From 71f975af437ae5316679752e8cdb72cd1c8b7b1d Mon Sep 17 00:00:00 2001 From: Aarushi Jain <142941703+aarushjain29@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:35:56 -0500 Subject: [PATCH] [ROCm][CI] Add MLA decode accuracy and determinism tests (#50480) Signed-off-by: aarushjain29 --- .buildkite/test-amd.yaml | 36 +++ .../attention/test_rocm_aiter_mla_decode.py | 225 ++++++++++++++++++ 2 files changed, 261 insertions(+) create mode 100644 tests/kernels/attention/test_rocm_aiter_mla_decode.py diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index caea1f7666..7b62fda8bb 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1530,6 +1530,24 @@ steps: - pytest -v -s tests/ir - pytest -v -s tests/kernels/ir +- label: Kernels MLA (MI300) # TBD + timeout_in_minutes: 180 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] + dind: false + agent_pool: mi300_1 + working_dir: "/vllm-workspace/" + source_file_dependencies: + - vllm/v1/attention/backends/mla/rocm_aiter_mla.py + - vllm/v1/attention/selector.py + - vllm/platforms/rocm.py + - vllm/_aiter_ops.py + - tests/kernels/attention/test_rocm_aiter_mla_decode.py + - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py + commands: + - rocm-smi + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py + - label: Kernels Attention Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -3387,6 +3405,24 @@ steps: - pytest -v -s tests/kernels/attention/test_attention_selector.py - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py +- label: Kernels MLA (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: + - vllm/v1/attention/backends/mla/rocm_aiter_mla.py + - vllm/v1/attention/selector.py + - vllm/platforms/rocm.py + - vllm/_aiter_ops.py + - tests/kernels/attention/test_rocm_aiter_mla_decode.py + - tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py + commands: + - rocm-smi + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py + - pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py + - label: Kernels Attention Test %N # TBD timeout_in_minutes: 180 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] diff --git a/tests/kernels/attention/test_rocm_aiter_mla_decode.py b/tests/kernels/attention/test_rocm_aiter_mla_decode.py new file mode 100644 index 0000000000..3a101468d2 --- /dev/null +++ b/tests/kernels/attention/test_rocm_aiter_mla_decode.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm kernel correctness tests for AITER MLA decode. + +Compares ``rocm_aiter_ops.mla_decode_fwd`` against a pure PyTorch reference +under the absorbed MLA formulation (DeepSeek-V3/V4) with varied batch sizes, +head counts, and sequence lengths. +""" + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.utils.torch_utils import set_random_seed + +_SKIP_NON_MI3XX = True +if current_platform.is_rocm(): + from vllm.platforms.rocm import on_mi3xx + + _SKIP_NON_MI3XX = not on_mi3xx() + +pytestmark = [ + pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm-specific tests"), + pytest.mark.skipif(_SKIP_NON_MI3XX, reason="MI300/MI350 ROCm only"), +] + +# DeepSeek-V3/V4 MLA config. +Q_HEAD_DIM = 576 # kv_lora_rank + qk_rope_head_dim +V_HEAD_DIM = 512 # kv_lora_rank +SM_SCALE = Q_HEAD_DIM**-0.5 + +NUM_HEADS = [16, 128] +BATCH_SIZES = [1, 4, 16] +KV_SEQ_LENS = [16, 256] + +ATOL, RTOL = 1.5e-2, 1e-2 + + +def _require_aiter(): + from vllm._aiter_ops import is_aiter_found_and_supported + + if not is_aiter_found_and_supported(): + pytest.skip("aiter is required on supported ROCm hardware for this test") + + +def _ref_mla_decode( + q: torch.Tensor, + kv_buffer: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, +) -> torch.Tensor: + """Pure PyTorch reference for absorbed MLA decode. + + K uses full q_head_dim (576) for scoring, V uses first v_head_dim (512) + for output. This mirrors the production kernel's absorbed formulation. + """ + batch_size, num_heads, _ = q.shape + output = torch.zeros( + batch_size, num_heads, V_HEAD_DIM, dtype=q.dtype, device=q.device + ) + + for b in range(batch_size): + start = kv_indptr[b].item() + end = kv_indptr[b + 1].item() + token_indices = kv_indices[start:end] + + k = kv_buffer[token_indices].float() + v = kv_buffer[token_indices, :V_HEAD_DIM].float() + + for h in range(num_heads): + q_h = q[b, h, :].float() + scores = torch.mv(k, q_h) * SM_SCALE + weights = torch.softmax(scores, dim=0) + output[b, h, :] = torch.mv(v.t(), weights).to(q.dtype) + + return output + + +def _make_inputs( + batch_size: int, + nhead: int, + kv_seq_len: int, + *, + contiguous_indices: bool = True, +): + """Build valid MLA decode inputs on the current CUDA device. + + When contiguous_indices=False, tokens are randomly scattered in a 2x pool + to simulate real paged allocation. + """ + from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAHelper + + actual_nhead = AiterMLAHelper.get_actual_mla_num_heads(nhead) + total_kv_tokens = batch_size * kv_seq_len + + pool_size = total_kv_tokens if contiguous_indices else total_kv_tokens * 2 + + q = torch.randn( + batch_size, actual_nhead, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + kv_buffer = torch.randn( + pool_size, 1, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + o = torch.zeros( + batch_size, actual_nhead, V_HEAD_DIM, dtype=torch.bfloat16, device="cuda" + ) + + qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") + kv_indptr = ( + torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") * kv_seq_len + ) + kv_last_page_lens = torch.ones(batch_size, dtype=torch.int32, device="cuda") + + if contiguous_indices: + kv_indices = torch.arange(0, total_kv_tokens, dtype=torch.int32, device="cuda") + else: + perm = torch.randperm(pool_size, device="cuda")[:total_kv_tokens] + kv_indices = perm.to(torch.int32) + + return { + "q": q, + "kv_buffer": kv_buffer, + "o": o, + "qo_indptr": qo_indptr, + "kv_indptr": kv_indptr, + "kv_indices": kv_indices, + "kv_last_page_lens": kv_last_page_lens, + "actual_nhead": actual_nhead, + } + + +def _run_kernel(inputs: dict) -> torch.Tensor: + """Run decode through the production rocm_aiter_ops path.""" + from vllm._aiter_ops import rocm_aiter_ops + + rocm_aiter_ops.mla_decode_fwd( + inputs["q"], + inputs["kv_buffer"], + inputs["o"], + SM_SCALE, + inputs["qo_indptr"], + 1, + inputs["kv_indptr"], + inputs["kv_indices"], + inputs["kv_last_page_lens"], + ) + return inputs["o"] + + +def _ref_output(inputs: dict) -> torch.Tensor: + """Compute reference output from the same inputs.""" + kv_flat = inputs["kv_buffer"].squeeze(1) + return _ref_mla_decode( + inputs["q"], kv_flat, inputs["kv_indptr"], inputs["kv_indices"] + ) + + +@pytest.mark.parametrize("nhead", NUM_HEADS) +@pytest.mark.parametrize("batch_size", BATCH_SIZES) +@pytest.mark.parametrize("kv_seq_len", KV_SEQ_LENS) +@torch.inference_mode() +def test_mla_decode_accuracy( + nhead: int, + batch_size: int, + kv_seq_len: int, +) -> None: + """BF16 decode accuracy vs PyTorch reference (contiguous indices).""" + _require_aiter() + set_random_seed(0) + + inputs = _make_inputs(batch_size, nhead, kv_seq_len) + output = _run_kernel(inputs) + output_ref = _ref_output(inputs) + + torch.testing.assert_close(output, output_ref, atol=ATOL, rtol=RTOL) + + +@pytest.mark.parametrize("nhead", NUM_HEADS) +@pytest.mark.parametrize("batch_size", [1, 4, 8]) +@pytest.mark.parametrize("kv_seq_len", [16, 128]) +@torch.inference_mode() +def test_mla_decode_noncontiguous_indices( + nhead: int, + batch_size: int, + kv_seq_len: int, +) -> None: + """BF16 decode accuracy with shuffled kv_indices (paged allocation).""" + _require_aiter() + set_random_seed(0) + + inputs = _make_inputs(batch_size, nhead, kv_seq_len, contiguous_indices=False) + output = _run_kernel(inputs) + output_ref = _ref_output(inputs) + + torch.testing.assert_close(output, output_ref, atol=ATOL, rtol=RTOL) + + +@torch.inference_mode() +def test_mla_decode_determinism() -> None: + """Repeated decode calls produce bitwise-identical results.""" + _require_aiter() + set_random_seed(0) + + inputs = _make_inputs(batch_size=4, nhead=128, kv_seq_len=16) + out_first = _run_kernel(inputs).clone() + + for _ in range(3): + inputs["o"] = torch.zeros_like(inputs["o"]) + out = _run_kernel(inputs) + torch.testing.assert_close(out, out_first, atol=0, rtol=0) + + +@torch.inference_mode() +def test_mla_decode_smoke() -> None: + """Basic sanity: correct shape, dtype, finite, non-zero.""" + _require_aiter() + set_random_seed(0) + + inputs = _make_inputs(batch_size=4, nhead=128, kv_seq_len=64) + output = _run_kernel(inputs) + + assert output.shape == (4, 128, V_HEAD_DIM) + assert output.dtype == torch.bfloat16 + assert torch.isfinite(output).all() + assert not torch.all(output == 0)