Reduce duplicated MLA config comment per code review

Move the full MLA kv-heads explanation to a single top-of-file note and
reference it concisely from each MLA config entry.

Signed-off-by: copilot

Author: Copilot

tests/_test_configs.py

@@ -57,6 +57,13 @@
 
 LONGROPE_FACTORS = [1.0] * (int(TINY_HEAD_DIM * 0.5) // 2)
 
+# NOTE (MLA models): Multi-head Latent Attention models (DeepSeek-V2/V3,
+# LongCat-Flash, ...) reconstruct full-head K/V from a shared latent, so they do
+# not use grouped-query attention.  Their tiny configs must set
+# num_key_value_heads == num_attention_heads; otherwise HuggingFace's repeat_kv()
+# in the SDPA path duplicates the already-full-head K/V tensors, producing a
+# head-count mismatch against the query.
+
 
 def _base_config(config_cls=None, **overrides) -> ArchitectureConfig:
     """Create a tiny ArchitectureConfig for graph-build and parity tests.
@@ -550,8 +557,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig:
     (
         "deepseek_v3",
         {
-            # MLA uses full-head K/V (no GQA); num_key_value_heads must equal
-            # num_attention_heads or HF's repeat_kv in SDPA inflates the KV heads.
+            # MLA: kv heads must equal attn heads (see MLA note near top of file).
             "num_key_value_heads": TINY_HEADS,
             "q_lora_rank": 32,
             "kv_lora_rank": 16,
@@ -584,8 +590,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig:
     (
         "deepseek_v2",
         {
-            # MLA uses full-head K/V (no GQA); num_key_value_heads must equal
-            # num_attention_heads or HF's repeat_kv in SDPA inflates the KV heads.
+            # MLA: kv heads must equal attn heads (see MLA note near top of file).
             "num_key_value_heads": TINY_HEADS,
             "q_lora_rank": 32,
             "kv_lora_rank": 16,
@@ -610,8 +615,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig:
     (
         "deepseek_v2_moe",
         {
-            # MLA uses full-head K/V (no GQA); num_key_value_heads must equal
-            # num_attention_heads or HF's repeat_kv in SDPA inflates the KV heads.
+            # MLA: kv heads must equal attn heads (see MLA note near top of file).
             "num_key_value_heads": TINY_HEADS,
             "q_lora_rank": 32,
             "kv_lora_rank": 16,
@@ -783,8 +787,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig:
         "longcat_flash",
         {
             "_config_cls": LongcatFlashConfig,
-            # MLA uses full-head K/V (no GQA); num_key_value_heads must equal
-            # num_attention_heads or HF's repeat_kv in SDPA inflates the KV heads.
+            # MLA: kv heads must equal attn heads (see MLA note near top of file).
             "num_key_value_heads": TINY_HEADS,
             "q_lora_rank": 16,
             "kv_lora_rank": 8,
@@ -1097,8 +1100,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig:
     (
         "deepseek_v2",
         {
-            # MLA uses full-head K/V (no GQA); num_key_value_heads must equal
-            # num_attention_heads or HF's repeat_kv in SDPA inflates the KV heads.
+            # MLA: kv heads must equal attn heads (see MLA note near top of file).
             "num_key_value_heads": TINY_HEADS,
             "q_lora_rank": 32,
             "kv_lora_rank": 16,