[ROCm] Gate Torch FP8 scaled-MM on architecture support (#51021)
Signed-off-by: sstamenk <strahinja.stamenkovic@amd.com> Signed-off-by: Strahinja Stamenkovic <strahinja.stamenkovic@amd.com> Signed-off-by: Strahinja Stamenkovic <sstamenk@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Strahinja Stamenkovic committed
Aug 18, 2026 at 19:54 UTC
ddbf826bee965ace7d2a68599fb45a69e79745d2
1 file changed
+26
-17
vllm/model_executor/kernels/linear/scaled_mm/pytorch.py
+26
-17
@@ -24,6 +24,27 @@ def _get_num_tokens(output_shape: list) -> int:
24
return math.prod(output_shape[:-1])
25
26
27
+def _rocm_torch_fp8_scaled_mm_supported() -> bool:
28
+ from vllm.platforms.rocm import on_gfx12x, on_gfx942, on_gfx950, on_gfx1250
29
+
30
+ return on_gfx942() or on_gfx950() or on_gfx12x() or on_gfx1250()
31
+
32
+
33
+def _supports_torch_fp8_scaled_mm() -> bool:
34
+ if current_platform.is_cpu():
35
+ return True
36
+ if current_platform.is_xpu():
37
+ return True
38
+ if not current_platform.is_cuda_alike():
39
+ return False
40
+
41
+ # TODO: Use torch.cuda.is_scaled_mm_supported once it is in the supported
42
+ # PyTorch baseline.
43
+ if current_platform.is_rocm():
44
+ return _rocm_torch_fp8_scaled_mm_supported()
45
+ return current_platform.supports_fp8()
46
+
47
+
48
class TorchFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
49
"""
50
Base class for FP8 linear kernels using Torch.
@@ -35,17 +56,10 @@ class TorchFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
56
def is_supported(
57
cls, compute_capability: int | None = None
58
) -> tuple[bool, str | None]:
38
- if not (
39
- current_platform.is_cuda_alike()
40
- or current_platform.is_cpu()
41
- or current_platform.is_xpu()
42
- ):
43
- return False, "requires ROCm, CUDA, CPU or XPU."
44
-
45
- if compute_capability is not None and compute_capability < 89:
46
- return False, "requires compute capability 89 and above."
59
+ if _supports_torch_fp8_scaled_mm():
60
+ return True, None
61
48
- return True, None
62
+ return False, "requires a platform with torch FP8 scaled-MM support."
63
64
def get_output_padding(self) -> int | None:
65
# Note: we pad the input because torch._scaled_mm is more performant
@@ -111,13 +125,8 @@ class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
125
if not current_platform.is_rocm():
126
return False, "requires ROCm."
127
114
- from vllm.platforms.rocm import get_cdna_version, on_rdna4
115
-
116
- if get_cdna_version() <= 2 and not on_rdna4():
117
- return False, "requires CDNA3+ or RDNA4"
118
-
119
- if compute_capability is not None and compute_capability < 94:
120
- return False, "requires compute capability 94 and above."
128
+ if not _supports_torch_fp8_scaled_mm():
129
+ return False, "requires platform with torch FP8 scaled-MM support."
130
131
return True, None
132