[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>
2026-09-29 02:08:31 -05:00
parent 0ae480ff1a
commit 69d5c81fcd
6 changed files with 209 additions and 68 deletions
@@ -4487,6 +4487,77 @@ steps:

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

+- label: ":amd: (MI355) mHC Kernels"
+  timeout_in_minutes: 50
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - csrc/
+  - vllm/
+  - tests/kernels/
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s kernels/mhc
+
+- label: ":amd: (MI355) FLA Ops Kernels"
+  timeout_in_minutes: 50
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - csrc/
+  - vllm/
+  - tests/kernels/
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s kernels/fla
+
+- label: ":amd: (MI355) Miscellaneous Kernels"
+  timeout_in_minutes: 55
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - csrc/
+  - vllm/
+  - tests/kernels/
+  - vllm/platforms/rocm.py
+  commands:
+  - pytest -v -s
+      kernels/test_apply_repetition_penalties.py
+      kernels/test_awq_int4_to_int8.py
+      kernels/test_bf16_skinny_gemm.py
+      kernels/test_cache_kernels.py
+      kernels/test_compressor_kv_cache.py
+      kernels/test_cp_gather_fp8.py
+      kernels/test_deepseek_v4_cpu_kernels.py
+      kernels/test_dsv41_mega_attn_layouts.py
+      kernels/test_engram.py
+      kernels/test_flashinfer_bf16_gemm.py
+      kernels/test_flex_attention.py
+      kernels/test_fp32_router_gemm.py
+      kernels/test_fused_deepseek_v32_norm_rope.py
+      kernels/test_fused_indexer_q_rope_quant.py
+      kernels/test_fused_inv_rope_fp8_quant.py
+      kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
+      kernels/test_fused_quant_activation.py
+      kernels/test_gate_linear_rocm_dispatch.py
+      kernels/test_gemm_rs_ar.py
+      kernels/test_kpool_decode_update_batched.py
+      kernels/test_minimax_m3_amd_ops.py
+      kernels/test_minimax_m3_sparse_attn_fp8_scale.py
+      kernels/test_onednn.py
+      kernels/test_relu2_fp8_quant.py
+      kernels/test_rocm_fp32_router_gemm.py
+      --deselect=tests/kernels/test_bf16_skinny_gemm.py::test_every_dsv3_routed_shape_is_instantiated
+
 - label: ":amd: (MI355 DPX) Attention DiffKV Kernels"
  timeout_in_minutes: 35
  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
@@ -5191,36 +5262,9 @@ steps:
  - unset VLLM_USE_V2_MODEL_RUNNER
  - uv pip install --system torchao==0.17.0
  - uv pip install --system conch-triton-kernels
-  - 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
+  - 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
  - if [ "$$BUILDKITE_PARALLEL_JOB" = "0" ]; then pytest -v -s rocm/test_moe_weight_replay.py; fi

-# - label: Quantized MoE Test (B200-MI355) # TBD
-#   timeout_in_minutes: 180
-#   mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
-#   dind: false
-#   agent_pool: mi355_1
-#   working_dir: "/vllm-workspace/"
-#   source_file_dependencies:
-#   - tests/quantization/test_gfx950_moe.py
-#   - vllm/model_executor/models/deepseek_v2.py
-#   - vllm/model_executor/models/gpt_oss.py
-#   - vllm/model_executor/models/llama4.py
-#   - vllm/model_executor/layers/fused_moe
-#   - vllm/model_executor/layers/quantization/compressed_tensors
-#   - vllm/model_executor/layers/quantization/modelopt.py
-#   - vllm/model_executor/layers/quantization/mxfp4.py
-#   - vllm/v1/attention/backends/triton_attn.py
-#   - vllm/v1/attention/backends/rocm_attn.py
-#   - vllm/v1/attention/backends/rocm_aiter_fa.py
-#   - vllm/v1/attention/backends/mla/
-#   - vllm/v1/attention/selector.py
-#   - vllm/model_executor/layers/layernorm.py
-#   - vllm/_aiter_ops.py
-#   - vllm/platforms/rocm.py
-#   - vllm/model_executor/model_loader/
-#   commands:
-#   - pytest -s -v tests/quantization/test_gfx950_moe.py
-
 #---------------------------------------------------------  mi355 · samplers  ----------------------------------------------------------#

 - label: ":amd: (MI355 DPX) Samplers Multimodal Beam Search"
@@ -6078,6 +6122,7 @@ steps:
  timeout_in_minutes: 45
  working_dir: /vllm-workspace/
  source_file_dependencies:
+  - tests/quantization/test_gfx950_moe.py
  - tests/quantization/test_rocm_moe.py
  - tests/utils.py
  - vllm/model_executor/models/deepseek_v2.py
@@ -6094,6 +6139,7 @@ steps:
  - vllm/envs.py
  commands:
  - pytest -v -s tests/quantization/test_rocm_moe.py
+  - pytest -v -s tests/quantization/test_gfx950_moe.py
  dind: false
  agent_pool: mi355_1
  num_gpus: 1
@@ -184,6 +184,20 @@ steps:
  - tests/kernels/
  commands:
    - pytest -v -s kernels/mhc
+  mirror:
+    amd:
+      label: ":amd: (MI355) mHC Kernels"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 50
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/
+      - vllm/
+      - tests/kernels/
+      - vllm/platforms/rocm.py

 - label: ":nvidia: (B200) FLA Ops Kernels"
  key: kernels-fla-ops-test-b200
@@ -196,6 +210,20 @@ steps:
  - tests/kernels/
  commands:
    - pytest -v -s kernels/fla
+  mirror:
+    amd:
+      label: ":amd: (MI355) FLA Ops Kernels"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 50
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/
+      - vllm/
+      - tests/kernels/
+      - vllm/platforms/rocm.py

 # Catch-all for test files at the tests/kernels root. This job collects
 # the whole root so new files are wired by default.
@@ -231,6 +259,50 @@ steps:
        --ignore=kernels/test_shuffle_rows.py
    # BROKEN on main, pending kernel fixes (B200):
    #   test_shuffle_rows.py (1: test_shuffle_rows_edge_cases)
+  mirror:
+    amd:
+      label: ":amd: (MI355) Miscellaneous Kernels"
+      dind: false
+      device: mi355_1
+      timeout_in_minutes: 55
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/
+      - vllm/
+      - tests/kernels/
+      - vllm/platforms/rocm.py
+      # Keep the ROCm exclusions from #50519: incompatible FlyDSL TurboQuant
+      # and a CUDA source check whose files are absent from the AMD test image.
+      commands:
+      - pytest -v -s
+          kernels/test_apply_repetition_penalties.py
+          kernels/test_awq_int4_to_int8.py
+          kernels/test_bf16_skinny_gemm.py
+          kernels/test_cache_kernels.py
+          kernels/test_compressor_kv_cache.py
+          kernels/test_cp_gather_fp8.py
+          kernels/test_deepseek_v4_cpu_kernels.py
+          kernels/test_dsv41_mega_attn_layouts.py
+          kernels/test_engram.py
+          kernels/test_flashinfer_bf16_gemm.py
+          kernels/test_flex_attention.py
+          kernels/test_fp32_router_gemm.py
+          kernels/test_fused_deepseek_v32_norm_rope.py
+          kernels/test_fused_indexer_q_rope_quant.py
+          kernels/test_fused_inv_rope_fp8_quant.py
+          kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py
+          kernels/test_fused_quant_activation.py
+          kernels/test_gate_linear_rocm_dispatch.py
+          kernels/test_gemm_rs_ar.py
+          kernels/test_kpool_decode_update_batched.py
+          kernels/test_minimax_m3_amd_ops.py
+          kernels/test_minimax_m3_sparse_attn_fp8_scale.py
+          kernels/test_onednn.py
+          kernels/test_relu2_fp8_quant.py
+          kernels/test_rocm_fp32_router_gemm.py
+          --deselect=tests/kernels/test_bf16_skinny_gemm.py::test_every_dsv3_routed_shape_is_instantiated

 - label: ":nvidia: (L4) Attention Kernels Shard %N"
  key: kernels-attention-test
@@ -725,6 +797,35 @@ steps:
    - pytest -v -s kernels/moe/test_flashinfer_moe.py
    - pytest -v -s kernels/moe/test_nvfp4_moe.py
    - pytest -v -s kernels/moe/test_ocp_mx_moe.py
+  mirror:
+    amd:
+      label: ":amd: (MI355 DPX) FP4 and INT4 MoE Kernels"
+      dind: false
+      device: mi355_dpx
+      num_devices: 1
+      timeout_in_minutes: 40
+      optional: false
+      working_dir: /vllm-workspace/tests
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - csrc/moe/
+      - csrc/rocm/
+      - vllm/model_executor/layers/fused_moe/
+      - vllm/model_executor/layers/quantization/
+      - vllm/_aiter_ops.py
+      - vllm/platforms/rocm.py
+      - vllm/envs.py
+      - vllm/config/
+      - tests/kernels/moe/test_rocm_aiter_moe.py
+      - tests/kernels/moe/test_ocp_mx_moe.py
+      - tests/kernels/moe/test_flydsl_moe.py
+      - tests/kernels/moe/utils.py
+      - tests/kernels/utils.py
+      commands:
+      - pytest -v -s kernels/moe/test_rocm_aiter_moe.py -k 'mi350_mxfp4 or quant_scheme_support_matches_gfx950'
+      - pytest -v -s kernels/moe/test_ocp_mx_moe.py::test_rocm_mxfp4_moe_oracle -k AITER_MXFP4
+      - pytest -v -s kernels/moe/test_flydsl_moe.py


 - label: ":nvidia: (H100) FusedMoE Layer Kernels"
@@ -847,34 +948,6 @@ steps:
  - pytest -v -s kernels/quantization/test_rocm_mxfp8_linear.py
  - pytest -v -s kernels/quantization/test_rocm_skinny_gemms.py -k wvsplitkrc

- label: ":amd: (MI355 DPX) FP4 and INT4 MoE Kernels"
-  key: amd-fp4-moe-kernels-mi355
-  dind: false
-  device: mi355_dpx
-  num_devices: 1
-  timeout_in_minutes: 40
-  working_dir: /vllm-workspace/tests
-  depends_on:
-  - image-build-amd
-  source_file_dependencies:
-  - csrc/moe/
-  - csrc/rocm/
-  - vllm/model_executor/layers/fused_moe/
-  - vllm/model_executor/layers/quantization/
-  - vllm/_aiter_ops.py
-  - vllm/platforms/rocm.py
-  - vllm/envs.py
-  - vllm/config/
-  - tests/kernels/moe/test_rocm_aiter_moe.py
-  - tests/kernels/moe/test_ocp_mx_moe.py
-  - tests/kernels/moe/test_flydsl_moe.py
-  - tests/kernels/moe/utils.py
-  - tests/kernels/utils.py
-  commands:
-  - pytest -v -s kernels/moe/test_rocm_aiter_moe.py -k 'mi350_mxfp4 or quant_scheme_support_matches_gfx950'
-  - pytest -v -s kernels/moe/test_ocp_mx_moe.py::test_rocm_mxfp4_moe_oracle -k AITER_MXFP4
-  - pytest -v -s kernels/moe/test_flydsl_moe.py
-
 - label: ":amd: (MI355 DPX) FP8 MoE Kernels"
  key: amd-fp8-moe-kernels-mi355
  dind: false
@@ -102,6 +102,7 @@ steps:
      depends_on:
      - image-build-amd
      source_file_dependencies:
+      - tests/quantization/test_gfx950_moe.py
      - tests/quantization/test_rocm_moe.py
      - tests/utils.py
      - vllm/model_executor/models/deepseek_v2.py
@@ -118,6 +119,7 @@ steps:
      - vllm/envs.py
      commands:
      - pytest -v -s tests/quantization/test_rocm_moe.py
+      - pytest -v -s tests/quantization/test_gfx950_moe.py

 - label: ":nvidia: (H200 MIG 35GB) Quantized Models"
  device: h200_35gb
@@ -269,6 +269,9 @@ def test_kda_overlap_configs_match_measured_table() -> None:
    }


+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="KDA CuTeDSL helper requires CUDA"
+)
@pytest.mark.parametrize(
    "compute_capability,supported",
    [((9, 0), False), ((10, 0), True), ((10, 3), True), ((12, 0), True)],
@@ -2182,7 +2182,9 @@ def test_v41_rocm_csa2_full_pipeline():
        )

    q_out, kv_out, index_q_out, index_q_scale_out, index_weights_out_out = result
-    assert torch.equal(q_out, q)
+    assert torch.equal(
+        q_out, q.view(num_tokens, attention.n_local_heads, attention.head_dim)
+    )
    assert torch.equal(kv_out, kv)
    assert torch.equal(index_q_out, index_q)
    assert torch.equal(index_q_scale_out, index_q_scale)
@@ -2273,7 +2275,9 @@ def test_v41_rocm_csa2_reindex_pipeline():
        )

    q_out, kv_out, index_q_out, index_q_scale_out, index_weights_out_out = result
-    assert torch.equal(q_out, q)
+    assert torch.equal(
+        q_out, q.view(num_tokens, attention.n_local_heads, attention.head_dim)
+    )
    assert torch.equal(kv_out, kv)
    assert torch.equal(index_q_out, index_q)
    assert torch.equal(index_q_scale_out, index_q_scale)
@@ -233,6 +233,26 @@ def _assert_eq(r_ref, r_kern):
    )


+def _cache_pool_bytes(kv_cache: torch.Tensor, pool_slot: int) -> torch.Tensor:
+    """Read a logical pool's K and scale bytes from the platform cache layout."""
+    page_size = kv_cache.shape[1]
+    head_dim = kv_cache.shape[2] - 4
+    page_idx, token_offset = divmod(pool_slot, page_size)
+    flat = kv_cache[page_idx].reshape(-1)
+    dims = torch.arange(head_dim, device=kv_cache.device)
+    if current_platform.is_rocm() and page_size > 1:
+        k_offsets = (
+            (token_offset // 16) * 16 * head_dim
+            + (dims // 16) * 16 * 16
+            + (token_offset % 16) * 16
+            + dims % 16
+        )
+    else:
+        k_offsets = token_offset * head_dim + dims
+    scale_offset = page_size * head_dim + 4 * token_offset
+    return torch.cat((flat[k_offsets], flat[scale_offset : scale_offset + 4]))
+
+
@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm required")
 def test_amd_prefill_writer_uses_preshuffled_cache_layout():
    from vllm.models.glm5next.amd.ops.kpool_compress import (
@@ -343,7 +363,7 @@ def test_decode_writer_matches_prefill_writer(pool_size, ring_pools):
        p
        for p in range(n_pools)
        if not torch.equal(
-            kv_prefill[p // page, p % page], kv_decode[p // page, p % page]
+            _cache_pool_bytes(kv_prefill, p), _cache_pool_bytes(kv_decode, p)
        )
    ]
    assert not differing, (
@@ -379,13 +399,6 @@ def test_rejected_draft_redo_needs_ring_slots(ring_pools):
    kv = torch.zeros_like(kv_ref)
    tail = torch.zeros(nblk, 2, ring, HEAD_DIM, dtype=torch.bfloat16, device=dev)

-    def pool_bytes(cache, p):
-        # Page layout: [page * HEAD_DIM bytes of K rows | page * 4 bytes of scales].
-        flat = cache[0].reshape(-1)
-        k_bytes = flat[p * HEAD_DIM : (p + 1) * HEAD_DIM]
-        s_bytes = flat[page * HEAD_DIM + 4 * p : page * HEAD_DIM + 4 * (p + 1)]
-        return torch.cat([k_bytes, s_bytes])
-
    def step(positions, keys, scores):
        pos = torch.tensor([positions], dtype=torch.int32, device=dev)
        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):
    )
    step([8, 9, 10, 11], k[8:12], score[8:12])  # all drafts rejected
    for p in (1, 2):
-        assert torch.equal(pool_bytes(kv, p), pool_bytes(kv_ref, p)), p
+        assert torch.equal(_cache_pool_bytes(kv, p), _cache_pool_bytes(kv_ref, p)), p

    # Draft 7 completes pool 1 and is rejected. With a one-pool ring, drafts
    # 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):
        torch.cat([score[6:7], draft_scores]),
    )
    step([7, 8, 9, 10], k[7:11], score[7:11])
-    pool1_ok = torch.equal(pool_bytes(kv, 1), pool_bytes(kv_ref, 1))
+    pool1_ok = torch.equal(_cache_pool_bytes(kv, 1), _cache_pool_bytes(kv_ref, 1))
    if ring_pools == 1:
        assert not pool1_ok, "expected a one-pool ring to corrupt pool 1"
    else: