@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 262 lines 9.55 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Regression test for causal masking in the AITER MLA verify MTP path.
4
5 When ``num_heads < 16`` the ROCm AITER MLA backend cannot use the ASM decode
6 kernel for a multi-token block, so ``AiterMLAImpl.forward_mqa`` routes verify
7 through ``mla_gluon``'s native 4-D MTP entry (``q`` shaped
8 ``[batch, qlen, nhead, dim]`` with ``use_2d_view=False``). aiter applies the
9 causal tail internally: position ``t`` attends KV ``[0, seq_len - qlen + t]``.
10
11 The test drives the real builder and ``forward_mqa`` over a multi-token decode
12 batch, intercepts the arguments actually handed to the Gluon kernel, and pins
13 that the 4-D MTP contract is used with per-request ``block_table`` and
14 ``cache_seqlens`` rather than a flattened per-token paged-KV view.
15 """
16
17 import types
18 from unittest.mock import patch
19
20 import pytest
21 import torch
22
23 from vllm._aiter_ops import is_aiter_found
24 from vllm.platforms import current_platform
25
26
27 def _on_rocm_with_aiter() -> bool:
28 return current_platform.is_rocm() and is_aiter_found()
29
30
31 # The production verify MTP path is gfx950-only. This test replaces the Gluon
32 # kernel with a spy and forces the architecture feature probe so the metadata
33 # regression remains testable on every ROCm AITER CI runner.
34 pytestmark = pytest.mark.skipif(
35 not _on_rocm_with_aiter(),
36 reason="ROCM_AITER_MLA verify MTP requires ROCm and AITER",
37 )
38
39 # Below the 16-head threshold that selects the MTP path. DeepSeek-V3 / R1 at
40 # TP=16 and Kimi-K3 at TP=8 both land here.
41 NUM_QUERY_HEADS = 12
42 KV_LORA_RANK = 512
43 QK_NOPE_HEAD_DIM = 128
44 QK_ROPE_HEAD_DIM = 64
45 V_HEAD_DIM = 128
46 HEAD_SIZE = KV_LORA_RANK + QK_ROPE_HEAD_DIM
47
48 # 1 + num_speculative_tokens, i.e. the verify block length.
49 QLEN = 8
50 # Committed context per request, including a fresh request at zero context
51 # where the whole KV range is the verify block itself.
52 CONTEXT_LENS = [0, 1, 37, 512]
53 # A cudagraph padding request: seq_len 0.
54 PADDING_ROWS = 1
55
56 # The Gluon path flattens the KV cache to one page per token.
57 PAGE_SIZE = 1
58 MAX_MODEL_LEN = 1024
59 MAX_NUM_SEQS = 8
60
61 # The builder resolves the verify routing and records it on the metadata, so
62 # the gfx950 feature probe has to be forced for the build as well as for the
63 # forward; ROCm CI also runs gfx942, where the real probe is False.
64 _GLUON_SUPPORTED_TARGET = (
65 "vllm.v1.attention.backends.mla.rocm_aiter_mla._gluon_mla_decode_supported"
66 )
67
68
69 def _seq_lens() -> list[int]:
70 return [c + QLEN for c in CONTEXT_LENS] + [0] * PADDING_ROWS
71
72
73 def _run_verify_block():
74 """Drive the real builder + forward_mqa over one multi-token verify block.
75
76 Returns ``(metadata, captured)`` where ``captured`` holds the Gluon kwargs.
77 """
78 from tests.v1.attention.utils import (
79 BatchSpec,
80 create_common_attn_metadata,
81 create_vllm_config,
82 )
83 from vllm.config import SpeculativeConfig
84 from vllm.config.vllm import set_current_vllm_config
85 from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAImpl
86 from vllm.v1.attention.backends.registry import AttentionBackendEnum
87 from vllm.v1.kv_cache_interface import MLAAttentionSpec
88 from vllm.v1.worker.workspace import init_workspace_manager
89
90 device = torch.device("cuda:0")
91 seq_lens = _seq_lens()
92 num_reqs = len(seq_lens)
93 # One flat page per token, arange block indices, plus room to spare.
94 num_gpu_blocks = num_reqs * max(seq_lens) + 256
95
96 vllm_config = create_vllm_config(
97 model_name="deepseek-ai/DeepSeek-R1",
98 max_model_len=MAX_MODEL_LEN,
99 num_gpu_blocks=num_gpu_blocks,
100 block_size=PAGE_SIZE,
101 max_num_seqs=MAX_NUM_SEQS,
102 max_num_batched_tokens=8192,
103 hf_config_override={"num_attention_heads": NUM_QUERY_HEADS},
104 )
105 # What raises reorder_batch_threshold to QLEN, so a QLEN-token block is
106 # classified as a decode instead of a prefill. ngram needs no draft model.
107 vllm_config.speculative_config = SpeculativeConfig(
108 method="ngram", num_speculative_tokens=QLEN - 1
109 )
110 vllm_config.model_config.get_num_attention_heads = types.MethodType( # type: ignore[method-assign]
111 lambda self, parallel_config, arch_config=None: NUM_QUERY_HEADS,
112 vllm_config.model_config,
113 )
114
115 spec = MLAAttentionSpec(
116 block_size=PAGE_SIZE,
117 num_kv_heads=1,
118 head_size=vllm_config.model_config.get_head_size(),
119 dtype=vllm_config.model_config.dtype,
120 cache_dtype_str="auto",
121 )
122
123 backend_cls = AttentionBackendEnum.ROCM_AITER_MLA.get_class()
124 builder_cls = backend_cls.get_builder_cls()
125 impl_cls = backend_cls.get_impl_cls()
126
127 # The builder reads layer.prefill_backend from static_forward_context; a
128 # stub with the attribute is enough for metadata construction.
129 layer_name = "placeholder"
130 vllm_config.compilation_config.static_forward_context[layer_name] = (
131 types.SimpleNamespace(
132 prefill_backend=torch.empty((1,)),
133 q_lora_rank=None,
134 kv_lora_rank=KV_LORA_RANK,
135 qk_nope_head_dim=QK_NOPE_HEAD_DIM,
136 qk_rope_head_dim=QK_ROPE_HEAD_DIM,
137 v_head_dim=V_HEAD_DIM,
138 )
139 )
140
141 init_workspace_manager(device)
142
143 batch_spec = BatchSpec(seq_lens=seq_lens, query_lens=[QLEN] * num_reqs)
144
145 captured: dict = {}
146
147 def spy(**kwargs):
148 captured.update(kwargs)
149 captured["q_nope"] = kwargs["q_nope"].detach().clone()
150 captured["page_table"] = kwargs["page_table"].detach().clone()
151 captured["seq_info"] = kwargs["seq_info"].detach().clone()
152
153 with (
154 patch(_GLUON_SUPPORTED_TARGET, lambda: True),
155 set_current_vllm_config(vllm_config),
156 ):
157 builder = builder_cls(spec, [layer_name], vllm_config, device)
158 common_attn_metadata = create_common_attn_metadata(
159 batch_spec, PAGE_SIZE, device, arange_block_indices=True
160 )
161 metadata = builder.build(
162 common_prefix_len=0, common_attn_metadata=common_attn_metadata
163 )
164
165 impl = impl_cls(
166 num_heads=NUM_QUERY_HEADS,
167 head_size=HEAD_SIZE,
168 scale=(QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM) ** -0.5,
169 num_kv_heads=1,
170 alibi_slopes=None,
171 sliding_window=None,
172 kv_cache_dtype="auto",
173 logits_soft_cap=None,
174 attn_type="decoder",
175 kv_sharing_target_layer_name=None,
176 q_lora_rank=None,
177 kv_lora_rank=KV_LORA_RANK,
178 qk_nope_head_dim=QK_NOPE_HEAD_DIM,
179 qk_rope_head_dim=QK_ROPE_HEAD_DIM,
180 qk_head_dim=QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM,
181 v_head_dim=V_HEAD_DIM,
182 kv_b_proj=None,
183 )
184
185 num_rows = num_reqs * QLEN
186 dtype = torch.bfloat16
187 q_nope = torch.zeros(
188 num_rows, NUM_QUERY_HEADS, KV_LORA_RANK, dtype=dtype, device=device
189 )
190 q_pe = torch.zeros(
191 num_rows, NUM_QUERY_HEADS, QK_ROPE_HEAD_DIM, dtype=dtype, device=device
192 )
193 kv_cache = torch.zeros(
194 num_gpu_blocks, PAGE_SIZE, HEAD_SIZE, dtype=dtype, device=device
195 )
196
197 # The Gluon kernel only reads the metadata this test is about, so a spy in
198 # its place keeps the assertions independent of the AITER build.
199 with (
200 patch(_GLUON_SUPPORTED_TARGET, lambda: True),
201 patch(
202 "vllm.v1.attention.backends.mla.rocm_aiter_mla._get_mla_gluon",
203 lambda: spy,
204 ),
205 ):
206 assert isinstance(impl, AiterMLAImpl)
207 impl.forward_mqa((q_nope, q_pe), kv_cache, metadata, layer=None)
208
209 return metadata, captured
210
211
212 def test_verify_mtp_uses_native_4d_gluon_entry():
213 """Verify uses 4-D q with per-request paged_kv metadata.
214
215 Regression guard: before the fix every flattened row of a request got that
216 request's whole paged-KV range, so verify position t attended to the draft
217 tokens after it. The native MTP path delegates causal masking to aiter.
218 """
219 metadata, captured = _run_verify_block()
220
221 decode = metadata.decode
222 assert decode is not None, "batch was not classified as a decode"
223 assert decode.max_qo_len == QLEN, (
224 f"expected a {QLEN}-token verify block, got max_qo_len="
225 f"{decode.max_qo_len}; the MTP path under test was not reached"
226 )
227 assert decode.use_gluon_verify, (
228 "the builder owns verify routing, so the metadata -- not the impl -- "
229 "has to select the Gluon MTP entry for a small-head bf16 verify block"
230 )
231 assert captured, "forward_mqa did not reach the Gluon kernel"
232
233 seq_lens = _seq_lens()
234 num_reqs = len(seq_lens)
235
236 q_nope = captured["q_nope"]
237 assert q_nope.dim() == 4, (
238 f"verify must pass 4-D q to mla_gluon, got {q_nope.dim()}-D shape "
239 f"{tuple(q_nope.shape)}"
240 )
241 assert q_nope.shape[:2] == (num_reqs, QLEN), (
242 f"expected q shape [num_reqs={num_reqs}, qlen={QLEN}, ...], "
243 f"got {tuple(q_nope.shape)}"
244 )
245 assert q_nope.shape[2] == NUM_QUERY_HEADS
246
247 assert captured["use_2d_view"] is False
248 page_table = captured["page_table"]
249 seq_info = captured["seq_info"]
250 assert page_table.dim() == 1, (
251 f"verify MTP passes the 1-D paged_kv_indices buffer, got {page_table.dim()}-D"
252 )
253 assert seq_info.dim() == 1 and seq_info.numel() == num_reqs + 1, (
254 f"verify MTP passes paged_kv_indptr [num_reqs + 1], got shape "
255 f"{tuple(seq_info.shape)}"
256 )
257 assert decode is not None and decode.paged_kv_indptr is not None
258 assert torch.equal(seq_info, decode.paged_kv_indptr[: num_reqs + 1])
259 assert torch.equal(
260 page_table[: int(seq_info[-1].item())],
261 decode.paged_kv_indices[: int(seq_info[-1].item())],
262 )