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 3, 14:04 UTC
[ROCm] Widen OAI Triton MoE capability range to include gfx12 (RDNA4) (#37826)
Signed-off-by: L.B.R. <lbr@mmonad.com> Co-authored-by: L.B.R. <lbr@mmonad.com>
L.B.R. committed
May 15, 2026 at 15:59 UTC
d792d993c11db085770de908d7bb8d95b291cbc9
1 file changed
+26
-22
vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
+26
-22
@@ -34,6 +34,30 @@ from ..utils import swiglu_limit_func
34
logger = init_logger(__name__)
35
36
37
+def _triton_kernel_moe_supports_current_device() -> bool:
38
+ # Shared device gate for the OAI Triton MoE expert classes.
39
+ # Platform-aware to avoid ROCm capability aliasing — cap (9, 0)
40
+ # matches both gfx90a (verified) and gfx906 (unverified), so we
41
+ # dispatch on gfx-string helpers instead of the cap tuple on ROCm.
42
+ p = current_platform
43
+ if p.is_cuda():
44
+ cap = p.get_device_capability()
45
+ # Keep the original `(9, 0) <= cap < (11, 0)` window on
46
+ # CUDA (covers Hopper SM90 and Blackwell SM100, excludes
47
+ # SM120) — this PR is ROCm-scoped and the broader CUDA
48
+ # range was not validated.
49
+ return cap is not None and (9, 0) <= (cap.major, cap.minor) < (11, 0)
50
+ if p.is_rocm():
51
+ from vllm.platforms.rocm import on_gfx1x, on_gfx9
52
+
53
+ # gfx9 family: gfx90a (MI200), gfx942/gfx950 (MI3xx);
54
+ # on_gfx9() already excludes gfx906/gfx908.
55
+ # gfx1x family: gfx11xx (RDNA3/3.5) and gfx12xx (RDNA4);
56
+ # on_gfx1x() excludes gfx10xx (RDNA1/RDNA2).
57
+ return on_gfx9() or on_gfx1x()
58
+ return False
59
+
60
+
61
def _patch_make_bitmatrix_metadata() -> None:
62
"""Monkey-patch make_bitmatrix_metadata to support non-power-of-2 top_k.
63
@@ -524,17 +548,7 @@ class BaseOAITritonExperts(mk.FusedMoEExpertsModular):
548
549
@staticmethod
550
def _supports_current_device() -> bool:
527
- p = current_platform
528
- if not p.is_cuda_alike():
529
- return False
530
- cap = p.get_device_capability()
531
- if cap is None:
532
- return False
533
- # (9,0) <= cap < (11,0) covers CUDA SM90 (Hopper), SM100+ (Blackwell)
534
- # and ROCm gfx942/gfx950 (which map to 9.4/9.5).
535
- if not has_triton_kernels():
536
- return False
537
- return (9, 0) <= (cap.major, cap.minor) < (11, 0)
551
+ return _triton_kernel_moe_supports_current_device() and has_triton_kernels()
552
553
@staticmethod
554
def _supports_no_act_and_mul() -> bool:
@@ -938,17 +952,7 @@ class OAITritonMxfp4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
952
953
@staticmethod
954
def _supports_current_device() -> bool:
941
- p = current_platform
942
- if not p.is_cuda_alike():
943
- return False
944
- cap = p.get_device_capability()
945
- if cap is None:
946
- return False
947
- # (9,0) <= cap < (11,0) covers CUDA SM90 (Hopper), SM100+ (Blackwell)
948
- # and ROCm gfx942/gfx950 (which map to 9.4/9.5).
949
- if not has_triton_kernels():
950
- return False
951
- return (9, 0) <= (cap.major, cap.minor) < (11, 0)
955
+ return _triton_kernel_moe_supports_current_device() and has_triton_kernels()
956
957
@staticmethod
958
def _supports_no_act_and_mul() -> bool: