[https://nvbugs/6443620][fix] Reserve overlap-scheduler slack in gen KV capacity for one-model spec decode

[dc3671]

Background

With the overlap scheduler enabled and one-model speculative decoding (MTP), the KV-cache-manager V2 scheduler grows a generation request's KV capacity by 1 + draft_len per scheduled step, computed from host-side bookkeeping. However, under overlap the device-side KV position runs one verification round ahead of the host (up to draft_len accepted-but-uncommitted tokens), and the MTP draft layers then append draft-token KV beyond that position. The device's write reach can therefore exceed the host-allocated capacity by up to draft_len tokens.

This overrun is harmless while it stays inside an already-allocated block, but whenever it crosses a tokens_per_block boundary, applyMLARopeAndAssignQKVKernelGeneration looks up a block ordinal whose block-offset entry is still the empty sentinel (-1). A -1 offset computes an address below the pool base, producing:

torch.AcceleratorError: CUDA error: an illegal memory access was encountered

(reported asynchronously at the next sampler-event sync; the GPU coredump shows a Warp MMU Fault on the FP8 KV-append STG in the MLA rope generation kernel).

Observed with DeepSeek-V4 attention-DP (DEP8) + MTP3 at 128k ISL on a disaggregated gen worker: with the prompt length near a block boundary, the first sliding-window boundary crossing after prefill faults ~20–30 decode iterations in, deterministically. Debug instrumentation caught the exact inconsistency: a step ran with device kv_len = 131071 while manager capacity was 131070 (the grow to 131074 landed one scheduling round later), and the draft-KV append touched block ordinal 512 whose offsets entry was still -1.

Notably, TP-sharded attention (TEP) only escaped this because its measured MTP acceptance was near zero in the same setup, so the device position never outran the allocation — the bug is acceptance-rate-dependent, not attention-DP-specific.

Summary

Grow generation KV capacity by 1 + 2 * draft_len instead of 1 + draft_len — one draft_len for the step's draft tokens, one as slack for the overlap scheduler's advance booking. revert_allocate_generation and extend_capacity_for_tokens are updated symmetrically.

Impact

• Cost: at most draft_len extra reserved KV tokens per in-flight generation request (transiently one extra block near block boundaries; freed at request completion). Measured kv_cache_util impact in the repro: +0.2% over a 1024-token generation.
• No effect when speculative decoding is disabled (draft_len == 0) or for two-model spec decode (which reserves via dummy draft tokens in the scheduler).

Validation

• Repro (DSV4 DEP8 + MTP3, bs1/rank, 128k ISL, disaggregated gen worker, overlap scheduler + CUDA graphs): faulted at decode iter 19–28 in five consecutive runs before the fix; with the fix, the same config runs the full benchmark to completion (775+ iters, zero IMA), verified both with and without a debug validator that checks every block-offset entry the MLA rope kernel will address.
• One-axis ablations confirmed the trigger is the overlap+MTP capacity under-reservation: MTP0 clean, CUDA-graphs-off still faults, lm-head-TP-in-ADP-off still faults, CTX-server count and benchmark client irrelevant.

Follow-up (2nd commit): reclaim the slack each iteration

Long-decode validation under MTP3 (128k ISL / 1k OSL, disaggregated gen worker) exposed a defect in the first commit: the extra draft_len slack was granted on every scheduled iteration, but update_resources only rewinds the rejected drafts (py_rewind_len = draft_len - num_accepted), which balances just the recurring 1 + draft_len part. Net capacity growth per iteration was 1 + accepted + draft_len while the sequence only grows 1 + accepted — the slack compounded, leaking draft_len tokens of KV capacity per generation iteration. After ~750 iterations (MTP3) the leak reached ~2.3k tokens (~9 blocks at tokens_per_block=256), overran the max_blocks_per_seq-sized host_kv_cache_block_offsets rows, and crashed the scheduler in kv_cache.resize():

ValueError: User-provided base page indices is too short

right as requests approached the end of their 1024-token decode (reproduced 6/6 on DEP8/DEP16 × batch 1/2/4; MTP0 is unaffected since draft_len == 0).

The second commit keeps the slack a constant offset over the request lifetime: every grant site (try_allocate_generation, the draft-manager prepare_resources path, and the CUDA-graph padding top-up in extend_capacity_for_tokens) records the granted slack in _pending_overlap_slack, revert_allocate_generation deducts it, and update_resources reclaims it alongside the rewind. Capacity at forward prep still always includes the fresh 1 + 2 * draft_len growth, so the IMA protection from the first commit is preserved, while steady-state net growth returns to 1 + accepted per iteration.

Validation of the follow-up: the DSV4 DEP8 + MTP3 128k/1k gen-worker benchmark that previously crashed at iter ~757 in all six swept configs now runs to completion (1027+ iterations, all requests reaching the full 1024-token OSL, zero resize errors), on top of the original repro staying IMA-free.

🤖 Generated with Claude Code