[Kernel Slimming] Migrate GPTQ-Marlin repack kernel to JIT (#18543)

Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>

python/sglang/jit_kernel/benchmark/bench_gptq_marlin_repack.py

@@ -0,0 +1,104 @@
+import os
+
+import torch
+import triton
+import triton.testing
+from sgl_kernel.scalar_type import scalar_types
+
+from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack as jit_fn
+from sglang.srt.layers.quantization.utils import gptq_quantize_weights, pack_rows
+
+try:
+    from sgl_kernel import gptq_marlin_repack as aot_fn
+
+    AOT_AVAILABLE = True
+except ImportError:
+    AOT_AVAILABLE = False
+
+IS_CI = (
+    os.getenv("CI", "false").lower() == "true"
+    or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
+)
+
+# Fixed problem dimensions
+SIZE_N = 4096
+NUM_BITS = 4
+QUANT_TYPE = scalar_types.uint4b8
+GROUP_SIZE = 128
+
+# Pre-compute quantized weight for each size_k in the sweep
+_cache = {}
+
+
+def _get_inputs(size_k):
+    if size_k not in _cache:
+        size_n = SIZE_N
+        b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
+        _, q_w, _, _, _ = gptq_quantize_weights(
+            b_weight, QUANT_TYPE, GROUP_SIZE, act_order=False
+        )
+        q_w_gptq = pack_rows(q_w, NUM_BITS, size_k, size_n)
+        sort_indices = torch.empty(0, dtype=torch.int, device="cuda")
+        _cache[size_k] = (q_w_gptq, sort_indices)
+    return _cache[size_k]
+
+
+def check_correctness():
+    if not AOT_AVAILABLE:
+        print("sgl_kernel AOT not available, skipping correctness check")
+        return
+    size_k = 4096
+    q_w_gptq, sort_indices = _get_inputs(size_k)
+    out_jit = jit_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS)
+    out_aot = aot_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS)
+    torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0)
+    print("Correctness check passed (JIT vs AOT)")
+
+
+if IS_CI:
+    k_range = [128, 1024, 4096]
+else:
+    k_range = [128, 256, 512, 1024, 2048, 4096, 8192]
+
+if AOT_AVAILABLE:
+    line_vals = ["jit", "aot"]
+    line_names = ["JIT Kernel", "AOT Kernel"]
+    styles = [("blue", "-"), ("green", "-")]
+else:
+    line_vals = ["jit"]
+    line_names = ["JIT Kernel"]
+    styles = [("blue", "-")]
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["size_k"],
+        x_vals=k_range,
+        line_arg="provider",
+        line_vals=line_vals,
+        line_names=line_names,
+        styles=styles,
+        ylabel="us",
+        plot_name="gptq-marlin-repack-performance",
+        args={},
+    )
+)
+def benchmark(size_k, provider):
+    q_w_gptq, sort_indices = _get_inputs(size_k)
+
+    quantiles = [0.5, 0.2, 0.8]
+
+    if provider == "jit":
+        fn = lambda: jit_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS)
+    elif provider == "aot":
+        fn = lambda: aot_fn(q_w_gptq, sort_indices, size_k, SIZE_N, NUM_BITS)
+    else:
+        raise ValueError(f"Unknown provider: {provider}")
+
+    ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    check_correctness()
+    benchmark.run(print_data=True)

python/sglang/jit_kernel/csrc/gemm/marlin/gptq_marlin_repack.cuh

@@ -0,0 +1,362 @@
+/*
+ * Modified by Neural Magic
+ * Copyright (C) Marlin.2024 Elias Frantar
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *         http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/*
+ * Adapted from https://github.com/IST-DASLab/marlin
+ */
+
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include "marlin.cuh"
+
+namespace device::marlin {
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+template <int const num_threads, int const num_bits, bool const has_perm>
+__global__ void gptq_marlin_repack_kernel(
+    uint32_t const* __restrict__ b_q_weight_ptr,
+    uint32_t const* __restrict__ perm_ptr,
+    uint32_t* __restrict__ out_ptr,
+    int size_k,
+    int size_n) {
+  return;
+}
+#else
+template <int const num_threads, int const num_bits, bool const has_perm>
+__global__ void gptq_marlin_repack_kernel(
+    uint32_t const* __restrict__ b_q_weight_ptr,
+    uint32_t const* __restrict__ perm_ptr,
+    uint32_t* __restrict__ out_ptr,
+    int size_k,
+    int size_n) {
+  constexpr int pack_factor = 32 / num_bits;
+
+  int k_tiles = size_k / tile_k_size;
+  int n_tiles = size_n / tile_n_size;
+  int block_k_tiles = div_ceil(k_tiles, gridDim.x);
+
+  auto start_k_tile = blockIdx.x * block_k_tiles;
+  if (start_k_tile >= k_tiles) {
+    return;
+  }
+
+  int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles);
+
+  // Wait until the next thread tile has been loaded to shared memory.
+  auto wait_for_stage = [&]() {
+    // We only have `stages - 2` active fetches since we are double buffering
+    // and can only issue the next fetch when it is guaranteed that the previous
+    // shared memory load is fully complete (as it may otherwise be
+    // overwritten).
+    cp_async_wait<repack_stages - 2>();
+    __syncthreads();
+  };
+
+  extern __shared__ int4 sh[];
+
+  constexpr int perm_size = tile_k_size / 4;
+
+  int4* sh_perm_ptr = sh;
+  int4* sh_pipe_ptr = sh_perm_ptr;
+  if constexpr (has_perm) {
+    sh_pipe_ptr += perm_size;
+  }
+
+  constexpr int tile_ints = tile_k_size / pack_factor;
+
+  constexpr int stage_n_threads = tile_n_size / 4;
+  constexpr int stage_k_threads = has_perm ? tile_k_size : tile_ints;
+  constexpr int stage_size = stage_k_threads * stage_n_threads;
+
+  auto load_perm_to_shared = [&](int k_tile_id) {
+    int first_k_int4 = (k_tile_id * tile_k_size) / 4;
+
+    int4 const* perm_int4_ptr = reinterpret_cast<int4 const*>(perm_ptr);
+
+    if (threadIdx.x < perm_size) {
+      sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x];
+    }
+    __syncthreads();
+  };
+
+  auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) {
+    if (n_tile_id >= n_tiles) {
+      cp_async_fence();
+      return;
+    }
+
+    int first_n = n_tile_id * tile_n_size;
+
+    int4* sh_ptr = sh_pipe_ptr + stage_size * pipe;
+
+    if constexpr (has_perm) {
+      if (threadIdx.x < stage_size) {
+        auto k_id = threadIdx.x / stage_n_threads;
+        auto n_id = threadIdx.x % stage_n_threads;
+
+        uint32_t const* sh_perm_int_ptr = reinterpret_cast<uint32_t const*>(sh_perm_ptr);
+
+        int src_k = sh_perm_int_ptr[k_id];
+        int src_k_packed = src_k / pack_factor;
+
+        cp_async4(
+            &sh_ptr[k_id * stage_n_threads + n_id],
+            reinterpret_cast<int4 const*>(&(b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)])));
+      }
+
+    } else {
+      if (threadIdx.x < stage_size) {
+        auto k_id = threadIdx.x / stage_n_threads;
+        auto n_id = threadIdx.x % stage_n_threads;
+
+        int first_k = k_tile_id * tile_k_size;
+        int first_k_packed = first_k / pack_factor;
+
+        cp_async4(
+            &sh_ptr[k_id * stage_n_threads + n_id],
+            reinterpret_cast<int4 const*>(&(b_q_weight_ptr[(first_k_packed + k_id) * size_n + first_n + (n_id * 4)])));
+      }
+    }
+
+    cp_async_fence();
+  };
+
+  auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) {
+    if (n_tile_id >= n_tiles) {
+      return;
+    }
+
+    auto warp_id = threadIdx.x / 32;
+    auto th_id = threadIdx.x % 32;
+
+    if (warp_id >= 4) {
+      return;
+    }
+
+    int tc_col = th_id / 4;
+    int tc_row = (th_id % 4) * 2;
+
+    constexpr int tc_offsets[4] = {0, 1, 8, 9};
+
+    int cur_n = warp_id * 16 + tc_col;
+
+    constexpr int sh_stride = 64;
+    constexpr uint32_t mask = (1 << num_bits) - 1;
+
+    int4* sh_stage_ptr = sh_pipe_ptr + stage_size * pipe;
+    uint32_t* sh_stage_int_ptr = reinterpret_cast<uint32_t*>(sh_stage_ptr);
+
+    uint32_t* sh_perm_int_ptr = reinterpret_cast<uint32_t*>(sh_perm_ptr);
+
+    uint32_t vals[8];
+
+    if constexpr (has_perm) {
+      for (int i = 0; i < 4; i++) {
+        int k_idx = tc_row + tc_offsets[i];
+
+        uint32_t src_k = sh_perm_int_ptr[k_idx];
+        uint32_t src_k_pos = src_k % pack_factor;
+
+        uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n];
+        uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask;
+
+        uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8];
+        uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask;
+
+        vals[i] = b1_cur_val;
+        vals[4 + i] = b2_cur_val;
+      }
+
+    } else {
+      uint32_t b1_vals[tile_ints];
+      uint32_t b2_vals[tile_ints];
+
+#pragma unroll
+      for (int i = 0; i < tile_ints; i++) {
+        b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i];
+        b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i];
+      }
+
+#pragma unroll
+      for (int i = 0; i < 4; i++) {
+        int cur_elem = tc_row + tc_offsets[i];
+        int cur_int = cur_elem / pack_factor;
+        int cur_pos = cur_elem % pack_factor;
+
+        vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask;
+        vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask;
+      }
+    }
+
+    constexpr int tile_size = tile_k_size * tile_n_size / pack_factor;
+    int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size;
+
+    // Result of:
+    // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h
+    if constexpr (num_bits == 4) {
+      constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7};
+
+      uint32_t res = 0;
+#pragma unroll
+      for (int i = 0; i < 8; i++) {
+        res |= vals[pack_idx[i]] << (i * 4);
+      }
+
+      out_ptr[out_offset + th_id * 4 + warp_id] = res;
+
+    } else {
+      constexpr int pack_idx[4] = {0, 2, 1, 3};
+
+      uint32_t res1 = 0;
+      uint32_t res2 = 0;
+#pragma unroll
+      for (int i = 0; i < 4; i++) {
+        res1 |= vals[pack_idx[i]] << (i * 8);
+        res2 |= vals[4 + pack_idx[i]] << (i * 8);
+      }
+
+      out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1;
+      out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2;
+    }
+  };
+
+  auto start_pipes = [&](int k_tile_id, int n_tile_id) {
+#pragma unroll
+    for (int pipe = 0; pipe < repack_stages - 1; pipe++) {
+      fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe);
+    }
+
+    wait_for_stage();
+  };
+#pragma unroll
+  for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) {
+    int n_tile_id = 0;
+
+    if constexpr (has_perm) {
+      load_perm_to_shared(k_tile_id);
+    }
+
+    start_pipes(k_tile_id, n_tile_id);
+
+    while (n_tile_id < n_tiles) {
+#pragma unroll
+      for (int pipe = 0; pipe < repack_stages; pipe++) {
+        fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id, n_tile_id + pipe + repack_stages - 1);
+        repack_tile(pipe, k_tile_id, n_tile_id + pipe);
+        wait_for_stage();
+      }
+      n_tile_id += repack_stages;
+    }
+  }
+}
+#endif
+
+}  // namespace device::marlin
+
+#define CALL_IF_REPACK(NUM_BITS, HAS_PERM)                                                                        \
+  else if (num_bits == NUM_BITS && has_perm == HAS_PERM) {                                                        \
+    host::RuntimeDeviceCheck(cudaFuncSetAttribute(                                                                \
+        device::marlin::gptq_marlin_repack_kernel<device::marlin::repack_threads, NUM_BITS, HAS_PERM>,            \
+        cudaFuncAttributeMaxDynamicSharedMemorySize,                                                              \
+        max_shared_mem));                                                                                         \
+    host::LaunchKernel(blocks, device::marlin::repack_threads, stream, static_cast<std::size_t>(max_shared_mem))( \
+        device::marlin::gptq_marlin_repack_kernel<device::marlin::repack_threads, NUM_BITS, HAS_PERM>,            \
+        b_q_weight_ptr,                                                                                           \
+        perm_ptr,                                                                                                 \
+        out_ptr,                                                                                                  \
+        size_k,                                                                                                   \
+        size_n);                                                                                                  \
+  }
+
+void gptq_marlin_repack(
+    tvm::ffi::TensorView b_q_weight,
+    tvm::ffi::TensorView perm,
+    tvm::ffi::TensorView out,
+    int64_t size_k,
+    int64_t size_n,
+    int64_t num_bits) {
+  using namespace host;
+
+  // Validate num_bits
+  RuntimeCheck(num_bits == 4 || num_bits == 8, "num_bits must be 4 or 8. Got = ", num_bits);
+  int const pack_factor = 32 / static_cast<int>(num_bits);
+
+  // Validate size alignment
+  RuntimeCheck(
+      size_k % device::marlin::tile_k_size == 0,
+      "size_k = ",
+      size_k,
+      " is not divisible by tile_k_size = ",
+      device::marlin::tile_k_size);
+  RuntimeCheck(
+      size_n % device::marlin::tile_n_size == 0,
+      "size_n = ",
+      size_n,
+      " is not divisible by tile_n_size = ",
+      device::marlin::tile_n_size);
+
+  // Validate b_q_weight
+  auto bqw_dim0 = SymbolicSize{"bqw_dim0"};
+  auto bqw_dim1 = SymbolicSize{"bqw_dim1"};
+  bqw_dim0.set_value(size_k / pack_factor);
+  bqw_dim1.set_value(size_n);
+  auto device_ = SymbolicDevice{};
+  device_.set_options<kDLCUDA>();
+  TensorMatcher({bqw_dim0, bqw_dim1}).with_dtype<int32_t>().with_device(device_).verify(b_q_weight);
+
+  // Validate out
+  auto out_dim0 = SymbolicSize{"out_dim0"};
+  auto out_dim1 = SymbolicSize{"out_dim1"};
+  out_dim0.set_value(size_k / device::marlin::tile_size);
+  out_dim1.set_value(size_n * device::marlin::tile_size / pack_factor);
+  TensorMatcher({out_dim0, out_dim1}).with_dtype<int32_t>().with_device(device_).verify(out);
+
+  // Detect if there is act_order
+  bool has_perm = perm.size(0) != 0;
+
+  // Get ptrs
+  uint32_t const* b_q_weight_ptr = reinterpret_cast<uint32_t const*>(b_q_weight.data_ptr());
+  uint32_t const* perm_ptr = reinterpret_cast<uint32_t const*>(perm.data_ptr());
+  uint32_t* out_ptr = reinterpret_cast<uint32_t*>(out.data_ptr());
+
+  // Get dev info
+  DLDevice dl_device = device_.unwrap();
+  int dev = dl_device.device_id;
+  cudaStream_t stream = LaunchKernel::resolve_device(dl_device);
+  int blocks;
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, dev));
+
+  int max_shared_mem = 0;
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev));
+  RuntimeCheck(max_shared_mem > 0, "max_shared_mem must be > 0");
+
+  if (false) {
+  }
+  CALL_IF_REPACK(4, false)
+  CALL_IF_REPACK(4, true)
+  CALL_IF_REPACK(8, false)
+  CALL_IF_REPACK(8, true)
+  else {
+    Panic("Unsupported repack config: num_bits = ", num_bits, ", has_perm = ", has_perm);
+  }
+}
+
+#undef CALL_IF_REPACK

python/sglang/jit_kernel/gptq_marlin_repack.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+# Constants matching device::marlin:: in marlin.cuh
+_TILE_SIZE = 16
+
+
+@cache_once
+def _jit_gptq_marlin_repack_module() -> Module:
+    return load_jit(
+        "gptq_marlin_repack",
+        cuda_files=["gemm/marlin/gptq_marlin_repack.cuh"],
+        cuda_wrappers=[("gptq_marlin_repack", "gptq_marlin_repack")],
+    )
+
+
+def gptq_marlin_repack(
+    b_q_weight: torch.Tensor,
+    perm: torch.Tensor,
+    size_k: int,
+    size_n: int,
+    num_bits: int,
+) -> torch.Tensor:
+    pack_factor = 32 // num_bits
+
+    # Allocate output tensor
+    out = torch.empty(
+        (size_k // _TILE_SIZE, size_n * _TILE_SIZE // pack_factor),
+        dtype=b_q_weight.dtype,
+        device=b_q_weight.device,
+    )
+
+    module = _jit_gptq_marlin_repack_module()
+    module.gptq_marlin_repack(b_q_weight, perm, out, size_k, size_n, num_bits)
+    return out

python/sglang/jit_kernel/tests/test_gptq_marlin_repack.py

@@ -0,0 +1,101 @@
+import pytest
+import torch
+from sgl_kernel import gptq_marlin_repack as aot_gptq_marlin_repack
+from sgl_kernel.scalar_type import scalar_types
+
+from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
+from sglang.srt.layers.quantization.utils import (
+    gptq_quantize_weights,
+    pack_rows,
+    sort_weights,
+)
+from sglang.test.test_marlin_utils import get_weight_perm, marlin_weights
+
+MARLIN_K_CHUNKS = [128]
+MARLIN_N_CHUNKS = [64, 256]
+
+MNK_FACTORS = [
+    (1, 1, 1),
+    (1, 4, 8),
+    (1, 7, 5),
+    (13, 17, 67),
+    (26, 37, 13),
+    (67, 13, 11),
+    (257, 13, 11),
+    (658, 13, 11),
+]
+
+
+@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
+@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS)
+@pytest.mark.parametrize("quant_type", [scalar_types.uint4b8])
+@pytest.mark.parametrize("group_size", [-1, 32, 64, 128])
+@pytest.mark.parametrize("act_order", [False, True])
+@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
+def test_gptq_marlin_repack(
+    k_chunk, n_chunk, quant_type, group_size, act_order, mnk_factors
+):
+    m_factor, n_factor, k_factor = mnk_factors
+
+    size_k = k_chunk * k_factor
+    size_n = n_chunk * n_factor
+
+    # Filter act_order
+    if act_order:
+        if group_size == -1:
+            return
+        if group_size == size_k:
+            return
+
+    # Normalize group_size
+    if group_size == -1:
+        group_size = size_k
+    assert group_size <= size_k
+
+    if size_k % group_size != 0:
+        pytest.skip("size_k must be divisible by group_size")
+
+    # Create input
+    b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
+
+    # Quantize (and apply act_order if provided)
+    w_ref, q_w, s, g_idx, rand_perm = gptq_quantize_weights(
+        b_weight, quant_type, group_size, act_order
+    )
+
+    q_w_gptq = pack_rows(q_w, quant_type.size_bits, size_k, size_n)
+
+    # For act_order, sort the "weights" and "g_idx" so that group ids are
+    # increasing
+    sort_indices = torch.empty(0, dtype=torch.int, device=b_weight.device)
+    if act_order:
+        q_w, g_idx, sort_indices = sort_weights(q_w, g_idx)
+
+    marlin_layout_perm = get_weight_perm(quant_type.size_bits)
+    q_w_marlin_ref = marlin_weights(
+        q_w, size_k, size_n, quant_type.size_bits, marlin_layout_perm
+    )
+
+    # Run JIT repack kernel
+    jit_output = gptq_marlin_repack(
+        q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits
+    )
+
+    # Run AOT repack kernel
+    aot_output = aot_gptq_marlin_repack(
+        q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits
+    )
+
+    torch.cuda.synchronize()
+
+    # JIT should match the reference (computed from CPU marlin_weights)
+    torch.testing.assert_close(jit_output, q_w_marlin_ref)
+
+    # JIT should produce bitwise identical results to AOT
+    torch.testing.assert_close(jit_output, aot_output, rtol=0, atol=0)
+
+
+if __name__ == "__main__":
+    import subprocess
+
+    subprocess.call(["pytest", "--tb=short", str(__file__)])

python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py

@@ -43,7 +43,7 @@ from sglang.srt.utils import is_cuda
 _is_cuda = is_cuda()
 
 if _is_cuda:
-    from sgl_kernel import gptq_marlin_repack
+    from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
 
 
 ScalarType, scalar_types = get_scalar_types()

python/sglang/srt/layers/quantization/gptq.py

@@ -61,7 +61,9 @@ if TYPE_CHECKING:
 _is_cuda = is_cuda()
 
 if _is_cuda:
-    from sgl_kernel import gptq_gemm, gptq_marlin_repack, gptq_shuffle
+    from sgl_kernel import gptq_gemm, gptq_shuffle
+
+    from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
 
 _is_npu = is_npu()
 

python/sglang/srt/layers/quantization/marlin_utils_fp8.py

@@ -17,9 +17,8 @@ from sglang.srt.utils import is_cuda
 
 _is_cuda = is_cuda()
 if _is_cuda:
-    from sgl_kernel import gptq_marlin_repack
-
     from sglang.jit_kernel.gptq_marlin import gptq_marlin_gemm
+    from sglang.jit_kernel.gptq_marlin_repack import gptq_marlin_repack
 
 ScalarType, scalar_types = get_scalar_types()