Skip to content

fix: use explicit head dim for Puzzletron KV cache #2485

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
devxMani wants to merge 1 commit into from
+38 −6
Closed

fix: use explicit head dim for Puzzletron KV cache #2485

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -357,7 +357,15 @@ def calculate_attention_memory(
):
seq_len = min(seq_len, attention_chunk_size)

kv_dim = calculate_kv_dim(attention_config.num_key_value_heads, n_head, n_embd)
lm_config = descriptor.get_language_model_config(model_config)

@claude claude Bot Sep 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[CRITICAL Algorithm] Reading head_dim from the teacher config overrides the swept candidate hidden size, regressing KV-cache estimates for every hidden-size pruning candidate.

n_embd here is not the teacher's hidden size — it is a swept candidate value. Tracing upstream:

  • calc_subblock_stats.py:319-321 builds model_hidden_sizes = model_hidden_sizes + [lm_config.hidden_size] ("add a teacher model hidden size"), i.e. a list of candidate hidden sizes plus the teacher's.
  • calc_subblock_stats.py:326 iterates product(batch_sizes, data_types, model_hidden_sizes).
  • calc_subblock_stats.py:345-346 passes n_embd=model_hidden_size (varies per iteration) and n_head=lm_config.num_attention_heads (fixed).
  • model_config stays the fixed teacher config throughout, so lm_config.head_dim is a constant across the whole sweep.

The old head_size = n_embd // n_head correctly tracked each candidate hidden size. The new code pins head_size to the teacher's head_dim for all candidates.

Concrete failure: teacher with hidden_size=5120, num_attention_heads=40, head_dim=128, num_key_value_heads=8. For the candidate model_hidden_size=2560, the correct head size is 2560 // 40 = 64 → kv_dim = 2*8*64 = 1024. This code yields head_dim=128 → kv_dim = 2048, a 2× overestimate of kv_cache_size, which flows straight into kv_cache_memory_mib and memory_mib. Those are the memory costs the Puzzletron search optimizes against, so every hidden-size-reduced candidate is now penalized as twice as KV-expensive as it is, skewing candidate selection. Note this hits plain Llama teachers too: modern LlamaConfig always populates head_dim, so getattr(lm_config, "head_dim", None) returns a value and the old derivation is never reached.

The fix for #2482 needs to distinguish the two cases: head_dim is an independent model property only when it is not simply hidden_size // num_attention_heads (Qwen3, gpt-oss, etc.). When it is derived, it must keep tracking n_embd:

lm_config = descriptor.get_language_model_config(model_config)
configured_head_dim = getattr(lm_config, "head_dim", None)
# Only an explicitly decoupled head_dim is a fixed model property; a derived one must
# track the candidate hidden size being swept in n_embd.
if configured_head_dim == lm_config.hidden_size // lm_config.num_attention_heads:
    configured_head_dim = None

kv_dim = calculate_kv_dim(
    attention_config.num_key_value_heads, n_head, n_embd, head_dim=configured_head_dim
)

A regression test that pins this would be valuable: assert kv_cache_memory_mib still halves when n_embd halves for a teacher whose head_dim == hidden_size // num_heads. The current unit tests only exercise calculate_kv_dim in isolation, so they cannot catch this.

@soksof soksof Sep 23, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On the CRITICAL finding in the automated review, I think the premise doesn't hold for how Puzzletron prunes width. EmbeddingPruningSpec.update_config_object rewrites only hidden_size; head_dim, num_attention_heads and num_key_value_heads are left as in the teacher. A hidden-size-pruned child therefore builds k_proj/v_proj with unchanged output dimensions, and its KV cache per token is 2 * num_key_value_heads * head_dim, independent of the candidate hidden size. Using the teacher's head_dim for every candidate in the sweep is the correct value, not an overestimate. The previous n_embd // n_head understated KV for width candidates too, since it made KV shrink with hidden size. The suggested workaround (fall back to n_embd // n_head when head_dim == hidden_size // num_attention_heads) would bring that back for Llama-style teachers. A useful regression test would be the opposite of the one proposed: kv_cache_memory_mib stays constant when n_embd is reduced with heads unchanged.

Separately, qk_head_dim does not exist on AttentionConfig on main but does on feature/puzzletron_v2, where calculate_additive_metrics already reads it, so a port of this fix to v2 could honour it there.

head_dim = getattr(attention_config, "qk_head_dim", None) or getattr(lm_config, "head_dim", None)

@claude claude Bot Sep 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT Correctness] attention_config.qk_head_dim does not exist, so this half of the expression is permanently dead.

AttentionConfig (modelopt/torch/puzzletron/block_config.py:201) is a @dataclass(frozen=True, kw_only=True) whose full field set — including the SubblockConfig base at line 126 — is:

no_op, replace_with_linear, sparsify, weights_precision,   # SubblockConfig
num_key_value_heads, llama4, mamba                          # AttentionConfig

There is no qk_head_dim anywhere in modelopt/torch/puzzletron/ (a repo-wide grep finds only this new line), and because the dataclass is frozen it cannot be injected at runtime either. So getattr(attention_config, "qk_head_dim", None) always evaluates to None, and the resolved head_dim always comes from lm_config.head_dim.

Why it matters: it reads as if per-subblock head dims are honored during the Puzzletron search, which is the case a reviewer would most want to trust here. Anyone later adding a real per-subblock head dim will assume this line already wires it up. It is also silently unreachable — no error surfaces, the value is just wrong-by-omission if a subblock ever does override the head dim.

Suggested fix: drop the getattr on attention_config and read the head dim from the language-model config only (see my top-level comment for the version that reuses the existing helper). If honoring a per-subblock override is actually intended, add the field to AttentionConfig and plumb it through to_dict/from_dict so it round-trips.

Separately, prefer is not None over or when chaining these: or treats a legitimately falsy 0 as "unset". Not reachable today, but it is the kind of guard that breaks quietly once a real field lands.


kv_dim = calculate_kv_dim(
attention_config.num_key_value_heads,
n_head,
n_embd,
head_dim=head_dim,
)
total_num_tokens = seq_len * (batch_size + prefill_queue_size)
kv_cache_size = total_num_tokens * kv_dim
query_prefill_size = seq_len * n_embd if allocate_prefill_query else 0
Expand Down
16 changes: 11 additions & 5 deletions modelopt/torch/puzzletron/utils/misc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,23 +35,29 @@
]


def calculate_kv_dim(num_key_value_heads: int, n_head: int, n_embd: int) -> int:
def calculate_kv_dim(
num_key_value_heads: int,
n_head: int,
n_embd: int,
head_dim: int | None = None,
) -> int:
"""Calculate the key-value dimension for grouped-query attention.

Args:
num_key_value_heads: Number of key-value heads.
n_head: Total number of attention heads.
n_embd: Embedding dimension.
head_dim: Explicit attention head dimension. Falls back to n_embd // n_head.

Returns:
Combined dimension for key and value tensors (2 * num_key_value_heads * head_size).
Combined dimension for key and value tensors
(2 * num_key_value_heads * head_dim).
"""
if num_key_value_heads is None:
return 0
head_size = n_embd // n_head
kv_dim = 2 * num_key_value_heads * head_size
return kv_dim

head_size = head_dim if head_dim is not None else n_embd // n_head
return 2 * num_key_value_heads * head_size

def raise_unknown_subblock_config_error(subblock_config: Any) -> None:
"""Raise an error for invalid subblock configuration types.
Expand Down
18 changes: 18 additions & 0 deletions tests/unit/torch/puzzletron/test_misc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
from modelopt.torch.puzzletron.utils.misc import calculate_kv_dim

@claude claude Bot Sep 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT Compatibility] This new file is missing the Apache-2.0 license header, which will fail the repo's pre-commit gate.

.pre-commit-config.yaml:91 runs the insert-license hook (insert-license-py) over types_or: [python, shell], and its exclude list names only a handful of specific files that carry third-party headers — tests/ is not excluded. The hook will rewrite this file to insert the header and exit non-zero, so CI pre-commit fails until the header is committed.

Every sibling in this directory already has it (tests/unit/torch/puzzletron/conftest.py:1, test_checkpoint_utils_hf.py:1). Prepend the standard block, matching the 2026 form used by the newer files:

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Running pre-commit run --all-files locally will add this for you (the PR description notes the change was made in github.dev without running anything locally).



def test_calculate_kv_dim_with_explicit_head_dim():
assert calculate_kv_dim(
num_key_value_heads=8,
n_head=32,
n_embd=5120,
head_dim=128,
) == 2048


def test_calculate_kv_dim_falls_back_to_hidden_size():
assert calculate_kv_dim(
num_key_value_heads=8,
n_head=32,
n_embd=4096,
) == 2048
Loading