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

[Kernel] Enable Helion backend for Kimi Delta-Attention (#32593)

Browse Source 
Co-authored-by: Ethan Che <eche@meta.com>
This commit is contained in:
ethche
2026-08-14 22:22:46 -07:00
committed by GitHub
co-authored by Ethan Che
parent 3adbbec2fd
commit aeee1562e6
16 changed files with 4248 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/server_args.py
+21 -21
View File
@@ -377,6 +377,7 @@ LINEAR_ATTN_KERNEL_BACKEND_CHOICES = [
"flashkda",
"nvidia_kda",
"ptx_kda",
"helion",
]
@@ -2571,7 +2572,7 @@ class ServerArgs:
linear_attn_backend: A[
str,
Arg(
help="The default kernel backend for linear attention (GDN/KDA). Can be overridden per-mode by --linear-attn-decode-backend and --linear-attn-prefill-backend.",
help="The default kernel backend for linear attention (GDN/KDA). Can be overridden per-mode by --linear-attn-decode-backend and --linear-attn-prefill-backend. The Helion backend is KDA-only.",
choices=LINEAR_ATTN_KERNEL_BACKEND_CHOICES,
),
NS("exec.mamba"),
@@ -2606,11 +2607,10 @@ class ServerArgs:
bool,
"Enable the ReplaySSM buffered output-only linear-attn decode kernel. "
"Primarily a GDN (scalar-gate) decode-bandwidth optimization (~1.2-1.5x "
"at batch >= 64). The unified kernel also supports KDA (per-K gate) and "
"is numerically correct, but KDA decode is SLOWER than the packed "
"baseline (the per-K g_cache is K x larger and the reconstruction "
"refolds the per-K decay every step), so it is not recommended for KDA "
"models. Requires the Triton linear-attn decode backend and "
"at batch >= 64). KDA uses its selected Triton or Helion implementation, "
"but its per-K gate ring is larger and ReplaySSM is typically slower "
"than packed KDA decode; benchmark before enabling it. Requires the "
"Triton linear-attn decode backend, or Helion for KDA, and "
"--mamba-radix-cache-strategy no_buffer (the default).",
NS("exec.mamba"),
] = False
@@ -6171,6 +6171,7 @@ class ServerArgs:
# Fixed in FlashInfer v0.6.7: flashinfer-ai/flashinfer#2810
if (
self.linear_attn_decode_backend is None
and self.linear_attn_backend != "helion"
and is_sm100_supported()
and self.mamba_ssm_dtype == "bfloat16"
# Stage 4: flashinfer's recurrent_kda compiles the state slot stride
@@ -6249,8 +6250,8 @@ class ServerArgs:
f"got CUDA {cuda_version or 'unknown'}"
)
# GDN ReplaySSM buffered decode guards. Runs on the Triton GDN decode
# backend. cuda-graph is supported (slice 1b: CUDA-graph-safe static
# ReplaySSM buffered decode guards. Runs on Triton, or Helion for KDA.
# cuda-graph is supported (slice 1b: CUDA-graph-safe static
# write-cursor buffers). The RADIX prefix cache is now supported (slice
# 2b: the decode kernel force-flushes the ring into temporal[slot] on
# the radix track boundary `seq_lens % mamba_track_interval == 0`, and
@@ -6264,10 +6265,10 @@ class ServerArgs:
# cursor of the donated/kept slot would not be reset there. Handling
# that donation path is a follow-up; for now require no_buffer.
if self.enable_linear_replayssm:
if decode != "triton":
if decode not in {"triton", "helion"}:
raise ValueError(
"--enable-linear-replayssm requires the Triton "
"linear-attn decode backend, got "
"--enable-linear-replayssm requires Triton, or Helion for "
"KDA, as the linear-attn decode backend; got "
f"--linear-attn-decode-backend={decode!r}."
)
from sglang.srt.arg_groups.overrides import (
@@ -8285,21 +8286,20 @@ class ServerArgs:
# The Mamba/KDA state is stored in envelope-strided views; only
# stride-audited kernels may read it (Stage 4 audit, per slot):
# - decode: triton; flashinfer (recurrent_kda compiles the state slot
# stride as a free int64 — natively strided); cutedsl (KDA fused
# sigmoid-gating update made stride-safe) on KDA-hybrid models only —
# cutedsl_gdn still compiles h0 against a contiguous dummy.
# - prefill: triton; flashkda (wrapper gathers/scatters a contiguous
# per-slot copy, external kernel never sees the pool); cutedsl
# (kernel_h compiles h0/ht with dynamic int64 strides), same
# KDA-only caveat.
# stride as a free int64); helion (specializes KDA state strides 0-3
# and rejects a non-unit innermost stride); cutedsl (KDA fused sigmoid-
# gating update is stride-safe) on KDA-hybrid models only.
# - prefill: triton; flashkda (the wrapper gathers/scatters a contiguous
# per-slot copy); helion; cutedsl (kernel_h compiles h0/ht with dynamic
# int64 strides), with the same KDA-only caveat.
# - mamba (mamba2/short-conv state): triton only.
# use_mla_backend() distinguishes the KDA-hybrid family (K3/KimiLinear
# are MLA-hybrid) from GDN models (GQA-hybrid) for the cutedsl caveat.
# are MLA-hybrid) from GDN models (GQA-hybrid) for the KDA-only caveat.
decode_allowed = {"triton", "flashinfer"}
prefill_allowed = {"triton", "flashkda"}
if self.use_mla_backend():
decode_allowed.add("cutedsl")
prefill_allowed.add("cutedsl")
decode_allowed.update({"cutedsl", "helion"})
prefill_allowed.update({"cutedsl", "helion"})
resolved_linear_decode = (
self.linear_attn_decode_backend or self.linear_attn_backend
)