Repository navigation
Conversation
…ulation (NVIDIA#2482) Signed-off-by: musi22 <shawrashmi7@gmail.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughWalkthroughThe change updates KV-dimension calculation to use configured head dimensions, adds configuration fallbacks, supports descriptors without language-model configuration methods, and adds tests for KV-cache metrics. ChangesKV dimension resolution
Priority: ➖ Normal Estimated code review effort: 3 (Moderate) | ~20 minutes Change: Bug fix · Severity of issue fixed: Medium 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
… 100% coverage Signed-off-by: musi22 <shawrashmi7@gmail.com>
musi22
changed the title
soksof
added a commit
to soksof/Model-Optimizer
that referenced
this pull request
Oct 25, 2026
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 <sokratis.sofianopoulos@gmail.com>
Contributor
Author
|
Closing in favor of #3009 which ports this fix to current main with co-authorship and additional regression tests. |
soksof
added a commit
to soksof/Model-Optimizer
that referenced
this pull request
Oct 25, 2026
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: GitHub <noreply@github.com>
soksof
added a commit
to soksof/Model-Optimizer
that referenced
this pull request
Oct 25, 2026
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>
kevalmorabia97
added a commit
that referenced
this pull request
Oct 26, 2026
#3009) ### What does this PR do? Type of change: Bug fix `calculate_kv_dim` computes the head size as `n_embd // n_head` and does not read `head_dim`, so `calculate_attention_memory` misprices the KV cache of every attention subblock in two cases: - models that set `head_dim` independently of `hidden_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) - every candidate with a pruned hidden size, for all models: subblock statistics are computed per candidate `hidden_size` with the parent's `num_attention_heads`, while pruning `hidden_size` leaves `head_dim` unchanged. For example, Qwen3-8B with `hidden_size` pruned to 3072 gives 3072 // 32 = 96 instead of 128, i.e. a KV cache 25% too low The head dimension is now read from the language-model config (`descriptor.get_language_model_config(model_config).head_dim`), with `n_embd // n_head` as the fallback when it is not set. `calculate_kv_dim` gets an optional `head_dim` argument, so existing callers are unchanged. This is the fix of #2484 ported to the current `main` (that PR was based on `feature/puzzletron_v2`, and `main` has since renamed `num_kv_heads` to `num_key_value_heads`), with @musi22 as co-author, as suggested in #2482. On `main` the KV-cache memory is the only consumer of `calculate_kv_dim`. On `feature/puzzletron_v2`, the attention FLOP estimate in `calculate_additive_metrics` still uses `qk_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 on `main`. ### Usage No API change. `calculate_kv_dim(num_key_value_heads, n_head, n_embd, head_dim=None)` keeps its previous behaviour when `head_dim` is not given. ### Testing - New `tests/unit/torch/puzzletron/test_kv_dim.py`, 5 cases, 0.2 s in total: - `calculate_kv_dim` with an explicit `head_dim` smaller (Mistral-Small) and larger (gpt-oss) than `n_embd // n_head`, and the fallback without `head_dim` - `calculate_attention_memory` with `head_dim=128` at `n_embd` 5120 and 3840: the KV-cache memory is the same in both cases, i.e. pruning `hidden_size` does not change it. `calculate_subblock_params` is patched to 0, since the weight count is not part of the KV cache and would need a real model - All 5 cases fail without the fix (the KV cache comes out as 0.156 and 0.117 MiB instead of 0.125) and pass with it - `pytest tests/unit/torch/puzzletron`: 49 passed, on CPU (torch 2.14, transformers 5.17, Python 3.13) - `ruff format` and `ruff check` (0.15.20) on the changed files: no new findings (the 4 `PLR0917` findings in `calc_subblock_params_and_memory.py` are present on `main` as well) - Not run: GPU tests, and an end-to-end Puzzletron search ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ (the new argument is optional; subblock statistics computed with earlier versions change for the affected models and should be recomputed) - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ (logic from #2484 by @musi22, credited as co-author; no new dependency) - Did you write any new necessary tests?: ✅ - Did you update [[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ (Bug Fixes, 0.49.0) - Did you get Claude approval on this PR?: ❌ (to be run after opening) ### Additional Information Fixes #2482. Supersedes #2484 and #2485. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **Bug Fixes** * KV-cache memory estimates now use the configured attention head size when available, improving accuracy for models with explicit head-size settings and when embedding dimensions are pruned. If no head size is configured, estimates continue to use the inferred value. * Recompute existing subblock statistics to reflect the corrected memory estimates and ensure pruning results use up-to-date values. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Sokratis Sofianopoulos <s_sofian@athenarc.gr> Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: musi22 <shawrashmi7@gmail.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Fixes #2482
Currently,
calculate_kv_dimcomputeshead_size = n_embd // n_head. For architectures wherehead_dimis decoupled from embedding division (such as Mistral Small 24B wherehead_dim = 128vs5120 // 32 = 160, or gpt-oss 20B wherehead_dim = 64vs2880 // 64 = 45), this leads to inaccurate KV-cache memory and additive metric calculations.Changes
modelopt/torch/puzzletron/utils/misc.py:head_dim: int | None = Noneparameter tocalculate_kv_dim.head_dimis provided, it is used directly; otherwise, it falls back ton_embd // n_headto maintain complete backward compatibility.modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py:calculate_additive_metricsandcalculate_attention_memory, resolvedeffective_head_dimfollowing precedence:subblock_config.qk_head_dim>model_config.head_dim(via_language_model_attr) > fallbackn_embd // n_head._language_model_attrhelper to safely handle configs whendescriptorlacksget_language_model_config.tests/unit/torch/puzzletron/test_kv_dim_calculation.pywith 6 unit tests validating Mistral Small 24B, gpt-oss 20B, default fallback, and precedence behavior.Testing
Summary by CodeRabbit
Bug Fixes
Tests