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

[KDA] Add FlashKDA prefill backend for safe-gate KDA linear attention (#29472)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-07-02 11:37:04 +08:00
committed by GitHub
co-authored by luoyuan.luo
parent 6eeb9871e5
commit eb23f9a86b
5 changed files with 499 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/linear/kda_backend.py
+16 -1
View File
@@ -60,6 +60,12 @@ class KDAKernelDispatcher:
if prefill_backend.is_triton():
self.extend_kernel = triton_kernel
elif prefill_backend.is_flashkda():
from sglang.srt.layers.attention.linear.kernels.kda_flashkda import (
FlashKDAKernel,
)
self.extend_kernel = FlashKDAKernel()
elif prefill_backend.is_cutedsl():
if not is_cuda():
raise ValueError("KDA CuTe DSL backend requires CUDA")
@@ -81,7 +87,8 @@ class KDAKernelDispatcher:
else:
raise ValueError(
f"Unsupported KDA prefill backend: {prefill_backend}. "
"KDA supports 'triton' or 'cutedsl' (cutedsl prefill needs SM100)."
"KDA supports 'triton', 'flashkda', or 'cutedsl' "
"(cutedsl prefill needs SM100)."
)
self.supports_packed_decode = getattr(
@@ -367,6 +374,14 @@ class KDAAttnBackend(MambaAttnBackendBase):
A_log=layer.A_log,
dt_bias=layer.dt_bias,
lower_bound=getattr(layer, "lower_bound", None),
extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
# target_verify / draft_extend_v2 also reach forward_extend; they must
# stay rollback-able, so a kernel that commits state in place (e.g.
# FlashKDA) must not run for them.
is_spec_decode=(
forward_batch.forward_mode.is_target_verify()
or forward_batch.forward_mode.is_draft_extend_v2()
),
)
return core_attn_out
python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py
+257
View File
@@ -0,0 +1,257 @@
from typing import Optional
import torch
from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
LinearAttnKernelBase,
)
# FlashKDA chunk size. Sequences shorter than this fall back to Triton.
_FLASHKDA_CHUNK_SIZE = 64
# FlashKDA's max sequence length, Batches whose longest sequence exceeds this
# fall back to Triton for the whole batch.
_FLASHKDA_MAX_SEQ_LEN = 2048
def _load_flash_kda():
"""Import the optional ``flash_kda`` CUTLASS module."""
try:
import flash_kda
except ImportError as e:
raise ImportError(
"The 'flashkda' KDA prefill backend requires the flash_kda module, "
"which is not installed. Install it from source:\n"
" pip install git+https://github.com/MoonshotAI/FlashKDA.git"
) from e
return flash_kda
def _triton_fallback(
q,
k,
v,
g,
beta,
ssm_states,
cache_indices,
query_start_loc,
A_log=None,
dt_bias=None,
lower_bound=None,
):
"""Fall back to the Triton chunk_kda kernel (handles all preprocessing).
`g` is the RAW gate; chunk_kda applies the gate activation internally when
A_log is provided, so A_log/dt_bias/lower_bound must be threaded through too
-- otherwise the fallback silently skips activation. chunk_kda updates the
ssm state in-place via cache_indices and returns only the output tensor.
"""
from sglang.srt.layers.attention.fla.kda import chunk_kda
return chunk_kda(
q=q,
k=k,
v=v,
g=g,
beta=beta,
initial_state=ssm_states,
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,
)
class FlashKDAKernel(LinearAttnKernelBase):
"""FlashKDA (MoonshotAI) fully-fused CUTLASS KDA prefill backend.
Wraps the external ``flash_kda`` package (https://github.com/MoonshotAI/FlashKDA).
FlashKDA fuses q/k L2 norm, beta sigmoid, and the KDA gate *inside* the
kernel, so we pass RAW tensors plus ``A_log``/``dt_bias``/``lower_bound``.
It is prefill-only, bf16, K == V == 128, HV == H (no GVA), and requires the
safe (bounded) gate (``lower_bound`` set). The non-safe path and sequences
outside [chunk_size, max_seq_len] fall back to Triton ``chunk_kda``.
Requires an SM90+ GPU with the ``flash_kda`` package.
"""
def decode(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
*,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
**kwargs,
) -> torch.Tensor:
raise NotImplementedError("FlashKDAKernel only supports prefill (extend)")
def extend(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
*,
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,
extend_seq_lens_cpu: Optional[list] = None,
is_spec_decode: bool = False,
**kwargs,
) -> torch.Tensor:
if self._should_fall_back(
lower_bound, is_spec_decode, query_start_loc, extend_seq_lens_cpu
):
return _triton_fallback(
q,
k,
v,
g,
beta,
ssm_states,
cache_indices,
query_start_loc,
A_log=A_log,
dt_bias=dt_bias,
lower_bound=lower_bound,
)
return self._flashkda_extend(
q,
k,
v,
g,
beta,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
A_log=A_log,
dt_bias=dt_bias,
lower_bound=lower_bound,
)
@staticmethod
def _should_fall_back(
lower_bound: Optional[float],
is_spec_decode: bool,
query_start_loc: torch.Tensor,
extend_seq_lens_cpu: Optional[list],
) -> bool:
"""Whether to use the Triton chunk_kda path instead of the fused kernel."""
# Safe-gate only: the fused kernel does not support the unbounded gate
# (-exp(A_log)*softplus); those models leave lower_bound unset.
if lower_bound is None:
return True
# FlashKDA writes the committed recurrent state back in place, so it is
# unsafe for speculative verify / draft-extend forwards (which must stay
# rollback-able). Those reach this backend through forward_extend, so
# gate them here rather than relying on the decode/target_verify stubs.
if is_spec_decode:
return True
# Short sequences (< chunk size) and long sequences (> the crossover
# where Triton's chunked prefill wins) are faster on Triton. Read the
# per-request lengths from the CPU-side extend_seq_lens to avoid a
# GPU->CPU sync on every layer; derive from query_start_loc (one sync)
# only if they are unavailable.
if extend_seq_lens_cpu is not None:
if torch.is_tensor(extend_seq_lens_cpu):
lo = int(extend_seq_lens_cpu.min())
hi = int(extend_seq_lens_cpu.max())
else:
lo = min(extend_seq_lens_cpu)
hi = max(extend_seq_lens_cpu)
else:
seq_lens = query_start_loc[1:] - query_start_loc[:-1]
lo_t, hi_t = torch.aminmax(seq_lens)
lo, hi = int(lo_t), int(hi_t)
return lo < _FLASHKDA_CHUNK_SIZE or hi > _FLASHKDA_MAX_SEQ_LEN
def _flashkda_extend(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
*,
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,
) -> torch.Tensor:
flash_kda = _load_flash_kda()
# Input shapes (varlen, B == 1, matching chunk_kda's contract):
# q, k = [1, packed_seq, H, K] v = [1, packed_seq, HV, V]
# g = [1, packed_seq, HV, K] beta = [1, packed_seq, H]
# flash_kda wants these 4D tensors directly and RAW (it fuses l2norm /
# beta sigmoid / gate activation in-kernel).
num_heads = q.shape[2]
head_dim = q.shape[3]
scale = head_dim**-0.5
q = q.contiguous()
k = k.contiguous()
v = v.contiguous()
g = g.contiguous()
# KimiDeltaAttention.forward already applies sigmoid to beta on the
# prefill path, but flash_kda expects beta LOGITS (it sigmoids
# internally). Invert back so the kernel recovers the intended value:
# sigmoid(logit(p)) == p. (triton/cuLA consume the post-sigmoid beta.)
beta = torch.logit(beta.float().clamp_(1e-7, 1.0 - 1e-7)).to(torch.bfloat16)
beta = beta.contiguous()
# flash_kda wants A_log [H] fp32 and dt_bias [H, K] fp32. The model
# stores A_log as [1, 1, H, 1] and dt_bias as 1D [H*K], so reshape both.
A_log = A_log.reshape(-1).float().contiguous()
if dt_bias is not None:
dt_bias = dt_bias.reshape(num_heads, -1).float().contiguous()
# cu_seqlens must be int64 for flash_kda (FLA casts to long).
cu_seqlens = query_start_loc.to(torch.int64)
# flash_kda varlen state is [N, H, V, K] -- the SAME layout as sglang's
# KDA pool, so no transpose is needed. Advanced indexing copies, so the
# final state is written back in-place below (matching chunk_kda).
initial_state = ssm_states[cache_indices].contiguous()
out_buf = torch.empty_like(v)
final_state = torch.empty_like(initial_state)
flash_kda.fwd(
q,
k,
v,
g,
beta,
scale,
out_buf,
A_log,
dt_bias,
lower_bound,
initial_state=initial_state,
final_state=final_state,
cu_seqlens=cu_seqlens,
)
ssm_states[cache_indices] = final_state
# out_buf is already [1, packed_seq, HV, V].
return out_buf
python/sglang/srt/layers/attention/linear/utils.py
+4
View File
@@ -16,6 +16,7 @@ class LinearAttnKernelBackend(Enum):
TRITON = "triton"
CUTEDSL = "cutedsl"
FLASHINFER = "flashinfer"
FLASHKDA = "flashkda"
CUSTOM = "custom"
@classmethod
@@ -31,6 +32,9 @@ class LinearAttnKernelBackend(Enum):
def is_flashinfer(self):
return self == LinearAttnKernelBackend.FLASHINFER
def is_flashkda(self):
return self == LinearAttnKernelBackend.FLASHKDA
def is_custom(self):
return self == LinearAttnKernelBackend.CUSTOM
python/sglang/srt/server_args.py
+21 -1
View File
@@ -315,7 +315,7 @@ MAMBA_RADIX_CACHE_STRATEGY_CHOICES = [
MAMBA_BACKEND_CHOICES = ["triton", "flashinfer"]
LINEAR_ATTN_KERNEL_BACKEND_CHOICES = ["triton", "cutedsl", "flashinfer"]
LINEAR_ATTN_KERNEL_BACKEND_CHOICES = ["triton", "cutedsl", "flashinfer", "flashkda"]
# Allow external code to add more choices
@@ -5133,6 +5133,26 @@ class ServerArgs:
# SM100+ FlashInfer GDN decode requires bf16 state; SM90 uses float32.
decode = self.linear_attn_decode_backend or self.linear_attn_backend
# FlashKDA is a prefill-only KDA kernel (no decode kernel) but shares the
# backend choice list, so guard it from being selected for decode: error
# on an explicit --linear-attn-decode-backend flashkda, and fall back to
# triton decode when it was only inherited from base=flashkda (prefill
# keeps FlashKDA).
if decode == "flashkda":
if self.linear_attn_decode_backend == "flashkda":
raise ValueError(
"--linear-attn-decode-backend flashkda is not supported: "
"FlashKDA is prefill-only. Use "
"--linear-attn-prefill-backend flashkda (decode stays on triton)."
)
self.linear_attn_decode_backend = "triton"
decode = "triton"
logger.info(
"FlashKDA is prefill-only; using triton for KDA decode "
"(FlashKDA stays on prefill)."
)
if (
decode == "flashinfer"
and self.mamba_ssm_dtype != "bfloat16"
test/registered/attention/test_kda_prefill_flashkda.py
+201
View File
@@ -0,0 +1,201 @@
"""Correctness test for the FlashKDA prefill backend (safe-gate KDA).
Validates ``FlashKDAKernel`` (the external ``flash_kda`` fused CUTLASS kernel)
against the production Triton ``chunk_kda`` safe-gate reference. FlashKDA only
implements the safe/bounded gate, so this exercises ``lower_bound=-5``.
Requires an SM90+ GPU and the ``flash_kda`` package
(``pip install git+https://github.com/MoonshotAI/FlashKDA.git``); skips
otherwise. Mirrors ``test_kda_prefill_cutedsl.py``.
"""
import pytest
import torch
from sglang.test.ci.ci_register import register_cuda_ci
# SM90+ single-GPU kernel-unit suite. Disabled in CI: flash_kda is not in the
# public runner image, and the module-level pytest.skip aborts non-zero under
# `python3 file.py`. Drop `disabled=` once flash_kda ships in the runner image;
# still runs locally / on internal CI where it is installed.
register_cuda_ci(
est_time=20,
stage="base-b",
runner_config="1-gpu-large",
disabled="flash_kda not in public CI runner image (only on Ant-internal PyPI)",
)
if not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 9):
pytest.skip(
"FlashKDA requires CUDA SM90+ (Hopper or newer).",
allow_module_level=True,
)
try:
import flash_kda # noqa: F401
except ImportError:
pytest.skip(
"flash_kda not installed "
"(pip install git+https://github.com/MoonshotAI/FlashKDA.git).",
allow_module_level=True,
)
from sglang.srt.layers.attention.fla.kda import chunk_kda # noqa: E402
from sglang.srt.layers.attention.linear.kernels.kda_flashkda import ( # noqa: E402
FlashKDAKernel,
)
LOWER_BOUND = -5.0
H, K, V = 16, 128, 128 # FlashKDA requires K == V == 128, HV == H
def _cos(a: torch.Tensor, b: torch.Tensor) -> float:
return torch.nn.functional.cosine_similarity(
a.float().flatten(), b.float().flatten(), dim=0
).item()
def _make_inputs(seq_lens):
n = len(seq_lens)
cu = torch.zeros(n + 1, device="cuda", dtype=torch.int32)
cu[1:] = torch.tensor(seq_lens, device="cuda").cumsum(0)
total = int(cu[-1].item())
idx = torch.arange(n, device="cuda", dtype=torch.int32)
return dict(
cu=cu,
idx=idx,
# RAW q/k (FlashKDA L2-norms internally; chunk_kda via use_qk_l2norm).
q=torch.randn(1, total, H, K, device="cuda", dtype=torch.bfloat16) * 0.5,
k=torch.randn(1, total, H, K, device="cuda", dtype=torch.bfloat16) * 0.5,
v=torch.randn(1, total, H, V, device="cuda", dtype=torch.bfloat16) * 0.5,
g=torch.randn(1, total, H, K, device="cuda", dtype=torch.bfloat16) * 0.5,
# post-sigmoid beta in [0.1, 0.9] (FlashKDA inverts to logits internally).
beta=(torch.rand(1, total, H, device="cuda") * 0.8 + 0.1).to(torch.bfloat16),
A_log=torch.randn(1, 1, H, 1, device="cuda", dtype=torch.float32) * 0.5,
dt_bias=torch.randn(H * K, device="cuda", dtype=torch.float32) * 0.1,
pool=torch.randn(n, H, V, K, device="cuda", dtype=torch.float32) * 0.1,
)
def _chunk_kda_ref(d, lower_bound):
"""Triton chunk_kda reference. chunk_kda mutates g/v and the state in place,
so feed clones; returns (output, updated_state_slots)."""
st = d["pool"].clone()
out = chunk_kda(
q=d["q"].clone(),
k=d["k"].clone(),
v=d["v"].clone(),
g=d["g"].clone(),
beta=d["beta"].clone(),
initial_state=st,
initial_state_indices=d["idx"],
use_qk_l2norm_in_kernel=True,
cu_seqlens=d["cu"],
A_log=d["A_log"],
dt_bias=d["dt_bias"],
lower_bound=lower_bound,
)
return out, st[d["idx"]]
@pytest.mark.parametrize("seq_lens", [[128], [128, 384, 512], [96] * 4])
def test_flashkda_matches_triton_safe_gate(seq_lens):
torch.manual_seed(len(seq_lens))
d = _make_inputs(seq_lens)
ref_out, ref_state = _chunk_kda_ref(d, LOWER_BOUND)
st_fk = d["pool"].clone()
out = FlashKDAKernel().extend(
d["q"].clone(),
d["k"].clone(),
d["v"].clone(),
d["g"].clone(),
d["beta"].clone(),
ssm_states=st_fk,
cache_indices=d["idx"],
query_start_loc=d["cu"],
A_log=d["A_log"],
dt_bias=d["dt_bias"],
lower_bound=LOWER_BOUND,
extend_seq_lens_cpu=seq_lens,
)
torch.cuda.synchronize()
assert torch.isfinite(out).all(), "FlashKDA output has non-finite values"
assert torch.isfinite(st_fk).all(), "FlashKDA final state has non-finite values"
# bf16 cross-implementation noise (chunk=16 CUTLASS vs chunk=64 Triton);
# measured cos ~0.985 output / ~0.9999 state on H20-3e and B200.
assert _cos(ref_out, out) > 0.95, f"output cos too low: {_cos(ref_out, out):.4f}"
assert (
_cos(ref_state, st_fk[d["idx"]]) > 0.99
), f"state cos too low: {_cos(ref_state, st_fk[d['idx']]):.4f}"
def test_flashkda_falls_back_without_lower_bound():
"""Unbounded gate (lower_bound=None): FlashKDA must route to the Triton
chunk_kda fallback (it only supports the safe gate)."""
torch.manual_seed(0)
d = _make_inputs([256])
ref_out, _ = _chunk_kda_ref(d, None)
st_fk = d["pool"].clone()
out = FlashKDAKernel().extend(
d["q"].clone(),
d["k"].clone(),
d["v"].clone(),
d["g"].clone(),
d["beta"].clone(),
ssm_states=st_fk,
cache_indices=d["idx"],
query_start_loc=d["cu"],
A_log=d["A_log"],
dt_bias=d["dt_bias"],
lower_bound=None,
extend_seq_lens_cpu=[256],
)
torch.cuda.synchronize()
assert torch.isfinite(out).all()
# Same Triton code path as the reference -> matches closely.
assert _cos(ref_out, out) > 0.999, f"fallback cos too low: {_cos(ref_out, out):.4f}"
def test_flashkda_spec_verify_falls_back():
"""Speculative verify / draft-extend (is_spec_decode=True) must route to the
Triton fallback even with a safe gate and in-range seqs -- FlashKDA commits
the recurrent state in place, which would break draft rollback."""
torch.manual_seed(0)
d = _make_inputs([256])
ref_out, _ = _chunk_kda_ref(d, LOWER_BOUND)
st_fk = d["pool"].clone()
out = FlashKDAKernel().extend(
d["q"].clone(),
d["k"].clone(),
d["v"].clone(),
d["g"].clone(),
d["beta"].clone(),
ssm_states=st_fk,
cache_indices=d["idx"],
query_start_loc=d["cu"],
A_log=d["A_log"],
dt_bias=d["dt_bias"],
lower_bound=LOWER_BOUND,
extend_seq_lens_cpu=[256],
is_spec_decode=True,
)
torch.cuda.synchronize()
assert torch.isfinite(out).all()
# Took the Triton fallback (not FlashKDA) -> matches chunk_kda closely. If
# FlashKDA had run, the cross-impl cos would be ~0.985 and this would fail.
assert (
_cos(ref_out, out) > 0.999
), f"spec-decode did not fall back: {_cos(ref_out, out):.4f}"
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v"]))