minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

[KDA] Fuse gate+cumsum and reuse chunk index for KDA (#23038)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-04-21 17:54:20 +08:00
committed by GitHub
co-authored by luoyuan.luo
parent 8589b92a89
commit 48daa831ea
10 changed files with 582 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/bench_linear_attention/bench_fused_gate_cumsum.py
+213
View File
@@ -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())
python/sglang/srt/layers/attention/fla/chunk.py
+4 -1
View File
@@ -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,
python/sglang/srt/layers/attention/fla/chunk_delta_h.py
+3 -5
View File
@@ -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
python/sglang/srt/layers/attention/fla/chunk_intra.py
+1
View File
@@ -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
python/sglang/srt/layers/attention/fla/cumsum.py
+9 -8
View File
@@ -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(
python/sglang/srt/layers/attention/fla/kda.py
+224 -133
View File
@@ -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
python/sglang/srt/layers/attention/linear/kda_backend.py
+3
View File
@@ -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
python/sglang/srt/layers/attention/linear/kernels/kda_triton.py
+8
View File
@@ -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,
)
python/sglang/srt/models/kimi_linear.py
+6 -5
View File
@@ -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)
test/registered/attention/test_kda_kernels.py
+111 -2
View File
@@ -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()