[KDA] Optimize prefill kernels with diagonal and recompute fuse (#24271)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

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(
+@triton.autotune(
-#     configs=[
+    # Single hardcoded config. The kernel writes ht (final state) back into
-#         triton.Config({"BV": BV}, num_warps=num_warps, num_stages=num_stages)
+    # initial_state in-place; with multiple configs, triton's autotune benchmark
-#         for num_warps in [2, 4]
+    # phase invokes the kernel many times for timing and corrupts the cache pool,
-#         for num_stages in [2, 3, 4]
+    # producing silently wrong output on the first user request. Restoring via
-#         for BV in [32, 64]
+    # `restore_value=["initial_state"]` works for unit tests but OOMs on
-#     ],
+    # production-scale models (e.g. Kimi-Linear-48B at default mem_fraction)
-#     key=["H", "K", "V", "BT", "USE_G"],
+    # because cloning the cache pool for each benchmark exceeds available memory.
-#     use_cuda_graph=use_cuda_graph,
+    # 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,
+        NT_BUCKET=(0 if NT <= 32 else (1 if NT <= 128 else 2)),
-        num_stages=2,
     )
     return h, v_new

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_qg1 = (b_q1 * b_gqn).to(tl.bfloat16)
-            b_Akk10 += tl.dot(b_k1 * b_gqn, b_kgt)
+            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_qg2 = (b_q2 * b_gqn2).to(tl.bfloat16)
-                b_kg2 = b_k2 * b_gqn2
+                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_qg3 = (b_q3 * b_gqn3).to(tl.bfloat16)
-                    b_kg3 = b_k3 * b_gqn3
+                    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,6 +354,87 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
         b_Akk31 = b_Akk31 * b_b3[:, None]
         b_Akk32 = b_Akk32 * b_b3[:, None]
 
+    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)
         )
@@ -318,9 +454,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
 
     ################################################################################
     # 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,22 +538,225 @@ 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))
+    if FUSE_RECOMPUTE:
-    p_Akk10 = tl.make_block_ptr(Akk, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0))
+        # 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)
+
+        # 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_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_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)
         )
@@ -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,14 +929,15 @@ def chunk_kda_fwd_intra(
     NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
     NC = triton.cdiv(BT, BC)
 
+    if fuse_diagonal:
         Aqk = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype)
-    # Akk must be zero-initialized - kernel only writes lower triangular
+    else:
-    Akk = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype)
+        Aqk = torch.empty(B, T, H, BT, device=k.device, dtype=k.dtype)
-    # Separate fp32 buffer for diagonal 16x16 blocks (for precision in solve_tril)
+    Akkd = torch.empty(B, T, H, BC, device=k.device, dtype=torch.float32)
-    Akkd = torch.zeros(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 Akkd (fp32)
-    # Step 1: compute diagonal blocks into Akk_diag (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)
@@ -625,8 +973,42 @@ def chunk_kda_fwd_intra(
                 sub_chunk_size=BC,
             )
 
-    # Step 2: Fused inter + solve_tril (works for both fixed-len and varlen)
+    # 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,
+            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,
+            USE_SAFE_GATE=safe_gate,
+            FUSE_RECOMPUTE=True,
+            FUSE_DIAGONAL=fuse_diagonal,
+        )
+        return w, u, None, kg, Aqk, None
+
+    # 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,
     )

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"],
 )

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 BK in [64]
-        for BV in [64, 128]
+        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