You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[TRTLLM-17518][refactor] Build EncoderExecutor from shared setup components instead of ModelEngine - #21410
EncoderExecutor now creates and owns an EncoderRunner directly. Shared setup functions in engine/model_setup.py support both executor paths. ModelEngine rejects encode_only and no longer accepts packed inputs in forward or warms up packed runners.
The change updates EncoderExecutor’s constructor and batch_forward API. Callers now read limits and pretrained_config from the executor. Verify that setup extraction preserves capacity resolution, compilation, refit, and cleanup behavior. The supplied summary reports bitwise-identical outputs and lower host overhead in tested configurations; these results are not independently verified here.
QA Engineer Review
Modified tests cover encoder creation and forwarding, runner graph configuration, ModelEngine rejection and refit behavior, encode limits and graph fallback, telemetry fallback, and integration encode behavior. The changed integration test is present in CI lists tests/integration/test_lists/test-db/l0_a100.yml and l0_l40s.yml, and in manual-QA list tests/integration/test_lists/qa/llm_function_core.txt. No test-list files changed.
The supplied CI report says one PR_Github run succeeded while its associated pipeline failed, and a second PR_Github run failed while its associated pipeline was unstable. Specific test failures are not supplied. Coverage verdict: needs follow-up.
Per-File QA Perspective
docs/source/torch/arch_overview.md — Verify that the architecture overview describes direct EncoderRunner setup and execution with PackedInputs.
tensorrt_llm/_torch/pyexecutor/encoder_executor.py — Verify factory setup, warmup, gather_context_logits forwarding, and cleanup after setup failure and shutdown.
tensorrt_llm/_torch/pyexecutor/engine/model_setup.py — Verify shared capacity resolution, model loading, compilation, attention-backend selection, and process setup for encoder and generation paths.
tensorrt_llm/_torch/pyexecutor/engine/runners/encoder.py — Verify EncoderRunner.create selects graph configuration only for EncodeCudaGraphConfig and passes runner settings correctly.
tensorrt_llm/_torch/pyexecutor/model_engine.py — Verify encode_only rejection, scheduled-request-only forwarding, and generation warmup, compilation, refit, and capacity behavior.
tensorrt_llm/_torch/pyexecutor/py_executor_creator.py — Verify encoder creation delegates to EncoderExecutor.create without constructing a ModelEngine.
tensorrt_llm/llmapi/llm.py — Verify encode limits and telemetry configuration use executor-owned fields.
tensorrt_llm/serve/openai_server.py — Verify the embedding batcher receives the executor’s resolved limits.
tests/integration/defs/accuracy/test_llm_api_pytorch_encode.py — The Qwen3 embedding test reads pretrained_config from the executor. The integration test is listed in CI and manual-QA lists.
tests/unittest/_torch/pyexecutor/engine/runners/test_common.py — Covers removal of the encoder runner’s registered ModelEngine initializer. This unit test is not an integration test-list entry.
tests/unittest/_torch/pyexecutor/engine/runners/test_encoder.py — Covers graph-configuration selection and runner arguments in EncoderRunner.create. This unit test is not an integration test-list entry.
tests/unittest/_torch/pyexecutor/engine/test_multimodal.py — Checks multimodal token-budget setup with patches redirected to shared setup. This unit test is not an integration test-list entry.
tests/unittest/_torch/pyexecutor/test_encoder_executor.py — Covers rejection of multi-GPU mappings and forwarding through the runner. This unit test is not an integration test-list entry.
tests/unittest/_torch/pyexecutor/test_model_engine.py — Covers encode_only rejection and refit compilation behavior; it removes the encoder graph-configuration test. This unit test is not an integration test-list entry.
tests/unittest/_torch/ray_orchestrator/multi_gpu/test_refit_compile_isolation.py — Updates a docstring reference only. This unit test is not an integration test-list entry.
tests/unittest/llmapi/test_llm_encode.py — Checks executor-owned batch limits and graph replay versus eager fallback. This unit test is not an integration test-list entry.
tests/unittest/llmapi/test_llm_telemetry.py — Covers executor-owned runtime configuration and fallback when it is missing. This unit test is not an integration test-list entry.
tests/unittest/_torch/executor/test_model_engine_initialization.py — Resets shared setup state in the fixture. This unit test is not an integration test-list entry.
tests/unittest/others/test_import_gpu_free.py — Checks the GPU-free import path while calling the relocated setup function. This unit test is not an integration test-list entry.
Description
LLM.encode() used to build a ModelEngine and route every step through ModelEngine.forward to an EncoderRunner, re-checking a dispatch that is fixed at initialization. EncoderExecutor now builds its own EncoderRunner and calls it directly, and the encode path no longer constructs ModelEngine at all.
To do that without duplicating setup code, the setup steps that do not depend on the runner family move out of ModelEngine into stateless functions in pyexecutor/engine/model_setup.py:
process and loading: init_process_setup, create_model_loader, prepare_loaded_model, resolve_attention_backend.
ModelEngine calls the same functions, so the PyExecutor path, the RL refit entry points and model_engine.model_loader are unchanged. The lifecycle that only generation needs (warmup scope, MoE load balancer and its three-phase collective shutdown, refit) stays in ModelEngine. EncoderExecutor.create loads, compiles and warms up the model, builds the runner with the new EncoderRunner.create, and on shutdown releases the loader, the runner's graphs and then the model, in the same order as ModelEngine's rank-local cleanup. ModelEngine now rejects encode_only.
LLM.encode(), the /v1/embeddings batcher and usage telemetry read max_batch_size, max_num_tokens, max_seq_len and pretrained_config from the executor. With no packed caller left, the packed-batch branches in ModelEngine.forward and _warmup_impl are removed.
Notes for review:
EncoderExecutor.__init__ takes the runner, the model loader and the resolved limits instead of a ModelEngine; EncoderExecutor.create builds them. batch_forward now takes only gather_context_logits, the one keyword its caller passes.
EncoderExecutor.create rejects more than one rank. The encode path runs in-process on one GPU; trtllm-serve embeddings and docs/source/features/embeddings.md already said so, but LLM(..., encode_only=True) did not check it.
Encode warmup no longer goes through ModelEngine.warmup or its warmup scope. With one rank moe_ep_size is 1, so no MoE load balancer exists and the MoE communication timeout guard covers no kernel; the scope has nothing to switch. The path also stops recording total_warmup_seconds, which nothing on it read.
Multimodal models are not in scope for encode_only. EncoderRunner consumes no multimodal inputs, and every documented and tested encode_only model is a text model, so the encode path no longer runs the multimodal encoder setup that ModelEngine ran for every model (MultimodalItemScheduler.maybe_create, setup_mm_encoder_attn_metadata, setup_mm_encoder_cuda_graphs).
Validated on H100 against the base commit with the same build:
outputs are bitwise identical for BERT and Qwen3-Embedding-0.6B (eager), BERT with encoder CUDA graphs, and BERT with TorchCompileConfig; TinyLlama context logits through PyExecutor are bitwise identical eager and with TorchCompileConfig;
capacity errors, copy_logits_to_host=False device views and the in-shutdown() memory release are unchanged;
a call trace shows that construction, the first call, the steps and shutdown of the encode path run no code from model_engine.py beyond importing it.
host overhead per encode() call drops by the removed ModelEngine forwarding (3.6–3.7 µs eager, 1.4 µs with CUDA graphs on BERT and Qwen3, batch size 1); no segment gets slower.
Test Coverage
tests/unittest/_torch/pyexecutor/test_encoder_executor.py (create rejects more than one rank)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
Any new dependencies have been scanned for license and vulnerabilities
…ut of ModelEngine
compile_model returns the compiled model, its caller and backend; ModelEngine keeps the public refit entry points.
Signed-off-by: Lori Ren <lorir@nvidia.com>
… of ModelEngine
These setup steps do not depend on the runner family, so EncoderExecutor can reuse them without ModelEngine.
Signed-off-by: Lori Ren <lorir@nvidia.com>
EncoderRunner.create builds the runner from llm_args and the resolved limits, so an executor can construct it without ModelEngine.
Signed-off-by: Lori Ren <lorir@nvidia.com>
lori-ren
changed the title
[TRTLLM-17518][refactor] Take ModelEngine out of EncoderExecutor's per-step path
[TRTLLM-17518][refactor] Build EncoderExecutor from shared setup components instead of ModelEngine
Nov 2, 2026
Encoder execution now creates and owns an EncoderRunner without using ModelEngine. Shared setup helpers handle model loading, capacity resolution, compilation, process setup, and backend selection. Encoding and serving code read runtime configuration and limits from EncoderExecutor.
Changes
Encoder runtime and shared model setup
Layer / File(s)
Summary
Shared model setup tensorrt_llm/_torch/pyexecutor/engine/model_setup.py
Adds helpers for capacity resolution, model loading and preparation, process setup, compilation, and attention-backend selection.
ModelEngine delegates setup and compilation to shared helpers, rejects encode_only configurations, and handles scheduled requests without a packed-input path. Tests reflect the updated setup and refit behavior.
EncoderExecutor.create performs setup and creates an EncoderRunner. The executor forwards packed inputs and manages runner and loader cleanup.
Encoder API and serving integration tensorrt_llm/llmapi/llm.py, tensorrt_llm/serve/openai_server.py, tests/integration/defs/accuracy/test_llm_api_pytorch_encode.py, tests/unittest/llmapi/*
Encoding, usage reporting, and embedding batch setup read runtime configuration and limits from EncoderExecutor.
sequenceDiagram
participant Creator as create_encoder_executor
participant Executor as EncoderExecutor.create
participant Setup as model_setup
participant Loader as ModelLoader
participant Runner as EncoderRunner.create
Creator->>Executor: pass model and distributed configuration
Executor->>Setup: initialize process and resolve capacities
Executor->>Setup: create loader and compile model
Setup->>Loader: load and prepare model
Executor->>Runner: pass model, limits, and runtime settings
Runner-->>Executor: return encoder runner
Executor-->>Creator: return encoder executor
Loading
Suggested reviewers: brnguyen2
Merge Risk: Low · up to c9c34
The encoder executor now builds its runner directly and no longer goes through the generation engine. The behavior looks correct, but a regression in loader cleanup after a failed startup, or in shutdown ordering, would currently go undetected. This is a small follow-up test, not a merge blocker.
Pre-merge checks: 5
Check name
Status
Explanation
Title check
The title follows the required ticket and type format and clearly describes the main refactor: building EncoderExecutor from shared setup components instead of ModelEngine.
Description check
The description explains the problem, solution, scope, lifecycle changes, validation results, and relevant test coverage. It also includes the required checklist and documents the encode-only API and …
Docstring Coverage
Docstring coverage is 82.86% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentati…
Linked Issues check
Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check
Check skipped because no linked issues were found for this pull request.
Full details: Description check
Explanation
The description explains the problem, solution, scope, lifecycle changes, validation results, and relevant test coverage. It also includes the required checklist and documents the encode-only API and single-rank limitations.
Full details: Docstring Coverage
Explanation
Docstring coverage is 82.86% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 35 functions across 18 files.
Finishing Touches
Generate unit tests (beta)
Create a new PR
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Warning
Some tools did not complete. Review the errors below.
ast-grep (0.51.1)
ast-grep could not parse rule config: /ast-grep-rules/ast-grep-essentials/rules/typescript/security/fetch-response-statustext-check-typescript.yml
Tip
🎯 Customize your risk assessments
Guide how CodeRabbit assesses merge risk and security architecture risk. Highlight sensitive changes and explain your security boundaries using Merge Risk Instructions and Security Risk Instructions in your review settings. Security architecture risk assessments require the Advanced plan or above.
✅ Automatic approvals, with your rules
Let CodeRabbit approve pull requests once required threads are resolved, pre-merge checks pass, and your risk limits are met. Enable Automatic approvals in review settings, with overrides for teams or file paths.
CodeRabbit Verify
Run CodeRabbit Verify to test the changes in this pull request.
Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements.
test_init_from_args_forwards_vision_gpu_normalize still patched the
function on model_engine, which no longer has it: the EncoderExecutor
refactor (PR NVIDIA#21410) moved it to engine/model_setup.py, so both cases
fail with AttributeError in the CPU stage. Patch it on model_setup, as
the other _init_from_args test in the file does.
Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
test_init_from_args_forwards_vision_gpu_normalize still patched the
function on model_engine, which no longer has it: the EncoderExecutor
refactor (PR NVIDIA#21410) moved it to engine/model_setup.py, so both cases
fail with AttributeError in the CPU stage. Patch it on model_setup, as
the other _init_from_args test in the file does.
Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Dev Engineer Review
EncoderExecutornow creates and owns anEncoderRunnerdirectly. Shared setup functions inengine/model_setup.pysupport both executor paths.ModelEnginerejectsencode_onlyand no longer accepts packed inputs inforwardor warms up packed runners.The change updates
EncoderExecutor’s constructor andbatch_forwardAPI. Callers now read limits andpretrained_configfrom the executor. Verify that setup extraction preserves capacity resolution, compilation, refit, and cleanup behavior. The supplied summary reports bitwise-identical outputs and lower host overhead in tested configurations; these results are not independently verified here.QA Engineer Review
Modified tests cover encoder creation and forwarding, runner graph configuration,
ModelEnginerejection and refit behavior, encode limits and graph fallback, telemetry fallback, and integration encode behavior. The changed integration test is present in CI liststests/integration/test_lists/test-db/l0_a100.ymlandl0_l40s.yml, and in manual-QA listtests/integration/test_lists/qa/llm_function_core.txt. No test-list files changed.The supplied CI report says one
PR_Githubrun succeeded while its associated pipeline failed, and a secondPR_Githubrun failed while its associated pipeline was unstable. Specific test failures are not supplied. Coverage verdict: needs follow-up.Per-File QA Perspective
docs/source/torch/arch_overview.md— Verify that the architecture overview describes directEncoderRunnersetup and execution withPackedInputs.tensorrt_llm/_torch/pyexecutor/encoder_executor.py— Verify factory setup, warmup,gather_context_logitsforwarding, and cleanup after setup failure and shutdown.tensorrt_llm/_torch/pyexecutor/engine/model_setup.py— Verify shared capacity resolution, model loading, compilation, attention-backend selection, and process setup for encoder and generation paths.tensorrt_llm/_torch/pyexecutor/engine/runners/encoder.py— VerifyEncoderRunner.createselects graph configuration only forEncodeCudaGraphConfigand passes runner settings correctly.tensorrt_llm/_torch/pyexecutor/model_engine.py— Verifyencode_onlyrejection, scheduled-request-only forwarding, and generation warmup, compilation, refit, and capacity behavior.tensorrt_llm/_torch/pyexecutor/py_executor_creator.py— Verify encoder creation delegates toEncoderExecutor.createwithout constructing aModelEngine.tensorrt_llm/llmapi/llm.py— Verify encode limits and telemetry configuration use executor-owned fields.tensorrt_llm/serve/openai_server.py— Verify the embedding batcher receives the executor’s resolved limits.tests/integration/defs/accuracy/test_llm_api_pytorch_encode.py— The Qwen3 embedding test readspretrained_configfrom the executor. The integration test is listed in CI and manual-QA lists.tests/unittest/_torch/pyexecutor/engine/runners/test_common.py— Covers removal of the encoder runner’s registeredModelEngineinitializer. This unit test is not an integration test-list entry.tests/unittest/_torch/pyexecutor/engine/runners/test_encoder.py— Covers graph-configuration selection and runner arguments inEncoderRunner.create. This unit test is not an integration test-list entry.tests/unittest/_torch/pyexecutor/engine/test_multimodal.py— Checks multimodal token-budget setup with patches redirected to shared setup. This unit test is not an integration test-list entry.tests/unittest/_torch/pyexecutor/test_encoder_executor.py— Covers rejection of multi-GPU mappings and forwarding through the runner. This unit test is not an integration test-list entry.tests/unittest/_torch/pyexecutor/test_model_engine.py— Coversencode_onlyrejection and refit compilation behavior; it removes the encoder graph-configuration test. This unit test is not an integration test-list entry.tests/unittest/_torch/ray_orchestrator/multi_gpu/test_refit_compile_isolation.py— Updates a docstring reference only. This unit test is not an integration test-list entry.tests/unittest/llmapi/test_llm_encode.py— Checks executor-owned batch limits and graph replay versus eager fallback. This unit test is not an integration test-list entry.tests/unittest/llmapi/test_llm_telemetry.py— Covers executor-owned runtime configuration and fallback when it is missing. This unit test is not an integration test-list entry.tests/unittest/_torch/executor/test_model_engine_initialization.py— Resets shared setup state in the fixture. This unit test is not an integration test-list entry.tests/unittest/others/test_import_gpu_free.py— Checks the GPU-free import path while calling the relocated setup function. This unit test is not an integration test-list entry.Description
LLM.encode()used to build aModelEngineand route every step throughModelEngine.forwardto anEncoderRunner, re-checking a dispatch that is fixed at initialization.EncoderExecutornow builds its ownEncoderRunnerand calls it directly, and the encode path no longer constructsModelEngineat all.To do that without duplicating setup code, the setup steps that do not depend on the runner family move out of
ModelEngineinto stateless functions inpyexecutor/engine/model_setup.py:infer_max_seq_len,resolve_max_seq_len,resolve_max_num_tokens;compile_model,apply_torch_compile,remove_torch_compile, userbuffers;init_process_setup,create_model_loader,prepare_loaded_model,resolve_attention_backend.ModelEnginecalls the same functions, so the PyExecutor path, the RL refit entry points andmodel_engine.model_loaderare unchanged. The lifecycle that only generation needs (warmup scope, MoE load balancer and its three-phase collective shutdown, refit) stays inModelEngine.EncoderExecutor.createloads, compiles and warms up the model, builds the runner with the newEncoderRunner.create, and on shutdown releases the loader, the runner's graphs and then the model, in the same order asModelEngine's rank-local cleanup.ModelEnginenow rejectsencode_only.LLM.encode(), the/v1/embeddingsbatcher and usage telemetry readmax_batch_size,max_num_tokens,max_seq_lenandpretrained_configfrom the executor. With no packed caller left, the packed-batch branches inModelEngine.forwardand_warmup_implare removed.Notes for review:
EncoderExecutor.__init__takes the runner, the model loader and the resolved limits instead of aModelEngine;EncoderExecutor.createbuilds them.batch_forwardnow takes onlygather_context_logits, the one keyword its caller passes.EncoderExecutor.createrejects more than one rank. The encode path runs in-process on one GPU;trtllm-serve embeddingsanddocs/source/features/embeddings.mdalready said so, butLLM(..., encode_only=True)did not check it.ModelEngine.warmupor its warmup scope. With one rankmoe_ep_sizeis 1, so no MoE load balancer exists and the MoE communication timeout guard covers no kernel; the scope has nothing to switch. The path also stops recordingtotal_warmup_seconds, which nothing on it read.encode_only.EncoderRunnerconsumes no multimodal inputs, and every documented and testedencode_onlymodel is a text model, so the encode path no longer runs the multimodal encoder setup thatModelEngineran for every model (MultimodalItemScheduler.maybe_create,setup_mm_encoder_attn_metadata,setup_mm_encoder_cuda_graphs).openai_server.pystill reads_encoder_executorin one place, so [TRTLLM-17272][feat] Add public LLM health queries and a private-access lint for the serving frontends #21393's private-access baseline is unchanged.Validated on H100 against the base commit with the same build:
TorchCompileConfig; TinyLlama context logits through PyExecutor are bitwise identical eager and withTorchCompileConfig;copy_logits_to_host=Falsedevice views and the in-shutdown()memory release are unchanged;model_engine.pybeyond importing it.encode()call drops by the removedModelEngineforwarding (3.6–3.7 µs eager, 1.4 µs with CUDA graphs on BERT and Qwen3, batch size 1); no segment gets slower.Test Coverage
tests/unittest/_torch/pyexecutor/test_encoder_executor.py(createrejects more than one rank)tests/unittest/_torch/pyexecutor/engine/runners/test_encoder.py(EncoderRunner.create)tests/unittest/_torch/pyexecutor/test_model_engine.py(ModelEnginerejectsencode_only; moved patch targets)tests/unittest/llmapi/test_llm_encode.pytests/unittest/llmapi/test_llm_telemetry.pytests/unittest/llmapi/apps/_test_openai_embeddings.pytests/integration/defs/accuracy/test_llm_api_pytorch_encode.py::TestEncoderEncodePR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.