[Taimys]

Overview

Avoid materializing dense initial states for CUDA Mamba2 continued-prefill calls containing one sequence with direct final-state writeback. State passing reads the initial state from the cache and preserves each chunk's entering state in the existing SSD workspace.

Fixes #62919.

Claims

• Removes the separate initial-state gather and masking operations from the eligible single-sequence path.
• Improves cached short-continuation throughput by 10.75% on the tested L20/Granite configuration, with 9.69% lower p50 TTFT and 9.58% lower p95 TTFT.
• Preserves exact numerical equality in the tested kernel cases and matching token outputs in the serving comparison.
• A mixed-prefill serving workload showed approximately unchanged throughput (+0.50%).

Validation

Correctness

.venv/bin/python -m pytest tests/kernels/mamba/test_mamba_ssm_ssd.py -q

385 passed, with no failures, errors or skips, on one NVIDIA L20.

Coverage includes FP16/BF16/FP32 state storage, multiple SSD shapes, in-place updates, shared/crossed cache slots, masked initial states, checkpoint/returned-state fallbacks, and protection against destructive autotuning. Numerical comparisons use rtol=0, atol=0.

The suite includes 18 gather-warmup regressions using strided block-table column views. They require compile-only startup warmup and reject new gather compilation during the subsequent SSD call. A separate negative control using stride-1-only warmup triggers the expected failure.

A compiled Granite mixed-prefill/decode run completed 48 requests with 1,116 exact layer-level comparisons across direct, checkpoint and multi-sequence paths, including 756 comparisons in mixed prefill/decode steps.

Serving performance

Baseline: unmodified 8407bc0bb6b0061dacb731ae3cc54b6035f112b3.

Candidate: the same revision plus this patch. Both variants use native extensions built from the baseline source; this patch changes no native code.

Environment:

• One NVIDIA L20, 48 GB, TP1.
• Granite-4.0-h-micro, revision 03a64d0592a0af98cf9e80e13415a711e825e3e1.
• Python 3.12.14, PyTorch 2.14+cu130, Triton 3.8.
• BF16 activations/state, FP8 Torch linear backend, FlashAttention 2.
• Model Runner V2, async scheduling, prefix caching, compilation and FULL_AND_PIECEWISE CUDA graphs.
• Server and client pinned to separate sets of eight CPU cores.

A custom streaming HTTP client measured two workloads:

• Short continuation: 4,224-token cached prefix + 32-token suffix, max_tokens=1, 32 concurrent single-prompt RPCs.
• Mixed prefill: the same shared prefix, eight prompts per RPC, four concurrent RPCs, suffix lengths [32, 600, 1600, 128, 32, 600, 1600, 128], and max_tokens=64. Long suffixes begin with a varying nonce to avoid additional cache hits across rounds.

Workload	Stock prompt/s	Candidate prompt/s	Change	TTFT p50, stock → candidate	TTFT p95, stock → candidate
Short continuation	55.422	61.383	+10.75%	561.530 → 507.096 ms	568.582 → 514.117 ms
Mixed prefill	18.976	19.071	+0.50%	333.848 → 327.068 ms	987.042 → 979.144 ms

Services ran serially in stock/candidate/candidate/stock order. Each variant had two independent processes, with three measured rounds per workload per process after warmup. Each round contained 64 prompts. Table values are medians of the six round-level statistics; TTFT percentiles are computed per round.

TTFT is measured per prompt from the start of its HTTP RPC to the first streamed token. Throughput is measured in prompts/s, including all eight prompts in each batched RPC.

No benchmark-specific SSD configurations were forced. Existing autotuned kernels tuned independently in each process. The candidate's indexed state-passing path used its fixed production configuration (BLOCK_SIZE=512), bypassing autotuning because trial launches could overwrite aliased input states.

Client concurrency and SSD batch size are distinct. Separate probes of the short-continuation workload observed B=1 SSD prefill calls. The mixed workload produced B=3 checkpoint-prefill calls, alongside B=1 calls and decode work. These probes ran separately; timed workers had no scheduling probes, shadow execution or profiler.

All 1,536 recorded responses (768 per variant) matched in token IDs and text at corresponding input positions across runs. Short outputs contained one token; mixed outputs reached the 64-token limit.

Serving configuration

MODEL points to a local snapshot of the model revision listed above. Both variants used the same server configuration:

VLLM_USE_V2_MODEL_RUNNER=1 \
OMP_NUM_THREADS=1 \
OPENBLAS_NUM_THREADS=1 \
.venv/bin/python -m vllm.entrypoints.openai.api_server \
  --model "$MODEL" \
  --served-model-name granite \
  --host 127.0.0.1 \
  --port 8276 \
  --dtype bfloat16 \
  --max-model-len 16384 \
  --max-num-seqs 32 \
  --max-num-batched-tokens 2048 \
  --kv-cache-memory-bytes 8589934592 \
  --gpu-memory-utilization 0.85 \
  --enable-prefix-caching \
  --block-size 528 \
  --async-scheduling \
  --quantization fp8_per_tensor \
  --kernel-config '{"linear_backend":"torch"}' \
  --compilation-config '{"cudagraph_capture_sizes":[1,4,8,16,32]}' \
  --enable-prompt-tokens-details \
  --disable-log-requests

Each process used separate vLLM and Inductor cache directories. Compiled Triton binaries were shared, while disk caching of autotuning results remained disabled.

The client used /v1/completions with temperature=0, streaming enabled and return_token_ids=true. It tracked the first token separately for each choice.index, including tokens with empty decoded text. The shared prefix was primed before workload warmup and measurement.

Details

Continued prefill currently materializes dense initial states through indexing and torch.where, even for short suffixes. This snapshot also protects the original state after final-state writeback.

For CUDA SSD calls containing one prefill sequence with direct final-state writeback, state passing now reads the indexed cache and stores each chunk's entering state in the existing workspace. Chunk scan consumes those states. The FP32 recurrence and final-state writeback arithmetic remain unchanged; cache and workspace dtypes must match.

Indexed state passing uses a fixed BLOCK_SIZE=512, bypassing autotuning to avoid destructive trial writes. Multi-sequence, checkpoint and returned-state calls retain snapshot semantics: CUDA uses the existing fused gather, while other platforms retain the PyTorch snapshot.

Gather warmup uses the existing compile-only registry and covers the alignment and stride classes produced by Mamba2 metadata. The existing SSD warmup also exercises the indexed direct path.

Scope and tradeoffs

The serving measurements cover one model/GPU configuration and two workloads. The mixed workload exercises checkpoint fallback together with single-sequence prefill and decode; it does not isolate short multi-sequence SSD calls.

Earlier prototype measurements identified host overhead in the snapshot fallback. Isolated API cost was not remeasured on the final patch, and the mixed serving result does not establish regression-free behavior for every fallback workload.

For batched completions, the tested baseline reports cached_tokens only for the last prompt. Per-prompt cache-hit counts were therefore asserted for the single-prompt workload; mixed-workload coverage was additionally checked through the independent scheduling probes.

Related work checked: #61497 addresses final-state writeback and dispatch overhead, #61501 introduces the opt-in FlashInfer scan, and #61528 targets piecewise graph capture. This change addresses the remaining initial-state snapshot in the default Triton SSD path.

Assisted by Astra.

[mergify]

Documentation preview: https://vllm--62922.org.readthedocs.build/en/62922/