From c924530a6384ead67dcd1b20203b4881d3f256a0 Mon Sep 17 00:00:00 2001 From: musi22 Date: Sun, 20 Sep 2026 18:35:58 +0530 Subject: [PATCH 1/2] fix(puzzletron): respect explicit head_dim in KV-cache dimension calculation (#2482) Signed-off-by: musi22 --- .../calc_subblock_params_and_memory.py | 29 +++- modelopt/torch/puzzletron/utils/misc.py | 15 +- .../puzzletron/test_kv_dim_calculation.py | 158 ++++++++++++++++++ 3 files changed, 196 insertions(+), 6 deletions(-) create mode 100644 tests/unit/torch/puzzletron/test_kv_dim_calculation.py diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py index 6e6b4ab8237..a87e6ae4d78 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py +++ b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py @@ -105,7 +105,14 @@ def calculate_additive_metrics( decode_flops = 2 * active_params * batch_size * decode_tokens if isinstance(subblock_config, AttentionConfig) and not subblock_config.no_op: - kv_dim = calculate_kv_dim(subblock_config.num_kv_heads, n_head, n_embd) + effective_head_dim = ( + subblock_config.qk_head_dim + if subblock_config.qk_head_dim is not None + else _language_model_attr(model_config, descriptor, "head_dim", None) + ) + kv_dim = calculate_kv_dim( + subblock_config.num_kv_heads, n_head, n_embd, head_dim=effective_head_dim + ) kv_cache_bytes_per_token = kv_dim * sizeof_dtype(kv_cache_dtype) query_heads = int(subblock_config.num_query_heads or n_head) head_dim = int(subblock_config.qk_head_dim or (n_embd // n_head)) @@ -170,8 +177,15 @@ def calculate_additive_metrics( } -def _language_model_attr(config: PretrainedConfig, descriptor: type[ModelDescriptor], name: str, default=None): - return getattr(descriptor.get_language_model_config(config), name, getattr(config, name, default)) +def _language_model_attr( + config: PretrainedConfig, descriptor: type[ModelDescriptor], name: str, default=None +): + lm_config = ( + descriptor.get_language_model_config(config) + if descriptor is not None and hasattr(descriptor, "get_language_model_config") + else config + ) + return getattr(lm_config, name, getattr(config, name, default)) def _configured_subblock( @@ -554,7 +568,14 @@ def calculate_attention_memory( if isinstance(sliding_window, int): seq_len = min(seq_len, sliding_window) - kv_dim = calculate_kv_dim(attention_config.num_kv_heads, n_head, n_embd) + effective_head_dim = ( + attention_config.qk_head_dim + if getattr(attention_config, "qk_head_dim", None) is not None + else _language_model_attr(model_config, descriptor, "head_dim", None) + ) + kv_dim = calculate_kv_dim( + attention_config.num_kv_heads, n_head, n_embd, head_dim=effective_head_dim + ) total_num_tokens = seq_len * batch_size kv_cache_size = total_num_tokens * kv_dim if num_params is None: diff --git a/modelopt/torch/puzzletron/utils/misc.py b/modelopt/torch/puzzletron/utils/misc.py index f4b667d827e..1176943eb3a 100644 --- a/modelopt/torch/puzzletron/utils/misc.py +++ b/modelopt/torch/puzzletron/utils/misc.py @@ -35,20 +35,31 @@ ] -def calculate_kv_dim(num_kv_heads: int, n_head: int, n_embd: int) -> int: +def calculate_kv_dim( + num_kv_heads: int, + n_head: int | None = None, + n_embd: int | None = None, + head_dim: int | None = None, +) -> int: """Calculate the key-value dimension for grouped-query attention. Args: num_kv_heads: Number of key-value heads. n_head: Total number of attention heads. n_embd: Embedding dimension. + head_dim: Explicit dimension per attention head. If provided, takes precedence over n_embd // n_head. Returns: Combined dimension for key and value tensors (2 * num_kv_heads * head_size). """ if num_kv_heads is None: return 0 - head_size = n_embd // n_head + if head_dim is not None: + head_size = head_dim + elif n_head is not None and n_embd is not None and n_head > 0: + head_size = n_embd // n_head + else: + head_size = 0 kv_dim = 2 * num_kv_heads * head_size return kv_dim diff --git a/tests/unit/torch/puzzletron/test_kv_dim_calculation.py b/tests/unit/torch/puzzletron/test_kv_dim_calculation.py new file mode 100644 index 00000000000..820ea97f2c1 --- /dev/null +++ b/tests/unit/torch/puzzletron/test_kv_dim_calculation.py @@ -0,0 +1,158 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: PLR0917 + +from types import SimpleNamespace + +import pytest +import torch + +from modelopt.torch.puzzletron.block_config import AttentionConfig +from modelopt.torch.puzzletron.subblock_stats.calc_subblock_params_and_memory import ( + calculate_additive_metrics, + calculate_attention_memory, +) +from modelopt.torch.puzzletron.utils.misc import calculate_kv_dim + + +class DummyDescriptor: + @staticmethod + def get_language_model_config(config): + return config + + +@pytest.mark.parametrize( + ("name", "n_embd", "n_head", "num_kv_heads", "head_dim", "expected_kv_dim"), + [ + # (a) Explicit head_dim smaller than inferred: Mistral Small 24B (128 < 160) + ("Mistral-Small-24B", 5120, 32, 8, 128, 2048), + # (b) Explicit head_dim larger than inferred: gpt-oss 20B (64 > 45) + ("gpt-oss-20b", 2880, 64, 8, 64, 1024), + # (c) Fallback when head_dim is None: inferred from n_embd // n_head + ("Llama-fallback", 4096, 32, 8, None, 2048), + ], +) +def test_calculate_kv_dim_arithmetic(name, n_embd, n_head, num_kv_heads, head_dim, expected_kv_dim): + """Test calculate_kv_dim respects explicit head_dim and falls back to n_embd // n_head.""" + kv_dim = calculate_kv_dim(num_kv_heads, n_head, n_embd, head_dim=head_dim) + assert kv_dim == expected_kv_dim + + +def test_calculate_attention_memory_with_model_head_dim(): + """Verify calculate_attention_memory uses model_config.head_dim over n_embd // n_head.""" + # Mistral Small 24B configuration: n_embd=5120, n_head=32, num_kv_heads=8, head_dim=128 + attention_config = AttentionConfig(num_query_heads=32, num_kv_heads=8) + model_config = SimpleNamespace(head_dim=128) + + batch_size = 2 + prefill_seq_len = 10 + generation_seq_len = 0 + seq_len = prefill_seq_len + generation_seq_len + dtype = torch.bfloat16 + dtype_bytes = 2 + + # Expected: kv_dim = 2 * 8 * 128 = 2048. kv_cache_size = 2 * 10 * 2048 = 40960 elements. + expected_kv_cache_bytes = batch_size * seq_len * (2 * 8 * 128) * dtype_bytes + expected_kv_cache_mib = expected_kv_cache_bytes / 2**20 + + mem = calculate_attention_memory( + attention_config, + model_config=model_config, + descriptor=DummyDescriptor, + batch_size=batch_size, + prefill_seq_len=prefill_seq_len, + generation_seq_len=generation_seq_len, + n_embd=5120, + n_head=32, + weights_dtype=dtype, + kv_cache_dtype=dtype, + num_params=0, + ) + + assert mem["kv_cache_memory_mib"] == pytest.approx(expected_kv_cache_mib) + + +def test_calculate_additive_metrics_with_model_head_dim(): + """Verify calculate_additive_metrics uses model_config.head_dim for kv_cache_bytes_per_token.""" + # gpt-oss 20B configuration: n_embd=2880, n_head=64, num_kv_heads=8, head_dim=64 + attention_config = AttentionConfig(num_query_heads=64, num_kv_heads=8) + model_config = SimpleNamespace(head_dim=64) + + metrics = calculate_additive_metrics( + attention_config, + model_config=model_config, + descriptor=DummyDescriptor, + batch_size=1, + prefill_seq_len=4, + generation_seq_len=1, + n_embd=2880, + n_head=64, + weights_dtype=torch.bfloat16, + kv_cache_dtype=torch.bfloat16, + num_params=0, + active_params=0, + ) + + # Expected: kv_dim = 2 * 8 * 64 = 1024. Bytes per token = 1024 * 2 = 2048. + expected_bytes_per_token = (2 * 8 * 64) * 2 + assert metrics["kv_cache_bytes_per_token"] == expected_bytes_per_token + + +def test_effective_head_dim_precedence(): + """Verify precedence: qk_head_dim > model_config.head_dim > (n_embd // n_head).""" + # 1. When qk_head_dim is set, it overrides model_config.head_dim and fallback + subblock_1 = AttentionConfig(num_query_heads=32, num_kv_heads=8, qk_head_dim=64) + model_config_1 = SimpleNamespace(head_dim=128) + m1 = calculate_additive_metrics( + subblock_1, + model_config=model_config_1, + descriptor=DummyDescriptor, + batch_size=1, + prefill_seq_len=4, + generation_seq_len=1, + n_embd=5120, + n_head=32, + weights_dtype=torch.bfloat16, + kv_cache_dtype=torch.bfloat16, + num_params=0, + active_params=0, + ) + assert m1["kv_cache_bytes_per_token"] == (2 * 8 * 64) * 2 + + # 2. When qk_head_dim is None, model_config.head_dim overrides fallback + subblock_2 = AttentionConfig(num_query_heads=32, num_kv_heads=8, qk_head_dim=None) + model_config_2 = SimpleNamespace(head_dim=128) + m2 = calculate_additive_metrics( + subblock_2, + model_config=model_config_2, + descriptor=DummyDescriptor, + batch_size=1, + prefill_seq_len=4, + generation_seq_len=1, + n_embd=5120, + n_head=32, + weights_dtype=torch.bfloat16, + kv_cache_dtype=torch.bfloat16, + num_params=0, + active_params=0, + ) + assert m2["kv_cache_bytes_per_token"] == (2 * 8 * 128) * 2 + + # 3. When neither is set, falls back to n_embd // n_head = 5120 // 32 = 160 + subblock_3 = AttentionConfig(num_query_heads=32, num_kv_heads=8, qk_head_dim=None) + model_config_3 = SimpleNamespace() + m3 = calculate_additive_metrics( + subblock_3, + model_config=model_config_3, + descriptor=DummyDescriptor, + batch_size=1, + prefill_seq_len=4, + generation_seq_len=1, + n_embd=5120, + n_head=32, + weights_dtype=torch.bfloat16, + kv_cache_dtype=torch.bfloat16, + num_params=0, + active_params=0, + ) + assert m3["kv_cache_bytes_per_token"] == (2 * 8 * 160) * 2 From c36dac7f05244dd01d155bb11e3ac0abb229aa8e Mon Sep 17 00:00:00 2001 From: musi22 Date: Sun, 20 Sep 2026 20:12:04 +0530 Subject: [PATCH 2/2] docs(puzzletron): add docstrings for configuration helpers to satisfy 100% coverage Signed-off-by: musi22 --- .../subblock_stats/calc_subblock_params_and_memory.py | 1 + tests/unit/torch/puzzletron/test_kv_dim_calculation.py | 3 +++ 2 files changed, 4 insertions(+) diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py index a87e6ae4d78..8df515d86db 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py +++ b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_params_and_memory.py @@ -180,6 +180,7 @@ def calculate_additive_metrics( def _language_model_attr( config: PretrainedConfig, descriptor: type[ModelDescriptor], name: str, default=None ): + """Retrieve an attribute from language model config or top-level config.""" lm_config = ( descriptor.get_language_model_config(config) if descriptor is not None and hasattr(descriptor, "get_language_model_config") diff --git a/tests/unit/torch/puzzletron/test_kv_dim_calculation.py b/tests/unit/torch/puzzletron/test_kv_dim_calculation.py index 820ea97f2c1..74197dfa73f 100644 --- a/tests/unit/torch/puzzletron/test_kv_dim_calculation.py +++ b/tests/unit/torch/puzzletron/test_kv_dim_calculation.py @@ -16,8 +16,11 @@ class DummyDescriptor: + """Dummy model descriptor for testing configuration extraction.""" + @staticmethod def get_language_model_config(config): + """Return the given model configuration directly.""" return config