diff --git a/.agents/specs/mamba2-ssd.md b/.agents/specs/mamba2-ssd.md index b33042335..75c18e3b9 100644 --- a/.agents/specs/mamba2-ssd.md +++ b/.agents/specs/mamba2-ssd.md @@ -283,12 +283,23 @@ FAIL, round 2 PASS). `KERNEL-SSM-MAMBA` stays `INVENTORIED`: this is a host reference, not generic Mamba support, and no lifecycle state moved, so it owes no `STATUS.md` / `BENCHMARKS.md` projection. No performance claim is made. -**Owed before the row can move:** W2 (the CUDA arm, byte-compared to these host -references, `compute-sanitizer` clean on dgx), W3 (the `MambaSpec` producer for -Mamba2 shapes), and the one missing decode refusal `SUBCASE` recorded in §8.2. - -**Next action:** dispatch a fresh implementer for W2 (CUDA), and fold the §8.2 -`SUBCASE` into that task since it touches the same suites. +**W2 (2026-08-13):** the CUDA arm for all three ops is implemented and gated — +see §8.3 for the equivalence contract it was written against, which supersedes +§6's "byte-compared to W1" exit criterion with a named reason (the two arms call +different libms, so a byte compare is not reachable; the primary gate is the +device output against the same double-precision reference at the same +upstream-ported tolerances). The §8.2 decode `SUBCASE` is CLOSED and re-proved. +`KERNEL-SSM-MAMBA` still stays `INVENTORIED`: this lands `src/`, `include/`, +`tests/` and this spec only, no lifecycle state moved, and no performance claim +is made. + +**Owed before the row can move:** W3 (the `MambaSpec` producer for Mamba2 +shapes), a fresh scoped review of W2, and the two residuals named in §8.3 — +routing the device-side `A < 0` and `state_indices` precondition checks through +the deferred error ring, and #547 (`ReferenceTierEligible` gating on +`UnifiedMemory()` where it needs `DeviceMemoryIsHostAddressable()`). + +**Next action:** dispatch a fresh scoped review of W2, then W3. ### 8.1 W1 progress (host references landed, awaiting a fresh scoped review) @@ -399,8 +410,22 @@ half of F4. Dropping `CheckMamba2ANegative` at `cpu_ops.cpp:1877` leaves chunk-scan twin at `:1633` reds. The guard itself is present, correct and reachable (a direct probe refuses `A` = `+1.0`, `0.0`, `-0.0` naming `A_log`, and accepts `-1e-30` and `-9.8e-45`), so this is a missing mutation-proof, not -a defect — **owed:** an "A must be negative" `SUBCASE` on the state-update -refusal case mirroring `tests/vt/test_ops_mamba2_ssd.cpp:900`. +a defect — ~~**owed:** an "A must be negative" `SUBCASE` on the state-update +refusal case mirroring `tests/vt/test_ops_mamba2_ssd.cpp:900`.~~ **CLOSED in W2** +(`tests/vt/test_ops_mamba2_state_update.cpp:658`). Re-proved here rather than +taken on report, as mutation **M9** of the W2 sweep: deleting +`CheckMamba2ANegative(A, "mamba2_state_update");` at `cpu_ops.cpp:1877` — the +call site is unique, the other two hits of that symbol being its definition at +`:1582` and the chunk-scan call at `:1633` — takes +`test_ops_mamba2_state_update -tc=mamba2 state update refuses the arms it does +not implement` from `1 passed | 0 failed`, `assertions: 11 | 11 passed`, +`Status: SUCCESS!` to `0 passed | 1 failed`, `assertions: 11 | 8 passed | +3 failed`, `Status: FAILURE!` — the three reds being `CHECK(threw)`, +`CHECK(msg.find("A_log") != npos)` and `CHECK_THROWS(...) did NOT throw at all` +at `:672`, `:673`, `:680`. The pristine binary was run under the identical +filter FIRST, so the filter is proved to select a non-zero assertion count +rather than nothing; `cpu_ops.cpp` was restored byte-for-byte and its md5 +re-asserted at `9ed9eb980c239eca37ec7d92bfe0e766`. One repo-wide test trap found while capturing the RED output, and worth carrying to any doctest suite: **doctest 2.5.2 `INFO` prints a `const char*` VARIABLE as @@ -416,6 +441,266 @@ itself: the Triton dots downcast their tile inputs (`ssd_chunk_state.py:283-285` device-vs-host comparison is a tolerance comparison at the activation dtype, not a byte compare. +### 8.3 W2 — the declared equivalence contract for the CUDA arm + +This section is the spec copy of the contract the W2 implementer decided +**before** writing the kernel and recorded in `src/vt/cuda/cuda_mamba2_ssd.cuh` +and in all three test headers. Its original copy was a staged blob lost with the +worktree; it is re-authored here from the recovery commit `fcdb7d824`, unchanged. +It is a contract, not a tolerance budget: nothing in it was renegotiated to make +a run pass. + +**1. f32 accumulation throughout; the upstream tile downcasts are deliberately +NOT mirrored.** Upstream downcasts its tiles before `tl.dot` — `b.to(x_ptr. +dtype.element_ty)` at `ssd_chunk_state.py:283-285`, `cb.to(...)` and +`prev_states.to(C_ptr.dtype.element_ty)` at `ssd_chunk_scan.py:266-269` and +`:359-363`. Those casts are the **input-precision requirement of `tl.dot`**, i.e. +of a tensor-core MMA, not a statement of the algorithm: every one of those tiles +is loaded with an explicit `.to(tl.float32)` and computed in f32 right up to the +instant it is fed to the MMA. These are scalar-FMA kernels with no MMA, so +mirroring the downcast would copy a constraint we do not have, and would be lossy +for nothing. + +**2. This is not the "too wide" deviation §7 warns about, and the distinction is +checkable.** A token gate cannot catch a dtype that is too wide +([[token-gates-cannot-see-dequant-fallbacks]]), so the claim is made about the +**memory format**, which is byte-for-byte the host arm's: every load and store +goes through the operand's own declared dtype (`M2Load` / `M2Store`); `states` +and `CB` are f32 because upstream pins them there (`states_in_fp32=True`, +`ssd_combined.py:100-102`; `output_dtype=torch.float32`, `:124`); and the +inter-chunk `passed` buffer is allocated at **`state_dtype`**, *not* at the host +reference's f32 working width, which §8.2 F9 explicitly flagged as a width W2 +must not inherit. No extra byte moves. Only the register precision of one product +differs, and it differs in the direction Triton itself takes wherever it is not +feeding an MMA. + +**3. Accumulation ORDER is part of the port.** Except in the gated norm's group +reduction — which is a block reduction, and says so at the kernel — every +accumulation runs in ONE thread, over the SAME index range in the SAME direction +as the host reference. That is deliberate: it leaves the elementary functions as +the *only* admitted source of divergence, so the derived bound has exactly one +term to account for. + +**4. A byte compare against the host arm is NOT reachable, and the downcasts are +not why.** The two arms call different libms — CUDA `expf` is documented at +≤ 2 ulp, glibc's at ≤ 0.5 — and the gated norm additionally reorders one +non-negative reduction. Everything else is identical by construction. §9's third +stop condition ("the device arm cannot reach the host reference byte-for-byte") +is therefore resolved as **not reachable for a named, non-defect reason**, and +the gap is kept open in the form below rather than closed by widening anything. + +**5. The primary gate is therefore NOT device-vs-host.** It is the device output +against the **same double-precision sequential reference** the host arm is held +to, at the **same upstream-ported tolerances**, on the **same inputs** — e.g. atol +8e-3 / rtol 5e-3 from `test_mamba_ssm_ssd.py:210-213` on the driver shapes. Both +arms are asserted against it in the same test case, so a failure separates +cleanly: device-only means a device defect; both means the cited upstream +threshold does not cover this shape, which is a `NEEDS_DECISION`, not a wider +tolerance. + +**6. The derived device-vs-host bar is `rtol(K) = 4·(K+2)·2⁻²⁴`** +(`ExpectDeviceMatchesHost` / `DerivedRtol`, `tests/vt/test_ops_mamba2_ssd.cpp`), +from 2.5 ulp of libm disagreement per decay factor through a product of at most +`K`, plus `(K-1)·u` of summation error. **No number was tuned and no tolerance +was widened**; `K` is the sequence length the comparison actually ran at. +Because a bar nobody audits is a false claim, every comparison logs the +**fraction of the budget actually used** through `MESSAGE` — not `INFO`, because +doctest prints `INFO` only on failure, so an `INFO` would have been invisible on +the green run that the claim rests on. + +**Two deviations recorded with the arm, carried forward deliberately:** + +- **Placement.** The kernels live in `src/vt/cuda/cuda_mamba2_ssd.cuh`, included + by `cuda_gdn.cu`, not in a new `.cu`. Same reason as W1's `cpu_ops.cpp` + placement (§8.1): a new library TU must be listed in the root `CMakeLists.txt`, + which `check-doc-checkpoint` classifies `user_usage` + `landing_page`, so any + new `src/vt/` file owes a `docs/USAGE.md` update that a kernel exposing no + command, config key or C-ABI entry point has nothing true to write (#515). +- **The device arm does not re-check `A < 0` or `state_indices` distinctness.** + Both operands are on-device; re-reading them costs a D2H plus a stream + synchronise per call — the same host tax the GDN prefill path was rebuilt to + remove — and makes the op uncapturable in a CUDA graph. This mirrors the policy + `cuda_gdn.cu:8-13` already states for exactly this case. The kernels stay + **memory safe** under a violation: an out-of-range `state_indices` slot writes + nothing at all rather than out of bounds. **Owed, not implemented here:** route + both through the deferred device error ring at `cuda_ops.cu:790-940`. + +**Not fixed here, filed as #547.** The W2 RED run SIGSEGV'd on all three +binaries. GB10 reports `Backend::UnifiedMemory() == true`, so +`ReferenceTierEligible(kCUDA)` is true and, with no native kernel registered, +`GetOp` installs the CPU host kernel as a `vt-cpu-ref` provider over `cudaMalloc` +pointers — which `include/vt/backend.h` already says are not host-dereferenceable +on GB10. `op_provider.cpp:515-526` gates on `UnifiedMemory()` where it needs +`DeviceMemoryIsHostAddressable()`. That is shared-seam semantics across three +backends, so it takes its own row. Every CUDA case here calls +`RequireNativeCudaProvider`, so a device arm can never be gated by running the +host arm twice. + +### 8.4 W2 evidence (gate host `promaxgb10-4ad8`, GB10 / sm_121a, 2026-08-13) + +Build recipe, both arms: `cmake -G Ninja -DVLLM_CPP_CUDA=ON +-DVLLM_CPP_CUDA_ARCHITECTURES=121a -DVLLM_CPP_CUTLASS_DIR=$HOME/cutlass-4.5.0 +-DVLLM_CPP_TRITON=ON -DVLLM_CPP_BUILD_TESTS=ON`. The configure log was READ, not +assumed: `cutlass-nvfp4: ENABLED`, `cutlass-fp8: ENABLED`, `marlin-nvfp4: +ENABLED`, `fa2: ENABLED for [121a]`, `CUTLASS found at ~/cutlass-4.5.0` — an +absent CUTLASS silently falls back and the arm would not be the shipped one. +**0 warnings** in both build logs. + +**Release, after the `origin/main` re-merge** — identical to the pre-merge +counts, so the merge moved nothing: + +| suite | test cases | assertions | status | +|---|---|---|---| +| `test_ops_mamba2_ssd` | 11 / 11 passed | 2069 / 2069 | `SUCCESS!` | +| `test_ops_mamba2_state_update` | 10 / 10 passed | 5965 / 5965 | `SUCCESS!` | +| `test_ops_mamba2_gated_norm` | 12 / 12 passed | 3723 / 3723 | `SUCCESS!` | + +**Debug arm** (`CMAKE_BUILD_TYPE=Debug`, so `NDEBUG` is OFF and every `assert` +in the tree is live; CXX `-g -O0`, CUDA `-g` and deliberately *not* `-G`, which +would disable device optimisation and change what was measured): the same +`11 / 2069`, `10 / 5965`, `12 / 3723`, all `SUCCESS!`, exit 0. This arm exists +because the gate build is `-O3 -DNDEBUG`, where an assert-abort defect stays +latent behind a green Release run. + +**`compute-sanitizer`, 8 runs, all `ERROR SUMMARY: 0 errors` / `EXIT=0`:** +`memcheck` on the ssd optional-arms + dtype-knobs case, on the continuous-batch +`initial_states` subcase, and on both decode suites' `*CUDA arm*` cases; +`initcheck` on three; `synccheck` on the gated norm, which is the one kernel +with a block reduction and `__syncthreads`. + +**Mutation sweep — 9 of 9 CAUGHT.** Each mutation patches one source file, is +rebuilt, and is run under a doctest `-tc` filter; the **pristine** binary is run +under the *identical* filter first, because a filter that selects no test case +makes doctest print `SUCCESS!` and an unverified filter would score a false +catch. Sources restored byte-for-byte after each, md5 re-asserted +(`cuda_mamba2_ssd.cuh` `cbb1f928f4b421bdea2e24476012eed2`, `cpu_ops.cpp` +`9ed9eb980c239eca37ec7d92bfe0e766`). + +| # | mutation | control assertions | mutant | +|---|---|---|---| +| M1 | drop the inter-chunk state term | 27 `SUCCESS!` | `FAILURE!` | +| M2 | ignore `initial_states` in state passing | 297 `SUCCESS!` | `FAILURE!` | +| M3 | read `states[c]` for `states[c-1]` | 297 `SUCCESS!` | `FAILURE!` | +| M4 | drop the `D` skip connection | 570 `SUCCESS!` | `FAILURE!` | +| M5 | ignore `state_indices` (slot = row) | 1318 `SUCCESS!` | `FAILURE!` | +| M6 | treat the NULL row as slot 0 | 1318 `SUCCESS!` | `FAILURE!` | +| M7 | whole-row variance instead of per-group | 9 `SUCCESS!` | `FAILURE!` | +| M8 | sigmoid instead of silu | 9 `SUCCESS!` | `FAILURE!` | +| M9 | drop `CheckMamba2ANegative` on decode (§8.2) | 11 `SUCCESS!` | `FAILURE!` | + +**Two mutations had to be REFORMULATED, and that is worth carrying.** The +obvious form of M1 (`if (!prev_zero)` → `if (false)`) and of M7 (passing +`1, hidden` for `n_groups, group_size`) do not COMPILE: the CUDA arm is built +`-Werror=all-warnings`, and nvcc raises `#550-D "prev_zero was set but never +used"` and `#177-D "group_size was declared but never referenced"` once the +mutation dead-codes the read. A mutation that will not build is not a caught +mutation and must not be scored as one. Both were rewritten to drop exactly the +same term while leaving every variable read — M1 multiplies the inter-chunk +product by `0.0f`, M7 passes `1, group_size * args.n_groups` (which *is* +`hidden`) — and both then failed as intended. + +**Full `ctest` on the gate host**, all 392 test targets built (777 ninja edges, +0 warnings), `ctest -j 1` — serial is required, not cautious: GB10 memory is +UNIFIED, so a parallel CUDA suite reserves HOST RAM and has OOM-rebooted this +box, and several suites starve under `-j` and red spuriously. Result: +**`98% tests passed, 10 tests failed out of 431`**, `CTEST_EXIT=8`, 53 min. + +None of the three mamba2 suites is among the failures. Nine of the ten match, by +name, an **independent same-box baseline** — another agent's full `ctest` on +`row/pool-device-key` finishing 80 minutes earlier, `98% tests passed, 9 tests +failed out of 437`: `test_serve_low_tools`, `test_linear_method`, +`test_glm4_moe_lite_paged_engine`, `test_capi` (SEGFAULT), `test_ops_gdn`, +`test_qwen3_apc_e2e`, `test_minicpm3_paged_engine`, `test_internlm2_paged_engine`, +`test_llama_paged_engine`. Two branches, two builds, the same nine. + +The tenth, **`test_minimax_h3` (SEGFAULT at 11.81 s)**, passed on that baseline +and was the one difference, so it was NOT dismissed. It is now **fully +attributed, and it is not this brick**. + +It was RE-RUN STANDALONE, serially, under the lock, on a box that had just +rebooted and was idle (`up 5 min`, load 0.12, no CUDA process resident) — so the +contention hypothesis this section previously recorded as "plausible, not +proven" is **REFUTED**: it reproduces with nobody else on the GPU. It is instead +**#486**, already filed and open: *"test_minimax_h3 is RED on dgx (GB10): +cudaFree invalid argument + SIGSEGV when two CUDA cases run in one process"*. +The standalone re-run reproduces that issue's recorded signature exactly, number +for number: + +- `minimax_h3: the WHOLE t2va path composes end to end` throws + `vt cuda: cudaFree: invalid argument` (`test_minimax_h3.cpp:3537`); +- `minimax_h3: an NVFP4 checkpoint loads into a runnable DiT` then SIGSEGVs + (`:3977`); +- `test cases: 38 | 36 passed | 2 failed | 41 skipped`, `assertions: 42724 | + 42724 passed | 0 failed` — the identical counts #486 records. + +#486 further records that each case passes ALONE and only crashes when the two +run in one process, i.e. cross-test CUDA state, and that an A/B in another +agent's tree already proved it was not THEIR change either. Its root cause is +tracked as **#516** — *"vllm::Pool() free list is keyed by size class with no +DEVICE in the key: a cudaMalloc'd block can be handed to a CPU DBuf"* — which is +exactly what `row/pool-device-key` is repairing, and therefore exactly why the +baseline branch passed a test that main-based branches fail. The difference was +the baseline carrying a FIX, not this branch carrying a defect. + +Two independent lines already pointed the same way and are kept because they +remain true: **no model or layer code calls these ops at all** — `grep` for +`vt::Mamba2ChunkScan`, `vt::Mamba2StateUpdate` and `vt::RmsNormGatedGroup` +outside `src/vt/` returns `include/vt/ops.h` declarations and the three unit +tests, nothing else — so the H3 path cannot reach a kernel this brick added; and +the only W2 delta it could see is three extra registrations in the op table. + +So **all ten ctest failures are pre-existing and tracked** (#486 / #516 for +`test_minimax_h3`, #233 for `test_glm4_moe_lite_paged_engine` "plus 4 more +pre-existing ctest failures", and the same-name baseline for the rest). None is +attributable to W2. The standalone log is `~/w2ssd-evidence/refail.log`. + +**All ten reproduce STANDALONE** on the idle box, one `ctest -R` per binary under +the lock — `CTEST_EXIT=8` for every one of `test_minimax_h3`, `test_ops_gdn`, +`test_capi`, `test_serve_low_tools`, `test_linear_method`, +`test_glm4_moe_lite_paged_engine`, `test_qwen3_apc_e2e`, +`test_minicpm3_paged_engine`, `test_internlm2_paged_engine`, +`test_llama_paged_engine`, `W2_REFAIL_DONE`. That matters in both directions: it +rules out "the 431-test run starved them", so none of the ten is a contention +flake to be waved through, and it confirms they are deterministic reds that +exist independently of this branch — which is what makes the same-name baseline +comparison sound rather than a coincidence of two noisy runs. + +**A live protocol defect found while doing this, and fixed.** The re-run had +been relaunched (not by this session) as +`flock -w 3600 $HOME/gpu.lock ./w2refail.sh` — an OUTER `flock` wrapping a +script that takes the SAME lock itself on its own fd. `flock` locks an open file +DESCRIPTION, so the inner acquisition blocks against the outer one held by its +own parent: a self-deadlock that **held the shared GPU mutex while making no +progress**, with three other agents' jobs queued behind it. Killed the stack, the +lock passed straight to a waiter, and the re-run was relaunched with its own +single acquisition. Anything wrapping a script that already locks must not lock +again. + +**The gate host rebooted TWICE during this window** (08:57 and ~09:30 CEST, +`up 14 min` then `up 5 min`), which is the documented GB10 unified-memory +OOM-reboot under multi-agent load, and it killed two queued attempts before the +third landed. Recorded because it is the environment every measurement on this +box is taken in, not as an excuse for any result above. + +**CI: the two Windows jobs are the `main` BASELINE, not this branch.** +`windows-msvc-cpu` and `windows-msvc-vulkan` fail on PR #566, and the baseline +was subtracted rather than assumed: **#580** (`row/ENG-ISSUE-TABLE-INTAKE`, a +records-only change) fails both at 22m15s / 22m6s, and **#576** — already +*merged* to `main` — failed both at 22m6s / 21m33s. Two fixes are in flight for +exactly these: **#583** "Invoke-Checked rejected the empty argument list every +no-arg test uses" (#512), which is the failing step here — *Build and execute +the native Windows CPU focused gate* — and **#578** "the Windows arm cannot +compile `test_backend_cross_device`" (#514, #540). This branch's diff since the +merge is one markdown file and cannot reach an MSVC build; no Windows-arm claim +is made or repaired here. + +**The derived bar is audited, not asserted.** Across the 55 device-vs-host +comparisons in a green run, the worst one used **7.66%** of `rtol(K) = +4·(K+2)·2⁻²⁴`; the driver shapes used 0.32% and 0.18%. For contrast the same +`MESSAGE` line under mutant M3 reads `used 962173% of its derived budget`. So +the bound is neither tuned down to the observed error nor wide enough to hide a +defect. + ## 9. Stop conditions - The chunked scan cannot be made to match the sequential double-precision diff --git a/src/vt/cuda/cuda_gdn.cu b/src/vt/cuda/cuda_gdn.cu index 0ad4eff66..5dcfb22e8 100644 --- a/src/vt/cuda/cuda_gdn.cu +++ b/src/vt/cuda/cuda_gdn.cu @@ -35,6 +35,17 @@ #include "vt/cuda/conv_update_fast.h" #include "vt/cuda/cuda_device_caps.h" #include "vt/cuda/cuda_gdn_internal.h" +// MAMBA2 / SSD IS NOT THE GATED DELTA RULE (.agents/specs/mamba2-ssd.md §0, §7): +// no delta-removal term, decay driven by A_log/dt, B/C shared across n_groups. +// The three kernels live in their own header and their own namespace and share +// nothing with the GDN/KDA code below; they are compiled into THIS translation +// unit because it is where the sibling SSM/linear-attention device arms already +// live, and because a new .cu has to be listed in the ROOT CMakeLists.txt, which +// check-doc-checkpoint classifies as `user_usage` and therefore charges a +// docs/USAGE.md update that a kernel exposing no command, config key or C-ABI +// entry point has nothing true to write (the same deviation W1 recorded for +// cpu_ops.cpp, mamba2-ssd.md §8.1). +#include "vt/cuda/cuda_mamba2_ssd.cuh" #include "vt/cuda/gdn_decode_fused.h" #include "vt/cuda/gdn_packed_decode_triton.h" #include "vt/cuda/gdn_prefill_conv.h" @@ -6641,6 +6652,17 @@ struct Registrar { RegisterOp(OpId::kIndexCopy, DeviceType::kCUDA, reinterpret_cast( static_cast(&IndexCopyKernelCuda))); + // Mamba2 / SSD device arm (mamba2-ssd.md W2, #496) — sibling ops, never a + // parameterisation of the GDN kernels above. + RegisterOp(OpId::kMamba2ChunkScan, DeviceType::kCUDA, + reinterpret_cast( + static_cast(&mamba2::Mamba2ChunkScanKernelCuda))); + RegisterOp(OpId::kMamba2StateUpdate, DeviceType::kCUDA, + reinterpret_cast( + static_cast(&mamba2::Mamba2StateUpdateKernelCuda))); + RegisterOp(OpId::kRmsNormGatedGroup, DeviceType::kCUDA, + reinterpret_cast( + static_cast(&mamba2::RmsNormGatedGroupKernelCuda))); } } registrar; diff --git a/src/vt/cuda/cuda_mamba2_ssd.cuh b/src/vt/cuda/cuda_mamba2_ssd.cuh new file mode 100644 index 000000000..b23586d7c --- /dev/null +++ b/src/vt/cuda/cuda_mamba2_ssd.cuh @@ -0,0 +1,578 @@ +// CUDA device arm of the Mamba2 / SSD selective-scan core. +// .agents/specs/mamba2-ssd.md W2, issue #496. +// +// Three kernels, each the DEVICE transcription of the CPU host reference that +// landed in W1 (src/vt/cpu/cpu_ops.cpp `Mamba2ChunkScanKernel`, +// `Mamba2StateUpdateKernel`, `RmsNormGatedGroupKernel`), which is itself the 1:1 +// port of the upstream path named on it at the pinned oracle `555967922` +// (vLLM 0.26.0.dev0): +// +// Mamba2ChunkScan <- ops/ssd_combined.py:27-235 (the 5-stage varlen pipeline) +// over ssd_chunk_state.py, ssd_state_passing.py, +// ssd_bmm.py and ssd_chunk_scan.py +// Mamba2StateUpdate <- ops/mamba_ssm.py:497+ `selective_state_update` +// RmsNormGatedGroup <- mamba_mixer2.py:100-149 `Mixer2RMSNormGated.forward_native` +// +// ─── THE DECLARED EQUIVALENCE CONTRACT (mamba2-ssd.md §8.3) ─────────────────── +// +// This arm keeps **f32 accumulation throughout** and does NOT mirror the tile +// downcasts in upstream's Triton kernels — `b = b.to(x_ptr.dtype.element_ty)` +// before `tl.dot` (ssd_chunk_state.py:283-285), and `cb.to(x_ptr.dtype.element_ty)` +// / `prev_states.to(C_ptr.dtype.element_ty)` (ssd_chunk_scan.py:266-269, +// :359-363). Those casts are the INPUT-PRECISION REQUIREMENT OF `tl.dot`, i.e. of +// a tensor-core MMA, not a statement of the algorithm: every one of those tiles is +// loaded with an explicit `.to(tl.float32)` and computed in f32 right up to the +// instant it is fed to the MMA. These are scalar-FMA kernels with no MMA, where +// the downcast would be lossy for nothing. +// +// That is NOT a "wider dtype" deviation, and the distinction matters because a +// token gate cannot catch a dtype that is too wide (.agents/porting.md): the +// MEMORY FORMAT here is byte-for-byte the host arm's. Every load and store goes +// through the operand's own declared dtype; `states` and `CB` are f32 because +// upstream pins them there (`states_in_fp32=True` ssd_combined.py:100-102, +// `output_dtype=torch.float32` :124); the inter-chunk `passed` buffer is allocated +// at `state_dtype` — NOT at the host reference's f32 working width, which W1 +// explicitly flagged as a width W2 must not inherit (cpu_ops.cpp, §8.2 F9). +// No extra byte moves; only the register precision of one product differs, and it +// differs in the direction Triton itself takes wherever it is not feeding an MMA. +// +// The consequence for the gate is stated in the test files: the two arms are NOT +// bit-identical, because they call different libms (`expf`/`log1pf`), and the +// gated norm additionally reorders one non-negative reduction. Both effects carry +// a DERIVED forward-error bound; neither is a tuned number. See +// tests/vt/test_ops_mamba2_ssd.cpp `DerivedRtol`. +// +// ─── ACCUMULATION ORDER IS PART OF THE PORT ────────────────────────────────── +// Except in the gated norm's group reduction (which is a block reduction, and +// says so), every accumulation below runs in ONE thread, over the SAME index +// range in the SAME direction as the host reference. That is deliberate: it +// leaves the elementary functions as the ONLY admitted source of divergence, so +// the derived bound has exactly one term to account for. +// +// ─── WHAT THIS ARM DOES NOT CHECK (named residual, §8.3) ───────────────────── +// The host arm validates two data-dependent preconditions by READING THE TENSORS: +// `A < 0` (`CheckMamba2ANegative`, cpu_ops.cpp) and the distinctness of +// `state_indices`. Both operands live on the DEVICE here, so re-checking them +// would cost a D2H copy plus a stream synchronise on every call — the same host +// tax the GDN prefill path was rebuilt to remove (`GdnArgs::query_start_loc_host`, +// include/vt/ops.h), and it would make the op uncapturable in a CUDA graph. This +// arm therefore mirrors the policy cuda_gdn.cu already states for exactly this +// case ("here bad metadata is unchecked -- correctness-grade; the M0.9 builder +// owns metadata integrity", cuda_gdn.cu:8-13). The device kernels remain MEMORY +// SAFE under a violation: an out-of-range `state_indices` slot writes nothing at +// all rather than out of bounds. Closing the gap needs the deferred device error +// ring cuda_ops.cu:790-940 already implements for embedding, and is recorded as +// owed rather than silently dropped. +#ifndef VT_CUDA_MAMBA2_SSD_CUH_ +#define VT_CUDA_MAMBA2_SSD_CUH_ + +#include +#include +#include + +#include +#include +#include + +#include "vt/ops.h" + +namespace vt::cuda::mamba2 { +namespace { + +constexpr int kM2Block = 256; + +void M2Check(cudaError_t err, const char* what) { + if (err != cudaSuccess) { + throw std::runtime_error(std::string("vt cuda mamba2: ") + what + ": " + + cudaGetErrorString(err)); + } +} + +cudaStream_t M2Stream(const Queue& q) { return static_cast(q.handle); } + +// Grid for a grid-stride loop over `n` items at kM2Block threads, capped so a +// launch stays reasonable on any element count. +unsigned M2Grid(int64_t n) { + const int64_t blocks = (n + kM2Block - 1) / kM2Block; + if (blocks < 1) return 1u; + return static_cast(blocks < 65535 ? blocks : 65535); +} + +// f32 load/store through the operand's OWN dtype — the memory format is the host +// arm's, only the arithmetic is f32 (`LoadF32`/`StoreF32`, cpu_ops.cpp). +__device__ inline float M2Load(const void* p, DType dt, int64_t i) { + if (dt == DType::kF32) return static_cast(p)[i]; + if (dt == DType::kF16) return __half2float(static_cast(p)[i]); + return __bfloat162float(static_cast(p)[i]); +} + +__device__ inline void M2Store(void* p, DType dt, int64_t i, float v) { + if (dt == DType::kF32) { + static_cast(p)[i] = v; + } else if (dt == DType::kF16) { + static_cast<__half*>(p)[i] = __float2half_rn(v); + } else { + static_cast<__nv_bfloat16*>(p)[i] = __float2bfloat16(v); // RNE, as host F32ToBF16 + } +} + +// The value `v` as it reads back after a store/load round trip through `dt` — +// the `state_dtype` / `input_dtype` cast points (`RoundThrough`, cpu_ops.cpp). +__device__ inline float M2RoundThrough(DType dt, float v) { + if (dt == DType::kF32) return v; + if (dt == DType::kF16) return __half2float(__float2half_rn(v)); + return __bfloat162float(__float2bfloat16(v)); +} + +// softplus, guarded exactly as upstream: `tl.where(dt <= 20.0, softplus(dt), dt)` +// (ssd_chunk_state.py:94; the same guard at csrc/cpu/mamba_kernels.hpp:177). +__device__ inline float M2Softplus(float v) { return v <= 20.0f ? log1pf(expf(v)) : v; } + +__device__ inline float M2Silu(float z) { return z / (1.0f + expf(-z)); } + +// ───────────────────────────────────────────────────────────────────────────── +// stage 1 — `_chunk_cumsum_fwd` (ssd_chunk_state.py:300-346). +// One thread per (h, c): the prefix sum over the chunk is SEQUENTIAL, in the host +// reference's order. Positions past a partial chunk hold dt = 0 (:104-107), so +// dA_cumsum[..., cs-1] is the chunk's TOTAL decay whatever its length. +// ───────────────────────────────────────────────────────────────────────────── +__global__ void M2CumsumKernel(float* dtv, float* dac, const void* dt_in, DType dt_dtype, + const float* A, const float* dbp, const int32_t* ccs, int64_t H, + int64_t nchunks, int64_t cs, bool softplus, float dt_min, + float dt_max) { + const int64_t total = H * nchunks; + for (int64_t r = blockIdx.x * blockDim.x + threadIdx.x; r < total; + r += static_cast(gridDim.x) * blockDim.x) { + const int64_t h = r / nchunks, c = r % nchunks; + const float a = A[h]; + const int64_t start = ccs[c], len = ccs[c + 1] - start; + const int64_t base = r * cs; + float acc = 0.0f; + for (int64_t i = 0; i < cs; ++i) { + float d = 0.0f; + if (i < len) { + d = M2Load(dt_in, dt_dtype, (start + i) * H + h); + if (dbp != nullptr) d += dbp[h]; + if (softplus) d = M2Softplus(d); + d = fminf(fmaxf(d, dt_min), dt_max); + } + dtv[base + i] = d; + acc += d * a; + dac[base + i] = acc; + } + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// stage 2 — `_chunk_state_fwd` (ssd_chunk_state.py:349-407). +// states[c,h,p,n] = sum_i x[i,h,p] * (B[i,g,n] * exp(min(dA_last - dA_i, 0)) * dt_i) +// f32 by upstream's own `states_in_fp32=True` (ssd_combined.py:100-102). One +// thread per (c,h,p,n), accumulating over i in the host reference's order. +// ───────────────────────────────────────────────────────────────────────────── +__global__ void M2ChunkStateKernel(float* states, const void* x, DType xdt, const void* B, + DType Bdt, const float* dtv, const float* dac, + const int32_t* ccs, int64_t nchunks, int64_t H, int64_t P, + int64_t G, int64_t N, int64_t cs, int64_t hpg) { + const int64_t total = nchunks * H * P * N; + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; + idx += static_cast(gridDim.x) * blockDim.x) { + const int64_t n = idx % N; + const int64_t p = (idx / N) % P; + const int64_t h = (idx / (N * P)) % H; + const int64_t c = idx / (N * P * H); + const int64_t g = h / hpg; + const int64_t start = ccs[c], len = ccs[c + 1] - start; + const int64_t dbase = (h * nchunks + c) * cs; + const float da_last = dac[dbase + cs - 1]; + float acc = 0.0f; + for (int64_t i = 0; i < len; ++i) { + // The `min(., 0)` is upstream's and is an algebraic no-op inside the + // enforced contract (`A < 0`, `dt >= 0` make dA_cumsum non-increasing over + // i); it is kept because upstream keeps it. + const float scale = expf(fminf(da_last - dac[dbase + i], 0.0f)) * dtv[dbase + i]; + if (scale == 0.0f) continue; + const float xv = M2Load(x, xdt, ((start + i) * H + h) * P + p); + const float bv = M2Load(B, Bdt, ((start + i) * G + g) * N + n) * scale; + acc += xv * bv; + } + states[((c * H + h) * P + p) * N + n] = acc; + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// stage 3 — `_state_passing_fwd` (ssd_state_passing.py:99-146). +// S_c = exp(dA_last[c]) * S_{c-1} + states[c], S_{-1} = initial_states[b] +// out[c] is the state AFTER chunk c (:90-97). +// +// THE RUNNING STATE STAYS F32 AND ONLY THE STORE ROUNDS: upstream carries +// `states` in f32 registers across the chunk loop and stores a `state_dtype` copy +// per chunk (:88-97) — it never reads that store back into the recurrence. +// `passed` is allocated at `state_dtype` here, NOT at the host reference's f32 +// working width (cpu_ops.cpp records that width as one W2 must not inherit). +// One thread per (b,h,p,n), sequential over chunks. +// ───────────────────────────────────────────────────────────────────────────── +__global__ void M2StatePassKernel(void* passed, DType sdt, void* final_states, const float* states, + const float* dac, const int32_t* lci, const void* init, + DType initdt, int64_t S, int64_t H, int64_t P, int64_t N, + int64_t nchunks, int64_t cs) { + const int64_t row = P * N; + const int64_t total = S * H * row; + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; + idx += static_cast(gridDim.x) * blockDim.x) { + const int64_t i = idx % row; + const int64_t h = (idx / row) % H; + const int64_t b = idx / (row * H); + const int64_t chunk_end = lci[b] + 1; + const int64_t chunk_start = b > 0 ? lci[b - 1] + 1 : 0; + float s = init != nullptr ? M2Load(init, initdt, (b * H + h) * row + i) : 0.0f; + for (int64_t c = chunk_start; c < chunk_end; ++c) { + const float decay = expf(dac[(h * nchunks + c) * cs + cs - 1]); + s = s * decay + states[(c * H + h) * row + i]; + M2Store(passed, sdt, (c * H + h) * row + i, s); + } + // `varlen_states = states[last_chunk_indices]` (ssd_combined.py:154). + M2Store(final_states, sdt, (b * H + h) * row + i, s); + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// stage 4 — `_bmm_chunk_fwd` (ssd_bmm.py:148-209). +// CB[c,g,i,j] = sum_n C[i,g,n] * B[j,g,n], f32 REGARDLESS of the activation dtype +// (`output_dtype=torch.float32`, ssd_combined.py:124). Only j <= i is ever read +// (IS_CAUSAL), so only j <= i is written — as in the host reference. +// ───────────────────────────────────────────────────────────────────────────── +__global__ void M2BmmKernel(float* cb, const void* Bp, DType Bdt, const void* Cp, DType Cdt, + const int32_t* ccs, int64_t nchunks, int64_t G, int64_t N, + int64_t cs) { + const int64_t total = nchunks * G * cs * cs; + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; + idx += static_cast(gridDim.x) * blockDim.x) { + const int64_t j = idx % cs; + const int64_t i = (idx / cs) % cs; + if (j > i) continue; + const int64_t g = (idx / (cs * cs)) % G; + const int64_t c = idx / (cs * cs * G); + const int64_t start = ccs[c], len = ccs[c + 1] - start; + if (i >= len) continue; + float acc = 0.0f; + for (int64_t n = 0; n < N; ++n) { + acc += M2Load(Cp, Cdt, ((start + i) * G + g) * N + n) * + M2Load(Bp, Bdt, ((start + j) * G + g) * N + n); + } + cb[((c * G + g) * cs + i) * cs + j] = acc; + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// stage 5 — `_chunk_scan_fwd` (ssd_chunk_scan.py:216-525). +// out_i = exp(dA_i) * (C_i . S_{c-1}) inter +// + sum_{j<=i} CB[i,j] * exp(min(dA_i - dA_j, 0)) * dt_j * x_j intra +// + D * x_i skip +// then `out *= z * sigmoid(z)` when z is given (:394-406). +// `S_{c-1}` is `initial_states[seq_idx[c]]` when this chunk opens a new sequence +// AND initial states were supplied, ZEROS when they were not (:236-250, :271-289) +// — that is what makes a sequence boundary INSIDE a physical chunk correct. +// One thread per (c,h,i,p); both accumulations run in the host arm's order. +// ───────────────────────────────────────────────────────────────────────────── +__global__ void M2ChunkScanKernel(void* out, DType odt, const void* x, DType xdt, const void* Cp, + DType Cdt, const float* cb, const float* dtv, const float* dac, + const void* passed, DType sdt, const void* init, DType initdt, + const float* D, bool d_has_hdim, const void* z, DType zdt, + const int32_t* ccs, const int32_t* sidx, int64_t nchunks, + int64_t H, int64_t P, int64_t G, int64_t N, int64_t cs, + int64_t hpg) { + const int64_t total = nchunks * H * cs * P; + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; + idx += static_cast(gridDim.x) * blockDim.x) { + const int64_t p = idx % P; + const int64_t i = (idx / P) % cs; + const int64_t h = (idx / (P * cs)) % H; + const int64_t c = idx / (P * cs * H); + const int64_t start = ccs[c], len = ccs[c + 1] - start; + if (i >= len) continue; + const int64_t g = h / hpg; + const int64_t dbase = (h * nchunks + c) * cs; + const int64_t row = P * N; + + const int32_t si = sidx[c]; + const int32_t si_prev = c >= 1 ? sidx[c - 1] : -1; + bool prev_zero = false; + const void* prevp = passed; + DType prevdt = sdt; + int64_t prevbase = ((c - 1) * H + h) * row; + if (si != si_prev) { + if (init != nullptr) { + prevp = init; + prevdt = initdt; + prevbase = (static_cast(si) * H + h) * row; + } else { + prev_zero = true; + } + } + + const float da_i = dac[dbase + i]; + const float scale_m = expf(da_i); + float acc = 0.0f; + if (!prev_zero) { + for (int64_t n = 0; n < N; ++n) { + acc += M2Load(Cp, Cdt, ((start + i) * G + g) * N + n) * + M2Load(prevp, prevdt, prevbase + p * N + n); + } + } + acc *= scale_m; + const float* cbc = cb + (c * G + g) * cs * cs; + for (int64_t j = 0; j <= i; ++j) { + const float w = cbc[i * cs + j] * expf(fminf(da_i - dac[dbase + j], 0.0f)) * dtv[dbase + j]; + acc += w * M2Load(x, xdt, ((start + j) * H + h) * P + p); + } + const float xi = M2Load(x, xdt, ((start + i) * H + h) * P + p); + if (D != nullptr) acc += (d_has_hdim ? D[h * P + p] : D[h]) * xi; + if (z != nullptr) acc *= M2Silu(M2Load(z, zdt, ((start + i) * H + h) * P + p)); + M2Store(out, odt, ((start + i) * H + h) * P + p, acc); + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// vt::Mamba2StateUpdate — `selective_state_update` (ops/mamba_ssm.py:497+) at the +// scalar-per-head shape (csrc/cpu/mamba_kernels.hpp:104-250). +// One thread per (b,h,p), sequential over n exactly as the host arm. +// +// The readout uses the F32 value, not the value re-read from the cache: the +// Triton kernel holds `state` in registers and computes `out = sum(state * C)` +// from them, storing the cache-width copy separately (mamba_ssm.py:433,451). +// ───────────────────────────────────────────────────────────────────────────── +__global__ void M2StateUpdateKernel(void* out, DType odt, void* state, DType sdt, const void* x, + DType xdt, const void* dtp, DType dtdt, const float* A, + const void* Bp, DType Bdt, const void* Cp, DType Cdt, + const float* D, const void* z, DType zdt, const float* dbp, + const int32_t* sidx, int64_t Nb, int64_t H, int64_t P, + int64_t G, int64_t N, int64_t S, int64_t hpg, + bool softplus) { + const int64_t total = Nb * H * P; + for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total; + idx += static_cast(gridDim.x) * blockDim.x) { + const int64_t p = idx % P; + const int64_t h = (idx / P) % H; + const int64_t b = idx / (P * H); + int64_t slot = b; + if (sidx != nullptr) { + // LOCAL ABI: index < 0 is the NULL row — its cache slot is untouched + // (`continue`, mamba_kernels.hpp:147) and its output row is zeroed, as + // GdnDecode already models it. + if (sidx[b] < 0) { + M2Store(out, odt, (b * H + h) * P + p, 0.0f); + continue; + } + slot = sidx[b]; + // The host arm REFUSES an out-of-range slot; a device kernel cannot throw, + // so it writes nothing at all rather than out of bounds (see the header + // note on unchecked device-side preconditions). + if (slot >= S) continue; + } + const int64_t gg = h / hpg; + float d = M2Load(dtp, dtdt, b * H + h); + if (dbp != nullptr) d += dbp[h]; + if (softplus) d = M2Softplus(d); + const float dA = expf(A[h] * d); + const float xv = M2Load(x, xdt, (b * H + h) * P + p); + const int64_t sbase = ((slot * H + h) * P + p) * N; + float y = 0.0f; + for (int64_t n = 0; n < N; ++n) { + const float bv = M2Load(Bp, Bdt, (b * G + gg) * N + n); + const float cv = M2Load(Cp, Cdt, (b * G + gg) * N + n); + const float sn = M2Load(state, sdt, sbase + n) * dA + bv * xv * d; + M2Store(state, sdt, sbase + n, sn); + y += sn * cv; + } + if (D != nullptr) y += D[h] * xv; + if (z != nullptr) y *= M2Silu(M2Load(z, zdt, (b * H + h) * P + p)); + M2Store(out, odt, (b * H + h) * P + p, y); + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// vt::RmsNormGatedGroup — `Mixer2RMSNormGated.forward_native` (mamba_mixer2.py:100-149). +// v = x * silu(f32(gate)) (:114) +// out = weight * dtype(x)( v * rsqrt(mean(v^2 over its group) + eps) ) (:136-141, :149) +// +// One BLOCK per (row, group). This is the ONE accumulation in this file whose +// order differs from the host arm's sequential sum: a per-group reduction is the +// whole shape of the op, and a block reduction is how it is done on device. The +// summands are all NON-NEGATIVE (they are squares), so there is no cancellation +// and the reordering carries the plain forward-error bound the tests state. +// ───────────────────────────────────────────────────────────────────────────── +__device__ inline float M2BlockReduceSum(float v) { + __shared__ float smem[kM2Block / 32]; + const int lane = static_cast(threadIdx.x) & 31; + const int warp = static_cast(threadIdx.x) >> 5; + for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffffu, v, off); + if (lane == 0) smem[warp] = v; + __syncthreads(); + if (threadIdx.x == 0) { + float total = 0.0f; + const int nwarps = static_cast(blockDim.x) >> 5; + for (int i = 0; i < nwarps; ++i) total += smem[i]; + smem[0] = total; + } + __syncthreads(); + const float out = smem[0]; + __syncthreads(); // smem is reused by the next (row, group) this block takes + return out; +} + +__global__ void M2GatedNormKernel(void* out, DType odt, const void* x, DType xdt, DType input_dt, + const void* gate, DType gdt, const void* w, DType wdt, + bool has_w, int64_t nblocks, int64_t hidden, int64_t n_groups, + int64_t group_size, float eps) { + for (int64_t blk = blockIdx.x; blk < nblocks; blk += gridDim.x) { + const int64_t r = blk / n_groups, gsel = blk % n_groups; + const int64_t off = r * hidden + gsel * group_size; + float part = 0.0f; + for (int64_t j = threadIdx.x; j < group_size; j += blockDim.x) { + // The gate is promoted to f32 BEFORE the silu (:114). + const float v = M2Load(x, xdt, off + j) * M2Silu(M2Load(gate, gdt, off + j)); + if (!has_w) { + // use_rms_norm == False: no parameter, no norm (:94-96, :115-116). + M2Store(out, odt, off + j, M2RoundThrough(input_dt, v)); + } + part += v * v; + } + if (!has_w) continue; + // f32 accumulation, NOT double: upstream reduces `x.pow(2).mean(-1)` in f32. + const float ss = M2BlockReduceSum(part); + // `rsqrt(variance + eps)` (:130, :141) — eps is INSIDE the square root. + // Written as 1/sqrt to match the host reference's rounding, not `rsqrtf`. + const float inv = 1.0f / sqrtf(ss / static_cast(group_size) + eps); + for (int64_t j = threadIdx.x; j < group_size; j += blockDim.x) { + const float v = M2Load(x, xdt, off + j) * M2Silu(M2Load(gate, gdt, off + j)); + const float normed = M2RoundThrough(input_dt, v * inv); + M2Store(out, odt, off + j, M2Load(w, wdt, gsel * group_size + j) * normed); + } + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// host launchers +// ───────────────────────────────────────────────────────────────────────────── + +void Mamba2ChunkScanKernelCuda(Queue& q, Tensor& out, Tensor& final_states, const Tensor& x, + const Tensor& dt_in, const Tensor& A, const Tensor& B, + const Tensor& C, const Tensor* D, const Tensor* z, + const Tensor* dt_bias, const Tensor* initial_states, + const Tensor& /*cu_seqlens*/, const Tensor& cu_chunk_seqlens, + const Tensor& last_chunk_indices, const Tensor& seq_idx, + const Mamba2Args& args) { + const int64_t T = x.shape[0], H = x.shape[1], P = x.shape[2]; + const int64_t G = B.shape[1], N = B.shape[2]; + const int64_t S = final_states.shape[0]; + const int64_t cs = args.chunk_size; + const int64_t nchunks = cu_chunk_seqlens.shape[0] - 1; + const int64_t hpg = H / G; // nheads_ngroups_ratio (ssd_chunk_state.py:238) + if (T == 0 || nchunks == 0) return; + + cudaStream_t s = M2Stream(q); + const int32_t* ccs = cu_chunk_seqlens.Ptr(); + const int32_t* lci = last_chunk_indices.Ptr(); + const int32_t* sidx = seq_idx.Ptr(); + const float* Ap = A.Ptr(); + const float* dbp = dt_bias != nullptr ? dt_bias->Ptr() : nullptr; + const float* Dp = D != nullptr ? D->Ptr() : nullptr; + const bool d_has_hdim = D != nullptr && D->rank == 2; + const DType sdt = final_states.dtype; // `state_dtype` (ssd_combined.py:46,119,176) + + // Per-call scratch on the stream's memory pool, exactly as the sibling varlen + // prefill KdaChunkPrefill does (cuda_gdn.cu). This is the PREFILL path; the + // decode kernel below allocates nothing. + const size_t n_cumsum = static_cast(H * nchunks * cs); + const size_t n_states = static_cast(nchunks * H * P * N); + const size_t n_cb = static_cast(nchunks * G * cs * cs); + const size_t state_elem = sdt == DType::kF32 ? 4u : 2u; + float* dtv = nullptr; + float* dac = nullptr; + float* states = nullptr; + float* cb = nullptr; + void* passed = nullptr; + M2Check(cudaMallocAsync(&dtv, n_cumsum * sizeof(float), s), "dtv alloc"); + M2Check(cudaMallocAsync(&dac, n_cumsum * sizeof(float), s), "dac alloc"); + M2Check(cudaMallocAsync(&states, n_states * sizeof(float), s), "states alloc"); + M2Check(cudaMallocAsync(&cb, n_cb * sizeof(float), s), "cb alloc"); + M2Check(cudaMallocAsync(&passed, n_states * state_elem, s), "passed alloc"); + // cudaMallocAsync hands back DIRTY pool memory. Every element of dtv/dac/states + // is written before it is read; `cb` and `passed` are written only where they + // are read (the causal triangle, and the chunks of a scheduled sequence), so + // they are zeroed rather than left to an initcheck report. + M2Check(cudaMemsetAsync(cb, 0, n_cb * sizeof(float), s), "cb zero"); + M2Check(cudaMemsetAsync(passed, 0, n_states * state_elem, s), "passed zero"); + + M2CumsumKernel<<>>( + dtv, dac, dt_in.data, dt_in.dtype, Ap, dbp, ccs, H, nchunks, cs, args.dt_softplus, + args.dt_min, args.dt_max); + M2ChunkStateKernel<<(n_states)), kM2Block, 0, s>>>( + states, x.data, x.dtype, B.data, B.dtype, dtv, dac, ccs, nchunks, H, P, G, N, cs, hpg); + M2StatePassKernel<<(S) * H * P * N), kM2Block, 0, s>>>( + passed, sdt, final_states.data, states, dac, lci, + initial_states != nullptr ? initial_states->data : nullptr, + initial_states != nullptr ? initial_states->dtype : DType::kF32, S, H, P, N, nchunks, cs); + M2BmmKernel<<(n_cb)), kM2Block, 0, s>>>( + cb, B.data, B.dtype, C.data, C.dtype, ccs, nchunks, G, N, cs); + M2ChunkScanKernel<<>>( + out.data, out.dtype, x.data, x.dtype, C.data, C.dtype, cb, dtv, dac, passed, sdt, + initial_states != nullptr ? initial_states->data : nullptr, + initial_states != nullptr ? initial_states->dtype : DType::kF32, Dp, d_has_hdim, + z != nullptr ? z->data : nullptr, z != nullptr ? z->dtype : DType::kF32, ccs, sidx, + nchunks, H, P, G, N, cs, hpg); + + const cudaError_t launched = cudaGetLastError(); + M2Check(cudaFreeAsync(dtv, s), "dtv free"); + M2Check(cudaFreeAsync(dac, s), "dac free"); + M2Check(cudaFreeAsync(states, s), "states free"); + M2Check(cudaFreeAsync(cb, s), "cb free"); + M2Check(cudaFreeAsync(passed, s), "passed free"); + M2Check(launched, "mamba2_chunk_scan launch"); +} + +void Mamba2StateUpdateKernelCuda(Queue& q, Tensor& out, Tensor& state, const Tensor& x, + const Tensor& dt_in, const Tensor& A, const Tensor& B, + const Tensor& C, const Tensor* D, const Tensor* z, + const Tensor* dt_bias, const Tensor* state_indices, + const Mamba2Args& args) { + const int64_t Nb = x.shape[0], H = x.shape[1], P = x.shape[2]; + const int64_t G = B.shape[1], N = B.shape[2]; + const int64_t S = state.shape[0]; + const int64_t hpg = H / G; + if (Nb == 0) return; + cudaStream_t s = M2Stream(q); + M2StateUpdateKernel<<>>( + out.data, out.dtype, state.data, state.dtype, x.data, x.dtype, dt_in.data, dt_in.dtype, + A.Ptr(), B.data, B.dtype, C.data, C.dtype, + D != nullptr ? D->Ptr() : nullptr, z != nullptr ? z->data : nullptr, + z != nullptr ? z->dtype : DType::kF32, dt_bias != nullptr ? dt_bias->Ptr() : nullptr, + state_indices != nullptr ? state_indices->Ptr() : nullptr, Nb, H, P, G, N, S, hpg, + args.dt_softplus); + M2Check(cudaGetLastError(), "mamba2_state_update launch"); +} + +void RmsNormGatedGroupKernelCuda(Queue& q, Tensor& out, const Tensor& x, const Tensor& gate, + const Tensor* weight, const RmsNormGatedGroupArgs& args) { + const int64_t hidden = x.shape[x.rank - 1]; + int64_t rows = 1; + for (int r = 0; r < x.rank - 1; ++r) rows *= x.shape[r]; + if (rows == 0 || hidden == 0) return; + const int64_t group_size = hidden / args.n_groups; + const int64_t nblocks = rows * args.n_groups; + // `input_dtype = x.dtype` (:113) is the width the normalized value is cast back + // to before the weight multiply (`self.weight * x.to(input_dtype)`, :149). + cudaStream_t s = M2Stream(q); + const unsigned grid = static_cast(nblocks < 65535 ? nblocks : 65535); + M2GatedNormKernel<<>>( + out.data, out.dtype, x.data, x.dtype, x.dtype, gate.data, gate.dtype, + weight != nullptr ? weight->data : nullptr, + weight != nullptr ? weight->dtype : DType::kF32, weight != nullptr, nblocks, hidden, + args.n_groups, group_size, args.eps); + M2Check(cudaGetLastError(), "rms_norm_gated_group launch"); +} + +} // namespace +} // namespace vt::cuda::mamba2 + +#endif // VT_CUDA_MAMBA2_SSD_CUH_ diff --git a/tests/vt/test_ops_mamba2_gated_norm.cpp b/tests/vt/test_ops_mamba2_gated_norm.cpp index 81e160911..acc7d66dd 100644 --- a/tests/vt/test_ops_mamba2_gated_norm.cpp +++ b/tests/vt/test_ops_mamba2_gated_norm.cpp @@ -543,3 +543,347 @@ TEST_CASE("mamba2 gated group norm refuses the arms it does not implement") { CHECK_THROWS(RunNorm(in, {rows, hidden}, 7, 1e-6f, DType::kF32, true)); } } + +// ═════════════════════════════════════════════════════════════════════════════ +// (10) THE CUDA ARM — .agents/specs/mamba2-ssd.md W2, issue #496. +// +// The declared equivalence contract is stated in full at the head of the CUDA +// section of tests/vt/test_ops_mamba2_ssd.cpp and in +// src/vt/cuda/cuda_mamba2_ssd.cuh. ONE thing is different for this op, and it is +// stated here rather than inherited: the group reduction is a BLOCK reduction on +// device and a sequential sum on host, so this arm admits a second source of +// divergence — summation ORDER — on top of the libm difference. Its summands are +// all squares, hence non-negative, so there is no cancellation and the +// reordering carries the plain forward-error bound: for a length-m sum, +// |fl_a - fl_b| <= 2(m-1)*u*sum|x| = 2(m-1)*u*sum(x) because sum|x| IS the sum. +// `DerivedRtol(group_size)` covers it with the same (K + 8)*u expression the +// other two suites use, K being the group's own length. Nothing is tuned. +// ═════════════════════════════════════════════════════════════════════════════ +#ifdef VLLM_CPP_CUDA + +#include +#include + +namespace { + +using vt::Backend; + +Backend* MaybeCuda() { + try { + return &vt::GetBackend(DeviceType::kCUDA); + } catch (const std::exception&) { + return nullptr; + } +} + +// A GREEN TEST DOES NOT PROVE THE DEVICE RAN IT. GB10 is +// `integrated && pageable_memory_access`, so `Backend::UnifiedMemory()` is TRUE +// and `ReferenceTierEligible(kCUDA)` with it: absent a native kernel, `GetOp` +// does not throw — it installs the CPU HOST kernel as a `kReferenceProviderName` +// provider and runs THAT over the device pointers (op_provider.h, "portable +// reference tier"), so every assertion below would pass while nothing ran on the +// GPU. Every CUDA case therefore asserts the SELECTED provider is native. These +// are EAGER dispatches, so the counters are populated +// ([[graph-replay-does-no-host-dispatch-counters-read-zero]]). +void RequireNativeCudaProvider(vt::OpId op, const std::string& what) { + const vt::OpProviderStats st = vt::GetOpProviderStats(op, DeviceType::kCUDA); + INFO(what << ": selected CUDA provider = " + << (st.last_selected != nullptr ? st.last_selected : "") + << "; process-wide reference-tier hits = " << vt::GetReferenceTierHits()); + REQUIRE(st.last_selected != nullptr); + CHECK(std::string(st.last_selected) != std::string(vt::kReferenceProviderName)); +} + +// `4*(K + 2)*u` — the bound derived at the head of the CUDA section of +// tests/vt/test_ops_mamba2_ssd.cpp, which for this op covers the reordering of a +// length-K non-negative reduction (2(K-1)*u, no cancellation because sum|x| IS +// the sum) plus the libm difference in silu's `expf`. +constexpr double kUnitRoundoff = 5.9604644775390625e-08; // 2^-24 +double DerivedRtol(int64_t K) { return 4.0 * static_cast(K + 2) * kUnitRoundoff; } + +void ExpectDeviceMatchesHost(const std::string& what, const std::vector& dev, + const std::vector& host, int64_t K) { + REQUIRE(dev.size() == host.size()); + REQUIRE(!dev.empty()); + double scale = 0.0; + for (float v : host) scale = std::max(scale, std::abs(static_cast(v))); + const double rtol = DerivedRtol(K); + const double atol = rtol * scale; + size_t bit_differing = 0, worst_i = 0; + double worst_ratio = -1.0, worst_diff = 0.0; + for (size_t i = 0; i < dev.size(); ++i) { + if (dev[i] != host[i]) ++bit_differing; + const double d = std::abs(static_cast(dev[i]) - static_cast(host[i])); + const double budget = atol + rtol * std::abs(static_cast(host[i])); + const double ratio = budget > 0.0 ? d / budget : (d > 0.0 ? 1e30 : 0.0); + if (!std::isfinite(static_cast(dev[i])) || ratio > worst_ratio) { + worst_ratio = ratio; + worst_i = i; + worst_diff = d; + if (!std::isfinite(static_cast(dev[i]))) break; + } + } + // MESSAGE, not INFO: doctest prints an INFO context only when an assertion in + // its scope FAILS, so the used slack has to be logged unconditionally for the + // derived bar to be auditable on the green run that matters. + MESSAGE(what << ": K=" << K << " rtol=" << rtol << " scale=" << scale << "; " + << bit_differing << " of " << dev.size() + << " elements differ in any bit; worst element [" << worst_i + << "] dev=" << dev[worst_i] << " host=" << host[worst_i] + << " |diff|=" << worst_diff << " used " << (worst_ratio * 100.0) + << "% of its derived budget"); + CHECK(std::isfinite(static_cast(dev[worst_i]))); + CHECK(worst_ratio <= 1.0); +} + +Tensor MakeTDev(void* data, DType dt, Device dev, const std::vector& shape) { + Tensor t; + t.data = data; + t.dtype = dt; + t.device = dev; + t.rank = static_cast(shape.size()); + int64_t stride = 1; + for (int i = t.rank - 1; i >= 0; --i) { + t.shape[i] = shape[static_cast(i)]; + t.stride[i] = stride; + stride *= shape[static_cast(i)]; + } + return t; +} + +class DBuf { + public: + DBuf(Backend& b, Queue& q, const void* host, size_t bytes) : b_(&b), bytes_(bytes) { + p_ = b.Alloc(bytes == 0 ? 1 : bytes); + if (host != nullptr && bytes > 0) b.Copy(q, p_, host, bytes); + } + ~DBuf() { + if (p_ != nullptr) b_->Free(p_); + } + DBuf(const DBuf&) = delete; + DBuf& operator=(const DBuf&) = delete; + void* get() const { return p_; } + void Download(Queue& q, void* dst) const { + if (bytes_ > 0) b_->Copy(q, dst, p_, bytes_); + b_->Synchronize(q); + } + + private: + Backend* b_; + void* p_ = nullptr; + size_t bytes_ = 0; +}; + +// The CUDA twin of RunNorm, argument for argument. +std::vector RunNormCuda(Backend& gpu, const NormInputs& in, + const std::vector& shape, int64_t n_groups, float eps, + DType dt, bool use_rms_norm, int64_t tp_world_size = 1, + DType weight_dt = DType::kF32, DType out_dt = kSameAsAct) { + Queue q = gpu.CreateQueue(); + const Device dev{DeviceType::kCUDA, 0}; + if (out_dt == kSameAsAct) out_dt = dt; + size_t n = 1; + for (int64_t d : shape) n *= static_cast(d); + const std::vector xb = Pack(in.x, dt); + const std::vector gb = Pack(in.gate, dt); + const std::vector wb = Pack(in.weight, weight_dt); + const size_t out_bytes = n * vt::SizeOf(out_dt); + + DBuf dx(gpu, q, xb.data(), xb.size()); + DBuf dg(gpu, q, gb.data(), gb.size()); + DBuf dw(gpu, q, wb.data(), wb.size()); + DBuf dout(gpu, q, nullptr, out_bytes); + + Tensor xt = MakeTDev(dx.get(), dt, dev, shape); + Tensor gt = MakeTDev(dg.get(), dt, dev, shape); + Tensor ot = MakeTDev(dout.get(), out_dt, dev, shape); + Tensor wt = MakeTDev(dw.get(), weight_dt, dev, {shape.back()}); + + RmsNormGatedGroupArgs args; + args.eps = eps; + args.n_groups = n_groups; + args.tp_world_size = tp_world_size; + vt::RmsNormGatedGroup(q, ot, xt, gt, use_rms_norm ? &wt : nullptr, args); + + std::vector ob(out_bytes); + dout.Download(q, ob.data()); + gpu.Synchronize(q); + gpu.DestroyQueue(q); + return Unpack(ob, n, out_dt); +} + +} // namespace + +TEST_CASE("mamba2 gated group norm CUDA arm matches forward_native") { + Backend* gpu = MaybeCuda(); + if (gpu == nullptr) { + MESSAGE("SKIP: no CUDA backend registered (CPU-only build/box)"); + return; + } + // Upstream's shapes (batch 8, seq 128, hidden 64, n_groups in {1,2,4}, + // test_mamba_mixer2.py:21-33) plus one at the driver's own width: the mamba + // layer's gated norm runs over `intermediate_size` with n_groups = 8 + // (mamba2-ssd.md §1.4), which is the case whose group_size is large enough for + // the block reduction to differ from the host's sequential sum at all. + struct Case { + int64_t rows, hidden, n_groups; + }; + const std::vector cases{{8 * 128, 64, 1}, {8 * 128, 64, 2}, + {8 * 128, 64, 4}, {64, 4096, 8}}; + const float eps = 1e-6f; + for (const Case& c : cases) { + const NormInputs in = GenerateNorm(c.rows, c.hidden, 0x9A17Eu); + const int64_t group_size = c.hidden / c.n_groups; + INFO("rows=" << c.rows << " hidden=" << c.hidden << " n_groups=" << c.n_groups); + + const std::vector ref = + GatedGroupNormRef(in.x, in.gate, &in.weight, c.rows, c.hidden, c.n_groups, eps); + const std::vector host = + RunNorm(in, {c.rows, c.hidden}, c.n_groups, eps, DType::kF32, true); + const std::vector dev = + RunNormCuda(*gpu, in, {c.rows, c.hidden}, c.n_groups, eps, DType::kF32, true); + RequireNativeCudaProvider(vt::OpId::kRmsNormGatedGroup, "forward_native shapes"); + // G1 — device vs the INDEPENDENT double reference at upstream's tolerance. + ExpectClose("device out f32", dev, ref, 5e-3, 1e-3); + // G2 — device vs host, at the derived bound for a length-group_size reduction. + ExpectDeviceMatchesHost("out f32 device vs host", dev, host, group_size); + + // bf16 activation arm (upstream runs float16; the vt `out` contract is + // f32/bf16). `input_dtype` is x's dtype (:113), so the reference casts + // through bf16 at :149 too. + const std::vector bref = + GatedGroupNormRef(in.x, in.gate, &in.weight, c.rows, c.hidden, c.n_groups, eps, + DType::kBF16); + ExpectClose("device out bf16", + RunNormCuda(*gpu, in, {c.rows, c.hidden}, c.n_groups, eps, DType::kBF16, true), + bref, 5e-2, 1e-2); + } +} + +// The two differences that make this a SIBLING of vt::RmsNormGated rather than a +// mode of it, pinned on device: the reduction extent (per GROUP, not per row) and +// the activation (silu, not sigmoid). +TEST_CASE("mamba2 gated group norm CUDA arm keeps both sibling differences") { + Backend* gpu = MaybeCuda(); + if (gpu == nullptr) { + MESSAGE("SKIP: no CUDA backend registered (CPU-only build/box)"); + return; + } + const int64_t rows = 32, hidden = 64; + const float eps = 1e-6f; + const NormInputs in = GenerateNorm(rows, hidden, 0x6A7Eu); + + SUBCASE("the reduction is per group, not per row") { + const std::vector g1 = + RunNormCuda(*gpu, in, {rows, hidden}, 1, eps, DType::kF32, true); + const std::vector g4 = + RunNormCuda(*gpu, in, {rows, hidden}, 4, eps, DType::kF32, true); + RequireNativeCudaProvider(vt::OpId::kRmsNormGatedGroup, "per-group reduction"); + REQUIRE(g1.size() == g4.size()); + double max_diff = 0.0; + for (size_t i = 0; i < g1.size(); ++i) + max_diff = std::max(max_diff, std::abs(static_cast(g1[i]) - g4[i])); + INFO("max|n_groups=1 - n_groups=4| on device = " << max_diff); + // A whole-row variance would make these IDENTICAL. They must not be. + CHECK(max_diff > 1e-3); + // ... and n_groups=4 is the one that matches a per-group double reference. + ExpectClose("device n_groups=4", g4, + GatedGroupNormRef(in.x, in.gate, &in.weight, rows, hidden, 4, eps), 5e-3, 1e-3); + } + + SUBCASE("the gate is silu, not sigmoid") { + // A sigmoid-gated reference must NOT match what the device produced. + const NormInputs& sig = in; + const std::vector dev = + RunNormCuda(*gpu, in, {rows, hidden}, 2, eps, DType::kF32, true); + // The reference a SIGMOID gate would give, written out here rather than + // parameterised, so it shares nothing with the op under test. + std::vector sigmoid_ref(static_cast(rows * hidden), 0.0); + { + const int64_t group_size = hidden / 2; + for (int64_t r = 0; r < rows; ++r) { + std::vector v(static_cast(hidden)); + for (int64_t j = 0; j < hidden; ++j) { + const double zv = sig.gate[static_cast(r * hidden + j)]; + v[static_cast(j)] = + static_cast(sig.x[static_cast(r * hidden + j)]) / + (1.0 + std::exp(-zv)); // SIGMOID, not silu + } + for (int64_t g = 0; g < 2; ++g) { + double ss = 0.0; + for (int64_t j = 0; j < group_size; ++j) { + const double t = v[static_cast(g * group_size + j)]; + ss += t * t; + } + const double inv = 1.0 / std::sqrt(ss / static_cast(group_size) + eps); + for (int64_t j = 0; j < group_size; ++j) { + const int64_t idx = g * group_size + j; + sigmoid_ref[static_cast(r * hidden + idx)] = + static_cast(sig.weight[static_cast(idx)]) * + v[static_cast(idx)] * inv; + } + } + } + } + double max_diff = 0.0; + for (size_t i = 0; i < dev.size(); ++i) + max_diff = std::max(max_diff, std::abs(static_cast(dev[i]) - sigmoid_ref[i])); + INFO("max|device - sigmoid-gated reference| = " << max_diff); + CHECK(max_diff > 1e-2); + } +} + +// `use_rms_norm == False`: no parameter, no norm, just the gated value cast back +// to the input dtype (mamba_mixer2.py:94-96, :115-116); and every leading dim is +// a row (`*prefix_dims, hidden_dim`, :136), so a rank-3 [T,H,D] input is the same +// computation as its flattened rank-2 view. +TEST_CASE("mamba2 gated group norm CUDA arm covers the no-weight and rank-3 arms") { + Backend* gpu = MaybeCuda(); + if (gpu == nullptr) { + MESSAGE("SKIP: no CUDA backend registered (CPU-only build/box)"); + return; + } + const float eps = 1e-6f; + + SUBCASE("no weight: the gated value, unnormalised") { + const int64_t rows = 40, hidden = 64; + const NormInputs in = GenerateNorm(rows, hidden, 0x1234u); + const std::vector ref = + GatedGroupNormRef(in.x, in.gate, nullptr, rows, hidden, 4, eps); + const std::vector dev = + RunNormCuda(*gpu, in, {rows, hidden}, 4, eps, DType::kF32, false); + RequireNativeCudaProvider(vt::OpId::kRmsNormGatedGroup, "no-weight arm"); + ExpectClose("device out (no weight)", dev, ref, 5e-3, 1e-3); + ExpectDeviceMatchesHost("no-weight device vs host", dev, + RunNorm(in, {rows, hidden}, 4, eps, DType::kF32, false), 1); + } + + SUBCASE("rank 3 [T,H,D] equals its flattened rank-2 view") { + const int64_t T = 12, Hh = 4, Dd = 32; + const int64_t rows = T * Hh; + const NormInputs in = GenerateNorm(rows, Dd, 0x5678u); + const std::vector flat = + RunNormCuda(*gpu, in, {rows, Dd}, 2, eps, DType::kF32, true); + const std::vector cube = + RunNormCuda(*gpu, in, {T, Hh, Dd}, 2, eps, DType::kF32, true); + REQUIRE(flat.size() == cube.size()); + for (size_t i = 0; i < flat.size(); ++i) CHECK(flat[i] == cube[i]); + } + + SUBCASE("the weight is read at the WEIGHT's dtype") { + // `Mixer2RMSNormGated.weight` is created at the MODEL dtype (:91), bf16 for + // every checkpoint that ships this layer. A kernel that read it as f32 would + // over-read a real allocation and shift every output — the W1 F1 finding + // (mamba2-ssd.md §8.2), re-pinned here for the device arm. + const int64_t rows = 32, hidden = 64; + NormInputs in = GenerateNorm(rows, hidden, 0x9999u); + for (auto& w : in.weight) w = vt::BF16ToF32(vt::F32ToBF16(w)); + const std::vector ref = + GatedGroupNormRef(in.x, in.gate, &in.weight, rows, hidden, 2, eps); + const std::vector dev = RunNormCuda(*gpu, in, {rows, hidden}, 2, eps, DType::kF32, + true, 1, DType::kBF16); + ExpectClose("device out, bf16 weight", dev, ref, 5e-3, 1e-3); + } +} + +#endif // VLLM_CPP_CUDA diff --git a/tests/vt/test_ops_mamba2_ssd.cpp b/tests/vt/test_ops_mamba2_ssd.cpp index e4a5eee7c..5c6777c9c 100644 --- a/tests/vt/test_ops_mamba2_ssd.cpp +++ b/tests/vt/test_ops_mamba2_ssd.cpp @@ -999,3 +999,581 @@ TEST_CASE("mamba2 chunk scan refuses the arms it does not implement") { nullptr, nullptr, cust, ccst, lcit, sit, args)); } } + +// ═════════════════════════════════════════════════════════════════════════════ +// (8) THE CUDA ARM — .agents/specs/mamba2-ssd.md W2, issue #496. +// +// ─── THE DECLARED EQUIVALENCE CONTRACT ────────────────────────────────────── +// The CUDA kernels keep f32 accumulation THROUGHOUT and do NOT mirror the tile +// downcasts inside upstream's Triton dots (`b.to(x_ptr.dtype.element_ty)` +// ssd_chunk_state.py:283-285; `cb.to(...)` / `prev_states.to(C_ptr.dtype...)` +// ssd_chunk_scan.py:266-269, :359-363). Those casts are the input-precision +// requirement of `tl.dot`, i.e. of a tensor-core MMA — every one of those tiles is +// loaded `.to(tl.float32)` and computed in f32 right up to the MMA. Our kernels +// are scalar-FMA kernels with no MMA, so the bar here is NOT a downcast-derived +// tolerance. The memory format is unchanged: every load and store goes through the +// operand's own dtype, and the inter-chunk `passed` buffer is allocated at +// `state_dtype`, not at the host reference's f32 working width (§8.2 F9). +// +// Two gates, in order of authority: +// +// G1 (PRIMARY, INDEPENDENT). The device outputs are held to the SAME +// double-precision sequential reference (`SequentialSsdRef`) at the SAME +// upstream-ported tolerances as the host arm. Every structural defect — a +// dropped inter-chunk term, `states[c]` for `states[c-1]`, ignored +// `initial_states`, a missing `D` skip — is an O(1) error and fails here, +// against a reference the kernel was not written from. +// +// G2 (DEVICE vs HOST, DERIVED). A BYTE COMPARE IS NOT REACHABLE, and the reason +// is exactly one thing: the two arms call different libms. CUDA's `expf` is +// documented to <= 2 ulp and glibc's to <= 0.5, and they disagree in the last +// ulp on some inputs, so `exp`/`log1p` alone put the two arms off each other by +// ulps that the recurrence then amplifies. Everything else is held IDENTICAL by +// construction: the device kernels accumulate every output element in ONE +// thread, over the same index range in the same direction as the host arm, so +// summation order is not a second source. `DerivedRtol` below propagates that +// one source, and only that. Nothing in it was tuned, and the slack actually +// USED is reported on every comparison — if it ever approached the bar, the bar +// would have stopped being a statement about libm and the finding would be a +// NEEDS_DECISION, not a wider tolerance. +// ═════════════════════════════════════════════════════════════════════════════ +#ifdef VLLM_CPP_CUDA + +#include +#include + +namespace { + +using vt::Backend; + +Backend* MaybeCuda() { + try { + return &vt::GetBackend(DeviceType::kCUDA); + } catch (const std::exception&) { + return nullptr; + } +} + +// A GREEN TEST DOES NOT PROVE THE DEVICE RAN IT — and on THIS box it very nearly +// proves the opposite. GB10 is `integrated && pageable_memory_access`, so +// `Backend::UnifiedMemory()` is TRUE (cuda_backend.cu Registrar) and therefore +// `ReferenceTierEligible(kCUDA)` is TRUE. Absent a native kernel, `GetOp` does +// not throw: it installs the CPU HOST kernel as a `kReferenceProviderName` +// provider and runs THAT over the device pointers (op_provider.h, "portable +// reference tier"). Every numeric assertion below would then pass — the device +// arm would be gated by running the host arm twice, the exact false-green shape +// of [[absent-hook-looks-like-armed-instrument]] and +// [[gate-comparing-shared-helper-proves-consistency-not-correctness]]. +// +// So every CUDA case asserts the SELECTED provider is native. These are EAGER +// dispatches rather than a captured graph, so the counters are genuinely +// populated ([[graph-replay-does-no-host-dispatch-counters-read-zero]]). +void RequireNativeCudaProvider(vt::OpId op, const std::string& what) { + const vt::OpProviderStats st = vt::GetOpProviderStats(op, DeviceType::kCUDA); + INFO(what << ": selected CUDA provider = " + << (st.last_selected != nullptr ? st.last_selected : "") + << "; process-wide reference-tier hits = " << vt::GetReferenceTierHits()); + REQUIRE(st.last_selected != nullptr); + CHECK(std::string(st.last_selected) != std::string(vt::kReferenceProviderName)); +} + +// f32 unit roundoff, and the bound the two arms are held to. +// +// A value that has run through a product of at most K decay factors and a +// length-K f32 summation carries, between the two arms: +// * <= 2.5 ulp of libm disagreement PER FACTOR — CUDA's `expf` is documented +// to <= 2 ulp and glibc's to <= 0.5 — so <= 2.5*K*u on the product, and +// * the standard (K-1)*u forward error of the summation itself, +// i.e. <= 3.5*K*u. `4*(K + 2)*u` is that, rounded up to integers. K is the +// case's own recurrence length; nothing here is fitted, and the slack actually +// USED is reported on every comparison so a bar that had stopped doing work +// would be visible rather than silently absorbing a defect. +constexpr double kUnitRoundoff = 5.9604644775390625e-08; // 2^-24 +double DerivedRtol(int64_t K) { return 4.0 * static_cast(K + 2) * kUnitRoundoff; } + +// atol is `rtol * max|host|` rather than 0: a bound proportional to |want| alone +// is vacuous for an element that is near zero through cancellation of O(max) +// terms, which this recurrence produces routinely. +void ExpectDeviceMatchesHost(const std::string& what, const std::vector& dev, + const std::vector& host, int64_t K) { + REQUIRE(dev.size() == host.size()); + REQUIRE(!dev.empty()); + double scale = 0.0; + for (float v : host) scale = std::max(scale, std::abs(static_cast(v))); + const double rtol = DerivedRtol(K); + const double atol = rtol * scale; + size_t bit_differing = 0, worst_i = 0; + double worst_ratio = -1.0, worst_diff = 0.0; + for (size_t i = 0; i < dev.size(); ++i) { + if (dev[i] != host[i]) ++bit_differing; + const double d = std::abs(static_cast(dev[i]) - static_cast(host[i])); + const double budget = atol + rtol * std::abs(static_cast(host[i])); + const double ratio = budget > 0.0 ? d / budget : (d > 0.0 ? 1e30 : 0.0); + if (!std::isfinite(static_cast(dev[i])) || ratio > worst_ratio) { + worst_ratio = ratio; + worst_i = i; + worst_diff = d; + if (!std::isfinite(static_cast(dev[i]))) break; + } + } + // MESSAGE, not INFO: doctest prints an INFO context only when an assertion in + // its scope FAILS, so a claim that the used slack "is reported" would be false + // on the green run that matters. This line is emitted unconditionally, which is + // what makes the derived bar auditable rather than merely asserted. + MESSAGE(what << ": K=" << K << " rtol=" << rtol << " scale=" << scale << "; " + << bit_differing << " of " << dev.size() + << " elements differ in any bit; worst element [" << worst_i + << "] dev=" << dev[worst_i] << " host=" << host[worst_i] + << " |diff|=" << worst_diff << " used " << (worst_ratio * 100.0) + << "% of its derived budget"); + CHECK(std::isfinite(static_cast(dev[worst_i]))); + CHECK(worst_ratio <= 1.0); +} + +Tensor MakeTDev(void* data, DType dt, Device dev, const std::vector& shape) { + Tensor t; + t.data = data; + t.dtype = dt; + t.device = dev; + t.rank = static_cast(shape.size()); + int64_t stride = 1; + for (int i = t.rank - 1; i >= 0; --i) { + t.shape[i] = shape[static_cast(i)]; + t.stride[i] = stride; + stride *= shape[static_cast(i)]; + } + return t; +} + +// Owning device allocation, uploaded from host bytes (or left zero-sized). +class DBuf { + public: + DBuf(Backend& b, Queue& q, const void* host, size_t bytes) : b_(&b), bytes_(bytes) { + p_ = b.Alloc(bytes == 0 ? 1 : bytes); + if (host != nullptr && bytes > 0) b.Copy(q, p_, host, bytes); + } + ~DBuf() { + if (p_ != nullptr) b_->Free(p_); + } + DBuf(const DBuf&) = delete; + DBuf& operator=(const DBuf&) = delete; + void* get() const { return p_; } + void Download(Queue& q, void* dst) const { + if (bytes_ > 0) b_->Copy(q, dst, p_, bytes_); + b_->Synchronize(q); + } + + private: + Backend* b_; + void* p_ = nullptr; + size_t bytes_ = 0; +}; + +// The CUDA twin of RunChunkScan, argument for argument. +RunOut RunChunkScanCuda(Backend& gpu, const Inputs& in, int64_t T, int64_t H, int64_t P, + int64_t G, int64_t N, const std::vector& cu_seqlens, + const std::vector* D, const std::vector* z, + const std::vector* dt_bias, + const std::vector* initial_states, const RunCfg& cfg) { + Queue q = gpu.CreateQueue(); + const Device dev{DeviceType::kCUDA, 0}; + const int64_t S = static_cast(cu_seqlens.size()) - 1; + ChunkMeta meta = ComputeVarlenChunkMetadata(cu_seqlens, cfg.chunk_size); + const int64_t nchunks = static_cast(meta.seq_idx.size()); + + const std::vector xb = Pack(in.x, cfg.act_dtype); + const std::vector dtb = Pack(in.dt, cfg.act_dtype); + const std::vector Bb = Pack(in.B, cfg.act_dtype); + const std::vector Cb = Pack(in.C, cfg.act_dtype); + const size_t out_bytes = static_cast(T * H * P) * vt::SizeOf(cfg.act_dtype); + const size_t fs_bytes = static_cast(S * H * P * N) * vt::SizeOf(cfg.state_dtype); + std::vector cus = cu_seqlens; + + DBuf dx(gpu, q, xb.data(), xb.size()); + DBuf ddt(gpu, q, dtb.data(), dtb.size()); + DBuf dA(gpu, q, in.A.data(), in.A.size() * sizeof(float)); + DBuf dB(gpu, q, Bb.data(), Bb.size()); + DBuf dC(gpu, q, Cb.data(), Cb.size()); + DBuf dout(gpu, q, nullptr, out_bytes); + DBuf dfs(gpu, q, nullptr, fs_bytes); + DBuf dcus(gpu, q, cus.data(), cus.size() * sizeof(int32_t)); + DBuf dccs(gpu, q, meta.cu_chunk_seqlens.data(), meta.cu_chunk_seqlens.size() * sizeof(int32_t)); + DBuf dlci(gpu, q, meta.last_chunk_indices.data(), + meta.last_chunk_indices.size() * sizeof(int32_t)); + DBuf dsi(gpu, q, meta.seq_idx.data(), meta.seq_idx.size() * sizeof(int32_t)); + + Tensor xt = MakeTDev(dx.get(), cfg.act_dtype, dev, {T, H, P}); + Tensor dtt = MakeTDev(ddt.get(), cfg.act_dtype, dev, {T, H}); + Tensor At = MakeTDev(dA.get(), DType::kF32, dev, {H}); + Tensor Bt = MakeTDev(dB.get(), cfg.act_dtype, dev, {T, G, N}); + Tensor Ct = MakeTDev(dC.get(), cfg.act_dtype, dev, {T, G, N}); + Tensor outt = MakeTDev(dout.get(), cfg.act_dtype, dev, {T, H, P}); + Tensor fst = MakeTDev(dfs.get(), cfg.state_dtype, dev, {S, H, P, N}); + Tensor cust = MakeTDev(dcus.get(), DType::kI32, dev, {S + 1}); + Tensor ccst = MakeTDev(dccs.get(), DType::kI32, dev, {nchunks + 1}); + Tensor lcit = MakeTDev(dlci.get(), DType::kI32, dev, {S}); + Tensor sit = MakeTDev(dsi.get(), DType::kI32, dev, {nchunks}); + + std::vector Dc; + std::unique_ptr dD; + Tensor Dt; + if (D != nullptr) { + Dc = *D; + dD = std::make_unique(gpu, q, Dc.data(), Dc.size() * sizeof(float)); + Dt = cfg.d_has_hdim ? MakeTDev(dD->get(), DType::kF32, dev, {H, P}) + : MakeTDev(dD->get(), DType::kF32, dev, {H}); + } + std::vector zb; + std::unique_ptr dz; + Tensor zt; + if (z != nullptr) { + zb = Pack(*z, cfg.act_dtype); + dz = std::make_unique(gpu, q, zb.data(), zb.size()); + zt = MakeTDev(dz->get(), cfg.act_dtype, dev, {T, H, P}); + } + std::vector dbc; + std::unique_ptr ddb; + Tensor dbt; + if (dt_bias != nullptr) { + dbc = *dt_bias; + ddb = std::make_unique(gpu, q, dbc.data(), dbc.size() * sizeof(float)); + dbt = MakeTDev(ddb->get(), DType::kF32, dev, {H}); + } + std::vector isb; + std::unique_ptr dis; + Tensor ist; + if (initial_states != nullptr) { + isb = Pack(*initial_states, cfg.state_dtype); + dis = std::make_unique(gpu, q, isb.data(), isb.size()); + ist = MakeTDev(dis->get(), cfg.state_dtype, dev, {S, H, P, N}); + } + + Mamba2Args args; + args.chunk_size = cfg.chunk_size; + args.dt_softplus = cfg.dt_softplus; + args.dt_min = cfg.dt_min; + args.dt_max = cfg.dt_max; + + vt::Mamba2ChunkScan(q, outt, fst, xt, dtt, At, Bt, Ct, D != nullptr ? &Dt : nullptr, + z != nullptr ? &zt : nullptr, dt_bias != nullptr ? &dbt : nullptr, + initial_states != nullptr ? &ist : nullptr, cust, ccst, lcit, sit, args); + + std::vector outb(out_bytes), fsb(fs_bytes); + dout.Download(q, outb.data()); + dfs.Download(q, fsb.data()); + gpu.Synchronize(q); + + RunOut r; + r.y = Unpack(outb, static_cast(T * H * P), cfg.act_dtype); + r.final_states = Unpack(fsb, static_cast(S * H * P * N), cfg.state_dtype); + gpu.DestroyQueue(q); + return r; +} + +} // namespace + +// G1 + G2 on the shapes the row exists for: Nemotron-3.5-Lightning-30B-A3B's +// mamba layer — nheads 64, headdim 64, dstate 128, ngroups 8, chunk_size 128, +// mamba_ssm_cache_dtype float32 (mamba2-ssd.md §1.4). T is 200 so `nchunks == 2` +// and the arm actually exercises inter-chunk state passing — the failure shape of +// [[h3-video-decode-temporal-and-tiling-compose]], and of §8.2 F6. +TEST_CASE("mamba2 chunk scan CUDA arm on the driver shapes") { + Backend* gpu = MaybeCuda(); + if (gpu == nullptr) { + MESSAGE("SKIP: no CUDA backend registered (CPU-only build/box)"); + return; + } + const int64_t T = 200, H = 64, P = 64, G = 8, N = 128, chunk = 128; + const Inputs in = GenerateInputs(T, H, P, G, N, 0x4E33Au); + const std::vector cu{0, static_cast(T)}; + REQUIRE(ComputeVarlenChunkMetadata(cu, chunk).seq_idx.size() > 1); + + std::mt19937 rng(0x77u); + std::normal_distribution nd(0.0f, 1.0f); + std::vector D(static_cast(H)); + for (auto& v : D) v = nd(rng); + + RunCfg cfg; + cfg.chunk_size = chunk; + const SeqRefOut ref = + SequentialSsdRef(in, T, H, P, G, N, cu, &D, false, nullptr, nullptr, nullptr, {}); + const RunOut host = RunChunkScan(in, T, H, P, G, N, cu, &D, nullptr, nullptr, nullptr, cfg); + const RunOut dev = + RunChunkScanCuda(*gpu, in, T, H, P, G, N, cu, &D, nullptr, nullptr, nullptr, cfg); + RequireNativeCudaProvider(vt::OpId::kMamba2ChunkScan, "driver shapes"); + + // G1 — the device arm against the INDEPENDENT double reference, at upstream's + // own f32 threshold for the chunked single-example test (atol 8e-3 / rtol 5e-3, + // test_mamba_ssm_ssd.py:210-213). THE HOST ARM IS HELD TO THE SAME NUMBERS on + // the same inputs, so a failure separates cleanly: device-only means a device + // defect, both means the cited threshold does not cover this shape and the + // finding is a NEEDS_DECISION rather than a wider tolerance. + ExpectClose("host y vs sequential double", host.y, ref.y, 8e-3, 5e-3); + ExpectClose("host final_states vs sequential double", host.final_states, ref.final_states, + 8e-3, 5e-3); + ExpectClose("device y vs sequential double", dev.y, ref.y, 8e-3, 5e-3); + ExpectClose("device final_states vs sequential double", dev.final_states, ref.final_states, + 8e-3, 5e-3); + // G2 — device vs host, at the derived libm bound. + ExpectDeviceMatchesHost("y device vs host", dev.y, host.y, T); + ExpectDeviceMatchesHost("final_states device vs host", dev.final_states, host.final_states, T); +} + +// The structural properties the chunked factorisation lives on, on device: +// chunk-size invariance, sequence boundaries INSIDE a physical chunk, and +// `initial_states`. Small shapes so the sweep stays cheap; the driver shapes are +// covered above. +TEST_CASE("mamba2 chunk scan CUDA arm holds the chunked-factorisation properties") { + Backend* gpu = MaybeCuda(); + if (gpu == nullptr) { + MESSAGE("SKIP: no CUDA backend registered (CPU-only build/box)"); + return; + } + + SUBCASE("invariant to chunk_size") { + const int64_t T = 300, H = 8, P = 16, G = 2, N = 32; + const Inputs in = GenerateInputs(T, H, P, G, N, 0xC0FFEEu); + const std::vector cu{0, static_cast(T)}; + const SeqRefOut ref = + SequentialSsdRef(in, T, H, P, G, N, cu, nullptr, false, nullptr, nullptr, nullptr, {}); + std::vector first_y, first_state; + for (int64_t chunk : {8, 16, 32, 64, 128}) { + RunCfg cfg; + cfg.chunk_size = chunk; + const int64_t nchunks = + static_cast(ComputeVarlenChunkMetadata(cu, chunk).seq_idx.size()); + INFO("chunk_size=" << chunk << " nchunks=" << nchunks); + REQUIRE(nchunks > 1); + const RunOut dev = + RunChunkScanCuda(*gpu, in, T, H, P, G, N, cu, nullptr, nullptr, nullptr, nullptr, cfg); + RequireNativeCudaProvider(vt::OpId::kMamba2ChunkScan, "chunk_size invariance"); + const RunOut host = + RunChunkScan(in, T, H, P, G, N, cu, nullptr, nullptr, nullptr, nullptr, cfg); + ExpectClose("device y vs sequential", dev.y, ref.y, 1e-2, 5e-3); + ExpectClose("device final_states vs sequential", dev.final_states, ref.final_states, 1e-2, + 5e-3); + ExpectDeviceMatchesHost("y device vs host", dev.y, host.y, T); + ExpectDeviceMatchesHost("final_states device vs host", dev.final_states, host.final_states, + T); + if (first_y.empty()) { + first_y = dev.y; + first_state = dev.final_states; + } else { + ExpectCloseF("device y vs chunk_size=8", dev.y, first_y, 1e-2, 5e-3); + ExpectCloseF("device final_states vs chunk_size=8", dev.final_states, first_state, 1e-2, + 5e-3); + } + } + } + + SUBCASE("continuous batches, with and without initial_states") { + struct Case { + std::vector lens; + int64_t chunk; + }; + const std::vector cases{ + {{64, 32}, 8}, {{4, 4, 4, 4}, 8}, {{5, 30, 1, 2}, 256}, + {{138, 225}, 128}, {{270, 88}, 8}, + }; + const int64_t H = 8, P = 16, G = 2, N = 16; + for (const Case& c : cases) { + std::vector cu{0}; + for (int32_t l : c.lens) cu.push_back(cu.back() + l); + const int64_t T = cu.back(); + const int64_t S = static_cast(c.lens.size()); + const int64_t maxlen = *std::max_element(c.lens.begin(), c.lens.end()); + const double atol = maxlen > 256 ? 1e-2 : 5e-3; + const Inputs in = GenerateInputs(T, H, P, G, N, 0xBEEF01u + static_cast(c.chunk)); + INFO("chunk=" << c.chunk << " nseq=" << S << " T=" << T); + RunCfg cfg; + cfg.chunk_size = c.chunk; + + // (a) fresh sequences — the `seq_idx[c] != seq_idx[c-1]` branch must take + // ZEROS as the previous state (ssd_chunk_scan.py:271-274). + { + const SeqRefOut ref = SequentialSsdRef(in, T, H, P, G, N, cu, nullptr, false, nullptr, + nullptr, nullptr, {}); + const RunOut dev = RunChunkScanCuda(*gpu, in, T, H, P, G, N, cu, nullptr, nullptr, + nullptr, nullptr, cfg); + RequireNativeCudaProvider(vt::OpId::kMamba2ChunkScan, "continuous batch (fresh)"); + const RunOut host = + RunChunkScan(in, T, H, P, G, N, cu, nullptr, nullptr, nullptr, nullptr, cfg); + ExpectClose("device y (fresh)", dev.y, ref.y, atol, 5e-3); + ExpectClose("device final_states (fresh)", dev.final_states, ref.final_states, atol, + 5e-3); + ExpectDeviceMatchesHost("y (fresh) device vs host", dev.y, host.y, maxlen); + ExpectDeviceMatchesHost("final_states (fresh) device vs host", dev.final_states, + host.final_states, maxlen); + } + // (b) with initial_states — the same branch must instead take + // initial_states[seq_idx[c]] (ssd_chunk_scan.py:236-243). + { + std::mt19937 rng(1234u); + std::normal_distribution nd(0.0f, 0.5f); + std::vector init(static_cast(S * H * P * N)); + for (auto& v : init) v = nd(rng); + const std::vector initd(init.begin(), init.end()); + const SeqRefOut ref = SequentialSsdRef(in, T, H, P, G, N, cu, nullptr, false, nullptr, + nullptr, &initd, {}); + const RunOut dev = + RunChunkScanCuda(*gpu, in, T, H, P, G, N, cu, nullptr, nullptr, nullptr, &init, cfg); + const RunOut host = + RunChunkScan(in, T, H, P, G, N, cu, nullptr, nullptr, nullptr, &init, cfg); + ExpectClose("device y (init states)", dev.y, ref.y, atol, 5e-3); + ExpectClose("device final_states (init states)", dev.final_states, ref.final_states, atol, + 5e-3); + ExpectDeviceMatchesHost("y (init) device vs host", dev.y, host.y, maxlen); + ExpectDeviceMatchesHost("final_states (init) device vs host", dev.final_states, + host.final_states, maxlen); + } + } + } +} + +// The optional arms and the dtype knobs, on device: D as [H] and [H,P], the z +// silu gate, dt_bias + dt_softplus, the dt_limit clamp, a bf16 activation stream +// and a bf16 SSM state. +TEST_CASE("mamba2 chunk scan CUDA arm covers the optional arms and the dtype knobs") { + Backend* gpu = MaybeCuda(); + if (gpu == nullptr) { + MESSAGE("SKIP: no CUDA backend registered (CPU-only build/box)"); + return; + } + const int64_t T = 100, H = 8, P = 12, G = 2, N = 16, chunk = 32; + const Inputs in = GenerateInputs(T, H, P, G, N, 0xD00D42u); + const std::vector cu{0, 40, static_cast(T)}; + const int64_t maxlen = 60; // the longer of the two sequences + + std::mt19937 rng(7u); + std::normal_distribution nd(0.0f, 1.0f); + std::uniform_real_distribution ud(0.0f, 1.0f); + std::vector d_head_scalar(static_cast(H)); + for (auto& v : d_head_scalar) v = nd(rng); + std::vector d_hdim(static_cast(H * P)); + for (auto& v : d_hdim) v = nd(rng); + std::vector z(static_cast(T * H * P)); + for (auto& v : z) v = nd(rng); + std::vector dt_bias(static_cast(H)); + for (auto& v : dt_bias) v = ud(rng) - 4.0f; + + SUBCASE("D as [H]") { + RunCfg cfg; + cfg.chunk_size = chunk; + const SeqRefOut ref = SequentialSsdRef(in, T, H, P, G, N, cu, &d_head_scalar, false, nullptr, + nullptr, nullptr, {}); + const RunOut dev = RunChunkScanCuda(*gpu, in, T, H, P, G, N, cu, &d_head_scalar, nullptr, + nullptr, nullptr, cfg); + RequireNativeCudaProvider(vt::OpId::kMamba2ChunkScan, "D as [H]"); + const RunOut host = + RunChunkScan(in, T, H, P, G, N, cu, &d_head_scalar, nullptr, nullptr, nullptr, cfg); + ExpectClose("device y", dev.y, ref.y, 5e-3, 5e-3); + ExpectDeviceMatchesHost("y device vs host", dev.y, host.y, maxlen); + } + SUBCASE("D as [H,P]") { + RunCfg cfg; + cfg.chunk_size = chunk; + cfg.d_has_hdim = true; + const SeqRefOut ref = + SequentialSsdRef(in, T, H, P, G, N, cu, &d_hdim, true, nullptr, nullptr, nullptr, {}); + const RunOut dev = + RunChunkScanCuda(*gpu, in, T, H, P, G, N, cu, &d_hdim, nullptr, nullptr, nullptr, cfg); + const RunOut host = + RunChunkScan(in, T, H, P, G, N, cu, &d_hdim, nullptr, nullptr, nullptr, cfg); + ExpectClose("device y", dev.y, ref.y, 5e-3, 5e-3); + ExpectDeviceMatchesHost("y device vs host", dev.y, host.y, maxlen); + } + SUBCASE("z silu gate") { + RunCfg cfg; + cfg.chunk_size = chunk; + const SeqRefOut ref = + SequentialSsdRef(in, T, H, P, G, N, cu, nullptr, false, &z, nullptr, nullptr, {}); + const RunOut dev = + RunChunkScanCuda(*gpu, in, T, H, P, G, N, cu, nullptr, &z, nullptr, nullptr, cfg); + const RunOut host = RunChunkScan(in, T, H, P, G, N, cu, nullptr, &z, nullptr, nullptr, cfg); + ExpectClose("device y", dev.y, ref.y, 5e-3, 5e-3); + ExpectDeviceMatchesHost("y device vs host", dev.y, host.y, maxlen); + } + SUBCASE("dt_bias + dt_softplus, then the dt_limit clamp") { + RefCfg rc; + rc.dt_softplus = true; + RunCfg cfg; + cfg.chunk_size = chunk; + cfg.dt_softplus = true; + { + const SeqRefOut ref = + SequentialSsdRef(in, T, H, P, G, N, cu, nullptr, false, nullptr, &dt_bias, nullptr, rc); + const RunOut dev = + RunChunkScanCuda(*gpu, in, T, H, P, G, N, cu, nullptr, nullptr, &dt_bias, nullptr, cfg); + const RunOut host = + RunChunkScan(in, T, H, P, G, N, cu, nullptr, nullptr, &dt_bias, nullptr, cfg); + ExpectClose("device y", dev.y, ref.y, 5e-3, 5e-3); + ExpectClose("device final_states", dev.final_states, ref.final_states, 5e-3, 5e-3); + ExpectDeviceMatchesHost("y device vs host", dev.y, host.y, maxlen); + } + RefCfg clamped = rc; + clamped.dt_min = 0.05; + clamped.dt_max = 0.10; + RunCfg ccfg = cfg; + ccfg.dt_min = 0.05f; + ccfg.dt_max = 0.10f; + const SeqRefOut cref = SequentialSsdRef(in, T, H, P, G, N, cu, nullptr, false, nullptr, + &dt_bias, nullptr, clamped); + const RunOut cdev = + RunChunkScanCuda(*gpu, in, T, H, P, G, N, cu, nullptr, nullptr, &dt_bias, nullptr, ccfg); + ExpectClose("device y (clamped)", cdev.y, cref.y, 5e-3, 5e-3); + // The clamp must actually BITE, or the comparison is trivial. + const SeqRefOut loose = + SequentialSsdRef(in, T, H, P, G, N, cu, nullptr, false, nullptr, &dt_bias, nullptr, rc); + double max_diff = 0.0; + for (size_t i = 0; i < loose.y.size(); ++i) + max_diff = std::max(max_diff, std::abs(loose.y[i] - cref.y[i])); + CHECK(max_diff > 1e-2); + } + SUBCASE("bf16 activations") { + Inputs bin = in; + RoundInputsTo(bin, DType::kBF16); + RunCfg cfg; + cfg.chunk_size = chunk; + cfg.act_dtype = DType::kBF16; + const SeqRefOut bref = + SequentialSsdRef(bin, T, H, P, G, N, cu, nullptr, false, nullptr, nullptr, nullptr, {}); + const RunOut dev = + RunChunkScanCuda(*gpu, bin, T, H, P, G, N, cu, nullptr, nullptr, nullptr, nullptr, cfg); + ExpectClose("device y bf16", dev.y, bref.y, 5e-2, 5e-2); + } + // `state_dtype` is a SEPARATE knob (ssd_combined.py:46,119,176), and it moves + // `out` as well as `final_states` because `_chunk_scan_fwd` reads the stored + // copy back (:249-250, :266-269). The device arm allocates its inter-chunk + // buffer at that width — NOT at the host reference's f32 working width (§8.2 F9) + // — so the two arms must agree on BOTH outputs. + SUBCASE("bf16 SSM state with f32 activations") { + const int64_t T2 = 96, H2 = 4, P2 = 8, G2 = 2, N2 = 16, chunk2 = 16; + const Inputs in2 = GenerateInputs(T2, H2, P2, G2, N2, 0x51A7Eu); + const std::vector cu2{0, static_cast(T2)}; + RunCfg cfg; + cfg.chunk_size = chunk2; + cfg.state_dtype = DType::kBF16; + const RunOut dev = RunChunkScanCuda(*gpu, in2, T2, H2, P2, G2, N2, cu2, nullptr, nullptr, + nullptr, nullptr, cfg); + RequireNativeCudaProvider(vt::OpId::kMamba2ChunkScan, "bf16 SSM state"); + const RunOut host = + RunChunkScan(in2, T2, H2, P2, G2, N2, cu2, nullptr, nullptr, nullptr, nullptr, cfg); + for (float v : dev.final_states) CHECK(vt::BF16ToF32(vt::F32ToBF16(v)) == v); + ExpectDeviceMatchesHost("y (bf16 state) device vs host", dev.y, host.y, T2); + ExpectDeviceMatchesHost("final_states (bf16 state) device vs host", dev.final_states, + host.final_states, T2); + // ... and the bf16 state really does move `out` on the device arm too. + RunCfg f32cfg = cfg; + f32cfg.state_dtype = DType::kF32; + const RunOut f32dev = RunChunkScanCuda(*gpu, in2, T2, H2, P2, G2, N2, cu2, nullptr, nullptr, + nullptr, nullptr, f32cfg); + double max_out_diff = 0.0; + for (size_t i = 0; i < dev.y.size(); ++i) + max_out_diff = + std::max(max_out_diff, std::abs(static_cast(dev.y[i]) - f32dev.y[i])); + INFO("device max|out(bf16 state) - out(f32 state)| = " << max_out_diff); + CHECK(max_out_diff > 1e-6); + } +} + +#endif // VLLM_CPP_CUDA diff --git a/tests/vt/test_ops_mamba2_state_update.cpp b/tests/vt/test_ops_mamba2_state_update.cpp index 81a2e9595..12ed32771 100644 --- a/tests/vt/test_ops_mamba2_state_update.cpp +++ b/tests/vt/test_ops_mamba2_state_update.cpp @@ -643,6 +643,50 @@ TEST_CASE("mamba2 state update refuses the arms it does not implement") { CHECK(msg.find("extra_groups_for_head_shards") != std::string::npos); } + // ── the precondition the decay rests on ─────────────────────────────────── + // `A = -exp(A_log)` (mamba_mixer2.py:456) is strictly negative for every finite + // `A_log`, and the decode step's `exp(A[h]*dt)` is a DECAY only while it is. + // Fed `A > 0` the step returns a silently GROWING recurrence where the caller + // asked for a decaying one, so the op refuses — `CheckMamba2ANegative`, + // src/vt/cpu/cpu_ops.cpp, the decode twin of the chunk-scan guard this case + // mirrors (test_ops_mamba2_ssd.cpp:900). + // + // WHY THIS SUBCASE EXISTS (mamba2-ssd.md §8.2, round-2 review correction). The + // guard was present, correct and reachable, and pinned by NOTHING: deleting the + // call left this whole suite green, while the same deletion on the chunk-scan + // twin reds. A guard no test can see is a guard the next refactor removes. + SUBCASE("A must be negative (A = -exp(A_log))") { + StepCfg cfg; + StepInputs bad = in; + bad.A[static_cast(H - 1)] = 1.0f; + std::vector raw(static_cast(Nb * H * P * N) * 4, 0); + bool threw = false; + std::string msg; + try { + RunStateUpdate(raw, Nb, bad, nullptr, Nb, H, P, G, N, cfg); + } catch (const std::exception& e) { + threw = true; + msg = e.what(); + } + INFO(msg); + CHECK(threw); + CHECK(msg.find("A_log") != std::string::npos); + + // Zero is refused too: `-exp(A_log)` is strictly negative for every finite + // A_log, and a flat (non-decaying) state is not what the layer asked for. + StepInputs zero = in; + zero.A[0] = 0.0f; + std::vector raw_zero(static_cast(Nb * H * P * N) * 4, 0); + CHECK_THROWS(RunStateUpdate(raw_zero, Nb, zero, nullptr, Nb, H, P, G, N, cfg)); + + // ... and a legitimately tiny negative A is ACCEPTED, so the guard is a + // sign test and not an accidental magnitude floor. + StepInputs tiny = in; + for (auto& v : tiny.A) v = -1e-30f; + std::vector raw_tiny(static_cast(Nb * H * P * N) * 4, 0); + CHECK_NOTHROW(RunStateUpdate(raw_tiny, Nb, tiny, nullptr, Nb, H, P, G, N, cfg)); + } + SUBCASE("a compact state must have one row per token") { StepCfg cfg; std::vector raw(static_cast((Nb + 1) * H * P * N) * 4, 0); @@ -682,3 +726,500 @@ TEST_CASE("mamba2 state update refuses the arms it does not implement") { CHECK_NOTHROW(RunStateUpdate(raw_null, S, in, &nulls, Nb, H, P, G, N, cfg)); } } + +// ═════════════════════════════════════════════════════════════════════════════ +// (6) THE CUDA ARM — .agents/specs/mamba2-ssd.md W2, issue #496. +// The declared equivalence contract is stated in full at the head of the CUDA +// section of tests/vt/test_ops_mamba2_ssd.cpp and in +// src/vt/cuda/cuda_mamba2_ssd.cuh: f32 accumulation throughout, no tile +// downcasts, identical memory format, and the same per-element accumulation +// order — leaving the libm difference as the only admitted source of divergence, +// which `DerivedRtol` bounds without a tuned number. +// +// The decode step's evidence is mostly EXACT rather than approximate, and +// deliberately so: which cache slot was written, which was left alone, and what a +// NULL row produced are byte facts, not tolerances. +// ═════════════════════════════════════════════════════════════════════════════ +#ifdef VLLM_CPP_CUDA + +#include +#include + +namespace { + +using vt::Backend; + +Backend* MaybeCuda() { + try { + return &vt::GetBackend(DeviceType::kCUDA); + } catch (const std::exception&) { + return nullptr; + } +} + +// A GREEN TEST DOES NOT PROVE THE DEVICE RAN IT. GB10 is +// `integrated && pageable_memory_access`, so `Backend::UnifiedMemory()` is TRUE +// and `ReferenceTierEligible(kCUDA)` with it: absent a native kernel, `GetOp` +// does not throw — it installs the CPU HOST kernel as a `kReferenceProviderName` +// provider and runs THAT over the device pointers (op_provider.h, "portable +// reference tier"), so every assertion below would pass while nothing ran on the +// GPU. Every CUDA case therefore asserts the SELECTED provider is native. These +// are EAGER dispatches, so the counters are populated +// ([[graph-replay-does-no-host-dispatch-counters-read-zero]]). +void RequireNativeCudaProvider(vt::OpId op, const std::string& what) { + const vt::OpProviderStats st = vt::GetOpProviderStats(op, DeviceType::kCUDA); + INFO(what << ": selected CUDA provider = " + << (st.last_selected != nullptr ? st.last_selected : "") + << "; process-wide reference-tier hits = " << vt::GetReferenceTierHits()); + REQUIRE(st.last_selected != nullptr); + CHECK(std::string(st.last_selected) != std::string(vt::kReferenceProviderName)); +} + +// `4*(K + 2)*u` — the bound derived at the head of the CUDA section of +// tests/vt/test_ops_mamba2_ssd.cpp: 2.5 ulp of glibc-vs-CUDA libm disagreement +// per decay factor through a product of at most K, plus the standard (K-1)*u +// forward error of a length-K f32 summation. +constexpr double kUnitRoundoff = 5.9604644775390625e-08; // 2^-24 +double DerivedRtol(int64_t K) { return 4.0 * static_cast(K + 2) * kUnitRoundoff; } + +void ExpectDeviceMatchesHost(const std::string& what, const std::vector& dev, + const std::vector& host, int64_t K) { + REQUIRE(dev.size() == host.size()); + REQUIRE(!dev.empty()); + double scale = 0.0; + for (float v : host) scale = std::max(scale, std::abs(static_cast(v))); + const double rtol = DerivedRtol(K); + const double atol = rtol * scale; + size_t bit_differing = 0, worst_i = 0; + double worst_ratio = -1.0, worst_diff = 0.0; + for (size_t i = 0; i < dev.size(); ++i) { + if (dev[i] != host[i]) ++bit_differing; + const double d = std::abs(static_cast(dev[i]) - static_cast(host[i])); + const double budget = atol + rtol * std::abs(static_cast(host[i])); + const double ratio = budget > 0.0 ? d / budget : (d > 0.0 ? 1e30 : 0.0); + if (!std::isfinite(static_cast(dev[i])) || ratio > worst_ratio) { + worst_ratio = ratio; + worst_i = i; + worst_diff = d; + if (!std::isfinite(static_cast(dev[i]))) break; + } + } + // MESSAGE, not INFO: doctest prints an INFO context only when an assertion in + // its scope FAILS, so the used slack has to be logged unconditionally for the + // derived bar to be auditable on the green run that matters. + MESSAGE(what << ": K=" << K << " rtol=" << rtol << " scale=" << scale << "; " + << bit_differing << " of " << dev.size() + << " elements differ in any bit; worst element [" << worst_i + << "] dev=" << dev[worst_i] << " host=" << host[worst_i] + << " |diff|=" << worst_diff << " used " << (worst_ratio * 100.0) + << "% of its derived budget"); + CHECK(std::isfinite(static_cast(dev[worst_i]))); + CHECK(worst_ratio <= 1.0); +} + +Tensor MakeTDev(void* data, DType dt, Device dev, const std::vector& shape) { + Tensor t; + t.data = data; + t.dtype = dt; + t.device = dev; + t.rank = static_cast(shape.size()); + int64_t stride = 1; + for (int i = t.rank - 1; i >= 0; --i) { + t.shape[i] = shape[static_cast(i)]; + t.stride[i] = stride; + stride *= shape[static_cast(i)]; + } + return t; +} + +class DBuf { + public: + DBuf(Backend& b, Queue& q, const void* host, size_t bytes) : b_(&b), bytes_(bytes) { + p_ = b.Alloc(bytes == 0 ? 1 : bytes); + if (host != nullptr && bytes > 0) b.Copy(q, p_, host, bytes); + } + ~DBuf() { + if (p_ != nullptr) b_->Free(p_); + } + DBuf(const DBuf&) = delete; + DBuf& operator=(const DBuf&) = delete; + void* get() const { return p_; } + void Download(Queue& q, void* dst) const { + if (bytes_ > 0) b_->Copy(q, dst, p_, bytes_); + b_->Synchronize(q); + } + + private: + Backend* b_; + void* p_ = nullptr; + size_t bytes_ = 0; +}; + +// The CUDA twin of RunStateUpdate. `state_raw` is updated IN PLACE, as on host. +std::vector RunStateUpdateCuda(Backend& gpu, std::vector& state_raw, int64_t S, + const StepInputs& in, const std::vector* slots, + int64_t Nb, int64_t H, int64_t P, int64_t G, int64_t N, + const StepCfg& cfg) { + Queue q = gpu.CreateQueue(); + const Device dev{DeviceType::kCUDA, 0}; + const std::vector xb = Pack(in.x, cfg.act_dtype); + const std::vector dtb = Pack(in.dt, cfg.act_dtype); + const std::vector Bb = Pack(in.B, cfg.act_dtype); + const std::vector Cb = Pack(in.C, cfg.act_dtype); + const std::vector zb = Pack(in.z, cfg.act_dtype); + const size_t out_bytes = static_cast(Nb * H * P) * vt::SizeOf(cfg.act_dtype); + + DBuf dstate(gpu, q, state_raw.data(), state_raw.size()); + DBuf dx(gpu, q, xb.data(), xb.size()); + DBuf ddt(gpu, q, dtb.data(), dtb.size()); + DBuf dA(gpu, q, in.A.data(), in.A.size() * sizeof(float)); + DBuf dB(gpu, q, Bb.data(), Bb.size()); + DBuf dC(gpu, q, Cb.data(), Cb.size()); + DBuf dD(gpu, q, in.D.data(), in.D.size() * sizeof(float)); + DBuf dz(gpu, q, zb.data(), zb.size()); + DBuf ddb(gpu, q, in.dt_bias.data(), in.dt_bias.size() * sizeof(float)); + DBuf dout(gpu, q, nullptr, out_bytes); + + Tensor st = MakeTDev(dstate.get(), cfg.state_dtype, dev, {S, H, P, N}); + Tensor xt = MakeTDev(dx.get(), cfg.act_dtype, dev, {Nb, H, P}); + Tensor dtt = MakeTDev(ddt.get(), cfg.act_dtype, dev, {Nb, H}); + Tensor At = MakeTDev(dA.get(), DType::kF32, dev, {H}); + Tensor Bt = MakeTDev(dB.get(), cfg.act_dtype, dev, {Nb, G, N}); + Tensor Ct = MakeTDev(dC.get(), cfg.act_dtype, dev, {Nb, G, N}); + Tensor Dt = MakeTDev(dD.get(), DType::kF32, dev, {H}); + Tensor zt = MakeTDev(dz.get(), cfg.act_dtype, dev, {Nb, H, P}); + Tensor dbt = MakeTDev(ddb.get(), DType::kF32, dev, {H}); + Tensor outt = MakeTDev(dout.get(), cfg.act_dtype, dev, {Nb, H, P}); + + std::vector idx; + std::unique_ptr didx; + Tensor idxt; + if (slots != nullptr) { + idx = *slots; + didx = std::make_unique(gpu, q, idx.data(), idx.size() * sizeof(int32_t)); + idxt = MakeTDev(didx->get(), DType::kI32, dev, {Nb}); + } + + Mamba2Args args; + args.dt_softplus = cfg.dt_softplus; + args.tp_world_size = cfg.tp_world_size; + vt::Mamba2StateUpdate(q, outt, st, xt, dtt, At, Bt, Ct, cfg.use_D ? &Dt : nullptr, + cfg.use_z ? &zt : nullptr, cfg.use_dt_bias ? &dbt : nullptr, + slots != nullptr ? &idxt : nullptr, args); + + std::vector outb(out_bytes); + dout.Download(q, outb.data()); + dstate.Download(q, state_raw.data()); + gpu.Synchronize(q); + gpu.DestroyQueue(q); + return Unpack(outb, static_cast(Nb * H * P), cfg.act_dtype); +} + +} // namespace + +TEST_CASE("mamba2 state update CUDA arm matches the reference recurrence") { + Backend* gpu = MaybeCuda(); + if (gpu == nullptr) { + MESSAGE("SKIP: no CUDA backend registered (CPU-only build/box)"); + return; + } + // The driver shapes: nheads 64, headdim 64, dstate 128, ngroups 8 + // (mamba2-ssd.md §1.4), plus the small sweep the host arm runs. + struct Shape { + int64_t Nb, H, P, G, N; + }; + const std::vector shapes{{4, 64, 64, 8, 128}, {3, 16, 64, 1, 16}, {3, 16, 64, 4, 64}}; + for (const Shape& sh : shapes) { + for (bool has_z : {false, true}) { + const StepInputs in = + GenerateStep(sh.Nb, sh.H, sh.P, sh.G, sh.N, 0x51A7Eu + static_cast(sh.N)); + std::mt19937 rng(9u); + std::normal_distribution nd(0.0f, 1.0f); + std::vector state0(static_cast(sh.Nb * sh.H * sh.P * sh.N)); + for (auto& v : state0) v = nd(rng); + const std::vector state0d(state0.begin(), state0.end()); + + StepCfg cfg; + cfg.use_z = has_z; + const RefStep ref = + SelectiveStateUpdateRef(state0d, in.x, in.dt, in.A, in.B, in.C, &in.D, + has_z ? &in.z : nullptr, &in.dt_bias, true, sh.Nb, sh.H, sh.P, + sh.G, sh.N); + INFO("Nb=" << sh.Nb << " H=" << sh.H << " P=" << sh.P << " G=" << sh.G << " N=" << sh.N + << " has_z=" << has_z); + + std::vector raw_dev = Pack(state0, DType::kF32); + const std::vector out_dev = + RunStateUpdateCuda(*gpu, raw_dev, sh.Nb, in, nullptr, sh.Nb, sh.H, sh.P, sh.G, sh.N, + cfg); + RequireNativeCudaProvider(vt::OpId::kMamba2StateUpdate, "plain decode step"); + std::vector raw_host = Pack(state0, DType::kF32); + const std::vector out_host = + RunStateUpdate(raw_host, sh.Nb, in, nullptr, sh.Nb, sh.H, sh.P, sh.G, sh.N, cfg); + + // G1 — the device arm against the INDEPENDENT double reference, at + // upstream's own f32 tolerance (rtol 3e-4 / atol 1e-3, test_mamba_ssm.py:828). + ExpectClose("device out", out_dev, ref.out, 1e-3, 3e-4); + ExpectClose("device state", Unpack(raw_dev, state0.size(), DType::kF32), ref.state, 1e-3, + 3e-4); + // G2 — device vs host. A decode step runs exactly ONE decay factor and one + // length-N summation over dstate, so K is N and the bound is at its + // tightest anywhere in this row. + ExpectDeviceMatchesHost("out device vs host", out_dev, out_host, sh.N); + ExpectDeviceMatchesHost("state device vs host", + Unpack(raw_dev, state0.size(), DType::kF32), + Unpack(raw_host, state0.size(), DType::kF32), sh.N); + } + } +} + +// SCATTERED CACHE SLOTS and the NULL row, on device. This is the case that pins +// `state_indices` being honoured at all: a kernel that wrote through the ROW +// index instead of the SLOT index, or that treated the NULL row as slot 0, passes +// nothing here. The evidence is BYTE-EXACT — no tolerance can hide a slot that +// was written when it should not have been. +TEST_CASE("mamba2 state update CUDA arm honours scattered slots and the NULL row") { + Backend* gpu = MaybeCuda(); + if (gpu == nullptr) { + MESSAGE("SKIP: no CUDA backend registered (CPU-only build/box)"); + return; + } + const int64_t H = 8, P = 32, N = 16, G = 2; + const int64_t real = 3, padding = 5, Nb = real + padding; + const int64_t S = 30; + + const StepInputs in = GenerateStep(Nb, H, P, G, N, 0xFACEu); + std::mt19937 rng(11u); + std::normal_distribution nd(0.0f, 1.0f); + std::vector cache(static_cast(S * H * P * N)); + for (auto& v : cache) v = nd(rng); + + const std::vector chosen{17, 2, 25}; + std::vector slots = chosen; + for (int64_t i = 0; i < padding; ++i) slots.push_back(-1); + + std::vector state0d(static_cast(real * H * P * N)); + for (int64_t b = 0; b < real; ++b) { + const size_t src = + static_cast(chosen[static_cast(b)]) * static_cast(H * P * N); + for (int64_t i = 0; i < H * P * N; ++i) + state0d[static_cast(b * H * P * N + i)] = cache[src + static_cast(i)]; + } + StepInputs real_in = in; + real_in.x.resize(static_cast(real * H * P)); + real_in.dt.resize(static_cast(real * H)); + real_in.B.resize(static_cast(real * G * N)); + real_in.C.resize(static_cast(real * G * N)); + real_in.z.resize(static_cast(real * H * P)); + const RefStep ref = + SelectiveStateUpdateRef(state0d, real_in.x, real_in.dt, real_in.A, real_in.B, real_in.C, + &real_in.D, &real_in.z, &real_in.dt_bias, true, real, H, P, G, N); + + StepCfg cfg; + cfg.use_z = true; + std::vector raw = Pack(cache, DType::kF32); + const std::vector before = raw; + const std::vector out = RunStateUpdateCuda(*gpu, raw, S, in, &slots, Nb, H, P, G, N, cfg); + RequireNativeCudaProvider(vt::OpId::kMamba2StateUpdate, "scattered slots"); + const std::vector after = Unpack(raw, cache.size(), DType::kF32); + + // (a) the three selected slots hold the advanced state. + std::vector got_state(static_cast(real * H * P * N)); + for (int64_t b = 0; b < real; ++b) { + const size_t src = + static_cast(chosen[static_cast(b)]) * static_cast(H * P * N); + for (int64_t i = 0; i < H * P * N; ++i) + got_state[static_cast(b * H * P * N + i)] = after[src + static_cast(i)]; + } + ExpectClose("device scattered state", got_state, ref.state, 1e-3, 3e-4); + + // (b) the three real output rows match. + std::vector got_out(out.begin(), out.begin() + static_cast(real * H * P)); + ExpectClose("device scattered out", got_out, ref.out, 1e-3, 3e-4); + + // (c) EVERY slot that was not selected is BYTE-identical — upstream's + // `torch.equal(state_before[unused], state[unused])` (test_mamba_ssm.py:800). + // Slot 0 is among them, so a kernel that mapped the NULL row onto slot 0 + // fails here and nowhere else. + size_t untouched_slots = 0; + for (int64_t s = 0; s < S; ++s) { + if (std::find(chosen.begin(), chosen.end(), static_cast(s)) != chosen.end()) continue; + ++untouched_slots; + const size_t off = static_cast(s) * static_cast(H * P * N) * 4; + CHECK(std::memcmp(before.data() + off, raw.data() + off, + static_cast(H * P * N) * 4) == 0); + } + CHECK(untouched_slots == static_cast(S - real)); + + // (d) the NULL rows write an EXACTLY zeroed output row. + for (int64_t b = real; b < Nb; ++b) + for (int64_t i = 0; i < H * P; ++i) + CHECK(out[static_cast(b * H * P + i)] == 0.0f); +} + +// DECODE == PREFILL, both arms on device. The two go through entirely different +// device code (a single-token step vs the 5-stage chunked pipeline), which is +// what makes the agreement evidence rather than a shared-helper tautology +// ([[gate-comparing-shared-helper-proves-consistency-not-correctness]]). +TEST_CASE("mamba2 state update CUDA arm reproduces the CUDA chunked prefill") { + Backend* gpu = MaybeCuda(); + if (gpu == nullptr) { + MESSAGE("SKIP: no CUDA backend registered (CPU-only build/box)"); + return; + } + const int64_t T = 40, H = 4, P = 8, G = 2, N = 16, chunk = 8; + std::mt19937 rng(0x1234u); + std::normal_distribution nd(0.0f, 1.0f); + std::uniform_real_distribution ud(0.0f, 1.0f); + + std::vector A(static_cast(H)); + for (auto& v : A) v = -ud(rng) - 1.0f; + std::vector D(static_cast(H)); + for (auto& v : D) v = nd(rng); + std::vector dt_bias(static_cast(H)); + for (auto& v : dt_bias) v = ud(rng) - 4.0f; + std::vector x(static_cast(T * H * P)); + for (auto& v : x) v = nd(rng); + std::vector dt(static_cast(T * H)); + for (auto& v : dt) v = nd(rng); + std::vector B(static_cast(T * G * N)); + for (auto& v : B) v = nd(rng); + std::vector C(static_cast(T * G * N)); + for (auto& v : C) v = nd(rng); + + // ── prefill arm, on device ── + std::vector y_prefill(static_cast(T * H * P), 0.0f); + std::vector fs(static_cast(H * P * N), 0.0f); + { + Queue q = gpu->CreateQueue(); + const Device dev{DeviceType::kCUDA, 0}; + std::vector cu{0, static_cast(T)}; + std::vector ccs{0}; + std::vector sidx; + for (int64_t pos = 0; pos < T; pos += chunk) { + ccs.push_back(static_cast(std::min(pos + chunk, T))); + sidx.push_back(0); + } + std::vector lci{static_cast(sidx.size()) - 1}; + + DBuf dx(*gpu, q, x.data(), x.size() * sizeof(float)); + DBuf ddt(*gpu, q, dt.data(), dt.size() * sizeof(float)); + DBuf dA(*gpu, q, A.data(), A.size() * sizeof(float)); + DBuf dB(*gpu, q, B.data(), B.size() * sizeof(float)); + DBuf dC(*gpu, q, C.data(), C.size() * sizeof(float)); + DBuf dD(*gpu, q, D.data(), D.size() * sizeof(float)); + DBuf ddb(*gpu, q, dt_bias.data(), dt_bias.size() * sizeof(float)); + DBuf dout(*gpu, q, nullptr, y_prefill.size() * sizeof(float)); + DBuf dfs(*gpu, q, nullptr, fs.size() * sizeof(float)); + DBuf dcu(*gpu, q, cu.data(), cu.size() * sizeof(int32_t)); + DBuf dccs(*gpu, q, ccs.data(), ccs.size() * sizeof(int32_t)); + DBuf dlci(*gpu, q, lci.data(), lci.size() * sizeof(int32_t)); + DBuf dsi(*gpu, q, sidx.data(), sidx.size() * sizeof(int32_t)); + + Tensor xt = MakeTDev(dx.get(), DType::kF32, dev, {T, H, P}); + Tensor dtt = MakeTDev(ddt.get(), DType::kF32, dev, {T, H}); + Tensor At = MakeTDev(dA.get(), DType::kF32, dev, {H}); + Tensor Bt = MakeTDev(dB.get(), DType::kF32, dev, {T, G, N}); + Tensor Ct = MakeTDev(dC.get(), DType::kF32, dev, {T, G, N}); + Tensor Dt = MakeTDev(dD.get(), DType::kF32, dev, {H}); + Tensor dbt = MakeTDev(ddb.get(), DType::kF32, dev, {H}); + Tensor outt = MakeTDev(dout.get(), DType::kF32, dev, {T, H, P}); + Tensor fst = MakeTDev(dfs.get(), DType::kF32, dev, {1, H, P, N}); + Tensor cust = MakeTDev(dcu.get(), DType::kI32, dev, {2}); + Tensor ccst = MakeTDev(dccs.get(), DType::kI32, dev, {static_cast(ccs.size())}); + Tensor lcit = MakeTDev(dlci.get(), DType::kI32, dev, {1}); + Tensor sit = MakeTDev(dsi.get(), DType::kI32, dev, {static_cast(sidx.size())}); + Mamba2Args args; + args.chunk_size = chunk; + args.dt_softplus = true; + vt::Mamba2ChunkScan(q, outt, fst, xt, dtt, At, Bt, Ct, &Dt, nullptr, &dbt, nullptr, cust, + ccst, lcit, sit, args); + dout.Download(q, y_prefill.data()); + dfs.Download(q, fs.data()); + gpu->Synchronize(q); + gpu->DestroyQueue(q); + } + RequireNativeCudaProvider(vt::OpId::kMamba2ChunkScan, "decode-vs-prefill: prefill arm"); + + // ── decode arm: T single-token device steps on one cache slot ── + std::vector raw(static_cast(H * P * N) * 4, 0); + std::vector y_decode(static_cast(T * H * P), 0.0f); + for (int64_t t = 0; t < T; ++t) { + StepInputs step; + step.A = A; + step.D = D; + step.dt_bias = dt_bias; + step.x.assign(x.begin() + static_cast(t * H * P), + x.begin() + static_cast((t + 1) * H * P)); + step.dt.assign(dt.begin() + static_cast(t * H), + dt.begin() + static_cast((t + 1) * H)); + step.B.assign(B.begin() + static_cast(t * G * N), + B.begin() + static_cast((t + 1) * G * N)); + step.C.assign(C.begin() + static_cast(t * G * N), + C.begin() + static_cast((t + 1) * G * N)); + step.z.assign(static_cast(H * P), 0.0f); + StepCfg cfg; + const std::vector o = + RunStateUpdateCuda(*gpu, raw, 1, step, nullptr, 1, H, P, G, N, cfg); + std::copy(o.begin(), o.end(), y_decode.begin() + static_cast(t * H * P)); + } + + RequireNativeCudaProvider(vt::OpId::kMamba2StateUpdate, "decode-vs-prefill: decode arm"); + ExpectCloseF("device decode vs device chunked prefill", y_decode, y_prefill, 5e-3, 5e-3); + ExpectCloseF("device final state", Unpack(raw, static_cast(H * P * N), DType::kF32), fs, + 5e-3, 5e-3); +} + +// The SSM cache dtype is its own knob (mamba_utils.py:73-81), and `out` does not +// depend on it at all: the readout uses the f32 state held in registers, not the +// value re-read from the cache (mamba_ssm.py:433,451). On device that is the +// EXACT, tolerance-free statement it is on host. +TEST_CASE("mamba2 state update CUDA arm reads out the f32 state, not the rounded cache") { + Backend* gpu = MaybeCuda(); + if (gpu == nullptr) { + MESSAGE("SKIP: no CUDA backend registered (CPU-only build/box)"); + return; + } + const int64_t Nb = 2, H = 4, P = 8, G = 2, N = 16; + const StepInputs in = GenerateStep(Nb, H, P, G, N, 0xDEC0DEu); + std::mt19937 rng(0x99u); + std::normal_distribution nd(0.0f, 1.0f); + std::vector state0(static_cast(Nb * H * P * N)); + for (auto& v : state0) v = vt::BF16ToF32(vt::F32ToBF16(nd(rng))); + + for (bool has_z : {false, true}) { + StepCfg f32cfg; + f32cfg.use_z = has_z; + StepCfg bf16cfg = f32cfg; + bf16cfg.state_dtype = DType::kBF16; + + std::vector raw_f32 = Pack(state0, DType::kF32); + const std::vector out_f32 = + RunStateUpdateCuda(*gpu, raw_f32, Nb, in, nullptr, Nb, H, P, G, N, f32cfg); + std::vector raw_bf16 = Pack(state0, DType::kBF16); + const std::vector out_bf16 = + RunStateUpdateCuda(*gpu, raw_bf16, Nb, in, nullptr, Nb, H, P, G, N, bf16cfg); + RequireNativeCudaProvider(vt::OpId::kMamba2StateUpdate, "cache-dtype independence"); + + REQUIRE(out_f32.size() == out_bf16.size()); + size_t differing = 0; + for (size_t i = 0; i < out_f32.size(); ++i) + if (out_f32[i] != out_bf16[i]) ++differing; + INFO("has_z=" << has_z << ": " << differing << " of " << out_f32.size() + << " device outputs move with the CACHE dtype"); + CHECK(differing == 0); + + // The cache rounding really is live, so the check above is not comparing two + // identical computations. + const std::vector after_f32 = Unpack(raw_f32, state0.size(), DType::kF32); + const std::vector after_bf16 = Unpack(raw_bf16, state0.size(), DType::kBF16); + for (float v : after_bf16) CHECK(vt::BF16ToF32(vt::F32ToBF16(v)) == v); + size_t rounded = 0; + for (size_t i = 0; i < after_f32.size(); ++i) + if (after_f32[i] != after_bf16[i]) ++rounded; + INFO("stored states the bf16 cache rounded: " << rounded); + REQUIRE(rounded > 0); + } +} + +#endif // VLLM_CPP_CUDA