[GDN] Fuse GDN kkt + solve_tril into one kernel (#21411)

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

python/sglang/srt/layers/attention/fla/chunk.py

@@ -8,19 +8,20 @@ import torch
 from einops import rearrange
 
 from sglang.srt.layers.attention.fla.chunk_delta_h import chunk_gated_delta_rule_fwd_h
+from sglang.srt.layers.attention.fla.chunk_fwd import chunk_gated_delta_rule_fwd_intra
 from sglang.srt.layers.attention.fla.chunk_o import chunk_fwd_o
-from sglang.srt.layers.attention.fla.chunk_scaled_dot_kkt import (
-    chunk_scaled_dot_kkt_fwd,
-)
 from sglang.srt.layers.attention.fla.cumsum import chunk_local_cumsum
+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.solve_tril import solve_tril
 from sglang.srt.layers.attention.fla.utils import (
     SUPPRESS_LEVEL,
     autocast_custom_fwd,
     input_guard,
 )
-from sglang.srt.layers.attention.fla.wy_fast import recompute_w_u_fwd
+
+CHUNK_SIZE = 64
 
 
 def chunk_gated_delta_rule_fwd(
@@ -33,21 +34,20 @@ def chunk_gated_delta_rule_fwd(
     initial_state: torch.Tensor,
     initial_state_indices: torch.Tensor,
     cu_seqlens: Optional[torch.LongTensor] = None,
+    chunk_indices: torch.LongTensor | None = None,
 ):
-    g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens)
+    g = chunk_local_cumsum(g, chunk_size=CHUNK_SIZE, cu_seqlens=cu_seqlens)
-    # obtain WY representation. u is actually the new v.
+
-    A = chunk_scaled_dot_kkt_fwd(
+    # fused kkt + solve_tril + recompute_w_u
-        k=k, beta=beta, g_cumsum=g, cu_seqlens=cu_seqlens, output_dtype=torch.float32
+    w, u, A = chunk_gated_delta_rule_fwd_intra(
-    )
-    A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype)
-    w, u = recompute_w_u_fwd(
         k=k,
         v=v,
+        g=g,
         beta=beta,
-        A=A,
-        g_cumsum=g,
         cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
     )
+
     h, v_new = chunk_gated_delta_rule_fwd_h(
         k=k,
         w=w,
@@ -97,6 +97,11 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
             q = l2norm_fwd(q)
             k = l2norm_fwd(k)
 
+        chunk_indices = (
+            prepare_chunk_indices(cu_seqlens, CHUNK_SIZE)
+            if cu_seqlens is not None
+            else None
+        )
         g, o, A, w, h, v_new = chunk_gated_delta_rule_fwd(
             q=q,
             k=k,
@@ -107,6 +112,7 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
             initial_state=initial_state,
             initial_state_indices=initial_state_indices,
             cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices,
         )
         return o.to(q.dtype), h
 

python/sglang/srt/layers/attention/fla/chunk_fwd.py

@@ -0,0 +1,416 @@
+# Adapted from https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/gated_delta_rule/chunk_fwd.py
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.layers.attention.fla.index import prepare_chunk_indices
+from sglang.srt.layers.attention.fla.op import safe_exp
+from sglang.srt.layers.attention.fla.utils import (
+    autotune_cache_kwargs,
+    is_tf32_supported,
+)
+from sglang.srt.layers.attention.fla.wy_fast import recompute_w_u_fwd
+
+# TF32 for the block-merge dot products (16x16 matmuls) is safe and ~2x faster on SM90.
+# The numerically sensitive forward-substitution uses scalar ops, not tl.dot.
+if is_tf32_supported:
+    _MERGE_DOT_PRECISION = tl.constexpr("tf32")
+else:
+    _MERGE_DOT_PRECISION = tl.constexpr("ieee")
+
+
+@triton.heuristics(
+    {
+        "USE_G": lambda args: args["g"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.autotune(
+    configs=[
+        triton.Config({"BK": BK}, num_warps=num_warps)
+        for BK in [32, 64]
+        for num_warps in [1, 2, 4]
+    ],
+    key=["H", "Hg", "K", "BC"],
+    **autotune_cache_kwargs,
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_gated_delta_rule_fwd_kkt_solve_kernel(
+    k,
+    g,
+    beta,
+    A,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    Hg: tl.constexpr,
+    K: tl.constexpr,
+    BT: tl.constexpr,
+    BC: tl.constexpr,
+    BK: tl.constexpr,
+    USE_G: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    """
+    Fused kernel: compute beta * K @ K^T (lower triangular) + solve_tril (I+A)^{-1} in one pass.
+
+    This kernel fuses chunk_scaled_dot_kkt_fwd and solve_tril into a single kernel,
+    avoiding the HBM round-trip for the intermediate A matrix.
+
+    Steps:
+    1. Compute all 10 lower-triangular [BC, BC] blocks of beta * K @ K^T in registers
+    2. Apply gate and beta scaling
+    3. Forward substitution on diagonal blocks
+    4. Block merge to get full (I+A)^{-1}
+    5. Write result to A (output)
+    """
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    i_b, i_h = i_bh // H, i_bh % H
+
+    if IS_VARLEN:
+        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
+            chunk_indices + i_t * 2 + 1
+        ).to(tl.int32)
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
+            cu_seqlens + i_n + 1
+        ).to(tl.int32)
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    if i_t * BT >= T:
+        return
+
+    i_tc0 = i_t * BT
+    i_tc1 = i_t * BT + BC
+    i_tc2 = i_t * BT + 2 * BC
+    i_tc3 = i_t * BT + 3 * BC
+
+    k += (bos * Hg + i_h // (H // Hg)) * K
+    A += (bos * H + i_h) * BT
+
+    o_i = tl.arange(0, BC)
+    m_tc0 = (i_tc0 + o_i) < T
+    m_tc1 = (i_tc1 + o_i) < T
+    m_tc2 = (i_tc2 + o_i) < T
+    m_tc3 = (i_tc3 + o_i) < T
+
+    # load beta for each sub-chunk
+    p_b0 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc0,), (BC,), (0,))
+    p_b1 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc1,), (BC,), (0,))
+    p_b2 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc2,), (BC,), (0,))
+    p_b3 = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_tc3,), (BC,), (0,))
+    b_b0 = tl.load(p_b0, boundary_check=(0,)).to(tl.float32)
+    b_b1 = tl.load(p_b1, boundary_check=(0,)).to(tl.float32)
+    b_b2 = tl.load(p_b2, boundary_check=(0,)).to(tl.float32)
+    b_b3 = tl.load(p_b3, boundary_check=(0,)).to(tl.float32)
+
+    # load gate if used
+    if USE_G:
+        p_g0 = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_tc0,), (BC,), (0,))
+        p_g1 = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_tc1,), (BC,), (0,))
+        p_g2 = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_tc2,), (BC,), (0,))
+        p_g3 = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_tc3,), (BC,), (0,))
+
+        b_g0 = tl.load(p_g0, boundary_check=(0,)).to(tl.float32)
+        b_g1 = tl.load(p_g1, boundary_check=(0,)).to(tl.float32)
+        b_g2 = tl.load(p_g2, boundary_check=(0,)).to(tl.float32)
+        b_g3 = tl.load(p_g3, boundary_check=(0,)).to(tl.float32)
+
+    ############################################################################
+    # Step 1: compute all 10 lower-triangular [BC, BC] blocks of K @ K^T
+    ############################################################################
+
+    # 4 diagonal blocks
+    b_A00 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A11 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A22 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A33 = tl.zeros([BC, BC], dtype=tl.float32)
+
+    # 6 off-diagonal blocks
+    b_A10 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A20 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A21 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A30 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A31 = tl.zeros([BC, BC], dtype=tl.float32)
+    b_A32 = tl.zeros([BC, BC], dtype=tl.float32)
+
+    for i_k in range(tl.cdiv(K, BK)):
+        p_k0 = tl.make_block_ptr(
+            k, (T, K), (Hg * K, 1), (i_tc0, i_k * BK), (BC, BK), (1, 0)
+        )
+        b_k0 = tl.load(p_k0, boundary_check=(0, 1))
+        # diagonal block 0
+        b_A00 += tl.dot(b_k0, tl.trans(b_k0))
+
+        if i_tc1 < T:
+            p_k1 = tl.make_block_ptr(
+                k, (T, K), (Hg * K, 1), (i_tc1, i_k * BK), (BC, BK), (1, 0)
+            )
+            b_k1 = tl.load(p_k1, boundary_check=(0, 1))
+            # diagonal block 1
+            b_A11 += tl.dot(b_k1, tl.trans(b_k1))
+            # off-diagonal (1,0)
+            b_A10 += tl.dot(b_k1, tl.trans(b_k0))
+
+            if i_tc2 < T:
+                p_k2 = tl.make_block_ptr(
+                    k, (T, K), (Hg * K, 1), (i_tc2, i_k * BK), (BC, BK), (1, 0)
+                )
+                b_k2 = tl.load(p_k2, boundary_check=(0, 1))
+                # diagonal block 2
+                b_A22 += tl.dot(b_k2, tl.trans(b_k2))
+                # off-diagonal (2,0), (2,1)
+                b_A20 += tl.dot(b_k2, tl.trans(b_k0))
+                b_A21 += tl.dot(b_k2, tl.trans(b_k1))
+
+                if i_tc3 < T:
+                    p_k3 = tl.make_block_ptr(
+                        k, (T, K), (Hg * K, 1), (i_tc3, i_k * BK), (BC, BK), (1, 0)
+                    )
+                    b_k3 = tl.load(p_k3, boundary_check=(0, 1))
+                    # diagonal block 3
+                    b_A33 += tl.dot(b_k3, tl.trans(b_k3))
+                    # off-diagonal (3,0), (3,1), (3,2)
+                    b_A30 += tl.dot(b_k3, tl.trans(b_k0))
+                    b_A31 += tl.dot(b_k3, tl.trans(b_k1))
+                    b_A32 += tl.dot(b_k3, tl.trans(b_k2))
+
+    ############################################################################
+    # Step 2: apply gate and beta scaling
+    ############################################################################
+
+    if USE_G:
+        # diagonal blocks: g_diff = g_i - g_j within sub-chunk
+        b_A00 *= safe_exp(b_g0[:, None] - b_g0[None, :])
+        b_A11 *= safe_exp(b_g1[:, None] - b_g1[None, :])
+        b_A22 *= safe_exp(b_g2[:, None] - b_g2[None, :])
+        b_A33 *= safe_exp(b_g3[:, None] - b_g3[None, :])
+
+        # off-diagonal blocks: g_diff = g_row - g_col (cross sub-chunk)
+        b_A10 *= safe_exp(b_g1[:, None] - b_g0[None, :])
+        b_A20 *= safe_exp(b_g2[:, None] - b_g0[None, :])
+        b_A21 *= safe_exp(b_g2[:, None] - b_g1[None, :])
+        b_A30 *= safe_exp(b_g3[:, None] - b_g0[None, :])
+        b_A31 *= safe_exp(b_g3[:, None] - b_g1[None, :])
+        b_A32 *= safe_exp(b_g3[:, None] - b_g2[None, :])
+
+    # apply beta to row dimension and mask
+    m_d = o_i[:, None] > o_i[None, :]
+    m_I = o_i[:, None] == o_i[None, :]
+
+    # diagonal blocks: strictly lower triangular within sub-chunk, scaled by beta
+    b_A00 = (
+        tl.where(m_d & (m_tc0[:, None] & m_tc0[None, :]), b_A00, 0.0) * b_b0[:, None]
+    )
+    b_A11 = (
+        tl.where(m_d & (m_tc1[:, None] & m_tc1[None, :]), b_A11, 0.0) * b_b1[:, None]
+    )
+    b_A22 = (
+        tl.where(m_d & (m_tc2[:, None] & m_tc2[None, :]), b_A22, 0.0) * b_b2[:, None]
+    )
+    b_A33 = (
+        tl.where(m_d & (m_tc3[:, None] & m_tc3[None, :]), b_A33, 0.0) * b_b3[:, None]
+    )
+
+    # off-diagonal blocks: full block, scaled by beta
+    b_A10 = b_A10 * b_b1[:, None]
+    b_A20 = b_A20 * b_b2[:, None]
+    b_A21 = b_A21 * b_b2[:, None]
+    b_A30 = b_A30 * b_b3[:, None]
+    b_A31 = b_A31 * b_b3[:, None]
+    b_A32 = b_A32 * b_b3[:, None]
+
+    ############################################################################
+    # Step 3: forward substitution on diagonal blocks -> (I + A_diag)^{-1}
+    #
+    # Same algorithm as solve_tril, but rows are extracted from in-register
+    # [BC, BC] tensor via tl.sum(tl.where(mask, tensor, 0), 0) instead of
+    # tl.load from HBM.
+    ############################################################################
+
+    b_Ai00 = -b_A00
+    b_Ai11 = -b_A11
+    b_Ai22 = -b_A22
+    b_Ai33 = -b_A33
+
+    for i in range(2, min(BC, T - i_tc0)):
+        b_a00 = tl.sum(tl.where((o_i == i)[:, None], -b_A00, 0.0), 0)
+        b_a00 = tl.where(o_i < i, b_a00, 0.0)
+        b_a00 = b_a00 + tl.sum(b_a00[:, None] * b_Ai00, 0)
+        b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00)
+    for i in range(2, min(BC, T - i_tc1)):
+        b_a11 = tl.sum(tl.where((o_i == i)[:, None], -b_A11, 0.0), 0)
+        b_a11 = tl.where(o_i < i, b_a11, 0.0)
+        b_a11 = b_a11 + tl.sum(b_a11[:, None] * b_Ai11, 0)
+        b_Ai11 = tl.where((o_i == i)[:, None], b_a11, b_Ai11)
+    for i in range(2, min(BC, T - i_tc2)):
+        b_a22 = tl.sum(tl.where((o_i == i)[:, None], -b_A22, 0.0), 0)
+        b_a22 = tl.where(o_i < i, b_a22, 0.0)
+        b_a22 = b_a22 + tl.sum(b_a22[:, None] * b_Ai22, 0)
+        b_Ai22 = tl.where((o_i == i)[:, None], b_a22, b_Ai22)
+    for i in range(2, min(BC, T - i_tc3)):
+        b_a33 = tl.sum(tl.where((o_i == i)[:, None], -b_A33, 0.0), 0)
+        b_a33 = tl.where(o_i < i, b_a33, 0.0)
+        b_a33 = b_a33 + tl.sum(b_a33[:, None] * b_Ai33, 0)
+        b_Ai33 = tl.where((o_i == i)[:, None], b_a33, b_Ai33)
+
+    b_Ai00 += m_I
+    b_Ai11 += m_I
+    b_Ai22 += m_I
+    b_Ai33 += m_I
+
+    ############################################################################
+    # Step 4: block merge -> full (I + A)^{-1}
+    ############################################################################
+
+    b_Ai10 = -tl.dot(
+        tl.dot(b_Ai11, b_A10, input_precision=_MERGE_DOT_PRECISION),
+        b_Ai00,
+        input_precision=_MERGE_DOT_PRECISION,
+    )
+    b_Ai21 = -tl.dot(
+        tl.dot(b_Ai22, b_A21, input_precision=_MERGE_DOT_PRECISION),
+        b_Ai11,
+        input_precision=_MERGE_DOT_PRECISION,
+    )
+    b_Ai32 = -tl.dot(
+        tl.dot(b_Ai33, b_A32, input_precision=_MERGE_DOT_PRECISION),
+        b_Ai22,
+        input_precision=_MERGE_DOT_PRECISION,
+    )
+
+    b_Ai20 = -tl.dot(
+        b_Ai22,
+        tl.dot(b_A20, b_Ai00, input_precision=_MERGE_DOT_PRECISION)
+        + tl.dot(b_A21, b_Ai10, input_precision=_MERGE_DOT_PRECISION),
+        input_precision=_MERGE_DOT_PRECISION,
+    )
+    b_Ai31 = -tl.dot(
+        b_Ai33,
+        tl.dot(b_A31, b_Ai11, input_precision=_MERGE_DOT_PRECISION)
+        + tl.dot(b_A32, b_Ai21, input_precision=_MERGE_DOT_PRECISION),
+        input_precision=_MERGE_DOT_PRECISION,
+    )
+    b_Ai30 = -tl.dot(
+        b_Ai33,
+        tl.dot(b_A30, b_Ai00, input_precision=_MERGE_DOT_PRECISION)
+        + tl.dot(b_A31, b_Ai10, input_precision=_MERGE_DOT_PRECISION)
+        + tl.dot(b_A32, b_Ai20, input_precision=_MERGE_DOT_PRECISION),
+        input_precision=_MERGE_DOT_PRECISION,
+    )
+
+    ############################################################################
+    # Step 5: store full (I + A)^{-1} to output A
+    ############################################################################
+
+    p_A00 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_tc0, 0), (BC, BC), (1, 0))
+    p_A10 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_tc1, 0), (BC, BC), (1, 0))
+    p_A11 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_tc1, BC), (BC, BC), (1, 0))
+    p_A20 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_tc2, 0), (BC, BC), (1, 0))
+    p_A21 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_tc2, BC), (BC, BC), (1, 0))
+    p_A22 = tl.make_block_ptr(
+        A, (T, BT), (H * BT, 1), (i_tc2, 2 * BC), (BC, BC), (1, 0)
+    )
+    p_A30 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_tc3, 0), (BC, BC), (1, 0))
+    p_A31 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_tc3, BC), (BC, BC), (1, 0))
+    p_A32 = tl.make_block_ptr(
+        A, (T, BT), (H * BT, 1), (i_tc3, 2 * BC), (BC, BC), (1, 0)
+    )
+    p_A33 = tl.make_block_ptr(
+        A, (T, BT), (H * BT, 1), (i_tc3, 3 * BC), (BC, BC), (1, 0)
+    )
+
+    tl.store(p_A00, b_Ai00.to(A.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_A10, b_Ai10.to(A.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_A11, b_Ai11.to(A.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_A20, b_Ai20.to(A.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_A21, b_Ai21.to(A.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_A22, b_Ai22.to(A.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_A30, b_Ai30.to(A.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_A31, b_Ai31.to(A.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_A32, b_Ai32.to(A.dtype.element_ty), boundary_check=(0, 1))
+    tl.store(p_A33, b_Ai33.to(A.dtype.element_ty), boundary_check=(0, 1))
+
+
+def chunk_gated_delta_rule_fwd_intra(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor | None = None,
+    beta: torch.Tensor | None = None,
+    cu_seqlens: torch.LongTensor | None = None,
+    chunk_size: int = 64,
+    chunk_indices: torch.LongTensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    r"""
+    GDN intra-chunk forward: fused kkt + solve_tril + recompute_w_u.
+
+    Equivalent to:
+        A = chunk_scaled_dot_kkt_fwd(k, g, beta, ...)       # kernel 1
+        A = solve_tril(A, ...)                                # kernel 2
+        w, u = recompute_w_u_fwd(k, v, beta, A, g, ...)      # kernel 3
+
+    Fuses kernels 1+2 into a single kernel, reducing from 3 to 2 kernel launches
+    and eliminating the HBM round-trip for the intermediate A matrix.
+
+    Args:
+        k (torch.Tensor):
+            The key tensor of shape `[B, T, H, K]`.
+        v (torch.Tensor):
+            The value tensor of shape `[B, T, H, V]`.
+        g (torch.Tensor):
+            The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`.
+        beta (torch.Tensor):
+            The beta tensor of shape `[B, T, H]`.
+        cu_seqlens (torch.LongTensor):
+            The cumulative sequence lengths. Default: `None`.
+        chunk_size (int):
+            The chunk size. Default: 64.
+        chunk_indices (torch.LongTensor):
+            Precomputed chunk indices. Default: `None`.
+
+    Returns:
+        w (torch.Tensor): shape `[B, T, H, K]`
+        u (torch.Tensor): shape `[B, T, H, V]`
+        A (torch.Tensor): shape `[B, T, H, BT]`, the solved (I+A)^{-1} matrix
+    """
+    B, T, Hg, K = k.shape
+    H = beta.shape[-1]
+    BT = chunk_size
+    BC = 16
+
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    # Step 1: fused kkt + solve_tril
+    A = torch.zeros(B, T, H, BT, device=k.device, dtype=k.dtype)
+    chunk_gated_delta_rule_fwd_kkt_solve_kernel[(NT, B * H)](
+        k=k,
+        g=g,
+        beta=beta,
+        A=A,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        Hg=Hg,
+        K=K,
+        BT=BT,
+        BC=BC,
+    )
+
+    # Step 2: recompute_w_u
+    w, u = recompute_w_u_fwd(
+        k=k,
+        v=v,
+        beta=beta,
+        A=A,
+        g_cumsum=g,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+    )
+    return w, u, A

python/sglang/srt/layers/attention/fla/fused_recurrent.py

@@ -940,3 +940,6 @@ def fused_recurrent_gated_delta_rule_update(
         retrieve_parent_token,
     )
     return o
+
+
+fused_recurrent_gdn = fused_recurrent_gated_delta_rule

python/sglang/srt/layers/attention/fla/utils.py

@@ -3,6 +3,7 @@
 
 import contextlib
 import functools
+import inspect
 import logging
 import os
 import sys
@@ -20,6 +21,16 @@ logger = logging.getLogger(__name__)
 
 COMPILER_MODE = os.getenv("FLA_COMPILER_MODE") == "1"
 FLA_CI_ENV = os.getenv("FLA_CI_ENV") == "1"
+FLA_CACHE_RESULTS = os.getenv("FLA_CACHE_RESULTS", "1") == "1"
+
+
+SUPPORTS_AUTOTUNE_CACHE = (
+    "cache_results" in inspect.signature(triton.autotune).parameters
+)
+
+autotune_cache_kwargs = (
+    {"cache_results": FLA_CACHE_RESULTS} if SUPPORTS_AUTOTUNE_CACHE else {}
+)
 
 
 @lru_cache(maxsize=1)
@@ -323,3 +334,6 @@ else:
 
     def custom_device_ctx(index: int):
         return torch.cuda.device(index)
+
+
+device_platform = get_available_device()

python/sglang/srt/layers/attention/fla/wy_fast.py

@@ -115,14 +115,14 @@ def recompute_w_u_fwd(
     g_cumsum: torch.Tensor,
     A: torch.Tensor,
     cu_seqlens: Optional[torch.LongTensor],
+    chunk_indices: torch.LongTensor | None = None,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     B, T, Hg, K, V = *k.shape, v.shape[-1]
     H = v.shape[-2]
     BT = A.shape[-1]
 
-    chunk_indices = (
+    if chunk_indices is None and cu_seqlens is not None:
-        prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+        chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
-    )
     NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
     BK = 64
     BV = 64