From 459c6c56d899b1c2f25b0c708348db91e2987b20 Mon Sep 17 00:00:00 2001
From: Mani <122438942+devxMani@users.noreply.github.com>
Date: Sun, 20 Sep 2026 18:58:44 +0530
Subject: [PATCH] fix: use explicit head dim for puzzletron kv cache

---
 .../calc_subblock_params_and_memory.py         | 10 +++++++++-
 modelopt/torch/puzzletron/utils/misc.py        | 16 +++++++++++-----
 tests/unit/torch/puzzletron/test_misc.py       | 18 ++++++++++++++++++
 3 files changed, 38 insertions(+), 6 deletions(-)
 create mode 100644 tests/unit/torch/puzzletron/test_misc.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 531f7a3f0a1..ea62933b38a 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
@@ -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)
+    head_dim = getattr(attention_config, "qk_head_dim", None) or getattr(lm_config, "head_dim", None)
+
+    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
diff --git a/modelopt/torch/puzzletron/utils/misc.py b/modelopt/torch/puzzletron/utils/misc.py
index 68751d1e07e..478d21dde5f 100644
--- a/modelopt/torch/puzzletron/utils/misc.py
+++ b/modelopt/torch/puzzletron/utils/misc.py
@@ -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.
diff --git a/tests/unit/torch/puzzletron/test_misc.py b/tests/unit/torch/puzzletron/test_misc.py
new file mode 100644
index 00000000000..4f66422052b
--- /dev/null
+++ b/tests/unit/torch/puzzletron/test_misc.py
@@ -0,0 +1,18 @@
+from modelopt.torch.puzzletron.utils.misc import calculate_kv_dim
+
+
+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
\ No newline at end of file