@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 107 lines 4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """ROCm custom op schema tests for AITER MLA decode.
4
5 ``opcheck`` verifies that the decode ops are registered and that their schemas
6 and fake implementations are consistent with the real kernels: fake-tensor
7 support for torch.compile tracing and ``mutates_args=["o"]`` in-place output
8 aliasing.
9 """
10
11 import pytest
12 import torch
13
14 from tests.kernels.utils import opcheck
15 from vllm.platforms import current_platform
16
17 pytestmark = pytest.mark.skipif(
18 not current_platform.is_rocm(), reason="ROCm-specific tests"
19 )
20
21 Q_HEAD_DIM = 576 # kv_lora_rank + qk_rope_head_dim
22 V_HEAD_DIM = 512 # kv_lora_rank
23
24
25 def _require_aiter():
26 from vllm._aiter_ops import is_aiter_found_and_supported
27 from vllm.platforms.rocm import get_cdna_version
28
29 if get_cdna_version() not in (3, 4):
30 pytest.skip("AITER MLA requires CDNA 3 or 4")
31
32 if not is_aiter_found_and_supported():
33 pytest.skip("aiter is required on supported ROCm hardware for this test")
34
35
36 @torch.inference_mode()
37 def test_mla_decode_fwd_op_schema() -> None:
38 """Opcheck validates registration, schema, fake-tensor, and ``o`` aliasing.
39
40 A single opcheck call covers that the op is registered/callable, that its
41 fake implementation matches the real op (torch.compile tracing), and that
42 the ``mutates_args=["o"]`` in-place output aliasing is declared correctly.
43 """
44 _require_aiter()
45 # Import ensures the custom op is registered.
46 from vllm._aiter_ops import rocm_aiter_ops # noqa: F401
47
48 batch_size, nhead = 4, 128
49
50 q = torch.randn(batch_size, nhead, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
51 kv_buffer = torch.randn(64, 1, 1, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
52 o = torch.zeros(batch_size, nhead, V_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
53 qo_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda")
54 kv_indptr = torch.arange(0, batch_size + 1, dtype=torch.int32, device="cuda") * 16
55 kv_indices = torch.arange(0, 64, dtype=torch.int32, device="cuda")
56 kv_last_page_lens = torch.ones(batch_size, dtype=torch.int32, device="cuda")
57
58 opcheck(
59 torch.ops.vllm.rocm_aiter_mla_decode_fwd,
60 (q, kv_buffer, o, qo_indptr, 1),
61 {
62 "kv_indptr": kv_indptr,
63 "kv_indices": kv_indices,
64 "kv_last_page_lens": kv_last_page_lens,
65 "sm_scale": Q_HEAD_DIM**-0.5,
66 "logit_cap": 0.0,
67 "q_scale": None,
68 "kv_scale": None,
69 "work_meta_data": None,
70 "work_indptr": None,
71 "work_info_set": None,
72 "reduce_indptr": None,
73 "reduce_final_map": None,
74 "reduce_partial_map": None,
75 },
76 )
77
78
79 @torch.inference_mode()
80 def test_mla_decode_fwd_lse_op_schema() -> None:
81 """Validate graph registration and mutation schema for LSE decode."""
82 _require_aiter()
83 # Import ensures the custom op is registered.
84 from vllm._aiter_ops import rocm_aiter_ops # noqa: F401
85
86 batch_size, nhead = 2, 16
87 q = torch.randn(batch_size, nhead, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
88 kv_buffer = torch.randn(32, 1, 1, Q_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
89 o = torch.zeros(batch_size, nhead, V_HEAD_DIM, dtype=torch.bfloat16, device="cuda")
90 qo_indptr = torch.arange(batch_size + 1, dtype=torch.int32, device="cuda")
91 kv_indptr = torch.arange(batch_size + 1, dtype=torch.int32, device="cuda") * 16
92 kv_indices = torch.arange(32, dtype=torch.int32, device="cuda")
93 kv_last_page_lens = torch.ones(batch_size, dtype=torch.int32, device="cuda")
94
95 opcheck(
96 torch.ops.vllm.rocm_aiter_mla_decode_fwd_lse,
97 (q, kv_buffer, o, qo_indptr, 1),
98 {
99 "kv_indptr": kv_indptr,
100 "kv_indices": kv_indices,
101 "kv_last_page_lens": kv_last_page_lens,
102 "sm_scale": Q_HEAD_DIM**-0.5,
103 "logit_cap": 0.0,
104 "q_scale": None,
105 "kv_scale": None,
106 },
107 )