hotfix: cross-warp race in checkpointing SSU kernel (mamba) (#3439)

<!-- .github/pull_request_template.md -->

## 📌 Description

Fixes a per-launch non-determinism bug in the CUDA `checkpointing_ssu`
kernel's `must_checkpoint=True` code path, caused by a missing
cross-warp barrier between `load_state_per_warp` (Phase 0) and
`replay_state_mma` (Phase 1).

**Root cause.** `load_state_per_warp` partitions M=DIM across warps —
warp W loads rows `[W*D_PER_CTA/4, (W+1)*D_PER_CTA/4)` of `smem.state`
via cp.async. `replay_state_mma` uses a `Layout<_1, _4>` tiled MMA (1
warp on M, 4 warps on N), so every warp reads the **full M=DIM extent**
of `smem.state` when forming its `frag_h` initial value. The `load_data`
tail used only `__syncwarp()` + `__pipeline_wait_prior(0)`, which
establishes visibility *within* a single warp but not *across* warps.
Result: warp 0's replay reads rows that warps 1/2/3 may not have
committed yet, picking up partial/stale smem and producing different
output every launch.

**Symptom (pre-fix).** Hashing the post-kernel `state`, `state_scale`,
and `out` across 5 launches with bit-identical inputs gave 5 distinct
hashes per config. `state_diff` (max abs delta between launches) up to
~13 in fp16 at `batch=99`, scaling roughly with how much `smem.state`
each warp's MMA touched. Sub-ULP for small batches/heads where the race
rarely fires before the consumer thread arrives — i.e. genuine race
timing, not arithmetic noise.

**Fix.** Hoist a single `__syncthreads()` from inside `ssu_nocheckpoint`
to the dispatch site in `checkpointing_ssu_kernel`, right before the `if
(must_checkpoint)` branch. One barrier now covers cross-warp visibility
for everything both branches consume:
- (a) load_data's per-warp-partitioned `smem.state`
- (b) `smem.x` (warp 2-loaded), `smem.z` (warp 3-loaded)
- (c) `compute_CB_scaled_2warp` writes (warps 0,1)
- (d) `compute_CB_old_2warp` writes (warps 2,3, no-checkpoint path)

Net barrier-count delta is **zero** — the `__syncthreads()` previously
inside `ssu_nocheckpoint` is just relocated, and `ssu_checkpoint` was
missing one. No new compilation flags, no smem layout changes, no
perf-relevant code path touched.

**Scope.** Only the generic kernel (`kernel_checkpointing_ssu.cuh`) is
affected. The 8-bit kernel (`kernel_checkpointing_ssu_8bit.cuh`, used
for `int8`/`fp8_e4m3fn` state) goes through a separate code path and was
already deterministic — verified empirically across `mw ∈ {8, 16}, np ∈
{8, 16}`, `philox ∈ {0, 5}`, all `prev_k` values.

## 🔍 Related Issues

Discovered while investigating intermittent failures in the batch-sweep
parity test added in #3431. The race is in the existing kernel — not
introduced by that PR — but the new sweep stresses it across enough
`(batch, heads_per_group)` configurations to expose it reliably. Landing
this hotfix should let #3431's CUDA-vs-Triton parity test go green.

Adjacent prior art: NVIDIA/TensorRT-LLM#14203 fixes a structurally
similar (but mechanically distinct) bug in the Triton replay kernel —
Triton alias analysis reordering writes ahead of reads on a
single-buffered `old_x`. Our CUDA kernel doesn't have that issue because
cp.async + explicit barriers preserve in-thread ordering; our bug is
purely cross-warp visibility.

## 🚀 Pull Request Checklist

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

## 🧪 Tests

- [x] Tests have been added:
`tests/mamba/test_checkpointing_ssu.py::test_checkpointing_ssu_determinism_across_launches`
— runs the kernel 5× with bit-identical inputs and asserts that the
`state`, `state_scale` (quantized path only), and `out` tensors hash to
a single value across all launches. Four parametrizations:
  - `fp16-no_checkpoint` (`prev_k=4`, no state write)
  - `fp16-checkpoint` (`prev_k=12`, state writeback)
  - `fp8-no_checkpoint`
  - `fp8-checkpoint`
- [x] All tests are passing (`unittest`, etc.). *(Local: full
`tests/mamba/test_checkpointing_ssu.py` — please trigger CI.)*

## Reviewer Notes

- **Why hash-equality instead of `assert_close(rtol, atol)`?** For
*self*-determinism (one kernel against itself, same inputs, same
compiled binary) there should be zero numerical noise — any difference
indicates a real race / uninitialized read / nondeterministic atomic. A
tolerance-based check would silently swallow sub-ULP variance that still
represents a correctness bug (some pre-fix configs had `state_diff ≈
1e-3` — small in magnitude but the symptom of the same race). See the
test docstring for details.
- **Performance impact**: a single extra `__syncthreads()` per CTA on
the dispatch path. Branch divergence is unchanged — `must_checkpoint` is
CTA-uniform (derived from broadcast `prev_k` + compile-time `NPREDICTED`
+ `MAX_WINDOW`).
- **Follow-up opportunity (not in this PR)**: the redundant per-warp
cp.async loads of `old_x` / `old_B` / scalar `old_dt` / `old_cumAdt` in
`load_data` were a defensive measure for the missing-sync regime. With
the dispatch-site barrier in place, these could be partitioned across
all 128 threads to reclaim the 3× redundant cp.async issue cost. Out of
scope for the hotfix.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* Unified GPU kernel synchronization to ensure correct cross-warp
shared-memory visibility, preventing inconsistent outcomes between
checkpointing and non-checkpointing executions.

* **Tests**
* Added a determinism regression test that validates bit-exact,
launch-to-launch consistency across multiple runs, data types, and
checkpointing configurations.

<!-- review_stack_entry_start -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3439?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->