@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 69d5c81fcd

[ROCm][CI] Mirror split kernel groups on MI355 and fix exposed tests (#58654)

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

Andreas Karatzas committed Sep 29, 2026 at 02:08 UTC 69d5c81fcd52f29a3b498dd4c403b088c91e4ee1

6 files changed +209 -68

.buildkite/test-amd.yaml

+74 -28

@@ -4487,6 +4487,77 @@ steps:

4487

4488 #---------------------------------------------------------- mi355 · kernels ----------------------------------------------------------#

4489

4490 +- label: ":amd: (MI355) mHC Kernels"

4491 + timeout_in_minutes: 50

4492 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]

4493 + dind: false

4494 + agent_pool: mi355_1

4495 + optional: true

4496 + working_dir: "/vllm-workspace/tests"

4497 + source_file_dependencies:

4498 + - csrc/

4499 + - vllm/

4500 + - tests/kernels/

4501 + - vllm/platforms/rocm.py

4502 + commands:

4503 + - pytest -v -s kernels/mhc

4504 +

4505 +- label: ":amd: (MI355) FLA Ops Kernels"

4506 + timeout_in_minutes: 50

4507 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]

4508 + dind: false

4509 + agent_pool: mi355_1

4510 + optional: true

4511 + working_dir: "/vllm-workspace/tests"

4512 + source_file_dependencies:

4513 + - csrc/

4514 + - vllm/

4515 + - tests/kernels/

4516 + - vllm/platforms/rocm.py

4517 + commands:

4518 + - pytest -v -s kernels/fla

4519 +

4520 +- label: ":amd: (MI355) Miscellaneous Kernels"

4521 + timeout_in_minutes: 55

4522 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]

4523 + dind: false

4524 + agent_pool: mi355_1

4525 + optional: true

4526 + working_dir: "/vllm-workspace/tests"

4527 + source_file_dependencies:

4528 + - csrc/

4529 + - vllm/

4530 + - tests/kernels/

4531 + - vllm/platforms/rocm.py

4532 + commands:

4533 + - pytest -v -s

4534 + kernels/test_apply_repetition_penalties.py

4535 + kernels/test_awq_int4_to_int8.py

4536 + kernels/test_bf16_skinny_gemm.py

4537 + kernels/test_cache_kernels.py

4538 + kernels/test_compressor_kv_cache.py

4539 + kernels/test_cp_gather_fp8.py

4540 + kernels/test_deepseek_v4_cpu_kernels.py

4541 + kernels/test_dsv41_mega_attn_layouts.py

4542 + kernels/test_engram.py

4543 + kernels/test_flashinfer_bf16_gemm.py

4544 + kernels/test_flex_attention.py

4545 + kernels/test_fp32_router_gemm.py

4546 + kernels/test_fused_deepseek_v32_norm_rope.py

4547 + kernels/test_fused_indexer_q_rope_quant.py

4548 + kernels/test_fused_inv_rope_fp8_quant.py

4549 + kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py

4550 + kernels/test_fused_quant_activation.py

4551 + kernels/test_gate_linear_rocm_dispatch.py

4552 + kernels/test_gemm_rs_ar.py

4553 + kernels/test_kpool_decode_update_batched.py

4554 + kernels/test_minimax_m3_amd_ops.py

4555 + kernels/test_minimax_m3_sparse_attn_fp8_scale.py

4556 + kernels/test_onednn.py

4557 + kernels/test_relu2_fp8_quant.py

4558 + kernels/test_rocm_fp32_router_gemm.py

4559 + --deselect=tests/kernels/test_bf16_skinny_gemm.py::test_every_dsv3_routed_shape_is_instantiated

4560 +

4561 - label: ":amd: (MI355 DPX) Attention DiffKV Kernels"

4562 timeout_in_minutes: 35

4563 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]

@@ -5191,36 +5262,9 @@ steps:

5262 - unset VLLM_USE_V2_MODEL_RUNNER

5263 - uv pip install --system torchao==0.17.0

5264 - uv pip install --system conch-triton-kernels

5194 - - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT

5265 + - VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py --ignore quantization/test_rocm_moe.py --ignore quantization/test_gfx950_moe.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT

5266 - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi

5267

5197 -# - label: Quantized MoE Test (B200-MI355) # TBD

5198 -# timeout_in_minutes: 180

5199 -# mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]

5200 -# dind: false

5201 -# agent_pool: mi355_1

5202 -# working_dir: "/vllm-workspace/"

5203 -# source_file_dependencies:

5204 -# - tests/quantization/test_gfx950_moe.py

5205 -# - vllm/model_executor/models/deepseek_v2.py

5206 -# - vllm/model_executor/models/gpt_oss.py

5207 -# - vllm/model_executor/models/llama4.py

5208 -# - vllm/model_executor/layers/fused_moe

5209 -# - vllm/model_executor/layers/quantization/compressed_tensors

5210 -# - vllm/model_executor/layers/quantization/modelopt.py

5211 -# - vllm/model_executor/layers/quantization/mxfp4.py

5212 -# - vllm/v1/attention/backends/triton_attn.py

5213 -# - vllm/v1/attention/backends/rocm_attn.py

5214 -# - vllm/v1/attention/backends/rocm_aiter_fa.py

5215 -# - vllm/v1/attention/backends/mla/

5216 -# - vllm/v1/attention/selector.py

5217 -# - vllm/model_executor/layers/layernorm.py

5218 -# - vllm/_aiter_ops.py

5219 -# - vllm/platforms/rocm.py

5220 -# - vllm/model_executor/model_loader/

5221 -# commands:

5222 -# - pytest -s -v tests/quantization/test_gfx950_moe.py

5223 -

5268 #--------------------------------------------------------- mi355 · samplers ----------------------------------------------------------#

5269

5270 - label: ":amd: (MI355 DPX) Samplers Multimodal Beam Search"

@@ -6078,6 +6122,7 @@ steps:

6122 timeout_in_minutes: 45

6123 working_dir: /vllm-workspace/

6124 source_file_dependencies:

6125 + - tests/quantization/test_gfx950_moe.py

6126 - tests/quantization/test_rocm_moe.py

6127 - tests/utils.py

6128 - vllm/model_executor/models/deepseek_v2.py

@@ -6094,6 +6139,7 @@ steps:

6139 - vllm/envs.py

6140 commands:

6141 - pytest -v -s tests/quantization/test_rocm_moe.py

6142 + - pytest -v -s tests/quantization/test_gfx950_moe.py

6143 dind: false

6144 agent_pool: mi355_1

6145 num_gpus: 1

.buildkite/test_areas/kernels.yaml

+101 -28

@@ -184,6 +184,20 @@ steps:

184 - tests/kernels/

185 commands:

186 - pytest -v -s kernels/mhc

187 + mirror:

188 + amd:

189 + label: ":amd: (MI355) mHC Kernels"

190 + dind: false

191 + device: mi355_1

192 + timeout_in_minutes: 50

193 + working_dir: "/vllm-workspace/tests"

194 + depends_on:

195 + - image-build-amd

196 + source_file_dependencies:

197 + - csrc/

198 + - vllm/

199 + - tests/kernels/

200 + - vllm/platforms/rocm.py

201

202 - label: ":nvidia: (B200) FLA Ops Kernels"

203 key: kernels-fla-ops-test-b200

@@ -196,6 +210,20 @@ steps:

210 - tests/kernels/

211 commands:

212 - pytest -v -s kernels/fla

213 + mirror:

214 + amd:

215 + label: ":amd: (MI355) FLA Ops Kernels"

216 + dind: false

217 + device: mi355_1

218 + timeout_in_minutes: 50

219 + working_dir: "/vllm-workspace/tests"

220 + depends_on:

221 + - image-build-amd

222 + source_file_dependencies:

223 + - csrc/

224 + - vllm/

225 + - tests/kernels/

226 + - vllm/platforms/rocm.py

227

228 # Catch-all for test files at the tests/kernels root. This job collects

229 # the whole root so new files are wired by default.

@@ -231,6 +259,50 @@ steps:

259 --ignore=kernels/test_shuffle_rows.py

260 # BROKEN on main, pending kernel fixes (B200):

261 # test_shuffle_rows.py (1: test_shuffle_rows_edge_cases)

262 + mirror:

263 + amd:

264 + label: ":amd: (MI355) Miscellaneous Kernels"

265 + dind: false

266 + device: mi355_1

267 + timeout_in_minutes: 55

268 + working_dir: "/vllm-workspace/tests"

269 + depends_on:

270 + - image-build-amd

271 + source_file_dependencies:

272 + - csrc/

273 + - vllm/

274 + - tests/kernels/

275 + - vllm/platforms/rocm.py

276 + # Keep the ROCm exclusions from #50519: incompatible FlyDSL TurboQuant

277 + # and a CUDA source check whose files are absent from the AMD test image.

278 + commands:

279 + - pytest -v -s

280 + kernels/test_apply_repetition_penalties.py

281 + kernels/test_awq_int4_to_int8.py

282 + kernels/test_bf16_skinny_gemm.py

283 + kernels/test_cache_kernels.py

284 + kernels/test_compressor_kv_cache.py

285 + kernels/test_cp_gather_fp8.py

286 + kernels/test_deepseek_v4_cpu_kernels.py

287 + kernels/test_dsv41_mega_attn_layouts.py

288 + kernels/test_engram.py

289 + kernels/test_flashinfer_bf16_gemm.py

290 + kernels/test_flex_attention.py

291 + kernels/test_fp32_router_gemm.py

292 + kernels/test_fused_deepseek_v32_norm_rope.py

293 + kernels/test_fused_indexer_q_rope_quant.py

294 + kernels/test_fused_inv_rope_fp8_quant.py

295 + kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py

296 + kernels/test_fused_quant_activation.py

297 + kernels/test_gate_linear_rocm_dispatch.py

298 + kernels/test_gemm_rs_ar.py

299 + kernels/test_kpool_decode_update_batched.py

300 + kernels/test_minimax_m3_amd_ops.py

301 + kernels/test_minimax_m3_sparse_attn_fp8_scale.py

302 + kernels/test_onednn.py

303 + kernels/test_relu2_fp8_quant.py

304 + kernels/test_rocm_fp32_router_gemm.py

305 + --deselect=tests/kernels/test_bf16_skinny_gemm.py::test_every_dsv3_routed_shape_is_instantiated

306

307 - label: ":nvidia: (L4) Attention Kernels Shard %N"

308 key: kernels-attention-test

@@ -725,6 +797,35 @@ steps:

797 - pytest -v -s kernels/moe/test_flashinfer_moe.py

798 - pytest -v -s kernels/moe/test_nvfp4_moe.py

799 - pytest -v -s kernels/moe/test_ocp_mx_moe.py

800 + mirror:

801 + amd:

802 + label: ":amd: (MI355 DPX) FP4 and INT4 MoE Kernels"

803 + dind: false

804 + device: mi355_dpx

805 + num_devices: 1

806 + timeout_in_minutes: 40

807 + optional: false

808 + working_dir: /vllm-workspace/tests

809 + depends_on:

810 + - image-build-amd

811 + source_file_dependencies:

812 + - csrc/moe/

813 + - csrc/rocm/

814 + - vllm/model_executor/layers/fused_moe/

815 + - vllm/model_executor/layers/quantization/

816 + - vllm/_aiter_ops.py

817 + - vllm/platforms/rocm.py

818 + - vllm/envs.py

819 + - vllm/config/

820 + - tests/kernels/moe/test_rocm_aiter_moe.py

821 + - tests/kernels/moe/test_ocp_mx_moe.py

822 + - tests/kernels/moe/test_flydsl_moe.py

823 + - tests/kernels/moe/utils.py

824 + - tests/kernels/utils.py

825 + commands:

826 + - pytest -v -s kernels/moe/test_rocm_aiter_moe.py -k 'mi350_mxfp4 or quant_scheme_support_matches_gfx950'

827 + - pytest -v -s kernels/moe/test_ocp_mx_moe.py::test_rocm_mxfp4_moe_oracle -k AITER_MXFP4

828 + - pytest -v -s kernels/moe/test_flydsl_moe.py

829

830

831 - label: ":nvidia: (H100) FusedMoE Layer Kernels"

@@ -847,34 +948,6 @@ steps:

948 - pytest -v -s kernels/quantization/test_rocm_mxfp8_linear.py

949 - pytest -v -s kernels/quantization/test_rocm_skinny_gemms.py -k wvsplitkrc

950

850 -- label: ":amd: (MI355 DPX) FP4 and INT4 MoE Kernels"

851 - key: amd-fp4-moe-kernels-mi355

852 - dind: false

853 - device: mi355_dpx

854 - num_devices: 1

855 - timeout_in_minutes: 40

856 - working_dir: /vllm-workspace/tests

857 - depends_on:

858 - - image-build-amd

859 - source_file_dependencies:

860 - - csrc/moe/

861 - - csrc/rocm/

862 - - vllm/model_executor/layers/fused_moe/

863 - - vllm/model_executor/layers/quantization/

864 - - vllm/_aiter_ops.py

865 - - vllm/platforms/rocm.py

866 - - vllm/envs.py

867 - - vllm/config/

868 - - tests/kernels/moe/test_rocm_aiter_moe.py

869 - - tests/kernels/moe/test_ocp_mx_moe.py

870 - - tests/kernels/moe/test_flydsl_moe.py

871 - - tests/kernels/moe/utils.py

872 - - tests/kernels/utils.py

873 - commands:

874 - - pytest -v -s kernels/moe/test_rocm_aiter_moe.py -k 'mi350_mxfp4 or quant_scheme_support_matches_gfx950'

875 - - pytest -v -s kernels/moe/test_ocp_mx_moe.py::test_rocm_mxfp4_moe_oracle -k AITER_MXFP4

876 - - pytest -v -s kernels/moe/test_flydsl_moe.py

877 -

951 - label: ":amd: (MI355 DPX) FP8 MoE Kernels"

952 key: amd-fp8-moe-kernels-mi355

953 dind: false

.buildkite/test_areas/quantization.yaml

@@ -102,6 +102,7 @@ steps:

102 depends_on:

103 - image-build-amd

104 source_file_dependencies:

105 + - tests/quantization/test_gfx950_moe.py

106 - tests/quantization/test_rocm_moe.py

107 - tests/utils.py

108 - vllm/model_executor/models/deepseek_v2.py

@@ -118,6 +119,7 @@ steps:

119 - vllm/envs.py

120 commands:

121 - pytest -v -s tests/quantization/test_rocm_moe.py

122 + - pytest -v -s tests/quantization/test_gfx950_moe.py

123

124 - label: ":nvidia: (H200 MIG 35GB) Quantized Models"

125 device: h200_35gb

tests/kernels/test_bf16_skinny_gemm.py

@@ -269,6 +269,9 @@ def test_kda_overlap_configs_match_measured_table() -> None:

269 }

270

271

272 +@pytest.mark.skipif(

273 + not current_platform.is_cuda(), reason="KDA CuTeDSL helper requires CUDA"

274 +)

275 @pytest.mark.parametrize(

276 "compute_capability,supported",

277 [((9, 0), False), ((10, 0), True), ((10, 3), True), ((12, 0), True)],

tests/kernels/test_compressor_kv_cache.py

+6 -2

@@ -2182,7 +2182,9 @@ def test_v41_rocm_csa2_full_pipeline():

2182 )

2183

2184 q_out, kv_out, index_q_out, index_q_scale_out, index_weights_out_out = result

2185 - assert torch.equal(q_out, q)

2185 + assert torch.equal(

2186 + q_out, q.view(num_tokens, attention.n_local_heads, attention.head_dim)

2187 + )

2188 assert torch.equal(kv_out, kv)

2189 assert torch.equal(index_q_out, index_q)

2190 assert torch.equal(index_q_scale_out, index_q_scale)

@@ -2273,7 +2275,9 @@ def test_v41_rocm_csa2_reindex_pipeline():

2275 )

2276

2277 q_out, kv_out, index_q_out, index_q_scale_out, index_weights_out_out = result

2276 - assert torch.equal(q_out, q)

2278 + assert torch.equal(

2279 + q_out, q.view(num_tokens, attention.n_local_heads, attention.head_dim)

2280 + )

2281 assert torch.equal(kv_out, kv)

2282 assert torch.equal(index_q_out, index_q)

2283 assert torch.equal(index_q_scale_out, index_q_scale)

tests/kernels/test_kpool_decode_update_batched.py

+23 -10

@@ -233,6 +233,26 @@ def _assert_eq(r_ref, r_kern):

233 )

234

235

236 +def _cache_pool_bytes(kv_cache: torch.Tensor, pool_slot: int) -> torch.Tensor:

237 + """Read a logical pool's K and scale bytes from the platform cache layout."""

238 + page_size = kv_cache.shape[1]

239 + head_dim = kv_cache.shape[2] - 4

240 + page_idx, token_offset = divmod(pool_slot, page_size)

241 + flat = kv_cache[page_idx].reshape(-1)

242 + dims = torch.arange(head_dim, device=kv_cache.device)

243 + if current_platform.is_rocm() and page_size > 1:

244 + k_offsets = (

245 + (token_offset // 16) * 16 * head_dim

246 + + (dims // 16) * 16 * 16

247 + + (token_offset % 16) * 16

248 + + dims % 16

249 + )

250 + else:

251 + k_offsets = token_offset * head_dim + dims

252 + scale_offset = page_size * head_dim + 4 * token_offset

253 + return torch.cat((flat[k_offsets], flat[scale_offset : scale_offset + 4]))

254 +

255 +

256 @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")

257 def test_amd_prefill_writer_uses_preshuffled_cache_layout():

258 from vllm.models.glm5next.amd.ops.kpool_compress import (

@@ -343,7 +363,7 @@ def test_decode_writer_matches_prefill_writer(pool_size, ring_pools):

363 p

364 for p in range(n_pools)

365 if not torch.equal(

346 - kv_prefill[p // page, p % page], kv_decode[p // page, p % page]

366 + _cache_pool_bytes(kv_prefill, p), _cache_pool_bytes(kv_decode, p)

367 )

368 ]

369 assert not differing, (

@@ -379,13 +399,6 @@ def test_rejected_draft_redo_needs_ring_slots(ring_pools):

399 kv = torch.zeros_like(kv_ref)

400 tail = torch.zeros(nblk, 2, ring, HEAD_DIM, dtype=torch.bfloat16, device=dev)

401

382 - def pool_bytes(cache, p):

383 - # Page layout: [page * HEAD_DIM bytes of K rows | page * 4 bytes of scales].

384 - flat = cache[0].reshape(-1)

385 - k_bytes = flat[p * HEAD_DIM : (p + 1) * HEAD_DIM]

386 - s_bytes = flat[page * HEAD_DIM + 4 * p : page * HEAD_DIM + 4 * (p + 1)]

387 - return torch.cat([k_bytes, s_bytes])

388 -

402 def step(positions, keys, scores):

403 pos = torch.tensor([positions], dtype=torch.int32, device=dev)

404 slots = [(p // pool) if p % pool == pool - 1 else -1 for p in positions]

@@ -415,7 +428,7 @@ def test_rejected_draft_redo_needs_ring_slots(ring_pools):

428 )

429 step([8, 9, 10, 11], k[8:12], score[8:12]) # all drafts rejected

430 for p in (1, 2):

418 - assert torch.equal(pool_bytes(kv, p), pool_bytes(kv_ref, p)), p

431 + assert torch.equal(_cache_pool_bytes(kv, p), _cache_pool_bytes(kv_ref, p)), p

432

433 # Draft 7 completes pool 1 and is rejected. With a one-pool ring, drafts

434 # 8 and 9 overwrite the slots of positions 4 and 5, which are read by the

@@ -430,7 +443,7 @@ def test_rejected_draft_redo_needs_ring_slots(ring_pools):

443 torch.cat([score[6:7], draft_scores]),

444 )

445 step([7, 8, 9, 10], k[7:11], score[7:11])

433 - pool1_ok = torch.equal(pool_bytes(kv, 1), pool_bytes(kv_ref, 1))

446 + pool1_ok = torch.equal(_cache_pool_bytes(kv, 1), _cache_pool_bytes(kv_ref, 1))

447 if ring_pools == 1:

448 assert not pool1_ok, "expected a one-pool ring to corrupt pool 1"

449 else: