[Bugfix] fix use_atomic_add support of marlin kernel when using v1 engine (#15946)

Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com>
This commit is contained in:
Jinzhen Lin
2025-04-06 11:04:22 +08:00
committed by GitHub
parent 13affc432d
commit 2fa66ef713
2 changed files with 6 additions and 2 deletions

View File

@@ -1785,7 +1785,7 @@ __global__ void Marlin(
<<<blocks, NUM_THREADS, max_shared_mem, stream>>>( \
A_ptr, B_ptr, C_ptr, C_tmp_ptr, s_ptr, zp_ptr, g_idx_ptr, \
num_groups, prob_m, prob_n, prob_k, lda, locks, \
use_atomic_add, use_fp32_reduce); \
part_use_atomic_add, use_fp32_reduce); \
} \
}
@@ -2215,6 +2215,10 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,
thread_m_blocks = exec_cfg.max_m_blocks;
}
// atomic add reduce have better performance only when m * n is small
bool part_use_atomic_add =
use_atomic_add && div_ceil(prob_m, 64) * prob_n <= 2048;
if (false) {
}
GPTQ_CALL_IF(vllm::kU4B8, 16, 4, 256)