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, 02:15 UTC
[feat] enable SM100 CUTLASS block scaled group gemm for smaller batch sizes (#20640)
Signed-off-by: Duncan Moss <djm.moss@gmail.com>
Duncan Moss committed
Jul 8, 2025 at 20:03 UTC
97abeb1daac6007526af435244d3f7047db272cd
2 files changed
+5
-7
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+4
-6
@@ -522,16 +522,14 @@ def cutlass_moe_fp4(a: torch.Tensor, a1_gscale: torch.Tensor,
522
return out.to(dtype=out_dtype)
523
524
525
-def _valid_cutlass_block_scaled_grouped_gemm(hidden_states: torch.Tensor,
526
- w1: torch.Tensor,
525
+def _valid_cutlass_block_scaled_grouped_gemm(w1: torch.Tensor,
526
w2: torch.Tensor) -> bool:
527
529
- def _valid_cutlass_block_scaled_grouped_gemm_shape(M: int, N: int, K: int):
530
- return M >= 128 and N % 128 == 0 and K % 128 == 0
528
+ def _valid_cutlass_block_scaled_grouped_gemm_shape(N: int, K: int):
529
+ return N % 128 == 0 and K % 128 == 0
530
532
- m = hidden_states.size(0)
531
_, K, N = w2.size()
534
- if not _valid_cutlass_block_scaled_grouped_gemm_shape(m, N, K):
532
+ if not _valid_cutlass_block_scaled_grouped_gemm_shape(N, K):
533
logger.debug(
534
"CutlassBlockScaledGroupedGemm disabled: unalinged problem size.")
535
return False
vllm/model_executor/layers/fused_moe/fused_moe.py
+1
-1
@@ -1180,7 +1180,7 @@ def fused_experts(
1180
apply_router_weight_on_input=apply_router_weight_on_input,
1181
)
1182
elif (allow_cutlass_block_scaled_grouped_gemm and use_fp8_w8a8
1183
- and _valid_cutlass_block_scaled_grouped_gemm(hidden_states, w1, w2)):
1183
+ and _valid_cutlass_block_scaled_grouped_gemm(w1, w2)):
1184
assert apply_router_weight_on_input is False
1185
return run_cutlass_block_scaled_fused_experts(
1186
a=hidden_states,