diff --git a/benchmarks/bench_kda_packed_decode.py b/benchmarks/bench_kda_packed_decode.py new file mode 100644 index 0000000..d9e13e7 --- /dev/null +++ b/benchmarks/bench_kda_packed_decode.py @@ -0,0 +1,264 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +"""bench_kda_packed_decode.py — micro-benchmark: packed vs non-packed KDA decode. + +Compares single-token (T=1) decode routes that share the same CuTe DSL kernel +body, so the delta isolates host-side orchestration + the q/k/v materialization +that a real caller pays. + +Two routes are timed, BOTH seeded from the same packed ``mixed_qkv`` (the form +a conv layer actually produces, e.g. in SGLang's KDA decode): + + 1. non-packed (realistic): the caller must turn ``mixed_qkv`` back into + separate contiguous q/k/v before calling ``kda_decode`` — i.e. a per-call + ``split + unsqueeze + contiguous``. The ``.contiguous()`` is a REAL copy + when ``N>1`` because the split leaves row stride = qkv_dim. This is the + cost the packed path is meant to remove. + 2. packed: ``cula.kda.kda_packed_decode`` feeds q/k/v as strided views + directly — no materialization, no ``.contiguous()`` copy. + +A third "non-packed (pre-split)" column is included as an oracle: it uses +q/k/v that were split once outside the timed loop (no per-call copy). This +shows the floor of the non-packed approach — i.e. how fast decode could be if +the caller already had separate contiguous q/k/v. packed matching or beating +the realistic column while staying near the oracle is the win. + +All routes are timed with CUDA events bracketing the full callable (host-side +view construction + cache lookup + kernel launch). The ``mixed_qkv`` is reused +across iterations; only the per-call split+contiguous is inside the timed +region for the realistic non-packed route. + +Usage: + python benchmarks/bench_kda_packed_decode.py + python benchmarks/bench_kda_packed_decode.py --batch-sizes 1 4 16 64 128 + python benchmarks/bench_kda_packed_decode.py --head-pairs 8:8 16:16 32:32 64:64 + python benchmarks/bench_kda_packed_decode.py --ncu +""" + +import argparse +import pathlib +import sys + +import torch + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from benchmarks.utils import benchmark_cuda_fn +from cula.kda import fused_sigmoid_gating_delta_rule_update as cula_fused +from cula.kda import kda_packed_decode + + +def make_inputs(N, H, HV, K, V, device="cuda", seed=42): + torch.manual_seed(seed) + q = torch.randn(N, H, K, device=device, dtype=torch.bfloat16) + k = torch.randn(N, H, K, device=device, dtype=torch.bfloat16) + v = torch.randn(N, HV, V, device=device, dtype=torch.bfloat16) + a = (torch.randn(N, HV, K, device=device, dtype=torch.float32) * 0.1).to(torch.bfloat16) + b = torch.randn(N, HV, device=device, dtype=torch.bfloat16) + A_log = -torch.rand(HV, device=device, dtype=torch.float32) * 2 + dt_bias = torch.randn(HV, K, device=device, dtype=torch.float32) * 0.1 + state = torch.randn(N, HV, V, K, device=device, dtype=torch.float32) * 0.01 + return q, k, v, a, b, A_log, dt_bias, state + + +def run_config(N, H, HV, K, V, warmup, rep): + device = "cuda" + scale = K**-0.5 + + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V, device) + + q_4d = q.unsqueeze(1).contiguous() + k_4d = k.unsqueeze(1).contiguous() + v_4d = v.unsqueeze(1).contiguous() + a_flat = a.reshape(N, 1, -1).contiguous() + b_3d = b.unsqueeze(1).contiguous() + mixed_qkv = torch.cat([q.view(N, -1), k.view(N, -1), v.view(N, -1)], dim=-1).contiguous() + qk_dim = H * K + v_dim = HV * V + indices = torch.arange(N, device=device, dtype=torch.int32) + + state_init = state.clone().contiguous() + + # non-packed (oracle): separate contiguous q/k/v split once outside the loop. + def call_oracle(state_buf): + return cula_fused( + A_log=A_log, + a=a_flat, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=q_4d, + k=k_4d, + v=v_4d, + b=b_3d, + initial_state_source=state_buf, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + is_kda=True, + state_layout="vk", + ) + + # non-packed (realistic): start from mixed_qkv, split+unsqueeze+contiguous + # every call — the cost a real SGLang caller pays when q/k/v are not kept + # pre-split. + def call_nonpacked(state_buf): + qq, kk, vv = mixed_qkv.split([qk_dim, qk_dim, v_dim], dim=-1) + return cula_fused( + A_log=A_log, + a=a_flat, + dt_bias=dt_bias, + softplus_beta=1.0, + softplus_threshold=20.0, + q=qq.view(N, 1, H, K).contiguous(), + k=kk.view(N, 1, H, K).contiguous(), + v=vv.view(N, 1, HV, V).contiguous(), + b=b_3d, + initial_state_source=state_buf, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + is_kda=True, + state_layout="vk", + ) + + # packed: kda_packed_decode route (mixed_qkv reused every iteration) + def call_packed(state_buf): + return kda_packed_decode( + mixed_qkv, + a_flat, + b_3d, + A_log=A_log, + dt_bias=dt_bias, + state=state_buf, + state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + state_layout="vk", + ) + + # Correctness sanity: packed must match the oracle. + state_ora = state_init.clone() + state_pck = state_init.clone() + with torch.no_grad(): + o_ora = call_oracle(state_ora) + o_pck = call_packed(state_pck) + out_diff = (o_ora.float() - o_pck.float()).abs().max().item() + state_diff = (state_ora.float() - state_pck.float()).abs().max().item() + + state_bench_ora = state_init.clone() + state_bench_non = state_init.clone() + state_bench_pck = state_init.clone() + + def setup_ora(): + state_bench_ora.copy_(state_init) + + def setup_non(): + state_bench_non.copy_(state_init) + + def setup_pck(): + state_bench_pck.copy_(state_init) + + with torch.no_grad(): + t_ora = benchmark_cuda_fn(lambda: call_oracle(state_bench_ora), setup_fn=setup_ora, warmup=warmup, rep=rep) + t_non = benchmark_cuda_fn(lambda: call_nonpacked(state_bench_non), setup_fn=setup_non, warmup=warmup, rep=rep) + t_pck = benchmark_cuda_fn(lambda: call_packed(state_bench_pck), setup_fn=setup_pck, warmup=warmup, rep=rep) + + # q/k/v bytes the non-packed path copies per call (split+contiguous): bf16. + qkv_dim = mixed_qkv.shape[1] + copy_bytes = 2 * (2 * H * K + HV * V) * N + + return { + "N": N, + "H": H, + "HV": HV, + "K": K, + "V": V, + "qkv_dim": qkv_dim, + "t_oracle_ms": t_ora, + "t_non_ms": t_non, + "t_packed_ms": t_pck, + "saved_vs_non_us": (t_non - t_pck) * 1e3, + "saved_vs_ora_us": (t_ora - t_pck) * 1e3, + "speedup_vs_non": t_non / t_pck if t_pck > 0 else float("inf"), + "out_diff": out_diff, + "state_diff": state_diff, + "copy_bytes": copy_bytes, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 4, 16, 64, 128]) + parser.add_argument("--Hs", type=int, nargs="+", default=[8, 16]) + parser.add_argument("--HV", type=int, default=16) + parser.add_argument( + "--head-pairs", + type=str, + nargs="+", + default=None, + help="Explicit H:HV pairs, e.g. --head-pairs 8:8 16:16 32:32 64:64", + ) + parser.add_argument("--K", type=int, default=128) + parser.add_argument("--V", type=int, default=128) + parser.add_argument("--warmup", type=int, default=30) + parser.add_argument("--rep", type=int, default=200) + parser.add_argument("--ncu", action="store_true") + args = parser.parse_args() + + if args.ncu: + args.warmup, args.rep = 1, 1 + + gpu = torch.cuda.get_device_name(0) + print(f"# cuLA Packed KDA Decode micro-bench [{gpu}]") + print(f"# K={args.K} V={args.V} warmup={args.warmup} rep={args.rep}") + print() + + if args.head_pairs is None: + head_pairs = [(H, args.HV if args.HV >= 2 * H else 2 * H) for H in args.Hs] + else: + head_pairs = [] + for item in args.head_pairs: + try: + h_str, hv_str = item.split(":", 1) + H, HV = int(h_str), int(hv_str) + except ValueError as exc: + raise ValueError(f"Invalid --head-pairs entry {item!r}; expected H:HV, e.g. 8:8") from exc + if H <= 0 or HV <= 0 or HV % H != 0: + raise ValueError(f"Invalid head pair H={H}, HV={HV}; expected positive values with HV % H == 0") + head_pairs.append((H, HV)) + + for H, HV in head_pairs: + print(f"## H={H} HV={HV} K={args.K} V={args.V}") + hdr = ( + f"{'N':>5} | {'qkv_dim':>7} | {'oracle (ms)':>12} | {'non-packed (ms)':>16} " + f"| {'cula_packed (ms)':>17} | {'save vs non (us)':>16} | {'speedup vs non':>15} " + f"| {'out_diff':>9} | {'state_diff':>10}" + ) + print(hdr) + print("-" * len(hdr)) + for N in args.batch_sizes: + r = run_config(N, H, HV, args.K, args.V, args.warmup, args.rep) + print( + f"{r['N']:>5} | {r['qkv_dim']:>7} | {r['t_oracle_ms']:>12.4f} | {r['t_non_ms']:>16.4f} " + f"| {r['t_packed_ms']:>17.4f} | {r['saved_vs_non_us']:>16.2f} | {r['speedup_vs_non']:>14.3f}x " + f"| {r['out_diff']:>9.2e} | {r['state_diff']:>10.2e}" + ) + print() + + +if __name__ == "__main__": + main() diff --git a/cula/kda/__init__.py b/cula/kda/__init__.py index 0b61b13..73f3321 100644 --- a/cula/kda/__init__.py +++ b/cula/kda/__init__.py @@ -20,6 +20,7 @@ "kda_decode_mtp", "kda_decode_mtp_recurrent", "kda_decode_mtp_recurrent_ws", + "kda_packed_decode", "fused_sigmoid_gating_delta_rule_update", "kda_prefill_hopper", "kda_prefill_hopper_opt", @@ -29,6 +30,8 @@ _LAZY = { "chunk_kda": ("cula.kda.chunk", "chunk_kda"), "kda_prefill_hopper": ("cula.kda.hopper_fused_fwd", "cula_kda_prefill"), + "kda_prefill_hopper_opt": ("cula.kda.hopper_fused_fwd_opt", "cula_kda_prefill_opt"), + "kda_prefill_hopper_auto": ("cula.kda.auto_route", "cula_kda_prefill_auto"), "kda_decode": ("cula.ops.kda.decode.cute", "kda_decode"), "kda_decode_mtp": ("cula.ops.kda.decode.mtp", "kda_decode_mtp"), "kda_decode_mtp_recurrent": ("cula.ops.kda.decode.mtp", "kda_decode_mtp_recurrent"), @@ -36,6 +39,7 @@ "cula.ops.kda.decode.mtp", "kda_decode_mtp_recurrent_ws", ), + "kda_packed_decode": ("cula.ops.kda.decode.cute", "kda_packed_decode"), "fused_sigmoid_gating_delta_rule_update": ( "cula.ops.kda.decode.cute", "fused_sigmoid_gating_delta_rule_update", diff --git a/cula/ops/kda/decode/cute.py b/cula/ops/kda/decode/cute.py index d84c77b..37c6025 100644 --- a/cula/ops/kda/decode/cute.py +++ b/cula/ops/kda/decode/cute.py @@ -285,6 +285,199 @@ def _try_fast_dense_decode( return o +def _try_fast_dense_packed_decode( + A_log: torch.Tensor, + dt_bias: torch.Tensor, + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + *, + N: int, + H: int, + HV: int, + K: int, + V: int, + initial_state_source: torch.Tensor, + initial_state_indices: torch.Tensor, + is_varlen_decode: bool, + cu_seqlens: torch.Tensor | None, + scale: float | None, + use_qk_l2norm_in_kernel: bool, + softplus_beta: float, + softplus_threshold: float, + out: torch.Tensor | None, + state_layout: str | None, +): + """Fast path for packed-QKV decode. + + Mirrors ``_try_fast_dense_decode`` but takes a packed ``mixed_qkv`` of shape + ``[N, qkv_dim]`` (``= [Q(H·K) | K(H·K) | V(HV·V)]``, head-major, last dim + contiguous) and constructs the q/k/v as strided views instead of requiring + three separate contiguous tensors — so no q/k/v materialization or + ``.contiguous()`` copy is needed. + + Returns ``None`` (falling back to the general path) when the inputs are not + already in the exact kernel-ready layout/dtype; never raises on shape/dtype + mismatches. + """ + if K != TILE_K or V % TILE_V_SMALL != 0 or V % TILE_V != 0: + return None + + qkv_dim = 2 * H * K + HV * V + if ( + mixed_qkv.ndim != 2 + or mixed_qkv.shape != (N, qkv_dim) + or mixed_qkv.stride(-1) != 1 + or mixed_qkv.stride(0) < qkv_dim + or mixed_qkv.device.type != "cuda" + or mixed_qkv.dtype != torch.bfloat16 + ): + return None + + if ( + initial_state_source.dtype != torch.float32 + or initial_state_indices.dtype != torch.int32 + or initial_state_indices.shape != (N,) + or A_log.dtype != torch.float32 + or dt_bias.dtype != torch.float32 + ): + return None + + if not ( + initial_state_source.is_contiguous() + and initial_state_indices.is_contiguous() + and A_log.is_contiguous() + and dt_bias.is_contiguous() + ): + return None + if A_log.numel() != HV or dt_bias.shape != (HV, K): + return None + + normalized_layout = "vk" if state_layout is None else str(state_layout).strip().lower() + if normalized_layout == "vk": + if initial_state_source.ndim != 4 or initial_state_source.shape[1:] != (HV, V, K): + return None + state_layout_is_kv = False + elif normalized_layout == "kv": + if initial_state_source.ndim != 4 or initial_state_source.shape[1:] != (HV, K, V): + return None + state_layout_is_kv = True + else: + return None + + if not a.is_contiguous() or a.device != mixed_qkv.device or a.dtype != torch.bfloat16: + return None + if is_varlen_decode: + # varlen kernel compiled a shape: (N, HV, K) -- 3D + if a.dim() == 3 and a.shape == (N, HV, K): + a_kernel = a + else: + return None + else: + # dense kernel compiled a shape: (N, 1, HV, K) -- 4D + if a.dim() == 4 and a.shape == (N, 1, HV, K): + a_kernel = a + elif a.dim() == 3 and a.shape == (N, 1, HV * K): + a_kernel = a.view(N, 1, HV, K) + elif a.dim() == 3 and a.shape == (N, HV, K): + a_kernel = a.unsqueeze(1) + else: + return None + + if b.device != mixed_qkv.device or b.dtype != torch.bfloat16 or not b.is_contiguous(): + return None + if is_varlen_decode: + # varlen b compiled: (N, HV) -- 2D + if b.dim() == 2 and b.shape == (N, HV): + b_kernel = b + else: + return None + else: + # dense b compiled: (N, 1, HV) -- 3D + if b.dim() == 3 and b.shape == (N, 1, HV): + b_kernel = b + elif b.dim() == 2 and b.shape == (N, HV): + b_kernel = b.unsqueeze(1) + else: + return None + + if scale is None: + scale = K**-0.5 + elif scale <= 0: + return None + + # Construct strided q/k/v views (row stride may include padding, last dim contiguous) + # and the matching output. Shapes align with the compile-time mock in the + # packed compile paths: varlen -> [1,N,...], dense -> [N,1,...]. + if is_varlen_decode: + q_view = mixed_qkv.narrow(1, 0, H * K).view(1, N, H, K) + k_view = mixed_qkv.narrow(1, H * K, H * K).view(1, N, H, K) + v_view = mixed_qkv.narrow(1, 2 * H * K, HV * V).view(1, N, HV, V) + o = _prepare_output_tensor(mixed_qkv, out, (1, N, HV, V)) + else: + q_view = mixed_qkv.narrow(1, 0, H * K).view(N, 1, H, K) + k_view = mixed_qkv.narrow(1, H * K, H * K).view(N, 1, H, K) + v_view = mixed_qkv.narrow(1, 2 * H * K, HV * V).view(N, 1, HV, V) + o = _prepare_output_tensor(mixed_qkv, out, (N, 1, HV, V)) + + if cu_seqlens is not None: + if cu_seqlens.dtype != torch.int32 or cu_seqlens.numel() != N + 1 or not cu_seqlens.is_contiguous(): + return None + cu_seqlens_to_use = cu_seqlens + else: + cache_key = (N, str(mixed_qkv.device)) + if cache_key not in _cu_seqlens_cache: + _cu_seqlens_cache[cache_key] = torch.arange(N + 1, dtype=torch.int32, device=mixed_qkv.device) + cu_seqlens_to_use = _cu_seqlens_cache[cache_key] + + use_small_batch = N < SMALL_BATCH_THRESHOLD + if is_varlen_decode: + dense_small_hv_parallel = False + else: + dense_small_hv_parallel_head_threshold = ( + N4_DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD if N <= 4 else DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD + ) + dense_small_hv_parallel = ( + use_small_batch and dense_small_hv_parallel_head_threshold >= H and N <= DENSE_SMALL_HV_PARALLEL_MAX_N + ) + num_blocks_per_state_small = _select_small_blocks_per_state(N, H, HV, V) + + compiled_kernel = _get_compiled_packed_kernel( + N, + H, + HV, + K, + V, + mixed_qkv.stride(0), + initial_state_source.shape[0], + use_small_batch, + is_varlen_decode, + scale=scale, + use_qk_l2norm=use_qk_l2norm_in_kernel, + state_layout_is_kv=state_layout_is_kv, + precomputed_decay_beta=False, + num_blocks_per_state_small=num_blocks_per_state_small, + dense_small_hv_parallel=dense_small_hv_parallel, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + ) + compiled_kernel( + cu_seqlens_to_use, + q_view, + k_view, + v_view, + a_kernel, + b_kernel, + A_log, + dt_bias, + initial_state_source, + initial_state_indices, + o, + _get_cached_stream(mixed_qkv.device), + ) + return o + + def _define_kernels(): """Define CuTe DSL kernels for KDA normal and varlen decode modes.""" @@ -1668,6 +1861,149 @@ def _get_compiled_kernel( return compiled_kernel +def _get_compiled_packed_kernel( + N, + H, + HV, + K, + V, + row_stride, + pool_size, + use_small_batch, + is_varlen_decode, + scale, + use_qk_l2norm, + state_layout_is_kv, + precomputed_decay_beta, + num_blocks_per_state_small, + dense_small_hv_parallel, + softplus_beta, + softplus_threshold, +): + """Get or lazily compile a packed-QKV kernel with static packed strides. + + The q/k/v mock tensors are strided views sliced from a packed + ``mixed_qkv`` mock, so their row stride is the compile-time constant + ``row_stride``. Runtime packed views with the same shape/stride can then + use the static layout specialization while still avoiding q/k/v + materialization. + """ + global _compiled_kernels + + qkv_dim = 2 * H * K + HV * V + key = ( + N, + H, + HV, + K, + V, + qkv_dim, + row_stride, + pool_size, + use_small_batch, + is_varlen_decode, + scale, + use_qk_l2norm, + state_layout_is_kv, + precomputed_decay_beta, + num_blocks_per_state_small, + dense_small_hv_parallel, + softplus_beta, + softplus_threshold, + "packed_qkv", + ) + if key in _compiled_kernels: + return _compiled_kernels[key] + + cu_seqlens = torch.zeros(N + 1, dtype=torch.int32, device="cuda") + if row_stride < qkv_dim: + raise ValueError(f"row_stride={row_stride} must be >= qkv_dim={qkv_dim}") + mixed_qkv = torch.zeros(N, row_stride, dtype=torch.bfloat16, device="cuda") + + if is_varlen_decode: + q = mixed_qkv.narrow(1, 0, H * K).view(1, N, H, K) + k = mixed_qkv.narrow(1, H * K, H * K).view(1, N, H, K) + v = mixed_qkv.narrow(1, 2 * H * K, HV * V).view(1, N, HV, V) + a = torch.zeros(N, HV, K, dtype=torch.bfloat16, device="cuda") + b = torch.zeros(N, HV, dtype=torch.bfloat16, device="cuda") + o = torch.zeros(1, N, HV, V, dtype=torch.bfloat16, device="cuda") + else: + q = mixed_qkv.narrow(1, 0, H * K).view(N, 1, H, K) + k = mixed_qkv.narrow(1, H * K, H * K).view(N, 1, H, K) + v = mixed_qkv.narrow(1, 2 * H * K, HV * V).view(N, 1, HV, V) + a = torch.zeros(N, 1, HV, K, dtype=torch.bfloat16, device="cuda") + b = torch.zeros(N, 1, HV, dtype=torch.bfloat16, device="cuda") + o = torch.zeros(N, 1, HV, V, dtype=torch.bfloat16, device="cuda") + + A_log = torch.zeros(HV, dtype=torch.float32, device="cuda") + dt_bias = torch.zeros(HV, K, dtype=torch.float32, device="cuda") + if state_layout_is_kv: + h0_source = torch.zeros(pool_size, HV, K, V, dtype=torch.float32, device="cuda") + else: + h0_source = torch.zeros(pool_size, HV, V, K, dtype=torch.float32, device="cuda") + h0_indices = torch.zeros(N, dtype=torch.int32, device="cuda") + + cu_seqlens_tensor = from_dlpack(cu_seqlens, assumed_align=16) + q_tensor = from_dlpack(q, assumed_align=16) + k_tensor = from_dlpack(k, assumed_align=16) + v_tensor = from_dlpack(v, assumed_align=16) + a_tensor = from_dlpack(a, assumed_align=16) + b_tensor = from_dlpack(b, assumed_align=16) + A_log_tensor = from_dlpack(A_log, assumed_align=16) + dt_bias_tensor = from_dlpack(dt_bias, assumed_align=16) + h0_source_tensor = from_dlpack(h0_source, assumed_align=16) + h0_indices_tensor = from_dlpack(h0_indices, assumed_align=16) + o_tensor = from_dlpack(o, assumed_align=16) + + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + run_small, run_small_varlen, run_large, run_large_varlen = _get_jit_functions() + if use_small_batch: + kernel_func = run_small_varlen if is_varlen_decode else run_small + else: + kernel_func = run_large_varlen if is_varlen_decode else run_large + + compiled_kernel = cute.compile( + kernel_func, + cu_seqlens_tensor, + q_tensor, + k_tensor, + v_tensor, + a_tensor, + b_tensor, + A_log_tensor, + dt_bias_tensor, + h0_source_tensor, + h0_indices_tensor, + o_tensor, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + scale=scale, + B=1 if is_varlen_decode else N, + T=N if is_varlen_decode else 1, + H=H, + K=K, + V=V, + HV=HV, + use_initial_state=True, + use_qk_l2norm=use_qk_l2norm, + state_layout_is_kv=state_layout_is_kv, + precomputed_decay_beta=precomputed_decay_beta, + num_blocks_per_state_small=num_blocks_per_state_small, + dense_small_hv_parallel=dense_small_hv_parallel, + stream=stream, + options="--enable-tvm-ffi --opt-level 1", + ) + + _compiled_kernels[key] = compiled_kernel + logger.info( + "CuTe DSL KDA packed static-stride kernel compiled: " + f"N={N}, H={H}, HV={HV}, K={K}, V={V}, qkv_dim={qkv_dim}, row_stride={row_stride}, pool_size={pool_size}, " + f"small_batch={use_small_batch}, varlen={is_varlen_decode}" + ) + return compiled_kernel + + def _normalize_A_log(A_log: torch.Tensor, HV: int) -> torch.Tensor: if A_log.numel() != HV: raise ValueError(f"Unexpected A_log shape: {A_log.shape}; expected numel={HV}") @@ -2094,3 +2430,227 @@ def kda_decode( ) return o + + +def kda_packed_decode( + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + *, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + state: torch.Tensor, + state_indices: torch.Tensor, + out: torch.Tensor | None = None, + scale: float | None = None, + use_qk_l2norm_in_kernel: bool = True, + softplus_beta: float = 1.0, + softplus_threshold: float = 20.0, + cu_seqlens: torch.Tensor | None = None, + state_layout: str = "vk", +) -> torch.Tensor: + """Packed-QKV variant of :func:`kda_decode`. + + Takes a packed ``mixed_qkv`` of shape ``[N, qkv_dim]`` laid out as + ``[Q(H·K) | K(H·K) | V(HV·V)]`` (head-major, last dim contiguous, row stride + >= qkv_dim) and feeds the existing CuTe DSL KDA decode kernel q/k/v as + strided views directly — avoiding the q/k/v materialization and the + ``.contiguous()`` copy that ``kda_decode`` performs. + + Numerics match ``kda_decode`` on the same inputs repacked, since the kernel + body is unchanged. The packed static-stride compile path builds q/k/v mock + tensors with the same row stride as runtime packed views, so the kernel + accepts the non-contiguous views without dynamic layout. + + Args: + mixed_qkv: ``[N, qkv_dim]`` bf16, last dim contiguous. + a: gate, dense ``(N,1,HV,K)`` / ``(N,HV,K)`` / varlen-compatible; see + ``_normalize_kda_a``. + b: ``(N,1,HV)`` (dense) or ``(N,HV)`` (varlen), bf16. + A_log: ``(HV,)`` fp32. + dt_bias: ``(HV,K)`` fp32. + state: ``(num_slots, HV, V, K)`` for ``state_layout='vk'`` or + ``(num_slots, HV, K, V)`` for ``'kv'``, fp32. + state_indices: ``(N,)`` int32, ``-1`` marks dummy slots. + out: optional preallocated output, dense ``(N,1,HV,V)`` or varlen + ``(1,N,HV,V)``, bf16, contiguous. + cu_seqlens: when not None, varlen decode (otherwise dense). + state_layout: ``"vk"`` (default) or ``"kv"``. + + Returns: + Output tensor of shape ``(N,1,HV,V)`` (dense) or ``(1,N,HV,V)`` + (varlen). ``state`` is updated in place. + """ + state_layout_canonical = _canonicalize_state_layout(state_layout) + state_layout_is_kv = state_layout_canonical == "kv" + + if state.dim() != 4: + raise ValueError(f"Unexpected state shape: {state.shape}; expected a 4D state tensor") + + if state_layout_is_kv: + pool_size, HV, K, V = state.shape + else: + pool_size, HV, V, K = state.shape + + if K != TILE_K: + raise ValueError(f"Current CuTe DSL KDA kernel requires K={TILE_K}, got K={K}") + if V % TILE_V_SMALL != 0 or V % TILE_V != 0: + raise ValueError(f"Current CuTe DSL KDA kernel requires V % {TILE_V_SMALL} == 0 and V % {TILE_V} == 0, got V={V}") + + if mixed_qkv.ndim != 2: + raise ValueError(f"mixed_qkv must be 2D [N, qkv_dim], got shape {mixed_qkv.shape}") + qkv_dim = mixed_qkv.shape[1] + qk_dim = qkv_dim - HV * V + if qk_dim % 2 != 0: + raise ValueError( + f"mixed_qkv q/k segment (qkv_dim - HV*V = {qk_dim}) must be even for equal Q|K, got qkv_dim={qkv_dim}" + ) + H = (qk_dim // 2) // K + if H <= 0 or qk_dim // 2 != H * K: + raise ValueError( + f"Inconsistent mixed_qkv layout: q/k segment {qk_dim // 2} not divisible by K={K}, got qkv_dim={qkv_dim}" + ) + if HV % H != 0: + raise ValueError(f"HV={HV} must be divisible by H={H} (grouped-head mapping)") + if qkv_dim != 2 * H * K + HV * V: + raise ValueError( + f"mixed_qkv qkv_dim={qkv_dim} != 2*H*K + HV*V = {2 * H * K + HV * V} for H={H}, HV={HV}, K={K}, V={V}" + ) + + N = state_indices.shape[0] + is_varlen_decode = cu_seqlens is not None + if mixed_qkv.shape[0] != N: + raise ValueError(f"mixed_qkv batch size {mixed_qkv.shape[0]} must match state_indices length {N}") + row_stride = mixed_qkv.stride(0) + if mixed_qkv.stride(-1) != 1 or row_stride < qkv_dim: + raise ValueError( + f"mixed_qkv must use packed-row layout with stride(-1)=1 and stride(0)>=qkv_dim={qkv_dim}; " + f"got stride={mixed_qkv.stride()}" + ) + + if scale is None: + scale = K**-0.5 + elif scale <= 0: + raise ValueError(f"scale must be positive, got {scale}") + + fast_out = _try_fast_dense_packed_decode( + A_log, + dt_bias, + mixed_qkv, + a, + b, + N=N, + H=H, + HV=HV, + K=K, + V=V, + initial_state_source=state, + initial_state_indices=state_indices, + is_varlen_decode=is_varlen_decode, + cu_seqlens=cu_seqlens, + scale=scale, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + out=out, + state_layout=state_layout_canonical, + ) + if fast_out is not None: + return fast_out + + # --- General path: normalize then launch the packed static-stride kernel. --- + # We deliberately do NOT route through kda_decode's slow path, which would + # .contiguous()-copy the q/k/v (defeating packed) and requires separate + # q/k/v tensors. We normalize the meta tensors (a/b/A_log/dt_bias/state/ + # indices) but feed q/k/v as strided views into the packed kernel directly. + h0_source, pool_size, _ = _normalize_state_source( + state, + N=N, + HV=HV, + K=K, + V=V, + device=mixed_qkv.device, + state_layout=state_layout_canonical, + ) + a_kernel = _normalize_kda_a(a, is_varlen_decode=is_varlen_decode, N=N, HV=HV, K=K) + a_kernel = a_kernel if a_kernel.is_contiguous() else a_kernel.contiguous() + if is_varlen_decode: + if b.dim() == 3: + b_kernel = b.squeeze(0) + else: + b_kernel = b + else: + if b.dim() == 2: + b_kernel = b.unsqueeze(1) + else: + b_kernel = b + b_kernel = b_kernel if b_kernel.is_contiguous() else b_kernel.contiguous() + A_log_n = _normalize_A_log(A_log, HV) + dt_bias_n = _normalize_dt_bias(dt_bias, HV, K) + indices_n = _normalize_state_indices(state_indices, N=N, pool_size=pool_size, device=mixed_qkv.device) + + if is_varlen_decode: + q_view = mixed_qkv.narrow(1, 0, H * K).view(1, N, H, K) + k_view = mixed_qkv.narrow(1, H * K, H * K).view(1, N, H, K) + v_view = mixed_qkv.narrow(1, 2 * H * K, HV * V).view(1, N, HV, V) + o = _prepare_output_tensor(mixed_qkv, out, (1, N, HV, V)) + else: + q_view = mixed_qkv.narrow(1, 0, H * K).view(N, 1, H, K) + k_view = mixed_qkv.narrow(1, H * K, H * K).view(N, 1, H, K) + v_view = mixed_qkv.narrow(1, 2 * H * K, HV * V).view(N, 1, HV, V) + o = _prepare_output_tensor(mixed_qkv, out, (N, 1, HV, V)) + + if cu_seqlens is not None: + cu_seqlens_to_use = cu_seqlens.contiguous() + else: + cache_key = (N, str(mixed_qkv.device)) + if cache_key not in _cu_seqlens_cache: + _cu_seqlens_cache[cache_key] = torch.arange(N + 1, dtype=torch.int32, device=mixed_qkv.device) + cu_seqlens_to_use = _cu_seqlens_cache[cache_key] + + use_small_batch = N < SMALL_BATCH_THRESHOLD + if is_varlen_decode: + dense_small_hv_parallel = False + else: + dense_small_hv_parallel_head_threshold = ( + N4_DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD if N <= 4 else DENSE_SMALL_HV_PARALLEL_HEAD_THRESHOLD + ) + dense_small_hv_parallel = ( + use_small_batch and dense_small_hv_parallel_head_threshold >= H and N <= DENSE_SMALL_HV_PARALLEL_MAX_N + ) + num_blocks_per_state_small = _select_small_blocks_per_state(N, H, HV, V) + + compiled_kernel = _get_compiled_packed_kernel( + N, + H, + HV, + K, + V, + row_stride, + pool_size, + use_small_batch, + is_varlen_decode, + scale=scale, + use_qk_l2norm=use_qk_l2norm_in_kernel, + state_layout_is_kv=state_layout_is_kv, + precomputed_decay_beta=False, + num_blocks_per_state_small=num_blocks_per_state_small, + dense_small_hv_parallel=dense_small_hv_parallel, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + ) + compiled_kernel( + cu_seqlens_to_use, + q_view, + k_view, + v_view, + a_kernel, + b_kernel, + A_log_n, + dt_bias_n, + h0_source, + indices_n, + o, + _get_cached_stream(mixed_qkv.device), + ) + return o diff --git a/tests/test_kda_packed_decode.py b/tests/test_kda_packed_decode.py new file mode 100644 index 0000000..259baed --- /dev/null +++ b/tests/test_kda_packed_decode.py @@ -0,0 +1,807 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# 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. + +""" +Unit tests for kda_packed_decode (packed-QKV CuTe DSL KDA decode kernel). + +Numerical ground truth = the existing non-packed ``kda_decode``: the same base +inputs (q, k, v, a, b, A_log, dt_bias, state) are repacked into a mixed_qkv of +shape ``[N, qkv_dim] = [Q(H·K) | K(H·K) | V(HV·V)]`` (head-major, last dim +contiguous) and ``kda_packed_decode`` is checked against ``kda_decode``. +""" + +import pathlib +import sys + +import pytest +import torch + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from cula.kda import kda_decode, kda_packed_decode +from tests.test_kda_decode import _assert_close, make_inputs # noqa: F401 + +K = 128 + + +def _pack_mixed_qkv(q, k, v, N): + """q,k: (N,H,K) bf16 ; v: (N,HV,V) bf16 -> mixed_qkv (N, qkv_dim) bf16.""" + return torch.cat([q.view(N, -1), k.view(N, -1), v.view(N, -1)], dim=-1) + + +# --------------------------------------------------------------------------- +# Dense: packed vs non-packed kda_decode +# --------------------------------------------------------------------------- +def _run_nonpacked_dense(q, k, v, a, b, A_log, dt_bias, state, scale): + N, H, Kd = q.shape + state_ref = state.clone() + o = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=torch.arange(N, device=q.device, dtype=torch.int32), + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + return o.squeeze(1), state_ref # (N,HV,V), (N,HV,V,K) + + +def _run_packed_dense(q, k, v, a, b, A_log, dt_bias, state, scale): + N, H, Kd = q.shape + mixed = _pack_mixed_qkv(q, k, v, N) + state_p = state.clone() + o = kda_packed_decode( + mixed, + a.unsqueeze(1).contiguous(), + b.unsqueeze(1).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=torch.arange(N, device=q.device, dtype=torch.int32), + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + return o.squeeze(1), state_p # (N,HV,V), (N,HV,V,K) + + +@pytest.mark.parametrize("N", [1, 2, 8, 16, 32, 64, 128]) +@pytest.mark.parametrize("H,HV", [(8, 16), (16, 32)]) +@pytest.mark.parametrize("V", [128, 256]) +def test_packed_dense(N, H, HV, V): + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + o_ref, state_ref = _run_nonpacked_dense(q, k, v, a, b, A_log, dt_bias, state, scale) + o_p, state_p = _run_packed_dense(q, k, v, a, b, A_log, dt_bias, state, scale) + + # Sanity: the packed should be ~identical to non-packed (same kernel, strided view). + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +# --------------------------------------------------------------------------- +# Varlen: packed vs non-packed kda_decode +# --------------------------------------------------------------------------- +def _run_nonpacked_varlen(q, k, v, a, b, A_log, dt_bias, state, scale): + N, H, Kd = q.shape + state_ref = state.clone() + o = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(0).contiguous(), + k=k.unsqueeze(0).contiguous(), + v=v.unsqueeze(0).contiguous(), + a=a.contiguous(), + b=b.contiguous(), + initial_state_source=state_ref, + initial_state_indices=torch.arange(N, device=q.device, dtype=torch.int32), + cu_seqlens=torch.arange(N + 1, device=q.device, dtype=torch.int32), + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + return o.squeeze(0), state_ref # (N,HV,V), (N,HV,V,K) + + +def _run_packed_varlen(q, k, v, a, b, A_log, dt_bias, state, scale): + N, H, Kd = q.shape + mixed = _pack_mixed_qkv(q, k, v, N) + state_p = state.clone() + o = kda_packed_decode( + mixed, + a.contiguous(), + b.contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=torch.arange(N, device=q.device, dtype=torch.int32), + cu_seqlens=torch.arange(N + 1, device=q.device, dtype=torch.int32), + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + return o.squeeze(0), state_p # (N,HV,V), (N,HV,V,K) + + +@pytest.mark.parametrize("N", [2, 8, 16, 32, 64, 128]) +@pytest.mark.parametrize("H,HV", [(8, 16), (16, 32)]) +@pytest.mark.parametrize("V", [128, 256]) +def test_packed_varlen(N, H, HV, V): + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + o_ref, state_ref = _run_nonpacked_varlen(q, k, v, a, b, A_log, dt_bias, state, scale) + o_p, state_p = _run_packed_varlen(q, k, v, a, b, A_log, dt_bias, state, scale) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +# --------------------------------------------------------------------------- +# Equal-head coverage: H == HV (group ratio = 1) +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("N", [1, 8, 64]) +@pytest.mark.parametrize("H", [8, 16, 32, 64]) +def test_packed_dense_equal_heads(N, H): + HV, V = H, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + o_ref, state_ref = _run_nonpacked_dense(q, k, v, a, b, A_log, dt_bias, state, scale) + o_p, state_p = _run_packed_dense(q, k, v, a, b, A_log, dt_bias, state, scale) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +@pytest.mark.parametrize("N", [2, 8, 64]) +@pytest.mark.parametrize("H", [8, 16, 32, 64]) +def test_packed_varlen_equal_heads(N, H): + HV, V = H, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + o_ref, state_ref = _run_nonpacked_varlen(q, k, v, a, b, A_log, dt_bias, state, scale) + o_p, state_p = _run_packed_varlen(q, k, v, a, b, A_log, dt_bias, state, scale) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +# --------------------------------------------------------------------------- +# -1 dummy slots: output 0, corresponding state untouched +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("N", [1, 2, 8]) +def test_packed_minus1_dummy(N): + H, HV, V = 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + # pool larger than N; mark the first token as dummy (-1). + pool = N + 2 + state_pool = torch.zeros(pool, HV, V, K, device="cuda", dtype=torch.float32) + real_idx = torch.arange(N, device="cuda", dtype=torch.int32) + 1 # use slots 1..N + state_pool[real_idx] = state + indices = real_idx.clone() + indices[0] = -1 # first batch token is a dummy + + mixed = _pack_mixed_qkv(q, k, v, N) + state_before = state_pool.clone() + o = kda_packed_decode( + mixed, + a.unsqueeze(1).contiguous(), + b.unsqueeze(1).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_pool, + state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + ).squeeze(1) # (N,HV,V) + + # Dummy tokens (pool_idx == -1) skip ALL kernel work — the kernel neither + # reads/writes their state slot nor writes their output row. So: + # - the dummy's pool slot is untouched, and + # - the dummy's output row is UNINITIALIZED (leave it out of checks). + dummy_slot = int(real_idx[0]) if indices[0] == -1 else None + assert dummy_slot is not None + assert torch.equal(state_pool[dummy_slot], state_before[dummy_slot]), "dummy state slot was modified" + + # Compare against non-packed using the SAME pool + SAME indices (no remap), + # so packed and non-packed write the same physical state slots. + state_ref_pool = state_before.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref_pool, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=True, + ).squeeze(1) + + if N > 1: + _assert_close("real output", o_ref[1:].float(), o[1:].float()) + _assert_close("real state", state_ref_pool[real_idx[1:]], state_pool[real_idx[1:]]) + else: + # N == 1: the only token is the dummy, so there is no real row to check. + # Just confirm the dummy slot's state survived untouched (already asserted). + pass + + +# --------------------------------------------------------------------------- +# KV state layout +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("is_varlen", [False, True]) +def test_packed_kv_state_layout(is_varlen): + N, H, HV, V = 8, 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state_vk = make_inputs(N, H, HV, K, V) + state_kv = state_vk.permute(0, 1, 3, 2).contiguous() # (N,HV,K,V) + + mixed = _pack_mixed_qkv(q, k, v, N) + indices = torch.arange(N, device="cuda", dtype=torch.int32) + cu_seqlens = torch.arange(N + 1, device="cuda", dtype=torch.int32) if is_varlen else None + + if is_varlen: + state_p = state_kv.clone() + o_kv = kda_packed_decode( + mixed, + a.contiguous(), + b.contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + state_layout="kv", + ).squeeze(0) + else: + state_p = state_kv.clone() + o_kv = kda_packed_decode( + mixed, + a.unsqueeze(1).contiguous(), + b.unsqueeze(1).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + scale=scale, + state_layout="kv", + ).squeeze(1) + + # Compare against vk-layout non-packed reference (same numerics). + state_vk_ref = state_vk.clone() + if is_varlen: + o_vk = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(0).contiguous(), + k=k.unsqueeze(0).contiguous(), + v=v.unsqueeze(0).contiguous(), + a=a.contiguous(), + b=b.contiguous(), + initial_state_source=state_vk_ref, + initial_state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + ).squeeze(0) + else: + o_vk = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_vk_ref, + initial_state_indices=indices, + scale=scale, + ).squeeze(1) + + _assert_close("kv output", o_kv.float(), o_vk.float()) + # state: kv layout transposed back to vk for comparison + _assert_close("kv state", state_vk_ref, state_p.permute(0, 1, 3, 2).contiguous()) + + +# --------------------------------------------------------------------------- +# No L2 norm path +# --------------------------------------------------------------------------- +def test_packed_no_l2norm(): + N, H, HV, V = 8, 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + mixed = _pack_mixed_qkv(q, k, v, N) + indices = torch.arange(N, device="cuda", dtype=torch.int32) + + state_p = state.clone() + o_p = kda_packed_decode( + mixed, + a.unsqueeze(1).contiguous(), + b.unsqueeze(1).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=False, + ).squeeze(1) + + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=indices, + scale=scale, + use_qk_l2norm_in_kernel=False, + ).squeeze(1) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +# --------------------------------------------------------------------------- +# Focused API compatibility and validation coverage +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("is_varlen", [False, True]) +def test_packed_explicit_out_scale_and_softplus(is_varlen): + N, H, HV, V = 8, 8, 16, 128 + scale = 0.25 + softplus_beta = 0.75 + softplus_threshold = 10.0 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + mixed = _pack_mixed_qkv(q, k, v, N) + indices = torch.arange(N, device="cuda", dtype=torch.int32) + cu_seqlens = torch.arange(N + 1, device="cuda", dtype=torch.int32) if is_varlen else None + + state_p = state.clone() + if is_varlen: + out = torch.empty(1, N, HV, V, device="cuda", dtype=torch.bfloat16) + o_p = kda_packed_decode( + mixed, + a.contiguous(), + b.contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + out=out, + cu_seqlens=cu_seqlens, + scale=scale, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + ) + assert o_p is out + o_p_cmp = o_p.squeeze(0) + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(0).contiguous(), + k=k.unsqueeze(0).contiguous(), + v=v.unsqueeze(0).contiguous(), + a=a.contiguous(), + b=b.contiguous(), + initial_state_source=state_ref, + initial_state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + ).squeeze(0) + else: + out = torch.empty(N, 1, HV, V, device="cuda", dtype=torch.bfloat16) + o_p = kda_packed_decode( + mixed, + a.unsqueeze(1).contiguous(), + b.unsqueeze(1).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + out=out, + scale=scale, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + ) + assert o_p is out + o_p_cmp = o_p.squeeze(1) + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=indices, + scale=scale, + softplus_beta=softplus_beta, + softplus_threshold=softplus_threshold, + ).squeeze(1) + + _assert_close("output", o_ref.float(), o_p_cmp.float()) + _assert_close("state", state_ref, state_p) + + +@pytest.mark.parametrize("is_varlen", [False, True]) +def test_packed_slow_path_compatible_a_b_shapes(is_varlen): + N, H, HV, V = 8, 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + mixed = _pack_mixed_qkv(q, k, v, N) + indices = torch.arange(N, device="cuda", dtype=torch.int32) + cu_seqlens = torch.arange(N + 1, device="cuda", dtype=torch.int32) if is_varlen else None + + state_p = state.clone() + if is_varlen: + # Fast path expects a=(N,HV,K), b=(N,HV); these compatible public + # shapes force the general normalization path. + o_p = kda_packed_decode( + mixed, + a.reshape(1, N, HV * K).contiguous(), + b.unsqueeze(0).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + ).squeeze(0) + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(0).contiguous(), + k=k.unsqueeze(0).contiguous(), + v=v.unsqueeze(0).contiguous(), + a=a.contiguous(), + b=b.contiguous(), + initial_state_source=state_ref, + initial_state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + ).squeeze(0) + else: + # Fast path does not accept dense a=(N,HV*K), so this covers the + # packed general path while preserving the same numerics. + o_p = kda_packed_decode( + mixed, + a.reshape(N, HV * K).contiguous(), + b.contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + scale=scale, + ).squeeze(1) + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=indices, + scale=scale, + ).squeeze(1) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +def test_packed_mixed_qkv_allows_padded_row_stride(): + N, H, HV, V = 8, 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + mixed = _pack_mixed_qkv(q, k, v, N) + padded = torch.empty(N, mixed.shape[1] + 8, device="cuda", dtype=torch.bfloat16) + mixed_padded_view = padded[:, : mixed.shape[1]] + mixed_padded_view.copy_(mixed) + assert mixed_padded_view.stride(-1) == 1 + assert mixed_padded_view.stride(0) != mixed.shape[1] + + indices = torch.arange(N, device="cuda", dtype=torch.int32) + state_p = state.clone() + o_p = kda_packed_decode( + mixed_padded_view, + a.unsqueeze(1).contiguous(), + b.unsqueeze(1).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + scale=scale, + ).squeeze(1) + + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=indices, + scale=scale, + ).squeeze(1) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +def test_packed_mixed_qkv_allows_padded_row_stride_n1(): + N, H, HV, V = 1, 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + + mixed = _pack_mixed_qkv(q, k, v, N) + padded = torch.empty(N, mixed.shape[1] + 8, device="cuda", dtype=torch.bfloat16) + mixed_padded_view = padded[:, : mixed.shape[1]] + mixed_padded_view.copy_(mixed) + assert mixed_padded_view.stride(-1) == 1 + assert mixed_padded_view.stride(0) != mixed.shape[1] + + indices = torch.arange(N, device="cuda", dtype=torch.int32) + state_p = state.clone() + o_p = kda_packed_decode( + mixed_padded_view, + a.unsqueeze(1).contiguous(), + b.unsqueeze(1).contiguous(), + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + scale=scale, + ).squeeze(1) + + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=indices, + scale=scale, + ).squeeze(1) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +@pytest.mark.parametrize("is_varlen", [False, True]) +def test_packed_general_path_accepts_noncontiguous_a_b(is_varlen): + N, H, HV, V = 4, 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V) + mixed = _pack_mixed_qkv(q, k, v, N) + + a_flat = a.reshape(N, HV * K) + a_padded = torch.empty(N, HV * K + 1, device="cuda", dtype=torch.bfloat16) + a_arg = a_padded[:, : HV * K] + a_arg.copy_(a_flat) + b_padded = torch.empty(N, HV + 1, device="cuda", dtype=torch.bfloat16) + b_arg = b_padded[:, :HV] + b_arg.copy_(b) + assert not a_arg.is_contiguous() + assert not b_arg.is_contiguous() + + state_p = state.clone() + if is_varlen: + o_p = kda_packed_decode( + mixed, + a_arg, + b_arg, + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=torch.arange(N, device="cuda", dtype=torch.int32), + cu_seqlens=torch.arange(N + 1, device="cuda", dtype=torch.int32), + scale=scale, + ).squeeze(0) + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(0).contiguous(), + k=k.unsqueeze(0).contiguous(), + v=v.unsqueeze(0).contiguous(), + a=a.contiguous(), + b=b.contiguous(), + initial_state_source=state_ref, + initial_state_indices=torch.arange(N, device="cuda", dtype=torch.int32), + cu_seqlens=torch.arange(N + 1, device="cuda", dtype=torch.int32), + scale=scale, + ).squeeze(0) + else: + o_p = kda_packed_decode( + mixed, + a_arg, + b_arg, + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=torch.arange(N, device="cuda", dtype=torch.int32), + scale=scale, + ).squeeze(1) + state_ref = state.clone() + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=torch.arange(N, device="cuda", dtype=torch.int32), + scale=scale, + ).squeeze(1) + + _assert_close("output", o_ref.float(), o_p.float()) + _assert_close("state", state_ref, state_p) + + +# --------------------------------------------------------------------------- +# Padded batch + CUDA graph capture/replay +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("is_varlen", [False, True]) +def test_packed_padded_batch(is_varlen): + N_real, N_pad, H, HV, V = 6, 8, 8, 16, 128 + scale = K**-0.5 + q, k, v, a, b, A_log, dt_bias, state_real = make_inputs(N_real, H, HV, K, V) + + # Pad to N_pad with extra rows; mark them as dummy (-1). + q_pad = torch.zeros(N_pad, H, K, device="cuda", dtype=torch.bfloat16) + k_pad = torch.zeros(N_pad, H, K, device="cuda", dtype=torch.bfloat16) + v_pad = torch.zeros(N_pad, HV, V, device="cuda", dtype=torch.bfloat16) + a_pad = torch.zeros(N_pad, HV, K, device="cuda", dtype=torch.bfloat16) + b_pad = torch.zeros(N_pad, HV, device="cuda", dtype=torch.bfloat16) + q_pad[:N_real], k_pad[:N_real], v_pad[:N_real] = q, k, v + a_pad[:N_real] = a + b_pad[:N_real] = b + + pool = N_pad + state_pool = torch.zeros(pool, HV, V, K, device="cuda", dtype=torch.float32) + indices = torch.arange(N_pad, device="cuda", dtype=torch.int32) + indices[N_real:] = -1 # padded slots are dummies + state_pool[:N_real] = state_real + + mixed = _pack_mixed_qkv(q_pad, k_pad, v_pad, N_pad) + cu_seqlens = torch.arange(N_pad + 1, device="cuda", dtype=torch.int32) if is_varlen else None + + # Forward args shared by eager + graph path + if is_varlen: + a_arg, b_arg = a_pad.contiguous(), b_pad.contiguous() + else: + a_arg, b_arg = a_pad.unsqueeze(1).contiguous(), b_pad.unsqueeze(1).contiguous() + + state_p = state_pool.clone() + o_p = kda_packed_decode( + mixed, + a_arg, + b_arg, + A_log=A_log, + dt_bias=dt_bias, + state=state_p, + state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + ) + if is_varlen: + o_real = o_p[:, :N_real] # (1,N_real,HV,V) + else: + o_real = o_p[:N_real] # (N_real,1,HV,V) + + # Compare real rows against non-packed run on the real-only inputs. + state_ref = state_real.clone() + if is_varlen: + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(0).contiguous(), + k=k.unsqueeze(0).contiguous(), + v=v.unsqueeze(0).contiguous(), + a=a.contiguous(), + b=b.contiguous(), + initial_state_source=state_ref, + initial_state_indices=torch.arange(N_real, device="cuda", dtype=torch.int32), + cu_seqlens=torch.arange(N_real + 1, device="cuda", dtype=torch.int32), + scale=scale, + ) # (1,N_real,HV,V) + else: + o_ref = kda_decode( + A_log=A_log, + dt_bias=dt_bias, + q=q.unsqueeze(1).contiguous(), + k=k.unsqueeze(1).contiguous(), + v=v.unsqueeze(1).contiguous(), + a=a.unsqueeze(1).contiguous(), + b=b.unsqueeze(1).contiguous(), + initial_state_source=state_ref, + initial_state_indices=torch.arange(N_real, device="cuda", dtype=torch.int32), + scale=scale, + ) # (N_real,1,HV,V) + + _assert_close("real output", o_ref.float(), o_real.float()) + # dummy (padded) state slots must be untouched; their output rows are + # UNINITIALIZED by the kernel (no write), so we don't assert on them. + assert torch.equal(state_p[N_real:], state_pool[N_real:]), "dummy state slots were modified" + + # CUDA graph capture + replay smoke (capture must not raise; replay must not raise). + # Run once eagerly first so the kernel is compiled and any warming cu_seqlens + # cache is populated before capture. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + state_g = state_pool.clone() + _ = kda_packed_decode( + mixed, + a_arg, + b_arg, + A_log=A_log, + dt_bias=dt_bias, + state=state_g, + state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + ) + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out_g2 = kda_packed_decode( + mixed, + a_arg, + b_arg, + A_log=A_log, + dt_bias=dt_bias, + state=state_g, + state_indices=indices, + cu_seqlens=cu_seqlens, + scale=scale, + ) + g.replay() + assert out_g2 is not None