diff --git a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py index 0c7f80f42..5feb33906 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py +++ b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py @@ -13,22 +13,30 @@ from sglang.srt.layers.attention.fla.index import ( prepare_chunk_offsets, ) from sglang.srt.layers.attention.fla.op import exp, safe_exp -from sglang.srt.layers.attention.fla.utils import is_nvidia_hopper +from sglang.srt.layers.attention.fla.utils import ( + autotune_cache_kwargs, + is_nvidia_hopper, +) NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8, 16] CHUNK_SIZE = 64 -# @triton.autotune( -# configs=[ -# triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages) -# for num_warps in [2, 4] -# for num_stages in [2, 3, 4] -# for BV in [32, 64] -# ], -# key=["H", "K", "V", "BT", "USE_G"], -# use_cuda_graph=use_cuda_graph, -# ) +@triton.autotune( + # Single hardcoded config. The kernel writes ht (final state) back into + # initial_state in-place; with multiple configs, triton's autotune benchmark + # phase invokes the kernel many times for timing and corrupts the cache pool, + # producing silently wrong output on the first user request. Restoring via + # `restore_value=["initial_state"]` works for unit tests but OOMs on + # production-scale models (e.g. Kimi-Linear-48B at default mem_fraction) + # because cloning the cache pool for each benchmark exceeds available memory. + # NT_BUCKET is kept in the autotune key for forward-compatibility (allows + # future per-bucket configs once the kernel is refactored to write final + # state to a separate output buffer). + configs=[triton.Config({"BV": 32}, num_warps=4, num_stages=2)], + key=["H", "K", "V", "BT", "USE_GK", "NT_BUCKET"], + **autotune_cache_kwargs, +) @triton.jit(do_not_specialize=["T"]) def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( k, @@ -55,6 +63,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64( INPLACE_UPDATE: tl.constexpr, SAVE_NEW_VALUE: tl.constexpr, IS_VARLEN: tl.constexpr, + NT_BUCKET: tl.constexpr, ): i_v, i_nh = tl.program_id(0), tl.program_id(1) i_n, i_h = i_nh // H, i_nh % H @@ -325,14 +334,12 @@ def chunk_gated_delta_rule_fwd_h( K=K, V=V, BT=BT, - BV=32, USE_G=g is not None, USE_GK=gk is not None, USE_INITIAL_STATE=initial_state is not None, INPLACE_UPDATE=True, SAVE_NEW_VALUE=v_new is not None, IS_VARLEN=cu_seqlens is not None, - num_warps=4, - num_stages=2, + NT_BUCKET=(0 if NT <= 32 else (1 if NT <= 128 else 2)), ) return h, v_new diff --git a/python/sglang/srt/layers/attention/fla/chunk_intra.py b/python/sglang/srt/layers/attention/fla/chunk_intra.py index 00b62727c..a8d35629e 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_intra.py +++ b/python/sglang/srt/layers/attention/fla/chunk_intra.py @@ -11,7 +11,7 @@ from sglang.srt.layers.attention.fla.chunk_intra_token_parallel import ( from sglang.srt.layers.attention.fla.index import ( prepare_chunk_indices, ) -from sglang.srt.layers.attention.fla.op import exp2, gather +from sglang.srt.layers.attention.fla.op import exp, exp2, gather from sglang.srt.layers.attention.fla.utils import ( autotune_cache_kwargs, is_gather_supported, @@ -36,11 +36,11 @@ else: ) @triton.autotune( configs=[ - triton.Config({"BK": BK}, num_warps=num_warps) + triton.Config({"BK": BK, "BV": 64}, num_warps=num_warps) for BK in [32, 64] for num_warps in [1, 2, 4] ], - key=["H", "K", "BC"], + key=["H", "K", "BC", "V", "FUSE_RECOMPUTE", "FUSE_DIAGONAL"], **autotune_cache_kwargs, ) @triton.jit(do_not_specialize=["T"]) @@ -53,16 +53,24 @@ def chunk_kda_fwd_kernel_inter_solve_fused( Akkd, Akk, scale, + v_in, + w_out, + u_out, + kg_out, cu_seqlens, chunk_indices, T, H: tl.constexpr, K: tl.constexpr, + V: tl.constexpr, BT: tl.constexpr, BC: tl.constexpr, BK: tl.constexpr, + BV: tl.constexpr, IS_VARLEN: tl.constexpr, USE_SAFE_GATE: tl.constexpr, + FUSE_RECOMPUTE: tl.constexpr, + FUSE_DIAGONAL: tl.constexpr, ): """ Fused kernel: compute inter-subchunk Akk + solve_tril in one pass. @@ -125,8 +133,19 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Aqk32 = tl.zeros([BC, BC], dtype=tl.float32) b_Akk32 = tl.zeros([BC, BC], dtype=tl.float32) + if FUSE_DIAGONAL: + b_Aqk_d0 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk_d0 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk_d1 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk_d1 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk_d2 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk_d2 = tl.zeros([BC, BC], dtype=tl.float32) + b_Aqk_d3 = tl.zeros([BC, BC], dtype=tl.float32) + b_Akk_d3 = tl.zeros([BC, BC], dtype=tl.float32) + m_tc0 = (i_tc0 + o_i) < T + ################################################################################ - # off-diagonal blocks + # off-diagonal blocks (+ optional diagonal blocks) ################################################################################ for i_k in range(tl.cdiv(K, BK)): o_k = i_k * BK + tl.arange(0, BK) @@ -141,6 +160,19 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_k0 = tl.load(p_k0, boundary_check=(0, 1)).to(tl.float32) b_g0 = tl.load(p_g0, boundary_check=(0, 1)).to(tl.float32) + if FUSE_DIAGONAL: + p_q0 = tl.make_block_ptr( + q, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) + ) + b_q0 = tl.load(p_q0, boundary_check=(0, 1)).to(tl.float32) + b_gn0 = tl.load(g + i_tc0 * H * K + o_k, mask=m_k, other=0).to(tl.float32) + b_gm0 = tl.clamp(b_g0 - b_gn0[None, :], -126.0, 126.0) + b_gq0 = tl.where(m_tc0[:, None], exp2(b_gm0), 0.0) + b_gk0 = tl.where(m_tc0[:, None], exp2(-b_gm0), 0.0) + b_kgt_d0 = tl.trans(b_k0 * b_gk0) + b_Aqk_d0 += tl.dot(b_q0 * b_gq0, b_kgt_d0) + b_Akk_d0 += tl.dot(b_k0 * b_gq0, b_kgt_d0) + if i_tc1 < T: p_q1 = tl.make_block_ptr( q, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) @@ -160,10 +192,19 @@ def chunk_kda_fwd_kernel_inter_solve_fused( # [BC, BK] b_gqn = tl.where(m_tc1[:, None], exp2(b_g1 - b_gn1[None, :]), 0) # [BK, BC] - b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0)) + b_kgt = tl.trans(b_k0 * exp2(b_gn1[None, :] - b_g0)).to(tl.bfloat16) # [BC, BC] - b_Aqk10 += tl.dot(b_q1 * b_gqn, b_kgt) - b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt) + b_qg1 = (b_q1 * b_gqn).to(tl.bfloat16) + b_kg1 = (b_k1 * b_gqn).to(tl.bfloat16) + b_Aqk10 += tl.dot(b_qg1, b_kgt) + b_Akk10 += tl.dot(b_kg1, b_kgt) + + if FUSE_DIAGONAL: + b_gm1_d = tl.clamp(b_gn1[None, :] - b_g1, -126.0, 126.0) + b_gk1_d = tl.where(m_tc1[:, None], exp2(b_gm1_d), 0.0) + b_kgt_d1 = tl.trans(b_k1 * b_gk1_d) + b_Aqk_d1 += tl.dot(b_q1 * b_gqn, b_kgt_d1) + b_Akk_d1 += tl.dot(b_k1 * b_gqn, b_kgt_d1) if i_tc2 < T: p_q2 = tl.make_block_ptr( @@ -185,18 +226,25 @@ def chunk_kda_fwd_kernel_inter_solve_fused( ) # [BC, BK] b_gqn2 = tl.where(m_tc2[:, None], exp2(b_g2 - b_gn2[None, :]), 0) - b_qg2 = b_q2 * b_gqn2 - b_kg2 = b_k2 * b_gqn2 + b_qg2 = (b_q2 * b_gqn2).to(tl.bfloat16) + b_kg2 = (b_k2 * b_gqn2).to(tl.bfloat16) # [BK, BC] - b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)) + b_kgt = tl.trans(b_k0 * exp2(b_gn2[None, :] - b_g0)).to(tl.bfloat16) b_Aqk20 += tl.dot(b_qg2, b_kgt) b_Akk20 += tl.dot(b_kg2, b_kgt) # [BC, BC] - b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)) + b_kgt = tl.trans(b_k1 * exp2(b_gn2[None, :] - b_g1)).to(tl.bfloat16) # [BC, BC] b_Aqk21 += tl.dot(b_qg2, b_kgt) b_Akk21 += tl.dot(b_kg2, b_kgt) + if FUSE_DIAGONAL: + b_gm2_d = tl.clamp(b_gn2[None, :] - b_g2, -126.0, 126.0) + b_gk2_d = tl.where(m_tc2[:, None], exp2(b_gm2_d), 0.0) + b_kgt_d2 = tl.trans(b_k2 * b_gk2_d) + b_Aqk_d2 += tl.dot(b_q2 * b_gqn2, b_kgt_d2) + b_Akk_d2 += tl.dot(b_k2 * b_gqn2, b_kgt_d2) + if i_tc3 < T: p_q3 = tl.make_block_ptr( q, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) @@ -217,24 +265,31 @@ def chunk_kda_fwd_kernel_inter_solve_fused( ) # [BC, BK] b_gqn3 = tl.where(m_tc3[:, None], exp2(b_g3 - b_gn3[None, :]), 0) - b_qg3 = b_q3 * b_gqn3 - b_kg3 = b_k3 * b_gqn3 + b_qg3 = (b_q3 * b_gqn3).to(tl.bfloat16) + b_kg3 = (b_k3 * b_gqn3).to(tl.bfloat16) # [BK, BC] - b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)) + b_kgt = tl.trans(b_k0 * exp2(b_gn3[None, :] - b_g0)).to(tl.bfloat16) # [BC, BC] b_Aqk30 += tl.dot(b_qg3, b_kgt) b_Akk30 += tl.dot(b_kg3, b_kgt) # [BK, BC] - b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)) + b_kgt = tl.trans(b_k1 * exp2(b_gn3[None, :] - b_g1)).to(tl.bfloat16) # [BC, BC] b_Aqk31 += tl.dot(b_qg3, b_kgt) b_Akk31 += tl.dot(b_kg3, b_kgt) # [BK, BC] - b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)) + b_kgt = tl.trans(b_k2 * exp2(b_gn3[None, :] - b_g2)).to(tl.bfloat16) # [BC, BC] b_Aqk32 += tl.dot(b_qg3, b_kgt) b_Akk32 += tl.dot(b_kg3, b_kgt) + if FUSE_DIAGONAL: + b_gm3_d = tl.clamp(b_gn3[None, :] - b_g3, -126.0, 126.0) + b_gk3_d = tl.where(m_tc3[:, None], exp2(b_gm3_d), 0.0) + b_kgt_d3 = tl.trans(b_k3 * b_gk3_d) + b_Aqk_d3 += tl.dot(b_q3 * b_gqn3, b_kgt_d3) + b_Akk_d3 += tl.dot(b_k3 * b_gqn3, b_kgt_d3) + ################################################################################ # save off-diagonal Aqk blocks and prepare Akk ################################################################################ @@ -299,28 +354,114 @@ def chunk_kda_fwd_kernel_inter_solve_fused( b_Akk31 = b_Akk31 * b_b3[:, None] b_Akk32 = b_Akk32 * b_b3[:, None] - p_Akk00 = tl.make_block_ptr( - Akkd, (T, BC), (H * BC, 1), (i_tc0, 0), (BC, BC), (1, 0) - ) - p_Akk11 = tl.make_block_ptr( - Akkd, (T, BC), (H * BC, 1), (i_tc1, 0), (BC, BC), (1, 0) - ) - p_Akk22 = tl.make_block_ptr( - Akkd, (T, BC), (H * BC, 1), (i_tc2, 0), (BC, BC), (1, 0) - ) - p_Akk33 = tl.make_block_ptr( - Akkd, (T, BC), (H * BC, 1), (i_tc3, 0), (BC, BC), (1, 0) - ) - b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) - b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) - b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) - b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) + if FUSE_DIAGONAL: + m_Aqk_diag = o_i[:, None] >= o_i[None, :] + m_Akk_diag = o_i[:, None] > o_i[None, :] + + b_Aqk_d0 = tl.where(m_Aqk_diag, b_Aqk_d0, 0.0) + b_Akk_d0 = tl.where(m_Akk_diag, b_Akk_d0, 0.0) + b_Aqk_d1 = tl.where(m_Aqk_diag, b_Aqk_d1, 0.0) + b_Akk_d1 = tl.where(m_Akk_diag, b_Akk_d1, 0.0) + b_Aqk_d2 = tl.where(m_Aqk_diag, b_Aqk_d2, 0.0) + b_Akk_d2 = tl.where(m_Akk_diag, b_Akk_d2, 0.0) + b_Aqk_d3 = tl.where(m_Aqk_diag, b_Aqk_d3, 0.0) + b_Akk_d3 = tl.where(m_Akk_diag, b_Akk_d3, 0.0) + + p_Aqk_d0 = tl.make_block_ptr( + Aqk, (T, BT), (H * BT, 1), (i_tc0, 0), (BC, BC), (1, 0) + ) + p_Aqk_d1 = tl.make_block_ptr( + Aqk, (T, BT), (H * BT, 1), (i_tc1, BC), (BC, BC), (1, 0) + ) + p_Aqk_d2 = tl.make_block_ptr( + Aqk, (T, BT), (H * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0) + ) + p_Aqk_d3 = tl.make_block_ptr( + Aqk, (T, BT), (H * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0) + ) + tl.store( + p_Aqk_d0, (b_Aqk_d0 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + tl.store( + p_Aqk_d1, (b_Aqk_d1 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + tl.store( + p_Aqk_d2, (b_Aqk_d2 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + tl.store( + p_Aqk_d3, (b_Aqk_d3 * scale).to(Aqk.dtype.element_ty), boundary_check=(0, 1) + ) + + p_bd0 = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_tc0,), (BC,), (0,) + ) + p_bd1 = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_tc1,), (BC,), (0,) + ) + p_bd2 = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_tc2,), (BC,), (0,) + ) + p_bd3 = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_tc3,), (BC,), (0,) + ) + b_bd0 = tl.load(p_bd0, boundary_check=(0,)).to(tl.float32) + b_bd1 = tl.load(p_bd1, boundary_check=(0,)).to(tl.float32) + b_bd2 = tl.load(p_bd2, boundary_check=(0,)).to(tl.float32) + b_bd3 = tl.load(p_bd3, boundary_check=(0,)).to(tl.float32) + b_Akk_d0 = b_Akk_d0 * b_bd0[:, None] + b_Akk_d1 = b_Akk_d1 * b_bd1[:, None] + b_Akk_d2 = b_Akk_d2 * b_bd2[:, None] + b_Akk_d3 = b_Akk_d3 * b_bd3[:, None] + + p_Akkd00 = tl.make_block_ptr( + Akkd, (T, BC), (H * BC, 1), (i_tc0, 0), (BC, BC), (1, 0) + ) + p_Akkd11 = tl.make_block_ptr( + Akkd, (T, BC), (H * BC, 1), (i_tc1, 0), (BC, BC), (1, 0) + ) + p_Akkd22 = tl.make_block_ptr( + Akkd, (T, BC), (H * BC, 1), (i_tc2, 0), (BC, BC), (1, 0) + ) + p_Akkd33 = tl.make_block_ptr( + Akkd, (T, BC), (H * BC, 1), (i_tc3, 0), (BC, BC), (1, 0) + ) + tl.store(p_Akkd00, b_Akk_d0.to(Akkd.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akkd11, b_Akk_d1.to(Akkd.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akkd22, b_Akk_d2.to(Akkd.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akkd33, b_Akk_d3.to(Akkd.dtype.element_ty), boundary_check=(0, 1)) + + b_Ai00 = b_Akk_d0 + b_Ai11 = b_Akk_d1 + b_Ai22 = b_Akk_d2 + b_Ai33 = b_Akk_d3 + else: + p_Akk00 = tl.make_block_ptr( + Akkd, (T, BC), (H * BC, 1), (i_tc0, 0), (BC, BC), (1, 0) + ) + p_Akk11 = tl.make_block_ptr( + Akkd, (T, BC), (H * BC, 1), (i_tc1, 0), (BC, BC), (1, 0) + ) + p_Akk22 = tl.make_block_ptr( + Akkd, (T, BC), (H * BC, 1), (i_tc2, 0), (BC, BC), (1, 0) + ) + p_Akk33 = tl.make_block_ptr( + Akkd, (T, BC), (H * BC, 1), (i_tc3, 0), (BC, BC), (1, 0) + ) + b_Ai00 = tl.load(p_Akk00, boundary_check=(0, 1)).to(tl.float32) + b_Ai11 = tl.load(p_Akk11, boundary_check=(0, 1)).to(tl.float32) + b_Ai22 = tl.load(p_Akk22, boundary_check=(0, 1)).to(tl.float32) + b_Ai33 = tl.load(p_Akk33, boundary_check=(0, 1)).to(tl.float32) ################################################################################ # forward substitution on diagonals + # Diagonal blocks are RAW (need substitution) when: + # - FUSE_DIAGONAL=True: blocks were computed fresh above as gated k·k. + # - FUSE_DIAGONAL=False with USE_SAFE_GATE=False: token_parallel wrote raw. + # They are pre-inverted only by the safe_gate diagonal kernel + # (USE_SAFE_GATE=True, FUSE_DIAGONAL=False). ################################################################################ - if not USE_SAFE_GATE: + if FUSE_DIAGONAL or not USE_SAFE_GATE: m_A = o_i[:, None] > o_i[None, :] m_I = o_i[:, None] == o_i[None, :] @@ -397,42 +538,245 @@ def chunk_kda_fwd_kernel_inter_solve_fused( ) ################################################################################ - # store full Akk_inv to Akk + # Output: store Akk_inv OR compute w, u, kg from registers ################################################################################ - p_Akk00 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc0, 0), (BC, BC), (1, 0)) - p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0)) - p_Akk11 = tl.make_block_ptr( - Akk, (T, BT), (H * BT, 1), (i_tc1, BC), (BC, BC), (1, 0) - ) - p_Akk20 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc2, 0), (BC, BC), (1, 0)) - p_Akk21 = tl.make_block_ptr( - Akk, (T, BT), (H * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) - ) - p_Akk22 = tl.make_block_ptr( - Akk, (T, BT), (H * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0) - ) - p_Akk30 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc3, 0), (BC, BC), (1, 0)) - p_Akk31 = tl.make_block_ptr( - Akk, (T, BT), (H * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) - ) - p_Akk32 = tl.make_block_ptr( - Akk, (T, BT), (H * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) - ) - p_Akk33 = tl.make_block_ptr( - Akk, (T, BT), (H * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0) - ) + if FUSE_RECOMPUTE: + # Cast A-inverse sub-blocks to input dtype for dot products + b_Ai00_h = b_Ai00.to(k.dtype.element_ty) + b_Ai10_h = b_Ai10.to(k.dtype.element_ty) + b_Ai11_h = b_Ai11.to(k.dtype.element_ty) + b_Ai20_h = b_Ai20.to(k.dtype.element_ty) + b_Ai21_h = b_Ai21.to(k.dtype.element_ty) + b_Ai22_h = b_Ai22.to(k.dtype.element_ty) + b_Ai30_h = b_Ai30.to(k.dtype.element_ty) + b_Ai31_h = b_Ai31.to(k.dtype.element_ty) + b_Ai32_h = b_Ai32.to(k.dtype.element_ty) + b_Ai33_h = b_Ai33.to(k.dtype.element_ty) - tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) - tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + # Load beta for all 4 sub-chunks + p_b0 = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_tc0,), (BC,), (0,) + ) + b_b0 = tl.load(p_b0, boundary_check=(0,)).to(tl.float32) + p_b1r = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_tc1,), (BC,), (0,) + ) + b_b1r = tl.load(p_b1r, boundary_check=(0,)).to(tl.float32) + p_b2r = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_tc2,), (BC,), (0,) + ) + b_b2r = tl.load(p_b2r, boundary_check=(0,)).to(tl.float32) + p_b3r = tl.make_block_ptr( + beta + bos * H + i_h, (T,), (H,), (i_tc3,), (BC,), (0,) + ) + b_b3r = tl.load(p_b3r, boundary_check=(0,)).to(tl.float32) + + # ---- u = A_inv @ (v * beta) ---- + v_base = v_in + (bos * H + i_h) * V + u_base = u_out + (bos * H + i_h) * V + for i_v in range(tl.cdiv(V, BV)): + p_v0 = tl.make_block_ptr( + v_base, (T, V), (H * V, 1), (i_tc0, i_v * BV), (BC, BV), (1, 0) + ) + p_v1 = tl.make_block_ptr( + v_base, (T, V), (H * V, 1), (i_tc1, i_v * BV), (BC, BV), (1, 0) + ) + p_v2 = tl.make_block_ptr( + v_base, (T, V), (H * V, 1), (i_tc2, i_v * BV), (BC, BV), (1, 0) + ) + p_v3 = tl.make_block_ptr( + v_base, (T, V), (H * V, 1), (i_tc3, i_v * BV), (BC, BV), (1, 0) + ) + + b_v0 = tl.load(p_v0, boundary_check=(0, 1)) + b_v1 = tl.load(p_v1, boundary_check=(0, 1)) + b_v2 = tl.load(p_v2, boundary_check=(0, 1)) + b_v3 = tl.load(p_v3, boundary_check=(0, 1)) + + b_vb0 = (b_v0 * b_b0[:, None]).to(b_v0.dtype) + b_vb1 = (b_v1 * b_b1r[:, None]).to(b_v1.dtype) + b_vb2 = (b_v2 * b_b2r[:, None]).to(b_v2.dtype) + b_vb3 = (b_v3 * b_b3r[:, None]).to(b_v3.dtype) + + b_u0 = tl.dot(b_Ai00_h, b_vb0) + b_u1 = tl.dot(b_Ai10_h, b_vb0) + tl.dot(b_Ai11_h, b_vb1) + b_u2 = ( + tl.dot(b_Ai20_h, b_vb0) + + tl.dot(b_Ai21_h, b_vb1) + + tl.dot(b_Ai22_h, b_vb2) + ) + b_u3 = ( + tl.dot(b_Ai30_h, b_vb0) + + tl.dot(b_Ai31_h, b_vb1) + + tl.dot(b_Ai32_h, b_vb2) + + tl.dot(b_Ai33_h, b_vb3) + ) + + p_u0 = tl.make_block_ptr( + u_base, (T, V), (H * V, 1), (i_tc0, i_v * BV), (BC, BV), (1, 0) + ) + p_u1 = tl.make_block_ptr( + u_base, (T, V), (H * V, 1), (i_tc1, i_v * BV), (BC, BV), (1, 0) + ) + p_u2 = tl.make_block_ptr( + u_base, (T, V), (H * V, 1), (i_tc2, i_v * BV), (BC, BV), (1, 0) + ) + p_u3 = tl.make_block_ptr( + u_base, (T, V), (H * V, 1), (i_tc3, i_v * BV), (BC, BV), (1, 0) + ) + tl.store(p_u0, b_u0.to(p_u0.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u1, b_u1.to(p_u1.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u2, b_u2.to(p_u2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_u3, b_u3.to(p_u3.dtype.element_ty), boundary_check=(0, 1)) + + # ---- w = A_inv @ (k * beta * exp(gk)), kg = k * exp(gn - gk) ---- + w_base = w_out + (bos * H + i_h) * K + kg_base = kg_out + (bos * H + i_h) * K + last_idx = min(i_t * BT + BT, T) - 1 + + for i_k in range(tl.cdiv(K, BK)): + o_k = i_k * BK + tl.arange(0, BK) + m_k = o_k < K + b_gn = tl.load(g + last_idx * H * K + o_k, mask=m_k, other=0.0).to( + tl.float32 + ) + + p_k0 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) + ) + p_k1 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) + p_k2 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) + p_k3 = tl.make_block_ptr( + k, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) + + p_gk0 = tl.make_block_ptr( + g, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) + ) + p_gk1 = tl.make_block_ptr( + g, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) + p_gk2 = tl.make_block_ptr( + g, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) + p_gk3 = tl.make_block_ptr( + g, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) + + b_k0r = tl.load(p_k0, boundary_check=(0, 1)) + b_k1r = tl.load(p_k1, boundary_check=(0, 1)) + b_k2r = tl.load(p_k2, boundary_check=(0, 1)) + b_k3r = tl.load(p_k3, boundary_check=(0, 1)) + + b_gk0r = tl.load(p_gk0, boundary_check=(0, 1)).to(tl.float32) + b_gk1r = tl.load(p_gk1, boundary_check=(0, 1)).to(tl.float32) + b_gk2r = tl.load(p_gk2, boundary_check=(0, 1)).to(tl.float32) + b_gk3r = tl.load(p_gk3, boundary_check=(0, 1)).to(tl.float32) + + b_kb0 = (b_k0r * b_b0[:, None] * exp(b_gk0r)).to(b_k0r.dtype) + b_kb1 = (b_k1r * b_b1r[:, None] * exp(b_gk1r)).to(b_k1r.dtype) + b_kb2 = (b_k2r * b_b2r[:, None] * exp(b_gk2r)).to(b_k2r.dtype) + b_kb3 = (b_k3r * b_b3r[:, None] * exp(b_gk3r)).to(b_k3r.dtype) + + b_w0 = tl.dot(b_Ai00_h, b_kb0) + b_w1 = tl.dot(b_Ai10_h, b_kb0) + tl.dot(b_Ai11_h, b_kb1) + b_w2 = ( + tl.dot(b_Ai20_h, b_kb0) + + tl.dot(b_Ai21_h, b_kb1) + + tl.dot(b_Ai22_h, b_kb2) + ) + b_w3 = ( + tl.dot(b_Ai30_h, b_kb0) + + tl.dot(b_Ai31_h, b_kb1) + + tl.dot(b_Ai32_h, b_kb2) + + tl.dot(b_Ai33_h, b_kb3) + ) + + p_w0 = tl.make_block_ptr( + w_base, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) + ) + p_w1 = tl.make_block_ptr( + w_base, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) + p_w2 = tl.make_block_ptr( + w_base, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) + p_w3 = tl.make_block_ptr( + w_base, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) + tl.store(p_w0, b_w0.to(p_w0.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w1, b_w1.to(p_w1.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w2, b_w2.to(p_w2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_w3, b_w3.to(p_w3.dtype.element_ty), boundary_check=(0, 1)) + + b_kg0 = b_k0r * exp(b_gn[None, :] - b_gk0r) + b_kg1 = b_k1r * exp(b_gn[None, :] - b_gk1r) + b_kg2 = b_k2r * exp(b_gn[None, :] - b_gk2r) + b_kg3 = b_k3r * exp(b_gn[None, :] - b_gk3r) + + p_kg0 = tl.make_block_ptr( + kg_base, (T, K), (H * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0) + ) + p_kg1 = tl.make_block_ptr( + kg_base, (T, K), (H * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0) + ) + p_kg2 = tl.make_block_ptr( + kg_base, (T, K), (H * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0) + ) + p_kg3 = tl.make_block_ptr( + kg_base, (T, K), (H * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0) + ) + tl.store(p_kg0, b_kg0.to(p_kg0.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_kg1, b_kg1.to(p_kg1.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_kg2, b_kg2.to(p_kg2.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_kg3, b_kg3.to(p_kg3.dtype.element_ty), boundary_check=(0, 1)) + else: + p_Akk00 = tl.make_block_ptr( + Akk, (T, BT), (H * BT, 1), (i_tc0, 0), (BC, BC), (1, 0) + ) + p_Akk10 = tl.make_block_ptr( + Akk, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0) + ) + p_Akk11 = tl.make_block_ptr( + Akk, (T, BT), (H * BT, 1), (i_tc1, BC), (BC, BC), (1, 0) + ) + p_Akk20 = tl.make_block_ptr( + Akk, (T, BT), (H * BT, 1), (i_tc2, 0), (BC, BC), (1, 0) + ) + p_Akk21 = tl.make_block_ptr( + Akk, (T, BT), (H * BT, 1), (i_tc2, BC), (BC, BC), (1, 0) + ) + p_Akk22 = tl.make_block_ptr( + Akk, (T, BT), (H * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0) + ) + p_Akk30 = tl.make_block_ptr( + Akk, (T, BT), (H * BT, 1), (i_tc3, 0), (BC, BC), (1, 0) + ) + p_Akk31 = tl.make_block_ptr( + Akk, (T, BT), (H * BT, 1), (i_tc3, BC), (BC, BC), (1, 0) + ) + p_Akk32 = tl.make_block_ptr( + Akk, (T, BT), (H * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0) + ) + p_Akk33 = tl.make_block_ptr( + Akk, (T, BT), (H * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0) + ) + + tl.store(p_Akk00, b_Ai00.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk10, b_Ai10.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk11, b_Ai11.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk20, b_Ai20.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk21, b_Ai21.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk22, b_Ai22.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk30, b_Ai30.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk31, b_Ai31.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk32, b_Ai32.to(Akk.dtype.element_ty), boundary_check=(0, 1)) + tl.store(p_Akk33, b_Ai33.to(Akk.dtype.element_ty), boundary_check=(0, 1)) @triton.heuristics( @@ -573,8 +917,11 @@ def chunk_kda_fwd_intra( chunk_indices: torch.LongTensor | None = None, safe_gate: bool = False, disable_recompute: bool = False, + fuse_recompute: bool = False, + fuse_diagonal: bool = False, ): B, T, H, K = k.shape + V = v.shape[-1] BT = chunk_size BC = 16 if chunk_indices is None and cu_seqlens is not None: @@ -582,51 +929,86 @@ def chunk_kda_fwd_intra( NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) NC = triton.cdiv(BT, BC) - Aqk = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype) - # Akk must be zero-initialized - kernel only writes lower triangular - Akk = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype) - # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril) - Akkd = torch.zeros(B, T, H, BC, device=k.device, dtype=torch.float32) + if fuse_diagonal: + Aqk = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype) + else: + Aqk = torch.empty(B, T, H, BT, device=k.device, dtype=k.dtype) + Akkd = torch.empty(B, T, H, BC, device=k.device, dtype=torch.float32) - # Step 1: Run token_parallel first to compute diagonal blocks into Akkd (fp32) - # Step 1: compute diagonal blocks into Akk_diag (fp32) - if safe_gate: - grid = (NT, NC, B * H) - BK = triton.next_power_of_2(K) - chunk_kda_fwd_kernel_intra_sub_chunk[grid]( + # Step 1: compute diagonal blocks into Akkd (fp32) + # When fuse_diagonal=True, diagonal blocks are computed inside inter_solve + if not fuse_diagonal: + if safe_gate: + grid = (NT, NC, B * H) + BK = triton.next_power_of_2(K) + chunk_kda_fwd_kernel_intra_sub_chunk[grid]( + q=q, + k=k, + g=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + T=T, + H=H, + K=K, + BT=BT, + BC=BC, + BK=BK, + USE_GATHER=is_gather_supported, + ) + else: + Aqk, Akkd = chunk_kda_fwd_intra_token_parallel( + q=q, + k=k, + gk=gk, + beta=beta, + Aqk=Aqk, + Akk=Akkd, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_size=BT, + sub_chunk_size=BC, + ) + + # Step 2: inter_solve (+ optional fused recompute) + grid = (NT, B * H) + + if fuse_recompute: + w = torch.empty_like(k) + u = torch.empty_like(v) + kg = torch.empty_like(k) + chunk_kda_fwd_kernel_inter_solve_fused[grid]( q=q, k=k, g=gk, beta=beta, Aqk=Aqk, - Akk=Akkd, + Akkd=Akkd, + Akk=k, # unused placeholder when FUSE_RECOMPUTE=True (dead branch) scale=scale, + v_in=v, + w_out=w, + u_out=u, + kg_out=kg, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, T=T, H=H, K=K, + V=V, BT=BT, BC=BC, - BK=BK, - USE_GATHER=is_gather_supported, - ) - else: - Aqk, Akkd = chunk_kda_fwd_intra_token_parallel( - q=q, - k=k, - gk=gk, - beta=beta, - Aqk=Aqk, - Akk=Akkd, - scale=scale, - cu_seqlens=cu_seqlens, - chunk_size=BT, - sub_chunk_size=BC, + USE_SAFE_GATE=safe_gate, + FUSE_RECOMPUTE=True, + FUSE_DIAGONAL=fuse_diagonal, ) + return w, u, None, kg, Aqk, None - # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen) - grid = (NT, B * H) + # Non-fused path: inter_solve stores Akk, then separate recompute + Akk = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype) chunk_kda_fwd_kernel_inter_solve_fused[grid]( q=q, k=k, @@ -636,15 +1018,24 @@ def chunk_kda_fwd_intra( Akkd=Akkd, Akk=Akk, scale=scale, + # v_in/w_out/u_out/kg_out unused when FUSE_RECOMPUTE=False (dead branch) + v_in=k, + w_out=k, + u_out=k, + kg_out=k, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, T=T, H=H, K=K, + V=0, BT=BT, BC=BC, USE_SAFE_GATE=safe_gate, + FUSE_RECOMPUTE=False, + FUSE_DIAGONAL=fuse_diagonal, ) + from sglang.srt.layers.attention.fla.kda import ( recompute_w_u_fwd as kda_recompute_w_u_fwd, ) diff --git a/python/sglang/srt/layers/attention/fla/cumsum.py b/python/sglang/srt/layers/attention/fla/cumsum.py index 0bc05ed88..911b78e00 100644 --- a/python/sglang/srt/layers/attention/fla/cumsum.py +++ b/python/sglang/srt/layers/attention/fla/cumsum.py @@ -70,9 +70,10 @@ def chunk_local_cumsum_scalar_kernel( @triton.autotune( configs=[ - triton.Config({"BS": BS}, num_warps=num_warps) + triton.Config({"BS": BS}, num_warps=num_warps, num_stages=num_stages) for BS in BS_LIST for num_warps in [2, 4, 8] + for num_stages in [2, 3, 4] ], key=["B", "H", "S", "BT", "IS_VARLEN", "REVERSE", "HAS_SCALE"], ) diff --git a/python/sglang/srt/layers/attention/fla/kda.py b/python/sglang/srt/layers/attention/fla/kda.py index f0ff3e449..2a6fd059b 100644 --- a/python/sglang/srt/layers/attention/fla/kda.py +++ b/python/sglang/srt/layers/attention/fla/kda.py @@ -17,10 +17,14 @@ from sglang.srt.layers.attention.fla.fused_norm_gate import layer_norm_gated_fwd from sglang.srt.layers.attention.fla.fused_recurrent import ( fused_recurrent_gated_delta_rule_fwd_kernel, ) -from sglang.srt.layers.attention.fla.index import prepare_chunk_indices +from sglang.srt.layers.attention.fla.index import ( + prepare_chunk_indices, +) from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd from sglang.srt.layers.attention.fla.op import exp, log -from sglang.srt.layers.attention.fla.utils import check_shared_mem +from sglang.srt.layers.attention.fla.utils import ( + check_shared_mem, +) BS_LIST = [32, 64] if check_shared_mem() else [16, 32] @@ -488,11 +492,13 @@ def chunk_kda_scaled_dot_kkt_fwd( @triton.autotune( configs=[ - triton.Config({}, num_warps=num_warps, num_stages=num_stages) + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) + for BK in [64, 128] + for BV in [64, 128] for num_warps in [2, 4, 8] for num_stages in [2, 3, 4] ], - key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"], + key=["H", "K", "V", "BT", "IS_VARLEN"], ) @triton.jit(do_not_specialize=["T"]) def recompute_w_u_fwd_kernel( @@ -650,8 +656,6 @@ def recompute_w_u_fwd( ) -> tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] BT = A.shape[-1] - BK = 64 - BV = 64 if chunk_indices is None and cu_seqlens is not None: chunk_indices = prepare_chunk_indices(cu_seqlens, BT) @@ -678,12 +682,10 @@ def recompute_w_u_fwd( K=K, V=V, BT=BT, - BK=BK, - BV=BV, STORE_QG=False, STORE_KG=kg is not None, IS_VARLEN=cu_seqlens is not None, - DOT_PRECISION="ieee", + DOT_PRECISION="tf32", ) return w, u, None, kg @@ -691,8 +693,8 @@ def recompute_w_u_fwd( @triton.autotune( configs=[ triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages) - for BK in [32, 64] - for BV in [64, 128] + for BK in [64] + for BV in [64] for num_warps in [2, 4, 8] for num_stages in [2, 3, 4] ], @@ -803,7 +805,7 @@ def chunk_gla_fwd_kernel_o( # [BT, BT] b_A = tl.load(p_A, boundary_check=(0, 1)) b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype) - b_o += tl.dot(b_A, b_v, allow_tf32=False) + b_o += tl.dot(b_A, b_v) tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) @@ -1065,7 +1067,23 @@ def chunk_kda_fwd( chunk_indices=chunk_indices, ) - # Fused: scaled_dot_kkt + solve_tril + recompute_w_u + # FUSE_DIAGONAL (fold diagonal-block compute into inter+solve) and + # FUSE_RECOMPUTE (also fold w/u/kg recompute) save kernel launches and HBM + # round-trips, but cost register footprint per CTA. Wins at small grid + # where launch overhead dominates; loses at large grid where the extra + # register pressure spills. Gate both on the same grid heuristic. + # Total CTAs in inter_solve_fused = NT * B * H_per_rank. For varlen, + # chunks don't cross sequence boundaries, so per-sequence ceil-divs sum to + # more than cdiv(total_tokens, chunk_size); use chunk_indices.shape[0] which + # already enumerates all (seq, chunk) pairs. + _NT_pr = ( + triton.cdiv(q.shape[1], chunk_size) + if cu_seqlens is None + else chunk_indices.shape[0] + ) + _H_pr = q.shape[-2] + _B = q.shape[0] + _small_grid = _B * _NT_pr * _H_pr <= 256 w, u, _, kg, Aqk, _ = chunk_kda_fwd_intra( q=q, k=k, @@ -1076,6 +1094,8 @@ def chunk_kda_fwd( cu_seqlens=cu_seqlens, chunk_size=chunk_size, chunk_indices=chunk_indices, + fuse_diagonal=_small_grid, + fuse_recompute=_small_grid, ) h, v_new = chunk_gated_delta_rule_fwd_h( @@ -1089,6 +1109,7 @@ def chunk_kda_fwd( chunk_indices=chunk_indices, ) del w, u, kg + o = chunk_gla_fwd_o_gk( q=q, v=v_new, @@ -1097,11 +1118,12 @@ def chunk_kda_fwd( h=h, o=v, scale=scale, - cu_seqlens=cu_seqlens, chunk_size=chunk_size, + cu_seqlens=cu_seqlens, chunk_indices=chunk_indices, ) del Aqk, v_new, h + return o