kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-02-16 07:53:55 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[None][fix] Remove overlap scheduler adjustment for max sequence length in create_py_executor function (#9229)

Browse Source 
Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com>
This commit is contained in:
Robin Kobus 2026-02-11 17:46:25 +01:00 committed by GitHub
parent c47ff4da43
commit 7a103035be
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 6 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
tensorrt_llm/_torch/pyexecutor/_util.py
View File

@ -273,10 +273,8 @@ class KvCacheCreator:
num_cache_blocks = 0
num_extra_tokens_per_seq = 1 # account for generated tokens
spec_cfg = self._speculative_config
if not self._llm_args.disable_overlap_scheduler:
num_extra_tokens_per_seq = num_extra_tokens_per_seq + 1
if spec_cfg is not None:
num_extra_tokens_per_seq += spec_cfg.max_total_draft_tokens
if not self._llm_args.disable_overlap_scheduler and spec_cfg is not None:
num_extra_tokens_per_seq += spec_cfg.max_total_draft_tokens
if spec_cfg is not None:
num_extra_tokens_per_seq += spec_cfg.max_total_draft_tokens

3
tensorrt_llm/_torch/pyexecutor/model_engine.py
View File

@ -457,8 +457,7 @@ class PyTorchModelEngine(ModelEngine):
# This way it can also be used for CUDA graphs.
if self.use_beam_search:
self.cache_indirection_attention = torch.zeros(
(self.batch_size, self.max_beam_width, self.max_seq_len +
(0 if self._disable_overlap_scheduler else 1)),
(self.batch_size, self.max_beam_width, self.max_seq_len),
device="cuda",
dtype=torch.int32)
else:

6
tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
View File

@ -436,10 +436,8 @@ def create_py_executor(
# PyTorchModelEngine modifies these fields, update them
model_engine_max_seq_len = model_engine.max_seq_len
net_max_seq_len = model_engine_max_seq_len
if not llm_args.disable_overlap_scheduler:
model_engine_max_seq_len = model_engine.max_seq_len + 1
if spec_config is not None:
model_engine_max_seq_len += spec_config.max_total_draft_tokens
if not llm_args.disable_overlap_scheduler and spec_config is not None:
model_engine_max_seq_len += spec_config.max_total_draft_tokens
if spec_config is not None:
model_engine_max_seq_len += get_num_extra_kv_tokens(spec_config)

2
tensorrt_llm/_torch/pyexecutor/sampler.py
View File

@ -1047,7 +1047,7 @@ class TorchSampler(Sampler[SampleStateTorch], AsyncWorkerMixin):
self.CACHE_INDIRECTION_SHAPE = (
self.max_num_sequences,
self.max_beam_width,
self.max_seq_len + (0 if args.disable_overlap_scheduler else 1),
self.max_seq_len,
)
self.LOGPROBS_SHAPE = (self.max_num_sequences, self.max_beam_width, self.max_tokens)
self.TOPK_LOGPROBS_SHAPE = (self.max_num_sequences, self.max_tokens, self.max_topk_logprobs)