Skip to content

[None][fix] Add KV cache V2 recompute pause path #14816

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
yizhang-nv wants to merge 1 commit into from
Closed

[None][fix] Add KV cache V2 recompute pause path #14816

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 60 additions & 6 deletions tensorrt_llm/_torch/pyexecutor/py_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,11 @@
# Default: "0" (only rank 0 prints, matching existing behavior).
PROFILE_LOG_RANKS_ENV_VAR_NAME = "TLLM_PROFILE_LOG_RANKS"

# C++ LlmRequest.pause() requires a prompt-length cap. Recompute pause should
# replay all generated tokens in PyTorch instead of inheriting TRT build-time
# max_input_len truncation.
_UNBOUNDED_PAUSE_MAX_INPUT_LEN = 0x7fffffff


class PPCommTag(IntEnum):
"""
Expand Down Expand Up @@ -1234,8 +1239,9 @@ def _update_iter_stats(self, stats, iter_latency_ms, num_completed_requests,
stats.inflight_batching_stats.num_context_requests = scheduled_batch.num_context_requests
stats.inflight_batching_stats.num_gen_requests = scheduled_batch.num_generation_requests
stats.inflight_batching_stats.num_scheduled_requests = stats.inflight_batching_stats.num_context_requests + stats.inflight_batching_stats.num_gen_requests
stats.inflight_batching_stats.num_paused_requests = len(
scheduled_batch.paused_requests)
paused_requests = (scheduled_batch.paused_requests +
scheduled_batch.recompute_paused_requests)
stats.inflight_batching_stats.num_paused_requests = len(paused_requests)
stats.inflight_batching_stats.avg_num_decoded_tokens_per_iter = 0
stats.inflight_batching_stats.micro_batch_id = micro_batch_id

Expand Down Expand Up @@ -1307,7 +1313,7 @@ def _update_iter_stats(self, stats, iter_latency_ms, num_completed_requests,
# RuntimeError on a mutated request.
num_ctx_kv_tokens = 0
for req in scheduled_batch.context_requests:
if getattr(req, "is_attention_dp_dummy", False):
if req.is_attention_dp_dummy:
continue
last_chunk = getattr(req, "py_last_context_chunk", None)
if last_chunk is not None and last_chunk[0] is not None:
Expand All @@ -1324,7 +1330,7 @@ def _update_iter_stats(self, stats, iter_latency_ms, num_completed_requests,
# summed across scheduled generation requests.
num_gen_kv_tokens = 0
for req in scheduled_batch.generation_requests:
if getattr(req, "is_attention_dp_dummy", False):
if req.is_attention_dp_dummy:
continue
try:
num_gen_kv_tokens += req.get_num_tokens(0)
Expand Down Expand Up @@ -1375,8 +1381,8 @@ def _update_iter_stats(self, stats, iter_latency_ms, num_completed_requests,
# requests — were decoding but got evicted back to the waiting
# pool for this iteration.
num_paused_kv_tokens = 0
for req in scheduled_batch.paused_requests:
if getattr(req, "is_attention_dp_dummy", False):
for req in paused_requests:
if req.is_attention_dp_dummy:
continue
try:
num_paused_kv_tokens += req.get_num_tokens(0)
Expand Down Expand Up @@ -1629,6 +1635,9 @@ def _executor_loop_pp(self):
self.scheduler.schedule_request(self.active_requests,
self.inflight_req_ids)

self._terminate_recompute_paused_requests(scheduled_batch)
self._pause_recompute_paused_requests(scheduled_batch)

# For requests that are fitting disagg gen init, also prepare resources for KV cache manager
if self.kv_cache_transceiver:
self._prepare_disagg_gen_init(
Expand Down Expand Up @@ -2428,6 +2437,9 @@ def _executor_loop(self):
self._revert_gen_alloc(scheduled_batch)
continue

self._terminate_recompute_paused_requests(scheduled_batch)
self._pause_recompute_paused_requests(scheduled_batch)

if not self._scheduler_manages_kv_suspend:
self._terminate_requests(scheduled_batch.paused_requests)
self._pause_requests(scheduled_batch.paused_requests)
Expand Down Expand Up @@ -2720,6 +2732,8 @@ def _executor_loop_overlap(self):
self._revert_gen_alloc(scheduled_batch)
continue

self._terminate_recompute_paused_requests(scheduled_batch)

if not self._scheduler_manages_kv_suspend:
self._terminate_requests(scheduled_batch.paused_requests)

Expand Down Expand Up @@ -2849,6 +2863,7 @@ def _executor_loop_overlap(self):

if not self._scheduler_manages_kv_suspend:
self._pause_requests(scheduled_batch.paused_requests)
self._pause_recompute_paused_requests(scheduled_batch)

if can_queue:
guided_decoder_failed_requests = None
Expand Down Expand Up @@ -3452,6 +3467,7 @@ def _schedule(self):
scheduled_requests.reset_context_requests(scheduled_context_requests)
scheduled_requests.generation_requests = scheduler_output.generation_requests
scheduled_requests.paused_requests = scheduler_output.paused_requests
scheduled_requests.recompute_paused_requests = scheduler_output.recompute_paused_requests

return scheduled_requests, scheduler_output.fitting_disagg_gen_init_requests, num_fitting

Expand Down Expand Up @@ -4488,6 +4504,44 @@ def _pause_requests(self, requests_to_pause):
for req in requests_to_pause:
req.pause(self.max_input_len)

def _pause_recompute_request(self, req):
req.pause(_UNBOUNDED_PAUSE_MAX_INPUT_LEN)
req.py_batch_idx = None
req.py_seq_slot = None
req.py_prompt_len = req.prompt_len
req.py_orig_prompt_len = req.prompt_len
req.py_max_new_tokens = req.max_new_tokens
req.seqlen_this_rank_cp = req.prompt_len
req.total_input_len_cp = req.prompt_len
req.py_draft_pages_allocated = 0
req.py_rewind_len = 0
req.py_draft_tokens = []
req.draft_tokens = []
req.py_last_context_chunk = (None, None)
req.py_last_draft_tokens = None
req.py_num_accepted_draft_tokens = 0
req.py_num_accepted_draft_tokens_indices = []
req.py_rewind_draft_token_separate_adjustment = 0
req.py_decoding_iter = 0
req.py_ctx_pre_resize_cap = None
req._cached_tokens = 0
req._cached_tokens_set = False

def _terminate_recompute_paused_requests(
self, scheduled_batch: ScheduledRequests):
requests = scheduled_batch.recompute_paused_requests
if not requests:
return
self._terminate_requests(requests)

def _pause_recompute_paused_requests(self,
scheduled_batch: ScheduledRequests):
requests = scheduled_batch.recompute_paused_requests
if not requests:
return
for req in requests:
self._pause_recompute_request(req)

def _add_inflight_ids(self, scheduled_requests: ScheduledRequests):
"""Add request IDs of current sampling requests to self.inflight_req_ids.

Expand Down
3 changes: 3 additions & 0 deletions tensorrt_llm/_torch/pyexecutor/resource_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2102,6 +2102,8 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
logger.info(
f"KV cache manager v2 host cache quota set to {host_quota / (1 << 30):.2f}GiB"
)
self.has_host_cache_tier = any(
isinstance(tier, HostCacheTierConfig) for tier in cache_tiers)

self.vocab_size = vocab_size

Expand Down Expand Up @@ -2133,6 +2135,7 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
cache_tiers=cache_tiers_gpu_only,
)
cache_tiers = cache_tiers_gpu_only
self.has_host_cache_tier = False
self.kv_cache_manager_py_config = config
self.impl = KVCacheManagerPy(config,
event_manager=self.event_manager)
Expand Down
78 changes: 67 additions & 11 deletions tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,18 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import dataclasses
import inspect
from abc import ABC, abstractmethod
Expand Down Expand Up @@ -49,16 +64,46 @@ def _call_with_optional_summary(
return fn(*args, cached_summary=cached_summary)


SchedulerOutput = namedtuple(
"SchedulerOutput",
[
"context_requests",
"generation_requests",
"paused_requests",
"fitting_disagg_gen_init_requests",
"num_fitting_requests",
],
)
class SchedulerOutput(
namedtuple(
"_SchedulerOutputBase",
[
"context_requests",
"generation_requests",
"paused_requests",
"fitting_disagg_gen_init_requests",
"num_fitting_requests",
"recompute_paused_requests",
],
)
):
"""Scheduler result.

``recompute_paused_requests`` is V2-only and defaults to an empty list so
existing V1 schedulers can keep constructing the original five-field
output.
"""

__slots__ = ()

def __new__(
cls,
context_requests: RequestList,
generation_requests: RequestList,
paused_requests: RequestList,
fitting_disagg_gen_init_requests: RequestList,
num_fitting_requests: int,
recompute_paused_requests: RequestList | None = None,
):
return super(SchedulerOutput, cls).__new__(
cls,
context_requests,
generation_requests,
paused_requests,
fitting_disagg_gen_init_requests,
num_fitting_requests,
[] if recompute_paused_requests is None else recompute_paused_requests,
)


class ScheduledRequests:
Expand All @@ -77,13 +122,16 @@ class ScheduledRequests:
generation_requests: RequestList
"""Requests that are in the generation phase."""
paused_requests: RequestList
"""Requests that are paused."""
"""Requests whose KV cache was suspended without resetting request state."""
recompute_paused_requests: RequestList
"""Requests that must release resources and restart from context."""

def __init__(self):
self.context_requests_chunking: RequestList = []
self.context_requests_last_chunk: RequestList = []
self.generation_requests: RequestList = []
self.paused_requests: RequestList = []
self.recompute_paused_requests: RequestList = []

@property
def is_generation_only(self) -> bool:
Expand Down Expand Up @@ -173,6 +221,8 @@ class SerializableSchedulerOutput:
int
] # request ids of fitting disaggregated generation initialization requests
num_fitting_requests: int # number of fitting requests
recompute_paused_requests: list[int] = dataclasses.field(default_factory=list)
"""Request ids of recompute-paused requests."""

@classmethod
def from_scheduler_result(
Expand All @@ -194,6 +244,9 @@ def from_scheduler_result(
req.request_id for req in fitting_disagg_gen_init_requests
],
num_fitting_requests=num_fitting_requests,
recompute_paused_requests=[
req.request_id for req in scheduled_requests.recompute_paused_requests
],
)

def to_scheduler_result(
Expand All @@ -213,6 +266,9 @@ def to_scheduler_result(
scheduled_requests.paused_requests = [
id_to_request[req_id] for req_id in self.paused_requests
]
scheduled_requests.recompute_paused_requests = [
id_to_request[req_id] for req_id in self.recompute_paused_requests
]
fitting_disagg_gen_init_requests = [
id_to_request[req_id] for req_id in self.fitting_disagg_gen_init_requests
]
Expand Down
Loading
Loading