[Feature]: Paged / KV-manager-resident storage for the DFlash draft context, replacing the dense buffer

[chungen04]

🚀 The feature, motivation and pitch

Problem

DFlash's draft model cross-attends over an accumulated bank of projected target hidden states. At upstream it is written as a dense, worst-case pre-allocated GPU buffer (tensorrt_llm/_torch/speculative/dflash.py,_lazy_init_ctx_buffers):

_ctx_k_buf / _ctx_v_buf : [max_batch, L_draft, max_ctx + block_size, n_kv, head_dim]  (bf16, x2 for K and V)

with max_ctx defaulting to min(runtime max_seq_len, draft max_position_embeddings). This reserves max_batch × max_ctx token-slots up front.

When launching the engine with model: MiniMax-M2.5-NVFP4, TP=2, max_batch_size=64,max_seq_len=131072 with draft model: 5 layers, 8 KV heads, head_dim 128:

DFlash: allocated ctx buffers: max_batch=64, max_ctx=131083, dtype=torch.bfloat16

• Per-token context cost: 2 × 5 × 4 × 128 × 2 B = 10 KB/token/rank
• Dense reservation: 64 × 131,091 ~ 8.39 M token-slots × 10 KB ~ 80 GB per rank
• On a 180 GB GPU, the buffer consumes much of the remaining memory. The target KV pool shrinks to the point that admission starves (utilization pins at 100%, new requests can't be scheduled).

Current design is also not able to retrieve the context on a prefix-cache hit.

vLLM's solution for reference

vLLM's DFlash implementation stores the draft context inside the standard paged KV cache system, so no dedicated buffer at all. Based on release v0.24.0:

• Draft layers are first-class members of the hybrid KV cache manager. Layer grouping explicitly accommodates "speculative decoding drafters that add extra layers to one attention type" (vllm/v1/core/kv_cache_utils.py).
• Projected context K/V is scattered into the draft layers' paged cache via slot mappings, with norm/RoPE (and optionally quantization) applied at store time: precompute_and_store_context_kv(states, positions, slot_mappings) (vllm/model_executor/models/qwen3_dflash.py,vllm/models/deepseek_v4/nvidia/dspark.py).
• Per-iteration metadata (context positions, context/query slot mappings) is built on the GPU from the block table by a single Triton kernel (copy_and_expand_dflash_inputs_kernel, vllm/v1/spec_decode/utils.py), making zero host-side table maintenance, also CUDA-graph-clean (vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py).
• The scheduler natively accounts for the draft context as ordinary lookahead (num_lookahead_tokens = num_spec_tokens + 1,vllm/v1/core/sched/scheduler.py).
• Because context lives in ordinary cached blocks, prefix-cache hits restore the draft context.

Constraints discovered when porting this to TensorRT-LLM

• Upstream flash_attn's paged read requires 256-token pages. DFlash's cross-attention currently calls flash_attn_with_kvcache. Its paged path asserts Paged KV cache block size must be divisible by 256, so it cannot consume the manager's tokens_per_block=32 blocks. (vLLM avoids this because its attention backends read 16/32-token pages natively.)
• Manager pools share one dtype. Manager-resident context inherits the target's KV dtype (e.g. fp8), so store-time quantization of the projected context is required.
• Scheduler accounting: once draft-layer context grows inside the manager, admission must budget it.

Proposed solution

Register the draft layers as spec layers of the target manager. Some machinery already exists:

• get_num_spec_layers / get_pp_layers append spec layers.
• _build_per_layer_num_kv_heads already builds heterogeneous per-layer head lists
• Multi-pool handles differing geometry. get_buffers(layer_idx) exposes per-layer paged views for the store path.

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.