[Perf] Fuse DSV4.1 input metadata preparation with Triton

[WoosukKwon]

Purpose

DSV4.1 prepares token-to-request mappings and flattened indexer decode metadata with chains of PyTorch operations on every model step. Replace those operations with Triton kernels that write directly into the existing buffers:

• Build token-to-request mappings from device query boundaries, including zero-length requests and padding. Device boundaries are necessary for DSpark adaptive verification, where CPU boundaries can be stale.
• Fuse SM100 varlen indexer sequence-length expansion, block-table expansion, request indices, per-request decode lengths, and padding initialization.
• Add focused correctness tests for flattened decode metadata and device token-to-request mappings, including CUDA-graph replay with stale CPU boundaries.

The replaced paths call Triton directly. No new JIT warmup-framework wrappers or registrations are added. The existing non-varlen indexer implementation is unchanged.

Duplicate-work checks: searched open PRs for DeepSeek V4 metadata, indexer metadata triton, token_to_req_indices, and metadata DSpark. No existing PR addresses this replacement. Related work is distinct: #56344 fuses attention math, #56254 changes sparse MQA-logit computation, #52628 enables ROCm fused draft-step updates, and #55692 hardens the downstream top-k block-table gather. No associated issue was supplied.

Test Plan

Tested against 9dcf6bf344caa7793bae0b45a7896d3f8e03a01a on four NVIDIA GB200 GPUs, with deepseek-ai/DeepSeek-V4.1-Flash, TP=4, and FP8 KV cache.

.venv/bin/python -m pytest \
  tests/v1/attention/test_kpool_tail_slot_mapping.py \
  tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py -q

.venv/bin/pre-commit run --files \
  vllm/v1/attention/backend.py \
  vllm/v1/attention/backends/mla/indexer.py \
  vllm/v1/attention/ops/metadata.py \
  tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py

Test Result

• 30 attention tests passed, including zero-length requests, strided block tables, padding, and CUDA-graph replay with changed device boundaries and stale CPU boundaries.
• All 24 metadata-builder cases matched baseline tensors exactly (rtol=atol=0).
• Exact live metadata comparisons passed on all four ranks for ordinary decoding, DSpark, and adaptive verification, covering prefill/mixed/decode.
• Pre-commit and targeted mypy 3.12 passed; git diff --check passed.

Metadata profile and microbenchmark

Measured with a local benchmark harness that is not included in this PR. CUPTI GPU timings of the five relevant builders (SWA, MLA ratios 1/2, indexer ratios 1/2), CUDA-graph replay with cold L2, batch size 16. These are metadata-only speedups.

Case	Before (µs)	After (µs)	Speedup
Prefill	168.29	118.21	1.42×
Chunked prefill	203.23	153.74	1.32×
Mixed	220.42	72.48	3.04×
Decode	191.10	46.27	4.13×
DSpark verification	172.96	56.86	3.04×
DSpark adaptive verification	167.73	50.37	3.33×
DSpark draft	173.42	55.14	3.15×
DSpark adaptive draft	174.75	58.18	3.00×

Batch-16 decode GPU launches decrease from 77 to 17 per builder suite. Across three profiled iterations, CPU aten::repeat_interleave events decrease from 54 to zero. A full-model baseline trace also confirmed these operations occur in the input-preparation path.

Batch-size-1 TPOT, 512 input / 128 output tokens

One repeated prompt, greedy generation, full decode CUDA graphs, compilation mode 0, FlashInfer autotuning disabled. Median of four interleaved samples per implementation on the same loaded model; engine first/last-token timestamps exclude prefill. All eight generated sequences match within each mode. The optimized adaptive run includes a retained 2.142 ms outlier.

Mode	Before (ms/token)	After (ms/token)	Reduction
No DSpark	6.091	5.736	5.82%
DSpark	1.657	1.407	15.07%
DSpark + adaptive	1.781	1.457	18.20%

SPEED-Bench: 8K input / 1K output, DSpark adaptive enabled

NVIDIA SPEED-Bench throughput_8k/low_entropy, 512 source rows, normalized to exactly 8,192 DSV4.1 chat-template tokens using padding/truncation, default thinking, greedy sampling, 1,024 forced output tokens. Two runs per implementation in baseline/optimized then optimized/baseline order; max(32, 2 × concurrency) requests per run, identical seed and request order. At concurrency 1024, the source dataset is repeated to produce 2,048 requests per run.

Same loaded model, V2 runner, Rust frontend, Python vllm bench serve client, five speculative tokens, full decode graphs, compilation mode 0, FlashInfer autotuning and prefix caching disabled. Fixed server limits: max_num_seqs=256, max_num_batched_tokens=16384, max_model_len=16384, GPU memory utilization 0.9. GPU clocks were not locked. Metadata methods were switched outside model CUDA graphs; no profiling or validation wrappers during timing.

Mean TTFT, including queueing

Concurrency	Before (ms)	After (ms)	Reduction
1	195.876	193.766	1.08%
8	380.581	382.415	-0.48%
128	5024.837	4989.004	0.71%
1024	128491.480	126614.766	1.46%

Mean TPOT, excluding TTFT

Concurrency	Before (ms/token)	After (ms/token)	Reduction
1	3.091	2.494	19.31%
8	4.910	4.794	2.35%
128	24.053	23.995	0.24%
1024	48.804	48.119	1.40%

Output throughput

Concurrency	Before (tok/s)	After (tok/s)	Speedup
1	304.9	373.0	1.223×
8	1432.7	1479.2	1.032×
128	4202.3	4224.8	1.005×
1024	4898.4	4966.1	1.014×

Latency means pool all requests across both repetitions; throughput is total output tokens / total measured duration. All 16 runs completed with exact requested token counts and zero preemptions. Concurrency-1024 TTFT includes substantial queueing behind the fixed 256 active-request limit. The concurrency-128 throughput difference is within observed run variation.

These TPOT values are not directly comparable to the earlier 512/128 test: context length, prompts, measurement boundary (streaming client vs engine), and aggregation differ. Long-prompt outputs also vary between baseline-only repetitions; universal generated-token parity is not claimed. Mean acceptance lengths before/after are 3.294/3.274, 3.177/3.153, 2.846/2.840, and 3.122/3.110 at the four concurrency levels.

GSM8K accuracy

Optimized implementation, full 1,319-question test set, repository five-shot prompts, default thinking/high effort, greedy sampling, seed 0, maximum 16,384 output tokens, context limit 32,768. Extract answers from final text after </think>; unfinished reasoning is invalid.

Mode	Correct	Accuracy	Truncated
No DSpark	1277 / 1319	96.82%	0
DSpark	1279 / 1319	96.97%	1
DSpark + adaptive	1277 / 1319	96.82%	0

These are single-run speculative-mode evaluations of the optimized implementation, not before/after accuracy-equivalence measurements.

AI assistance: OpenAI Codex implemented the change and ran the reported checks and evaluations under the submitter's direction. This PR is a draft; human review of every changed line and human-run validation remain required before marking it ready.

[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.

[WoosukKwon]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88471 for commit e84929af3cea.