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 8, 15:15 UTC
[ROCm][Bugfix] Resolve Dynamo tracing crash from amdsmi calls in on_gfx* arch detection (#34108)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Feb 9, 2026 at 22:50 UTC
4cde2e015944495e6bd650a4415cfb342bd73cfb
1 file changed
+27
-35
vllm/platforms/rocm.py
+27
-35
@@ -101,12 +101,10 @@ def _query_gcn_arch_from_amdsmi() -> str:
101
raise RuntimeError("amdsmi did not return valid GCN arch")
102
103
104
-@cache
105
-def _get_gcn_arch_via_amdsmi() -> str:
104
+def _get_gcn_arch() -> str:
105
"""
107
- Get the GCN architecture name using amdsmi instead of torch.cuda.
108
- This avoids initializing CUDA, which is important for Ray workers
109
- that need to set CUDA_VISIBLE_DEVICES after importing vLLM.
106
+ Get GCN arch via amdsmi (no CUDA init), fallback to torch.cuda.
107
+ Called once at module level; result stored in _GCN_ARCH.
108
"""
109
try:
110
return _query_gcn_arch_from_amdsmi()
@@ -121,34 +119,36 @@ def _get_gcn_arch_via_amdsmi() -> str:
119
return torch.cuda.get_device_properties("cuda").gcnArchName
120
121
124
-@cache
122
+# Resolve once at module load. Uses amdsmi (no CUDA init) so Ray workers
123
+# can still set CUDA_VISIBLE_DEVICES after import.
124
+# These are plain Python bools — fully torch.compile/Dynamo safe.
125
+_GCN_ARCH = _get_gcn_arch()
126
+
127
+_ON_GFX1X = any(arch in _GCN_ARCH for arch in ["gfx11", "gfx12"])
128
+_ON_MI3XX = any(arch in _GCN_ARCH for arch in ["gfx942", "gfx950"])
129
+_ON_GFX9 = any(arch in _GCN_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
130
+_ON_GFX942 = "gfx942" in _GCN_ARCH
131
+_ON_GFX950 = "gfx950" in _GCN_ARCH
132
+
133
+
134
def on_gfx1x() -> bool:
126
- GPU_ARCH = _get_gcn_arch_via_amdsmi()
127
- return any(arch in GPU_ARCH for arch in ["gfx11", "gfx12"])
135
+ return _ON_GFX1X
136
137
130
-@cache
138
def on_mi3xx() -> bool:
132
- GPU_ARCH = _get_gcn_arch_via_amdsmi()
133
- return any(arch in GPU_ARCH for arch in ["gfx942", "gfx950"])
139
+ return _ON_MI3XX
140
141
136
-@cache
142
def on_gfx9() -> bool:
138
- GPU_ARCH = _get_gcn_arch_via_amdsmi()
139
- return any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
143
+ return _ON_GFX9
144
145
142
-@cache
146
def on_gfx942() -> bool:
144
- GPU_ARCH = _get_gcn_arch_via_amdsmi()
145
- return any(arch in GPU_ARCH for arch in ["gfx942"])
147
+ return _ON_GFX942
148
149
148
-@cache
150
def on_gfx950() -> bool:
150
- GPU_ARCH = _get_gcn_arch_via_amdsmi()
151
- return any(arch in GPU_ARCH for arch in ["gfx950"])
151
+ return _ON_GFX950
152
153
154
@cache
@@ -163,13 +163,9 @@ def use_rocm_custom_paged_attention(
163
alibi_slopes: torch.Tensor | None = None,
164
sinks: torch.Tensor | None = None,
165
) -> bool:
166
- GPU_ARCH = _get_gcn_arch_via_amdsmi()
167
- ON_GFX9 = any(arch in GPU_ARCH for arch in ["gfx90a", "gfx942", "gfx950"])
168
- ON_GFX11_GFX12 = any(arch in GPU_ARCH for arch in ["gfx11", "gfx12"])
169
-
166
# custom paged attn always supported on V0. On V1, requires sliding window
167
# disabled due to observed numerical discrepancy.
172
- if ON_GFX9:
168
+ if _ON_GFX9:
169
return (
170
(sliding_window == 0 or sliding_window == (-1, -1))
171
and (qtype == torch.half or qtype == torch.bfloat16)
@@ -183,7 +179,7 @@ def use_rocm_custom_paged_attention(
179
180
else:
181
return (
186
- ON_GFX11_GFX12
182
+ _ON_GFX1X
183
and (sliding_window == 0 or sliding_window == (-1, -1))
184
and (qtype == torch.half or qtype == torch.bfloat16)
185
and head_size == 128
@@ -611,18 +607,16 @@ class RocmPlatform(Platform):
607
608
@classmethod
609
def supports_mx(cls) -> bool:
614
- gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
615
- return any(gfx in gcn_arch for gfx in ["gfx95"])
610
+ return any(gfx in _GCN_ARCH for gfx in ["gfx95"])
611
612
@classmethod
613
def supports_fp8(cls) -> bool:
619
- gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
620
- return any(gfx in gcn_arch for gfx in ["gfx94", "gfx95", "gfx12"])
614
+ return any(gfx in _GCN_ARCH for gfx in ["gfx94", "gfx95", "gfx12"])
615
616
@classmethod
617
def is_fp8_fnuz(cls) -> bool:
618
# only device 0 is checked, this assumes MI300 platforms are homogeneous
625
- return "gfx94" in torch.cuda.get_device_properties(0).gcnArchName
619
+ return "gfx94" in _GCN_ARCH
620
621
@classmethod
622
def fp8_dtype(cls) -> torch.dtype:
@@ -634,9 +628,7 @@ class RocmPlatform(Platform):
628
@classmethod
629
def use_custom_allreduce(cls) -> bool:
630
# We only enable custom allreduce for MI300 series
637
- gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
638
- supported_archs = ["gfx94", "gfx95"]
639
- return any(gfx in gcn_arch for gfx in supported_archs)
631
+ return any(gfx in _GCN_ARCH for gfx in ["gfx94", "gfx95"])
632
633
@classmethod
634
def opaque_attention_op(cls) -> bool:
@@ -644,7 +636,7 @@ class RocmPlatform(Platform):
636
637
@classmethod
638
def is_navi(cls) -> bool:
647
- return "gfx1" in torch.cuda.get_device_properties(0).gcnArchName
639
+ return "gfx1" in _GCN_ARCH
640
641
@classmethod
642
def get_static_graph_wrapper_cls(cls) -> str: