minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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)