align shape suffix between attn and linear attn

.claude/CLAUDE.md

@@ -109,7 +109,7 @@ this for the comments and docstrings you are adding or rewriting.
 - **Shape-suffix tensor names.** In model code, name tensors with shape
   suffixes (Noam Shazeer convention:
   https://medium.com/@NoamShazeer/shape-suffixes-good-coding-style-f836e72e24fd),
-  e.g. `x_BLD`, `q_BLNH`, `out_TNH`. Capital-letter suffixes denote *logical*
+  e.g. `x_BLD`, `q_BLHK`, `out_THV`. Capital-letter suffixes denote *logical*
   tensor dimensions, not a physical sharding layout -- a name like
   `routed_input_RD` keeps the same suffix whether or not `R` is a local shard
   under EP/SP. Letters are scoped per module, not global: give each module that

CONTRIBUTING.md

@@ -76,7 +76,7 @@ When appropriate, one should consider
   - To add a unit test, put it in the [tests](tests/) folder and follow the existing test files.
   - To add a GPU integration test, add a configuration to the matching module in [torchtitan_recipes/tests](torchtitan_recipes/tests/) and a new `OverrideDefinitions` naming it in [integration_tests](tests/integration_tests/). These suites name a full Trainer configuration per run.
 - Updating [README](README.md) and writing a new note in the [docs](docs/) folder on installation and usage, similar to [float8.md](torchtitan/components/quantization/float8.md).
-- Following the tensor shape-suffix naming convention for new model code (e.g. `x_BLD`, `q_BLNH`, `out_TNH`), with a per-module legend comment as in [attention.py](torchtitan/models/common/attention.py). Capital suffixes name logical tensor dimensions (not sharding layout) and are scoped per file.
+- Following the tensor shape-suffix naming convention for new model code (e.g. `x_BLD`, `q_BLHK`, `out_THV`), with a per-module legend comment as in [attention.py](torchtitan/models/common/attention.py). Capital suffixes name logical tensor dimensions (not sharding layout) and are scoped per file.
 - Adding a new file with benchmark results in [benchmarks](benchmarks) folder.
 - Creating GitHub issues for things that cannot be addressed at the moment.
 - Writing a post on [PyTorch Forums](https://discuss.pytorch.org/c/distributed/torchtitan/44) and linking to it.

tests/unit_tests/cpu/test_flex_attention.py

@@ -29,53 +29,53 @@ def _mask(seq_len: int, batch_size: int):
             device="cpu",
         )
 
-    def test_tnh_layout(self) -> None:
+    def test_thk_thv_layout(self) -> None:
         num_tokens, num_heads, head_dim = 8, 4, 16
-        q_TNH = torch.randn(num_tokens, num_heads, head_dim)
-        k_TNH = torch.randn_like(q_TNH)
-        v_TNH = torch.randn_like(q_TNH)
+        q_THK = torch.randn(num_tokens, num_heads, head_dim)
+        k_THK = torch.randn_like(q_THK)
+        v_THV = torch.randn_like(q_THK)
 
-        def kernel(q_BNTH, k_BNTH, v_BNTH, **kwargs):
-            self.assertEqual(q_BNTH.shape, (1, num_heads, num_tokens, head_dim))
-            self.assertEqual(k_BNTH.shape, q_BNTH.shape)
-            self.assertEqual(v_BNTH.shape, q_BNTH.shape)
-            lse_BNT = torch.randn(1, num_heads, num_tokens)
-            return q_BNTH, SimpleNamespace(lse=lse_BNT)
+        def kernel(q_1HTK, k_1HTK, v_1HTV, **kwargs):
+            self.assertEqual(q_1HTK.shape, (1, num_heads, num_tokens, head_dim))
+            self.assertEqual(k_1HTK.shape, q_1HTK.shape)
+            self.assertEqual(v_1HTV.shape, q_1HTK.shape)
+            lse_1HT = torch.randn(1, num_heads, num_tokens)
+            return q_1HTK, SimpleNamespace(lse=lse_1HT)
 
         with patch.object(FlexAttention, "compiled_flex_attn", side_effect=kernel):
-            out_TNH = self.attention(
-                q_TNH,
-                k_TNH,
-                v_TNH,
+            out_THV = self.attention(
+                q_THK,
+                k_THK,
+                v_THV,
                 attention_masks=self._mask(num_tokens, 1),
             )
 
-        torch.testing.assert_close(out_TNH, q_TNH)
+        torch.testing.assert_close(out_THV, q_THK)
 
-    def test_tnh_out_transform_layout(self) -> None:
+    def test_thk_thv_out_transform_layout(self) -> None:
         num_tokens, num_heads, head_dim = 8, 4, 16
-        q_TNH = torch.randn(num_tokens, num_heads, head_dim)
-        expected_lse_TN = torch.randn(num_tokens, num_heads)
+        q_THK = torch.randn(num_tokens, num_heads, head_dim)
+        expected_lse_TH = torch.randn(num_tokens, num_heads)
 
-        def kernel(q_BNTH, k_BNTH, v_BNTH, **kwargs):
-            return q_BNTH, SimpleNamespace(
-                lse=expected_lse_TN.transpose(0, 1).unsqueeze(0)
+        def kernel(q_1HTK, k_1HTK, v_1HTV, **kwargs):
+            return q_1HTK, SimpleNamespace(
+                lse=expected_lse_TH.transpose(0, 1).unsqueeze(0)
             )
 
-        def out_transform(out_TNH, lse_TN):
-            torch.testing.assert_close(lse_TN, expected_lse_TN)
-            return out_TNH
+        def out_transform(out_THV, lse_TH):
+            torch.testing.assert_close(lse_TH, expected_lse_TH)
+            return out_THV
 
         with patch.object(FlexAttention, "compiled_flex_attn", side_effect=kernel):
-            out_TNH = self.attention(
-                q_TNH,
-                q_TNH,
-                q_TNH,
+            out_THV = self.attention(
+                q_THK,
+                q_THK,
+                q_THK,
                 attention_masks=self._mask(num_tokens, 1),
                 out_transform=out_transform,
             )
 
-        torch.testing.assert_close(out_TNH, q_TNH)
+        torch.testing.assert_close(out_THV, q_THK)
 
 
 if __name__ == "__main__":

tests/unit_tests/cpu/test_fused_qkv.py

@@ -141,22 +141,22 @@ def test_forward_outputs_are_contiguous_and_correct(self):
         fused = _build_fused()
         num_tokens = 6
         x_TD = torch.randn(num_tokens, _DIM)
-        xq_TNH, xk_TNH, xv_TNH = fused(x_TD)
+        xq_THK, xk_THK, xv_THV = fused(x_TD)
 
-        self.assertEqual(xq_TNH.shape, (num_tokens, _N_HEADS, _HEAD_DIM))
-        self.assertEqual(xk_TNH.shape, (num_tokens, _N_KV_HEADS, _HEAD_DIM))
-        self.assertEqual(xv_TNH.shape, (num_tokens, _N_KV_HEADS, _HEAD_DIM))
-        self.assertTrue(xq_TNH.is_contiguous())
-        self.assertTrue(xk_TNH.is_contiguous())
-        self.assertTrue(xv_TNH.is_contiguous())
+        self.assertEqual(xq_THK.shape, (num_tokens, _N_HEADS, _HEAD_DIM))
+        self.assertEqual(xk_THK.shape, (num_tokens, _N_KV_HEADS, _HEAD_DIM))
+        self.assertEqual(xv_THV.shape, (num_tokens, _N_KV_HEADS, _HEAD_DIM))
+        self.assertTrue(xq_THK.is_contiguous())
+        self.assertTrue(xk_THK.is_contiguous())
+        self.assertTrue(xv_THV.is_contiguous())
 
         sd = fused.state_dict()
-        ref_q_TNH = (x_TD @ sd["wq.weight"].T).view(num_tokens, _N_HEADS, _HEAD_DIM)
-        ref_k_TNH = (x_TD @ sd["wk.weight"].T).view(num_tokens, _N_KV_HEADS, _HEAD_DIM)
-        ref_v_TNH = (x_TD @ sd["wv.weight"].T).view(num_tokens, _N_KV_HEADS, _HEAD_DIM)
-        torch.testing.assert_close(xq_TNH, ref_q_TNH)
-        torch.testing.assert_close(xk_TNH, ref_k_TNH)
-        torch.testing.assert_close(xv_TNH, ref_v_TNH)
+        ref_q_THK = (x_TD @ sd["wq.weight"].T).view(num_tokens, _N_HEADS, _HEAD_DIM)
+        ref_k_THK = (x_TD @ sd["wk.weight"].T).view(num_tokens, _N_KV_HEADS, _HEAD_DIM)
+        ref_v_THV = (x_TD @ sd["wv.weight"].T).view(num_tokens, _N_KV_HEADS, _HEAD_DIM)
+        torch.testing.assert_close(xq_THK, ref_q_THK)
+        torch.testing.assert_close(xk_THK, ref_k_THK)
+        torch.testing.assert_close(xv_THV, ref_v_THV)
 
     def test_raw_pointer_read_needs_contiguous(self):
         """A consumer reading the base pointer with contiguous head-major strides

tests/unit_tests/cpu/test_model_td_layout.py

@@ -20,38 +20,38 @@
 
 
 class _AttentionOutput(nn.Module):
-    def forward(self, q_TNH, k_TNH, v_TNH, *, out_transform=None, **kwargs):
-        num_q_heads = q_TNH.shape[1]
-        num_v_heads = v_TNH.shape[1]
-        out_TNH = v_TNH.repeat_interleave(num_q_heads // num_v_heads, dim=1)
+    def forward(self, q_THK, k_THK, v_THV, *, out_transform=None, **kwargs):
+        num_q_heads = q_THK.shape[1]
+        num_v_heads = v_THV.shape[1]
+        out_THV = v_THV.repeat_interleave(num_q_heads // num_v_heads, dim=1)
         if out_transform is not None:
-            lse_TN = torch.zeros(
-                q_TNH.shape[:2], device=q_TNH.device, dtype=q_TNH.dtype
+            lse_TH = torch.zeros(
+                q_THK.shape[:2], device=q_THK.device, dtype=q_THK.dtype
             )
-            out_TNH = out_transform(out_TNH, lse_TN)
-        return out_TNH
+            out_THV = out_transform(out_THV, lse_TH)
+        return out_THV
 
 
 class TestModelTDLayout(unittest.TestCase):
-    def test_sdpa_preserves_blnh_shape(self):
+    def test_sdpa_preserves_blhv_shape(self):
         attention = ScaledDotProductAttention.Config().build()
-        q_BLNH = torch.randn(2, 8, 4, 16)
-        k_BLNH = torch.randn(2, 8, 2, 16)
-        v_BLNH = torch.randn(2, 8, 2, 16)
+        q_BLHK = torch.randn(2, 8, 4, 16)
+        k_BLHK = torch.randn(2, 8, 2, 16)
+        v_BLHV = torch.randn(2, 8, 2, 16)
 
-        out_BLNH = attention(q_BLNH, k_BLNH, v_BLNH, enable_gqa=True)
+        out_BLHV = attention(q_BLHK, k_BLHK, v_BLHV, enable_gqa=True)
 
-        self.assertEqual(out_BLNH.shape, q_BLNH.shape)
+        self.assertEqual(out_BLHV.shape, q_BLHK.shape)
 
-    def test_graph_trainer_sdpa_preserves_tnh_shape(self):
+    def test_graph_trainer_sdpa_preserves_thv_shape(self):
         attention = GraphTrainerScaledDotProductAttention.Config().build()
-        q_TNH = torch.randn(8, 4, 16)
-        k_TNH = torch.randn(8, 2, 16)
-        v_TNH = torch.randn(8, 2, 16)
+        q_THK = torch.randn(8, 4, 16)
+        k_THK = torch.randn(8, 2, 16)
+        v_THV = torch.randn(8, 2, 16)
 
-        out_TNH = attention(q_TNH, k_TNH, v_TNH, enable_gqa=True)
+        out_THV = attention(q_THK, k_THK, v_THV, enable_gqa=True)
 
-        self.assertEqual(out_TNH.shape, q_TNH.shape)
+        self.assertEqual(out_THV.shape, q_THK.shape)
 
     def test_gpt_oss_attention_preserves_td_shape(self):
         config = gptoss_configs["debugmodel"]("standard", "varlen")

tests/unit_tests/cpu/test_varlen_attention.py

@@ -5,7 +5,8 @@
 # LICENSE file in the root directory of this source tree.
 
 # Shape suffix legend:
-#   T = packed tokens, N = attention heads, H = head dimension, D = model dim
+#   T = packed tokens, H = attention heads, K = query/key head dimension,
+#   V = value head dimension, D = model dimension
 
 import unittest
 from unittest.mock import patch
@@ -63,11 +64,11 @@ def test_gqa_preserves_td_shape(self):
         positions_T = torch.tensor([0, 1, 0, 1, 2, 3])
         metadata = create_varlen_metadata_for_document(positions_T)
 
-        def _identity_varlen(q_TNH, k_TNH, v_TNH, *args, **kwargs):
-            self.assertEqual(q_TNH.ndim, 3)
-            self.assertEqual(k_TNH.ndim, 3)
-            self.assertEqual(v_TNH.ndim, 3)
-            return q_TNH
+        def _identity_varlen(q_THK, k_THK, v_THV, *args, **kwargs):
+            self.assertEqual(q_THK.ndim, 3)
+            self.assertEqual(k_THK.ndim, 3)
+            self.assertEqual(v_THV.ndim, 3)
+            return q_THK
 
         with patch(
             "torchtitan.models.common.attention._varlen_attn",
@@ -77,7 +78,7 @@ def _identity_varlen(q_TNH, k_TNH, v_TNH, *args, **kwargs):
 
         self.assertEqual(out_TD.shape, x_TD.shape)
 
-    def test_tnh_sharding_uses_varlen_argument_names(self):
+    def test_thk_thv_sharding_uses_varlen_argument_names(self):
         from torchtitan.models.llama3 import llama3_configs
         from torchtitan.models.llama3.sharding import set_llama3_sharding_config
 
@@ -88,45 +89,45 @@ def test_tnh_sharding_uses_varlen_argument_names(self):
         assert sharding is not None
         self.assertEqual(
             set(sharding.in_src_shardings or {}),
-            {"q_TNH", "k_TNH", "v_TNH"},
+            {"q_THK", "k_THK", "v_THV"},
         )
-        q_layout = (sharding.in_src_shardings or {})["q_TNH"]
-        k_dst_layout = (sharding.in_dst_shardings or {})["k_TNH"]
+        q_layout = (sharding.in_src_shardings or {})["q_THK"]
+        k_dst_layout = (sharding.in_dst_shardings or {})["k_THK"]
         axis_types = _per_axis_types(q_layout)
         self.assertEqual(axis_types[MeshAxisName.DP], spmd.S(0))
         self.assertEqual(axis_types[MeshAxisName.CP], spmd.S(0))
         self.assertEqual(axis_types[MeshAxisName.TP], spmd.S(1))
         self.assertEqual(_per_axis_types(k_dst_layout)[MeshAxisName.CP], spmd.R)
 
-    def test_out_transform_receives_tn_lse(self):
+    def test_out_transform_receives_th_lse(self):
         num_tokens, num_heads, head_dim = 5, 2, 4
-        q_TNH = torch.randn(num_tokens, num_heads, head_dim)
+        q_THK = torch.randn(num_tokens, num_heads, head_dim)
         positions_T = torch.tensor([0, 1, 0, 1, 2])
         metadata = create_varlen_metadata_for_document(positions_T)
         inner_attention = VarlenAttention.Config().build()
 
         def _varlen_with_lse(q, k, v, *args, **kwargs):
-            lse_NT = torch.randn(num_heads, num_tokens)
-            return q, lse_NT
+            lse_HT = torch.randn(num_heads, num_tokens)
+            return q, lse_HT
 
-        def _check_shapes(out_TNH, lse_TN):
-            self.assertEqual(out_TNH.shape, q_TNH.shape)
-            self.assertEqual(lse_TN.shape, (num_tokens, num_heads))
-            return out_TNH
+        def _check_shapes(out_THV, lse_TH):
+            self.assertEqual(out_THV.shape, q_THK.shape)
+            self.assertEqual(lse_TH.shape, (num_tokens, num_heads))
+            return out_THV
 
         with patch(
             "torchtitan.models.common.attention._varlen_attn",
             side_effect=_varlen_with_lse,
         ):
-            out_TNH = inner_attention(
-                q_TNH,
-                q_TNH,
-                q_TNH,
+            out_THV = inner_attention(
+                q_THK,
+                q_THK,
+                q_THK,
                 attention_masks=metadata,
                 out_transform=_check_shapes,
             )
 
-        self.assertEqual(out_TNH.shape, q_TNH.shape)
+        self.assertEqual(out_THV.shape, q_THK.shape)
 
     def test_llama_decoder_preserves_td_shape(self):
         from torchtitan.models.llama3 import llama3_configs
@@ -138,8 +139,8 @@ def test_llama_decoder_preserves_td_shape(self):
         positions_T = torch.tensor([0, 1, 0, 1, 2, 3])
         metadata = model.get_attention_masks(positions_T)
 
-        def _identity_varlen(q_TNH, k_TNH, v_TNH, *args, **kwargs):
-            return q_TNH
+        def _identity_varlen(q_THK, k_THK, v_THV, *args, **kwargs):
+            return q_THK
 
         with patch(
             "torchtitan.models.common.attention._varlen_attn",

tests/unit_tests/cpu/test_vision_attention.py

@@ -19,9 +19,9 @@ def __init__(self) -> None:
         super().__init__()
         self.input_shape: torch.Size | None = None
 
-    def forward(self, q_TNH, k_TNH, v_TNH, *, attention_masks):
-        self.input_shape = q_TNH.shape
-        return q_TNH
+    def forward(self, q_THK, k_THK, v_THV, *, attention_masks):
+        self.input_shape = q_THK.shape
+        return q_THK
 
 
 class TestVisionAttention(unittest.TestCase):