diff --git a/benchmarks/kda_prefill_autotune.py b/benchmarks/kda_prefill_autotune.py deleted file mode 100644 index 76eb45208..000000000 --- a/benchmarks/kda_prefill_autotune.py +++ /dev/null @@ -1,267 +0,0 @@ -"""Multi-shape autotuning entry point for KDA prefill helpers.""" - -from __future__ import annotations - -import argparse -from itertools import pairwise - -from examples.linear.kda_prefill import _chunk_output -from examples.linear.kda_prefill import _chunk_state -from examples.linear.kda_prefill import _chunk_state_varlen -from examples.linear.kda_prefill import _intra_matrices_wide -from examples.linear.kda_prefill import _intra_matrices_wide_forward -from examples.linear.kda_prefill import _intra_solve -from examples.linear.kda_prefill import _intra_solve_recompute -from examples.linear.kda_prefill import _intra_solve_recompute_newton -from examples.linear.kda_prefill import _recompute_u -from examples.linear.kda_prefill import _recompute_w_kg -from examples.linear.kda_prefill import prepare_chunk_indices -from examples.linear.kda_prefill import prepare_chunk_offsets -import torch - - -def _matrix_args( - sequence_length: int, - varlen: bool = False, -) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - float, - bool, -]: - batch, heads, key_dim = 1, 16, 128 - q = torch.nn.functional.normalize( - torch.randn( - batch, - sequence_length, - heads, - key_dim, - device="cuda", - ), - dim=-1, - ).to(torch.bfloat16) - k = torch.nn.functional.normalize( - torch.randn_like(q, dtype=torch.float32), - dim=-1, - ).to(torch.bfloat16) - g_step = ( - -torch.rand( - batch, - sequence_length, - heads, - key_dim, - device="cuda", - ) - * 0.1 - ) - beta = torch.rand(batch, sequence_length, heads, device="cuda") * 0.1 - if varlen: - if sequence_length == 512: - lengths = [129, 383] - elif sequence_length == 8192: - lengths = [513, 1023, 2049, 4607] - else: - first = sequence_length // 3 + 1 - lengths = [first, sequence_length - first] - cu_seqlens = torch.tensor( - [0, *torch.tensor(lengths).cumsum(0).tolist()], - device="cuda", - dtype=torch.int32, - ) - chunk_indices = prepare_chunk_indices(cu_seqlens) - else: - lengths = [sequence_length] - cu_seqlens = torch.empty(0, device="cuda", dtype=torch.int32) - chunk_indices = torch.empty(0, 2, device="cuda", dtype=torch.long) - g = torch.empty_like(g_step) - sequence_begin = 0 - for length in lengths: - for chunk_begin in range(sequence_begin, sequence_begin + length, 64): - chunk_end = min(chunk_begin + 64, sequence_begin + length) - g[:, chunk_begin:chunk_end] = ( - torch.cumsum( - g_step[:, chunk_begin:chunk_end], - dim=1, - ) - * 1.4426950216293335 - ) - sequence_begin += length - return q, k, g, beta, cu_seqlens, chunk_indices, key_dim**-0.5, varlen - - -def _kernel_args( - kernel_name: str, - sequence_length: int, - varlen: bool, - newton_schulz: bool, -) -> tuple[object, ...]: - matrix_args = _matrix_args(sequence_length, varlen) - preinvert_diagonal = not newton_schulz - if kernel_name == "matrix": - return (*matrix_args, preinvert_diagonal, newton_schulz) - - q, k, g, beta, cu_seqlens, chunk_indices, _, is_varlen = matrix_args - v = torch.randn_like(q) - matrix_kernel = ( - _intra_matrices_wide_forward if preinvert_diagonal else _intra_matrices_wide - ) - aqk, akk = matrix_kernel( - *matrix_args, - preinvert_diagonal, - newton_schulz, - ) - qg = (q.float() * (q.size(-1) ** -0.5) * torch.exp2(g)).to(q.dtype) - wk = (k.float() * beta[..., None] * torch.exp2(g)).to(k.dtype) - kg = torch.empty_like(k) - if is_varlen: - sequence_offsets = cu_seqlens.tolist() - else: - sequence_offsets = [0, q.size(1)] - chunk_decays = [] - for sequence_begin, sequence_end in pairwise(sequence_offsets): - for chunk_begin in range(sequence_begin, sequence_end, 64): - chunk_end = min(chunk_begin + 64, sequence_end) - chunk_decays.append(torch.exp2(g[0, chunk_end - 1])) - kg[:, chunk_begin:chunk_end] = ( - k[:, chunk_begin:chunk_end].float() - * torch.exp2( - g[:, chunk_end - 1 : chunk_end] - g[:, chunk_begin:chunk_end] - ) - ).to(k.dtype) - if kernel_name == "fused": - return ( - akk, - wk, - v, - beta, - cu_seqlens, - chunk_indices, - is_varlen, - newton_schulz, - preinvert_diagonal, - ) - solve_args = (akk, k, cu_seqlens, chunk_indices, is_varlen) - if kernel_name == "solve": - return solve_args - - inverse = _intra_solve(*solve_args) - if kernel_name == "u": - return v, beta, inverse, cu_seqlens, chunk_indices, is_varlen - if kernel_name == "w": - return k, g, beta, inverse, cu_seqlens, chunk_indices, is_varlen - if kernel_name in {"state", "output"}: - solve_kernel = ( - _intra_solve_recompute_newton if newton_schulz else _intra_solve_recompute - ) - w, u = solve_kernel( - akk, - wk, - v, - beta, - cu_seqlens, - chunk_indices, - is_varlen, - newton_schulz, - preinvert_diagonal, - ) - if is_varlen: - num_sequences = cu_seqlens.size(0) - 1 - chunk_offsets = prepare_chunk_offsets(cu_seqlens) - else: - num_sequences = q.size(0) - chunk_offsets = torch.empty(0, device="cuda", dtype=torch.long) - initial_state = torch.zeros( - num_sequences, - q.size(2), - v.size(3), - k.size(3), - device="cuda", - dtype=torch.float32, - ) - initial_state_indices = torch.arange( - num_sequences, - device="cuda", - dtype=torch.int32, - ) - state_args = ( - kg, - w, - u, - torch.stack(chunk_decays), - initial_state, - initial_state_indices, - cu_seqlens, - chunk_indices, - chunk_offsets, - is_varlen, - ) - if kernel_name == "state": - return state_args - state_kernel = _chunk_state_varlen if is_varlen else _chunk_state - h, v_new = state_kernel(*state_args) - return ( - qg, - v_new, - aqk, - h, - v.clone(), - cu_seqlens, - chunk_indices, - is_varlen, - ) - raise AssertionError(f"unsupported kernel: {kernel_name}") - - -def main() -> None: - parser = argparse.ArgumentParser() - parser.add_argument("--sequence-lengths", type=int, nargs="+", default=[512, 8192]) - parser.add_argument("--cache-tag", default="kda-prefill-matrix-gb200-v1") - parser.add_argument("--varlen", action="store_true") - parser.add_argument("--newton-schulz", action="store_true") - parser.add_argument( - "--kernel", - choices=["matrix", "solve", "fused", "u", "w", "state", "output"], - default="matrix", - ) - args = parser.parse_args() - - fused_kernel = ( - _intra_solve_recompute_newton if args.newton_schulz else _intra_solve_recompute - ) - state_kernel = _chunk_state_varlen if args.varlen else _chunk_state - matrix_kernel = ( - _intra_matrices_wide if args.newton_schulz else _intra_matrices_wide_forward - ) - kernels = { - "matrix": matrix_kernel, - "solve": _intra_solve, - "fused": fused_kernel, - "u": _recompute_u, - "w": _recompute_w_kg, - "state": state_kernel, - "output": _chunk_output, - } - winner = kernels[args.kernel].autotune_multi( - [ - _kernel_args( - args.kernel, - length, - args.varlen, - args.newton_schulz, - ) - for length in args.sequence_lengths - ], - aggregation="geomean", - relative_to=None, - cache_tag=args.cache_tag, - force=True, - ) - print(f"Multi-shape winner: {winner}") - - -if __name__ == "__main__": - main() diff --git a/examples/linear/kda_packed_decode.py b/examples/linear/kda_packed_decode.py deleted file mode 100644 index 1d10b0835..000000000 --- a/examples/linear/kda_packed_decode.py +++ /dev/null @@ -1,995 +0,0 @@ -"""Helion versus SGLang packed Kimi Delta Attention decode. - -This module implements the exact callable contract of SGLang's default -``fused_recurrent_kda_packed_decode`` kernel and provides a correctness and -latency comparison against that kernel loaded from an SGLang checkout. - -The production Kimi-Linear-48B-A3B shape is K=V=128 with 32 global heads. -Tensor parallelism shards those heads, so the default benchmark uses TP=2 and -16 local heads. Activations are bfloat16 while the recurrent state is float32. -This published-model default selects SGLang's in-tree Triton decode. Explicitly -using a bfloat16 state on SM100 makes SGLang auto-select external FlashInfer and -is a different baseline. - -Run from the Helion repository root: - - python -m examples.linear.kda_packed_decode - python -m examples.linear.kda_packed_decode --tp-sizes 1 2 4 8 - HELION_PRINT_OUTPUT_CODE=1 python -m examples.linear.kda_packed_decode -""" - -from __future__ import annotations - -import argparse -from dataclasses import dataclass -import importlib.util -import inspect -from pathlib import Path -import statistics -import sys -from types import ModuleType -from typing import Callable -from typing import Literal -from typing import cast - -import torch - -import helion -import helion.language as hl - -KIMI_GLOBAL_HEADS = 32 -KIMI_HEAD_K_DIM = 128 -KIMI_HEAD_V_DIM = 128 -KIMI_KDA_LAYERS = 20 -SOFTPLUS_THRESHOLD = 20.0 - - -# CUDA x traverses V tiles, keeping the x grid at 16 across all supported TP sizes. -_KDA_CONFIG = helion.Config( - block_sizes=[8], - loop_orders=[[2, 1, 0]], - num_warps=1, - num_stages=1, - indexing="pointer", - pid_type="xyz", -) - - -@helion.kernel( - static_shapes=False, - config=_KDA_CONFIG, -) -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: 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 - - gate_input = a[i_b, i_hv * K + k_offsets].float() - gate_input = gate_input + dt_bias[i_hv * K + k_offsets].float() - gate_exp = torch.exp(gate_input) - softplus = torch.where( - gate_input <= 20.0, - torch.log(1.0 + gate_exp), - gate_input, - ) - A_log_value = A_log[i_hv].float() - A = torch.exp(A_log_value) - beta = torch.sigmoid(b[i_b, i_hv].float()) - log_decay = -A * softplus - - state = initial_state[state_index, i_hv, tile_v.index, k_offsets].float() - decay = torch.exp(log_decay) - state = state * decay[None, :] - - k = mixed_qkv[i_b, k_input_offsets].float() - if use_qk_l2norm_in_kernel: - k = k / torch.sqrt((k * k).sum() + 1e-6) - 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 = q / torch.sqrt((q * q).sum() + 1e-6) - 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 - - -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, -) -> 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. - * ``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. - """ - _validate_packed_decode_inputs( - mixed_qkv, - a, - b, - A_log, - dt_bias, - initial_state, - out, - ssm_state_indices, - ) - result = _helion_fused_recurrent_kda_packed_decode( - mixed_qkv, - a, - b, - A_log, - dt_bias, - scale, - initial_state, - out, - ssm_state_indices, - use_qk_l2norm_in_kernel, - ) - return result, initial_state - - -def torch_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, -) -> tuple[torch.Tensor, torch.Tensor]: - """Independent PyTorch reference with the same mutating contract.""" - B, H, HV, K, V = _validate_packed_decode_inputs( - mixed_qkv, - a, - b, - A_log, - dt_bias, - initial_state, - out, - ssm_state_indices, - ) - - q_end = H * K - k_end = 2 * H * K - q = mixed_qkv[:, :q_end].reshape(B, H, K).float() - k = mixed_qkv[:, q_end:k_end].reshape(B, H, K).float() - v = mixed_qkv[:, k_end:].reshape(B, HV, V).float() - if HV != H: - repeat = HV // H - q = q.repeat_interleave(repeat, dim=1) - k = k.repeat_interleave(repeat, dim=1) - - if use_qk_l2norm_in_kernel: - 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 * scale - - gate_input = a.reshape(B, HV, K).float() + dt_bias.reshape(1, HV, K).float() - softplus = torch.where( - gate_input <= SOFTPLUS_THRESHOLD, - torch.log(1.0 + torch.exp(gate_input)), - gate_input, - ) - log_decay = -torch.exp(A_log.reshape(1, HV, 1).float()) * softplus - beta = torch.sigmoid(b.float()) - - valid = ssm_state_indices >= 0 - safe_indices = torch.where(valid, ssm_state_indices, 0).long() - state = initial_state.index_select(0, safe_indices).float() - state = state * torch.exp(log_decay)[:, :, None, :] - value_residual = v - (state * k[:, :, None, :]).sum(-1) - value_residual = value_residual * beta[:, :, None] - state = state + value_residual[:, :, :, None] * k[:, :, None, :] - output = (state * q[:, :, None, :]).sum(-1) - - out[:, 0] = torch.where(valid[:, None, None], output, 0.0).to(out.dtype) - if valid.any(): - initial_state[safe_indices[valid]] = state[valid].to(initial_state.dtype) - return out, initial_state - - -PackedDecode = Callable[ - [ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - float, - torch.Tensor, - torch.Tensor, - torch.Tensor, - bool, - ], - tuple[torch.Tensor, torch.Tensor], -] -PackedDecodeArgs = tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - float, - torch.Tensor, - torch.Tensor, - torch.Tensor, - bool, -] - - -@dataclass -class KDAInputs: - 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 = True - - def args(self) -> PackedDecodeArgs: - return ( - self.mixed_qkv, - self.a, - self.b, - self.A_log, - self.dt_bias, - self.scale, - self.initial_state, - self.out, - self.ssm_state_indices, - self.use_qk_l2norm_in_kernel, - ) - - def clone_mutable(self) -> KDAInputs: - return KDAInputs( - mixed_qkv=self.mixed_qkv, - a=self.a, - b=self.b, - A_log=self.A_log, - dt_bias=self.dt_bias, - scale=self.scale, - initial_state=self.initial_state.clone(), - out=torch.empty_like(self.out), - ssm_state_indices=self.ssm_state_indices, - use_qk_l2norm_in_kernel=self.use_qk_l2norm_in_kernel, - ) - - -def make_kda_inputs( - B: int, - H: int, - HV: int, - K: int, - V: int, - *, - device: torch.device | str = "cuda", - activation_dtype: torch.dtype = torch.bfloat16, - state_dtype: torch.dtype = torch.float32, - pool_size: int | None = None, - seed: int = 42, - padded: bool = False, -) -> KDAInputs: - """Create stable, production-layout packed KDA decode inputs.""" - if HV % H != 0: - raise ValueError(f"HV={HV} must be divisible by H={H}") - if pool_size is None: - pool_size = B + 16 - if pool_size < B: - raise ValueError(f"pool_size={pool_size} must be at least B={B}") - - generator = torch.Generator(device=device).manual_seed(seed) - qkv_dim = 2 * H * K + HV * V - mixed_qkv = ( - torch.randn( - B, - qkv_dim, - device=device, - dtype=activation_dtype, - generator=generator, - ) - * 0.1 - ) - a = ( - torch.randn( - B, - HV * K, - device=device, - dtype=activation_dtype, - generator=generator, - ) - * 0.5 - - 1.0 - ) - b = ( - torch.randn( - B, - HV, - device=device, - dtype=activation_dtype, - generator=generator, - ) - * 0.5 - ) - A_log = ( - torch.randn(HV, device=device, dtype=torch.float32, generator=generator) * 0.2 - ) - dt_bias = ( - torch.randn(HV * K, device=device, dtype=torch.float32, generator=generator) - * 0.1 - ) - initial_state = ( - torch.randn( - pool_size, - HV, - V, - K, - device=device, - dtype=state_dtype, - generator=generator, - ) - * 0.01 - ) - ssm_state_indices = torch.arange(B, device=device, dtype=torch.int32) - if padded: - ssm_state_indices[1::2] = -1 - out = torch.empty(B, 1, HV, V, device=device, dtype=activation_dtype) - return KDAInputs( - mixed_qkv=mixed_qkv.contiguous(), - a=a.contiguous(), - b=b.contiguous(), - A_log=A_log.contiguous(), - dt_bias=dt_bias.contiguous(), - scale=K**-0.5, - initial_state=initial_state.contiguous(), - out=out, - ssm_state_indices=ssm_state_indices, - ) - - -def _install_namespace(name: str, path: Path | None = None) -> None: - module = ModuleType(name) - module.__path__ = [] if path is None else [str(path)] # type: ignore[attr-defined] - sys.modules[name] = module - - -def _load_module(name: str, path: Path) -> ModuleType: - spec = importlib.util.spec_from_file_location(name, path) - if spec is None or spec.loader is None: - raise ImportError(f"Unable to load {name} from {path}") - module = importlib.util.module_from_spec(spec) - sys.modules[name] = module - spec.loader.exec_module(module) - return module - - -def load_sglang_packed_decode(sglang_root: Path) -> PackedDecode: - """Load the exact SGLang baseline without importing the full server package.""" - fla_dir = ( - sglang_root / "python" / "sglang" / "kernels" / "ops" / "attention" / "fla" - ) - baseline_path = fla_dir / "fused_recurrent.py" - op_path = fla_dir / "op.py" - if not baseline_path.is_file() or not op_path.is_file(): - raise FileNotFoundError( - f"Expected SGLang KDA sources under {fla_dir}; " - "pass --sglang-root explicitly." - ) - - package_paths = { - "sglang": sglang_root / "python" / "sglang", - "sglang.kernels": sglang_root / "python" / "sglang" / "kernels", - "sglang.kernels.ops": sglang_root / "python" / "sglang" / "kernels" / "ops", - "sglang.kernels.ops.attention": fla_dir.parent, - "sglang.kernels.ops.attention.fla": fla_dir, - } - for name, path in package_paths.items(): - _install_namespace(name, path) - - utils_name = "sglang.kernels.ops.attention.fla.utils" - utils = ModuleType(utils_name) - utils.input_guard = lambda fn: fn # type: ignore[attr-defined] - utils.is_gather_supported = hasattr( # type: ignore[attr-defined] - __import__("triton.language", fromlist=["gather"]), "gather" - ) - sys.modules[utils_name] = utils - - _load_module("sglang.kernels.ops.attention.fla.op", op_path) - module = _load_module( - "sglang.kernels.ops.attention.fla.fused_recurrent", baseline_path - ) - baseline = module.fused_recurrent_kda_packed_decode - if not callable(baseline): - raise TypeError(f"Unexpected baseline object: {baseline!r}") - return cast("PackedDecode", baseline) - - -def assert_matching_signatures(baseline: PackedDecode) -> None: - """Check argument names, order, kinds, and defaults against SGLang.""" - expected = inspect.signature(baseline) - actual = inspect.signature(helion_fused_recurrent_kda_packed_decode) - expected_params = list(expected.parameters.values()) - actual_params = list(actual.parameters.values()) - if len(expected_params) != len(actual_params): - raise AssertionError(f"Signature length mismatch: {actual} != {expected}") - for actual_param, expected_param in zip( - actual_params, expected_params, strict=True - ): - if ( - actual_param.name != expected_param.name - or actual_param.kind != expected_param.kind - or actual_param.default != expected_param.default - ): - raise AssertionError( - f"Signature mismatch at {actual_param.name}: {actual} != {expected}" - ) - - -def _max_abs_diff(actual: torch.Tensor, expected: torch.Tensor) -> float: - return float((actual.float() - expected.float()).abs().max()) - - -def check_correctness( - baseline: PackedDecode, - inputs: KDAInputs, - *, - atol: float = 2e-2, - rtol: float = 1e-2, -) -> tuple[float, float]: - """Check reference, output/state values, mutations, aliases, and padding.""" - original_state = inputs.initial_state.clone() - original_readonly = tuple( - tensor.clone() - for tensor in ( - inputs.mixed_qkv, - inputs.a, - inputs.b, - inputs.A_log, - inputs.dt_bias, - inputs.ssm_state_indices, - ) - ) - - reference_inputs = inputs.clone_mutable() - baseline_inputs = inputs.clone_mutable() - helion_inputs = inputs.clone_mutable() - - reference_result = torch_fused_recurrent_kda_packed_decode(*reference_inputs.args()) - baseline_result = baseline(*baseline_inputs.args()) - helion_result = helion_fused_recurrent_kda_packed_decode(*helion_inputs.args()) - torch.cuda.synchronize() - - for result, call_inputs, name in ( - (reference_result, reference_inputs, "reference"), - (baseline_result, baseline_inputs, "SGLang"), - (helion_result, helion_inputs, "Helion"), - ): - if result[0].data_ptr() != call_inputs.out.data_ptr(): - raise AssertionError(f"{name} did not return the supplied out tensor") - if result[1].data_ptr() != call_inputs.initial_state.data_ptr(): - raise AssertionError(f"{name} did not return the supplied state tensor") - - torch.testing.assert_close( - baseline_inputs.out, reference_inputs.out, atol=atol, rtol=rtol - ) - torch.testing.assert_close( - baseline_inputs.initial_state, - reference_inputs.initial_state, - atol=atol, - rtol=rtol, - ) - torch.testing.assert_close( - helion_inputs.out, baseline_inputs.out, atol=atol, rtol=rtol - ) - torch.testing.assert_close( - helion_inputs.initial_state, - baseline_inputs.initial_state, - atol=atol, - rtol=rtol, - ) - - valid_indices = inputs.ssm_state_indices[inputs.ssm_state_indices >= 0].long() - touched = torch.zeros( - original_state.shape[0], dtype=torch.bool, device=original_state.device - ) - touched[valid_indices] = True - if not torch.equal(helion_inputs.initial_state[~touched], original_state[~touched]): - raise AssertionError("Helion modified an unselected state-cache row") - invalid = inputs.ssm_state_indices < 0 - if invalid.any() and torch.count_nonzero(helion_inputs.out[invalid]) != 0: - raise AssertionError("Helion did not zero output for a padded state index") - - for current, original, name in zip( - ( - inputs.mixed_qkv, - inputs.a, - inputs.b, - inputs.A_log, - inputs.dt_bias, - inputs.ssm_state_indices, - ), - original_readonly, - ("mixed_qkv", "a", "b", "A_log", "dt_bias", "ssm_state_indices"), - strict=True, - ): - if not torch.equal(current, original): - raise AssertionError(f"Read-only input {name} was modified") - - return ( - _max_abs_diff(helion_inputs.out, baseline_inputs.out), - _max_abs_diff( - helion_inputs.initial_state[valid_indices], - baseline_inputs.initial_state[valid_indices], - ), - ) - - -def benchmark_one( - baseline: PackedDecode, - inputs: KDAInputs, - *, - warmup_ms: int, - rep_ms: int, - rounds: int, -) -> tuple[float, float]: - """Return median SGLang and Helion latency in milliseconds.""" - from triton.testing import do_bench - - baseline_inputs = inputs.clone_mutable() - helion_inputs = inputs.clone_mutable() - - def run_baseline() -> None: - baseline(*baseline_inputs.args()) - - def run_helion() -> None: - helion_fused_recurrent_kda_packed_decode(*helion_inputs.args()) - - run_baseline() - run_helion() - torch.cuda.synchronize() - samples: dict[str, list[float]] = {"SGLang": [], "Helion": []} - benchmarks = (("SGLang", run_baseline), ("Helion", run_helion)) - for round_index in range(rounds): - round_benchmarks = benchmarks if round_index % 2 == 0 else benchmarks[::-1] - for name, fn in round_benchmarks: - samples[name].append( - cast("float", do_bench(fn, warmup=warmup_ms, rep=rep_ms)) - ) - return statistics.median(samples["SGLang"]), statistics.median(samples["Helion"]) - - -def autotune_decode_shapes( - batch_sizes: list[int], - *, - local_heads: int, - activation_dtype: torch.dtype, - state_dtype: torch.dtype, - seed: int, - aggregation: Literal["geomean", "max"], -) -> helion.Config: - """Select one config using equal-weight relative latency across batch sizes.""" - arg_sets = [ - make_kda_inputs( - batch_size, - local_heads, - local_heads, - KIMI_HEAD_K_DIM, - KIMI_HEAD_V_DIM, - activation_dtype=activation_dtype, - state_dtype=state_dtype, - seed=seed + index, - ).args() - for index, batch_size in enumerate(batch_sizes) - ] - cache_tag = ( - "kda-packed-decode-v1-" - f"h{local_heads}-b{'-'.join(map(str, batch_sizes))}-" - f"{str(activation_dtype).removeprefix('torch.')}-" - f"{str(state_dtype).removeprefix('torch.')}" - ) - return _helion_fused_recurrent_kda_packed_decode.autotune_multi( - arg_sets, - aggregation=aggregation, - relative_to="default", - cache_tag=cache_tag, - force=True, - ) - - -def _dtype(name: str) -> torch.dtype: - return { - "float16": torch.float16, - "bfloat16": torch.bfloat16, - "float32": torch.float32, - }[name] - - -def _default_sglang_root() -> Path: - return Path(__file__).resolve().parents[3] / "sglang" - - -def main() -> None: - parser = argparse.ArgumentParser( - description="Compare Helion with SGLang's default packed KDA decode kernel." - ) - parser.add_argument("--sglang-root", type=Path, default=_default_sglang_root()) - parser.add_argument( - "--mode", choices=("all", "correctness", "bench"), default="all" - ) - parser.add_argument( - "--batch-sizes", type=int, nargs="+", default=[1, 4, 16, 64, 128, 256] - ) - parser.add_argument( - "--multi-autotune", - action="store_true", - help="Jointly autotune one config before running the selected mode.", - ) - parser.add_argument("--tune-batch-sizes", type=int, nargs="+", default=[1, 256]) - parser.add_argument( - "--tune-aggregation", choices=("geomean", "max"), default="geomean" - ) - parser.add_argument( - "--tp-sizes", - type=int, - nargs="+", - default=[2], - help="Tensor-parallel sizes; local heads are 32 / TP (default: 2).", - ) - parser.add_argument( - "--activation-dtype", - choices=("float16", "bfloat16", "float32"), - default="bfloat16", - ) - parser.add_argument( - "--state-dtype", - choices=("float16", "bfloat16", "float32"), - default="float32", - help="Kimi-Linear's default recurrent-state dtype is float32.", - ) - parser.add_argument("--warmup-ms", type=int, default=100) - parser.add_argument("--rep-ms", type=int, default=500) - parser.add_argument("--rounds", type=int, default=3) - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--atol", type=float, default=2e-2) - parser.add_argument("--rtol", type=float, default=1e-2) - args = parser.parse_args() - - if not torch.cuda.is_available(): - raise SystemExit("CUDA is required for the SGLang and Helion kernels.") - if args.rounds < 1: - raise SystemExit("--rounds must be at least 1.") - for tp_size in args.tp_sizes: - if KIMI_GLOBAL_HEADS % tp_size != 0: - raise SystemExit( - f"TP={tp_size} does not divide {KIMI_GLOBAL_HEADS} Kimi heads." - ) - - baseline = load_sglang_packed_decode(args.sglang_root.resolve()) - assert_matching_signatures(baseline) - - activation_dtype = _dtype(args.activation_dtype) - state_dtype = _dtype(args.state_dtype) - device = torch.device("cuda") - capability = torch.cuda.get_device_capability(device) - print( - f"Device: {torch.cuda.get_device_name(device)} " - f"(SM{capability[0]}{capability[1]})" - ) - print(f"SGLang source: {args.sglang_root.resolve()}") - print(f"Helion config: {_KDA_CONFIG}") - print( - f"Contract: activation={activation_dtype}, state={state_dtype}, " - f"K={KIMI_HEAD_K_DIM}, V={KIMI_HEAD_V_DIM}, " - "raw gate/beta logits, q/k L2 normalization" - ) - if capability[0] >= 10 and state_dtype is torch.bfloat16: - print( - "Note: this still compares the explicit Triton source; SGLang's " - "SM100 server default for a bfloat16 state is external FlashInfer." - ) - - if args.multi_autotune: - if len(args.tp_sizes) != 1: - raise SystemExit("--multi-autotune requires exactly one --tp-sizes value.") - local_heads = KIMI_GLOBAL_HEADS // args.tp_sizes[0] - print( - "Joint autotune: " - f"B={args.tune_batch_sizes}, H={local_heads}, " - f"aggregation={args.tune_aggregation}, relative_to=default" - ) - winner = autotune_decode_shapes( - args.tune_batch_sizes, - local_heads=local_heads, - activation_dtype=activation_dtype, - state_dtype=state_dtype, - seed=args.seed, - aggregation=args.tune_aggregation, - ) - print(f"Joint autotune winner: {winner}") - - if args.mode in ("all", "correctness"): - print("\nCorrectness and mutation contract") - for tp_size in args.tp_sizes: - local_heads = KIMI_GLOBAL_HEADS // tp_size - for batch_size in args.batch_sizes: - inputs = make_kda_inputs( - batch_size, - local_heads, - local_heads, - KIMI_HEAD_K_DIM, - KIMI_HEAD_V_DIM, - device=device, - activation_dtype=activation_dtype, - state_dtype=state_dtype, - seed=args.seed, - ) - output_diff, state_diff = check_correctness( - baseline, inputs, atol=args.atol, rtol=args.rtol - ) - print( - f" TP={tp_size} B={batch_size:>3} H={local_heads:>2}: " - f"PASS output_max={output_diff:.3e} " - f"state_max={state_diff:.3e}" - ) - - padded_inputs = make_kda_inputs( - 4, - local_heads, - local_heads, - KIMI_HEAD_K_DIM, - KIMI_HEAD_V_DIM, - device=device, - activation_dtype=activation_dtype, - state_dtype=state_dtype, - seed=args.seed + 1, - padded=True, - ) - check_correctness(baseline, padded_inputs, atol=args.atol, rtol=args.rtol) - print(f" TP={tp_size} padded cache indices: PASS") - - grouped_inputs = make_kda_inputs( - 4, - 8, - 16, - KIMI_HEAD_K_DIM, - KIMI_HEAD_V_DIM, - device=device, - activation_dtype=activation_dtype, - state_dtype=state_dtype, - seed=args.seed + 2, - ) - check_correctness(baseline, grouped_inputs, atol=args.atol, rtol=args.rtol) - print(" grouped heads H=8, HV=16: PASS") - - if args.mode in ("all", "bench"): - print("\nLatency (microseconds, lower is better)") - print(" TP B H | SGLang Helion speedup saved/layer saved/20 layers") - print(" " + "-" * 72) - for tp_size in args.tp_sizes: - local_heads = KIMI_GLOBAL_HEADS // tp_size - for batch_size in args.batch_sizes: - inputs = make_kda_inputs( - batch_size, - local_heads, - local_heads, - KIMI_HEAD_K_DIM, - KIMI_HEAD_V_DIM, - device=device, - activation_dtype=activation_dtype, - state_dtype=state_dtype, - seed=args.seed, - ) - baseline_ms, helion_ms = benchmark_one( - baseline, - inputs, - warmup_ms=args.warmup_ms, - rep_ms=args.rep_ms, - rounds=args.rounds, - ) - baseline_us = baseline_ms * 1000 - helion_us = helion_ms * 1000 - saved_us = baseline_us - helion_us - speedup = baseline_us / helion_us - print( - f" {tp_size:>2} {batch_size:>4} {local_heads:>3} | " - f"{baseline_us:>7.1f} {helion_us:>8.1f} " - f"{speedup:>7.2f}x {saved_us:>11.1f} " - f"{saved_us * KIMI_KDA_LAYERS:>15.1f}" - ) - - -if __name__ == "__main__": - main() diff --git a/examples/linear/kda_prefill.py b/examples/linear/kda_prefill.py deleted file mode 100644 index 0ede29e6f..000000000 --- a/examples/linear/kda_prefill.py +++ /dev/null @@ -1,2379 +0,0 @@ -"""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 torch - -import helion -import helion.language as hl - -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 - - -_CHUNK_INDICES_CACHE: list[tuple[torch.Tensor, int, torch.Tensor]] = [] -_CHUNK_OFFSETS_CACHE: list[tuple[torch.Tensor, int, torch.Tensor]] = [] - - -_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) -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=[[2, 3, 1, 0]], - num_warps=1, - num_stages=1, - indexing="pointer", -) - - -_GATE_VARLEN_CONFIG = helion.Config( - block_sizes=[16], - loop_orders=[[2, 1, 0]], - num_warps=1, - 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.exp(a_log.float()) - if use_lower_bound: - return lower_bound * torch.sigmoid(a * raw_gate) - softplus = torch.where( - raw_gate < SOFTPLUS_THRESHOLD, - torch.log(1.0 + torch.exp(raw_gate)), - raw_gate, - ) - return -a * softplus - - -@helion.kernel(static_shapes=False, config=_GATE_FIXED_CONFIG) -def _gate_cumsum_fixed( - g: torch.Tensor, - a_log: torch.Tensor, - dt_bias: torch.Tensor, - 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] -) -> torch.Tensor: - """Gate activation and chunk-local cumsum for equal-length batches.""" - B = g.size(0) - T = g.size(1) - H = hl.specialize(g.size(2)) - K = hl.specialize(g.size(3)) - hl.specialize( - ( - g.stride(1), - g.stride(2), - g.stride(3), - a_log.stride(0), - dt_bias.stride(0), - ) - ) - - out = torch.empty_like(g, dtype=torch.float32) - g_rows = g.view(B * T * H, K) - out_rows = out.view(B * T * H, K) - chunks = (T + CHUNK_SIZE - 1) // CHUNK_SIZE - block_k = hl.register_block_size(16, K) - - for tile_b, tile_chunk, tile_h, tile_k in hl.tile( - [B, chunks, H, K], - block_size=[1, 1, 1, block_k], - ): - time = hl.arange(64) - token = tile_chunk.id * CHUNK_SIZE + time - valid = token < T - row = (tile_b.id * T + 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.cumsum(value, dim=0) * scale - hl.store( - out_rows, - [row[:, None], tile_k.index[None, :]], - value, - extra_mask=valid[:, None], - ) - - return out - - -@helion.kernel(static_shapes=False, config=_GATE_VARLEN_CONFIG) -def _gate_cumsum_varlen( - g: torch.Tensor, - a_log: torch.Tensor, - dt_bias: torch.Tensor, - cu_seqlens: torch.Tensor, - chunk_indices: torch.Tensor, - 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] -) -> torch.Tensor: - """Gate activation and chunk-local cumsum for packed ragged sequences.""" - T = g.size(1) - H = hl.specialize(g.size(2)) - K = hl.specialize(g.size(3)) - chunks = chunk_indices.size(0) - hl.specialize( - ( - g.stride(1), - g.stride(2), - g.stride(3), - 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) - g_rows = g.view(T * H, K) - out_rows = out.view(T * H, K) - block_k = hl.register_block_size(16, K) - - for tile_chunk, tile_h, tile_k in hl.tile( - [chunks, H, K], - block_size=[1, 1, block_k], - ): - time = hl.arange(64) - 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() - 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.cumsum(value, dim=0) * scale - hl.store( - out_rows, - [row[:, None], tile_k.index[None, :]], - value, - extra_mask=valid[:, None], - ) - - return out - - -@helion.kernel(static_shapes=False, config=_GATE_FIXED_CONFIG) -def _gate_cumsum_operands_fixed( - g: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - beta: torch.Tensor, - a_log: torch.Tensor, - dt_bias: 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] -) -> 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 = (T + CHUNK_SIZE - 1) // CHUNK_SIZE - 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), - ) - ) - - 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( - [B * 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_b, tile_chunk, tile_h, tile_k in hl.tile( - [B, chunks, H, K], - block_size=[1, 1, 1, block_k], - ): - time = hl.arange(64) - token = tile_chunk.id * CHUNK_SIZE + time - valid = token < T - row = (tile_b.id * T + 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) - global_chunk = tile_b.id * chunks + tile_chunk.id - 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, - [global_chunk, tile_h.id, tile_k.index], - torch.exp2(last_gate), - ) - - return out, qg, wk, kg, chunk_decay - - -@helion.kernel(static_shapes=False, config=_GATE_VARLEN_CONFIG) -def _gate_cumsum_operands_varlen( - 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] -) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, -]: - """Packed cumulative gates with rounded Q/K operands.""" - T = g.size(1) - H = hl.specialize(g.size(2)) - K = hl.specialize(g.size(3)) - chunks = chunk_indices.size(0) - 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( - [chunks, H, K], - dtype=torch.float32, - device=g.device, - ) - g_rows = g.view(T * H, K) - q_rows = q.view(T * H, K) - k_rows = k.view(T * H, K) - beta_rows = beta.view(T * H) - out_rows = out.view(T * H, K) - qg_rows = qg.view(T * H, K) - wk_rows = wk.view(T * H, K) - kg_rows = kg.view(T * H, K) - block_k = hl.register_block_size(16, K) - - for tile_chunk, tile_h, tile_k in hl.tile( - [chunks, H, K], - block_size=[1, 1, block_k], - ): - time = hl.arange(64) - 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() - 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 - - -def prepare_chunk_indices( - cu_seqlens: torch.Tensor, - chunk_size: int = CHUNK_SIZE, -) -> torch.Tensor: - """Build the same ``(sequence, local_chunk)`` map used by SGLang.""" - for index, (cached_cu_seqlens, cached_chunk_size, result) in enumerate( - _CHUNK_INDICES_CACHE - ): - if cu_seqlens is cached_cu_seqlens and chunk_size == cached_chunk_size: - _CHUNK_INDICES_CACHE.append(_CHUNK_INDICES_CACHE.pop(index)) - return result - - lengths = cu_seqlens[1:] - cu_seqlens[:-1] - chunk_counts = torch.div( - lengths + chunk_size - 1, - chunk_size, - rounding_mode="floor", - ) - local_chunks = torch.cat( - [ - torch.arange(count, device=cu_seqlens.device, dtype=cu_seqlens.dtype) - for count in chunk_counts.tolist() - ] - ) - result = torch.stack( - [local_chunks.eq(0).cumsum(0) - 1, local_chunks], - dim=1, - ) - _CHUNK_INDICES_CACHE.append((cu_seqlens, chunk_size, result)) - if len(_CHUNK_INDICES_CACHE) > 4: - _CHUNK_INDICES_CACHE.pop(0) - return result - - -def prepare_chunk_offsets( - cu_seqlens: torch.Tensor, - chunk_size: int = CHUNK_SIZE, -) -> torch.Tensor: - """Return the packed output offset for each ragged sequence.""" - for index, (cached_cu_seqlens, cached_chunk_size, result) in enumerate( - _CHUNK_OFFSETS_CACHE - ): - if cu_seqlens is cached_cu_seqlens and chunk_size == cached_chunk_size: - _CHUNK_OFFSETS_CACHE.append(_CHUNK_OFFSETS_CACHE.pop(index)) - return result - - lengths = cu_seqlens[1:] - cu_seqlens[:-1] - chunk_counts = torch.div( - lengths + chunk_size - 1, - chunk_size, - rounding_mode="floor", - ) - result = torch.cat([cu_seqlens.new_zeros(1), chunk_counts]).cumsum(0) - _CHUNK_OFFSETS_CACHE.append((cu_seqlens, chunk_size, result)) - if len(_CHUNK_OFFSETS_CACHE) > 4: - _CHUNK_OFFSETS_CACHE.pop(0) - return result - - -def gate_chunk_cumsum( - g: torch.Tensor, - *, - 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, - scale: float = RCP_LN2, -) -> torch.Tensor: - """Apply KDA gate preprocessing with SGLang-compatible option semantics.""" - 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 - - if cu_seqlens is None: - return _gate_cumsum_fixed( - g, - flat_a_log, - flat_bias, - scale, - lower_bound_value, - activate, - has_bias, - use_lower_bound, - ) - - 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) - return _gate_cumsum_varlen( - g, - flat_a_log, - flat_bias, - cu_seqlens, - chunk_indices, - scale, - lower_bound_value, - activate, - has_bias, - use_lower_bound, - ) - - -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 - - if cu_seqlens is None: - return _gate_cumsum_operands_fixed( - g, - q, - k, - beta, - flat_a_log, - flat_bias, - gate_scale, - q_scale, - lower_bound_value, - activate, - has_bias, - use_lower_bound, - ) - - 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) - return _gate_cumsum_operands_varlen( - g, - q, - k, - beta, - flat_a_log, - flat_bias, - cu_seqlens, - chunk_indices, - gate_scale, - q_scale, - lower_bound_value, - activate, - has_bias, - use_lower_bound, - ) - - -_INTRA_MATRIX_CONFIG = helion.Config( - block_sizes=[32], - loop_orders=[[2, 1, 0]], - num_warps=1, - num_stages=2, - indexing="pointer", -) - - -_INTRA_MATRIX_FORWARD_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) -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] - preinvert_diagonal: hl.constexpr = False, # pyrefly: ignore[bad-function-definition] - newton_schulz: hl.constexpr = False, # 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 - 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, # pyrefly: ignore[bad-argument-type] - ).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, - ) - q_off = (q_row * torch.exp2(g_row - g_anchor[None, :])).to( - torch.bfloat16 - ) - k_off = (k_row * torch.exp2(g_row - g_anchor[None, :])).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_anchor[None, :], - -126.0, - 126.0, - ) - q_diag = q_row * torch.exp2(diag_delta) - k_diag_fwd = k_row * torch.exp2(diag_delta) - k_diag_bwd = k_row * torch.exp2(-diag_delta) - 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, - ) - if preinvert_diagonal: - diagonal_matrix = _invert_lower_16( - diagonal_matrix, - newton_schulz, - ) - 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 - - -_intra_matrices_wide_forward = helion.kernel( - static_shapes=False, - config=_INTRA_MATRIX_FORWARD_CONFIG, -)(_intra_matrices_wide.fn) - - -@helion.kernel(static_shapes=False, config=_INTRA_MATRIX_CONFIG) -def _intra_matrices( - 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 causal QK and beta-scaled KK blocks for each KDA chunk.""" - 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.zeros([B, T, H, CHUNK_SIZE], dtype=q.dtype, device=q.device) - akk = torch.zeros([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, tile_col_block in hl.tile( - [total_chunks, H, 4, 4], - block_size=[1, 1, 1, 1], - ): - if tile_col_block.id <= tile_row_block.id: - 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 - row_token = chunk_begin + tile_row_block.id * 16 + lane - col_token = chunk_begin + tile_col_block.id * 16 + lane - row_valid = row_token < end - col_valid = col_token < end - row = row_token * H + tile_h.id - col = col_token * H + tile_h.id - anchor_row = chunk_begin + tile_row_block.id * 16 - anchor = anchor_row * H + tile_h.id - aqk_value = hl.zeros([16, 16], dtype=torch.float32) - akk_value = 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() - k_col = hl.load( - k_rows, - [col[:, None], tile_k.index[None, :]], - extra_mask=col_valid[:, None], - ).float() - g_col = hl.load( - g_rows, - [col[:, None], tile_k.index[None, :]], - extra_mask=col_valid[:, None], - ).float() - g_anchor = hl.load( - g_rows, - [anchor, tile_k.index], - extra_mask=anchor_row < end, # pyrefly: ignore[bad-argument-type] - ).float() - - if tile_row_block.id == tile_col_block.id: - delta_row = torch.clamp( - g_row - g_anchor[None, :], - -126.0, - 126.0, - ) - delta_col = torch.clamp( - g_anchor[None, :] - g_col, - -126.0, - 126.0, - ) - q_scaled = (q_row * torch.exp2(delta_row)).to(torch.bfloat16) - k_scaled = (k_row * torch.exp2(delta_row)).to(torch.bfloat16) - k_col_scaled = (k_col * torch.exp2(delta_col)).to(torch.bfloat16) - else: - q_scaled = (q_row * torch.exp2(g_row - g_anchor[None, :])).to( - torch.bfloat16 - ) - k_scaled = (k_row * torch.exp2(g_row - g_anchor[None, :])).to( - torch.bfloat16 - ) - k_col_scaled = (k_col * torch.exp2(g_anchor[None, :] - g_col)).to( - torch.bfloat16 - ) - - aqk_value = hl.dot( - q_scaled, - k_col_scaled.T, - acc=aqk_value, - out_dtype=torch.float32, - ) - akk_value = hl.dot( - k_scaled, - k_col_scaled.T, - acc=akk_value, - out_dtype=torch.float32, - ) - - row_local = tile_row_block.id * 16 + lane - col_local = tile_col_block.id * 16 + lane - aqk_mask = row_valid[:, None] & col_valid[None, :] - akk_mask = row_valid[:, None] & col_valid[None, :] - aqk_causal = row_local[:, None] >= col_local[None, :] - akk_causal = row_local[:, None] > col_local[None, :] - row_beta = hl.load( - beta_rows, - [row], - extra_mask=row_valid, - ).float() - hl.store( - aqk_rows, - [row[:, None], col_local[None, :]], - torch.where(aqk_causal, aqk_value * scale, 0.0), - extra_mask=aqk_mask, - ) - hl.store( - akk_rows, - [row[:, None], col_local[None, :]], - torch.where(akk_causal, akk_value * row_beta[:, None], 0.0), - extra_mask=akk_mask, - ) - - 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 _invert_lower_16_newton_schulz(matrix: torch.Tensor) -> torch.Tensor: - lane = hl.arange(16) - strictly_lower = lane[:, None] > lane[None, :] - diagonal = lane[:, None] == lane[None, :] - lower = torch.where(strictly_lower, matrix, 0.0) - system = (lower + diagonal.float()).to(torch.bfloat16) - inverse = diagonal.float() - lower - for _ in range(3): - inverse_bf16 = inverse.to(torch.bfloat16) - correction = hl.dot(system, inverse_bf16, out_dtype=torch.float32) - inverse = 2.0 * inverse_bf16.float() - hl.dot( - inverse_bf16, - correction.to(torch.bfloat16), - out_dtype=torch.float32, - ) - inverse = torch.where(strictly_lower | diagonal, inverse, 0.0) - return inverse - - -def _invert_lower_16( - matrix: torch.Tensor, - newton_schulz: hl.constexpr, -) -> torch.Tensor: - if newton_schulz: - return _invert_lower_16_newton_schulz(matrix) - return _invert_lower_16_forward_substitution(matrix) - - -_INTRA_SOLVE_CONFIG = helion.Config( - loop_orders=[[1, 0]], - num_warps=2, - num_stages=2, - indexing="pointer", -) - - -@helion.kernel(static_shapes=False, config=_INTRA_SOLVE_CONFIG) -def _intra_solve( - akk: torch.Tensor, - output_template: torch.Tensor, - cu_seqlens: torch.Tensor, - chunk_indices: torch.Tensor, - is_varlen: hl.constexpr, # pyrefly: ignore[bad-function-definition] -) -> torch.Tensor: - """Invert the 64x64 unit-lower KDA system as four 16x16 blocks.""" - B = akk.size(0) - T = akk.size(1) - H = hl.specialize(akk.size(2)) - 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((akk.stride(1), akk.stride(2), akk.stride(3))) - - inverse = torch.empty( - [B, T, H, CHUNK_SIZE], - dtype=output_template.dtype, - device=akk.device, - ) - akk_rows = akk.view(B * T * H, CHUNK_SIZE) - inverse_rows = inverse.view(B * T * H, CHUNK_SIZE) - - 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() - - i00 = _invert_lower_16_forward_substitution(m00) - i11 = _invert_lower_16_forward_substitution(m11) - i22 = _invert_lower_16_forward_substitution(m22) - i33 = _invert_lower_16_forward_substitution(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, - ) - - hl.store( - inverse_rows, - [flat0[:, None], col0[None, :]], - i00, - extra_mask=valid0[:, None] & valid0[None, :], - ) - hl.store( - inverse_rows, - [flat1[:, None], col0[None, :]], - i10, - extra_mask=valid1[:, None] & valid0[None, :], - ) - hl.store( - inverse_rows, - [flat1[:, None], col1[None, :]], - i11, - extra_mask=valid1[:, None] & valid1[None, :], - ) - hl.store( - inverse_rows, - [flat2[:, None], col0[None, :]], - i20, - extra_mask=valid2[:, None] & valid0[None, :], - ) - hl.store( - inverse_rows, - [flat2[:, None], col1[None, :]], - i21, - extra_mask=valid2[:, None] & valid1[None, :], - ) - hl.store( - inverse_rows, - [flat2[:, None], col2[None, :]], - i22, - extra_mask=valid2[:, None] & valid2[None, :], - ) - hl.store( - inverse_rows, - [flat3[:, None], col0[None, :]], - i30, - extra_mask=valid3[:, None] & valid0[None, :], - ) - hl.store( - inverse_rows, - [flat3[:, None], col1[None, :]], - i31, - extra_mask=valid3[:, None] & valid1[None, :], - ) - hl.store( - inverse_rows, - [flat3[:, None], col2[None, :]], - i32, - extra_mask=valid3[:, None] & valid2[None, :], - ) - hl.store( - inverse_rows, - [flat3[:, None], col3[None, :]], - i33, - extra_mask=valid3[:, None] & valid3[None, :], - ) - - return inverse - - -_RECOMPUTE_U_CONFIG = helion.Config( - block_sizes=[128], - loop_orders=[[0, 1, 2, 3]], - static_ranges=[True], - num_warps=2, - num_stages=2, - indexing="pointer", -) - - -@helion.kernel(static_shapes=False, config=_RECOMPUTE_U_CONFIG) -def _recompute_u( - v: torch.Tensor, - beta: torch.Tensor, - inverse: torch.Tensor, - cu_seqlens: torch.Tensor, - chunk_indices: torch.Tensor, - is_varlen: hl.constexpr, # pyrefly: ignore[bad-function-definition] -) -> torch.Tensor: - B = v.size(0) - T = v.size(1) - H = hl.specialize(v.size(2)) - 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 - u = torch.empty_like(v) - v_rows = v.view(B * T * H, V) - beta_rows = beta.view(B * T * H) - inverse_rows = inverse.view(B * T * H, CHUNK_SIZE) - u_rows = u.view(B * T * H, V) - block_v = hl.register_block_size(32, V) - - for tile_chunk, tile_h, tile_row_block, tile_v in hl.tile( - [total_chunks, H, 4, V], - block_size=[1, 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() - 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 - row_token = chunk_begin + tile_row_block.id * 16 + lane - row_valid = row_token < end - row = row_token * H + tile_h.id - value = hl.zeros([16, tile_v], dtype=torch.float32) - for source_block in range(4): - if source_block <= tile_row_block.id: - source_token = chunk_begin + source_block * 16 + lane - source_valid = source_token < end - source = source_token * H + tile_h.id - inv = hl.load( - inverse_rows, - [row[:, None], (source_block * 16 + lane)[None, :]], - extra_mask=row_valid[:, None] & source_valid[None, :], - ) - source_v = hl.load( - v_rows, - [source[:, None], tile_v.index[None, :]], - extra_mask=source_valid[:, None], - ) - source_beta = hl.load( - beta_rows, - [source], - extra_mask=source_valid, - ) - value = hl.dot( - inv, - (source_v * source_beta[:, None]).to(v.dtype), - acc=value, - out_dtype=torch.float32, - ) - hl.store( - u_rows, - [row[:, None], tile_v.index[None, :]], - value, - extra_mask=row_valid[:, None], - ) - return u - - -_RECOMPUTE_W_CONFIG = helion.Config( - block_sizes=[128], - loop_orders=[[1, 2, 0, 3]], - l2_groupings=[4], - static_ranges=[True], - num_warps=4, - num_stages=2, - indexing="pointer", -) - - -@helion.kernel(static_shapes=False, config=_RECOMPUTE_W_CONFIG) -def _recompute_w_kg( - k: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - inverse: torch.Tensor, - cu_seqlens: torch.Tensor, - chunk_indices: torch.Tensor, - is_varlen: hl.constexpr, # pyrefly: ignore[bad-function-definition] -) -> tuple[torch.Tensor, torch.Tensor]: - B = k.size(0) - T = k.size(1) - H = hl.specialize(k.size(2)) - K = hl.specialize(k.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 - w = torch.empty_like(k) - kg = torch.empty_like(k) - k_rows = k.view(B * T * H, K) - g_rows = g.view(B * T * H, K) - beta_rows = beta.view(B * T * H) - inverse_rows = inverse.view(B * T * H, CHUNK_SIZE) - w_rows = w.view(B * T * H, K) - kg_rows = kg.view(B * T * H, K) - block_k = hl.register_block_size(32, K) - - for tile_chunk, tile_h, tile_row_block, tile_k in hl.tile( - [total_chunks, H, 4, K], - block_size=[1, 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 - - lane = hl.arange(16) - chunk_begin = begin + local_chunk * CHUNK_SIZE - row_token = chunk_begin + tile_row_block.id * 16 + lane - row_valid = row_token < end - row = row_token * H + tile_h.id - value = hl.zeros([16, tile_k], dtype=torch.float32) - for source_block in range(4): - if source_block <= tile_row_block.id: - source_token = chunk_begin + source_block * 16 + lane - source_valid = source_token < end - source = source_token * H + tile_h.id - inv = hl.load( - inverse_rows, - [row[:, None], (source_block * 16 + lane)[None, :]], - extra_mask=row_valid[:, None] & source_valid[None, :], - ) - source_k = hl.load( - k_rows, - [source[:, None], tile_k.index[None, :]], - extra_mask=source_valid[:, None], - ) - source_g = hl.load( - g_rows, - [source[:, None], tile_k.index[None, :]], - extra_mask=source_valid[:, None], - ).float() - source_beta = hl.load( - beta_rows, - [source], - extra_mask=source_valid, - ).float() - weighted_k = ( - source_k.float() * source_beta[:, None] * torch.exp2(source_g) - ).to(k.dtype) - value = hl.dot( - inv, - weighted_k, - acc=value, - out_dtype=torch.float32, - ) - - row_k = hl.load( - k_rows, - [row[:, None], tile_k.index[None, :]], - extra_mask=row_valid[:, None], - ).float() - row_g = hl.load( - g_rows, - [row[:, None], tile_k.index[None, :]], - extra_mask=row_valid[:, None], - ).float() - if is_varlen: - chunk_end = ( - end.new_full( # pyrefly: ignore[missing-attribute] - [], CHUNK_SIZE - ) - + chunk_begin - ) - last_token = ( - torch.minimum( - chunk_end, - end, # pyrefly: ignore[bad-argument-type] - ) - - 1 - ) - else: - last_token = min(chunk_begin + CHUNK_SIZE, end) - 1 - last = last_token * H + tile_h.id - last_g = g_rows[last, tile_k.index].float() - kg_value = row_k * torch.exp2(last_g[None, :] - row_g) - hl.store( - w_rows, - [row[:, None], tile_k.index[None, :]], - value, - extra_mask=row_valid[:, None], - ) - hl.store( - kg_rows, - [row[:, None], tile_k.index[None, :]], - kg_value, - extra_mask=row_valid[:, None], - ) - return w, kg - - -_FORWARD_SOLVE_RECOMPUTE_CONFIG = helion.Config( - block_sizes=[64, 64], - loop_orders=[[0, 1]], - num_warps=1, - num_stages=3, - indexing="pointer", -) - - -_NEWTON_SOLVE_RECOMPUTE_CONFIG = helion.Config( - block_sizes=[64, 64], - loop_orders=[[0, 1]], - num_warps=2, - num_stages=3, - indexing="pointer", -) - - -@helion.kernel(static_shapes=False, config=_FORWARD_SOLVE_RECOMPUTE_CONFIG) -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] - newton_schulz: hl.constexpr = False, # pyrefly: ignore[bad-function-definition] - diagonal_preinverted: hl.constexpr = False, # pyrefly: ignore[bad-function-definition] -) -> tuple[torch.Tensor, torch.Tensor]: - """Solve the 64x64 system and emit W and U from pre-scaled operands. - - Forward substitution preserves the Triton baseline's floating-point order. - Newton-Schulz is an explicitly selected, algebraically equivalent fast path. - """ - 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 - 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() - - if diagonal_preinverted: - i00 = m00 - i11 = m11 - i22 = m22 - i33 = m33 - else: - i00 = _invert_lower_16(m00, newton_schulz) - i11 = _invert_lower_16(m11, newton_schulz) - i22 = _invert_lower_16(m22, newton_schulz) - i33 = _invert_lower_16(m33, newton_schulz) - 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, - ) - i00 = i00.to(wk.dtype) - i10 = i10.to(wk.dtype) - i11 = i11.to(wk.dtype) - i20 = i20.to(wk.dtype) - i21 = i21.to(wk.dtype) - i22 = i22.to(wk.dtype) - i30 = i30.to(wk.dtype) - i31 = i31.to(wk.dtype) - i32 = i32.to(wk.dtype) - i33 = i33.to(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 = hl.dot(i00, vb0, out_dtype=torch.float32) - u1 = hl.dot(i10, vb0, out_dtype=torch.float32) + hl.dot( - i11, vb1, out_dtype=torch.float32 - ) - u2 = ( - hl.dot(i20, vb0, out_dtype=torch.float32) - + hl.dot(i21, vb1, out_dtype=torch.float32) - + hl.dot(i22, vb2, out_dtype=torch.float32) - ) - u3 = ( - hl.dot(i30, vb0, out_dtype=torch.float32) - + hl.dot(i31, vb1, out_dtype=torch.float32) - + hl.dot(i32, vb2, out_dtype=torch.float32) - + hl.dot(i33, vb3, out_dtype=torch.float32) - ) - 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 = hl.dot(i00, wk0, out_dtype=torch.float32) - w1 = hl.dot(i10, wk0, out_dtype=torch.float32) + hl.dot( - i11, wk1, out_dtype=torch.float32 - ) - w2 = ( - hl.dot(i20, wk0, out_dtype=torch.float32) - + hl.dot(i21, wk1, out_dtype=torch.float32) - + hl.dot(i22, wk2, out_dtype=torch.float32) - ) - w3 = ( - hl.dot(i30, wk0, out_dtype=torch.float32) - + hl.dot(i31, wk1, out_dtype=torch.float32) - + hl.dot(i32, wk2, out_dtype=torch.float32) - + hl.dot(i33, wk3, out_dtype=torch.float32) - ) - 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 - - -_intra_solve_recompute_newton = helion.kernel( - static_shapes=False, - config=_NEWTON_SOLVE_RECOMPUTE_CONFIG, -)(_intra_solve_recompute.fn) - - -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, - newton_schulz: bool = False, -) -> 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) - 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) - - preinvert_diagonal = not newton_schulz - matrix_kernel = ( - _intra_matrices_wide_forward if preinvert_diagonal else _intra_matrices_wide - ) - aqk, akk = matrix_kernel( - q, - k, - g, - beta, - metadata, - chunk_indices, - scale, - is_varlen, - preinvert_diagonal, - newton_schulz, - ) - solve_kernel = ( - _intra_solve_recompute_newton if newton_schulz else _intra_solve_recompute - ) - w, u = solve_kernel( - akk, - wk, - v, - beta, - metadata, - chunk_indices, - is_varlen, - newton_schulz, - preinvert_diagonal, - ) - return w, u, kg, aqk - - -_STATE_FIXED_CONFIG = helion.Config( - block_sizes=[16], - num_warps=8, - num_stages=3, - indexing="pointer", -) - - -_STATE_VARLEN_CONFIG = helion.Config( - block_sizes=[16], - loop_orders=[[1, 2, 0]], - num_warps=4, - num_stages=3, - indexing="pointer", -) - - -@helion.kernel(static_shapes=False, config=_STATE_FIXED_CONFIG) -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, -)(_chunk_state.fn) - - -_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) -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) - - 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``. - - ``newton_schulz=True`` in ``kwargs`` enables a faster inverse with a - different floating-point operation order. The default uses Triton's - forward-substitution order. - """ - 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") - newton_schulz = bool(kwargs.get("newton_schulz")) - - 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) 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, - newton_schulz, - ) - - is_varlen = cu_seqlens is not None - if is_varlen: - chunk_offsets = prepare_chunk_offsets(cu_seqlens) - 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 = _chunk_state_varlen if is_varlen else _chunk_state - 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 - - -def main() -> None: - """Compile the front-end kernels for the production Kimi-Linear shape.""" - device = torch.device("cuda") - q = torch.randn(1, 512, 16, 128, device=device, dtype=torch.bfloat16) - k = torch.randn_like(q) - g = torch.randn_like(q) - a_log = torch.randn(16, device=device) - dt_bias = torch.randn(16 * 128, device=device) - _l2norm_qk(q, k) - gate_chunk_cumsum( - g, - a_log=a_log, - dt_bias=dt_bias, - cu_seqlens=None, - ) - - -if __name__ == "__main__": - main() diff --git a/examples/linear/kda_prefill_sglang_benchmarking.md b/examples/linear/kda_prefill_sglang_benchmarking.md deleted file mode 100644 index 21657f749..000000000 --- a/examples/linear/kda_prefill_sglang_benchmarking.md +++ /dev/null @@ -1,594 +0,0 @@ -# KDA Prefill: Helion and SGLang Benchmarking Notes - -This note records the Kimi Delta Attention (KDA) prefill implementation, the -baseline used for comparison, and the commands used on two GB200 GPUs. The -`sitecustomize.py` integration is an A/B harness, not a proposed SGLang API. - -## Baseline and dispatch - -For `moonshotai/Kimi-Linear-48B-A3B-Instruct`, SGLang's default linear-attention -backend is `triton`. The prefill/extend call chain is: - -```text -KDAAttnBackend.forward_extend - -> KDAKernelDispatcher.extend - -> TritonKDAKernel.extend - -> sglang.kernels.ops.attention.fla.kda.chunk_kda -``` - -The dispatcher and wrapper live in: - -```text -python/sglang/srt/layers/attention/linear/kda_backend.py -python/sglang/srt/layers/attention/linear/kernels/kda_triton.py -``` - -The kernels used by the default path are in-tree under: - -```text -python/sglang/kernels/ops/attention/fla/kda.py -python/sglang/kernels/ops/attention/fla/chunk_intra.py -python/sglang/kernels/ops/attention/fla/chunk_delta_h.py -``` - -This code was copied from Flash Linear Attention and adapted through vLLM, but -it is vendored into SGLang and does not import an external FLA package at -runtime. - -SGLang also has two non-default KDA prefill implementations: - -- `cutedsl`: in-tree Blackwell kernels, selected explicitly with - `--linear-attn-prefill-backend cutedsl`. It requires SM100 and K=128, and the - current wrapper cannot return intermediate states for the default - `extra_buffer` radix-cache strategy. -- `flashkda`: an optional wrapper around the external MoonshotAI `flash_kda` - CUTLASS package. Unsupported shapes, gates, speculative extend, and requests - for intermediate states fall back to the Triton baseline. - -FlashInfer is not a KDA prefill baseline. Its SGLang wrapper implements KDA -decode and verify only and explicitly leaves extend/prefill to Triton or CuTe -DSL. - -## Production contract - -The TP=2 shape for Kimi Linear 48B is: - -```text -global heads: 32 -local heads/GPU: 16 -K: 128 -V: 128 -chunk size: 64 -activation/output: bfloat16 -recurrent state: float32 -``` - -SGLang's normal Kimi prefill/extend path concatenates each request's new tokens -and always passes `query_start_loc` as `cu_seqlens`. The default hot-path -specialization is therefore **packed varlen + forward substitution + FP32 -state**. Fixed-length input remains supported for direct callers and isolated -microbenchmarks. Newton-Schulz is opt-in and can be combined with either input -layout. - -The Helion public entry point has the same parameter order and defaults as -SGLang's `chunk_kda`: - -```python -chunk_kda( - q, - k, - v, - g, - beta, - scale=None, - initial_state=None, - initial_state_indices=None, - use_qk_l2norm_in_kernel=False, - cu_seqlens=None, - A_log=None, - dt_bias=None, - lower_bound=None, - output_intermediate_states=False, - **kwargs, -) -``` - -It preserves the corresponding numerical and mutation contract: - -- fixed-length and packed variable-length input; -- raw gates with `A_log`, optional `dt_bias`, and optional safe-gate - `lower_bound`, or already activated gates when `A_log` is absent; -- optional in-kernel Q/K L2 normalization; -- indexed FP32 or BF16 state pools in `[slot, H, V, K]` layout; -- in-place output through `v` and in-place final-state updates; -- optional intermediate states with the same fixed/packed layouts; -- partial chunks, FP16 or BF16 activations, K up to 256, and arbitrary V. - -The implementation and focused tests are: - -```text -examples/linear/kda_prefill.py -test/test_kda_prefill.py -``` - -## Implementation - -The production path launches six Helion-generated kernels: - -1. Q/K L2 normalization. -2. Gate activation and chunk-local cumulative sum. -3. Chunk-local QK and KKT matrices. -4. Fused triangular solve plus U, W, and gated-K recomputation. -5. Recurrent chunk-state update. -6. Output projection and in-place output store. - -The important fusion is step 4. Keeping solve and recomputation separate was -slower for every production sequence length tested. - -Each numerical path has one selected configuration for all observed sequence -lengths. Fixed and packed state propagation also use separate configurations -because they have materially different load patterns: - -| Stage/path | Blocks | Loop order | Warps | Stages | Indexing | -|---|---|---|---:|---:|---| -| matrices, forward substitution | 32 | `[1, 2, 0]` | 1 | 2 | pointer | -| matrices, Newton-Schulz | 32 | `[2, 1, 0]` | 1 | 2 | pointer | -| solve/recompute, forward substitution | 64, 64 | `[0, 1]` | 1 | 3 | pointer | -| solve/recompute, Newton-Schulz | 64, 64 | `[0, 1]` | 2 | 3 | pointer | -| state, fixed length | 16 | default | 8 | 3 | pointer | -| state, packed varlen | 16 | `[1, 2, 0]` | 4 | 3 | pointer | -| output | 128 | `[1, 2, 0]` | 2 | 4 | pointer | - -All kernels use `static_shapes=False`. H, K, V, inner layouts, dtypes, -fixed/varlen mode, and optional numerical modes remain specialized. Total -sequence length and every leading stride derived from it stay dynamic, matching -the Triton baseline's `do_not_specialize=["T"]` behavior. Generated launch -code takes runtime T, grid sizes, and leading strides while retaining constexpr -H/K/V. - -## Tuning findings - -Multi-shape autotuning and focused sweeps used one geometric-mean objective, -with production points at T=512 and T=8192 and validation at T=2048. - -- The prologue now emits the cumulative gate, rounded Q/K operands, gated K, - and each chunk's final decay in one pass. This removes repeated exponentials - and avoids rereading the last gate in state propagation. -- Forward substitution distributes the four independent 16x16 diagonal - inversions over the matrix CTAs. The solve consumes those pre-inverted - blocks. This was bitwise identical for the matrix, W, and U intermediates and - reduced full-pipeline latency by 13.5%, 12.3%, and 3.4% at - T=512/2048/8192 relative to doing the same inversions in the solve CTA. -- Moving Newton-Schulz inversion into the matrix stage was neutral or slower, - so that path keeps inversion in the solve CTA. Three BF16 MMA-style - refinements match the CuTe DSL structure and outperform FP32 refinements - while remaining within the recurrent-reference tolerance. -- Separate solve sweeps selected one warp for forward substitution and two for - Newton-Schulz. Both use 64x64 registered blocks and three pipeline stages. -- Fixed state propagation benefits from eight warps: 11.320, 36.134, and - 162.226 us at T=512/2048/8192, versus 11.552, 37.091, and 172.232 us with four - warps. Packed varlen instead prefers four warps and loop order `[1, 2, 0]`. -- Matrix tuning selected block 32, one warp, and two stages, with distinct loop - orders for the two inverse paths. -- Tensor descriptors were emitted for the contiguous normalization loads, but - that variant was 3-4x slower. Gathered matrix/output addresses did not lower - to descriptors. Pointer indexing is retained. -- L2 grouping was useful only for the now-unused split W helper. Eviction hints, - persistent PID, `maxnreg`, range unrolling, and range staging were neutral or - slower on the production kernels. Four range stages regressed performance. -- Warp specialization reached the Triton operation but failed in the SM100 - compiler for the matrix kernel, so it is not enabled. -- A broad random fused search found configurations with minute-long compile - outliers and a worse geometric mean. Those were rejected because compile - behavior is part of the acceptance criteria. - -The autotuning entry point is: - -```text -benchmarks/kda_prefill_autotune.py -``` - -For example: - -```bash -python benchmarks/kda_prefill_autotune.py \ - --kernel fused \ - --varlen \ - --sequence-lengths 512 8192 \ - --cache-tag kda-prefill-fused-gb200 -``` - -## Microbenchmark results - -Environment: - -```text -Helion: 7ba4dc7070252d110a03d3f4bcb575d53f5e9699 + this worktree -SGLang: d0b9689805232d8ab37789121cbc3b766b5c723e + benchmark changes -Torch: 2.11.0+cu130 -Triton: 3.6.0 -GPU: NVIDIA GB200, SM100, 152 SMs -``` - -SGLang's existing CuTe DSL prefill benchmark was extended in place to include -Helion. With BF16, H=16, K=V=128, one fixed sequence, and CUDA-graph timing: - -| T | Triton (ms) | CuTe DSL (ms) | Helion forward-substitution (ms) | Triton/Helion | Helion Newton-Schulz (ms) | Triton/Helion-NS | -|---:|---:|---:|---:|---:|---:|---:| -| 512 | 0.067 | 0.031 | 0.042 | 1.62x | 0.039 | 1.74x | -| 2048 | 0.176 | 0.099 | 0.107 | 1.64x | 0.103 | 1.68x | -| 8192 | 0.674 | 0.363 | 0.379 | 1.78x | 0.367 | 1.82x | - -The default forward-substitution path retains the Triton baseline's inverse -operation order. Newton-Schulz is algebraically equivalent and passes the same -recurrent-reference tolerances, but changes floating-point ordering. It is -therefore opt-in through the ``newton_schulz=True`` keyword and is reported as a -separate result rather than as a replacement for the default baseline. - -CuTe DSL remains faster on this fixed-shape GB200 microbenchmark. The default -benchmark command is: - -```bash -cd /path/to/sglang -python benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py \ - --mode bench \ - --dtype bfloat16 \ - --num-heads 16 \ - --seq-lens 512 2048 8192 \ - --helion-root /path/to/helion-multi-autotune -``` - -Add ``--helion-newton-schulz`` to select and clearly label the experimental -Newton-Schulz specialization. - -A paired packed-varlen benchmark using raw gates, internal Q/K normalization, -H=16, K=V=128, and ragged non-aligned lengths produced: - -| Total T | Triton (ms) | Helion forward (ms) | Triton/Helion | Helion NS (ms) | Triton/Helion-NS | -|---:|---:|---:|---:|---:|---:| -| 512 | 0.067957 | 0.052078 | 1.305x | 0.047767 | 1.423x | -| 2048 | 0.157140 | 0.126784 | 1.239x | 0.122787 | 1.280x | -| 8192 | 0.532932 | 0.438449 | 1.215x | 0.430456 | 1.238x | - -### BF16 decode tuning - -The decode kernel now has separate fixed configurations for FP32 and BF16 -recurrent state. The public parameter order, return tuple, aliases, in-place -state/output mutations, padding behavior, and numerical operations are shared -by both paths. The FP32 configuration is unchanged. - -One multi-shape search evaluated 1,752 configurations on B=1 and B=256 with an -equal-weight geometric mean of latency ratios. The selected BF16 configuration -uses a V tile of 16, one warp, four stages, flat PID, L2 grouping 16, and the -exact indexing and eviction fields recorded in -`examples/linear/kda_packed_decode.py`. It improved the raw kernel as follows: - -| B | Prior config (us) | BF16 config (us) | Ratio | -|---:|---:|---:|---:| -| 1 | 11.360 | 8.288 | 0.730x | -| 256 | 72.224 | 51.680 | 0.716x | - -The joint objective was `0.7225x`, or a `1.384x` raw-kernel speedup. The search -command was: - -```bash -cd /home/eche/local/helion-multi-autotune -python -m examples.linear.kda_packed_decode \ - --sglang-root /home/eche/local/sglang \ - --mode correctness \ - --batch-sizes 1 \ - --tp-sizes 2 \ - --activation-dtype bfloat16 \ - --state-dtype bfloat16 \ - --multi-autotune \ - --tune-batch-sizes 1 256 \ - --tune-aggregation geomean -``` - -SGLang's existing FlashInfer decode benchmark was extended with an optional -paired Helion baseline. It validates both Helion variants against packed -Triton, alternates timing order for three rounds, and reports medians. Direct -timing includes Python validation, output allocation, and dispatch, which hide -most of the raw gain at small B. The tuned kernel remained `1.68-1.83x` faster -than FlashInfer and was 8% faster than the old Helion config at B=256: - -```bash -cd /home/eche/local/sglang -python benchmark/bench_linear_attention/bench_kda_flashinfer_mtp.py \ - --mode bench \ - --task decode \ - --dtype bfloat16 \ - --batch-sizes 1 4 16 32 64 128 256 \ - --num-q-heads 16 \ - --num-v-heads 16 \ - --helion-root /home/eche/local/helion \ - --helion-baseline-root /home/eche/local/helion-multi-autotune -``` - -Generated Triton confirmed the selected block, PID, stage, L2, and eviction -choices. The requested tensor-descriptor indexing entries were downgraded to -pointer loads because these accesses are gathered through runtime state indices; -no descriptor instructions were emitted. The exact autotuned config is retained. - -A follow-up inline Triton experiment isolated only the recurrent-state access. -The pointer and TMA specializations shared launch geometry and all arithmetic; -TMA used a `[1, 1, block_v, 128]` descriptor box with the state slot loaded from -`ssm_state_indices` on device. Descriptor load/store was bitwise identical to -the pointer variant, confirming that dynamic slot coordinates are supported. -It was not profitable: - -| Block V | Warps | B=1 TMA/pointer | B=256 TMA/pointer | Geomean | -|---:|---:|---:|---:|---:| -| 16 | 1 | 1.796x | 1.433x | 1.605x | -| 32 | 1 | 1.856x | 1.052x | 1.398x | -| 32 | 2 | 2.016x | 1.273x | 1.602x | -| 64 | 8 | 1.522x | 1.255x | 1.382x | -| 128 | 8 | 1.675x | 1.638x | 1.656x | - -Issuing the transfer before the independent gate exponentials improved overlap, -but the best matched result still regressed by 26.8% at B=1 and 9.5% at B=256. -The production kernel therefore retains coalesced pointer state access. The -reproducer is `benchmarks/kda_decode_tma_probe.py`. - -The fresh TP=2 BF16 end-to-end comparison used default FlashInfer decode and -Triton prefill as the baseline. The Helion row selected Triton only as the -packed-decode integration carrier, then substituted the tuned Helion decode; -prefill remained Triton. Each server ran one 20-request warmup followed by three -100-request samples of the workload documented below. - -| Configuration | Input tok/s samples | Median | Geomean | vs baseline | -|---|---:|---:|---:|---:| -| Default FlashInfer decode | 26,952 / 25,937 / 28,201 | 26,952 | 27,014 | baseline | -| Tuned Helion decode | 27,275 / 27,874 / 27,286 | 27,286 | 27,477 | +1.24% median, +1.71% geomean | - -Every matched run generated identical text and token lengths, processed 247,525 -input and 485 output tokens, and had zero request errors. Logs, JSONL outputs, -the paired microbenchmark, and generated code are under: - -```text -/home/eche/results/kda-bf16-decode-tuning-20260724 -/home/eche/results/kda-bf16-decode-tuned-e2e-20260724 -``` - -## Dynamic-shape and startup behavior - -After one T=512 compile, previously unseen T=2048 and T=8192 calls reused the -same six bound kernels in 0.00225 and 0.00209 seconds and emitted no additional -code. A TP=2 server emitted exactly 12 output-code events, six per worker, all -before readiness and none during the serving benchmarks. - -Cold compilation remains a cost. A local first compile took about 13.35 -seconds. Across comparable TP=2 launches, Helion reached readiness roughly -10-12 seconds later than Triton. This is substantially better than specializing -and compiling for every new T, but it is not compile-time parity with the -existing Triton path. - -## End-to-end A/B result - -The serving workload used two GB200 GPUs, TP=2, FP32 recurrent state, Triton -decode, 100 deterministic random prompts derived with the SGLang benchmarker, -target input length 4096 with range ratio 0.25, output length 8, concurrency 4, -seed 4242, and a cache flush before each measured run. Each run processed -247,525 input tokens and 485 output tokens with zero errors. - -| Backend | Run 1 input tok/s | Run 2 input tok/s | Geomean | -|---|---:|---:|---:| -| Triton prefill | 26,575 | 27,597 | 27,081 | -| Helion prefill | 27,542 | 28,957 | 28,241 | - -The clean-run geomean improvement is **4.28%**. One earlier Helion sample had a -non-reproducing 4.38-second scheduler stall and is preserved in the result -directory as `helion-prefill-n100.jsonl`; there was no Helion code generation -during the stall. It is excluded from the clean-run geomean rather than hidden. - -Results and logs from this machine are under: - -```text -/home/eche/results/kda-prefill-ab-20260723 -``` - -### Path-isolated follow-up - -After tuning the forward-substitution and Newton-Schulz paths separately, a -fresh TP=2 run kept FP32 state and Triton decode fixed. Each clean row is one -100-request sample after a 20-request warmup: - -| Prefill path | Input tok/s | Duration (s) | Mean TTFT (ms) | P99 TTFT (ms) | vs Triton | -|---|---:|---:|---:|---:|---:| -| Triton | 26,874.99 | 9.21 | 218.15 | 349.65 | baseline | -| Helion forward substitution | 26,487.65 | 9.34 | 252.91 | 367.74 | -1.44% | -| Helion Newton-Schulz | 25,926.43 | 9.55 | 237.83 | 402.78 | -3.53% | - -The first forward and first Newton measured samples each hit the same isolated -1.7-second p99 TTFT scheduler stall. They are preserved as the non-`r2` files -and excluded from the table. With only one clean serving sample per path, the -small differences above are not robust enough to override the paired kernel -measurements or the earlier multi-run serving result. - -All rows processed the same 247,525 input and 485 output tokens with no request -errors. Forward substitution generated byte-identical text to Triton. -Newton-Schulz preserved every output length but changed 90 of 100 generated -texts, which is consistent with its intentionally different floating-point -order. Both Helion servers emitted 12 generated-code events before readiness -and none during measurement. The complete artifacts are under: - -```text -/home/eche/results/kda-prefill-structural-20260724 -``` - -## End-to-end reproduction - -Set paths and offline mode: - -```bash -export HELION_ROOT=/path/to/helion-multi-autotune -export SGLANG_ROOT=/path/to/sglang -export MODEL=/path/to/Kimi-Linear-48B-A3B-Instruct/snapshot -export DATASET=/path/to/ShareGPT_V3_unfiltered_cleaned_split.json -export RESULTS=/path/to/results/kda-prefill-ab -export PYTHONPATH="${HELION_ROOT}/scripts/kda_sglang_ab:${HELION_ROOT}:${SGLANG_ROOT}/python" -export HF_HUB_OFFLINE=1 -export TRANSFORMERS_OFFLINE=1 -mkdir -p "${RESULTS}" -``` - -Launch the default Triton baseline: - -```bash -cd "${SGLANG_ROOT}" -unset SGLANG_KDA_HELION_PREFILL -python -m sglang.launch_server \ - --model-path "${MODEL}" \ - --trust-remote-code \ - --tp-size 2 \ - --linear-attn-backend triton \ - --linear-attn-decode-backend triton \ - --linear-attn-prefill-backend triton \ - --mamba-ssm-dtype float32 \ - --disable-custom-all-reduce \ - --port 30000 \ - 2>&1 | tee "${RESULTS}/triton-server.log" -``` - -For the numerical-contract Helion path, use the same command with: - -```bash -export SGLANG_KDA_HELION_PREFILL=1 -unset HELION_KDA_PREFILL_NEWTON_SCHULZ -``` - -Set `HELION_KDA_PREFILL_NEWTON_SCHULZ=1` in addition to the substitution flag -to benchmark the separately tuned Newton-Schulz path. - -With only this flag, the hook in `scripts/kda_sglang_ab/sitecustomize.py` -replaces only `kda_triton.chunk_kda`; decode and all other model kernels remain -unchanged. `SGLANG_KDA_HELION_DECODE=1` independently enables the packed-decode -substitution. Confirm the corresponding two `first ... call` markers, one per -TP worker, before using a Helion result. - -After the server reports readiness, run one short workload to settle the full -serving path, then measure: - -```bash -python -m sglang.benchmark.serving \ - --backend sglang \ - --host 127.0.0.1 \ - --port 30000 \ - --dataset-name random \ - --dataset-path "${DATASET}" \ - --model "${MODEL}" \ - --num-prompts 100 \ - --random-input-len 4096 \ - --random-output-len 8 \ - --random-range-ratio 0.25 \ - --max-concurrency 4 \ - --seed 4242 \ - --temperature 0 \ - --flush-cache \ - --warmup-requests 1 \ - --output-details \ - --disable-tqdm \ - --output-file "${RESULTS}/result.jsonl" -``` - -## Verification - -Run the focused contract tests before benchmarking: - -```bash -cd "${HELION_ROOT}" -pytest test/test_kda_prefill.py -x -vv -s -``` - -The final run passed all four tests in 36.01 seconds. A wider direct packed test -at T=73, H=2, K=256, and V=80 had maximum output error `3.052e-5`, maximum -intermediate-state error `1.707e-4`, finite outputs, and exact preservation of -untouched state-pool rows. A Newton-Schulz packed-varlen safe-gate check at -K=V=32 preserved the in-place and untouched-slot contracts with maximum output -and state errors of `1.831e-4` and `7.321e-4`. - -SGLang's benchmark correctness mode also passed both Helion paths and CuTe DSL -against `fused_recurrent_kda` at T=128, 192, 256, 512, and 1024 with H=32 and -K=V=128. Both Helion paths had maximum output error `4.88e-4`; the largest -final-state errors were `4.39e-3` for forward substitution and `4.71e-3` for -Newton-Schulz. - -## Default-backend FP32/BF16 matrix - -On 2026-07-24, the TP=2 serving benchmark was repeated without the base or -prefill backend arguments. SGLang resolved the default backends as follows: - -```text -FP32 state: decode=triton, prefill=triton -BF16 state: decode=flashinfer, prefill=triton -``` - -Helion decode requires `TritonKDAKernel` as its integration point. Therefore, -the BF16 decode-only and combined rows explicitly selected -`--linear-attn-decode-backend triton`, then replaced its packed-decode callable -with Helion. Their performance comparison is still against the flag-free -FlashInfer baseline. BF16 prefill-only retained default FlashInfer decode. - -Each configuration ran one 20-request warmup followed by three deterministic -100-request samples. Each sample processed 247,525 input and 485 output tokens -with no request errors. The primary statistic is the three-run median because -the serving benchmark occasionally contains multi-second scheduling stalls. - -| State | Configuration | Input tok/s samples | Median | vs default | Geomean | -|---|---|---|---:|---:|---:| -| FP32 | default Triton/Triton | 24,129 / 26,784 / 25,279 | 25,279 | baseline | 25,374 | -| FP32 | Helion decode | 27,449 / 29,436 / 28,704 | 28,704 | +13.55% | 28,518 | -| FP32 | Helion prefill | 16,588 / 26,729 / 25,833 | 25,833 | +2.19% | 22,541 | -| FP32 | Helion decode + prefill | 29,113 / 29,148 / 29,636 | 29,148 | +15.31% | 29,298 | -| BF16 | default FlashInfer/Triton | 25,618 / 27,369 / 27,437 | 27,369 | baseline | 26,795 | -| BF16 | Helion decode | 25,594 / 27,795 / 27,175 | 27,175 | -0.71% | 26,838 | -| BF16 | Helion prefill | 27,169 / 28,700 / 27,496 | 27,496 | +0.46% | 27,781 | -| BF16 | Helion decode + prefill | 29,196 / 28,937 / 29,420 | 29,196 | +6.67% | 29,183 | - -The first FP32 Helion-prefill sample reproduced the previously observed -one-time long-run stall. It emitted no new Helion code and subsequent samples -were normal. It is retained in both the raw samples and geomean; the median is -reported as the robust primary result. The small BF16 prefill-only difference -is within observed run-to-run noise, while the combined improvements are much -larger and stable across all three samples. - -BF16 state also nearly doubled recurrent-state capacity at the same memory -budget: `max_mamba_cache_size` increased from 2,605 to 5,040 and -`max_running_requests` from 521 to 1,008. The complete logs and JSONL outputs -are saved under: - -```text -/home/eche/results/kda-default-matrix-20260724 -``` - -## ShareGPT FP32-state occupancy sweep - -A flag-free default baseline and the three Helion substitutions were measured -with 128 fixed-seed ShareGPT prompts, 64 output tokens per prompt, FP32 state, -TP=2, and concurrency 1/4/16/32/64/128. No linear-attention backend arguments -were supplied; SGLang resolved the baseline to Triton decode and Triton prefill. -Each cache-flushed cell processed 31,609 input and 8,192 output tokens with 128 -successful requests and no nonempty errors. - -| Configuration | C=1 | C=4 | C=16 | C=32 | C=64 | C=128 | Ratio geomean | -|---|---:|---:|---:|---:|---:|---:|---:| -| Triton baseline | 848.3 | 2,427.6 | 7,206.0 | 12,780.0 | 19,239.6 | 30,858.0 | 1.0000x | -| Helion decode | 876.7 | 2,535.2 | 7,339.8 | 12,992.4 | 19,886.1 | 30,915.2 | 1.0247x | -| Helion prefill | 948.6 | 2,543.5 | 7,562.7 | 13,067.4 | 19,592.1 | 31,292.5 | 1.0445x | -| Helion decode + prefill | 942.0 | 2,469.7 | 7,296.1 | 12,964.2 | 19,299.5 | 30,430.8 | 1.0233x | - -Values are total input-plus-output tokens per second. The geomean speedups were -2.47% for decode, 4.45% for prefill, and 2.33% for both. Prefill-only reduced -mean-TTFT geomean by 9.95%; at concurrency 1 it improved total throughput by -11.82% and reduced mean TTFT from 156.1 to 117.7 ms. The combined row was not -additive in this one-run sweep, so its small differences at higher occupancy -should be treated as serving variance rather than a kernel interaction. - -The first baseline C=128 sample and a later baseline C=64 repeat each hit an -isolated scheduler stall. Both are preserved; the table uses the clean original -C=64 and clean repeated C=128. Full logs, JSONL files, per-cell percentages, and -the outlier policy are under: - -```text -/home/eche/results/kda-sharegpt-occupancy-fp32-20260724 -``` diff --git a/scripts/kda_sglang_ab/sitecustomize.py b/scripts/kda_sglang_ab/sitecustomize.py deleted file mode 100644 index 4563c1bd1..000000000 --- a/scripts/kda_sglang_ab/sitecustomize.py +++ /dev/null @@ -1,85 +0,0 @@ -"""Opt-in Helion KDA substitutions for SGLang A/B benchmarks.""" - -from __future__ import annotations - -import importlib.util -import os -from pathlib import Path -import sys -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from types import ModuleType - - -def _load_source_module(name: str, path: Path) -> ModuleType: - spec = importlib.util.spec_from_file_location(name, path) - if spec is None or spec.loader is None: - raise ImportError(f"cannot load {path}") - module = importlib.util.module_from_spec(spec) - sys.modules[name] = module - spec.loader.exec_module(module) - return module - - -enable_decode = ( - os.environ.get("SGLANG_KDA_HELION_DECODE") == "1" - or os.environ.get("SGLANG_KDA_HELION") == "1" -) -enable_prefill = os.environ.get("SGLANG_KDA_HELION_PREFILL") == "1" - -if enable_decode or enable_prefill: - from sglang.srt.layers.attention.linear.kernels import kda_triton - -if enable_decode: - decode_root = os.environ.get("HELION_KDA_DECODE_ROOT") - if decode_root is None: - from examples.linear.kda_packed_decode import ( - helion_fused_recurrent_kda_packed_decode, - ) - else: - decode_module = _load_source_module( - "_helion_kda_decode_ab", - Path(decode_root) / "examples/linear/kda_packed_decode.py", - ) - helion_fused_recurrent_kda_packed_decode = ( - decode_module.helion_fused_recurrent_kda_packed_decode - ) - - _decode_called = False - - def _helion_packed_decode(*args: object, **kwargs: object) -> object: - global _decode_called - if not _decode_called: - print("[helion-kda-ab] first packed decode call", flush=True) - _decode_called = True - return helion_fused_recurrent_kda_packed_decode(*args, **kwargs) - - kda_triton.fused_recurrent_kda_packed_decode = _helion_packed_decode - print("[helion-kda-ab] installed packed decode substitution", flush=True) - -if enable_prefill: - prefill_root = os.environ.get("HELION_KDA_PREFILL_ROOT") - if prefill_root is None: - from examples.linear.kda_prefill import chunk_kda as helion_chunk_kda - else: - prefill_module = _load_source_module( - "_helion_kda_prefill_ab", - Path(prefill_root) / "examples/linear/kda_prefill.py", - ) - helion_chunk_kda = prefill_module.chunk_kda - newton_schulz = os.environ.get("HELION_KDA_PREFILL_NEWTON_SCHULZ") == "1" - - _prefill_called = False - - def _helion_chunk_kda(*args: object, **kwargs: object) -> object: - global _prefill_called - if not _prefill_called: - inverse = "newton-schulz" if newton_schulz else "forward-substitution" - print(f"[helion-kda-ab] first prefill call ({inverse})", flush=True) - _prefill_called = True - kwargs["newton_schulz"] = newton_schulz - return helion_chunk_kda(*args, **kwargs) - - kda_triton.chunk_kda = _helion_chunk_kda - print("[helion-kda-ab] installed prefill substitution", flush=True) diff --git a/test/test_kda_packed_decode.py b/test/test_kda_packed_decode.py deleted file mode 100644 index e8201a322..000000000 --- a/test/test_kda_packed_decode.py +++ /dev/null @@ -1,95 +0,0 @@ -from __future__ import annotations - -import inspect - -from examples.linear.kda_packed_decode import _helion_fused_recurrent_kda_packed_decode -from examples.linear.kda_packed_decode import helion_fused_recurrent_kda_packed_decode -from examples.linear.kda_packed_decode import make_kda_inputs -from examples.linear.kda_packed_decode import torch_fused_recurrent_kda_packed_decode -import torch - -from helion._testing import DEVICE -from helion._testing import RefEagerTestDisabled -from helion._testing import TestCase -from helion._testing import onlyBackends - - -@onlyBackends(["triton"]) -class TestKdaPackedDecode(RefEagerTestDisabled, TestCase): - def test_uses_one_config_for_all_shapes(self) -> None: - self.assertEqual(len(_helion_fused_recurrent_kda_packed_decode.configs), 1) - self.assertFalse( - _helion_fused_recurrent_kda_packed_decode.settings.static_shapes - ) - config = _helion_fused_recurrent_kda_packed_decode.configs[0] - self.assertEqual(config.block_sizes, [8]) - self.assertEqual(config.loop_orders, [[2, 1, 0]]) - self.assertEqual(config.pid_type, "xyz") - - def test_public_signature(self) -> None: - signature = inspect.signature(helion_fused_recurrent_kda_packed_decode) - self.assertEqual( - list(signature.parameters), - [ - "mixed_qkv", - "a", - "b", - "A_log", - "dt_bias", - "scale", - "initial_state", - "out", - "ssm_state_indices", - "use_qk_l2norm_in_kernel", - ], - ) - self.assertIs(signature.parameters["use_qk_l2norm_in_kernel"].default, False) - - def test_values_mutations_aliases_and_padding(self) -> None: - for use_qk_l2norm_in_kernel in (False, True): - with self.subTest(use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel): - inputs = make_kda_inputs( - 3, - 2, - 4, - 128, - 128, - device=DEVICE, - pool_size=7, - seed=123, - ) - inputs.ssm_state_indices.copy_( - torch.tensor([5, -1, 2], device=DEVICE, dtype=torch.int32) - ) - inputs.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel - original_state = inputs.initial_state.clone() - - reference_inputs = inputs.clone_mutable() - actual_inputs = inputs.clone_mutable() - torch_fused_recurrent_kda_packed_decode(*reference_inputs.args()) - result = helion_fused_recurrent_kda_packed_decode(*actual_inputs.args()) - - self.assertEqual(result[0].data_ptr(), actual_inputs.out.data_ptr()) - self.assertEqual( - result[1].data_ptr(), actual_inputs.initial_state.data_ptr() - ) - torch.testing.assert_close( - actual_inputs.out, reference_inputs.out, atol=2e-2, rtol=1e-2 - ) - torch.testing.assert_close( - actual_inputs.initial_state, - reference_inputs.initial_state, - atol=2e-2, - rtol=1e-2, - ) - self.assertEqual(torch.count_nonzero(actual_inputs.out[1]).item(), 0) - - untouched = torch.tensor( - [0, 1, 3, 4, 6], device=DEVICE, dtype=torch.long - ) - self.assertTrue( - torch.equal( - actual_inputs.initial_state[untouched], - original_state[untouched], - ) - ) diff --git a/test/test_kda_prefill.py b/test/test_kda_prefill.py deleted file mode 100644 index 80cd2418c..000000000 --- a/test/test_kda_prefill.py +++ /dev/null @@ -1,278 +0,0 @@ -from __future__ import annotations - -import inspect -import math - -from examples.linear.kda_prefill import chunk_kda -import torch - -from helion._testing import DEVICE -from helion._testing import RefEagerTestDisabled -from helion._testing import TestCase -from helion._testing import onlyBackends - - -def _torch_chunk_kda( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - initial_state: torch.Tensor, - initial_state_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]: - if scale is None: - scale = k.shape[-1] ** -0.5 - if use_qk_l2norm_in_kernel: - q = ( - q.float() / torch.sqrt((q.float().square()).sum(-1, keepdim=True) + 1e-6) - ).to(q.dtype) - k = ( - k.float() / torch.sqrt((k.float().square()).sum(-1, keepdim=True) + 1e-6) - ).to(k.dtype) - - if A_log is not None: - raw = g.float() - if dt_bias is not None: - raw = raw + dt_bias.reshape(g.shape[2], g.shape[3])[None, None] - a = torch.exp(A_log.reshape(-1).float())[None, None, :, None] - if lower_bound is None: - gate = -a * torch.nn.functional.softplus(raw) - else: - gate = lower_bound * torch.sigmoid(a * raw) - else: - gate = g.float() - - output = torch.empty_like(v) - h_chunks: list[torch.Tensor] = [] - if cu_seqlens is None: - sequences = [ - (batch, batch * q.shape[1], (batch + 1) * q.shape[1]) - for batch in range(q.shape[0]) - ] - else: - offsets = cu_seqlens.cpu().tolist() - sequences = [ - (sequence, offsets[sequence], offsets[sequence + 1]) - for sequence in range(len(offsets) - 1) - ] - - for sequence, begin, end in sequences: - state_index = int(initial_state_indices[sequence].item()) - state = initial_state[state_index].float() - sequence_chunks: list[torch.Tensor] = [] - for flat_token in range(begin, end): - if (flat_token - begin) % 64 == 0: - sequence_chunks.append(state.to(q.dtype)) - if cu_seqlens is None: - batch = sequence - token = flat_token - begin - else: - batch = 0 - token = flat_token - q_t = q[batch, token].float() - k_t = k[batch, token].float() - v_t = v[batch, token].float() - beta_t = beta[batch, token].float() - state = state * torch.exp(gate[batch, token])[:, None, :] - prediction = (state * k_t[:, None, :]).sum(-1) - residual = (v_t - prediction) * beta_t[:, None] - state = state + residual[:, :, None] * k_t[:, None, :] - output[batch, token] = ( - (state * (q_t * scale)[:, None, :]).sum(-1).to(output.dtype) - ) - initial_state[state_index] = state.to(initial_state.dtype) - h_chunks.extend(sequence_chunks) - - h = torch.stack(h_chunks).unsqueeze(0) - if cu_seqlens is None: - chunks_per_batch = math.ceil(q.shape[1] / 64) - h = h.reshape(q.shape[0], chunks_per_batch, *h.shape[2:]) - return output, h - - -@onlyBackends(["triton"]) -class TestKdaPrefill(RefEagerTestDisabled, TestCase): - def test_public_signature(self) -> None: - signature = inspect.signature(chunk_kda) - expected_names = [ - "q", - "k", - "v", - "g", - "beta", - "scale", - "initial_state", - "initial_state_indices", - "use_qk_l2norm_in_kernel", - "cu_seqlens", - "A_log", - "dt_bias", - "lower_bound", - "output_intermediate_states", - "kwargs", - ] - self.assertEqual(list(signature.parameters), expected_names) - self.assertEqual( - [parameter.kind for parameter in signature.parameters.values()], - [inspect.Parameter.POSITIONAL_OR_KEYWORD] * 14 - + [inspect.Parameter.VAR_KEYWORD], - ) - self.assertEqual( - [parameter.default for parameter in signature.parameters.values()], - [inspect.Parameter.empty] * 5 - + [None, None, None, False, None, None, None, None, False] - + [inspect.Parameter.empty], - ) - - def test_fixed_partial_chunk_and_state_pool(self) -> None: - torch.manual_seed(123) - B, T, H, K, V = 2, 17, 2, 32, 32 - q = torch.randn(B, T, H, K, device=DEVICE, dtype=torch.bfloat16) - k = torch.randn_like(q) - v = torch.randn(B, T, H, V, device=DEVICE, dtype=torch.bfloat16) * 0.1 - g = torch.randn_like(q) * 0.2 - beta = torch.rand(B, T, H, device=DEVICE) - a_log = torch.full([H], -2.0, device=DEVICE) - dt_bias = torch.zeros(H * K, device=DEVICE) - indices = torch.tensor([3, 1], device=DEVICE, dtype=torch.int32) - initial = torch.randn(5, H, V, K, device=DEVICE) * 0.01 - - reference_state = initial.clone() - actual_state = initial.clone() - expected, expected_h = _torch_chunk_kda( - q, - k, - v, - g, - beta, - reference_state, - indices, - use_qk_l2norm_in_kernel=True, - A_log=a_log, - dt_bias=dt_bias, - ) - output_buffer = v.clone() - actual, actual_h = chunk_kda( - q, - k, - output_buffer, - g, - beta, - initial_state=actual_state, - initial_state_indices=indices, - use_qk_l2norm_in_kernel=True, - A_log=a_log, - dt_bias=dt_bias, - output_intermediate_states=True, - ) - - self.assertEqual(actual.data_ptr(), output_buffer.data_ptr()) - torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) - torch.testing.assert_close(actual_h, expected_h, atol=2e-2, rtol=2e-2) - torch.testing.assert_close(actual_state, reference_state, atol=2e-2, rtol=2e-2) - self.assertTrue(torch.equal(actual_state[0], initial[0])) - self.assertTrue(torch.equal(actual_state[2], initial[2])) - self.assertTrue(torch.equal(actual_state[4], initial[4])) - - def test_varlen_safe_gate(self) -> None: - torch.manual_seed(456) - lengths = [1, 15, 17] - T, H, K, V = sum(lengths), 2, 32, 32 - cu_seqlens = torch.tensor( - [0, *torch.tensor(lengths).cumsum(0).tolist()], - device=DEVICE, - dtype=torch.int32, - ) - q = torch.randn(1, T, H, K, device=DEVICE, dtype=torch.bfloat16) - k = torch.randn_like(q) - v = torch.randn(1, T, H, V, device=DEVICE, dtype=torch.bfloat16) * 0.1 - g = torch.randn_like(q) * 0.2 - beta = torch.rand(1, T, H, device=DEVICE) - a_log = torch.full([H], -2.0, device=DEVICE) - indices = torch.tensor([4, 1, 3], device=DEVICE, dtype=torch.int32) - initial = torch.randn(6, H, V, K, device=DEVICE) * 0.01 - - reference_state = initial.clone() - actual_state = initial.clone() - expected, expected_h = _torch_chunk_kda( - q, - k, - v, - g, - beta, - reference_state, - indices, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens, - A_log=a_log, - lower_bound=-0.01, - ) - actual, actual_h = chunk_kda( - q, - k, - v.clone(), - g, - beta, - initial_state=actual_state, - initial_state_indices=indices, - use_qk_l2norm_in_kernel=True, - cu_seqlens=cu_seqlens, - A_log=a_log, - lower_bound=-0.01, - output_intermediate_states=True, - ) - - torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) - torch.testing.assert_close(actual_h, expected_h, atol=2e-2, rtol=2e-2) - torch.testing.assert_close(actual_state, reference_state, atol=2e-2, rtol=2e-2) - - def test_fp16_preactivated_gate(self) -> None: - torch.manual_seed(789) - B, T, H, K, V = 1, 17, 1, 32, 32 - q = torch.nn.functional.normalize( - torch.randn(B, T, H, K, device=DEVICE), dim=-1 - ).half() - k = torch.nn.functional.normalize( - torch.randn(B, T, H, K, device=DEVICE), dim=-1 - ).half() - v = torch.randn(B, T, H, V, device=DEVICE, dtype=torch.float16) * 0.1 - g = -torch.rand(B, T, H, K, device=DEVICE) * 0.01 - beta = torch.rand(B, T, H, device=DEVICE) - indices = torch.tensor([1], device=DEVICE, dtype=torch.int32) - initial = torch.randn(3, H, V, K, device=DEVICE, dtype=torch.bfloat16) * 0.01 - - reference_state = initial.clone() - actual_state = initial.clone() - expected, expected_h = _torch_chunk_kda( - q, - k, - v, - g, - beta, - reference_state, - indices, - ) - actual, actual_h = chunk_kda( - q, - k, - v.clone(), - g, - beta, - initial_state=actual_state, - initial_state_indices=indices, - output_intermediate_states=True, - ) - - self.assertEqual(actual.dtype, torch.float16) - self.assertEqual(actual_h.dtype, torch.float16) - torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) - torch.testing.assert_close(actual_h, expected_h, atol=2e-2, rtol=2e-2) - torch.testing.assert_close(actual_state, reference_state, atol=2e-2, rtol=2e-2)