[None][feat] Clean up ngram auto mode, add max_concurrency to configs #6676
Conversation
📝 WalkthroughWalkthroughThis update introduces a concurrency-aware heuristic for automatically configuring speculative decoding in the PyTorch backend. It adds a new function to generate default speculative configs based on batch size, updates class constructors and logic to support a Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant LLMAPI
participant suggest_spec_config
participant Drafter
participant Scheduler
User->>LLMAPI: Request model with AUTO speculative config
LLMAPI->>suggest_spec_config: Call with max_batch_size
suggest_spec_config-->>LLMAPI: Return NGramDecodingConfig
LLMAPI->>Drafter: Initialize with max_concurrency
LLMAPI->>Scheduler: Pass spec_config and drafter
Scheduler->>Drafter: should_use_spec_decode(requests)
Drafter-->>Scheduler: Return True/False based on max_concurrency
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~15 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Note ⚡️ Unit Test Generation is now available in beta!Learn more here, or try it out under "Finishing Touches" below. 📜 Recent review detailsConfiguration used: .coderabbit.yaml 📒 Files selected for processing (9)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (7)
🧰 Additional context used🪛 Ruff (0.12.2)tensorrt_llm/_torch/speculative/auto_heuristic.py1-1: Undefined name (F821) ⏰ 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)
🔇 Additional comments (1)
✨ 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 comments)
Other keywords and placeholders
Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
tensorrt_llm/llmapi/llm_args.py (1)
539-543: Fix docstring formatting per coding standards.The simplified docstring content looks good and aligns with the PR objective of reducing excessive documentation. However, the docstring format needs a blank line between the summary and description.
Apply this diff to fix the docstring formatting:
""" Configuration for auto speculative decoding. + This config will automatically select a good, draft-model free speculation algorithm with some heuristic. Attributes that are inherited from the base class are ignored. """
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (9)
tensorrt_llm/_torch/pyexecutor/py_executor.py(1 hunks)tensorrt_llm/_torch/pyexecutor/py_executor_creator.py(1 hunks)tensorrt_llm/_torch/speculative/__init__.py(2 hunks)tensorrt_llm/_torch/speculative/auto_heuristic.py(1 hunks)tensorrt_llm/_torch/speculative/drafter.py(2 hunks)tensorrt_llm/_torch/speculative/model_drafter.py(1 hunks)tensorrt_llm/_torch/speculative/ngram.py(1 hunks)tensorrt_llm/llmapi/llm.py(2 hunks)tensorrt_llm/llmapi/llm_args.py(3 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.py: The code developed for TensorRT-LLM should conform to Python 3.8+.
Indent Python code with 4 spaces. Do not use tabs.
Always maintain the namespace when importing in Python, even if only one class or function from a module is used.
Python filenames should use snake_case (e.g., some_file.py).
Python classes should use PascalCase (e.g., class SomeClass).
Python functions and methods should use snake_case (e.g., def my_awesome_function():).
Python local variables should use snake_case. Prefix k for variable names that start with a number (e.g., k_99th_percentile = ...).
Python global variables should use upper snake_case and prefix G (e.g., G_MY_GLOBAL = ...).
Python constants should use upper snake_case (e.g., MY_CONSTANT = ...).
Avoid shadowing variables declared in an outer scope in Python.
Initialize all externally visible members of a class in the constructor in Python.
For interfaces that may be used outside a file, prefer docstrings over comments in Python.
Comments in Python should be reserved for code within a function, or interfaces that are local to a file.
Use Google style docstrings for classes and functions in Python, which can be parsed by Sphinx.
Attributes and variables in Python can be documented inline; attribute docstrings will be rendered under the docstring for the class.
Avoid using reflection in Python when functionality can be easily achieved without it.
When using try-except blocks in Python, limit the except to the smallest set of errors possible.
When using try-except blocks to handle multiple possible variable types in Python, keep the body of the try as small as possible, using the else block to implement the logic.
Files:
tensorrt_llm/_torch/speculative/__init__.pytensorrt_llm/_torch/speculative/drafter.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytensorrt_llm/_torch/speculative/model_drafter.pytensorrt_llm/_torch/speculative/auto_heuristic.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/llmapi/llm.pytensorrt_llm/_torch/speculative/ngram.py
**/*.{cpp,h,hpp,cc,cxx,cu,py}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
All TensorRT-LLM Open Source Software code should contain an NVIDIA copyright header that includes the current year. This includes .cpp, .h, .cu, .py, and any other source files which are compiled or interpreted.
Files:
tensorrt_llm/_torch/speculative/__init__.pytensorrt_llm/_torch/speculative/drafter.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytensorrt_llm/_torch/speculative/model_drafter.pytensorrt_llm/_torch/speculative/auto_heuristic.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/llmapi/llm.pytensorrt_llm/_torch/speculative/ngram.py
🧠 Learnings (6)
📚 Learning: in tensorrt_llm/executor/worker.py, the lora adapter cache optimization logic that checks `is_adapte...
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/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.py
📚 Learning: applies to **/*.py : the code developed for tensorrt-llm should conform to python 3.8+....
Learnt from: CR
PR: NVIDIA/TensorRT-LLM#0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-08-06T08:45:40.701Z
Learning: Applies to **/*.py : The code developed for TensorRT-LLM should conform to Python 3.8+.
Applied to files:
tensorrt_llm/llmapi/llm_args.pytensorrt_llm/llmapi/llm.py
📚 Learning: in tensorrt-llm's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()...
Learnt from: yechank-nvidia
PR: NVIDIA/TensorRT-LLM#6254
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:1201-1204
Timestamp: 2025-07-22T09:22:14.726Z
Learning: In TensorRT-LLM's multimodal processing pipeline, shared tensor recovery using `from_shared_tensor()` is only needed during the context phase. Generation requests reuse the already-recovered tensor data and only need to call `strip_for_generation()` to remove unnecessary multimodal data while preserving the recovered tensors. This avoids redundant tensor recovery operations during generation.
Applied to files:
tensorrt_llm/llmapi/llm.py
📚 Learning: in tensorrt-llm, test files (files under tests/ directories) do not require nvidia copyright headers...
Learnt from: galagam
PR: NVIDIA/TensorRT-LLM#6487
File: tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_trtllm_bench.py:1-12
Timestamp: 2025-08-06T13:58:07.506Z
Learning: In TensorRT-LLM, test files (files under tests/ directories) do not require NVIDIA copyright headers, unlike production source code files. Test files typically start directly with imports, docstrings, or code.
Applied to files:
tensorrt_llm/llmapi/llm.py
📚 Learning: in tensorrt-llm testing, it's common to have both cli flow tests (test_cli_flow.py) and pytorch api ...
Learnt from: moraxu
PR: NVIDIA/TensorRT-LLM#6303
File: tests/integration/test_lists/qa/examples_test_list.txt:494-494
Timestamp: 2025-07-28T17:06:08.621Z
Learning: In TensorRT-LLM testing, it's common to have both CLI flow tests (test_cli_flow.py) and PyTorch API tests (test_llm_api_pytorch.py) for the same model. These serve different purposes: CLI flow tests validate the traditional command-line workflow, while PyTorch API tests validate the newer LLM API backend. Both are legitimate and should coexist.
Applied to files:
tensorrt_llm/llmapi/llm.py
📚 Learning: in tensorrt-llm, examples directory can have different dependency versions than the root requirement...
Learnt from: yibinl-nvidia
PR: NVIDIA/TensorRT-LLM#6506
File: examples/models/core/mixtral/requirements.txt:3-3
Timestamp: 2025-08-01T15:14:45.673Z
Learning: In TensorRT-LLM, examples directory can have different dependency versions than the root requirements.txt file. Version conflicts between root and examples dependencies are acceptable because examples are designed to be standalone and self-contained.
Applied to files:
tensorrt_llm/llmapi/llm.py
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/speculative/auto_heuristic.py
1-1: Undefined name DecodingBaseConfig
(F821)
tensorrt_llm/llmapi/llm_args.py
546-548: 1 blank line required between summary line and description
(D205)
⏰ 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 (12)
tensorrt_llm/_torch/speculative/__init__.py (1)
1-1: LGTM! Clean integration of the new heuristic function.The import and export of
suggest_spec_configproperly exposes the centralized heuristic logic for automatic speculative decoding configuration, aligning with the PR's goal of moving heuristic logic to a dedicated module.Also applies to: 27-27
tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
874-877: LGTM! Helpful clarifying comment.This comment clearly explains the conditional behavior of
_prepare_draft_requests()based on speculation state, improving code readability and maintainability.tensorrt_llm/_torch/speculative/ngram.py (1)
170-170: LGTM! Proper integration with concurrency gating.The constructor correctly passes
spec_config.max_concurrencyto the parentDrafterclass, enabling concurrency-based control of speculative decoding for NGram drafters.tensorrt_llm/_torch/speculative/model_drafter.py (1)
49-50: LGTM! Consistent integration with concurrency gating.The constructor properly passes
spec_config.max_concurrencyto the parentDrafterclass, maintaining consistency with other drafter implementations and enabling concurrency-based speculation control.tensorrt_llm/_torch/speculative/drafter.py (1)
12-13: LGTM! Clean addition of concurrency configuration.The
__init__method properly accepts and stores themax_concurrencyparameter, enabling concurrency-based control of speculative decoding across drafter implementations.tensorrt_llm/_torch/speculative/auto_heuristic.py (1)
9-17: LGTM! Solid heuristic implementation.The batch-size-based heuristic logic is well-designed:
- Smaller batches (≤4) get longer draft lengths (max_draft_len=5) for better speculation
- Larger batches (>4) get shorter draft lengths (max_draft_len=3) to reduce overhead
- max_concurrency=32 effectively disables speculation for large batch sizes as documented
- Local import avoids potential circular dependency issues
tensorrt_llm/llmapi/llm.py (2)
35-36: LGTM! Clean import removal.Correctly removed the unused
NGramDecodingConfigimport since the manual heuristic logic has been replaced with the centralizedsuggest_spec_configfunction.
1019-1021: Excellent refactoring to centralize speculation heuristics.The replacement of manual heuristic logic with the centralized
suggest_spec_configfunction improves maintainability and consistency. The local import pattern avoids potential circular dependencies.Note: The case change from "AUTO" to "Auto" should be consistent with changes in the
AutoDecodingConfigclass.tensorrt_llm/_torch/pyexecutor/py_executor_creator.py (1)
172-178: LGTM! Excellent defensive compatibility check.This addition properly handles the interaction between speculative decoding and the overlap scheduler:
- Defensive checks prevent unnecessary modifications when overlap scheduler is already disabled
- Uses the appropriate
support_overlap_scheduler()method to determine compatibility- Provides clear warning message when automatic adjustment occurs
- Prevents potential runtime issues from incompatible feature combinations
tensorrt_llm/llmapi/llm_args.py (3)
344-347: LGTM! Clean addition of concurrency control.The new
max_concurrencyfield is well-implemented with proper typing, clear documentation, and appropriate default value. The PyTorch-specific comment helps clarify its scope of usage.
549-549: LGTM! Consistent with mapping change.The
decoding_typechange from"AUTO"to"Auto"maintains consistency with the updatedfrom_dictmapping and follows a more conventional casing pattern.
551-552: LGTM! Proper backend restriction implementation.The
supports_backendmethod correctly restrictsAutoDecodingConfigto the PyTorch backend only, which aligns with the auto speculation functionality being PyTorch-specific. The implementation follows the established pattern used by other decoding configuration classes.
mikeiovine
force-pushed
the
fix-ngram-auto
branch
from
533a6fe to
d8c9081
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
tensorrt_llm/_torch/speculative/auto_heuristic.py (1)
1-1: Fix the return type annotation and import.The return type annotation uses
"DecodingBaseConfig"but the function actually returnsNGramDecodingConfig. Additionally,DecodingBaseConfigis not imported, causing the static analysis error.Apply this diff to fix the type annotation:
+from tensorrt_llm.llmapi.llm_args import NGramDecodingConfig + -def suggest_spec_config(max_batch_size: int) -> "DecodingBaseConfig": +def suggest_spec_config(max_batch_size: int) -> NGramDecodingConfig:
🧹 Nitpick comments (1)
tensorrt_llm/_torch/speculative/auto_heuristic.py (1)
9-17: Consider moving the import to module level.The import is currently inside the function, which works but is unconventional. Moving it to the top would improve readability and follows standard Python conventions.
+from tensorrt_llm.llmapi.llm_args import NGramDecodingConfig + def suggest_spec_config(max_batch_size: int) -> NGramDecodingConfig: """ Suggests a reasonable draft model free speculation scheme. Used when the user specifies spec_mode == AUTO. For now, we always use an ngram scheme that gets disabled at BS>=32. """ - from tensorrt_llm.llmapi.llm_args import NGramDecodingConfig return NGramDecodingConfig(
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (9)
tensorrt_llm/_torch/pyexecutor/py_executor.py(1 hunks)tensorrt_llm/_torch/pyexecutor/py_executor_creator.py(1 hunks)tensorrt_llm/_torch/speculative/__init__.py(2 hunks)tensorrt_llm/_torch/speculative/auto_heuristic.py(1 hunks)tensorrt_llm/_torch/speculative/drafter.py(2 hunks)tensorrt_llm/_torch/speculative/model_drafter.py(1 hunks)tensorrt_llm/_torch/speculative/ngram.py(1 hunks)tensorrt_llm/llmapi/llm.py(2 hunks)tensorrt_llm/llmapi/llm_args.py(2 hunks)
✅ Files skipped from review due to trivial changes (1)
- tensorrt_llm/_torch/pyexecutor/py_executor.py
🚧 Files skipped from review as they are similar to previous changes (6)
- tensorrt_llm/_torch/speculative/init.py
- tensorrt_llm/llmapi/llm.py
- tensorrt_llm/_torch/speculative/ngram.py
- tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
- tensorrt_llm/_torch/speculative/model_drafter.py
- tensorrt_llm/llmapi/llm_args.py
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: in tensorrt_llm/executor/worker.py, the lora adapter cache optimization logic that checks `is_adapte...
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/_torch/speculative/drafter.py
🪛 Ruff (0.12.2)
tensorrt_llm/_torch/speculative/auto_heuristic.py
1-1: Undefined name DecodingBaseConfig
(F821)
⏰ 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 (3)
tensorrt_llm/_torch/speculative/drafter.py (2)
12-13: LGTM! Constructor correctly initializes concurrency parameter.The
__init__method properly accepts and stores the optionalmax_concurrencyparameter, enabling concurrency-aware gating in the drafter.
31-32: LGTM! Concurrency gating logic is now correct.The logic correctly disables speculation when the number of requests reaches or exceeds
max_concurrency, which aligns with the PR objectives. This addresses the previous review concern about inverted logic.tensorrt_llm/_torch/speculative/auto_heuristic.py (1)
10-17: LGTM! Heuristic parameters are well-tuned.The heuristic logic appropriately adjusts parameters based on batch size:
- Smaller batches (≤4) get more aggressive speculation with longer drafts
- Larger batches get shorter drafts to balance performance
- Fixed
max_concurrency=32effectively disables speculation at high batch sizes as intendedThe boolean flags are consistently set to enable optimal ngram behavior.
mikeiovine
force-pushed
the
fix-ngram-auto
branch
from
d8c9081 to
95b2c55
Compare
|
/bot run |
|
PR_Github #14334 [ run ] triggered by Bot |
|
PR_Github #14334 [ run ] completed with state |
…onfig Signed-off-by: Mike Iovine <[email protected]>
mikeiovine
force-pushed
the
fix-ngram-auto
branch
from
95b2c55 to
27c401e
Compare
|
/bot reuse-pipeline |
|
PR_Github #14494 [ reuse-pipeline ] triggered by Bot |
|
PR_Github #14494 [ reuse-pipeline ] completed with state |
…NVIDIA#6676) Signed-off-by: Mike Iovine <[email protected]>
Description
Address some followups from the initial ngram auto PR.
_torch/speculative/auto_heuristic.py.NGramDecodingconfig is self-explanatory. We also should not be leaking these details in user-facing docstrings.is_auto_heuristicinNGramDecodingConfigmax_concurrencyfield toDecodingBaseConfig.Drafterwill now disable speculation if the number of active requests goes above this number. If nomax_concurrencyis given, speculation will always be on (default).should_use_spec_decodeshould be used to turn ngram off.Test Coverage
Existing tests cover ngram and auto. Tried to add a new accuracy tests but there are CUDA graph bugs to address. Opened https://nvbugspro.nvidia.com/bug/5441438 to follow up.
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-listparameter 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.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip 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-pipelineReuse 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.
Summary by CodeRabbit
Summary by CodeRabbit
New Features
Improvements
Bug Fixes
Documentation