[Bugfix] Initialize DeepGemmQuantScaleFMT oracle lazily; bound QuantFP8 UE8M0 packed path to group_size 128 (#51359)
Signed-off-by: BabyDrangoner <148877251+BabyDrangoner@users.noreply.github.com> Signed-off-by: kiroxu <148877251+BabyDrangoner@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
kiroxu committed
Aug 13, 2026 at 02:27 UTC
19d61b1273a56e0ebd8a9160bffb2e75dccad82b
3 files changed
+30
-2
tests/kernels/quantization/test_fp8_quant_group.py
+15
@@ -73,6 +73,21 @@ def test_quantfp8_group_functionality(
73
if is_divisible:
74
x_quant_cuda, scales_cuda = quant_op.forward_cuda(x.clone())
75
assert x_quant_cuda.shape == x.shape
76
+
77
+ if scales_cuda.dtype == torch.int32:
78
+ # DeepGEMM UE8M0 fast path (group_size == 128 on Blackwell with
79
+ # the UE8M0 oracle): scales are packed 4-per-int32 in a
80
+ # TMA-aligned layout. Unpack and compare against native.
81
+ assert use_ue8m0 and group_size == 128
82
+ packed_groups = (expected_num_groups + 3) // 4
83
+ assert scales_cuda.shape == (batch_size, packed_groups)
84
+ unpacked = (
85
+ scales_cuda.contiguous()
86
+ .view(torch.uint8)
87
+ .reshape(batch_size, -1)[:, :expected_num_groups]
88
+ )
89
+ scales_cuda = torch.pow(2.0, unpacked.float() - 127)
90
+
91
assert scales_cuda.shape == (batch_size, expected_num_groups)
92
93
# Verify CUDA/native consistency
vllm/model_executor/layers/quantization/input_quant_fp8.py
+2
@@ -94,6 +94,8 @@ class QuantFP8(CustomOp):
94
self.is_group_quant
95
and self.use_ue8m0
96
and self.use_deep_gemm_supported
97
+ and not self.static
98
+ and self.group_size == 128
99
and (DeepGemmQuantScaleFMT.from_oracle() == DeepGemmQuantScaleFMT.UE8M0)
100
):
101
return fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
vllm/utils/deep_gemm.py
+13
-2
@@ -84,9 +84,20 @@ class DeepGemmQuantScaleFMT(Enum):
84
85
@classmethod
86
def from_oracle(cls) -> "DeepGemmQuantScaleFMT":
87
- """Return the pre-initialized oracle decision"""
87
+ """Return the oracle decision, initializing it on first use.
88
+
89
+ The cache is normally populated by ``_lazy_init()`` (e.g. during
90
+ engine startup), but standalone consumers such as ``QuantFP8`` with an
91
+ explicit ``use_ue8m0=True`` can reach this before any DeepGEMM kernel
92
+ wrapper has run. Resolve the DeepGEMM symbols and initialize the
93
+ decision here instead of asserting; without DeepGEMM this yields
94
+ FLOAT32, matching ``is_deep_gemm_e8m0_used()``.
95
+ """
96
cached = getattr(cls, "_oracle_cache", None)
89
- assert cached is not None, "DeepGemmQuantScaleFMT oracle cache not initialized"
97
+ if cached is None:
98
+ _lazy_init()
99
+ cls.init_oracle_cache()
100
+ cached = cls._oracle_cache # type: ignore[attr-defined]
101
return cached
102
103