| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Standard (autoregressive) decode correctness for the Mamba2 ReplaySSM kernels. |
| 4 | |
| 5 | ReplaySSM caches the recent SSM inputs ``(x, dt, B)`` in a small ring buffer and |
| 6 | reconstructs / reads out the recurrent state on the fly, writing the full state |
| 7 | back to HBM only when the buffer flushes. This file checks that, over a |
| 8 | multi-step decode, the ReplaySSM output_only kernel reproduces the exact SSM |
| 9 | recurrence. |
| 10 | |
| 11 | For every step we assert, against trusted oracles driven one token at a time: |
| 12 | |
| 13 | * the ReplaySSM output matches its pure-PyTorch reference, which models the |
| 14 | kernel's exact arithmetic (including its bf16 reconstruction), at every |
| 15 | precision, |
| 16 | * when the state is fp32, the output also matches the baseline decode kernel |
| 17 | (``selective_state_update``); at bf16 state the baseline legitimately |
| 18 | differs -- it downcasts the fp32 state to bf16 every step, while ReplaySSM |
| 19 | accumulates a whole buffer in fp32 and is the more accurate path, |
| 20 | * when state and activations are both fp32, the output also matches the exact |
| 21 | elementwise ``selective_state_update_ref``, |
| 22 | * the cached inputs match the reference cache management, |
| 23 | * the checkpoint state matches the reference (and the baseline at fp32 state). |
| 24 | |
| 25 | State and activation/buffer precision are swept independently. Nemotron-3 |
| 26 | defaults to fp32 SSM state (``mamba_ssm_cache_dtype=float32``), but bf16 state is |
| 27 | also supported; the buffer dtype follows the activation dtype and ``dt_cache`` |
| 28 | is always fp32. |
| 29 | """ |
| 30 | |
| 31 | import pytest |
| 32 | import torch |
| 33 | |
| 34 | from tests.kernels.mamba.utils import ( |
| 35 | allocate_update_caches, |
| 36 | selective_state_update_ref, |
| 37 | selective_state_update_replayssm_output_only_ref, |
| 38 | ) |
| 39 | from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update |
| 40 | from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501 |
| 41 | selective_state_update_replayssm_output_only, |
| 42 | ) |
| 43 | from vllm.utils.torch_utils import set_random_seed |
| 44 | from vllm.v1.attention.backends.utils import NULL_BLOCK_ID |
| 45 | |
| 46 | |
| 47 | def _tolerances(dtype: torch.dtype) -> tuple[float, float]: |
| 48 | # fp32 demands true fp32 parity. A correct fp32 reconstruction (tl.dot with |
| 49 | # input_precision="tf32x3"/"ieee") matches the elementwise baseline to |
| 50 | # ~1e-5, while a TF32 reconstruction drifts to ~1e-2. atol=1e-3 sits between, |
| 51 | # so it flags TF32 degradation yet passes a correct fp32 kernel. bf16 stays |
| 52 | # loose (bf16 rounding dominates and is unaffected by the matmul precision). |
| 53 | if dtype == torch.float32: |
| 54 | return 1e-4, 1e-3 |
| 55 | return 6e-2, 2e-1 |
| 56 | |
| 57 | |
| 58 | def _tied_A(nheads: int, headdim: int, dstate: int, device: str) -> torch.Tensor: |
| 59 | A = -torch.rand(nheads, device=device) - 1.0 |
| 60 | return A.view(nheads, 1, 1).expand(nheads, headdim, dstate) |
| 61 | |
| 62 | |
| 63 | def _tied_dt( |
| 64 | batch: int, |
| 65 | nheads: int, |
| 66 | headdim: int, |
| 67 | device: str, |
| 68 | dtype: torch.dtype, |
| 69 | ) -> torch.Tensor: |
| 70 | dt = torch.randn(batch, nheads, device=device, dtype=dtype) |
| 71 | return dt.unsqueeze(-1).expand(batch, nheads, headdim) |
| 72 | |
| 73 | |
| 74 | def _tied_dt_bias(nheads: int, headdim: int, device: str) -> torch.Tensor: |
| 75 | dt_bias = torch.rand(nheads, device=device) - 4.0 |
| 76 | return dt_bias.view(nheads, 1).expand(nheads, headdim) |
| 77 | |
| 78 | |
| 79 | def _run_standard_decode( |
| 80 | *, |
| 81 | state_dtype: torch.dtype, |
| 82 | act_dtype: torch.dtype, |
| 83 | batch: int, |
| 84 | nheads: int, |
| 85 | headdim: int, |
| 86 | ngroups: int, |
| 87 | dstate: int, |
| 88 | max_cache_len: int, |
| 89 | num_steps: int, |
| 90 | has_z: bool, |
| 91 | dt_softplus: bool, |
| 92 | use_dt_bias: bool, |
| 93 | desync_write_pos: bool = False, |
| 94 | seed: int = 0, |
| 95 | ) -> None: |
| 96 | """Drive ``num_steps`` decode steps and check every path against the |
| 97 | trusted recurrence. ``state_dtype`` is the recurrent-state precision; |
| 98 | ``act_dtype`` is the activation/buffer precision.""" |
| 99 | device = "cuda" |
| 100 | both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 |
| 101 | rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16) |
| 102 | set_random_seed(seed) |
| 103 | |
| 104 | # One state copy per path; all start identical at the same precision. |
| 105 | state0 = torch.randn( |
| 106 | batch, nheads, headdim, dstate, dtype=state_dtype, device=device |
| 107 | ) |
| 108 | state_anchor = state0.clone() |
| 109 | state_baseline = state0.clone() |
| 110 | state_cached = state0.clone() |
| 111 | state_ref = state0.clone() |
| 112 | |
| 113 | A = _tied_A(nheads, headdim, dstate, device) |
| 114 | dt_bias = _tied_dt_bias(nheads, headdim, device) if use_dt_bias else None |
| 115 | D = torch.randn(nheads, headdim, device=device) |
| 116 | |
| 117 | # Caches follow the activation dtype (dt_cache is forced to fp32 inside). |
| 118 | x_cache, dt_cache, B_cache, _ = allocate_update_caches( |
| 119 | batch, |
| 120 | nheads, |
| 121 | ngroups, |
| 122 | headdim, |
| 123 | dstate, |
| 124 | max_cache_len, |
| 125 | device, |
| 126 | act_dtype, |
| 127 | act_dtype, |
| 128 | ) |
| 129 | x_cache_ref, dt_cache_ref, B_cache_ref, _ = allocate_update_caches( |
| 130 | batch, |
| 131 | nheads, |
| 132 | ngroups, |
| 133 | headdim, |
| 134 | dstate, |
| 135 | max_cache_len, |
| 136 | device, |
| 137 | act_dtype, |
| 138 | act_dtype, |
| 139 | ) |
| 140 | bc_pre = torch.empty( |
| 141 | batch, ngroups, max_cache_len, device=device, dtype=torch.float32 |
| 142 | ) |
| 143 | |
| 144 | if desync_write_pos: |
| 145 | # Rows start at different ring positions so they flush on different |
| 146 | # steps, exercising per-row write-position handling. |
| 147 | write_pos = ( |
| 148 | torch.arange(batch, device=device, dtype=torch.int32) % max_cache_len |
| 149 | ) |
| 150 | else: |
| 151 | write_pos = torch.zeros(batch, dtype=torch.int32, device=device) |
| 152 | |
| 153 | for _ in range(num_steps): |
| 154 | x = torch.randn(batch, nheads, headdim, device=device, dtype=act_dtype) |
| 155 | dt = _tied_dt(batch, nheads, headdim, device, act_dtype) |
| 156 | B = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) |
| 157 | C = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) |
| 158 | z = torch.randn_like(x) if has_z else None |
| 159 | is_flush = write_pos == max_cache_len - 1 |
| 160 | |
| 161 | # Trusted recurrence (mutates state_anchor in place, returns output). |
| 162 | out_anchor = selective_state_update_ref( |
| 163 | state_anchor, |
| 164 | x, |
| 165 | dt, |
| 166 | A, |
| 167 | B, |
| 168 | C, |
| 169 | D=D, |
| 170 | z=z, |
| 171 | dt_bias=dt_bias, |
| 172 | dt_softplus=dt_softplus, |
| 173 | ) |
| 174 | |
| 175 | # Upstream baseline decode kernel. |
| 176 | out_baseline = torch.empty_like(x) |
| 177 | selective_state_update( |
| 178 | state_baseline, |
| 179 | x, |
| 180 | dt, |
| 181 | A, |
| 182 | B, |
| 183 | C, |
| 184 | D=D, |
| 185 | z=z, |
| 186 | dt_bias=dt_bias, |
| 187 | dt_softplus=dt_softplus, |
| 188 | out=out_baseline, |
| 189 | ) |
| 190 | |
| 191 | # ReplaySSM kernel under test + its pure-PyTorch reference. |
| 192 | out_cached = torch.empty_like(x) |
| 193 | common = dict( |
| 194 | D=D, |
| 195 | z=z, |
| 196 | dt_bias=dt_bias, |
| 197 | dt_softplus=dt_softplus, |
| 198 | x_cache=x_cache, |
| 199 | dt_cache=dt_cache, |
| 200 | B_cache=B_cache, |
| 201 | write_pos=write_pos, |
| 202 | is_flush=is_flush, |
| 203 | max_cache_len=max_cache_len, |
| 204 | out=out_cached, |
| 205 | ) |
| 206 | selective_state_update_replayssm_output_only( |
| 207 | state_cached, x, dt, A, B, C, bc_pre=bc_pre, **common |
| 208 | ) |
| 209 | out_ref = selective_state_update_replayssm_output_only_ref( |
| 210 | state_ref, |
| 211 | x, |
| 212 | dt, |
| 213 | A, |
| 214 | B, |
| 215 | C, |
| 216 | D=D, |
| 217 | z=z, |
| 218 | dt_bias=dt_bias, |
| 219 | dt_softplus=dt_softplus, |
| 220 | x_cache=x_cache_ref, |
| 221 | dt_cache=dt_cache_ref, |
| 222 | B_cache=B_cache_ref, |
| 223 | write_pos=write_pos, |
| 224 | max_cache_len=max_cache_len, |
| 225 | ) |
| 226 | |
| 227 | # The reference models the kernel's exact arithmetic (including its bf16 |
| 228 | # reconstruction), so the kernel must match it tightly at every |
| 229 | # precision. At fp32 this also flags any TF32 reconstruction drift. |
| 230 | torch.testing.assert_close(out_cached, out_ref, rtol=rtol, atol=atol) |
| 231 | # When the STATE is fp32 the baseline decode kernel is a valid oracle (it |
| 232 | # does not downcast the state per step), so ReplaySSM must match it. At |
| 233 | # bf16 state the baseline legitimately differs: it downcasts the fp32 |
| 234 | # state to bf16 every step, while ReplaySSM accumulates a whole buffer in |
| 235 | # fp32 and is the MORE accurate path -- so it is not a tight oracle there. |
| 236 | if state_dtype == torch.float32: |
| 237 | torch.testing.assert_close(out_cached, out_baseline, rtol=rtol, atol=atol) |
| 238 | # The exact elementwise reference is valid only when state AND |
| 239 | # activations are fp32; otherwise it downcasts the state (at readout for |
| 240 | # bf16 activations, or per step for bf16 state). |
| 241 | if both_fp32: |
| 242 | torch.testing.assert_close(out_cached, out_anchor, rtol=rtol, atol=atol) |
| 243 | |
| 244 | # Cached inputs match the reference cache management. |
| 245 | torch.testing.assert_close(x_cache, x_cache_ref, rtol=rtol, atol=atol) |
| 246 | torch.testing.assert_close(dt_cache, dt_cache_ref, rtol=rtol, atol=atol) |
| 247 | torch.testing.assert_close(B_cache, B_cache_ref, rtol=rtol, atol=atol) |
| 248 | |
| 249 | # Checkpoint state at flush matches the reference (and, when the state is |
| 250 | # fp32, the baseline kernel). |
| 251 | if bool(is_flush.any()): |
| 252 | torch.testing.assert_close( |
| 253 | state_cached[is_flush], state_ref[is_flush], rtol=rtol, atol=atol |
| 254 | ) |
| 255 | if state_dtype == torch.float32: |
| 256 | torch.testing.assert_close( |
| 257 | state_cached[is_flush], |
| 258 | state_baseline[is_flush], |
| 259 | rtol=rtol, |
| 260 | atol=atol, |
| 261 | ) |
| 262 | |
| 263 | write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1) |
| 264 | |
| 265 | |
| 266 | # State/activation precisions. fp32 state is the default; bf16/fp16 are the |
| 267 | # reduced-footprint configs. fp16 appears both as an activation dtype (fully-fp16 |
| 268 | # model sfp16_afp16, or fp16 act over fp32 state s32_afp16) and as a state dtype |
| 269 | # under bf16 activations (sfp16_a16): fp16 has a finer mantissa than bf16 at the |
| 270 | # same 2 bytes, so it is a more accurate state at no extra footprint. We still |
| 271 | # skip fp16 state under fp32 activations (the unused low-state/high-act mix). |
| 272 | _PRECISIONS = [ |
| 273 | pytest.param((torch.float32, torch.float32), id="s32_a32"), |
| 274 | pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), |
| 275 | pytest.param((torch.bfloat16, torch.bfloat16), id="s16_a16"), |
| 276 | pytest.param((torch.float32, torch.float16), id="s32_afp16"), |
| 277 | pytest.param((torch.float16, torch.float16), id="sfp16_afp16"), |
| 278 | pytest.param((torch.float16, torch.bfloat16), id="sfp16_a16"), |
| 279 | ] |
| 280 | # Small synthetic shapes for the full axis sweep (compile fast). |
| 281 | _SMALL_GEOMETRIES = [ |
| 282 | pytest.param((8, 64, 64, 4), id="small"), |
| 283 | pytest.param((4, 64, 16, 1), id="tiny"), |
| 284 | ] |
| 285 | # Production Mamba2 shapes (nheads, headdim, dstate, ngroups), TP=1. |
| 286 | _REAL_GEOMETRIES = [ |
| 287 | pytest.param((96, 80, 128, 8), id="nano4b"), |
| 288 | pytest.param((128, 64, 128, 8), id="super120b"), |
| 289 | pytest.param((256, 64, 128, 8), id="ultra550b"), |
| 290 | ] |
| 291 | |
| 292 | |
| 293 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") |
| 294 | @pytest.mark.parametrize("precision", _PRECISIONS) |
| 295 | @pytest.mark.parametrize("max_cache_len", [1, 4, 16]) |
| 296 | @pytest.mark.parametrize("geometry", _SMALL_GEOMETRIES) |
| 297 | @pytest.mark.parametrize("has_z", [False, True]) |
| 298 | def test_replayssm_standard_decode_matches_reference( |
| 299 | precision: tuple[torch.dtype, torch.dtype], |
| 300 | max_cache_len: int, |
| 301 | geometry: tuple[int, int, int, int], |
| 302 | has_z: bool, |
| 303 | ): |
| 304 | state_dtype, act_dtype = precision |
| 305 | nheads, headdim, dstate, ngroups = geometry |
| 306 | _run_standard_decode( |
| 307 | state_dtype=state_dtype, |
| 308 | act_dtype=act_dtype, |
| 309 | batch=4, |
| 310 | nheads=nheads, |
| 311 | headdim=headdim, |
| 312 | ngroups=ngroups, |
| 313 | dstate=dstate, |
| 314 | max_cache_len=max_cache_len, |
| 315 | num_steps=2 * max_cache_len + 1, |
| 316 | has_z=has_z, |
| 317 | dt_softplus=True, |
| 318 | use_dt_bias=True, |
| 319 | ) |
| 320 | |
| 321 | |
| 322 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") |
| 323 | @pytest.mark.parametrize("precision", _PRECISIONS) |
| 324 | @pytest.mark.parametrize("geometry", _REAL_GEOMETRIES) |
| 325 | def test_replayssm_standard_decode_real_geometry( |
| 326 | precision: tuple[torch.dtype, torch.dtype], |
| 327 | geometry: tuple[int, int, int, int], |
| 328 | ): |
| 329 | # Production Mamba2 shapes for the Nemotron-3 family (Nano-4B / Super-120B / |
| 330 | # Ultra-550B), at the production buffer length (8). All three precisions, |
| 331 | # including the bf16 state case. |
| 332 | state_dtype, act_dtype = precision |
| 333 | nheads, headdim, dstate, ngroups = geometry |
| 334 | _run_standard_decode( |
| 335 | state_dtype=state_dtype, |
| 336 | act_dtype=act_dtype, |
| 337 | batch=4, |
| 338 | nheads=nheads, |
| 339 | headdim=headdim, |
| 340 | ngroups=ngroups, |
| 341 | dstate=dstate, |
| 342 | max_cache_len=8, |
| 343 | num_steps=17, |
| 344 | has_z=True, |
| 345 | dt_softplus=True, |
| 346 | use_dt_bias=True, |
| 347 | ) |
| 348 | |
| 349 | |
| 350 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") |
| 351 | @pytest.mark.parametrize( |
| 352 | "precision", |
| 353 | [ |
| 354 | pytest.param((torch.float32, torch.float32), id="s32_a32"), |
| 355 | pytest.param((torch.float32, torch.bfloat16), id="s32_a16"), |
| 356 | pytest.param((torch.float32, torch.float16), id="s32_afp16"), |
| 357 | ], |
| 358 | ) |
| 359 | def test_replayssm_standard_decode_desync_write_pos( |
| 360 | precision: tuple[torch.dtype, torch.dtype], |
| 361 | ): |
| 362 | # Rows start at staggered ring positions, so they flush on different steps |
| 363 | # and hold genuinely different cached histories. |
| 364 | state_dtype, act_dtype = precision |
| 365 | _run_standard_decode( |
| 366 | state_dtype=state_dtype, |
| 367 | act_dtype=act_dtype, |
| 368 | batch=4, |
| 369 | nheads=8, |
| 370 | headdim=64, |
| 371 | ngroups=4, |
| 372 | dstate=64, |
| 373 | max_cache_len=4, |
| 374 | num_steps=12, |
| 375 | has_z=True, |
| 376 | dt_softplus=True, |
| 377 | use_dt_bias=True, |
| 378 | desync_write_pos=True, |
| 379 | ) |
| 380 | |
| 381 | |
| 382 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") |
| 383 | @pytest.mark.parametrize("precision", _PRECISIONS) |
| 384 | @pytest.mark.parametrize("with_padding", [False, True]) |
| 385 | def test_replayssm_standard_decode_with_batch_indices( |
| 386 | precision: tuple[torch.dtype, torch.dtype], |
| 387 | with_padding: bool, |
| 388 | ): |
| 389 | # Sparse state allocation via state_batch_indices, with NULL_BLOCK_ID |
| 390 | # padding rows. The pure-PyTorch references do not model the sparse |
| 391 | # gather, so the anchor here is the upstream baseline decode kernel. |
| 392 | state_dtype, act_dtype = precision |
| 393 | device = "cuda" |
| 394 | both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 |
| 395 | rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16) |
| 396 | set_random_seed(0) |
| 397 | |
| 398 | batch = 3 |
| 399 | padding = 2 if with_padding else 0 |
| 400 | padded_batch = batch + padding |
| 401 | total_state_slots = 16 |
| 402 | nheads = 4 |
| 403 | ngroups = 2 |
| 404 | headdim = 64 |
| 405 | dstate = 16 |
| 406 | max_cache_len = 4 |
| 407 | num_steps = 2 * max_cache_len |
| 408 | |
| 409 | state = torch.randn( |
| 410 | total_state_slots, nheads, headdim, dstate, dtype=state_dtype, device=device |
| 411 | ) |
| 412 | state_baseline = state.clone() |
| 413 | state_cached = state.clone() |
| 414 | state_before = state.clone() |
| 415 | |
| 416 | state_indices = ( |
| 417 | torch.randperm(total_state_slots - 1, device=device)[:batch] + 1 |
| 418 | ).to(torch.int32) |
| 419 | state_batch_indices = torch.cat( |
| 420 | [ |
| 421 | state_indices, |
| 422 | torch.full((padding,), NULL_BLOCK_ID, dtype=torch.int32, device=device), |
| 423 | ] |
| 424 | ) |
| 425 | unused_states = torch.ones(total_state_slots, dtype=torch.bool, device=device) |
| 426 | unused_states[state_indices] = False |
| 427 | |
| 428 | A = _tied_A(nheads, headdim, dstate, device) |
| 429 | dt_bias = _tied_dt_bias(nheads, headdim, device) |
| 430 | D = torch.randn(nheads, headdim, device=device) |
| 431 | x_cache = torch.zeros( |
| 432 | total_state_slots, |
| 433 | nheads, |
| 434 | max_cache_len, |
| 435 | headdim, |
| 436 | device=device, |
| 437 | dtype=act_dtype, |
| 438 | ) |
| 439 | dt_cache = torch.zeros( |
| 440 | total_state_slots, nheads, max_cache_len, device=device, dtype=torch.float32 |
| 441 | ) |
| 442 | B_cache = torch.zeros( |
| 443 | total_state_slots, |
| 444 | ngroups, |
| 445 | max_cache_len, |
| 446 | dstate, |
| 447 | device=device, |
| 448 | dtype=act_dtype, |
| 449 | ) |
| 450 | bc_pre = torch.empty( |
| 451 | padded_batch, ngroups, max_cache_len, device=device, dtype=torch.float32 |
| 452 | ) |
| 453 | write_pos = torch.zeros(padded_batch, dtype=torch.int32, device=device) |
| 454 | |
| 455 | for _ in range(num_steps): |
| 456 | x = torch.randn(padded_batch, nheads, headdim, device=device, dtype=act_dtype) |
| 457 | dt = _tied_dt(padded_batch, nheads, headdim, device, act_dtype) |
| 458 | B = torch.randn(padded_batch, ngroups, dstate, device=device, dtype=act_dtype) |
| 459 | C = torch.randn(padded_batch, ngroups, dstate, device=device, dtype=act_dtype) |
| 460 | z = torch.randn_like(x) |
| 461 | is_flush = write_pos == max_cache_len - 1 |
| 462 | |
| 463 | out_baseline = torch.empty_like(x) |
| 464 | selective_state_update( |
| 465 | state_baseline, |
| 466 | x, |
| 467 | dt, |
| 468 | A, |
| 469 | B, |
| 470 | C, |
| 471 | D=D, |
| 472 | z=z, |
| 473 | dt_bias=dt_bias, |
| 474 | dt_softplus=True, |
| 475 | state_batch_indices=state_batch_indices, |
| 476 | out=out_baseline, |
| 477 | ) |
| 478 | |
| 479 | out_cached = torch.full_like(x, 42) |
| 480 | common = dict( |
| 481 | D=D, |
| 482 | z=z, |
| 483 | dt_bias=dt_bias, |
| 484 | dt_softplus=True, |
| 485 | x_cache=x_cache, |
| 486 | dt_cache=dt_cache, |
| 487 | B_cache=B_cache, |
| 488 | write_pos=write_pos, |
| 489 | is_flush=is_flush, |
| 490 | max_cache_len=max_cache_len, |
| 491 | state_batch_indices=state_batch_indices, |
| 492 | out=out_cached, |
| 493 | ) |
| 494 | selective_state_update_replayssm_output_only( |
| 495 | state_cached, x, dt, A, B, C, bc_pre=bc_pre, **common |
| 496 | ) |
| 497 | |
| 498 | torch.testing.assert_close( |
| 499 | out_cached[:batch], out_baseline[:batch], rtol=rtol, atol=atol |
| 500 | ) |
| 501 | if with_padding: |
| 502 | assert torch.equal( |
| 503 | out_cached[batch:], torch.full_like(out_cached[batch:], 42) |
| 504 | ) |
| 505 | |
| 506 | if bool(is_flush[:batch].all()): |
| 507 | torch.testing.assert_close( |
| 508 | state_cached[state_indices], |
| 509 | state_baseline[state_indices], |
| 510 | rtol=rtol, |
| 511 | atol=atol, |
| 512 | ) |
| 513 | |
| 514 | write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1) |
| 515 | |
| 516 | assert torch.equal(state_cached[unused_states], state_before[unused_states]) |
| 517 | assert torch.equal(state_baseline[unused_states], state_before[unused_states]) |
| 518 | |
| 519 | |
| 520 | # Geometries with (nheads, ngroups) both divisible by the tp below. |
| 521 | _TP_GEOMETRIES = [ |
| 522 | pytest.param((8, 64, 64, 4), id="small"), |
| 523 | pytest.param((96, 80, 128, 8), id="nano4b"), |
| 524 | ] |
| 525 | |
| 526 | |
| 527 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device") |
| 528 | @pytest.mark.parametrize("precision", _PRECISIONS) |
| 529 | @pytest.mark.parametrize("geometry", _TP_GEOMETRIES) |
| 530 | @pytest.mark.parametrize("tp", [2]) |
| 531 | def test_replayssm_standard_decode_tp_head_shard_equivalence( |
| 532 | precision: tuple[torch.dtype, torch.dtype], |
| 533 | geometry: tuple[int, int, int, int], |
| 534 | tp: int, |
| 535 | ): |
| 536 | """Tensor-parallel correctness at the kernel boundary (single GPU). |
| 537 | |
| 538 | The kernel has no cross-rank communication, so head sharding must be exactly |
| 539 | separable: one run over all heads must equal concatenating ``tp`` independent |
| 540 | per-rank runs on ``nheads // tp`` heads and ``ngroups // tp`` groups. This |
| 541 | guards the per-rank divisors and group->head mapping the state-shape wiring |
| 542 | relies on. The real TP1==TP2 engine check lives in the v1/e2e suite.""" |
| 543 | state_dtype, act_dtype = precision |
| 544 | nheads, headdim, dstate, ngroups = geometry |
| 545 | assert nheads % tp == 0 and ngroups % tp == 0 |
| 546 | device = "cuda" |
| 547 | both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32 |
| 548 | rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16) |
| 549 | set_random_seed(0) |
| 550 | |
| 551 | batch = 4 |
| 552 | max_cache_len = 4 |
| 553 | num_steps = 2 * max_cache_len + 1 |
| 554 | nh_s = nheads // tp |
| 555 | ng_s = ngroups // tp |
| 556 | |
| 557 | # Shards slice the tied params; never .contiguous() -- that would drop the |
| 558 | # stride-0 broadcast the kernel's TIE_HDIM asserts require. |
| 559 | A = _tied_A(nheads, headdim, dstate, device) |
| 560 | dt_bias = _tied_dt_bias(nheads, headdim, device) |
| 561 | D = torch.randn(nheads, headdim, device=device) |
| 562 | |
| 563 | state_full = torch.randn( |
| 564 | batch, nheads, headdim, dstate, dtype=state_dtype, device=device |
| 565 | ) |
| 566 | x_cache, dt_cache, B_cache, _ = allocate_update_caches( |
| 567 | batch, |
| 568 | nheads, |
| 569 | ngroups, |
| 570 | headdim, |
| 571 | dstate, |
| 572 | max_cache_len, |
| 573 | device, |
| 574 | act_dtype, |
| 575 | act_dtype, |
| 576 | ) |
| 577 | bc_pre = torch.empty( |
| 578 | batch, ngroups, max_cache_len, device=device, dtype=torch.float32 |
| 579 | ) |
| 580 | |
| 581 | # Per-rank shards, each seeded from the matching head slice so all start equal. |
| 582 | shard_state = [] |
| 583 | shard_caches = [] |
| 584 | for r in range(tp): |
| 585 | h0 = r * nh_s |
| 586 | shard_state.append(state_full[:, h0 : h0 + nh_s].contiguous()) |
| 587 | xc, dtc, bc, _ = allocate_update_caches( |
| 588 | batch, |
| 589 | nh_s, |
| 590 | ng_s, |
| 591 | headdim, |
| 592 | dstate, |
| 593 | max_cache_len, |
| 594 | device, |
| 595 | act_dtype, |
| 596 | act_dtype, |
| 597 | ) |
| 598 | bcp = torch.empty( |
| 599 | batch, ng_s, max_cache_len, device=device, dtype=torch.float32 |
| 600 | ) |
| 601 | shard_caches.append((xc, dtc, bc, bcp)) |
| 602 | |
| 603 | write_pos = torch.zeros(batch, dtype=torch.int32, device=device) |
| 604 | for _ in range(num_steps): |
| 605 | x = torch.randn(batch, nheads, headdim, device=device, dtype=act_dtype) |
| 606 | dt = _tied_dt(batch, nheads, headdim, device, act_dtype) |
| 607 | B = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) |
| 608 | C = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype) |
| 609 | z = torch.randn_like(x) |
| 610 | is_flush = write_pos == max_cache_len - 1 |
| 611 | |
| 612 | out_full = torch.empty_like(x) |
| 613 | selective_state_update_replayssm_output_only( |
| 614 | state_full, |
| 615 | x, |
| 616 | dt, |
| 617 | A, |
| 618 | B, |
| 619 | C, |
| 620 | D=D, |
| 621 | z=z, |
| 622 | dt_bias=dt_bias, |
| 623 | dt_softplus=True, |
| 624 | x_cache=x_cache, |
| 625 | dt_cache=dt_cache, |
| 626 | B_cache=B_cache, |
| 627 | bc_pre=bc_pre, |
| 628 | write_pos=write_pos, |
| 629 | is_flush=is_flush, |
| 630 | max_cache_len=max_cache_len, |
| 631 | out=out_full, |
| 632 | ) |
| 633 | |
| 634 | for r in range(tp): |
| 635 | h0 = r * nh_s |
| 636 | g0 = r * ng_s |
| 637 | xc, dtc, bc, bcp = shard_caches[r] |
| 638 | out_shard = torch.empty( |
| 639 | batch, nh_s, headdim, device=device, dtype=act_dtype |
| 640 | ) |
| 641 | selective_state_update_replayssm_output_only( |
| 642 | shard_state[r], |
| 643 | x[:, h0 : h0 + nh_s].contiguous(), |
| 644 | dt[:, h0 : h0 + nh_s], |
| 645 | A[h0 : h0 + nh_s], |
| 646 | B[:, g0 : g0 + ng_s].contiguous(), |
| 647 | C[:, g0 : g0 + ng_s].contiguous(), |
| 648 | D=D[h0 : h0 + nh_s].contiguous(), |
| 649 | z=z[:, h0 : h0 + nh_s].contiguous(), |
| 650 | dt_bias=dt_bias[h0 : h0 + nh_s], |
| 651 | dt_softplus=True, |
| 652 | x_cache=xc, |
| 653 | dt_cache=dtc, |
| 654 | B_cache=bc, |
| 655 | bc_pre=bcp, |
| 656 | write_pos=write_pos, |
| 657 | is_flush=is_flush, |
| 658 | max_cache_len=max_cache_len, |
| 659 | out=out_shard, |
| 660 | ) |
| 661 | |
| 662 | torch.testing.assert_close( |
| 663 | out_full[:, h0 : h0 + nh_s], out_shard, rtol=rtol, atol=atol |
| 664 | ) |
| 665 | if bool(is_flush.any()): |
| 666 | torch.testing.assert_close( |
| 667 | state_full[:, h0 : h0 + nh_s][is_flush], |
| 668 | shard_state[r][is_flush], |
| 669 | rtol=rtol, |
| 670 | atol=atol, |
| 671 | ) |
| 672 | |
| 673 | write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1) |