@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / cc6dafaef2
Fork 0

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 10, 14:00 UTC

[Perf][Kernels] Enable FlashInfer DeepGEMM swapAB on SM90 (for W8A8 Linear Op) (#29213)

Signed-off-by: Kate Cheng <yunhsuanc@nvidia.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com> Co-authored-by: Jhao-Ting Chen <jhaotingc@nvidia.com>

Kate Cheng committed Jan 7, 2026 at 07:53 UTC cc6dafaef2bc6b0a3ea0d80efc93985fce6bc8b9
4 files changed +257 -1
tests/kernels/quantization/test_block_fp8.py
+51
@@ -24,6 +24,10 @@ from vllm.utils.deep_gemm import (
24 per_block_cast_to_fp8,
25 should_use_deepgemm_for_fp8_linear,
26 )
27 +from vllm.utils.flashinfer import (
28 + flashinfer_fp8_blockscale_gemm,
29 + has_flashinfer_fp8_blockscale_gemm,
30 +)
31 from vllm.utils.import_utils import has_deep_gemm
32
33 if current_platform.get_device_capability() < (9, 0):
@@ -205,3 +209,50 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
209 torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))
210 ) / torch.mean(torch.abs(ref_out.to(torch.float32)))
211 assert rel_diff < 0.001
212 +
213 +
214 +@pytest.mark.skipif(
215 + current_platform.is_fp8_fnuz(),
216 + reason="This platform supports e4m3fnuz, not e4m3fn.",
217 +)
218 +@pytest.mark.parametrize(
219 + "M,N,K,block_size,out_dtype,seed",
220 + itertools.product(M, N, K, BLOCK_SIZE, OUT_DTYPES, SEEDS),
221 +)
222 +@torch.inference_mode()
223 +def test_w8a8_block_fp8_flashinfer_matmul(M, N, K, block_size, out_dtype, seed):
224 + if not has_flashinfer_fp8_blockscale_gemm():
225 + pytest.skip(
226 + "FlashInfer block GEMM not available (requires SM90+ and FlashInfer)"
227 + )
228 + # only aligned sizes
229 + if K % 128 != 0 or N % 64 != 0:
230 + pytest.skip(f"Skipping test; invalid size {M}, {N}, {K}")
231 +
232 + torch.manual_seed(seed)
233 + fp8_info = torch.finfo(torch.float8_e4m3fn)
234 + fp8_max = fp8_info.max
235 +
236 + A_bf16 = (torch.rand(M, K, dtype=torch.bfloat16) - 0.5) * 2 * fp8_max
237 + B_bf16 = (torch.rand(N, K, dtype=torch.bfloat16) - 0.5) * 2 * fp8_max
238 +
239 + A_fp8, As_fp8 = per_token_group_quant_fp8(A_bf16, block_size[1], use_ue8m0=False)
240 + B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_bf16, block_size, use_ue8m0=False)
241 +
242 + As = As_fp8.to(torch.float32)
243 + Bs = Bs_fp8.to(torch.float32)
244 +
245 + ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
246 +
247 + out = flashinfer_fp8_blockscale_gemm(
248 + input=A_bf16,
249 + weight=B_fp8,
250 + input_scale=None,
251 + weight_scale=Bs,
252 + out_dtype=out_dtype,
253 + )
254 +
255 + rel_diff = torch.mean(
256 + torch.abs(out.to(torch.bfloat16) - ref_out.to(torch.bfloat16))
257 + ) / torch.mean(torch.abs(ref_out.to(torch.bfloat16)))
258 + assert rel_diff < 0.001
vllm/envs.py
+6
@@ -168,6 +168,7 @@ if TYPE_CHECKING:
168 "relax",
169 ] = "relax"
170 VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True
171 + VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = False
172 VLLM_USE_FLASHINFER_MOE_FP16: bool = False
173 VLLM_USE_FLASHINFER_MOE_FP8: bool = False
174 VLLM_USE_FLASHINFER_MOE_FP4: bool = False
@@ -1206,6 +1207,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
1207 "VLLM_USE_FUSED_MOE_GROUPED_TOPK": lambda: bool(
1208 int(os.getenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "1"))
1209 ),
1210 + # Allow use of FlashInfer FP8 block-scale GEMM for linear layers.
1211 + # This uses TensorRT-LLM kernels and requires SM90+ (Hopper).
1212 + "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool(
1213 + int(os.getenv("VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER", "0"))
1214 + ),
1215 # Allow use of FlashInfer MoE kernels for fused moe ops.
1216 "VLLM_USE_FLASHINFER_MOE_FP16": lambda: bool(
1217 int(os.getenv("VLLM_USE_FLASHINFER_MOE_FP16", "0"))
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+143 -1
@@ -38,6 +38,11 @@ from vllm.utils.deep_gemm import (
38 should_use_deepgemm_for_fp8_linear,
39 transform_sf_into_required_layout,
40 )
41 +from vllm.utils.flashinfer import (
42 + flashinfer_fp8_blockscale_gemm,
43 + is_flashinfer_fp8_blockscale_gemm_supported,
44 + should_use_flashinfer_for_blockscale_fp8_gemm,
45 +)
46 from vllm.utils.torch_utils import direct_register_custom_op
47
48 logger = init_logger(__name__)
@@ -229,6 +234,112 @@ direct_register_custom_op(
234 )
235
236
237 +def _flashinfer_fp8_blockscale_gemm_impl(
238 + input: torch.Tensor,
239 + weight: torch.Tensor,
240 + weight_scale: torch.Tensor,
241 + group_size: int,
242 + use_deep_gemm_e8m0: bool,
243 +) -> torch.Tensor:
244 + """
245 + Conditional FlashInfer FP8 blockscale GEMM with batch-size-dependent selection.
246 +
247 + This function switches between two optimized kernels based on the input batch size:
248 + - For small batches (M < 32): Uses FlashInfer's DeepGEMM swapAB optimization.
249 + - For larger batches (M >= 32): Uses the official DeepGEMM kernel.
250 +
251 + The conditional logic must use torch.cond() instead of a simple if-else statement
252 + to maintain compatibility with torch.compile graph compilation.
253 +
254 + This batch-size-dependent selection is essential for maintaining model accuracy.
255 + Benchmarks on GSM8K show a significant accuracy gap (88% vs 95%) for DeepSeek-V3.1
256 + when using FlashInfer's DeepGEMM on M>=32. The M < 32 strategy fixes the accurracy
257 + drop.
258 +
259 + Args:
260 + input: Input tensor of shape (batch_size, input_dim) in FP8 format
261 + weight: Weight tensor of shape (output_dim, input_dim) in FP8 format
262 + weight_scale: Scale factors for weight quantization (per-group)
263 + group_size: Quantization group size for the weight tensor
264 + use_deep_gemm_e8m0: Whether to use the E8M0 format in DeepGEMM quantization
265 +
266 + Returns:
267 + Output tensor of shape (batch_size, output_dim) in bfloat16 format
268 + """
269 +
270 + def run_flashinfer_deepgemm_swapAB(
271 + input: torch.Tensor,
272 + weight: torch.Tensor,
273 + weight_scale: torch.Tensor,
274 + ) -> torch.Tensor:
275 + return flashinfer_fp8_blockscale_gemm(
276 + input=input,
277 + weight=weight,
278 + weight_scale=weight_scale,
279 + out_dtype=torch.bfloat16,
280 + )
281 +
282 + def run_deepgemm(
283 + input: torch.Tensor,
284 + weight: torch.Tensor,
285 + weight_scale: torch.Tensor,
286 + ) -> torch.Tensor:
287 + q_input, input_scale = per_token_group_quant_fp8(
288 + input,
289 + group_size=group_size,
290 + column_major_scales=True,
291 + use_ue8m0=use_deep_gemm_e8m0,
292 + )
293 + output = torch.empty(
294 + (q_input.shape[0], weight.shape[0]),
295 + dtype=torch.bfloat16,
296 + device=q_input.device,
297 + )
298 + fp8_gemm_nt(
299 + (q_input, input_scale),
300 + (weight, weight_scale),
301 + output,
302 + is_deep_gemm_e8m0_used=use_deep_gemm_e8m0,
303 + )
304 + return output
305 +
306 + condition = input.shape[0] < 32
307 +
308 + # PyTorch's torch.compile cannot handle input-dependent control flow in standard
309 + # Python conditionals. torch.cond() explicitly registers both code paths in the
310 + # computation graph, allowing torch.compile to capture both branches.
311 + # without torch.cond, the M < 32 condition won't be able to be captured by torch
312 + # compile
313 + return torch.cond(
314 + condition,
315 + run_flashinfer_deepgemm_swapAB,
316 + run_deepgemm,
317 + (input, weight, weight_scale),
318 + )
319 +
320 +
321 +def _flashinfer_fp8_blockscale_gemm_fake(
322 + input: torch.Tensor,
323 + weight: torch.Tensor,
324 + weight_scale: torch.Tensor,
325 + group_size: int,
326 + use_deep_gemm_e8m0: bool,
327 +) -> torch.Tensor:
328 + """
329 + Required fake/meta implementation for torch.compile graph tracing.
330 + """
331 + return torch.empty(
332 + input.shape[0], weight.shape[0], dtype=torch.bfloat16, device=input.device
333 + )
334 +
335 +
336 +direct_register_custom_op(
337 + "flashinfer_fp8_blockscale_gemm",
338 + _flashinfer_fp8_blockscale_gemm_impl,
339 + fake_impl=_flashinfer_fp8_blockscale_gemm_fake,
340 +)
341 +
342 +
343 # TODO fix ROCm->Triton custom path:
344 # https://github.com/vllm-project/vllm/issues/14397
345 class W8A8BlockFp8LinearOp:
@@ -249,6 +360,7 @@ class W8A8BlockFp8LinearOp:
360 self.is_deep_gemm_supported = is_deep_gemm_supported()
361 self.is_hopper = current_platform.is_device_capability(90)
362 self.use_deep_gemm_e8m0 = is_deep_gemm_e8m0_used()
363 + self.is_flashinfer_supported = is_flashinfer_fp8_blockscale_gemm_supported()
364
365 # Get the correct blockscale mul and input quant operations.
366 # We can't use _dispatch_w8a8_blockscale_op to figure out if we want
@@ -284,7 +396,14 @@ class W8A8BlockFp8LinearOp:
396 output_shape = [*input.shape[:-1], weight.shape[0]]
397 output_dtype = input.dtype
398
287 - if should_use_deepgemm_for_fp8_linear(
399 + if should_use_flashinfer_for_blockscale_fp8_gemm(
400 + self.is_flashinfer_supported, output_dtype, input_2d, weight
401 + ) and should_use_deepgemm_for_fp8_linear(
402 + output_dtype, weight, self.is_deep_gemm_supported
403 + ):
404 + output = self._run_flashinfer(input_2d, weight, weight_scale)
405 +
406 + elif should_use_deepgemm_for_fp8_linear(
407 output_dtype, weight, self.is_deep_gemm_supported
408 ):
409 output = self._run_deepgemm(input_2d, weight, weight_scale)
@@ -412,6 +531,29 @@ class W8A8BlockFp8LinearOp:
531 input_2d.dtype,
532 )
533
534 + def _run_flashinfer(
535 + self,
536 + input_2d: torch.Tensor,
537 + weight: torch.Tensor,
538 + weight_scale: torch.Tensor,
539 + ) -> torch.Tensor:
540 + """
541 + Run FlashInfer FP8 block-scale GEMM.
542 +
543 + This backend uses TensorRT-LLM's FP8 block-scale GEMM kernels
544 + and supports FP8+FP8 (W8A8 full quantization) on SM90+ (Hopper).
545 + """
546 + # Now call FlashInfer with BF16 input + FP8 weight, input will be
547 + # quantized with FlashInfer kernel (W8A8)
548 + output = torch.ops.vllm.flashinfer_fp8_blockscale_gemm(
549 + input=input_2d, # BF16 input
550 + weight=weight, # FP8 weight
551 + weight_scale=weight_scale, # Weight scales
552 + group_size=self.act_quant_group_shape.col,
553 + use_deep_gemm_e8m0=self.use_deep_gemm_e8m0,
554 + )
555 + return output
556 +
557 def _dispatch_w8a8_blockscale_op(
558 self,
559 use_cutlass: bool,
vllm/utils/flashinfer.py
+57
@@ -540,6 +540,59 @@ def flashinfer_scaled_fp8_mm(
540 return output
541
542
543 +flashinfer_fp8_blockscale_gemm = _lazy_import_wrapper(
544 + "flashinfer.gemm", "fp8_blockscale_gemm_sm90"
545 +)
546 +
547 +
548 +@functools.cache
549 +def has_flashinfer_fp8_blockscale_gemm() -> bool:
550 + """Return `True` if FlashInfer block-scale FP8 GEMM is available."""
551 + return (
552 + has_flashinfer()
553 + and current_platform.is_device_capability(90)
554 + and hasattr(_get_submodule("flashinfer.gemm"), "fp8_blockscale_gemm_sm90")
555 + )
556 +
557 +
558 +@functools.cache
559 +def is_flashinfer_fp8_blockscale_gemm_supported() -> bool:
560 + """Return `True` if FlashInfer block-scale FP8 GEMM is supported."""
561 + return (
562 + envs.VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER
563 + and has_flashinfer_fp8_blockscale_gemm()
564 + )
565 +
566 +
567 +def should_use_flashinfer_for_blockscale_fp8_gemm(
568 + is_flashinfer_supported: bool,
569 + output_dtype: torch.dtype,
570 + input: torch.Tensor,
571 + weight: torch.Tensor,
572 +):
573 + if not is_flashinfer_supported:
574 + return False
575 +
576 + # Verify DeepGEMM N/K dims requirements
577 + # NOTE: Also synchronized with test_w8a8_block_fp8_deep_gemm_matmul
578 + # test inside kernels/quatization/test_block_fp8.py
579 + N_MULTIPLE = 64
580 + K_MULTIPLE = 128
581 +
582 + weight_dtype = weight.dtype
583 + input_dtype = input.dtype
584 +
585 + should_use_flashinfer = (
586 + output_dtype == torch.bfloat16
587 + and input_dtype == torch.bfloat16
588 + and weight_dtype == torch.float8_e4m3fn
589 + and weight.shape[0] % N_MULTIPLE == 0
590 + and weight.shape[1] % K_MULTIPLE == 0
591 + )
592 +
593 + return should_use_flashinfer
594 +
595 +
596 __all__ = [
597 "has_flashinfer",
598 "flashinfer_trtllm_fp8_block_scale_moe",
@@ -556,10 +609,14 @@ __all__ = [
609 "has_flashinfer_all2all",
610 "has_flashinfer_cutlass_fused_moe",
611 "has_flashinfer_cutedsl_grouped_gemm_nt_masked",
612 + "has_flashinfer_fp8_blockscale_gemm",
613 "has_nvidia_artifactory",
614 "supports_trtllm_attention",
615 "can_use_trtllm_attention",
616 "use_trtllm_attention",
617 "flashinfer_scaled_fp4_mm",
618 "flashinfer_scaled_fp8_mm",
619 + "flashinfer_fp8_blockscale_gemm",
620 + "should_use_flashinfer_for_blockscale_fp8_gemm",
621 + "is_flashinfer_fp8_blockscale_gemm_supported",
622 ]