[None][perf] Preserve default V2 KV cache pool sizing - #16783
Conversation
|
/bot run |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe PR makes KV-cache warmup configuration conditional on an explicitly provided average sequence length, preserves DeepSeek-V4 fallback behavior, and adds Gemma 4 NVFP4 B200 performance-sanity coverage for pre-merge CI. ChangesKV cache typical sequence handling
B200 performance-sanity coverage
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant KvCacheConfig
participant KVCacheManagerV2
participant WarmupConstraints
KvCacheConfig->>KVCacheManagerV2: provide avg_seq_len or unset value
KVCacheManagerV2->>KVCacheManagerV2: resolve typical sequence length
KVCacheManagerV2->>WarmupConstraints: build warmup step and constraints when configured
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@jenkins/L0_Test.groovy`:
- Around line 4896-4898: Update the perfMode predicate near the stage
configuration logic to recognize stage keys containing "-PerfSanity-" in
addition to "-Perf-". Ensure both DGX_B200-PyTorch-PerfSanity-1 and
DGX_B200-PyTorch-PerfSanity-Post-Merge-1 receive performance mode and run with
--perf.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 591c9744-f352-4e46-a3fd-399cb1acaf56
📒 Files selected for processing (8)
jenkins/L0_Test.groovytensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytests/integration/defs/perf/_model_paths.pytests/integration/test_lists/test-db/l0_b200_perf_sanity.ymltests/scripts/perf-sanity/aggregated/gemma4_26b_a4b_nvfp4_blackwell.yamltests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.pytests/unittest/_torch/executor/test_kv_cache_manager_v2.py
|
PR_Github #61222 [ run ] triggered by Bot. Commit: |
|
/bot run |
|
PR_Github #61224 [ run ] triggered by Bot. Commit: |
|
PR_Github/16783-0b93b6f #61222 was force-killed by a newer pipeline run. |
|
PR_Github #61224 [ run ] completed with state
|
chzblych
left a comment
There was a problem hiding this comment.
Approved for the CI changes.
0b93b6f to
291449c
Compare
|
/bot run |
|
PR_Github #61341 [ run ] completed with state
|
SimengLiu-nv
left a comment
There was a problem hiding this comment.
The kvcache changes LGTM.
|
/bot run |
|
PR_Github #61424 [ run ] triggered by Bot. Commit: |
|
PR_Github #61424 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61500 [ run ] triggered by Bot. Commit: |
|
PR_Github #61500 [ run ] completed with state
|
| linux_distribution_name: ubuntu* | ||
| cpu: x86_64 | ||
| terms: | ||
| stage: pre_merge |
There was a problem hiding this comment.
Do we need to put this in pre-merge? What could be the cost? I would suggest only adding post merge perf test.
There was a problem hiding this comment.
Given how often we've been negatively impacting perf for these models (10x reduction happened twice in 2 weeks), I would prefer to keep it in pre-merge. Removed the post-merge entry instead. We can always revisit if this ends up being a major bottleneck. It's a "small" model so I don't expect it to make a noticeable dent in CI.
There was a problem hiding this comment.
Why the perf regression happen so frequent? Which PR causes it? Does these perf regressions only affect Gemm4?
There was a problem hiding this comment.
The regression happens so frequently because:
- there are no perf regression tests gating this in pre-merge.
- all KV cache manager v2 related changes seem to assume they only affect DSv4 / specific models.
The 2 PRs that regressed perf that I know of:
- [https://nvbugs/6438658][fix] Fix KV cache estimation capacity #16545 (what this PR is trying to fix)
- [TRTLLM-14138][perf] Make FlashInfer decode plans sync-free with host-built page tables #16073 is the prior regression for FlashInfer (which Gemma4 has to use).
291449c to
05d2aad
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@jenkins/L0_Test.groovy`:
- Around line 4960-4961: Add a Jenkins entry for
DGX_B200-PyTorch-PerfSanity-Post-Merge-1 alongside
DGX_B200-PyTorch-PerfSanity-1, reusing the same auto:dgx-b200-flex platform,
l0_b200_perf_sanity test list, and single-GPU settings while selecting the
post_merge stage.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7a405d65-79b3-4941-a78d-78e95534a9ee
📒 Files selected for processing (4)
jenkins/L0_Test.groovytests/integration/defs/perf/_model_paths.pytests/integration/test_lists/test-db/l0_b200_perf_sanity.ymltests/scripts/perf-sanity/aggregated/gemma4_26b_a4b_nvfp4_blackwell.yaml
🚧 Files skipped from review as they are similar to previous changes (3)
- tests/integration/test_lists/test-db/l0_b200_perf_sanity.yml
- tests/integration/defs/perf/_model_paths.py
- tests/scripts/perf-sanity/aggregated/gemma4_26b_a4b_nvfp4_blackwell.yaml
05d2aad to
7d2d538
Compare
|
/bot run |
|
PR_Github #61682 [ run ] triggered by Bot. Commit: |
|
PR_Github #61682 [ run ] completed with state
|
| return ( | ||
| kv_cache_config.avg_seq_len | ||
| if kv_cache_config.avg_seq_len is not None | ||
| else self.max_seq_len |
There was a problem hiding this comment.
Why only dsv4 need fallback?
There was a problem hiding this comment.
It's not so much that DSv4 needs a fallback, as it is that this PR moved DSv4 specific logic into the base class.
So what this commit does is introduce the _get_typical_seq_len helper that does the expected thing for DSv4, but also restores the behavior prior to that PR for the base class (which Gemma 4 uses).
* Why? Using `max_seq_len` as the default typical sequence length over- allocates Gemma4's full-attention pool and exhausts its SWA pool. This prevents CUDA graph padding from allocating dummy requests and forces eager decode, leading to hugely degraded performance. * What? Only derive generic warmup constraints from an explicitly configured `avg_seq_len`, preserving allocator-derived pool sizing by default. Keep DeepSeek-V4's max-length fallback in its specialized cache manager. This restores the behavior for Gemma4 prior to `697738c1f6` introducing the new sizing logic for DeepSeek-V4. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
7d2d538 to
2cb0720
Compare
|
/bot run |
|
PR_Github #61942 [ run ] triggered by Bot. Commit: |
|
PR_Github #61942 [ run ] completed with state
|
|
/bot run |
|
PR_Github #62007 [ run ] triggered by Bot. Commit: |
|
PR_Github #62007 [ run ] completed with state
|
Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
Dev Engineer Review
KVCacheManagerV2._build_base_configso warmuptypical_stepandconstraintsare derived only whenkv_cache_config.avg_seq_lenis explicitly provided (via_get_typical_seq_len). Theavg_seq_len > max_seq_lenvalidation is now conditional onavg_seq_lenbeing set, preserving the prior “allocator fallback” behavior whenavg_seq_lenis unset.DeepseekV4CacheManager._get_typical_seq_len(...)to keep DeepSeek-V4 pool-sizing semantics: returnkv_cache_config.avg_seq_lenwhen set; otherwise fall back toself.max_seq_len.tests/integration/defs/perf/_model_paths.py:gemma_4_26b_a4b_nvfp4 -> gemma/nvidia-Gemma-4-26B-A4B-NVFP4.tests/scripts/perf-sanity/aggregated/gemma4_26b_a4b_nvfp4_blackwell.yaml(gemma4_26b_a4b_nvfp4_tp1_1k1k, CUDA graph padding enabled, KV cachefp8, single GPU).tests/integration/test_lists/test-db/l0_b200_perf_sanity.yml: adds a pre_merge run forperf/test_perf_sanity.py::test_e2e[aggr_upload-gemma4_26b_a4b_nvfp4_blackwell-gemma4_26b_a4b_nvfp4_tp1_1k1k] TIMEOUT (60)under*b200*,ubuntu*,x86_64,backend: pytorch,orchestrator: mpi,system_gpu_count: 1.jenkins/L0_Test.groovy: addedDGX_B200-PyTorch-PerfSanity-1(single-GPU,auto:dgx-b200-flex, runningl0_b200_perf_sanity).tests/unittest/_torch/executor/test_kv_cache_manager_v2.pywarmup/constraint assertions to reflect the new gating behavior:test_default_uses_allocator_fallback(defaults keeptypical_step is Noneandconstraints == []).test_avg_seq_len_builds_warmup_constraintsto validate warmuptypical_stepandconstraintswhenavg_seq_lenis set.test_extra_tokens_are_in_context_capacityto drive config viaKvCacheConfig(avg_seq_len=264).tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py::test_typical_seq_len_preserves_deepseek_v4_fallbackto verify DeepSeek-V4 fallback tomax_seq_lenwhenavg_seq_lenis unset.QA Engineer Review
Test files changed (integration/perf-regression)
tests/integration/test_lists/test-db/l0_b200_perf_sanity.ymlstage: pre_mergeonly:perf/test_perf_sanity.py::test_e2e[aggr_upload-gemma4_26b_a4b_nvfp4_blackwell-gemma4_26b_a4b_nvfp4_tp1_1k1k] TIMEOUT (60)gpu: '*b200*',ubuntu*,x86_64,backend: pytorch,orchestrator: mpi,system_gpu_count: 1.tests/integration/defs/perf/_model_paths.pygemma_4_26b_a4b_nvfp4.tests/scripts/perf-sanity/aggregated/gemma4_26b_a4b_nvfp4_blackwell.yamlgemma4_26b_a4b_nvfp4_tp1_1k1kand client profilecon64_iter3_1k1k.Verdict: needs follow-up (no CBTS coverage availability provided in the provided context).
Test code changed (unit)
tests/unittest/_torch/executor/test_kv_cache_manager_v2.pytest_default_uses_allocator_fallbacktest_avg_seq_len_builds_warmup_constraintstest_extra_tokens_are_in_context_capacity(drives viaKvCacheConfig(avg_seq_len=264))test_builds_warmup_constraintstests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.pytest_typical_seq_len_preserves_deepseek_v4_fallbackCoverage linkage: unit tests are not represented in
tests/integration/test_lists/test-db/; the Gemma4 end-to-end regression coverage is exercised vial0_b200_perf_sanity.yml.Description
Using
max_seq_lenas the default typical sequence length over-allocates Gemma4's full-attention pool and exhausts its SWA pool.
This prevents CUDA graph padding from allocating dummy requests
and forces eager decode, leading to hugely degraded performance.
Only derive generic warmup constraints from an explicitly configured
avg_seq_len, preserving allocator-derived pool sizing by default.Keep DeepSeek-V4's max-length fallback in its specialized cache manager.
This restores the behavior for Gemma4 prior to
697738c1f6introducingthe new sizing logic for DeepSeek-V4, and adds an E2E perf regression test
for it.
Test Coverage
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)
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.