[Puzzletron] Fix KV-cache dimension calculation for models with custom head_dim (#2482)

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,16 @@ 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
+):
+    """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")
+        else config
+    )
+    return getattr(lm_config, name, getattr(config, name, default))
 
 
 def _configured_subblock(
@@ -554,7 +569,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:

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
 

tests/unit/torch/puzzletron/test_kv_dim_calculation.py

@@ -0,0 +1,161 @@
+# 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:
+    """Dummy model descriptor for testing configuration extraction."""
+
+    @staticmethod
+    def get_language_model_config(config):
+        """Return the given model configuration directly."""
+        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