A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix][ROCm] Fix FP8 per-tensor scale rank mismatch causing Inductor assertion failure (#44912)
Signed-off-by: nehmathe2 <nehmathe2@gmail.com> Signed-off-by: Divakar Verma <divakar.verma@amd.com> Signed-off-by: nehmathe <nehmathe@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Divakar Verma <divakar.verma@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
nehmathe2 committed
Jun 16, 2026 at 20:17 UTC
2785a5e0e6eefe54c42307189f1deeb648e0a5b4
1 file changed
+6
vllm/model_executor/kernels/linear/scaled_mm/pytorch.py
+6
@@ -78,6 +78,12 @@ class PerTensorTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
78
bias: torch.Tensor | None,
79
output_shape: list,
80
) -> torch.Tensor:
81
+ # torch._scaled_mm under torch.compile does not support 0-D scales
82
+ if As.dim() == 0:
83
+ As = As.view(1)
84
+ if Bs.dim() == 0:
85
+ Bs = Bs.view(1)
86
+
87
output = torch._scaled_mm(
88
A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias
89
)