feat: add explanation on `hidden_size = num_attention_heads * dim_per_head`

README.md

@@ -27,7 +27,8 @@ learning notes only and are **not** imported by the app.
 1. **Random matrix** — sample even-width Q (and K) tensors, apply numpy RoPE,
    inspect heatmaps, pairwise 2D rotation, `QK^T`, and additive sinusoidal PE.
 2. **Real model** — lazy-load an ungated Llama-like checkpoint (default
-   `HuggingFaceTB/SmolLM2-135M`), take `embed_tokens`, first-layer `q_proj` /
+   `HuggingFaceTB/SmolLM2-135M`; `HuggingFaceM4/tiny-random-LlamaForCausalLM`
+   is included for a very small test model), take `embed_tokens`, first-layer `q_proj` /
    `k_proj` (GQA-aware), and compare educational numpy RoPE (`llama` pairing)
    to the model's `rotary_emb`. No Hugging Face token is required. Gated models
    are not used.

app.py

@@ -14,7 +14,7 @@ from src.extract import (
     MODEL_CHOICES,
     expand_kv_heads,
     extract_from_model,
-    get_model_head_dim,
+    get_model_dimensions,
     random_qk,
     select_head,
 )
@@ -93,18 +93,28 @@ def _safe_slider_max(n: int) -> int:
 
 def update_dimension(source: str, model_name: str):
     if source.startswith("Random"):
-        return gr.update(minimum=4, maximum=128, value=32, step=2, interactive=True)
+        return (
+            gr.update(minimum=4, maximum=128, value=32, step=2, interactive=True),
+            gr.update(value=32),
+            gr.update(value=1),
+            gr.update(value=32),
+        )
     try:
-        model_dim = get_model_head_dim(model_name)
-        return gr.update(
-            minimum=4,
-            maximum=max(128, model_dim),
-            value=model_dim,
-            step=2,
-            interactive=False,
+        total_dim, n_heads, head_dim = get_model_dimensions(model_name)
+        return (
+            gr.update(
+                minimum=4,
+                maximum=max(128, head_dim),
+                value=head_dim,
+                step=2,
+                interactive=False,
+            ),
+            gr.update(value=total_dim),
+            gr.update(value=n_heads),
+            gr.update(value=head_dim),
         )
     except Exception:
-        return gr.update()
+        return gr.update(), gr.update(), gr.update(), gr.update()
 
 
 def compute(
@@ -306,6 +316,15 @@ with gr.Blocks(title="RoPE Explorer") as demo:
                 seq_len = gr.Slider(2, MAX_SEQ_LEN, value=16, step=1, label="Sequence length")
                 dim = gr.Slider(4, 128, value=32, step=2, label="Dimension (even; per attention head)")
                 seed = gr.Number(value=42, label="Seed", precision=0)
+            with gr.Row():
+                total_dim = gr.Number(value=32, label="Total dimension", precision=0, interactive=False)
+                attention_heads = gr.Number(value=1, label="Attention heads", precision=0, interactive=False)
+                head_dim = gr.Number(value=32, label="Dimension per attention head", precision=0, interactive=False)
+            gr.Markdown(
+                "**Why these numbers differ:** `total dimension = attention heads × dimension per head`. "
+                "RoPE rotates each query/key head separately, so its Dimension slider uses "
+                "the per-head value, not the model's total dimension."
+            )
             base = gr.Number(
                 value=10000,
                 label="RoPE base (overridden by config.rope_theta for real models)",
@@ -355,12 +374,12 @@ with gr.Blocks(title="RoPE Explorer") as demo:
     source.change(
         update_dimension,
         inputs=[source, model_name],
-        outputs=[dim],
+        outputs=[dim, total_dim, attention_heads, head_dim],
     )
     model_name.change(
         update_dimension,
         inputs=[source, model_name],
-        outputs=[dim],
+        outputs=[dim, total_dim, attention_heads, head_dim],
     )
 
     bulk_inputs = [state, which, head, mod_2pi]

src/extract.py

@@ -108,20 +108,33 @@ def _require_hf():
     return torch, AutoModelForCausalLM, AutoTokenizer
 
 
-def get_model_head_dim(model_name: str) -> int:
-    """Return the Q/K dimension per attention head without loading model weights."""
+def _get_model_config(model_name: str):
     if model_name not in _config_cache:
         try:
             from transformers import AutoConfig
         except ImportError as exc:
             raise RuntimeError("Real-model mode needs `transformers`.") from exc
         _config_cache[model_name] = AutoConfig.from_pretrained(model_name)
+    return _config_cache[model_name]
 
-    config = _config_cache[model_name]
+
+def get_model_dimensions(model_name: str) -> tuple[int, int, int]:
+    """Return total hidden size, attention heads, and per-head Q/K size."""
+    config = _get_model_config(model_name)
+    total_dim = int(config.hidden_size)
+    n_heads = int(config.num_attention_heads)
     configured_head_dim = getattr(config, "head_dim", None)
-    if configured_head_dim is not None:
-        return int(configured_head_dim)
-    return int(config.hidden_size) // int(config.num_attention_heads)
+    head_dim = (
+        int(configured_head_dim)
+        if configured_head_dim is not None
+        else total_dim // n_heads
+    )
+    return total_dim, n_heads, head_dim
+
+
+def get_model_head_dim(model_name: str) -> int:
+    """Return the Q/K dimension per attention head without loading model weights."""
+    return get_model_dimensions(model_name)[2]
 
 
 def get_model(model_name: str):