@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 213 lines 7.29 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Regression test for AITER MLA persistent decode metadata dtypes.
4
5 For the gfx950 fp8/fp8 nhead=32 qlen=1 fold path, the split/reduce metadata
6 layout depends on the q/kv element size. The builder must forward dtype_q/dtype_kv
7 to ``get_mla_metadata_v1``; omitting them lays out the work for the wrong dtype
8 and corrupts decode output. The test pins the builder's metadata to a golden
9 recomputed at runtime with the explicit correct dtypes.
10 """
11
12 import types
13 from unittest.mock import patch
14
15 import pytest
16 import torch
17
18 from vllm._aiter_ops import is_aiter_found
19 from vllm.platforms import current_platform
20
21
22 def _on_gfx950() -> bool:
23 if not (current_platform.is_rocm() and is_aiter_found()):
24 return False
25 from vllm.platforms.rocm import on_gfx950
26
27 return on_gfx950()
28
29
30 pytestmark = pytest.mark.skipif(
31 not _on_gfx950(),
32 reason="AITER MLA fp8 persistent decode metadata is gfx950-only",
33 )
34
35 # The fold path that the bug corrupted: fp8 query + fp8 KV-cache, 32 query
36 # heads, single-token decode, batch 128, context 8192, page_size 1.
37 NUM_QUERY_HEADS = 32
38 DECODE_QLEN = 1
39 BATCH_SIZE = 128
40 CONTEXT_LEN = 8192
41 PAGE_SIZE = 1
42
43 # On the fp8 KV-cache path the builder forwards the platform fp8 dtype (aiter
44 # dtypes.fp8) for both q and kv. Mirror it via current_platform.fp8_dtype()
45 # instead of hardcoding a literal (see #47276).
46 EXPECTED_Q_DTYPE = current_platform.fp8_dtype()
47 EXPECTED_KV_DTYPE = current_platform.fp8_dtype()
48
49 # The split/reduce content tensors filled by get_mla_metadata_v1. work_meta_data
50 # is excluded: it holds raw device pointers, never equal across allocations.
51 _CONTENT_METADATA_FIELDS = (
52 "work_indptr",
53 "work_info_set",
54 "reduce_indptr",
55 "reduce_final_map",
56 "reduce_partial_map",
57 )
58
59 # The builder's get_mla_metadata_v1 call passes 6 input args then 6 output
60 # buffers (see AiterMLAMetadataBuilder._build_decode). Output order -> field.
61 _NUM_INPUT_ARGS = 6
62 _OUTPUT_ARG_FIELDS = (
63 "work_meta_data",
64 "work_info_set",
65 "work_indptr",
66 "reduce_indptr",
67 "reduce_final_map",
68 "reduce_partial_map",
69 )
70
71
72 def _build_decode_metadata():
73 """Build AITER MLA decode metadata for the fp8/fp8 nhead=32 fold path.
74
75 Returns ``(metadata, captured)`` where ``captured`` records the positional
76 args/kwargs the builder passed to ``get_mla_metadata_v1``, so the golden can
77 be recomputed from the identical inputs.
78 """
79 from tests.v1.attention.utils import (
80 BatchSpec,
81 create_common_attn_metadata,
82 create_vllm_config,
83 )
84 from vllm.config.vllm import set_current_vllm_config
85 from vllm.model_executor.layers.attention.mla_attention import get_mla_dims
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
92 vllm_config = create_vllm_config(
93 model_name="deepseek-ai/DeepSeek-R1",
94 max_model_len=CONTEXT_LEN,
95 # One flat page per token (page_size=1); +buffer for the null block.
96 num_gpu_blocks=BATCH_SIZE * CONTEXT_LEN + 200,
97 block_size=PAGE_SIZE,
98 max_num_seqs=BATCH_SIZE,
99 max_num_batched_tokens=8192,
100 hf_config_override={"num_attention_heads": NUM_QUERY_HEADS},
101 )
102 vllm_config.cache_config.cache_dtype = "fp8"
103
104 spec = MLAAttentionSpec(
105 block_size=PAGE_SIZE,
106 num_kv_heads=1,
107 head_size=vllm_config.model_config.get_head_size(),
108 dtype=vllm_config.model_config.dtype,
109 cache_dtype_str="fp8",
110 )
111
112 builder_cls = AttentionBackendEnum.ROCM_AITER_MLA.get_class().get_builder_cls()
113
114 # The builder reads prefill_backend and the MLA latent dims off the layer,
115 # so the stub carries both; the dims come from the model config to stay in
116 # step with the checkpoint.
117 mla_dims = get_mla_dims(vllm_config.model_config)
118 layer_name = "placeholder"
119 vllm_config.compilation_config.static_forward_context[layer_name] = (
120 types.SimpleNamespace(
121 prefill_backend=torch.empty((1,)),
122 q_lora_rank=mla_dims.q_lora_rank,
123 kv_lora_rank=mla_dims.kv_lora_rank,
124 qk_nope_head_dim=mla_dims.qk_nope_head_dim,
125 qk_rope_head_dim=mla_dims.qk_rope_head_dim,
126 v_head_dim=mla_dims.v_head_dim,
127 )
128 )
129
130 init_workspace_manager(device)
131
132 batch_spec = BatchSpec(
133 seq_lens=[CONTEXT_LEN] * BATCH_SIZE,
134 query_lens=[DECODE_QLEN] * BATCH_SIZE,
135 )
136
137 captured: dict = {}
138
139 with set_current_vllm_config(vllm_config):
140 builder = builder_cls(spec, [layer_name], vllm_config, device)
141 common_attn_metadata = create_common_attn_metadata(
142 batch_spec, PAGE_SIZE, device, arange_block_indices=True
143 )
144
145 import aiter
146
147 real_get_mla_metadata_v1 = aiter.get_mla_metadata_v1
148
149 def spy(*args, **kwargs):
150 captured["args"] = args
151 captured["kwargs"] = dict(kwargs)
152 return real_get_mla_metadata_v1(*args, **kwargs)
153
154 with patch("aiter.get_mla_metadata_v1", spy):
155 metadata = builder.build(
156 common_prefix_len=0,
157 common_attn_metadata=common_attn_metadata,
158 )
159
160 return metadata, captured
161
162
163 def _compute_golden_metadata(captured: dict) -> dict[str, torch.Tensor]:
164 """Recompute the persistent metadata with explicit fp8/bf16 dtypes.
165
166 Replays ``get_mla_metadata_v1`` on the builder's exact input tensors with
167 fresh output buffers and the explicitly-correct dtypes. This reference must
168 match the builder's output when the fix is in place.
169 """
170 import aiter
171
172 args = captured["args"]
173 inputs = args[:_NUM_INPUT_ARGS]
174 # Fresh copies so the golden does not alias the builder's persistent buffers.
175 fresh_outputs = [arg.clone() for arg in args[_NUM_INPUT_ARGS:]]
176
177 golden_kwargs = dict(captured["kwargs"])
178 golden_kwargs["dtype_q"] = EXPECTED_Q_DTYPE
179 golden_kwargs["dtype_kv"] = EXPECTED_KV_DTYPE
180
181 aiter.get_mla_metadata_v1(*inputs, *fresh_outputs, **golden_kwargs)
182
183 return dict(zip(_OUTPUT_ARG_FIELDS, fresh_outputs))
184
185
186 def test_persistent_decode_metadata_matches_fp8_golden():
187 """The builder's metadata must match the dtype-correct golden.
188
189 Regression guard: the fixed builder forwards fp8/bf16 dtypes so its
190 split/reduce metadata matches the golden recomputed with those explicit
191 dtypes. Dropping the dtypes (the original bug) produces a different layout
192 and fails this test.
193 """
194 metadata, captured = _build_decode_metadata()
195
196 # qlen=1 must take the persistent-metadata path for this to be meaningful.
197 assert metadata.decode is not None
198 assert metadata.decode.has_persistent_metadata
199 assert metadata.work_meta_data is not None
200
201 golden = _compute_golden_metadata(captured)
202
203 mismatched = [
204 name
205 for name in _CONTENT_METADATA_FIELDS
206 if getattr(metadata, name).shape != golden[name].shape
207 or not torch.equal(getattr(metadata, name), golden[name])
208 ]
209 assert not mismatched, (
210 "AITER MLA persistent decode metadata does not match the fp8/bf16 "
211 f"golden for fields {mismatched}; the builder must forward "
212 "dtype_q/dtype_kv to get_mla_metadata_v1."
213 )