[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>
Oxana Korzh committed
Sep 23, 2026 at 10:25 UTC
c4cfd7007bad6f1125ef938feb656707b3ccab76
4 files changed
+102
.buildkite/test-amd.yaml
+42
@@ -1818,6 +1818,48 @@ steps:
1818
commands:
1819
- pytest -v -s models/kimi_k3/test_kda.py models/kimi_k3/test_kda_metadata.py
1820
1821
+- label: ":amd: (MI355) Kimi K3" # TBD
1822
+ timeout_in_minutes: 60
1823
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
1824
+ dind: false
1825
+ agent_pool: mi355_1
1826
+ num_gpus: 1
1827
+ optional: true
1828
+ working_dir: "/vllm-workspace/tests"
1829
+ source_file_dependencies:
1830
+ - vllm/models/kimi_k3/
1831
+ - csrc/libtorch_stable/kimi_k3/
1832
+ - tests/models/kimi_k3/
1833
+ - vllm/platforms/rocm.py
1834
+ commands:
1835
+ # Mirrors the ":amd: (MI355) Kimi K3" step generated from
1836
+ # .buildkite/test_areas/models_basic.yaml. The ROCm KDA/AttnRes kernels are
1837
+ # gated on gfx950, so this has to be MI355. The two kernels/ files the NVIDIA
1838
+ # step runs are CUDA-only (they import cuda.bindings.driver), so they are
1839
+ # deliberately left out here.
1840
+ - pytest -v -s models/kimi_k3
1841
+
1842
+- label: ":amd: (MI355) Kimi K3 Multi-GPU" # TBD
1843
+ timeout_in_minutes: 30
1844
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
1845
+ dind: false
1846
+ agent_pool: mi355_4
1847
+ num_gpus: 4
1848
+ optional: true
1849
+ working_dir: "/vllm-workspace/tests"
1850
+ source_file_dependencies:
1851
+ - vllm/models/kimi_k3/
1852
+ - tests/models/kimi_k3/
1853
+ - vllm/platforms/rocm.py
1854
+ commands:
1855
+ # The single-GPU step above skips every @multi_gpu_test case. `-m "distributed"`
1856
+ # selects exactly those: multi_gpu_test() applies pytest.mark.distributed
1857
+ # alongside its skipif (tests/utils.py:2080), so this picks up a new
1858
+ # multi-GPU test anywhere in the directory without editing this list.
1859
+ # Four GPUs run the two tp4 cases in test_amd_latent_moe_runner.py; the tp8
1860
+ # case and the 8/16-GPU cases in test_latent_moe_tail.py still skip.
1861
+ - pytest -v -s -m "distributed" models/kimi_k3
1862
+
1863
- label: ":amd: (MI300) Mamba Kernels" # TBD
1864
timeout_in_minutes: 55
1865
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
.buildkite/test_areas/models_basic.yaml
+17
@@ -109,6 +109,23 @@ steps:
109
commands:
110
# The native NVIDIA Kimi K3 kernels require the SM100 family.
111
- pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py
112
+ mirror:
113
+ amd:
114
+ label: ":amd: (MI355) Kimi K3"
115
+ dind: false
116
+ device: mi355_1
117
+ num_devices: 1
118
+ timeout_in_minutes: 60
119
+ working_dir: "/vllm-workspace/tests"
120
+ depends_on:
121
+ - image-build-amd
122
+ source_file_dependencies:
123
+ - vllm/platforms/rocm.py
124
+ commands:
125
+ # The ROCm KDA/AttnRes kernels are gated on gfx950, so this has to be MI355.
126
+ # The two kernels/ files the NVIDIA job adds are CUDA-only (they import
127
+ # cuda.bindings.driver), so they are deliberately left out here.
128
+ - pytest -v -s models/kimi_k3
129
130
- label: ":nvidia: (H200) GLM5Next"
131
key: glm5next-unit-tests
tests/models/kimi_k3/test_attn_res.py
+23
@@ -17,6 +17,26 @@ EPS = 1e-5
17
attn_res_module = importlib.import_module("vllm.models.kimi_k3.nvidia.ops.attn_res")
18
19
20
+def _on_rocm_below_10() -> bool:
21
+ if not current_platform.is_rocm():
22
+ return False
23
+ from vllm.platforms.rocm import get_rocm_version
24
+
25
+ return (get_rocm_version() or (0,)) < (10,)
26
+
27
+
28
+# The Triton bundled with ROCm < 10 (3.7.x) crashes in the AMD
29
+# CanonicalizePointers pass on the kernel's tl.where over pointer tensors. Only
30
+# the num_blocks > 0 loop contains it. ROCm serves AttnRes from
31
+# vllm/models/kimi_k3/amd/ops/attn_res.py instead, so nothing real is lost.
32
+_OLD_ROCM = _on_rocm_below_10()
33
+
34
+
35
+def _skip_on_old_rocm(num_blocks: int) -> None:
36
+ if _OLD_ROCM and num_blocks > 0:
37
+ pytest.skip("Triton on ROCm < 10 cannot compile this kernel's pointer select")
38
+
39
+
40
def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor:
41
storage = torch.randn(
42
*shape[:-1],
@@ -79,6 +99,7 @@ def test_attn_res(
99
):
100
if backend == "nvidia" and not current_platform.is_device_capability_family(100):
101
pytest.skip("NVIDIA AttnRes requires the SM100 family")
102
+ _skip_on_old_rocm(num_blocks)
103
104
prefix = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding)
105
delta = (
@@ -133,6 +154,7 @@ def test_attn_res(
154
155
@pytest.mark.parametrize("num_blocks", range(MAX_BLOCKS + 1))
156
def test_attn_res_block_counts(num_blocks: int):
157
+ _skip_on_old_rocm(num_blocks)
158
prefix = torch.randn(1, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
159
blocks = torch.randn(
160
1, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
@@ -169,6 +191,7 @@ def test_attn_res_block_counts(num_blocks: int):
191
192
193
def test_attn_res_without_output_norm():
194
+ _skip_on_old_rocm(MAX_BLOCKS)
195
prefix = torch.randn(7, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
196
delta = torch.randn_like(prefix)
197
blocks = torch.randn(
vllm/platforms/rocm.py
+20
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import importlib.metadata
5
import os
6
import platform
7
from datetime import timedelta
@@ -361,6 +362,25 @@ def get_cdna_version() -> int:
362
return 0
363
364
365
+@cache
366
+def get_rocm_version() -> tuple[int, ...] | None:
367
+ """Return the installed ROCm release as (major, minor, patch), or None."""
368
+ # ROCm 10+ ships as the `rocm` pip SDK; older releases install to /opt/rocm.
369
+ try:
370
+ version = importlib.metadata.version("rocm")
371
+ except importlib.metadata.PackageNotFoundError:
372
+ rocm_path = os.environ.get("ROCM_PATH", "/opt/rocm")
373
+ try:
374
+ with open(os.path.join(rocm_path, ".info", "version")) as f:
375
+ version = f.read()
376
+ except OSError:
377
+ return None
378
+ match = re.match(r"\s*(\d+)\.(\d+)(?:\.(\d+))?", version)
379
+ if match is None:
380
+ return None
381
+ return tuple(int(part) for part in match.groups() if part is not None)
382
+
383
+
384
# Enable HIP online tuning early, before hipBLASLt initializes.
385
# Turn on hipBLASLt online tuning if use AITER hipBLASLt GEMM.
386
if (