[Puzzletron] Fix KV-cache dimension calculation for models with custom head_dim (#2482)fix(puzzletron): respect explicit head_dim in KV-cache dimension calc…

[musi22]

Description

Fixes #2482

Currently, calculate_kv_dim computes head_size = n_embd // n_head. For architectures where head_dim is decoupled from embedding division (such as Mistral Small 24B where head_dim = 128 vs 5120 // 32 = 160, or gpt-oss 20B where head_dim = 64 vs 2880 // 64 = 45), this leads to inaccurate KV-cache memory and additive metric calculations.

Changes

1. modelopt/torch/puzzletron/utils/misc.py:
   • Added optional head_dim: int | None = None parameter to calculate_kv_dim.
   • If head_dim is provided, it is used directly; otherwise, it falls back to n_embd // n_head to maintain complete backward compatibility.
2. modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py:
   • In both calculate_additive_metrics and calculate_attention_memory, resolved effective_head_dim following precedence: subblock_config.qk_head_dim > model_config.head_dim (via _language_model_attr) > fallback n_embd // n_head.
   • Hardened _language_model_attr helper to safely handle configs when descriptor lacks get_language_model_config.
3. Tests:
   • Added tests/unit/torch/puzzletron/test_kv_dim_calculation.py with 6 unit tests validating Mistral Small 24B, gpt-oss 20B, default fallback, and precedence behavior.

Testing

• Ran unit test suite:

  pytest tests/unit/torch/puzzletron/test_kv_dim_calculation.py tests/unit/torch/puzzletron/test_additive_subblock_metrics.py tests/unit/torch/puzzletron/test_attention_runtime_memory.py

  Result: 20 passed.
• Ruff checks and format passed.

Summary by CodeRabbit

• Bug Fixes

  • Improved attention KV-dimension calculations by honoring configured query/key head dimensions when available.
  • Added reliable fallbacks for models that provide general head-dimension or embedding configuration.
  • Improved KV-cache memory and bytes-per-token estimates across supported model configurations.

• Tests

  • Added coverage for explicit and inferred head dimensions, configuration precedence, and attention-memory calculations.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: abaa9560-4d50-415d-94b8-180e62f7aab8

📥 Commits

Reviewing files that changed from the base of the PR and between c924530 and c36dac7.

📒 Files selected for processing (2)

• modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py
• tests/unit/torch/puzzletron/test_kv_dim_calculation.py

🚧 Files skipped from review as they are similar to previous changes (2)

• tests/unit/torch/puzzletron/test_kv_dim_calculation.py
• modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Walkthrough

Walkthrough

The 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.

Changes

KV dimension resolution

Layer / File(s)	Summary
Head dimension resolution and metric integration modelopt/torch/puzzletron/utils/misc.py, modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py	calculate_kv_dim accepts optional dimensions and prioritizes head_dim. Attention metrics use qk_head_dim, model head_dim, or inferred dimensions. Configuration lookup supports descriptors without get_language_model_config.
KV metric coverage tests/unit/torch/puzzletron/test_kv_dim_calculation.py	Tests cover explicit and inferred dimensions, KV-cache memory, additive metrics, and dimension precedence.

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)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main change: correcting KV-cache dimension calculations for models with an explicit custom head_dim. It is overly long and contains duplicated title text, but it rem…
Linked Issues check	✅ Passed	The PR meets the coding requirements in #2482. calculate_kv_dim accepts head_dim and uses n_embd // n_head when no explicit dimension exists. calculate_attention_memory and `calculate_additive…
Out of Scope Changes check	✅ Passed	The changes remain within #2482. The _language_model_attr safeguard supports model configuration lookup for the required head_dim resolution. The compatibility update and the added tests directly …
Docstring Coverage	✅ Passed	Docstring coverage is 90.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 10 functions across 3 files.
Security Anti-Patterns	✅ Passed	The pull request adds only KV-dimension logic, documentation, and tests. The authoritative diff adds no prohibited torch.load(..., weights_only=False), numpy.load(..., allow_pickle=True), hardcode…

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

---

_{Comment @coderabbitai help to get the list of available commands.}