From 0a697f58717de9551d074e89d2843a00ab58e255 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 9 Aug 2026 10:44:43 +0000 Subject: [PATCH] arch(scan): dispatch_scan.cuh Phase 1/Phase 2 separation in GDN chunk loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Direct translation of CCCL dispatch_scan.cuh (1469 lines) architecture: CCCL dispatch_scan has two kernels: 1. DeviceScanInitKernel — initializes tile_state (parallelizable) 2. DeviceScanKernel — sequential scan using tile_state propagation Our _torch_chunk_gated_delta_rule now separates: Phase 1 (init, parallelizable): pre-compute ALL chunk-local attn matrices attn_i[c] = q[c] @ k[c].T * decay[c] — does NOT depend on state Also pre-compute g.exp() and clamped g once, outside loop Phase 2 (scan, sequential): only state-dependent ops in the loop v_prime, v_new, attn_inter, core_out, state update This matches CCCL's insight: everything that doesn't need tile_state should be computed before the scan kernel, not interleaved with it. --- qwen3_6_scripts/qwen3_5.py | 33 +++++++++++++++++++++++++-------- 1 file changed, 25 insertions(+), 8 deletions(-) diff --git a/qwen3_6_scripts/qwen3_5.py b/qwen3_6_scripts/qwen3_5.py index a5f8edef..6f13eead 100644 --- a/qwen3_6_scripts/qwen3_5.py +++ b/qwen3_6_scripts/qwen3_5.py @@ -218,17 +218,34 @@ 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): - q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] - attn_i = (_ix_matmul(q_i, k_i.transpose(-1, -2)) * decay_mask[:, :, i]).masked_fill_(mask_upper2, 0) + # dispatch_scan.cuh Phase 1: pre-compute ALL chunk-local attention matrices + # outside the state loop. attn_i[c] only depends on q, k, decay_mask — NOT state. + # This is the CCCL "init kernel" pattern: compute everything possible + # before the sequential scan kernel that needs tile_state propagation. + num_chunks = total_len // chunk_size + attn_i_all = torch.empty( + batch, num_heads, num_chunks, chunk_size, chunk_size, + dtype=value.dtype, device=value.device) + for i in range(num_chunks): + attn_i_all[:, :, i] = ( + _ix_matmul(query[:, :, i], key[:, :, i].transpose(-1, -2)) + * decay_mask[:, :, i] + ).masked_fill_(mask_upper2, 0) + + # dispatch_scan.cuh Phase 2: sequential state propagation (scan kernel). + # Only state-dependent ops remain in this loop. + g_exp_cache = g.clamp(-20, 20).exp() # pre-compute once + g_clamped = g.clamp(-20, 20) # keep raw clamped g for difference computation + for i in range(num_chunks): v_prime = _ix_matmul(k_cumdecay[:, :, i], last_state) - v_new = v_i - v_prime - attn_inter = _ix_matmul(q_i * g[:, :, i, :, None].clamp(-20, 20).exp(), last_state) - core_out[:, :, i] = attn_inter + _ix_matmul(attn_i, v_new) + v_new = value[:, :, i] - v_prime + attn_inter = _ix_matmul(query[:, :, i] * g_exp_cache[:, :, i, :, None], last_state) + core_out[:, :, i] = attn_inter + _ix_matmul(attn_i_all[:, :, i], v_new) + # State update uses difference form: exp(g[-1] - g[:]) to avoid division last_state = ( - last_state * g[:, :, i, -1, None, None].clamp(-20, 20).exp() + last_state * g_exp_cache[:, :, i, -1, None, None] + _ix_matmul( - (k_i * (g[:, :, i, -1, None] - g[:, :, i]).clamp(-20, 20).exp()[..., None]) + (key[:, :, i] * (g_clamped[:, :, i, -1, None] - g_clamped[:, :, i]).exp()[..., None]) .transpose(-1, -2), v_new) )