[ROCm][CI] Add MI355 TP2 AR-RMS and B200 fusion mirrors (#58284)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Andreas Karatzas committed
Sep 29, 2026 at 02:10 UTC
78fc1e6274d326c869432e4dfde8dc5b4c6e85a8
4 files changed
+165
-1
.buildkite/test-amd.yaml
+50
@@ -3228,6 +3228,56 @@ steps:
3228
commands:
3229
- pytest -v -s compile/h100/test_startup.py
3230
3231
+- label: ":amd: (MI355) Fusion E2E TP2 AR-RMS FP8 Config Sweep"
3232
+ timeout_in_minutes: 45
3233
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3234
+ dind: false
3235
+ agent_pool: mi355_2
3236
+ num_gpus: 2
3237
+ optional: true
3238
+ working_dir: "/vllm-workspace/"
3239
+ source_file_dependencies:
3240
+ - vllm/compilation/
3241
+ - vllm/model_executor/
3242
+ - vllm/distributed/
3243
+ - vllm/v1/attention/
3244
+ - vllm/_aiter_ops.py
3245
+ - vllm/platforms/rocm.py
3246
+ - tests/compile/fusions_e2e/
3247
+ - tests/utils.py
3248
+ - tests/conftest.py
3249
+ - csrc/quantization/
3250
+ commands:
3251
+ # Existing MI300 jobs own BF16; TP2 Quick owns the excluded Llama cell.
3252
+ - amd-smi
3253
+ - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and Llama-3 and ROCM_AITER_UNIFIED_ATTN and (dynamo_partition or +quant_fp8 or +rms_norm)'
3254
+
3255
+- label: ":amd: (MI355) Fusion E2E TP2 FP8 BF16 rewrites"
3256
+ timeout_in_minutes: 45
3257
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3258
+ dind: false
3259
+ agent_pool: mi355_2
3260
+ num_gpus: 2
3261
+ optional: true
3262
+ working_dir: "/vllm-workspace/"
3263
+ source_file_dependencies:
3264
+ - vllm/compilation/
3265
+ - vllm/model_executor/
3266
+ - vllm/distributed/
3267
+ - vllm/v1/attention/
3268
+ - vllm/_aiter_ops.py
3269
+ - vllm/platforms/rocm.py
3270
+ - tests/compile/fusions_e2e/
3271
+ - tests/utils.py
3272
+ - tests/conftest.py
3273
+ - csrc/quantization/
3274
+ commands:
3275
+ # Dummy-weight rewrite coverage; trained FP8 AR/RMS equivalence is unresolved.
3276
+ # TODO: FP4 output-quant and trained MXFP4 AR/RMS fusion remain unsupported.
3277
+ - amd-smi
3278
+ - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and qwen3 and ROCM_AITER_UNIFIED_ATTN and +quant_fp8 and (dynamo_partition or not +rms_norm)'
3279
+ - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_moe.py
3280
+
3281
- label: ":amd: (MI355) Fusion and Compile FP8"
3282
timeout_in_minutes: 45
3283
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
.buildkite/test_areas/compile.yaml
+52
@@ -267,6 +267,31 @@ steps:
267
- nvidia-smi
268
# Run just llama3 (fp8 & bf16) for all config combinations
269
- pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "llama-3"
270
+ mirror:
271
+ amd:
272
+ label: ":amd: (MI355) Fusion E2E TP2 AR-RMS FP8 Config Sweep"
273
+ dind: false
274
+ device: mi355_2
275
+ num_devices: 2
276
+ timeout_in_minutes: 45
277
+ working_dir: "/vllm-workspace/"
278
+ depends_on:
279
+ - image-build-amd
280
+ source_file_dependencies:
281
+ - vllm/compilation/
282
+ - vllm/model_executor/
283
+ - vllm/distributed/
284
+ - vllm/v1/attention/
285
+ - vllm/_aiter_ops.py
286
+ - vllm/platforms/rocm.py
287
+ - tests/compile/fusions_e2e/
288
+ - tests/utils.py
289
+ - tests/conftest.py
290
+ - csrc/quantization/
291
+ commands:
292
+ # Existing MI300 jobs own BF16; TP2 Quick owns the excluded Llama cell.
293
+ - amd-smi
294
+ - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and Llama-3 and ROCM_AITER_UNIFIED_ATTN and (dynamo_partition or +quant_fp8 or +rms_norm)'
295
296
- label: ":nvidia: (H100) Fusion E2E TP2 AsyncTP Config Sweep"
297
key: fusion-e2e-tp2-asynctp-config-sweep-h100
@@ -334,3 +359,30 @@ steps:
359
# for ar-rms-quant-fp4, also sweep llama3
360
- pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "(FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))) or Llama-3.1-8B-Instruct-FP4"
361
- pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))"
362
+ mirror:
363
+ amd:
364
+ label: ":amd: (MI355) Fusion E2E TP2 FP8 BF16 rewrites"
365
+ dind: false
366
+ device: mi355_2
367
+ num_devices: 2
368
+ timeout_in_minutes: 45
369
+ working_dir: "/vllm-workspace/"
370
+ depends_on:
371
+ - image-build-amd
372
+ source_file_dependencies:
373
+ - vllm/compilation/
374
+ - vllm/model_executor/
375
+ - vllm/distributed/
376
+ - vllm/v1/attention/
377
+ - vllm/_aiter_ops.py
378
+ - vllm/platforms/rocm.py
379
+ - tests/compile/fusions_e2e/
380
+ - tests/utils.py
381
+ - tests/conftest.py
382
+ - csrc/quantization/
383
+ commands:
384
+ # Dummy-weight rewrite coverage; trained FP8 AR/RMS equivalence is unresolved.
385
+ # TODO: FP4 output-quant and trained MXFP4 AR/RMS fusion remain unsupported.
386
+ - amd-smi
387
+ - VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_CUSTOM_AR=1 pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k 'test_tp2_ar_rms_fp8_fusions and qwen3 and ROCM_AITER_UNIFIED_ATTN and +quant_fp8 and (dynamo_partition or not +rms_norm)'
388
+ - pytest -v -s tests/compile/fusions_e2e/test_tp2_rocm_moe.py
tests/compile/fusions_e2e/test_tp2_ar_rms.py
+3
-1
@@ -111,7 +111,9 @@ 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:
114
+ if current_platform.is_rocm() and (
115
+ "llama-3" in model_name_lower or "qwen" in model_name_lower
116
+ ):
117
model_kwargs.update(
118
gpu_memory_utilization=0.01,
119
kv_cache_memory_bytes=64 * 1024**2,
tests/compile/fusions_e2e/test_tp2_rocm_moe.py
new
+60
@@ -0,0 +1,60 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""ROCm MoE coverage for the TP2 GEMM/collective fusion contract."""
4
+
5
+import pytest
6
+
7
+from tests.utils import multi_gpu_test
8
+from vllm.config import PassConfig
9
+from vllm.platforms import current_platform
10
+
11
+from .common import INDUCTOR_GRAPH_PARTITION
12
+from .models import TRITON_ATTN, qwen3_a3b
13
+
14
+
15
+@multi_gpu_test(num_gpus=2)
16
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm AsyncTP contract")
17
+@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION[:1])
18
+@pytest.mark.parametrize("attn_backend", [TRITON_ATTN])
19
+def test_tp2_rocm_moe_async_tp_bf16_fusions(
20
+ inductor_graph_partition, attn_backend, run_e2e_fusion_test
21
+):
22
+ """MoE dense projections retain the Qwen SP/AsyncTP rewrite counts."""
23
+ n_layers = 4
24
+ run_e2e_fusion_test(
25
+ qwen3_a3b.model_name,
26
+ qwen3_a3b.matches(n_layers),
27
+ dict(
28
+ hf_overrides=dict(
29
+ num_hidden_layers=n_layers,
30
+ hidden_size=512,
31
+ intermediate_size=1024,
32
+ moe_intermediate_size=128,
33
+ num_attention_heads=8,
34
+ num_key_value_heads=2,
35
+ head_dim=64,
36
+ num_experts=8,
37
+ num_experts_per_tok=2,
38
+ ),
39
+ load_format="dummy",
40
+ max_model_len=1024,
41
+ kv_cache_memory_bytes=64 * 1024**2,
42
+ gpu_memory_utilization=0.01,
43
+ max_num_batched_tokens=1024,
44
+ max_num_seqs=32,
45
+ ),
46
+ attn_backend,
47
+ dict(
48
+ use_inductor_graph_partition=inductor_graph_partition,
49
+ custom_ops=["-rms_norm"],
50
+ pass_config=PassConfig(
51
+ enable_qk_norm_rope_fusion=True,
52
+ enable_sp=True,
53
+ fuse_gemm_comms=True,
54
+ fuse_allreduce_rms=False,
55
+ sp_min_token_num=512,
56
+ ),
57
+ ),
58
+ ["norm_rope_fusion", "sequence_parallel", "async_tp"],
59
+ tp_size=2,
60
+ )