-
Notifications
You must be signed in to change notification settings - Fork 1.7k
[None][chore] Create PyExecutor from TorchLlmArgs Part 1 #7105
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.
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
[None][chore] Create PyExecutor from TorchLlmArgs Part 1 #7105
Conversation
📝 WalkthroughWalkthroughPropagates Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant L as LLM (_build_model)
participant E as GenerationExecutor.create
participant P as GenerationExecutorProxy
participant W as GenerationExecutorWorker
participant PT as PyTorch Executor
participant TRT as C++/AutoDeploy Engine
L->>E: create(engine, ..., hf_model_dir, tokenizer, llm_args)
E->>P: spawn proxy (worker_kwargs include hf_model_dir, tokenizer, llm_args)
Note right of P #DDDDDD: GC threshold read from llm_args.garbage_collection_gen0_threshold
P->>W: start worker with worker_kwargs
alt llm_args provided (PyTorch path)
W->>W: cfg = llm_args.get_executor_config(hf_model_dir, tokenizer)
W->>PT: _create_py_executor(cfg, lora_config, llm_args.garbage_collection_gen0_threshold, ...)
PT-->>W: Py executor ready
else fallback (C++/autodeploy)
W->>TRT: _create_engine(engine, executor_config, ...)
TRT-->>W: Engine ready
end
W-->>P: worker ready
P-->>L: executor handle returned
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
Status, Documentation and Community
|
/bot run |
PR_Github #15985 [ run ] triggered by Bot |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/executor/worker.py (2)
121-154
: UnboundLocalError risk: local assignment to executor_config shadows enclosing variable in _create_engine.
executor_config = tllm.ExecutorConfig(1)
makesexecutor_config
local; the earlierif executor_config is None
then reads the local before assignment at runtime (Ruff F823). Also,self._executor_config
is never set when a default is created.Fix by introducing a local
ec
and persisting it:- def _create_engine(comm_ranks, device_ids): - if executor_config is None: - executor_config = tllm.ExecutorConfig(1) - - executor_config.logits_post_processor_config = tllm.LogitsPostProcessorConfig( + def _create_engine(comm_ranks, device_ids): + ec = executor_config or tllm.ExecutorConfig(1) + # Persist so code paths relying on self._executor_config work + self._executor_config = ec + + ec.logits_post_processor_config = tllm.LogitsPostProcessorConfig( processor_batched=batched_logits_processor, replicate=False) - executor_config.parallel_config = tllm.ParallelConfig( + ec.parallel_config = tllm.ParallelConfig( participant_ids=comm_ranks, device_ids=device_ids) if isinstance(engine, Engine): return tllm.Executor(engine.engine, json.dumps(engine.config.to_dict(), cls=ConfigEncoder), tllm.ModelType.DECODER_ONLY, - executor_config=executor_config, + executor_config=ec, managed_weights=engine.managed_weights) - if not hasattr(executor_config, "backend"): - return tllm.Executor(engine, tllm.ModelType.DECODER_ONLY, - executor_config) + if not hasattr(ec, "backend"): + return tllm.Executor(engine, tllm.ModelType.DECODER_ONLY, ec) args = { - "executor_config": executor_config, - "checkpoint_dir": executor_config.hf_model_dir, + "executor_config": ec, + "checkpoint_dir": ec.hf_model_dir, } - if executor_config.backend == "_autodeploy": + if ec.backend == "_autodeploy": from tensorrt_llm._torch.auto_deploy.shim.ad_executor import \ create_autodeploy_executor create_executor = create_autodeploy_executor else: raise ValueError( - f"Unsupported backend config: {executor_config.backend}") + f"Unsupported backend config: {ec.backend}") return create_executor(**args)
188-199
: LoRA manager initialization conditional uses ctor parameter, not the resolved backend flag.Using
getattr(executor_config, "backend", "") == "pytorch"
misses the PyTorch path when the ctor param is None. Switch toself._is_pytorch_backend
.- if getattr(executor_config, "backend", - "") == "pytorch" and lora_config is not None: + if self._is_pytorch_backend and lora_config is not None: from tensorrt_llm._torch.pyexecutor.resource_manager import \ ResourceManagerType peft_cache_manager = self.engine.resource_manager.resource_managers.get( ResourceManagerType.PEFT_CACHE_MANAGER) self._lora_manager = LoraManager( cpp_peft_cache_manager=peft_cache_manager.impl) lora_model_config = self.engine.model_engine.lora_model_config assert lora_model_config is not None self._lora_model_config = lora_model_config
🧹 Nitpick comments (6)
tensorrt_llm/llmapi/llm_args.py (2)
2364-2376
: Minor duplication: max_beam_width set twice.
max_beam_width
is passed into_ExecutorConfig(...)
and then later assigned again toexecutor_config.max_beam_width
. One source of truth is enough.Apply this diff to drop the redundant assignment later (keep the constructor version):
- executor_config = _ExecutorConfig( - max_beam_width=self.max_beam_width, + executor_config = _ExecutorConfig( + max_beam_width=self.max_beam_width, scheduler_config=PybindMirror.maybe_to_pybind( self.scheduler_config), max_batch_size=self.max_batch_size, max_num_tokens=self.max_num_tokens, gather_generation_logits=self.gather_generation_logits, fail_fast_on_attention_window_too_large=getattr( self, 'fail_fast_on_attention_window_too_large', False), )And remove the later reassignment in Lines 2404–2405 (see the next comment).
2403-2405
: Remove redundant max_beam_width override.This duplicates the value already set in the constructor.
- executor_config.enable_chunked_context = self.enable_chunked_prefill - executor_config.max_beam_width = self.max_beam_width + executor_config.enable_chunked_context = self.enable_chunked_prefilltensorrt_llm/llmapi/llm.py (1)
978-982
: Replace debug prints with logger.debug gated by debug flag.Unconditional stdout prints are noisy for library consumers and CI logs.
- print("---- self._executor_cls is: {}".format(self._executor_cls), - flush=True) - print("---- self._engine_dir is: {}".format(self._engine_dir), - flush=True) + if enable_llm_debug(): + logger.debug(f"self._executor_cls: {self._executor_cls}") + logger.debug(f"self._engine_dir: {self._engine_dir}")tensorrt_llm/executor/proxy.py (1)
89-92
: Initialize GC threshold from llm_args safely and avoid double dict lookup.Current code is correct but slightly brittle/readability could improve. Also consider the case where
customized_gc_thresholds(None)
might not be a no-op.- self.garbage_collection_gen0_threshold = worker_kwargs[ - "llm_args"].garbage_collection_gen0_threshold if worker_kwargs.get( - "llm_args", None) is not None else None + llm_args = worker_kwargs.get("llm_args") + self.garbage_collection_gen0_threshold = ( + llm_args.garbage_collection_gen0_threshold if llm_args is not None + else None + )tensorrt_llm/executor/worker.py (2)
155-166
: Device selection: guard zero CUDA device case.If
torch.cuda.device_count()
is 0, this will throw. If CPU isn’t a supported environment, consider raising a clear error early; otherwise guard it.- device_id = self.global_rank % torch.cuda.device_count() - torch.cuda.set_device(device_id) + num_devs = torch.cuda.device_count() + if num_devs == 0: + raise RuntimeError("No CUDA devices available for executor worker.") + device_id = self.global_rank % num_devs + torch.cuda.set_device(device_id)
791-801
: Replace debug print with logger and debug flag.Same rationale as in llm.py; printing in workers pollutes logs.
- print("---- worker_cls is: {}".format(worker_cls), flush=True) + if enable_llm_debug(): + logger.debug(f"worker_cls: {worker_cls}")
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (5)
tensorrt_llm/executor/executor.py
(7 hunks)tensorrt_llm/executor/proxy.py
(1 hunks)tensorrt_llm/executor/worker.py
(8 hunks)tensorrt_llm/llmapi/llm.py
(3 hunks)tensorrt_llm/llmapi/llm_args.py
(3 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.py
: Python code must target Python 3.8+
Python indentation: 4 spaces, no tabs
Maintain module namespace in imports (from package.subpackage import foo; then use foo.SomeClass())
Python file names use snake_case
Python class names use PascalCase
Python functions/methods and local variables use snake_case; variables starting with a number get k_ prefix (e.g., k_99th_percentile)
Global variables use G_ prefixed UPPER_SNAKE_CASE (e.g., G_MY_GLOBAL)
Constants use UPPER_SNAKE_CASE in Python
Avoid shadowing variables from outer scopes in Python
Initialize all externally visible members of a Python class in init
Prefer docstrings for interfaces used outside a file; comments for local code
Use Google-style docstrings for classes and functions (Sphinx-parsable)
Document attributes/variables inline with short docstrings
Avoid reflection when simple alternatives exist (e.g., prefer explicit parameters over dict(**locals()))
In try/except, catch the narrowest exceptions possible
For duck-typing with try/except, keep try body minimal and put logic in else
Files:
tensorrt_llm/executor/executor.py
tensorrt_llm/llmapi/llm_args.py
tensorrt_llm/executor/proxy.py
tensorrt_llm/executor/worker.py
tensorrt_llm/llmapi/llm.py
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Prepend NVIDIA copyright header (current year) to all source files
Files:
tensorrt_llm/executor/executor.py
tensorrt_llm/llmapi/llm_args.py
tensorrt_llm/executor/proxy.py
tensorrt_llm/executor/worker.py
tensorrt_llm/llmapi/llm.py
🪛 Ruff (0.12.2)
tensorrt_llm/executor/worker.py
122-122: Local variable executor_config
referenced before assignment
(F823)
🔇 Additional comments (16)
tensorrt_llm/llmapi/llm_args.py (3)
47-49
: Import of _GuidedDecodingConfig is correct and scoped.Bringing
_GuidedDecodingConfig
into this module makes sense given the new get_executor_config wiring. No concerns here.
60-61
: Tokenizer helpers import alignment looks good.Importing both
_llguidance_tokenizer_info
and_xgrammar_tokenizer_info
here (instead of exporting them from llm.py) fits the new responsibility boundary.
2414-2417
: Script to locate thesuggest_spec_config
definition is running. I’ll analyze the function signature once I have the results.tensorrt_llm/llmapi/llm.py (3)
40-41
: Public surface: narrowed tokenizer exports are fine.Dropping
_llguidance_tokenizer_info
here and keeping_xgrammar_tokenizer_info
is consistent with moving guided decoding wiring into llm_args.
960-960
: Assert on TorchLlmArgs is good.This keeps the PyTorch path honest during the migration.
985-1000
: LLM now delegates config creation to worker via llm_args; OK but update docs/changelog.Passing
executor_config=None
withhf_model_dir
andllm_args
is the right direction for centralizing configuration. Ensure API docs mention this change so downstream users do not expect to pass anExecutorConfig
for PyTorch.Would you like me to draft a brief upgrade note covering:
- Removal of
garbage_collection_gen0_threshold
from executor creation entrypoints- New
llm_args
-driven initialization for PyTorch?tensorrt_llm/executor/proxy.py (2)
96-97
: is_llm_executor=False in worker_kwargs is correct.Proxy hosts the “main” executor API; workers run as subordinates. This override makes sense.
156-161
: No changes required for GC threshold guardThe
customized_gc_thresholds
context manager already treats aNone
(or any falsy)gen0_threshold
as a no-op—it never callsgc.set_threshold
whengen0_threshold
isNone
(or0
), and simply yields without error before and after the block. Therefore, wrapping it in anullcontext
is purely optional and not required for correctness.tensorrt_llm/executor/executor.py (4)
24-24
: Importing TorchLlmArgs for type hints is fine.No behavior change here.
382-388
: Forwarding llm_args and hf_model_dir to workers: good.This plumbs the Torch path without disturbing TRT behavior.
405-406
: Propagating is_llm_executor consistently across creation paths.Looks consistent across MPI reuse, single-process, and default proxy branches.
Also applies to: 417-417, 429-429, 439-439
358-360
: GenerationExecutor.create signature change is backward-compatible—no code updates requiredAdding the new
hf_model_dir
andllm_args
parameters as trailing defaults does not break any existing positional call sites, and no invocation of the removedgarbage_collection_gen0_threshold
keyword was found. All tests and internal code continue to callGenerationExecutor.create(engine_dir)with no additional arguments, which still works as before.
• No call sites pass
garbage_collection_gen0_threshold
tocreate()
• All existing invocations rely on the default values forhf_model_dir
andllm_args
Please update the public API documentation to:
- Document the new
hf_model_dir: Optional[Path] = None
andllm_args: Optional[TorchLlmArgs] = None
parameters- Remove references to
garbage_collection_gen0_threshold
from thecreate()
APItensorrt_llm/executor/worker.py (4)
63-65
: Constructor additions (hf_model_dir, llm_args) are appropriate.This is the right place to collect inputs for both the Torch and Autodeploy branches.
85-87
: Backend detection is fine.
self._is_pytorch_backend
is a simple and reliable switch.
457-467
: Overlap scheduler guard: solid validation.Good to explicitly reject context-only requests with overlap enabled on the PyTorch path when disaggregated.
470-490
: Default max_tokens deduction relies on executor_config.max_seq_len; ensure it’s set.With the above fixes persisting
self._executor_config
,_deduce_max_tokens
will work. Ifmax_seq_len
isn’t set by the Torch path, consider populating it duringget_executor_config
or returning a clearer error here.Would you like me to thread
max_seq_len
into the Torchexecutor_config
earlier to make this deduction robust?
e63d49d
to
ad49221
Compare
PR_Github #15985 [ run ] completed with state |
/bot run |
PR_Github #15988 [ run ] triggered by Bot |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/executor/worker.py (2)
123-156
: Bug: Unbound local due to executor_config assignment inside nested _create_engineexecutor_config is assigned in _create_engine, making it a local and triggering “referenced before assignment” (ruff F823). Use a separate local variable to avoid shadowing the outer parameter.
Apply:
- def _create_engine(comm_ranks, device_ids): - if executor_config is None: - executor_config = tllm.ExecutorConfig(1) - - executor_config.logits_post_processor_config = tllm.LogitsPostProcessorConfig( + def _create_engine(comm_ranks, device_ids): + local_executor_config = executor_config or tllm.ExecutorConfig(1) + + local_executor_config.logits_post_processor_config = tllm.LogitsPostProcessorConfig( processor_batched=batched_logits_processor, replicate=False) - executor_config.parallel_config = tllm.ParallelConfig( + local_executor_config.parallel_config = tllm.ParallelConfig( participant_ids=comm_ranks, device_ids=device_ids) if isinstance(engine, Engine): return tllm.Executor(engine.engine, json.dumps(engine.config.to_dict(), cls=ConfigEncoder), tllm.ModelType.DECODER_ONLY, - executor_config=executor_config, + executor_config=local_executor_config, managed_weights=engine.managed_weights) - if not hasattr(executor_config, "backend"): + if not hasattr(local_executor_config, "backend"): return tllm.Executor(engine, tllm.ModelType.DECODER_ONLY, - executor_config) + local_executor_config) args = { - "executor_config": executor_config, - "checkpoint_dir": executor_config.hf_model_dir, + "executor_config": local_executor_config, + "checkpoint_dir": local_executor_config.hf_model_dir, } - if executor_config.backend == "_autodeploy": + if local_executor_config.backend == "_autodeploy": from tensorrt_llm._torch.auto_deploy.shim.ad_executor import \ create_autodeploy_executor create_executor = create_autodeploy_executor else: raise ValueError( - f"Unsupported backend config: {executor_config.backend}") + f"Unsupported backend config: {local_executor_config.backend}") return create_executor(**args)
190-201
: LoRA manager init for PyTorch path still checks executor_config backendIn the Torch path, the outer executor_config is often None. This branch won’t run, and LoRA optimizations won’t initialize. Prefer the already-computed self._is_pytorch_backend (or inspect the engine type).
Apply:
- if getattr(executor_config, "backend", - "") == "pytorch" and lora_config is not None: + if self._is_pytorch_backend and lora_config is not None: from tensorrt_llm._torch.pyexecutor.resource_manager import \ ResourceManagerType peft_cache_manager = self.engine.resource_manager.resource_managers.get( ResourceManagerType.PEFT_CACHE_MANAGER) self._lora_manager = LoraManager( cpp_peft_cache_manager=peft_cache_manager.impl) lora_model_config = self.engine.model_engine.lora_model_config assert lora_model_config is not None self._lora_model_config = lora_model_config
🧹 Nitpick comments (4)
tensorrt_llm/executor/proxy.py (1)
89-91
: Safer extraction of GC threshold from llm_argsMinor robustness/readability tweak: avoid the chained dict access and use getattr on a local. This prevents accidental AttributeError if a different llm_args type is introduced later.
Apply:
- self.garbage_collection_gen0_threshold = worker_kwargs[ - "llm_args"].garbage_collection_gen0_threshold if worker_kwargs.get( - "llm_args", None) is not None else None + llm_args = worker_kwargs.get("llm_args") + self.garbage_collection_gen0_threshold = getattr( + llm_args, "garbage_collection_gen0_threshold", None + )tensorrt_llm/executor/worker.py (3)
63-65
: Constructor now accepts hf_model_dir and llm_args — document new paramsAdd Google-style docstring entries for hf_model_dir and llm_args (type TorchLlmArgs), including how they steer backend selection and checkpoint_dir for PyExecutor.
157-164
: CUDA device selection — add a quick availability guardIf this ever runs in an environment where CUDA is not initialized, torch.cuda.device_count() can be 0. A simple assert or early error message improves debuggability.
- device_id = self.global_rank % torch.cuda.device_count() + assert torch.cuda.is_available() and torch.cuda.device_count() > 0, \ + "CUDA is not available or no devices found." + device_id = self.global_rank % torch.cuda.device_count()
459-468
: Overlap scheduler restriction: clear error earlyGood defensive check to prevent unsupported context-only requests when overlap is enabled in PyTorch. Consider hinting how to disable overlap in the error text for quicker user recovery (e.g., set TorchLlmArgs.disable_overlap_scheduler=True).
- raise ValueError( - "Context only requests are not supported in pytorch backend when overlap is enabled." - ) + raise ValueError( + "Context-only requests are not supported in the PyTorch backend when overlap scheduling is enabled. " + "Set TorchLlmArgs.disable_overlap_scheduler=True or disable disaggregated serving." + )
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (5)
tensorrt_llm/executor/executor.py
(7 hunks)tensorrt_llm/executor/proxy.py
(1 hunks)tensorrt_llm/executor/worker.py
(8 hunks)tensorrt_llm/llmapi/llm.py
(3 hunks)tensorrt_llm/llmapi/llm_args.py
(3 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/llmapi/llm_args.py
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.py
: Python code must target Python 3.8+
Python indentation: 4 spaces, no tabs
Maintain module namespace in imports (from package.subpackage import foo; then use foo.SomeClass())
Python file names use snake_case
Python class names use PascalCase
Python functions/methods and local variables use snake_case; variables starting with a number get k_ prefix (e.g., k_99th_percentile)
Global variables use G_ prefixed UPPER_SNAKE_CASE (e.g., G_MY_GLOBAL)
Constants use UPPER_SNAKE_CASE in Python
Avoid shadowing variables from outer scopes in Python
Initialize all externally visible members of a Python class in init
Prefer docstrings for interfaces used outside a file; comments for local code
Use Google-style docstrings for classes and functions (Sphinx-parsable)
Document attributes/variables inline with short docstrings
Avoid reflection when simple alternatives exist (e.g., prefer explicit parameters over dict(**locals()))
In try/except, catch the narrowest exceptions possible
For duck-typing with try/except, keep try body minimal and put logic in else
Files:
tensorrt_llm/executor/proxy.py
tensorrt_llm/executor/worker.py
tensorrt_llm/executor/executor.py
tensorrt_llm/llmapi/llm.py
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Prepend NVIDIA copyright header (current year) to all source files
Files:
tensorrt_llm/executor/proxy.py
tensorrt_llm/executor/worker.py
tensorrt_llm/executor/executor.py
tensorrt_llm/llmapi/llm.py
🧠 Learnings (1)
📚 Learning: 2025-07-17T09:01:27.402Z
Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.402Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
Applied to files:
tensorrt_llm/executor/worker.py
🪛 Ruff (0.12.2)
tensorrt_llm/executor/worker.py
124-124: Local variable executor_config
referenced before assignment
(F823)
🔇 Additional comments (10)
tensorrt_llm/executor/proxy.py (1)
96-96
: Explicitly setting is_llm_executor=False for worker kwargs is correctThis aligns the proxy/worker split of iteration-result responsibilities and avoids double-initializing result queues.
tensorrt_llm/executor/executor.py (2)
382-388
: Propagating hf_model_dir and llm_args to worker is the right placeThreading these through worker_kwargs centralizes backend selection in the worker. LGTM.
405-406
: Consistent propagation of is_llm_executorPassing is_llm_executor through each creation branch keeps stats/kv-events iteration queues consistent. Looks good.
Also applies to: 416-417, 428-429, 439-440
tensorrt_llm/llmapi/llm.py (2)
960-961
: Asserting TorchLlmArgs on PyTorch pathThe explicit assert prevents accidental mixing of Trt/AAD args with Torch backend. Good safeguard.
967-981
: Verify PyTorch executor path for residual_executor_config
referencesI wasn’t able to identify a separate PyTorch-specific executor implementation in
tensorrt_llm/llmapi
or any branch that clearly omitsexecutor_config
in itscreate
call. Please manually confirm that, in the pure‐PyTorch code path (where you expectexecutor_config=None
), no downstream logic still reads or assumesself._executor_config
is set on the LLM object.tensorrt_llm/executor/worker.py (5)
85-87
: Backend flag derived from llm_args is clearUsing llm_args to decide PyTorch mode is straightforward and avoids fragile executor_config checks. LGTM.
94-122
: Persisting executor_config from llm_args fixes downstream max_tokens deductionCreating executor_config via llm_args.get_executor_config(hf_model_dir) and persisting to self._executor_config addresses previous failures when deducing default max_tokens.
165-168
: Engine selection toggle (llm_args vs. legacy) is correctClear split between PyTorch path (llm_args) and legacy/C++ path. LGTM.
672-674
: worker_main signature updated to accept hf_model_dir and llm_argsMatches the proxy/executor plumbing. LGTM.
94-122
: ExecutorConfig in Torch Path Already Exposes max_seq_len; cp_size Defaults to 1 (No Further Changes Needed)I checked the PyTorch executor config in
tensorrt_llm/_torch/pyexecutor/config.py
and found that theExecutorConfig
dataclass directly declaresmax_seq_len
(defaulting toNone
) alongside other core fields. It also carries aparallel_config
object—however, the Torch backend currently does not use cross-processor (cp) parallelism, socp_size
is implicitly 1 and appears only in the mapping when pipeline parallelism is active. No explicitcp_size
field is needed for the single-process Torch path. As a result:
ExecutorConfig.max_seq_len
is available for downstream deduce_max_tokens logic without modification.parallel_config.cp_size
would default to 1 even if queried, matching Torch’s no‐cp paradigm.No changes to
_create_py_executor
are required.
PR_Github #15988 [ run ] completed with state |
/bot run |
PR_Github #16030 [ run ] triggered by Bot |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 0
♻️ Duplicate comments (1)
tensorrt_llm/llmapi/llm_args.py (1)
2377-2448
: Harden get_executor_config: kv-cache None guard, guided-decoding backend checks, tokenizer presence, docstringOverall direction is right. A few correctness hardening tweaks recommended.
- FORCE_DETERMINISTIC kv-cache guard:
- If executor_config.kv_cache_config is unset, the toggles will raise. Add a guard/initializer. This is a repeat of an earlier comment.
- llguidance must be PyTorch-only:
- Add a backend check to fail fast when used on TRT. This is a repeat of an earlier comment.
- Require tokenizer when guided decoding is enabled (xgrammar/llguidance):
- Today, if tokenizer is None, _xgrammar/_llguidance helpers will raise later. Provide a clearer error early.
- Add a brief docstring for get_executor_config.
Apply the following minimal diff:
- def get_executor_config(self, - _hf_model_dir: Optional[Path] = None - ) -> _ExecutorConfig: + def get_executor_config(self, + _hf_model_dir: Optional[Path] = None + ) -> _ExecutorConfig: + """Build an ExecutorConfig from TorchLlmArgs. + This centralizes PyTorch backend runtime configuration, including + scheduler, KV/PEFT/guided decoding configs, chunked prefill, mapping, + speculative decoding, and mm-encoder-only flag. + """ @@ - if self.kv_cache_config is not None: - executor_config.kv_cache_config = PybindMirror.maybe_to_pybind( - self.kv_cache_config) + if self.kv_cache_config is not None: + executor_config.kv_cache_config = PybindMirror.maybe_to_pybind( + self.kv_cache_config) if os.getenv("FORCE_DETERMINISTIC", "0") == "1": # Disable KV cache reuse for deterministic mode - executor_config.kv_cache_config.enable_block_reuse = False - executor_config.kv_cache_config.enable_partial_reuse = False + if getattr(executor_config, "kv_cache_config", None) is None: + executor_config.kv_cache_config = _KvCacheConfig() + executor_config.kv_cache_config.enable_block_reuse = False + executor_config.kv_cache_config.enable_partial_reuse = False @@ - if self.guided_decoding_backend == 'xgrammar': + if self.guided_decoding_backend == 'xgrammar': + if self.tokenizer is None: + raise ValueError("guided_decoding_backend requires a tokenizer; got None.") executor_config.guided_decoding_config = _GuidedDecodingConfig( backend=_GuidedDecodingConfig.GuidedDecodingBackend.XGRAMMAR, **_xgrammar_tokenizer_info(self.tokenizer)) - elif self.guided_decoding_backend == 'llguidance': + elif self.guided_decoding_backend == 'llguidance': + if self.backend != "pytorch": + raise ValueError("llguidance guided decoding is supported only on the PyTorch backend.") + if self.tokenizer is None: + raise ValueError("guided_decoding_backend requires a tokenizer; got None.") executor_config.guided_decoding_config = _GuidedDecodingConfig( backend=_GuidedDecodingConfig.GuidedDecodingBackend.LLGUIDANCE, **_llguidance_tokenizer_info(self.tokenizer))Optional verification regarding AUTO speculation:
- suggest_spec_config(max_batch_size) may need a non-None value. If max_batch_size can be None here, either guard or confirm suggest_spec_config handles None.
#!/bin/bash # Inspect suggest_spec_config signature/usage rg -nP 'def\s+suggest_spec_config\(' tensorrt_llm/_torch -C2 rg -nP 'suggest_spec_config\s*\(' tensorrt_llm/_torch -C2
🧹 Nitpick comments (7)
tensorrt_llm/llmapi/mm_encoder.py (3)
1-6
: Missing NVIDIA copyright headerPer repo guidelines, prepend the current-year NVIDIA copyright header to all source files.
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + from pathlib import Path from typing import Any, List, Literal, Optional, Sequence, Union
58-62
: Replace assert with explicit runtime checkAsserts can be stripped with Python optimization flags and are not ideal for user-facing type errors. Use an explicit check with a clear exception.
- assert isinstance(self.args, TorchLlmArgs) - # Update the tokenizer in TorchLlmArgs, so it can be used in GenerationExecutorWorker to init executor_config + if not isinstance(self.args, TorchLlmArgs): + raise TypeError(f"Expected TorchLlmArgs for PyTorch backend, got {type(self.args).__name__}") + # Update the tokenizer in TorchLlmArgs, so it can be used in GenerationExecutorWorker to init executor_config self.args.set_tokenizer(self.tokenizer) self.args.set_mm_encoder_only(True)
74-78
: Implement minimal argument validation or track with an issuePlaceholder validator currently does nothing. At minimum, reject TRT-only or generation-oriented knobs that don’t apply to a pure encoder (e.g., guided decoding, streaming/logprobs). If not in this PR, please track with a follow-up.
Example minimal guard:
def _validate_mm_args_for_torch_backend(self, kwargs: dict) -> None: - """Validate that users don't pass LLM-specific arguments when using MultimodalEncoder (PyTorch). - Placeholder for now. - """ + """Validate MM encoder kwargs for PyTorch backend.""" + disallowed = {"guided_decoding_backend", "return_generation_logits", "prompt_logprobs", "logprobs"} + bad = sorted(set(kwargs).intersection(disallowed)) + if bad: + raise ValueError(f"Unsupported args for MultimodalEncoder: {bad}")tensorrt_llm/llmapi/llm.py (1)
960-963
: Prefer explicit runtime type check over assertSame reasoning as in MultimodalEncoder: use a runtime check for clearer errors and to avoid being stripped by -O.
- assert isinstance(self.args, TorchLlmArgs) + if not isinstance(self.args, TorchLlmArgs): + raise TypeError(f"Expected TorchLlmArgs for PyTorch backend, got {type(self.args).__name__}") # Update the tokenizer in TorchLlmArgs, so it can be used in GenerationExecutorWorker to init executor_config self.args.set_tokenizer(self.tokenizer)tensorrt_llm/llmapi/llm_args.py (3)
60-61
: Importing _llguidance_tokenizer_info here is appropriate, but add a backend guardUsing llguidance should be restricted to the PyTorch backend. See suggested guard below in get_executor_config for correctness.
1825-1827
: Type hints and docstring for set_tokenizerPublicly accessible helper should be annotated and documented for clarity.
- def set_tokenizer(self, tokenizer): - self.tokenizer = tokenizer + def set_tokenizer(self, tokenizer: Optional[TokenizerBase | PreTrainedTokenizerBase]) -> None: + """Assign tokenizer after object construction (used by PyTorch path).""" + self.tokenizer = tokenizer
2374-2376
: Type hints and docstring for set_mm_encoder_onlyMinor polish to match style and avoid untyped public API.
- def set_mm_encoder_only(self, mm_encoder_only): - self.mm_encoder_only = mm_encoder_only + def set_mm_encoder_only(self, mm_encoder_only: bool) -> None: + """Toggle vision-encoder-only execution mode.""" + self.mm_encoder_only = mm_encoder_only
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (4)
tensorrt_llm/executor/worker.py
(8 hunks)tensorrt_llm/llmapi/llm.py
(3 hunks)tensorrt_llm/llmapi/llm_args.py
(5 hunks)tensorrt_llm/llmapi/mm_encoder.py
(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/executor/worker.py
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py
: Python code must target Python 3.8+
Python indentation: 4 spaces, no tabs
Maintain module namespace in imports (from package.subpackage import foo; then use foo.SomeClass())
Python file names use snake_case
Python class names use PascalCase
Python functions/methods and local variables use snake_case; variables starting with a number get k_ prefix (e.g., k_99th_percentile)
Global variables use G_ prefixed UPPER_SNAKE_CASE (e.g., G_MY_GLOBAL)
Constants use UPPER_SNAKE_CASE in Python
Avoid shadowing variables from outer scopes in Python
Initialize all externally visible members of a Python class in init
Prefer docstrings for interfaces used outside a file; comments for local code
Use Google-style docstrings for classes and functions (Sphinx-parsable)
Document attributes/variables inline with short docstrings
Avoid reflection when simple alternatives exist (e.g., prefer explicit parameters over dict(**locals()))
In try/except, catch the narrowest exceptions possible
For duck-typing with try/except, keep try body minimal and put logic in else
Files:
tensorrt_llm/llmapi/mm_encoder.py
tensorrt_llm/llmapi/llm.py
tensorrt_llm/llmapi/llm_args.py
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Prepend NVIDIA copyright header (current year) to all source files
Files:
tensorrt_llm/llmapi/mm_encoder.py
tensorrt_llm/llmapi/llm.py
tensorrt_llm/llmapi/llm_args.py
🧬 Code graph analysis (3)
tensorrt_llm/llmapi/mm_encoder.py (3)
tensorrt_llm/llmapi/llm_args.py (6)
TorchLlmArgs
(2035-2515)set_tokenizer
(1825-1826)set_mm_encoder_only
(2374-2375)parallel_config
(1341-1342)world_size
(248-260)world_size
(269-276)tensorrt_llm/llmapi/llm.py (2)
tokenizer
(682-686)tokenizer
(689-690)tensorrt_llm/llmapi/mpi_session.py (1)
external_mpi_comm_available
(57-66)
tensorrt_llm/llmapi/llm.py (3)
tensorrt_llm/llmapi/tokenizer.py (2)
TokenizerBase
(24-25)_xgrammar_tokenizer_info
(293-326)tensorrt_llm/llmapi/llm_args.py (2)
TorchLlmArgs
(2035-2515)set_tokenizer
(1825-1826)tensorrt_llm/executor/executor.py (1)
create
(346-439)
tensorrt_llm/llmapi/llm_args.py (3)
tensorrt_llm/llmapi/llm.py (2)
tokenizer
(682-686)tokenizer
(689-690)tensorrt_llm/llmapi/tokenizer.py (4)
TokenizerBase
(24-25)_llguidance_tokenizer_info
(329-333)_xgrammar_tokenizer_info
(293-326)tokenizer_factory
(270-290)tensorrt_llm/_torch/pyexecutor/config.py (1)
update_executor_config
(129-172)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (5)
tensorrt_llm/llmapi/mm_encoder.py (1)
63-73
: Confirm GenerationExecutor.create arguments for MultimodalEncoderWe’ve traced through both proxy and worker paths and confirmed that the MultimodalEncoder’s call to
GenerationExecutor.create(…)
is safe and behavior-preserving:• executor_config=None → in
GenerationExecutorWorker.__init__
,_create_py_executor
assertsexecutor_config is None
and then callsexecutor_config = llm_args.get_executor_config(hf_model_dir)so the config is built from
llm_args
, matching all other LLM paths.
• return_logits/default postproc → the staticcreate
method uses defaultreturn_logits=False
and constructspostproc_worker_config = PostprocWorkerConfig() # num_postprocess_workers=0so no logits are returned and no postprocessing workers are spawned unless explicitly enabled.
• is_llm_executor=True →
– In the proxy (main) process, passingis_llm_executor=True
tosuper().__init__
enables iteration‐stats and kv‐events collection via_maybe_initialize_iteration_results
.
– In worker processes, the proxy overridesworker_kwargs["is_llm_executor"]=False
, so workers behave identically to non‐LLM executors.In addition, a grep of all
GenerationExecutor.create(
call sites shows no user code relying on any other signature parameters, and themm_encoder_only
flag is correctly propagated throughTorchLlmArgs
into both Python and C++ backends.No regressions detected. Feel free to remove or update the TODO on
is_llm_executor
once its stats‐collection semantics are documented.tensorrt_llm/llmapi/llm.py (2)
40-40
: Import surface cleanup is correctDropping _llguidance_tokenizer_info from this module aligns with moving guided decoding wiring into llm_args.get_executor_config(). No issues spotted.
967-983
: ExecutorConfig usage centralized—no further changes neededI’ve confirmed that there are no remaining direct
_ExecutorConfig(...)
constructions in the PyTorch execution code—only the centralized builder intensorrt_llm/llmapi/llm_args.py
at lines 2357–2361. All backend executor instantiations now useexecutor_config=None
withllm_args
, so this refactor is fully applied.tensorrt_llm/llmapi/llm_args.py (2)
47-49
: SchedulerConfig/GuidedDecodingConfig import syncBringing _GuidedDecodingConfig into this module is consistent with centralizing executor_config construction. No concerns.
2176-2181
: mm_encoder_only flag addition looks goodField name, default=False, and description align with the MM encoder use-case.
PR_Github #16030 [ run ] completed with state |
/bot run |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/executor/worker.py (1)
182-190
: LoRA manager init is gated off the ctor arg, not the actual backend.Using getattr(executor_config, "backend", "") may be "", causing this block to be skipped in the PyTorch path and breaking LoRA in PyTorch. Gate on self._is_pytorch_backend instead.
- if getattr(executor_config, "backend", - "") == "pytorch" and lora_config is not None: + if self._is_pytorch_backend and lora_config is not None: from tensorrt_llm._torch.pyexecutor.resource_manager import \ ResourceManagerType peft_cache_manager = self.engine.resource_manager.resource_managers.get( ResourceManagerType.PEFT_CACHE_MANAGER) self._lora_manager = LoraManager( cpp_peft_cache_manager=peft_cache_manager.impl)
♻️ Duplicate comments (2)
tensorrt_llm/llmapi/llm_args.py (2)
1845-1849
: FORCE_DETERMINISTIC can None-deref executor_config.kv_cache_config.This dereferences executor_config.kv_cache_config without ensuring it exists. Initialize or guard before toggling.
- if os.getenv("FORCE_DETERMINISTIC", "0") == "1": - # Disable KV cache reuse for deterministic mode - executor_config.kv_cache_config.enable_block_reuse = False - executor_config.kv_cache_config.enable_partial_reuse = False + if os.getenv("FORCE_DETERMINISTIC", "0") == "1": + # Disable KV cache reuse for deterministic mode + if getattr(executor_config, "kv_cache_config", None) is None: + executor_config.kv_cache_config = _KvCacheConfig() + executor_config.kv_cache_config.enable_block_reuse = False + executor_config.kv_cache_config.enable_partial_reuse = False
1854-1865
: Guard guided decoding: require tokenizer; llguidance must be PyTorch-only.
- When guided_decoding_backend is set, a tokenizer is mandatory; currently None would raise later with a less clear error.
- llguidance backend should fail fast unless backend == "pytorch".
- if self.guided_decoding_backend == 'xgrammar': + if self.guided_decoding_backend in ('xgrammar', 'llguidance'): + if self.tokenizer is None: + raise ValueError( + "guided_decoding_backend requires a tokenizer; set LlmArgs.tokenizer or call set_tokenizer(...)." + ) + if self.guided_decoding_backend == 'xgrammar': executor_config.guided_decoding_config = _GuidedDecodingConfig( backend=_GuidedDecodingConfig.GuidedDecodingBackend.XGRAMMAR, **_xgrammar_tokenizer_info(self.tokenizer)) elif self.guided_decoding_backend == 'llguidance': + if self.backend != "pytorch": + raise ValueError("llguidance backend is supported only on the PyTorch backend.") executor_config.guided_decoding_config = _GuidedDecodingConfig( backend=_GuidedDecodingConfig.GuidedDecodingBackend.LLGUIDANCE, **_llguidance_tokenizer_info(self.tokenizer))
🧹 Nitpick comments (8)
tensorrt_llm/llmapi/llm_args.py (5)
1825-1827
: Normalize and type-hint set_tokenizer for consistency.Call tokenizer_factory to keep the invariant that self.tokenizer is normalized across entry points, and add type hints/docstring.
- def set_tokenizer(self, tokenizer): - self.tokenizer = tokenizer + def set_tokenizer( + self, + tokenizer: Optional[Union[str, Path, TokenizerBase, PreTrainedTokenizerBase]], + ) -> None: + """Assign and normalize tokenizer using tokenizer_factory for consistency.""" + self.tokenizer = tokenizer_factory( + tokenizer, + trust_remote_code=self.trust_remote_code, + use_fast=self.tokenizer_mode != 'slow', + )
1868-1868
: Nit: duplicate assignment of max_beam_width.max_beam_width is already provided in the _ExecutorConfig constructor above; this reassignment is redundant.
- executor_config.max_beam_width = self.max_beam_width
2247-2253
: mm_encoder_only should be a bool, not Optional[bool].The default is False and no None semantics are used. Tighten the type.
- mm_encoder_only: Optional[bool] = Field( - default=False, + mm_encoder_only: bool = Field( + default=False, description= "Only load/execute the vision encoder part of the full model.", status="prototype", )
2445-2447
: Type-hint and document set_mm_encoder_only.Keep the public surface consistent and self-documenting.
- def set_mm_encoder_only(self, mm_encoder_only): - self.mm_encoder_only = mm_encoder_only + def set_mm_encoder_only(self, mm_encoder_only: bool) -> None: + """Toggle vision-encoder-only mode for PyTorch backend.""" + self.mm_encoder_only = bool(mm_encoder_only)
2448-2454
: Optional: drop redundant mm_encoder_only set if base passes it to update.If you adopt the earlier change to pass mm_encoder_only into update_executor_config, this override can simply return super().get_executor_config(...) without reassigning the flag.
def get_executor_config(self, _hf_model_dir: Optional[Path] = None ) -> _ExecutorConfig: executor_config = super().get_executor_config(_hf_model_dir) - executor_config.mm_encoder_only = self.mm_encoder_only return executor_config
tensorrt_llm/executor/worker.py (3)
63-65
: Constructor API extension LGTM; consider documenting the new args.hf_model_dir and llm_args are now first-class; add a short docstring to init for discoverability.
94-102
: Guard against zero CUDA devices and compute device_id robustly.Avoid modulo by zero in environments without visible GPUs, and make the intent explicit.
- def _get_comm_ranks_device_id(): - device_id = self.global_rank % torch.cuda.device_count() + def _get_comm_ranks_device_id(): + dc = torch.cuda.device_count() + assert dc > 0, "No CUDA devices available." + device_id = self.global_rank % dc torch.cuda.set_device(device_id) # Make sure C++ executor would use same devices/ranks as py_executor global_rank = global_mpi_rank() comm_ranks = mpi_comm().allgather(global_rank) device_ids = mpi_comm().allgather(device_id) return comm_ranks, device_ids
451-461
: Overlap scheduler + disaggregated: UX improvement suggestion.The error is correct; consider hinting the knob name in the message (disable_overlap_scheduler) to reduce support churn.
- raise ValueError( - "Context only requests are not supported in pytorch backend when overlap is enabled." - ) + raise ValueError( + "Context-only requests are not supported on the PyTorch backend when overlap is enabled. " + "Set TorchLlmArgs.disable_overlap_scheduler=True or disable disaggregated serving." + )
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (3)
tensorrt_llm/executor/worker.py
(7 hunks)tensorrt_llm/llmapi/llm.py
(3 hunks)tensorrt_llm/llmapi/llm_args.py
(5 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/llmapi/llm.py
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py
: Python code must target Python 3.8+
Python indentation: 4 spaces, no tabs
Maintain module namespace in imports (from package.subpackage import foo; then use foo.SomeClass())
Python file names use snake_case
Python class names use PascalCase
Python functions/methods and local variables use snake_case; variables starting with a number get k_ prefix (e.g., k_99th_percentile)
Global variables use G_ prefixed UPPER_SNAKE_CASE (e.g., G_MY_GLOBAL)
Constants use UPPER_SNAKE_CASE in Python
Avoid shadowing variables from outer scopes in Python
Initialize all externally visible members of a Python class in init
Prefer docstrings for interfaces used outside a file; comments for local code
Use Google-style docstrings for classes and functions (Sphinx-parsable)
Document attributes/variables inline with short docstrings
Avoid reflection when simple alternatives exist (e.g., prefer explicit parameters over dict(**locals()))
In try/except, catch the narrowest exceptions possible
For duck-typing with try/except, keep try body minimal and put logic in else
Files:
tensorrt_llm/llmapi/llm_args.py
tensorrt_llm/executor/worker.py
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Prepend NVIDIA copyright header (current year) to all source files
Files:
tensorrt_llm/llmapi/llm_args.py
tensorrt_llm/executor/worker.py
🧠 Learnings (1)
📚 Learning: 2025-07-17T09:01:27.402Z
Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.402Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
Applied to files:
tensorrt_llm/executor/worker.py
🧬 Code graph analysis (2)
tensorrt_llm/llmapi/llm_args.py (4)
tensorrt_llm/llmapi/llm.py (2)
tokenizer
(682-686)tokenizer
(689-690)tensorrt_llm/llmapi/tokenizer.py (4)
TokenizerBase
(24-25)_llguidance_tokenizer_info
(329-333)_xgrammar_tokenizer_info
(293-326)tokenizer_factory
(270-290)tensorrt_llm/_torch/pyexecutor/config.py (1)
update_executor_config
(129-172)tensorrt_llm/_torch/speculative/auto_heuristic.py (1)
suggest_spec_config
(1-17)
tensorrt_llm/executor/worker.py (2)
tensorrt_llm/llmapi/llm_args.py (5)
PybindMirror
(564-708)TorchLlmArgs
(2106-2521)get_executor_config
(1828-1897)get_executor_config
(2448-2453)parallel_config
(1341-1342)tensorrt_llm/_torch/pyexecutor/py_executor_creator.py (1)
create_py_executor
(205-473)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (6)
tensorrt_llm/llmapi/llm_args.py (2)
47-49
: Imports LGTM.Adding _SchedulerConfig and _GuidedDecodingConfig aligns with new guided-decoding wiring.
60-61
: Tokenizer helpers import LGTM.Consistent with the helper usage below.
tensorrt_llm/executor/worker.py (4)
85-87
: Backend detection LGTM.Clear split between PyTorch and others.
157-160
: Engine path selection LGTM.Clear routing: llm_args drives the PyTorch/autodeploy creation; else, legacy C++ flow.
464-485
: Default max_tokens deduction LGTM given persisted executor_config.With the persistence fixes above, this logic should be reliable across backends.
784-794
: worker_main wiring LGTM.New args are passed through cleanly to the worker.
PR_Github #16061 [ run ] triggered by Bot |
PR_Github #16061 [ run ] completed with state |
/bot run |
PR_Github #16087 [ run ] triggered by Bot |
PR_Github #16087 [ run ] completed with state |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 0
♻️ Duplicate comments (3)
tensorrt_llm/executor/worker.py (3)
182-193
: Use self._is_pytorch_backend for LoRA manager pathRely on the already computed boolean instead of introspecting executor_config; reduces brittleness and handles None cleanly. Note: There’s a known LoRA adapter CPU cache race (July 17 learning); unchanged here, just flagging awareness.
Apply this diff:
- if getattr(self._executor_config, "backend", - "") == "pytorch" and lora_config is not None: + if self._is_pytorch_backend and lora_config is not None:
103-111
: Do not assert on executor_config; ignore it with a warning and derive from llm_argsAsserting here will break legacy call sites that still pass a non-None executor_config during this transitional PR (reflected by repeated pipeline failures). Keep the already-added persistence to self._executor_config.
Apply this diff:
def _create_py_executor(executor_config): - assert executor_config is None, "expect an empty executor_config is _create_py_executor" - executor_config = llm_args.get_executor_config(hf_model_dir) - # Persist so downstream code (e.g., default max_tokens deduction) has access - self._executor_config = executor_config + if executor_config is not None: + logger.warning( + "Ignoring provided executor_config; deriving from llm_args for PyTorch backend." + ) + executor_config = llm_args.get_executor_config(hf_model_dir) + # Persist so downstream code (e.g., default max_tokens deduction) has access + self._executor_config = executor_config executor_config.logits_post_processor_config = tllm.LogitsPostProcessorConfig( processor_batched=batched_logits_processor, replicate=False) comm_ranks, device_ids = _get_comm_ranks_device_id()
136-156
: Persist executor_config on the TRT/C++ path to unblock max_tokens deductionWhen using the TRT path, self._executor_config remains None, and _deduce_max_tokens() can fail when max_tokens is not provided. Persist it after wiring parallel_config. Also, replace the assert with an explicit exception.
Apply this diff:
def _create_engine(executor_config): if executor_config is None: executor_config = tllm.ExecutorConfig(1) executor_config.logits_post_processor_config = tllm.LogitsPostProcessorConfig( processor_batched=batched_logits_processor, replicate=False) comm_ranks, device_ids = _get_comm_ranks_device_id() executor_config.parallel_config = tllm.ParallelConfig( participant_ids=comm_ranks, device_ids=device_ids) + # Persist so downstream code (e.g., default max_tokens deduction) has access + self._executor_config = executor_config @@ - assert not hasattr(executor_config, "backend") + if hasattr(executor_config, "backend"): + raise ValueError( + "C++/TRT path expects a plain tllm.ExecutorConfig without a 'backend' attribute." + ) return tllm.Executor(engine, tllm.ModelType.DECODER_ONLY, executor_config)
🧹 Nitpick comments (6)
tensorrt_llm/executor/worker.py (6)
63-65
: New parameters in init: clarify contract and docstringsYou’ve added hf_model_dir and llm_args, which is good for the new construction flow. However, since the PyTorch path will ignore any provided executor_config later, the constructor should document that contract to avoid confusion. Minimal docstring addition recommended.
Would you like me to add concise Google-style docstrings for the new parameters in both GenerationExecutorWorker.init and worker_main?
94-102
: Guard against missing CUDA before set_devicedevice_id = global_rank % torch.cuda.device_count() will crash on systems without visible GPUs (mod by zero). Add a clear guard to fail fast with an actionable error.
Apply this diff:
def _get_comm_ranks_device_id(): - device_id = self.global_rank % torch.cuda.device_count() - torch.cuda.set_device(device_id) + if not torch.cuda.is_available() or torch.cuda.device_count() == 0: + raise RuntimeError( + "CUDA device not available. TensorRT-LLM worker requires at least one CUDA device." + ) + device_id = self.global_rank % torch.cuda.device_count() + torch.cuda.set_device(device_id) # Make sure C++ executor would use same devices/ranks as py_executor global_rank = global_mpi_rank() comm_ranks = mpi_comm().allgather(global_rank) device_ids = mpi_comm().allgather(device_id) return comm_ranks, device_ids
117-120
: Prefer explicit exception over assert for runtime validationAsserts may be stripped with -O. Use a ValueError for user-supplied configuration validation and fix the message grammar.
Apply this diff:
-assert hasattr( - executor_config, "backend" -), "executor_config should be with backend in _create_py_executor" +if not hasattr(executor_config, "backend"): + raise ValueError( + "executor_config is expected to have a 'backend' attribute in _create_py_executor" + )
126-126
: Guard optional GC threshold before passingIf garbage_collection_gen0_threshold can be None, avoid passing it to the creator. Pass only when set.
Apply this diff:
- args[ - "garbage_collection_gen0_threshold"] = llm_args.garbage_collection_gen0_threshold + if llm_args.garbage_collection_gen0_threshold is not None: + args["garbage_collection_gen0_threshold"] = ( + llm_args.garbage_collection_gen0_threshold + )Also verify the expected parameter name in create_py_executor/create_autodeploy_executor signatures across Torch paths.
157-159
: Backend dispatch is correct; minor naming nitThe dispatch correctly routes to the Torch/Autodeploy path when llm_args is provided. Minor nit: _create_py_executor also handles "_autodeploy", so consider renaming it to reflect both paths (e.g., _create_torch_or_autodeploy_executor) for clarity.
451-461
: Safer probe for kv_cache_transceiver attributeSome executors may not define kv_cache_transceiver. Use getattr to avoid AttributeError when overlap is enabled.
Apply this diff:
- is_disaggregated = self.engine.kv_cache_transceiver is not None + is_disaggregated = getattr(self.engine, "kv_cache_transceiver", None) is not None
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (1)
tensorrt_llm/executor/worker.py
(8 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py
: Python code must target Python 3.8+
Python indentation: 4 spaces, no tabs
Maintain module namespace in imports (from package.subpackage import foo; then use foo.SomeClass())
Python file names use snake_case
Python class names use PascalCase
Python functions/methods and local variables use snake_case; variables starting with a number get k_ prefix (e.g., k_99th_percentile)
Global variables use G_ prefixed UPPER_SNAKE_CASE (e.g., G_MY_GLOBAL)
Constants use UPPER_SNAKE_CASE in Python
Avoid shadowing variables from outer scopes in Python
Initialize all externally visible members of a Python class in init
Prefer docstrings for interfaces used outside a file; comments for local code
Use Google-style docstrings for classes and functions (Sphinx-parsable)
Document attributes/variables inline with short docstrings
Avoid reflection when simple alternatives exist (e.g., prefer explicit parameters over dict(**locals()))
In try/except, catch the narrowest exceptions possible
For duck-typing with try/except, keep try body minimal and put logic in else
Files:
tensorrt_llm/executor/worker.py
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Prepend NVIDIA copyright header (current year) to all source files
Files:
tensorrt_llm/executor/worker.py
🧠 Learnings (1)
📚 Learning: 2025-07-17T09:01:27.402Z
Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.402Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
Applied to files:
tensorrt_llm/executor/worker.py
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (4)
tensorrt_llm/executor/worker.py (4)
21-21
: Add TorchLlmArgs/PybindMirror import — LGTMImport is correct and aligns with downstream usage (PybindMirror at Line 501 and TorchLlmArgs across the worker).
85-87
: Backend detection via llm_args — LGTMDeriving _is_pytorch_backend from llm_args.backend centralizes backend checks and avoids scattered string comparisons.
793-794
: Plumbing hf_model_dir/llm_args into worker init — LGTMThe values are correctly forwarded into the worker for backend selection and configuration.
664-666
: I’ve requested inspection of theworker_kwargs
definition around theGenerationExecutorWorker
call and a search for any remainingexecutor_config
usage. Running these will confirm defaults and remove reliance on the old argument.
/bot run |
PR_Github #16096 [ run ] triggered by Bot |
/bot run |
PR_Github #16096 [ run ] completed with state |
PR_Github #16112 [ run ] triggered by Bot |
PR_Github #16112 [ run ] completed with state |
PR_Github #16353 [ run ] triggered by Bot |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 0
♻️ Duplicate comments (5)
tensorrt_llm/llmapi/llm_args.py (3)
1859-1862
: Potential None deref when FORCE_DETERMINISTIC=1 and kv_cache_config is None.When
self.kv_cache_config
is None, Line 1857 setsexecutor_config.kv_cache_config = None
. The subsequent lines 1861-1862 would then raise an AttributeError when trying to accessenable_block_reuse
andenable_partial_reuse
.Apply this defensive fix:
if os.getenv("FORCE_DETERMINISTIC", "0") == "1": # Disable KV cache reuse for deterministic mode + if executor_config.kv_cache_config is None: + executor_config.kv_cache_config = PybindMirror.maybe_to_pybind(KvCacheConfig()) executor_config.kv_cache_config.enable_block_reuse = False executor_config.kv_cache_config.enable_partial_reuse = False
1873-1877
: Consider adding backend validation for llguidance.The docstring states llguidance is PyTorch-only, but there's no runtime check to enforce this restriction. Consider adding a backend guard.
elif self.guided_decoding_backend == 'llguidance': + if self.backend not in ["pytorch", "_autodeploy"]: + raise ValueError("llguidance backend is supported only on PyTorch backend.") assert tokenizer is not None executor_config.guided_decoding_config = _GuidedDecodingConfig( backend=_GuidedDecodingConfig.GuidedDecodingBackend.LLGUIDANCE, **_llguidance_tokenizer_info(tokenizer))
1894-1896
: Handle None max_batch_size before calling suggest_spec_config.
executor_config.max_batch_size
can be None, and passing None tosuggest_spec_config
will raise an error.if spec_config is not None and spec_config.decoding_type == "AUTO": from tensorrt_llm._torch.speculative import suggest_spec_config - spec_config = suggest_spec_config(max_batch_size) + effective_batch_size = max_batch_size or self.max_batch_size or 1 + spec_config = suggest_spec_config(int(effective_batch_size))tensorrt_llm/executor/worker.py (2)
105-110
: Replace assertion with warning for better backward compatibility.The assertion could break existing code paths where
executor_config
might be passed. A warning would be more graceful.def _create_py_executor(executor_config): - assert executor_config is None, "expect an empty executor_config is _create_py_executor" + if executor_config is not None: + logger.warning("Ignoring provided executor_config; deriving from llm_args for PyTorch backend.") executor_config = llm_args.get_executor_config( hf_model_dir, tokenizer) # Persist so downstream code (e.g., default max_tokens deduction) has access self._executor_config = executor_config
139-147
: Persist executor_config for C++ path to avoid downstream issues.
self._executor_config
can remain None in the C++ path, which could break_deduce_max_tokens()
whenrequest.sampling_params.max_tokens
is unset.def _create_engine(executor_config): if executor_config is None: executor_config = tllm.ExecutorConfig(1) executor_config.logits_post_processor_config = tllm.LogitsPostProcessorConfig( processor_batched=batched_logits_processor, replicate=False) comm_ranks, device_ids = _get_comm_ranks_device_id() executor_config.parallel_config = tllm.ParallelConfig( participant_ids=comm_ranks, device_ids=device_ids) + # Persist so downstream code (e.g., default max_tokens deduction) has access + self._executor_config = executor_config
🧹 Nitpick comments (1)
tensorrt_llm/executor/worker.py (1)
96-103
: Consider extracting device configuration logic.The device setup logic inside a nested function could be extracted to improve readability and testability.
+ def _get_comm_ranks_and_device_ids(self): + """Get communication ranks and device IDs for all processes.""" + device_id = self.global_rank % torch.cuda.device_count() + torch.cuda.set_device(device_id) + # Make sure C++ executor would use same devices/ranks as py_executor + global_rank = global_mpi_rank() + comm_ranks = mpi_comm().allgather(global_rank) + device_ids = mpi_comm().allgather(device_id) + return comm_ranks, device_ids + def _get_comm_ranks_device_id(): - device_id = self.global_rank % torch.cuda.device_count() - torch.cuda.set_device(device_id) - # Make sure C++ executor would use same devices/ranks as py_executor - global_rank = global_mpi_rank() - comm_ranks = mpi_comm().allgather(global_rank) - device_ids = mpi_comm().allgather(device_id) - return comm_ranks, device_ids + return self._get_comm_ranks_and_device_ids()
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (8)
tensorrt_llm/executor/executor.py
(7 hunks)tensorrt_llm/executor/proxy.py
(1 hunks)tensorrt_llm/executor/worker.py
(8 hunks)tensorrt_llm/llmapi/llm.py
(3 hunks)tensorrt_llm/llmapi/llm_args.py
(5 hunks)tensorrt_llm/llmapi/mm_encoder.py
(2 hunks)tests/unittest/api_stability/references/llm.yaml
(1 hunks)tests/unittest/llmapi/test_llm_args.py
(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (4)
- tensorrt_llm/executor/proxy.py
- tests/unittest/api_stability/references/llm.yaml
- tensorrt_llm/llmapi/mm_encoder.py
- tensorrt_llm/executor/executor.py
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py
: Python code must target Python 3.8+
Python indentation: 4 spaces, no tabs
Maintain module namespace in imports (from package.subpackage import foo; then use foo.SomeClass())
Python file names use snake_case
Python class names use PascalCase
Python functions/methods and local variables use snake_case; variables starting with a number get k_ prefix (e.g., k_99th_percentile)
Global variables use G_ prefixed UPPER_SNAKE_CASE (e.g., G_MY_GLOBAL)
Constants use UPPER_SNAKE_CASE in Python
Avoid shadowing variables from outer scopes in Python
Initialize all externally visible members of a Python class in init
Prefer docstrings for interfaces used outside a file; comments for local code
Use Google-style docstrings for classes and functions (Sphinx-parsable)
Document attributes/variables inline with short docstrings
Avoid reflection when simple alternatives exist (e.g., prefer explicit parameters over dict(**locals()))
In try/except, catch the narrowest exceptions possible
For duck-typing with try/except, keep try body minimal and put logic in else
Files:
tests/unittest/llmapi/test_llm_args.py
tensorrt_llm/llmapi/llm.py
tensorrt_llm/llmapi/llm_args.py
tensorrt_llm/executor/worker.py
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Prepend NVIDIA copyright header (current year) to all source files
Files:
tests/unittest/llmapi/test_llm_args.py
tensorrt_llm/llmapi/llm.py
tensorrt_llm/llmapi/llm_args.py
tensorrt_llm/executor/worker.py
🧠 Learnings (1)
📚 Learning: 2025-07-17T09:01:27.402Z
Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.402Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
Applied to files:
tensorrt_llm/executor/worker.py
🧬 Code graph analysis (4)
tests/unittest/llmapi/test_llm_args.py (2)
tensorrt_llm/llmapi/llm_args.py (2)
get_executor_config
(1840-1913)get_executor_config
(2461-2468)tensorrt_llm/llmapi/llm.py (2)
tokenizer
(691-695)tokenizer
(698-699)
tensorrt_llm/llmapi/llm.py (1)
tensorrt_llm/llmapi/tokenizer.py (2)
TokenizerBase
(24-25)_xgrammar_tokenizer_info
(293-326)
tensorrt_llm/llmapi/llm_args.py (4)
tensorrt_llm/llmapi/tokenizer.py (4)
TokenizerBase
(24-25)_llguidance_tokenizer_info
(329-333)_xgrammar_tokenizer_info
(293-326)tokenizer_factory
(270-290)tensorrt_llm/_torch/pyexecutor/config.py (1)
update_executor_config
(129-172)tensorrt_llm/_torch/speculative/auto_heuristic.py (1)
suggest_spec_config
(1-17)tensorrt_llm/_torch/auto_deploy/llm_args.py (1)
get_pytorch_backend_config
(330-333)
tensorrt_llm/executor/worker.py (3)
tensorrt_llm/llmapi/llm_args.py (4)
TorchLlmArgs
(2122-2536)get_executor_config
(1840-1913)get_executor_config
(2461-2468)parallel_config
(1356-1357)tensorrt_llm/llmapi/tokenizer.py (1)
TokenizerBase
(24-25)tensorrt_llm/_torch/pyexecutor/py_executor_creator.py (1)
create_py_executor
(205-476)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (8)
tests/unittest/llmapi/test_llm_args.py (1)
441-446
: LGTM!The test correctly uses the public API
llm.args.get_executor_config(llm._hf_model_dir, llm.tokenizer)
instead of accessing private fields, which is a good practice.tensorrt_llm/llmapi/llm.py (3)
40-40
: Import cleanup needed.The import
_llguidance_tokenizer_info
has been removed from the public surface but it's being used in line 1876. Ensure the newget_executor_config
path is handling this correctly.
973-988
: LGTM!The change to set
executor_config=None
and passhf_model_dir
,tokenizer
, andllm_args
for the PyTorch backend properly aligns with the new architecture where the executor configuration is derived fromllm_args
.
990-1008
: LGTM!The validation method correctly identifies TensorRT-specific arguments and provides clear guidance to users about using the appropriate backend import. This helps prevent runtime errors from incompatible argument combinations.
tensorrt_llm/llmapi/llm_args.py (2)
2263-2268
: LGTM!The new
mm_encoder_only
field is properly documented and integrated. It provides a clear control for loading only the vision encoder part of multimodal models.
2461-2468
: LGTM!The override correctly propagates the
mm_encoder_only
flag to the executor config, ensuring the PyTorch backend can properly handle vision-only loading.tensorrt_llm/executor/worker.py (2)
160-162
: LGTM!The executor selection logic correctly chooses between PyTorch and C++ backends based on the presence of
llm_args
.
454-463
: LGTM!The overlap scheduler validation correctly checks the
llm_args
configuration and ensures that context-only requests are not supported when overlap is enabled in the PyTorch backend.
PR_Github #16353 [ run ] completed with state |
Signed-off-by: leslie-fang25 <[email protected]>
Signed-off-by: leslie-fang25 <[email protected]>
Signed-off-by: leslie-fang25 <[email protected]>
Signed-off-by: leslie-fang25 <[email protected]>
Signed-off-by: leslie-fang25 <[email protected]>
Signed-off-by: leslie-fang25 <[email protected]>
Signed-off-by: leslie-fang25 <[email protected]>
Signed-off-by: leslie-fang25 <[email protected]>
dceb34a
to
e9d9d99
Compare
/bot run |
PR_Github #16378 [ run ] triggered by Bot |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/executor/worker.py (1)
105-137
: Remove hard assert on executor_config being None; warn and ignore instead.Legacy callers may still forward a non-None executor_config. Prefer a warning and derive from llm_args.
Apply:
- def _create_py_executor(executor_config): - assert executor_config is None, "expect an empty executor_config is _create_py_executor" + def _create_py_executor(executor_config): + if executor_config is not None: + logger.warning( + "Ignoring provided executor_config; deriving from llm_args for PyTorch backend.") executor_config = llm_args.get_executor_config( hf_model_dir, tokenizer) # Persist so downstream code (e.g., default max_tokens deduction) has access self._executor_config = executor_config
♻️ Duplicate comments (3)
tensorrt_llm/llmapi/llm_args.py (1)
1840-1914
: Fix potential None deref for kv_cache_config in deterministic mode; guard llguidance to PyTorch; handle AUTO spec_config when max_batch_size is None.
- If FORCE_DETERMINISTIC=1 and self.kv_cache_config is None, dereferencing executor_config.kv_cache_config will raise. Initialize it first.
- llguidance backend is PyTorch-only per docs; add a backend guard to fail fast when misused. [duplicate of earlier bot comment.]
- suggest_spec_config(max_batch_size) will raise if max_batch_size is None; choose an effective default.
Apply:
if self.kv_cache_config is not None: executor_config.kv_cache_config = PybindMirror.maybe_to_pybind( self.kv_cache_config) - if os.getenv("FORCE_DETERMINISTIC", "0") == "1": + if os.getenv("FORCE_DETERMINISTIC", "0") == "1": # Disable KV cache reuse for deterministic mode - executor_config.kv_cache_config.enable_block_reuse = False - executor_config.kv_cache_config.enable_partial_reuse = False + if getattr(executor_config, "kv_cache_config", None) is None: + executor_config.kv_cache_config = _KvCacheConfig() + executor_config.kv_cache_config.enable_block_reuse = False + executor_config.kv_cache_config.enable_partial_reuse = False @@ - elif self.guided_decoding_backend == 'llguidance': + elif self.guided_decoding_backend == 'llguidance': assert tokenizer is not None + if self.backend != "pytorch": + raise ValueError( + "llguidance backend is supported only on the PyTorch backend.") executor_config.guided_decoding_config = _GuidedDecodingConfig( backend=_GuidedDecodingConfig.GuidedDecodingBackend.LLGUIDANCE, **_llguidance_tokenizer_info(tokenizer)) @@ - if spec_config is not None and spec_config.decoding_type == "AUTO": + if spec_config is not None and spec_config.decoding_type == "AUTO": from tensorrt_llm._torch.speculative import suggest_spec_config - spec_config = suggest_spec_config(max_batch_size) + effective_bs = max_batch_size or self.max_batch_size or 1 + spec_config = suggest_spec_config(int(effective_bs))tensorrt_llm/executor/executor.py (1)
359-362
: API surface change: create() adds hf_model_dir/tokenizer/llm_args and removed GC arg — sweep call sites and docs.Great direction routing via llm_args. Ensure no remaining call sites still pass garbage_collection_gen0_threshold and update docs accordingly. [duplicate of a prior review bot reminder.]
Run:
#!/bin/bash # 1) Old arg in Python call sites (should be gone) rg -nP "garbage_collection_gen0_threshold\s*=" --type=py -C3 # 2) Any references in non-Python files rg -n "garbage_collection_gen0_threshold" -g '!**/*.py' -C2 # 3) All GenerationExecutor.create call sites for unexpected kwargs rg -nP --type=py -C2 '\bGenerationExecutor\.create\s*\('tensorrt_llm/executor/worker.py (1)
139-159
: Persist executor_config in C++ path to support default max_tokens deduction.When creating a default tllm.ExecutorConfig (None passed), store it on self to unblock _deduce_max_tokens.
Apply:
def _create_engine(executor_config): if executor_config is None: - executor_config = tllm.ExecutorConfig(1) + executor_config = tllm.ExecutorConfig(1) + # Persist so downstream code (e.g., default max_tokens deduction) has access + self._executor_config = executor_config
🧹 Nitpick comments (10)
tensorrt_llm/llmapi/llm_args.py (4)
1-3
: Add NVIDIA copyright header.Per repository guidelines, prepend the current-year NVIDIA copyright header to all Python sources.
Apply to this file as well.
1898-1911
: Pass mm_encoder_only via update_executor_config instead of setting afterward.Keeping this invariant in one place reduces duplication. Optionally pass mm_encoder_only from TorchLlmArgs by forwarding getattr(self, "mm_encoder_only", False) here; then the Torch override can drop the manual assignment.
Apply:
update_executor_config( executor_config, backend=self.backend, pytorch_backend_config=self.get_pytorch_backend_config() if self.backend in ["pytorch", "_autodeploy"] else None, mapping=self.parallel_config.to_mapping(), speculative_config=spec_config, hf_model_dir=_hf_model_dir, max_input_len=self.max_input_len, max_seq_len=self.max_seq_len, + mm_encoder_only=getattr(self, "mm_encoder_only", False), checkpoint_format=None if self.backend == "_autodeploy" else self.checkpoint_format, checkpoint_loader=None if self.backend == "_autodeploy" else self.checkpoint_loader)
If you adopt this, you can simplify TorchLlmArgs.get_executor_config (see below).
2263-2269
: New mm_encoder_only flag: good addition; document stability.The field is clearly marked prototype and documented. Consider adding a brief note in TorchLlmArgs docstring or API docs about interactions with downstream components (e.g., disabling text decoder).
2461-2469
: Avoid redundant mm_encoder_only assignment if passed through update_executor_config.If you implement the earlier refactor to pass mm_encoder_only into update_executor_config, this explicit assignment becomes redundant and risks drift. Otherwise, keeping it here is fine.
- executor_config = super().get_executor_config(_hf_model_dir, tokenizer) - executor_config.mm_encoder_only = self.mm_encoder_only - return executor_config + executor_config = super().get_executor_config(_hf_model_dir, tokenizer) + # mm_encoder_only already populated via update_executor_config + return executor_configtensorrt_llm/llmapi/mm_encoder.py (2)
1-3
: Add NVIDIA copyright header.Please prepend the standard NVIDIA copyright header.
61-72
: Pass batched_logits_processor to executor.create for parity with LLM.This class is using the PyTorch path; while embeddings usually don't need logits processors, passing through args.batched_logits_processor keeps parity with _TorchLLM._build_model and avoids surprises.
self._executor = self._executor_cls.create( self._engine_dir, - executor_config=None, + executor_config=None, + batched_logits_processor=self.args.batched_logits_processor, model_world_size=self.args.parallel_config.world_size, mpi_session=self.mpi_session, reuse_mpi_comm=external_mpi_comm_available( self.args.parallel_config.world_size), is_llm_executor=True, # TODO: check if this is correct or needed hf_model_dir=self._hf_model_dir, tokenizer=self.tokenizer, llm_args=self.args)tensorrt_llm/executor/executor.py (1)
1-3
: Add NVIDIA copyright header.Please prepend the standard NVIDIA copyright header.
tensorrt_llm/llmapi/llm.py (1)
1-3
: Add NVIDIA copyright header.Please prepend the standard NVIDIA copyright header.
tensorrt_llm/executor/worker.py (2)
1-3
: Add NVIDIA copyright header.Please prepend the standard NVIDIA copyright header.
96-104
: Guard device_count==0 to avoid ZeroDivisionError in CPU-only contexts.Optional: if torch.cuda.device_count() == 0, set device_id=0 or raise a clearer error. This avoids modulo-by-zero.
- device_id = self.global_rank % torch.cuda.device_count() + ndev = torch.cuda.device_count() + if ndev == 0: + raise RuntimeError("No CUDA devices visible; TensorRT-LLM requires GPU.") + device_id = self.global_rank % ndev
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (8)
tensorrt_llm/executor/executor.py
(7 hunks)tensorrt_llm/executor/proxy.py
(1 hunks)tensorrt_llm/executor/worker.py
(8 hunks)tensorrt_llm/llmapi/llm.py
(3 hunks)tensorrt_llm/llmapi/llm_args.py
(5 hunks)tensorrt_llm/llmapi/mm_encoder.py
(2 hunks)tests/unittest/api_stability/references/llm.yaml
(1 hunks)tests/unittest/llmapi/test_llm_args.py
(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (3)
- tests/unittest/api_stability/references/llm.yaml
- tensorrt_llm/executor/proxy.py
- tests/unittest/llmapi/test_llm_args.py
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
**/*.py
: Python code must target Python 3.8+
Python indentation: 4 spaces, no tabs
Maintain module namespace in imports (from package.subpackage import foo; then use foo.SomeClass())
Python file names use snake_case
Python class names use PascalCase
Python functions/methods and local variables use snake_case; variables starting with a number get k_ prefix (e.g., k_99th_percentile)
Global variables use G_ prefixed UPPER_SNAKE_CASE (e.g., G_MY_GLOBAL)
Constants use UPPER_SNAKE_CASE in Python
Avoid shadowing variables from outer scopes in Python
Initialize all externally visible members of a Python class in init
Prefer docstrings for interfaces used outside a file; comments for local code
Use Google-style docstrings for classes and functions (Sphinx-parsable)
Document attributes/variables inline with short docstrings
Avoid reflection when simple alternatives exist (e.g., prefer explicit parameters over dict(**locals()))
In try/except, catch the narrowest exceptions possible
For duck-typing with try/except, keep try body minimal and put logic in else
Files:
tensorrt_llm/llmapi/llm.py
tensorrt_llm/executor/worker.py
tensorrt_llm/llmapi/mm_encoder.py
tensorrt_llm/llmapi/llm_args.py
tensorrt_llm/executor/executor.py
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
Prepend NVIDIA copyright header (current year) to all source files
Files:
tensorrt_llm/llmapi/llm.py
tensorrt_llm/executor/worker.py
tensorrt_llm/llmapi/mm_encoder.py
tensorrt_llm/llmapi/llm_args.py
tensorrt_llm/executor/executor.py
🧠 Learnings (1)
📚 Learning: 2025-07-17T09:01:27.402Z
Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.402Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.
Applied to files:
tensorrt_llm/executor/worker.py
🧬 Code graph analysis (5)
tensorrt_llm/llmapi/llm.py (2)
tensorrt_llm/_torch/auto_deploy/models/factory.py (1)
tokenizer
(48-50)tensorrt_llm/llmapi/tokenizer.py (2)
TokenizerBase
(24-25)_xgrammar_tokenizer_info
(293-326)
tensorrt_llm/executor/worker.py (3)
tensorrt_llm/llmapi/llm_args.py (5)
PybindMirror
(579-723)TorchLlmArgs
(2122-2536)get_executor_config
(1840-1913)get_executor_config
(2461-2468)parallel_config
(1356-1357)tensorrt_llm/llmapi/tokenizer.py (1)
TokenizerBase
(24-25)tensorrt_llm/_torch/pyexecutor/py_executor_creator.py (1)
create_py_executor
(205-476)
tensorrt_llm/llmapi/mm_encoder.py (4)
tensorrt_llm/llmapi/llm_args.py (3)
TorchLlmArgs
(2122-2536)world_size
(248-260)world_size
(269-276)tensorrt_llm/executor/executor.py (1)
create
(347-442)tensorrt_llm/llmapi/mpi_session.py (1)
external_mpi_comm_available
(57-66)tensorrt_llm/llmapi/llm.py (2)
tokenizer
(691-695)tokenizer
(698-699)
tensorrt_llm/llmapi/llm_args.py (3)
tensorrt_llm/llmapi/llm.py (2)
tokenizer
(691-695)tokenizer
(698-699)tensorrt_llm/llmapi/tokenizer.py (4)
TokenizerBase
(24-25)_llguidance_tokenizer_info
(329-333)_xgrammar_tokenizer_info
(293-326)tokenizer_factory
(270-290)tensorrt_llm/_torch/pyexecutor/config.py (1)
update_executor_config
(129-172)
tensorrt_llm/executor/executor.py (3)
tensorrt_llm/llmapi/llm_args.py (1)
TorchLlmArgs
(2122-2536)tensorrt_llm/llmapi/llm.py (2)
tokenizer
(691-695)tokenizer
(698-699)tensorrt_llm/llmapi/tokenizer.py (1)
TokenizerBase
(24-25)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Pre-commit Check
🔇 Additional comments (9)
tensorrt_llm/llmapi/llm_args.py (1)
60-61
: Tokenizer helpers import looks good.Importing _llguidance_tokenizer_info and _xgrammar_tokenizer_info from tokenizer centralizes backend-specific details. No issues.
tensorrt_llm/llmapi/mm_encoder.py (1)
58-60
: mm_encoder_only flag set early: LGTM.Setting TorchLlmArgs.mm_encoder_only before executor creation ensures the worker builds the correct executor_config.
tensorrt_llm/executor/executor.py (2)
384-391
: Propagating hf_model_dir/tokenizer/llm_args via worker_kwargs: LGTM.This cleanly carries state into the worker.
408-409
: is_llm_executor propagation maintained: LGTM.No functional concerns.
Also applies to: 419-419, 431-432
tensorrt_llm/llmapi/llm.py (1)
969-989
: PyTorch executor creation via llm_args: LGTM.Switching to executor_config=None + hf_model_dir/tokenizer/llm_args aligns with the new worker flow. return_logits parity preserved.
tensorrt_llm/executor/worker.py (4)
86-89
: _is_pytorch_backend detection: LGTM.Using llm_args is a reliable signal for the new path.
185-196
: LoRA manager hookup for PyTorch path: LGTM.Conditional PEFT manager retrieval is correct when backend is PyTorch and LoRA is enabled.
454-463
: Overlap scheduler vs. disaggregated context-only requests: LGTM.The guard prevents an unsupported combination; error message is clear.
652-670
: Ensure new parameters are forwarded in proxy creationI’ve confirmed in
executor.py
thatworker_kwargs
is built with the new keys:
- It includes
"hf_model_dir": hf_model_dir
,"tokenizer": tokenizer
, and"llm_args": llm_args
when callingGenerationExecutorProxy
.In
proxy.py
, these keys are received inworker_kwargs
and then passed through when spawning workers viaworker_main
, sinceworker_kwargs
is ultimately forwarded intact (with added IPC queues, etc.) into_start_executor_workers
.No missing forwarding paths were detected. All proxies correctly include
hf_model_dir
,tokenizer
, andllm_args
in the arguments forworker_main
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM
PR_Github #16378 [ run ] completed with state |
/bot run |
PR_Github #16429 [ run ] triggered by Bot |
PR_Github #16429 [ run ] completed with state |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM
Summary by CodeRabbit
New Features
Changes
Tests
Description
This is the first part of a refactor to build
PyExecutor
usingTorchLlmArgs
instead ofExecutorConfig
. This PR includes changes of:Replacing the dependency of
ExecutorConfig
withTorchLlmArgs
before creation ofGenerationExecutorWorker
.Creating
ExecutorConfig
insideGenerationExecutorWorker
to buildPyExecutor
for PyTorch Backend. This step will be further refactored in a subsequent PR.Test Coverage
GitHub Bot Help
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...
Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]
to print this help message.See details below for each supported subcommand.
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]
Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id
(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test
(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast
(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test
(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"
(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"
(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"
(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test
(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test
(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test
(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge
(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"
(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log
(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug
(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-list
parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.md
and the
scripts/test_to_stage_mapping.py
helper.kill
kill
Kill all running builds associated with pull request.
skip
skip --comment COMMENT
Skip testing for latest commit on pull request.
--comment "Reason for skipping build/test"
is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.