Repository navigation
[Puzzletron] Use the model's head_dim for the KV-cache memory estimate - #3009
Conversation
|
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:
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Head-dimension contract and callers modelopt/torch/puzzletron/utils/misc.py, modelopt/torch/puzzletron/pruning/kv_heads_pruning_mixin.py, modelopt/torch/puzzletron/pruning/pruning_utils.py, modelopt/torch/puzzletron/tools/bypassed_training/child_init.py |
calculate_kv_dim accepts an optional head_dim. The new get_attention_head_dim helper returns the configured value or derives it from hidden_size // num_attention_heads. Pruning and child initialization use the shared helper, and the previous local helpers are removed. |
Attention memory and validation modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py, tests/unit/torch/puzzletron/test_kv_dim.py, tests/gpu/torch/puzzletron/test_puzzletron.py, CHANGELOG.rst |
calculate_attention_memory passes the configured head dimension to calculate_kv_dim. Unit tests cover explicit and inferred dimensions and memory estimates. GPU test expectations change for several models. The changelog notes that earlier subblock statistics must be recomputed. |
Priority:
Estimated code review effort: 2 (Simple) | ~10 minutes
Change: Bug fix · Severity of issue fixed: Medium
Possibly related issues
- [Puzzletron] KV-cache memory computes head size as n_embd // n_head instead of reading head_dim #2482: The issue describes incorrect KV-cache sizing when
head_dimdiffers fromhidden_size // num_attention_heads, which this change addresses.
Possibly related PRs
- NVIDIA/Model-Optimizer#2484: Changes the KV-dimension calculation and attention-memory path for the same head-dimension mismatch.
- NVIDIA/Model-Optimizer#2485: Changes the helper contract and uses the configured head dimension for attention KV-cache memory.
Suggested reviewers: danielkorzekwa, separius
Merge Risk: 
39c43
-
No confirmed merge-blocking issue was identified in the reviewed change.
Pre-merge checks: 
-
Check name Status Explanation Resolution Docstring Coverage Docstring coverage is 58.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 12 functions across 7 files. Write docstrings for the functions missing them to satisfy the coverage threshold. -
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title names the primary change: Puzzletron's KV-cache memory estimate now uses the model's head_dim. The raw summary and objectives confirm this. The changeset adds an optional head_dim parameter … Linked Issues check [ #2482] Requirement: the KV-cache head size must come from the model's head_dim. When head_dim is not set, it must fall back to n_embd // n_head. At head 39c4323, calculate_attention_memory in calc_su…Out of Scope Changes check The whole-PR diff (merge base 3788881 to head 39c4323) changes CHANGELOG.rst (one bug-fix entry), misc.py (optional head_dim argument and get_attention_head_dim helper), calc_subblock_params_and_mem… Security Anti-Patterns PASS. The reviewed modelopt changes add only head-dimension resolution and KV-dimension calculation. The PR adds no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded t… Full details: Title check
-
Explanation
The title names the primary change: Puzzletron's KV-cache memory estimate now uses the model's head_dim. The raw summary and objectives confirm this. The changeset adds an optional head_dim parameter to calculate_kv_dim, adds get_attention_head_dim, and updates callers and tests to use the model's head dimension.
Full details: Linked Issues check
-
Explanation
[
#2482] Requirement: the KV-cache head size must come from the model's head_dim. When head_dim is not set, it must fall back to n_embd // n_head. At head 39c4323, calculate_attention_memory in calc_subblock_params_and_memory.py calls get_attention_head_dim(descriptor.get_language_model_config(model_config)). misc.py get_attention_head_dim returns lm_config.head_dim when set, else hidden_size // num_attention_heads of the parent config. The result is passed to calculate_kv_dim as head_dim. calculate_kv_dim keeps kv_dim = 2 * num_key_value_heads * head_size. It uses head_dim when given, else n_embd // n_head. This covers the Mistral-Small-24B (128 vs 160) and gpt-oss-20b (64 vs 45) cases in the issue. Pruned candidate hidden_size values use the parent config for the head dimension, so pruning hidden_size does not change the KV cache. The file tests/unit/torch/puzzletron/test_kv_dim.py at head covers: explicit head_dim smaller and larger than n_embd // n_head, the no-head_dim fallback, and KV-cache memory at n_embd 5120 and 3840 for both an explicit head_dim and the inferred parent head dimension. This closes the gap noted in the previous review. The issue's optional qk_head_dim path exists only on feature/puzzletron_v2 and is not on main, so it is not required here. Limits: the tests were not re-executed in this review. The direct calculate_kv_dim call in the issue's reproduction, with three arguments, still returns the old value. The fix is supplied through the new optional head_dim argument, which the caller now passes.
Full details: Out of Scope Changes check
-
Explanation
The whole-PR diff (merge base 3788881 to head 39c4323) changes CHANGELOG.rst (one bug-fix entry), misc.py (optional head_dim argument and get_attention_head_dim helper), calc_subblock_params_and_memory.py (head-dimension lookup in calculate_attention_memory), tests/unit/torch/puzzletron/test_kv_dim.py (regression tests), tests/gpu/torch/puzzletron/test_puzzletron.py (updated expected teacher-memory values), and the pruning and child-init files. The pruning and child-init files are: kv_heads_pruning_mixin.py, pruning_utils.py, and child_init.py. Per the change summary, these replace the separate _lm_head_dim and _get_head_dim helpers with the shared get_attention_head_dim. The commit message describes this as the same head-dimension resolution. This is a supporting consolidation for the
#2482fix, not an unrelated change. Changed expected values in test_puzzletron.py follow from the corrected KV-cache estimate. Limits: the kv_heads_pruning_mixin.py and child_init.py changes were checked only through the change summary and the previous review, not through a direct diff at head. No unrelated functional change was found in the inspected code.
Full details: Security Anti-Patterns
-
Explanation
PASS. The reviewed modelopt changes add only head-dimension resolution and KV-dimension calculation. The PR adds no torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcoded trust_remote_code=True, eval(), exec(), or # nosec usage. The diff changes no pyproject.toml or requirements.txt dependency files. Existing torch.load(p) and descriptor-controlled trust_remote_code usage are unchanged and do not match a newly introduced prohibited pattern.
-
-
Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
|
/claude review |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
Thanks @soksof for porting the fix from #2484 to main and adding the additional regression tests, and thank you @kevalmorabia97! The changes look great to me. |
|
Just a small heads-up @soksof: the DCO check is currently flagging an author vs sign-off email mismatch on the commits (\s_sofian@athenarc.gr\ vs \sokratis.sofianopoulos@gmail.com). A quick rebase with matching author/sign-off emails will get the check green for merge. |
50cb893 to
5857ba9
Compare
calculate_kv_dim computed the head size as n_embd // n_head, so calculate_attention_memory mispriced the KV cache of every attention subblock whenever the model sets head_dim independently (Mistral-Small, GPT-OSS, Gemma-3), and for every candidate with a pruned hidden size, since subblock statistics are computed per candidate hidden_size with the parent's num_attention_heads. Pruning hidden_size leaves head_dim unchanged, so the head dimension is now read from the language-model config, with n_embd // n_head as the fallback when it is not set. Fixes NVIDIA#2482. Based on the fix in NVIDIA#2484, ported to the current main. Co-authored-by: musi22 <shawrashmi7@gmail.com> Signed-off-by: Sokratis Sofianopoulos <s_sofian@athenarc.gr>
5857ba9 to
1b9ec16
Compare
Thanks for the heads-up! I squashed the two commits into one (1b9ec16) with matching author and sign-off emails, and the DCO check is green now. I mixed up my institutional and my gmail! |
|
/claude review |
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (bedrock-claude-opus-5-5) — DM the bot to share feedback.
Approving: this is a small, correct fix. The KV-cache estimate now reads head_dim from the language-model config and falls back to n_embd // n_head when it isn't set, and new tests cover both paths.
No action needed.
No action needed:
calculate_attention_memoryalready receivesmodel_configanddescriptor. Readinghead_dimfrom the parent config gives the right value for candidates with a prunedhidden_size, because pruning doesn't changehead_dim.- Existing callers keep their behaviour, since the new
head_dimargument is optional. test_kv_dim.pycovers an explicithead_dimsmaller and larger thann_embd // n_head, the fallback, and the case where pruningn_embddoesn't change the KV cache. Stubbingcalculate_subblock_paramsthere is reasonable.- The new test file uses the standard NVIDIA SPDX header. The ported logic comes from in-repo PR #2484, and its author is credited as co-author.
- Nit:
test_kv_dim.pyhas no trailing newline. - MLA-style configs aren't modelled correctly either way, e.g. DeepSeek's
head_dimis the rope dimension. That was already true before this change and is out of scope here.
|
@soksof can you fix all CI failures so we can merge this fix? |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review — 1 IMPORTANT, 1 SUGGESTION
Scope: full review (the trigger comment was a bare /claude review). All 4 changed files (+76/-3) reviewed: modelopt/torch/puzzletron/utils/misc.py, modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py, tests/unit/torch/puzzletron/test_kv_dim.py, CHANGELOG.rst. I also read the callers (calc_subblock_stats.py), the existing head-size resolvers (pruning_utils.py, child_init.py), and kv_heads_pruning_mixin.py.
Findings
| Severity | Count |
|---|---|
| CRITICAL | 0 |
| IMPORTANT | 1 |
| SUGGESTION | 1 |
[IMPORTANT Compatibility] calc_subblock_params_and_memory.py:360-363 — inline comment. The core fix is right for models that publish head_dim, but the fallback path still divides the candidate hidden size: when head_dim is absent or None, calculate_kv_dim computes n_embd // n_head, and at this call site n_embd is the per-candidate model_hidden_size from the model_hidden_sizes loop in calc_subblock_stats.py while n_head is always the teacher's num_attention_heads. So the second failure mode the PR describes — "every candidate with a pruned hidden size, for all models" — persists for configs that don't set head_dim, and the new comment ("Pruning hidden_size leaves head_dim unchanged") promises an invariant the fallback breaks.
It also diverges from the two resolvers the repo already has for this exact quantity — pruning_utils._lm_head_dim (pruning_utils.py:70) and child_init._get_head_dim (child_init.py:1228) — both of which fall back to the parent lm_config.hidden_size // lm_config.num_attention_heads. Since _lm_head_dim is what kv_heads_pruning_mixin.py:83 uses to size the weights actually written into the child checkpoint, the memory model and the produced checkpoint can resolve different head sizes for the same candidate. Reusing _lm_head_dim(model_config, descriptor) here fixes both and shrinks the diff (also matches the "extend an existing mechanism rather than add a parallel one" rule in CLAUDE.md).
[SUGGESTION] tests/unit/torch/puzzletron/test_kv_dim.py:65 has no trailing newline. ruff-format runs as a pre-commit hook here and will rewrite the file, so the pre-commit CI job should fail on it as committed — worth fixing in the same push even though it changes nothing functionally.
What checked out
head_dimis read from the teacher config, which is loaded once incalc_subblock_stats.pyand never mutated for the candidate hidden size, so for models that sethead_dimthe KV-cache estimate is now correctly independent of the prunedn_embd. The Mistral-Small / gpt-oss / Gemma-3 numbers in the description match.- The new argument is optional and the existing single in-repo caller is the one being updated, so no behavioural change for external callers of
calculate_kv_dim. getattr(..., "head_dim", None)handles both the missing-attribute and the explicit-Nonecase (the latter matters —calc_runtime_stats.py:64deliberately passeshead_dim=NonetoLlamaConfig).- The tests fail without the fix and the monkeypatch of
calculate_subblock_paramsis sound (the real call is positional, solambda *args: 0matches). Keeping the weight count out of a KV-cache test is the right call. CHANGELOG.rstentry is in the right sub-section, user-facing, and correctly tells users to recompute existing subblock statistics.- MLA-style configs (DeepSeek, where
head_dimis the rope dimension) are still not modelled, but that predates this change and is out of scope.
Risk
Low. The change is small, well-tested, and strictly improves accuracy over main; the one IMPORTANT finding is an incomplete-fix gap on the fallback branch rather than a new defect, and the suggested fix is a one-liner.
🤖 Generated with Claude Code
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #3009 +/- ##
==========================================
+ Coverage 77.51% 79.14% +1.62%
==========================================
Files 661 661
Lines 74937 75165 +228
==========================================
+ Hits 58086 59487 +1401
+ Misses 16851 15678 -1173
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
When head_dim is not set, the KV-cache estimate now divides the parent's hidden_size by its head count instead of the candidate's, by reusing lm_head_dim (made public), the helper that sizes the attention weights of the pruned checkpoint. Also normalize the line endings of the new test. Signed-off-by: Sokratis Sofianopoulos <s_sofian@athenarc.gr>
|
@kevalmorabia97 Fixed: the test file had Windows line endings, which failed ruff-format. The review finding is addressed in the same push |
|
/claude review |
|
/ok to test 76a45b5 |
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Findings: CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 2 (both posted inline, both non-blocking)
Scope. Full review — 6 changed files (+93/-11). Reviewed all 5 source/test files in modelopt/ and tests/; CHANGELOG.rst checked against the repo's changelog policy. Nothing deliberately skipped.
Correctness: the fix traces cleanly.
calculate_kv_dimgains an optional trailinghead_dim, so the old three-positional-argument call stays valid — backward compatible, and the grep confirmscalc_subblock_params_and_memory.py:363is the only production caller onmain, matching the PR description.- The key question for this fix is whether the
model_configreachingcalculate_attention_memoryis really the parent config, as the new comment claims — if it were a per-candidate config, the no-head_dimfallback would silently reintroduce the bug it is meant to fix. It is:calc_subblock_stats.py:338-345passes the same teachermodel_configobject on every iteration of themodel_hidden_sizes× dtype × batch-size product while onlyn_embd=model_hidden_sizevaries, andcalculate_subblock_paramsdeep-copies before mutating, so the object is never aliased or mutated under the loop.lm_head_dim(model_config, descriptor)therefore returns the invariant head dim in both the explicit-head_dimand the fallback case. head_dimis correctly treated as a non-pruning dimension:AttentionConfig(block_config.py:201-221) carries onlynum_key_value_heads/llama4/mamba, andblock_config_to_layer_overridescannot introduce a head size — so taking it from the parent while takingnum_key_value_headsfrom the candidate is the right split. The weight-count term is unaffected, sincecalculate_subblock_paramsbuilds a meta model from the real config and picks uphead_dimon its own.- No circular-import risk from the new
subblock_stats→pruningedge (nothing underpruning/oranymodel/importssubblock_stats), andpuzzletron/__init__.pyalready imports both subpackages, so there is no added import cost either. - Mode/state, export, and plugin-laziness categories are not touched by this diff — no mode registration, no
modelopt_stateschema change, no new public API beyond thelm_head_dimrename of a previously private symbol.
Tests. The new test_kv_dim.py earns its place: test_attention_kv_cache_independent_of_pruned_hidden_size parametrizes over both n_embd 5120 and 3840 and asserts the KV-cache term is identical, which is exactly the invariant the bug violated — it fails before the fix and cannot pass by restating the implementation. The module-level monkeypatch.setattr(calc, "calculate_subblock_params", ...) resolves correctly against the module global, and zeroing it is justified (the weight count is not part of the KV cache and would need a real model). CPU-only and ~0.2 s, so it belongs in tests/unit. There is no pre-existing test file for calc_subblock_params_and_memory to extend, so a new file is the right choice under the "be stingy" rule.
Most impactful observation (inline, non-blocking): the runtime objective still has the same root cause via a different mechanism. calc_runtime_stats.create_benchmark_model builds its synthetic Llama with head_dim=None and hidden_size=n_embd (calc_runtime_stats.py:58-70), which forces precisely the n_embd // n_head head size this PR removes from the memory path. When runtime_stats_enabled=True, the benchmarked attention has the wrong head dim for the same two populations the PR fixes — models with an explicit head_dim, and every pruned-hidden_size candidate — skewing the latency numbers that feed the MIP in the same direction the memory numbers were skewed. Correctly out of scope for this PR (different file, different objective, needs a vLLM run to validate), but worth a note on #2482 or a follow-up issue so it is not lost now that the memory side is resolved. The second suggestion is a de-duplication opportunity: child_init.py:1228 _get_head_dim already implements the same rule, and the new public name lm_head_dim reads as the lm_head projection rather than the attention head dim.
Risk: low. Tightly scoped bug fix, backward-compatible signature, no API or state-schema change, and the behavior change is correctly called out as invalidating previously computed subblock statistics in both the PR body and the CHANGELOG.rst entry — which is appropriately placed under 0.49.0 Bug Fixes and written for external users. The CHANGELOG and tests/unit/torch/export/test_export_diffusers.py hunks visible in a two-dot diff against the origin/main tip are unrelated drift from the shallow checkout, not part of this PR's six files.
🤖 Generated with Claude Code
…ctations Move the head-dim resolver to utils.misc.get_attention_head_dim and use it for the KV-cache estimate, KV-heads pruning and child init, replacing lm_head_dim and child_init._get_head_dim. The toy test models keep the parent's head_dim (64 for gpt-oss, 128 for the others except Qwen2.5), so the corrected KV-cache estimate raises the expected teacher memory in test_puzzletron. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
|
/ok to test 39c4323 |
What does this PR do?
Type of change: Bug fix
calculate_kv_dimcomputes the head size asn_embd // n_headand does not readhead_dim, socalculate_attention_memorymisprices the KV cache of every attention subblock in two cases:head_dimindependently ofhidden_size / num_attention_heads, e.g. Mistral-Small-24B (128 against 160, 25% too high), gpt-oss-20b (64 against 45, 30% too low) and Gemma-3-12B (256 against 240)hidden_sizewith the parent'snum_attention_heads, while pruninghidden_sizeleaveshead_dimunchanged. For example, Qwen3-8B withhidden_sizepruned to 3072 gives 3072 // 32 = 96 instead of 128, i.e. a KV cache 25% too lowThe head dimension is now read from the language-model config (
descriptor.get_language_model_config(model_config).head_dim), withn_embd // n_headas the fallback when it is not set.calculate_kv_dimgets an optionalhead_dimargument, so existing callers are unchanged.This is the fix of #2484 ported to the current
main(that PR was based onfeature/puzzletron_v2, andmainhas since renamednum_kv_headstonum_key_value_heads), with @musi22 as co-author, as suggested in #2482. Onmainthe KV-cache memory is the only consumer ofcalculate_kv_dim. Onfeature/puzzletron_v2, the attention FLOP estimate incalculate_additive_metricsstill usesqk_head_dim or (n_embd // n_head)and should use the same resolution; this is not included here, given that the function does not exist onmain.Usage
No API change.
calculate_kv_dim(num_key_value_heads, n_head, n_embd, head_dim=None)keeps its previous behaviour whenhead_dimis not given.Testing
tests/unit/torch/puzzletron/test_kv_dim.py, 5 cases, 0.2 s in total:calculate_kv_dimwith an explicithead_dimsmaller (Mistral-Small) and larger (gpt-oss) thann_embd // n_head, and the fallback withouthead_dimcalculate_attention_memorywithhead_dim=128atn_embd5120 and 3840: the KV-cache memory is the same in both cases, i.e. pruninghidden_sizedoes not change it.calculate_subblock_paramsis patched to 0, since the weight count is not part of the KV cache and would need a real modelpytest tests/unit/torch/puzzletron: 49 passed, on CPU (torch 2.14, transformers 5.17, Python 3.13)ruff formatandruff check(0.15.20) on the changed files: no new findings (the 4PLR0917findings incalc_subblock_params_and_memory.pyare present onmainas well)Before your PR is "Ready for review"
CONTRIBUTING.md: ✅ (logic from [Puzzletron] Fix KV-cache dimension calculation for models with custom head_dim (#2482) #2484 by @musi22, credited as co-author; no new dependency)Additional Information
Fixes #2482. Supersedes #2484 and #2485.
Summary by CodeRabbit