diff --git a/.agents/specs/mamba2-ssd.md b/.agents/specs/mamba2-ssd.md index b33042335..a9bf103fd 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 @@ -412,9 +437,507 @@ reported `1: worst element ...` instead of naming the tensor. The labels are Named residuals unchanged from §2: the CUDA arm (W2), `n_groups` TP sharding, spec-decode temporal state, ReplaySSM and Mamba v1. One more, from the port itself: the Triton dots downcast their tile inputs (`ssd_chunk_state.py:283-285`, -`ssd_chunk_scan.py:266-269`) where this host reference stays f32, so W2's -device-vs-host comparison is a tolerance comparison at the activation dtype, not -a byte compare. +`ssd_chunk_scan.py:266-269`) where this host reference stays f32. + +**Corrected by §8.3, which supersedes this paragraph.** As originally written it +predicted that W2 would mirror those downcasts, and therefore that W2's +device-vs-host comparison would be "a tolerance comparison at the activation +dtype, not a byte compare". W2 did NOT mirror them — §8.3 point 1 records the +decision and why — so **both arms stay f32** and the dtype is not the reason a +byte compare is out of reach. The reasons are **libm and FMA contraction** +(§8.3 points 4 and 4b). Only the conclusion survived; its stated cause did not, +and two different causes for one fact in one spec is the drift this note closes. + +### 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: order is the *amplifying* source, and +pinning it holds the derived bound to the two terms named in points 4 and 4b. + +**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. §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. + +**4b. The second reason is FMA CONTRACTION, and the original derivation omitted +it.** Corrected here after a fresh review of PR #566 (finding F1, MEDIUM); +points 3 and 4 as first written claimed the elementary functions were the *only* +admitted source, and that was false as written. Host C++ is pinned +`-ffp-contract=off` (`CMakeLists.txt:41-56`) precisely so `a*b + c` keeps two +roundings. **Nothing passes `--fmad=false` to nvcc** — `grep -rn fmad +CMakeLists.txt cmake/` returns nothing — so `cuda_mamba2_ssd.cuh` compiles at +nvcc's **default `--fmad=true`**, and every `acc += a*b` in it is a +single-rounding `fma` whose host twin is not. This project has *measured* that +exact idiom: `.agents/benchmark-record.md:532` records a pre-rounded `v²` +differing by ≤ 1 ulp from the nvcc-`fmad` form and flipping a near-tie at token +108. `CMakeLists.txt:41-56` carves CUDA out of the contraction policy on the +grounds that "GPU parity tests compare GPU-vs-GPU"; **G2 is precisely the case +that carve-out does not cover**. + +Nothing was hidden empirically — the worst audited comparison used 7.66% of the +bound and the driver shapes 0.32% / 0.18% — so this is a **derivation-accuracy** +defect, not a numerical failure. It is repaired by carrying the term, not by +turning the flag off: + +- **`-fmad=false` on the TU: REJECTED.** nvcc takes the flag per *translation + unit*. `cuda_mamba2_ssd.cuh` is a **header**, included by `cuda_gdn.cu:48`, so + the only ways to apply it are (a) de-contract every GDN kernel in that TU — a + measured hot decode path ([[gdn-packed-bridge-closes-31pct-decode-gap]]) — or + (b) split a new `src/vt/` TU, which the *Placement* deviation below already + records as blocked on #515. Slowing a shipped kernel to make a bound's prose + true is the wrong trade. +- **Carrying the term: TAKEN.** The bound moves from `4·(K+2)·u` to + `5·(K+2)·u`; the arithmetic is in point 6. + +**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) = 5·(K+2)·2⁻²⁴`** +(`ExpectDeviceMatchesHost` / `DerivedRtol`, in all three suites). Three terms, +one per admitted source: + +| term | bound | source | +|---|---|---| +| libm | `2.5·K·u` | ≤ 2.5 ulp per decay factor (CUDA `expf` ≤ 2, glibc ≤ 0.5) through a product of at most `K` | +| summation | `(K-1)·u` | the standard forward error of a length-`K` f32 sum, which is what *amplifies* the libm difference | +| contraction | `K·u` | the `K` product roundings the host keeps under `-ffp-contract=off` and the device's `fma` does not (point 4b) | + +Total `≤ 2.5·K·u + (K-1)·u + K·u = 4.5·K·u − u`, and `5·(K+2)·u` covers it +for every `K ≥ 0` because `5K + 10 ≥ 4.5K − 1` reduces to `0.5K + 11 ≥ 0`. + +**The previous `4·(K+2)·u` did not, and that is why the constant moved.** +`4.5K − 1 ≤ 4K + 8` reduces to `K ≤ 18`, so the old bound was provable only up +to `K = 18` — while the driver-shapes case runs at `K = T = 200` +(`test_ops_mamba2_ssd.cpp`, §1.4 shapes). The constant changed because the +**derivation gained a term the build actually emits**, not because a run needed +slack: re-scaling §8.4's audit by `4/5`, the worst of the 55 comparisons goes +from 7.66% to **6.13%** of budget, the driver shapes from 0.32%/0.18% to +**0.26%/0.14%**, and mutant M3 from 962173% to **769738%** — still caught by +four orders of magnitude. **No number is tuned**; `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 re-checks NONE of the host arm's metadata VALUE checks.** + Corrected here after the fresh review of PR #566 (finding F2, MEDIUM), which + established both that the list was longer than "`A < 0` and `state_indices` + distinctness" and that the memory-safety claim attached to it was broader than + the kernels guarantee. + + The **shared** validator (`Mamba2ChunkScan` / `Mamba2StateUpdate`, + `src/vt/ops.cpp`) checks metadata **shape, dtype and device only** + (`CheckI32Meta`, `ops.cpp:1717-1723`). Every **value** check lives in the host + kernel, where the data is host-readable (`cpu_ops.cpp:1622-1648`). The device + arm therefore does not check, in full: + + 1. `A < 0` (`CheckMamba2ANegative`) — value-only + 2. `state_indices` distinctness (§8.2 F8) — value-only + 3. the `cu_chunk_seqlens` tiling of `[0,T)` — **memory-unsafe** + 4. per-chunk `0 < len ≤ chunk_size` — **memory-unsafe** + 5. `seq_idx[c] ∈ [0,S)` — **memory-unsafe**, now CLAMPED + 6. `0 ≤ last_chunk_indices[b] < nchunks` (§8.2 **F7**) — **memory-unsafe**, now CLAMPED + + The reason for dropping them is unchanged and still holds for the *reads*: + the operands are on-device, so re-reading them costs a D2H plus a stream + synchronise per call — the host tax the GDN prefill path was rebuilt to remove + — and makes the op uncapturable in a CUDA graph, mirroring the policy + `cuda_gdn.cu:8-13` states for exactly this case. + + **But that reason does not reach 5 and 6, and the review was right that the + file was internally inconsistent.** Both values are already **in device + registers** at their use sites, and the decode kernel has always clamped its + `state_indices` slot for free on exactly that basis. Unclamped, `lci[b] ≥ + nchunks` makes the `passed` store an out-of-bounds **write** past the + `cudaMallocAsync` allocation, and a `seq_idx[c] ∉ [0,S)` an out-of-bounds + **read** of `initial_states` (with `seq_idx[0] < 0` additionally indexing + `passed` at chunk −1). Both are now **clamped in registers**, which is free, + matches the decode kernel, and matters because W4 is about to become the first + caller of these ops. The clamps do **not** restore the checks: out-of-contract + metadata still yields a **wrong answer**, now with a defined shape ("the chunk + loop stops at `nchunks`", "that chunk opens with a zero previous state"). They + bound only *where* it is read from. Pinned by a device-only case + (`test_ops_mamba2_ssd.cpp`, "clamps out-of-contract metadata in registers") + against in-contract reference runs, so the assertions are exact. + + **3 and 4 are NOT clamped and the arm is NOT memory-safe under them** — a + garbage `ccs` yields a `start`/`len` that index `x`/`B`/`C`/`z`/`out` out of + bounds in every stage, and bounding that needs `T` at each use site plus a + clamp inside the inner loops, which is not free. That is stated as unsafe + rather than folded into a blanket "memory safe" claim. **Owed, not implemented + here:** route 1–4 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 | `n_groups = 1` in the kernel, launcher grid unchanged (see below) | 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. + +**M7's label is narrower than "whole-row variance", and the fresh review of +#566 was right to say so (finding F3, LOW).** The reformulated M7 forces +`n_groups = 1` in the *kernel* while the launcher's `nblocks = rows * +args.n_groups` is left alone, so blocks `blk ≥ rows` compute `r = blk / 1 ≥ rows` +and read and write past the tensor. The device mutant is therefore +**memory-unsafe**, and fails partly for that rather than purely on the +whole-row-variance arithmetic. The guarantee IS pinned — the reviewer ran a +clean CPU twin of the same mutation, with no grid mismatch, and it reds on the +intended assertion — so only the label was wrong, and it is corrected in the +table above. A device mutant that is also memory-unsafe is a weaker instrument +than one that is not, and is recorded as such rather than re-scored. + +**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 is the one difference, so it is NOT dismissed. What is established: **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 can see +is three extra registrations in the op table. What is NOT yet established is the +positive cause. + +One contention fact IS established, and it invalidates the "idle box" premise +this run was read under: an unrelated job (`~/work/marlin442`, files touched +03:14 → 04:20, spanning the whole 03:51-04:44 ctest window) serialises on +**`/tmp/gpu.lock`, not `$HOME/gpu.lock`** — a different file, so the shared +mutex did not exclude it and it was on the GPU while this suite ran. +`test_minimax_h3` loads a ~41 GB model on a box whose memory is UNIFIED, which +is precisely the documented OOM failure mode. That is a plausible cause, not a +proven one, and it is recorded as the former. + +A standalone serial re-run of all ten is queued and writes +`~/w2ssd/refail.log`; it is **PENDING on a named external resource** — +`$HOME/gpu.lock` has been held ~2.3 h by an unrelated benchmark series with +three jobs ahead of it. That re-run, not this paragraph, settles the +attribution, and it is owed before the fresh review closes. **`~/w2ssd` on the +gate host is deliberately left in place for it** (`rm -rf ~/w2ssd` once +`refail.log` is read). + +**Update, 06:50 CEST: `dgx.casa` went unreachable** (`No route to host`, ping +100% loss) while that re-run was still queued, so its state is +**`REMOTE_UNVERIFIED`** — unknown is neither absence nor success. The box has +OOM-rebooted before under exactly the unified-memory pressure described above, +which is suggestive and is *not* offered as proof of anything. Whoever picks +this up: `~/w2ssd/refail.log` either completed or did not, and re-running +`~/w2ssd/w2refail.sh` is cheap. Every result recorded above this line was +captured and read BEFORE the box went away. + +**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. + +These percentages are **against the bound as it stood at this run**, +`4·(K+2)·2⁻²⁴`. §8.5 corrects the derivation and the bound is now +`5·(K+2)·2⁻²⁴`, against which the same measurements read 6.13%, 0.26%/0.14% and +769738%. The numbers above are left as captured rather than restated, because +they are what the run produced; §8.3 point 6 carries the conversion. + +### 8.5 W2 tightening pass (fresh review of PR #566 returned PASS + 5 findings) + +`row/KERNEL-SSM-MAMBA-SSD-W2-FIX`, branched from `1e819144e` with `origin/main` +merged. The review's **verdict was PASS**; it verified the equivalence contract's +central MMA reasoning, reconstructed the reformulated mutations on CPU twins, +and confirmed the grep, the recovered-byte md5s and the CI baseline subtraction. +None of that was re-done. What changed: + +- **F1 (MEDIUM, derivation accuracy)** — repaired by carrying the FMA-contraction + term, not by disabling contraction. `-fmad=false` was rejected on the grounds + in §8.3 point 4b: it is a per-TU flag on a header included by `cuda_gdn.cu`, so + it would de-contract a measured hot decode path. `DerivedRtol` is + `5·(K+2)·2⁻²⁴` in all three suites; §8.3 point 6 carries the arithmetic and the + re-scaled audit. +- **F2 (MEDIUM, over-broad memory-safety claim)** — repaired by taking BOTH + halves the finding offered: the two free register-local clamps AND the narrowed + claim. The header and §8.3's second deviation now enumerate all six dropped + value checks, mark which are memory-unsafe, and state plainly that the + `cu_chunk_seqlens` tiling and length checks are **NOT** clamped and **NOT** + safe. `M2ChunkScanKernel` gained an `S` parameter for the `seq_idx` clamp. +- **F3, F4 (LOW, record accuracy)** — M7's label corrected in §8.4 with the grid + mismatch stated; §8.2's superseded downcast-tolerance sentence reconciled + against §8.3. +- **F5 (LOW)** — taken. The five per-call `cudaMallocAsync` scratch buffers are + held by an `M2Scratch` scope guard, so a throw on the Nth no longer leaks the + N-1 before it. `Release()` also frees all five before reporting, where the + open-coded sequence it replaces leaked the remainder on a mid-sequence failure. + +**Evidence, CPU box `promaxgb10` worktree host, `df -h /` = 85% used (67G free) +before and after every result below.** + +| suite | test cases | assertions | status | +|---|---|---|---| +| `test_ops_mamba2_ssd` | 8 / 8 passed | 1175 / 1175 | `SUCCESS!` | +| `test_ops_mamba2_state_update` | 6 / 6 passed | 2469 / 2469 | `SUCCESS!` | +| `test_ops_mamba2_gated_norm` | 9 / 9 passed | 2107 / 2107 | `SUCCESS!` | + +Identical to the pre-change CPU counts, as expected: every code change is inside +`#ifdef VLLM_CPP_CUDA` or in the `.cuh`. `Status:` was read, not `assertions:` +alone ([[doctest-assertions-line-hides-thrown-cases]]). + +Clean Release rebuild of every target — **817 / 817 ninja edges, 0 warnings, 0 +errors** at `-Wall -Wextra -Werror`, exit 0 — then **full `ctest -j 4`: +`100% tests passed, 0 tests failed out of 403`**, `CTEST_EXIT=0`, 22.35 s, +2 skipped (`test_modelopt_mixed_precision_checkpoint`, `test_voxtral_e2e`). This +is the CPU-only lane, so it is a much smaller and much faster gate than §8.4's +431-test GPU-host run and is **not** a substitute for it; none of §8.4's ten +failures is reachable here. `scripts/agent-preflight.sh --staged` returned +`RC=0`. + +**The CUDA arm was COMPILED by the operator on `dgx.casa`; it was not RUN.** +`dgx.casa` returned at **06:57 UTC** (it had rebooted; uptime was 0), which +supersedes §8.4's `REMOTE_UNVERIFIED` note on the host being away. The +implementer's box has no `nvcc` and no GPU, so the compile was **operator-run, +not implementer-run**, and is recorded as the operator's result rather than +folded into the implementer's evidence. The branch was transferred by +`git archive` — never `rsync`, which has previously overwritten goldens into +false passes — and built with real nvcc: + +``` +NVCC: cuda_13.0.r13.0 +CUTLASS found at ~/cutlass-4.5.0; enabling sm120a NVFP4 cutlass GEMM +Marlin NVFP4 W4A16 MoE GEMM enabled (vendored) for [121a] +FlashAttention-2 prefill/decode: ENABLED for arch(es) [121a] +CONFIGURE_EXIT=0 BUILD_EXIT=0 WARNINGS=0 ENOSPC=0 +``` + +Zero nvcc errors and **zero warnings under the project's `-Werror` flags**, disk +unchanged at 64G either side (so not the stale-binary false-green shape), and +the three fast-path features read **out of the configure log** rather than +assumed — an absent CUTLASS also exits 0, so "the build succeeded" alone would +have proved nothing ([[dgx-build-fast-path-verification]]). This closes the axis +the shim check below could only approximate: `cuda_mamba2_ssd.cuh` now compiles +through `cuda_gdn.cu` with real nvcc at the arch it ships on, which is what +actually retires the risk that the new `S` kernel parameter or the `M2Scratch` +guard broke the device build. + +**Attribute it to the operator, not to CI.** `cuda-fat-build` has still never +completed on this branch (see the cancellation note below), so no CI job has +compiled this code. + +The two substitutes below were run BEFORE that compile existed. They are kept +because they are what the implementer could establish unaided, and because the +second one is not superseded by any compile — but neither is offered as the +device gate: + +1. **`.cuh` compile + arity check.** The header was compiled by the host compiler + at `-std=c++20 -Wall -Wextra -Werror` against minimal CUDA shims, with each + `Kernel<<>>(args)` rewritten to `M2Sink(cfg), Kernel(args)` — which drops + the launch configuration while PRESERVING the argument-count and + argument-type check on all 7 launches. `EXIT 0`. Proved ARMED rather than + vacuous by deleting the `S` argument the F2 clamp added to the chunk-scan + launch on a scratch copy: `too few arguments to function M2ChunkScanKernel`, + exit 1. This checks C++, not PTX or device semantics. +2. **F2 clamp CPU twin.** The two index computations were transcribed and walked + over the device case's own shape (`S=H=nchunks=4`, `row=128`, `passed` + allocation 2048 elements), clamped and unclamped. Unclamped and out of + contract, every claimed hole reproduced: `lci[3] = nchunks` forms index + **2432** (an out-of-bounds WRITE), `lci[1] = -9` forms **−4096**, + `seq_idx >= S` reaches **463232** into a 2048-element `initial_states`, and + `seq_idx = {-1,…}` forms **−512** into `passed` — the `c == 0`, + `si == si_prev` hole. Clamped, all of them land in `[0, 1920] ⊂ [0, 2048)`, + the `lci` clamp reproduces the in-contract index range exactly, and both + `seq_idx` violations read no previous state at all — which is what the new + device case compares against an in-contract zero-init run. + +**Still `omitted_gates` — and the blocker is now the GPU LOCK, not the host.** +`dgx.casa` is back, so these are no longer waiting on a machine; they are +waiting on `$HOME/gpu.lock`, currently held by other coordinators' jobs with +8-hour timeouts. A host being reachable is not the same as the GPU being +available, and running any of these against a contended GPU would reproduce +exactly the undetected-contention defect §8.4 already records. Owed: + +1. the three suites' CUDA arms (Release and Debug) — **execution**, which the + operator's compile does not supply; +2. `compute-sanitizer memcheck` on the new "clamps out-of-contract metadata in + registers" case. This is the one that actually proves the F2 clamps: a green + run without it is necessary and not sufficient, because an out-of-bounds + write into a `cudaMallocAsync` pool commonly does not fault; +3. a re-run of the 9-mutation sweep against the moved bound `5·(K+2)·u` — the + §8.4 sweep was scored against `4·(K+2)·u`, and a widened bound is precisely + the change that could stop a mutation reddening; +4. §8.4's `~/w2ssd/refail.log` attribution, still `REMOTE_UNVERIFIED`. + +Item 3 is the one a reader is most likely to assume is safe. It is not assumed +here: the re-scaled margins in §8.3 point 6 make it very likely to hold, and +"very likely" is not a gate result. + +**CI on this branch is `REMOTE_UNVERIFIED`, and the distinction matters.** Every +GitHub Actions run for PR #592 -- all four SHAs, both the `ci` and `containers` +workflows -- ended `conclusion: cancelled`, never `failure` and never `success`. +`gh pr checks` renders a cancelled job as `fail`, so the PR reads as 16 reds +that are not reds; the run-level conclusion is what settles it. Three of the +four were cancelled by my own subsequent pushes, which is ordinary. The fourth +was not: run `31676434945` on the head SHA `f5c589f5e` sat queued from 07:06:49, +started at ~07:42, and was cancelled at **07:44:45-07:44:55** together with +**every other run in the repository** -- 20 of 20 across 7 branches, including +`row/LTX25-L9C-REGISTER-GATE`, `row/GATE-GPU-LOCK-WRAPPER` and +`row/FIX-HF-SNAPSHOT-ORDER-551`. `windows-msvc-cpu` had already passed two steps +when it was killed mid-build. A simultaneous repo-wide cancellation is an +Actions-side event, not a verdict on any diff, and it is recorded as unknown +rather than as either absence or success. A re-run was triggered at 07:47 and +was still `queued` at 07:50 when this was written; **whatever it reports, not +this paragraph, is the CI result.** + +**The rendering trap is worth carrying beyond this row, because two people hit +it independently from opposite ends.** `gh pr checks` prints a CANCELLED job as +`fail`. The implementer's watcher reported 16 failures on #592 and the +operator's reported 20, and the true count of failures in both was **zero**. A +per-check listing therefore cannot distinguish "this branch is red" from "the +whole runner pool was killed"; only the RUN-level `conclusion` field can +(`gh run view --json status,conclusion` → `cancelled`). Anyone subtracting +a CI baseline on this repo should read run conclusions, not check rows, or they +will attribute an infrastructure event to a diff. + +**No CI claim is made or repaired here.** In particular the two Windows jobs +remain the `main` baseline described in §8.4 (#512 now fixed by #583, #514, +#584) -- subtracted, not inherited -- and nothing above should be read as +evidence that this branch's CI passed. ## 9. Stop conditions 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..f7df33ae7 --- /dev/null +++ b/src/vt/cuda/cuda_mamba2_ssd.cuh @@ -0,0 +1,692 @@ +// 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. TWO effects are admitted, both structural and both carrying a +// DERIVED forward-error bound; neither is a tuned number. See +// tests/vt/test_ops_mamba2_ssd.cpp `DerivedRtol`. +// +// ─── ADMITTED SOURCE 1: THE ELEMENTARY FUNCTIONS ───────────────────────────── +// The two arms call different libms (`expf`/`log1pf`) — CUDA's `expf` is +// documented to <= 2 ulp and glibc's to <= 0.5. +// +// ─── ADMITTED SOURCE 2: FMA CONTRACTION ────────────────────────────────────── +// Host C++ is pinned `-ffp-contract=off` (CMakeLists.txt:41-56) precisely so +// `a*b + c` keeps two roundings. NOTHING passes `--fmad=false` to nvcc, so this +// header compiles at nvcc's DEFAULT `--fmad=true` and every `acc += a*b` below +// (:290, :335, :364, :434, :442, :501, and the gated norm's `part += v*v` at +// :554) is a SINGLE-rounding `fma` whose host twin is not. That is measured, not +// theoretical: .agents/benchmark-record.md:532 records a pre-rounded `v²` +// differing by <= 1 ulp from this exact nvcc-`fmad` idiom and flipping a +// near-tie. CMakeLists.txt:41-56 carves CUDA out of the contraction policy on +// the grounds that "GPU parity tests compare GPU-vs-GPU"; the device-vs-host +// comparison is exactly the case that carve-out does not cover, so the bound +// carries the term instead of the build removing it. +// +// `-fmad=false` was REJECTED, not overlooked. nvcc takes it per TRANSLATION +// UNIT and this is a header, included by cuda_gdn.cu:48 — so applying it means +// either de-contracting every GDN kernel in that TU (a measured hot decode +// path) or splitting a new `src/vt/` TU, which §8.3 already records as blocked +// on #515. Slowing an unrelated shipped kernel to make a bound's prose true is +// the wrong trade; widening the bound by the term the build actually emits is +// the right one. `DerivedRtol` is `5·(K+2)·u`, not `4·(K+2)·u`, and §8.3 shows +// the arithmetic. +// +// ─── 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 — order +// is the AMPLIFYING source and pinning it keeps the bound to the two terms +// above. It does not, on its own, make the libm the only one. +// +// ─── WHAT THIS ARM DOES NOT CHECK (named residual, §8.3) ───────────────────── +// The SHARED validator checks metadata SHAPE, DTYPE and DEVICE only +// (`CheckI32Meta`, ops.cpp:1717-1723). Every VALUE check lives in the host +// kernel, which reads the tensors (cpu_ops.cpp:1622-1648): `A < 0` +// (`CheckMamba2ANegative`), `state_indices` distinctness, the `cu_chunk_seqlens` +// tiling, per-chunk length bounds, `seq_idx[c] ∈ [0,S)`, and both halves of +// `0 <= last_chunk_indices[b] < nchunks`. This arm re-checks NONE of them: the +// operands live on the DEVICE, so reading them costs a D2H copy plus a stream +// synchronise per 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 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). Closing the gap needs the deferred device error ring +// cuda_ops.cu:790-940 already implements for embedding, and is owed. +// +// MEMORY SAFETY IS A NARROWER CLAIM AND IS MADE SEPARATELY, because dropping a +// check whose consequence is a WRONG NUMBER is correctness-grade while dropping +// one whose consequence is an out-of-bounds ACCESS is not: +// +// CLAMPED, therefore memory-safe under a violation — +// * `last_chunk_indices[b] >= nchunks` would make M2Store(passed, ...) at +// :336 write PAST the cudaMallocAsync allocation; `lci[b-1] < -1` would +// make `states[...]` at :335 read before it. Clamped at :319-331. +// * `seq_idx[c] ∉ [0,S)` would read `initial_states` out of bounds at :435, +// and a `seq_idx[0] < 0` would additionally make `si == si_prev` at c == 0 +// and index `passed` at chunk -1. Clamped at :404-419. +// * an out-of-range `state_indices` slot writes nothing at all. Clamped at +// :486, as it always has been. +// The clamps are the reason the D2H argument above does not apply to these: +// the values are ALREADY IN REGISTERS at their use sites, so bounding them +// costs nothing and needs no host round trip. They do NOT restore the checks — +// out-of-contract metadata still produces a WRONG ANSWER, now with a defined +// shape. They bound only WHERE it is read from. Both are pinned by device-only +// cases in tests/vt/test_ops_mamba2_ssd.cpp rather than asserted here. +// +// NOT CLAMPED, therefore NOT memory-safe under a violation — +// * the `cu_chunk_seqlens` tiling and per-chunk length checks. `start` and +// `len` derived from a garbage `ccs` index x/B/C/z/out out of bounds in +// every stage. Bounding these needs T at each use site and a clamp in the +// inner loops, which is not free; it is owed with the error ring above. +// VALUE-ONLY, in-bounds wrong number — +// * `A < 0` and `state_indices` distinctness. +#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); } + +// Scope guard for the prefill path's per-call scratch. `M2Check` throws, so a +// failure on the Nth `cudaMallocAsync` would otherwise leak the N-1 before it. +// The happy path calls `Release()`, which frees on the stream and CHECKS each +// free exactly as the open-coded sequence it replaces did; the destructor is the +// unwinding path only, and cannot throw. +class M2Scratch { + public: + explicit M2Scratch(cudaStream_t s) : s_(s) {} + M2Scratch(const M2Scratch&) = delete; + M2Scratch& operator=(const M2Scratch&) = delete; + ~M2Scratch() { + for (int i = 0; i < n_; ++i) static_cast(cudaFreeAsync(p_[i], s_)); + } + void* Alloc(size_t bytes, const char* what) { + // Refuse rather than overrun if a sixth buffer is ever added: a silent + // overflow here would be the exact defect class the guard exists to remove. + if (n_ >= kMax) throw std::runtime_error("vt cuda mamba2: scratch slots exhausted"); + void* p = nullptr; + M2Check(cudaMallocAsync(&p, bytes, s_), what); + p_[n_++] = p; + return p; + } + void Release() { + const int n = n_; + n_ = 0; + // Free ALL of them before reporting, so a mid-sequence failure does not + // leak the remainder the way the open-coded `M2Check(cudaFreeAsync(...))` + // sequence this replaces did; then throw on the first error seen. + cudaError_t first = cudaSuccess; + for (int i = 0; i < n; ++i) { + const cudaError_t e = cudaFreeAsync(p_[i], s_); + if (first == cudaSuccess) first = e; + } + M2Check(first, "scratch free"); + } + + private: + static constexpr int kMax = 5; + cudaStream_t s_; + void* p_[kMax] = {}; + int n_ = 0; +}; + +// 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); + // REGISTER-LOCAL MEMORY-SAFETY CLAMP (see the header's "what this arm does + // not check"). The host arm value-checks `0 <= last_chunk_indices[b] < + // nchunks` (cpu_ops.cpp); this arm cannot. Unclamped, `lci[b] >= nchunks` + // makes the `M2Store(passed, ...)` below write PAST the cudaMallocAsync + // allocation, and `lci[b-1] < -1` makes `states[...]` read before it. Both + // bounds are already in registers, so the D2H argument for dropping the + // check does not reach them and this costs nothing. It does NOT restore the + // check: out-of-contract metadata still yields a wrong number, now with a + // DEFINED shape — the chunk loop stops at nchunks and starts at 0. + int64_t chunk_end = lci[b] + 1; + int64_t chunk_start = b > 0 ? lci[b - 1] + 1 : 0; + if (chunk_end > nchunks) chunk_end = nchunks; + if (chunk_start < 0) chunk_start = 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 S, + 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; + // REGISTER-LOCAL MEMORY-SAFETY CLAMP, as in M2StatePassKernel above. The + // host arm value-checks `seq_idx[c] in [0,S)` (cpu_ops.cpp); unclamped, an + // out-of-range `si` reads `initial_states` out of bounds at `prevbase`, and + // a `sidx[0] < 0` makes `si == si_prev` at c == 0 and indexes `passed` at + // chunk -1. In-contract `si` is ALWAYS in range, so every in-contract path + // below is bit-identical to the unclamped form; out of contract the answer + // is still wrong, now with the defined shape "this chunk opens with a zero + // previous state". + const bool si_ok = si >= 0 && static_cast(si) < S; + bool prev_zero = false; + const void* prevp = passed; + DType prevdt = sdt; + int64_t prevbase = ((c - 1) * H + h) * row; + if (!si_ok) { + prev_zero = true; + } else 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; + // `M2Check` THROWS, so the five allocations below are held by a scope guard: + // without it a failure on the 3rd leaks the 1st and 2nd. The guard is released + // once the explicit frees at the end of the happy path have run. + M2Scratch scratch(s); + float* dtv = static_cast(scratch.Alloc(n_cumsum * sizeof(float), "dtv alloc")); + float* dac = static_cast(scratch.Alloc(n_cumsum * sizeof(float), "dac alloc")); + float* states = static_cast(scratch.Alloc(n_states * sizeof(float), "states alloc")); + float* cb = static_cast(scratch.Alloc(n_cb * sizeof(float), "cb alloc")); + void* passed = scratch.Alloc(n_states * state_elem, "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, S, cs, hpg); + + const cudaError_t launched = cudaGetLastError(); + scratch.Release(); // frees all five on the stream, reporting the first error + 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..a90ef15d2 100644 --- a/tests/vt/test_ops_mamba2_gated_norm.cpp +++ b/tests/vt/test_ops_mamba2_gated_norm.cpp @@ -543,3 +543,352 @@ 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, including the FMA-contraction term: `part += +// v * v` is one nvcc-`fmad` rounding on device and two under the host's +// `-ffp-contract=off`. ONE further 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 THIRD source of +// divergence — summation ORDER — on top of libm and contraction. 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. +// That is 2(K-1)*u for the reorder plus K*u for the contraction plus the silu +// `expf` difference, comfortably inside the `5*(K + 10)*u` that +// `DerivedRtol(group_size)` — the same expression the other two suites use — +// gives, 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)); +} + +// `5*(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 K*u for the contraction of `part += v * v` on the device side +// only, plus the libm difference in silu's `expf`. +constexpr double kUnitRoundoff = 5.9604644775390625e-08; // 2^-24 +double DerivedRtol(int64_t K) { return 5.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..37add73fd 100644 --- a/tests/vt/test_ops_mamba2_ssd.cpp +++ b/tests/vt/test_ops_mamba2_ssd.cpp @@ -999,3 +999,750 @@ 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, for TWO +// named reasons and no others: +// +// (a) LIBM. The two arms call different ones. 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. +// (b) FMA CONTRACTION. Host C++ is pinned `-ffp-contract=off` +// (CMakeLists.txt:41-56), so `a*b + c` keeps two roundings; nothing +// passes `--fmad=false` to nvcc, so the device arm compiles at the +// DEFAULT `--fmad=true` and every `acc += a*b` is a single-rounding +// `fma` whose host twin is not. `-fmad=false` was weighed and rejected +// at src/vt/cuda/cuda_mamba2_ssd.cuh — it is a per-TU flag on a header +// included by a hot GDN TU — so the bound carries the term. +// +// Summation ORDER is not a third source: the device kernels accumulate every +// output element in ONE thread, over the same index range in the same +// direction as the host arm. `DerivedRtol` below propagates (a) and (b), and +// only those. 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 contraction 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 sum of products carries, between the two arms, THREE terms: +// +// libm <= 2.5 ulp PER FACTOR — CUDA `expf` <= 2 ulp, glibc <= 0.5 — +// through a product of at most K, i.e. <= 2.5*K*u; +// summation the standard (K-1)*u forward error of a length-K f32 sum, +// which is what AMPLIFIES the libm difference in the inputs; +// contraction the device arm's `acc += a*b` is ONE nvcc-`fmad` rounding and +// the host arm's is TWO (`-ffp-contract=off`, CMakeLists.txt:55), +// so the host carries K product roundings the device does not: +// <= K*u. See the FMA-contraction note in +// src/vt/cuda/cuda_mamba2_ssd.cuh for why the flag is not simply +// turned off instead. +// +// Total <= 2.5*K*u + (K-1)*u + K*u = 4.5*K*u - u. `5*(K + 2)*u` is that, +// rounded up to integers, and it covers the model for every K >= 0 because +// 5K + 10 >= 4.5K - 1 reduces to 0.5K + 11 >= 0. +// +// THE OLD `4*(K + 2)*u` DID NOT. It omitted the contraction term, and +// 4.5K - 1 <= 4K + 8 holds only for K <= 18 — while the driver-shapes case runs +// at K = T = 200. The constant moved because the DERIVATION gained a term the +// build actually emits, not because a run needed slack: at 4*(K+2) the worst of +// the 55 audited comparisons used 7.66% of budget, so at 5*(K+2) it uses 6.13%, +// and mutant M3's 962173% becomes 769738% — still caught by four orders of +// magnitude (§8.4). +// +// 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 5.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. +// +// `lci_override` / `sidx_override` replace the metadata `ComputeVarlenChunkMetadata` +// derived, and exist for ONE caller: the clamp case below, which must reach the +// device kernels with metadata the host arm refuses. Every other caller passes +// nullptr and gets the derived metadata unchanged. +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, + const std::vector* lci_override = nullptr, + const std::vector* sidx_override = nullptr) { + 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()); + if (lci_override != nullptr) meta.last_chunk_indices = *lci_override; + if (sidx_override != nullptr) meta.seq_idx = *sidx_override; + + 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); + } +} + +// ───────────────────────────────────────────────────────────────────────────── +// OUT-OF-CONTRACT METADATA, DEVICE ONLY. +// +// The shared validator checks metadata SHAPE / DTYPE / DEVICE only +// (`CheckI32Meta`, ops.cpp); every VALUE check lives in the HOST kernel, which +// reads the tensors (cpu_ops.cpp). The device arm therefore runs with NONE of +// them. src/vt/cuda/cuda_mamba2_ssd.cuh enumerates all six and says which are +// memory-unsafe; this case covers the TWO that are both memory-unsafe and +// bounded by a register-local clamp: +// * `last_chunk_indices[b] >= nchunks` makes the `passed` store run past its +// `cudaMallocAsync` allocation, and +// * a `seq_idx[c]` outside `[0,S)` reads `initial_states` out of bounds, with +// `seq_idx[0] < 0` additionally indexing `passed` at chunk -1. +// It PINS the clamped behaviour so it is a test, not a sentence in a header. +// The `cu_chunk_seqlens` checks are NOT clamped and are NOT covered here; the +// header states that as an open memory-safety gap rather than hiding it. +// +// THIS CASE HAS NO HOST TWIN, deliberately: the host kernel REFUSES both inputs, +// which is exactly why the device arm needed the clamp. Each clamp is pinned +// against an IN-CONTRACT reference run whose result the clamp is defined to +// reproduce, so the assertions are EXACT — a tolerance here would be wide enough +// to hide the defect it is meant to catch. +// +// What this case does NOT establish is memory safety itself. An out-of-bounds +// write into a `cudaMallocAsync` pool very often does not fault, so a green run +// is necessary and not sufficient; `compute-sanitizer memcheck` on this case is +// what proves it, and is recorded as owed in mamba2-ssd.md §8.5. +TEST_CASE("mamba2 chunk scan CUDA arm clamps out-of-contract metadata in registers") { + Backend* gpu = MaybeCuda(); + if (gpu == nullptr) { + MESSAGE("SKIP: no CUDA backend registered (CPU-only build/box)"); + return; + } + // Four one-chunk sequences, so `nchunks == S == 4` and every chunk opens a new + // sequence. That makes the in-contract `seq_idx` {0,1,2,3} and the in-contract + // `last_chunk_indices` {0,1,2,3}. + const int64_t H = 4, P = 8, G = 2, N = 16, chunk = 16; + const std::vector cu{0, 16, 32, 48, 64}; + const int64_t T = cu.back(), S = static_cast(cu.size()) - 1; + const Inputs in = GenerateInputs(T, H, P, G, N, 0xC1A47u); + const ChunkMeta meta = ComputeVarlenChunkMetadata(cu, chunk); + const int64_t nchunks = static_cast(meta.seq_idx.size()); + REQUIRE(nchunks == S); + // A braced init-list cannot appear inside a doctest macro — the preprocessor + // splits it on the commas — so the expectation is named first. + const std::vector in_contract{0, 1, 2, 3}; + REQUIRE(meta.last_chunk_indices == in_contract); + REQUIRE(meta.seq_idx == in_contract); + + std::mt19937 rng(0xB0B0u); + std::normal_distribution nd(0.0f, 0.5f); + std::vector init(static_cast(S * H * P * N)); + for (auto& v : init) v = nd(rng); + + RunCfg cfg; + cfg.chunk_size = chunk; + + auto max_abs_diff = [](const std::vector& a, const std::vector& b) { + REQUIRE(a.size() == b.size()); + double worst = 0.0; + for (size_t i = 0; i < a.size(); ++i) + worst = std::max(worst, std::abs(static_cast(a[i]) - static_cast(b[i]))); + return worst; + }; + auto count_differing = [](const std::vector& a, const std::vector& b) { + REQUIRE(a.size() == b.size()); + size_t n = 0; + for (size_t i = 0; i < a.size(); ++i) + if (a[i] != b[i]) ++n; + return n; + }; + + SUBCASE("last_chunk_indices past the end stops the chunk loop at nchunks") { + // In contract, the last sequence's `lci` is `nchunks - 1`. The clamp is + // defined to make anything >= nchunks behave as nchunks - 1, so these two + // runs must agree BIT FOR BIT — in `final_states` as much as in `y`, since + // the clamp is inside the state-passing kernel that produces both. + const RunOut ref = + RunChunkScanCuda(*gpu, in, T, H, P, G, N, cu, nullptr, nullptr, nullptr, &init, cfg); + RequireNativeCudaProvider(vt::OpId::kMamba2ChunkScan, "lci clamp reference"); + for (int32_t past : {static_cast(nchunks), static_cast(nchunks + 41)}) { + std::vector lci = meta.last_chunk_indices; + lci.back() = past; + const RunOut got = RunChunkScanCuda(*gpu, in, T, H, P, G, N, cu, nullptr, nullptr, nullptr, + &init, cfg, &lci, nullptr); + INFO("last_chunk_indices.back() = " << past << " (nchunks = " << nchunks << ")"); + CHECK(count_differing(got.y, ref.y) == 0); + CHECK(count_differing(got.final_states, ref.final_states) == 0); + } + } + + SUBCASE("a seq_idx outside [0,S) opens the chunk with a ZERO previous state") { + // The reference is an IN-CONTRACT run with NO initial states: `seq_idx` + // {0,1,2,3} makes every chunk open a new sequence, and with `initial_states + // == nullptr` every one of them therefore opens from zero. That is exactly + // what the clamp is defined to do for an out-of-range index, so `y` must + // match bit for bit. `final_states` is NOT compared: the state-passing + // kernel reads `initial_states` by `b`, not by `seq_idx`, so it legitimately + // differs between a run that was given initial states and one that was not. + const RunOut zero_ref = + RunChunkScanCuda(*gpu, in, T, H, P, G, N, cu, nullptr, nullptr, nullptr, nullptr, cfg); + RequireNativeCudaProvider(vt::OpId::kMamba2ChunkScan, "seq_idx clamp reference"); + // ... and the pin is not vacuous: with the SAME metadata in contract, the + // initial states genuinely move `y`, so "matches the zero-init run" is a + // real statement rather than two identical computations + // ([[gate-comparing-shared-helper-proves-consistency-not-correctness]]). + const RunOut init_ref = + RunChunkScanCuda(*gpu, in, T, H, P, G, N, cu, nullptr, nullptr, nullptr, &init, cfg); + INFO("max|y(initial_states) - y(none)| = " << max_abs_diff(init_ref.y, zero_ref.y)); + REQUIRE(max_abs_diff(init_ref.y, zero_ref.y) > 1e-6); + + // Over the top: every index >= S, which unclamped indexes `initial_states` + // out of bounds. Under the bottom: every index negative, which unclamped + // ALSO makes `si == si_prev` at c == 0 and indexes `passed` at chunk -1. + const std::vector> bad{ + {static_cast(S), static_cast(S + 1), static_cast(S + 2), + static_cast(S + 900)}, + {-1, -1, -1, -1}, + }; + for (const std::vector& sidx : bad) { + REQUIRE(static_cast(sidx.size()) == nchunks); + const RunOut got = RunChunkScanCuda(*gpu, in, T, H, P, G, N, cu, nullptr, nullptr, nullptr, + &init, cfg, nullptr, &sidx); + INFO("seq_idx = {" << sidx[0] << ", " << sidx[1] << ", " << sidx[2] << ", " << sidx[3] + << "}, S = " << S); + CHECK(count_differing(got.y, zero_ref.y) == 0); + } + } +} + +#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..d4e6f9bf1 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,503 @@ 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 TWO admitted sources of divergence, the libm difference and +// nvcc's default `--fmad=true` contraction of `y += sn * cv`, both of 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)); +} + +// `5*(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, plus K*u for the K product +// roundings the host arm's `-ffp-contract=off` keeps and the device arm's +// nvcc-`fmad` `fma` does not (`y += sn * cv`, cuda_mamba2_ssd.cuh). +constexpr double kUnitRoundoff = 5.9604644775390625e-08; // 2^-24 +double DerivedRtol(int64_t K) { return 5.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