diff --git a/docs/api/kda.rst b/docs/api/kda.rst index e7f27af61aa..cc8f638de27 100644 --- a/docs/api/kda.rst +++ b/docs/api/kda.rst @@ -52,6 +52,16 @@ The public ``recurrent_kda`` API supports standard decode with one token per sequence (``T=1``) and packed speculative decode with two or more tokens per sequence (``T>=2``). +``fused_kda_decode`` combines width-four causal convolution, one KDA update, +and gated RMSNorm for ``T=1``. ``packed_fused_kda_decode`` extends the same +pipeline to packed ragged speculative decode, retaining one recurrent +checkpoint per token and an extended rolling convolution window. Its optional +``t1_state_indices`` argument is a T=1-only fast path containing caller-resolved +per-row cache slots. T>1 always uses the ``state_indices[N, T]``, +``query_start_loc``, and ``num_accepted_tokens`` metadata with the multitoken +kernel; passing ``t1_state_indices`` for T>1 is rejected. No T>1 kernel +behavior or performance is changed by this optimization. + Pass ``backend="cake"`` to select the exported Cake backend. On SM100-family SM100a (B200/GB200) and SM103a (B300/GB300) devices, its D128 ``T=1..6`` family with in-kernel QK normalization exports 25 frozen CUDA bodies: @@ -167,6 +177,7 @@ small ``B * HV * T``. Requires SM90+ for the WY path and ``K = V = 128``; :toctree: ../generated fused_kda_decode + packed_fused_kda_decode packed_kda_decode recurrent_kda diff --git a/flashinfer/__init__.py b/flashinfer/__init__.py index 2b47c542495..8235e6d34c6 100644 --- a/flashinfer/__init__.py +++ b/flashinfer/__init__.py @@ -218,6 +218,7 @@ def __dir__() -> list[str]: from .kda import recurrent_kda as recurrent_kda from .kda_decode import fused_kda_decode as fused_kda_decode from .kda_decode import packed_kda_decode as packed_kda_decode +from .kda_decode import packed_fused_kda_decode from .cake_minimax_h3 import MiniMaxH3Mxfp8PreAttention as MiniMaxH3Mxfp8PreAttention from .mla import BatchMLAPagedAttentionWrapper as BatchMLAPagedAttentionWrapper from . import mhc as mhc diff --git a/flashinfer/kda_decode.py b/flashinfer/kda_decode.py index bb74b959ff8..2aa3513ae1c 100644 --- a/flashinfer/kda_decode.py +++ b/flashinfer/kda_decode.py @@ -32,6 +32,7 @@ from .api_logging import flashinfer_api from .trace.templates.kda import ( + packed_fused_kda_decode_trace, fused_kda_decode_trace, packed_kda_decode_trace, recurrent_kda_decode_trace, @@ -48,6 +49,16 @@ _FUSED_KDA_DECODE_AVAILABLE = False from .kda_kernels import run_packed_kda_decode as _run_packed_kda_decode + +try: + from .kda_kernels.fused_kda_decode_multitoken import ( + run_fused_kda_decode_multitoken as _run_packed_fused_kda_decode, + ) + + _FUSED_KDA_DECODE_PACKED_AVAILABLE = True +except (ImportError, RuntimeError): + _run_packed_fused_kda_decode = None + _FUSED_KDA_DECODE_PACKED_AVAILABLE = False from .kda_kernels import run_recurrent_kda as _run_recurrent_kda try: @@ -605,6 +616,227 @@ def fused_kda_decode( ) +def _validate_packed_t1_remap_inputs(x, query_start_loc, state_indices): + """Validate remap metadata structure without reading device values.""" + if query_start_loc.dtype != torch.int32: + raise TypeError( + f"query_start_loc must have dtype torch.int32, got {query_start_loc.dtype}" + ) + if state_indices.dtype != torch.int32: + raise TypeError( + f"state_indices must have dtype torch.int32, got {state_indices.dtype}" + ) + if state_indices.ndim != 2 or state_indices.shape[1] != 1: + raise ValueError("state_indices must have shape [N, 1]") + if ( + query_start_loc.ndim != 1 + or query_start_loc.shape[0] != state_indices.shape[0] + 1 + ): + raise ValueError("query_start_loc must have shape [state_indices.shape[0] + 1]") + if not query_start_loc.is_contiguous(): + raise ValueError("query_start_loc must be contiguous") + if not state_indices.is_contiguous(): + raise ValueError("state_indices must be contiguous") + if query_start_loc.device != x.device or state_indices.device != x.device: + raise ValueError("T=1 remap metadata must be on the same device as x") + + +@flashinfer_api(trace=packed_fused_kda_decode_trace) +def packed_fused_kda_decode( + x: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + raw_gate: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + state_indices: torch.Tensor, + state: torch.Tensor, + output_gate: torch.Tensor, + norm_weight: torch.Tensor, + lower_bound: Optional[float] = -5.0, + norm_eps: float = 1e-5, + output: Optional[torch.Tensor] = None, + *, + query_start_loc: torch.Tensor, + num_accepted_tokens: torch.Tensor, + t1_state_indices: Optional[torch.Tensor] = None, +) -> torch.Tensor: + r"""Run packed T>=1 fused KDA with per-token cache checkpoints. + + Tokens use a packed ragged layout described by ``query_start_loc``. + ``state_indices`` has shape ``[N, T]``, where T is the maximum verification + length. ``num_accepted_tokens[n] - 1`` selects the source checkpoint and + convolution-history offset for sequence ``n``; token ``t`` writes its + recurrent checkpoint to ``state_indices[n, t]``. The convolution cache is + one rolling window of length ``T + 2`` in ``state_indices[n, 0]``. + + This facade is CuTe-only: T=1 uses the CuTe DSL implementation underlying + ``fused_kda_decode`` (without Cake or backend selection), including + bfloat16-state and softplus-gate support. The optional + ``t1_state_indices`` optimization is strictly T=1-only; passing it for + T>1 fails closed. T>1 uses the SM10x CuTe DSL packed backend, which + requires FP32 + recurrent state and a finite negative ``lower_bound``. Both paths require + head dimension 128, convolution width four, and 12, 24, 32, 48, or 96 + heads. Other tensors follow :func:`fused_kda_decode`, with the packed row + count supplied by ``x``. + + ``query_start_loc`` must be contiguous int32 of shape ``[N+1]``, start at + zero, and contain nondecreasing offsets within ``x`` with active lengths no + larger than T. + ``num_accepted_tokens`` must be contiguous int32 of shape ``[N]`` with + values in ``[1, T]``. For T=1, ``t1_state_indices`` may optionally provide + the already-resolved contiguous per-row cache slots with shape + ``[num_rows]``. When supplied, it is used by the direct fused T=1 + specialization and ``query_start_loc`` is retained only as structural + metadata; its values are not read by the host or kernel. This is intended + for CUDA-Graph-stable caller-owned buffers updated in-place. + Active recurrent destinations must be positive cache + slots and must not alias destinations from another active sequence; + non-positive or zero-length rows are null rows and do not mutate either + cache. In T=1, packed rows beyond ``query_start_loc[-1]`` are trailing + unused capacity: they receive zero output and do not mutate either cache. + These value constraints are caller-owned so CUDA Graph replay never performs + a device-to-host validation sync. + + Args: + x: + Packed QKV projection with shape ``[num_rows, 3 * H * 128]`` and + dtype bfloat16. + weight: + Depthwise convolution weights with shape ``[3, 4, H * 128]`` and + dtype float32. + conv_state: + Paged bfloat16 convolution cache. T=1 uses history length three; + T>1 uses the extended rolling history length ``T + 2``. + raw_gate: + Raw recurrence gate with shape ``[1, num_rows, H, 128]`` and dtype + bfloat16. + raw_beta: + Raw delta-rule learning-rate logits with shape + ``[1, num_rows, H]`` and dtype bfloat16. + A_log: + Log decay parameter with ``H`` elements and dtype float32. + dt_bias: + Per-channel decay bias with ``H * 128`` elements and dtype float32. + state_indices: + Contiguous int32 cache slots with shape ``[N, T]``. For T>1, + column ``t`` is the recurrent checkpoint destination for token + ``t``. Column zero also selects the rolling convolution cache. + query_start_loc: + Contiguous int32 packed-row offsets with shape ``[N + 1]``. + num_accepted_tokens: + Contiguous int32 accepted-token counts with shape ``[N]``. + t1_state_indices: + Optional contiguous int32 per-packed-row cache slots with shape + ``[num_rows]``. T=1 only; null and trailing rows must be zero or + non-positive. The tensor is never allocated or populated by this + function. + state: + Paged recurrent state with shape ``[num_slots, H, 128, 128]``. + T=1 accepts float32 or bfloat16; T>1 requires float32. + output_gate: + Gated RMSNorm logits with shape ``[num_rows, H, 128]`` or + ``[1, num_rows, H, 128]`` and dtype bfloat16. + norm_weight: + RMSNorm weight with 128 elements and dtype float32. + lower_bound: + Negative recurrence-gate lower bound. T=1 also accepts ``None`` + for the softplus gate; T>1 requires a finite negative value. + norm_eps: + Non-negative RMSNorm epsilon. + output: + Optional preallocated contiguous bfloat16 output with shape + ``[1, num_rows, H, 128]``. + + Returns: + The packed bfloat16 output with shape ``[1, num_rows, H, 128]``. + """ + if state_indices.ndim < 2: + raise ValueError("state_indices must have shape [N, T]") + if t1_state_indices is not None: + if state_indices.ndim != 2 or state_indices.shape[1] != 1: + raise ValueError( + "t1_state_indices is only supported for T=1; T>1 uses the " + "multitoken kernel unchanged" + ) + if ( + t1_state_indices.ndim != 1 + or t1_state_indices.shape != (x.shape[0],) + or t1_state_indices.dtype != torch.int32 + or not t1_state_indices.is_contiguous() + or t1_state_indices.device != x.device + ): + raise ValueError( + "t1_state_indices must be contiguous int32 [num_rows] on x.device" + ) + if _run_fused_kda_decode is None: + raise NotImplementedError("fused KDA decode backend is unavailable") + _validate_packed_t1_remap_inputs(x, query_start_loc, state_indices) + return _run_fused_kda_decode( + x=x, + weight=weight, + conv_state=conv_state, + raw_gate=raw_gate, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + state_indices=t1_state_indices, + state=state, + output_gate=output_gate, + norm_weight=norm_weight, + lower_bound=lower_bound, + norm_eps=norm_eps, + output=output, + ) + if state_indices.ndim == 2 and state_indices.shape[1] == 1: + if _run_fused_kda_decode is None: + raise NotImplementedError("fused KDA decode backend is unavailable") + _validate_packed_t1_remap_inputs(x, query_start_loc, state_indices) + return _run_fused_kda_decode( + x=x, + weight=weight, + conv_state=conv_state, + raw_gate=raw_gate, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + # T=1 does not need a remap buffer. The direct ABI argument is + # unused by the packed constexpr specialization; pass the first + # packed destination as the compatible [N] view. + state_indices=state_indices[:, 0], + query_start_loc=query_start_loc, + packed_state_indices=state_indices, + state=state, + output_gate=output_gate, + norm_weight=norm_weight, + lower_bound=lower_bound, + norm_eps=norm_eps, + output=output, + ) + if _run_packed_fused_kda_decode is None: + raise NotImplementedError("packed fused KDA decode backend is unavailable") + return _run_packed_fused_kda_decode( + x=x, + weight=weight, + conv_state=conv_state, + raw_gate=raw_gate, + raw_beta=raw_beta, + A_log=A_log, + dt_bias=dt_bias, + state_indices=state_indices, + query_start_loc=query_start_loc, + num_accepted_tokens=num_accepted_tokens, + state=state, + output_gate=output_gate, + norm_weight=norm_weight, + lower_bound=lower_bound, + norm_eps=norm_eps, + output=output, + ) + + # Cached trivial metadata for the frozen mode's uniform batched form: the # arange cu_seqlens / slot indices are content-stable per (device, B, T), so # building them once keeps the hot path allocation- and launch-free (required diff --git a/flashinfer/kda_kernels/__init__.py b/flashinfer/kda_kernels/__init__.py index 0de29901c5e..05ec0932562 100644 --- a/flashinfer/kda_kernels/__init__.py +++ b/flashinfer/kda_kernels/__init__.py @@ -26,6 +26,7 @@ - run_recurrent_kda: Recurrent KDA standard decode and speculative decode backend - run_fused_kda_decode: Fused Kimi K3 conv, recurrent KDA, and RMSNorm backend - run_packed_kda_decode: Packed Kimi K3 T=1 recurrent decode backend +- run_fused_kda_decode_multitoken: Packed T>=1 fused KDA backend The ``run_*`` names above are canonical. The bare ``recurrent_kda``, ``fused_kda_decode`` and ``packed_kda_decode`` aliases are kept for existing @@ -61,6 +62,14 @@ # ``run_recurrent_kda`` in ``recurrent_kda.py``); import it by module path # only for tests and benchmarks. +try: + from .fused_kda_decode_multitoken import run_fused_kda_decode_multitoken + + fused_kda_decode_multitoken = run_fused_kda_decode_multitoken +except (ImportError, RuntimeError): + run_fused_kda_decode_multitoken = None # type: ignore + fused_kda_decode_multitoken = None # type: ignore + try: if _torch.cuda.is_available(): from ..cute_dsl.availability import is_cute_dsl_arch_supported as _dsl_arch_ok @@ -112,6 +121,7 @@ __all__ = [ "run_fused_kda_decode", + "run_fused_kda_decode_multitoken", "run_packed_kda_decode", "run_recurrent_kda", ] diff --git a/flashinfer/kda_kernels/fused_kda_decode.py b/flashinfer/kda_kernels/fused_kda_decode.py index 1c68f851f6a..b48538a6b1b 100644 --- a/flashinfer/kda_kernels/fused_kda_decode.py +++ b/flashinfer/kda_kernels/fused_kda_decode.py @@ -88,10 +88,13 @@ def _fused_kda_decode_kernel( A_log, raw_beta, state_indices, + query_start_loc, + packed_state_indices, state, output_gate, norm_weight, output, + packed_t1: cutlass.Constexpr, state_is_bf16: cutlass.Constexpr, use_lower_bound: cutlass.Constexpr, lower_bound: cutlass.Constexpr, @@ -140,10 +143,38 @@ def _fused_kda_decode_kernel( cute.nvgpu.CopyUniversalOp(), BF16, num_bits_per_copy=64 ) - # Slot zero is a read-only null slot. Non-positive indices produce a zero - # output and suppress both cache updates while still following a uniform - # control-flow path through the CTA. - requested_slot = state_indices[row_idx] + # Slot zero is a read-only null slot. Packed T=1 resolves the sequence + # checkpoint in this kernel so the public wrapper has no metadata launch. + # Initialize before staged control flow; CuTe's Python staging requires + # values assigned in constexpr branches to have a definition visible on + # both paths. The direct assignment below is eliminated for packed T=1. + requested_slot = cutlass.Int32(0) + if cutlass.const_expr(packed_t1): + num_sequences = packed_state_indices.shape[0] + row = cutlass.Int32(row_idx) + # For valid T=1 metadata, qsl[r + 1] == r + 1 proves that every + # preceding sequence has length one, so row r is sequence r. Keep the + # endpoint load device-side: graph replay may change metadata. + identity = False + if row < num_sequences: + identity = query_start_loc[row + 1] == row + 1 + if identity: + requested_slot = packed_state_indices[row, 0] + else: + lo = cutlass.Int32(0) + hi = num_sequences + while lo < hi: + mid = (lo + hi) // 2 + if query_start_loc[mid + 1] <= row: + lo = mid + 1 + else: + hi = mid + if lo < num_sequences: + requested_slot = packed_state_indices[lo, 0] + else: + requested_slot = cutlass.Int32(0) + else: + requested_slot = state_indices[row_idx] is_live = requested_slot > 0 slot = cute.max(requested_slot, cutlass.Int32(0)) @@ -402,14 +433,17 @@ def _fused_kda_decode_launch( output_gate, norm_weight, output, + query_start_loc, + packed_state_indices, stream: cuda.CUstream, state_is_bf16: cutlass.Constexpr, + packed_t1: cutlass.Constexpr, use_lower_bound: cutlass.Constexpr, lower_bound: cutlass.Constexpr, norm_eps: cutlass.Constexpr, ): num_heads = A_log.shape[0] - num_rows = state_indices.shape[0] + num_rows = x.shape[0] hidden_size = num_heads * _HEAD_DIM qkv_size = 3 * hidden_size @@ -496,10 +530,13 @@ def _fused_kda_decode_launch( A_log, beta_layout, state_indices, + query_start_loc, + packed_state_indices, state_layout, output_gate_layout, norm_weight_layout, output_layout, + packed_t1, state_is_bf16, use_lower_bound, lower_bound, @@ -572,17 +609,28 @@ def compact(shape, dtype): raw_beta, compact((num_heads,), F32), compact((hidden_size,), F32), - compact((num_rows,), cutlass.Int32), + # This direct-ABI placeholder is also marshalled for packed T=1, + # where its runtime length is N rather than num_rows. Keep its + # extent independent so TVM-FFI does not validate an unrelated + # row-count symbol before the constexpr-dead direct argument is + # eliminated. + compact((cute.sym_int(),), cutlass.Int32), state, output_gate, compact((_HEAD_DIM,), F32), compact((1, num_rows, num_heads, _HEAD_DIM), BF16), + cute.runtime.make_fake_tensor( + cutlass.Int32, (cute.sym_int(),), (1,), assumed_align=4 + ), + cute.runtime.make_fake_tensor( + cutlass.Int32, (cute.sym_int(), 1), (1, 1), assumed_align=4 + ), cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), ) @functools.cache -def _get_compiled_kernel(state_dtype, lower_bound, norm_eps): +def _get_compiled_kernel(state_dtype, lower_bound, norm_eps, packed_t1=False): """Get a shape-dynamic specialization from the two-level CuTe DSL cache.""" state_is_bf16 = state_dtype == torch.bfloat16 compile_state_dtype = BF16 if state_is_bf16 else F32 @@ -597,14 +645,16 @@ def _get_compiled_kernel(state_dtype, lower_bound, norm_eps): kernel_name = ( f"d128_w4_{gate_name}_state{state_name}" f"_eps{str(float(norm_eps)).replace('.', '_').replace('-', 'm')}" + f"_packed{int(packed_t1)}" ) - return build_and_load_cute_dsl_kernel( + kernel = build_and_load_cute_dsl_kernel( _CUTE_DSL_MODULE, kernel_name, lambda: cute.compile( _fused_kda_decode_launch, *_make_compile_inputs(compile_state_dtype), state_is_bf16, + packed_t1, use_lower_bound, compile_lower_bound, float(norm_eps), @@ -612,6 +662,17 @@ def _get_compiled_kernel(state_dtype, lower_bound, norm_eps): ), extra_key_files=_SOURCE_FILES, ) + if packed_t1: + return kernel + + # The launch ABI always carries both metadata tensors. Direct T=1 uses + # views of state_indices as rank-compatible, unread placeholders; the + # constexpr specialization removes their device access. + def direct_kernel(*args): + state_indices = args[7] + kernel(*args, state_indices, state_indices.reshape(-1, 1)) + + return direct_kernel def _check_cuda_tensor(name, tensor, dtype): @@ -739,6 +800,8 @@ def run_fused_kda_decode( *, backend: Literal["cute-dsl", "cake", "auto"] = "cute-dsl", state_indices_mode: CakeFusedKDADecodeStateIndicesMode | None = None, + query_start_loc: torch.Tensor | None = None, + packed_state_indices: torch.Tensor | None = None, ) -> torch.Tensor: """Run fused Kimi decode and update both cache tensors in-place.""" if backend not in ("cute-dsl", "cake", "auto"): @@ -820,7 +883,36 @@ def run_fused_kda_decode( raise ValueError("A_log must be contiguous with shape [H]") if dt_bias.shape != (hidden_size,) or not dt_bias.is_contiguous(): raise ValueError("dt_bias must be contiguous with shape [H * 128]") - if state_indices.shape != (num_rows,) or not state_indices.is_contiguous(): + if (query_start_loc is None) != (packed_state_indices is None): + raise ValueError( + "query_start_loc and packed_state_indices must be provided together" + ) + packed_t1 = query_start_loc is not None + if packed_t1: + if packed_state_indices is None: + raise ValueError("packed_state_indices is required for packed T=1") + if ( + query_start_loc.ndim != 1 + or packed_state_indices.ndim != 2 + or packed_state_indices.shape[1] != 1 + or query_start_loc.shape[0] != packed_state_indices.shape[0] + 1 + or not query_start_loc.is_contiguous() + or not packed_state_indices.is_contiguous() + ): + raise ValueError("invalid packed T=1 metadata shapes or strides") + if ( + query_start_loc.dtype != torch.int32 + or packed_state_indices.dtype != torch.int32 + ): + raise TypeError("packed T=1 metadata must have dtype torch.int32") + if ( + query_start_loc.device != x.device + or packed_state_indices.device != x.device + ): + raise ValueError("packed T=1 metadata must be on the same device as x") + if state_indices.shape not in ((num_rows,), (packed_state_indices.shape[0],)): + raise ValueError("state_indices must have shape [num_rows] or [N]") + elif state_indices.shape != (num_rows,) or not state_indices.is_contiguous(): raise ValueError("state_indices must be contiguous with shape [num_rows]") if ( state.ndim != 4 @@ -869,6 +961,9 @@ def run_fused_kda_decode( "output must be contiguous with shape [1, num_rows, H, 128]" ) + if packed_t1 and backend != "cute-dsl": + raise ValueError("packed T=1 metadata requires backend='cute-dsl'") + cake_variant = None if backend != "cute-dsl" and state_indices_mode is not None: cake_variant = _select_cake_variant( @@ -907,8 +1002,10 @@ def run_fused_kda_decode( "dtype, scalar configuration, and state_indices_mode" ) - kernel = _get_compiled_kernel(state.dtype, lower_bound, float(norm_eps)) - kernel( + kernel = _get_compiled_kernel( + state.dtype, lower_bound, float(norm_eps), packed_t1=packed_t1 + ) + args = ( x, weight, conv_state, @@ -922,4 +1019,8 @@ def run_fused_kda_decode( norm_weight, output, ) + if packed_t1: + kernel(*args, query_start_loc, packed_state_indices) + else: + kernel(*args) return output diff --git a/flashinfer/kda_kernels/fused_kda_decode_multitoken.py b/flashinfer/kda_kernels/fused_kda_decode_multitoken.py new file mode 100644 index 00000000000..62bc5e2920b --- /dev/null +++ b/flashinfer/kda_kernels/fused_kda_decode_multitoken.py @@ -0,0 +1,1059 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""Fused KDA multi-token decode (conv + SiLU, gated delta-rule, gated RMSNorm). + +Architecture -- one coherent design, one GPU kernel body, one launch: + + * The D x D recurrent state of one (sequence n, head h) pair is held in + registers across the whole packed token loop, so the checkpoint pool is + read once and only the mandatory per-token checkpoints are written back. + * One CTA owns one sequence/head pair at ordinary grids. Low grids split the + head's value rows over a four- or eight-CTA cluster. Thread (rg, kq) owns + rows `rg*R ..` and k positions `jj*(4*KQ) + kq*4 + m`; state traffic stays + fully coalesced and reductions over k use log2(KQ)-step butterflies. + * Value rows are independent in the delta rule. Split ranks exchange only + their per-token sums of squares for the gated RMSNorm denominator. + * (SPLIT, R, KQ) are selected from live SM count and sequence-head geometry. + * Each split rank computes the front-end into a small shared staging buffer. + All ranks rendezvous after loading accepted convolution history, before + partitioned rolling-window stores, so source/destination aliasing is safe. +""" + +import functools +import math +from pathlib import Path + +import cutlass +import cutlass.utils +from cutlass import cute +import cuda.bindings.driver as cuda +import torch +import tvm_ffi # noqa: F401 -- TVM FFI supplies the active PyTorch stream + +from ..jit.cute_dsl_core import build_and_load_cute_dsl_kernel +from ..norm import utils as norm_utils + +_D = 128 +_W = 4 +_L2_EPS = 1.0e-6 +_LOG2E = 1.4426950408889634 # log2(e); exp(x) == exp2(x*log2e) +_CUTE_DSL_MODULE = "fused_kda_decode_multitoken" +_SOURCE_FILES = ( + str(Path(__file__).resolve()), + str(Path(norm_utils.__file__).resolve()), +) + + +@cute.jit +def _fexp(x): + # fast exp(x) via ex2.approx.ftz + return cute.math.exp2(x * _LOG2E, fastmath=True) + + +@cute.jit +def _fsigmoid(x): + # fast sigmoid = 1/(1+exp(-x)) using ex2.approx.ftz + rcp.approx + return cute.arch.rcp_approx(1.0 + cute.math.exp2((-x) * _LOG2E, fastmath=True)) + + +def _align_up(value, alignment): + return (value + alignment - 1) // alignment * alignment + + +def _required_smem_bytes(T, R, KQ, NCH, SPLIT=1): + """Return the statically allocated shared memory for one packed CTA.""" + threads = (((_D // SPLIT) // NCH) // R) * KQ + warps = threads // 32 + offset = 0 + for size, alignment in ( + (T * 3 * _D * 2, 16), + (T * _D * 4, 16), + (T * _D * 4, 16), + (T * 4 * 4, 16), + (T * warps * 4, 16), + (T * 4, 16), + (T * 4, 16), + (8, 8), + (SPLIT * T * 4, 16), + ): + offset = _align_up(offset, alignment) + size + return offset + + +def _svec(base_addr, elem_off, dtype=None): + """4 consecutive SMEM elements as an aligned tile (one wide LDS).""" + dtype = cutlass.Float32 if dtype is None else dtype + w = dtype.width // 8 + return cute.make_tensor( + cute.make_ptr( + dtype, + base_addr + elem_off * w, + cutlass.AddressSpace.smem, + assumed_align=4 * w, + ), + cute.make_layout(4), + ) + + +@cute.kernel +def _kda_kernel( + gX: cute.Tensor, + gW: cute.Tensor, + gConv: cute.Tensor, + gGate: cute.Tensor, + gBeta: cute.Tensor, + gAlog: cute.Tensor, + gDtb: cute.Tensor, + gIdx: cute.Tensor, + gQueryStart: cute.Tensor, + gAccepted: cute.Tensor, + gState: cute.Tensor, + gOG: cute.Tensor, + gNW: cute.Tensor, + gOut: cute.Tensor, + T: cutlass.Constexpr, + H: cutlass.Constexpr, + lb: cutlass.Constexpr, + eps: cutlass.Constexpr, + R: cutlass.Constexpr, + KQ: cutlass.Constexpr, + SPLIT: cutlass.Constexpr, + NCH: cutlass.Constexpr, +): + RB = _D // SPLIT + RC = RB // NCH + JJ = _D // (KQ * 4) + NT = (RC // R) * KQ + NWARP = NT // 32 + LOGKQ = int(math.log2(KQ)) + + tid, _, _ = cute.arch.thread_idx() + bx, n, _ = cute.arch.block_idx() + h = bx // SPLIT + rank = bx % SPLIT + row0 = rank * RB + + dim = H * _D + scale = 1.0 / math.sqrt(_D) + + smem = cutlass.utils.SmemAllocator() + sQKV = smem.allocate_tensor( + cutlass.BFloat16, cute.make_layout(T * 3 * _D), byte_alignment=16 + ) + sDec = smem.allocate_tensor( + cutlass.Float32, cute.make_layout(T * _D), byte_alignment=16 + ) + sOut = smem.allocate_tensor( + cutlass.Float32, cute.make_layout(T * RB), byte_alignment=16 + ) + sRed = smem.allocate_tensor( + cutlass.Float32, cute.make_layout(T * 4), byte_alignment=16 + ) + sRed2 = smem.allocate_tensor( + cutlass.Float32, cute.make_layout(T * NWARP), byte_alignment=16 + ) + sBeta = smem.allocate_tensor( + cutlass.Float32, cute.make_layout(T), byte_alignment=16 + ) + sTot = smem.allocate_tensor(cutlass.Float32, cute.make_layout(T), byte_alignment=16) + # per-token partial sums-of-squares exchanged with the cluster peers + mbar = smem.allocate_array(cutlass.Int64, 1, byte_alignment=8) + mail_ptr = smem.allocate_array(cutlass.Float32, SPLIT * T, byte_alignment=16) + sMail = cute.make_tensor(mail_ptr, cute.make_layout(SPLIT * T)) + + rg = tid // KQ # row group: owns rows rg*R .. rg*R+R-1 + kq = tid % KQ # lane group inside the row + warp = tid // 32 + lane = tid % 32 + + bos = gQueryStart[n] + eos = gQueryStart[n + 1] + seq_len = eos - bos + accepted = cute.min( + cute.max(gAccepted[n] - cutlass.Int32(1), cutlass.Int32(0)), + cutlass.Int32(T - 1), + ) + conv_slot = gIdx[n, 0] + state_slot = gIdx[n, accepted] + livef = ( + cute.min( + cute.max(conv_slot.to(cutlass.Float32), cutlass.Float32(0.0)), + cutlass.Float32(1.0), + ) + * cute.min( + cute.max(state_slot.to(cutlass.Float32), cutlass.Float32(0.0)), + cutlass.Float32(1.0), + ) + * cute.min( + cute.max(seq_len.to(cutlass.Float32), cutlass.Float32(0.0)), + cutlass.Float32(1.0), + ) + ) + src = cute.max(state_slot, cutlass.Int32(0)) + conv_src = cute.max(conv_slot, cutlass.Int32(0)) + live = livef > 0.5 + slots = [gIdx[n, t] for t in range(T)] + p_state = gState.iterator.toint() + state_page = gState.stride[0] + + if cutlass.const_expr(SPLIT > 1): + # Arm the cluster mailbox up front: at kernel entry the peers are still + # in lockstep, so this rendezvous is nearly free, whereas the same + # rendezvous placed at the reduction would sit on the critical path. + if tid == 0: + cute.arch.mbarrier_init(mbar, 1) + cute.arch.mbarrier_init_fence() + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + # When the CTA's rows fit in a single register chunk there is no chunk loop + # to overlap global latency with, so the two loads that are only consumed + # at the very end (output gate, norm weight) and the checkpoint read itself + # are issued up front, before the front end. + NEP = (T * RB + NT - 1) // NT + if cutlass.const_expr(NCH == 1): + ogv = [] + nwv = [] + for blk in cutlass.range_constexpr(NEP): + pidx = cute.min(blk * NT + tid, cutlass.Int32(T * RB - 1)) + pvv = row0 + pidx % RB + og = cutlass.Float32(0.0) + if seq_len > 0: + token = cute.min(pidx // RB, seq_len - 1) + og = _fsigmoid(gOG[bos + token, h * _D + pvv].to(cutlass.Float32)) + ogv.append(og) + nwv.append(gNW[pvv]) + + src_base = p_state + src.to(cutlass.Int64) * (state_page * 4) + if cutlass.const_expr(NCH == 1): + sfrg0 = [ + cute.make_rmem_tensor(cute.make_layout((4, JJ)), cutlass.Float32) + for _ in range(R) + ] + lane_off0 = [ + ((h * _D + row0 + rg * R + r) * _D + kq * 4).to(cutlass.Int64) * 4 + for r in range(R) + ] + for r in cutlass.range_constexpr(R): + cute.autovec_copy( + cute.make_tensor( + cute.make_ptr( + cutlass.Float32, + src_base + lane_off0[r], + cutlass.AddressSpace.gmem, + assumed_align=16, + ), + cute.make_layout((4, JJ), stride=(1, KQ * 4)), + ), + sfrg0[r], + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.EVICT_FIRST, + ) + + # ========== stage 1: causal conv + SiLU, decay, beta, checkpoints ======== + # 4*D "channel tasks": 3*D convolution channels then D decay channels. + for cb in cutlass.range_constexpr((4 * _D + NT - 1) // NT): + c = cb * NT + tid + ci = cutlass.Int32(0) + cd = cutlass.Int32(0) + chan = cutlass.Int32(0) + hh0 = cutlass.Float32(0.0) + hh1 = cutlass.Float32(0.0) + hh2 = cutlass.Float32(0.0) + tp0 = cutlass.Float32(0.0) + tp1 = cutlass.Float32(0.0) + tp2 = cutlass.Float32(0.0) + tp3 = cutlass.Float32(0.0) + if c < 3 * _D: + ci = c // _D + cd = c % _D + chan = ci * dim + h * _D + cd + hh0 = gConv[conv_src, chan, accepted].to(cutlass.Float32) + hh1 = gConv[conv_src, chan, accepted + 1].to(cutlass.Float32) + hh2 = gConv[conv_src, chan, accepted + 2].to(cutlass.Float32) + + tp0 = gW[ci, 0, h * _D + cd] + tp1 = gW[ci, 1, h * _D + cd] + tp2 = gW[ci, 2, h * _D + cd] + tp3 = gW[ci, 3, h * _D + cd] + + if cutlass.const_expr(SPLIT > 1): + # All cluster ranks read the shared convolution source. No rank may + # overwrite the rolling window until every peer has captured its + # accepted history in registers. + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + + if c < 3 * _D: + # the conv history bytes are partitioned over the cluster ranks + mine = ((c * SPLIT) // (3 * _D)) == rank + + if live: + if mine: + gConv[conv_src, chan, 0] = hh1.to(cutlass.BFloat16) + gConv[conv_src, chan, 1] = hh2.to(cutlass.BFloat16) + + for t in cutlass.range_constexpr(T): + if t < seq_len: + xv = gX[bos + t, chan].to(cutlass.Float32) + mixed = tp0 * hh0 + tp1 * hh1 + tp2 * hh2 + tp3 * xv + mixed = mixed * _fsigmoid(mixed) + sQKV[t * (3 * _D) + c] = mixed.to(cutlass.BFloat16) + hh0 = hh1 + hh1 = hh2 + hh2 = xv + + if live: + if mine: + gConv[conv_src, chan, t + 2] = xv.to(cutlass.BFloat16) + elif c < 4 * _D: + dd = c - 3 * _D + a_h = _fexp(gAlog[h]) + dtb = gDtb[h * _D + dd] + for t in cutlass.range_constexpr(T): + if t < seq_len: + gi = gGate[bos + t, h * _D + dd].to(cutlass.Float32) + dtb + sDec[t * _D + dd] = _fexp(lb * _fsigmoid(a_h * gi)) + + if tid < T: + if tid < seq_len: + sBeta[tid] = _fsigmoid(gBeta[bos + tid, h].to(cutlass.Float32)) + + cute.arch.barrier() + + # ---- L2 norms of q, k and their cross term: one warp per token, each + # ---- lane owning 4 channels, so the whole 128-wide reduction is a single + # ---- butterfly and the token loop below reads 3 ready-made scalars + # ---- instead of re-reducing 12 partials per token. + a_qkv = sQKV.iterator.toint() + for nb in cutlass.range_constexpr((T * 32 + NT - 1) // NT): + idx = nb * NT + tid + if idx < T * 32: + t = idx // 32 + lane_idx = idx % 32 + if t < seq_len: + nq4 = cute.make_rmem_tensor(cute.make_layout(4), cutlass.BFloat16) + nk4 = cute.make_rmem_tensor(cute.make_layout(4), cutlass.BFloat16) + cute.autovec_copy( + _svec(a_qkv, t * (3 * _D) + lane_idx * 4, cutlass.BFloat16), + nq4, + ) + cute.autovec_copy( + _svec( + a_qkv, + t * (3 * _D) + _D + lane_idx * 4, + cutlass.BFloat16, + ), + nk4, + ) + sqq = cutlass.Float32(0.0) + skk = cutlass.Float32(0.0) + skq = cutlass.Float32(0.0) + for m in cutlass.range_constexpr(4): + qv = nq4[m].to(cutlass.Float32) + kv = nk4[m].to(cutlass.Float32) + sqq = sqq + qv * qv + skk = skk + kv * kv + skq = skq + kv * qv + for off in cutlass.range_constexpr(5): + sqq = sqq + cute.arch.shuffle_sync_bfly(sqq, 1 << off) + skk = skk + cute.arch.shuffle_sync_bfly(skk, 1 << off) + skq = skq + cute.arch.shuffle_sync_bfly(skq, 1 << off) + if lane_idx == 0: + sRed[t * 4] = ( + cute.rsqrt(sqq + _L2_EPS, approx=True, ftz=True) * scale + ) + sRed[t * 4 + 1] = cute.rsqrt(skk + _L2_EPS, approx=True, ftz=True) + sRed[t * 4 + 2] = skq + cute.arch.barrier() + + # ========== stage 2: register-resident gated delta-rule recurrence ======= + # The RB rows this CTA owns are walked in NCH register chunks of RC rows. + # NCH trades register residency for occupancy: the chunk state is what sits + # in registers, so raising NCH shrinks the per-thread tile and lets more + # CTAs (more memory-level parallelism) sit on an SM at once. Rows are + # independent under the delta rule, so the chunk order is immaterial. + a_dec = sDec.iterator.toint() + rtot = [cutlass.Float32(0.0) for _ in range(T)] + + for ch in cutlass.range_constexpr(NCH): + if cutlass.const_expr(NCH == 1): + sfrg = sfrg0 + lane_off = lane_off0 + else: + sfrg = [ + cute.make_rmem_tensor(cute.make_layout((4, JJ)), cutlass.Float32) + for _ in range(R) + ] + lane_off = [ + ((h * _D + row0 + ch * RC + rg * R + r) * _D + kq * 4).to(cutlass.Int64) + * 4 + for r in range(R) + ] + for r in cutlass.range_constexpr(R): + gtile = cute.make_tensor( + cute.make_ptr( + cutlass.Float32, + src_base + lane_off[r], + cutlass.AddressSpace.gmem, + assumed_align=16, + ), + cute.make_layout((4, JJ), stride=(1, KQ * 4)), + ) + cute.autovec_copy( + gtile, + sfrg[r], + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.EVICT_FIRST, + ) + + for t in cutlass.range_constexpr(T): + if t < seq_len: + qb = t * (3 * _D) + db = t * _D + rn_q = sRed[t * 4] + rn_k = sRed[t * 4 + 1] + bkq = sRed[t * 4 + 2] + bt = sBeta[t] + + # One pass over the state: decay it and contract it with BOTH the + # raw key and the raw query. Because the update is s + d*K, the + # recurrent readout is + d*, so a single reduction per + # token suffices. + acc = [cutlass.Float32(0.0) for _ in range(R)] + aq = [cutlass.Float32(0.0) for _ in range(R)] + kreg = [ + cute.make_rmem_tensor(cute.make_layout(4), cutlass.BFloat16) + for _ in range(JJ) + ] + for jj in cutlass.range_constexpr(JJ): + k0 = jj * (KQ * 4) + kq * 4 + d4 = cute.make_rmem_tensor(cute.make_layout(4), cutlass.Float32) + q4 = cute.make_rmem_tensor(cute.make_layout(4), cutlass.BFloat16) + cute.autovec_copy(_svec(a_dec, db + k0), d4) + cute.autovec_copy( + _svec(a_qkv, qb + _D + k0, cutlass.BFloat16), kreg[jj] + ) + cute.autovec_copy(_svec(a_qkv, qb + k0, cutlass.BFloat16), q4) + for m in cutlass.range_constexpr(4): + dc = d4[m] + kv = kreg[jj][m].to(cutlass.Float32) + qv = q4[m].to(cutlass.Float32) + for r in cutlass.range_constexpr(R): + s = sfrg[r][m, jj] * dc + sfrg[r][m, jj] = s + acc[r] = acc[r] + s * kv + aq[r] = aq[r] + s * qv + + d2 = [] + rec = [] + for r in cutlass.range_constexpr(R): + a = acc[r] + b = aq[r] + for off in cutlass.range_constexpr(LOGKQ): + a = a + cute.arch.shuffle_sync_bfly(a, 1 << off) + b = b + cute.arch.shuffle_sync_bfly(b, 1 << off) + dd2 = ( + ( + sQKV[qb + 2 * _D + row0 + ch * RC + rg * R + r].to( + cutlass.Float32 + ) + - a * rn_k + ) + * bt + * rn_k + ) + d2.append(dd2) + rec.append((b + dd2 * bkq) * rn_q) + + for jj in cutlass.range_constexpr(JJ): + for m in cutlass.range_constexpr(4): + kv = kreg[jj][m].to(cutlass.Float32) + for r in cutlass.range_constexpr(R): + sfrg[r][m, jj] = sfrg[r][m, jj] + d2[r] * kv + + rsum = cutlass.Float32(0.0) + for r in cutlass.range_constexpr(R): + # Match the composed vLLM path and the legacy fused kernel: + # recurrent output is materialized as BF16 before RMSNorm. + x = rec[r].to(cutlass.BFloat16).to(cutlass.Float32) * livef + rsum = rsum + x * x + if kq == 0: + sOut[t * RB + ch * RC + rg * R + r] = x + # rsum is uniform inside each KQ-lane group, so only the groups of + # the warp still have to be combined to get this warp's row sum. + for off in cutlass.range_constexpr(LOGKQ, 5): + rsum = rsum + cute.arch.shuffle_sync_bfly(rsum, 1 << off) + rtot[t] = rtot[t] + rsum + + if live: + if slots[t] > 0: + dst_base = p_state + slots[t].to(cutlass.Int64) * ( + state_page * 4 + ) + for r in cutlass.range_constexpr(R): + gwtile = cute.make_tensor( + cute.make_ptr( + cutlass.Float32, + dst_base + lane_off[r], + cutlass.AddressSpace.gmem, + assumed_align=16, + ), + cute.make_layout((4, JJ), stride=(1, KQ * 4)), + ) + cute.autovec_copy( + sfrg[r], + gwtile, + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.EVICT_FIRST, + ) + + if lane == 0: + for t in cutlass.range_constexpr(T): + sRed2[t * NWARP + warp] = rtot[t] + + cute.arch.barrier() + + # ---- collapse the per-warp partials to one per token, then (when the + # ---- head is spread over a cluster) add in the peer ranks' partials + # ---- straight out of distributed shared memory. + if cutlass.const_expr(SPLIT > 1): + if tid == 0: + cute.arch.mbarrier_arrive_and_expect_tx(mbar, SPLIT * T * 4) + cute.arch.barrier() + if tid < T: + tot = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(NWARP): + tot = tot + sRed2[tid * NWARP + w] + slot = mail_ptr + (rank * T + tid) + for p in cutlass.range_constexpr(SPLIT): + norm_utils.store_shared_remote(tot, slot, mbar, cutlass.Int32(p)) + cute.arch.fence_acq_rel_cta() + if tid == 0: + cute.arch.mbarrier_wait(mbar, 0) + cute.arch.barrier() + if tid < T: + tot = cutlass.Float32(0.0) + for p in cutlass.range_constexpr(SPLIT): + tot = tot + sMail[p * T + tid] + sTot[tid] = cute.rsqrt(tot * (1.0 / _D) + eps, approx=True, ftz=True) + cute.arch.barrier() + else: + if tid < T: + tot = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(NWARP): + tot = tot + sRed2[tid * NWARP + w] + sTot[tid] = cute.rsqrt(tot * (1.0 / _D) + eps, approx=True, ftz=True) + cute.arch.barrier() + + # ========== stage 3: gated RMSNorm epilogue ============================= + # one thread per (token, row) so the bf16 output store is fully coalesced + if cutlass.const_expr(NCH == 1): + for blk in cutlass.range_constexpr(NEP): + idx = blk * NT + tid + if idx < T * RB: + token = idx // RB + if token < seq_len: + gOut[bos + token, h * _D + row0 + idx % RB] = ( + sOut[idx] * sTot[token] * nwv[blk] * ogv[blk] + ).to(cutlass.BFloat16) + else: + for blk in cutlass.range_constexpr(NEP): + idx = blk * NT + tid + if idx < T * RB: + ete = idx // RB + evv = row0 + idx % RB + if ete < seq_len: + eog = _fsigmoid(gOG[bos + ete, h * _D + evv].to(cutlass.Float32)) + gOut[bos + ete, h * _D + evv] = ( + sOut[idx] * sTot[ete] * gNW[evv] * eog + ).to(cutlass.BFloat16) + + +@cute.jit +def _kda_launch( + m_x, + m_w, + m_conv, + m_gate, + m_beta, + m_alog, + m_dtb, + m_idx, + m_query_start, + m_accepted, + m_state, + m_og, + m_nw, + m_out, + x_row_stride: cutlass.Int64, + beta_row_stride: cutlass.Int64, + output_gate_row_stride: cutlass.Int64, + stream: cuda.CUstream, + T: cutlass.Constexpr, + H: cutlass.Constexpr, + lb: cutlass.Constexpr, + eps: cutlass.Constexpr, + R: cutlass.Constexpr, + KQ: cutlass.Constexpr, + bpm: cutlass.Constexpr, + SPLIT: cutlass.Constexpr, + NCH: cutlass.Constexpr, +): + p_x = m_x.iterator.toint() + p_w = m_w.iterator.toint() + p_gate = m_gate.iterator.toint() + p_beta = m_beta.iterator.toint() + p_alog = m_alog.iterator.toint() + p_dtb = m_dtb.iterator.toint() + p_idx = m_idx.iterator.toint() + p_og = m_og.iterator.toint() + p_nw = m_nw.iterator.toint() + p_out = m_out.iterator.toint() + dim = H * _D + BF = cutlass.BFloat16 + F32 = cutlass.Float32 + I32 = cutlass.Int32 + GM = cutlass.AddressSpace.gmem + N = m_idx.shape[0] + num_rows = m_x.shape[0] + + def _t(dtype, addr, shape, stride): + return cute.make_tensor( + cute.make_ptr(dtype, addr, GM, assumed_align=16), + cute.make_layout(shape, stride=stride), + ) + + gX = _t(BF, p_x, (num_rows, 3 * dim), (x_row_stride, 1)) + gW = _t(F32, p_w, (3, _W, dim), (_W * dim, dim, 1)) + gConv = m_conv + gGate = _t(BF, p_gate, (num_rows, dim), (dim, 1)) + gBeta = _t(BF, p_beta, (num_rows, H), (beta_row_stride, 1)) + gAlog = _t(F32, p_alog, (H,), (1,)) + gDtb = _t(F32, p_dtb, (dim,), (1,)) + gIdx = _t(I32, p_idx, (N, T), (T, 1)) + gQueryStart = m_query_start + gAccepted = m_accepted + gOG = _t(BF, p_og, (num_rows, dim), (output_gate_row_stride, 1)) + gNW = _t(F32, p_nw, (_D,), (1,)) + gOut = _t(BF, p_out, (num_rows, dim), (dim, 1)) + + NTHR = (((_D // SPLIT) // NCH) // R) * KQ + kernel = _kda_kernel( + gX, + gW, + gConv, + gGate, + gBeta, + gAlog, + gDtb, + gIdx, + gQueryStart, + gAccepted, + m_state, + gOG, + gNW, + gOut, + T, + H, + lb, + eps, + R, + KQ, + SPLIT, + NCH, + ) + if cutlass.const_expr(SPLIT > 1): + kernel.launch( + grid=(H * SPLIT, N, 1), + block=(NTHR, 1, 1), + cluster=(SPLIT, 1, 1), + stream=stream, + ) + else: + kernel.launch( + grid=(H, N, 1), + block=(NTHR, 1, 1), + min_blocks_per_mp=bpm, + stream=stream, + ) + + +# (rows per thread, k-lane groups, min CTAs per SM, row chunks per CTA). +# NCH walks the CTA's rows in register chunks: raising it shrinks the +# per-thread tile and buys occupancy (memory-level parallelism on the +# checkpoint stream) at the price of re-reading the staged decay/key/query +# once per chunk and token, so the chunk count is bounded by the token loop. +_TILE_TALL = (2, 16, 1, 1) +_TILE_T1 = (2, 16, 4, 4) +_TILE_WIDE_SHORT = (1, 16, 4, 4) +_TILE_WIDE_LONG = (2, 16, 3, 2) +# Split tiles must keep NTHR == 4*D. The pre-store cluster rendezvous is safe +# because every convolution-history channel is loaded in that single pass. +_TILE_SPLIT = { + 4: (1, 16, 1, 1), + 8: (1, 32, 1, 1), +} + + +def _pick_split(sequence_heads, sm_count): + if sequence_heads * 8 <= sm_count: + return 8 + if sequence_heads * 4 <= sm_count: + return 4 + return 1 + + +def _compact_fake(shape, dtype): + return cute.runtime.make_fake_compact_tensor( + dtype, + shape, + assumed_align=16, + stride_order=tuple(reversed(range(len(shape)))), + ) + + +@functools.cache +def _get_compiled_kernel( + T, + H, + lower_bound, + norm_eps, + R, + KQ, + bpm, + NCH, + SPLIT, +): + dim = H * _D + sequences = cute.sym_int() + boundaries = cute.sym_int() + rows = cute.sym_int() + slots = cute.sym_int() + fake_inputs = ( + cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (rows, 3 * dim), + (cute.sym_int64(), 1), + assumed_align=16, + ), + _compact_fake((3, _W, dim), cutlass.Float32), + cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (slots, 3 * dim, T + _W - 2), + (cute.sym_int64(divisibility=16), 1, 3 * dim), + assumed_align=16, + ), + _compact_fake((1, rows, H, _D), cutlass.BFloat16), + cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (1, rows, H), + (cute.sym_int64(), cute.sym_int64(), 1), + assumed_align=16, + ), + _compact_fake((H,), cutlass.Float32), + _compact_fake((dim,), cutlass.Float32), + _compact_fake((sequences, T), cutlass.Int32), + _compact_fake((boundaries,), cutlass.Int32), + _compact_fake((sequences,), cutlass.Int32), + cute.runtime.make_fake_tensor( + cutlass.Float32, + (slots, H, _D, _D), + (cute.sym_int64(divisibility=16), _D * _D, _D, 1), + assumed_align=16, + ), + cute.runtime.make_fake_tensor( + cutlass.BFloat16, + (rows, H, _D), + (cute.sym_int64(), _D, 1), + assumed_align=16, + ), + _compact_fake((_D,), cutlass.Float32), + _compact_fake((1, rows, H, _D), cutlass.BFloat16), + ) + kernel_name = ( + f"d128_w4_h{H}_t{T}" + f"_lb{str(lower_bound).replace('.', '_').replace('-', 'm')}" + f"_eps{str(norm_eps).replace('.', '_').replace('-', 'm')}" + f"_r{R}_kq{KQ}_bpm{bpm}_nch{NCH}_split{SPLIT}" + ) + return build_and_load_cute_dsl_kernel( + _CUTE_DSL_MODULE, + kernel_name, + lambda: cute.compile( + _kda_launch, + *fake_inputs, + 1, + 1, + 1, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + T, + H, + lower_bound, + norm_eps, + R, + KQ, + bpm, + SPLIT, + NCH, + options="--enable-tvm-ffi --generate-line-info", + ), + extra_key_files=_SOURCE_FILES, + ) + + +def _run_fused_kda_decode_multitoken( + x, + weight, + conv_state, + raw_gate, + raw_beta, + A_log, + dt_bias, + state_indices, + query_start_loc, + num_accepted_tokens, + state, + output_gate, + norm_weight, + lower_bound, + norm_eps, + out, +): + key = ( + state_indices.shape[1], + A_log.shape[0], + float(lower_bound), + float(norm_eps), + ) + T, H, lower_bound, norm_eps = key + N = state_indices.shape[0] + sm_count = torch.cuda.get_device_properties(x.device).multi_processor_count + SPLIT = _pick_split(N * H, sm_count) + if T == 1: + SPLIT = 1 + R, KQ, bpm, NCH = _TILE_T1 + elif SPLIT > 1: + R, KQ, bpm, NCH = _TILE_SPLIT[SPLIT] + elif 2 * sm_count > N * H: + R, KQ, bpm, NCH = _TILE_TALL + elif T <= 3: + R, KQ, bpm, NCH = _TILE_WIDE_SHORT + else: + R, KQ, bpm, NCH = _TILE_WIDE_LONG + properties = torch.cuda.get_device_properties(x.device) + required_smem = _required_smem_bytes(T, R, KQ, NCH, SPLIT) + if required_smem > properties.shared_memory_per_block_optin: + raise ValueError( + f"T={T} requires {required_smem} bytes of shared memory, but " + f"{x.device} supports {properties.shared_memory_per_block_optin} bytes per block" + ) + entry = _get_compiled_kernel( + T, + H, + lower_bound, + norm_eps, + R, + KQ, + bpm, + NCH, + SPLIT, + ) + + entry( + x, + weight, + conv_state, + raw_gate, + raw_beta, + A_log, + dt_bias, + state_indices, + query_start_loc, + num_accepted_tokens, + state, + output_gate, + norm_weight, + out, + x.stride(0), + raw_beta.stride(1), + output_gate.stride(0), + ) + + +def _check_cuda_tensor(name, tensor, dtype): + if not isinstance(tensor, torch.Tensor) or not tensor.is_cuda: + raise ValueError(f"{name} must be a CUDA tensor") + if tensor.dtype != dtype: + raise TypeError(f"{name} must have dtype {dtype}, got {tensor.dtype}") + + +@torch.no_grad() +def run_fused_kda_decode_multitoken( + x, + weight, + conv_state, + raw_gate, + raw_beta, + A_log, + dt_bias, + state_indices, + query_start_loc, + num_accepted_tokens, + state, + output_gate, + norm_weight, + lower_bound=-5.0, + norm_eps=1e-5, + output=None, +): + """Run packed fused KDA and write one recurrent checkpoint per token. + + ``query_start_loc`` supplies each sequence's packed token range and + ``num_accepted_tokens`` selects the source checkpoint and convolution + history offset from the preceding verification window. Zero-length or + non-positive rows do not mutate the cache. Destination stores are guarded + per token as a final safety net. + """ + _check_cuda_tensor("x", x, torch.bfloat16) + _check_cuda_tensor("weight", weight, torch.float32) + _check_cuda_tensor("conv_state", conv_state, torch.bfloat16) + _check_cuda_tensor("raw_gate", raw_gate, torch.bfloat16) + _check_cuda_tensor("raw_beta", raw_beta, torch.bfloat16) + _check_cuda_tensor("A_log", A_log, torch.float32) + _check_cuda_tensor("dt_bias", dt_bias, torch.float32) + _check_cuda_tensor("state_indices", state_indices, torch.int32) + _check_cuda_tensor("query_start_loc", query_start_loc, torch.int32) + _check_cuda_tensor("num_accepted_tokens", num_accepted_tokens, torch.int32) + _check_cuda_tensor("state", state, torch.float32) + _check_cuda_tensor("output_gate", output_gate, torch.bfloat16) + _check_cuda_tensor("norm_weight", norm_weight, torch.float32) + + major, _ = torch.cuda.get_device_capability(x.device) + if major != 10: + raise NotImplementedError("packed fused KDA T>1 currently requires SM10x") + + if state_indices.ndim != 2 or not state_indices.is_contiguous(): + raise ValueError("state_indices must be contiguous with shape [N, T]") + num_sequences, num_tokens = state_indices.shape + if num_sequences <= 0 or num_tokens <= 0: + raise ValueError("state_indices must contain at least one sequence and token") + if x.ndim != 2 or x.shape[0] > num_sequences * num_tokens: + raise ValueError("x must have at most N * T packed rows") + if x.shape[1] % (3 * _D) or x.stride(1) != 1: + raise ValueError("x must have a contiguous channel dimension of 3 * H * 128") + num_heads = x.shape[1] // (3 * _D) + hidden_size = num_heads * _D + if num_heads not in (12, 24, 32, 48, 96): + raise ValueError("H must be one of 12, 24, 32, 48, or 96") + if weight.shape != (3, _W, hidden_size) or not weight.is_contiguous(): + raise ValueError("weight must be contiguous with shape [3, 4, H * 128]") + if ( + conv_state.ndim != 3 + or conv_state.shape[1:] != (3 * hidden_size, num_tokens + _W - 2) + or conv_state.stride(1) != 1 + or conv_state.stride(2) != 3 * hidden_size + ): + raise ValueError("conv_state must use the extended paged SD cache layout") + num_rows = x.shape[0] + if raw_gate.shape != (1, num_rows, num_heads, _D): + raise ValueError("raw_gate must have shape [1, num_rows, H, 128]") + if not raw_gate.is_contiguous(): + raise ValueError("raw_gate must be contiguous") + if raw_beta.shape != (1, num_rows, num_heads) or raw_beta.stride(2) != 1: + raise ValueError("raw_beta must have shape [1, num_rows, H]") + if ( + query_start_loc.shape != (num_sequences + 1,) + or not query_start_loc.is_contiguous() + ): + raise ValueError("query_start_loc must be contiguous with shape [N + 1]") + if ( + num_accepted_tokens.shape != (num_sequences,) + or not num_accepted_tokens.is_contiguous() + ): + raise ValueError("num_accepted_tokens must be contiguous with shape [N]") + if A_log.shape != (num_heads,) or not A_log.is_contiguous(): + raise ValueError("A_log must be contiguous with shape [H]") + if dt_bias.shape != (hidden_size,) or not dt_bias.is_contiguous(): + raise ValueError("dt_bias must be contiguous with shape [H * 128]") + if ( + state.ndim != 4 + or state.shape[1:] != (num_heads, _D, _D) + or state.stride()[1:] != (_D * _D, _D, 1) + ): + raise ValueError("state must have shape [slots, H, 128, 128]") + if conv_state.shape[0] != state.shape[0]: + raise ValueError("conv_state and state must have matching slot counts") + if output_gate.ndim == 4: + if output_gate.shape[0] != 1: + raise ValueError("4D output_gate must have leading dimension one") + output_gate = output_gate[0] + if output_gate.shape != (num_rows, num_heads, _D): + raise ValueError("output_gate must have shape [num_rows, H, 128]") + if output_gate.stride(2) != 1 or output_gate.stride(1) != _D: + raise ValueError("output_gate must be contiguous within each token") + if norm_weight.shape != (_D,) or not norm_weight.is_contiguous(): + raise ValueError("norm_weight must be contiguous with shape [128]") + if lower_bound is None or not math.isfinite(lower_bound) or lower_bound >= 0: + raise ValueError("lower_bound must be a finite negative float") + if not math.isfinite(norm_eps) or norm_eps < 0: + raise ValueError("norm_eps must be finite and non-negative") + + tensors = ( + weight, + conv_state, + raw_gate, + raw_beta, + A_log, + dt_bias, + state_indices, + query_start_loc, + num_accepted_tokens, + state, + output_gate, + norm_weight, + ) + if any(tensor.device != x.device for tensor in tensors): + raise ValueError("all inputs must be on the same device") + expected_output_shape = (1, num_rows, num_heads, _D) + if output is None: + output = torch.empty( + expected_output_shape, dtype=torch.bfloat16, device=x.device + ) + else: + _check_cuda_tensor("output", output, torch.bfloat16) + if output.shape != expected_output_shape or not output.is_contiguous(): + raise ValueError( + "output must be contiguous with shape [1, num_rows, H, 128]" + ) + + _run_fused_kda_decode_multitoken( + x, + weight, + conv_state, + raw_gate, + raw_beta, + A_log, + dt_bias, + state_indices, + query_start_loc, + num_accepted_tokens, + state, + output_gate, + norm_weight, + float(lower_bound), + float(norm_eps), + output, + ) + return output diff --git a/flashinfer/trace/templates/kda.py b/flashinfer/trace/templates/kda.py index eec73943881..2e78aca148a 100644 --- a/flashinfer/trace/templates/kda.py +++ b/flashinfer/trace/templates/kda.py @@ -283,3 +283,69 @@ def _recurrent_kda_template( ], tags=["stage:decode", "status:verified"], ) + + +packed_fused_kda_decode_trace = TraceTemplate( + op_type="kda", + name_prefix="packed_fused_kda_decode", + description=( + "Packed ragged T>=1 Kimi width-four causal convolution, recurrent KDA, " + "and gated RMSNorm with one recurrent checkpoint per token." + ), + axes={ + "num_sequences": Var(description="Number of independent sequences."), + "num_tokens": Const(description="Maximum tokens per sequence.", abbrev="t"), + "num_rows": Var(description="Total packed tokens."), + "num_sequence_offsets": Var( + description="Number of packed sequence boundaries." + ), + "num_heads": Const(description="Number of KDA heads.", abbrev="h"), + "head_dim": Const(description="KDA head dimension.", abbrev="d"), + "singleton": Const(description="Leading singleton dimension.", abbrev=""), + "projection_groups": Const( + description="Packed QKV projection groups.", abbrev="" + ), + "hidden_size": Const(description="Channels in one projection.", abbrev=""), + "qkv_width": Const(description="Packed QKV width.", abbrev=""), + "conv_width": Const(description="Depthwise convolution width.", abbrev=""), + "conv_history": Const(description="Cached convolution history.", abbrev=""), + "num_slots": Var(description="Number of cache slots."), + }, + inputs={ + "x": Tensor(["num_rows", "qkv_width"]), + "weight": Tensor(["projection_groups", "conv_width", "hidden_size"]), + "conv_state": Tensor(["num_slots", "qkv_width", "conv_history"]), + "raw_gate": Tensor(["singleton", "num_rows", "num_heads", "head_dim"]), + "raw_beta": Tensor(["singleton", "num_rows", "num_heads"]), + "A_log": Tensor(["num_heads"]), + "dt_bias": Tensor(["hidden_size"]), + "state_indices": Tensor(["num_sequences", "num_tokens"]), + "query_start_loc": Tensor(["num_sequence_offsets"]), + "num_accepted_tokens": Tensor(["num_sequences"]), + "t1_state_indices": Tensor(["num_rows"], optional=True), + "state": Tensor(["num_slots", "num_heads", "head_dim", "head_dim"]), + "output_gate": Tensor(["num_rows", "num_heads", "head_dim"]), + "norm_weight": Tensor(["head_dim"]), + "lower_bound": Scalar("float32", optional=True), + "norm_eps": Scalar("float32", optional=True), + }, + outputs={ + "output": Tensor( + ["singleton", "num_rows", "num_heads", "head_dim"], dtype_from="x" + ), + }, + constraints=[ + "num_rows <= num_sequences * num_tokens", + "num_sequence_offsets == num_sequences + 1", + "num_tokens >= 1", + "qkv_width == 3 * num_heads * head_dim", + "hidden_size == num_heads * head_dim", + "singleton == 1", + "projection_groups == 3", + "head_dim == 128", + "num_heads in (12, 24, 32, 48, 96)", + "conv_width == 4", + "conv_history == num_tokens + 2", + ], + tags=["stage:decode", "status:experimental"], +) diff --git a/tests/kda/test_cake_fused_kda_decode_dispatch.py b/tests/kda/test_cake_fused_kda_decode_dispatch.py index 10f944406a5..42064aa7d22 100644 --- a/tests/kda/test_cake_fused_kda_decode_dispatch.py +++ b/tests/kda/test_cake_fused_kda_decode_dispatch.py @@ -425,7 +425,7 @@ def _patch_validated_run(monkeypatch, cake_variant): monkeypatch.setattr( fused, "_get_compiled_kernel", - lambda *args: lambda *kernel_args: fallback_calls.append(kernel_args), + lambda *args, **kwargs: lambda *kernel_args: fallback_calls.append(kernel_args), ) return cake_calls, fallback_calls diff --git a/tests/kda/test_fused_kda_decode.py b/tests/kda/test_fused_kda_decode.py index aa23b5bd375..23e1d747369 100644 --- a/tests/kda/test_fused_kda_decode.py +++ b/tests/kda/test_fused_kda_decode.py @@ -22,11 +22,13 @@ from flashinfer.kda_decode import ( _FUSED_KDA_DECODE_AVAILABLE, fused_kda_decode, + packed_fused_kda_decode, ) _has_fused_kda_decode = _FUSED_KDA_DECODE_AVAILABLE except ImportError: fused_kda_decode = None + packed_fused_kda_decode = None _has_fused_kda_decode = False @@ -65,15 +67,17 @@ def _require_fused_kda_decode(): pytest.skip("fused_kda_decode is unavailable (missing CuTe DSL dependencies)") -def _page_strides(num_heads, state_dtype): - conv_slot_bytes = 3 * num_heads * 128 * 3 * 2 +def _page_strides(num_heads, state_dtype, conv_history=3): + conv_slot_bytes = 3 * num_heads * 128 * conv_history * 2 state_element_bytes = torch.empty((), dtype=state_dtype).element_size() state_slot_bytes = num_heads * 128 * 128 * state_element_bytes page_bytes = conv_slot_bytes + state_slot_bytes return page_bytes // 2, page_bytes // state_element_bytes -def _make_inputs(num_heads, num_rows, seed=42, state_dtype=torch.float32): +def _make_inputs( + num_heads, num_rows, seed=42, state_dtype=torch.float32, conv_history=3 +): device = torch.device("cuda") num_slots = num_rows + 1 hidden_size = num_heads * 128 @@ -88,14 +92,18 @@ def randn(shape, dtype=torch.float32): x = x_storage[:, : 3 * hidden_size] weight = 0.1 * randn((3, 4, hidden_size)) - conv_slot_stride, state_slot_stride = _page_strides(num_heads, state_dtype) + conv_slot_stride, state_slot_stride = _page_strides( + num_heads, state_dtype, conv_history + ) conv_state = torch.empty_strided( - (num_slots, 3 * hidden_size, 3), + (num_slots, 3 * hidden_size, conv_history), (conv_slot_stride, 1, 3 * hidden_size), dtype=torch.bfloat16, device=device, ) - conv_state.copy_(0.1 * randn((num_slots, 3 * hidden_size, 3), torch.bfloat16)) + conv_state.copy_( + 0.1 * randn((num_slots, 3 * hidden_size, conv_history), torch.bfloat16) + ) state = torch.empty_strided( (num_slots, num_heads, 128, 128), (state_slot_stride, 128 * 128, 128, 1), @@ -125,6 +133,16 @@ def randn(shape, dtype=torch.float32): } +def _fused_kda_decode(**kwargs): + assert fused_kda_decode is not None + return fused_kda_decode(**kwargs) + + +def _packed_fused_kda_decode(**kwargs): + assert packed_fused_kda_decode is not None + return packed_fused_kda_decode(**kwargs) + + def _clone_strided(tensor): clone = torch.empty_strided( tensor.shape, tensor.stride(), dtype=tensor.dtype, device=tensor.device @@ -134,7 +152,7 @@ def _clone_strided(tensor): @torch.no_grad() -def _reference(inputs, conv_state, state, lower_bound=-5.0): +def _reference(inputs, conv_state, state, lower_bound: float | None = -5.0): x = inputs["x"] num_rows = x.shape[0] num_heads = inputs["A_log"].numel() @@ -178,6 +196,94 @@ def _reference(inputs, conv_state, state, lower_bound=-5.0): return output.unsqueeze(0).to(torch.bfloat16) +@torch.no_grad() +def _packed_reference(inputs, conv_state, state): + indices = inputs["state_indices"] + query_start_loc = inputs["query_start_loc"].cpu().tolist() + accepted_tokens = inputs["num_accepted_tokens"].cpu().tolist() + num_sequences, max_tokens = indices.shape + num_heads = inputs["A_log"].numel() + hidden_size = num_heads * 128 + taps = inputs["weight"].permute(0, 2, 1).reshape(3 * hidden_size, 4).float() + output = torch.zeros( + (1, inputs["x"].shape[0], num_heads, 128), + dtype=torch.bfloat16, + device=state.device, + ) + decay_a = inputs["A_log"].exp().view(num_heads, 1) + + for sequence in range(num_sequences): + begin, end = query_start_loc[sequence : sequence + 2] + sequence_length = end - begin + conv_slot = int(indices[sequence, 0]) + if sequence_length <= 0 or conv_slot <= 0: + continue + accepted = accepted_tokens[sequence] - 1 + state_slot = int(indices[sequence, accepted]) + history = conv_state[conv_slot, :, accepted : accepted + 3].float() + current_state = state[state_slot].float().clone() + conv_state[conv_slot, :, 0] = history[:, 1].to(torch.bfloat16) + conv_state[conv_slot, :, 1] = history[:, 2].to(torch.bfloat16) + + for token in range(sequence_length): + row = begin + token + current_x = inputs["x"][row].float() + window = torch.cat((history, current_x.unsqueeze(-1)), dim=-1) + history = window[:, 1:] + conv_state[conv_slot, :, token + 2] = current_x.to(torch.bfloat16) + + mixed = F.silu((window * taps).sum(-1)).to(torch.bfloat16) + query, key, value = mixed.float().view(3, num_heads, 128).unbind(0) + query *= torch.rsqrt(query.square().sum(-1, keepdim=True) + 1e-6) + query *= 128**-0.5 + key *= torch.rsqrt(key.square().sum(-1, keepdim=True) + 1e-6) + raw_gate = inputs["raw_gate"][0, row].float() + raw_gate += inputs["dt_bias"].view(num_heads, 128) + decay = torch.exp(-5.0 * torch.sigmoid(decay_a * raw_gate)) + current_state *= decay.unsqueeze(-2) + state_key = torch.einsum("hvk,hk->hv", current_state, key) + delta = value - state_key + beta = inputs["raw_beta"][0, row].float().sigmoid() + current_state += (delta * beta.unsqueeze(-1)).unsqueeze(-1) * key.unsqueeze( + -2 + ) + destination = int(indices[sequence, token]) + if destination > 0: + state[destination].copy_(current_state) + recurrent = torch.einsum("hvk,hk->hv", current_state, query) + output[0, row] = recurrent.to(torch.bfloat16) + + output_float = output.float() + inverse_rms = torch.rsqrt(output_float.square().mean(-1, keepdim=True) + 1e-5) + return ( + output_float + * inverse_rms + * inputs["norm_weight"] + * inputs["output_gate"].float().sigmoid().unsqueeze(0) + ).to(torch.bfloat16) + + +def _make_packed_inputs(num_heads, num_sequences, num_tokens, seed=42): + inputs = _make_inputs( + num_heads, + num_sequences * num_tokens, + seed=seed, + conv_history=num_tokens + 2, + ) + inputs["state_indices"] = inputs["state_indices"].reshape(num_sequences, num_tokens) + inputs["query_start_loc"] = torch.arange( + 0, + num_sequences * num_tokens + 1, + num_tokens, + dtype=torch.int32, + device=torch.device("cuda"), + ) + inputs["num_accepted_tokens"] = torch.ones( + num_sequences, dtype=torch.int32, device=torch.device("cuda") + ) + return inputs + + @pytest.mark.parametrize( ("num_heads", "num_rows"), [ @@ -193,7 +299,7 @@ def test_fused_kda_decode_kimi_k3_shapes(num_heads, num_rows): actual_state = _clone_strided(inputs["state"]) expected = _reference(inputs, reference_conv_state, reference_state) - actual = fused_kda_decode( + actual = _fused_kda_decode( x=inputs["x"], weight=inputs["weight"], conv_state=actual_conv_state, @@ -230,7 +336,7 @@ def test_fused_kda_decode_kimi_linear_softplus(num_rows, state_dtype): "state": actual_state, "lower_bound": None, } - actual = fused_kda_decode(**kwargs) + actual = _fused_kda_decode(**kwargs) torch.testing.assert_close(actual, expected, rtol=3e-2, atol=2e-2) torch.testing.assert_close(actual_conv_state, reference_conv_state, rtol=0, atol=0) @@ -256,10 +362,10 @@ def test_fused_kda_decode_cuda_graph(num_heads, lower_bound, state_dtype): # Compile and allocate before capture. Replays then contain only the fused # kernel and preserve the in-place cache update semantics. - fused_kda_decode(**kwargs) + _fused_kda_decode(**kwargs) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - captured_output = fused_kda_decode(**kwargs) + captured_output = _fused_kda_decode(**kwargs) graph.replay() torch.cuda.synchronize() @@ -276,7 +382,7 @@ def test_fused_kda_decode_null_slots(state_dtype): conv_state_before = _clone_strided(inputs["conv_state"]) state_before = _clone_strided(inputs["state"]) - output = fused_kda_decode(**inputs) + output = _fused_kda_decode(**inputs) torch.testing.assert_close(output[:, :2], torch.zeros_like(output[:, :2])) torch.testing.assert_close(inputs["conv_state"][0], conv_state_before[0]) @@ -295,7 +401,48 @@ def test_fused_kda_decode_requires_matching_cache_slots(): ValueError, match="conv_state and state must have the same number of cache slots", ): - fused_kda_decode(**inputs) + _fused_kda_decode(**inputs) + + +def test_fused_kda_decode_compile_abi_orders_flags_and_metadata(monkeypatch): + import importlib + + module = importlib.import_module("flashinfer.kda_kernels.fused_kda_decode") + compiled = [] + + monkeypatch.setattr(module, "_make_compile_inputs", lambda dtype: (dtype,)) + + def fake_compile(function, *args, **kwargs): + del function, kwargs + compiled.append(args) + return "compiled" + + def fake_build(module_name, kernel_name, compile_fn, extra_key_files): + del module_name, kernel_name, extra_key_files + compile_fn() + return "loaded" + + monkeypatch.setattr(module.cute, "compile", fake_compile) + monkeypatch.setattr(module, "build_and_load_cute_dsl_kernel", fake_build) + module._get_compiled_kernel.cache_clear() + try: + for state_dtype in (torch.float32, torch.bfloat16): + for packed_t1 in (False, True): + module._get_compiled_kernel( + state_dtype, -5.0, 1e-5, packed_t1=packed_t1 + ) + finally: + module._get_compiled_kernel.cache_clear() + + assert len(compiled) == 4 + # _make_compile_inputs contributes one sentinel; constexpr flags follow + # the launch signature: state_is_bf16, packed_t1. + assert {(args[-5], args[-4]) for args in compiled} == { + (False, False), + (False, True), + (True, False), + (True, True), + } def test_fused_kda_decode_cache_key_distinguishes_state_and_gate(monkeypatch): @@ -322,3 +469,919 @@ def record_kernel_name(module_name, kernel_name, compile_fn, extra_key_files): kernel_module._get_compiled_kernel.cache_clear() assert len(kernel_names) == len(set(kernel_names)) == 4 + + +@pytest.mark.parametrize( + ("num_heads", "num_sequences", "num_tokens"), + [ + pytest.param(12, 1, 1, id="h12-n1-t1"), + pytest.param(12, 2, 1, id="h12-n2-t1-dynamic-batch"), + pytest.param(12, 2, 2, id="h12-n2-t2"), + pytest.param(24, 4, 3, id="h24-n4-t3"), + pytest.param(32, 2, 4, id="h32-n2-t4"), + pytest.param(48, 2, 8, id="h48-n2-t8"), + pytest.param(96, 1, 3, id="h96-n1-t3"), + ], +) +def test_packed_fused_kda_decode(num_heads, num_sequences, num_tokens): + inputs = _make_packed_inputs(num_heads, num_sequences, num_tokens) + reference_conv = _clone_strided(inputs["conv_state"]) + reference_state = _clone_strided(inputs["state"]) + actual_conv = _clone_strided(inputs["conv_state"]) + actual_state = _clone_strided(inputs["state"]) + + expected = _packed_reference(inputs, reference_conv, reference_state) + actual = _packed_fused_kda_decode( + **{**inputs, "conv_state": actual_conv, "state": actual_state} + ) + + torch.testing.assert_close(actual, expected, rtol=3e-2, atol=2e-2) + torch.testing.assert_close(actual_conv, reference_conv, rtol=0, atol=0) + torch.testing.assert_close(actual_state, reference_state, rtol=3e-2, atol=2e-3) + + +def test_packed_fused_kda_decode_t1_pre_resolved_matches_direct_and_graph_replay(): + inputs = _make_inputs(12, 4, state_dtype=torch.float32) + inputs["state_indices"] = torch.tensor( + [[1], [2], [3], [4]], dtype=torch.int32, device="cuda" + ) + inputs["query_start_loc"] = torch.tensor( + [0, 1, 2, 2, 2], dtype=torch.int32, device="cuda" + ) + inputs["num_accepted_tokens"] = torch.ones(4, dtype=torch.int32, device="cuda") + packed_slots = torch.tensor([1, 2, 0, 0], dtype=torch.int32, device="cuda") + + def clone_inputs(): + return { + key: _clone_strided(value) if isinstance(value, torch.Tensor) else value + for key, value in inputs.items() + } + + direct_inputs = clone_inputs() + direct_inputs.pop("query_start_loc") + direct_inputs.pop("num_accepted_tokens") + direct_inputs["state_indices"] = packed_slots.clone() + direct = _fused_kda_decode(**direct_inputs) + packed_inputs = clone_inputs() + packed_output = torch.empty_like(direct) + packed = _packed_fused_kda_decode( + **packed_inputs, + t1_state_indices=packed_slots, + output=packed_output, + ) + assert torch.equal(packed, direct) + assert torch.equal(packed_inputs["conv_state"], direct_inputs["conv_state"]) + assert torch.equal(packed_inputs["state"], direct_inputs["state"]) + assert torch.equal(packed[0, 2:], torch.zeros_like(packed[0, 2:])) + + packed_conv_seed = _clone_strided(inputs["conv_state"]) + packed_state_seed = _clone_strided(inputs["state"]) + graph_output = torch.empty_like(direct) + graph_inputs = clone_inputs() + graph_inputs["conv_state"].copy_(packed_conv_seed) + graph_inputs["state"].copy_(packed_state_seed) + # Capture through a closure so the caller-owned slot tensor is read at + # replay time, exactly as it is in the serving CUDA-graph path. + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _packed_fused_kda_decode( + **graph_inputs, + t1_state_indices=packed_slots, + output=graph_output, + ) + packed_slots.copy_(torch.tensor([2, 1, 0, 0], device="cuda")) + graph_inputs["conv_state"].copy_(packed_conv_seed) + graph_inputs["state"].copy_(packed_state_seed) + graph.replay() + torch.cuda.synchronize() + + expected_inputs = clone_inputs() + expected_inputs.pop("query_start_loc") + expected_inputs.pop("num_accepted_tokens") + expected_inputs["state_indices"] = packed_slots.clone() + expected = _fused_kda_decode(**expected_inputs) + assert torch.equal(graph_output, expected) + assert torch.equal(graph_inputs["conv_state"], expected_inputs["conv_state"]) + assert torch.equal(graph_inputs["state"], expected_inputs["state"]) + + +def test_packed_fused_kda_decode_t1_uses_legacy_modes(): + inputs = _make_inputs(12, 2, state_dtype=torch.bfloat16) + inputs["state_indices"] = inputs["state_indices"].reshape(2, 1) + inputs["query_start_loc"] = torch.arange( + 3, dtype=torch.int32, device=torch.device("cuda") + ) + inputs["num_accepted_tokens"] = torch.ones( + 2, dtype=torch.int32, device=torch.device("cuda") + ) + legacy_conv = _clone_strided(inputs["conv_state"]) + legacy_state = _clone_strided(inputs["state"]) + packed_conv = _clone_strided(inputs["conv_state"]) + packed_state = _clone_strided(inputs["state"]) + + legacy_inputs = { + key: value + for key, value in inputs.items() + if key not in ("query_start_loc", "num_accepted_tokens") + } + legacy = _fused_kda_decode( + **{ + **legacy_inputs, + "conv_state": legacy_conv, + "state": legacy_state, + "state_indices": inputs["state_indices"][:, 0], + "lower_bound": None, + } + ) + packed = _packed_fused_kda_decode( + **{ + **inputs, + "conv_state": packed_conv, + "state": packed_state, + "lower_bound": None, + } + ) + + assert torch.equal(packed, legacy) + assert torch.equal(packed_conv, legacy_conv) + assert torch.equal(packed_state, legacy_state) + + +@pytest.mark.parametrize( + "query_start_loc", + [ + pytest.param([0, 0, 1, 2], id="empty-first"), + pytest.param([0, 1, 1, 2], id="empty-middle"), + pytest.param([0, 1, 2, 2], id="empty-last"), + ], +) +def test_packed_fused_kda_decode_t1_empty_sequence(query_start_loc): + inputs = _make_inputs(12, 3, state_dtype=torch.bfloat16) + for key in ("x", "raw_gate", "raw_beta", "output_gate"): + row_dimension = 1 if key in ("raw_gate", "raw_beta") else 0 + inputs[key] = inputs[key].narrow(row_dimension, 0, 2) + inputs["state_indices"] = torch.arange( + 1, 4, dtype=torch.int32, device=torch.device("cuda") + ).reshape(3, 1) + inputs["query_start_loc"] = torch.tensor( + query_start_loc, dtype=torch.int32, device=torch.device("cuda") + ) + inputs["num_accepted_tokens"] = torch.ones( + 3, dtype=torch.int32, device=torch.device("cuda") + ) + live_sequences = [ + sequence + for sequence in range(3) + if query_start_loc[sequence] != query_start_loc[sequence + 1] + ] + legacy_indices = inputs["state_indices"][live_sequences, 0].contiguous() + legacy_conv = _clone_strided(inputs["conv_state"]) + legacy_state = _clone_strided(inputs["state"]) + packed_conv = _clone_strided(inputs["conv_state"]) + packed_state = _clone_strided(inputs["state"]) + + legacy = _fused_kda_decode( + **{ + **{ + key: value + for key, value in inputs.items() + if key + not in ( + "query_start_loc", + "num_accepted_tokens", + "state_indices", + ) + }, + "conv_state": legacy_conv, + "state": legacy_state, + "state_indices": legacy_indices, + "lower_bound": None, + } + ) + packed_output = torch.empty_like(legacy) + packed_kwargs = { + **inputs, + "conv_state": packed_conv, + "state": packed_state, + "lower_bound": None, + "output": packed_output, + } + _packed_fused_kda_decode(**packed_kwargs) + packed_conv.copy_(inputs["conv_state"]) + packed_state.copy_(inputs["state"]) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _packed_fused_kda_decode(**packed_kwargs) + packed_conv.copy_(inputs["conv_state"]) + packed_state.copy_(inputs["state"]) + graph.replay() + torch.cuda.synchronize() + + assert torch.equal(packed_output, legacy) + assert torch.equal(packed_conv, legacy_conv) + assert torch.equal(packed_state, legacy_state) + + +@pytest.mark.parametrize( + ("state_dtype", "lower_bound", "error", "match"), + [ + pytest.param( + torch.bfloat16, + -5.0, + TypeError, + "state must have dtype torch.float32", + id="bf16-state", + ), + pytest.param( + torch.float32, + None, + ValueError, + "lower_bound must be a finite negative float", + id="softplus", + ), + ], +) +def test_packed_fused_kda_decode_t2_rejects_legacy_only_modes( + state_dtype, lower_bound, error, match +): + inputs = _make_inputs(12, 2, state_dtype=state_dtype, conv_history=4) + inputs["state_indices"] = inputs["state_indices"].reshape(1, 2) + inputs["query_start_loc"] = torch.tensor( + [0, 2], dtype=torch.int32, device=torch.device("cuda") + ) + inputs["num_accepted_tokens"] = torch.ones( + 1, dtype=torch.int32, device=torch.device("cuda") + ) + + with pytest.raises(error, match=match): + _packed_fused_kda_decode(**inputs, lower_bound=lower_bound) + + +def test_packed_fused_kda_decode_rejects_unsupported_arch(monkeypatch): + import importlib + + module = importlib.import_module( + "flashinfer.kda_kernels.fused_kda_decode_multitoken" + ) + inputs = _make_packed_inputs(12, 1, 2) + monkeypatch.setattr( + module.torch.cuda, "get_device_capability", lambda _device: (9, 0) + ) + + with pytest.raises(NotImplementedError, match="requires SM10x"): + _packed_fused_kda_decode(**inputs) + + +def test_packed_fused_kda_decode_shared_memory_limit(): + import importlib + + module = importlib.import_module( + "flashinfer.kda_kernels.fused_kda_decode_multitoken" + ) + limit = torch.cuda.get_device_properties(0).shared_memory_per_block_optin + supported = 1 + while ( + module._required_smem_bytes( + supported + 1, *module._TILE_TALL[:2], module._TILE_TALL[3] + ) + <= limit + ): + supported += 1 + + assert supported >= 96 + assert ( + module._required_smem_bytes( + supported, *module._TILE_TALL[:2], module._TILE_TALL[3] + ) + <= limit + ) + assert ( + module._required_smem_bytes( + supported + 1, *module._TILE_TALL[:2], module._TILE_TALL[3] + ) + > limit + ) + + +def test_packed_fused_kda_decode_split_tiles_load_frontend_in_one_pass(): + import importlib + + module = importlib.import_module( + "flashinfer.kda_kernels.fused_kda_decode_multitoken" + ) + for split, (rows, k_lanes, _, chunks) in module._TILE_SPLIT.items(): + threads = (((128 // split) // chunks) // rows) * k_lanes + assert threads == 4 * 128 + + +def test_packed_fused_kda_decode_null_row(): + inputs = _make_packed_inputs(num_heads=12, num_sequences=2, num_tokens=3) + inputs["state_indices"][0].zero_() + reference_conv = _clone_strided(inputs["conv_state"]) + reference_state = _clone_strided(inputs["state"]) + actual_conv = _clone_strided(inputs["conv_state"]) + actual_state = _clone_strided(inputs["state"]) + + expected = _packed_reference(inputs, reference_conv, reference_state) + actual = _packed_fused_kda_decode( + **{**inputs, "conv_state": actual_conv, "state": actual_state} + ) + + torch.testing.assert_close(actual, expected, rtol=3e-2, atol=2e-2) + assert torch.count_nonzero(actual[:, :3]) == 0 + torch.testing.assert_close(actual_conv, reference_conv, rtol=0, atol=0) + torch.testing.assert_close(actual_state, reference_state, rtol=3e-2, atol=2e-3) + + +def test_packed_fused_kda_decode_ragged_rows_and_padding(): + inputs = _make_packed_inputs(num_heads=12, num_sequences=3, num_tokens=4) + inputs["query_start_loc"].copy_( + torch.tensor([0, 4, 6, 6], dtype=torch.int32, device=torch.device("cuda")) + ) + inputs["num_accepted_tokens"].copy_( + torch.tensor([1, 2, 1], dtype=torch.int32, device=torch.device("cuda")) + ) + reference_conv = _clone_strided(inputs["conv_state"]) + reference_state = _clone_strided(inputs["state"]) + actual_conv = _clone_strided(inputs["conv_state"]) + actual_state = _clone_strided(inputs["state"]) + output = torch.full( + (1, inputs["x"].shape[0], 12, 128), + torch.nan, + dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + + expected = _packed_reference(inputs, reference_conv, reference_state) + actual = _packed_fused_kda_decode( + **{ + **inputs, + "conv_state": actual_conv, + "state": actual_state, + "output": output, + } + ) + + torch.testing.assert_close(actual[:, :6], expected[:, :6], rtol=3e-2, atol=2e-2) + assert torch.isnan(actual[:, 6:]).all() + torch.testing.assert_close(actual_conv, reference_conv, rtol=0, atol=0) + torch.testing.assert_close(actual_state, reference_state, rtol=3e-2, atol=2e-3) + + +@pytest.mark.parametrize("num_tokens", [3, 7, 8]) +def test_packed_fused_kda_decode_two_verification_windows(num_tokens): + first = _make_packed_inputs(12, 2, num_tokens, seed=41) + reference_conv = _clone_strided(first["conv_state"]) + reference_state = _clone_strided(first["state"]) + actual_conv = _clone_strided(first["conv_state"]) + actual_state = _clone_strided(first["state"]) + + expected_first = _packed_reference(first, reference_conv, reference_state) + actual_first = _packed_fused_kda_decode( + **{**first, "conv_state": actual_conv, "state": actual_state} + ) + torch.testing.assert_close(actual_conv, reference_conv, rtol=0, atol=0) + torch.testing.assert_close(actual_state, reference_state, rtol=3e-2, atol=2e-3) + torch.testing.assert_close(actual_first, expected_first, rtol=3e-2, atol=2e-2) + + second = _make_packed_inputs(12, 2, num_tokens, seed=83) + second["state_indices"].copy_(first["state_indices"]) + second["num_accepted_tokens"].copy_( + torch.tensor([2, num_tokens], dtype=torch.int32, device=torch.device("cuda")) + ) + expected_second = _packed_reference(second, reference_conv, reference_state) + actual_second = _packed_fused_kda_decode( + **{**second, "conv_state": actual_conv, "state": actual_state} + ) + + torch.testing.assert_close(actual_second, expected_second, rtol=3e-2, atol=2e-2) + torch.testing.assert_close(actual_conv, reference_conv, rtol=0, atol=0) + torch.testing.assert_close(actual_state, reference_state, rtol=3e-2, atol=2e-3) + + +def test_packed_fused_kda_decode_all_accepted_offsets(): + num_tokens = 8 + num_sequences = 32 + first = _make_packed_inputs(12, num_sequences, num_tokens, seed=107) + reference_conv = _clone_strided(first["conv_state"]) + reference_state = _clone_strided(first["state"]) + actual_conv = _clone_strided(first["conv_state"]) + actual_state = _clone_strided(first["state"]) + + _packed_reference(first, reference_conv, reference_state) + _packed_fused_kda_decode( + **{**first, "conv_state": actual_conv, "state": actual_state} + ) + + second = _make_packed_inputs(12, num_sequences, num_tokens, seed=109) + second["state_indices"].copy_(first["state_indices"]) + second["num_accepted_tokens"].copy_( + torch.arange(num_sequences, dtype=torch.int32, device=torch.device("cuda")) + % num_tokens + + 1 + ) + expected = _packed_reference(second, reference_conv, reference_state) + actual = _packed_fused_kda_decode( + **{**second, "conv_state": actual_conv, "state": actual_state} + ) + + torch.testing.assert_close(actual, expected, rtol=3e-2, atol=2e-2) + torch.testing.assert_close(actual_conv, reference_conv, rtol=0, atol=0) + torch.testing.assert_close(actual_state, reference_state, rtol=3e-2, atol=2e-3) + + +@pytest.mark.parametrize(("num_sequences", "num_tokens"), [(1, 7), (2, 8)]) +def test_packed_fused_kda_decode_cluster_determinism(num_sequences, num_tokens): + inputs = _make_packed_inputs(12, num_sequences, num_tokens, seed=127) + reference_conv = _clone_strided(inputs["conv_state"]) + reference_state = _clone_strided(inputs["state"]) + expected = _packed_reference(inputs, reference_conv, reference_state) + first = None + + for _ in range(8): + actual_conv = _clone_strided(inputs["conv_state"]) + actual_state = _clone_strided(inputs["state"]) + actual = _packed_fused_kda_decode( + **{**inputs, "conv_state": actual_conv, "state": actual_state} + ) + torch.cuda.synchronize() + + torch.testing.assert_close(actual, expected, rtol=3e-2, atol=2e-2) + torch.testing.assert_close(actual_conv, reference_conv, rtol=0, atol=0) + torch.testing.assert_close(actual_state, reference_state, rtol=3e-2, atol=2e-3) + if first is None: + first = (actual.clone(), actual_conv.clone(), actual_state.clone()) + else: + assert torch.equal(actual, first[0]) + assert torch.equal(actual_conv, first[1]) + assert torch.equal(actual_state, first[2]) + + +def test_packed_fused_kda_decode_cuda_graph(): + inputs = _make_packed_inputs(num_heads=12, num_sequences=2, num_tokens=3) + expected_conv = _clone_strided(inputs["conv_state"]) + expected_state = _clone_strided(inputs["state"]) + expected = _packed_reference(inputs, expected_conv, expected_state) + output = torch.empty_like(expected) + kwargs = {**inputs, "output": output} + + _packed_fused_kda_decode(**kwargs) + # Restore the exact seeded inputs rather than relying on capture-time state. + fresh = _make_packed_inputs(12, 2, 3) + inputs["conv_state"].copy_(fresh["conv_state"]) + inputs["state"].copy_(fresh["state"]) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _packed_fused_kda_decode(**kwargs) + inputs["conv_state"].copy_(fresh["conv_state"]) + inputs["state"].copy_(fresh["state"]) + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close(output, expected, rtol=3e-2, atol=2e-2) + torch.testing.assert_close(inputs["conv_state"], expected_conv, rtol=0, atol=0) + torch.testing.assert_close(inputs["state"], expected_state, rtol=3e-2, atol=2e-3) + + +def _t1_inputs(num_heads, state_dtype): + inputs = _make_inputs(num_heads, 4, state_dtype=state_dtype) + # Capture a fixed three-row graph while retaining five total cache slots 0..4; four are positive. + for key in ("x", "raw_gate", "raw_beta", "output_gate"): + dimension = 1 if key in ("raw_gate", "raw_beta") else 0 + inputs[key] = inputs[key].narrow(dimension, 0, 3) + inputs["state_indices"] = torch.empty((3, 1), dtype=torch.int32, device="cuda") + inputs["query_start_loc"] = torch.empty(4, dtype=torch.int32, device="cuda") + inputs["num_accepted_tokens"] = torch.ones(3, dtype=torch.int32, device="cuda") + return inputs + + +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("lower_bound", [-5.0, None]) +def test_packed_fused_kda_decode_t1_graph_replay_remaps_rows(state_dtype, lower_bound): + inputs = _t1_inputs(12, state_dtype) + original_conv = _clone_strided(inputs["conv_state"]) + original_state = _clone_strided(inputs["state"]) + inputs["query_start_loc"].copy_(torch.tensor([0, 1, 2, 3], device="cuda")) + inputs["state_indices"].copy_(torch.tensor([[1], [2], [3]], device="cuda")) + output = torch.empty((1, 3, 12, 128), dtype=torch.bfloat16, device="cuda") + kwargs = {**inputs, "output": output, "lower_bound": lower_bound} + + _packed_fused_kda_decode(**kwargs) + inputs["conv_state"].copy_(original_conv) + inputs["state"].copy_(original_state) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _packed_fused_kda_decode(**kwargs) + + trials = [ + ([0, 1, 2, 3], [1, 2, 3], [1, 2, 3]), + ([0, 1, 1, 2], [4, 2, 1], [4, 1]), + ([0, 1, 2, 3], [3, 4, 2], [3, 4, 2]), + ] + for qsl_values, slots, expected_slots in trials: + assert len(qsl_values) == 4 and max(qsl_values) <= 3 + assert all(0 < slot < inputs["state"].shape[0] for slot in slots) + inputs["conv_state"].copy_(original_conv) + inputs["state"].copy_(original_state) + inputs["query_start_loc"].copy_(torch.tensor(qsl_values, device="cuda")) + inputs["state_indices"].copy_(torch.tensor(slots, device="cuda").view(3, 1)) + output.fill_(torch.nan) + + live = len(expected_slots) + oracle = { + **inputs, + "x": inputs["x"][:live], + "raw_gate": inputs["raw_gate"][:, :live], + "raw_beta": inputs["raw_beta"][:, :live], + "output_gate": inputs["output_gate"][:live], + "state_indices": torch.tensor(expected_slots, device="cuda"), + } + expected_conv = _clone_strided(original_conv) + expected_state = _clone_strided(original_state) + expected = _reference(oracle, expected_conv, expected_state, lower_bound) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output[:, :live], expected, rtol=3e-2, atol=2e-2) + assert torch.equal(output[:, live:], torch.zeros_like(output[:, live:])) + torch.testing.assert_close(inputs["conv_state"], expected_conv, rtol=0, atol=0) + torch.testing.assert_close( + inputs["state"], expected_state, rtol=3e-2, atol=2e-3 + ) + for slot in range(inputs["state"].shape[0]): + if slot not in expected_slots: + assert torch.equal(inputs["state"][slot], original_state[slot]) + assert torch.equal(inputs["conv_state"][slot], original_conv[slot]) + + +@pytest.mark.parametrize("num_rows", [129, 257]) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("lower_bound", [-5.0, None]) +def test_packed_fused_kda_decode_t1_tail_bound_and_dual_role( + num_rows, state_dtype, lower_bound +): + # role 2 is both the live sequence CTA for row 0 and a tail CTA for + # output row 2. The independent literal oracle catches an unconditional + # tail return; the large capacities catch using x.shape[0] (the vector + # lane extent) instead of the transformed output row extent. + inputs = _make_inputs(12, num_rows, state_dtype=state_dtype) + inputs["state_indices"] = torch.tensor( + [[0], [0], [7]], dtype=torch.int32, device="cuda" + ) + inputs["query_start_loc"] = torch.tensor( + [0, 0, 0, 1], dtype=torch.int32, device="cuda" + ) + inputs["num_accepted_tokens"] = torch.ones(3, dtype=torch.int32, device="cuda") + before_conv = _clone_strided(inputs["conv_state"]) + before_state = _clone_strided(inputs["state"]) + output = torch.full( + (1, num_rows, 12, 128), torch.nan, dtype=torch.bfloat16, device="cuda" + ) + actual = _packed_fused_kda_decode( + **{**inputs, "output": output, "lower_bound": lower_bound} + ) + + # Only sequence 2 is live, and it owns packed row 0. Every other output + # row is padding and must be exactly zero, including the rows beyond 4. + oracle = { + **inputs, + "x": inputs["x"][:1], + "raw_gate": inputs["raw_gate"][:, :1], + "raw_beta": inputs["raw_beta"][:, :1], + "output_gate": inputs["output_gate"][:1], + "state_indices": torch.tensor([7], dtype=torch.int32, device="cuda"), + } + expected_conv = _clone_strided(before_conv) + expected_state = _clone_strided(before_state) + expected = _reference(oracle, expected_conv, expected_state, lower_bound) + torch.testing.assert_close(actual[:, :1], expected, rtol=3e-2, atol=2e-2) + assert torch.equal(actual[:, 1:], torch.zeros_like(actual[:, 1:])) + torch.testing.assert_close(inputs["conv_state"], expected_conv, rtol=0, atol=0) + torch.testing.assert_close(inputs["state"], expected_state, rtol=3e-2, atol=2e-3) + for slot in range(inputs["state"].shape[0]): + if slot != 7: + assert torch.equal(inputs["state"][slot], before_state[slot]) + assert torch.equal(inputs["conv_state"][slot], before_conv[slot]) + + +def test_packed_fused_kda_decode_t1_all_empty_preserves_positive_slots(): + inputs = _t1_inputs(12, torch.bfloat16) + before_conv = _clone_strided(inputs["conv_state"]) + before_state = _clone_strided(inputs["state"]) + inputs["query_start_loc"].copy_(torch.tensor([0, 0, 0, 0], device="cuda")) + inputs["state_indices"].copy_(torch.tensor([[1], [2], [3]], device="cuda")) + output = _packed_fused_kda_decode(**inputs) + assert torch.equal(output, torch.zeros_like(output)) + torch.testing.assert_close(inputs["conv_state"], before_conv, rtol=0, atol=0) + torch.testing.assert_close(inputs["state"], before_state, rtol=0, atol=0) + + +@pytest.mark.parametrize( + ("num_sequences", "qsl_values", "slots", "expected_slots"), + [ + (4, [0, 1, 2, 2, 2], [1, 2, 3, 4], [1, 2]), + (3, [0, 1, 1, 2], [4, 2, 1], [4, 1]), + (4, [0, 0, 0, 0, 0], [1, 2, 3, 4], []), + ], +) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("lower_bound", [-5.0, None]) +def test_packed_fused_kda_decode_t1_eager_valid_padding( + num_sequences, qsl_values, slots, expected_slots, state_dtype, lower_bound +): + inputs = _make_inputs(12, 4, state_dtype=state_dtype) + inputs["state_indices"] = torch.tensor( + slots, dtype=torch.int32, device="cuda" + ).view(num_sequences, 1) + inputs["query_start_loc"] = torch.tensor( + qsl_values, dtype=torch.int32, device="cuda" + ) + inputs["num_accepted_tokens"] = torch.ones( + num_sequences, dtype=torch.int32, device="cuda" + ) + before_conv = _clone_strided(inputs["conv_state"]) + before_state = _clone_strided(inputs["state"]) + output = torch.full((1, 4, 12, 128), torch.nan, dtype=torch.bfloat16, device="cuda") + output = _packed_fused_kda_decode( + **{**inputs, "output": output, "lower_bound": lower_bound} + ) + assert all(0 < slot < inputs["state"].shape[0] for slot in slots) + assert len(expected_slots) == inputs["query_start_loc"][-1].item() + if expected_slots: + oracle = { + **inputs, + "x": inputs["x"][: len(expected_slots)], + "raw_gate": inputs["raw_gate"][:, : len(expected_slots)], + "raw_beta": inputs["raw_beta"][:, : len(expected_slots)], + "output_gate": inputs["output_gate"][: len(expected_slots)], + "state_indices": torch.tensor(expected_slots, device="cuda"), + } + expected_conv = _clone_strided(before_conv) + expected_state = _clone_strided(before_state) + expected = _reference(oracle, expected_conv, expected_state, lower_bound) + torch.testing.assert_close( + output[:, : len(expected_slots)], expected, rtol=3e-2, atol=2e-2 + ) + torch.testing.assert_close(inputs["conv_state"], expected_conv, rtol=0, atol=0) + torch.testing.assert_close( + inputs["state"], expected_state, rtol=3e-2, atol=2e-3 + ) + else: + assert torch.equal(output, torch.zeros_like(output)) + assert torch.equal(inputs["conv_state"], before_conv) + assert torch.equal(inputs["state"], before_state) + assert torch.equal( + output[:, len(expected_slots) :], + torch.zeros_like(output[:, len(expected_slots) :]), + ) + for slot in range(inputs["state"].shape[0]): + if slot not in expected_slots: + assert torch.equal(inputs["state"][slot], before_state[slot]) + assert torch.equal(inputs["conv_state"][slot], before_conv[slot]) + + +def test_packed_fused_kda_decode_cuda_graph_dynamic_metadata(): + inputs = _make_packed_inputs(num_heads=12, num_sequences=3, num_tokens=4) + seeded_conv = _clone_strided(inputs["conv_state"]) + seeded_state = _clone_strided(inputs["state"]) + output = torch.full( + (1, inputs["x"].shape[0], 12, 128), + torch.nan, + dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + kwargs = {**inputs, "output": output} + + _packed_fused_kda_decode(**kwargs) + inputs["conv_state"].copy_(seeded_conv) + inputs["state"].copy_(seeded_state) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _packed_fused_kda_decode(**kwargs) + + inputs["conv_state"].copy_(seeded_conv) + inputs["state"].copy_(seeded_state) + output.fill_(torch.nan) + inputs["query_start_loc"].copy_( + torch.tensor([0, 4, 6, 6], dtype=torch.int32, device=torch.device("cuda")) + ) + inputs["num_accepted_tokens"].copy_( + torch.tensor([1, 2, 1], dtype=torch.int32, device=torch.device("cuda")) + ) + expected_conv = _clone_strided(seeded_conv) + expected_state = _clone_strided(seeded_state) + expected = _packed_reference(inputs, expected_conv, expected_state) + + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close(output[:, :6], expected[:, :6], rtol=3e-2, atol=2e-2) + assert torch.isnan(output[:, 6:]).all() + torch.testing.assert_close(inputs["conv_state"], expected_conv, rtol=0, atol=0) + torch.testing.assert_close(inputs["state"], expected_state, rtol=3e-2, atol=2e-3) + + +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("lower_bound", [-5.0, None]) +def test_packed_fused_kda_decode_t1_trailing_empty_sequences_beyond_rows( + state_dtype, lower_bound +): + # The live row is valid, while three trailing sequences have qsl==rows. + # This specifically guards the old sequence-CTA input load at row_idx=qsl[s] + # when N > num_rows and total_tokens == num_rows. + inputs = _make_inputs(12, 4, state_dtype=state_dtype) + for key in ("x", "raw_gate", "raw_beta", "output_gate"): + dimension = 1 if key in ("raw_gate", "raw_beta") else 0 + inputs[key] = inputs[key].narrow(dimension, 0, 1) + inputs["query_start_loc"] = torch.tensor( + [0, 1, 1, 1, 1], dtype=torch.int32, device="cuda" + ) + inputs["num_accepted_tokens"] = torch.ones(4, dtype=torch.int32, device="cuda") + inputs["state_indices"] = torch.tensor( + [[1], [2], [3], [4]], dtype=torch.int32, device="cuda" + ) + before_conv = _clone_strided(inputs["conv_state"]) + before_state = _clone_strided(inputs["state"]) + output = torch.full((1, 1, 12, 128), torch.nan, dtype=torch.bfloat16, device="cuda") + actual = _packed_fused_kda_decode( + **{**inputs, "output": output, "lower_bound": lower_bound} + ) + + oracle = { + **inputs, + "state_indices": torch.tensor([1], dtype=torch.int32, device="cuda"), + } + expected_conv = _clone_strided(before_conv) + expected_state = _clone_strided(before_state) + expected = _reference(oracle, expected_conv, expected_state, lower_bound) + torch.testing.assert_close(actual, expected, rtol=3e-2, atol=2e-2) + torch.testing.assert_close(inputs["conv_state"], expected_conv, rtol=0, atol=0) + torch.testing.assert_close(inputs["state"], expected_state, rtol=3e-2, atol=2e-3) + for slot in (0, 2, 3, 4): + assert torch.equal(inputs["state"][slot], before_state[slot]) + assert torch.equal(inputs["conv_state"][slot], before_conv[slot]) + + +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("lower_bound", [-5.0, None]) +@pytest.mark.parametrize("live_slot", [0, -1, 4]) +def test_packed_fused_kda_decode_t1_late_active_dual_role_null_slots( + state_dtype, lower_bound, live_slot +): + # Empty sequences precede the live sequence. Null slots must not touch + # caches; a positive live slot still owns packed row zero. + inputs = _make_inputs(12, 4, state_dtype=state_dtype) + for key in ("x", "raw_gate", "raw_beta", "output_gate"): + dimension = 1 if key in ("raw_gate", "raw_beta") else 0 + inputs[key] = inputs[key].narrow(dimension, 0, 3) + inputs["query_start_loc"] = torch.tensor( + [0, 0, 0, 1], dtype=torch.int32, device="cuda" + ) + inputs["num_accepted_tokens"] = torch.ones(3, dtype=torch.int32, device="cuda") + inputs["state_indices"] = torch.tensor( + [[0], [-1], [live_slot]], dtype=torch.int32, device="cuda" + ) + assert inputs["x"].shape[0] == 3 + assert inputs["query_start_loc"].tolist() == [0, 0, 0, 1] + assert -1 <= live_slot < inputs["state"].shape[0] + before_conv = _clone_strided(inputs["conv_state"]) + before_state = _clone_strided(inputs["state"]) + output = torch.full((1, 3, 12, 128), torch.nan, dtype=torch.bfloat16, device="cuda") + actual = _packed_fused_kda_decode( + **{**inputs, "output": output, "lower_bound": lower_bound} + ) + + if live_slot > 0: + oracle = { + **inputs, + "x": inputs["x"][:1], + "raw_gate": inputs["raw_gate"][:, :1], + "raw_beta": inputs["raw_beta"][:, :1], + "output_gate": inputs["output_gate"][:1], + "state_indices": torch.tensor( + [live_slot], dtype=torch.int32, device="cuda" + ), + } + expected_conv = _clone_strided(before_conv) + expected_state = _clone_strided(before_state) + expected = _reference(oracle, expected_conv, expected_state, lower_bound) + torch.testing.assert_close(actual[:, :1], expected, rtol=3e-2, atol=2e-2) + torch.testing.assert_close(inputs["conv_state"], expected_conv, rtol=0, atol=0) + torch.testing.assert_close( + inputs["state"], expected_state, rtol=3e-2, atol=2e-3 + ) + else: + assert torch.equal(actual, torch.zeros_like(actual)) + assert torch.equal(inputs["conv_state"], before_conv) + assert torch.equal(inputs["state"], before_state) + assert torch.equal(actual[:, 1:], torch.zeros_like(actual[:, 1:])) + for slot in range(5): + if slot != live_slot: + assert torch.equal(inputs["state"][slot], before_state[slot]) + assert torch.equal(inputs["conv_state"][slot], before_conv[slot]) + + +def test_packed_fused_kda_decode_t1_graph_changed_metadata_null_to_active(): + inputs = _make_inputs(12, 4, state_dtype=torch.float32) + for key in ("x", "raw_gate", "raw_beta", "output_gate"): + dimension = 1 if key in ("raw_gate", "raw_beta") else 0 + inputs[key] = inputs[key].narrow(dimension, 0, 3) + inputs["query_start_loc"] = torch.tensor( + [0, 1, 2, 3, 3], dtype=torch.int32, device="cuda" + ) + inputs["num_accepted_tokens"] = torch.ones(4, dtype=torch.int32, device="cuda") + inputs["state_indices"] = torch.tensor( + [[1], [2], [3], [4]], dtype=torch.int32, device="cuda" + ) + assert inputs["x"].shape[0] == 3 + assert inputs["query_start_loc"].tolist() == [0, 1, 2, 3, 3] + assert all(0 < slot < inputs["state"].shape[0] for slot in [1, 2, 3, 4]) + output = torch.full((1, 3, 12, 128), torch.nan, dtype=torch.bfloat16, device="cuda") + kwargs = {**inputs, "output": output, "lower_bound": -5.0} + _packed_fused_kda_decode(**kwargs) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + _packed_fused_kda_decode(**kwargs) + seeded_conv = _clone_strided(inputs["conv_state"]) + seeded_state = _clone_strided(inputs["state"]) + + # all-active -> trailing-empty -> late-active, including null/negative + # slots. Each replay uses a fresh literal oracle and must clear padding. + trials = [ + ([0, 1, 2, 3, 3], [1, 2, 3, 4], [1, 2, 3]), + ([0, 1, 1, 1, 1], [4, 2, 3, 1], [4]), + ([0, 0, 0, 1, 1], [1, 2, 4, 3], [4]), + ([0, 0, 0, 1, 1], [1, 2, 0, 3], []), + ([0, 1, 2, 3, 3], [1, 2, 3, 4], [1, 2, 3]), + ] + for qsl, slots, expected_slots in trials: + assert len(qsl) == 5 and len(slots) == 4 + assert qsl[-1] <= inputs["x"].shape[0] + assert all(-1 <= slot < inputs["state"].shape[0] for slot in slots) + inputs["conv_state"].copy_(seeded_conv) + inputs["state"].copy_(seeded_state) + inputs["query_start_loc"].copy_(torch.tensor(qsl, device="cuda")) + inputs["state_indices"].copy_(torch.tensor(slots, device="cuda").view(4, 1)) + output.fill_(torch.nan) + expected_conv = _clone_strided(seeded_conv) + expected_state = _clone_strided(seeded_state) + if expected_slots: + oracle = { + **inputs, + "x": inputs["x"][: len(expected_slots)], + "raw_gate": inputs["raw_gate"][:, : len(expected_slots)], + "raw_beta": inputs["raw_beta"][:, : len(expected_slots)], + "output_gate": inputs["output_gate"][: len(expected_slots)], + "state_indices": torch.tensor(expected_slots, device="cuda"), + } + expected = _reference(oracle, expected_conv, expected_state, -5.0) + else: + expected = None + graph.replay() + torch.cuda.synchronize() + if expected is not None: + torch.testing.assert_close( + output[:, : len(expected_slots)], expected, rtol=3e-2, atol=2e-2 + ) + else: + assert torch.equal(output, torch.zeros_like(output)) + assert torch.equal( + output[:, len(expected_slots) :], + torch.zeros_like(output[:, len(expected_slots) :]), + ) + torch.testing.assert_close(inputs["conv_state"], expected_conv, rtol=0, atol=0) + torch.testing.assert_close( + inputs["state"], expected_state, rtol=3e-2, atol=2e-3 + ) + for slot in range(5): + if slot not in expected_slots: + assert torch.equal(inputs["state"][slot], seeded_state[slot]) + assert torch.equal(inputs["conv_state"][slot], seeded_conv[slot]) + + +def test_packed_fused_kda_decode_cache_key(monkeypatch): + import importlib + + module = importlib.import_module( + "flashinfer.kda_kernels.fused_kda_decode_multitoken" + ) + kernel_names = [] + + def record_kernel_name(module_name, kernel_name, compile_fn, extra_key_files): + del module_name, compile_fn, extra_key_files + kernel_names.append(kernel_name) + return kernel_name + + monkeypatch.setattr(module, "build_and_load_cute_dsl_kernel", record_kernel_name) + module._get_compiled_kernel.cache_clear() + base = [3, 12, -5.0, 1e-5, 2, 16, 1, 1, 1] + alternatives = [4, 24, -4.0, 2e-5, 1, 32, 2, 2, 4] + try: + module._get_compiled_kernel(*base) + for index, alternative in enumerate(alternatives): + args = list(base) + args[index] = alternative + module._get_compiled_kernel(*args) + finally: + module._get_compiled_kernel.cache_clear() + + assert len(kernel_names) == len(set(kernel_names)) == 1 + len(alternatives) diff --git a/tests/kda/test_packed_fused_kda_decode_api.py b/tests/kda/test_packed_fused_kda_decode_api.py new file mode 100644 index 00000000000..54b9ae5439c --- /dev/null +++ b/tests/kda/test_packed_fused_kda_decode_api.py @@ -0,0 +1,209 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +from typing import Any, cast + +import pytest +import torch # pyright: ignore[reportMissingImports] + +import flashinfer.kda_decode as kda_decode + + +def _arguments(num_sequences=4, num_rows=4): + return { + "x": torch.empty((num_rows, 1)), + "weight": torch.empty((1,)), + "conv_state": torch.empty((1,)), + "raw_gate": torch.empty((1,)), + "raw_beta": torch.empty((1,)), + "A_log": torch.empty((1,)), + "dt_bias": torch.empty((1,)), + "state_indices": torch.tensor( + [[11 + 11 * index] for index in range(num_sequences)], dtype=torch.int32 + ), + "state": torch.empty((1,)), + "output_gate": torch.empty((1,)), + "norm_weight": torch.empty((1,)), + "query_start_loc": torch.arange(num_sequences + 1, dtype=torch.int32), + "num_accepted_tokens": torch.ones(num_sequences, dtype=torch.int32), + } + + +def test_packed_fused_kda_decode_signature_and_public_name(): + packed = inspect.signature(kda_decode.packed_fused_kda_decode) + legacy = inspect.signature(kda_decode.fused_kda_decode) + packed_parameters = list(packed.parameters.values()) + legacy_parameters = list(legacy.parameters.values()) + assert [parameter.name for parameter in packed_parameters[:14]] == [ + parameter.name for parameter in legacy_parameters[:14] + ] + assert [parameter.kind for parameter in packed_parameters[:14]] == [ + parameter.kind for parameter in legacy_parameters[:14] + ] + assert [parameter.default for parameter in packed_parameters[:14]] == [ + parameter.default for parameter in legacy_parameters[:14] + ] + assert [parameter.name for parameter in packed_parameters[14:]] == [ + "query_start_loc", + "num_accepted_tokens", + "t1_state_indices", + ] + assert all( + parameter.kind is inspect.Parameter.KEYWORD_ONLY + for parameter in packed_parameters[14:] + ) + assert all( + parameter.default is inspect.Parameter.empty + for parameter in packed_parameters[14:16] + ) + assert packed_parameters[16].default is None + assert not hasattr(kda_decode, "fused_kda_decode_packed") + flashinfer = __import__("flashinfer") + assert not hasattr(flashinfer, "fused_kda_decode_packed") + assert hasattr(flashinfer, "packed_fused_kda_decode") + + +def _capture_metadata(monkeypatch, arguments): + captured = {} + + def fake_backend(**kwargs): + captured.update(kwargs) + return "output" + + monkeypatch.setattr(kda_decode, "_run_fused_kda_decode", fake_backend) + assert kda_decode.packed_fused_kda_decode(**arguments) == "output" + return captured + + +def test_packed_fused_kda_decode_t2_rejects_t1_state_indices(monkeypatch): + arguments = _arguments(num_sequences=4, num_rows=4) + arguments["state_indices"] = torch.ones((4, 2), dtype=torch.int32) + arguments["t1_state_indices"] = torch.ones((4,), dtype=torch.int32) + monkeypatch.setattr(kda_decode, "_run_fused_kda_decode", lambda **kwargs: "bad") + with pytest.raises(ValueError, match="only supported for T=1"): + kda_decode.packed_fused_kda_decode(**arguments) + + +def test_packed_fused_kda_decode_t2_forwards_none_lower_bound(monkeypatch): + captured = {} + + def fake_backend(**kwargs): + captured.update(kwargs) + return "output" + + monkeypatch.setattr(kda_decode, "_run_packed_fused_kda_decode", fake_backend) + arguments = _arguments(num_sequences=4, num_rows=4) + arguments["state_indices"] = torch.ones((4, 2), dtype=torch.int32) + arguments["query_start_loc"] = torch.arange(5, dtype=torch.int32) + arguments["num_accepted_tokens"] = torch.ones(4, dtype=torch.int32) + assert kda_decode.packed_fused_kda_decode(**arguments, lower_bound=None) == "output" + assert captured["lower_bound"] is None + + +def test_packed_fused_kda_decode_t1_forwards_metadata_without_remap(monkeypatch): + arguments = _arguments(num_sequences=3, num_rows=5) + arguments["state_indices"] = torch.tensor([[11], [22], [33]], dtype=torch.int32) + arguments["query_start_loc"] = torch.tensor([0, 1, 1, 2], dtype=torch.int32) + captured = _capture_metadata(monkeypatch, arguments) + assert captured["query_start_loc"] is arguments["query_start_loc"] + assert captured["packed_state_indices"] is arguments["state_indices"] + assert ( + captured["state_indices"].data_ptr() + == arguments["state_indices"][:, 0].data_ptr() + ) + assert captured["state_indices"].shape == (3,) + + +def test_packed_fused_kda_decode_t1_forwards_pre_resolved_rows(monkeypatch): + arguments = _arguments(num_sequences=3, num_rows=5) + arguments["state_indices"] = torch.tensor([[11], [22], [33]], dtype=torch.int32) + arguments["query_start_loc"] = torch.tensor([0, 1, 1, 2], dtype=torch.int32) + arguments["t1_state_indices"] = torch.tensor([11, 0, 33, 0, 0], dtype=torch.int32) + captured = _capture_metadata(monkeypatch, arguments) + assert captured["state_indices"] is arguments["t1_state_indices"] + assert "query_start_loc" not in captured + assert "packed_state_indices" not in captured + + +@pytest.mark.parametrize( + ("field", "value", "error"), + [ + ( + "query_start_loc", + torch.tensor([0, 1, 2, 3, 4], dtype=torch.int64), + "dtype", + ), + ("state_indices", torch.ones((4, 1), dtype=torch.int64), "dtype"), + ("query_start_loc", torch.arange(4, dtype=torch.int32), "shape"), + ( + "query_start_loc", + torch.arange(10, dtype=torch.int32)[::2], + "contiguous", + ), + ], +) +def test_packed_fused_kda_decode_t1_validates_remap_structure( + monkeypatch, field, value, error +): + arguments = _arguments() + arguments[field] = value + monkeypatch.setattr(kda_decode, "_run_fused_kda_decode", lambda **kwargs: "output") + with pytest.raises((TypeError, ValueError), match=error): + kda_decode.packed_fused_kda_decode(**arguments) + + +def test_packed_fused_kda_decode_cache_getters_are_cached(monkeypatch): + kernel_module = __import__( + "flashinfer.kda_kernels.fused_kda_decode_multitoken", fromlist=["module"] + ) + calls = [] + + def fake_build(*args, **kwargs): + del kwargs + calls.append(args) + return object() + + monkeypatch.setattr(kernel_module, "build_and_load_cute_dsl_kernel", fake_build) + assert not hasattr(kernel_module, "_get_t1_remap_kernel") + assert not hasattr(kernel_module, "remap_packed_t1_state_indices") + assert hasattr(kernel_module._get_compiled_kernel, "cache_info") + assert hasattr(kernel_module._get_compiled_kernel, "cache_clear") + kernel_module._get_compiled_kernel.cache_clear() + try: + first = kernel_module._get_compiled_kernel(3, 12, -5.0, 1e-5, 2, 16, 1, 1, 1) + second = kernel_module._get_compiled_kernel(3, 12, -5.0, 1e-5, 2, 16, 1, 1, 1) + assert first is second + assert len(calls) == 1 + finally: + kernel_module._get_compiled_kernel.cache_clear() + + +def test_packed_fused_kda_decode_metadata_is_required(): + arguments = _arguments() + positional = [ + arguments.get(name, parameter.default) + for name, parameter in list( + inspect.signature(kda_decode.fused_kda_decode).parameters.items() + )[:14] + ] + with pytest.raises(TypeError): + kda_decode.packed_fused_kda_decode(*positional) # type: ignore[call-arg] + with pytest.raises(TypeError): + packed = cast(Any, kda_decode.packed_fused_kda_decode) + packed( + *positional, + arguments["query_start_loc"], + arguments["num_accepted_tokens"], + ) diff --git a/tests/trace/example.py b/tests/trace/example.py index 6391f02afcd..459e057757a 100644 --- a/tests/trace/example.py +++ b/tests/trace/example.py @@ -33,6 +33,7 @@ recurrent_kda_q8_v16_d128.json packed_kda_decode_h12_d128.json fused_kda_decode_h12_d128.json +packed_fused_kda_decode_t3_h12_d128.json gemm_bf16_N256_K7168.json gemm_bf16_N4096_K4096.json gemm_fp4_N2048_K7168_block_size16.json @@ -1079,6 +1080,40 @@ def example_dsv41_fp4_cache(): fk_norm_weight, ) +# ── packed fused Kimi K3 speculative decode (conv + KDA + gated RMSNorm) ──── +fkp_N, fkp_T = 2, 3 +fkp_rows = fkp_N * fkp_T +fkp_slots = fkp_rows + 1 +fkp_x = torch.randn(fkp_rows, 3 * fk_hidden, dtype=torch.bfloat16, device=device) +fkp_conv_storage = torch.zeros( + fkp_slots, fkp_T + 2, 3 * fk_hidden, dtype=torch.bfloat16, device=device +) +fkp_conv_state = fkp_conv_storage.transpose(1, 2) +fkp_raw_gate = torch.randn(1, fkp_rows, fk_H, fk_D, dtype=torch.bfloat16, device=device) +fkp_raw_beta = torch.randn(1, fkp_rows, fk_H, dtype=torch.bfloat16, device=device) +fkp_indices = torch.arange(fkp_rows, 0, -1, dtype=torch.int32, device=device).reshape( + fkp_N, fkp_T +) +fkp_query_start = torch.arange(0, fkp_rows + 1, fkp_T, dtype=torch.int32, device=device) +fkp_accepted = torch.ones(fkp_N, dtype=torch.int32, device=device) +fkp_state = torch.zeros(fkp_slots, fk_H, fk_D, fk_D, dtype=torch.float32, device=device) +fkp_output_gate = torch.randn(fkp_rows, fk_H, fk_D, dtype=torch.bfloat16, device=device) +flashinfer.kda_decode.packed_fused_kda_decode( + fkp_x, + fk_weight, + fkp_conv_state, + fkp_raw_gate, + fkp_raw_beta, + fk_A_log, + fk_dt_bias, + fkp_indices, + fkp_state, + fkp_output_gate, + fk_norm_weight, + query_start_loc=fkp_query_start, + num_accepted_tokens=fkp_accepted, +) + # ── AlphaMoE fused router (SM100/SM103) ────────────────────────────────────── _alpha_router_logits = torch.randn(32, 512, dtype=torch.float32, device=device) _alpha_router_cc = torch.cuda.get_device_capability(device) diff --git a/tests/trace/fi_trace_out/packed_fused_kda_decode_t3_h12_d128.json b/tests/trace/fi_trace_out/packed_fused_kda_decode_t3_h12_d128.json new file mode 100644 index 00000000000..20f0e890485 --- /dev/null +++ b/tests/trace/fi_trace_out/packed_fused_kda_decode_t3_h12_d128.json @@ -0,0 +1,204 @@ +{ + "name": "packed_fused_kda_decode_t3_h12_d128", + "description": "Packed ragged T>=1 Kimi width-four causal convolution, recurrent KDA, and gated RMSNorm with one recurrent checkpoint per token.", + "op_type": "kda", + "tags": [ + "fi_api:flashinfer.kda_decode.packed_fused_kda_decode", + "stage:decode", + "status:experimental" + ], + "axes": { + "num_sequences": { + "type": "var", + "description": "Number of independent sequences." + }, + "num_tokens": { + "type": "const", + "value": 3, + "description": "Maximum tokens per sequence." + }, + "num_rows": { + "type": "var", + "description": "Total packed tokens." + }, + "num_sequence_offsets": { + "type": "var", + "description": "Number of packed sequence boundaries." + }, + "num_heads": { + "type": "const", + "value": 12, + "description": "Number of KDA heads." + }, + "head_dim": { + "type": "const", + "value": 128, + "description": "KDA head dimension." + }, + "singleton": { + "type": "const", + "value": 1, + "description": "Leading singleton dimension." + }, + "projection_groups": { + "type": "const", + "value": 3, + "description": "Packed QKV projection groups." + }, + "hidden_size": { + "type": "const", + "value": 1536, + "description": "Channels in one projection." + }, + "qkv_width": { + "type": "const", + "value": 4608, + "description": "Packed QKV width." + }, + "conv_width": { + "type": "const", + "value": 4, + "description": "Depthwise convolution width." + }, + "conv_history": { + "type": "const", + "value": 5, + "description": "Cached convolution history." + }, + "num_slots": { + "type": "var", + "description": "Number of cache slots." + } + }, + "constraints": [ + "num_rows <= num_sequences * num_tokens", + "num_sequence_offsets == num_sequences + 1", + "num_tokens >= 1", + "qkv_width == 3 * num_heads * head_dim", + "hidden_size == num_heads * head_dim", + "singleton == 1", + "projection_groups == 3", + "head_dim == 128", + "num_heads in (12, 24, 32, 48, 96)", + "conv_width == 4", + "conv_history == num_tokens + 2" + ], + "inputs": { + "x": { + "shape": [ + "num_rows", + "qkv_width" + ], + "dtype": "bfloat16" + }, + "weight": { + "shape": [ + "projection_groups", + "conv_width", + "hidden_size" + ], + "dtype": "float32" + }, + "conv_state": { + "shape": [ + "num_slots", + "qkv_width", + "conv_history" + ], + "dtype": "bfloat16" + }, + "raw_gate": { + "shape": [ + "singleton", + "num_rows", + "num_heads", + "head_dim" + ], + "dtype": "bfloat16" + }, + "raw_beta": { + "shape": [ + "singleton", + "num_rows", + "num_heads" + ], + "dtype": "bfloat16" + }, + "A_log": { + "shape": [ + "num_heads" + ], + "dtype": "float32" + }, + "dt_bias": { + "shape": [ + "hidden_size" + ], + "dtype": "float32" + }, + "state_indices": { + "shape": [ + "num_sequences", + "num_tokens" + ], + "dtype": "int32" + }, + "query_start_loc": { + "shape": [ + "num_sequence_offsets" + ], + "dtype": "int32" + }, + "num_accepted_tokens": { + "shape": [ + "num_sequences" + ], + "dtype": "int32" + }, + "state": { + "shape": [ + "num_slots", + "num_heads", + "head_dim", + "head_dim" + ], + "dtype": "float32" + }, + "output_gate": { + "shape": [ + "num_rows", + "num_heads", + "head_dim" + ], + "dtype": "bfloat16" + }, + "norm_weight": { + "shape": [ + "head_dim" + ], + "dtype": "float32" + }, + "lower_bound": { + "shape": null, + "dtype": "float32", + "optional": true + }, + "norm_eps": { + "shape": null, + "dtype": "float32", + "optional": true + } + }, + "outputs": { + "output": { + "shape": [ + "singleton", + "num_rows", + "num_heads", + "head_dim" + ], + "dtype": "bfloat16" + } + }, + "check": "def standard_check(\n reference_outputs: Any,\n actual_outputs: Any,\n *,\n rtol: Optional[float] = None,\n atol: Optional[float] = None,\n max_mismatch_pct: float = 0.0,\n min_cos_sim: Optional[float] = 1.0 - 1e-3,\n) -> bool:\n \"\"\"Default trace correctness check used when a template does not override it.\"\"\"\n from flashinfer.trace import default_check\n\n return default_check(\n reference_outputs,\n actual_outputs,\n rtol=rtol,\n atol=atol,\n max_mismatch_pct=max_mismatch_pct,\n min_cos_sim=min_cos_sim,\n )\n" +}