[CI][ROCm] Add an MI355 Kimi-K3 unit test group (#58012)

Signed-off-by: Oxana Korzh <okorzh@amd.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>

.buildkite/test-amd.yaml

@@ -1818,6 +1818,48 @@ steps:
   commands:
   - pytest -v -s models/kimi_k3/test_kda.py models/kimi_k3/test_kda_metadata.py
 
+- label: ":amd: (MI355) Kimi K3" # TBD
+  timeout_in_minutes: 60
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_1
+  num_gpus: 1
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/models/kimi_k3/
+  - csrc/libtorch_stable/kimi_k3/
+  - tests/models/kimi_k3/
+  - vllm/platforms/rocm.py
+  commands:
+  # Mirrors the ":amd: (MI355) Kimi K3" step generated from
+  # .buildkite/test_areas/models_basic.yaml. The ROCm KDA/AttnRes kernels are
+  # gated on gfx950, so this has to be MI355. The two kernels/ files the NVIDIA
+  # step runs are CUDA-only (they import cuda.bindings.driver), so they are
+  # deliberately left out here.
+  - pytest -v -s models/kimi_k3
+
+- label: ":amd: (MI355) Kimi K3 Multi-GPU" # TBD
+  timeout_in_minutes: 30
+  mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
+  dind: false
+  agent_pool: mi355_4
+  num_gpus: 4
+  optional: true
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+  - vllm/models/kimi_k3/
+  - tests/models/kimi_k3/
+  - vllm/platforms/rocm.py
+  commands:
+  # The single-GPU step above skips every @multi_gpu_test case. `-m "distributed"`
+  # selects exactly those: multi_gpu_test() applies pytest.mark.distributed
+  # alongside its skipif (tests/utils.py:2080), so this picks up a new
+  # multi-GPU test anywhere in the directory without editing this list.
+  # Four GPUs run the two tp4 cases in test_amd_latent_moe_runner.py; the tp8
+  # case and the 8/16-GPU cases in test_latent_moe_tail.py still skip.
+  - pytest -v -s -m "distributed" models/kimi_k3
+
 - label: ":amd: (MI300) Mamba Kernels" # TBD
   timeout_in_minutes: 55
   mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]

.buildkite/test_areas/models_basic.yaml

@@ -109,6 +109,23 @@ steps:
   commands:
     # The native NVIDIA Kimi K3 kernels require the SM100 family.
     - pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py
+  mirror:
+    amd:
+      label: ":amd: (MI355) Kimi K3"
+      dind: false
+      device: mi355_1
+      num_devices: 1
+      timeout_in_minutes: 60
+      working_dir: "/vllm-workspace/tests"
+      depends_on:
+      - image-build-amd
+      source_file_dependencies:
+      - vllm/platforms/rocm.py
+      commands:
+      # The ROCm KDA/AttnRes kernels are gated on gfx950, so this has to be MI355.
+      # The two kernels/ files the NVIDIA job adds are CUDA-only (they import
+      # cuda.bindings.driver), so they are deliberately left out here.
+      - pytest -v -s models/kimi_k3
 
 - label: ":nvidia: (H200) GLM5Next"
   key: glm5next-unit-tests

tests/models/kimi_k3/test_attn_res.py

@@ -17,6 +17,26 @@ EPS = 1e-5
 attn_res_module = importlib.import_module("vllm.models.kimi_k3.nvidia.ops.attn_res")
 
 
+def _on_rocm_below_10() -> bool:
+    if not current_platform.is_rocm():
+        return False
+    from vllm.platforms.rocm import get_rocm_version
+
+    return (get_rocm_version() or (0,)) < (10,)
+
+
+# The Triton bundled with ROCm < 10 (3.7.x) crashes in the AMD
+# CanonicalizePointers pass on the kernel's tl.where over pointer tensors. Only
+# the num_blocks > 0 loop contains it. ROCm serves AttnRes from
+# vllm/models/kimi_k3/amd/ops/attn_res.py instead, so nothing real is lost.
+_OLD_ROCM = _on_rocm_below_10()
+
+
+def _skip_on_old_rocm(num_blocks: int) -> None:
+    if _OLD_ROCM and num_blocks > 0:
+        pytest.skip("Triton on ROCm < 10 cannot compile this kernel's pointer select")
+
+
 def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor:
     storage = torch.randn(
         *shape[:-1],
@@ -79,6 +99,7 @@ def test_attn_res(
 ):
     if backend == "nvidia" and not current_platform.is_device_capability_family(100):
         pytest.skip("NVIDIA AttnRes requires the SM100 family")
+    _skip_on_old_rocm(num_blocks)
 
     prefix = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding)
     delta = (
@@ -133,6 +154,7 @@ def test_attn_res(
 
 @pytest.mark.parametrize("num_blocks", range(MAX_BLOCKS + 1))
 def test_attn_res_block_counts(num_blocks: int):
+    _skip_on_old_rocm(num_blocks)
     prefix = torch.randn(1, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
     blocks = torch.randn(
         1, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
@@ -169,6 +191,7 @@ def test_attn_res_block_counts(num_blocks: int):
 
 
 def test_attn_res_without_output_norm():
+    _skip_on_old_rocm(MAX_BLOCKS)
     prefix = torch.randn(7, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
     delta = torch.randn_like(prefix)
     blocks = torch.randn(

vllm/platforms/rocm.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import importlib.metadata
 import os
 import platform
 from datetime import timedelta
@@ -361,6 +362,25 @@ def get_cdna_version() -> int:
     return 0
 
 
+@cache
+def get_rocm_version() -> tuple[int, ...] | None:
+    """Return the installed ROCm release as (major, minor, patch), or None."""
+    # ROCm 10+ ships as the `rocm` pip SDK; older releases install to /opt/rocm.
+    try:
+        version = importlib.metadata.version("rocm")
+    except importlib.metadata.PackageNotFoundError:
+        rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm")
+        try:
+            with open(os.path.join(rocm_path, ".info", "version")) as f:
+                version = f.read()
+        except OSError:
+            return None
+    match = re.match(r"\s*(\d+)\.(\d+)(?:\.(\d+))?", version)
+    if match is None:
+        return None
+    return tuple(int(part) for part in match.groups() if part is not None)
+
+
 # Enable HIP online tuning early, before hipBLASLt initializes.
 # Turn on hipBLASLt online tuning if use AITER hipBLASLt GEMM.
 if (