[Perf] Mamba2 selective_state_update up to 2x faster on B200 via 8x1 launch config for dstate 128 (+7.5% Nemotron-3-Super serving)

[TarzanZhao]

Summary

The Mamba2 decode kernel selective_state_update is memory bound: each step it reads and writes every request's fp32 SSM state. For 64 < dstate <= 128 it launches with BLOCK_SIZE_M=4, num_warps=4, so each thread handles only 4 elements. That is too few loads in flight to fill HBM, and on B200 the kernel stays near 3.1 TB/s.

This PR changes that one entry to BLOCK_SIZE_M=8, num_warps=1, 32 elements per thread. On one B200 the kernel is 1.19x, 1.86x and 2.02x faster at batch 8, 128 and 512 (3.1 -> 6.3 TB/s at batch 512), with bitwise-identical outputs. Nemotron-3-Super on 8x B200 serves +7.5% output tokens/s. H200 numbers are in the comment below.

Speedup Result

Kernel alone, one B200, CUDA graph timing with cold L2 (headdim 64, dstate 128, fp32 state, bf16 inputs, 16 heads per GPU as in Nemotron-3-Super TP8):

batch	baseline, us	this PR, us	speedup
8	4.77	4.02	1.19x
128	44.7	24.1	1.86x
512	173.5	85.7	2.02x

Other launch configs at the same batches, all bitwise identical to the baseline. Of the 20 configs I tried, (8, 1) is the fastest at batch 128 and 512, and 4% behind the fastest, (8, 4), at batch 8.

(BLOCK_SIZE_M, num_warps)	batch 8, us	batch 128, us	batch 512, us
(4, 4), baseline	4.77	44.7	173.5
(4, 1)	4.26	36.6	138.1
(8, 1), this PR	4.02	24.1	85.7
(16, 1)	4.04	24.2	86.1
(16, 4)	3.96	25.6	94.7

End to end, 8x B200 (TP8), sglang.bench_serving, ShareGPT, 1000 prompts, request rate inf:

	baseline (567d5925fe)	this PR	change
output token throughput, mean of 3	10,505 tok/s (10,499 / 10,516 / 10,499)	11,289 tok/s (11,283 / 11,262 / 11,321)	+7.5%
peak GPU memory per GPU, max of 3	164.2 GB	163.9 GB

Correctness Verification

I ran the baseline and this branch as servers and compared greedy completions for the same requests: the token id and logprob of every generated token. I also ran the kernel alone with both configs from the same inputs and initial state, and compared its output and updated state bit for bit. I set the tolerance before the runs: exact equality, because the baseline server reproduces its outputs bit for bit across restarts. The scripts are not part of this PR, and I can share them.

Every compared value is identical.

recorded	baseline vs this PR	tolerance
greedy token ids and logprobs, 64 ShareGPT prompts x 128 tokens, 3 server runs	identical	exact
same, 8 prompts of 3.2k to 3.5k tokens x 128 tokens, 3 server runs	identical	exact
kernel output and updated state, 5 head counts x 184 batch sizes (920 points)	bitwise identical	exact

Checklist

• Format: ruff format and ruff check --select=F401,F821,UP037 (the pre-commit rules) pass on the changed file.
• Unit tests: no new test. test/registered/layers/mamba/test_mamba_ssm.py covers this path and passes, 340 passed on B200.
• Documentation: none needed. The change is internal to a kernel launcher.
• Accuracy and speed results: above.
• Code style: one-value change inside the existing table, no new comments.

Scope: this affects every model that runs Mamba2 decode through the default Triton selective_state_update with 64 < dstate <= 128, on any GPU. I measured only B200. For the same shape, vLLM's tuned configs for H100 and H200 mostly pick 1 or 2 warps rather than 4x4, which points the same way but is not a measurement.

Details: hardware, model, full command, traces

Hardware. 8x NVIDIA B200 (183 GB, sm_100), driver 580.126.20, CUDA 13.0, torch 2.13.0+cu130, triton 3.7.1, flashinfer 0.6.18, sglang at 567d5925fe.

Model. nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16, TP8. 88 layers: 40 Mamba2, 40 MoE, 8 attention. Per rank the Mamba2 layers have 16 heads, headdim 64, dstate 128, and an fp32 SSM cache (mamba_ssm_cache_dtype in the checkpoint config).

Kernel measurement.

• Shapes follow the model: state pool [pool, H, 64, 128] fp32 with random state_batch_indices; x [bs, H, 64] and B/C [bs, 1, 128] are strided views of one tensor; dt, A, D and dt_bias are stride-0 broadcasts.
• Timing: max(20, K) calls captured in one CUDA graph, 20 replays timed with CUDA events, median of 5 repeats. Call i reads state slot set (i mod K); the K sets are disjoint and together hold at least 3x the 126.5 MiB L2, so every call reads state from HBM, as in real decode where each layer's state is touched once per step. Without this rotation, batches whose state fits in L2 looked up to 2.6x faster than they are.
• Speed tables: 16 heads, batch 8, 128, 512, 20 launch configs. Bitwise check (Correctness table): heads 8, 16, 32, 64, 128; batch 1 to 64, then every 8 up to 1024.
• The config is switched through a temporary override in the launcher, so both configs run the same compiled kernel code.

End-to-end measurement.

• Measured on base 567d5925fe with only this one-line change applied. The branch is rebased onto ec070ec8c8; mamba_ssm.py is identical at both, and test_mamba_ssm.py passes on the rebased branch.
• Both versions run the same server command, including --cuda-graph-max-bs-decode 1024. That flag is not part of this PR; without it, decode batches above 512 run without CUDA graphs on this workload, which hides kernel differences.
• 3 interleaved runs per version on one node, warm caches, ShareGPT 1000 prompts, seed 42.
• Throughput is measured over the Benchmark duration window of bench_serving: first request sent to last response, excluding server startup, dataset loading and warmup.

Server, both versions:

python -m sglang.launch_server --model-path nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16 \
  --tp 8 --trust-remote-code --disable-radix-cache \
  --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 50}' \
  --cuda-graph-max-bs-decode 1024

Client:

python -m sglang.bench_serving --dataset-name sharegpt --num-prompts 1000 --seed 42 --request-rate inf --flush-cache

Settings tried and rejected.

• --mamba-backend flashinfer: similar end-to-end speed, but greedy outputs are not bitwise identical to the Triton kernel.

Traces (torch.profiler, 30 decode steps at batch ~360, ranks 0 and 7; open in https://ui.perfetto.dev):

• baseline: (local: commits/c1-flag-cuda-graph-max-bs-decode-1024/trace/decode_graph_rank*.trace.json.gz)

🤖 Generated with Claude Code

[TarzanZhao]

Same kernel measured on H200 (sm_90), for reference alongside the B200 numbers above.

Kernel alone, one H200, CUDA graph timing with cold L2 (headdim 64, dstate 128, fp32 state, bf16 inputs, 16 heads per GPU; x, B and C strided views of one tensor and dt / A / dt_bias head-expanded as in the model; the calls in each graph rotate over disjoint state slots so every call reads at least 180 MiB of state, 3x the 60 MiB L2; median of 5 replays):

batch	baseline, us	this PR, us	speedup
8	5.67	4.73	1.20x
128	47.1	37.3	1.26x
512	180.1	139.6	1.29x

The same launch configs as in the description, at the same batches, all bitwise identical to the baseline (outputs and updated state). (8, 1) is 7% behind the fastest at batch 8; the fastest at batch 128; the fastest at batch 512.

(BLOCK_SIZE_M, num_warps)	batch 8, us	batch 128, us	batch 512, us
(4, 4), baseline	5.67	47.1	180.1
(4, 1)	5.42	42.9	160.7
(8, 1), this PR	4.73	37.3	139.6
(16, 1)	5.40	38.4	142.9
(16, 4)	4.44	38.1	144.0