diff --git a/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py b/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py index 72a3b84..8debafa 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py @@ -153,8 +153,8 @@ class FlashInferFp8DeepGEMMDynamicBlockScaledKernel( Conditional FlashInfer / DeepGEMM FP8 block-scaled GEMM. Dispatches between two kernels based on input batch size: - - Small batches (M < 32): FlashInfer's swapAB trick for better utilisation. - - Large batches (M >= 32): DeepGEMM for peak throughput. + - Decode-sized batches (M <= 32): FlashInfer's fused BF16-to-FP8 path. + - Larger batches (M > 32): DeepGEMM for peak throughput. apply_input_quant is False because FlashInfer accepts BF16 input and handles FP8 conversion internally. The DeepGEMM branch therefore @@ -240,8 +240,8 @@ def _dynamic_flashinfer_deepgemm_blockscale_gemm_impl( Conditional FlashInfer FP8 blockscale GEMM with batch-size-dependent selection. This function switches between two optimized kernels based on the input batch size: - - For small batches (M < 32): Uses FlashInfer's DeepGEMM swapAB optimization. - - For larger batches (M >= 32): Uses the official DeepGEMM kernel. + - For decode-sized batches (M <= 32): Uses FlashInfer's fused FP8 path. + - For larger batches (M > 32): Uses the official DeepGEMM kernel. The conditional logic must use torch.cond() instead of a simple if-else statement to maintain compatibility with torch.compile graph compilation. @@ -301,12 +301,12 @@ def _dynamic_flashinfer_deepgemm_blockscale_gemm_impl( if envs.VLLM_BATCH_INVARIANT: return run_deepgemm(input, weight, weight_scale) - condition = input.shape[0] < 32 + condition = input.shape[0] <= 32 # PyTorch's torch.compile cannot handle input-dependent control flow in standard # Python conditionals. torch.cond() explicitly registers both code paths in the # computation graph, allowing torch.compile to capture both branches. - # without torch.cond, the M < 32 condition won't be able to be captured by torch + # without torch.cond, the M <= 32 condition won't be able to be captured by torch # compile return torch.cond( condition,