[None][perf] Use +64 batch sizes for padding-enabled CUDA graphs#12895
Conversation
📝 WalkthroughWalkthroughModified the batch size generation logic in Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes 🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/llmapi/llm_args.py (1)
169-187:⚠️ Potential issue | 🟠 MajorClamp padding-mode batch sizes to
max_batch_sizebefore returning.On Line 181-Line 183, filtering/sorting now runs only in the
elsebranch.
Withenable_padding=True,batch_sizescan exceedmax_batch_size(e.g.,max_batch_size=64still keeps values up to 128), which can lead to unnecessary graph capture and inconsistent config behavior.🔧 Proposed fix
if enable_padding: # Start with [1, 2, 4, 8, 16, 24, ..., 128] (multiples of 8) batch_sizes = [1, 2, 4] + [i * 8 for i in range(1, 17)] # Sliding 64: extend by increments of 64 up to max_batch_size while batch_sizes[-1] + 64 <= max_batch_size: batch_sizes.append(batch_sizes[-1] + 64) else: batch_sizes = list(range(1, 32)) + [32, 64, 128] # Add powers of 2 up to max_batch_size batch_sizes += [ 2**i for i in range(8, math.ceil(math.log(max_batch_size, 2))) ] - # Filter and sort batch sizes - batch_sizes = sorted( - [size for size in batch_sizes if size <= max_batch_size]) + # Filter and sort batch sizes for both branches + batch_sizes = sorted(size for size in batch_sizes if size <= max_batch_size) # Add max_batch_size if not already included - if max_batch_size != batch_sizes[-1]: + if not batch_sizes or max_batch_size != batch_sizes[-1]: batch_sizes.append(max_batch_size)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/llmapi/llm_args.py` around lines 169 - 187, The enable_padding branch for building batch_sizes can produce values > max_batch_size; update the block that builds batch_sizes when enable_padding is True (the batch_sizes list creation) to clamp/filter values to <= max_batch_size and sort/unique them just like the else branch does, ensuring batch_sizes = sorted([s for s in batch_sizes if s <= max_batch_size]) before the final check that appends max_batch_size if missing; keep the subsequent append-of-max_batch_size logic unchanged.
🧹 Nitpick comments (1)
tensorrt_llm/llmapi/llm_args.py (1)
169-189: Add regression coverage for padding schedule edge cases.Please add tests for
enable_padding=Truewith at least:max_batch_size=64,129, and320to assert all values are<= max_batch_size, sorted, and includemax_batch_size. This will prevent regressions from branch-local filtering changes.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/llmapi/llm_args.py` around lines 169 - 189, Add regression tests covering the padding schedule branch where enable_padding=True by calling the batch-size generator (the function that takes enable_padding and max_batch_size and returns batch_sizes) with max_batch_size values 64, 129, and 320; for each case assert that every returned value in batch_sizes is <= max_batch_size, that batch_sizes is sorted (non-decreasing), and that max_batch_size is present in the returned list; target the branch that builds batch_sizes using the enable_padding path (referencing the enable_padding parameter, max_batch_size variable, and the returned batch_sizes) so future changes to the padding schedule are validated.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 169-187: The enable_padding branch for building batch_sizes can
produce values > max_batch_size; update the block that builds batch_sizes when
enable_padding is True (the batch_sizes list creation) to clamp/filter values to
<= max_batch_size and sort/unique them just like the else branch does, ensuring
batch_sizes = sorted([s for s in batch_sizes if s <= max_batch_size]) before the
final check that appends max_batch_size if missing; keep the subsequent
append-of-max_batch_size logic unchanged.
---
Nitpick comments:
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 169-189: Add regression tests covering the padding schedule branch
where enable_padding=True by calling the batch-size generator (the function that
takes enable_padding and max_batch_size and returns batch_sizes) with
max_batch_size values 64, 129, and 320; for each case assert that every returned
value in batch_sizes is <= max_batch_size, that batch_sizes is sorted
(non-decreasing), and that max_batch_size is present in the returned list;
target the branch that builds batch_sizes using the enable_padding path
(referencing the enable_padding parameter, max_batch_size variable, and the
returned batch_sizes) so future changes to the padding schedule are validated.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 5920502e-9591-4b1c-983a-406885dd4e30
📒 Files selected for processing (1)
tensorrt_llm/llmapi/llm_args.py
|
/bot run |
|
PR_Github #42571 [ run ] triggered by Bot. Commit: |
|
PR_Github #42571 [ run ] completed with state
|
|
/bot run |
|
PR_Github #42656 [ run ] triggered by Bot. Commit: |
|
@yijingl-nvidia can you put a summary on the performance validation results on the PR description? Thanks. |
|
PR_Github #42656 [ run ] completed with state
|
|
/bot run |
|
PR_Github #42731 [ run ] triggered by Bot. Commit: |
|
PR_Github #42731 [ run ] completed with state
|
Thanks for the suggestion. Edited the summary to include a brief summary of our performance findings. |
|
/bot run |
|
PR_Github #43059 [ run ] triggered by Bot. Commit: |
|
PR_Github #43059 [ run ] completed with state
|
a557b8b to
f77a4d2
Compare
|
/bot run |
|
PR_Github #44494 [ run ] triggered by Bot. Commit: |
|
PR_Github #44494 [ run ] completed with state
|
f77a4d2 to
e4fa2e3
Compare
|
/bot run |
|
PR_Github #44781 [ run ] triggered by Bot. Commit: |
|
PR_Github #44781 [ run ] completed with state
|
|
/bot run |
|
PR_Github #44822 [ run ] triggered by Bot. Commit: |
When enable_padding=True, replace the sparse powers-of-2 schedule (256, 512, 1024, 2048) with uniform +64 increments after the initial [1,2,4,8,...,128] base, giving denser coverage (192, 256, 320, …, max_batch_size) and reducing padding waste at intermediate batch sizes. Signed-off-by: Yijing Li <[email protected]>
Move filter/sort outside the if/else so sizes exceeding max_batch_size are dropped in the enable_padding=True branch as well. Add guard for empty list before the max_batch_size append. Add regression tests for edge cases: max_batch_size=64, 129, 320. Signed-off-by: Yijing Li <[email protected]>
e4fa2e3 to
6be963b
Compare
|
/bot run |
|
PR_Github #44826 [ run ] triggered by Bot. Commit: |
|
PR_Github #44826 [ run ] completed with state
|
|
/bot run |
|
PR_Github #44902 [ run ] triggered by Bot. Commit: |
|
PR_Github #44902 [ run ] completed with state
|
|
/bot run |
|
PR_Github #44990 [ run ] triggered by Bot. Commit: |
|
PR_Github #44990 [ run ] completed with state
|
|
/bot run |
|
PR_Github #45101 [ run ] triggered by Bot. Commit: |
|
PR_Github #45101 [ run ] completed with state |
|
Corresponding blog post at #13393 |
…DIA#12895) Signed-off-by: Yijing Li <[email protected]>
Description
When enable_padding=True, replace the sparse powers-of-2 schedule (256, 512, 1024, 2048) with uniform +64 increments after the initial [1,2,4,8,...,128] base, giving denser coverage (192, 256, 320, …, max_batch_size) and reducing padding waste at intermediate batch sizes.
To test how denser batch sizes improve output throughput, we’ve tested on four major models, DeepSeek R1 and V3.2, Kimi K2 and Qwen 3, and in both aggregate and disagg modes.
A denser set of batch sizes comes with the cost of more GPU memory. Each CUDA graph metadata takes about 10 MiB for DeepSeek R1. For +64 batch sizes, it increases from 220MiB to 490 MiB per GPU. For +8 batch sizes, it goes further to 2.6 GiB.
From the experiments, we found that denser batch sizes led to smaller drop on output throughput. Both +64 and +8 batch sizes improve the metrics. +64 batch sizes improve by up to 1.3x for agg and 1.5x for disagg for some large concurrencies.
+8 batch sizes improves a bit more than +64 on most models. However. We did notice a regression of +8 batch sizes on DeepSeek V3.2. We suspect DeepSeek V3.2 uses DSA which accelerates the attention computation but makes the CUDA graph more complex to store, thereby costing more than 2x GPU memory to save. This squeezes out available space for KV cache, reducing server capacity. Each GPU in the +8 experiment has 5 GiB fewer KV cache memory than the +64.
On the downside, +64 batch sizes increase server startup time. TRTLLM repeats the graph capture twice. 1st as a dry run to measure how much GPU memory is used at runtime to estimate how much KV cache we can safely allocate. 2nd as the actual startup with the full KV cache allocated. Tested on GB200, DeepSeek R1 TP=4, IFB, max batch size 2048, it increases CUDA graph capture time by 1.4x and total LLM init time (measured as the time of
get_llm()) by 1.17x. Future remedy can be to reduce the number of CUDA graph captures in the dry run phase and use an estimate based on the memory usage of the remaining CUDA graphs captured.We concluded by suggesting landing the change to use +64 batch sizes for the next release. We are also writing a blog post with the data and guidance on how to set CUDA graph batch sizes for improving performance. Future work would include testing more options like +16, and designing a way of automatically setting the batch sizes to reach the best balance of memory cost and performance gains.
Test Coverage
CI
PR 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)
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.Summary by CodeRabbit