diff --git a/python/sglang/srt/layers/attention/fla/chunk.py b/python/sglang/srt/layers/attention/fla/chunk.py index 28fc166f4..f715afdc8 100644 --- a/python/sglang/srt/layers/attention/fla/chunk.py +++ b/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) - # obtain WY representation. u is actually the new v. - A = chunk_scaled_dot_kkt_fwd( - k=k, beta=beta, g_cumsum=g, cu_seqlens=cu_seqlens, output_dtype=torch.float32 - ) - A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype) - w, u = recompute_w_u_fwd( + g = chunk_local_cumsum(g, chunk_size=CHUNK_SIZE, cu_seqlens=cu_seqlens) + + # fused kkt + solve_tril + recompute_w_u + w, u, A = chunk_gated_delta_rule_fwd_intra( 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 diff --git a/python/sglang/srt/layers/attention/fla/chunk_fwd.py b/python/sglang/srt/layers/attention/fla/chunk_fwd.py new file mode 100644 index 000000000..432a274cd --- /dev/null +++ b/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 diff --git a/python/sglang/srt/layers/attention/fla/fused_recurrent.py b/python/sglang/srt/layers/attention/fla/fused_recurrent.py index 44e42e2d6..f110770c0 100644 --- a/python/sglang/srt/layers/attention/fla/fused_recurrent.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/fla/utils.py b/python/sglang/srt/layers/attention/fla/utils.py index af6ca3d6e..4154a3c52 100644 --- a/python/sglang/srt/layers/attention/fla/utils.py +++ b/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() diff --git a/python/sglang/srt/layers/attention/fla/wy_fast.py b/python/sglang/srt/layers/attention/fla/wy_fast.py index 757e56210..980a475cc 100644 --- a/python/sglang/srt/layers/attention/fla/wy_fast.py +++ b/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 = ( - prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None - ) + 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) BK = 64 BV = 64