Mamba-2 SSD fused layer via VGpuWarp — design sketch

Status: DESIGN (not built). The profiler (MESH_PROFILE_MAMBA=1, model_mamba.rs) must first turn the inferred ~177 ms/token non-matmul cost into a measured MATMUL / CONV / SSD_STEP split before committing. This is the multi-session lever after the GPU scan (3×) + q8 f32 (1.38×) landed codestral at 4.5 tok/s.

S0 RESULT (measured 2026-07-03, MESH_PROFILE_MAMBA=1, codestral f32)

[mamba-prof] matmul=65ms (35%) conv=11ms (6%) ssd_step=108ms (59%) accounted=184ms of a ~222ms token → 83% of the token is in the layer forward (cost is NOT outside it — not coordinator/tokenizer/prefill; the fused layer is aimed right). SSD_STEP is the dominant 59% — but the scan kernel is only 8192 lanes × 128 tiny FMAs + 1 exp each (~µs), so the 108ms ≈ 1.7ms/layer is almost ENTIRELY the per-layer htod(scratch)/dtoh(y)/ SYNC around the scan + the host gated-norm. i.e. the ping-pong, confirmed. REFINED PRIORITY: S1 should target the SSD_STEP first — (a) device gated-RMSNorm so y never dtoh's to host, (b) keep the scan's inputs (conv→x/B/C) and output (y) resident so there's no htod/dtoh per layer. That single change attacks the 59%. matmul (35%) is next (resident GEMV chaining); conv (6%) is cheap but must go on-device to keep the chain resident.

S1 RESULT (measured 2026-07-03 on dedicated gnosis-codestral-fs) — NO WIN, S0 hypothesis FALSIFIED

S1 (device gated-norm + resident scan I/O, MESH_GPU_SSM_FUSED, kernel CONFIRMED compiled+engaged) gave ssd_step = 108ms fused, IDENTICAL to non-fused. Removing the per-layer dtoh(y) + host gated-norm changed NOTHING → those were NOT the ssd_step cost. So the S0 read ("scan kernel is µs-cheap; the 108ms is the host↔GPU sync ping-pong") is WRONG. The ssd_step's 108ms (~1.7ms/layer) is inside the SCAN itself: the htod(scratch dt/x/B/C) + the scan KERNEL + its resident-h traffic (4MB/layer read+write = 512MB/token) + whatever sync the launch/htod_sync forces — NOT the dtoh/norm. RE-AIM before S2/S3: finer-grained profiler INSIDE mamba2_ssd_scan (time the htod_sync vs the kernel launch vs any dtoh separately) to find which. Likely culprits: (a) the 8192-thread kernel is LOW-OCCUPANCY on the L4 (~7000 cores; need ~30k threads to hide the h-read latency → the d_state=128 loop stalls on each global h read) → optimize occupancy / stage h in shared/registers; (b) htod_sync of the scratch is synchronous per layer. The CUDA-graph idea (S3) still helps launch overhead but NOT if the kernel itself is the wall. tok/s was 4.06 (fused) vs ~4.5 baseline — a slight REGRESSION (the device-norm + resident bookkeeping added cost without removing the real bottleneck). Keep MESH_GPU_ SSM_FUSED OFF by default (it's a loss). DEDICATED SLOT WIN: gnosis-codestral-fs (separate GPU service, args=openai-server, only codestral, no swap) serves codestral reliably ISOLATED from the mesh 22B-contention that was 500'ing prod — the real daily-driver-reliability fix.

S2 RESULT (measured 2026-07-03, gnosis-codestral-fs) — ALSO NO WIN → the KERNEL is the cost

S2 (htod-pack: 4 htod_sync/layer → 1, 256→64 syncs/token) gave ssd_step = 107.5ms, unchanged. So NEITHER removing the dtoh/norm (S1) NOR cutting the htod sync-count (S2) moved the 108ms. By elimination the 108ms is the scan KERNEL launch+execution, not the host↔GPU transfers/syncs. (Caveat: the finer [mamba-scan-prof] htod/launch/tail split I added in S2 did NOT print — either the fused path bypasses the instrumented helper or the print cadence never fired — so this is elimination, not a direct kernel-time reading. FIX the profiler to confirm.) LIKELY MECHANISM: the mamba2_ssd_scan kernel's resident-h access is uncoalesced — thread c (channel) sweeps n=0..d_state, reading h[c*d_state+n]; consecutive threads at the same n are d_state=128 floats (512B) apart → every h read is a separate cache line → ~32× effective-bandwidth waste (4 useful bytes per 128B line). 8192×128 uncoalesced R+W ≈ 16GB/token effective / ~300GB/s ≈ ~100ms — matches. NEXT LEVER (S3): coalesce the h layout (store h as [n][channel] so consecutive threads hit consecutive addresses) +/or raise occupancy (8192 threads ≈ 9% of the L4; split d_state across threads). tok/s stuck at ~4.1-4.5 across S0/S1/S2 (the ssd_step is the wall). Parity held throughout (coherent code). DELIVERED VALUE UNCHANGED: codestral correct + 4.5 tok/s + dedicated reliable slot (gnosis-codestral-fs).

The problem (profiling-established)

A codestral token ≈ 222 ms. Only ~45 ms is math (f32 GEMV). The other ~177 ms is coordination: the SSM forward is CPU-orchestrated, so per layer it does

in_proj(GPU) → dtoh+SYNC → conv(CPU) → htod → ssd_scan(GPU) → dtoh+SYNC
             → gated_norm(CPU) → htod → out_proj(GPU) → dtoh+SYNC

≈ 3 host↔GPU syncs × 64 layers ≈ 192 round-trips/token. Neither side is saturated — it's latency-bound ping-pong. Mirror-skip being neutral proved the cost is sync count, not bytes moved. The fix is to keep the layer on-device so the boundary count drops from ~192 to O(1).

Why VGpuWarp is the right abstraction

VGpuWarp (Lean SSOT Gnosis/VGpuWarp.lean + Rust warp.rs) already gives:

• fork/fold = diverge/reconverge — the execution algebra for a warp program.
• WarpBatch lane-mask, batched == solo bit-identical (proved, 0 axioms) — the correctness guarantee that lets us fuse without drifting the argmax.
• The batched width-axis forward primitives (matvec_resident_batch, RWKV time_mix7_batch) — one resident weight read serves the wave.

The fused layer = express the SSD layer as one VGpuWarp program that lowers to a single replayable device graph, instead of Rust calling kernels and syncing between each.

Caveat: the vCPU resumable-inference ISA (StreamingEngine) targets the isolate/CPU substrate (resume-anywhere, checkpoint-KV). This is a device backend — VGpuWarp as the execution model, a new lowering, NOT reuse of StreamingEngine.

Lane mapping (batch=1 daily-driver)

The SSD scan is already warp-shaped: num_heads·head_dim = 128·64 = 8192 independent lanes, each sweeping d_state = 128. That is the wave; fork over (head, head_dim), fold the d_state reduction. The existing mamba2_ssd_scan.cu already does exactly this (one thread per channel) — so the scan kernel is done; the win is not re-launching around it.

NOTE: WarpBatch (one-weight-read-serves-N-requests) is a THROUGHPUT lever — pays off at batch>1. For a single daily-driver user it does NOT cut latency; the win here is on-device coherence (fewer syncs), and the within-token 8192-lane parallelism, not cross-request.

The concrete mechanism: capture the token as ONE CUDA graph

A VGpuWarp program is a graph. The lowering target is CUDA Graphs: capture the per-token forward once, replay every token → the ~192 per-op launch+sync costs collapse to one graph launch. Blocker: conv and gated RMSNorm are host ops today, so the graph would break at each. They must become device ops for the graph to be end-to-end. Then a token = [embed] → 64×[in_proj·conv·scan·gatednorm·out_proj·residual] → [finalnorm·lm_head], all on-device, one replay.

What has to become a device op (the actual work)

1. causal_conv1d_step → device kernel — trivial: per-channel FIR over conv_dim=10240 channels × d_conv=4 taps + a rolling conv buffer resident in VRAM (like the scan's h). ~one warp per channel-block. (Currently model_mamba.rs host loop.)
2. gated RMSNorm (MambaRMSNormGated) → device kernel — y·=silu(z), then rms over d_inner=8192, ×norm_w. One block-reduce. (Currently host, inside mamba2_ssd_step.)
3. Chain the GEMVs without dtoh — q8_gemv already writes its output to a resident device buffer; today matmul() dtoh's it back to a host Vec. Add a matmul_resident(out_dev, x_dev, ...) that leaves the result on-device so conv/scan/norm read it directly. This is the crux: the activations (proj_out, conv_out, y, out_proj, residual x) stay in VRAM for the whole layer.
4. Residual add + pre-norm on device — cheap element-wise / block-reduce kernels.

Then the only host↔GPU crossings per token are: upload the 1 embedding row (or keep the embed table resident + index on-device), and dtoh the final logits. ~2 syncs/token, not ~192.

Staged, measurable plan (do NOT big-bang)

Each stage rebuilds + measures tok/s + PARITY (correct code, argmax unchanged):

• S0 (done here): per-section profiler → confirm SSD_STEP + the matmul-sync gap are the cost. If a surprise (e.g. mesh coordinator round-trip, not the layer) dominates, re-aim.
• S1: matmul_resident + conv device kernel — fuse in_proj→conv→scan without the two dtoh/htod between them (activations resident). Measure: expect the biggest single drop.
• S2: gated-norm device kernel + out_proj resident — the whole layer body on-device.
• S3: CUDA-graph-capture the per-token forward (or the per-layer sub-block) — kill the residual launch overhead. Replay.
• S4 (optional): chunked prefill (chunk_size=256 SSD duality) — parallel prompt ingest; fixes the tok/s-under-reports-decode metric + the long-context (O(1)-state) story.

Gate every stage on the batched==solo bit-identical discipline: the on-device path must match the CPU reference argmax (the /debug/forward + coding-prompt parity we already use).

Realistic ceiling

Codestral = 64 layers (2× mistral's 32) + a huge in_proj, so mistral's 14 tok/s is not physically reachable. Collapsing ~177 ms of coordination toward the ~45 ms math floor + some residual per-op cost → ~7-8 tok/s is the honest target, with the SSM's flat-at- long-context edge (O(1) state) on top — the daily-driver win transformers can't match.

Files

• src/model_mamba.rs — forward_one (the CPU orchestration to fuse) + the profiler.
• src/cuda/mamba2_ssd_scan.cu — the scan kernel (done; the warp shape).
• src/cuda/q8_gemv.cu + cuda_matmul.rs — GEMV (needs a _resident output variant).
• open-source/gnosis VGpuWarp SSOT (Gnosis/VGpuWarp.lean, warp.rs) — the algebra.