NemotronH runs the PREFILL mamba kernels at decode: vt::Mamba2StateUpdate is registered, gated, and has ZERO model callers

[localai-bot]

Found by a source audit of the decode gap against the pinned vLLM (5559679229bc). Nothing here is profiled — every figure is arithmetic from the released config.json or a hand count from call sites, and each is labelled as such.

The divergence

vLLM decode (vllm/model_executor/layers/mamba/mamba_mixer2.py, if has_decode: :981) calls the single-step recurrent kernels:

• causal_conv1d_update(...) :1013
• selective_state_update(...) :1087

Both take state_batch_indices / conv_state_indices and read+write the cache in place at the slot.

We call the chunk-scan prefill kernels for BOTH prefill and decode — vt::CausalConv1dFwd + vt::Mamba2ChunkScan wrapped in GatherNemotronHState / ScatterNemotronHState (src/vllm/model_executor/models/nemotron_h_device.cpp:1439,:1459, as of PR #1289).

The decode kernels are already ported and called by nothing

grep -rn Mamba2StateUpdate src/vllm/   ->  0
grep -rn Mamba2StateUpdate src/vt/     -> 15

• CUDA registration src/vt/cuda/cuda_gdn.cu:6672 -> src/vt/cuda/cuda_mamba2_ssd.cuh::Mamba2StateUpdateKernelCuda:648
• CPU arm src/vt/cpu/cpu_ops.cpp:3408
• Spec .agents/specs/mamba2-ssd.md §W2, landed 2026-08-13

vt::CausalConv1dUpdate does have callers (qwen3_5.cpp:4744, kimi_linear_device.cpp:1811) — just not NemotronH.

This is a live AGENTS.md "Nothing lands dead" case: a capability that landed, is gated, and is reachable from no production entry point.

Cost, at T=1, chunk_size=128, nchunks=1

Arithmetic from cuda_mamba2_ssd.cuh::Mamba2ChunkScanKernelCuda:577-645, not measured:

per mamba layer per token	chunk scan (ours)	state update (vLLM shape)
kernel launches	5	1
cudaMallocAsync/FreeAsync	5 + 5	0
cudaMemsetAsync	2, of 2.50 MiB	0
scratch allocated	4.56 MiB	0
M2ChunkScanKernel grid	524,288 elems = 128x one token	4,096 elems

x23 layers = 105 MiB scratch, 57.5 MiB zeroed, 230 driver alloc/free, 115 launches per token in the SSD alone. Plus gather/scatter state churn the decode kernels do not perform: the SSM page is 2.00 MiB/layer f32, read+written by gather and scatter and again by the final_states copy-back, roughly +414 MiB/token.

The tree already warns about this

src/vllm/model_executor/models/qwen3_5.cpp:4730-4746:

Passing the state indices to the op eliminates the per-request gather+scatter -- the two host<->device copies per sequence per layer that dominate the decode memcpy tax.

NemotronH does the thing that comment warns against.

Confirm or refute

A/B a decode-only window with vt::Mamba2StateUpdate + vt::CausalConv1dUpdate on the nd rows against the current chunk-scan arm, same binary. Refuted if per-token time moves less than 3%.

Semantics look drop-in: Mamba2StateUpdateKernelCuda:657 computes hpg = H/G, so n_groups=8 is handled, and it takes state_indices, so the gather/scatter goes with it.

Not verified: whether the state-update arm is numerically equivalent to the chunk scan at T=1 for n_groups=8. .agents/specs/mamba2-ssd.md §8.3 declares an equivalence contract; only the launcher signature was read, not the gate evidence.

Context

This is residual work AFTER the two host-side arms land — A2-Q1 (#1289) and A2-Q2b. Those two account for the measured 718.2x on their own: they re-expand 1.24e9 elements per token (2.49 GB) on the CPU. This issue is what remains once they are gone.