minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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/kernels/ops/attention/helion/__init__.py
+5
View File
@@ -0,0 +1,5 @@
"""Helion attention kernels."""
# K3 exposes 12 local value heads at TP=8. Lower value-head counts share the
# same small-head decode regime.
KDA_SMALL_VALUE_HEAD_THRESHOLD = 12
python/sglang/kernels/ops/attention/helion/kda_decode.py
+384
View File
@@ -0,0 +1,384 @@
"""Helion implementation of SGLang's packed KDA decode contract."""
from __future__ import annotations
import helion
import helion.language as hl
import torch
from sglang.kernels.ops.attention.helion import KDA_SMALL_VALUE_HEAD_THRESHOLD
# SGLang initializes torch.distributed, but this kernel has no collectives.
_IGNORED_WARNINGS = [helion.exc.ProcessGroupNameNotFound]
_LOG2_E = 1.4426950408889634
# Tile V on the CUDA x axis so tensor-parallel head counts do not change the
# grid width.
_KDA_CONFIG = helion.Config(
block_sizes=[8],
loop_orders=[[2, 1, 0]],
num_warps=1,
num_stages=1,
indexing="pointer",
pid_type="xyz",
)
_KDA_BF16_CONFIG = helion.Config(
atomic_indexing=[],
block_sizes=[16],
indexing="pointer",
l2_groupings=[16],
# Policies are positional in the traced load order. Retune them if the
# decode body gains, loses, or reorders loads.
load_eviction_policies=[
"",
"last",
"first",
"last",
"first",
"first",
"last",
"",
"last",
],
loop_orders=[[1, 2, 0]],
num_stages=1,
num_warps=1,
pid_type="flat",
range_flattens=[None],
range_multi_buffers=[None],
range_num_stages=[],
range_unroll_factors=[0],
)
# The bounded sigmoid gate has lower ALU and register pressure than the
# unbounded softplus gate, allowing small-head BF16 decode to use a wider V
# tile. The same tile regresses the unbounded path.
_KDA_BF16_SMALL_HEAD_CONFIG = helion.Config(
block_sizes=[32],
loop_orders=[[2, 1, 0]],
num_warps=1,
num_stages=1,
indexing="pointer",
pid_type="xyz",
)
def _helion_fused_recurrent_kda_packed_decode_body(
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
scale: float,
lower_bound: float,
initial_state: torch.Tensor,
out: torch.Tensor,
ssm_state_indices: torch.Tensor,
use_qk_l2norm_in_kernel: hl.constexpr = False, # pyrefly: ignore[bad-function-definition]
use_fast_rsqrt: hl.constexpr = False, # pyrefly: ignore[bad-function-definition]
use_lower_bound: hl.constexpr = False, # pyrefly: ignore[bad-function-definition]
) -> torch.Tensor:
"""Fused packed KDA decode body; mutates ``initial_state`` and ``out``."""
B = mixed_qkv.size(0)
HV = hl.specialize(initial_state.size(-3))
V = hl.specialize(initial_state.size(-2))
K = hl.specialize(initial_state.size(-1))
H = hl.specialize((mixed_qkv.size(1) - HV * V) // (2 * K))
heads_per_q = HV // H
hl.specialize(
(
mixed_qkv.stride(0),
mixed_qkv.stride(1),
a.stride(0),
a.stride(1),
b.stride(0),
b.stride(1),
A_log.stride(0),
dt_bias.stride(0),
initial_state.stride(0),
initial_state.stride(1),
initial_state.stride(2),
initial_state.stride(3),
out.stride(0),
out.stride(1),
out.stride(2),
out.stride(3),
ssm_state_indices.stride(0),
)
)
block_v = hl.register_block_size(1, V)
for tile_b, tile_hv, tile_v in hl.tile([B, HV, V], block_size=[1, 1, block_v]):
k_offsets = hl.arange(K)
i_b = tile_b.id
i_hv = tile_hv.id
i_h = i_hv // heads_per_q
state_index = ssm_state_indices[i_b].long()
if state_index < 0:
out[i_b, 0, i_hv, tile_v] = 0.0
else:
q_offsets = i_h * K + k_offsets
k_input_offsets = H * K + i_h * K + k_offsets
v_offsets = 2 * H * K + i_hv * V + tile_v.index
raw_gate = a[i_b, i_hv * K + k_offsets].float()
raw_gate = raw_gate + dt_bias[i_hv * K + k_offsets].float()
A_log_value = A_log[i_hv].float()
A = torch.exp2(A_log_value * _LOG2_E)
if use_lower_bound:
log_decay = lower_bound * torch.sigmoid(A * raw_gate)
else:
gate_exp = torch.exp2(raw_gate * _LOG2_E)
softplus = torch.where(
raw_gate <= 20.0,
torch.log(1.0 + gate_exp),
raw_gate,
)
log_decay = -A * softplus
beta = torch.sigmoid(b[i_b, i_hv].float())
state = initial_state[state_index, i_hv, tile_v.index, k_offsets].float()
decay = torch.exp2(log_decay * _LOG2_E)
state = state * decay[None, :]
k = mixed_qkv[i_b, k_input_offsets].float()
if use_qk_l2norm_in_kernel:
k_norm = (k * k).sum() + 1e-6
if use_fast_rsqrt:
k = k * torch.rsqrt(k_norm)
else:
k = k / torch.sqrt(k_norm)
v = mixed_qkv[i_b, v_offsets].float()
value_residual = v - (state * k[None, :]).sum(-1)
value_residual = value_residual * beta
state = state + value_residual[:, None] * k[None, :]
q = mixed_qkv[i_b, q_offsets].float()
if use_qk_l2norm_in_kernel:
q_norm = (q * q).sum() + 1e-6
if use_fast_rsqrt:
q = q * torch.rsqrt(q_norm)
else:
q = q / torch.sqrt(q_norm)
q = q * scale
output = (state * q[None, :]).sum(-1)
out[i_b, 0, i_hv, tile_v] = output.to(out.dtype)
initial_state[state_index, i_hv, tile_v.index, k_offsets] = state
return out
_helion_fused_recurrent_kda_packed_decode = helion.kernel(
_helion_fused_recurrent_kda_packed_decode_body,
static_shapes=False,
config=_KDA_CONFIG,
ignore_warnings=_IGNORED_WARNINGS,
)
_helion_fused_recurrent_kda_packed_decode_bf16 = helion.kernel(
_helion_fused_recurrent_kda_packed_decode_body,
static_shapes=False,
config=_KDA_BF16_CONFIG,
ignore_warnings=_IGNORED_WARNINGS,
)
_helion_fused_recurrent_kda_packed_decode_bf16_small_head = helion.kernel(
_helion_fused_recurrent_kda_packed_decode_body,
static_shapes=False,
config=_KDA_BF16_SMALL_HEAD_CONFIG,
ignore_warnings=_IGNORED_WARNINGS,
)
def _select_decode_kernel(
*,
is_bf16_state: bool,
num_v_heads: int,
use_lower_bound: bool,
) -> helion.Kernel:
if (
is_bf16_state
and use_lower_bound
and num_v_heads <= KDA_SMALL_VALUE_HEAD_THRESHOLD
):
return _helion_fused_recurrent_kda_packed_decode_bf16_small_head
if is_bf16_state:
return _helion_fused_recurrent_kda_packed_decode_bf16
return _helion_fused_recurrent_kda_packed_decode
def validate_packed_decode_inputs(
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
initial_state: torch.Tensor,
out: torch.Tensor,
ssm_state_indices: torch.Tensor,
) -> tuple[int, int, int, int, int]:
"""Apply the shape and layout checks from SGLang's packed wrapper."""
if mixed_qkv.ndim != 2:
raise ValueError(
f"`mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.ndim})."
)
if mixed_qkv.stride(-1) != 1:
raise ValueError("`mixed_qkv` must be contiguous in the last dim.")
if a.ndim != 2 or b.ndim != 2:
raise ValueError(
f"`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim={b.ndim})."
)
if a.stride(-1) != 1 or b.stride(-1) != 1:
raise ValueError("`a`/`b` must be contiguous in the last dim.")
if A_log.ndim != 1 or dt_bias.ndim != 1:
raise ValueError("`A_log`/`dt_bias` must be 1D tensors.")
if A_log.stride(0) != 1 or dt_bias.stride(0) != 1:
raise ValueError("`A_log`/`dt_bias` must be contiguous.")
if ssm_state_indices.ndim != 1:
raise ValueError(
"`ssm_state_indices` must be 1D for packed decode "
f"(got ndim={ssm_state_indices.ndim})."
)
if not out.is_contiguous():
raise ValueError("`out` must be contiguous.")
device = mixed_qkv.device
if any(
tensor.device != device
for tensor in (
a,
b,
A_log,
dt_bias,
initial_state,
out,
ssm_state_indices,
)
):
raise ValueError("All inputs must be on the same device.")
B = mixed_qkv.shape[0]
if a.shape[0] != B or b.shape[0] != B:
raise ValueError(
"Mismatched batch sizes: "
f"mixed_qkv.shape[0]={B}, a.shape[0]={a.shape[0]}, "
f"b.shape[0]={b.shape[0]}."
)
if ssm_state_indices.shape[0] != B:
raise ValueError(
f"`ssm_state_indices` must have shape [B] "
f"(got {tuple(ssm_state_indices.shape)}; expected ({B},))."
)
if initial_state.ndim != 4:
raise ValueError(
f"`initial_state` must be a 4D tensor (got ndim={initial_state.ndim})."
)
if initial_state.stride(-1) != 1:
raise ValueError("`initial_state` must be contiguous in the last dim.")
HV, V, K = initial_state.shape[-3:]
if a.shape[1] != HV * K:
raise ValueError(
f"`a` must have shape [B, HV*K] with HV={HV}, K={K} "
f"(got a.shape={tuple(a.shape)})."
)
if b.shape[1] != HV:
raise ValueError(
f"`b` must have shape [B, HV] with HV={HV} (got b.shape={tuple(b.shape)})."
)
if A_log.numel() != HV:
raise ValueError(f"`A_log` must have {HV} elements (got {A_log.numel()}).")
if dt_bias.numel() != HV * K:
raise ValueError(
f"`dt_bias` must have {HV * K} elements (got {dt_bias.numel()})."
)
if out.shape != (B, 1, HV, V):
raise ValueError(
f"`out` must have shape {(B, 1, HV, V)} (got out.shape={tuple(out.shape)})."
)
qkv_dim = mixed_qkv.shape[1]
qk_dim = qkv_dim - HV * V
if qk_dim <= 0 or qk_dim % 2 != 0:
raise ValueError(
f"Invalid packed `mixed_qkv` last dim={qkv_dim} for HV={HV}, V={V}."
)
q_dim = qk_dim // 2
if q_dim % K != 0:
raise ValueError(
f"Invalid packed Q size {q_dim}: must be divisible by K={K}. "
"KDA packed decode requires num_q_heads == num_k_heads and "
"head_q_dim == head_k_dim."
)
H = q_dim // K
if H <= 0 or HV % H != 0:
raise ValueError(
f"Invalid head config inferred from mixed_qkv: H={H}, HV={HV}."
)
return B, H, HV, K, V
def helion_fused_recurrent_kda_packed_decode(
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
scale: float,
initial_state: torch.Tensor,
out: torch.Tensor,
ssm_state_indices: torch.Tensor,
use_qk_l2norm_in_kernel: bool = False,
lower_bound: float | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Helion implementation of SGLang's packed KDA decode contract.
Inputs, mutations, padding semantics, and outputs match
``fused_recurrent_kda_packed_decode``:
* ``mixed_qkv`` is ``[B, 2*H*K + HV*V]`` after the short convolution.
* ``a`` and ``b`` are raw forget-gate and beta logits.
* ``lower_bound`` selects the bounded sigmoid decay used by safe-gate KDA.
* ``initial_state`` is ``[num_slots, HV, V, K]`` and is updated in place.
* ``ssm_state_indices == -1`` writes a zero output and leaves state untouched.
* ``out`` is ``[B, 1, HV, V]`` and is written in place.
* The return is the same ``(out, initial_state)`` object pair supplied by the
caller.
"""
_, _, num_v_heads, _, _ = validate_packed_decode_inputs(
mixed_qkv,
a,
b,
A_log,
dt_bias,
initial_state,
out,
ssm_state_indices,
)
use_lower_bound = lower_bound is not None
is_bf16_state = initial_state.dtype is torch.bfloat16
kernel = _select_decode_kernel(
is_bf16_state=is_bf16_state,
num_v_heads=num_v_heads,
use_lower_bound=use_lower_bound,
)
lower_bound_value = 0.0 if lower_bound is None else lower_bound
result = kernel(
mixed_qkv,
a,
b,
A_log,
dt_bias,
scale,
lower_bound_value,
initial_state,
out,
ssm_state_indices,
use_qk_l2norm_in_kernel,
is_bf16_state,
use_lower_bound,
)
return result, initial_state
python/sglang/kernels/ops/attention/helion/kda_prefill.py
+1399
View File
File diff suppressed because it is too large Load Diff
python/sglang/kernels/ops/attention/helion/kda_replayssm.py
+780
View File
@@ -0,0 +1,780 @@
"""Helion ReplaySSM decode for Kimi Delta Attention.
The public :func:`helion_fused_recurrent_kda_replayssm_decode` mirrors
``fused_recurrent_linear_replayssm_decode(..., is_kda=True)`` from
``sglang.kernels.ops.attention.fla.fused_recurrent_linear_replayssm``. That
Triton kernel is gate-generic (GDN scalar gate / KDA per-K gate); this module
implements the KDA path only, and additionally supports the bounded gate,
which the Triton ReplaySSM kernel does not expose.
"""
from __future__ import annotations
import helion
import helion.language as hl
import torch
from sglang.kernels.ops.attention.helion import KDA_SMALL_VALUE_HEAD_THRESHOLD
from sglang.kernels.ops.attention.helion.kda_decode import (
validate_packed_decode_inputs,
)
# SGLang initializes torch.distributed, but this kernel has no collectives.
_IGNORED_WARNINGS = [helion.exc.ProcessGroupNameNotFound]
_LOG2_E = 1.4426950408889634
# Indexing and eviction-policy lists are positional in the traced load order.
# Retune them if the ReplaySSM body gains, loses, or reorders loads.
_KDA_REPLAYSSM_FP32_CONFIG = helion.Config(
atomic_indexing=[],
block_sizes=[32, 128],
indexing=[
"tensor_descriptor",
"tensor_descriptor",
"tensor_descriptor",
"tensor_descriptor",
"tensor_descriptor",
"pointer",
"pointer",
"tensor_descriptor",
"pointer",
"tensor_descriptor",
"tensor_descriptor",
"tensor_descriptor",
"pointer",
"tensor_descriptor",
"pointer",
"tensor_descriptor",
"pointer",
"pointer",
"tensor_descriptor",
"tensor_descriptor",
"pointer",
"pointer",
"tensor_descriptor",
"tensor_descriptor",
"tensor_descriptor",
"pointer",
"pointer",
"tensor_descriptor",
"tensor_descriptor",
"tensor_descriptor",
"tensor_descriptor",
],
l2_groupings=[4],
load_eviction_policies=[
"last",
"last",
"first",
"first",
"",
"first",
"",
"first",
"first",
"",
"last",
"first",
"last",
"first",
"last",
"first",
"",
"first",
"last",
"",
"last",
"first",
"last",
"first",
"first",
],
loop_orders=[[1, 2, 0]],
num_warps=1,
num_stages=1,
pid_type="flat",
range_flattens=[None, None],
range_multi_buffers=[None, False],
range_num_stages=[0, 3],
range_unroll_factors=[0, 0],
)
_KDA_REPLAYSSM_BF16_CONFIG = helion.Config(
atomic_indexing=[],
block_sizes=[64, 64],
indexing=[
"pointer",
"pointer",
"tensor_descriptor",
"pointer",
"tensor_descriptor",
"pointer",
"tensor_descriptor",
"pointer",
"pointer",
"tensor_descriptor",
"tensor_descriptor",
"pointer",
"tensor_descriptor",
"tensor_descriptor",
"pointer",
"pointer",
"tensor_descriptor",
"tensor_descriptor",
"pointer",
"tensor_descriptor",
"tensor_descriptor",
"pointer",
"tensor_descriptor",
"pointer",
"pointer",
"pointer",
"pointer",
"pointer",
"pointer",
"pointer",
"pointer",
],
l2_groupings=[2],
load_eviction_policies=[
"first",
"last",
"last",
"first",
"",
"",
"last",
"last",
"first",
"first",
"first",
"",
"last",
"first",
"last",
"first",
"",
"first",
"",
"first",
"first",
"",
"",
"first",
"last",
],
loop_orders=[[1, 2, 0]],
num_stages=1,
num_warps=1,
pid_type="flat",
range_flattens=[None, False],
range_multi_buffers=[None, False],
range_num_stages=[0, 0],
range_unroll_factors=[0, 4],
)
# Small-head BF16 ReplaySSM uses the FP32 tile schedule with direct PID order.
_KDA_REPLAYSSM_BF16_SMALL_HEAD_CONFIG = helion.Config.from_dict(
{**_KDA_REPLAYSSM_FP32_CONFIG, "l2_groupings": [1]}
)
def _log_decay(
*,
raw_gate: torch.Tensor,
decay_rate: torch.Tensor,
lower_bound: float,
use_lower_bound: hl.constexpr,
) -> torch.Tensor:
if use_lower_bound:
return lower_bound * torch.sigmoid(decay_rate * raw_gate)
gate_exp = torch.exp2(raw_gate * _LOG2_E)
softplus = torch.where(
raw_gate <= 20.0,
torch.log(1.0 + gate_exp),
raw_gate,
)
return -decay_rate * softplus
def _load_log_decay(
*,
a: torch.Tensor,
dt_bias: torch.Tensor,
batch_index: torch.Tensor,
value_head: torch.Tensor,
key_indices: torch.Tensor,
key_valid: torch.Tensor,
key_dim: int,
decay_rate: torch.Tensor,
lower_bound: float,
use_lower_bound: hl.constexpr,
) -> torch.Tensor:
"""Load the current raw gate and apply the selected KDA gate contract."""
raw_gate = hl.load(
a,
[batch_index, value_head * key_dim + key_indices],
extra_mask=key_valid,
).float()
raw_gate = (
raw_gate
+ hl.load(
dt_bias,
[value_head * key_dim + key_indices],
extra_mask=key_valid,
).float()
)
return _log_decay(
raw_gate=raw_gate,
decay_rate=decay_rate,
lower_bound=lower_bound,
use_lower_bound=use_lower_bound,
)
def _helion_fused_recurrent_kda_replayssm_decode_body(
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
scale: float,
initial_state: torch.Tensor,
d_cache: torch.Tensor,
k_cache: torch.Tensor,
g_cache: torch.Tensor,
out: torch.Tensor,
ssm_state_indices: torch.Tensor,
write_pos: torch.Tensor,
force_flush: torch.Tensor | None,
lower_bound: float,
cache_block: hl.constexpr, # pyrefly: ignore[bad-function-definition]
use_qk_l2norm_in_kernel: hl.constexpr, # pyrefly: ignore[bad-function-definition]
use_lower_bound: hl.constexpr, # pyrefly: ignore[bad-function-definition]
) -> torch.Tensor:
"""Reconstruct the buffered state, emit one token, and flush if needed."""
B = mixed_qkv.size(0)
HV = hl.specialize(initial_state.size(1))
V = hl.specialize(initial_state.size(2))
K = hl.specialize(initial_state.size(3))
H = hl.specialize(k_cache.size(1))
cache_length = hl.specialize(d_cache.size(2))
heads_per_q = HV // H
# Keep storage setup inline so the generated host wrapper does not call
# Python helpers on every eager invocation.
state_strides = hl.specialize(
(
initial_state.stride(0),
initial_state.stride(1),
initial_state.stride(2),
initial_state.stride(3),
)
)
d_strides = hl.specialize(
(
d_cache.stride(0),
d_cache.stride(1),
d_cache.stride(2),
d_cache.stride(3),
)
)
k_strides = hl.specialize(
(
k_cache.stride(0),
k_cache.stride(1),
k_cache.stride(2),
k_cache.stride(3),
)
)
g_strides = hl.specialize(
(
g_cache.stride(0),
g_cache.stride(1),
g_cache.stride(2),
g_cache.stride(3),
)
)
state_storage_size = (
(initial_state.size(0) - 1) * state_strides[0]
+ (HV - 1) * state_strides[1]
+ (V - 1) * state_strides[2]
+ (K - 1) * state_strides[3]
+ 1
)
d_storage_size = (
(d_cache.size(0) - 1) * d_strides[0]
+ (HV - 1) * d_strides[1]
+ (cache_length - 1) * d_strides[2]
+ (V - 1) * d_strides[3]
+ 1
)
k_storage_size = (
(k_cache.size(0) - 1) * k_strides[0]
+ (H - 1) * k_strides[1]
+ (cache_length - 1) * k_strides[2]
+ (K - 1) * k_strides[3]
+ 1
)
g_storage_size = (
(g_cache.size(0) - 1) * g_strides[0]
+ (HV - 1) * g_strides[1]
+ (cache_length - 1) * g_strides[2]
+ (K - 1) * g_strides[3]
+ 1
)
# Omitting storage_offset preserves each source view's existing offset,
# which is required for envelope and page-major cache views.
state_storage = initial_state.as_strided([state_storage_size], [1])
d_storage = d_cache.as_strided([d_storage_size], [1])
k_storage = k_cache.as_strided([k_storage_size], [1])
g_storage = g_cache.as_strided([g_storage_size], [1])
hl.specialize(
(
mixed_qkv.stride(0),
mixed_qkv.stride(1),
a.stride(0),
a.stride(1),
b.stride(0),
b.stride(1),
A_log.stride(0),
dt_bias.stride(0),
out.stride(0),
out.stride(1),
out.stride(2),
out.stride(3),
ssm_state_indices.stride(0),
write_pos.stride(0),
)
)
block_v = hl.register_block_size(16, V)
block_k = hl.register_block_size(16, K)
for tile_b, tile_hv, tile_v in hl.tile([B, HV, V], block_size=[1, 1, block_v]):
i_b = tile_b.id
i_hv = tile_hv.id
i_h = i_hv // heads_per_q
state_index = ssm_state_indices[i_b].long()
v_valid = tile_v.index < V
if state_index < 0:
hl.store(
out,
[i_b, 0, i_hv, tile_v.index],
hl.zeros([tile_v], dtype=out.dtype),
extra_mask=v_valid,
)
else:
cursor = write_pos[i_b].long()
is_flush = cursor == cache_length - 1
if force_flush is not None:
is_flush = is_flush | (force_flush[i_b] != 0)
should_append = is_flush == 0
cache_positions = hl.arange(cache_block)
cache_valid = cache_positions < cursor
d_offsets = (
state_index * d_strides[0]
+ i_hv * d_strides[1]
+ cache_positions[:, None] * d_strides[2]
+ tile_v.index[None, :] * d_strides[3]
)
d_values = hl.load(
d_storage,
[d_offsets],
extra_mask=cache_valid[:, None] & v_valid[None, :],
).float()
# Advanced loads preserve tensor-axis order: d_cache contributes
# [L, V], while the reconstruction dot consumes [V, L].
d_dot = d_values.T.to(out.dtype)
full_k = hl.arange(K)
q_offsets = i_h * K + full_k
full_k_offsets = H * K + i_h * K + full_k
q_full = mixed_qkv[i_b, q_offsets].float()
k_full = mixed_qkv[i_b, full_k_offsets].float()
if use_qk_l2norm_in_kernel:
q_rnorm = 1.0 / torch.sqrt((q_full * q_full).sum() + 1e-6)
k_rnorm = 1.0 / torch.sqrt((k_full * k_full).sum() + 1e-6)
else:
q_rnorm = 1.0
k_rnorm = 1.0
value_offsets = 2 * H * K + i_hv * V + tile_v.index
value = hl.load(
mixed_qkv,
[i_b, value_offsets],
extra_mask=v_valid,
).float()
# ReplaySSM Triton rounds beta through the input dtype before FP32
# accumulation; packed KDA intentionally keeps beta in FP32.
beta = torch.sigmoid(b[i_b, i_hv].float()).to(b.dtype).float()
A = torch.exp2(A_log[i_hv].float() * _LOG2_E)
state_q = hl.zeros([tile_v], dtype=torch.float32)
state_k = hl.zeros([tile_v], dtype=torch.float32)
current_kq = hl.zeros([], dtype=torch.float32)
# Reconstruct each K tile directly from the checkpoint and ring.
# The full state is never materialized outside registers.
for tile_k in hl.tile(K, block_size=block_k):
k_valid = tile_k.index < K
q_value = hl.load(
mixed_qkv,
[i_b, i_h * K + tile_k.index],
extra_mask=k_valid,
).float()
k_value = hl.load(
mixed_qkv,
[i_b, H * K + i_h * K + tile_k.index],
extra_mask=k_valid,
).float()
q_value = q_value * q_rnorm
k_value = k_value * k_rnorm
q_scaled = q_value * scale
current_kq = current_kq + (k_value * q_scaled).sum()
cache_mask = cache_valid[:, None] & k_valid[None, :]
g_offsets = (
state_index * g_strides[0]
+ i_hv * g_strides[1]
+ cache_positions[:, None] * g_strides[2]
+ tile_k.index[None, :] * g_strides[3]
)
cached_k_offsets = (
state_index * k_strides[0]
+ i_h * k_strides[1]
+ cache_positions[:, None] * k_strides[2]
+ tile_k.index[None, :] * k_strides[3]
)
state_offsets = (
state_index * state_strides[0]
+ i_hv * state_strides[1]
+ tile_v.index[:, None] * state_strides[2]
+ tile_k.index[None, :] * state_strides[3]
)
state_mask = v_valid[:, None] & k_valid[None, :]
cached_g = hl.load(
g_storage,
[g_offsets],
extra_mask=cache_mask,
).float()
gate_prefix = torch.cumsum(cached_g, dim=0)
gate_total = cached_g.sum(0)
replay_decay = torch.where(
cache_mask,
torch.exp2((gate_total[None, :] - gate_prefix) * _LOG2_E),
0.0,
)
total_decay = torch.exp2(gate_total * _LOG2_E)
cached_k = hl.load(
k_storage,
[cached_k_offsets],
extra_mask=cache_mask,
).float()
cached_k = (cached_k * replay_decay).to(out.dtype)
state = hl.load(
state_storage,
[state_offsets],
extra_mask=state_mask,
).float()
state = state * total_decay[None, :]
state = state + hl.dot(
d_dot,
cached_k,
out_dtype=torch.float32,
)
current_gate = _load_log_decay(
a=a,
dt_bias=dt_bias,
batch_index=i_b,
value_head=i_hv,
key_indices=tile_k.index,
key_valid=k_valid,
key_dim=K,
decay_rate=A,
lower_bound=lower_bound,
use_lower_bound=use_lower_bound,
)
current_decay = torch.exp2(current_gate * _LOG2_E)
q_effective = q_scaled * current_decay
k_effective = k_value * current_decay
state_q = state_q + (state * q_effective[None, :]).sum(-1)
state_k = state_k + (state * k_effective[None, :]).sum(-1)
if should_append:
if tile_v.id == 0:
if i_hv == i_h * heads_per_q:
current_k_offsets = (
state_index * k_strides[0]
+ i_h * k_strides[1]
+ cursor * k_strides[2]
+ tile_k.index * k_strides[3]
)
hl.store(
k_storage,
[current_k_offsets],
k_value.to(k_cache.dtype),
extra_mask=(cursor < cache_length) & k_valid,
)
current_g_offsets = (
state_index * g_strides[0]
+ i_hv * g_strides[1]
+ cursor * g_strides[2]
+ tile_k.index * g_strides[3]
)
hl.store(
g_storage,
[current_g_offsets],
current_gate,
extra_mask=(cursor < cache_length) & k_valid,
)
delta = beta * (value - state_k)
output = state_q + delta * current_kq
hl.store(
out,
[i_b, 0, i_hv, tile_v.index],
output.to(out.dtype),
extra_mask=v_valid,
)
if is_flush:
# Reconstruct again now that the current delta is available,
# then fold the current rank-one update into the checkpoint.
# Helion cannot lower the cumsum through a device helper without
# creating a separate unsupported scan subgraph.
for tile_k in hl.tile(K, block_size=block_k):
k_valid = tile_k.index < K
k_value = hl.load(
mixed_qkv,
[i_b, H * K + i_h * K + tile_k.index],
extra_mask=k_valid,
).float()
k_value = k_value * k_rnorm
cache_mask = cache_valid[:, None] & k_valid[None, :]
g_offsets = (
state_index * g_strides[0]
+ i_hv * g_strides[1]
+ cache_positions[:, None] * g_strides[2]
+ tile_k.index[None, :] * g_strides[3]
)
cached_k_offsets = (
state_index * k_strides[0]
+ i_h * k_strides[1]
+ cache_positions[:, None] * k_strides[2]
+ tile_k.index[None, :] * k_strides[3]
)
state_offsets = (
state_index * state_strides[0]
+ i_hv * state_strides[1]
+ tile_v.index[:, None] * state_strides[2]
+ tile_k.index[None, :] * state_strides[3]
)
state_mask = v_valid[:, None] & k_valid[None, :]
cached_g = hl.load(
g_storage,
[g_offsets],
extra_mask=cache_mask,
).float()
gate_prefix = torch.cumsum(cached_g, dim=0)
gate_total = cached_g.sum(0)
replay_decay = torch.where(
cache_mask,
torch.exp2((gate_total[None, :] - gate_prefix) * _LOG2_E),
0.0,
)
total_decay = torch.exp2(gate_total * _LOG2_E)
cached_k = hl.load(
k_storage,
[cached_k_offsets],
extra_mask=cache_mask,
).float()
cached_k = (cached_k * replay_decay).to(out.dtype)
state = hl.load(
state_storage,
[state_offsets],
extra_mask=state_mask,
).float()
state = state * total_decay[None, :]
state = state + hl.dot(
d_dot,
cached_k,
out_dtype=torch.float32,
)
current_gate = _load_log_decay(
a=a,
dt_bias=dt_bias,
batch_index=i_b,
value_head=i_hv,
key_indices=tile_k.index,
key_valid=k_valid,
key_dim=K,
decay_rate=A,
lower_bound=lower_bound,
use_lower_bound=use_lower_bound,
)
current_decay = torch.exp2(current_gate * _LOG2_E)
state = state * current_decay[None, :]
state = state + delta[:, None] * k_value[None, :]
hl.store(
state_storage,
[state_offsets],
state.to(initial_state.dtype),
extra_mask=v_valid[:, None] & k_valid[None, :],
)
else:
current_d_offsets = (
state_index * d_strides[0]
+ i_hv * d_strides[1]
+ cursor * d_strides[2]
+ tile_v.index * d_strides[3]
)
hl.store(
d_storage,
[current_d_offsets],
delta.to(d_cache.dtype),
extra_mask=(cursor < cache_length) & v_valid,
)
return out
_helion_fused_recurrent_kda_replayssm_decode_fp32 = helion.kernel(
_helion_fused_recurrent_kda_replayssm_decode_body,
static_shapes=False,
config=_KDA_REPLAYSSM_FP32_CONFIG,
ignore_warnings=_IGNORED_WARNINGS,
)
_helion_fused_recurrent_kda_replayssm_decode_bf16 = helion.kernel(
_helion_fused_recurrent_kda_replayssm_decode_body,
static_shapes=False,
config=_KDA_REPLAYSSM_BF16_CONFIG,
ignore_warnings=_IGNORED_WARNINGS,
)
_helion_fused_recurrent_kda_replayssm_decode_bf16_small_head = helion.kernel(
_helion_fused_recurrent_kda_replayssm_decode_body,
static_shapes=False,
config=_KDA_REPLAYSSM_BF16_SMALL_HEAD_CONFIG,
ignore_warnings=_IGNORED_WARNINGS,
)
def _select_replayssm_decode_kernel(
*,
is_bf16_state: bool,
num_v_heads: int,
) -> helion.Kernel:
if is_bf16_state and num_v_heads <= KDA_SMALL_VALUE_HEAD_THRESHOLD:
return _helion_fused_recurrent_kda_replayssm_decode_bf16_small_head
if is_bf16_state:
return _helion_fused_recurrent_kda_replayssm_decode_bf16
return _helion_fused_recurrent_kda_replayssm_decode_fp32
def helion_fused_recurrent_kda_replayssm_decode(
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
scale: float,
initial_state: torch.Tensor,
d_cache: torch.Tensor,
k_cache: torch.Tensor,
g_cache: torch.Tensor,
out: torch.Tensor,
ssm_state_indices: torch.Tensor,
write_pos: torch.Tensor,
force_flush: torch.Tensor | None = None,
use_qk_l2norm_in_kernel: bool = False,
lower_bound: float | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Run one buffered KDA decode step using caller-owned ReplaySSM state.
Allocates nothing persistent: the caller owns ``d_cache`` / ``k_cache`` /
``g_cache`` and is responsible for advancing ``write_pos`` modulo the ring
length after a non-flush step and resetting it to zero after a natural or
forced flush. ``initial_state`` is both the checkpoint read (h0) and the
flush-only checkpoint write (ht), in place.
"""
batch = mixed_qkv.size(0)
if a.ndim not in (2, 3) or not a.is_contiguous():
raise ValueError("KDA `a` must be a contiguous 2D or 3D tensor.")
if dt_bias.ndim not in (1, 2) or not dt_bias.is_contiguous():
raise ValueError("KDA `dt_bias` must be a contiguous 1D or 2D tensor.")
flat_a = a.view(batch, -1)
flat_dt_bias = dt_bias.view(-1)
_, num_q_heads, num_v_heads, key_dim, value_dim = validate_packed_decode_inputs(
mixed_qkv,
flat_a,
b,
A_log,
flat_dt_bias,
initial_state,
out,
ssm_state_indices,
)
if write_pos.ndim != 1 or write_pos.dtype is not torch.int32:
raise ValueError("`write_pos` must be a 1D int32 tensor.")
if write_pos.shape != (batch,):
raise ValueError(f"`write_pos` must have shape {(batch,)}.")
if force_flush is not None and (
force_flush.ndim != 1
or force_flush.dtype is not torch.int32
or force_flush.shape != (batch,)
):
raise ValueError("`force_flush` must be a length-B int32 tensor or None.")
cache_length = d_cache.size(2)
if cache_length < 1:
raise ValueError("ReplaySSM cache length must be at least 1.")
if d_cache.shape[1:] != (num_v_heads, cache_length, value_dim):
raise ValueError("`d_cache` must have shape [slots, HV, L, V].")
if k_cache.shape[1:] != (num_q_heads, cache_length, key_dim):
raise ValueError("`k_cache` must have shape [slots, H, L, K].")
if g_cache.shape[1:] != (num_v_heads, cache_length, key_dim):
raise ValueError("`g_cache` must have shape [slots, HV, L, K].")
if g_cache.dtype is not torch.float32:
raise ValueError("`g_cache` must have dtype torch.float32.")
device = mixed_qkv.device
if any(
tensor.device != device for tensor in (d_cache, k_cache, g_cache, write_pos)
):
raise ValueError("ReplaySSM inputs must be on the same device.")
if force_flush is not None and force_flush.device != device:
raise ValueError("`force_flush` must be on the same device as the inputs.")
cache_block = helion.next_power_of_2(max(16, cache_length))
use_lower_bound = lower_bound is not None
kernel = _select_replayssm_decode_kernel(
is_bf16_state=initial_state.dtype is torch.bfloat16,
num_v_heads=num_v_heads,
)
result = kernel(
mixed_qkv,
flat_a,
b,
A_log,
flat_dt_bias,
scale,
initial_state,
d_cache,
k_cache,
g_cache,
out,
ssm_state_indices,
write_pos,
force_flush,
0.0 if lower_bound is None else lower_bound,
cache_block,
use_qk_l2norm_in_kernel,
use_lower_bound,
)
return result, initial_state
python/sglang/srt/layers/attention/linear/gdn_backend.py
+8
View File
@@ -137,6 +137,10 @@ class GDNKernelDispatcher:
flashinfer_kernel = FlashInferGDNKernel()
self.decode_kernel = flashinfer_kernel
elif decode_backend.is_helion():
raise ValueError(
"The Helion linear-attention backend supports KDA only, not GDN."
)
else:
raise ValueError(f"Unsupported GDN decode backend: {decode_backend}")
@@ -176,6 +180,10 @@ class GDNKernelDispatcher:
flashinfer_kernel = FlashInferGDNKernel()
self.extend_kernel = flashinfer_kernel
elif prefill_backend.is_helion():
raise ValueError(
"The Helion linear-attention backend supports KDA only, not GDN."
)
else:
raise ValueError(f"Unsupported GDN prefill backend: {prefill_backend}")
python/sglang/srt/layers/attention/linear/kda_backend.py
+23 -3
View File
@@ -47,9 +47,25 @@ class KDAKernelDispatcher:
):
self.verify_backend = verify_backend
triton_kernel = TritonKDAKernel()
helion_kernel = None
if decode_backend.is_helion() or prefill_backend.is_helion():
if not is_cuda():
raise ValueError("KDA Helion backend requires CUDA")
from sglang.srt.layers.attention.linear.kernels.kda_helion import (
HelionKDAKernel,
)
helion_kernel = HelionKDAKernel(
triton_fallback=triton_kernel,
enable_decode=decode_backend.is_helion(),
enable_prefill=prefill_backend.is_helion(),
)
if decode_backend.is_triton():
self.decode_kernel = triton_kernel
elif decode_backend.is_helion():
assert helion_kernel is not None
self.decode_kernel = helion_kernel
elif decode_backend.is_cutedsl():
if not is_cuda():
raise ValueError("KDA CuTe DSL backend requires CUDA")
@@ -71,7 +87,7 @@ class KDAKernelDispatcher:
else:
raise ValueError(
f"Unsupported KDA decode backend: {decode_backend}. "
"KDA supports 'triton', 'cutedsl', or 'flashinfer'."
"KDA supports 'triton', 'helion', 'cutedsl', or 'flashinfer'."
)
# target_verify kernel, selected via --linear-attn-verify-backend (defaults
@@ -110,6 +126,9 @@ class KDAKernelDispatcher:
if prefill_backend.is_triton():
self.extend_kernel = triton_kernel
elif prefill_backend.is_helion():
assert helion_kernel is not None
self.extend_kernel = helion_kernel
elif prefill_backend.is_flashkda():
from sglang.srt.layers.attention.linear.kernels.kda_flashkda import (
FlashKDAKernel,
@@ -166,8 +185,9 @@ class KDAKernelDispatcher:
else:
raise ValueError(
f"Unsupported KDA prefill backend: {prefill_backend}. "
"KDA supports 'triton', 'flashkda', 'cutedsl', 'nvidia_kda', or "
"'ptx_kda' (cutedsl/nvidia_kda prefill need SM100, ptx_kda SM103)."
"KDA supports 'triton', 'helion', 'flashkda', 'cutedsl', "
"'nvidia_kda', or 'ptx_kda' (cutedsl/nvidia_kda prefill need "
"SM100, ptx_kda SM103)."
)
self.supports_packed_decode = getattr(
python/sglang/srt/layers/attention/linear/kernels/kda_helion.py
+191
View File
@@ -0,0 +1,191 @@
"""Helion backend for Kimi Delta Attention."""
from __future__ import annotations
import torch
from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
LinearAttnKernelBase,
)
class HelionKDAKernel(LinearAttnKernelBase):
"""KDA packed decode and prefill implemented with Helion kernels.
The generic decode interface delegates to Triton, and the dispatcher routes
speculative target verification directly to Triton. The one-token decode
and ReplaySSM paths use :meth:`packed_decode`, while prefill uses
:meth:`extend`.
"""
supports_packed_decode = True
def __init__(
self,
triton_fallback: TritonKDAKernel | None = None,
*,
enable_decode: bool = True,
enable_prefill: bool = True,
) -> None:
self.supports_packed_decode = enable_decode
self._packed_decode = None
self._replayssm_decode = None
self._chunk_kda = None
if enable_decode or enable_prefill:
try:
import helion # noqa: F401
except ModuleNotFoundError as error:
if error.name != "helion":
raise
raise ImportError(
"The Helion package is required when a KDA backend is set to "
"Helion. Install it with: pip install helion==1.4.0"
) from None
if enable_decode:
from sglang.kernels.ops.attention.helion.kda_decode import (
helion_fused_recurrent_kda_packed_decode,
)
from sglang.kernels.ops.attention.helion.kda_replayssm import (
helion_fused_recurrent_kda_replayssm_decode,
)
self._packed_decode = helion_fused_recurrent_kda_packed_decode
self._replayssm_decode = helion_fused_recurrent_kda_replayssm_decode
if enable_prefill:
from sglang.kernels.ops.attention.helion.kda_prefill import chunk_kda
self._chunk_kda = chunk_kda
self._triton = triton_fallback or TritonKDAKernel()
def packed_decode(
self,
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
*,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
scale: float,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
num_v_heads: int,
head_v_dim: int,
lower_bound: float | None = None,
**kwargs,
) -> torch.Tensor:
assert self._packed_decode is not None
batch_size = mixed_qkv.shape[0]
out = mixed_qkv.new_empty(batch_size, 1, num_v_heads, head_v_dim)
replayssm_d = kwargs.get("replayssm_d")
replayssm_k = kwargs.get("replayssm_k")
replayssm_g = kwargs.get("replayssm_g")
replayssm_write_pos = kwargs.get("replayssm_write_pos")
if (
replayssm_d is not None
and replayssm_k is not None
and replayssm_g is not None
and replayssm_write_pos is not None
):
assert self._replayssm_decode is not None
self._replayssm_decode(
mixed_qkv=mixed_qkv,
a=a.reshape(batch_size, num_v_heads, -1).contiguous(),
b=b.reshape(batch_size, num_v_heads).contiguous(),
A_log=A_log.reshape(-1),
dt_bias=dt_bias.reshape(num_v_heads, -1).contiguous(),
scale=scale,
initial_state=ssm_states,
d_cache=replayssm_d,
k_cache=replayssm_k,
g_cache=replayssm_g,
out=out,
ssm_state_indices=cache_indices,
write_pos=replayssm_write_pos,
force_flush=kwargs.get("replayssm_force_flush"),
use_qk_l2norm_in_kernel=True,
lower_bound=lower_bound,
)
return out.transpose(0, 1)
if a.ndim != 2:
a = a.reshape(batch_size, -1)
if b.ndim != 2:
b = b.reshape(batch_size, -1)
self._packed_decode(
mixed_qkv=mixed_qkv,
a=a,
b=b,
A_log=A_log.reshape(-1),
dt_bias=dt_bias.reshape(-1),
scale=scale,
initial_state=ssm_states,
out=out,
ssm_state_indices=cache_indices,
use_qk_l2norm_in_kernel=True,
lower_bound=lower_bound,
)
return out.transpose(0, 1)
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:
return self._triton.decode(
q,
k,
v,
a,
b,
A_log=A_log,
dt_bias=dt_bias,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
**kwargs,
)
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: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
lower_bound: float | None = None,
return_intermediate_states: bool = False,
**kwargs,
) -> torch.Tensor:
assert self._chunk_kda is not None
return self._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,
output_intermediate_states=return_intermediate_states,
)
python/sglang/srt/layers/attention/linear/utils.py
+4
View File
@@ -20,6 +20,7 @@ class LinearAttnKernelBackend(Enum):
FLASHKDA = "flashkda"
NVIDIA_KDA = "nvidia_kda"
PTX_KDA = "ptx_kda"
HELION = "helion"
CUSTOM = "custom"
@classmethod
@@ -47,6 +48,9 @@ class LinearAttnKernelBackend(Enum):
def is_ptx_kda(self):
return self == LinearAttnKernelBackend.PTX_KDA
def is_helion(self):
return self == LinearAttnKernelBackend.HELION
def is_custom(self):
return self == LinearAttnKernelBackend.CUSTOM
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
)