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, 18:45 UTC
[Bugfix] Fix RMS norm + quant fusion on DeepGEMM UE8M0 path for B200 (#40552)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com>
Lucas Kabela committed
Apr 22, 2026 at 15:04 UTC
b8401a9bf462cbbdcd99a1c6521af8301ad8fa1b
2 files changed
+22
tests/compile/passes/test_fusion.py
+21
@@ -51,6 +51,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
51
)
52
from vllm.platforms import current_platform
53
from vllm.utils.deep_gemm import (
54
+ is_deep_gemm_e8m0_used,
55
is_deep_gemm_supported,
56
)
57
@@ -317,6 +318,26 @@ def test_fusion_rmsnorm_quant(
318
):
319
pytest.skip("Unsupported group shape 64 for CUTLASS/DeepGemm")
320
321
+ # TODO(quant-rms-fusion): DeepGEMM UE8M0 activation quant on B200 lowers
322
+ # to a packed int32-scale op (per_token_group_quant_fp8_packed_for_deepgemm),
323
+ # but the rms+quant fusion pattern only matches the fp32-scale variant, so
324
+ # the fused output gets a mismatched scale layout and produces NaN. Only
325
+ # reproduces on bf16 (DeepGEMM UE8M0 on B200 is bf16-only).
326
+ # To re-enable: make rms_norm_per_block_quant emit packed UE8M0 scales
327
+ # and extend the fusion pattern to rewrite the packed activation quant.
328
+ deepgemm_kernels = (
329
+ DeepGemmFp8BlockScaledMMKernel,
330
+ FlashInferFp8DeepGEMMDynamicBlockScaledKernel,
331
+ )
332
+ if (
333
+ dtype == torch.bfloat16
334
+ and force_kernel in deepgemm_kernels
335
+ and is_deep_gemm_e8m0_used()
336
+ ):
337
+ pytest.skip(
338
+ "rms+quant fusion does not yet match the packed UE8M0 DeepGEMM path"
339
+ )
340
+
341
custom_ops = []
342
if enable_rms_norm_custom_op:
343
custom_ops.append("+rms_norm")
tests/utils.py
+1
@@ -1826,6 +1826,7 @@ class TestFP8Layer(torch.nn.Module):
1826
self.weight = torch.rand(weight_shape).to(dtype=FP8_DTYPE)
1827
self.input_scale = None
1828
self.weight_scale = None
1829
+ self.weight_block_size = [block_size, block_size]
1830
if transpose_weights:
1831
self.weight = self.weight.t()
1832
else: