From 479d5f9abe8111d3450f36834b42d57a31f314e7 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Sat, 11 Jul 2026 20:09:43 -0700 Subject: [PATCH 1/2] [None][fix] Reserve overlap-scheduler slack in gen KV capacity for one-model spec decode Under the overlap scheduler with one-model speculative decoding (MTP), the device-side KV position runs one verification round ahead of host bookkeeping (up to draft_len accepted-but-uncommitted tokens), and the MTP draft layers then append draft-token KV beyond that position. The V2 scheduler grows generation KV capacity by only 1 + draft_len from the host view, so the device write reach can exceed the allocated capacity by up to draft_len tokens. Whenever that overrun crosses a tokens_per_block boundary, applyMLARopeAndAssignQKVKernelGeneration addresses a block ordinal whose block-offset entry is still the empty sentinel (-1), producing an illegal memory access. Observed with DSV4-Pro DEP8 + MTP3 at 128k ISL (disagg gen worker): with the prompt ending block-aligned, the first window-slide boundary crossing after prefill faults deterministically ~130 tokens into decode. TEP escaped only because its MTP acceptance rate happened to be near zero, keeping the device from outrunning the allocation. Grow generation capacity by 1 + 2 * draft_len instead (one draft_len for the step's draft tokens, one for the overlap advance slack), and mirror the doubled growth in revert_allocate_generation and extend_capacity_for_tokens. Costs at most draft_len extra reserved tokens per request (transiently one extra block near boundaries); max_blocks_per_seq headroom already covers it. Validated: DEP8+MTP3 128k gen_only repro previously faulting at decode iter 19-28 now runs the full benchmark clean (775+ iters, both with and without debug instrumentation). Signed-off-by: Zhenhuan Chen --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 20 +++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 5d22a2d395f2..757199c56b88 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -1811,9 +1811,18 @@ def _effective_draft_len(self, req: LlmRequest) -> int: def _required_gen_capacity(self, req: LlmRequest, current_capacity: int) -> int: """Compute generation KV cache capacity for a request. - Grows *current_capacity* by 1 + draft tokens. + Grows *current_capacity* by 1 + draft tokens, plus another draft_len + of slack for one-model speculative decoding: under the overlap + scheduler the device-side KV position runs one verification round + ahead of host bookkeeping (up to draft_len accepted-but-uncommitted + tokens), and the MTP draft layers then append draft KV beyond that + position. Without the slack, the draft KV append can address a block + ordinal the host has not allocated yet whenever the overrun crosses a + tokens_per_block boundary — an illegal memory access in the MLA rope + generation kernel (observed with DSV4 DEP+MTP3 at 128k, where the + first window-slide boundary crossing after prefill faults). """ - return current_capacity + 1 + self._effective_draft_len(req) + return current_capacity + 1 + 2 * self._effective_draft_len(req) def try_allocate_generation(self, req: LlmRequest) -> bool: """Try to allocate one additional KV cache slot for a generation request. @@ -1853,7 +1862,8 @@ def revert_allocate_generation(self, req: LlmRequest) -> None: draft_len = self._allocated_draft_lens.pop( req.py_request_id, self._effective_draft_len(req) ) - reverted_cap = kv_cache.capacity - 1 - draft_len + # Mirror the 1 + 2 * draft_len growth in _required_gen_capacity. + reverted_cap = kv_cache.capacity - 1 - 2 * draft_len if reverted_cap < 0: return if not kv_cache.resize(reverted_cap): @@ -2059,7 +2069,9 @@ def extend_capacity_for_tokens(self, request: LlmRequest) -> None: if allocated is None: return current_draft_len = get_draft_token_length(request) - delta = current_draft_len - allocated + # Growth is 1 + 2 * draft_len (see _required_gen_capacity), so the + # padding delta scales by 2 as well. + delta = 2 * (current_draft_len - allocated) if delta <= 0: return kv_cache = self.kv_cache_map[request.py_request_id] From 683a71e9e08b7059133bea89cd1e7feadc0b8328 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Sun, 12 Jul 2026 01:36:53 -0700 Subject: [PATCH 2/2] [None][fix] Reclaim per-iter overlap-scheduler KV slack in kv_cache_manager_v2 3e7fa61d73 grew generation KV capacity by 1 + 2 * draft_len per scheduled iteration (the extra draft_len reserves overlap-scheduler slack for one-model spec decode), but update_resources only trims py_rewind_len = draft_len - num_accepted, which balances just the recurring 1 + draft_len part. Net capacity growth per iteration was therefore 1 + accepted + draft_len while the sequence only grows 1 + accepted: the slack compounded, leaking draft_len tokens of capacity per iteration. Under MTP3 at 128k ISL / 1k OSL the leak reaches ~2.3k tokens (~9 blocks at tokens_per_block=256) by the time requests near the end of decode, overrunning max_blocks_per_seq sizing of host_kv_cache_block_offsets and crashing the scheduler in kv_cache.resize() with 'User-provided base page indices is too short' (then the executor error-broadcast hangs the whole disagg instance). MTP0 (draft_len=0) is unaffected, matching observations. Track the slack granted per request in _pending_overlap_slack (allocation, draft-manager prepare_resources, and CUDA-graph padding top-up all record it; revert_allocate_generation deducts it) and reclaim it in update_resources alongside the rewind. The slack thus stays a constant offset over the request lifetime: capacity at forward prep still always includes the fresh 1 + 2 * draft_len growth, so the IMA protection from 3e7fa61d73 is preserved, while steady-state net growth returns to 1 + accepted per iteration. Signed-off-by: Zhenhuan Chen --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 53 ++++++++++++++++++- 1 file changed, 51 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 757199c56b88..38676dc77f76 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -942,6 +942,16 @@ def append_to_kv_heads_per_layer( # unbounded capacity growth. self._allocated_draft_lens: dict[int, int] = {} + # Overlap-scheduler slack (the extra draft_len in + # _required_gen_capacity's 1 + 2 * draft_len growth) granted to each + # request and not yet reclaimed. The sampler's py_rewind_len only + # covers the rejected part of the recurring 1 + draft_len growth, so + # update_resources must additionally trim this slack — otherwise it + # compounds, leaking draft_len tokens of capacity per generation + # iteration until the request's block count overruns + # max_blocks_per_seq ("User-provided base page indices is too short"). + self._pending_overlap_slack: dict[int, int] = {} + # Defensive cap for get_num_available_tokens: when host cache is # enabled, clamp_max_seq_len_for_mem may return a value that spans # both GPU and host tiers. Storing the explicit max_tokens (if set) @@ -1821,6 +1831,11 @@ def _required_gen_capacity(self, req: LlmRequest, current_capacity: int) -> int: tokens_per_block boundary — an illegal memory access in the MLA rope generation kernel (observed with DSV4 DEP+MTP3 at 128k, where the first window-slide boundary crossing after prefill faults). + + The slack must stay a *constant* offset over the request's lifetime: + py_rewind_len only rewinds the rejected part of the recurring + 1 + draft_len growth, so every call site records the extra draft_len + in _pending_overlap_slack and update_resources trims it back. """ return current_capacity + 1 + 2 * self._effective_draft_len(req) @@ -1841,7 +1856,13 @@ def try_allocate_generation(self, req: LlmRequest) -> bool: draft_len = self._effective_draft_len(req) self._allocated_draft_lens[req.py_request_id] = draft_len - return kv_cache.resize(self._required_gen_capacity(req, kv_cache.capacity)) + if not kv_cache.resize(self._required_gen_capacity(req, kv_cache.capacity)): + return False + if draft_len > 0: + self._pending_overlap_slack[req.py_request_id] = ( + self._pending_overlap_slack.get(req.py_request_id, 0) + draft_len + ) + return True def revert_allocate_generation(self, req: LlmRequest) -> None: """Undo the capacity growth from try_allocate_generation. @@ -1866,6 +1887,15 @@ def revert_allocate_generation(self, req: LlmRequest) -> None: reverted_cap = kv_cache.capacity - 1 - 2 * draft_len if reverted_cap < 0: return + # The reverted growth included draft_len of overlap slack; deduct it + # so update_resources does not trim slack that no longer exists. + if draft_len > 0: + req_id = req.py_request_id + remaining = self._pending_overlap_slack.get(req_id, 0) - draft_len + if remaining > 0: + self._pending_overlap_slack[req_id] = remaining + else: + self._pending_overlap_slack.pop(req_id, None) if not kv_cache.resize(reverted_cap): raise RuntimeError( f"Failed to revert KV cache capacity for request " @@ -2074,6 +2104,11 @@ def extend_capacity_for_tokens(self, request: LlmRequest) -> None: delta = 2 * (current_draft_len - allocated) if delta <= 0: return + # Half of the delta tops up the overlap slack; record it so + # update_resources reclaims the full slack for this iteration. + self._pending_overlap_slack[request.py_request_id] = self._pending_overlap_slack.get( + request.py_request_id, 0 + ) + (current_draft_len - allocated) kv_cache = self.kv_cache_map[request.py_request_id] new_capacity = kv_cache.capacity + delta success = kv_cache.resize(new_capacity) @@ -2171,6 +2206,11 @@ def _prepare_draft_resources(self, scheduled_batch: ScheduledRequests): f"Draft KV cache generation resize failed for request " f"{req.py_request_id}: could not resize to {new_cap} tokens" ) + slack = self._effective_draft_len(req) + if slack > 0: + self._pending_overlap_slack[req.py_request_id] = ( + self._pending_overlap_slack.get(req.py_request_id, 0) + slack + ) def _augment_tokens_for_block_reuse( self, tokens: Sequence[int], req: LlmRequest, start: int = 0, end: int | None = None @@ -2779,6 +2819,7 @@ def release_index_slot(self, request_id: int) -> None: def free_resources(self, request: LlmRequest, pin_on_release: bool = False): self._allocated_draft_lens.pop(request.py_request_id, None) + self._pending_overlap_slack.pop(request.py_request_id, None) kv_cache = self.kv_cache_map.pop(request.py_request_id, None) if kv_cache is None: self.impl.clear_stats_excluded(request.py_request_id) @@ -3068,10 +3109,18 @@ def update_resources( # will be resumed by the scheduler on the next iteration. if not kv_cache.is_active: continue + # Reclaim this iteration's overlap slack together with the + # rejected-draft rewind; without this the constant slack in + # _required_gen_capacity compounds by draft_len every iteration + # and eventually overruns max_blocks_per_seq. + overlap_slack = self._pending_overlap_slack.pop(req.py_request_id, 0) new_capacity = ( None if req.state in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT) - else kv_cache.capacity - req.py_rewind_len + else max( + kv_cache.capacity - req.py_rewind_len - overlap_slack, + req.max_beam_num_tokens - 1, + ) ) success = kv_cache.resize(new_capacity, req.max_beam_num_tokens - 1) if not success: