diff --git a/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py b/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py index f1d5bdabd..bdb52b1c3 100644 --- a/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py +++ b/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py @@ -16,7 +16,10 @@ from sglang.srt.layers.attention.fla.utils import ( CHUNK_SIZE = 64 -# This kernel handles K blocks in a for loop to minimize register spills +# This kernel handles K blocks in a for loop to minimize register spills. +# Time is the OUTER loop; K blocks are processed in two inner phases per step: +# Phase 1: store h to output, accumulate v_correction = sum_k(w_k @ h_k^T) +# Phase 2: update h = gate * h + k^T @ v_gated, save to scratch (initial_state) @triton.autotune( configs=[triton.Config({"BV": 64}, num_warps=8, num_stages=2)], key=["H", "K", "V", "BT", "USE_GK", "USE_INITIAL_STATE", "NT_BUCKET"], @@ -110,67 +113,104 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64_k_loop( if INPLACE_UPDATE: ht = ht + i_h * V * K - # Explicit K loop here to reduce register pressure - for k_start in range(0, K, 64): - # [BV, BK] - b_h1 = tl.zeros([BV, 64], dtype=tl.float32) + # main recurrence — time is the outer loop + for i_t in range(NT): + ######################################################################## + # Phase 1: store h to output, compute v_new = u - sum_k(w_k @ h_k^T) + ######################################################################## + b_v_corr = tl.zeros([BT, BV], dtype=tl.float32) + for k_blk in range(0, K, 64): + # Load h: from initial_state (i_t==0) or scratch (i_t>0) + if i_t == 0: + if USE_INITIAL_STATE: + p_hs = tl.make_block_ptr( + h0, (V, K), (K, 1), (i_v * BV, k_blk), (BV, 64), (1, 0) + ) + b_h = tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32) + else: + b_h = tl.zeros([BV, 64], dtype=tl.float32) + else: + p_hs = tl.make_block_ptr( + ht, (V, K), (K, 1), (i_v * BV, k_blk), (BV, 64), (1, 0) + ) + b_h = tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32) - # load initial state - if USE_INITIAL_STATE: - p_h0_1 = tl.make_block_ptr( - h0, (V, K), (K, 1), (i_v * BV, k_start), (BV, 64), (1, 0) - ) - b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32) - - # main recurrence - for i_t in range(NT): - p_h1 = tl.make_block_ptr( + # Store pre-update h to output + p_ho = tl.make_block_ptr( h + i_t * stride_h, (V, K), (K, 1), - (i_v * BV, k_start), + (i_v * BV, k_blk), (BV, 64), (1, 0), ) - tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_ho, b_h.to(p_ho.dtype.element_ty), boundary_check=(0, 1)) - b_w = w_desc.load([i_t * BT, k_start]) - b_v = tl.dot(b_w, tl.trans(b_h1).to(b_w.dtype)) - b_v = v_desc.load([i_t * BT, i_v * BV]) - b_v + # Accumulate correction: w_k @ h_k^T + b_w = w_desc.load([i_t * BT, k_blk]) + b_v_corr += tl.dot(b_w, tl.trans(b_h).to(b_w.dtype)) - if SAVE_NEW_VALUE: - v_new_desc.store([i_t * BT, i_v * BV], b_v.to(v_new.dtype.element_ty)) + # v_new = u - correction + b_v = v_desc.load([i_t * BT, i_v * BV]) - b_v_corr - last_idx = min((i_t + 1) * BT, T) - 1 - if USE_G: - b_g_last = tl.load(g + bos * H + last_idx * H + i_h) - p_g = tl.make_block_ptr( - g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + if SAVE_NEW_VALUE: + v_new_desc.store([i_t * BT, i_v * BV], b_v.to(v_new.dtype.element_ty)) + + # Apply gate to v + last_idx = min((i_t + 1) * BT, T) - 1 + if USE_G: + b_g_last = tl.load(g + bos * H + last_idx * H + i_h) + p_g = tl.make_block_ptr( + g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,) + ) + b_g = tl.load(p_g, boundary_check=(0,)) + b_v = b_v * tl.expand_dims(safe_exp(b_g_last - b_g), 1) + b_g_last = exp(b_g_last) + + b_v = b_v.to(k.dtype.element_ty) + + ######################################################################## + # Phase 2: reload h, apply gate, update h += k^T @ v, save to scratch + ######################################################################## + for k_blk in range(0, K, 64): + # Reload h (same source as Phase 1) + if i_t == 0: + if USE_INITIAL_STATE: + p_hs = tl.make_block_ptr( + h0, (V, K), (K, 1), (i_v * BV, k_blk), (BV, 64), (1, 0) + ) + b_h = tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32) + else: + b_h = tl.zeros([BV, 64], dtype=tl.float32) + else: + p_hs = tl.make_block_ptr( + ht, (V, K), (K, 1), (i_v * BV, k_blk), (BV, 64), (1, 0) ) - b_g = tl.load(p_g, boundary_check=(0,)) - b_v = b_v * safe_exp(b_g_last - b_g)[:, None] - b_g_last = exp(b_g_last) - b_h1 = b_h1 * b_g_last + b_h = tl.load(p_hs, boundary_check=(0, 1)).to(tl.float32) + + # Gate decay on h + if USE_G: + b_h = b_h * b_g_last if USE_GK: - o_k1 = tl.arange(0, 64) + k_start + o_k1 = tl.arange(0, 64) + k_blk b_gk_last1 = tl.load( gk + (bos + last_idx) * H * K + i_h * K + o_k1, mask=(o_k1 < K), other=0.0, ) - b_h1 *= exp(b_gk_last1)[None, :] - b_v = b_v.to(k.dtype.element_ty) + b_h *= tl.expand_dims(exp(b_gk_last1), 0) - b_k = tl.trans(k_desc.load([i_t * BT, k_start])) - b_h1 += tl.trans(tl.dot(b_k, b_v)) + # Delta update: h += k^T @ v + b_k = tl.trans(k_desc.load([i_t * BT, k_blk])) + b_h += tl.trans(tl.dot(b_k, b_v)) - # epilogue - if INPLACE_UPDATE: - p_ht = tl.make_block_ptr( - ht, (V, K), (K, 1), (i_v * BV, k_start), (BV, 64), (1, 0) - ) - tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1)) + # Save updated h to scratch (initial_state) for next time step + if INPLACE_UPDATE: + p_hs = tl.make_block_ptr( + ht, (V, K), (K, 1), (i_v * BV, k_blk), (BV, 64), (1, 0) + ) + tl.store(p_hs, b_h.to(p_hs.dtype.element_ty), boundary_check=(0, 1)) def chunk_gated_delta_rule_fwd_h( diff --git a/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_fwd.py b/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_fwd.py index 3e54712dd..c2299a2f0 100644 --- a/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_fwd.py +++ b/python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_fwd.py @@ -233,6 +233,16 @@ def chunk_gated_delta_rule_fwd_kkt_solve_kernel_low_reg( ) tl.store(p_Ai_ij, b_Ai_ij.to(A.dtype.element_ty), boundary_check=(0, 1)) + # Clean up scratch slots: Pass 2 stored raw A_ij blocks in the upper-triangular + # part of row i_tc0 (cols BC..3*BC). These must be zeroed because + # recompute_w_u_fwd reads the full BT×BT block. + b_zero = tl.zeros([BC, BC], dtype=tl.float32) + for sc in tl.static_range(1, BT // BC): + p_scratch = tl.make_block_ptr( + A, (T, BT), (H * BT, 1), (i_tc0, sc * BC), (BC, BC), (1, 0) + ) + tl.store(p_scratch, b_zero.to(A.dtype.element_ty), boundary_check=(0, 1)) + def chunk_gated_delta_rule_fwd_intra( k: torch.Tensor, diff --git a/test/registered/attention/test_chunk_gated_delta_rule.py b/test/registered/attention/test_chunk_gated_delta_rule.py index d496d679a..c487f2577 100644 --- a/test/registered/attention/test_chunk_gated_delta_rule.py +++ b/test/registered/attention/test_chunk_gated_delta_rule.py @@ -210,6 +210,16 @@ class TestChunkGatedDeltaRule(unittest.TestCase): def test_large_pool(self): self._check_shape(B=4, T_per_seq=128, H=16, K=128, V=128, pool_size=512) + # ------------------------------------------------------------------ + # Long prompts (many chunks; regression test for cross-chunk errors) + # ------------------------------------------------------------------ + def test_long_prompt(self): + for B, T_per_seq in [(1, 1024), (1, 1536), (1, 2048), (2, 1024)]: + with self.subTest(T_per_seq=T_per_seq): + self._check_shape( + B=B, T_per_seq=T_per_seq, H=16, K=128, V=128, pool_size=32 + ) + # ------------------------------------------------------------------ # Combined stress # ------------------------------------------------------------------