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, 16:00 UTC
[Bugfix] fix use_atomic_add support of marlin kernel when using v1 engine (#15946)
Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com>
Jinzhen Lin committed
Apr 6, 2025 at 11:04 UTC
2fa66ef713a62f96c24135ceffd70f1bbafea4e4
2 files changed
+6
-2
csrc/quantization/gptq_marlin/gptq_marlin.cu
+5
-1
@@ -1785,7 +1785,7 @@ __global__ void Marlin(
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, \
1787
num_groups, prob_m, prob_n, prob_k, lda, locks, \
1788
- use_atomic_add, use_fp32_reduce); \
1788
+ part_use_atomic_add, use_fp32_reduce); \
1789
} \
1790
}
1791
@@ -2215,6 +2215,10 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
2215
thread_m_blocks = exec_cfg.max_m_blocks;
2216
}
2217
2218
+ // atomic add reduce have better performance only when m * n is small
2219
+ bool part_use_atomic_add =
2220
+ use_atomic_add && div_ceil(prob_m, 64) * prob_n <= 2048;
2221
+
2222
if (false) {
2223
}
2224
GPTQ_CALL_IF(vllm::kU4B8, 16, 4, 256)
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+1
-1
@@ -305,7 +305,7 @@ def should_use_atomic_add_reduce(m: int, n: int, k: int, device: torch.device,
305
306
# the performance of atomicAdd is better than global reduce
307
# only when m*n is small and k is large
308
- return max(m, 64) * n < 64 * 2048 and k >= 2048
308
+ return n < 2048 and k >= 2048
309
310
311
def apply_gptq_marlin_linear(