[Bugfix] Guard CUDA-only rms_norm_per_block_quant in FUSED_OPS for non-CUDA builds (#47296)
Signed-off-by: Tsvika Shapira <tsvika@moonmath.ai> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Tsvika Shapira committed
Jul 9, 2026 at 17:10 UTC
429f4057487076d569f4a80413738462513dbe5a
1 file changed
+7
-12
vllm/compilation/passes/fusion/rms_quant_fusion.py
+7
-12
@@ -128,19 +128,14 @@ FUSED_OPS: dict[FusedRMSQuantKey, OpOverload] = {
128
FusedRMSQuantKey(
129
kFp8DynamicTokenSym, True
130
): torch.ops._C.rms_norm_dynamic_per_token_quant.default, # noqa: E501
131
- FusedRMSQuantKey(
132
- kFp8Dynamic128Sym, False
133
- ): torch.ops._C.rms_norm_per_block_quant.default, # noqa: E501
134
- FusedRMSQuantKey(
135
- kFp8Dynamic128Sym, True
136
- ): torch.ops._C.rms_norm_per_block_quant.default, # noqa: E501
137
- FusedRMSQuantKey(
138
- kFp8Dynamic64Sym, False
139
- ): torch.ops._C.rms_norm_per_block_quant.default, # noqa: E501
140
- FusedRMSQuantKey(
141
- kFp8Dynamic64Sym, True
142
- ): torch.ops._C.rms_norm_per_block_quant.default, # noqa: E501
131
}
132
+# rms_norm_per_block_quant is CUDA-only; guard it like per_token_group_fp8_quant above.
133
+if hasattr(torch.ops._C, "rms_norm_per_block_quant"):
134
+ _rms_norm_per_block_quant = torch.ops._C.rms_norm_per_block_quant.default
135
+ FUSED_OPS[FusedRMSQuantKey(kFp8Dynamic128Sym, False)] = _rms_norm_per_block_quant # noqa: E501
136
+ FUSED_OPS[FusedRMSQuantKey(kFp8Dynamic128Sym, True)] = _rms_norm_per_block_quant # noqa: E501
137
+ FUSED_OPS[FusedRMSQuantKey(kFp8Dynamic64Sym, False)] = _rms_norm_per_block_quant # noqa: E501
138
+ FUSED_OPS[FusedRMSQuantKey(kFp8Dynamic64Sym, True)] = _rms_norm_per_block_quant # noqa: E501
139
140
141
class RMSNormQuantPattern: