[Perf][Attention] Pin MLA chunked-context metadata tensors so H2D copies are truly non-blocking

[zixi-qi]

Purpose

Fix a hidden host-side stall in the MLA chunked-prefill metadata build that serializes CPU metadata prep with GPU execution on every scheduler step.

In MLACommonMetadataBuilder.build(), the chunked-context section creates three small CPU tensors and ships them to the GPU with .to(device, non_blocking=True) when assembling ChunkedContextMetadata:

• cu_seq_lens_cpu — allocated with pin_memory=True ✅
• chunk_starts — pageable ❌
• token_to_seq_tensor_cpu — pageable ❌

non_blocking=True is only actually asynchronous when the host source is pinned. From pageable memory, CUDA makes the copy host-synchronous: the CPU blocks inside cudaMemcpyAsync until the GPU stream has drained all previously launched kernels. During chunked-prefill serving of long prompts this happens on the hot path of every step, so the metadata build stalls until the previous step's forward pass finishes.

Measured on Kimi-K2.5-NVFP4 (GB200 ×4, --enforce-eager, --max-num-batched-tokens 4096, 80k-token prompts, concurrency 16): the metadata build took 67–101 ms per 4096-token step — the single largest per-step cost — while the actual tensor work in that section is microseconds (1 prefill, 1–2 context chunks). Section timers confirmed ~100% of it was the chunked-context block, i.e. the hidden stream drain.

The fix pins the remaining pageable tensors so the H2D copies are truly asynchronous (.pin_memory() on chunk_starts, pin_memory=True on token_to_seq_tensor_cpu), and — per review feedback — the same fix on the DCP path (.pin_memory() on local_chunk_starts, the only pageable non_blocking=True source there; padded_local_cu_chunk_seq_lens_cpu was already pinned). No behavior change.

Duplicate check: searched open PRs for pin_memory, token_to_seq, chunk_starts, chunked prefill pinned, non_blocking pageable — no open PR addresses this (closest hits are unrelated: #44149 is LoRA adapter device detection, #34393 is MLA gather-kernel unification).

Test Plan

A/B benchmark on a 4× GB200 node, identical server and workload, only the pin change differing. Kimi-K2.5 uses standard MLA, so long-context chunked prefill goes through the fixed ChunkedContextMetadata build. Output length is 1 so the measurement isolates prefill:

vllm serve nvidia/Kimi-K2.5-NVFP4 --trust-remote-code \
  --tensor-parallel-size 4 --enable-expert-parallel \
  --enforce-eager --max-model-len 40960 \
  --max-num-batched-tokens 4096 --max-num-seqs 64 \
  --no-enable-prefix-caching --kv-cache-dtype fp8

vllm bench serve --model nvidia/Kimi-K2.5-NVFP4 --trust-remote-code \
  --backend openai-chat --endpoint /v1/chat/completions \
  --base-url http://127.0.0.1:8200 \
  --dataset-name random --random-input-len 32768 --random-output-len 1 \
  --random-range-ratio 0 --num-prompts 48 --max-concurrency 16 \
  --request-rate inf --ignore-eos

With --max-num-batched-tokens 4096 and 32k-token prompts, every prefill step after the first runs the chunked-context metadata build, exercising the fixed path.

Lint: ruff check / ruff format --check pass on the changed file.

Test Result

48/48 requests succeeded in both runs (1.57M input tokens each, prefill-only):

Metric	Before (pageable)	After (pinned)	Δ
Benchmark duration	92.39 s	89.41 s	−3.2%
Input token throughput	17,038 tok/s	17,607 tok/s	+3.3%
Mean TTFT	27,293 ms	26,339 ms	−3.5%
Median TTFT	25,570 ms	25,007 ms	−2.2%
P99 TTFT	43,328 ms	40,985 ms	−5.4%

At 384 prefill steps per run (48 × 32k / 4096), the −3.0 s duration delta is ~8 ms saved per step in this configuration; the per-step stall grows with context length (the original 80k-prompt measurement showed 67–101 ms metadata builds), and in mixed prefill+decode serving the same stall additionally shows up as decode ITL tail latency, so end-to-end gains there are larger.

---

This PR was developed with AI assistance (Claude Code); the changed lines and benchmark results have been reviewed by the human submitter.

[ivanium]

👍 Makes sense to me. Do we need to fix the DCP path too? cc @GirasoleY @njhill

[ivanium]

Should we apply the same fix to DCP path too?

[zixi-qi]

Added!

[njhill]

Wow this is a huge speedup!

[zixi-qi]

Wow this is a huge speedup!

TBH the benchmark setting is a bit biased in favor of this change, in a realistic compute bound prefill only eager run the improvement should be much smaller. But I left the existing benchmark numbers to demonstrate that the fix has taken effect

[WoosukKwon]

Holding since I found it confusing

[zixi-qi]

Holding since I found it confusing

Thanks @WoosukKwon for the review! The changed code is actually not invoked by DeepSeek V4 so initial benchmark result are most likely due to performance variance from running decode in eager mode. Reran the benchmark with Kimi + prefill only workload and updated the PR description

[WoosukKwon]

Thanks for the update!