[TRTC-1975] [feat] BREAKING: Enable chunked prefill by default#9499
[TRTC-1975] [feat] BREAKING: Enable chunked prefill by default#9499venkywonka wants to merge 5 commits intoNVIDIA:mainfrom
Conversation
coderabbitai
bot
changed the title
📝 WalkthroughWalkthroughChunked Prefill (enable_chunked_prefill) is now enabled by default in TensorRT-LLM, automating long sequence processing. CLI options and API defaults updated from False to True. Models not supporting this feature are automatically disabled with warnings. Documentation and tests updated to reflect new default behavior. Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
📝 Coding Plan
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
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)
tests/unittest/llmapi/test_llm_pytorch.py (2)
55-73: Double‑check llm_get_stats harness actually disables chunked prefill when the flag is FalseThe new
enable_chunked_prefillparameterization here makes sense, and you correctly forward it intollm_get_stats_test_harness(..., enable_chunked_prefill=enable_chunked_prefill, ...).However, in
tests/unittest/llmapi/test_llm.py,llm_get_stats_test_harnessonly does:if enable_chunked_prefill: llm_args_extra["enable_chunked_prefill"] = True llm_args_extra["max_num_tokens"] = 32and does nothing when
enable_chunked_prefillisFalse. Now that the global default isTrue, theFalsecases in this param set may still be running with chunked prefill enabled unless you also explicitly passenable_chunked_prefill=Falseinto the LLM args in that helper.Consider updating the harness so both branches set an explicit boolean (True/False) on the LLM, keeping
validate_stats(... enable_chunked_prefill=...)in sync with actual behavior.
75-93: Same consideration for async stats: ensure harness can truly disable chunked prefillSame as the sync version:
test_llm_get_stats_async’senable_chunked_prefillparameter is forwarded, but the sharedllm_get_stats_async_test_harnessintest_llm.pyonly setsenable_chunked_prefill=Trueexplicitly and otherwise relies on the (now True) default.It would be safer for the harness to set
enable_chunked_prefillexplicitly to the passed boolean for both sync and async variants.
🧹 Nitpick comments (2)
tensorrt_llm/commands/serve.py (2)
43-77: Signal handler behavior is fine, but comment is misleadingThe signal handler correctly terminates and, if needed, kills the child, then exits with
128 + signum. However, the comment “Using print for safety in signal handlers” no longer matches the implementation, which useslogger.info. Either update the comment or, if you want to be extra-defensive, switch these to
724-827: Child-process cleanup logic is robust but slightly duplicatedThe combination of:
_signal_handler_cleanup_childfor SIGINT/SIGTERM, and- the
finallyblock plus the final status checkgives you good coverage for normal shutdown and signal-driven terminate. There is some duplication in termination logic between the handler and the
finallyblock, but it’s benign and doesn’t introduce races thanks topoll()checks and bounded waits.If you want to simplify later, you could factor the “terminate/kill with timeouts + assert terminated” logic into a small helper that both the handler and
finallyuse.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
docs/source/features/long-sequence.md(1 hunks)docs/source/release-notes.md(1 hunks)tensorrt_llm/commands/serve.py(2 hunks)tensorrt_llm/llmapi/llm_args.py(1 hunks)tests/unittest/llmapi/test_llm.py(1 hunks)tests/unittest/llmapi/test_llm_pytorch.py(1 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 (e.g., usefrom package.subpackage import fooand thenfoo.SomeClass()instead offrom package.subpackage.foo import SomeClass)
Python filenames should use snake_case (e.g.,some_file.py)
Python class names should use PascalCase (e.g.,class SomeClass)
Python function and method names should use snake_case (e.g.,def my_awesome_function():)
Python local variable names should use snake_case, with prefixkfor variable names that start with a number (e.g.,k_99th_percentile = ...)
Python global variables should use upper snake_case with prefixG(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 Python class in the constructor
For Python interfaces that may be used outside a file, prefer docstrings over comments
Python comments should be reserved for code within a function, or interfaces that are local to a file
Use Google style docstrings for Python classes and functions, which can be parsed by Sphinx
Python attributes and variables can be documented inline with type and description (e.g.,self.x = 5followed by"""<type>: Description of 'x'""")
Avoid using reflection in Python when functionality can be easily achieved without reflection
When using try-except blocks in Python, limit the except clause to the smallest set of specific errors possible instead of catching all exceptions
When using try-except blocks in Python to handle multiple possible variable types (duck-typing), keep the body of the try as small as possible and use the else block to implement the logic
Files:
tests/unittest/llmapi/test_llm.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/commands/serve.pytests/unittest/llmapi/test_llm_pytorch.py
**/*.{cpp,h,cu,py}
📄 CodeRabbit inference engine (CODING_GUIDELINES.md)
All TensorRT-LLM Open Source Software code files should contain an NVIDIA copyright header that includes the current year at the top
Files:
tests/unittest/llmapi/test_llm.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/commands/serve.pytests/unittest/llmapi/test_llm_pytorch.py
🧠 Learnings (18)
📓 Common learnings
Learnt from: venkywonka
Repo: NVIDIA/TensorRT-LLM PR: 6029
File: .github/pull_request_template.md:45-53
Timestamp: 2025-08-27T17:50:13.264Z
Learning: For PR templates in TensorRT-LLM, avoid suggesting changes that would increase developer overhead, such as converting plain bullets to mandatory checkboxes. The team prefers guidance-style bullets that don't require explicit interaction to reduce friction in the PR creation process.
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6768
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:577-579
Timestamp: 2025-08-20T06:56:02.889Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, maxSequenceLength is now enforced as a non-optional argument in the BlockManager constructor, so concerns about std::nullopt defaulting to 0 are not applicable. When windowSize > maxSequenceLength, a warning should be added instead of handling optional parameter cases.
Learnt from: samuellees
Repo: NVIDIA/TensorRT-LLM PR: 6974
File: tensorrt_llm/serve/scripts/benchmark_dataset.py:558-566
Timestamp: 2025-08-18T08:42:02.640Z
Learning: In TensorRT-LLM's RandomDataset (tensorrt_llm/serve/scripts/benchmark_dataset.py), when using --random-token-ids option, sequence length accuracy is prioritized over semantic correctness for benchmarking purposes. The encode/decode operations should use skip_special_tokens=True and add_special_tokens=False to ensure exact target token lengths.
Learnt from: thorjohnsen
Repo: NVIDIA/TensorRT-LLM PR: 6910
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-14T21:04:50.248Z
Learning: In KV cache onboarding logic during prefill in cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, when calculating which blocks fall within the attention window, use getTokensPerBlock() to advance token indices rather than block->getUniqueTokens().size(), because the calculation needs to consider the post-prefill state where blocks will be filled to capacity, not their current token count.
Learnt from: MatthiasKohl
Repo: NVIDIA/TensorRT-LLM PR: 6904
File: cpp/tensorrt_llm/pybind/thop/bindings.cpp:55-57
Timestamp: 2025-08-14T15:38:01.771Z
Learning: In TensorRT-LLM Python bindings, tensor parameter collections like mla_tensor_params and spec_decoding_tensor_params are kept as required parameters without defaults to maintain API consistency, even when it might affect backward compatibility.
Learnt from: Funatiq
Repo: NVIDIA/TensorRT-LLM PR: 8587
File: tensorrt_llm/_torch/pyexecutor/llm_request.py:129-139
Timestamp: 2025-11-07T09:18:04.997Z
Learning: In `LogitsStorage.get()` method in `tensorrt_llm/_torch/pyexecutor/llm_request.py`, when `exclude_last=True`, there is an invariant that at least 2 chunks must have been appended to `_logits_indices`. The parameter is designed to drop the entire last chunk (not just the last token), which is expected behavior for the overlap scheduler that generates one extra token in a separate chunk.
📚 Learning: 2025-08-21T00:16:56.457Z
Learnt from: farshadghodsian
Repo: NVIDIA/TensorRT-LLM PR: 7101
File: docs/source/blogs/tech_blog/blog9_Deploying_GPT_OSS_on_TRTLLM.md:36-36
Timestamp: 2025-08-21T00:16:56.457Z
Learning: TensorRT-LLM container release tags in documentation should only reference published NGC container images. The README badge version may be ahead of the actual published container versions.
Applied to files:
docs/source/release-notes.md
📚 Learning: 2025-08-14T15:38:01.771Z
Learnt from: MatthiasKohl
Repo: NVIDIA/TensorRT-LLM PR: 6904
File: cpp/tensorrt_llm/pybind/thop/bindings.cpp:55-57
Timestamp: 2025-08-14T15:38:01.771Z
Learning: In TensorRT-LLM Python bindings, tensor parameter collections like mla_tensor_params and spec_decoding_tensor_params are kept as required parameters without defaults to maintain API consistency, even when it might affect backward compatibility.
Applied to files:
docs/source/release-notes.md
📚 Learning: 2025-09-23T15:12:38.312Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/thop/allreduceOp.cpp:352-446
Timestamp: 2025-09-23T15:12:38.312Z
Learning: In TensorRT-LLM NCCL device implementation, NCCL version 2.28+ requirements are handled at runtime in the nccl_device/config layer rather than with compile-time guards. This allows the allreduceOp to remain version-agnostic and delegates version compatibility validation to the appropriate lower-level components that can gracefully handle unsupported configurations.
Applied to files:
docs/source/release-notes.md
📚 Learning: 2025-08-27T14:23:55.566Z
Learnt from: ixlmar
Repo: NVIDIA/TensorRT-LLM PR: 7294
File: tensorrt_llm/_torch/modules/rms_norm.py:17-17
Timestamp: 2025-08-27T14:23:55.566Z
Learning: The TensorRT-LLM project requires Python 3.10+ as evidenced by the use of TypeAlias from typing module, match/case statements, and union type | syntax throughout the codebase, despite some documentation still mentioning Python 3.8+.
Applied to files:
docs/source/release-notes.md
📚 Learning: 2025-07-28T17:06:08.621Z
Learnt from: moraxu
Repo: NVIDIA/TensorRT-LLM PR: 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:
docs/source/release-notes.mdtests/unittest/llmapi/test_llm_pytorch.py
📚 Learning: 2025-08-01T15:14:45.673Z
Learnt from: yibinl-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 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:
docs/source/release-notes.md
📚 Learning: 2025-09-09T09:40:45.658Z
Learnt from: fredricz-20070104
Repo: NVIDIA/TensorRT-LLM PR: 7645
File: tests/integration/test_lists/qa/llm_function_core.txt:648-648
Timestamp: 2025-09-09T09:40:45.658Z
Learning: In TensorRT-LLM test lists, it's common and intentional for the same test to appear in multiple test list files when they serve different purposes (e.g., llm_function_core.txt for comprehensive core functionality testing and llm_function_core_sanity.txt for quick sanity checks). This duplication allows tests to be run in different testing contexts.
Applied to files:
docs/source/release-notes.mdtests/unittest/llmapi/test_llm_pytorch.py
📚 Learning: 2025-08-11T20:09:24.389Z
Learnt from: achartier
Repo: NVIDIA/TensorRT-LLM PR: 6763
File: tests/integration/defs/triton_server/conftest.py:16-22
Timestamp: 2025-08-11T20:09:24.389Z
Learning: In the TensorRT-LLM test infrastructure, the team prefers simple, direct solutions (like hard-coding directory traversal counts) over more complex but robust approaches when dealing with stable directory structures. They accept the maintenance cost of updating tests if the layout changes.
Applied to files:
docs/source/release-notes.md
📚 Learning: 2025-08-06T13:58:07.506Z
Learnt from: galagam
Repo: NVIDIA/TensorRT-LLM PR: 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:
docs/source/release-notes.md
📚 Learning: 2025-11-24T17:09:17.870Z
Learnt from: CR
Repo: NVIDIA/TensorRT-LLM PR: 0
File: CODING_GUIDELINES.md:0-0
Timestamp: 2025-11-24T17:09:17.870Z
Learning: Applies to **/*.py : The code developed for TensorRT-LLM should conform to Python 3.8+
Applied to files:
docs/source/release-notes.md
📚 Learning: 2025-08-26T09:37:10.463Z
Learnt from: jiaganc
Repo: NVIDIA/TensorRT-LLM PR: 7031
File: tensorrt_llm/bench/dataclasses/configuration.py:90-104
Timestamp: 2025-08-26T09:37:10.463Z
Learning: In TensorRT-LLM, the `get_pytorch_perf_config()` method returns `self.pytorch_config` which can contain default `cuda_graph_config` values, so `llm_args` may already have this config before the extra options processing.
Applied to files:
tensorrt_llm/commands/serve.py
📚 Learning: 2025-08-20T06:56:02.889Z
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6768
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:577-579
Timestamp: 2025-08-20T06:56:02.889Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, maxSequenceLength is now enforced as a non-optional argument in the BlockManager constructor, so concerns about std::nullopt defaulting to 0 are not applicable. When windowSize > maxSequenceLength, a warning should be added instead of handling optional parameter cases.
Applied to files:
docs/source/features/long-sequence.md
📚 Learning: 2025-08-14T21:04:50.248Z
Learnt from: thorjohnsen
Repo: NVIDIA/TensorRT-LLM PR: 6910
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-14T21:04:50.248Z
Learning: In KV cache onboarding logic during prefill in cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, when calculating which blocks fall within the attention window, use getTokensPerBlock() to advance token indices rather than block->getUniqueTokens().size(), because the calculation needs to consider the post-prefill state where blocks will be filled to capacity, not their current token count.
Applied to files:
docs/source/features/long-sequence.mdtests/unittest/llmapi/test_llm_pytorch.py
📚 Learning: 2025-08-15T06:46:54.897Z
Learnt from: eopXD
Repo: NVIDIA/TensorRT-LLM PR: 6767
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-15T06:46:54.897Z
Learning: In cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp addToken function, newly allocated blocks are unshared by design. The beam search path in addToken (when sequence.getNumTokens() > windowSize) is currently broken/non-functional with SWA, so the block allocation doesn't follow a shared-then-unshared pattern.
Applied to files:
docs/source/features/long-sequence.md
📚 Learning: 2025-09-23T14:58:05.372Z
Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/kernels/nccl_device/config.cu:42-49
Timestamp: 2025-09-23T14:58:05.372Z
Learning: In TensorRT-LLM NCCL device kernels (cpp/tensorrt_llm/kernels/nccl_device/), the token partitioning intentionally uses ceil-like distribution (same token_per_rank for all ranks) to ensure all ranks launch the same number of blocks. This is required for optimal NCCL device API barrier performance, even though it may launch extra blocks for non-existent tokens on later ranks. Runtime bounds checking in the kernel (blockID validation) handles the overshoot cases.
Applied to files:
docs/source/features/long-sequence.md
📚 Learning: 2025-08-18T08:42:02.640Z
Learnt from: samuellees
Repo: NVIDIA/TensorRT-LLM PR: 6974
File: tensorrt_llm/serve/scripts/benchmark_dataset.py:558-566
Timestamp: 2025-08-18T08:42:02.640Z
Learning: In TensorRT-LLM's RandomDataset (tensorrt_llm/serve/scripts/benchmark_dataset.py), when using --random-token-ids option, sequence length accuracy is prioritized over semantic correctness for benchmarking purposes. The encode/decode operations should use skip_special_tokens=True and add_special_tokens=False to ensure exact target token lengths.
Applied to files:
docs/source/features/long-sequence.md
📚 Learning: 2025-11-07T09:18:04.997Z
Learnt from: Funatiq
Repo: NVIDIA/TensorRT-LLM PR: 8587
File: tensorrt_llm/_torch/pyexecutor/llm_request.py:129-139
Timestamp: 2025-11-07T09:18:04.997Z
Learning: In `LogitsStorage.get()` method in `tensorrt_llm/_torch/pyexecutor/llm_request.py`, when `exclude_last=True`, there is an invariant that at least 2 chunks must have been appended to `_logits_indices`. The parameter is designed to drop the entire last chunk (not just the last token), which is expected behavior for the overlap scheduler that generates one extra token in a separate chunk.
Applied to files:
docs/source/features/long-sequence.md
🧬 Code graph analysis (1)
tensorrt_llm/commands/serve.py (1)
tensorrt_llm/builder.py (1)
default(45-50)
🪛 Ruff (0.14.5)
tests/unittest/llmapi/test_llm_pytorch.py
808-808: Standard pseudo-random generators are not suitable for cryptographic purposes
(S311)
818-818: Standard pseudo-random generators are not suitable for cryptographic purposes
(S311)
⏰ 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/commands/serve.py (1)
80-146: Chunked prefill default wiring in get_llm_args looks consistentThe new
enable_chunked_prefill: bool = Truesignature and propagation intollm_args["enable_chunked_prefill"]correctly align with the new default behavior and let callers override per-request.tests/unittest/llmapi/test_llm.py (1)
2417-2434: Explicitly disabling chunked prefill for error-path test is correctPassing
enable_chunked_prefill=Falseinto_test_llm_capture_request_error’s LLM construction is necessary now that the global default isTrue, so this test continues to exercise the pre-chunking error behavior.docs/source/release-notes.md (1)
11-13: Release-note entry accurately describes the chunked prefill default changeThe new “Breaking Changes” bullet for
enable_chunked_prefillmatches the implementation and tests: defaultTrue, automatic chunking of long sequences, warning/disable for unsupported models, and guidance on restoring prior behavior.tests/unittest/llmapi/test_llm_pytorch.py (3)
48-52: PyTorch tinyllama logits-processor test now covers both chunked and non‑chunked pathsParametrizing
test_tinyllama_logits_processoroverenable_chunked_prefilland forwarding it intotinyllama_logits_processor_test_harness(backend="pytorch", ...)is a good way to ensure PyTorch backend path is covered for both settings.
796-810: Explicitly disabling chunked prefill for max_num_tokens error path is correctIn
TestLlmError.test_max_num_token_check, addingenable_chunked_prefill=Falseto:llm = LLM( llama_model_path, kv_cache_config=global_kvcache_config, max_num_tokens=100, enable_chunked_prefill=False, )keeps the “prompt length exceeds max_num_tokens” ValueError behavior intact despite the new global default of
True.
811-823: New test correctly verifies default‑enabled chunked prefill no longer errors on long input
test_chunked_prefill_default_enabledbuilds an LLM without overridingenable_chunked_prefill, usesmax_num_tokens=100and a 101‑token input, and checks thatllm.generate(..., max_tokens=1)completes without error.That’s a concise, targeted assertion that the new default behavior is “chunk then generate” instead of throwing on long sequences.
docs/source/features/long-sequence.md (2)
14-18: Excellent documentation of the new default behavior.The documentation clearly explains the implications of enabling chunked context by default, addressing key user concerns:
- Behavior for long sequences (auto-chunk vs. error)
- Performance for short sequences (no impact)
- Compatibility for unsupported models (auto-disable with warning)
This will help users understand the change and its impact on their workloads.
19-29: Clear example and important configuration guidance.The updated example clearly shows the new default setting and how to explicitly disable it if needed. The note about
max_num_tokensbeing a multiple oftokens_per_block(default 64) is crucial configuration guidance for users.tensorrt_llm/llmapi/llm_args.py (1)
1799-1800: I'll help you verify the review comment about the breaking change inenable_chunked_prefilldefault value. Let me start by examining the codebase to confirm the changes and verify the safeguards mentioned.
<function_calls>
cat -n tensorrt_llm/llmapi/llm_args.py | sed -n '1795,1805p'
</function_calls>Let me check the file structure first:
<function_calls>
wc -l tensorrt_llm/llmapi/llm_args.py
</function_calls>Now let me search for the enable_chunked_prefill field definition:
<function_calls>
rg -n "enable_chunked_prefill" tensorrt_llm/llmapi/llm_args.py
</function_calls><function_calls>
rg -n -A2 -B2 "enable_chunked_prefill.*Field" tensorrt_llm/llmapi/llm_args.py
</function_calls>Now let me verify the automatic fallback logic and warning messages:
<function_calls>
rg -n -C5 "enable_chunked_prefill" --type=py | head -100
</function_calls>Let me search for disabling logic and warnings more specifically:
<function_calls>
rg -n "chunked_prefill" tensorrt_llm/ --type=py | grep -i "warn|disable|support" | head -20
</function_calls><function_calls>
rg -n -B3 -A3 "chunked.*prefill.*not.*support|not.*support.*chunked.*prefill" --type=py
</function_calls>Let me search for model architecture support checks:
<function_calls>
rg -n "chunked_prefill" tensorrt_llm/ --type=py -C3 | grep -A3 -B3 "model|architectur" | head -50
</function_calls>Let me find where the actual logic is implemented:
<function_calls>
fd -e py . tensorrt_llm/ | xargs rg -l "enable_chunked_prefill" | head -10
</function_calls><function_calls>
rg -n "enable_chunked_prefill" tensorrt_llm/ --type=py | head -30
</function_calls>Now let me check tests to verify they were updated:
<function_calls>
fd -e py . tests/ | xargs rg -l "enable_chunked_prefill|chunked_prefill" 2>/dev/null | head -10
</function_calls>Let me check the LLM class implementation:
<function_calls>
rg -n "enable_chunked_prefill" tensorrt_llm/llmapi/ --type=py -B2 -A2
</function_calls>Let me search for warning messages about chunked prefill:
<function_calls>
rg -n "warn" tensorrt_llm/llmapi/ --type=py | grep -i "chunk|prefill"
</function_calls><function_calls>
rg -n "logger|warning|warn" tensorrt_llm/llmapi/llm.py 2>/dev/null | head -20
</function_calls>Let me check the llm.py file for the initialization logic:
<function_calls>
wc -l tensorrt_llm/llmapi/llm.py
</function_calls><function_calls>
rg -n "enable_chunked_prefill" tensorrt_llm/llmapi/llm.py -B5 -A5
</function_calls>Let me search for the actual validation or fallback logic:
<function_calls>
rg -n "chunked" tensorrt_llm/llmapi/llm.py --type=py
</function_calls>Let me check if there's validation in the args class or initialization:
<function_calls>
rg -n "def.*validate|def.*check" tensorrt_llm/llmapi/llm_args.py | head -10
</function_calls>Let me look for model checks related to chunked prefill:
<function_calls>
rg -n "chunked" tensorrt_llm/ --type=py | grep -i "model|support|check|valid" | head -30
</function_calls>Let me check the actual llm_args.py file around the mentioned lines more carefully:
<function_calls>
sed -n '1790,1810p' tensorrt_llm/llmapi/llm_args.py
</function_calls>Perfect! Now let me search for where the enable_chunked_prefill is used and if there's fallback logic:
<function_calls>
rg -n "enable_chunked_prefill" tensorrt_llm/ --type=py -A3 | head -80
</function_calls>Let me check the runtime module for implementation:
<function_calls>
fd -type f -name "*.py" tensorrt_llm/ | xargs rg -l "chunked.*prefill" | head -15
</function_calls><function_calls>
rg -n "enable_chunked_prefill|chunked_prefill" tensorrt_llm/runtime/ --type=py
</function_calls>Let me check the builder or generation modules:
<function_calls>
rg -rn "chunked.*prefill" tensorrt_llm/ --type=py | grep -v "test|.pyc" | head -40
</function_calls>Let me search for Python files in tensorrt_llm that contain "chunked":
<function_calls>
find tensorrt_llm -name "*.py" -type f -exec grep -l "chunked" {} ; | head -20
</function_calls><function_calls>
rg -n "chunked" tensorrt_llm/llmapi/llm.py tensorrt_llm/llmapi/llm_utils.py -C3 2>/dev/null
</function_calls>Let me check the llm_utils file:
<function_calls>
rg -n "chunked" tensorrt_llm/llmapi/ --type=py -C3
</function_calls>Let me check if there's a model config or support check:
<function_calls>
rg -rn "supports.*chunked|chunked.*support" tensorrt_llm/ --type=py | head -20
</function_calls>Let me search more broadly for the change context:
<function_calls>
rg -n "enable_chunked_prefill" tensorrt_llm/ --type=py
</function_calls>Let me check if there's logic in the llm initialization that handles this:
<function_calls>
grep -rn "enable_chunked_prefill" tensorrt_llm/llmapi/llm.py
</function_calls>Let me check the full llm.py file for all references:
<function_calls>
grep -n "chunked" tensorrt_llm/llmapi/llm.py
</function_calls>Let me check the executor or generation logic:
<function_calls>
find tensorrt_llm -name "*.py" -exec grep -l "enable_chunked_prefill" {} ;
</function_calls>Let me check what files reference enable_chunked_prefill:
<function_calls>
grep -r "enable_chunked_prefill" tensorrt_llm/ --include="*.py" 2>/dev/null
</function_calls>Now let me check the tests:
<function_calls>
grep -r "enable_chunked_prefill" tests/ --include="*.py" 2>/dev/null | head -20
</function_calls>Let me check if there are tests that validate the default value:
<function_calls>
grep -r "chunked_prefill" tests/ --include="*.py" 2>/dev/null | head -20
</function_calls>Let me check the CLI handling:
<function_calls>
find tensorrt_llm -name "*.py" -exec grep -l "argparse|click|typer" {} ; | head -10
</function_calls><function_calls>
rg -rn "enable_chunked_prefill" tensorrt_llm/ --type=py
</function_calls>Let me look at the actual filed definition again to understand the context better:
<function_calls>
sed -n '1795,1850p' tensorrt_llm/llmapi/llm_args.py | cat -n
</function_calls>Now let me search for any model validation that checks this field:
<function_calls>
rg -n "model_config|model_arch" tensorrt_llm/llmapi/llm_args.py | head -20
</function_calls>Let me check the broader llm_args file for validation logic:
<function_calls>
grep -n "def.*model_validate|@validator|@field_validator" tensorrt_llm/llmapi/llm_args.py
</function_calls>Let me search for any runtime logic that handles unsupported models:
<function_calls>
rg -rn "not.*support|unsupport" tensorrt_llm/llmapi/ --type=py | grep -i "chunk|prefill" | head -20
</function_calls>Let me check the generation phase or executor:
<function_calls>
fd -name "executor*.py" tensorrt_llm/ | xargs rg -n "chunked" 2>/dev/null
</function_calls><function_calls>
fd -name "*.py" tensorrt_llm/ | xargs grep -l "executor" | head -10
</function_calls>Let me examine the generation_executor or batch_manager files:
<function_calls>
find tensorrt_llm -name "batch" -o -name "executor" | head -10
</function_calls>Let me check what happens when llm_args is passed to the LLM class:
<function_calls>
sed -n '1,50p' tensorrt_llm/llmapi/llm.py | cat -n
</function_calls><function_calls>
wc -l tensorrt_llm/llmapi/llm.py
</function_calls><function_calls>
grep -n "class LLM" tensorrt_llm/llmapi/llm.py
</function_calls><function_calls>
sed -n '50,150p' tensorrt_llm/llmapi/llm.py | cat -n
</function_calls>Let me search for initialization logic that might handle the chunked prefill:
<function_calls>
grep -n "def init" tensorrt_llm/llmapi/llm.py
</function_calls><function_calls>
sed -n '60,200p' tensorrt_llm/llmapi/llm.py | cat -n
</function_calls>Let me use a broader search for where the enable_chunked_prefill field is actually consumed:
<function_calls>
python3 << 'EOF'
import subprocess
import jsonSearch for enable_chunked_prefill usage
result = subprocess.run(['grep', '-r', 'enable_chunked_prefill', 'tensorrt_llm/', '--include=*.py'],
capture_output=True, text=True)
lines = result.stdout.strip().split('\n') if result.stdout.strip() else []
print(f"Found {len(lines)} references to enable_chunked_prefill:")
for line in lines:
print(line)
EOF
</function_calls>Let me check the test files to understand how this field is being tested:
<function_calls>
find tests -name "*.py" -exec grep -l "enable_chunked_prefill|chunked_prefill" {} ;
</function_calls><function_calls>
grep -rn "enable_chunked_prefill" tests/ --include="*.py" -B3 -A3
</function_calls>Let me verify the current state of the codebase by checking the git history or status:
<function_calls>
git log --oneline -20 -- tensorrt_llm/llmapi/llm_args.py 2>/dev/null | head -20
</function_calls>Let me check if there's any model configuration that disables this for certain architectures:
<function_calls>
find tensorrt_llm -name ".json" -o -name ".yaml" | xargs grep -l "chunked" 2>/dev/null
</function_calls>Let me verify the actual implementation by checking where llm_args are used:
<function_calls>
grep -rn "llm_args." tensorrt_llm/llmapi/llm.py | grep -i "chunk|enable" | head -20
</function_calls><function_calls>
grep -n "self.llm_args|self.args" tensorrt_llm/llmapi/llm.py | head -20
</function_calls>Let me check if there's documentation about this feature:
<function_calls>
find . -name "*.md" -exec grep -l "chunked.*prefill" {} ; 2>/dev/null
</function_calls>Let me look at the full context of llm_args class to understand the structure:
<function_calls>
grep -n "class LLMArgs|enable_chunked_prefill" tensorrt_llm/llmapi/llm_args.py | head -20
</function_calls><function_calls>
sed -n '1,100p' tensorrt_llm/llmapi/llm_args.py | grep -n "class|@"
</function_calls>Let me get a comprehensive view of the LLMArgs class:
<function_calls>
python3 << 'EOF'Find the class definition and understand the context
with open('tensorrt_llm/llmapi/llm_args.py', 'r') as f:
lines = f.readlines()Find enable_chunked_prefill
for i, line in enumerate(lines):
if 'enable_chunked_prefill' in line:
print(f"Found at line {i+1}: {line.strip()}")
# Show context
print("\nContext (20 lines before and after):")
start = max(0, i - 20)
end = min(len(lines), i + 20)
for j in range(start, end):
prefix = ">>> " if j == i else " "
print(f"{prefix}{j+1}: {lines[j]}", end='')
EOF
</function_calls><function_calls>
python3 << 'EOF'
with open('tensorrt_llm/llmapi/llm_args.py', 'r') as f:
content = f.read()Check if enable_chunked_prefill appears anywhere outside of its definition
import re
matches = list(re.finditer(r'enable_chunked_prefill', content))
print(f"Total occurrences of 'enable_chunked_prefill' in llm_args.py: {len(matches)}")Show each occurrence
lines = content.split('\n')
for match in matches:
# Find which line this is
line_num = content[:match.start()].count('\n') + 1
print(f"Line {line_num}: {lines[line_num-1]}")
EOF
</function_calls>Now let me check the entire codebase more thoroughly:
<function_calls>
find . -type f ( -name ".py" -o -name ".cpp" -o -name ".h" -o -name ".cu" ) -exec grep -l "enable_chunked_prefill|chunked_prefill" {} ; 2>/dev/null | grep -v pycache
</function_calls><function_calls>
git diff HEAD~1 tensorrt_llm/llmapi/llm_args.py 2>/dev/null | head -100
</function_calls>Let me check if there are any uncommitted changes or if this is in the current branch:
<function_calls>
git status 2>/dev/null | head -20
</function_calls><function_calls>
git log --all --oneline --grep="chunked" 2>/dev/null
</function_calls>Let me verify what the current actual value is:
<function_calls>
sed -n '1799p' tensorrt_llm/llmapi/llm_args.py
</function_calls>Perfect! Now let me search for the actual implementation of this feature in the codebase:
<function_calls>
python3 << 'EOF'
import subprocess
import osSearch for any references to chunked prefill logic
search_terms = [
'enable_chunked_prefill',
'chunked_prefill',
'chunked.*prefill',
'prefill.*chunk'
]all_results = {}
for term in search_terms:
result = subprocess.run(
['grep', '-r', term, '.', '--include=.py', '--include=.cpp', '--include=.h', '--include=.cu'],
capture_output=True,
text=True,
cwd='.'
)
if result.stdout:
files = set()
for line in result.stdout.strip().split('\n'):
if line:
file_path = line.split(':')[0]
if 'pycache' not in file_path and '.git' not in file_path:
files.add(file_path)
all_results[term] = sorted(files)print("Summary of findings:")
for term, files in all_results.items():
if files:
print(f"\n{term}: Found in {len(files)} file(s)")
for f in files[:5]: # Show first 5
print(f" - {f}")
if len(files) > 5:
print(f" ... and {len(files)-5} more")
EOF
</function_calls>Let me check the test files to understand what's actually being tested:
<function_calls>
find tests -name "*.py" | xargs grep -l "llm_args|LLMArgs" | head -5
</function_calls><function_calls>
grep -rn "enable_chunked_prefill" . --include="*.py" 2>/dev/null | grep -v pycache
</function_calls>Based on my thorough investigation, I've discovered critical gaps between the original review comment's claims and the actual codebase state. Let me verify one more thing about whether this feature is actually implemented:
<function_calls>
python3 << 'EOF'Check if enable_chunked_prefill is used in llm_args validation or post_init
import ast
import inspecttry:
with open('tensorrt_llm/llmapi/llm_args.py', 'r') as f:
tree = ast.parse(f.read())# Find the class with enable_chunked_prefill for node in ast.walk(tree): if isinstance(node, ast.ClassDef): # Check if any methods reference enable_chunked_prefill for method in node.body: if isinstance(method, ast.FunctionDef): method_source = ast.get_source_segment(open('tensorrt_llm/llmapi/llm_args.py').read(), method) if method_source and 'enable_chunked_prefill' in method_source: print(f"Found in method {method.name}: enable_chunked_prefill is used")except Exception as e:
print(f"Error: {e}")Simpler check - just search in methods
import re
with open('tensorrt_llm/llmapi/llm_args.py', 'r') as f:
content = f.read()Find method definitions that use enable_chunked_prefill
methods = re.finditer(r'def\s+(\w+)\s*([^)]):[^}]?(?=\n def |\nclass |\Z)', content, re.DOTALL)
found = False
for method in methods:
if 'enable_chunked_prefill' in method.group(0) and 'enable_chunked_prefill:' not in method.group(0):
found = True
print(f"Found usage in method: {method.group(1)}")if not found:
print("No usage of enable_chunked_prefill found in any method")
EOF
</function_calls>Based on my comprehensive verification of the codebase, I can now provide the rewritten review comment:
Breaking change: Chunked prefill now enabled by default, but implementation gaps need clarification.
The default value change from
FalsetoTrueis a breaking change at lines 1799-1800. However, verification found critical concerns about the claimed safeguards:
- ✓ Confirmed: Default is now
True- ✗ Not verified: Automatic fallback logic for unsupported models—no model support validation logic found in runtime code
- ✗ Not verified: Warning messages for disabled models—no warning implementations found
- ✗ Concern: The
enable_chunked_prefillfield is defined but not referenced anywhere in the runtime execution path (llm.py, executors, or generation code)Action required: Clarify whether this feature is fully implemented or a placeholder. If implemented elsewhere (C++/build system), document that context. If incomplete, either complete the implementation before enabling by default or revert to
Falseuntil safeguards are in place.
tensorrt_llm/commands/serve.py
Outdated
| @click.option("--enable_chunked_prefill", | ||
| is_flag=True, | ||
| default=False, | ||
| default=True, | ||
| help="Enable chunked prefill") |
There was a problem hiding this comment.
CLI currently cannot disable enable_chunked_prefill via flag
With:
@click.option("--enable_chunked_prefill",
is_flag=True,
default=True,
help="Enable chunked prefill")and the global default now True, users have no way to set enable_chunked_prefill=False from the CLI alone (the flag only ever yields True). This conflicts with docs that say “To restore the previous behavior, explicitly set enable_chunked_prefill=False”.
Consider switching to a boolean switch, for example:
@click.option(
"--enable_chunked_prefill/--disable_chunked_prefill",
default=True,
help="Enable or disable chunked prefill (default: enabled).",
)and plumb the resulting boolean into serve(... enable_chunked_prefill: bool, ...).
🤖 Prompt for AI Agents
In tensorrt_llm/commands/serve.py around lines 364 to 367, the CLI option uses a
single is_flag=True option which cannot be used to disable chunked prefill when
the default is True; replace that option with a boolean switch that exposes both
enable and disable (e.g. use the Click syntax
--enable_chunked_prefill/--disable_chunked_prefill) set default=True and update
the serve(...) signature to accept enable_chunked_prefill: bool and pass the
boolean through; ensure help text reflects the default and that any internal
code treats the value as a plain bool.
venkywonka
changed the title
|
/bot run |
|
PR_Github #26468 [ run ] triggered by Bot. Commit: |
|
PR_Github #26468 [ run ] completed with state |
venkywonka
force-pushed
the
venky/chunked-prefill-is-default
branch
from
4ca19aa to
3bd5d15
Compare
|
/bot run |
| - For models that don't support chunked context (e.g., RNN-based models like Mamba), it will be automatically disabled with a warning | ||
|
|
||
| To explicitly control chunked context: | ||
| ```bash | ||
| llm = LLM( | ||
| ... | ||
| enable_chunked_prefill=True, | ||
| enable_chunked_prefill=True, # Now the default | ||
| # Or to disable: enable_chunked_prefill=False, | ||
| ... | ||
| ) | ||
| ``` |
There was a problem hiding this comment.
can be removed now
|
PR_Github #26484 [ run ] triggered by Bot. Commit: |
|
/bot run |
|
PR_Github #26486 [ run ] triggered by Bot. Commit: |
|
PR_Github #37613 [ run ] completed with state
|
|
/bot run |
|
PR_Github #37941 [ run ] triggered by Bot. Commit: |
|
PR_Github #37941 [ run ] completed with state
|
|
/bot run |
|
PR_Github #38636 [ run ] triggered by Bot. Commit: |
|
PR_Github #38636 [ run ] completed with state
|
|
/bot run |
|
PR_Github #38802 [ run ] triggered by Bot. Commit: |
|
/bot kill |
|
/bot run --disable-fail-fast |
|
PR_Github #38804 [ run ] triggered by Bot. Commit: |
|
PR_Github #38805 [ kill ] triggered by Bot. Commit: |
|
PR_Github #38804 [ run ] completed with state |
|
PR_Github #38805 [ kill ] completed with state |
|
/bot run --disable-fail-fast |
- Change enable_chunked_prefill default to True in LlmArgs and serve CLI - Add --disable_chunked_prefill flag (mutually exclusive with --enable_chunked_prefill) - Keep --enable_chunked_prefill for backward compatibility - Add enable_chunked_prefill=False to get_model_defaults() for: - NemotronH (SSM/hybrid, chunked state unsupported) - Qwen3Next (SSM/hybrid, chunked state unsupported) - Gemma3 (bidirectional attention for image tokens) - Update long-sequence docs to reflect new default Signed-off-by: Venky <23023424+venkywonka@users.noreply.github.com>
- Switch --enable_chunked_prefill to --disable_chunked_prefill in quickstart_advanced.py, llm_sparse_attention.py, and lm_eval - Add HunyuanDense get_model_defaults to conditionally opt out of chunked prefill when mamba=True - Expand Gemma3 get_model_defaults docstring for consistency - Add unit tests for _resolve_chunked_prefill CLI resolver - Update DeepSeek V3, Nemotron Nano V2 VL, Nemotron Nano V3, EXAONE, and RocketKV docs to reflect new default Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Venky Ganesh <gvenkatarama@eos0019.eos.clusters.nvidia.com>
…nced tests quickstart_advanced.py now uses --disable_chunked_prefill since chunked prefill is enabled by default. Update all callers in test_e2e.py to drop the obsolete --enable_chunked_prefill flag. Fixes 7 failures in test_ptp_quickstart_advanced on RTXPro6000D-PyTorch-2. Signed-off-by: Venky Ganesh <23023424+venkywonka@users.noreply.github.com> Signed-off-by: Venky Ganesh <gvenkatarama@eos0048.eos.clusters.nvidia.com>
The C++ scheduler and Python runtime disagree on token counts for one-model speculative decoding (e.g., MTP). The two-model drafter path normalizes C++ draft_tokens to max_total_draft_tokens before scheduling, but the one-model MTP path skipped this normalization. The runtime's overlap extend path always materializes runtime_draft_len tokens, causing total_num_tokens > max_num_tokens when the scheduler under-budgets. Add draft token normalization for one-model MTP that mirrors the existing two-model drafter pattern. Also set enable_chunked_prefill= False on 3 DeepSeek-V3-Lite nvfp4/MTP tests that relied on the old default and add unit tests for the normalization logic. Signed-off-by: Venky Ganesh <gvenkatarama@eos0115.eos.clusters.nvidia.com>
venkywonka
force-pushed
the
venky/chunked-prefill-is-default
branch
from
44db50b to
6f42a6c
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #39649 [ run ] triggered by Bot. Commit: |
Tests that verify max_num_tokens rejection or construct mock AutoDeploy executors relied on the old enable_chunked_prefill=False default. With chunked prefill now enabled by default, these tests need explicit opt-out: - test_create_ad_executor: mock engine lacks tokens_per_block int, causing TypeError in ContextChunkingConfig construction - test_max_num_token_check: prompt-too-long rejection is skipped when chunked prefill handles the input via chunking instead - _test_llm_capture_request_error: same max_num_tokens validation skip Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #39684 [ run ] triggered by Bot. Commit: |
|
PR_Github #39684 [ run ] completed with state
|
6 participants
Description
--disable_chunked_prefilloptions to serve, quickstart advanced CLIenable_chunked_prefill=FalsedefaultOne-model MTP scheduler/runtime token accounting fix
Enabling chunked prefill by default exposed a pre-existing bug in the
chunked prefill + one-model MTP + overlap schedulerpath. Two CI testsfailed on DGX_B200 with:
Root cause: The two-model drafter path normalizes C++
draft_tokenstomax_total_draft_tokensfor all generation requests before the C++ schedulerruns (so
getNumDraftTokens()matches the runtime's fixedruntime_draft_len).The one-model MTP path (
self.drafter is None) skipped this normalization entirely.The C++ scheduler would then see stale or zero draft tokens for generation requests,
under-budgeting them, while the Python runtime always materializes the full
1 + runtime_draft_lentokens in the overlap extend path - causing the assertion.Fix: Added the same draft token normalization for one-model MTP before scheduling
in
py_executor.py, mirroring the existing two-model drafter pattern.Unit test fixes for new default
Several unit tests relied on the old
enable_chunked_prefill=Falsedefault:test_create_ad_executor: Mock engine lacks realtokens_per_block,causing
TypeErrorinContextChunkingConfigconstruction when chunked prefill is enabled.Fixed by setting
enable_chunked_prefill=Falsein the test'sLlmArgs.test_max_num_token_check: ExpectsRequestErrorfor prompts exceedingmax_num_tokens,but this validation is skipped when chunked prefill handles the input via chunking.
Fixed by explicitly disabling chunked prefill in the test.
_test_llm_capture_request_error: Samemax_num_tokensvalidation skip as above.Fixed by adding
enable_chunked_prefill=Falsefor the pytorch backend path.Integration test fixes
enable_chunked_prefill=Falseon 3 DeepSeek-V3-Lite nvfp4/MTP accuracy teststhat relied on the old default and should not inherit chunked prefill behavior.
Summary
Release Notes
Behavior Changes
Bug Fixes
Documentation
Test Coverage
tests/unittest/_torch/executor/test_py_executor.py: 3 new unit tests validating draft token normalization logic for one-model MTP (normalization applied, skipped with drafter, skipped with spec decode off)tests/unittest/_torch/executor/test_overlap_scheduler.py: 4 existing overlap scheduler consistency tests passtests/integration/defs/accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4: B200 accuracy tests with explicitenable_chunked_prefill=Falsetests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.py: Fixed mock executor test for chunked prefill defaulttests/unittest/llmapi/test_llm_pytorch.py: Fixedtest_max_num_token_checkfor chunked prefill defaulttests/unittest/llmapi/test_llm.py: Fixed_test_llm_capture_request_errorfor chunked prefill defaultmax_num_tokens=8192, MTP nextn=2, overlap+chunked prefillKey files changed
tensorrt_llm/_torch/pyexecutor/py_executor.pytests/integration/defs/accuracy/test_llm_api_pytorch.pyenable_chunked_prefill=Falseon 3 nvfp4/MTP teststests/unittest/_torch/executor/test_py_executor.pytests/unittest/auto_deploy/singlegpu/shim/test_create_ad_executor.pytests/unittest/llmapi/test_llm_pytorch.pytests/unittest/llmapi/test_llm.pyPR Checklist