[Perf][Hybrid] 3D-grid tiling of the state-copy Triton kernels

[fuscof-ibm]

Purpose

Follow-up to #48110.

#48110 optimized the copy of the states for hybrid models using 64bit load/stores. However the state copy is still underutilizing the HBM memory at small batch and the kernel requires 8B-aligned state tensors.

This PR further improves the performance by adding a 3D grid and lift the hard 8B alignment precondition by using a head/body/tail pattern. Those are the main changes:

1. Tiles the temporal copy across CTAs. The uint64 body range of each temporal state is partitioned across TEMPORAL_TILES CTAs along a new program_id(2) axis. Grid goes from (num_reqs, total_states) to (num_reqs, total_states, TEMPORAL_TILES). Small-batch runs (reqs=1..16) now fill the SMs instead of leaving them idle.
2. Extends the 64bit load/store path to the SD conv state. SD Conv copies previously went through a separate byte-wise route; they now share _copy_mamba_state_block and the same fused launch, with tile_idx > 0 early-returning so conv still runs single-CTA (it's ~60–80 KiB per block, compared to 1-4MB of the temporal part).
3. Drops the hard 8B alignment precondition. _memcpy_u64_tiled now splits the copy into head (≤7 B to align dst to 8 B) / body (u64-vectorized) / tail (≤7 B). The setup-time asserts become one-time warning_once calls — misalignment is a throughput hint, not a correctness failure.
4. Refactors the shared body into _memcpy_u64_tiled so postprocess_mamba_fused_kernel and precopy_mamba_align_fused_kernel so both kernel can use it and the 64 aligned copies can be tested in isolation.
5. Tunes TEMPORAL_TILES = 16 based on a sweep over {8, 16, 32} on GB200 with Qwen3.5-9B and Nemotron-3-Super-120B-A12B.

Files touched: vllm/v1/worker/mamba_utils.py, plus two test files.

Net Effect: 3D-grid tiling lifts small-batch kernel bandwidth by 6–7.7× at reqs=1 and pins post-copy throughput at ~72–84% of GB200 HBM3e peak bandwidth across reqs ∈ {1..128} (vs. main's 11–79% swing). End-to-end on Qwen3.5-9B + MTP + prefix caching: +1.1–1.8% output tok/s and −0.7 to −1.6% median TPOT / ITL / E2EL across conc ∈ {16, 32, 64}, with MTP acceptance unchanged and no regressions.

Test Plan

Unit tests (new):

python -m pytest tests/kernels/mamba/test_memcpy_u64_tiled.py -v
python -m pytest tests/kernels/mamba/test_precopy_mamba_align.py -v

• test_memcpy_u64_tiled.py: byte-identity check of _memcpy_u64_tiled vs. a torch slice copy across copy_size ∈ [0, 4 KiB], all sub-8B src/dst alignment offsets, and NUM_TILES ∈ {1, 16}. Guards head/body/tail correctness and the tile partitioning at the boundaries.
• test_precopy_mamba_align.py: existing V1-vs-V2 copy-spec equivalence test, now parametrized over TEMPORAL_TILES ∈ {1, 16}.

Integration tests (existing hybrid coverage):

python -m pytest tests/models/language/generation/test_hybrid.py -v

Microbenchmark

Reused the benchmarks/bench_copy_mamba_state_block.py developed for #48110 and extended to support a) multiple models and b) the 3D grid launch. In particular we use 36a924e for 3D grids (as in this PR) and 5cce1bc for 2D grid (as in main).

Sweep num_reqs ∈ {1, 4, 8, 16, 32, 64, 128} on GB200 to confirm tiling improves small-batch throughput without regressing at large batch.

End-to-end

Exact same benchmark as in #48110, but using MRv2.

For the main we use

vllm bench serve against vllm serve with MTP and prefix caching enabled (which exercises the state copies):

VLLM_USE_V2_MODEL_RUNNER=1 \
vllm serve \
    --model Qwen/Qwen3.5-9B \
    -tp 1 -pp 1 -dp 1 \
    --language-model-only \
    --reasoning-parser qwen3 \
    --speculative-config '{"method":"qwen3_next_mtp","num_speculative_tokens":2}' \
    --enable-prefix-caching

Benchmark workload:

    vllm bench serve \
        --dataset-name random \
        --seed 42 \
        --temperature 0 \
        --random-input-len "$INPUT_LEN" \
        --random-output-len "$OUTPUT_LEN" \
        --num-prompts "$NUM_PROMPTS" \
        --request-rate inf \
        --max-concurrency "$conc" \
        --ignore-eos \
        --percentile-metrics "ttft,tpot,itl,e2el" \
        --metric-percentiles "50,90,99"

Workload: --dataset-name random, INPUT_LEN=500, OUTPUT_LEN=5330, NUM_PROMPTS=500.

The workload is chosen to exercise the _copy_mamba_state_block which happens when a block is crossed.
For Qwen3.5-9B the block size is 533. The output length is chosen to be 10x the block size.

Hardware: NVIDIA GB200. Concurrency sweep {16,32,64}. At each concurrency we ran the benchmark 4 times per branch and discarded the 1st run to strip warmup/JIT effects; reported numbers aggregate the remaining 3 runs

Test Result

Unit Tests

Unit and integration tests: all passing

Microbenchmark

Raw results:

microbenchmarks_results.txt

Throughput

Peak = ~8 TB/s for GB200 HBM3e

Qwen3.5-9B

reqs	MAIN (GB/s)	PR (GB/s)	Gain	MAIN % peak	PR % peak
1	585.7	4519.1	7.72×	7.3%	56.5%
4	2073.5	5683.7	2.74×	25.9%	71.0%
8	3610.2	6265.6	1.74×	45.1%	78.3%
16	5360.5	6521.5	1.22×	67.0%	81.5%
32	6292.4	6659.7	1.06×	78.7%	83.2%
64	5925.7	6713.5	1.13×	74.1%	83.9%
128	6216.5	6748.3	1.09×	77.7%	84.4%

On Qwen3.5-9B the PR delivers > 70% of the peak starting from reqs=4 while main requires reqs=23 to achieve that value.

Nemotron-3-Super-120B-A12B-NVFP4

reqs	MAIN (GB/s)	PR (GB/s)	Gain	MAIN % peak	PR % peak
1	949.9	5740.9	6.04×	11.9%	71.8%
4	3107.1	6320.0	2.03×	38.8%	79.0%
8	4918.8	6504.7	1.32×	61.5%	81.3%
16	6524.3	6617.2	1.01×	81.6%	82.7%
32	5790.2	6669.8	1.15×	72.4%	83.4%
64	6263.5	6692.7	1.07×	78.3%	83.7%
128	6315.3	6702.9	1.06×	78.9%	83.8%

Unlike main, on Nemotron-3-Super this PR achieves an almost flat curve starting from reqs=1.

---

Kernel Latency

Qwen3.5-9B — Kernel Latency (μs)

reqs	MAIN	PR	Speedup
1	178.58	23.15	7.71×
4	201.78	73.61	2.74×
8	231.78	133.55	1.74×
16	312.20	256.62	1.22×
32	531.92	502.58	1.06×
64	1129.67	997.11	1.13×
128	2153.65	1983.94	1.09×

This PR is provides lower latencies at all regimes, with larger gains below reqs=16`.

Nemotron-3-Super-120B-A12B-NVFP4 — Kernel Latency (μs)

reqs	MAIN	PR	Speedup
1	360.15	59.59	6.04×
4	440.41	216.52	2.03×
8	556.40	420.74	1.32×
16	838.95	827.17	1.01×
32	1890.63	1641.31	1.15×
64	3495.56	3271.35	1.07×

End-to-end

Raw results:

tiled_16.txt
tiled_32.txt
tiled_64.txt

Values below are mean ± stddev across 3 runs (sample stddev, n=3).
Δ rows compare PR to main; for latency metrics negative is better, for throughput positive is better.

Metric	Branch	conc=16	conc=32	conc=64
Output tok/s	main	6219.19 ± 10.28	9744.67 ± 5.37	13848.81 ± 37.01
Output tok/s	PR	6331.49 ± 9.81	9864.17 ± 10.30	14004.27 ± 20.58
Δ throughput		+1.81 %	+1.23 %	+1.12 %
Median TPOT (ms)	main	2.407 ± 0.006	3.027 ± 0.006	4.113 ± 0.025
Median TPOT (ms)	PR	2.373 ± 0.006	2.987 ± 0.006	4.063 ± 0.032
Δ median TPOT		−1.39 %	−1.32 %	−1.22 %
Median ITL (ms)	main	7.003 ± 0.006	8.623 ± 0.012	11.773 ± 0.032
Median ITL (ms)	PR	6.890 ± 0.010	8.517 ± 0.015	11.693 ± 0.021
Δ median ITL		−1.62 %	−1.24 %	−0.68 %
Mean E2EL (ms)	main	13224.01 ± 15.90	16889.25 ± 7.79	23644.57 ± 32.34
Mean E2EL (ms)	PR	13010.79 ± 8.04	16680.60 ± 25.10	23416.33 ± 40.85
Δ mean E2EL		−1.61 %	−1.24 %	−0.97 %
Acceptance rate (%)	main	98.68 ± 0.05	98.61 ± 0.03	98.60 ± 0.09
Acceptance rate (%)	PR	98.67 ± 0.03	98.51 ± 0.08	98.62 ± 0.05

Summary

• Consistent ~1.1–1.8% throughput improvement across all concurrency levels, well above the ~0.1–0.3% run-to-run noise (visible from the stddevs).
• Latency reductions track throughput gains: TPOT, ITL, and E2EL all improve by ~0.7–1.6%, monotonically decreasing as concurrency rises.
• Gain shrinks as concurrency grows: at conc=16 tiling helps more when the 2D grid is smaller
• Speculative decoding is unaffected: acceptance rate stays ~98.5–98.7% and acceptance length stays 2.97 on both branches — this is purely a postprocess optimization with no semantic effect on drafts.
• No regressions on any metric at any concurrency level.

Duplicate-work check

Per AGENTS.md:

gh pr list --repo vllm-project/vllm --state open --search "postprocess_mamba"
gh pr list --repo vllm-project/vllm --state open --search "mamba tiled"
gh pr list --repo vllm-project/vllm --state open --search "_copy_mamba_state_block"

This PR is a strict follow-on to the earlier postprocess_mamba uint64 vectorization PR — same file, same kernel, but pushes past the single-CTA-per-copy ceiling and drops the alignment precondition. Not duplicating any open work.
Tested on MRv2.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[ZJY0516]

@claude review

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @fuscof-ibm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[tdoublep]

LGTM

[tdoublep]

/ci retry

[github-actions]

✅ No earlier CI build exists for this PR. Use /ci run first.

[tdoublep]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83111 for commit 6748a62bea87.

[fuscof-ibm]

/ci run

[tdoublep]

/ci run

[github-actions]

✅ Triggered Buildkite CI #83145 for commit fee7490c6f1e.