From aeee1562e62d41ed68aa3d8fd96e72e42b14958c Mon Sep 17 00:00:00 2001 From: ethche Date: Sat, 15 Aug 2026 01:22:46 -0400 Subject: [PATCH] [Kernel] Enable Helion backend for Kimi Delta-Attention (#32593) Co-authored-by: Ethan Che --- .../bench_kda_decode.py | 75 +- .../bench_kda_flashinfer_mtp.py | 46 +- .../bench_kda_prefill_cutedsl.py | 66 +- .../kernels/ops/attention/helion/__init__.py | 5 + .../ops/attention/helion/kda_decode.py | 384 +++++ .../ops/attention/helion/kda_prefill.py | 1399 +++++++++++++++++ .../ops/attention/helion/kda_replayssm.py | 780 +++++++++ .../layers/attention/linear/gdn_backend.py | 8 + .../layers/attention/linear/kda_backend.py | 26 +- .../attention/linear/kernels/kda_helion.py | 191 +++ .../srt/layers/attention/linear/utils.py | 4 + python/sglang/srt/server_args.py | 42 +- .../kernels/ops/attention/test_kda_helion.py | 1026 ++++++++++++ .../test_gdn_prefill_backend_policy.py | 13 + .../attention/test_kda_helion_dispatcher.py | 197 +++ .../test_page_major_backend_allowlist.py | 24 +- 16 files changed, 4248 insertions(+), 38 deletions(-) create mode 100644 python/sglang/kernels/ops/attention/helion/__init__.py create mode 100644 python/sglang/kernels/ops/attention/helion/kda_decode.py create mode 100644 python/sglang/kernels/ops/attention/helion/kda_prefill.py create mode 100644 python/sglang/kernels/ops/attention/helion/kda_replayssm.py create mode 100644 python/sglang/srt/layers/attention/linear/kernels/kda_helion.py create mode 100644 test/registered/kernels/ops/attention/test_kda_helion.py create mode 100644 test/registered/unit/layers/attention/test_kda_helion_dispatcher.py diff --git a/benchmark/bench_linear_attention/bench_kda_decode.py b/benchmark/bench_linear_attention/bench_kda_decode.py index c06d254c8..f5846a387 100644 --- a/benchmark/bench_linear_attention/bench_kda_decode.py +++ b/benchmark/bench_linear_attention/bench_kda_decode.py @@ -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). diff --git a/benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py b/benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py index 31e1e1718..5e988392e 100644 --- a/benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py +++ b/benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py @@ -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(): diff --git a/benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py b/benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py index 9b438d015..aac599ffe 100644 --- a/benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py +++ b/benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py @@ -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) diff --git a/python/sglang/kernels/ops/attention/helion/__init__.py b/python/sglang/kernels/ops/attention/helion/__init__.py new file mode 100644 index 000000000..749da3d74 --- /dev/null +++ b/python/sglang/kernels/ops/attention/helion/__init__.py @@ -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 diff --git a/python/sglang/kernels/ops/attention/helion/kda_decode.py b/python/sglang/kernels/ops/attention/helion/kda_decode.py new file mode 100644 index 000000000..8bda2e945 --- /dev/null +++ b/python/sglang/kernels/ops/attention/helion/kda_decode.py @@ -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 diff --git a/python/sglang/kernels/ops/attention/helion/kda_prefill.py b/python/sglang/kernels/ops/attention/helion/kda_prefill.py new file mode 100644 index 000000000..931b370fc --- /dev/null +++ b/python/sglang/kernels/ops/attention/helion/kda_prefill.py @@ -0,0 +1,1399 @@ +"""Helion kernels for SGLang's Kimi Delta Attention prefill path. + +The public :func:`chunk_kda` entry point in this module is intended to match +``sglang.kernels.ops.attention.fla.kda.chunk_kda``. KDA uses 64-token chunks +and keeps the cumulative per-key decay in base-2 logarithm space. +""" + +from __future__ import annotations + +import helion +import helion.language as hl +import torch + +from sglang.kernels.ops.attention.fla.index import ( + prepare_chunk_indices, + prepare_chunk_offsets, +) + +CHUNK_SIZE = 64 +# Rounded identically to flash-linear-attention/SGLang before FP32 multiply. +RCP_LN2 = 1.4426950216293335 +L2_NORM_EPS = 1e-6 +SOFTPLUS_THRESHOLD = 20.0 + +# SGLang initializes torch.distributed, but these kernels have no collectives. +_IGNORED_WARNINGS = [helion.exc.ProcessGroupNameNotFound] +# K3 exposes 12 local heads at TP=8. Lower head counts share the same +# low-occupancy packed-varlen state-propagation regime. +_PREFILL_SMALL_HEAD_THRESHOLD = 12 + + +_L2_NORM_CONFIG = helion.Config( + block_sizes=[8], + num_warps=4, + num_stages=2, + indexing="pointer", +) + + +@helion.kernel( + static_shapes=False, + config=_L2_NORM_CONFIG, + ignore_warnings=_IGNORED_WARNINGS, +) +def _l2norm_qk( + q: torch.Tensor, + k: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Normalize Q and K rows with Triton's FP32 accumulation contract.""" + B = q.size(0) + T = q.size(1) + H = hl.specialize(q.size(2)) + K = hl.specialize(q.size(3)) + hl.specialize( + ( + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(1), + k.stride(2), + k.stride(3), + ) + ) + + q_out = torch.empty_like(q) + k_out = torch.empty_like(k) + q_rows = q.view(B * T * H, K) + k_rows = k.view(B * T * H, K) + q_out_rows = q_out.view(B * T * H, K) + k_out_rows = k_out.view(B * T * H, K) + block_rows = hl.register_block_size(1, 16) + + for tile_rows in hl.tile(B * T * H, block_size=block_rows): + q_value = q_rows[tile_rows, :].float() + k_value = k_rows[tile_rows, :].float() + q_norm = torch.sqrt((q_value * q_value).sum(-1) + L2_NORM_EPS) + k_norm = torch.sqrt((k_value * k_value).sum(-1) + L2_NORM_EPS) + q_out_rows[tile_rows, :] = (q_value / q_norm[:, None]).to(q.dtype) + k_out_rows[tile_rows, :] = (k_value / k_norm[:, None]).to(k.dtype) + + return q_out, k_out + + +_GATE_FIXED_CONFIG = helion.Config( + block_sizes=[16], + loop_orders=[[1, 2, 0]], + num_warps=1, + num_stages=1, + indexing="pointer", +) + + +# Fixed chunks use arithmetic indexing; varlen chunks load sequence metadata. +# Their generated kernels have different occupancy and cache-policy optima. +# Eviction policies are positional in the traced load order. Retune them if +# the shared gate body gains, loses, or reorders loads. +_GATE_VARLEN_CONFIG = helion.Config( + block_sizes=[16], + # The sixth generated load streams the gate tile once per program. + load_eviction_policies=["", "", "", "", "", "first"] + [""] * 11, + loop_orders=[[2, 1, 0]], + num_warps=8, + num_stages=1, + indexing="pointer", +) + + +def _activate_gate( + raw_gate: torch.Tensor, + a_log: torch.Tensor, + lower_bound: float, + use_lower_bound: hl.constexpr, +) -> torch.Tensor: + a = torch.exp2(a_log.float() * RCP_LN2) + if use_lower_bound: + return lower_bound * torch.sigmoid(a * raw_gate) + softplus = torch.where( + raw_gate < SOFTPLUS_THRESHOLD, + torch.log(1.0 + torch.exp2(raw_gate * RCP_LN2)), + raw_gate, + ) + return -a * softplus + + +@helion.kernel( + static_shapes=False, + config=_GATE_FIXED_CONFIG, + ignore_warnings=_IGNORED_WARNINGS, +) +def _gate_cumsum_operands( + g: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + beta: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + cu_seqlens: torch.Tensor, + chunk_indices: torch.Tensor, + gate_scale: float, + q_scale: float, + lower_bound: float, + activate: hl.constexpr, # pyrefly: ignore[bad-function-definition] + has_bias: hl.constexpr, # pyrefly: ignore[bad-function-definition] + use_lower_bound: hl.constexpr, # pyrefly: ignore[bad-function-definition] + is_varlen: hl.constexpr, # pyrefly: ignore[bad-function-definition] +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Compute cumulative gates and rounded Q/K operands in one pass.""" + B = g.size(0) + T = g.size(1) + H = hl.specialize(g.size(2)) + K = hl.specialize(g.size(3)) + chunks_per_batch = (T + CHUNK_SIZE - 1) // CHUNK_SIZE + total_chunks = chunk_indices.size(0) if is_varlen else B * chunks_per_batch + hl.specialize( + ( + g.stride(1), + g.stride(2), + g.stride(3), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(1), + k.stride(2), + k.stride(3), + beta.stride(1), + beta.stride(2), + a_log.stride(0), + dt_bias.stride(0), + cu_seqlens.stride(0), + chunk_indices.stride(0), + chunk_indices.stride(1), + ) + ) + + out = torch.empty_like(g, dtype=torch.float32) + qg = torch.empty_like(q) + wk = torch.empty_like(k) + kg = torch.empty_like(k) + chunk_decay = torch.empty( + [total_chunks, H, K], + dtype=torch.float32, + device=g.device, + ) + g_rows = g.view(B * T * H, K) + q_rows = q.view(B * T * H, K) + k_rows = k.view(B * T * H, K) + beta_rows = beta.view(B * T * H) + out_rows = out.view(B * T * H, K) + qg_rows = qg.view(B * T * H, K) + wk_rows = wk.view(B * T * H, K) + kg_rows = kg.view(B * T * H, K) + block_k = hl.register_block_size(16, K) + + for tile_chunk, tile_h, tile_k in hl.tile( + [total_chunks, H, K], + block_size=[1, 1, block_k], + ): + if is_varlen: + sequence = chunk_indices[tile_chunk.id, 0].long() + local_chunk = chunk_indices[tile_chunk.id, 1].long() + begin = cu_seqlens[sequence].long() + end = cu_seqlens[sequence + 1].long() + else: + sequence = tile_chunk.id // chunks_per_batch + local_chunk = tile_chunk.id % chunks_per_batch + begin = sequence * T + end = begin + T + time = hl.arange(64) + token = begin + local_chunk * CHUNK_SIZE + time + valid = token < end + row = token * H + tile_h.id + value = hl.load( + g_rows, + [row[:, None], tile_k.index[None, :]], + extra_mask=valid[:, None], + ).float() + if activate: + if has_bias: + value = value + dt_bias[tile_h.id * K + tile_k.index].float()[None, :] + value = _activate_gate( + value, + a_log[tile_h.id], + lower_bound, + use_lower_bound, + ) + value = torch.where(valid[:, None], value, 0.0) + value = torch.cumsum(value, dim=0) * gate_scale + q_value = hl.load( + q_rows, + [row[:, None], tile_k.index[None, :]], + extra_mask=valid[:, None], + ).float() + k_value = hl.load( + k_rows, + [row[:, None], tile_k.index[None, :]], + extra_mask=valid[:, None], + ).float() + beta_value = hl.load(beta_rows, [row], extra_mask=valid).float() + last_gate = torch.where(time[:, None] == 63, value, 0.0).sum(0) + gate_value = torch.exp2(value) + hl.store( + out_rows, + [row[:, None], tile_k.index[None, :]], + value, + extra_mask=valid[:, None], + ) + hl.store( + qg_rows, + [row[:, None], tile_k.index[None, :]], + (q_value * q_scale * gate_value).to(q.dtype), + extra_mask=valid[:, None], + ) + hl.store( + wk_rows, + [row[:, None], tile_k.index[None, :]], + (k_value * beta_value[:, None] * gate_value).to(k.dtype), + extra_mask=valid[:, None], + ) + hl.store( + kg_rows, + [row[:, None], tile_k.index[None, :]], + (k_value * torch.exp2(last_gate[None, :] - value)).to(k.dtype), + extra_mask=valid[:, None], + ) + hl.store( + chunk_decay, + [tile_chunk.id, tile_h.id, tile_k.index], + torch.exp2(last_gate), + ) + + return out, qg, wk, kg, chunk_decay + + +_gate_cumsum_operands_varlen = helion.kernel( + static_shapes=False, + config=_GATE_VARLEN_CONFIG, + ignore_warnings=_IGNORED_WARNINGS, +)(_gate_cumsum_operands.fn) + + +def gate_chunk_cumsum_operands( + g: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + beta: torch.Tensor, + *, + q_scale: float, + a_log: torch.Tensor | None, + dt_bias: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, + chunk_indices: torch.Tensor | None = None, + lower_bound: float | None = None, + gate_scale: float = RCP_LN2, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Gate preprocessing plus rounded operands reused by later stages.""" + flat_a_log = ( + a_log.reshape(-1) + if a_log is not None + else torch.empty(1, device=g.device, dtype=torch.float32) + ) + flat_bias = ( + dt_bias.reshape(-1) + if dt_bias is not None + else torch.empty(1, device=g.device, dtype=torch.float32) + ) + activate = a_log is not None + has_bias = dt_bias is not None + use_lower_bound = lower_bound is not None + lower_bound_value = 0.0 if lower_bound is None else lower_bound + + is_varlen = cu_seqlens is not None + if is_varlen: + if g.size(0) != 1: + raise ValueError("varlen KDA requires batch size 1") + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, CHUNK_SIZE) + metadata = cu_seqlens + gate_kernel = _gate_cumsum_operands_varlen + else: + metadata = torch.empty(0, device=g.device, dtype=torch.int32) + chunk_indices = torch.empty(0, 2, device=g.device, dtype=torch.long) + gate_kernel = _gate_cumsum_operands + + return gate_kernel( + g, + q, + k, + beta, + flat_a_log, + flat_bias, + metadata, + chunk_indices, + gate_scale, + q_scale, + lower_bound_value, + activate, + has_bias, + use_lower_bound, + is_varlen, + ) + + +_INTRA_MATRIX_CONFIG = helion.Config( + block_sizes=[32], + loop_orders=[[1, 2, 0]], + num_warps=1, + num_stages=2, + indexing="pointer", +) + + +@helion.kernel( + static_shapes=False, + config=_INTRA_MATRIX_CONFIG, + ignore_warnings=_IGNORED_WARNINGS, +) +def _intra_matrices_wide( + q: torch.Tensor, + k: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + cu_seqlens: torch.Tensor, + chunk_indices: torch.Tensor, + scale: float, + is_varlen: hl.constexpr, # pyrefly: ignore[bad-function-definition] +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute a full 16x64 causal matrix row per CTA.""" + B = q.size(0) + T = q.size(1) + H = hl.specialize(q.size(2)) + K = hl.specialize(q.size(3)) + chunks_per_batch = (T + CHUNK_SIZE - 1) // CHUNK_SIZE + total_chunks = chunk_indices.size(0) if is_varlen else B * chunks_per_batch + hl.specialize( + ( + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(1), + k.stride(2), + k.stride(3), + g.stride(1), + g.stride(2), + g.stride(3), + beta.stride(1), + beta.stride(2), + ) + ) + + aqk = torch.empty([B, T, H, CHUNK_SIZE], dtype=q.dtype, device=q.device) + akk = torch.empty([B, T, H, CHUNK_SIZE], dtype=torch.float32, device=q.device) + q_rows = q.view(B * T * H, K) + k_rows = k.view(B * T * H, K) + g_rows = g.view(B * T * H, K) + beta_rows = beta.view(B * T * H) + aqk_rows = aqk.view(B * T * H, CHUNK_SIZE) + akk_rows = akk.view(B * T * H, CHUNK_SIZE) + block_k = hl.register_block_size(32, K) + + for tile_chunk, tile_h, tile_row_block in hl.tile( + [total_chunks, H, 4], + block_size=[1, 1, 1], + ): + if is_varlen: + sequence = chunk_indices[tile_chunk.id, 0].long() + local_chunk = chunk_indices[tile_chunk.id, 1].long() + begin = cu_seqlens[sequence].long() + end = cu_seqlens[sequence + 1].long() + else: + sequence = tile_chunk.id // chunks_per_batch + local_chunk = tile_chunk.id % chunks_per_batch + begin = sequence * T + end = begin + T + + row_lane = hl.arange(16) + col_lane = hl.arange(64) + chunk_begin = begin + local_chunk * CHUNK_SIZE + row_local = tile_row_block.id * 16 + row_lane + row_token = chunk_begin + row_local + col_token = chunk_begin + col_lane + row_valid = row_token < end + col_valid = col_token < end + block_causal = col_lane < (tile_row_block.id + 1) * 16 + row = row_token * H + tile_h.id + col = col_token * H + tile_h.id + anchor_token = chunk_begin + tile_row_block.id * 16 + anchor = anchor_token * H + tile_h.id + # Varlen ``end`` is a loaded scalar tensor; fixed ``end`` is symbolic. + if is_varlen: + diag_anchor_token = torch.minimum(anchor_token + 8, end - 1) + else: + diag_anchor_token = min(anchor_token + 8, end - 1) + diag_anchor = diag_anchor_token * H + tile_h.id + aqk_off = hl.zeros([16, 64], dtype=torch.float32) + akk_off = hl.zeros([16, 64], dtype=torch.float32) + aqk_diag = hl.zeros([16, 16], dtype=torch.float32) + akk_diag = hl.zeros([16, 16], dtype=torch.float32) + + for tile_k in hl.tile(K, block_size=block_k): + q_row = hl.load( + q_rows, + [row[:, None], tile_k.index[None, :]], + extra_mask=row_valid[:, None], + ).float() + k_row = hl.load( + k_rows, + [row[:, None], tile_k.index[None, :]], + extra_mask=row_valid[:, None], + ).float() + g_row = hl.load( + g_rows, + [row[:, None], tile_k.index[None, :]], + extra_mask=row_valid[:, None], + ).float() + g_anchor = hl.load( + g_rows, + [anchor, tile_k.index], + extra_mask=anchor_token < end, + ).float() + g_diag_anchor = hl.load( + g_rows, + [diag_anchor, tile_k.index], + extra_mask=diag_anchor_token < end, + ).float() + if tile_row_block.id > 0: + k_col = hl.load( + k_rows, + [col[:, None], tile_k.index[None, :]], + extra_mask=col_valid[:, None] & block_causal[:, None], + ).float() + g_col = hl.load( + g_rows, + [col[:, None], tile_k.index[None, :]], + extra_mask=col_valid[:, None] & block_causal[:, None], + ).float() + off_col = col_lane < tile_row_block.id * 16 + off_col_delta = torch.where( + off_col[:, None], + g_anchor[None, :] - g_col, + 0.0, + ) + off_row_delta = g_row - g_anchor[None, :] + # Masked rows load zero; cap only their positive overflow edge. + # Valid KDA cumulative gates are non-increasing in each chunk. + off_row_delta = torch.clamp(off_row_delta, max=126.0) + off_row_factor = torch.exp2(off_row_delta) + q_off = (q_row * off_row_factor).to(torch.bfloat16) + k_off = (k_row * off_row_factor).to(torch.bfloat16) + k_col_off = (k_col * torch.exp2(off_col_delta)).to(torch.bfloat16) + aqk_off = hl.dot( + q_off, + k_col_off.T, + acc=aqk_off, + out_dtype=torch.float32, + ) + akk_off = hl.dot( + k_off, + k_col_off.T, + acc=akk_off, + out_dtype=torch.float32, + ) + + diag_delta = torch.clamp( + g_row - g_diag_anchor[None, :], + -126.0, + 126.0, + ) + diag_forward_factor = torch.exp2(diag_delta) + diag_backward_factor = torch.exp2(-diag_delta) + q_diag = q_row * diag_forward_factor + k_diag_fwd = k_row * diag_forward_factor + k_diag_bwd = k_row * diag_backward_factor + aqk_diag = hl.dot( + q_diag, + k_diag_bwd.T, + acc=aqk_diag, + out_dtype=torch.float32, + ) + akk_diag = hl.dot( + k_diag_fwd, + k_diag_bwd.T, + acc=akk_diag, + out_dtype=torch.float32, + ) + + causal = row_local[:, None] >= col_lane[None, :] + strictly_causal = row_local[:, None] > col_lane[None, :] + row_beta = hl.load( + beta_rows, + [row], + extra_mask=row_valid, + ).float() + hl.store( + aqk_rows, + [row[:, None], col_lane[None, :]], + torch.where(causal & col_valid[None, :], aqk_off * scale, 0.0), + extra_mask=row_valid[:, None], + ) + hl.store( + akk_rows, + [row[:, None], col_lane[None, :]], + torch.where( + strictly_causal & col_valid[None, :], + akk_off * row_beta[:, None], + 0.0, + ), + extra_mask=row_valid[:, None], + ) + diag_col = tile_row_block.id * 16 + row_lane + diag_causal = row_lane[:, None] >= row_lane[None, :] + diag_strict = row_lane[:, None] > row_lane[None, :] + diagonal_matrix = torch.where( + diag_strict & row_valid[None, :], + akk_diag * row_beta[:, None], + 0.0, + ) + diagonal_matrix = _invert_lower_16_forward_substitution(diagonal_matrix) + hl.store( + aqk_rows, + [row[:, None], diag_col[None, :]], + torch.where(diag_causal & row_valid[None, :], aqk_diag * scale, 0.0), + extra_mask=row_valid[:, None], + ) + hl.store( + akk_rows, + [row[:, None], diag_col[None, :]], + diagonal_matrix, + extra_mask=row_valid[:, None], + ) + + return aqk, akk + + +def _invert_lower_16_forward_substitution(matrix: torch.Tensor) -> torch.Tensor: + lane = hl.arange(16) + strictly_lower = lane[:, None] > lane[None, :] + inverse = -torch.where(strictly_lower, matrix, 0.0) + for row in range(2, 16): + value = -torch.where((lane == row)[:, None], matrix, 0.0).sum(0) + value = torch.where(lane < row, value, 0.0) + value = value + (value[:, None] * inverse).sum(0) + inverse = torch.where((lane == row)[:, None], value[None, :], inverse) + return inverse + (lane[:, None] == lane[None, :]).float() + + +def _assemble_lower_64_inverse( + matrix: tuple[torch.Tensor, ...], + output_dtype: torch.dtype, +) -> tuple[torch.Tensor, ...]: + """Assemble the 64x64 inverse from pre-inverted 16x16 diagonal blocks.""" + m00, m10, m11, m20, m21, m22, m30, m31, m32, m33 = matrix + i00, i11, i22, i33 = m00, m11, m22, m33 + + i10 = -hl.dot( + hl.dot(i11, m10, out_dtype=torch.float32), + i00, + out_dtype=torch.float32, + ) + i21 = -hl.dot( + hl.dot(i22, m21, out_dtype=torch.float32), + i11, + out_dtype=torch.float32, + ) + i32 = -hl.dot( + hl.dot(i33, m32, out_dtype=torch.float32), + i22, + out_dtype=torch.float32, + ) + i20 = -hl.dot( + i22, + hl.dot(m20, i00, out_dtype=torch.float32) + + hl.dot(m21, i10, out_dtype=torch.float32), + out_dtype=torch.float32, + ) + i31 = -hl.dot( + i33, + hl.dot(m31, i11, out_dtype=torch.float32) + + hl.dot(m32, i21, out_dtype=torch.float32), + out_dtype=torch.float32, + ) + i30 = -hl.dot( + i33, + hl.dot(m30, i00, out_dtype=torch.float32) + + hl.dot(m31, i10, out_dtype=torch.float32) + + hl.dot(m32, i20, out_dtype=torch.float32), + out_dtype=torch.float32, + ) + return tuple( + block.to(output_dtype) + for block in (i00, i10, i11, i20, i21, i22, i30, i31, i32, i33) + ) + + +def _apply_lower_64_blocks( + matrix: tuple[torch.Tensor, ...], + rhs: tuple[torch.Tensor, ...], +) -> tuple[torch.Tensor, ...]: + """Apply a 64x64 lower-triangular matrix to four 16-row blocks.""" + i00, i10, i11, i20, i21, i22, i30, i31, i32, i33 = matrix + x0, x1, x2, x3 = rhs + y0 = hl.dot(i00, x0, out_dtype=torch.float32) + y1 = hl.dot(i10, x0, out_dtype=torch.float32) + hl.dot( + i11, x1, out_dtype=torch.float32 + ) + y2 = ( + hl.dot(i20, x0, out_dtype=torch.float32) + + hl.dot(i21, x1, out_dtype=torch.float32) + + hl.dot(i22, x2, out_dtype=torch.float32) + ) + y3 = ( + hl.dot(i30, x0, out_dtype=torch.float32) + + hl.dot(i31, x1, out_dtype=torch.float32) + + hl.dot(i32, x2, out_dtype=torch.float32) + + hl.dot(i33, x3, out_dtype=torch.float32) + ) + return y0, y1, y2, y3 + + +_SOLVE_RECOMPUTE_CONFIG = helion.Config( + block_sizes=[64, 64], + loop_orders=[[0, 1]], + num_warps=1, + num_stages=3, + indexing="pointer", +) + + +@helion.kernel( + static_shapes=False, + config=_SOLVE_RECOMPUTE_CONFIG, + ignore_warnings=_IGNORED_WARNINGS, +) +def _intra_solve_recompute( + akk: torch.Tensor, + wk: torch.Tensor, + v: torch.Tensor, + beta: torch.Tensor, + cu_seqlens: torch.Tensor, + chunk_indices: torch.Tensor, + is_varlen: hl.constexpr, # pyrefly: ignore[bad-function-definition] +) -> tuple[torch.Tensor, torch.Tensor]: + """Solve the 64x64 system and emit W and U from pre-scaled operands.""" + B = wk.size(0) + T = wk.size(1) + H = hl.specialize(wk.size(2)) + K = hl.specialize(wk.size(3)) + V = hl.specialize(v.size(3)) + chunks_per_batch = (T + CHUNK_SIZE - 1) // CHUNK_SIZE + total_chunks = chunk_indices.size(0) if is_varlen else B * chunks_per_batch + # Each CTA owns one (chunk, head) and loads a complete column tile before + # overwriting it, so these output aliases do not introduce cross-CTA races. + w = wk + u = v + akk_rows = akk.view(B * T * H, CHUNK_SIZE) + wk_rows = wk.view(B * T * H, K) + v_rows = v.view(B * T * H, V) + beta_rows = beta.view(B * T * H) + w_rows = w.view(B * T * H, K) + u_rows = u.view(B * T * H, V) + block_v = hl.register_block_size(32, V) + block_k = hl.register_block_size(32, K) + + for tile_chunk, tile_h in hl.tile( + [total_chunks, H], + block_size=[1, 1], + ): + if is_varlen: + sequence = chunk_indices[tile_chunk.id, 0].long() + local_chunk = chunk_indices[tile_chunk.id, 1].long() + begin = cu_seqlens[sequence].long() + end = cu_seqlens[sequence + 1].long() + else: + sequence = tile_chunk.id // chunks_per_batch + local_chunk = tile_chunk.id % chunks_per_batch + begin = sequence * T + end = begin + T + + lane = hl.arange(16) + chunk_begin = begin + local_chunk * CHUNK_SIZE + row0 = chunk_begin + lane + row1 = chunk_begin + 16 + lane + row2 = chunk_begin + 32 + lane + row3 = chunk_begin + 48 + lane + valid0 = row0 < end + valid1 = row1 < end + valid2 = row2 < end + valid3 = row3 < end + flat0 = row0 * H + tile_h.id + flat1 = row1 * H + tile_h.id + flat2 = row2 * H + tile_h.id + flat3 = row3 * H + tile_h.id + col0 = lane + col1 = 16 + lane + col2 = 32 + lane + col3 = 48 + lane + + m00 = hl.load( + akk_rows, + [flat0[:, None], col0[None, :]], + extra_mask=valid0[:, None] & valid0[None, :], + ).float() + m10 = hl.load( + akk_rows, + [flat1[:, None], col0[None, :]], + extra_mask=valid1[:, None] & valid0[None, :], + ).float() + m11 = hl.load( + akk_rows, + [flat1[:, None], col1[None, :]], + extra_mask=valid1[:, None] & valid1[None, :], + ).float() + m20 = hl.load( + akk_rows, + [flat2[:, None], col0[None, :]], + extra_mask=valid2[:, None] & valid0[None, :], + ).float() + m21 = hl.load( + akk_rows, + [flat2[:, None], col1[None, :]], + extra_mask=valid2[:, None] & valid1[None, :], + ).float() + m22 = hl.load( + akk_rows, + [flat2[:, None], col2[None, :]], + extra_mask=valid2[:, None] & valid2[None, :], + ).float() + m30 = hl.load( + akk_rows, + [flat3[:, None], col0[None, :]], + extra_mask=valid3[:, None] & valid0[None, :], + ).float() + m31 = hl.load( + akk_rows, + [flat3[:, None], col1[None, :]], + extra_mask=valid3[:, None] & valid1[None, :], + ).float() + m32 = hl.load( + akk_rows, + [flat3[:, None], col2[None, :]], + extra_mask=valid3[:, None] & valid2[None, :], + ).float() + m33 = hl.load( + akk_rows, + [flat3[:, None], col3[None, :]], + extra_mask=valid3[:, None] & valid3[None, :], + ).float() + + inverse = _assemble_lower_64_inverse( + (m00, m10, m11, m20, m21, m22, m30, m31, m32, m33), + wk.dtype, + ) + + beta0 = hl.load(beta_rows, [flat0], extra_mask=valid0).float() + beta1 = hl.load(beta_rows, [flat1], extra_mask=valid1).float() + beta2 = hl.load(beta_rows, [flat2], extra_mask=valid2).float() + beta3 = hl.load(beta_rows, [flat3], extra_mask=valid3).float() + for tile_v in hl.tile(V, block_size=block_v): + v0 = hl.load( + v_rows, + [flat0[:, None], tile_v.index[None, :]], + extra_mask=valid0[:, None], + ) + v1 = hl.load( + v_rows, + [flat1[:, None], tile_v.index[None, :]], + extra_mask=valid1[:, None], + ) + v2 = hl.load( + v_rows, + [flat2[:, None], tile_v.index[None, :]], + extra_mask=valid2[:, None], + ) + v3 = hl.load( + v_rows, + [flat3[:, None], tile_v.index[None, :]], + extra_mask=valid3[:, None], + ) + vb0 = (v0 * beta0[:, None]).to(v.dtype) + vb1 = (v1 * beta1[:, None]).to(v.dtype) + vb2 = (v2 * beta2[:, None]).to(v.dtype) + vb3 = (v3 * beta3[:, None]).to(v.dtype) + u0, u1, u2, u3 = _apply_lower_64_blocks( + inverse, + (vb0, vb1, vb2, vb3), + ) + hl.store( + u_rows, + [flat0[:, None], tile_v.index[None, :]], + u0, + extra_mask=valid0[:, None], + ) + hl.store( + u_rows, + [flat1[:, None], tile_v.index[None, :]], + u1, + extra_mask=valid1[:, None], + ) + hl.store( + u_rows, + [flat2[:, None], tile_v.index[None, :]], + u2, + extra_mask=valid2[:, None], + ) + hl.store( + u_rows, + [flat3[:, None], tile_v.index[None, :]], + u3, + extra_mask=valid3[:, None], + ) + + for tile_k in hl.tile(K, block_size=block_k): + wk0 = hl.load( + wk_rows, + [flat0[:, None], tile_k.index[None, :]], + extra_mask=valid0[:, None], + ) + wk1 = hl.load( + wk_rows, + [flat1[:, None], tile_k.index[None, :]], + extra_mask=valid1[:, None], + ) + wk2 = hl.load( + wk_rows, + [flat2[:, None], tile_k.index[None, :]], + extra_mask=valid2[:, None], + ) + wk3 = hl.load( + wk_rows, + [flat3[:, None], tile_k.index[None, :]], + extra_mask=valid3[:, None], + ) + w0, w1, w2, w3 = _apply_lower_64_blocks( + inverse, + (wk0, wk1, wk2, wk3), + ) + hl.store( + w_rows, + [flat0[:, None], tile_k.index[None, :]], + w0, + extra_mask=valid0[:, None], + ) + hl.store( + w_rows, + [flat1[:, None], tile_k.index[None, :]], + w1, + extra_mask=valid1[:, None], + ) + hl.store( + w_rows, + [flat2[:, None], tile_k.index[None, :]], + w2, + extra_mask=valid2[:, None], + ) + hl.store( + w_rows, + [flat3[:, None], tile_k.index[None, :]], + w3, + extra_mask=valid3[:, None], + ) + + return w, u + + +def chunk_kda_fwd_intra( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + wk: torch.Tensor, + kg: torch.Tensor, + scale: float, + cu_seqlens: torch.Tensor | None, + chunk_indices: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Helion equivalent of SGLang's intra-chunk KDA preparation.""" + is_varlen = cu_seqlens is not None + if is_varlen: + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, CHUNK_SIZE) + metadata = cu_seqlens + else: + metadata = torch.empty(0, device=q.device, dtype=torch.int32) + chunk_indices = torch.empty(0, 2, device=q.device, dtype=torch.long) + + aqk, akk = _intra_matrices_wide( + q, + k, + g, + beta, + metadata, + chunk_indices, + scale, + is_varlen, + ) + w, u = _intra_solve_recompute( + akk, + wk, + v, + beta, + metadata, + chunk_indices, + is_varlen, + ) + return w, u, kg, aqk + + +_STATE_FIXED_CONFIG = helion.Config( + block_sizes=[16], + num_warps=8, + num_stages=3, + indexing="pointer", +) + + +# Eviction policies are positional in the traced load order. Retune both +# varlen configs if the shared state body gains, loses, or reorders loads. +_STATE_VARLEN_CONFIG = helion.Config( + atomic_indexing=[], + block_sizes=[64], + indexing=[ + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "tensor_descriptor", + "pointer", + "pointer", + "pointer", + "pointer", + "tensor_descriptor", + "pointer", + "pointer", + ], + l2_groupings=[4], + load_eviction_policies=[ + "", + "", + "", + "last", + "last", + "first", + "first", + "first", + "first", + ], + loop_orders=[[0, 2, 1]], + num_stages=2, + num_warps=4, + pid_type="flat", + range_flattens=[None, None], + range_multi_buffers=[None, False], + range_num_stages=[], + range_unroll_factors=[0, 2], +) + +# Packed small-head workloads benefit from a wider V tile during state propagation. +_STATE_VARLEN_SMALL_HEAD_CONFIG = helion.Config( + atomic_indexing=[], + block_sizes=[32], + indexing=["pointer"] * 12, + l2_groupings=[1], + load_eviction_policies=["", "", "", "last", "first", "", "", "", "first"], + loop_orders=[[1, 2, 0]], + num_stages=3, + num_warps=8, + pid_type="flat", + range_flattens=[None, None], + range_multi_buffers=[None, True], + range_num_stages=[], + range_unroll_factors=[0, 0], +) + + +@helion.kernel( + static_shapes=False, + config=_STATE_FIXED_CONFIG, + ignore_warnings=_IGNORED_WARNINGS, +) +def _chunk_state( + kg: torch.Tensor, + w: torch.Tensor, + u: torch.Tensor, + chunk_decay: torch.Tensor, + initial_state: torch.Tensor, + initial_state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + chunk_indices: torch.Tensor, + chunk_offsets: torch.Tensor, + is_varlen: hl.constexpr, # pyrefly: ignore[bad-function-definition] +) -> tuple[torch.Tensor, torch.Tensor]: + """Propagate KDA state between chunks and update the state pool in place.""" + B = kg.size(0) + T = kg.size(1) + H = hl.specialize(kg.size(2)) + K = hl.specialize(kg.size(3)) + V = hl.specialize(u.size(3)) + N = cu_seqlens.size(0) - 1 if is_varlen else B + chunks_per_batch = (T + CHUNK_SIZE - 1) // CHUNK_SIZE + total_chunks = chunk_indices.size(0) if is_varlen else chunks_per_batch + hl.specialize( + ( + kg.stride(1), + kg.stride(2), + kg.stride(3), + w.stride(1), + w.stride(2), + w.stride(3), + u.stride(1), + u.stride(2), + u.stride(3), + chunk_decay.stride(0), + chunk_decay.stride(1), + chunk_decay.stride(2), + initial_state.stride(0), + initial_state.stride(1), + initial_state.stride(2), + initial_state.stride(3), + initial_state_indices.stride(0), + ) + ) + + h = torch.empty( + [B, total_chunks, H, V, K], + dtype=kg.dtype, + device=kg.device, + ) + v_new = u + kg_rows = kg.view(B * T * H, K) + w_rows = w.view(B * T * H, K) + u_rows = u.view(B * T * H, V) + decay_rows = chunk_decay.view(-1, K) + v_new_rows = v_new.view(B * T * H, V) + h_rows = h.view(B * total_chunks * H, V, K) + block_v = hl.register_block_size(1, V) + + for tile_sequence, tile_h, tile_v in hl.tile( + [N, H, V], + block_size=[1, 1, block_v], + ): + if is_varlen: + begin = cu_seqlens[tile_sequence.id].long() + end = cu_seqlens[tile_sequence.id + 1].long() + output_offset = chunk_offsets[tile_sequence.id].long() + else: + begin = tile_sequence.id * T + end = begin + T + output_offset = tile_sequence.id * chunks_per_batch + sequence_length = end - begin + state_index = initial_state_indices[tile_sequence.id].long() + state = initial_state[ + state_index, + tile_h.id, + tile_v.index, + :, + ].float() + + for token_tile in hl.tile(sequence_length, block_size=64): + global_chunk = output_offset + token_tile.id + h_rows[ + global_chunk * H + tile_h.id, + tile_v, + :, + ] = state.to(h.dtype) + token = begin + token_tile.index + valid = token < end + row = token * H + tile_h.id + w_value = hl.load( + w_rows, + [row[:, None], hl.arange(K)[None, :]], + extra_mask=valid[:, None], + ) + residual = -hl.dot( + w_value, + state.T.to(w.dtype), + out_dtype=torch.float32, + ) + residual = residual + u_rows[row, tile_v].float() + v_new_rows[row, tile_v] = residual.to(v_new.dtype) + decay = decay_rows[global_chunk * H + tile_h.id, :] + state = state * decay[None, :] + kg_value = hl.load( + kg_rows, + [row[:, None], hl.arange(K)[None, :]], + extra_mask=valid[:, None], + ) + state = state + hl.dot( + residual.T.to(kg.dtype), + kg_value, + out_dtype=torch.float32, + ) + + initial_state[ + state_index, + tile_h.id, + tile_v.index, + :, + ] = state.to(initial_state.dtype) + + return h, v_new + + +_chunk_state_varlen = helion.kernel( + static_shapes=False, + config=_STATE_VARLEN_CONFIG, + ignore_warnings=_IGNORED_WARNINGS, +)(_chunk_state.fn) +_chunk_state_varlen_small_head = helion.kernel( + static_shapes=False, + config=_STATE_VARLEN_SMALL_HEAD_CONFIG, + ignore_warnings=_IGNORED_WARNINGS, +)(_chunk_state.fn) + + +def _select_state_kernel(*, is_varlen: bool, num_heads: int) -> helion.Kernel: + if not is_varlen: + return _chunk_state + if num_heads <= _PREFILL_SMALL_HEAD_THRESHOLD: + return _chunk_state_varlen_small_head + return _chunk_state_varlen + + +_OUTPUT_CONFIG = helion.Config( + block_sizes=[128], + loop_orders=[[1, 2, 0]], + l2_groupings=[32], + num_warps=2, + num_stages=4, + indexing="pointer", +) + + +@helion.kernel( + static_shapes=False, + config=_OUTPUT_CONFIG, + ignore_warnings=_IGNORED_WARNINGS, +) +def _chunk_output( + qg: torch.Tensor, + v_new: torch.Tensor, + aqk: torch.Tensor, + h: torch.Tensor, + out: torch.Tensor, + cu_seqlens: torch.Tensor, + chunk_indices: torch.Tensor, + is_varlen: hl.constexpr, # pyrefly: ignore[bad-function-definition] +) -> torch.Tensor: + """Compose inter-chunk state output and causal intra-chunk output.""" + B = qg.size(0) + T = qg.size(1) + H = hl.specialize(qg.size(2)) + K = hl.specialize(qg.size(3)) + V = hl.specialize(v_new.size(3)) + chunks_per_batch = (T + CHUNK_SIZE - 1) // CHUNK_SIZE + total_chunks = chunk_indices.size(0) if is_varlen else B * chunks_per_batch + h_chunks = h.size(1) + hl.specialize( + ( + qg.stride(1), + qg.stride(2), + qg.stride(3), + v_new.stride(1), + v_new.stride(2), + v_new.stride(3), + aqk.stride(1), + aqk.stride(2), + aqk.stride(3), + h.stride(1), + h.stride(2), + h.stride(3), + h.stride(4), + out.stride(1), + out.stride(2), + out.stride(3), + ) + ) + + qg_rows = qg.view(B * T * H, K) + v_rows = v_new.view(B * T * H, V) + aqk_rows = aqk.view(B * T * H, CHUNK_SIZE) + h_rows = h.view(B * h_chunks * H, V, K) + out_rows = out.view(B * T * H, V) + block_v = hl.register_block_size(32, V) + + # ``out`` may alias ``v_new``. Each CTA loads its complete input tile before + # overwriting that same tile, and no other CTA reads it. + for tile_chunk, tile_h, tile_v in hl.tile( + [total_chunks, H, V], + block_size=[1, 1, block_v], + ): + if is_varlen: + sequence = chunk_indices[tile_chunk.id, 0].long() + local_chunk = chunk_indices[tile_chunk.id, 1].long() + begin = cu_seqlens[sequence].long() + end = cu_seqlens[sequence + 1].long() + h_chunk = tile_chunk.id + else: + sequence = tile_chunk.id // chunks_per_batch + local_chunk = tile_chunk.id % chunks_per_batch + begin = sequence * T + end = begin + T + h_chunk = tile_chunk.id + + lane = hl.arange(64) + token = begin + local_chunk * CHUNK_SIZE + lane + valid = token < end + row = token * H + tile_h.id + qg_value = hl.load( + qg_rows, + [row[:, None], hl.arange(K)[None, :]], + extra_mask=valid[:, None], + ) + h_value = h_rows[ + h_chunk * H + tile_h.id, + tile_v, + :, + ] + output = hl.dot( + qg_value, + h_value.T, + out_dtype=torch.float32, + ) + a_value = hl.load( + aqk_rows, + [row[:, None], lane[None, :]], + extra_mask=valid[:, None], + ) + v_value = hl.load( + v_rows, + [row, tile_v], + extra_mask=valid[:, None], + ) + output = hl.dot( + a_value.to(v_new.dtype), + v_value, + acc=output, + out_dtype=torch.float32, + ) + hl.store( + out_rows, + [row, tile_v], + output, + extra_mask=valid[:, None], + ) + + return out + + +def chunk_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + initial_state_indices: torch.Tensor | None = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + lower_bound: float | None = None, + output_intermediate_states: bool = False, + **kwargs: object, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Match the public forward contract of SGLang's Triton ``chunk_kda``.""" + if scale is None: + scale = k.shape[-1] ** -0.5 + if initial_state is None or initial_state_indices is None: + raise ValueError("KDA prefill requires an indexed initial-state pool") + + q = q.contiguous() + k = k.contiguous() + if use_qk_l2norm_in_kernel: + q, k = _l2norm_qk(q, k) + v = v.contiguous() + g = g.contiguous() + beta = beta.contiguous() + chunk_indices = ( + prepare_chunk_indices(cu_seqlens, CHUNK_SIZE) + if cu_seqlens is not None + else None + ) + g, qg, wk, kg, chunk_decay = gate_chunk_cumsum_operands( + g, + q, + k, + beta, + q_scale=scale, + a_log=A_log, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=lower_bound, + ) + w, u, kg, aqk = chunk_kda_fwd_intra( + q, + k, + v, + g, + beta, + wk, + kg, + scale, + cu_seqlens, + chunk_indices, + ) + + is_varlen = cu_seqlens is not None + if is_varlen: + chunk_offsets = prepare_chunk_offsets(cu_seqlens, CHUNK_SIZE) + metadata = cu_seqlens + else: + metadata = torch.empty(0, device=q.device, dtype=torch.int32) + chunk_offsets = torch.empty(0, device=q.device, dtype=torch.long) + state_kernel = _select_state_kernel(is_varlen=is_varlen, num_heads=q.size(2)) + h, v_new = state_kernel( + kg, + w, + u, + chunk_decay, + initial_state, + initial_state_indices, + metadata, + ( + chunk_indices + if chunk_indices is not None + else torch.empty(0, 2, device=q.device, dtype=torch.long) + ), + chunk_offsets, + is_varlen, + ) + if chunk_indices is None: + chunk_indices = torch.empty(0, 2, device=q.device, dtype=torch.long) + output = _chunk_output( + qg, + v_new, + aqk, + h, + v, + metadata, + chunk_indices, + is_varlen, + ) + if output_intermediate_states: + return output, h + return output diff --git a/python/sglang/kernels/ops/attention/helion/kda_replayssm.py b/python/sglang/kernels/ops/attention/helion/kda_replayssm.py new file mode 100644 index 000000000..3323c2c6c --- /dev/null +++ b/python/sglang/kernels/ops/attention/helion/kda_replayssm.py @@ -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 diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 6b49fae3b..89f014703 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -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}") diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index 06fd7dd78..e09984bec 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/python/sglang/srt/layers/attention/linear/kda_backend.py @@ -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( diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_helion.py b/python/sglang/srt/layers/attention/linear/kernels/kda_helion.py new file mode 100644 index 000000000..4e6bf31b7 --- /dev/null +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_helion.py @@ -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, + ) diff --git a/python/sglang/srt/layers/attention/linear/utils.py b/python/sglang/srt/layers/attention/linear/utils.py index 99bdc07f1..681f589c6 100644 --- a/python/sglang/srt/layers/attention/linear/utils.py +++ b/python/sglang/srt/layers/attention/linear/utils.py @@ -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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 0a9bbf2cb..9d444ad9d 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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 ) diff --git a/test/registered/kernels/ops/attention/test_kda_helion.py b/test/registered/kernels/ops/attention/test_kda_helion.py new file mode 100644 index 000000000..b2894cb61 --- /dev/null +++ b/test/registered/kernels/ops/attention/test_kda_helion.py @@ -0,0 +1,1026 @@ +from __future__ import annotations + +import inspect +import sys + +import pytest +import torch + +from sglang.kernels.ops.attention.fla.fused_recurrent import ( + fused_recurrent_kda_packed_decode, +) +from sglang.kernels.ops.attention.fla.fused_recurrent_linear_replayssm import ( + fused_recurrent_linear_replayssm_decode, +) +from sglang.kernels.ops.attention.fla.kda import chunk_kda as triton_chunk_kda +from sglang.test.ci.ci_register import register_cuda_ci + +try: + import helion # noqa: F401 +except ModuleNotFoundError as error: + if error.name != "helion": + raise + HELION_AVAILABLE = False +else: + HELION_AVAILABLE = True + from sglang.kernels.ops.attention.helion.kda_decode import ( + helion_fused_recurrent_kda_packed_decode, + ) + from sglang.kernels.ops.attention.helion.kda_prefill import ( + _intra_matrices_wide, + ) + from sglang.kernels.ops.attention.helion.kda_prefill import ( + chunk_kda as helion_chunk_kda, + ) + from sglang.kernels.ops.attention.helion.kda_replayssm import ( + helion_fused_recurrent_kda_replayssm_decode, + ) + +register_cuda_ci(est_time=180, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +pytestmark = pytest.mark.skipif( + not HELION_AVAILABLE, + reason="helion is not installed", +) + +_DECODE_STATE_ATOL = { + torch.float32: 1e-5, + torch.bfloat16: 2e-3, + torch.float16: 5e-4, +} + + +def test_public_signatures_match_triton() -> None: + decode_signature = inspect.signature(fused_recurrent_kda_packed_decode) + helion_decode_signature = inspect.signature( + helion_fused_recurrent_kda_packed_decode + ) + decode_parameters = list(decode_signature.parameters.values()) + helion_decode_parameters = list(helion_decode_signature.parameters.values()) + assert [parameter.name for parameter in helion_decode_parameters] == [ + parameter.name for parameter in decode_parameters + ] + assert [parameter.kind for parameter in helion_decode_parameters] == [ + parameter.kind for parameter in decode_parameters + ] + assert [parameter.default for parameter in helion_decode_parameters] == [ + parameter.default for parameter in decode_parameters + ] + + replay_parameters = list( + inspect.signature(fused_recurrent_linear_replayssm_decode).parameters.values() + ) + helion_replay_parameters = list( + inspect.signature( + helion_fused_recurrent_kda_replayssm_decode + ).parameters.values() + ) + shared_replay_parameter_count = 15 + assert ( + helion_replay_parameters[:shared_replay_parameter_count] + == replay_parameters[:shared_replay_parameter_count] + ) + assert [ + parameter.name + for parameter in helion_replay_parameters[shared_replay_parameter_count:] + ] == ["lower_bound"] + + prefill_signature = inspect.signature(triton_chunk_kda) + helion_prefill_signature = inspect.signature(helion_chunk_kda) + assert list(helion_prefill_signature.parameters) == list( + prefill_signature.parameters + ) + assert [ + parameter.kind for parameter in helion_prefill_signature.parameters.values() + ] == [parameter.kind for parameter in prefill_signature.parameters.values()] + assert [ + parameter.default for parameter in helion_prefill_signature.parameters.values() + ] == [parameter.default for parameter in prefill_signature.parameters.values()] + + +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16, torch.float16]) +def test_packed_decode_contract(state_dtype: torch.dtype) -> None: + torch.manual_seed(123) + batch, q_heads, v_heads, key_dim, value_dim = 3, 2, 4, 128, 128 + pool_size = 7 + mixed_qkv = torch.randn( + batch, + 2 * q_heads * key_dim + v_heads * value_dim, + device="cuda", + dtype=torch.bfloat16, + ) + gate = torch.randn(batch, v_heads * key_dim, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(batch, v_heads, device="cuda", dtype=torch.bfloat16) + a_log = torch.randn(v_heads, device="cuda", dtype=torch.float32) + dt_bias = torch.randn(v_heads * key_dim, device="cuda", dtype=torch.float32) + state = ( + torch.randn( + pool_size, + v_heads, + value_dim, + key_dim, + device="cuda", + dtype=state_dtype, + ) + * 0.01 + ) + indices = torch.tensor([5, -1, 2], device="cuda", dtype=torch.int32) + triton_state = state.clone() + helion_state = state.clone() + triton_out = mixed_qkv.new_empty(batch, 1, v_heads, value_dim) + helion_out = torch.empty_like(triton_out) + + fused_recurrent_kda_packed_decode( + mixed_qkv, + gate, + beta, + a_log, + dt_bias, + key_dim**-0.5, + triton_state, + triton_out, + indices, + True, + ) + result, result_state = helion_fused_recurrent_kda_packed_decode( + mixed_qkv, + gate, + beta, + a_log, + dt_bias, + key_dim**-0.5, + helion_state, + helion_out, + indices, + True, + ) + + assert result.data_ptr() == helion_out.data_ptr() + assert result_state.data_ptr() == helion_state.data_ptr() + torch.testing.assert_close(helion_out, triton_out, atol=1e-4, rtol=1e-4) + torch.testing.assert_close( + helion_state, + triton_state, + atol=_DECODE_STATE_ATOL[state_dtype], + rtol=1e-4, + ) + assert torch.count_nonzero(helion_out[1]).item() == 0 + untouched = torch.tensor([0, 1, 3, 4, 6], device="cuda") + assert torch.equal(helion_state[untouched], state[untouched]) + + +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +def test_packed_decode_lower_bound_contract(state_dtype: torch.dtype) -> None: + torch.manual_seed(321) + batch, q_heads, v_heads, key_dim, value_dim = 3, 2, 4, 128, 128 + pool_size = 7 + mixed_qkv = torch.randn( + batch, + 2 * q_heads * key_dim + v_heads * value_dim, + device="cuda", + dtype=torch.bfloat16, + ) + gate = torch.randn(batch, v_heads * key_dim, device="cuda", dtype=torch.bfloat16) + beta = torch.randn(batch, v_heads, device="cuda", dtype=torch.bfloat16) + a_log = torch.randn(v_heads, device="cuda", dtype=torch.float32) + dt_bias = torch.randn(v_heads * key_dim, device="cuda", dtype=torch.float32) + state = ( + torch.randn( + pool_size, + v_heads, + value_dim, + key_dim, + device="cuda", + dtype=state_dtype, + ) + * 0.01 + ) + indices = torch.tensor([5, -1, 2], device="cuda", dtype=torch.int32) + reference_state = state.clone() + helion_state = state.clone() + reference_out = mixed_qkv.new_zeros(batch, 1, v_heads, value_dim) + helion_out = torch.empty_like(reference_out) + scale = key_dim**-0.5 + lower_bound = -5.0 + + heads_per_q = v_heads // q_heads + q, k, v = mixed_qkv.split( + [q_heads * key_dim, q_heads * key_dim, v_heads * value_dim], dim=-1 + ) + q = q.float().view(batch, q_heads, key_dim) + k = k.float().view(batch, q_heads, key_dim) + q = q / torch.sqrt((q * q).sum(-1, keepdim=True) + 1e-6) + k = k / torch.sqrt((k * k).sum(-1, keepdim=True) + 1e-6) + q = q.repeat_interleave(heads_per_q, dim=1) + k = k.repeat_interleave(heads_per_q, dim=1) + v = v.float().view(batch, v_heads, value_dim) + raw_gate = gate.float().view(batch, v_heads, key_dim) + raw_gate = raw_gate + dt_bias.view(1, v_heads, key_dim) + A = torch.exp(a_log.float()).view(1, v_heads, 1) + decay = torch.exp(lower_bound * torch.sigmoid(A * raw_gate)) + beta_value = torch.sigmoid(beta.float()) + for batch_idx, state_idx in enumerate(indices.tolist()): + if state_idx < 0: + continue + current_state = reference_state[state_idx].float() + current_state = current_state * decay[batch_idx, :, None, :] + residual = v[batch_idx] - (current_state * k[batch_idx, :, None, :]).sum(-1) + residual = residual * beta_value[batch_idx, :, None] + current_state = current_state + residual[..., None] * k[batch_idx, :, None, :] + reference_out[batch_idx, 0] = ( + current_state * (q[batch_idx] * scale)[:, None, :] + ).sum(-1) + reference_state[state_idx] = current_state + + result, result_state = helion_fused_recurrent_kda_packed_decode( + mixed_qkv, + gate, + beta, + a_log, + dt_bias, + scale, + helion_state, + helion_out, + indices, + True, + lower_bound, + ) + + assert result.data_ptr() == helion_out.data_ptr() + assert result_state.data_ptr() == helion_state.data_ptr() + torch.testing.assert_close(helion_out, reference_out, atol=1e-4, rtol=1e-4) + torch.testing.assert_close( + helion_state, + reference_state, + atol=_DECODE_STATE_ATOL[state_dtype], + rtol=1e-4, + ) + assert torch.count_nonzero(helion_out[1]).item() == 0 + + +# `v_heads` selects the tuned config: <= KDA_SMALL_VALUE_HEAD_THRESHOLD picks the +# small-head bf16 schedule, above it the wide bf16 one. Cover both. +@pytest.mark.parametrize( + ("state_dtype", "lower_bound", "v_heads"), + [ + (torch.float32, None, 4), + (torch.bfloat16, None, 4), + (torch.bfloat16, -5.0, 4), + (torch.bfloat16, None, 16), + ], + ids=["fp32", "bf16-small-head", "bf16-small-head-lower-bound", "bf16"], +) +def test_replayssm_decode_contract( + state_dtype: torch.dtype, lower_bound: float | None, v_heads: int +) -> None: + """Match Triton ring writes, forced flushes, and natural flushes.""" + batch, q_heads, key_dim, value_dim = 3, 2, 128, 128 + cache_length, pool_size = 4, 5 + scale = key_dim**-0.5 + torch.manual_seed(721) + a_log = torch.randn(v_heads, device="cuda", dtype=torch.float32) * 0.3 + dt_bias = torch.randn(v_heads, key_dim, device="cuda", dtype=torch.float32) * 0.1 + initial = torch.randn( + pool_size, + v_heads, + value_dim, + key_dim, + device="cuda", + dtype=state_dtype, + ) + triton_state = initial.clone() + helion_state = initial.clone() + triton_d = torch.zeros( + pool_size, + v_heads, + cache_length, + value_dim, + device="cuda", + dtype=state_dtype, + ) + helion_d = triton_d.clone() + triton_k = torch.zeros( + pool_size, + q_heads, + cache_length, + key_dim, + device="cuda", + dtype=state_dtype, + ) + helion_k = triton_k.clone() + triton_g = torch.zeros( + pool_size, + v_heads, + cache_length, + key_dim, + device="cuda", + dtype=torch.float32, + ) + helion_g = triton_g.clone() + indices = torch.tensor([3, -1, 1], device="cuda", dtype=torch.int32) + write_pos = torch.zeros(batch, device="cuda", dtype=torch.int32) + + for step in range(7): + generator = torch.Generator(device="cuda").manual_seed(900 + step) + mixed_qkv = torch.randn( + batch, + 2 * q_heads * key_dim + v_heads * value_dim, + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + gate = ( + torch.randn( + batch, + v_heads, + key_dim, + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + * 0.5 + ) + beta = torch.randn( + batch, + v_heads, + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + triton_out = torch.empty( + batch, 1, v_heads, value_dim, device="cuda", dtype=torch.bfloat16 + ) + helion_out = torch.empty_like(triton_out) + force_flush = ( + torch.ones(batch, device="cuda", dtype=torch.int32) if step == 2 else None + ) + + # Triton ReplaySSM is the protocol oracle for the unbounded gate. Its + # bounded-gate path is unsupported, so that case uses the full recurrent + # update as a stronger state oracle and checks the checkpoint on flushes. + if lower_bound is None: + fused_recurrent_linear_replayssm_decode( + mixed_qkv=mixed_qkv, + a=gate, + b=beta, + A_log=a_log, + dt_bias=dt_bias, + scale=scale, + initial_state=triton_state, + d_cache=triton_d, + k_cache=triton_k, + g_cache=triton_g, + out=triton_out, + ssm_state_indices=indices, + write_pos=write_pos, + force_flush=force_flush, + use_qk_l2norm_in_kernel=True, + is_kda=True, + nk=2, + ) + else: + fused_recurrent_kda_packed_decode( + mixed_qkv=mixed_qkv, + a=gate.view(batch, -1), + b=beta, + A_log=a_log, + dt_bias=dt_bias.view(-1), + scale=scale, + initial_state=triton_state, + out=triton_out, + ssm_state_indices=indices, + use_qk_l2norm_in_kernel=True, + lower_bound=lower_bound, + ) + helion_fused_recurrent_kda_replayssm_decode( + mixed_qkv=mixed_qkv, + a=gate, + b=beta, + A_log=a_log, + dt_bias=dt_bias, + scale=scale, + initial_state=helion_state, + d_cache=helion_d, + k_cache=helion_k, + g_cache=helion_g, + out=helion_out, + ssm_state_indices=indices, + write_pos=write_pos, + force_flush=force_flush, + use_qk_l2norm_in_kernel=True, + lower_bound=lower_bound, + ) + + torch.testing.assert_close(helion_out, triton_out, atol=5e-4, rtol=1e-2) + assert torch.count_nonzero(helion_out[1]).item() == 0 + is_flush = force_flush is not None or write_pos[0].item() == cache_length - 1 + if is_flush: + torch.testing.assert_close( + helion_state.float(), + triton_state.float(), + atol=4e-3, + rtol=1e-2, + ) + write_pos.zero_() + else: + write_pos.add_(1) + + assert torch.equal(helion_state[4], initial[4]) + + +@pytest.mark.parametrize( + ("write_pos_values", "force_flush_values", "flushed_rows"), + [ + ([0, 2, 3], None, [False, False, True]), + ([1, 2, 1], [1, 0, 1], [True, False, True]), + ], + ids=["divergent-natural-flush", "mixed-forced-flush"], +) +def test_replayssm_per_row_flush_contract( + write_pos_values: list[int], + force_flush_values: list[int] | None, + flushed_rows: list[bool], +) -> None: + """Keep each row's cursor and partial-ring flush decision independent.""" + batch, q_heads, v_heads, key_dim, value_dim = 3, 2, 4, 128, 128 + cache_length = 4 + torch.manual_seed(977) + mixed_qkv = torch.randn( + batch, + 2 * q_heads * key_dim + v_heads * value_dim, + device="cuda", + dtype=torch.bfloat16, + ) + gate = torch.randn( + batch, + v_heads, + key_dim, + device="cuda", + dtype=torch.bfloat16, + ) + beta = torch.randn(batch, v_heads, device="cuda", dtype=torch.bfloat16) + a_log = torch.randn(v_heads, device="cuda", dtype=torch.float32) * 0.2 + dt_bias = torch.randn(v_heads, key_dim, device="cuda", dtype=torch.float32) * 0.1 + initial = ( + torch.randn( + batch, + v_heads, + value_dim, + key_dim, + device="cuda", + dtype=torch.float32, + ) + * 0.02 + ) + d_cache = ( + torch.randn( + batch, + v_heads, + cache_length, + value_dim, + device="cuda", + dtype=torch.float32, + ) + * 0.02 + ) + k_cache = torch.randn( + batch, + q_heads, + cache_length, + key_dim, + device="cuda", + dtype=torch.float32, + ) + g_cache = ( + -torch.rand( + batch, + v_heads, + cache_length, + key_dim, + device="cuda", + dtype=torch.float32, + ) + * 0.1 + ) + indices = torch.arange(batch, device="cuda", dtype=torch.int32) + write_pos = torch.tensor(write_pos_values, device="cuda", dtype=torch.int32) + force_flush = ( + None + if force_flush_values is None + else torch.tensor(force_flush_values, device="cuda", dtype=torch.int32) + ) + + triton_state = initial.clone() + helion_state = initial.clone() + triton_d, helion_d = d_cache.clone(), d_cache.clone() + triton_k, helion_k = k_cache.clone(), k_cache.clone() + triton_g, helion_g = g_cache.clone(), g_cache.clone() + triton_out = torch.empty( + batch, 1, v_heads, value_dim, device="cuda", dtype=torch.bfloat16 + ) + helion_out = torch.empty_like(triton_out) + + common_args = dict( + mixed_qkv=mixed_qkv, + a=gate, + b=beta, + A_log=a_log, + dt_bias=dt_bias, + scale=key_dim**-0.5, + ssm_state_indices=indices, + write_pos=write_pos, + force_flush=force_flush, + use_qk_l2norm_in_kernel=True, + ) + fused_recurrent_linear_replayssm_decode( + **common_args, + initial_state=triton_state, + d_cache=triton_d, + k_cache=triton_k, + g_cache=triton_g, + out=triton_out, + is_kda=True, + nk=2, + ) + helion_fused_recurrent_kda_replayssm_decode( + **common_args, + initial_state=helion_state, + d_cache=helion_d, + k_cache=helion_k, + g_cache=helion_g, + out=helion_out, + ) + + torch.testing.assert_close(helion_out, triton_out, atol=5e-4, rtol=1e-2) + torch.testing.assert_close(helion_state, triton_state, atol=2e-3, rtol=1e-2) + torch.testing.assert_close(helion_d, triton_d, atol=2e-3, rtol=1e-2) + torch.testing.assert_close(helion_k, triton_k, atol=2e-3, rtol=1e-2) + torch.testing.assert_close(helion_g, triton_g, atol=1e-5, rtol=1e-4) + for row, flushed in enumerate(flushed_rows): + if flushed: + assert not torch.equal(helion_state[row], initial[row]) + else: + assert torch.equal(helion_state[row], initial[row]) + + +def test_replayssm_cuda_graph_replay_with_strided_state() -> None: + """Keep cursor branches dynamic and preserve envelope-strided state I/O.""" + batch, q_heads, v_heads, key_dim, value_dim = 2, 2, 4, 128, 128 + cache_length, pool_size = 4, 3 + state_size = v_heads * value_dim * key_dim + slot_stride = state_size + 257 + storage = torch.empty(pool_size * slot_stride, device="cuda", dtype=torch.float32) + state = torch.as_strided( + storage, + (pool_size, v_heads, value_dim, key_dim), + (slot_stride, value_dim * key_dim, key_dim, 1), + ) + torch.manual_seed(811) + initial = torch.randn_like(state) + state.copy_(initial) + mixed_qkv = torch.randn( + batch, + 2 * q_heads * key_dim + v_heads * value_dim, + device="cuda", + dtype=torch.bfloat16, + ) + gate = torch.randn( + batch, + v_heads, + key_dim, + device="cuda", + dtype=torch.bfloat16, + ) + beta = torch.randn(batch, v_heads, device="cuda", dtype=torch.bfloat16) + a_log = torch.randn(v_heads, device="cuda", dtype=torch.float32) + dt_bias = torch.randn(v_heads, key_dim, device="cuda", dtype=torch.float32) + d_cache = torch.zeros( + pool_size, + v_heads, + cache_length, + value_dim, + device="cuda", + dtype=torch.float32, + ) + k_cache = torch.zeros( + pool_size, + q_heads, + cache_length, + key_dim, + device="cuda", + dtype=torch.float32, + ) + g_cache = torch.zeros( + pool_size, + v_heads, + cache_length, + key_dim, + device="cuda", + dtype=torch.float32, + ) + indices = torch.tensor([2, 0], device="cuda", dtype=torch.int32) + write_pos = torch.zeros(batch, device="cuda", dtype=torch.int32) + force_flush = torch.zeros(batch, device="cuda", dtype=torch.int32) + output = torch.empty( + batch, 1, v_heads, value_dim, device="cuda", dtype=torch.bfloat16 + ) + + def run_helion() -> None: + helion_fused_recurrent_kda_replayssm_decode( + mixed_qkv=mixed_qkv, + a=gate, + b=beta, + A_log=a_log, + dt_bias=dt_bias, + scale=key_dim**-0.5, + initial_state=state, + d_cache=d_cache, + k_cache=k_cache, + g_cache=g_cache, + out=output, + ssm_state_indices=indices, + write_pos=write_pos, + force_flush=force_flush, + use_qk_l2norm_in_kernel=True, + ) + + run_helion() + torch.cuda.synchronize() + state.copy_(initial) + d_cache.zero_() + k_cache.zero_() + g_cache.zero_() + write_pos.fill_(1) + force_flush.zero_() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run_helion() + + state.copy_(initial) + d_cache.zero_() + k_cache.zero_() + g_cache.zero_() + write_pos.fill_(1) + force_flush.copy_(torch.tensor([1, 0], device="cuda", dtype=torch.int32)) + reference_state = initial.clone() + reference_out = torch.empty_like(output) + fused_recurrent_linear_replayssm_decode( + mixed_qkv=mixed_qkv, + a=gate, + b=beta, + A_log=a_log, + dt_bias=dt_bias, + scale=key_dim**-0.5, + initial_state=reference_state, + d_cache=d_cache.clone(), + k_cache=k_cache.clone(), + g_cache=g_cache.clone(), + out=reference_out, + ssm_state_indices=indices, + write_pos=write_pos, + force_flush=force_flush, + use_qk_l2norm_in_kernel=True, + is_kda=True, + nk=2, + ) + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close(output, reference_out, atol=5e-4, rtol=1e-2) + torch.testing.assert_close(state, reference_state, atol=2e-3, rtol=1e-3) + + +def _compare_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate: torch.Tensor, + beta: torch.Tensor, + state: torch.Tensor, + indices: torch.Tensor, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + A_log: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + lower_bound: float | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch, tokens, heads, key_dim = q.shape + value_dim = v.size(-1) + if scale is None: + scale = key_dim**-0.5 + + reference_q = q.float() + reference_k = k.float() + if use_qk_l2norm_in_kernel: + reference_q = reference_q / torch.sqrt( + (reference_q * reference_q).sum(-1, keepdim=True) + 1e-6 + ) + reference_k = reference_k / torch.sqrt( + (reference_k * reference_k).sum(-1, keepdim=True) + 1e-6 + ) + reference_q = reference_q.to(q.dtype).float() + reference_k = reference_k.to(k.dtype).float() + + reference_gate = gate.float() + if A_log is not None: + if dt_bias is not None: + reference_gate = reference_gate + dt_bias.view(1, 1, heads, key_dim) + a = torch.exp(A_log.float()).view(1, 1, heads, 1) + if lower_bound is not None: + reference_gate = lower_bound * torch.sigmoid(a * reference_gate) + else: + reference_gate = -a * torch.nn.functional.softplus(reference_gate) + + reference_state = state.clone() + reference_out = torch.empty_like(v) + q_rows = reference_q.view(batch * tokens, heads, key_dim) + k_rows = reference_k.view(batch * tokens, heads, key_dim) + v_rows = v.view(batch * tokens, heads, value_dim).float() + gate_rows = reference_gate.view(batch * tokens, heads, key_dim) + beta_rows = beta.view(batch * tokens, heads).float() + out_rows = reference_out.view(batch * tokens, heads, value_dim) + + if cu_seqlens is None: + sequence_bounds = [ + (sequence * tokens, (sequence + 1) * tokens) for sequence in range(batch) + ] + chunks_per_sequence = (tokens + 63) // 64 + reference_chunks = torch.empty( + batch, + chunks_per_sequence, + heads, + value_dim, + key_dim, + device=q.device, + dtype=v.dtype, + ) + else: + offsets = cu_seqlens.tolist() + sequence_bounds = list(zip(offsets, offsets[1:])) + total_chunks = sum((end - begin + 63) // 64 for begin, end in sequence_bounds) + reference_chunks = torch.empty( + 1, + total_chunks, + heads, + value_dim, + key_dim, + device=q.device, + dtype=v.dtype, + ) + + global_chunk = 0 + for sequence, (begin, end) in enumerate(sequence_bounds): + state_index = indices[sequence].item() + current_state = reference_state[state_index].float() + for local_chunk, chunk_begin in enumerate(range(begin, end, 64)): + chunk_index = local_chunk if cu_seqlens is None else global_chunk + chunk_batch = sequence if cu_seqlens is None else 0 + reference_chunks[chunk_batch, chunk_index] = current_state.to(v.dtype) + if cu_seqlens is not None: + global_chunk += 1 + for token in range(chunk_begin, min(chunk_begin + 64, end)): + current_state = current_state * torch.exp(gate_rows[token])[:, None, :] + residual = v_rows[token] - ( + current_state * k_rows[token][:, None, :] + ).sum(-1) + residual = residual * beta_rows[token][:, None] + current_state = current_state + ( + residual[:, :, None] * k_rows[token][:, None, :] + ) + output = (current_state * (q_rows[token] * scale)[:, None, :]).sum(-1) + out_rows[token] = output.to(v.dtype) + reference_state[state_index] = current_state.to(state.dtype) + + helion_state = state.clone() + helion_v = v.clone() + helion_out, helion_chunks = helion_chunk_kda( + q, + k, + helion_v, + gate, + beta, + initial_state=helion_state, + initial_state_indices=indices, + output_intermediate_states=True, + scale=scale, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + cu_seqlens=cu_seqlens, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + ) + + assert helion_out.data_ptr() == helion_v.data_ptr() + torch.testing.assert_close(helion_out, reference_out, atol=2e-2, rtol=2e-2) + torch.testing.assert_close(helion_chunks, reference_chunks, atol=2e-2, rtol=2e-2) + torch.testing.assert_close(helion_state, reference_state, atol=2e-2, rtol=2e-2) + return helion_out, helion_chunks, helion_state + + +def test_fixed_partial_prefill_and_state_pool_contract() -> None: + torch.manual_seed(789) + batch, tokens, heads, key_dim, value_dim = 2, 17, 2, 32, 32 + q = torch.randn(batch, tokens, heads, key_dim, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn( + batch, tokens, heads, value_dim, device="cuda", dtype=torch.bfloat16 + ) + gate = torch.randn_like(q) * 0.2 + beta = torch.rand(batch, tokens, heads, device="cuda") + a_log = torch.full([heads], -2.0, device="cuda") + dt_bias = torch.zeros(heads * key_dim, device="cuda") + indices = torch.tensor([3, 1], device="cuda", dtype=torch.int32) + state = torch.randn(5, heads, value_dim, key_dim, device="cuda") * 0.01 + + _, _, helion_state = _compare_prefill( + q, + k, + v, + gate, + beta, + state, + indices, + use_qk_l2norm_in_kernel=True, + A_log=a_log, + dt_bias=dt_bias, + ) + + untouched = torch.tensor([0, 2, 4], device="cuda") + assert torch.equal(helion_state[untouched], state[untouched]) + + +def test_prefill_uses_stable_subchunk_gates() -> None: + torch.manual_seed(1117) + tokens, heads, key_dim, value_dim = 64, 1, 32, 32 + q = torch.nn.functional.normalize( + torch.randn(1, tokens, heads, key_dim, device="cuda"), dim=-1 + ).bfloat16() + k = torch.nn.functional.normalize( + torch.randn(1, tokens, heads, key_dim, device="cuda"), dim=-1 + ).bfloat16() + v = torch.randn(1, tokens, heads, value_dim, device="cuda").bfloat16() + # A chunk-global reference would form exp2(63 * 2 * RCP_LN2), which + # overflows FP32. The 16-token anchors keep every matrix factor finite. + gate = torch.full( + (1, tokens, heads, key_dim), -2.0, device="cuda", dtype=torch.float32 + ) + beta = torch.full((1, tokens, heads), 0.5, device="cuda") + cu_seqlens = torch.tensor([0, tokens], device="cuda", dtype=torch.int32) + indices = torch.zeros(1, device="cuda", dtype=torch.int32) + state = torch.zeros(1, heads, value_dim, key_dim, device="cuda") + + output, chunks, final_state = _compare_prefill( + q, + k, + v, + gate, + beta, + state, + indices, + cu_seqlens=cu_seqlens, + ) + assert torch.isfinite(output).all() + assert torch.isfinite(chunks).all() + assert torch.isfinite(final_state).all() + + +@pytest.mark.parametrize("is_varlen", [False, True]) +def test_prefill_diagonal_uses_midpoint_gate_anchor(is_varlen: bool) -> None: + """Prevent leading-edge anchoring from saturating the +/-126 gate clamp.""" + tokens, heads, key_dim = 16, 1, 32 + q = torch.full( + (1, tokens, heads, key_dim), + key_dim**-0.5, + device="cuda", + dtype=torch.bfloat16, + ) + k = q.clone() + cumulative_gate = -10.0 * torch.arange(tokens, device="cuda", dtype=torch.float32) + gate = cumulative_gate.view(1, tokens, 1, 1).expand_as(q).float() + beta = torch.ones(1, tokens, heads, device="cuda") + if is_varlen: + metadata = torch.tensor([0, tokens], device="cuda", dtype=torch.int32) + chunk_indices = torch.tensor([[0, 0]], device="cuda", dtype=torch.int32) + else: + metadata = torch.empty(0, device="cuda", dtype=torch.int32) + chunk_indices = torch.empty(0, 2, device="cuda", dtype=torch.int32) + + aqk, _ = _intra_matrices_wide( + q, + k, + gate, + beta, + metadata, + chunk_indices, + 1.0, + is_varlen=is_varlen, + ) + + qk = q[0, :, 0].float() @ k[0, :, 0].float().T + gate_delta = cumulative_gate[:, None] - cumulative_gate[None, :] + causal = ( + torch.arange(tokens, device="cuda")[:, None] + >= torch.arange(tokens, device="cuda")[None, :] + ) + expected = torch.where(causal, qk * torch.exp2(gate_delta), 0.0) + actual = aqk[0, :, 0, :tokens].float() + + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, expected, atol=5e-3, rtol=5e-3) + + +def test_fp16_preactivated_gate_with_bf16_state_contract() -> None: + torch.manual_seed(1213) + batch, tokens, heads, key_dim, value_dim = 1, 17, 1, 32, 32 + q = torch.nn.functional.normalize( + torch.randn(batch, tokens, heads, key_dim, device="cuda"), dim=-1 + ).half() + k = torch.nn.functional.normalize( + torch.randn(batch, tokens, heads, key_dim, device="cuda"), dim=-1 + ).half() + v = torch.randn(batch, tokens, heads, value_dim, device="cuda", dtype=torch.float16) + gate = -torch.rand(batch, tokens, heads, key_dim, device="cuda") * 0.01 + beta = torch.rand(batch, tokens, heads, device="cuda") + indices = torch.tensor([1], device="cuda", dtype=torch.int32) + state = ( + torch.randn( + 3, + heads, + value_dim, + key_dim, + device="cuda", + dtype=torch.bfloat16, + ) + * 0.01 + ) + + output, chunks, _ = _compare_prefill( + q, + k, + v, + gate, + beta, + state, + indices, + ) + assert output.dtype == torch.float16 + assert chunks.dtype == torch.float16 + + +@pytest.mark.parametrize( + ("state_dtype", "lower_bound"), + [ + (torch.float32, None), + (torch.bfloat16, -5.0), + (torch.float16, None), + ], + ids=["fp32", "bf16-lower-bound", "fp16"], +) +def test_packed_varlen_prefill_contract( + state_dtype: torch.dtype, lower_bound: float | None +) -> None: + torch.manual_seed(456) + lengths = [65, 31] + tokens, heads, key_dim, value_dim = sum(lengths), 2, 128, 128 + q = torch.randn(1, tokens, heads, key_dim, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(1, tokens, heads, value_dim, device="cuda", dtype=torch.bfloat16) + gate = torch.randn_like(q) + beta = torch.sigmoid( + torch.randn(1, tokens, heads, device="cuda", dtype=torch.float32) + ) + a_log = torch.randn(heads, device="cuda", dtype=torch.float32) + dt_bias = torch.randn(heads * key_dim, device="cuda", dtype=torch.float32) + cu_seqlens = torch.tensor([0, lengths[0], tokens], device="cuda", dtype=torch.int32) + indices = torch.tensor([3, 1], device="cuda", dtype=torch.int32) + state = ( + torch.randn( + 5, + heads, + value_dim, + key_dim, + device="cuda", + dtype=state_dtype, + ) + * 0.01 + ) + _compare_prefill( + q, + k, + v, + gate, + beta, + state, + indices, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + A_log=a_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"])) diff --git a/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py b/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py index 63240dac7..4f5c6314b 100644 --- a/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py +++ b/test/registered/unit/layers/attention/test_gdn_prefill_backend_policy.py @@ -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() diff --git a/test/registered/unit/layers/attention/test_kda_helion_dispatcher.py b/test/registered/unit/layers/attention/test_kda_helion_dispatcher.py new file mode 100644 index 000000000..269d26e5d --- /dev/null +++ b/test/registered/unit/layers/attention/test_kda_helion_dispatcher.py @@ -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() diff --git a/test/registered/unit/server_args/test_page_major_backend_allowlist.py b/test/registered/unit/server_args/test_page_major_backend_allowlist.py index 2550e70b1..924ad155b 100644 --- a/test/registered/unit/server_args/test_page_major_backend_allowlist.py +++ b/test/registered/unit/server_args/test_page_major_backend_allowlist.py @@ -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()