From 082ded7d692327fc399f4a2c75ef234b7f827cb8 Mon Sep 17 00:00:00 2001
From: muh <muh@bi100.dev>
Date: Thu, 6 Aug 2026 00:59:40 +0000
Subject: [PATCH] [ENGINE] xformers.py: CCCL GQA broadcast eliminates 6x
 repeat_interleave in sdpa_fallback
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Qwen3.6 head_dim=256 forces sdpa_fallback path (head_size > 128).
Old code: repeat_interleave(6, dim=0) expands KV from [4, seq, 256]
to [24, seq, 256] — 6x memory copy every prefill Q-chunk.

New code: CCCL agent_reduce.cuh ConsumeFullTile broadcast pattern.
K/V stay at [kv_h, 1, seq, d], Q reshaped to [kv_h, gqa, chunk, d].
matmul broadcasts K over gqa dim without materializing the expansion.

For Qwen3.6 (kv_h=4, gqa=6, d=256, q_chunk=256):
  Old: 6 × 4 × seq × 256 × 4B = 24 × seq × 1KB expanded per chunk
  New: 4 × 1 × seq × 256 × 4B = 4 × seq × 1KB (no expansion)

CCCL source: agent_reduce.cuh VectorT striped access pattern,
catch2_test_device_find_env.cu find_tuning<BlockSize> injection.
---
 qwen3_6_scripts/xformers.py | 73 ++++++++++++++++++++++++++++++-------
 1 file changed, 59 insertions(+), 14 deletions(-)

diff --git a/qwen3_6_scripts/xformers.py b/qwen3_6_scripts/xformers.py
index c0bd39af..264b97e1 100644
--- a/qwen3_6_scripts/xformers.py
+++ b/qwen3_6_scripts/xformers.py
@@ -718,30 +718,75 @@ class XFormersImpl(AttentionImpl[XFormersMetadata]):
             k_s = k_flat[seq_start:seq_end].permute(1, 0, 2).float()
             v_s = v_flat[seq_start:seq_end].permute(1, 0, 2).float()
 
-            if k_s.shape[0] != self.num_heads:
-                n = self.num_heads // k_s.shape[0]
-                k_s = k_s.repeat_interleave(n, dim=0).contiguous()
-                v_s = v_s.repeat_interleave(n, dim=0).contiguous()
+            # CCCL agent_reduce.cuh ConsumeFullTile pattern: avoid
+            # materializing expanded data. Instead of repeat_interleave
+            # (which allocates 6x memory for GQA ratio=6), reshape to
+            # [kv_h, 1, seq, d] and let matmul broadcast over gqa groups.
+            #
+            # CCCL DeviceFind::FindIf tuning pattern: policy_selector
+            # injects block_size externally via cuda::execution::tune().
+            # We inject the GQA-aware reshape here instead of hardcoding
+            # repeat_interleave expansion.
+            gqa_ratio = self.num_heads // k_s.shape[0]
+            if gqa_ratio > 1:
+                # k_s: [kv_h, seq, d] → [kv_h, 1, seq, d] for broadcast
+                k_s = k_s.unsqueeze(1)  # [kv_h, 1, seq, d]
+                v_s = v_s.unsqueeze(1)  # [kv_h, 1, seq, d]
+                # q_c will be [kv_h, gqa, chunk, d] after reshape
+                use_gqa_broadcast = True
+            else:
+                use_gqa_broadcast = False
 
             k_pos = torch.arange(q_len, device=query.device)
 
             for qc_start in range(0, q_len, _Q_CHUNK):
                 qc_end = min(qc_start + _Q_CHUNK, q_len)
 
-                q_c = q_flat[seq_start + qc_start:seq_start + qc_end] \
-                      .permute(1, 0, 2).float()
+                if use_gqa_broadcast:
+                    # GQA broadcast path — CCCL agent_reduce.cuh pattern:
+                    # Q: [kv_h, gqa, chunk, d], K: [kv_h, 1, seq, d]
+                    # matmul broadcasts K over gqa dim without materializing.
+                    # Saves 6x memory vs repeat_interleave for Qwen3.6 (ratio=6).
+                    q_c = (q_flat[seq_start + qc_start:seq_start + qc_end]
+                           .float()
+                           .view(-1, self.num_kv_heads, gqa_ratio, self.head_size)
+                           .permute(1, 2, 0, 3))  # [kv_h, gqa, chunk, d]
 
-                attn_w = torch.matmul(q_c, k_s.transpose(-2, -1)) * self.scale
+                    # [kv_h, gqa, chunk, seq] via broadcast
+                    attn_w = torch.matmul(
+                        q_c, k_s.transpose(-2, -1)) * self.scale
 
-                qc_q_pos = torch.arange(qc_start, qc_end, device=query.device)
-                mask = k_pos.unsqueeze(0) > qc_q_pos.unsqueeze(1)
-                attn_w = attn_w.masked_fill(mask.unsqueeze(0), float("-inf"))
+                    qc_q_pos = torch.arange(qc_start, qc_end, device=query.device)
+                    mask = k_pos.unsqueeze(0) > qc_q_pos.unsqueeze(1)
+                    attn_w = attn_w.masked_fill(
+                        mask.unsqueeze(0).unsqueeze(0), float("-inf"))
 
-                attn_w = torch.softmax(attn_w, dim=-1)
-                out_c = torch.matmul(attn_w, v_s).to(orig_dtype)
+                    attn_w = torch.softmax(attn_w, dim=-1)
+                    # [kv_h, gqa, chunk, d]
+                    out_c = torch.matmul(attn_w, v_s).to(orig_dtype)
+                    # → [chunk, kv_h * gqa, d] = [chunk, num_heads, d]
+                    out_c = (out_c.permute(2, 0, 1, 3)
+                             .contiguous()
+                             .view(-1, self.num_heads, self.head_size))
+                    output[seq_start + qc_start:seq_start + qc_end] = out_c
+                else:
+                    # Non-GQA path (kv_heads == num_heads)
+                    q_c = q_flat[seq_start + qc_start:seq_start + qc_end] \
+                          .permute(1, 0, 2).float()
 
-                output[seq_start + qc_start:seq_start + qc_end] = (
-                    out_c.permute(1, 0, 2))
+                    attn_w = torch.matmul(
+                        q_c, k_s.transpose(-2, -1)) * self.scale
+
+                    qc_q_pos = torch.arange(qc_start, qc_end, device=query.device)
+                    mask = k_pos.unsqueeze(0) > qc_q_pos.unsqueeze(1)
+                    attn_w = attn_w.masked_fill(
+                        mask.unsqueeze(0), float("-inf"))
+
+                    attn_w = torch.softmax(attn_w, dim=-1)
+                    out_c = torch.matmul(attn_w, v_s).to(orig_dtype)
+
+                    output[seq_start + qc_start:seq_start + qc_end] = (
+                        out_c.permute(1, 0, 2))
 
             seq_start = seq_end