diff --git a/benchmark/bench_linear_attention/bench_fused_gate_cumsum.py b/benchmark/bench_linear_attention/bench_fused_gate_cumsum.py new file mode 100644 index 000000000..1b2c105c3 --- /dev/null +++ b/benchmark/bench_linear_attention/bench_fused_gate_cumsum.py @@ -0,0 +1,213 @@ +""" +Benchmark: Fused Gate+Cumsum vs Separate Gate + Cumsum. + +Compares two paths: + - Separate: torch gate activation -> chunk_local_cumsum (2 steps) + - Fused: kda_gate_chunk_cumsum (single kernel) + +Both produce the same output: cumsum of gate-activated g. + +Usage: + python bench_fused_gate_cumsum.py + python bench_fused_gate_cumsum.py --batch-sizes 4 16 64 128 + python bench_fused_gate_cumsum.py --seq-lens 64 128 256 512 1024 +""" + +import argparse +import os +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python")) + +import torch +import triton + +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.kda import kda_gate_chunk_cumsum + +CHUNK_SIZE = 64 + + +def make_inputs( + B: int, + T_per_seq: int, + H: int, + K: int, + device: str, + dtype: torch.dtype, + seed: int = 42, +): + T = B * T_per_seq + torch.manual_seed(seed) + + # Raw gate: [1, T_total, H, K] (varlen format, before activation) + raw_g = torch.randn(1, T, H, K, dtype=dtype, device=device) + + # A_log: [H] (per-head log-scale parameter) + A_log = torch.randn(H, dtype=torch.float32, device=device) * 0.5 + + # dt_bias: [H*K] (per-head bias, flat) + dt_bias = torch.randn(H * K, dtype=torch.float32, device=device) * 0.1 + + # cu_seqlens for varlen mode + cu_seqlens = torch.arange( + 0, (B + 1) * T_per_seq, T_per_seq, dtype=torch.long, device=device + ) + + return dict( + raw_g=raw_g, + A_log=A_log, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + B=B, + T=T, + T_per_seq=T_per_seq, + H=H, + K=K, + ) + + +def run_ref(inp): + """Separate path: torch gate activation -> chunk_local_cumsum.""" + raw_g = inp["raw_g"] # [1, T, H, K] + A_log = inp["A_log"] # [H] + dt_bias = inp["dt_bias"] # [H*K] + cu_seqlens = inp["cu_seqlens"] + H, K = inp["H"], inp["K"] + + # Step 1: gate activation using torch ops + g_float = raw_g.float() + if dt_bias is not None: + g_float = g_float + dt_bias.float().view(1, 1, H, K) + g_activated = -torch.exp( + A_log.float().view(1, 1, H, 1) + ) * torch.nn.functional.softplus(g_float) + + # Step 2: chunk-local cumsum + chunk_indices = prepare_chunk_indices(cu_seqlens, CHUNK_SIZE) + g_cumsum = chunk_local_cumsum( + g_activated, + chunk_size=CHUNK_SIZE, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return g_cumsum + + +def run_fused(inp): + """Fused path: kda_gate_chunk_cumsum (single kernel).""" + raw_g = inp["raw_g"] + A_log = inp["A_log"] + dt_bias = inp["dt_bias"] + cu_seqlens = inp["cu_seqlens"] + + chunk_indices = prepare_chunk_indices(cu_seqlens, CHUNK_SIZE) + g_cumsum = kda_gate_chunk_cumsum( + raw_g, + A_log=A_log, + chunk_size=CHUNK_SIZE, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return g_cumsum + + +def verify_correctness(inp): + """Verify fused and separate paths produce the same output.""" + out_separate = run_ref(inp) + out_fused = run_fused(inp) + + max_diff = (out_separate - out_fused).abs().max().item() + rel_diff = max_diff / (out_separate.abs().mean().item() + 1e-8) + return max_diff, rel_diff + + +def bench_shape(B, H, T_per_seq, K, device, dtype): + T = B * T_per_seq + inp = make_inputs(B, T_per_seq, H, K, device, dtype) + + # Warmup (includes triton compilation) + for _ in range(5): + run_ref(inp) + run_fused(inp) + torch.cuda.synchronize() + + ms_sep, ms_sep_lo, ms_sep_hi = triton.testing.do_bench( + lambda: run_ref(inp), quantiles=[0.5, 0.2, 0.8], warmup=50, rep=200 + ) + ms_fused, ms_fused_lo, ms_fused_hi = triton.testing.do_bench( + lambda: run_fused(inp), quantiles=[0.5, 0.2, 0.8], warmup=50, rep=200 + ) + + speedup = ms_sep / ms_fused if ms_fused > 0 else 0 + saved_us = (ms_sep - ms_fused) * 1000 # microseconds + + print( + f" {B:>5} {H:>3} {T_per_seq:>6} {T:>7} | " + f"{ms_sep:>8.3f} {ms_fused:>8.3f} | " + f"{speedup:>6.2f}x {saved_us:>+8.1f}us" + ) + + +def main(): + parser = argparse.ArgumentParser( + description="Benchmark: Fused vs Separate Gate+Cumsum" + ) + parser.add_argument("--dtype", choices=["bfloat16", "float16"], default="bfloat16") + parser.add_argument("--head-size-k", type=int, default=128) + parser.add_argument("--num-heads", type=int, nargs="+", default=[16]) + parser.add_argument( + "--batch-sizes", type=int, nargs="+", default=[4, 8, 16, 32, 64, 128] + ) + parser.add_argument( + "--seq-lens", type=int, nargs="+", default=[64, 128, 256, 512, 1024] + ) + args = parser.parse_args() + + device = "cuda" + dtype = getattr(torch, args.dtype) + K = args.head_size_k + + cap = torch.cuda.get_device_capability() + dev_name = torch.cuda.get_device_name() + print(f"Device: {dev_name} (SM {cap[0]}{cap[1]})") + print() + + # Correctness check + print("=" * 80) + print("Correctness verification") + print("=" * 80) + for H in args.num_heads: + inp = make_inputs(16, 256, H, K, device, dtype) + max_diff, rel_diff = verify_correctness(inp) + print( + f" H={H:>3}, B=16, T/seq=256: " + f"max_diff={max_diff:.2e}, rel_diff={rel_diff:.2e} " + f"{'PASS' if max_diff < 1e-3 else 'FAIL'}" + ) + print() + + # Performance benchmark + print("=" * 80) + print("Performance: Separate (gate+cumsum) vs Fused (single kernel)") + print("=" * 80) + print(f" Config: K={K}, chunk_size={CHUNK_SIZE}, dtype={dtype}") + print( + f" {'B':>5} {'H':>3} {'T/seq':>6} {'T_tot':>7} | " + f"{'sep(ms)':>8} {'fuse(ms)':>8} | " + f"{'speedup':>6} {'saved':>9}" + ) + print(" " + "-" * 73) + + for H in args.num_heads: + for B in args.batch_sizes: + for T_per_seq in args.seq_lens: + bench_shape(B, H, T_per_seq, K, device, dtype) + if len(args.num_heads) > 1: + print() + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/python/sglang/srt/layers/attention/fla/chunk.py b/python/sglang/srt/layers/attention/fla/chunk.py index f715afdc8..23d8d58fe 100644 --- a/python/sglang/srt/layers/attention/fla/chunk.py +++ b/python/sglang/srt/layers/attention/fla/chunk.py @@ -36,7 +36,9 @@ def chunk_gated_delta_rule_fwd( cu_seqlens: Optional[torch.LongTensor] = None, chunk_indices: torch.LongTensor | None = None, ): - g = chunk_local_cumsum(g, chunk_size=CHUNK_SIZE, cu_seqlens=cu_seqlens) + g = chunk_local_cumsum( + g, chunk_size=CHUNK_SIZE, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices + ) # fused kkt + solve_tril + recompute_w_u w, u, A = chunk_gated_delta_rule_fwd_intra( @@ -56,6 +58,7 @@ def chunk_gated_delta_rule_fwd( initial_state=initial_state, initial_state_indices=initial_state_indices, cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, ) o = chunk_fwd_o( q=q, 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 7e5fd53fc..0c7f80f42 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py +++ b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py @@ -281,16 +281,14 @@ def chunk_gated_delta_rule_fwd_h( initial_state_indices: Optional[torch.Tensor] = None, save_new_value: bool = True, cu_seqlens: Optional[torch.LongTensor] = None, + chunk_indices: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B, T, Hg, K, V = *k.shape, u.shape[-1] H = u.shape[-2] BT = CHUNK_SIZE - chunk_indices = ( - prepare_chunk_indices(cu_seqlens, CHUNK_SIZE) - 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, CHUNK_SIZE) # N: the actual number of sequences in the batch with either equal or variable lengths if cu_seqlens is None: N, NT, chunk_offsets = B, triton.cdiv(T, BT), None diff --git a/python/sglang/srt/layers/attention/fla/chunk_intra.py b/python/sglang/srt/layers/attention/fla/chunk_intra.py index 344de6117..00b62727c 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_intra.py +++ b/python/sglang/srt/layers/attention/fla/chunk_intra.py @@ -657,5 +657,6 @@ def chunk_kda_fwd_intra( q=q if disable_recompute else None, gk=gk, cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, ) return w, u, qg, kg, Aqk, Akk diff --git a/python/sglang/srt/layers/attention/fla/cumsum.py b/python/sglang/srt/layers/attention/fla/cumsum.py index 39d2f4722..0bc05ed88 100644 --- a/python/sglang/srt/layers/attention/fla/cumsum.py +++ b/python/sglang/srt/layers/attention/fla/cumsum.py @@ -163,6 +163,7 @@ def chunk_local_cumsum_scalar( cu_seqlens: Optional[torch.Tensor] = None, head_first: bool = False, output_dtype: Optional[torch.dtype] = torch.float, + chunk_indices: Optional[torch.LongTensor] = None, ) -> torch.Tensor: if head_first: B, H, T = g.shape @@ -172,9 +173,8 @@ def chunk_local_cumsum_scalar( chunk_size.bit_length() - 1 ), "chunk_size must be a power of 2" BT = chunk_size - 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) g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) grid = (NT, B * H) @@ -206,17 +206,15 @@ def chunk_local_cumsum_vector( cu_seqlens: Optional[torch.Tensor] = None, head_first: bool = False, output_dtype: Optional[torch.dtype] = torch.float, + chunk_indices: Optional[torch.LongTensor] = None, ) -> torch.Tensor: if head_first: B, H, T, S = g.shape else: B, T, H, S = g.shape BT = chunk_size - chunk_indices = ( - prepare_chunk_indices(cu_seqlens, chunk_size) - 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) assert chunk_size == 2 ** ( chunk_size.bit_length() - 1 @@ -258,6 +256,7 @@ def chunk_local_cumsum( cu_seqlens: Optional[torch.Tensor] = None, head_first: bool = False, output_dtype: Optional[torch.dtype] = torch.float, + chunk_indices: Optional[torch.LongTensor] = None, **kwargs, ) -> torch.Tensor: if cu_seqlens is not None: @@ -273,6 +272,7 @@ def chunk_local_cumsum( cu_seqlens=cu_seqlens, head_first=head_first, output_dtype=output_dtype, + chunk_indices=chunk_indices, ) elif len(g.shape) == 4: return chunk_local_cumsum_vector( @@ -283,6 +283,7 @@ def chunk_local_cumsum( cu_seqlens=cu_seqlens, head_first=head_first, output_dtype=output_dtype, + chunk_indices=chunk_indices, ) else: raise ValueError( diff --git a/python/sglang/srt/layers/attention/fla/kda.py b/python/sglang/srt/layers/attention/fla/kda.py index a8d5cb405..f0ff3e449 100644 --- a/python/sglang/srt/layers/attention/fla/kda.py +++ b/python/sglang/srt/layers/attention/fla/kda.py @@ -4,6 +4,8 @@ # the following copyright notice: # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang +from typing import Optional + import torch import triton import triton.language as tl @@ -18,10 +20,9 @@ from sglang.srt.layers.attention.fla.fused_recurrent import ( 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 is_amd +from sglang.srt.layers.attention.fla.utils import check_shared_mem -BT_LIST_AUTOTUNE = [32, 64, 128] -NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32] +BS_LIST = [32, 64] if check_shared_mem() else [16, 32] def cdiv(a: int, b: int) -> int: @@ -645,15 +646,15 @@ def recompute_w_u_fwd( q: torch.Tensor | None = None, gk: torch.Tensor | None = None, cu_seqlens: torch.LongTensor | None = None, + chunk_indices: torch.LongTensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: B, T, H, K, V = *k.shape, v.shape[-1] BT = A.shape[-1] BK = 64 BV = 64 - 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 = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) w = torch.empty_like(k) @@ -816,15 +817,13 @@ def chunk_gla_fwd_o_gk( scale: float, cu_seqlens: torch.LongTensor | None = None, chunk_size: int = 64, + chunk_indices: torch.LongTensor | None = None, ): B, T, H, K, V = *q.shape, v.shape[-1] BT = chunk_size - chunk_indices = ( - prepare_chunk_indices(cu_seqlens, chunk_size) - 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 = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) def grid(meta): @@ -850,6 +849,178 @@ def chunk_gla_fwd_o_gk( return o +@triton.jit +def softplus_fwd(x): + """Standard softplus: log(1 + exp(x)), with linear approx for large x.""" + return tl.where(x < 20.0, log(1.0 + exp(x)), x) + + +@triton.heuristics( + { + "HAS_BIAS": lambda args: args["dt_bias"] is not None, + "HAS_SCALE": lambda args: args["scale"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + "USE_LOWER_BOUND": lambda args: args["lower_bound"] is not None, + } +) +@triton.autotune( + configs=[ + triton.Config({"BS": BS}, num_warps=num_warps) + for BS in BS_LIST + for num_warps in [2, 4, 8] + ], + key=["H", "S", "BT", "IS_VARLEN"], +) +@triton.jit(do_not_specialize=["T"]) +def kda_gate_chunk_cumsum_vector_kernel( + s, + A_log, + dt_bias, + o, + scale, + cu_seqlens, + chunk_indices, + lower_bound, + T, + H: tl.constexpr, + S: tl.constexpr, + BT: tl.constexpr, + BS: tl.constexpr, + HAS_BIAS: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, + USE_LOWER_BOUND: tl.constexpr, +): + i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2) + 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 + + p_s = tl.make_block_ptr( + s + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + p_o = tl.make_block_ptr( + o + (bos * H + i_h) * S, + (T, S), + (H * S, 1), + (i_t * BT, i_s * BS), + (BT, BS), + (1, 0), + ) + # [BT, BS] + b_s = tl.load(p_s, boundary_check=(0, 1)).to(tl.float32) + + if HAS_BIAS: + p_b = tl.make_block_ptr( + dt_bias + i_h * S, + (S,), + (1,), + (i_s * BS,), + (BS,), + (0,), + ) + b_bias = tl.load(p_b, boundary_check=(0,)).to(tl.float32) + b_s = b_s + b_bias[None, :] + + b_A = tl.load(A_log + i_h).to(tl.float32) + if not USE_LOWER_BOUND: + # Standard gate: -exp(A_log) * softplus(g + bias) + b_gate = -exp(b_A) * softplus_fwd(b_s) + else: + # Safe gate: lower_bound * sigmoid(exp(A_log) * (g + bias)) + b_gate = lower_bound * tl.sigmoid(exp(b_A) * b_s) + + # Chunk-local cumulative sum + b_o = tl.cumsum(b_gate, axis=0) + + if HAS_SCALE: + b_o *= scale + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1)) + + +def kda_gate_chunk_cumsum( + g: torch.Tensor, + A_log: torch.Tensor, + chunk_size: int, + scale: float = None, + dt_bias: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, + output_dtype: Optional[torch.dtype] = torch.float, + chunk_indices: Optional[torch.LongTensor] = None, + lower_bound: Optional[float] = None, +) -> torch.Tensor: + """ + Fused KDA gate activation + chunk-local cumulative sum. + + Combines two memory-bound kernels into one: + 1. Gate activation: g = -exp(A_log) * softplus(raw_g + dt_bias) + 2. Chunk-local cumsum along the time axis + + Args: + g: Raw gate tensor of shape [B, T, H, K] (before activation). + A_log: Per-head log-scale parameter, [H] elements (any shape, numel=H). + chunk_size: Chunk size for cumsum (must be power of 2). + scale: Optional scale factor applied to output. + dt_bias: Optional per-head bias, flat [H*K] elements. + cu_seqlens: Cumulative sequence lengths for variable-length input. + output_dtype: Output dtype (default float32). + chunk_indices: Pre-computed chunk indices for varlen mode. + lower_bound: If set, use safe gate: lower_bound * sigmoid(exp(A_log) * g). + + Returns: + Cumulative-summed gated tensor of shape [B, T, H, K]. + """ + if cu_seqlens is not None: + assert ( + g.shape[0] == 1 + ), "Only batch size 1 is supported when cu_seqlens are provided" + assert len(g.shape) == 4 + B, T, H, S = g.shape + BT = chunk_size + if chunk_indices is None and cu_seqlens is not None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices) + assert chunk_size == 2 ** ( + chunk_size.bit_length() - 1 + ), "chunk_size must be a power of 2" + + g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype) + + def grid(meta): + return (cdiv(meta["S"], meta["BS"]), NT, B * H) + + kda_gate_chunk_cumsum_vector_kernel[grid]( + s=g_org, + A_log=A_log, + dt_bias=dt_bias, + o=g, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=lower_bound, + T=T, + H=H, + S=S, + BT=BT, + ) + return g + + def chunk_kda_fwd( q: torch.Tensor, k: torch.Tensor, @@ -859,10 +1030,40 @@ def chunk_kda_fwd( scale: float, initial_state: torch.Tensor, initial_state_indices: torch.Tensor, - cu_seqlens: torch.LongTensor | None = None, + cu_seqlens: Optional[torch.LongTensor] = None, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + lower_bound: Optional[float] = None, ): chunk_size = 64 - g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens) + # Pre-compute chunk indices once and thread through all downstream kernels. + # Without this, each of the 4 callees would recompute independently. + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + + if A_log is not None: + # Fused: gate activation + chunk-local cumsum in one kernel. + # g is raw gate (before activation); A_log, dt_bias drive the activation. + g = kda_gate_chunk_cumsum( + g, + A_log=A_log, + chunk_size=chunk_size, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=lower_bound, + ) + else: + # g is already gate-activated by caller; just do cumsum. + g = chunk_local_cumsum( + g, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) # Fused: scaled_dot_kkt + solve_tril + recompute_w_u w, u, _, kg, Aqk, _ = chunk_kda_fwd_intra( @@ -874,6 +1075,7 @@ def chunk_kda_fwd( scale=scale, cu_seqlens=cu_seqlens, chunk_size=chunk_size, + chunk_indices=chunk_indices, ) h, v_new = chunk_gated_delta_rule_fwd_h( @@ -884,6 +1086,7 @@ def chunk_kda_fwd( initial_state=initial_state, initial_state_indices=initial_state_indices, cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, ) del w, u, kg o = chunk_gla_fwd_o_gk( @@ -896,6 +1099,7 @@ def chunk_kda_fwd( scale=scale, cu_seqlens=cu_seqlens, chunk_size=chunk_size, + chunk_indices=chunk_indices, ) del Aqk, v_new, h return o @@ -911,7 +1115,10 @@ def chunk_kda( initial_state: torch.Tensor = None, initial_state_indices: torch.Tensor = None, use_qk_l2norm_in_kernel: bool = False, - cu_seqlens: torch.LongTensor | None = None, + cu_seqlens: Optional[torch.LongTensor] = None, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + lower_bound: Optional[float] = None, **kwargs, ): if scale is None: @@ -931,124 +1138,8 @@ def chunk_kda( initial_state=initial_state, initial_state_indices=initial_state_indices, cu_seqlens=cu_seqlens, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, ) return o - - -@triton.autotune( - configs=[ - triton.Config({"BT": bt}, num_warps=nw, num_stages=ns) - for bt in BT_LIST_AUTOTUNE - for nw in NUM_WARPS_AUTOTUNE - for ns in [2, 3] - ], - key=["H", "D"], -) -@triton.jit -def kda_gate_fwd_kernel( - g, - A, - y, - g_bias, - beta: tl.constexpr, - threshold: tl.constexpr, - T, - H, - D: tl.constexpr, - BT: tl.constexpr, - BD: tl.constexpr, - HAS_BIAS: tl.constexpr, -): - i_t, i_h = tl.program_id(0), tl.program_id(1) - n_t = i_t * BT - - b_a = tl.load(A + i_h).to(tl.float32) - b_a = -tl.exp(b_a) - - stride_row = H * D - stride_col = 1 - - g_ptr = tl.make_block_ptr( - base=g + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) - - y_ptr = tl.make_block_ptr( - base=y + i_h * D, - shape=(T, D), - strides=(stride_row, stride_col), - offsets=(n_t, 0), - block_shape=(BT, BD), - order=(1, 0), - ) - - b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32) - - if HAS_BIAS: - n_d = tl.arange(0, BD) - bias_mask = n_d < D - b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to( - tl.float32 - ) - b_g = b_g + b_bias[None, :] - - # softplus(x, beta) = (1/beta) * log(1 + exp(beta * x)) - # When beta * x > threshold, use linear approximation x - # Use threshold to switch to linear when beta*x > threshold - g_scaled = b_g * beta - use_linear = g_scaled > threshold - sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled))) - b_y = b_a * sp - - tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1)) - - -def fused_kda_gate( - g: torch.Tensor, - A: torch.Tensor, - head_k_dim: int, - g_bias: torch.Tensor | None = None, - beta: float = 1.0, - threshold: float = 20.0, -) -> torch.Tensor: - """ - Forward pass for KDA gate: - input g: [..., H*D] - param A: [H] or [1, 1, H, 1] - beta: softplus beta parameter - threshold: softplus threshold parameter - return : [..., H, D] - """ - orig_shape = g.shape[:-1] - - g = g.view(-1, g.shape[-1]) - T = g.shape[0] - HD = g.shape[1] - H = A.numel() - assert H * head_k_dim == HD - - y = torch.empty_like(g, dtype=torch.float32) - - def grid(meta): - return (cdiv(T, meta["BT"]), H) - - kda_gate_fwd_kernel[grid]( - g, - A, - y, - g_bias, - beta, - threshold, - T, - H, - head_k_dim, - BD=next_power_of_2(head_k_dim), - HAS_BIAS=g_bias is not None, - ) - - y = y.view(*orig_shape, H, head_k_dim) - return y diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index fb27462b0..bb77eaafe 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/python/sglang/srt/layers/attention/linear/kda_backend.py @@ -252,6 +252,9 @@ class KDAAttnBackend(MambaAttnBackendBase): ssm_states=ssm_states, cache_indices=cache_indices, query_start_loc=query_start_loc, + A_log=layer.A_log, + dt_bias=layer.dt_bias, + lower_bound=getattr(layer, "lower_bound", None), ) return core_attn_out diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py b/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py index 2ab8a5e0b..2d6d54d9e 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py @@ -1,3 +1,5 @@ +from typing import Optional + import torch from sglang.srt.layers.attention.linear.kernels.kernel_backend import ( @@ -58,6 +60,9 @@ class TritonKDAKernel(LinearAttnKernelBase): ssm_states: torch.Tensor, cache_indices: torch.Tensor, query_start_loc: torch.Tensor, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + lower_bound: Optional[float] = None, **kwargs, ) -> torch.Tensor: return chunk_kda( @@ -70,4 +75,7 @@ class TritonKDAKernel(LinearAttnKernelBase): initial_state_indices=cache_indices, use_qk_l2norm_in_kernel=True, cu_seqlens=query_start_loc, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, ) diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index 795e38b20..fbd8ebbfe 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -15,7 +15,6 @@ from sglang.srt.distributed import ( ) from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.layers.attention.fla.fused_norm_gate import FusedRMSNormGated -from sglang.srt.layers.attention.fla.kda import fused_kda_gate from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( @@ -382,11 +381,13 @@ class KimiDeltaAttention(nn.Module): hidden_states ) - # fused_kda_gate is fused to KimiLinearAttentionBackend with decode + # For prefill: raw gate is passed to chunk_kda_fwd, which fuses gate + # activation with chunk_local_cumsum (kda_gate_chunk_cumsum kernel). + # For decode: gate activation is handled inside fused_recurrent kernel. if not forward_batch.forward_mode.is_decode(): - forget_gate = fused_kda_gate( - forget_gate, self.A_log, self.head_dim, g_bias=self.dt_bias - ) + forget_gate = forget_gate.unflatten( + -1, (-1, self.head_dim) + ) # [T, H*K] -> [T, H, K] beta = beta.float().sigmoid() forget_gate = forget_gate.unsqueeze(0) beta = beta.unsqueeze(0) diff --git a/test/registered/attention/test_kda_kernels.py b/test/registered/attention/test_kda_kernels.py index f3b18ad59..51f4e9db1 100644 --- a/test/registered/attention/test_kda_kernels.py +++ b/test/registered/attention/test_kda_kernels.py @@ -2,10 +2,15 @@ import unittest import torch +from sglang.srt.layers.attention.fla.cumsum import chunk_local_cumsum from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import ( fused_sigmoid_gating_delta_rule_update, ) -from sglang.srt.layers.attention.fla.kda import fused_kda_gate, fused_recurrent_kda +from sglang.srt.layers.attention.fla.index import prepare_chunk_indices +from sglang.srt.layers.attention.fla.kda import ( + fused_recurrent_kda, + kda_gate_chunk_cumsum, +) from sglang.test.ci.ci_register import register_cuda_ci register_cuda_ci(est_time=8, suite="stage-b-test-1-gpu-large") @@ -95,7 +100,15 @@ class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase): def run_kda(self): b = self.beta.float().sigmoid() - g = fused_kda_gate(self.a, self.A_log, self.head_dim, g_bias=self.dt_bias) + # Reference gate activation using torch ops: + # g = -exp(A_log) * softplus(raw_g + dt_bias) + H, K = self.local_num_heads, self.head_dim + raw_g = self.a.float() # [1, T, H*K] + if self.dt_bias is not None: + raw_g = raw_g + self.dt_bias.float() + g = -torch.exp( + self.A_log.float().view(1, 1, H, 1) + ) * torch.nn.functional.softplus(raw_g.view(1, -1, H, K)) initial_state = self.ssm_states[self.cache_indices].clone() core_attn_out, last_state = fused_recurrent_kda( q=self.q, @@ -121,5 +134,101 @@ class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase): self.assertTrue(torch.allclose(last_state, last_state_ref)) +@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA") +class TestKDAGateChunkCumsum(unittest.TestCase): + """Test kda_gate_chunk_cumsum against torch reference (gate activation + cumsum).""" + + CHUNK_SIZE = 64 + + def _ref_gate_cumsum(self, raw_g, A_log, dt_bias, cu_seqlens, chunk_size): + """Reference: torch gate activation then chunk_local_cumsum.""" + B, T, H, K = raw_g.shape + g = raw_g.float() + if dt_bias is not None: + g = g + dt_bias.float().view(1, 1, H, K) + g = -torch.exp(A_log.float().view(1, 1, H, 1)) * torch.nn.functional.softplus(g) + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, chunk_size) + if cu_seqlens is not None + else None + ) + return chunk_local_cumsum( + g, chunk_size=chunk_size, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices + ) + + def _run_case(self, B, T_per_seq, H, K, use_bias, use_varlen): + T = B * T_per_seq + torch.manual_seed(42) + raw_g = torch.randn(1, T, H, K, dtype=torch.bfloat16, device="cuda") + A_log = torch.randn(H, dtype=torch.float32, device="cuda") * 0.5 + dt_bias = ( + torch.randn(H * K, dtype=torch.float32, device="cuda") * 0.1 + if use_bias + else None + ) + cu_seqlens = ( + torch.arange( + 0, (B + 1) * T_per_seq, T_per_seq, dtype=torch.long, device="cuda" + ) + if use_varlen + else None + ) + + out_fused = kda_gate_chunk_cumsum( + raw_g, + A_log=A_log, + chunk_size=self.CHUNK_SIZE, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + ) + out_ref = self._ref_gate_cumsum( + raw_g, A_log, dt_bias, cu_seqlens, self.CHUNK_SIZE + ) + + max_diff = (out_fused - out_ref).abs().max().item() + rel_diff = max_diff / (out_ref.abs().mean().item() + 1e-8) + return max_diff, rel_diff + + def test_varlen_with_bias(self): + max_diff, rel_diff = self._run_case( + B=4, T_per_seq=256, H=16, K=128, use_bias=True, use_varlen=True + ) + self.assertLess( + max_diff, 1e-3, f"max_diff={max_diff:.2e}, rel_diff={rel_diff:.2e}" + ) + + def test_varlen_no_bias(self): + max_diff, rel_diff = self._run_case( + B=4, T_per_seq=256, H=16, K=128, use_bias=False, use_varlen=True + ) + self.assertLess( + max_diff, 1e-3, f"max_diff={max_diff:.2e}, rel_diff={rel_diff:.2e}" + ) + + def test_fixed_len_with_bias(self): + max_diff, rel_diff = self._run_case( + B=4, T_per_seq=256, H=16, K=128, use_bias=True, use_varlen=False + ) + self.assertLess( + max_diff, 1e-3, f"max_diff={max_diff:.2e}, rel_diff={rel_diff:.2e}" + ) + + def test_single_seq_long(self): + max_diff, rel_diff = self._run_case( + B=1, T_per_seq=2048, H=16, K=128, use_bias=True, use_varlen=True + ) + self.assertLess( + max_diff, 1e-3, f"max_diff={max_diff:.2e}, rel_diff={rel_diff:.2e}" + ) + + def test_small_head_dim(self): + max_diff, rel_diff = self._run_case( + B=4, T_per_seq=128, H=8, K=64, use_bias=True, use_varlen=True + ) + self.assertLess( + max_diff, 1e-3, f"max_diff={max_diff:.2e}, rel_diff={rel_diff:.2e}" + ) + + if __name__ == "__main__": unittest.main()