@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 248 lines 9.75 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Prefill == decode equivalence for the Mamba2 ReplaySSM kernels.
4
5 The SSM recurrence is path-independent: vLLM's chunked prefill (the SSD kernel
6 ``mamba_chunk_scan_combined_varlen``) and step-by-step decode must produce the
7 same per-position outputs and final state over a sequence. This file feeds one
8 set of inputs through the production dt flow (raw dt + softplus + a per-head
9 dt_bias, applied inside each kernel) to:
10
11 * the exact fp32 step recurrence (``selective_state_update_ref``) -- the
12 ground truth,
13 * the chunked prefill kernel,
14 * the baseline decode kernel,
15 * the ReplaySSM output_only decode kernel,
16
17 and checks all of them agree. Prefill (a chunked scan) and decode (a step
18 recurrence) are different code paths, so they differ numerically: the chunked
19 scan carries ~2e-2 (fp32) / ~4e-2 (bf16) vs the exact recurrence, far above the
20 near-exact decode. We therefore anchor every path on the exact recurrence at
21 SSD-level tolerances (the same regime as ``test_mamba_ssm_ssd.py``), which the
22 chunked scan sets, keyed off the activation dtype.
23
24 State and activation/buffer precision are swept independently, including the
25 fp32-state + bf16-activation production config.
26 """
27
28 import pytest
29 import torch
30
31 from tests.kernels.mamba.utils import selective_state_update_ref
32 from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update
33 from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501
34 selective_state_update_replayssm_output_only,
35 )
36 from vllm.model_executor.layers.mamba.ops.ssd_combined import (
37 mamba_chunk_scan_combined_varlen,
38 )
39 from vllm.platforms import current_platform
40 from vllm.utils.torch_utils import set_random_seed
41 from vllm.v1.attention.backends.mamba2_attn import compute_varlen_chunk_metadata
42
43
44 def _prefill_tolerances(act_dtype: torch.dtype) -> tuple[float, float]:
45 # The chunked prefill scan, not the decode, sets these: it carries ~2e-2
46 # (fp32) / ~6e-2 (bf16) vs the exact recurrence, while ReplaySSM decode is
47 # near-exact (~2e-6 fp32). Keyed off the activation dtype (the outputs are
48 # in act_dtype). Same regime as test_mamba_ssm_ssd.py.
49 if act_dtype == torch.float32:
50 return 1e-2, 3e-2
51 return 6e-2, 1e-1
52
53
54 def _run_prefill_decode_equivalence(
55 *,
56 state_dtype: torch.dtype,
57 act_dtype: torch.dtype,
58 nheads: int,
59 headdim: int,
60 ngroups: int,
61 dstate: int,
62 seqlen: int,
63 chunk_size: int,
64 max_cache_len: int,
65 seed: int = 0,
66 ) -> None:
67 """Prefill the whole sequence and decode it step by step; check both match
68 the exact fp32 recurrence (and each other). All paths use the production dt
69 flow (raw dt + softplus + a per-head dt_bias), so this also checks prefill
70 and decode apply the softplus/bias preprocessing consistently. ``state_dtype``
71 is the recurrent-state precision; ``act_dtype`` the activation/buffer one."""
72 device = current_platform.device_type
73 rtol, atol = _prefill_tolerances(act_dtype)
74 set_random_seed(seed)
75
76 # Production dt flow: raw dt + a per-head dt_bias (~-4 keeps softplus(dt +
77 # bias) small and well-conditioned). dt_bias is (nheads,) for prefill and
78 # (nheads, headdim) for the decode kernels/reference.
79 A = -torch.exp(torch.rand(nheads, device=device, dtype=act_dtype))
80 dt = torch.randn(seqlen, nheads, device=device, dtype=act_dtype)
81 dt_bias = torch.rand(nheads, device=device, dtype=act_dtype) - 4
82 dt_bias_hd = dt_bias.view(nheads, 1).expand(nheads, headdim)
83 X = torch.randn(seqlen, nheads, headdim, device=device, dtype=act_dtype)
84 B = torch.randn(seqlen, ngroups, dstate, device=device, dtype=act_dtype)
85 C = torch.randn(seqlen, ngroups, dstate, device=device, dtype=act_dtype)
86 A_bcast = A.view(nheads, 1, 1).expand(nheads, headdim, dstate)
87
88 # Chunked prefill over the whole sequence (implicit batch=1, varlen). The
89 # kernel always returns the final state in fp32, so no state_dtype plumbing.
90 cu_seqlens = torch.tensor((0, seqlen), device=device).cumsum(0).to(torch.int32)
91 cu_chunk_seqlens, last_chunk_indices, seq_idx = compute_varlen_chunk_metadata(
92 cu_seqlens, chunk_size
93 )
94 y_prefill = torch.empty(seqlen, nheads, headdim, device=device, dtype=act_dtype)
95 final_state_prefill = mamba_chunk_scan_combined_varlen(
96 X,
97 dt,
98 A,
99 B,
100 C,
101 chunk_size,
102 cu_seqlens=cu_seqlens,
103 cu_chunk_seqlens=cu_chunk_seqlens,
104 last_chunk_indices=last_chunk_indices,
105 seq_idx=seq_idx,
106 out=y_prefill,
107 D=None,
108 dt_bias=dt_bias,
109 dt_softplus=True,
110 )
111
112 # Step paths: exact fp32 recurrence (ground truth), baseline, ReplaySSM.
113 # State follows state_dtype; caches follow act_dtype (dt_cache is fp32).
114 state_ref = torch.zeros(
115 1, nheads, headdim, dstate, device=device, dtype=torch.float32
116 )
117 state_base = torch.zeros(
118 1, nheads, headdim, dstate, device=device, dtype=state_dtype
119 )
120 state_dec = torch.zeros(
121 1, nheads, headdim, dstate, device=device, dtype=state_dtype
122 )
123 x_cache = torch.zeros(
124 1, nheads, max_cache_len, headdim, device=device, dtype=act_dtype
125 )
126 dt_cache = torch.zeros(1, nheads, max_cache_len, device=device, dtype=torch.float32)
127 B_cache = torch.zeros(
128 1, ngroups, max_cache_len, dstate, device=device, dtype=act_dtype
129 )
130 bc_pre = torch.empty(1, ngroups, max_cache_len, device=device, dtype=torch.float32)
131 write_pos = torch.zeros(1, dtype=torch.int32, device=device)
132 # No skip connection (D=0) on any path here; the D!=0 path is covered by the
133 # standard-decode suite. The baseline kernel needs a D tensor, not None.
134 D_zero = torch.zeros(nheads, headdim, device=device)
135
136 y_ref = torch.empty(seqlen, nheads, headdim, device=device, dtype=torch.float32)
137 y_base = torch.empty(seqlen, nheads, headdim, device=device, dtype=act_dtype)
138 y_dec = torch.empty(seqlen, nheads, headdim, device=device, dtype=act_dtype)
139 for t in range(seqlen):
140 dt_t = dt[t].view(1, nheads, 1).expand(1, nheads, headdim)
141 is_flush = write_pos == max_cache_len - 1
142
143 y_ref[t] = selective_state_update_ref(
144 state_ref,
145 X[t : t + 1].float(),
146 dt_t.float(),
147 A_bcast.float(),
148 B[t : t + 1].float(),
149 C[t : t + 1].float(),
150 dt_bias=dt_bias_hd.float(),
151 dt_softplus=True,
152 )[0]
153
154 out_b = torch.empty(1, nheads, headdim, device=device, dtype=act_dtype)
155 selective_state_update(
156 state_base,
157 X[t : t + 1],
158 dt_t,
159 A_bcast,
160 B[t : t + 1],
161 C[t : t + 1],
162 D=D_zero,
163 dt_bias=dt_bias_hd,
164 dt_softplus=True,
165 out=out_b,
166 )
167 y_base[t] = out_b[0]
168
169 out_d = torch.empty(1, nheads, headdim, device=device, dtype=act_dtype)
170 common = dict(
171 dt_bias=dt_bias_hd,
172 dt_softplus=True,
173 x_cache=x_cache,
174 dt_cache=dt_cache,
175 B_cache=B_cache,
176 write_pos=write_pos,
177 is_flush=is_flush,
178 max_cache_len=max_cache_len,
179 out=out_d,
180 )
181 selective_state_update_replayssm_output_only(
182 state_dec,
183 X[t : t + 1],
184 dt_t,
185 A_bcast,
186 B[t : t + 1],
187 C[t : t + 1],
188 bc_pre=bc_pre,
189 **common,
190 )
191 y_dec[t] = out_d[0]
192
193 write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1)
194
195 # Every path computes the same recurrence; anchor each on the fp32 truth.
196 torch.testing.assert_close(y_prefill.float(), y_ref, rtol=rtol, atol=atol)
197 torch.testing.assert_close(y_base.float(), y_ref, rtol=rtol, atol=atol)
198 torch.testing.assert_close(y_dec.float(), y_ref, rtol=rtol, atol=atol)
199 # Headline: ReplaySSM decode matches the chunked prefill directly.
200 torch.testing.assert_close(y_dec.float(), y_prefill.float(), rtol=rtol, atol=atol)
201 # Final state too (the recurrence ends in state_ref after the loop).
202 torch.testing.assert_close(
203 final_state_prefill[0].float(), state_ref[0], rtol=rtol, atol=atol
204 )
205
206
207 @pytest.mark.skipif(
208 not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
209 reason="Mamba2 ReplaySSM kernels require a CUDA-alike or XPU device.",
210 )
211 @pytest.mark.parametrize(
212 "precision",
213 # fp32 state is the default; bf16/fp16 are reduced-footprint configs. fp16
214 # appears as an activation dtype (s32_afp16, sfp16_afp16) and as a finer-
215 # mantissa state under bf16 activations (sfp16_a16).
216 [
217 pytest.param((torch.float32, torch.float32), id="s32_a32"),
218 pytest.param((torch.float32, torch.bfloat16), id="s32_a16"),
219 pytest.param((torch.bfloat16, torch.bfloat16), id="s16_a16"),
220 pytest.param((torch.float32, torch.float16), id="s32_afp16"),
221 pytest.param((torch.float16, torch.float16), id="sfp16_afp16"),
222 pytest.param((torch.float16, torch.bfloat16), id="sfp16_a16"),
223 ],
224 )
225 @pytest.mark.parametrize(
226 "geometry", # (nheads, headdim, dstate, ngroups)
227 [
228 pytest.param((8, 64, 64, 2), id="small"),
229 pytest.param((96, 80, 128, 8), id="nano4b"),
230 ],
231 )
232 def test_replayssm_prefill_decode_equivalence(
233 precision: tuple[torch.dtype, torch.dtype],
234 geometry: tuple[int, int, int, int],
235 ):
236 state_dtype, act_dtype = precision
237 nheads, headdim, dstate, ngroups = geometry
238 _run_prefill_decode_equivalence(
239 state_dtype=state_dtype,
240 act_dtype=act_dtype,
241 nheads=nheads,
242 headdim=headdim,
243 ngroups=ngroups,
244 dstate=dstate,
245 seqlen=16,
246 chunk_size=8,
247 max_cache_len=4,
248 )