@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e9f167129c

[CI][ROCM] Add the Fusion E2E TP2 Quick group on MI355, and the AITER MLA fix it needs (#58369)

Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com> Signed-off-by: Andreas Karatzas <andreas.karatzas@protonmail.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Andreas Karatzas <andreas.karatzas@protonmail.com> Co-authored-by: Codex <noreply@openai.com>

stefankoncarevic committed Sep 24, 2026 at 04:01 UTC e9f167129c09d5330805c2042c16329e9b218c5a
7 files changed +220 -16
.buildkite/test-amd.yaml
+41
@@ -3773,6 +3773,47 @@ steps:
3773 - pytest -v -s tests/distributed/test_rocm_quick_reduce.py
3774 - pytest -v -s tests/distributed/test_quick_all_reduce.py
3775
3776 +- label: ":amd: (MI355) Fusion E2E TP2 Quick" # TBD
3777 + timeout_in_minutes: 60
3778 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3779 + dind: false
3780 + agent_pool: mi355_2
3781 + num_gpus: 2
3782 + optional: true
3783 + working_dir: "/vllm-workspace/"
3784 + source_file_dependencies:
3785 + - csrc/quantization/
3786 + - vllm/compilation/
3787 + - vllm/distributed/
3788 + - vllm/model_executor/
3789 + - vllm/v1/attention/
3790 + - tests/compile/fusions_e2e/
3791 + - vllm/_aiter_ops.py
3792 + - vllm/platforms/rocm.py
3793 + commands:
3794 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1
3795 + - >-
3796 + VLLM_ROCM_USE_AITER=1
3797 + VLLM_ROCM_USE_AITER_CUSTOM_AR=1
3798 + pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py
3799 + -k 'test_tp2_ar_rms_fp8_fusions and inductor_partition and
3800 + not +quant_fp8 and not +rms_norm and Llama-3 and
3801 + ROCM_AITER_UNIFIED_ATTN'
3802 + - >-
3803 + VLLM_ROCM_USE_AITER=1
3804 + VLLM_ROCM_USE_AITER_CUSTOM_AR=1
3805 + pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py
3806 + -k 'test_tp2_ar_rms_fp8_fusions and inductor_partition and
3807 + +quant_fp8 and +rms_norm and qwen3 and
3808 + ROCM_AITER_UNIFIED_ATTN'
3809 + - >-
3810 + VLLM_ROCM_USE_AITER=1
3811 + VLLM_ROCM_USE_AITER_CUSTOM_AR=1
3812 + pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py
3813 + -k 'test_tp2_ar_rms_fp8_fusions and inductor_partition and
3814 + not +quant_fp8 and not +rms_norm and DeepSeek-Coder and
3815 + ROCM_AITER_MLA'
3816 +
3817 #-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------#
3818
3819 - label: ":amd: (MI355) Entrypoints Integration (API Server) %N"
.buildkite/test_areas/compile.yaml
+41
@@ -206,6 +206,47 @@ steps:
206 # Run all models and attn backends but only Inductor partition and native custom ops
207 - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))"
208 - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))"
209 + mirror:
210 + amd:
211 + label: ":amd: (MI355) Fusion E2E TP2 Quick"
212 + dind: false
213 + device: mi355_2
214 + timeout_in_minutes: 60
215 + depends_on:
216 + - image-build-amd
217 + working_dir: /vllm-workspace/
218 + source_file_dependencies:
219 + - csrc/quantization/
220 + - vllm/compilation/
221 + - vllm/distributed/
222 + - vllm/model_executor/
223 + - vllm/v1/attention/
224 + - tests/compile/fusions_e2e/
225 + - vllm/_aiter_ops.py
226 + - vllm/platforms/rocm.py
227 + commands:
228 + - export VLLM_TEST_CLEAN_GPU_MEMORY=1
229 + - >-
230 + VLLM_ROCM_USE_AITER=1
231 + VLLM_ROCM_USE_AITER_CUSTOM_AR=1
232 + pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py
233 + -k 'test_tp2_ar_rms_fp8_fusions and inductor_partition and
234 + not +quant_fp8 and not +rms_norm and Llama-3 and
235 + ROCM_AITER_UNIFIED_ATTN'
236 + - >-
237 + VLLM_ROCM_USE_AITER=1
238 + VLLM_ROCM_USE_AITER_CUSTOM_AR=1
239 + pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py
240 + -k 'test_tp2_ar_rms_fp8_fusions and inductor_partition and
241 + +quant_fp8 and +rms_norm and qwen3 and
242 + ROCM_AITER_UNIFIED_ATTN'
243 + - >-
244 + VLLM_ROCM_USE_AITER=1
245 + VLLM_ROCM_USE_AITER_CUSTOM_AR=1
246 + pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py
247 + -k 'test_tp2_ar_rms_fp8_fusions and inductor_partition and
248 + not +quant_fp8 and not +rms_norm and DeepSeek-Coder and
249 + ROCM_AITER_MLA'
250
251 - label: ":nvidia: (H100) Fusion E2E TP2 AR-RMS Config Sweep"
252 key: fusion-e2e-tp2-ar-rms-config-sweep-h100
tests/compile/fusions_e2e/models.py
+9
@@ -44,6 +44,15 @@ ROCM_AITER_UNIFIED_ATTN = pytest.param(
44 ),
45 )
46
47 +ROCM_AITER_MLA_ATTN = pytest.param(
48 + AttentionBackendCase(backend=AttentionBackendEnum.ROCM_AITER_MLA),
49 + id="ROCM_AITER_MLA",
50 + marks=pytest.mark.skipif(
51 + not is_aiter_found_and_supported(),
52 + reason="ROCM_AITER_MLA only for AMD when AITER is installed",
53 + ),
54 +)
55 +
56 FLASHINFER_MLA_ATTN = pytest.param(
57 AttentionBackendCase(backend=AttentionBackendEnum.FLASHINFER_MLA),
58 id="FLASHINFER_MLA",
tests/compile/fusions_e2e/test_tp2_ar_rms.py
+57 -8
@@ -19,6 +19,7 @@ from .common import (
19 from .models import (
20 FLASHINFER_ATTN,
21 FLASHINFER_MLA_ATTN,
22 + ROCM_AITER_MLA_ATTN,
23 ROCM_AITER_UNIFIED_ATTN,
24 ROCM_ATTN,
25 TRITON_ATTN,
@@ -40,25 +41,35 @@ pytestmark = pytest.mark.skipif(
41 )
42
43
43 -@multi_gpu_test(num_gpus=2)
44 -@pytest.mark.parametrize(
45 - "model_name, matches_fn, model_kwargs, hf_overrides",
46 - # qwen3 & dsv3 should still fuse AR+rms even though group quant is not yet supported
44 +# qwen3 & dsv3 should still fuse AR+rms even though group quant is not yet
45 +# supported. ROCm carries only the payloads its two backends support.
46 +TP2_FP8_MODELS = (
47 [
48 llama3_8b_fp8,
49 llama4_scout_fp8,
50 qwen3_a3b_fp8,
51 deepseek_coder_v2_lite_fp8,
52 deepseek_v3_fp8,
53 - ],
53 + ]
54 + if current_platform.is_cuda()
55 + else [llama3_8b_fp8, qwen3_a3b_fp8, deepseek_coder_v2_lite_fp8]
56 +)
57 +
58 +TP2_FP8_ATTN_BACKENDS = (
59 + [TRITON_ATTN, FLASHINFER_ATTN, FLASHINFER_MLA_ATTN]
60 + if current_platform.is_cuda()
61 + else [ROCM_AITER_UNIFIED_ATTN, ROCM_AITER_MLA_ATTN]
62 )
63 +
64 +
65 +@multi_gpu_test(num_gpus=2)
66 @pytest.mark.parametrize(
56 - "attn_backend", [TRITON_ATTN, FLASHINFER_ATTN, FLASHINFER_MLA_ATTN]
67 + "model_name, matches_fn, model_kwargs, hf_overrides", TP2_FP8_MODELS
68 )
69 +@pytest.mark.parametrize("attn_backend", TP2_FP8_ATTN_BACKENDS)
70 @pytest.mark.parametrize("n_layers", [4])
71 @pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
72 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
61 -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test CUDA")
73 def test_tp2_ar_rms_fp8_fusions(
74 model_name: str,
75 matches_fn: Callable[[int], Matches],
@@ -69,10 +80,27 @@ def test_tp2_ar_rms_fp8_fusions(
80 custom_ops: str,
81 inductor_graph_partition: bool,
82 run_e2e_fusion_test,
83 + monkeypatch,
84 ):
85 matches = matches_fn(n_layers)
86
75 - block_fp8 = "qwen" in model_name.lower() or "deepseek" in model_name.lower()
87 + model_name_lower = model_name.lower()
88 + # DeepSeek-Coder on AITER MLA is the one static-FP8 payload here, validated
89 + # against the native QuantFP8 pattern; every other DeepSeek stays group
90 + # quantized as on CUDA.
91 + rocm_static_fp8_mla = (
92 + current_platform.is_rocm()
93 + and "deepseek-coder" in model_name_lower
94 + and attn_backend.backend.name == "ROCM_AITER_MLA"
95 + )
96 + block_fp8 = "qwen" in model_name_lower or "deepseek" in model_name_lower
97 + if rocm_static_fp8_mla:
98 + block_fp8 = False
99 + if custom_ops != "-quant_fp8,-rms_norm":
100 + pytest.skip("ROCm static-FP8 MLA parity uses native quant and RMSNorm")
101 + # AITER's MoE GEMM rejects this reduced dummy model shape; keep AITER
102 + # for MLA and the all-reduce, and let the MoE run on Triton.
103 + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "0")
104 if block_fp8 and "-quant_fp8" in custom_ops:
105 # This is why config forces +quant_fp8 by default
106 pytest.skip("native QuantFP8 matching not supported for group quant")
@@ -83,6 +111,13 @@ def test_tp2_ar_rms_fp8_fusions(
111 model_kwargs["max_model_len"] = 1024
112 model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
113 model_kwargs["disable_custom_all_reduce"] = False
114 + if current_platform.is_rocm() and "llama-3" in model_name_lower:
115 + model_kwargs.update(
116 + gpu_memory_utilization=0.01,
117 + kv_cache_memory_bytes=64 * 1024**2,
118 + max_num_batched_tokens=1024,
119 + max_num_seqs=32,
120 + )
121
122 compilation_config = dict(
123 use_inductor_graph_partition=inductor_graph_partition,
@@ -103,6 +138,19 @@ def test_tp2_ar_rms_fp8_fusions(
138 "ar_rms_fusion",
139 ]
140
141 + use_aiter = current_platform.is_rocm()
142 + if use_aiter:
143 + matches = matches._replace(aiter_ar_rms_fusion=matches.ar_rms_fusion)
144 + # The AITER all-reduce pass registers the AR+RMS(+group-quant) patterns
145 + # before its AR+RMS-only ones, so the standalone AITER RMS-quant pass
146 + # does not own the post-AR sites here. ROCm's contract is native
147 + # QK-norm+RoPE and attention-output quant fusion plus AITER AR+RMS.
148 + matches_check = [
149 + "norm_rope_fusion",
150 + "attn_quant_fusion",
151 + "aiter_ar_rms_fusion",
152 + ]
153 +
154 run_e2e_fusion_test(
155 model_name,
156 matches,
@@ -111,6 +159,7 @@ def test_tp2_ar_rms_fp8_fusions(
159 compilation_config,
160 matches_check,
161 tp_size=2,
162 + use_aiter=use_aiter,
163 )
164
165
tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py
+19
@@ -151,6 +151,25 @@ def test_decode_and_verify_are_disjoint(kv_cache_dtype, num_heads):
151 )
152
153
154 +@pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES)
155 +def test_oversized_kv_cache_refuses_gluon(gluon_available, kv_cache_dtype):
156 + """Past 2 GiB per layer Gluon drops its KV bounds mask and reads out of range.
157 +
158 + ``None`` is the profiling case, before the cache is sized: nothing to
159 + refuse yet, so the usual route stands.
160 + """
161 + bound = rocm_aiter_mla._GLUON_MAX_KV_CACHE_BYTES
162 + for kv_cache_bytes in (None, bound):
163 + assert AiterMLAHelper.use_gluon_decode(8, 1, kv_cache_dtype, kv_cache_bytes)
164 + assert AiterMLAHelper.use_gluon_verify(
165 + 8, 8, kv_cache_dtype, kv_cache_bytes=kv_cache_bytes
166 + )
167 + assert not AiterMLAHelper.use_gluon_decode(8, 1, kv_cache_dtype, bound + 1)
168 + assert not AiterMLAHelper.use_gluon_verify(
169 + 8, 8, kv_cache_dtype, kv_cache_bytes=bound + 1
170 + )
171 +
172 +
173 @pytest.mark.parametrize("num_heads", [1, 2, 3, 5, 6, 7, 8, 9, 12, 15])
174 def test_padded_query_is_contiguous(num_heads):
175 """asm_mla.cu:805 requires Q.is_contiguous().
tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
+3 -1
@@ -107,6 +107,8 @@ def _builder(
107 _dcp_verify_buffers=None,
108 _graph_seq_lens=None,
109 _kv_cache_dtype_str=kv_cache_dtype,
110 + # Unsized as during profiling, so the Gluon KV bound stays out of the way.
111 + _kv_cache_bytes=None,
112 paged_kv_last_page_len=torch.ones(max_decode_rows, dtype=torch.int32),
113 paged_kv_indices=torch.empty(1024, dtype=torch.int32),
114 paged_kv_indptr=torch.empty(max_decode_rows + 1, dtype=torch.int32),
@@ -585,7 +587,7 @@ def test_mtp_builder_init_sizes_native_fp8_metadata(
587 get_num_attention_heads=lambda parallel_config: num_heads,
588 ),
589 scheduler_config=SimpleNamespace(max_num_seqs=2),
588 - cache_config=SimpleNamespace(cache_dtype="fp8_e4m3"),
590 + cache_config=SimpleNamespace(cache_dtype="fp8_e4m3", num_gpu_blocks=None),
591 compilation_config=SimpleNamespace(
592 cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False),
593 # Empty: the per-layer head-count probe finds no attention layer and
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+50 -7
@@ -178,6 +178,30 @@ def _gluon_mla_decode_supported() -> bool:
178 return on_gfx950()
179
180
181 +# Past 2 GiB, mla_gluon swaps its masked buffer_load for an unmasked
182 +# global_load that reads outside the cache and aborts the process. Mirrors the
183 +# kernel's own within_2gb bound, so the fallback engages exactly where it does.
184 +_GLUON_MAX_KV_CACHE_BYTES = 1 << 31
185 +
186 +
187 +def _gluon_kv_cache_in_bounds(kv_cache_bytes: int | None) -> bool:
188 + """Whether this layer's KV cache keeps Gluon on its bounds-checked path.
189 +
190 + The cache is one flat tensor per layer, so the bound disqualifies the layer
191 + outright; no batch shape brings an oversized cache back in range. ``None``
192 + is the not-yet-sized case during profiling, before any kernel runs.
193 + """
194 + if kv_cache_bytes is None or kv_cache_bytes <= _GLUON_MAX_KV_CACHE_BYTES:
195 + return True
196 + logger.warning_once(
197 + "KV cache is %.1f GiB per layer, past the 2 GiB where the Gluon MLA "
198 + "kernel drops its KV bounds mask; using the padded ASM decode instead. "
199 + "Lower --gpu-memory-utilization or --max-model-len to get Gluon back.",
200 + kv_cache_bytes / (1 << 30),
201 + )
202 + return False
203 +
204 +
205 @functools.lru_cache(maxsize=1)
206 def _get_segmented_mla_decode():
207 """Load AITER's segmented MLA decode with unreduced partial output."""
@@ -564,6 +588,13 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
588 # _build_decode needs the cache dtype to pick the decode kernel; keep
589 # the normalized string instead of dropping it at the end of __init__.
590 self._kv_cache_dtype_str = kv_cache_dtype_str
591 + # Sized per layer, matching the flat KV tensor the kernel is handed.
592 + num_gpu_blocks = vllm_config.cache_config.num_gpu_blocks
593 + self._kv_cache_bytes = (
594 + None
595 + if num_gpu_blocks is None
596 + else num_gpu_blocks * kv_cache_spec.page_size_bytes
597 + )
598 # MLAAttention quantizes decode Q to FP8 before calling this backend
599 # whenever the KV cache is FP8 and supports_quant_query_input is true.
600 q_dtype = (
@@ -1065,6 +1096,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1096 self._decode_num_heads,
1097 int(max_qo_len),
1098 self._kv_cache_dtype_str,
1099 + self._kv_cache_bytes,
1100 )
1101 use_gluon_verify = AiterMLAHelper.use_gluon_verify(
1102 self._decode_num_heads,
@@ -1072,6 +1104,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1104 self._kv_cache_dtype_str,
1105 self.dcp_world_size,
1106 causal,
1107 + self._kv_cache_bytes,
1108 )
1109 use_segmented_dcp_verify = (
1110 self._supports_segmented_dcp_verify and max_qo_len > 1
@@ -1481,7 +1514,12 @@ class AiterMLAHelper:
1514 return AiterMLAHelper._get_mla_unpadded_heads(num_heads, lse)
1515
1516 @staticmethod
1484 - def use_gluon_decode(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool:
1517 + def use_gluon_decode(
1518 + num_heads: int,
1519 + max_qo_len: int,
1520 + kv_cache_dtype: str,
1521 + kv_cache_bytes: int | None = None,
1522 + ) -> bool:
1523 # Small-head (<16) single-token decode takes either the Gluon kernel or
1524 # the padded asm persistent decode, selected by
1525 # VLLM_ROCM_AITER_MLA_ASM_PADDING and the arch (Gluon is gfx950 only).
@@ -1497,10 +1535,12 @@ class AiterMLAHelper:
1535 mode = _aiter_mla_small_head_mode()
1536 if mode == "asm":
1537 return False
1500 - gluon_supported = _gluon_mla_decode_supported()
1501 - if mode == "gluon":
1502 - return gluon_supported
1503 - return m % num_heads == 0 and gluon_supported
1538 + if not _gluon_mla_decode_supported():
1539 + return False
1540 + if mode != "gluon" and m % num_heads != 0:
1541 + return False
1542 + # Last, so the size warning only fires where Gluon would have run.
1543 + return _gluon_kv_cache_in_bounds(kv_cache_bytes)
1544
1545 @staticmethod
1546 def use_gluon_verify(
@@ -1509,6 +1549,7 @@ class AiterMLAHelper:
1549 kv_cache_dtype: str,
1550 dcp_world_size: int = 1,
1551 causal: bool = True,
1552 + kv_cache_bytes: int | None = None,
1553 ) -> bool:
1554 """Whether a small-head multi-token verify uses native Gluon MTP.
1555
@@ -1536,8 +1577,10 @@ class AiterMLAHelper:
1577 return False
1578 if not _gluon_mla_decode_supported():
1579 return False
1539 - # Same arch and mode gating as use_gluon_decode.
1540 - return _aiter_mla_small_head_mode() != "asm"
1580 + # Same arch, mode and cache-size gating as use_gluon_decode.
1581 + if _aiter_mla_small_head_mode() == "asm":
1582 + return False
1583 + return _gluon_kv_cache_in_bounds(kv_cache_bytes)
1584
1585 @staticmethod
1586 def dcp_local_verify_row_lens(