perf(deltanet): CCCL block_scan RAKING_MEMOIZE — precompute all exp() outside loop
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.
This commit is contained in:
@@ -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
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user