From be630106b22b7060ade1e842b704a83772704a24 Mon Sep 17 00:00:00 2001
From: project6 <dev@project6.com>
Date: Fri, 7 Aug 2026 09:08:53 +0000
Subject: [PATCH] =?UTF-8?q?perf(deltanet):=20CCCL=20block=5Fscan=20RAKING?=
 =?UTF-8?q?=5FMEMOIZE=20=E2=80=94=20precompute=20all=20exp()=20outside=20l?=
 =?UTF-8?q?oop?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Full translation of cub/block/block_scan.cuh BLOCK_SCAN_RAKING_MEMOIZE strategy
to _torch_chunk_gated_delta_rule cross-chunk scan loop:

CCCL RAKING_MEMOIZE: 'preserve upsweep segment values in registers while
performing warp-synchronous scan, allowing downsweep not to re-read from
shared memory.'

Translation: precompute g_exp_full, g_last_exp, g_diff_exp tensors outside
the sequential cross-chunk loop. Loop body now uses indexed lookups into
precomputed tensors instead of calling exp() 3 times per chunk iteration.

For seq_len=100K with chunk_size=64: 1562 chunks × 3 exp() = 4686 exp() calls
eliminated from the hot loop. Replaced with 3 bulk exp() + tensor indexing.

Memory tradeoff (same as RAKING_MEMOIZE's register pressure):
+3 tensors of shape (batch, heads, num_chunks, chunk_size) float32
= 3 × 1 × 6 × 1562 × 64 × 4B ≈ 7MB (negligible vs 16GB model weights)

Also inherits overflow_cast protection: g is clamped to [-20,20] before
exp(), so precomputed values stay in safe float32 range.
---
 qwen3_6_scripts/qwen3_5.py | 23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py
index 0afe9120..e1adbddf 100644
--- a/qwen3_6_scripts/qwen3_5.py
+++ b/qwen3_6_scripts/qwen3_5.py
@@ -263,16 +263,31 @@ def _torch_chunk_gated_delta_rule(
         torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device),
         diagonal=1)
 
-    for i in range(total_len // chunk_size):
+    # CCCL block_scan.cuh BLOCK_SCAN_RAKING_MEMOIZE strategy:
+    # Precompute all per-chunk exp values outside the loop, eliminating
+    # redundant exp() inside the sequential cross-chunk scan.
+    # RAKING_MEMOIZE: "preserve upsweep segment values in registers while
+    # performing warp-synchronous scan, allowing downsweep not to re-read."
+    num_chunks = total_len // chunk_size
+    # g shape: (batch, heads, num_chunks, chunk_size)
+    # g_exp_full[i] = exp(g[:,:,i,:]) for attn_inter computation
+    g_exp_full = g.exp()  # (batch, heads, num_chunks, chunk_size)
+    # g_last_exp[i] = exp(g[:,:,i,-1]) for state decay
+    g_last_exp = g_exp_full[:, :, :, -1]  # (batch, heads, num_chunks)
+    # g_diff_exp[i] = exp(g[:,:,i,-1] - g[:,:,i,:]) for k_i weighting
+    g_diff_exp = (g[:, :, :, -1:] - g).exp()  # (batch, heads, num_chunks, chunk_size)
+
+    for i in range(num_chunks):
         q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i]
         attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0)
         v_prime = k_cumdecay[:, :, i] @ last_state
         v_new = v_i - v_prime
-        attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_state
+        # Use precomputed exp (MEMOIZE: no redundant exp in loop body)
+        attn_inter = (q_i * g_exp_full[:, :, i, :, None]) @ last_state
         core_out[:, :, i] = attn_inter + attn_i @ v_new
         last_state = (
-            last_state * g[:, :, i, -1, None, None].exp()
-            + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None])
+            last_state * g_last_exp[:, :, i, None, None]
+            + (k_i * g_diff_exp[:, :, i, :, None])
             .transpose(-1, -2) @ v_new
         )