[muse_spark][perf] Split the packed qkv inside the q/k region

tests/unit_tests/gpu/models/muse_spark/test_attention_compile.py

@@ -33,21 +33,24 @@ def test_attention_region_matches_eager(region, layer_id):
     attention = layer.attention.build().cuda()
     num_tokens, num_heads, num_kv_heads = 64, attention.n_heads, attention.n_kv_heads
     head_dim = attention.head_dim
-    query_TNH = torch.randn(
-        num_tokens, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
-    )
     if region == "muse_spark_attention_qk":
         positions_T = torch.arange(num_tokens, device="cuda")
 
-        def fn(query_TNH, key_TKH):
-            return attention._prepare_query_and_key(query_TNH, key_TKH, positions_T)
+        def fn(qkv):
+            return attention._prepare_query_key_value(qkv, positions_T)
 
-        key_TKH = torch.randn(
-            num_tokens, num_kv_heads, head_dim, device="cuda", dtype=torch.bfloat16
+        qkv = torch.randn(
+            num_tokens,
+            (num_heads + 2 * num_kv_heads) * head_dim,
+            device="cuda",
+            dtype=torch.bfloat16,
         )
-        inputs_ = (query_TNH, key_TKH)
+        inputs_ = (qkv,)
     else:
         fn = attention._flatten_attention_output
+        query_TNH = torch.randn(
+            num_tokens, num_heads, head_dim, device="cuda", dtype=torch.bfloat16
+        )
         gate_TO = torch.randn(
             num_tokens, num_heads * head_dim, device="cuda", dtype=torch.bfloat16
         )

torchtitan/models/common/attention/attention.py

@@ -147,6 +147,15 @@ def __init__(self, config: Config):
         self.heads_per_kv = config.n_heads // config.n_kv_heads
         self.r_dim = self.heads_per_kv + 2
 
+    def forward(
+        self, x: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # Fused QKV: single matmul, then reshape and split along R dim.
+        qkv = self.wqkv(x)
+        # The split below copies the wqkv projection output with bare ops.
+        remat.recompute_needs_tensor(qkv)
+        return self.split(qkv)
+
     @spmd.local_map(
         out_types=(
             (
@@ -156,15 +165,17 @@ def __init__(self, config: Config):
         )
         * 3
     )
-    def forward(
-        self, x: torch.Tensor
+    def split(
+        self, qkv: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        # Fused QKV: single matmul, then reshape and split along R dim.
+        """Split the ``wqkv`` output into contiguous q, k and v.
+
+        Separate from ``forward`` so a compiled region can split ``wqkv(x)`` and fuse the
+        copies. An eager caller other than ``forward`` calls
+        ``remat.recompute_needs_tensor(qkv)`` first, as ``forward`` does.
+        """
         # [T, n_kv_heads * R * head_dim] -> [T, n_kv_heads, R, head_dim]
         # Use -1 for n_kv_heads so TP sharding is handled automatically.
-        qkv = self.wqkv(x)
-        # The split below copies the wqkv projection output with bare ops.
-        remat.recompute_needs_tensor(qkv)
         num_tokens = qkv.shape[0]
         with spmd.local():  # TODO(pianpwk): same QKV:S(1) unflatten case handled by even sharding
             qkv = qkv.view(num_tokens, -1, self.r_dim, self.head_dim)

torchtitan/models/muse_spark/model.py

@@ -119,13 +119,13 @@ def forward(
         attention_metadata: FlexAttentionMetadata | VarlenAttentionMetadata | None,
         positions: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        query_TNH, key_TKH, value_TKH = self.qkv_linear(x_TD)
+        qkv = self.qkv_linear.wqkv(x_TD)
         output_gate_TO = None
         if self.output_gate is not None:
             output_gate_TO = self.output_gate(x_TD)
 
-        remat.recompute_needs_tensor(query_TNH, key_TKH)
-        query_TNH, key_TKH = self._prepare_query_and_key(query_TNH, key_TKH, positions)
+        remat.recompute_needs_tensor(qkv)
+        query_TNH, key_TKH, value_TKH = self._prepare_query_key_value(qkv, positions)
 
         output_TNH = remat.region(
             self.inner_attention,
@@ -155,17 +155,20 @@ def forward(
         recompile_limit=16,
         options={"triton.autotune_pointwise": False},
     )
-    def _prepare_query_and_key(
+    def _prepare_query_key_value(
         self,
-        query_TNH: torch.Tensor,
-        key_TKH: torch.Tensor,
+        qkv: torch.Tensor,
         positions: torch.Tensor | None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Return the query and key as attention reads them. Each step runs only if configured.
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Split qkv and return q, k, v as attention reads them. Steps after the split run only if configured.
 
-        query: q_norm -> * scale_query_by -> * position temperature -> RoPE
-        key:   k_norm -> RoPE
+        query: split -> q_norm -> * scale_query_by -> * position temperature -> RoPE
+        key:   split -> k_norm -> RoPE
+        value: split
         """
+        # Split here, not in eager: the q/k copies fuse into the norms, and the backward
+        # writes the packed qkv grad directly instead of a cat.
+        query_TNH, key_TKH, value_TKH = self.qkv_linear.split(qkv)
         if self.q_norm is not None or self.k_norm is not None:
             if self.q_norm is None or self.k_norm is None:
                 raise ValueError("MuseSparkAttention needs both q_norm and k_norm.")
@@ -188,7 +191,7 @@ def _prepare_query_and_key(
 
         if self.rope is not None:
             query_TNH, key_TKH = self.rope(query_TNH, key_TKH, positions)
-        return query_TNH, key_TKH
+        return query_TNH, key_TKH, value_TKH
 
     @local_compile("muse_spark_attention_output", batch_invariant=False)
     def _flatten_attention_output(