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 4, 21:15 UTC
[Bugfix][Kernel] nvfp4 cutlass MoE: fix nvfp4 experts quant out-of-bounds read for expert counts not divisible by 4 or 16 (#40351)
Signed-off-by: Jakub Zakrzewski <jzakrzewski@nvidia.com>
Jakub Zakrzewski committed
Apr 21, 2026 at 21:06 UTC
6fbec8ed473f418a9c20d5b6a4d56486544139da
1 file changed
+6
-2
csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu
+6
-2
@@ -277,7 +277,9 @@ void quant_impl(void* output, void* output_scale, void* input,
277
(totalWorkSize + block.x * grid.x - 1) / (block.x * grid.x);
278
if (blockRepeat > 1) {
279
size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
280
- if (n_experts >= 4) {
280
+ // The shared-memory vectorized offset load only handles full 4-expert
281
+ // chunks. Use the scalar specialization for the remainder cases.
282
+ if (n_experts >= 4 && n_experts % 4 == 0) {
283
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
284
<<<grid, block, shared_mem_size, stream>>>(
285
m_topk, k, reinterpret_cast<T*>(input),
@@ -299,7 +301,9 @@ void quant_impl(void* output, void* output_scale, void* input,
301
n_experts);
302
}
303
} else {
302
- if (n_experts >= 16) {
304
+ // The low-latency vectorized expert lookup only handles full 16-expert
305
+ // chunks. Fall back to the scalar lookup path for the remainder cases.
306
+ if (n_experts >= 16 && n_experts % 16 == 0) {
307
cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
308
<<<grid, block, 0, stream>>>(
309
m_topk, k, reinterpret_cast<T*>(input),