From 3e5f59ab32c3ccd245bd79bedd8dcabfde1b815c Mon Sep 17 00:00:00 2001 From: Felipe Mello Date: Tue, 20 Oct 2026 01:35:44 -0700 Subject: [PATCH] [muse_spark][perf] Split the packed qkv inside the q/k region QKVLinear split wqkv's output and made three contiguous copies in eager before Muse Spark's q/k region. The split now lives in QKVLinear.split, and Muse Spark passes the packed qkv into its region, so the copies fuse into the region's kernels. Other models call forward = wqkv + split and run the same ops as before. Muse Spark 4-layer unit, GB300, CUDA graphs: decode -0.9% to -2.4% (T=1..256); train 16k -0.14% by kernel time (glue 1.69 -> 1.21 ms per unit; wall rows are min of 2 with ~0.5 ms noise). Outputs and grads bitwise identical. --- .../muse_spark/test_attention_compile.py | 19 ++++++++------ .../models/common/attention/attention.py | 23 ++++++++++++----- torchtitan/models/muse_spark/model.py | 25 +++++++++++-------- 3 files changed, 42 insertions(+), 25 deletions(-) diff --git a/tests/unit_tests/gpu/models/muse_spark/test_attention_compile.py b/tests/unit_tests/gpu/models/muse_spark/test_attention_compile.py index 1497b88fb37..da2d0f00653 100644 --- a/tests/unit_tests/gpu/models/muse_spark/test_attention_compile.py +++ b/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 ) diff --git a/torchtitan/models/common/attention/attention.py b/torchtitan/models/common/attention/attention.py index ad23ce5b43b..8b37deb4f83 100644 --- a/torchtitan/models/common/attention/attention.py +++ b/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) diff --git a/torchtitan/models/muse_spark/model.py b/torchtitan/models/muse_spark/model.py index 576e2ad7d6e..fe59d1d460d 100644 --- a/torchtitan/models/muse_spark/model.py +++ b/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(