[ROCm][CI][The Rock 10] Fix (MI355) Quantized Models failure on The Rock 10 with Triton 3.8.x (#54849)
Signed-off-by: Randall Smith <Randall.Smith@amd.com> Co-authored-by: Shanshan Shen <467638484@qq.com>
rasmith committed
Sep 17, 2026 at 15:20 UTC
acc2ed2a5f1fbe79575a220fc743a80c2d068ce8
1 file changed
+14
-2
vllm/model_executor/kernels/linear/mxfp8/rocm_native.py
+14
-2
@@ -119,6 +119,18 @@ def _mxfp8_dot_scaled_linear(
119
return out
120
121
122
+# Triton 3.8 enables TRITON_HIP_USE_ASYNC_COPY by default on gfx950; its extra LDS
123
+# buffer leaves no room for num_stages=3 at BLOCK_K=256.
124
+# TODO(rasmith)(Rohan138): Remove the 3.8 check once
125
+# https://github.com/vllm-project/vllm/pull/50605 merges.
126
+_BK256_STAGES = 3
127
+if triton.__version__.startswith("3.8") and current_platform.is_rocm():
128
+ from vllm.platforms.rocm import on_gfx950
129
+
130
+ if on_gfx950():
131
+ _BK256_STAGES = 2
132
+
133
+
134
def _select_cfg(M, N, K):
135
"""(BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) — graph-tuned on gfx950.
136
@@ -166,7 +178,7 @@ def _select_cfg(M, N, K):
178
# 3.6; on triton 3.7 its large BLOCK_M register/LDS footprint spills or hits
179
# "out of resources", so use 128x128x256 -- within the known-good footprint.)
180
if M >= 4096 and K >= 1024 and K % 256 == 0 and occ >= 256:
169
- return 128, 128, 256, 8, 3
181
+ return 128, 128, 256, 8, _BK256_STAGES
182
return 128, 128, 128, 8, 3
183
# large-K (K >= 2048). BLOCK_K is K-divisibility-guarded (the K-loop is unmasked):
184
# served large-K is 2048/6144 (%256==0), but fall back to 128 (always divides, since
@@ -186,7 +198,7 @@ def _select_cfg(M, N, K):
198
# Covers the qkv-class local N=1536 (TP=8 qkv / TP=4 shared_gate_up) and the deep-K
199
# / very-large-M shapes.
200
if K % 256 == 0 and (1280 < N <= 1536 or (occ >= 128 and (K >= 4096 or M >= 4096))):
189
- return 128, 128, 256, 8, 3
201
+ return 128, 128, 256, 8, _BK256_STAGES
202
# small local-N (e.g. TP=8 shared_gate_up N=768): a 64-wide BLOCK_N doubles the
203
# N-tile count -> better CU fill than 128x128 at this mid-large M (~1.4x there).
204
if N <= 1024 and K % 256 == 0: