@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 346 lines 11.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Mamba2 ReplaySSM decode write-position derivation in
4 BaseMambaAttentionMetadataBuilder: write_pos and is_flush computed from the
5 per-request ring origin (replayssm_decode_base) and num_computed.
6 """
7
8 from dataclasses import dataclass
9
10 import pytest
11 import torch
12
13 from tests.v1.attention.utils import (
14 BatchSpec,
15 MockMambaBuilder,
16 create_common_attn_metadata,
17 create_vllm_config,
18 )
19 from vllm.config import SpeculativeConfig
20 from vllm.config.compilation import CUDAGraphMode
21 from vllm.config.mamba import MambaBackendEnum
22 from vllm.v1.kv_cache_interface import MambaSpec
23
24 BLOCK_SIZE = 16
25 DEVICE = torch.device("cpu")
26
27
28 @dataclass
29 class ReplaySSMBuildCase:
30 """A decode batch and its expected per-row write_pos / is_flush.
31
32 num_computed = seq_len - query_len; write_pos =
33 (num_computed - decode_base) % buffer_len; is_flush = write_pos ==
34 buffer_len - 1 (or a forced one-token flush when num_computed < decode_base).
35 """
36
37 seq_lens: list[int]
38 query_lens: list[int]
39 is_prefilling: list[bool]
40 decode_base: list[int]
41 buffer_len: int
42 expected_write_pos: list[int]
43 expected_is_flush: list[int]
44 mamba_cache_mode: str = "none"
45
46
47 REPLAYSSM_BUILD_CASES = {
48 # decode_base == num_prompt (fresh request).
49 "fresh_decode": ReplaySSMBuildCase(
50 seq_lens=[106],
51 query_lens=[1],
52 is_prefilling=[False],
53 decode_base=[100],
54 buffer_len=16,
55 expected_write_pos=[5],
56 expected_is_flush=[0],
57 ),
58 # decode_base > num_prompt anchors write_pos at the resume point.
59 "resumed_reanchors_to_zero": ReplaySSMBuildCase(
60 seq_lens=[106],
61 query_lens=[1],
62 is_prefilling=[False],
63 decode_base=[105],
64 buffer_len=16,
65 expected_write_pos=[0],
66 expected_is_flush=[0],
67 ),
68 # write_pos == buffer_len - 1 flushes.
69 "flush_boundary": ReplaySSMBuildCase(
70 seq_lens=[116],
71 query_lens=[1],
72 is_prefilling=[False],
73 decode_base=[100],
74 buffer_len=16,
75 expected_write_pos=[15],
76 expected_is_flush=[1],
77 ),
78 # Resumed request landing on a flush boundary.
79 "resumed_flush_boundary": ReplaySSMBuildCase(
80 seq_lens=[121],
81 query_lens=[1],
82 is_prefilling=[False],
83 decode_base=[105],
84 buffer_len=16,
85 expected_write_pos=[15],
86 expected_is_flush=[1],
87 ),
88 # Per-row write_pos / is_flush are independent.
89 "mixed_rows": ReplaySSMBuildCase(
90 seq_lens=[104, 106, 216],
91 query_lens=[1, 1, 1],
92 is_prefilling=[False, False, False],
93 decode_base=[100, 105, 200],
94 buffer_len=16,
95 expected_write_pos=[3, 0, 15],
96 expected_is_flush=[0, 0, 1],
97 ),
98 # write_pos wraps within the buffer (6 % 4 == 2).
99 "small_buffer_wrap": ReplaySSMBuildCase(
100 seq_lens=[112],
101 query_lens=[1],
102 is_prefilling=[False],
103 decode_base=[105],
104 buffer_len=4,
105 expected_write_pos=[2],
106 expected_is_flush=[0],
107 ),
108 # Single-token prefill-as-decode still in the prompt (num_computed <
109 # decode_base): forced one-token flush.
110 "leftover_prompt_one_token_flush": ReplaySSMBuildCase(
111 seq_lens=[100],
112 query_lens=[1],
113 is_prefilling=[True],
114 decode_base=[100],
115 buffer_len=16,
116 expected_write_pos=[0],
117 expected_is_flush=[1],
118 ),
119 # Align mode (block_size 16). Past the first boundary the ring re-anchors at
120 # the block start: num_computed 117 -> block_start 112, write_pos 5 (vs 1 in
121 # none mode).
122 "align_reanchor_past_boundary": ReplaySSMBuildCase(
123 seq_lens=[118],
124 query_lens=[1],
125 is_prefilling=[False],
126 decode_base=[100],
127 buffer_len=16,
128 expected_write_pos=[5],
129 expected_is_flush=[0],
130 mamba_cache_mode="align",
131 ),
132 # First-block boundary (num_computed+1 == 112) forces a flush even though
133 # write_pos (11) != buffer_len - 1.
134 "align_first_block_boundary_flush": ReplaySSMBuildCase(
135 seq_lens=[112],
136 query_lens=[1],
137 is_prefilling=[False],
138 decode_base=[100],
139 buffer_len=16,
140 expected_write_pos=[11],
141 expected_is_flush=[1],
142 mamba_cache_mode="align",
143 ),
144 # First step of a new block re-anchors write_pos to 0.
145 "align_new_block_start_zero": ReplaySSMBuildCase(
146 seq_lens=[113],
147 query_lens=[1],
148 is_prefilling=[False],
149 decode_base=[100],
150 buffer_len=16,
151 expected_write_pos=[0],
152 expected_is_flush=[0],
153 mamba_cache_mode="align",
154 ),
155 # block_size % buffer_len == 0: a later boundary lands on write_pos ==
156 # buffer_len - 1, so the boundary flush coincides with the natural flush.
157 "align_boundary_coincides_natural_flush": ReplaySSMBuildCase(
158 seq_lens=[128],
159 query_lens=[1],
160 is_prefilling=[False],
161 decode_base=[100],
162 buffer_len=16,
163 expected_write_pos=[15],
164 expected_is_flush=[1],
165 mamba_cache_mode="align",
166 ),
167 # block_size % buffer_len != 0 (buffer_len 6): the boundary step still flushes
168 # although write_pos (3) != buffer_len - 1.
169 "align_unaligned_buffer_forces_flush": ReplaySSMBuildCase(
170 seq_lens=[128],
171 query_lens=[1],
172 is_prefilling=[False],
173 decode_base=[100],
174 buffer_len=6,
175 expected_write_pos=[3],
176 expected_is_flush=[1],
177 mamba_cache_mode="align",
178 ),
179 # Per-row independence in align mode: partial-block / new-block / boundary.
180 "align_mixed_rows": ReplaySSMBuildCase(
181 seq_lens=[105, 113, 112],
182 query_lens=[1, 1, 1],
183 is_prefilling=[False, False, False],
184 decode_base=[100, 100, 100],
185 buffer_len=16,
186 expected_write_pos=[4, 0, 11],
187 expected_is_flush=[0, 0, 1],
188 mamba_cache_mode="align",
189 ),
190 }
191
192
193 def _make_mamba_spec(
194 buffer_len: int,
195 mamba_backend: MambaBackendEnum,
196 num_speculative_tokens: int = 0,
197 ) -> MambaSpec:
198 ring_buffer_len = buffer_len + (
199 1 + num_speculative_tokens
200 if mamba_backend == MambaBackendEnum.FLASHINFER
201 else 0
202 )
203 shapes = (
204 (1, 1),
205 (1, 1, 1),
206 (1, ring_buffer_len, 1),
207 (1, ring_buffer_len),
208 (1, ring_buffer_len, 1),
209 )
210 return MambaSpec(
211 block_size=BLOCK_SIZE,
212 shapes=shapes,
213 dtypes=(torch.float32,),
214 )
215
216
217 def _create_replayssm_builder(
218 buffer_len: int,
219 mamba_cache_mode: str = "none",
220 *,
221 mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON,
222 num_speculative_tokens: int = 0,
223 ) -> MockMambaBuilder:
224 vllm_config = create_vllm_config(
225 model_name="Qwen/Qwen3.5-0.8B", block_size=BLOCK_SIZE
226 )
227 # Set the flags after construction to skip validate_mamba_cached_kernel
228 # (it requires a real SupportsReplaySSM model) on the mock model.
229 vllm_config.cache_config.use_replayssm = True
230 vllm_config.cache_config.replayssm_buffer_len = buffer_len
231 vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
232 vllm_config.mamba_config.backend = mamba_backend
233 if num_speculative_tokens > 0:
234 # Triton ReplaySSM requires exact synchronous CPU metadata.
235 vllm_config.scheduler_config.async_scheduling = False
236 vllm_config.speculative_config = SpeculativeConfig(
237 method="ngram",
238 num_speculative_tokens=num_speculative_tokens,
239 )
240 return MockMambaBuilder(
241 _make_mamba_spec(buffer_len, mamba_backend, num_speculative_tokens),
242 ["layer0"],
243 vllm_config,
244 DEVICE,
245 )
246
247
248 def _build(
249 builder: MockMambaBuilder,
250 case: ReplaySSMBuildCase,
251 num_accepted_tokens: torch.Tensor | None = None,
252 ):
253 batch = BatchSpec(seq_lens=case.seq_lens, query_lens=case.query_lens)
254 common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace(
255 is_prefilling=torch.tensor(case.is_prefilling, dtype=torch.bool),
256 replayssm_decode_base_cpu=torch.tensor(case.decode_base, dtype=torch.int32),
257 )
258 return builder.build(0, common, num_accepted_tokens=num_accepted_tokens)
259
260
261 @pytest.mark.parametrize(
262 "case", REPLAYSSM_BUILD_CASES.values(), ids=REPLAYSSM_BUILD_CASES.keys()
263 )
264 def test_replayssm_write_pos(case: ReplaySSMBuildCase):
265 builder = _create_replayssm_builder(case.buffer_len, case.mamba_cache_mode)
266 meta = _build(builder, case)
267
268 assert meta.write_pos_d is not None
269 assert meta.is_flush_d is not None
270 n = len(case.expected_write_pos)
271 assert meta.write_pos_d[:n].tolist() == case.expected_write_pos
272 assert meta.is_flush_d[:n].tolist() == case.expected_is_flush
273
274
275 def test_resumed_request_differs_from_fresh():
276 """Same token count, different decode_base: fresh (base 100) -> write_pos 5,
277 resumed (base 105) -> write_pos 0."""
278 builder = _create_replayssm_builder(16)
279 batch = BatchSpec(seq_lens=[106, 106], query_lens=[1, 1])
280 common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace(
281 is_prefilling=torch.tensor([False, False]),
282 replayssm_decode_base_cpu=torch.tensor([100, 105], dtype=torch.int32),
283 )
284 meta = builder.build(0, common)
285
286 assert meta.write_pos_d.tolist()[:2] == [5, 0]
287 assert meta.is_flush_d.tolist()[:2] == [0, 0]
288
289
290 def test_spec_decode_single_token_chunk_synthesizes_acceptance_metadata():
291 builder = _create_replayssm_builder(16, num_speculative_tokens=3)
292 case = REPLAYSSM_BUILD_CASES["leftover_prompt_one_token_flush"]
293
294 meta = _build(builder, case)
295
296 assert meta.query_start_loc_d is not None
297 assert meta.query_start_loc_d.tolist() == [0, 1]
298 assert meta.num_accepted_tokens is not None
299 assert meta.num_accepted_tokens.tolist() == [1]
300
301
302 def test_flashinfer_replayssm_state_indices_are_stable_for_full_cudagraph():
303 checkpointing_ssu = pytest.importorskip("flashinfer.mamba.checkpointing_ssu")
304 if not hasattr(checkpointing_ssu, "allocate_checkpointing_ssu_scratch"):
305 pytest.skip("requires FlashInfer ReplaySSM autotuning support")
306
307 builder = _create_replayssm_builder(
308 16,
309 mamba_backend=MambaBackendEnum.FLASHINFER,
310 num_speculative_tokens=3,
311 )
312 builder.compilation_config.cudagraph_mode = CUDAGraphMode.FULL
313
314 first = _build(
315 builder,
316 ReplaySSMBuildCase(
317 seq_lens=[106, 106],
318 query_lens=[1, 1],
319 is_prefilling=[False, False],
320 decode_base=[100, 100],
321 buffer_len=16,
322 expected_write_pos=[],
323 expected_is_flush=[],
324 ),
325 )
326 first_indices = first.replayssm_state_indices_d
327 assert first_indices is not None
328 assert first_indices.is_contiguous()
329 first_ptr = first_indices.data_ptr()
330
331 second = _build(
332 builder,
333 ReplaySSMBuildCase(
334 seq_lens=[122, 122],
335 query_lens=[1, 1],
336 is_prefilling=[False, False],
337 decode_base=[116, 116],
338 buffer_len=16,
339 expected_write_pos=[],
340 expected_is_flush=[],
341 ),
342 )
343 second_indices = second.replayssm_state_indices_d
344 assert second_indices is not None
345 assert second_indices.data_ptr() == first_ptr
346 assert torch.equal(second_indices, second.state_indices_tensor_d[:, 0])