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
benchmark/bench_linear_attention/bench_kda_decode.py
+70 -5
View File
@@ -4,6 +4,7 @@ Benchmark & Correctness: KDA Packed Decode vs Baseline Decode.
Compares:
- Baseline: split(mixed_qkv) -> view -> fused_sigmoid_gating_delta_rule_update(is_kda=True)
- Packed: fused_recurrent_kda_packed_decode (single fused kernel)
- Helion: helion_fused_recurrent_kda_packed_decode
Differences from the GDN packed decode benchmark:
- KDA gate ``a`` is per-K with shape ``[B, HV * K]`` (instead of ``[B, HV]``).
@@ -28,6 +29,9 @@ from sglang.kernels.ops.attention.fla.fused_recurrent import (
from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
fused_sigmoid_gating_delta_rule_update,
)
from sglang.kernels.ops.attention.helion.kda_decode import (
helion_fused_recurrent_kda_packed_decode,
)
def make_inputs(
@@ -131,6 +135,27 @@ def run_packed(inp):
return out.transpose(0, 1), ssm_states
def run_helion(inp):
"""Helion path: same packed decode contract and output layout."""
B, HV, V = inp["B"], inp["HV"], inp["V"]
ssm_states = inp["ssm_states"].clone()
out = inp["mixed_qkv"].new_empty(B, 1, HV, V)
helion_fused_recurrent_kda_packed_decode(
mixed_qkv=inp["mixed_qkv"],
a=inp["a"],
b=inp["b"],
A_log=inp["A_log"],
dt_bias=inp["dt_bias"],
scale=inp["K"] ** -0.5,
initial_state=ssm_states,
out=out,
ssm_state_indices=inp["cache_indices"],
use_qk_l2norm_in_kernel=True,
)
return out.transpose(0, 1), ssm_states
def check_correctness(B, H, HV, K, V, pool_size, device, dtype, seed=42):
"""Run correctness check for a single config. Returns True if PASS."""
tag = f"B={B:>4} H={H:>2} HV={HV:>2} K={K:>3} V={V:>3} pool={pool_size:>4}"
@@ -139,6 +164,7 @@ def check_correctness(B, H, HV, K, V, pool_size, device, dtype, seed=42):
o_baseline, state_baseline = run_baseline(inp)
o_packed, state_packed = run_packed(inp)
o_helion, state_helion = run_helion(inp)
atol = 2e-2 if dtype != torch.float32 else 1e-4
rtol = 1e-2 if dtype != torch.float32 else 1e-4
@@ -157,6 +183,19 @@ def check_correctness(B, H, HV, K, V, pool_size, device, dtype, seed=42):
atol, rtol * state_baseline[indices].float().abs().max().item()
)
helion_out_diff = (o_helion.float() - o_packed.float()).abs().max().item()
helion_state_diff = (
(state_helion[indices].float() - state_packed[indices].float())
.abs()
.max()
.item()
)
helion_ok = helion_out_diff <= max(
atol, rtol * o_packed.float().abs().max().item()
) and helion_state_diff <= max(
atol, rtol * state_packed[indices].float().abs().max().item()
)
passed = output_ok and state_ok
if passed:
print(
@@ -166,7 +205,12 @@ def check_correctness(B, H, HV, K, V, pool_size, device, dtype, seed=42):
print(
f" [FAIL] {tag} out max_diff={out_diff:.6f}, state max_diff={st_diff:.6f}"
)
return passed
print(
f" [{'PASS' if helion_ok else 'FAIL'}] Helion vs packed {tag} "
f"(out max_diff={helion_out_diff:.2e}, "
f"state max_diff={helion_state_diff:.2e})"
)
return passed and helion_ok
def bench_shape(B, H, HV, K, V, pool_size, device, dtype):
@@ -213,6 +257,20 @@ def bench_shape(B, H, HV, K, V, pool_size, device, dtype):
use_qk_l2norm_in_kernel=True,
)
def fn_helion():
helion_fused_recurrent_kda_packed_decode(
mixed_qkv=inp["mixed_qkv"],
a=inp["a"],
b=inp["b"],
A_log=inp["A_log"],
dt_bias=inp["dt_bias"],
scale=K**-0.5,
initial_state=inp["ssm_states"],
out=out_buf,
ssm_state_indices=inp["cache_indices"],
use_qk_l2norm_in_kernel=True,
)
# Intentionally wall-clock CUDA-event timing, not the shared do_bench /
# do_bench_cudagraph util: ~2/3 of the packed win is eager CPU dispatch
# (split + 3x unflatten + extra launch), which graph capture / L2-flush
@@ -222,6 +280,7 @@ def bench_shape(B, H, HV, K, V, pool_size, device, dtype):
for _ in range(warmup):
fn_baseline()
fn_packed()
fn_helion()
torch.cuda.synchronize()
def _time(fn):
@@ -236,7 +295,7 @@ def bench_shape(B, H, HV, K, V, pool_size, device, dtype):
ms_baseline = _time(fn_baseline)
ms_packed = _time(fn_packed)
ms_helion = _time(fn_helion)
speedup = ms_baseline / ms_packed if ms_packed > 0 else float("inf")
saved_us = (ms_baseline - ms_packed) * 1000
@@ -245,7 +304,9 @@ def bench_shape(B, H, HV, K, V, pool_size, device, dtype):
f"{ms_baseline * 1000:>10.1f} | "
f"{ms_packed * 1000:>10.1f} | "
f"{speedup:>7.2f}x | "
f"{saved_us:>+9.1f}"
f"{saved_us:>+9.1f} | "
f"{ms_helion * 1000:>10.1f} | "
f"{ms_packed / ms_helion:>10.2f}x"
)
@@ -289,8 +350,10 @@ def run_correctness(device, dtype):
)
o_baseline, _ = run_baseline(inp)
o_packed, _ = run_packed(inp)
o_helion, _ = run_helion(inp)
try:
torch.testing.assert_close(o_packed, o_baseline, atol=2e-2, rtol=1e-2)
torch.testing.assert_close(o_helion, o_packed, atol=2e-2, rtol=1e-2)
print(" [PASS] PAD_SLOT_ID=-1 handling")
except AssertionError as e:
print(f" [FAIL] PAD_SLOT_ID=-1 handling: {e}")
@@ -323,9 +386,11 @@ def run_benchmark(device, dtype, args):
f"{'base (us)':>10} | "
f"{'packed (us)':>10} | "
f"{'speedup':>8} | "
f"{'saved (us)':>10}"
f"{'saved (us)':>10} | "
f"{'Helion (us)':>10} | "
f"{'packed/H':>11}"
)
print(" " + "-" * 80)
print(" " + "-" * 116)
for B, H, HV in bench_configs:
# Packed kernel requires HV % H == 0 (GVA / grouped query layout).
benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py
+43 -3
View File
@@ -28,6 +28,9 @@ import argparse
import torch
from sglang.kernels.ops.attention.helion.kda_decode import (
helion_fused_recurrent_kda_packed_decode,
)
from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
@@ -59,6 +62,9 @@ def make_decode_inputs(B, H, HV, K, V, pool_size, device, dtype, seed=42):
ssm = torch.randn(pool_size, HV, V, K, device=device, dtype=dtype) * 0.01
cache_indices = torch.arange(B, device=device, dtype=torch.int32)
qsl = torch.arange(B + 1, device=device, dtype=torch.int32)
mixed_qkv = torch.cat(
(q.reshape(B, H * K), k.reshape(B, H * K), v.reshape(B, HV * V)), dim=-1
)
return dict(
q=q.contiguous(),
k=k.contiguous(),
@@ -70,6 +76,7 @@ def make_decode_inputs(B, H, HV, K, V, pool_size, device, dtype, seed=42):
ssm=ssm.contiguous(),
cache_indices=cache_indices,
qsl=qsl,
mixed_qkv=mixed_qkv.contiguous(),
B=B,
H=H,
HV=HV,
@@ -139,6 +146,23 @@ def call_decode(kernel, inp, ssm):
return out.reshape(inp["B"], inp["HV"], inp["V"]).float()
def call_helion(inp, ssm):
out = inp["mixed_qkv"].new_empty(inp["B"], 1, inp["HV"], inp["V"])
helion_fused_recurrent_kda_packed_decode(
mixed_qkv=inp["mixed_qkv"],
a=inp["a"],
b=inp["b"],
A_log=inp["A_log"],
dt_bias=inp["dt_bias"],
scale=inp["K"] ** -0.5,
initial_state=ssm,
out=out,
ssm_state_indices=inp["cache_indices"],
use_qk_l2norm_in_kernel=True,
)
return out.reshape(inp["B"], inp["HV"], inp["V"]).float()
def call_verify(kernel, inp, ssm, intermediate_states):
out = kernel.target_verify(
A_log=inp["A_log"],
@@ -186,11 +210,14 @@ def run(task, fi, tri, device, dtype, args):
)
print("=" * 92)
hdr = "B" if not is_verify else "B(xT)"
print(
header = (
f" {hdr:>6} {'H':>3} {'HV':>3} | {'triton(us)':>11} | "
f"{'flashinfer(us)':>14} | {'speedup':>8} | {'out_max_diff':>12}"
)
print(" " + "-" * 86)
if not is_verify:
header += f" | {'helion(us)':>11} | {'FI/H':>8} | {'H max diff':>10}"
print(header)
print(" " + "-" * (124 if not is_verify else 86))
for B in args.batch_sizes:
for H in args.num_q_heads:
@@ -235,10 +262,23 @@ def run(task, fi, tri, device, dtype, args):
)
fi_us = f"{ms_fi * 1000:>14.1f}" if fi is not None else f"{'skip':>14}"
sp = f"{speed:>7.2f}x" if fi is not None else f"{'-':>8}"
print(
line = (
f" {B:>6} {H:>3} {HV:>3} | {ms_tri * 1000:>11.1f} | "
f"{fi_us} | {sp} | {diff:>12}"
)
if not is_verify:
helion_state = inp["ssm"].clone()
o_helion = call_helion(inp, helion_state)
helion_diff = (o_helion - o_tri).abs().max().item()
ms_helion = _time(lambda: call_helion(inp, helion_state))
fi_helion = (
f"{ms_fi / ms_helion:>7.2f}x" if fi is not None else f"{'-':>8}"
)
line += (
f" | {ms_helion * 1000:>11.1f} | {fi_helion} | "
f"{helion_diff:>10.2e}"
)
print(line)
def main():
benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py
+63 -3
View File
@@ -4,6 +4,7 @@ Benchmark & Correctness: Triton KDA vs CuTeDSL KDA (prefill, SM100 Blackwell).
Compares:
- Triton: sglang's chunk_kda (FLA chunkwise gated delta rule, per-channel gate)
- CuteDSL: kda_blackwell pipeline (fused Triton prologue -> kkt_inv_uw -> h -> o)
- Helion: sglang's Helion chunk_kda
KDA differs from GDN by a PER-CHANNEL decay gate (g is [T, H, K], not scalar).
The cutedsl pipeline externalizes the per-channel decay into five pre-scaled
@@ -30,6 +31,9 @@ import torch
import torch.nn.functional as F
from sglang.kernels.ops.attention.fla.kda import chunk_kda, fused_recurrent_kda
from sglang.kernels.ops.attention.helion.kda_prefill import (
chunk_kda as helion_chunk_kda,
)
from sglang.kernels.ops.attention.linear.kda_blackwell import prepare_metadata
from sglang.kernels.ops.attention.linear.kda_blackwell.kernel_h import (
kda_h_cutedsl,
@@ -132,6 +136,7 @@ def cutedsl_buffers(inp, num_sms, device):
total=total,
num_sms=num_sms,
h0=torch.zeros(1, H, V, K, device=device, dtype=torch.float32),
state_indices=torch.zeros(1, device=device, dtype=torch.int32),
U=torch.empty(pad_t, H, V, device=device, dtype=torch.bfloat16),
W=torch.empty(pad_t, H, K, device=device, dtype=torch.bfloat16),
V_new=torch.empty(pad_t, H, V, device=device, dtype=torch.bfloat16),
@@ -172,6 +177,7 @@ def run_cutedsl_pipeline(inp, buf, scale):
buf["ht"],
buf["cu"],
buf["co"],
buf["state_indices"],
)
kda_o_cutedsl(
qg,
@@ -219,7 +225,35 @@ def check_shape(T, H, K, V, device, dtype, num_sms):
print(
f" [{status}] {tag} | o_err {o_err:.2e} state_err {s_err:.2e} finite={finite}"
)
return ok
helion_state = torch.zeros(1, H, V, K, device=device, dtype=torch.float32)
helion_o = helion_chunk_kda(
q=inp["q"],
k=inp["k"],
v=inp["v"].clone(),
g=inp["g_act"],
beta=inp["beta"],
scale=scale,
initial_state=helion_state,
initial_state_indices=torch.zeros(1, device=device, dtype=torch.int32),
use_qk_l2norm_in_kernel=False,
cu_seqlens=None,
A_log=None,
dt_bias=None,
lower_bound=None,
)
helion_finite = bool(
torch.isfinite(helion_o).all() and torch.isfinite(helion_state).all()
)
helion_o_err = (helion_o[0].float() - o_ref.float()).abs().max().item()
helion_s_err = (helion_state.float() - state_ref.float()).abs().max().item()
helion_ok = helion_finite and helion_o_err < 1e-2 and helion_s_err < 5e-2
print(
f" [{'PASS' if helion_ok else 'FAIL'}] Helion {tag} | "
f"o_err {helion_o_err:.2e} state_err {helion_s_err:.2e} "
f"finite={helion_finite}"
)
return ok and helion_ok
# ---------------------------------------------------------------------------
@@ -264,15 +298,37 @@ def bench_shape(T, H, K, V, device, dtype, num_sms):
def fn_cutedsl():
run_cutedsl_pipeline(inp, buf, scale)
helion_v = v.clone()
helion_state = torch.zeros(1, H, V, K, device=device, dtype=torch.float32)
def fn_helion():
helion_chunk_kda(
q=q,
k=k,
v=helion_v,
g=g_act,
beta=beta,
scale=scale,
initial_state=helion_state,
initial_state_indices=idx,
use_qk_l2norm_in_kernel=False,
cu_seqlens=None,
A_log=None,
dt_bias=None,
lower_bound=None,
)
quantiles = [0.5, 0.2, 0.8]
fn_triton()
fn_cutedsl()
fn_helion()
torch.cuda.synchronize()
ms_triton, _, _ = triton.testing.do_bench_cudagraph(fn_triton, quantiles=quantiles)
ms_cutedsl, _, _ = triton.testing.do_bench_cudagraph(
fn_cutedsl, quantiles=quantiles
)
ms_helion, _, _ = triton.testing.do_bench_cudagraph(fn_helion, quantiles=quantiles)
flops = kda_flops(T, H, K, V)
mem_bytes = kda_bytes(T, H, K, V, 1, dtype)
@@ -281,7 +337,10 @@ def bench_shape(T, H, K, V, device, dtype, num_sms):
f" {H:>3} {T:>7} | "
f"{ms_triton:>8.3f} {flops / ms_triton / 1e9:>7.2f} {mem_bytes / ms_triton / 1e9:>7.2f} | "
f"{ms_cutedsl:>8.3f} {flops / ms_cutedsl / 1e9:>7.2f} {mem_bytes / ms_cutedsl / 1e9:>7.2f} | "
f"{speedup:>7.2f}x"
f"{speedup:>7.2f}x | "
f"{ms_helion:>8.3f} {flops / ms_helion / 1e9:>7.2f} "
f"{mem_bytes / ms_helion / 1e9:>7.2f} | "
f"{ms_triton / ms_helion:>7.2f}x"
)
@@ -312,8 +371,9 @@ def run_benchmark(device, dtype, args, num_sms):
f" {'H':>3} {'T':>7} | "
f"{'tri(ms)':>8} {'TFLOP':>7} {'TB/s':>7} | "
f"{'cute(ms)':>8} {'TFLOP':>7} {'TB/s':>7} | {'speedup':>8}"
f" | {'helion(ms)':>10} {'TFLOP':>7} {'TB/s':>7} | {'tri/hel':>8}"
)
print(" " + "-" * 84)
print(" " + "-" * 126)
for H in args.num_heads:
for T in args.seq_lens:
bench_shape(T, H, 128, 128, device, dtype, num_sms)
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
)
test/registered/kernels/ops/attention/test_kda_helion.py
+1026
View File
File diff suppressed because it is too large Load Diff
test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py
+13
View File
@@ -217,6 +217,19 @@ class TestFlashInferGDNPrefillBackendPolicy(unittest.TestCase):
tree_verify.assert_called_once()
flashinfer_kernel.target_verify.assert_not_called()
def test_helion_backend_reports_kda_only(self):
cases = (
(LinearAttnKernelBackend.HELION, LinearAttnKernelBackend.TRITON),
(LinearAttnKernelBackend.TRITON, LinearAttnKernelBackend.HELION),
)
for decode_backend, prefill_backend in cases:
with self.subTest(
decode_backend=decode_backend,
prefill_backend=prefill_backend,
):
with self.assertRaisesRegex(ValueError, "supports KDA only"):
GDNKernelDispatcher(decode_backend, prefill_backend)
if __name__ == "__main__":
unittest.main()
test/registered/unit/layers/attention/test_kda_helion_dispatcher.py
+197
View File
@@ -0,0 +1,197 @@
import unittest
from unittest.mock import ANY, MagicMock, patch
import torch
from sglang.srt.layers.attention.linear.kda_backend import KDAKernelDispatcher
from sglang.srt.layers.attention.linear.kernels.kda_helion import HelionKDAKernel
from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
from sglang.srt.layers.attention.linear.utils import LinearAttnKernelBackend
from sglang.srt.server_args import ServerArgs
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
class TestHelionKDADispatcher(unittest.TestCase):
def _make_dispatcher(self, decode_backend, prefill_backend):
helion_kernel = MagicMock(supports_packed_decode=True)
with (
patch(
"sglang.srt.layers.attention.linear.kda_backend.is_cuda",
return_value=True,
),
patch(
"sglang.srt.layers.attention.linear.kernels.kda_helion."
"HelionKDAKernel",
return_value=helion_kernel,
) as constructor,
):
dispatcher = KDAKernelDispatcher(
decode_backend=decode_backend,
prefill_backend=prefill_backend,
verify_backend=LinearAttnKernelBackend.TRITON,
)
return dispatcher, helion_kernel, constructor
def test_combined_backend_reuses_adapter_and_keeps_triton_verify(self):
dispatcher, helion_kernel, constructor = self._make_dispatcher(
LinearAttnKernelBackend.HELION,
LinearAttnKernelBackend.HELION,
)
constructor.assert_called_once_with(
triton_fallback=ANY,
enable_decode=True,
enable_prefill=True,
)
self.assertIs(dispatcher.decode_kernel, helion_kernel)
self.assertIs(dispatcher.extend_kernel, helion_kernel)
self.assertIsInstance(dispatcher.verify_kernel, TritonKDAKernel)
self.assertTrue(dispatcher.supports_packed_decode)
def test_decode_only_keeps_triton_prefill_and_verify(self):
dispatcher, helion_kernel, constructor = self._make_dispatcher(
LinearAttnKernelBackend.HELION,
LinearAttnKernelBackend.TRITON,
)
constructor.assert_called_once_with(
triton_fallback=ANY,
enable_decode=True,
enable_prefill=False,
)
self.assertIs(dispatcher.decode_kernel, helion_kernel)
self.assertIsInstance(dispatcher.extend_kernel, TritonKDAKernel)
self.assertIsInstance(dispatcher.verify_kernel, TritonKDAKernel)
def test_prefill_only_keeps_triton_decode_and_verify(self):
dispatcher, helion_kernel, constructor = self._make_dispatcher(
LinearAttnKernelBackend.TRITON,
LinearAttnKernelBackend.HELION,
)
constructor.assert_called_once_with(
triton_fallback=ANY,
enable_decode=False,
enable_prefill=True,
)
self.assertIsInstance(dispatcher.decode_kernel, TritonKDAKernel)
self.assertIs(dispatcher.extend_kernel, helion_kernel)
self.assertIsInstance(dispatcher.verify_kernel, TritonKDAKernel)
def test_enum_recognizes_helion(self):
backend = LinearAttnKernelBackend("helion")
self.assertIs(backend, LinearAttnKernelBackend.HELION)
self.assertTrue(backend.is_helion())
def test_replayssm_decode_uses_native_helion_kernel(self):
kernel = HelionKDAKernel.__new__(HelionKDAKernel)
kernel._packed_decode = MagicMock()
kernel._replayssm_decode = MagicMock()
kernel._triton = MagicMock()
mixed_qkv = torch.empty(2, 20)
a = torch.empty(2, 8)
b = torch.empty(2, 1)
state = torch.empty(2, 1, 4, 8)
indices = torch.arange(2, dtype=torch.int32)
force_flush = torch.zeros(2, dtype=torch.int32)
replay_args = {
"replayssm_d": torch.empty(2, 1, 4, 4),
"replayssm_k": torch.empty(2, 1, 4, 8),
"replayssm_g": torch.empty(2, 1, 4, 8),
"replayssm_write_pos": torch.zeros(2, dtype=torch.int32),
"replayssm_force_flush": force_flush,
}
result = kernel.packed_decode(
mixed_qkv,
a,
b,
A_log=torch.empty(1),
dt_bias=torch.empty(8),
scale=0.5,
ssm_states=state,
cache_indices=indices,
num_v_heads=1,
head_v_dim=4,
lower_bound=-5.0,
**replay_args,
)
self.assertEqual(result.shape, (1, 2, 1, 4))
kernel._packed_decode.assert_not_called()
kernel._replayssm_decode.assert_called_once()
self.assertIs(
kernel._replayssm_decode.call_args.kwargs["force_flush"], force_flush
)
self.assertEqual(kernel._replayssm_decode.call_args.kwargs["lower_bound"], -5.0)
kernel._triton.packed_decode.assert_not_called()
def test_packed_decode_forwards_lower_bound(self):
kernel = HelionKDAKernel.__new__(HelionKDAKernel)
kernel._packed_decode = MagicMock()
kernel._triton = MagicMock()
mixed_qkv = torch.empty(2, 16)
a = torch.empty(2, 8)
b = torch.empty(2, 1)
a_log = torch.empty(1)
dt_bias = torch.empty(8)
state = torch.empty(2, 1, 4, 8)
indices = torch.arange(2, dtype=torch.int32)
kernel.packed_decode(
mixed_qkv,
a,
b,
A_log=a_log,
dt_bias=dt_bias,
scale=0.5,
ssm_states=state,
cache_indices=indices,
num_v_heads=1,
head_v_dim=4,
lower_bound=-5.0,
)
kernel._packed_decode.assert_called_once()
self.assertEqual(kernel._packed_decode.call_args.kwargs["lower_bound"], -5.0)
def test_replayssm_accepts_helion_and_rejects_other_backends(self):
with (
patch("sglang.srt.server_args.is_sm100_supported", return_value=False),
patch("sglang.srt.server_args.is_cuda", return_value=False),
):
helion_args = ServerArgs(
model_path="dummy",
linear_attn_decode_backend="helion",
enable_linear_replayssm=True,
)
helion_args._handle_linear_attn_backend()
flashinfer_args = ServerArgs(
model_path="dummy",
linear_attn_decode_backend="flashinfer",
enable_linear_replayssm=True,
)
with self.assertRaisesRegex(ValueError, "Triton, or Helion"):
flashinfer_args._handle_linear_attn_backend()
def test_explicit_base_backend_is_not_replaced_by_flashinfer(self):
args = ServerArgs(
model_path="dummy",
linear_attn_backend="helion",
mamba_ssm_dtype="bfloat16",
)
with (
patch("sglang.srt.server_args.is_sm100_supported", return_value=True),
patch("sglang.srt.server_args.is_cuda", return_value=False),
):
args._handle_linear_attn_backend()
self.assertIsNone(args.linear_attn_decode_backend)
self.assertEqual(args.linear_attn_backend, "helion")
if __name__ == "__main__":
unittest.main()
test/registered/unit/server_args/test_page_major_backend_allowlist.py
+21 -3
View File
@@ -37,7 +37,14 @@ from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _accepts(backend: str, *, use_mla: bool, unified: bool = True) -> bool:
def _accepts(
backend: str,
*,
use_mla: bool,
unified: bool = True,
linear_decode: str | None = None,
linear_prefill: str | None = None,
) -> bool:
"""Run just `_handle_page_major_kv_layout` against a minimal stand-in.
ServerArgs' real constructor pulls in a model config; this exercises the
@@ -53,8 +60,8 @@ def _accepts(backend: str, *, use_mla: bool, unified: bool = True) -> bool:
"prefill_attention_backend": None,
"decode_attention_backend": None,
"linear_attn_backend": "triton",
"linear_attn_decode_backend": None,
"linear_attn_prefill_backend": None,
"linear_attn_decode_backend": linear_decode,
"linear_attn_prefill_backend": linear_prefill,
"mamba_backend": "triton",
}.items():
object.__setattr__(sa, name, value)
@@ -115,6 +122,17 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
f"{backend} has no dense-id remapping and must be rejected",
)
def test_helion_linear_attention_is_kda_only(self):
for unified in (True, False):
for phase in ("decode", "prefill"):
kwargs = {f"linear_{phase}": "helion"}
self.assertTrue(
_accepts("triton", use_mla=True, unified=unified, **kwargs)
)
self.assertFalse(
_accepts("triton", use_mla=False, unified=unified, **kwargs)
)
if __name__ == "__main__":
unittest.main()