From ea8ef75e0f530d9f88a3f1dd3723acc7de4beed1 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 12 Aug 2026 16:37:24 +0000 Subject: [PATCH] spec(KERNEL-SSM-MAMBA): the Mamba2 SSD core was never built, only the GDN arm (#496) FOLLOWING_AGENTS_PROTOCOL The row is titled "General Mamba selective scan, causal convolution, SSD, and linear attention kernels". Only the second and fourth clauses were ever built. Everything landed under it is the GDN arm: the gated-delta-rule recurrence, its causal conv, its state gather/scatter, the sm_120 tiling campaign. The matrix row already conceded it in passing, four times, and the concession is easy to read as a scoping note rather than as an absent capability: "these GDN tests do not prove generic Mamba support". They do not, and the distance is not small. Gated delta rule carries the (I - beta k k^T) removal term and a per-head scalar decay; Mamba2's SSD is a diagonally-decayed gated linear recurrence driven by A_log / dt / D, with B and C shared across n_groups head groups. include/vt/ops.h has kCausalConv1dFwd, kCausalConv1dUpdate, kCausalConv1dSpecUpdate, kRmsNormGated, kGdnPrefill, kGdnDecode and kKdaGatedDeltaRule, and no chunked state-space scan and no selective state update at all. kGdnPrefill cannot be reshaped into one. What this costs: ten INVENTORIED model rows wait on this single kernel -- NemotronH, Mamba2, GraniteMoeHybrid, FalconH1, Zamba2, Jamba, Lfm2, Lfm2Moe, Plamo2, OlmoHybrid. The immediate driver is nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4 (23 mamba / 6 attention / 23 MoE, n_groups=8, ssm_state_size=128, chunk_size=128), which is 20.1 GiB and fits one GB10. The pinned oracle already runs that arch -- registry.py:179 -> nemotron_h.py, MambaMixer2 at nemotron_h.py:373, NemotronHMoEDecoderLayer at nemotron_h.py:317, ModelOptMixedPrecisionConfig at modelopt.py:2280 -- so the gap is entirely ours, and the SSD core is the part not reachable by composing what we already have. The spec is host-reference-first, mirroring the KDA and DeepSeek-V4 DSA lanes: W1 lands portable CPU references for the chunked scan, the single-token selective update and the silu-gated GROUP RMS norm (a sibling of our sigmoid kRmsNormGated, not a parameter of it), gated against a from-first-principles double-precision reference. Two tests upstream does not have carry the weight, because chunk cumsum, state passing and the scan combine are each plausible and jointly wrong in ways one chunk cannot see: chunked-must-equal-sequential, and chunk-boundary invariance across chunk_size in {8,16,32,64,128}. No lifecycle state moves and no measurement is claimed. The row stays INVENTORIED: the spec is committed, the implementation has not started, and W1 goes to a fresh implementer. FEATURES.md gains a "Inventoried but blocked" row because the honest reason differs from the other three entries there -- Nemotron-H fits the pool, so it is capability-blocked, not hardware-blocked, and the page should not imply we are waiting on a bigger box. Records and docs only, no product code touched. Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-opus-5 [Claude Code] --- .agents/kernel-matrix.md | 2 +- .agents/model-matrix.md | 2 +- .agents/roadmap_v1.md | 1 + .agents/specs/mamba2-ssd.md | 299 ++++++++++++++++++++++++++++++++++++ docs/FEATURES.md | 1 + 5 files changed, 303 insertions(+), 2 deletions(-) create mode 100644 .agents/specs/mamba2-ssd.md diff --git a/.agents/kernel-matrix.md b/.agents/kernel-matrix.md index 844a20c76..45bf9bddf 100644 --- a/.agents/kernel-matrix.md +++ b/.agents/kernel-matrix.md @@ -154,7 +154,7 @@ host/sched. Detail: state `KERNEL-FA2-GQA-SWAP-FLIP`. | `KERNEL-GDN-PACKED-DECODE` | Pure non-spec packed GDN recurrent decode for FP16/BF16/F32 | default and dispatch `vllm/envs.py:117,1123-1125`, `qwen_gdn_linear_attn.py:1286-1298,1644-1695`; packed body/launch `fla/ops/fused_recurrent.py:255-478`; upstream test `tests/kernels/test_fused_recurrent_packed_decode.py:13-98` | public [API](../include/vt/ops.h#L518), [validation](../src/vt/ops.cpp#L1466), portable [CPU recurrence](../src/vt/cpu/cpu_ops.cpp), registered [CUDA kernel](../src/vt/cuda/cuda_gdn.cu#L156), **register-resident perf kernel** [GdnPackedDecodeRegTileKernel](../src/vt/cuda/cuda_gdn.cu#L1183) + [rollback flag](../src/vt/cuda/gdn_packed_reg_tile.h) + [CPU flag test](../tests/vt/test_gdn_packed_reg_tile.cpp), exact [model dispatch](../src/vllm/model_executor/models/qwen3_5.cpp#L1555), [request-identity slot pool](../src/vllm/v1/worker/gpu/runner.cpp), [trace finalizer](../tools/bench/finalize_gdn_packed_trace.py#L218), production [component driver](../scripts/dgx-gdn-packed-component.sh), and [component finalizer](../tools/bench/gdn_packed_component.py) | Clean `f18ca23`/`9ad8fb7` close G0/G1; `f344dec` closes W1D2/G2; `7ff713e` + `24cea4f` close W1D3 structure; the identity-keyed slot fix at `c172336` proved the c16 repair (model gates 235/235, `--diagnostic-c16` 3/3). **W1D3 CLOSES on EQUIVALENCE** over eight sealed component roots + the 8-pair locked c16 A/B (**−0.205% ± 0.30, <1σ**) + the 24-window trace attribution (packed GPU-cheaper, kernel compute −1.30..−1.58%/step, no attributable packed-side cost): no STABLE regression on any axis; every failing axis is a sign-flipping band-edge statistic of a true-zero effect. The eighth (first 22-leg: cold-discard pair + 5 reps) seal `complete-failed` at **38/40 + 8/8 memory**, stability clean, `validation_error=None`, paired-consistency PASS at BOTH c2/c16 (the two fails are c2 `median_tpot_ms` 0.9899 and c2 pooled `p99_ttft_ms` 0.8464). Packed stays the **default** (exact-upstream semantics, 48-launch reduction); `VT_GDN_PACKED_DECODE=0` is the rollback; **no `complete-pass` marker exists and NO speed credit is claimed**. [Component tests](../tests/tools/test_gdn_packed_component.py#L1) **79/79**, all tools **162/162**; [runner test](../tests/vllm/v1/worker/test_runner.cpp) 8/8. Correctness blast radius in [BENCHMARKS](../docs/BENCHMARKS.md); closure [ledger](parity-ledger.md#L469). qkvz (`KERNEL-GEMM-BF16` W2) is UNBLOCKED and the exact grid is authorized. **2026-07-16 PERF LEVER (test-first, CPU-gated, DGX-pending):** the named +2.06 ms/step recurrence-tiling lever is ported into `GdnPackedDecodeRegTileKernel` — one warp per `[BV=32,BK]` tile, state block REGISTER-RESIDENT (no shared-state round-trip, no cross-warp `__shfl` reduction, no `__syncthreads`; mirrors vLLM FLA `num_warps=1`/`num_stages=3` `fused_recurrent.py:256-336`), default-on behind `VT_GDN_PACKED_REG_TILE` (=0 → legacy bit-for-bit, same binary), selected for `bv==32 && dk in {32,128}`. Boundary-fixture bit-exactness preserved (sequential per-row Dk reduction). **PROOF FAILED (dgx 54f0541): oracle boundary FAIL + c16 700.5 vs 793.6 tok/s → default flipped OFF (`309c218`).** **2026-07-16 RESOLVED via the sanctioned vendored Triton cubin (`CLAIM-GDN-DECODE-TRITON`):** Phase-1 cuobjdump MEASURED the codegen cause — vLLM FLA decode cubin REG:205/0-spill vs the naive hand port REG:255+STACK:48 (spills); DECISION = vendored cubin (`gdn_decode_h48`, 27B-only, `VT_GDN_PACKED_DECODE_TRITON` default OFF, hand kernel stays default). DGX gates: AOT op test 28/28, full GDN 49/49, oracle boundary 12/12, **27B model gate 235/235 token-exact with the Triton path ON**, memcheck 0/0. c16 A/B: triton [817.51, 821.06, 822.55] vs legacy [813.77, 815.62, 815.30] tok/s — paired mean **+5.48 tok/s (+0.67%)**, monotone (+3.74/+5.44/+7.25), 3/3 pairs positive; mean TPOT triton [161.04, 160.49, 160.35] vs legacy [162.09, 161.65, 161.93] = **-1.26 ms (-0.78%)** (median TPOT -1.13 ms); w0 cold-discard (triton 821.48/160.44) excluded. Anchors: [shim](../triton_kernels/fused_recurrent_packed_decode.py), [vendored cubin](../src/vt/cuda/triton_aot_vendored/sm_121a/gdn_decode_h48.h), [launcher `TryTritonPackedDecode`](../src/vt/cuda/cuda_gdn.cu), [default-ON flag header](../src/vt/cuda/gdn_packed_decode_triton.h) + [CPU flag test](../tests/vt/test_gdn_packed_decode_triton.cpp), [AOT test](../tests/vt/test_ops_gdn.cpp). ACCEPTANCE MET (oracle PASS + consistent c16 TPOT improvement + no throughput regression). **2026-07-16 DEFAULT FLIP ON (`CLAIM-GDN-DECODE-TRITON-FLIP`):** per MIRROR policy (the vendored kernel IS vLLM's exact token-identical FLA kernel, run by default upstream) `VT_GDN_PACKED_DECODE_TRITON` flipped OFF→ON, joining the sibling GDN Triton kernels; `=0` is the same-binary rollback. Test-first default-ON pure-header predicate + CPU flag test (RED→GREEN 10/10). **35B: no specialization** — excluded at the model level by the dense-only `ShouldUsePackedGdnDecode` (`qwen3_5.cpp:49`) and the launcher guard rejects its `Hv=32` shape (`cmake` H=32=35B) anyway (clean fallback; a 35B cubin would be dead code). Flip gates ALL EIGHT PASS exit 0 (root `~/work/vllm.cpp-gdn-decode-triton-flip` `gates.verdict`, `-DVLLM_CPP_TRITON=ON` + CUTLASS-4.5.0/nvcc-13.0, CUTLASS/FA2 configure-log lines verified, one flock): 27B DEFAULT (Triton) **235/235** + `=0` rollback **235/235**; 35B DEFAULT **315/315** + `=0` rollback **315/315** (inert); AOT op test **28/28** (default fires cubin, `=0` fires legacy); full GDN **49/49 (2,343/2,343)**; oracle boundary **12/12**; memcheck **28/28, 0 errors**. No new A/B (9dd7d3f's +5.48 tok/s / −1.26 ms TPOT stands). The next binding grid runs the Triton decode path by default; no separate flip speed credit is claimed. | [packed decode](specs/gdn-packed-decode.md) | `DONE` | `e47b4d6` | | `KERNEL-GDN-AOT-BF16` | GDN prefill/decode and BF16-output Triton/CuTe fast-path coverage | backend resolution `qwen_gdn_linear_attn.py:150-411`; FLA kernels `fla/ops/chunk_delta_h.py:43-371`; FlashInfer GDN entry `flashinfer-ref/gdn_prefill.py:37-87` | [canonical AOT contract](../cmake/TritonAOTKernels.cmake#L47), [BF16 manifest bases](../src/vt/cuda/triton_aot_vendored/sm_121a/MANIFEST#L21), [typed dispatch](../src/vt/cuda/cuda_gdn.cu#L2910), [27B-only default/f32 fallback](../src/vllm/model_executor/models/qwen3_5.cpp#L1555) | Existing AOT/safety/native gates are green. The BF16/f32 component is **1.007989×**, 16/20 timing and 2/4 memory. Immutable `3f256ab` is the exact v0.25 denominator: c16 total throughput passes at **1.027889×**, but mean TPOT/ITL remains **0.987450× normalized**; c1-c8 total and decode-shaped axes fail. Diagnostic cross-profiler totals are local `GdnDecodeFusedKernel` 19.101 s/73,578 calls and vLLM fused recurrent 28.659 s/70,848 calls; they are not a valid speed ratio. The independently gateable pure-decode mismatch is now owned by `KERNEL-GDN-PACKED-DECODE`; this row retains prefill AOT, BF16 default and scratch/teardown scope. Every 35B path stays f32; strict component/pool/teardown remain open | [inventory](specs/kernel-family-inventory.md) | `READY` | - | | `KERNEL-GDN-SCRATCH` | Stream-ordered GDN AOT scratch allocation, reuse, and capture safety | FLA/FlashInfer runtime path `qwen_gdn_linear_attn.py:1073-1161`; CuTe Blackwell scheduler `flashinfer-ref/gdn_kernels/blackwell/gated_delta_net_chunked.py:177-182,379-846` | [stream-owned pool](../src/vt/cuda/cuda_gdn.cu#L2662), [queue cleanup](../src/vt/cuda/cuda_gdn.cu#L3641), [test poison hook](../src/vt/cuda/cuda_gdn.cu#L3679) | [same-stream poison/reuse/growth/pool-off assertions](../tests/vt/test_ops_gdn.cpp#L1494); all 11 buffers are filled with `0xff` before reuse and dispatch counters reject hand fallback; integrated current-main GB10 graph gate, sanitizer, trace and pool A/B pending | [inventory](specs/kernel-family-inventory.md) | `READY` | - | -| `KERNEL-SSM-MAMBA` | General Mamba selective scan, causal convolution, SSD, and linear attention kernels | selective scan source `CMakeLists.txt:401`, `csrc/libtorch_stable/mamba/selective_scan_fwd.cu:394-820`; upstream mamba tests; decode conv-update FLA Triton `causal_conv1d.py:15-192` @ `e24d1b24` | only Qwen GDN-specific convolution/state path, including indexed cache gather/scatter [cuda_gdn.cu:156](../src/vt/cuda/cuda_gdn.cu#L156) and GDN dispatch [cuda_gdn.cu:3695](../src/vt/cuda/cuda_gdn.cu#L3695); **decode conv-update decode-fast** [CausalConv1dUpdateFastKernel + TryLaunchConvUpdateFast](../src/vt/cuda/cuda_gdn.cu) + [flag header](../src/vt/cuda/conv_update_fast.h); **prefill conv-fwd register-window** [CausalConv1dFwdRegKernel + LaunchConvFwdReg](../src/vt/cuda/cuda_gdn.cu) + **fused post-conv per-V-head split** [GdnPostConvSplitKernel](../src/vt/cuda/cuda_gdn.cu) + **fused post-conv fast (128-thread + 128-bit V copy)** [GdnPostConvFastKernel + GdnVecCopy8](../src/vt/cuda/cuda_gdn.cu) + [prefill flag header](../src/vt/cuda/gdn_prefill_conv.h) + [CPU flag test](../tests/vt/test_gdn_prefill_conv.cpp) | Indexed BF16/F32 cache-boundary + i8-mask CUDA tests [test_ops_gdn.cpp:994](../tests/vt/test_ops_gdn.cpp#L994), focused 7/7 memcheck and W1 structural trace pass; these GDN tests do not prove generic Mamba support. **2026-07-18 (`CLAIM-CONV-UPDATE-FAST-1`): decode conv-update kernel-efficiency sub-lever landed, DEFAULT ON** (the c16-trace scan lever #5; row stays `INVENTORIED` — broad Mamba support unchanged). `CausalConv1dUpdateFastKernel` behind `VT_CONV_UPDATE_FAST` (default ON / `=0` rollback) is BIT-IDENTICAL (0-ulp) to the shipped `CausalConv1dUpdateKernel`: same float op order, only a 2D grid (removes two int64 div/mod per thread) + a WIDTH-templated register-cached state row reused for the conv accumulation AND the roll (mirrors FLA's per-`KERNEL_WIDTH` `col0..col3` register specialization). DGX-PROVEN (flock, clean `-Werror`, CUTLASS+FA2 hard-verified): `test_ops_gdn` conv-update decode-fast `fast==shipped` **BYTE-EXACT (0-ulp)** on `out`+`conv_state` [330/330](../tests/vt/test_ops_gdn.cpp) (k∈{3,4,5}, bf16+f32 state, ±bias, silu/identity, compact+scattered cache incl. NULL-block); full GDN 51/51; [CPU flag test](../tests/vt/test_conv_update_fast.cpp) 10/10. Isolated nsys pure-kernel 27B c16 shape (batch=16 conv_dim=10240 k=4 bf16): shipped 7,072 ns vs fast 3,680 ns median = **1.92×** (clears ≥1.3× bar). Full default set (all fast kernels ON) 27B 235/235 + 35B 315/315; `=0` rollback arms 235/235 + 315/315. `benchmark_binding=false`; binding grid re-measures. **2026-07-18 (`CLAIM-GDN-PREFILL-CONV-1`): PREFILL conv-fwd + fused post-conv kernel-efficiency sub-lever landed** (row stays `INVENTORIED`). `CausalConv1dFwdRegKernel` (`VT_CONV_REG` DEFAULT ON / `=0`→tiled) is the register-resident sliding-window mirror of vLLM FLA `_causal_conv1d_fwd_kernel` (`causal_conv1d.py:397-452`): per-channel weights preloaded to registers, `(k-1)`-tap register sliding window (each x loaded once, coalesced), token-axis chunked over grid.z for `n<=4` low-batch prefill. `GdnPostConvSplitKernel` (`VT_GDN_POSTCONV_SPLIT` OPT-IN) mirrors `_fused_post_conv_kernel` grid `(cdiv(L,BLOCK_T), H+HV)` (per-V-head). BIT-IDENTICAL (0-ulp) to the shipped tiled/megablock kernels — same tap-order f32 sum + q/k L2-norm byte-for-byte. DGX-PROVEN: byte-exact reg==tiled + split==megablock 268 GPU assertions + full GDN 3081/3081; memcheck 0; 27B **235/235** + 35B **315/315** (final defaults reg ON/split OFF). nsys 35B (`~/work/prefill-attr-conv-35b`): conv c1 −4.7% / c6 −7.3% (bandwidth-bound) ⇒ DEFAULT ON; post-conv split c1 −3.8% / c6 +4.7% (near-neutral — GdnPostConv is q/k-L2-norm-dominated) ⇒ OPT-IN. TTFT c1 −0.39% within run-noise (conv ~2.5% of GPU). Residual vLLM conv gap = bf16 traffic (VT_GDN_IN_BF16, task #40 sibling), not structure. `benchmark_binding=false`. **2026-07-19 (`CLAIM-GDN-POSTCONV-FAST-1`): fused post-conv FAST sub-lever landed, DEFAULT ON** (row stays `INVENTORIED`). A fresh production-path nsys (`--cuda-graph-trace=node`, VLLM_CPP_TRITON=ON, input-1024) confirmed the GDN chunk **compute** runs the vendored FLA Triton AOT cubins by default (`chunk_gated_delta_rule_fwd_kernel_h_blockdim64` / `chunk_fwd_kernel_o` / `recompute_w_u_fwd_kernel` / `chunk_scaled_dot_kkt_fwd_kernel` = at FLA parity), so the #1 remaining NON-AOT GDN kernel on BOTH models is the fused post-conv prep. The split (above) measured neutral/slower, so `GdnPostConvFastKernel` (`VT_GDN_POSTCONV_FAST` DEFAULT ON / `=0`→megablock) keeps the megablock grid `(T,Hk+1)` but makes two BYTE-IDENTICAL changes for Dk==Dv==128: 128 threads/block (the 128-wide L2-norm tree is the 256-wide tree minus a leading `+0` step; better reduction occupancy) + a 128-bit-staged V copy (`GdnVecCopy8`: raw `int4` when dtypes match, else the SAME `__bfloat162float`/`__float2bfloat16` converts). No arithmetic reordered; grounded in FLA `_fused_post_conv_kernel` (`fused_gdn_prefill_post_conv.py:57-149`). Isolated nsys per-call: 27B **175.6→133.0 ms (−24.3%)**, 35B **93.5→70.3 ms (−24.8%)**; in-situ TTFT (input-1024, 3 reps): 27B c1 −1.14% / c2 −1.31%, 35B c1 −0.72% / c2 −0.99% (all reps positive). BIT-EXACT: 27B 235/235 + 35B 315/315 on default(fast)+`=0`(megablock) arms; CPU flag test 28/28; clean `-Werror`. Evidence `dgx:~/work/vllm.cpp-gdn-chunk/prof`. `benchmark_binding=false` | [prefill conv spec](specs/gdn-prefill-conv-reg-2026-07-18.md); [conv-update decode-fast spec](specs/conv-update-decode-fast-2026-07-18.md); [inventory](specs/kernel-family-inventory.md) | `INVENTORIED` | - | +| `KERNEL-SSM-MAMBA` | General Mamba selective scan, causal convolution, SSD, and linear attention kernels | selective scan source `CMakeLists.txt:401`, `csrc/libtorch_stable/mamba/selective_scan_fwd.cu:394-820`; upstream mamba tests; decode conv-update FLA Triton `causal_conv1d.py:15-192` @ `e24d1b24` | only Qwen GDN-specific convolution/state path, including indexed cache gather/scatter [cuda_gdn.cu:156](../src/vt/cuda/cuda_gdn.cu#L156) and GDN dispatch [cuda_gdn.cu:3695](../src/vt/cuda/cuda_gdn.cu#L3695); **decode conv-update decode-fast** [CausalConv1dUpdateFastKernel + TryLaunchConvUpdateFast](../src/vt/cuda/cuda_gdn.cu) + [flag header](../src/vt/cuda/conv_update_fast.h); **prefill conv-fwd register-window** [CausalConv1dFwdRegKernel + LaunchConvFwdReg](../src/vt/cuda/cuda_gdn.cu) + **fused post-conv per-V-head split** [GdnPostConvSplitKernel](../src/vt/cuda/cuda_gdn.cu) + **fused post-conv fast (128-thread + 128-bit V copy)** [GdnPostConvFastKernel + GdnVecCopy8](../src/vt/cuda/cuda_gdn.cu) + [prefill flag header](../src/vt/cuda/gdn_prefill_conv.h) + [CPU flag test](../tests/vt/test_gdn_prefill_conv.cpp) | Indexed BF16/F32 cache-boundary + i8-mask CUDA tests [test_ops_gdn.cpp:994](../tests/vt/test_ops_gdn.cpp#L994), focused 7/7 memcheck and W1 structural trace pass; these GDN tests do not prove generic Mamba support. **2026-07-18 (`CLAIM-CONV-UPDATE-FAST-1`): decode conv-update kernel-efficiency sub-lever landed, DEFAULT ON** (the c16-trace scan lever #5; row stays `INVENTORIED` — broad Mamba support unchanged). `CausalConv1dUpdateFastKernel` behind `VT_CONV_UPDATE_FAST` (default ON / `=0` rollback) is BIT-IDENTICAL (0-ulp) to the shipped `CausalConv1dUpdateKernel`: same float op order, only a 2D grid (removes two int64 div/mod per thread) + a WIDTH-templated register-cached state row reused for the conv accumulation AND the roll (mirrors FLA's per-`KERNEL_WIDTH` `col0..col3` register specialization). DGX-PROVEN (flock, clean `-Werror`, CUTLASS+FA2 hard-verified): `test_ops_gdn` conv-update decode-fast `fast==shipped` **BYTE-EXACT (0-ulp)** on `out`+`conv_state` [330/330](../tests/vt/test_ops_gdn.cpp) (k∈{3,4,5}, bf16+f32 state, ±bias, silu/identity, compact+scattered cache incl. NULL-block); full GDN 51/51; [CPU flag test](../tests/vt/test_conv_update_fast.cpp) 10/10. Isolated nsys pure-kernel 27B c16 shape (batch=16 conv_dim=10240 k=4 bf16): shipped 7,072 ns vs fast 3,680 ns median = **1.92×** (clears ≥1.3× bar). Full default set (all fast kernels ON) 27B 235/235 + 35B 315/315; `=0` rollback arms 235/235 + 315/315. `benchmark_binding=false`; binding grid re-measures. **2026-07-18 (`CLAIM-GDN-PREFILL-CONV-1`): PREFILL conv-fwd + fused post-conv kernel-efficiency sub-lever landed** (row stays `INVENTORIED`). `CausalConv1dFwdRegKernel` (`VT_CONV_REG` DEFAULT ON / `=0`→tiled) is the register-resident sliding-window mirror of vLLM FLA `_causal_conv1d_fwd_kernel` (`causal_conv1d.py:397-452`): per-channel weights preloaded to registers, `(k-1)`-tap register sliding window (each x loaded once, coalesced), token-axis chunked over grid.z for `n<=4` low-batch prefill. `GdnPostConvSplitKernel` (`VT_GDN_POSTCONV_SPLIT` OPT-IN) mirrors `_fused_post_conv_kernel` grid `(cdiv(L,BLOCK_T), H+HV)` (per-V-head). BIT-IDENTICAL (0-ulp) to the shipped tiled/megablock kernels — same tap-order f32 sum + q/k L2-norm byte-for-byte. DGX-PROVEN: byte-exact reg==tiled + split==megablock 268 GPU assertions + full GDN 3081/3081; memcheck 0; 27B **235/235** + 35B **315/315** (final defaults reg ON/split OFF). nsys 35B (`~/work/prefill-attr-conv-35b`): conv c1 −4.7% / c6 −7.3% (bandwidth-bound) ⇒ DEFAULT ON; post-conv split c1 −3.8% / c6 +4.7% (near-neutral — GdnPostConv is q/k-L2-norm-dominated) ⇒ OPT-IN. TTFT c1 −0.39% within run-noise (conv ~2.5% of GPU). Residual vLLM conv gap = bf16 traffic (VT_GDN_IN_BF16, task #40 sibling), not structure. `benchmark_binding=false`. **2026-07-19 (`CLAIM-GDN-POSTCONV-FAST-1`): fused post-conv FAST sub-lever landed, DEFAULT ON** (row stays `INVENTORIED`). A fresh production-path nsys (`--cuda-graph-trace=node`, VLLM_CPP_TRITON=ON, input-1024) confirmed the GDN chunk **compute** runs the vendored FLA Triton AOT cubins by default (`chunk_gated_delta_rule_fwd_kernel_h_blockdim64` / `chunk_fwd_kernel_o` / `recompute_w_u_fwd_kernel` / `chunk_scaled_dot_kkt_fwd_kernel` = at FLA parity), so the #1 remaining NON-AOT GDN kernel on BOTH models is the fused post-conv prep. The split (above) measured neutral/slower, so `GdnPostConvFastKernel` (`VT_GDN_POSTCONV_FAST` DEFAULT ON / `=0`→megablock) keeps the megablock grid `(T,Hk+1)` but makes two BYTE-IDENTICAL changes for Dk==Dv==128: 128 threads/block (the 128-wide L2-norm tree is the 256-wide tree minus a leading `+0` step; better reduction occupancy) + a 128-bit-staged V copy (`GdnVecCopy8`: raw `int4` when dtypes match, else the SAME `__bfloat162float`/`__float2bfloat16` converts). No arithmetic reordered; grounded in FLA `_fused_post_conv_kernel` (`fused_gdn_prefill_post_conv.py:57-149`). Isolated nsys per-call: 27B **175.6→133.0 ms (−24.3%)**, 35B **93.5→70.3 ms (−24.8%)**; in-situ TTFT (input-1024, 3 reps): 27B c1 −1.14% / c2 −1.31%, 35B c1 −0.72% / c2 −0.99% (all reps positive). BIT-EXACT: 27B 235/235 + 35B 315/315 on default(fast)+`=0`(megablock) arms; CPU flag test 28/28; clean `-Werror`. Evidence `dgx:~/work/vllm.cpp-gdn-chunk/prof`. `benchmark_binding=false`. **2026-08-12 (#496): the SSD half of this row's title is named and specced.** Everything above is the GDN arm; the generic Mamba2 selective scan has no local symbol at all — `include/vt/ops.h` has `kCausalConv1dFwd/Update/SpecUpdate`, `kRmsNormGated`, `kGdnPrefill/Decode`, `kKdaGatedDeltaRule` and **no** chunked state-space scan and no selective state update. Gated delta rule is not Mamba2: GDN carries the `(I − β kkᵀ)` removal term and a per-head scalar decay, Mamba2's SSD is a diagonally-decayed gated linear recurrence driven by `A_log`/`dt`/`D` with `B`/`C` shared across `n_groups`. Owed: `mamba_chunk_scan_combined_varlen` (`mamba/ops/ssd_combined.py:157-235` @ `555967922`), `selective_state_update` (`mamba/ops/mamba_ssm.py:497`), and the silu-gated GROUP RMS norm (`mamba_mixer2.py:100-149`, a sibling of our sigmoid `kRmsNormGated`, not a parameter of it). Blocks 10 `INVENTORIED` model rows; immediate driver is `NemotronHForCausalLM` (Nemotron-3.5-Lightning-30B-A3B-NVFP4: 23 mamba / 6 attention / 23 MoE, `n_groups=8`, `ssm_state_size=128`, `chunk_size=128`). Row stays `INVENTORIED` — spec committed, implementation not started | [Mamba2 SSD spec](specs/mamba2-ssd.md); [prefill conv spec](specs/gdn-prefill-conv-reg-2026-07-18.md); [conv-update decode-fast spec](specs/conv-update-decode-fast-2026-07-18.md); [inventory](specs/kernel-family-inventory.md) | `INVENTORIED` | - | | `KERNEL-SAMPLING` | Sampler reductions, top-k/top-p, probabilities, penalties, masks, RNG | stable sampler/top-k `CMakeLists.txt:399-400`; upstream `tests/v1/sample/test_sampler.py`, `test_topk_topp_sampler.py` | [cuda_sample.cu:638](../src/vt/cuda/cuda_sample.cu#L638) | [sampling tests](../tests/vt/test_ops_sample.cpp#L50) | [inventory](specs/kernel-family-inventory.md) | `ANCHOR-BACKFILL` | - | | `KERNEL-COLLECTIVES` | Custom/quick/FlashInfer all-reduce and fused collective kernels | custom all-reduce `CMakeLists.txt:405`; runtime `device_communicators/custom_all_reduce.py:109-264`; upstream distributed tests | - | - | [inventory](specs/kernel-family-inventory.md) | `INVENTORIED` | - | | `KERNEL-SPEC-DECODE` | Speculative-decode expand/copy, rejection, acceptance, and draft sampling kernels | Triton setup kernels `vllm/v1/spec_decode/utils.py:306-565`; proposer path `llm_base_proposer.py:67-1315` | - | - | [inventory](specs/kernel-family-inventory.md) | `INVENTORIED` | - | diff --git a/.agents/model-matrix.md b/.agents/model-matrix.md index f29df210e..9ae59607a 100644 --- a/.agents/model-matrix.md +++ b/.agents/model-matrix.md @@ -245,7 +245,7 @@ Transformers compatibility is capability-driven and excluded from finite counts. | `MODEL-TEXT-mimo-v2-mi-mo-v2-flash-for-causal-lm` | `MiMoV2FlashForCausalLM` | `registry.py:166`; `vllm/model_executor/models/mimo_v2.py::MiMoV2FlashForCausalLM` | causal generation / text | model loader/forward; paged attention/KV; FusedMoE/grouped GEMM; sliding-window attention | ☐ required | `INVENTORIED` | none | unassigned | | `MODEL-TEXT-mimo-v2-mi-mo-v2-for-causal-lm` | `MiMoV2ForCausalLM` | `registry.py:167`; `vllm/model_executor/models/mimo_v2.py::MiMoV2ForCausalLM` | causal generation / text | model loader/forward; paged attention/KV; FusedMoE/grouped GEMM; sliding-window attention | ☐ required | `INVENTORIED` | none | unassigned | | `MODEL-TEXT-nemotron-nemotron-for-causal-lm` | `NemotronForCausalLM` | `registry.py:168`; `vllm/model_executor/models/nemotron.py::NemotronForCausalLM` | causal generation / text | model loader/forward; paged attention/KV | ☐ required | `INVENTORIED` | none | unassigned | -| `MODEL-TEXT-nemotron-h-nemotron-hfor-causal-lm` | `NemotronHForCausalLM`, `NemotronHPuzzleForCausalLM` | `registry.py:169-170`; `vllm/model_executor/models/nemotron_h.py::NemotronHForCausalLM` | causal generation / text | model loader/forward; paged attention/KV; FusedMoE/grouped GEMM; Mamba/SSM state; sliding-window attention | ☐ required | `INVENTORIED` | none | unassigned | +| `MODEL-TEXT-nemotron-h-nemotron-hfor-causal-lm` | `NemotronHForCausalLM`, `NemotronHPuzzleForCausalLM` | `registry.py:169-170`; `vllm/model_executor/models/nemotron_h.py::NemotronHForCausalLM` | causal generation / text | model loader/forward; paged attention/KV; FusedMoE/grouped GEMM; Mamba/SSM state; sliding-window attention. **2026-08-12: BLOCKED on `KERNEL-SSM-MAMBA` [#496](https://github.com/mudler/vllm.cpp/issues/496)** ([spec](specs/mamba2-ssd.md)) — the Mamba2 SSD core is unported. Beyond that kernel this arch owes three further things, none of which exist locally: non-gated `relu²` MoE (every grouped-MoE op we have is SwiGLU-shaped), ModelOpt `MIXED_PRECISION` per-module loading (`quantization/modelopt.py:2280` — NVFP4 W4A16 g16 experts and FP8 W8A8 mamba projections in one checkpoint), and the DeepSeek-style MTP head. Driver checkpoint `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4` (20.1 GiB, fits one GB10) | ☐ required | `INVENTORIED` | none | unassigned | | `MODEL-TEXT-olmo-olmo-for-causal-lm` | `OlmoForCausalLM` | `registry.py:171`; `vllm/model_executor/models/olmo.py::OlmoForCausalLM` | causal generation / text | model loader/forward; paged attention/KV | ☐ required | `INVENTORIED` | none | unassigned | | `MODEL-TEXT-olmo2-olmo2-for-causal-lm` | `Olmo2ForCausalLM`, `Olmo3ForCausalLM` | `registry.py:172-173`; `vllm/model_executor/models/olmo2.py::Olmo2ForCausalLM` | causal generation / text | OLMo-2-0425-1B dense LANDED (W0-W4). Impl: `include/vllm/model_executor/models/olmo2.h` + `src/vllm/model_executor/models/{olmo2,olmo2_weights,olmo2_registry}.cpp` (TWO `REGISTER_VLLM_MODEL` for `Olmo2ForCausalLM`+`Olmo3ForCausalLM`, one shared factory; reuses the shared dense glue). ZERO new compute kernel — the two "distinctives" both reduce to WIRING over landed ops: (1) PURE POST-NORM (`norm_after`, `olmo2.py:261-277`) = standalone `vt::RmsNorm` on each sublayer OUTPUT + plain `vt::Add` residual re-join, NO pre-norm (the GLM-4 standalone-output-norm op, pre-norms dropped); (2) FULL-WIDTH QK-norm (`olmo2.py:113-117,160-172`) = two standalone `vt::RmsNorm` over `[T,q_size]`/`[T,kv_size]` before NeoX rope (NOT the per-head `kAttnQkNormRope`). Reuse: SiLU SwiGLU (`kSiluAndMul`), NeoX RoPE (`RopeFromCache`/`RopeNeox`), GQA/MHA paged glue (`dense_attn_block.h`), merged qkv/gate_up loader. Checkpoint is F32 on-disk → loader downcasts f32→bf16 (round-to-nearest-even, `vt::F32ToBF16`) to match vLLM-bf16. UNTIED lm_head (`tie_word_embeddings:false`). Loader 179 tensors, zero missing/unmapped. Tokenizer: ONE shared-TU touch — `tokenizer.cpp` accepts the OLMo-2 Split `behavior=Removed,invert=true` encoding (equiv. to `Isolated` for the full-cover cl100k regex → `kLlama3`); guarded additive branch, no other checkpoint uses it. NO-BOS-verified real ByteLevel gate. Runs EAGER (bf16, no decode graph). Gate: 16/16 vs vLLM 0.25.0 (STRICT 13/16 + near-tie 3/16, max gap 0.094 nats, 0 forward-divergent) — speed pending. **OLMo-3 W5 (batch3) IMPLEMENTED** (guarded additive edits to `olmo2.{h,cpp,weights}`, diff-inert for OLMo-2 — re-run OLMo-2 gate 16/16 UNCHANGED): per-layer interleaved routing off `config.layer_types` — sliding_attention layers use plain NeoX rope (theta 500000) + finite window (masked at the FA kernel, inert for short contexts), full_attention layers use a precomputed YaRN cos/sin cache (get_rope yarn: factor 8, original 8192, mscale=yarn_get_mscale(8)=1.2079=config attention_factor, indexed by real positions); dtype-aware loader (OLMo-3 is BF16 on-disk vs OLMo-2 F32). **BUT the pinned vLLM 0.25.0 oracle CANNOT run `allenai/OLMo-3-1025-7B`** — its transformers version predates OLMo-3's nested per-layer-type rope schema: `olmo2.py:143` does `rope_parameters["rope_theta"]` → `KeyError: 'rope_theta'` (top-level rope_theta not folded into standardized rope_parameters), and forcing it surfaces `Unrecognized keys {'sliding_attention','full_attention'}` → `TypeError: unhashable type: 'dict'`. So there is NO pinned-oracle SACRED bar for OLMo-3 (DEP-blocked, spec D5); our engine LOADS + RUNS it (bootstrap loaded clean). W5 SACRED gate pending an oracle that can construct OLMo-3's rope config. | ✅ [sweep-olmo2](specs/sweep-olmo2.md) | `PARTIAL` | `test_olmo2_paged_engine` 16/16 (dgx, 92 assertions, batch3 re-run UNCHANGED); `test_olmo3_paged_engine` present (oracle-blocked, skips); loader 179 tensors zero-unmapped; registry resolves both arch strings; registration `src/vllm/model_executor/models/olmo2_registry.cpp:126`; test `tests/vllm/models/test_model_registry.cpp:118` | `CLAIM-SWEEP-OLMO2` (Claude Code opus-4-8) | | `MODEL-TEXT-olmo-hybrid-olmo-hybrid-for-causal-lm` | `OlmoHybridForCausalLM` | `registry.py:174`; `vllm/model_executor/models/olmo_hybrid.py::OlmoHybridForCausalLM` | causal generation / text | model loader/forward; paged attention/KV; Mamba/SSM state; GDN/linear-attention state | ☐ required | `INVENTORIED` | none | unassigned | diff --git a/.agents/roadmap_v1.md b/.agents/roadmap_v1.md index c50732274..eb5a46133 100644 --- a/.agents/roadmap_v1.md +++ b/.agents/roadmap_v1.md @@ -84,6 +84,7 @@ issue is not yet placed. Keyed record: update in place, never append. | [#382](https://github.com/mudler/vllm.cpp/issues/382) | `KERNEL-ATTN-PAGED` | decode-opt attention kernel is head_dim-256 only; head_dim 128 (Qwen3-dense, Llama, Mistral) falls to the block kernel | perf | | [#206](https://github.com/mudler/vllm.cpp/issues/206) | `KERNEL-SSM-MAMBA` | RTX 5070 Ti: close Qwen3.5-4B TTFT, TPOT and VRAM gaps vs vLLM — owns the sm_120 post-conv token tile and the K=4 causal-conv arm (PR #155) | feature | | [#305](https://github.com/mudler/vllm.cpp/issues/305) | `KERNEL-SSM-MAMBA` | GDN causal-conv: the `conv_state` initial-state read races the final-state write across blocks (`VT_CONV_REG` + exact chunks, both default ON) | bug | +| [#496](https://github.com/mudler/vllm.cpp/issues/496) | `KERNEL-SSM-MAMBA` | Mamba2 SSD core is unported: the row carries only the GDN arm, so every hybrid-SSM arch (Nemotron-H/3.5, Granite-4, Zamba2, Falcon-H1, Mamba2) is blocked | feature | | [#352](https://github.com/mudler/vllm.cpp/issues/352) | `KV-EVENTS` | The batch envelope is never published and `kv_cache_report_mode` does not exist — events can only be observed by reaching into the `BlockPool` | feature | | [#353](https://github.com/mudler/vllm.cpp/issues/353) | `KV-EVENTS` | `KVEventsConfig` has no `__post_init__`: an enabled config with an unset publisher reports `unknown event publisher ''` instead of resolving to zmq | bug | | [#396](https://github.com/mudler/vllm.cpp/issues/396) | `KV-EXTERNAL-CACHE` | `test_lmcache_connector` data race under TSan: `MockLmcacheServer` writes non-atomic `listen_fd_` before joining its accept thread | bug | diff --git a/.agents/specs/mamba2-ssd.md b/.agents/specs/mamba2-ssd.md new file mode 100644 index 000000000..630e05b59 --- /dev/null +++ b/.agents/specs/mamba2-ssd.md @@ -0,0 +1,299 @@ +# Mamba2 SSD — the generic selective-scan core `KERNEL-SSM-MAMBA` never got + +**Claim:** `CLAIM-KERNEL-SSM-MAMBA-SSD`. **Kernel row:** `KERNEL-SSM-MAMBA` +(existing, stays `INVENTORIED` at this spec commit — see §8). +**Issue:** [#496](https://github.com/mudler/vllm.cpp/issues/496). + +**Model rows it unblocks (all `INVENTORIED`, all waiting on this one kernel):** +`MODEL-TEXT-nemotron-h-nemotron-hfor-causal-lm` (model-matrix.md:248, the +immediate driver), `MODEL-TEXT-mamba2-mamba2-for-causal-lm` (:234), +`MODEL-TEXT-granitemoehybrid-granite-moe-hybrid-for-causal-lm` (:214), +`MODEL-TEXT-falcon-h1-falcon-h1-for-causal-lm` (:194), +`MODEL-TEXT-zamba2-zamba2-for-causal-lm` (:280), +`MODEL-TEXT-jamba-jamba-for-causal-lm` (:227), +`MODEL-TEXT-lfm2-lfm2-for-causal-lm` (:229), +`MODEL-TEXT-lfm2-moe-lfm2-moe-for-causal-lm` (:230), +`MODEL-TEXT-plamo2-plamo2-for-causal-lm` (:264), +`MODEL-TEXT-olmo-hybrid-olmo-hybrid-for-causal-lm` (:251). + +**Base:** `main` HEAD `e1087a8812c9b7d96fca5a813981f378fcace638`. +**Pinned oracle:** `${VLLM_SOURCE}` = `/home/mudler/_git/vllm` @ +`5559679229bc961848b121ccdeaa8fa5d79bec98` (vLLM 0.26.0.dev0). + +**Signal (honest, up front):** this is a **host-reference-first kernel brick**, +mirroring the KDA lane ([kda-kernel-delta.md](kda-kernel-delta.md)) and the +DeepSeek-V4 DSA lane. W1 lands the SSD numerics as portable CPU references +gated against a from-first-principles double-precision reference and against +the sequential recurrence they are supposed to be algebraically equal to. The +W1 gate is **host-reference + structural review, NOT a dumped-oracle rel-L2**. +The device (CUDA) arm is W2, and the real end-to-end token gate is a NAMED +residual owned by the model row, not by this one. + +--- + +## 0. Scope (headline verdict) + +`KERNEL-SSM-MAMBA` is named "General Mamba selective scan, causal convolution, +SSD, and linear attention kernels" (kernel-matrix.md:157). Only the second and +fourth clauses were ever built. Everything landed under the row is the **GDN** +arm: the gated-delta-rule recurrence, its causal conv, its state gather/scatter +and the sm_120 tiling campaign. The row's own matrix text concedes it — +*"these GDN tests do not prove generic Mamba support"* — and repeats +"generic Mamba lifecycle stays `INVENTORIED`" at four checkpoints +(kernel-matrix.md:167-209). + +**Gated delta rule is not Mamba2.** GDN carries a delta-removal term +`(I − β kkᵀ)` and a per-head scalar decay; Mamba2's SSD is a diagonally-decayed +gated linear recurrence with **no** removal term, driven by `A_log`, a +per-token `dt` and a skip `D`, with `B`/`C` shared across `n_groups` head +groups. `kGdnPrefill` cannot be reshaped into it. This row owns exactly two new +numerical objects: + +1. **the chunked SSD scan** (`mamba_chunk_scan_combined_varlen`) — the varlen + prefill path: chunk cumulative decay, chunk-local state, inter-chunk state + passing, and the chunk-scan output combine; +2. **the single-token selective state update** (`selective_state_update`) — the + decode path, including the scattered-cache-slot indexing our GDN decode + already models. + +Plus one small third thing that is genuinely different from what we have: + +3. **the silu-gated GROUP RMS norm** (`Mixer2RMSNormGated.forward_native`, + `mamba_mixer2.py:100-149`) — `x · silu(gate)` followed by an RMS norm over + `group_size = intermediate_size / n_groups` slices. Our `kRmsNormGated` is + the GDN/KDA **sigmoid** gate over the whole row; the activation and the + reduction extent both differ, so it is a sibling op, not a parameter. + +**Out of scope, explicitly.** The Nemotron-3.5-Lightning model port — non-gated +`relu²` MoE grouped GEMM, ModelOpt `MIXED_PRECISION` per-module loading, the +MTP head, `layers_block_type` layer dispatch — is a separate row on +`MODEL-TEXT-nemotron-h-nemotron-hfor-causal-lm`. So is TP sharding of +`n_groups` (`mamba_v2_sharded_weight_loader`, `mamba_mixer2.py:174-236`); +this row lands `tp_world_size == 1` and states the shard seam it leaves open. +So is ReplaySSM (`use_replayssm`, `selective_state_update_replayssm_output_only`) +and Mamba **v1** (`mamba_mixer.py`), neither of which any scoped model needs. + +## 1. Upstream chain (`file:line` @ `555967922`) + +### 1.1 The layer — `vllm/model_executor/layers/mamba/mamba_mixer2.py` + +| What | Anchor | +|---|---| +| `Mixer2RMSNormGated` silu-gated group RMS norm | `:69-172` (native `:100-149`) | +| `MambaMixer2.__init__` (proj split, `A_log`/`D`/`dt_bias`) | `:250-547` | +| `conv_ssm_forward` — the prefill/decode split | `:687-1104` | +| decode/prefill token split (`num_decodes`, `num_prefills`) | `:738-790` | +| varlen prefill call site | `:870` (and warmup `:654`) | +| decode call site | `:1087` | +| `get_state_dtype` / `get_state_shape` | `:1105-1141` | +| `mamba_type` → `MambaAttentionBackendEnum.MAMBA2` | `:1142-1144` | + +### 1.2 The SSD ops — `vllm/model_executor/layers/mamba/ops/` + +| Kernel | File | Role | +|---|---|---| +| `mamba_chunk_scan_combined_varlen` | `ssd_combined.py:157-235` | varlen entry; `_mamba_chunk_scan_combined_fwd` at `:27-156` is the 5-stage pipeline | +| `_chunk_cumsum_fwd` + `_chunk_state_fwd` | `ssd_chunk_state.py` (407 L) | per-chunk `dA_cumsum` (with `dt_softplus`, `dt_limit`) and chunk-local state | +| `_state_passing_fwd` | `ssd_state_passing.py` (146 L) | inter-chunk recurrence, consumes `initial_states`, honours `seq_idx` boundaries | +| `_bmm_chunk_fwd` | `ssd_bmm.py` (209 L) | `CB = C·Bᵀ` per chunk, f32 accumulate (`ssd_combined.py:124`) | +| `_chunk_scan_fwd` | `ssd_chunk_scan.py` (525 L) | combines intra-chunk attention-like term + inter-chunk state term + `D` skip | +| `selective_state_update` | `mamba_ssm.py:497+` | decode; `state_batch_indices`/`null_block_id` scattered slots | +| `_layer_norm_fwd` (gated) | `layernorm_gated.py` (172 L) | the fused form of §0.3 | + +Contract points that must be mirrored, not re-derived: +`chunk_size` must be a power of 2 (`ssd_combined.py:48`); `seq_idx` is +per-chunk, `seq_idx.shape == (nchunks,)` (`:60-61`, `:189`); `initial_states` +is `(batch, nheads, headdim, dstate)` (`:79`, `:194`); `CB` accumulates in +**f32** regardless of activation dtype (`:124`); the SSM state dtype is a +separate knob from the activation dtype (`state_dtype`, `:46,119,176`). + +### 1.3 State layout — `vllm/model_executor/layers/mamba/mamba_utils.py` + +`mamba2_state_shape` (`:174-199`) returns exactly two shapes: +`conv_state = (conv_dim/tp, conv_kernel − 1 + num_spec)` where +`conv_dim = intermediate_size + 2·n_groups·state_size`, and +`temporal_state = (num_heads/tp, head_dim, state_size)`. +`mamba2_state_dtype` (`:73-81`) gives `(conv_dtype, ssm_dtype)` independently, +which is the same conv-then-temporal ordered `MambaSpec` pair our runner already +allocates (porting-inventory.md:109). + +### 1.4 The config that drives it + +`nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4` (`model_type: nemotron_h`, +52 layers = 23 mamba / 6 attention / 23 moe): `mamba_num_heads=64`, +`mamba_head_dim=64`, `n_groups=8`, `ssm_state_size=128`, `conv_kernel=4`, +`chunk_size=128`, `mamba_hidden_act=silu`, `use_conv_bias=true`, +`use_bias=false`, `mamba_ssm_cache_dtype=float32`. Wired at +`nemotron_h.py:373-389`. + +## 2. Our baseline — reuse vs new (our `file:line`) + +### REUSE (landed; the SSD recurrence rides on these unchanged) + +- `MambaSpec`, the ordered conv-then-temporal state pair, and its exact + `page_size_bytes` — `include/vllm/v1/kv_cache_interface.h`, + `src/vllm/v1/kv_cache_interface.cpp:102-127`. +- Hybrid KV coordinator and the per-group managers (full-attn left→right + + Mamba right→left single recurrent state) — + `src/vllm/v1/core/kv_cache_coordinator.cpp`, + `single_type_kv_cache_manager.cpp:652`. +- All three causal-conv arms: `vt::CausalConv1dFwd`, `CausalConv1dUpdate`, + `CausalConv1dSpecUpdate` (`include/vt/ops.h:95-97`) — Mamba2's conv is the + same op with the same persistent `conv_state`; only the channel split around + it differs. +- Scattered state slots: `kGdnStateGather` / `kGdnStateScatter` + (`include/vt/ops.h:186-187`) — the `state_batch_indices` / `null_block_id` + semantics of `selective_state_update` are the ones GDN decode already models + (`ops.h:1856`). +- `kQkvSplit`-style projection splitting, `kMatmulBT`, `kSiluAndMul`. + +### NEW (this brick) + +- `vt::Mamba2ChunkScan` — varlen chunked SSD prefill. +- `vt::Mamba2StateUpdate` — single-token decode selective update. +- `vt::RmsNormGatedGroup` — silu-gated, `n_groups`-wise RMS norm (§0.3). + +### NEW, but NOT this brick (named residuals) + +- The CUDA arm of all three (W2) — W1 is CPU host references only. +- `n_groups` TP sharding + `extra_groups_for_head_shards` + (`mamba_utils.py:187`). +- Spec-decode `num_spec > 0` state rows; the conv side already supports them + (`CausalConv1dSpecUpdate`), the temporal side does not. +- ReplaySSM, Mamba v1, `mamba_cache_mode=align` prefix retention + (already an open T1 item, porting-inventory.md:80). + +## 3. Port map (upstream → local) + +| Upstream | Local (new) | Note | +|---|---|---| +| `ssd_chunk_state.py::_chunk_cumsum_fwd` | `src/vt/cpu/cpu_mamba2_ssd.cpp` | `dt_softplus` + `dt_limit` clamp, then `dA_cumsum` per chunk | +| `ssd_chunk_state.py::_chunk_state_fwd` | same TU | chunk-local `(nheads, headdim, dstate)` | +| `ssd_state_passing.py::_state_passing_fwd` | same TU | consumes `initial_states`, resets on `seq_idx` change | +| `ssd_bmm.py::_bmm_chunk_fwd` | same TU | **f32 accumulate**, non-negotiable | +| `ssd_chunk_scan.py::_chunk_scan_fwd` | same TU | intra + inter + `D` skip | +| `ssd_combined.py::mamba_chunk_scan_combined_varlen` | `vt::Mamba2ChunkScan` in `src/vt/ops.cpp` + `include/vt/ops.h` | the 5 stages above, in upstream order | +| `mamba_ssm.py::selective_state_update` | `vt::Mamba2StateUpdate` | scattered slots + NULL row, mirroring `GdnDecode` | +| `mamba_mixer2.py:100-149` | `vt::RmsNormGatedGroup` | silu gate + group RMS | +| `mamba_utils.py::mamba2_state_shape` | `src/vllm/v1/kv_cache_spec_registry.cpp` | a second `MambaSpec` producer; no new spec type | + +Every new symbol carries the `file:line` it was ported from, per +[porting.md](../porting.md). Nothing here is written from scratch, so nothing +is owed to porting-inventory §9. + +## 4. Tests to port + +From `tests/kernels/mamba/` @ `555967922`, parameters, dtypes and tolerances +preserved, harness adaptation documented where unavoidable: + +| Upstream | Local | Covers | +|---|---|---| +| `test_mamba_ssm_ssd.py` | `tests/vt/test_ops_mamba2_ssd.cpp` | the chunked scan across `chunk_size`, `n_groups`, `seq_idx` boundaries, `initial_states` | +| `test_mamba_ssm.py` | `tests/vt/test_ops_mamba2_state_update.cpp` | `selective_state_update` incl. scattered/NULL slots | +| `test_mamba_mixer2.py` | `tests/vt/test_ops_mamba2_gated_norm.cpp` | the gated group RMS norm | + +Plus two tests upstream does not have, because our gate is not a dumped oracle: + +- **Chunked == sequential.** The SSD chunked scan must equal a naive + per-token recurrence written independently in double precision. This is the + test that actually catches a wrong `dA_cumsum` or a dropped inter-chunk term, + and it is the reason this row can gate without the GPU. +- **Chunk-boundary invariance.** The same sequence scanned at + `chunk_size ∈ {8, 16, 32, 64, 128}` must agree to the f32 tolerance. A + state-passing defect is invisible at one chunk size and loud across five — + the same failure shape as [[h3-video-decode-temporal-and-tiling-compose]], + where the gates ran below one chunk and saw nothing. + +**RED first.** Each test is committed failing for the intended reason with the +red output captured, before the implementation. The reviewer mutates the +claimed guarantee in a scratch copy — drop the `initial_states` term, widen +`dt_limit`, swap the `CB` accumulator to bf16 — and proves each mutation is +caught. `Approx` comparisons use `.scale(0.0)` +([[doctest-approx-scale-term-floor]]); pass/fail is read from the `Status` +line, not the `assertions:` count ([[doctest-assertions-line-hides-thrown-cases]]). + +## 5. Gates + +**W1 (host reference).** Focused: the three new test binaries green, with the +chunked==sequential and chunk-boundary-invariance arms passing at the stated f32 +tolerance. Full: `ctest` clean, re-run **serially** for the known parallel-flaky +set ([[flaky-under-parallel-ctest]]), clean `-Werror` on a **clean rebuild** +([[incremental-build-masks-werror]]), and a Debug arm so asserts are live +([[release-gate-masks-asserts]]). + +**W2 (device).** CUDA arm byte-compared against the W1 host reference on the +Nemotron-3.5 shapes (`nheads=64, headdim=64, dstate=128, ngroups=8, chunk=128`), +`compute-sanitizer` clean, on dgx under `flock $HOME/gpu.lock` with +`local-ai-worker` parked ([[localai-worker-down-on-dgx]]). + +**Not this row's gate.** A token-exact end-to-end comparison against the pinned +vLLM needs the model port, so it is stated here as a residual and owned by +`MODEL-TEXT-nemotron-h-nemotron-hfor-causal-lm`. When it runs, the oracle +identity is asserted before any number is believed — `vllm.__version__` + +flashinfer, abort on mismatch ([[oracle-identity-must-be-asserted]]) — and the +fixture is checked to be the checkpoint the changed path actually loads +([[sacred-27b-gate-loads-wrong-checkpoint]]). + +**No performance claim is made in this row.** The SSD kernel's throughput +against vLLM's Triton pipeline is a separate measurement with its own +both-arms `nsys` requirement ([[profile-both-arms-before-choosing-a-lever]]); +a host reference is not a speed result and will not be reported as one. + +## 6. W-breakdown + +| W | Content | Exit | +|---|---|---| +| W0 | This spec. Upstream chain read end to end, contract points extracted, reuse-vs-new settled | spec committed, issue linked in three places | +| W1 | RED tests → `vt::Mamba2ChunkScan`, `vt::Mamba2StateUpdate`, `vt::RmsNormGatedGroup` CPU references → focused green → full gate → fresh scoped review | host references gated, reviewed by a fresh agent | +| W2 | CUDA arm for all three, byte-compared to W1 | device green + `compute-sanitizer` clean on dgx | +| W3 | `MambaSpec` producer for Mamba2 shapes; het-KV group construction proven on the Nemotron-H layer pattern | spec/shape tests green; hands off to the model row | + +W1 is the load-bearing one. W0 and W3 are small; W2 is bounded by the shapes W1 +already gates. + +## 7. Risks / decisions + +- **The chunked scan is where correctness hides.** Chunk cumsum, state passing + and the scan combine are individually plausible and jointly wrong in ways a + single-chunk test cannot see. Mitigated by the two extra tests in §4, which + are the point of the design, not extras. +- **f32 discipline.** `CB` accumulates in f32 upstream (`ssd_combined.py:124`) + and the SSM state has its own dtype knob. Going *wider* than the oracle is + numerically correct and therefore invisible to a token gate while doubling + traffic ([[token-gates-cannot-see-dequant-fallbacks]]), so the memory format + is checked against the oracle explicitly, per [porting.md](../porting.md), + and every f32 buffer on this path carries a one-line reason. +- **A bf16 output arm can absorb a real reduction-order defect** + ([[bf16-store-absorbs-reduction-order-defects]]). Every new test sweeps an + **f32 output arm** as well. +- **`n_groups` sharding is deferred, not forgotten.** W1 asserts + `tp_world_size == 1` and refuses otherwise with a message naming + `extra_groups_for_head_shards`, rather than silently computing a wrong split. + An arm that is not implemented is refused, never discovered later. +- **Do not "reuse" GDN by parameterising it.** The two recurrences differ in + structure, not in constants. A shared entry point would make both harder to + gate. Sibling ops, one shared state layout. + +## 8. Now + +**State at this commit:** spec committed, implementation **not started**. +`KERNEL-SSM-MAMBA` stays `INVENTORIED` — this commit changes no lifecycle +state, so it owes no `STATUS.md` / `BENCHMARKS.md` projection. The row moves +when a fresh implementer claims W1 against issue #496. + +**Next action:** dispatch a fresh implementer for W1 with the +[implementer contract](../prompts/), scoped to `src/vt/cpu/cpu_mamba2_ssd.cpp`, +`include/vt/ops.h`, `src/vt/ops.cpp` and the three new test TUs, RED first. + +## 9. Stop conditions + +- The chunked scan cannot be made to match the sequential double-precision + recurrence at any chunk size → stop, report `NEEDS_DECISION`, do not widen + the tolerance to pass. A tolerance widened to make a scan agree with itself + is the defect, not the gate. +- A required upstream contract point is ambiguous in source *and* the oracle + cannot be run to settle it → `NEEDS_CONTEXT`; do not guess a mirrored + behaviour ([mirror, never ask how a feature should behave] applies to product + decisions, not to unread source — read the source first). +- The device arm cannot reach the host reference byte-for-byte → keep the gap + open, name the next traceable hypothesis, and never record it as a ceiling. diff --git a/docs/FEATURES.md b/docs/FEATURES.md index cfa5f6f14..b5b7e2c81 100644 --- a/docs/FEATURES.md +++ b/docs/FEATURES.md @@ -163,6 +163,7 @@ Enumerated in `.agents/model-matrix.md`, not registered, no runnable GB10 gate: | `DeepseekV3ForCausalLM`, `DeepseekV32ForCausalLM` | DeepSeek-V3 / V3.2 | 671B, ~642 GiB fp8 vs 119 GiB unified; V3.2 also DSA-indexer dep-blocked | | `GlmMoeDsaForCausalLM` | GLM-5 (DSA) | ~1404 GiB bf16; dep-blocked (GLM-5.x is DeepSeek-V3.2 verbatim) | | `MiniMaxM2ForCausalLM` | MiniMax-M2 | ~230B, ~428 GiB bf16, ~4x over the unified pool | +| `NemotronHForCausalLM` | Nemotron-H / Nemotron-3.5-Lightning-30B-A3B | 20.1 GiB NVFP4 fits the pool, so capability-blocked, not HW-blocked: Mamba2 SSD is unported ([#496](https://github.com/mudler/vllm.cpp/issues/496)), plus non-gated `relu²` MoE and ModelOpt `MIXED_PRECISION` | 27 of the 31 registered text-generation architectures carry a passing correctness gate today; the rest are honestly marked scaffold or blocked above.