diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index d2168c3d2..8a4784d8b 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/python/sglang/srt/layers/attention/linear/kda_backend.py @@ -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 diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py b/python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py new file mode 100644 index 000000000..2720dd228 --- /dev/null +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py @@ -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 diff --git a/python/sglang/srt/layers/attention/linear/utils.py b/python/sglang/srt/layers/attention/linear/utils.py index 13d2b271b..2d7bb7858 100644 --- a/python/sglang/srt/layers/attention/linear/utils.py +++ b/python/sglang/srt/layers/attention/linear/utils.py @@ -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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index e6bc49376..7d29ef345 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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" diff --git a/test/registered/attention/test_kda_prefill_flashkda.py b/test/registered/attention/test_kda_prefill_flashkda.py new file mode 100644 index 000000000..954d1dadc --- /dev/null +++ b/test/registered/attention/test_kda_prefill_flashkda.py @@ -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"]))