A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 17:15 UTC
[Kernel] allow non-contiguous input for marlin kernel (#14658)
Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com>
Jinzhen Lin committed
Mar 24, 2025 at 21:21 UTC
6b3cc75be07d7f7458f4a649eb3da10bdd2f9517
2 files changed
+73
-20
csrc/quantization/gptq_marlin/gptq_marlin.cu
+28
-20
@@ -42,7 +42,7 @@ namespace marlin {
42
__global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
43
int const* __restrict__ perm_int_ptr,
44
int4* __restrict__ out_int4_ptr, int size_m,
45
- int size_k, int block_rows) {}
45
+ int size_k, int lda, int block_rows) {}
46
47
template <typename scalar_t, // compute dtype, half or nv_float16
48
const vllm::ScalarTypeId w_type_id, // weight ScalarType id
@@ -459,7 +459,7 @@ __device__ inline void barrier_release(int* lock, bool reset = false) {
459
__global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
460
int const* __restrict__ perm_int_ptr,
461
int4* __restrict__ out_int4_ptr, int size_m,
462
- int size_k, int block_rows) {
462
+ int size_k, int lda, int block_rows) {
463
int start_row = block_rows * blockIdx.x;
464
int finish_row = start_row + block_rows;
465
if (finish_row > size_m) {
@@ -467,16 +467,19 @@ __global__ void permute_cols_kernel(int4 const* __restrict__ a_int4_ptr,
467
}
468
int cur_block_rows = finish_row - start_row;
469
470
- int row_stride = size_k * sizeof(half) / 16;
470
+ int input_row_stride = lda * sizeof(half) / 16;
471
+ int output_row_stride = size_k * sizeof(half) / 16;
472
473
auto permute_row = [&](int row) {
474
int iters = size_k / default_threads;
475
int rest = size_k % default_threads;
476
476
- int offset = row * row_stride;
477
+ int input_offset = row * input_row_stride;
478
+ int output_offset = row * output_row_stride;
479
478
- half const* a_row_half = reinterpret_cast<half const*>(a_int4_ptr + offset);
479
- half* out_half = reinterpret_cast<half*>(out_int4_ptr + offset);
480
+ half const* a_row_half =
481
+ reinterpret_cast<half const*>(a_int4_ptr + input_offset);
482
+ half* out_half = reinterpret_cast<half*>(out_int4_ptr + output_offset);
483
484
int base_k = 0;
485
@@ -537,6 +540,7 @@ __global__ void Marlin(
540
int prob_m, // batch dimension m
541
int prob_n, // output dimension n
542
int prob_k, // reduction dimension k
543
+ int lda, // A.stride(0), equal to prob_k is A is contiguous
544
int* locks, // extra global storage for barrier synchronization
545
bool use_atomic_add, // whether to use atomic add to reduce
546
bool use_fp32_reduce // whether to use fp32 global reduce
@@ -600,7 +604,7 @@ __global__ void Marlin(
604
// We can easily implement parallel problem execution by just remapping
605
// indices and advancing global pointers
606
if (slice_col_par >= n_tiles) {
603
- A += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_k / 8;
607
+ A += (slice_col_par / n_tiles) * 16 * thread_m_blocks * lda / 8;
608
C += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_n / 8;
609
locks += (slice_col_par / n_tiles) * n_tiles;
610
slice_col = slice_col_par % n_tiles;
@@ -631,7 +635,7 @@ __global__ void Marlin(
635
}
636
}
637
if (slice_col == n_tiles) {
634
- A += 16 * thread_m_blocks * prob_k / 8;
638
+ A += 16 * thread_m_blocks * lda / 8;
639
C += 16 * thread_m_blocks * prob_n / 8;
640
locks += n_tiles;
641
slice_col = 0;
@@ -643,7 +647,7 @@ __global__ void Marlin(
647
// A sizes/strides
648
649
// stride of the A matrix in global memory
646
- int a_gl_stride = prob_k / 8;
650
+ int a_gl_stride = lda / 8;
651
// stride of an A matrix tile in shared memory
652
constexpr int a_sh_stride = 16 * thread_k_blocks / 8;
653
// delta between subsequent A tiles in global memory
@@ -1780,8 +1784,8 @@ __global__ void Marlin(
1784
HAS_ZP, GROUP_BLOCKS, IS_ZP_FLOAT> \
1785
<<<blocks, NUM_THREADS, max_shared_mem, stream>>>( \
1786
A_ptr, B_ptr, C_ptr, C_tmp_ptr, s_ptr, zp_ptr, g_idx_ptr, \
1783
- num_groups, prob_m, prob_n, prob_k, locks, use_atomic_add, \
1784
- use_fp32_reduce); \
1787
+ num_groups, prob_m, prob_n, prob_k, lda, locks, \
1788
+ use_atomic_add, use_fp32_reduce); \
1789
} \
1790
}
1791
@@ -2071,7 +2075,7 @@ exec_config_t determine_thread_config(int prob_m, int prob_n, int prob_k,
2075
template <typename scalar_t>
2076
void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
2077
void* zp, void* g_idx, void* perm, void* a_tmp, int prob_m,
2074
- int prob_n, int prob_k, void* workspace,
2078
+ int prob_n, int prob_k, int lda, void* workspace,
2079
vllm::ScalarType const& q_type, bool has_act_order,
2080
bool is_k_full, bool has_zp, int num_groups, int group_size,
2081
int dev, cudaStream_t stream, int thread_k, int thread_n,
@@ -2184,8 +2188,9 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
2188
// Permute A columns
2189
int block_rows = div_ceil(prob_m, blocks);
2190
permute_cols_kernel<<<blocks, default_threads, 0, stream>>>(
2187
- A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, block_rows);
2191
+ A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, lda, block_rows);
2192
A_ptr = a_tmp_ptr;
2193
+ lda = prob_k;
2194
}
2195
2196
// If we have a full K, then we can run the non-act-order version of Marlin
@@ -2244,7 +2249,7 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
2249
", num_bits = ", num_bits);
2250
}
2251
2247
- A_ptr += 16 * thread_m_blocks * (prob_k / 8) * par;
2252
+ A_ptr += 16 * thread_m_blocks * (lda / 8) * par;
2253
C_ptr += 16 * thread_m_blocks * (prob_n / 8) * par;
2254
}
2255
}
@@ -2300,7 +2305,10 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
2305
2306
// Verify device and strides
2307
TORCH_CHECK(a.device().is_cuda(), "A is not on GPU");
2303
- TORCH_CHECK(a.is_contiguous(), "A is not contiguous");
2308
+ TORCH_CHECK(a.stride(1) == 1, "A.stride(1) is not 1");
2309
+ // We use int4 (16 bytes) to load A, so A must aligned to 16 bytes
2310
+ TORCH_CHECK(a.stride(0) % 8 == 0, "A.stride(0) must divisible by 8");
2311
+ TORCH_CHECK(((uint64_t)a.data_ptr()) % 16 == 0, "A must aligned to 16 bytes");
2312
2313
TORCH_CHECK(b_q_weight.device().is_cuda(), "b_q_weight is not on GPU");
2314
TORCH_CHECK(b_q_weight.is_contiguous(), "b_q_weight is not contiguous");
@@ -2432,7 +2440,7 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
2440
a.data_ptr<at::Half>(), b_q_weight.data_ptr(), c.data_ptr<at::Half>(),
2441
c_tmp.data_ptr<float>(), b_scales.data_ptr<at::Half>(),
2442
b_zeros.data_ptr(), g_idx.data_ptr(), perm.data_ptr(),
2435
- a_tmp.data_ptr<at::Half>(), size_m, size_n, size_k,
2443
+ a_tmp.data_ptr<at::Half>(), size_m, size_n, size_k, a.stride(0),
2444
workspace.data_ptr(), b_q_type, has_act_order, is_k_full, has_zp,
2445
num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),
2446
thread_k, thread_n, sms, marlin::max_par, use_atomic_add,
@@ -2443,10 +2451,10 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,
2451
c.data_ptr<at::BFloat16>(), c_tmp.data_ptr<float>(),
2452
b_scales.data_ptr<at::BFloat16>(), b_zeros.data_ptr(), g_idx.data_ptr(),
2453
perm.data_ptr(), a_tmp.data_ptr<at::BFloat16>(), size_m, size_n, size_k,
2446
- workspace.data_ptr(), b_q_type, has_act_order, is_k_full, has_zp,
2447
- num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),
2448
- thread_k, thread_n, sms, marlin::max_par, use_atomic_add,
2449
- use_fp32_reduce, is_zp_float);
2454
+ a.stride(0), workspace.data_ptr(), b_q_type, has_act_order, is_k_full,
2455
+ has_zp, num_groups, group_size, dev,
2456
+ at::cuda::getCurrentCUDAStream(dev), thread_k, thread_n, sms,
2457
+ marlin::max_par, use_atomic_add, use_fp32_reduce, is_zp_float);
2458
} else {
2459
TORCH_CHECK(false, "gpt_marlin_gemm only supports bfloat16 and float16");
2460
}
tests/kernels/test_marlin_gemm.py
+45
@@ -606,6 +606,51 @@ def test_marlin_qqq_gemm(
606
assert max_diff < 0.04
607
608
609
+def test_marlin_gemm_subset_input():
610
+ quant_type = scalar_types.uint4b8
611
+ group_size = 128
612
+
613
+ size_m, size_k, size_n = 32, 1024, 2048
614
+ big_m = size_m * 2
615
+ big_k = size_k * 2
616
+
617
+ a_input = rand_data((big_m, big_k))[8:size_m + 8, 8:size_k + 8]
618
+ b_weight = rand_data((size_k, size_n))
619
+
620
+ w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize(
621
+ b_weight, quant_type, group_size, False)
622
+
623
+ marlin_zp = marlin_make_empty_g_idx(marlin_s.device)
624
+ workspace = MarlinWorkspace(size_n, GPTQ_MARLIN_MIN_THREAD_N,
625
+ GPTQ_MARLIN_MAX_PARALLEL)
626
+
627
+ output = ops.gptq_marlin_gemm(
628
+ a_input,
629
+ marlin_q_w,
630
+ marlin_s,
631
+ marlin_zp,
632
+ g_idx,
633
+ sort_indices,
634
+ workspace.scratch,
635
+ quant_type,
636
+ a_input.shape[0],
637
+ b_weight.shape[1],
638
+ a_input.shape[1],
639
+ is_k_full=True,
640
+ has_zp=False,
641
+ use_atomic_add=False,
642
+ use_fp32_reduce=True,
643
+ is_zp_float=False,
644
+ )
645
+ output_ref = torch.matmul(a_input, w_ref)
646
+
647
+ torch.cuda.synchronize()
648
+
649
+ max_diff = compute_max_diff(output, output_ref)
650
+
651
+ assert max_diff < 0.04
652
+
653
+
654
def test_marlin_gemm_opcheck():
655
size_m = 2048
656
size_n = 4096