Skip to content

[https://nvbugs/6438658][fix] Fix KV cache estimation capacity#16545

Draft
jiaganc wants to merge 1 commit into
NVIDIA:mainfrom
jiaganc:codex/fix-kv-estimation-cuda-graph-warmup
Draft

[https://nvbugs/6438658][fix] Fix KV cache estimation capacity#16545
jiaganc wants to merge 1 commit into
NVIDIA:mainfrom
jiaganc:codex/fix-kv-estimation-cuda-graph-warmup

Conversation

@jiaganc

@jiaganc jiaganc commented Jul 17, 2026

Copy link
Copy Markdown
Collaborator

@coderabbitai summary

Description

Fix KV cache capacity estimation so the resulting cache can satisfy runtime warmup and resume requirements.

Without these constraints, pool-ratio estimation can allocate a pool layout that cannot run CUDA graph warmup. Ignoring resume utilization during distributed quota synchronization can also remove the reserved headroom and overestimate usable token capacity.

  • Preserve max_util_for_resume headroom when synchronizing KV cache capacity across distributed ranks.
  • Create the generic V2 cache configuration, standard layers, and estimation constraints in KVCacheManagerV2.
  • Cover the full CUDA graph generation warmup batch and the per-iteration context token budget when estimating pool ratios.
  • Let specialized cache managers update the base configuration; DeepSeek V4 supplies its custom layers and typical step while preserving the generic constraints.

Test Coverage

  • Added and updated unit coverage for distributed resume headroom, generic estimation constraints, DeepSeek V4 typical-step configuration, and max_util_for_resume validation.
  • pre-commit passed for all changed files.
  • python3 -m py_compile passed for the changed Python files.
  • Focused pytest execution was not available locally because the environment lacks compatible compiled tensorrt_llm.bindings.

PR Checklist

  • 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, comment /bot help.

Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
@jiaganc
jiaganc force-pushed the codex/fix-kv-estimation-cuda-graph-warmup branch from 084262f to 9917bcf Compare July 17, 2026 14:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant