From c4cfd7007bad6f1125ef938feb656707b3ccab76 Mon Sep 17 00:00:00 2001 From: Oxana Korzh Date: Wed, 23 Sep 2026 10:25:51 -0600 Subject: [PATCH] [CI][ROCm] Add an MI355 Kimi-K3 unit test group (#58012) Signed-off-by: Oxana Korzh Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Andreas Karatzas --- .buildkite/test-amd.yaml | 42 +++++++++++++++++++++++++ .buildkite/test_areas/models_basic.yaml | 17 ++++++++++ tests/models/kimi_k3/test_attn_res.py | 23 ++++++++++++++ vllm/platforms/rocm.py | 20 ++++++++++++ 4 files changed, 102 insertions(+) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index c95d22d33d..8bd434a355 100644 --- a/.buildkite/test-amd.yaml +++ b/.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] diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 930c2dd590..6f570b3459 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.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 diff --git a/tests/models/kimi_k3/test_attn_res.py b/tests/models/kimi_k3/test_attn_res.py index 574f7a5925..366abe0221 100644 --- a/tests/models/kimi_k3/test_attn_res.py +++ b/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( diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 9089d025e0..844fe4884b 100644 --- a/vllm/platforms/rocm.py +++ b/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 (