diff --git a/benchmarks/bench_packed_kda_decode.py b/benchmarks/bench_packed_kda_decode.py new file mode 100644 index 00000000000..c4c53cc2abb --- /dev/null +++ b/benchmarks/bench_packed_kda_decode.py @@ -0,0 +1,318 @@ +# 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. + +"""Benchmark the packed-input CuTe KDA T=1 decode kernel on B200.""" + +import argparse +import json +import statistics +from pathlib import Path + +import torch + +from flashinfer.kda_kernels.packed_kda_decode_cute import _select_tile_v +from flashinfer.kda_kernels.packed_kda_decode_cute import ( + run_packed_kda_decode_cute, +) +from flashinfer.testing import bench_gpu_time + + +HEADS = 12 +HEAD_DIM = 128 +MIXED_WIDTH = 3 * HEADS * HEAD_DIM +GATE_WIDTH = HEADS * HEAD_DIM +MIXED_STRIDE = 6144 +STATE_ELEMENTS = HEADS * HEAD_DIM * HEAD_DIM +STATE_PADDING = 256 +DEFAULT_BATCHES = (1, 8, 16, 31, 32, 64, 128, 256, 512) + + +def _state_view(storage, slots, slot_stride): + return storage.as_strided( + (slots, HEADS, HEAD_DIM, HEAD_DIM), + (slot_stride, HEAD_DIM * HEAD_DIM, HEAD_DIM, 1), + ) + + +def _make_case(batch, device, seed): + generator = torch.Generator(device=device).manual_seed(seed) + mixed_storage = torch.randn( + batch, + MIXED_STRIDE, + dtype=torch.bfloat16, + device=device, + generator=generator, + ).mul_(0.25) + raw_gate = torch.randn( + batch, + GATE_WIDTH, + dtype=torch.bfloat16, + device=device, + generator=generator, + ).mul_(0.25) + raw_beta = torch.randn( + batch, + HEADS, + dtype=torch.bfloat16, + device=device, + generator=generator, + ) + A_log = torch.empty(HEADS, dtype=torch.float32, device=device) + A_log.uniform_(-2.0, -0.1, generator=generator) + dt_bias = torch.randn( + GATE_WIDTH, + dtype=torch.float32, + device=device, + generator=generator, + ).mul_(0.1) + + slots = batch + 1 + slot_stride = STATE_ELEMENTS + STATE_PADDING + state_storage = torch.randn( + slots * slot_stride, + dtype=torch.bfloat16, + device=device, + generator=generator, + ).mul_(0.02) + state = _state_view(state_storage, slots, slot_stride) + state_indices = torch.arange( + batch, + 0, + -1, + dtype=torch.int32, + device=device, + ) + output = torch.empty( + batch, + 1, + HEADS, + HEAD_DIM, + dtype=torch.bfloat16, + device=device, + ) + return { + "batch": batch, + "mixed_qkv": mixed_storage[:, :MIXED_WIDTH], + "raw_gate": raw_gate, + "raw_beta": raw_beta, + "A_log": A_log, + "dt_bias": dt_bias, + "state_storage": state_storage, + "state": state, + "state_indices": state_indices, + "output": output, + "initial_state_storage": state_storage.clone(), + "slots": slots, + "slot_stride": slot_stride, + } + + +def _run(case, tile_v): + return run_packed_kda_decode_cute( + case["mixed_qkv"], + case["raw_gate"], + case["raw_beta"], + case["A_log"], + case["dt_bias"], + case["state"], + case["state_indices"], + output=case["output"], + tile_v=tile_v, + ) + + +def _restore(case): + case["state_storage"].copy_(case["initial_state_storage"]) + case["output"].zero_() + + +def _reference(case): + batch = case["batch"] + packed = case["mixed_qkv"].float().reshape(batch, 3, HEADS, HEAD_DIM) + q_raw = packed[:, 0] + k_raw = packed[:, 1] + q = ( + q_raw + * torch.rsqrt((q_raw * q_raw).sum(dim=-1, keepdim=True) + 1.0e-6) + * (HEAD_DIM**-0.5) + ) + k = k_raw * torch.rsqrt((k_raw * k_raw).sum(dim=-1, keepdim=True) + 1.0e-6) + value = packed[:, 2] + gate = case["raw_gate"].float().reshape(batch, HEADS, HEAD_DIM) + gate = gate + case["dt_bias"].reshape(HEADS, HEAD_DIM) + decay = torch.exp( + -5.0 * torch.sigmoid(torch.exp(case["A_log"])[None, :, None] * gate) + ) + beta = torch.sigmoid(case["raw_beta"].float()) + + indices = case["state_indices"].long() + reference_storage = case["initial_state_storage"].clone() + reference_state = _state_view( + reference_storage, + case["slots"], + case["slot_stride"], + ) + selected = reference_state.index_select(0, indices).float() + decayed = selected * decay[:, :, None, :] + prediction = torch.einsum("bhvk,bhk->bhv", decayed, k) + delta = (value - prediction) * beta[:, :, None] + updated = decayed + delta[:, :, :, None] * k[:, :, None, :] + projected = torch.einsum("bhvk,bhk->bhv", updated, q) + reference_state.index_copy_(0, indices, updated.to(torch.bfloat16)) + return projected.to(torch.bfloat16).unsqueeze(1), reference_state + + +def _check(case, tile_v): + expected_output, expected_state = _reference(case) + _restore(case) + result = _run(case, tile_v) + torch.cuda.synchronize() + torch.testing.assert_close( + result, + expected_output, + atol=1.0e-2, + rtol=1.0e-2, + check_dtype=False, + ) + torch.testing.assert_close( + case["state"], + expected_state, + atol=1.0e-2, + rtol=1.0e-2, + check_dtype=False, + ) + output_error = float((result.float() - expected_output.float()).abs().max()) + state_error = float((case["state"].float() - expected_state.float()).abs().max()) + _restore(case) + torch.cuda.synchronize() + return output_error, state_error + + +def _capture(case, tile_v): + _restore(case) + torch.cuda.synchronize() + stream = torch.cuda.Stream(device=case["state"].device) + stream.wait_stream(torch.cuda.current_stream(case["state"].device)) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + _run(case, tile_v) + torch.cuda.synchronize() + _restore(case) + torch.cuda.synchronize() + return graph + + +def _logical_bytes(batch): + bf16 = 2 + state = batch * STATE_ELEMENTS * bf16 * 2 + inputs = batch * (MIXED_WIDTH + GATE_WIDTH + HEADS) * bf16 + parameters = (HEADS + GATE_WIDTH) * 4 + output = batch * HEADS * HEAD_DIM * bf16 + indices = batch * 4 + return state + inputs + parameters + output + indices + + +def _parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--batch-size", + type=int, + nargs="+", + default=list(DEFAULT_BATCHES), + ) + parser.add_argument( + "--mode", + choices=("direct", "cuda_graph", "both"), + default="both", + ) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iterations", type=int, default=100) + parser.add_argument("--cold-l2", action="store_true") + parser.add_argument("--tile-v", type=int, choices=(8, 16, 32, 64, 128)) + parser.add_argument("--seed", type=int, default=20260805) + parser.add_argument("--json", type=Path) + args = parser.parse_args() + if any(batch <= 0 or batch > 65535 for batch in args.batch_size): + parser.error("batch sizes must be in [1, 65535]") + if args.warmup < 0 or args.iterations <= 0: + parser.error("warmup must be non-negative and iterations must be positive") + return args + + +def main(): + args = _parse_args() + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + device = torch.device("cuda") + if torch.cuda.get_device_capability(device) != (10, 0): + raise RuntimeError("packed-input CuTe KDA decode requires exact CC 10.0") + + modes = ("direct", "cuda_graph") if args.mode == "both" else (args.mode,) + rows = [] + print(f"{'mode':<10} {'B':>5} {'tile':>6} {'median_us':>12} {'logical_TB/s':>14}") + for ordinal, batch in enumerate(args.batch_size): + case = _make_case(batch, device, args.seed + ordinal) + # tile_v=None lets the per-batch policy pick the kernel shape; + # _select_tile_v is what it will choose (display only). + tile_v = args.tile_v + display_tile = tile_v or _select_tile_v(batch) + output_error, state_error = _check(case, tile_v) + for mode in modes: + run = ( + (lambda: _run(case, tile_v)) + if mode == "direct" + else _capture(case, tile_v).replay + ) + samples_ms = bench_gpu_time( + run, + enable_cupti=True, + cold_l2_cache=args.cold_l2, + use_cuda_graph=False, + dry_run_iters=args.warmup, + repeat_iters=args.iterations, + ) + median_ms = float(statistics.median(samples_ms)) + logical_tbps = _logical_bytes(batch) / median_ms / 1.0e9 + row = { + "mode": mode, + "batch_size": batch, + "tile_v": display_tile, + "median_us": median_ms * 1000.0, + "logical_tb_per_second": logical_tbps, + "output_max_abs": output_error, + "state_max_abs": state_error, + "samples_ms": [float(value) for value in samples_ms], + } + rows.append(row) + print( + f"{mode:<10} {batch:>5} {display_tile:>6} " + f"{row['median_us']:>12.4f} {logical_tbps:>14.4f}" + ) + torch.cuda.empty_cache() + + if args.json is not None: + report = { + "device": torch.cuda.get_device_name(device), + "compute_capability": list(torch.cuda.get_device_capability(device)), + "warmup": args.warmup, + "iterations": args.iterations, + "cold_l2": args.cold_l2, + "rows": rows, + } + args.json.parent.mkdir(parents=True, exist_ok=True) + args.json.write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/flashinfer/kda_kernels/__init__.py b/flashinfer/kda_kernels/__init__.py index fca6b3ecd21..46c5dafe42d 100644 --- a/flashinfer/kda_kernels/__init__.py +++ b/flashinfer/kda_kernels/__init__.py @@ -37,6 +37,12 @@ run_fused_kda_decode = None # type: ignore fused_kda_decode = None # type: ignore +# NOTE: flashinfer.kda_kernels.packed_kda_decode_cute is an internal +# implementation module, not public API. Its kernels back the T=1 fast path +# of the public ``flashinfer.recurrent_kda`` operation (see +# ``run_recurrent_kda`` in ``recurrent_kda.py``); import it by module path +# only for tests and benchmarks. + try: if _torch.cuda.is_available(): from ..cute_dsl.utils import is_cute_dsl_arch_supported as _dsl_arch_ok diff --git a/flashinfer/kda_kernels/packed_kda_decode_cute.py b/flashinfer/kda_kernels/packed_kda_decode_cute.py new file mode 100644 index 00000000000..44a694ba971 --- /dev/null +++ b/flashinfer/kda_kernels/packed_kda_decode_cute.py @@ -0,0 +1,2613 @@ +""" +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. + +CuTe-DSL backend for serving-native packed Kimi K3 T=1 recurrent decode. + +INTERNAL IMPLEMENTATION -- not public API. The supported way to reach these +kernels is the public ``flashinfer.recurrent_kda`` operation, whose T=1 fast +path dispatches eligible decode calls here (toggle with +``FLASHINFER_KDA_T1_FAST_PATH=0``). ``run_packed_kda_decode_cute`` is kept +for tests and benchmarks of the packed-input entry point. + +Same numerical contract and tensor layouts as the exported Cake backend in +``flashinfer.kda_kernels.packed_kda_decode``: packed bf16 QKV rows, raw gate +and beta logits, ``scale = 1/sqrt(128)``, L2 epsilon ``1e-6``, +``lower_bound = -5``, bf16 state pool updated in place, and ``-1`` state +indices producing zero output without touching the pool. + +Kernel design (ported from the GDN wide-vec T=1 kernel in +``flashinfer.gdn_kernels.gdn_decode_bf16_state`` and the register-tile +recurrent KDA kernel in ``flashinfer.kda_kernels.recurrent_kda``): + +- CTAs are ``num_groups`` groups of 16 lanes; each lane owns a contiguous + eight-element K vector (LDG.128 / STG.128 on the bf16 state). Grid is + ``num_v_tiles x H x B`` linearised; each CTA owns one ``tile_v`` V-row + tile of one ``(batch, head)`` pair. +- Barrier-free staging: each warp independently loads q/k/gate (four + elements per lane across 32 lanes), reduces the L2 norms with a full-warp + butterfly, computes the per-channel decay, and shuffle-transposes into the + 16-lane-by-8 register layout. +- State stays packed bf16 in registers; unpack is a shift/mask bit trick and + repack a single ``cvt.rn.bf16x2.f32`` (full-rate ALU instead of the + conversion pipe). FMAs use packed F32x2 pairs on SM100+. +- Two main-loop variants, selected per batch size by ``_select_config``: + a register-prefetch kernel (double-buffered LDG.128, CTA width tunable + down to 32 threads for small-batch latency), and a cp.async kernel that + streams state rows through a shared-memory ring (``n_stages`` deep, + ``cp.async.cg`` L1 bypass) so the in-flight read volume is not bounded by + the register file — the DRAM-bound large-batch regime runs there. +""" + +import functools +import os +from typing import Optional + +import cutlass +import cutlass.cute as cute +import cuda.bindings.driver as cuda +import torch +from cutlass._mlir.dialects import arith as mlir_arith +from cutlass._mlir.dialects import llvm as mlir_llvm +from cutlass.cute.runtime import make_fake_stream +import tvm_ffi # noqa: F401 -- TVM FFI required for zero-overhead dispatch + +from ..jit.cpp_ext import is_cuda_version_at_least +from ..utils import get_compute_capability + +_HEADS = 12 +_HEAD_DIM = 128 +_MIXED_WIDTH = 3 * _HEADS * _HEAD_DIM +_SCALE = float(_HEAD_DIM) ** -0.5 +_EPS = 1.0e-6 +_LOWER_BOUND = -5.0 +_LOG2_E = 1.4426950408889634 + +_LANES_PER_ROW = 16 # 16 lanes cover K=128 with 8 bf16 each +_ELEMS_PER_LANE = 8 +_NUM_THREADS = 128 +_NUM_GROUPS = _NUM_THREADS // _LANES_PER_ROW # 8 +_VALS_PER_LANE = _HEAD_DIM // 32 # 4: warp-staged q/k/gate elements per lane + + +@cute.jit +def _bf16x2_to_f32x2(u): + """Unpack one register holding two bf16 into two f32 via bit ops. + + bf16 widens to f32 by appending 16 zero mantissa bits, so a shift and a + mask on the packed word replace two CVT instructions (full-rate ALU + instead of the conversion pipe — the same trick the Cake-generated + kernel uses via inline PTX). Element 0 sits in the low half. + """ + lo = cutlass.Float32( + mlir_arith.bitcast(cutlass.Float32.mlir_type, (u << 16).ir_value()) + ) + hi = cutlass.Float32( + mlir_arith.bitcast(cutlass.Float32.mlir_type, (u & -65536).ir_value()) + ) + return lo, hi + + +@cute.jit +def _f32x2_to_bf16x2(h0, h1): + """Pack two f32 into one register of two bf16 with a single CVT. + + ``cvt.rn.bf16x2.f32`` rounds both halves to nearest-even — bitwise + identical to two scalar BFloat16() conversions (and to CUDA's + ``__float22bfloat162_rn``) — replacing two conversion-pipe CVTs plus a + register merge. ``h0`` lands in the low half (element 0). + """ + packed = mlir_llvm.inline_asm( + cutlass.Int32.mlir_type, + [ + cutlass.Float32(h1).ir_value(), + cutlass.Float32(h0).ir_value(), + ], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(packed) + + +def _cp_async_bf16x8_cg(base_addr_i64, bf16_elem_offset, smem_addr_i32, l2_hint=0): + """cp.async.cg, 16 B (8 bf16), L1 bypass (state is a stream-once input). + + Same inline-PTX shape as the GDN wy kernel helpers: the u32 element + offset is widened and scaled into the u64 global address in-asm. + ``l2_hint`` of 128/256 adds an ``.L2::B`` allocation hint — each 8 KiB + chunk densely covers its sectors, so full-sector allocation can cut DRAM + request overhead (constexpr; compiled in). + """ + hint = f".L2::{int(l2_hint)}B" if l2_hint else "" + r = mlir_llvm.inline_asm( + cutlass.Int32.mlir_type, + [ + smem_addr_i32.ir_value(), + base_addr_i64.ir_value(), + bf16_elem_offset.ir_value(), + ], + "{ .reg .u64 _a; mad.wide.u32 _a, $3, 2, $2;" + f" cp.async.cg.shared.global{hint} [$1], [_a], 16;" + " mov.u32 $0, 0; }", + "=r,r,l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(r) + + +def _l2_policy_evict_last(): + """64-bit L2 cache policy: evict_last (created once per thread).""" + r = mlir_llvm.inline_asm( + cutlass.Int64.mlir_type, + [], + "createpolicy.fractional.L2::evict_last.b64 $0, 1.0;", + "=l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int64(r) + + +def _l2_policy_evict_first(): + r = mlir_llvm.inline_asm( + cutlass.Int64.mlir_type, + [], + "createpolicy.fractional.L2::evict_first.b64 $0, 1.0;", + "=l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int64(r) + + +def _stg_hint_v4_b32(gmem_addr_i64, u0, u1, u2, u3, policy_i64): + """st.global.L2::cache_hint.v4.b32 — state write with an L2 policy. + + evict_last keeps the freshly written state resident in L2 so its + write-back drains after kernel-end (overlapping the next kernel in a + serving pipeline) instead of competing with the in-kernel read stream. + """ + r = mlir_llvm.inline_asm( + cutlass.Int32.mlir_type, + [ + gmem_addr_i64.ir_value(), + u0.ir_value(), + u1.ir_value(), + u2.ir_value(), + u3.ir_value(), + policy_i64.ir_value(), + ], + "{ st.global.L2::cache_hint.v4.b32 [$1], {$2, $3, $4, $5}, $6;" + " mov.u32 $0, 0; }", + "=r,l,r,r,r,r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(r) + + +def _cp_async_bf16x8_cg_hint( + base_addr_i64, bf16_elem_offset, smem_addr_i32, policy_i64 +): + """cp.async.cg with an L2 cache policy (evict_first for the read stream).""" + r = mlir_llvm.inline_asm( + cutlass.Int32.mlir_type, + [ + smem_addr_i32.ir_value(), + base_addr_i64.ir_value(), + bf16_elem_offset.ir_value(), + policy_i64.ir_value(), + ], + "{ .reg .u64 _a; mad.wide.u32 _a, $3, 2, $2;" + " cp.async.cg.shared.global.L2::cache_hint [$1], [_a], 16, $4;" + " mov.u32 $0, 0; }", + "=r,r,l,r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(r) + + +def _stg_v4_b32(gmem_addr_i64, u0, u1, u2, u3): + """st.global.v4.b32 at a raw byte address (default cache policy).""" + r = mlir_llvm.inline_asm( + cutlass.Int32.mlir_type, + [ + gmem_addr_i64.ir_value(), + u0.ir_value(), + u1.ir_value(), + u2.ir_value(), + u3.ir_value(), + ], + "{ st.global.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }", + "=r,l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(r) + + +def _stg_cg_v4_b32(gmem_addr_i64, u0, u1, u2, u3): + """st.global.cg.v4.b32 — stream the state write past L1 (write-once).""" + r = mlir_llvm.inline_asm( + cutlass.Int32.mlir_type, + [ + gmem_addr_i64.ir_value(), + u0.ir_value(), + u1.ir_value(), + u2.ir_value(), + u3.ir_value(), + ], + "{ st.global.cg.v4.b32 [$1], {$2, $3, $4, $5}; mov.u32 $0, 0; }", + "=r,l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(r) + + +def _cp_async_commit_group(): + r = mlir_llvm.inline_asm( + cutlass.Int32.mlir_type, + [], + "{ cp.async.commit_group; mov.u32 $0, 0; }", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(r) + + +def _cp_async_wait_group_n(n_const): + """Wait until at most ``n_const`` cp.async groups remain in flight. + + The ``~{memory}`` clobber matters in the barrier-free private ring: + without it the compiler may hoist the subsequent smem loads above the + wait (the barriered path is protected by bar.sync instead). + """ + r = mlir_llvm.inline_asm( + cutlass.Int32.mlir_type, + [], + f"{{ cp.async.wait_group {int(n_const)}; mov.u32 $0, 0; }}", + "=r,~{memory}", + has_side_effects=True, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(r) + + +def _cp_async_bulk_store(gmem_addr_i64, smem_addr_i32, size_bytes): + """One shared->global bulk store (async proxy, bulk_group completion). + + The issuing warp's lane 0 pushes its warp's contiguous smem region out + through the TMA path — state writes leave the SM without occupying + LSU/L1TEX wavefronts. + """ + r = mlir_llvm.inline_asm( + cutlass.Int32.mlir_type, + [gmem_addr_i64.ir_value(), smem_addr_i32.ir_value()], + "{ cp.async.bulk.global.shared::cta.bulk_group [$1], [$2], " + f"{int(size_bytes)};" + " mov.u32 $0, 0; }", + "=r,l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(r) + + +def _cp_async_bulk_read(smem_addr_i32, gmem_addr_i64, mbar_addr_i32, size_bytes): + """One global->shared bulk read completing on an mbarrier (TMA-1D).""" + r = mlir_llvm.inline_asm( + cutlass.Int32.mlir_type, + [ + smem_addr_i32.ir_value(), + gmem_addr_i64.ir_value(), + mbar_addr_i32.ir_value(), + ], + "{ cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes" + f" [$1], [$2], {int(size_bytes)}, [$3];" + " mov.u32 $0, 0; }", + "=r,r,l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(r) + + +def _cp_async_bulk_commit(): + r = mlir_llvm.inline_asm( + cutlass.Int32.mlir_type, + [], + "{ cp.async.bulk.commit_group; mov.u32 $0, 0; }", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(r) + + +def _cp_async_bulk_wait_read(n_const): + """Wait until at most n bulk groups have unread smem sources.""" + r = mlir_llvm.inline_asm( + cutlass.Int32.mlir_type, + [], + f"{{ cp.async.bulk.wait_group.read {int(n_const)}; mov.u32 $0, 0; }}", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(r) + + +def _cp_async_bulk_wait(n_const): + r = mlir_llvm.inline_asm( + cutlass.Int32.mlir_type, + [], + f"{{ cp.async.bulk.wait_group {int(n_const)}; mov.u32 $0, 0; }}", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(r) + + +def _fence_proxy_async(): + """Order generic-proxy smem writes before async-proxy bulk reads.""" + r = mlir_llvm.inline_asm( + cutlass.Int32.mlir_type, + [], + "{ fence.proxy.async.shared::cta; mov.u32 $0, 0; }", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(r) + + +def _sync_warp(): + r = mlir_llvm.inline_asm( + cutlass.Int32.mlir_type, + [], + "{ bar.warp.sync 0xffffffff; mov.u32 $0, 0; }", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=mlir_llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(r) + + +@cute.kernel +def _kda_packed_t1_kernel( + state: cute.Tensor, # [pool, H, V, K] bf16, padded slot stride allowed + q: cute.Tensor, # [B, H, K] bf16 strided view of mixed_qkv + k: cute.Tensor, # [B, H, K] bf16 strided view of mixed_qkv + v: cute.Tensor, # [B, H, V] bf16 strided view of mixed_qkv + g: cute.Tensor, # [B, H, K] bf16 strided view of raw_gate + beta: cute.Tensor, # [B, H] bf16 (row stride dynamic) + A_log: cute.Tensor, # [H] f32 + dt_bias: cute.Tensor, # [H, K] f32 + o: cute.Tensor, # [B, H, V] bf16 contiguous + state_indices: cute.Tensor, # [B] i32 + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + lb_log2e: cutlass.Constexpr[float], + precomputed: cutlass.Constexpr[bool], # lower_bound * log2(e) + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + num_v_tiles: cutlass.Constexpr[int], + ilp_rows: cutlass.Constexpr[int], + num_groups: cutlass.Constexpr[int], + use_packed_fma: cutlass.Constexpr[bool], + evict_first_state: cutlass.Constexpr[bool], +): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + + # Grid: (num_v_tiles x H x B) linearised, i_v fastest so consecutive CTAs + # touch adjacent V tiles of the same state slot (L2 locality). + i_v = bidx % num_v_tiles + tmp = bidx // num_v_tiles + i_h = tmp % H + i_n = tmp // H + + lane = tidx % 32 + k_lane = tidx % _LANES_PER_ROW + group_idx = tidx // _LANES_PER_ROW + + ROWS_PER_GROUP: cutlass.Constexpr[int] = tile_v // num_groups + ITERS: cutlass.Constexpr[int] = ROWS_PER_GROUP // ilp_rows + + vec: cutlass.Constexpr[int] = _ELEMS_PER_LANE + vals: cutlass.Constexpr[int] = _VALS_PER_LANE + + # Plain-Python unroll indices: the DSL stages `for x in range(...)` loops, + # whose loop variable cannot index Python lists of tiles/accumulators. + # Iterating tuples keeps the loop in Python (fully unrolled at trace). + ROWS = tuple(range(ilp_rows)) + VECI = tuple(range(vec)) + VECP = tuple(range(0, vec, 2)) + VALSI = tuple(range(vals)) + ITERI = tuple(range(ITERS)) + + raw_slot = state_indices[i_n] + + # All register tensors are allocated up front, before any dynamic branch + # (allocas must not live inside scf.if regions; same structure as the GDN + # wide-vec kernel and the register-tile recurrent KDA kernel). + q_src = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.Float32 + ) + k_src = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.Float32 + ) + d_src = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.Float32 + ) + q_bf16 = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.BFloat16 + ) + k_bf16 = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.BFloat16 + ) + g_bf16 = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.BFloat16 + ) + dtb_f32 = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.Float32 + ) + r_q = cute.make_rmem_tensor(cute.make_layout((vec,), stride=(1,)), cutlass.Float32) + r_k = cute.make_rmem_tensor(cute.make_layout((vec,), stride=(1,)), cutlass.Float32) + r_d = cute.make_rmem_tensor(cute.make_layout((vec,), stride=(1,)), cutlass.Float32) + # Double-buffered bf16 staging: while iteration `it` computes, iteration + # `it+1`'s state rows and v values are already in flight (the kernel is + # otherwise latency-bound — nothing saturates at base clocks). + r_hb = [ + [ + cute.make_rmem_tensor( + cute.make_layout((vec,), stride=(1,)), cutlass.BFloat16 + ) + for _ in ROWS + ] + for _ in range(2) + ] + r_hb32 = [[cute.recast_tensor(t, cutlass.Int32) for t in bufs] for bufs in r_hb] + r_v_bf16 = [ + cute.make_rmem_tensor( + cute.make_layout((ilp_rows,), stride=(1,)), cutlass.BFloat16 + ) + for _ in range(2) + ] + r_o_bf16 = cute.make_rmem_tensor( + cute.make_layout((ilp_rows,), stride=(1,)), cutlass.BFloat16 + ) + + slot_dead = (raw_slot < 0) | (cutlass.Int64(raw_slot) >= state.shape[0]) + if slot_dead: + # Inactive CUDA-graph padding row: zero this CTA's output rows and do + # not touch the state pool. Uniform per CTA (no divergence hazards). + for r in ROWS: + r_o_bf16[r] = cutlass.BFloat16(0.0) + if k_lane == 0: + for it in ITERI: + vb = i_v * tile_v + group_idx * ROWS_PER_GROUP + it * ilp_rows + ot = cute.local_tile(o, (1, 1, ilp_rows), (i_n, i_h, vb // ilp_rows)) + cute.autovec_copy(r_o_bf16, ot) + else: + # Issue iteration 0's state loads FIRST: they depend only on the slot + # index, and the whole q/k/gate staging chain below (LDG + norm + # butterflies + shuffles) then executes while they are in flight. + # Mirrors the GDN wide-vec tile_v=32 hoist. + h_slot = state[(cutlass.Int64(raw_slot), i_h, None, None)] + v_base = i_v * tile_v + group_idx * ROWS_PER_GROUP + + ht_of = [ + [ + cute.local_tile(h_slot, (1, vec), (v_base + it * ilp_rows + r, k_lane)) + for r in ROWS + ] + for it in ITERI + ] + vt_of = [ + cute.local_tile( + v, (1, 1, ilp_rows), (i_n, i_h, (v_base + it * ilp_rows) // ilp_rows) + ) + for it in ITERI + ] + + for r in ROWS: + if cutlass.const_expr(evict_first_state): + cute.autovec_copy( + ht_of[0][r], + r_hb[0][r], + l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.EVICT_FIRST, + ) + else: + cute.autovec_copy(ht_of[0][r], r_hb[0][r]) + cute.autovec_copy(vt_of[0], r_v_bf16[0]) + + # ------------------------------------------------------------------ + # Warp-local staging: 32 lanes load q/k/gate (4 elems each), butterfly + # the L2 norms, compute per-channel decay, then shuffle-transpose into + # the (k_lane, 8) register layout every group needs. + # ------------------------------------------------------------------ + q_tile = cute.local_tile(q, (1, 1, vals), (i_n, i_h, lane)) + k_tile = cute.local_tile(k, (1, 1, vals), (i_n, i_h, lane)) + g_tile = cute.local_tile(g, (1, 1, vals), (i_n, i_h, lane)) + dtb_tile = cute.local_tile(dt_bias, (1, vals), (i_h, lane)) + cute.autovec_copy(q_tile, q_bf16) + cute.autovec_copy(k_tile, k_bf16) + cute.autovec_copy(g_tile, g_bf16) + cute.autovec_copy(dtb_tile, dtb_f32) + + if cutlass.const_expr(precomputed): + # Pre-computed convention: beta arrives already sigmoided. + a_exp = cutlass.Float32(0.0) + r_beta = cutlass.Float32(beta[(i_n, i_h)]) + else: + a_exp = cute.exp(cutlass.Float32(A_log[i_h]), fastmath=True) + b_logit = cutlass.Float32(beta[(i_n, i_h)]) + r_beta = cute.rcp( + cute.exp(-b_logit, fastmath=True) + 1.0, approx=True, ftz=True + ) + + sum_q = cutlass.Float32(0.0) + sum_k = cutlass.Float32(0.0) + for i in VALSI: + q_val = cutlass.Float32(q_bf16[i]) + k_val = cutlass.Float32(k_bf16[i]) + q_src[i] = q_val + k_src[i] = k_val + sum_q += q_val * q_val + sum_k += k_val * k_val + if cutlass.const_expr(precomputed): + # g is the log-space decay: d = exp(g). + d_src[i] = cute.exp2( + cutlass.Float32(g_bf16[i]) * _LOG2_E, fastmath=True + ) + else: + gate_x = cutlass.Float32(g_bf16[i]) + dtb_f32[i] + sig = cute.rcp( + cute.exp(-(a_exp * gate_x), fastmath=True) + 1.0, + approx=True, + ftz=True, + ) + d_src[i] = cute.exp2(lb_log2e * sig, fastmath=True) + + for offset in [16, 8, 4, 2, 1]: + sum_q += cute.arch.shuffle_sync_bfly(sum_q, offset=offset, mask=0xFFFFFFFF) + sum_k += cute.arch.shuffle_sync_bfly(sum_k, offset=offset, mask=0xFFFFFFFF) + inv_q = cute.rsqrt(sum_q + eps, fastmath=True) * scale + inv_k = cute.rsqrt(sum_k + eps, fastmath=True) + + for i in VECI: + source_lane = 2 * k_lane + i // vals + source_value = i % vals + r_q[i] = ( + cute.arch.shuffle_sync( + q_src[source_value], offset=source_lane, mask=0xFFFFFFFF + ) + * inv_q + ) + r_k[i] = ( + cute.arch.shuffle_sync( + k_src[source_value], offset=source_lane, mask=0xFFFFFFFF + ) + * inv_k + ) + r_d[i] = cute.arch.shuffle_sync( + d_src[source_value], offset=source_lane, mask=0xFFFFFFFF + ) + + # ------------------------------------------------------------------ + # Main loop: ilp_rows V rows in registers per iteration, with the + # next iteration's loads issued before the current compute chain. + # ------------------------------------------------------------------ + for it in ITERI: + vb = v_base + it * ilp_rows + cur = it % 2 + if cutlass.const_expr(it + 1 < ITERS): + for r in ROWS: + if cutlass.const_expr(evict_first_state): + cute.autovec_copy( + ht_of[it + 1][r], + r_hb[1 - cur][r], + l1c_evict_priority=( + cute.nvgpu.CacheEvictionPriority.EVICT_FIRST + ), + ) + else: + cute.autovec_copy(ht_of[it + 1][r], r_hb[1 - cur][r]) + cute.autovec_copy(vt_of[it + 1], r_v_bf16[1 - cur]) + ht = ht_of[it] + + # Pass 1: s = (h * decay) . k, reduced over K. The fp32 decayed + # state is transient — state stays packed bf16 in registers (the + # Cake register-economy trick: ~32 fewer live registers, so more + # CTAs fit per SM; the decay product is recomputed identically in + # pass 2). + s_e = [cutlass.Float32(0.0) for _ in ROWS] + s_o = [cutlass.Float32(0.0) for _ in ROWS] + for i in VECP: + for r in ROWS: + h0, h1 = _bf16x2_to_f32x2(r_hb32[cur][r][i // 2]) + if cutlass.const_expr(use_packed_fma): + h0, h1 = cute.arch.fma_packed_f32x2( + src_a=(h0, h1), + src_b=(r_d[i], r_d[i + 1]), + src_c=(cutlass.Float32(0.0), cutlass.Float32(0.0)), + ) + s_e[r], s_o[r] = cute.arch.fma_packed_f32x2( + src_a=(h0, h1), + src_b=(r_k[i], r_k[i + 1]), + src_c=(s_e[r], s_o[r]), + ) + else: + h0 = h0 * r_d[i] + h1 = h1 * r_d[i + 1] + s_e[r] = s_e[r] + h0 * r_k[i] + s_o[r] = s_o[r] + h1 * r_k[i + 1] + s = [s_e[r] + s_o[r] for r in ROWS] + for offset in [8, 4, 2, 1]: + for r in ROWS: + s[r] += cute.arch.shuffle_sync_bfly( + s[r], offset=offset, mask=0xFFFFFFFF + ) + + # Delta rule: vn = (v - s) * beta. + vn = [(cutlass.Float32(r_v_bf16[cur][r]) - s[r]) * r_beta for r in ROWS] + + # Pass 2: recompute h*decay, apply the rank-1 update h += k * vn, + # accumulate o = (h_new * q), and pack h_new straight back into + # the bf16 staging registers (which then feed the STG.128s). + o_e = [cutlass.Float32(0.0) for _ in ROWS] + o_o = [cutlass.Float32(0.0) for _ in ROWS] + for i in VECP: + for r in ROWS: + h0, h1 = _bf16x2_to_f32x2(r_hb32[cur][r][i // 2]) + if cutlass.const_expr(use_packed_fma): + h0, h1 = cute.arch.fma_packed_f32x2( + src_a=(h0, h1), + src_b=(r_d[i], r_d[i + 1]), + src_c=(cutlass.Float32(0.0), cutlass.Float32(0.0)), + ) + h0, h1 = cute.arch.fma_packed_f32x2( + src_a=(r_k[i], r_k[i + 1]), + src_b=(vn[r], vn[r]), + src_c=(h0, h1), + ) + o_e[r], o_o[r] = cute.arch.fma_packed_f32x2( + src_a=(h0, h1), + src_b=(r_q[i], r_q[i + 1]), + src_c=(o_e[r], o_o[r]), + ) + else: + h0 = h0 * r_d[i] + h1 = h1 * r_d[i + 1] + h0 = r_k[i] * vn[r] + h0 + h1 = r_k[i + 1] * vn[r] + h1 + o_e[r] = o_e[r] + h0 * r_q[i] + o_o[r] = o_o[r] + h1 * r_q[i + 1] + r_hb32[cur][r][i // 2] = _f32x2_to_bf16x2(h0, h1) + + # State write-back: h_new is already packed bf16; the STG.128s + # issue before the output butterfly so they drain during the + # shuffle reduction. + for r in ROWS: + cute.autovec_copy(r_hb[cur][r], ht[r]) + + o_val = [o_e[r] + o_o[r] for r in ROWS] + for offset in [8, 4, 2, 1]: + for r in ROWS: + o_val[r] += cute.arch.shuffle_sync_bfly( + o_val[r], offset=offset, mask=0xFFFFFFFF + ) + + if k_lane == 0: + for r in ROWS: + r_o_bf16[r] = cutlass.BFloat16(o_val[r]) + ot = cute.local_tile(o, (1, 1, ilp_rows), (i_n, i_h, vb // ilp_rows)) + cute.autovec_copy(r_o_bf16, ot) + + +@cute.kernel +def _kda_packed_t1_smem_kernel( + state: cute.Tensor, # [pool, H, V, K] bf16, padded slot stride allowed + q: cute.Tensor, # [B, H, K] bf16 strided view of mixed_qkv + k: cute.Tensor, # [B, H, K] bf16 strided view of mixed_qkv + v: cute.Tensor, # [B, H, V] bf16 strided view of mixed_qkv + g: cute.Tensor, # [B, H, K] bf16 strided view of raw_gate + beta: cute.Tensor, # [B, H] bf16 (row stride dynamic) + A_log: cute.Tensor, # [H] f32 + dt_bias: cute.Tensor, # [H, K] f32 + o: cute.Tensor, # [B, H, V] bf16 contiguous + state_indices: cute.Tensor, # [B] i32 + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + lb_log2e: cutlass.Constexpr[float], + precomputed: cutlass.Constexpr[bool], + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + num_v_tiles: cutlass.Constexpr[int], + ilp_rows: cutlass.Constexpr[int], + n_stages: cutlass.Constexpr[int], + chunk_rows: cutlass.Constexpr[int], + cp_l2_hint: cutlass.Constexpr[int], + heads_per_cta: cutlass.Constexpr[int], + bulk_store: cutlass.Constexpr[bool], + tma_read: cutlass.Constexpr[bool], + l2_policy_mode: cutlass.Constexpr[int], + private_ring: cutlass.Constexpr[bool], + use_packed_fma: cutlass.Constexpr[bool], +): + """cp.async-pipelined variant for the DRAM-bound large-batch regime. + + State rows stream through a shared-memory ring buffer via cp.async.cg + (L1 bypass): unlike register prefetching, in-flight bytes are not bounded + by the register file, so each CTA keeps ``(n_stages - 1) * 8 KiB`` of + reads outstanding while computing — the standard GEMM mainloop shape. + Requires a 16-byte-aligned state pool (cp.async constraint); the + register-prefetch kernel remains the fallback for odd pools. + """ + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + + # Multiple heads per CTA (heads are adjacent compact 32 KiB blocks in + # the slot, so the ring streams one contiguous span and only the q/k/gate + # staging re-runs at head boundaries). Requires tile_v == V for HPC > 1. + HPC: cutlass.Constexpr[int] = heads_per_cta + HB: cutlass.Constexpr[int] = H // HPC + i_v = bidx % num_v_tiles + tmp = bidx // num_v_tiles + i_h0 = (tmp % HB) * HPC + i_n = tmp // HB + + lane = tidx % 32 + k_lane = tidx % _LANES_PER_ROW + group_idx = tidx // _LANES_PER_ROW + + # One pipeline chunk covers SUBI compute sub-iterations of + # ilp_rows rows per group (bigger chunks halve the wait+barrier count + # without touching the register profile). + CHUNK_ROWS: cutlass.Constexpr[int] = chunk_rows + ITERS_H: cutlass.Constexpr[int] = tile_v // CHUNK_ROWS + ITERS: cutlass.Constexpr[int] = HPC * ITERS_H + SUB: cutlass.Constexpr[int] = CHUNK_ROWS // (_NUM_GROUPS * ilp_rows) + + vec: cutlass.Constexpr[int] = _ELEMS_PER_LANE + vals: cutlass.Constexpr[int] = _VALS_PER_LANE + + ROWS = tuple(range(ilp_rows)) + VECI = tuple(range(vec)) + VECP = tuple(range(0, vec, 2)) + VALSI = tuple(range(vals)) + ITERI = tuple(range(ITERS)) + ITERI_H = tuple(range(ITERS_H)) + SUBI = tuple(range(SUB)) + + # cp.async issue geometry: 128 threads cover one chunk with + # CHUNK_ROWS * 2 / 128 sixteen-byte copies each. + THREADS_PER_ROW: cutlass.Constexpr[int] = _NUM_THREADS // CHUNK_ROWS + COL_SPAN: cutlass.Constexpr[int] = K // THREADS_PER_ROW + crow = tidx // THREADS_PER_ROW + ccol = (tidx % THREADS_PER_ROW) * COL_SPAN + CPJ = tuple(range(COL_SPAN // vec)) + + # Bulk-store geometry: with SUB == 1 each warp's rows are contiguous + # (2 groups x ilp_rows rows = one span), so its write-back is a single + # shared->global bulk copy issued by lane 0. + warp_idx = tidx // 32 + WARP_ROWS: cutlass.Constexpr[int] = 2 * ilp_rows + + smem = cutlass.utils.SmemAllocator() + sH = smem.allocate_tensor( + cutlass.BFloat16, + cute.make_layout((n_stages, CHUNK_ROWS, K), stride=(CHUNK_ROWS * K, K, 1)), + 16, + ) + sMbar = smem.allocate_tensor( + cutlass.Int64, cute.make_layout((2 * n_stages,), stride=(1,)), 8 + ) + CHUNK_BYTES: cutlass.Constexpr[int] = CHUNK_ROWS * K * 2 + + raw_slot = state_indices[i_n] + + q_src = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.Float32 + ) + k_src = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.Float32 + ) + d_src = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.Float32 + ) + q_bf16 = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.BFloat16 + ) + k_bf16 = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.BFloat16 + ) + g_bf16 = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.BFloat16 + ) + dtb_f32 = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.Float32 + ) + r_q = cute.make_rmem_tensor(cute.make_layout((vec,), stride=(1,)), cutlass.Float32) + r_k = cute.make_rmem_tensor(cute.make_layout((vec,), stride=(1,)), cutlass.Float32) + r_d = cute.make_rmem_tensor(cute.make_layout((vec,), stride=(1,)), cutlass.Float32) + r_hb = [ + [ + [ + cute.make_rmem_tensor( + cute.make_layout((vec,), stride=(1,)), cutlass.BFloat16 + ) + for _ in ROWS + ] + for _ in SUBI + ] + for _ in range(2) + ] + r_hb32 = [ + [[cute.recast_tensor(t, cutlass.Int32) for t in bufs] for bufs in par] + for par in r_hb + ] + r_v_all = [ + cute.make_rmem_tensor( + cute.make_layout((ilp_rows,), stride=(1,)), cutlass.BFloat16 + ) + for _ in range(ITERS_H * SUB) + ] + r_o_bf16 = cute.make_rmem_tensor( + cute.make_layout((ilp_rows,), stride=(1,)), cutlass.BFloat16 + ) + + slot_dead = (raw_slot < 0) | (cutlass.Int64(raw_slot) >= state.shape[0]) + if slot_dead: + for r in ROWS: + r_o_bf16[r] = cutlass.BFloat16(0.0) + if k_lane == 0: + for it in ITERI: + for sub in SUBI: + vb = ( + i_v * tile_v + + (it % ITERS_H) * CHUNK_ROWS + + (sub * _NUM_GROUPS + group_idx) * ilp_rows + ) + ot = cute.local_tile( + o, + (1, 1, ilp_rows), + (i_n, i_h0 + it // ITERS_H, vb // ilp_rows), + ) + cute.autovec_copy(r_o_bf16, ot) + else: + h_slots = [ + state[(cutlass.Int64(raw_slot), i_h0 + hh, None, None)] for hh in range(HPC) + ] + h_base = h_slots[0].iterator.toint() + sh_base = cutlass.Int32(sH.iterator.toint()) + + # L2 eviction steering: state writes pin as evict_last so their + # write-back drains after kernel-end (in serving it overlaps the next + # kernel); the read stream optionally marks evict_first. + pol_w = cutlass.Int64(0) + pol_r = cutlass.Int64(0) + if cutlass.const_expr(l2_policy_mode >= 1): + pol_w = _l2_policy_evict_last() + if cutlass.const_expr(l2_policy_mode >= 2): + pol_r = _l2_policy_evict_first() + + # Prologue: put the first D chunks in flight before the q/k/gate + # staging math, which then hides their latency. In bulk-store mode one + # ring slot stays spare so refills never wait on the newest bulk read. + D: cutlass.Constexpr[int] = n_stages - (2 if bulk_store else 1) + if cutlass.const_expr(tma_read): + # TMA-1D ring: one bulk read per chunk, full/empty mbarriers + # instead of LDGSTS groups + CTA barriers. All slots prefill. + if tidx == 0: + for s in range(n_stages): + cute.arch.mbarrier_init(sMbar.iterator + s, 1) + cute.arch.mbarrier_init(sMbar.iterator + n_stages + s, _NUM_THREADS) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + if tidx == 0: + for c in range(min(n_stages, ITERS)): + cute.arch.mbarrier_arrive_and_expect_tx( + sMbar.iterator + c, CHUNK_BYTES + ) + _cp_async_bulk_read( + sh_base + c * CHUNK_BYTES, + h_base + cutlass.Int64((i_v * tile_v + c * CHUNK_ROWS) * K * 2), + cutlass.Int32(sMbar.iterator.toint()) + c * 8, + CHUNK_BYTES, + ) + else: + for c in range(min(D, ITERS)): + grow_base = i_v * tile_v + c * CHUNK_ROWS + crow + for j in CPJ: + elem_off = grow_base * K + ccol + j * vec + smem_byte = ( + (c % n_stages) * CHUNK_ROWS * K + crow * K + ccol + j * vec + ) * 2 + _cp_async_bf16x8_cg( + h_base, elem_off, sh_base + smem_byte, cp_l2_hint + ) + _cp_async_commit_group() + + # ------------------------------------------------------------------ + # Pipelined main loop: wait chunk -> barrier -> refill ring -> compute. + # ------------------------------------------------------------------ + for it in ITERI: + if cutlass.const_expr(it % ITERS_H == 0): + # Head-boundary staging: q/k/gate/decay/beta/kq and this + # head's v values. For heads after the first, the raw q/k/g/ + # dt loads were prefetched two chunks early into the (dead) + # staging buffers, so only the staging math runs here. + i_hh = i_h0 + it // ITERS_H + if cutlass.const_expr(it == 0): + q_tile = cute.local_tile(q, (1, 1, vals), (i_n, i_hh, lane)) + k_tile = cute.local_tile(k, (1, 1, vals), (i_n, i_hh, lane)) + g_tile = cute.local_tile(g, (1, 1, vals), (i_n, i_hh, lane)) + dtb_tile = cute.local_tile(dt_bias, (1, vals), (i_hh, lane)) + cute.autovec_copy(q_tile, q_bf16) + cute.autovec_copy(k_tile, k_bf16) + cute.autovec_copy(g_tile, g_bf16) + cute.autovec_copy(dtb_tile, dtb_f32) + + if cutlass.const_expr(precomputed): + # Pre-computed convention: beta arrives already sigmoided. + a_exp = cutlass.Float32(0.0) + r_beta = cutlass.Float32(beta[(i_n, i_hh)]) + else: + a_exp = cute.exp(cutlass.Float32(A_log[i_hh]), fastmath=True) + b_logit = cutlass.Float32(beta[(i_n, i_hh)]) + r_beta = cute.rcp( + cute.exp(-b_logit, fastmath=True) + 1.0, approx=True, ftz=True + ) + + sum_q = cutlass.Float32(0.0) + sum_k = cutlass.Float32(0.0) + for i in VALSI: + q_val = cutlass.Float32(q_bf16[i]) + k_val = cutlass.Float32(k_bf16[i]) + q_src[i] = q_val + k_src[i] = k_val + sum_q += q_val * q_val + sum_k += k_val * k_val + if cutlass.const_expr(precomputed): + # g is the log-space decay: d = exp(g). + d_src[i] = cute.exp2( + cutlass.Float32(g_bf16[i]) * _LOG2_E, fastmath=True + ) + else: + gate_x = cutlass.Float32(g_bf16[i]) + dtb_f32[i] + sig = cute.rcp( + cute.exp(-(a_exp * gate_x), fastmath=True) + 1.0, + approx=True, + ftz=True, + ) + d_src[i] = cute.exp2(lb_log2e * sig, fastmath=True) + + for offset in [16, 8, 4, 2, 1]: + sum_q += cute.arch.shuffle_sync_bfly( + sum_q, offset=offset, mask=0xFFFFFFFF + ) + sum_k += cute.arch.shuffle_sync_bfly( + sum_k, offset=offset, mask=0xFFFFFFFF + ) + inv_q = cute.rsqrt(sum_q + eps, fastmath=True) * scale + inv_k = cute.rsqrt(sum_k + eps, fastmath=True) + + for i in VECI: + source_lane = 2 * k_lane + i // vals + source_value = i % vals + r_q[i] = ( + cute.arch.shuffle_sync( + q_src[source_value], offset=source_lane, mask=0xFFFFFFFF + ) + * inv_q + ) + r_k[i] = ( + cute.arch.shuffle_sync( + k_src[source_value], offset=source_lane, mask=0xFFFFFFFF + ) + * inv_k + ) + r_d[i] = cute.arch.shuffle_sync( + d_src[source_value], offset=source_lane, mask=0xFFFFFFFF + ) + + # kq = (normalized q) . (normalized k), used by the fused output form + # o = (h*d).q + vn * kq — one warp butterfly at staging instead of a + # second dependent reduction tree in every iteration. + kq_p = cutlass.Float32(0.0) + for i in VALSI: + kq_p += q_src[i] * k_src[i] + for offset in [16, 8, 4, 2, 1]: + kq_p += cute.arch.shuffle_sync_bfly( + kq_p, offset=offset, mask=0xFFFFFFFF + ) + r_kq = kq_p * (inv_q * inv_k) + + for itl in ITERI_H: + for sub in SUBI: + vb_it = ( + i_v * tile_v + + itl * CHUNK_ROWS + + (sub * _NUM_GROUPS + group_idx) * ilp_rows + ) + vt = cute.local_tile( + v, (1, 1, ilp_rows), (i_n, i_hh, vb_it // ilp_rows) + ) + cute.autovec_copy(vt, r_v_all[itl * SUB + sub]) + + if cutlass.const_expr(private_ring): + # Per-thread deep wait: own bytes of chunk it+1 resident, so + # the prefetched LDS below needs no CTA barrier — except + # while PROLOGUE chunks are being consumed: the prologue uses + # the shared crow/ccol issue geometry (better warp + # coalescing), so chunks 0..D-1 arrive via other threads' + # copies and each needs a publish after the matching wait. + _cp_async_wait_group_n( + max(0, min(D - 2, ITERS - 2 - it)) if it + 1 < ITERS else 0 + ) + if cutlass.const_expr(it <= D - 2): + cute.arch.barrier() + if cutlass.const_expr(it + D < ITERS): + c = it + D + for sub2 in SUBI: + for r in ROWS: + lrow = (sub2 * _NUM_GROUPS + group_idx) * ilp_rows + r + elem_off = ( + i_v * tile_v + c * CHUNK_ROWS + lrow + ) * K + k_lane * vec + smem_byte = ( + ((c % n_stages) * CHUNK_ROWS + lrow) * K + k_lane * vec + ) * 2 + _cp_async_bf16x8_cg( + h_base, elem_off, sh_base + smem_byte, cp_l2_hint + ) + _cp_async_commit_group() + elif cutlass.const_expr(tma_read): + cute.arch.mbarrier_wait( + sMbar.iterator + (it % n_stages), (it // n_stages) & 1 + ) + else: + # Deep wait: chunk it+1 must also be resident so its LDS can + # issue this iteration and cover a full compute phase. + _cp_async_wait_group_n( + max(0, min(D - 2, ITERS - 2 - it)) if it + 1 < ITERS else 0 + ) + if cutlass.const_expr(bulk_store): + if cutlass.const_expr(it >= 2): + # The slot refilled below was bulk-stored two + # iterations ago; allow the newest bulk group to stay + # in flight. + _cp_async_bulk_wait_read(1) + elif cutlass.const_expr(it == 1): + _cp_async_bulk_wait_read(0) + cute.arch.barrier() + + if cutlass.const_expr(it + D < ITERS): + c = it + D + grow_base = i_v * tile_v + c * CHUNK_ROWS + crow + for j in CPJ: + elem_off = grow_base * K + ccol + j * vec + smem_byte = ( + (c % n_stages) * CHUNK_ROWS * K + crow * K + ccol + j * vec + ) * 2 + if cutlass.const_expr(l2_policy_mode >= 2): + _cp_async_bf16x8_cg_hint( + h_base, elem_off, sh_base + smem_byte, pol_r + ) + else: + _cp_async_bf16x8_cg( + h_base, elem_off, sh_base + smem_byte, cp_l2_hint + ) + _cp_async_commit_group() + + h_slot = h_slots[it // ITERS_H] + cur = it % 2 + if cutlass.const_expr(it % ITERS_H == ITERS_H - 2 and it + 2 < ITERS): + # Prefetch the NEXT head's raw staging inputs into the dead + # bf16 buffers; their latency rides the remaining two chunks. + i_hn = i_h0 + it // ITERS_H + 1 + q_tile = cute.local_tile(q, (1, 1, vals), (i_n, i_hn, lane)) + k_tile = cute.local_tile(k, (1, 1, vals), (i_n, i_hn, lane)) + g_tile = cute.local_tile(g, (1, 1, vals), (i_n, i_hn, lane)) + dtb_tile = cute.local_tile(dt_bias, (1, vals), (i_hn, lane)) + cute.autovec_copy(q_tile, q_bf16) + cute.autovec_copy(k_tile, k_bf16) + cute.autovec_copy(g_tile, g_bf16) + cute.autovec_copy(dtb_tile, dtb_f32) + if cutlass.const_expr(it == 0): + stage0 = sH[(0, None, None)] + for sub in SUBI: + for r in ROWS: + hs = cute.local_tile( + stage0, + (1, vec), + ( + (sub * _NUM_GROUPS + group_idx) * ilp_rows + r, + k_lane, + ), + ) + cute.autovec_copy(hs, r_hb[0][sub][r]) + if cutlass.const_expr(it + 1 < ITERS): + stage_n = sH[((it + 1) % n_stages, None, None)] + for sub in SUBI: + for r in ROWS: + hs = cute.local_tile( + stage_n, + (1, vec), + ( + (sub * _NUM_GROUPS + group_idx) * ilp_rows + r, + k_lane, + ), + ) + cute.autovec_copy(hs, r_hb[1 - cur][sub][r]) + for sub in SUBI: + vb = ( + i_v * tile_v + + (it % ITERS_H) * CHUNK_ROWS + + (sub * _NUM_GROUPS + group_idx) * ilp_rows + ) + if cutlass.const_expr(tma_read and sub == SUB - 1): + # Slot consumed into registers: free it and let thread 0 + # refill immediately — the bulk read streams during the + # remaining compute of this iteration. + cute.arch.mbarrier_arrive( + sMbar.iterator + n_stages + (it % n_stages) + ) + if cutlass.const_expr(it + n_stages < ITERS): + if tidx == 0: + cute.arch.mbarrier_wait( + sMbar.iterator + n_stages + (it % n_stages), + (it // n_stages) & 1, + ) + cute.arch.mbarrier_arrive_and_expect_tx( + sMbar.iterator + (it % n_stages), CHUNK_BYTES + ) + _cp_async_bulk_read( + sh_base + (it % n_stages) * CHUNK_BYTES, + h_base + + cutlass.Int64( + (i_v * tile_v + (it + n_stages) * CHUNK_ROWS) + * K + * 2 + ), + cutlass.Int32(sMbar.iterator.toint()) + + (it % n_stages) * 8, + CHUNK_BYTES, + ) + # Pass 1 (fused): s = (h*d).k and od = (h*d).q in one sweep; + # their butterflies interleave (independent trees), and pass 2 + # no longer carries an output reduction. + s_e = [cutlass.Float32(0.0) for _ in ROWS] + s_o = [cutlass.Float32(0.0) for _ in ROWS] + o_e = [cutlass.Float32(0.0) for _ in ROWS] + o_o = [cutlass.Float32(0.0) for _ in ROWS] + for i in VECP: + for r in ROWS: + h0, h1 = _bf16x2_to_f32x2(r_hb32[cur][sub][r][i // 2]) + if cutlass.const_expr(use_packed_fma): + h0, h1 = cute.arch.fma_packed_f32x2( + src_a=(h0, h1), + src_b=(r_d[i], r_d[i + 1]), + src_c=(cutlass.Float32(0.0), cutlass.Float32(0.0)), + ) + s_e[r], s_o[r] = cute.arch.fma_packed_f32x2( + src_a=(h0, h1), + src_b=(r_k[i], r_k[i + 1]), + src_c=(s_e[r], s_o[r]), + ) + o_e[r], o_o[r] = cute.arch.fma_packed_f32x2( + src_a=(h0, h1), + src_b=(r_q[i], r_q[i + 1]), + src_c=(o_e[r], o_o[r]), + ) + else: + h0 = h0 * r_d[i] + h1 = h1 * r_d[i + 1] + s_e[r] = s_e[r] + h0 * r_k[i] + s_o[r] = s_o[r] + h1 * r_k[i + 1] + o_e[r] = o_e[r] + h0 * r_q[i] + o_o[r] = o_o[r] + h1 * r_q[i + 1] + s_red = [s_e[r] + s_o[r] for r in ROWS] + od = [o_e[r] + o_o[r] for r in ROWS] + for offset in [8, 4, 2, 1]: + for r in ROWS: + s_red[r] += cute.arch.shuffle_sync_bfly( + s_red[r], offset=offset, mask=0xFFFFFFFF + ) + od[r] += cute.arch.shuffle_sync_bfly( + od[r], offset=offset, mask=0xFFFFFFFF + ) + + vn = [ + (cutlass.Float32(r_v_all[(it % ITERS_H) * SUB + sub][r]) - s_red[r]) + * r_beta + for r in ROWS + ] + + # Output is already reduced: o = (h*d).q + vn * kq. + if k_lane == 0: + for r in ROWS: + r_o_bf16[r] = cutlass.BFloat16(od[r] + vn[r] * r_kq) + ot = cute.local_tile( + o, (1, 1, ilp_rows), (i_n, i_hh, vb // ilp_rows) + ) + cute.autovec_copy(r_o_bf16, ot) + + # Pass 2 (store only): h_new = h*decay + k*vn, repack bf16. + for i in VECP: + for r in ROWS: + h0, h1 = _bf16x2_to_f32x2(r_hb32[cur][sub][r][i // 2]) + if cutlass.const_expr(use_packed_fma): + h0, h1 = cute.arch.fma_packed_f32x2( + src_a=(h0, h1), + src_b=(r_d[i], r_d[i + 1]), + src_c=(cutlass.Float32(0.0), cutlass.Float32(0.0)), + ) + h0, h1 = cute.arch.fma_packed_f32x2( + src_a=(r_k[i], r_k[i + 1]), + src_b=(vn[r], vn[r]), + src_c=(h0, h1), + ) + else: + h0 = h0 * r_d[i] + h1 = h1 * r_d[i + 1] + h0 = r_k[i] * vn[r] + h0 + h1 = r_k[i + 1] * vn[r] + h1 + r_hb32[cur][sub][r][i // 2] = _f32x2_to_bf16x2(h0, h1) + + if cutlass.const_expr(bulk_store): + # h_new goes back into this chunk's smem slot (warp-local + # rows), then one bulk store per warp pushes it out via the + # async proxy. + stage_v = sH[(it % n_stages, None, None)] + for r in ROWS: + hs_w = cute.local_tile( + stage_v, + (1, vec), + ((sub * _NUM_GROUPS + group_idx) * ilp_rows + r, k_lane), + ) + cute.autovec_copy(r_hb[cur][sub][r], hs_w) + _sync_warp() + if lane == 0: + _fence_proxy_async() + row0 = it * CHUNK_ROWS + warp_idx * WARP_ROWS + g_addr = h_base + cutlass.Int64((i_v * tile_v + row0) * K * 2) + s_addr = sh_base + ( + ((it % n_stages) * CHUNK_ROWS + warp_idx * WARP_ROWS) + * K + * 2 + ) + _cp_async_bulk_store(g_addr, s_addr, WARP_ROWS * K * 2) + _cp_async_bulk_commit() + else: + if cutlass.const_expr(l2_policy_mode >= 1): + for r in ROWS: + st_addr = h_base + cutlass.Int64( + ((it // ITERS_H) * V * K + (vb + r) * K + k_lane * vec) + * 2 + ) + _stg_hint_v4_b32( + st_addr, + r_hb32[cur][sub][r][0], + r_hb32[cur][sub][r][1], + r_hb32[cur][sub][r][2], + r_hb32[cur][sub][r][3], + pol_w, + ) + else: + for r in ROWS: + ht = cute.local_tile(h_slot, (1, vec), (vb + r, k_lane)) + cute.autovec_copy(r_hb[cur][sub][r], ht) + + if cutlass.const_expr(bulk_store): + # Writes must be complete (not just smem-read) before exit. + _cp_async_bulk_wait(0) + + +@cute.kernel +def _kda_packed_t1_persist_kernel( + state: cute.Tensor, # [pool, H, V, K] bf16, padded slot stride allowed + q: cute.Tensor, # [B, H, K] bf16 strided view of mixed_qkv + k: cute.Tensor, # [B, H, K] bf16 strided view of mixed_qkv + v: cute.Tensor, # [B, H, V] bf16 strided view of mixed_qkv + g: cute.Tensor, # [B, H, K] bf16 strided view of raw_gate + beta: cute.Tensor, # [B, H] bf16 (row stride dynamic) + A_log: cute.Tensor, # [H] f32 + dt_bias: cute.Tensor, # [H, K] f32 + o: cute.Tensor, # [B, H, V] bf16 contiguous + state_indices: cute.Tensor, # [B] i32 + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + lb_log2e: cutlass.Constexpr[float], + precomputed: cutlass.Constexpr[bool], + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + n_stages: cutlass.Constexpr[int], + cp_l2_hint: cutlass.Constexpr[int], + use_packed_fma: cutlass.Constexpr[bool], +): + """Persistent v2: the champion pipeline with items streamed per CTA. + + Composition of every proven win — fused output form, LDS double-buffer, + barrier-free private cp.async ring — with a grid-stride item loop whose + ring NEVER drains: chunk look-ahead and staging loads cross item + boundaries, so the per-item fill bubble (the dominant remaining stall) + is paid once per CTA instead of once per item. Requires + ITERS_H % n_stages == 0 (slot/parity arithmetic stays compile-time) and + a 16-byte-aligned pool. One item = one whole (batch, head) head block. + """ + ilp_rows: cutlass.Constexpr[int] = 2 + CHUNK_ROWS: cutlass.Constexpr[int] = 16 + ITERS: cutlass.Constexpr[int] = V // CHUNK_ROWS # 8 chunks per item + P: cutlass.Constexpr[int] = n_stages - 1 # in-flight depth (2-cold refill) + + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + gdim, _, _ = cute.arch.grid_dim() + + B = cutlass.Int32(q.shape[0]) + num_items = B * H + + lane = tidx % 32 + k_lane = tidx % _LANES_PER_ROW + group_idx = tidx // _LANES_PER_ROW + + vec: cutlass.Constexpr[int] = _ELEMS_PER_LANE + vals: cutlass.Constexpr[int] = _VALS_PER_LANE + ROWS = tuple(range(ilp_rows)) + VECI = tuple(range(vec)) + VECP = tuple(range(0, vec, 2)) + VALSI = tuple(range(vals)) + ITERI = tuple(range(ITERS)) + + smem = cutlass.utils.SmemAllocator() + sH = smem.allocate_tensor( + cutlass.BFloat16, + cute.make_layout((n_stages, CHUNK_ROWS, K), stride=(CHUNK_ROWS * K, K, 1)), + 16, + ) + sh_base = cutlass.Int32(sH.iterator.toint()) + + q_src = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.Float32 + ) + k_src = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.Float32 + ) + d_src = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.Float32 + ) + q_bf16 = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.BFloat16 + ) + k_bf16 = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.BFloat16 + ) + g_bf16 = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.BFloat16 + ) + dtb_f32 = cute.make_rmem_tensor( + cute.make_layout((vals,), stride=(1,)), cutlass.Float32 + ) + r_q = cute.make_rmem_tensor(cute.make_layout((vec,), stride=(1,)), cutlass.Float32) + r_k = cute.make_rmem_tensor(cute.make_layout((vec,), stride=(1,)), cutlass.Float32) + r_d = cute.make_rmem_tensor(cute.make_layout((vec,), stride=(1,)), cutlass.Float32) + r_hb = [ + [ + cute.make_rmem_tensor( + cute.make_layout((vec,), stride=(1,)), cutlass.BFloat16 + ) + for _ in ROWS + ] + for _ in range(2) + ] + r_hb32 = [[cute.recast_tensor(t, cutlass.Int32) for t in bufs] for bufs in r_hb] + r_v_all = [ + cute.make_rmem_tensor( + cute.make_layout((ilp_rows,), stride=(1,)), cutlass.BFloat16 + ) + for _ in ITERI + ] + r_o_bf16 = cute.make_rmem_tensor( + cute.make_layout((ilp_rows,), stride=(1,)), cutlass.BFloat16 + ) + + # Item metadata: current + look-ahead (byte base of the 32 KiB head + # block; ok flags gate compute and refills of inactive/out-of-range). + item_cur = cutlass.Int32(bidx) + i_n_cur = item_cur // H + i_h_cur = item_cur % H + slot_cur = state_indices[i_n_cur] + base_cur = cutlass.Int64(0) + ok_cur = (slot_cur >= 0) & (cutlass.Int64(slot_cur) < state.shape[0]) + if ok_cur: + base_cur = state[ + (cutlass.Int64(slot_cur), i_h_cur, None, None) + ].iterator.toint() + + item_nxt = item_cur + gdim + i_n_nxt = item_nxt // H + i_h_nxt = item_nxt % H + ok_nxt = item_nxt < num_items + base_nxt = cutlass.Int64(0) + if ok_nxt: + slot_nxt = state_indices[i_n_nxt] + if (slot_nxt >= 0) & (cutlass.Int64(slot_nxt) < state.shape[0]): + base_nxt = state[ + (cutlass.Int64(slot_nxt), i_h_nxt, None, None) + ].iterator.toint() + else: + ok_nxt = False + + # Prologue: first P chunks of item 0 (empty commits when inactive) — + # self-service geometry, each thread fetches only the bytes it consumes. + for c in range(P): + if ok_cur: + for r in ROWS: + lrow = group_idx * ilp_rows + r + elem_off = (c * CHUNK_ROWS + lrow) * K + k_lane * vec + smem_byte = ( + ((c % n_stages) * CHUNK_ROWS + lrow) * K + k_lane * vec + ) * 2 + _cp_async_bf16x8_cg(base_cur, elem_off, sh_base + smem_byte, cp_l2_hint) + _cp_async_commit_group() + + n_my_items = (num_items - item_cur + gdim - 1) // gdim + for _j in cutlass.range(n_my_items, unroll=1): + # ------------------------------------------------------------------ + # Per-item staging (q/k/g/dt loads land here for the FIRST item; + # later items' loads were prefetched during the previous item and + # only the math runs). The ring keeps streaming underneath. + # ------------------------------------------------------------------ + if _j == 0: + q_tile = cute.local_tile(q, (1, 1, vals), (i_n_cur, i_h_cur, lane)) + k_tile = cute.local_tile(k, (1, 1, vals), (i_n_cur, i_h_cur, lane)) + g_tile = cute.local_tile(g, (1, 1, vals), (i_n_cur, i_h_cur, lane)) + dtb_tile = cute.local_tile(dt_bias, (1, vals), (i_h_cur, lane)) + cute.autovec_copy(q_tile, q_bf16) + cute.autovec_copy(k_tile, k_bf16) + cute.autovec_copy(g_tile, g_bf16) + cute.autovec_copy(dtb_tile, dtb_f32) + + if cutlass.const_expr(precomputed): + # Pre-computed convention: beta arrives already sigmoided. + a_exp = cutlass.Float32(0.0) + r_beta = cutlass.Float32(beta[(i_n_cur, i_h_cur)]) + else: + a_exp = cute.exp(cutlass.Float32(A_log[i_h_cur]), fastmath=True) + b_logit = cutlass.Float32(beta[(i_n_cur, i_h_cur)]) + r_beta = cute.rcp( + cute.exp(-b_logit, fastmath=True) + 1.0, approx=True, ftz=True + ) + + sum_q = cutlass.Float32(0.0) + sum_k = cutlass.Float32(0.0) + for i in VALSI: + q_val = cutlass.Float32(q_bf16[i]) + k_val = cutlass.Float32(k_bf16[i]) + q_src[i] = q_val + k_src[i] = k_val + sum_q += q_val * q_val + sum_k += k_val * k_val + if cutlass.const_expr(precomputed): + # g is the log-space decay: d = exp(g). + d_src[i] = cute.exp2( + cutlass.Float32(g_bf16[i]) * _LOG2_E, fastmath=True + ) + else: + gate_x = cutlass.Float32(g_bf16[i]) + dtb_f32[i] + sig = cute.rcp( + cute.exp(-(a_exp * gate_x), fastmath=True) + 1.0, + approx=True, + ftz=True, + ) + d_src[i] = cute.exp2(lb_log2e * sig, fastmath=True) + + kq_p = cutlass.Float32(0.0) + for i in VALSI: + kq_p += q_src[i] * k_src[i] + for offset in [16, 8, 4, 2, 1]: + sum_q += cute.arch.shuffle_sync_bfly(sum_q, offset=offset, mask=0xFFFFFFFF) + sum_k += cute.arch.shuffle_sync_bfly(sum_k, offset=offset, mask=0xFFFFFFFF) + kq_p += cute.arch.shuffle_sync_bfly(kq_p, offset=offset, mask=0xFFFFFFFF) + inv_q = cute.rsqrt(sum_q + eps, fastmath=True) * scale + inv_k = cute.rsqrt(sum_k + eps, fastmath=True) + r_kq = kq_p * (inv_q * inv_k) + + for i in VECI: + source_lane = 2 * k_lane + i // vals + source_value = i % vals + r_q[i] = ( + cute.arch.shuffle_sync( + q_src[source_value], offset=source_lane, mask=0xFFFFFFFF + ) + * inv_q + ) + r_k[i] = ( + cute.arch.shuffle_sync( + k_src[source_value], offset=source_lane, mask=0xFFFFFFFF + ) + * inv_k + ) + r_d[i] = cute.arch.shuffle_sync( + d_src[source_value], offset=source_lane, mask=0xFFFFFFFF + ) + + for it in ITERI: + vt = cute.local_tile( + v, + (1, 1, ilp_rows), + ( + i_n_cur, + i_h_cur, + (it * CHUNK_ROWS + group_idx * ilp_rows) // ilp_rows, + ), + ) + cute.autovec_copy(vt, r_v_all[it]) + + # ------------------------------------------------------------------ + # 8 chunks; refills and the LDS prefetch cross into the next item. + # ------------------------------------------------------------------ + for it in ITERI: + cur = it % 2 + # Per-thread wait: own bytes of chunks p and p+1 resident + # (steady-state constant count; tail passes trivially). + _cp_async_wait_group_n(P - 2 if P >= 2 else 0) + + # Refill chunk p+P (slot (p-1) % n: two iterations cold). + if cutlass.const_expr(it + P < ITERS): + c = it + P + if ok_cur: + for r in ROWS: + lrow = group_idx * ilp_rows + r + elem_off = (c * CHUNK_ROWS + lrow) * K + k_lane * vec + smem_byte = ( + (((it + P) % n_stages) * CHUNK_ROWS + lrow) * K + + k_lane * vec + ) * 2 + _cp_async_bf16x8_cg( + base_cur, elem_off, sh_base + smem_byte, cp_l2_hint + ) + _cp_async_commit_group() + else: + c = it + P - ITERS + if ok_nxt: + for r in ROWS: + lrow = group_idx * ilp_rows + r + elem_off = (c * CHUNK_ROWS + lrow) * K + k_lane * vec + smem_byte = ( + (((it + P) % n_stages) * CHUNK_ROWS + lrow) * K + + k_lane * vec + ) * 2 + _cp_async_bf16x8_cg( + base_nxt, elem_off, sh_base + smem_byte, cp_l2_hint + ) + _cp_async_commit_group() + + # First chunk of the very first item lands in buffer 0 here. + if cutlass.const_expr(it == 0): + if _j == 0: + stage0 = sH[(0, None, None)] + for r in ROWS: + hs = cute.local_tile( + stage0, (1, vec), (group_idx * ilp_rows + r, k_lane) + ) + cute.autovec_copy(hs, r_hb[0][r]) + # LDS prefetch of chunk p+1 (crosses into the next item at it=7; + # parity (it+1)%2 stays compile-time since ITERS is even). + stage_n = sH[((it + 1) % n_stages, None, None)] + for r in ROWS: + hs = cute.local_tile( + stage_n, (1, vec), (group_idx * ilp_rows + r, k_lane) + ) + cute.autovec_copy(hs, r_hb[1 - cur][r]) + + # Staging prefetch for the next item, mid-item (dead buffers). + if cutlass.const_expr(it == ITERS - 2): + if ok_nxt: + q_tile = cute.local_tile(q, (1, 1, vals), (i_n_nxt, i_h_nxt, lane)) + k_tile = cute.local_tile(k, (1, 1, vals), (i_n_nxt, i_h_nxt, lane)) + g_tile = cute.local_tile(g, (1, 1, vals), (i_n_nxt, i_h_nxt, lane)) + dtb_tile = cute.local_tile(dt_bias, (1, vals), (i_h_nxt, lane)) + cute.autovec_copy(q_tile, q_bf16) + cute.autovec_copy(k_tile, k_bf16) + cute.autovec_copy(g_tile, g_bf16) + cute.autovec_copy(dtb_tile, dtb_f32) + + vb = it * CHUNK_ROWS + group_idx * ilp_rows + hb32_cur = r_hb32[cur] + r_v_it = r_v_all[it] + if ok_cur: + # Pass 1 (fused): s = (h*d).k and od = (h*d).q in one sweep. + s_e = [cutlass.Float32(0.0) for _ in ROWS] + s_o = [cutlass.Float32(0.0) for _ in ROWS] + o_e = [cutlass.Float32(0.0) for _ in ROWS] + o_o = [cutlass.Float32(0.0) for _ in ROWS] + for i in VECP: + for r in ROWS: + h0, h1 = _bf16x2_to_f32x2(hb32_cur[r][i // 2]) + h0, h1 = cute.arch.fma_packed_f32x2( + src_a=(h0, h1), + src_b=(r_d[i], r_d[i + 1]), + src_c=(cutlass.Float32(0.0), cutlass.Float32(0.0)), + ) + s_e[r], s_o[r] = cute.arch.fma_packed_f32x2( + src_a=(h0, h1), + src_b=(r_k[i], r_k[i + 1]), + src_c=(s_e[r], s_o[r]), + ) + o_e[r], o_o[r] = cute.arch.fma_packed_f32x2( + src_a=(h0, h1), + src_b=(r_q[i], r_q[i + 1]), + src_c=(o_e[r], o_o[r]), + ) + s = [s_e[r] + s_o[r] for r in ROWS] + od = [o_e[r] + o_o[r] for r in ROWS] + for offset in [8, 4, 2, 1]: + for r in ROWS: + s[r] += cute.arch.shuffle_sync_bfly( + s[r], offset=offset, mask=0xFFFFFFFF + ) + od[r] += cute.arch.shuffle_sync_bfly( + od[r], offset=offset, mask=0xFFFFFFFF + ) + + vn = [(cutlass.Float32(r_v_it[r]) - s[r]) * r_beta for r in ROWS] + + if k_lane == 0: + for r in ROWS: + r_o_bf16[r] = cutlass.BFloat16(od[r] + vn[r] * r_kq) + ot = cute.local_tile( + o, (1, 1, ilp_rows), (i_n_cur, i_h_cur, vb // ilp_rows) + ) + cute.autovec_copy(r_o_bf16, ot) + + # Pass 2 (store only): h_new = h*d + k*vn, repack, STG.128. + for i in VECP: + for r in ROWS: + h0, h1 = _bf16x2_to_f32x2(hb32_cur[r][i // 2]) + h0, h1 = cute.arch.fma_packed_f32x2( + src_a=(h0, h1), + src_b=(r_d[i], r_d[i + 1]), + src_c=(cutlass.Float32(0.0), cutlass.Float32(0.0)), + ) + h0, h1 = cute.arch.fma_packed_f32x2( + src_a=(r_k[i], r_k[i + 1]), + src_b=(vn[r], vn[r]), + src_c=(h0, h1), + ) + hb32_cur[r][i // 2] = _f32x2_to_bf16x2(h0, h1) + + for r in ROWS: + st_addr = base_cur + cutlass.Int64( + ((vb + r) * K + k_lane * vec) * 2 + ) + _stg_v4_b32( + st_addr, + hb32_cur[r][0], + hb32_cur[r][1], + hb32_cur[r][2], + hb32_cur[r][3], + ) + else: + for r in ROWS: + r_o_bf16[r] = cutlass.BFloat16(0.0) + if k_lane == 0: + ot = cute.local_tile( + o, (1, 1, ilp_rows), (i_n_cur, i_h_cur, vb // ilp_rows) + ) + cute.autovec_copy(r_o_bf16, ot) + + # Rotate item metadata. + item_cur = item_nxt + i_n_cur = i_n_nxt + i_h_cur = i_h_nxt + base_cur = base_nxt + ok_cur = ok_nxt + item_nxt = item_cur + gdim + i_n_nxt = item_nxt // H + i_h_nxt = item_nxt % H + ok_nxt = item_nxt < num_items + base_nxt = cutlass.Int64(0) + if ok_nxt: + slot_n2 = state_indices[i_n_nxt] + if (slot_n2 >= 0) & (cutlass.Int64(slot_n2) < state.shape[0]): + base_nxt = state[ + (cutlass.Int64(slot_n2), i_h_nxt, None, None) + ].iterator.toint() + else: + ok_nxt = False + + +@cute.jit +def _kda_packed_t1_persist_launch( + state: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + g: cute.Tensor, + beta: cute.Tensor, + A_log: cute.Tensor, + dt_bias: cute.Tensor, + o: cute.Tensor, + state_indices: cute.Tensor, + grid_ctas: cutlass.Int32, + stream: cuda.CUstream, + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + lb_log2e: cutlass.Constexpr[float], + precomputed: cutlass.Constexpr[bool], + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + n_stages: cutlass.Constexpr[int], + cp_l2_hint: cutlass.Constexpr[int], + use_packed_fma: cutlass.Constexpr[bool], +): + smem_bytes = n_stages * 16 * K * 2 + 128 + _kda_packed_t1_persist_kernel( + state, + q, + k, + v, + g, + beta, + A_log, + dt_bias, + o, + state_indices, + scale, + eps, + lb_log2e, + precomputed, + H, + K, + V, + n_stages, + cp_l2_hint, + use_packed_fma, + ).launch( + grid=[grid_ctas, 1, 1], + block=[_NUM_THREADS, 1, 1], + smem=smem_bytes, + stream=stream, + ) + + +@cute.jit +def _kda_packed_t1_smem_launch( + state: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + g: cute.Tensor, + beta: cute.Tensor, + A_log: cute.Tensor, + dt_bias: cute.Tensor, + o: cute.Tensor, + state_indices: cute.Tensor, + stream: cuda.CUstream, + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + lb_log2e: cutlass.Constexpr[float], + precomputed: cutlass.Constexpr[bool], + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + ilp_rows: cutlass.Constexpr[int], + n_stages: cutlass.Constexpr[int], + chunk_rows: cutlass.Constexpr[int], + cp_l2_hint: cutlass.Constexpr[int], + heads_per_cta: cutlass.Constexpr[int], + bulk_store: cutlass.Constexpr[bool], + tma_read: cutlass.Constexpr[bool], + l2_policy_mode: cutlass.Constexpr[int], + private_ring: cutlass.Constexpr[bool], + min_blocks: cutlass.Constexpr[int], + use_packed_fma: cutlass.Constexpr[bool], +): + num_v_tiles: cutlass.Constexpr[int] = V // tile_v + B = cute.size(q.shape[0]) + grid_size = B * (H // heads_per_cta) * num_v_tiles + smem_bytes = n_stages * chunk_rows * K * 2 + 2 * n_stages * 8 + 128 + _kda_packed_t1_smem_kernel( + state, + q, + k, + v, + g, + beta, + A_log, + dt_bias, + o, + state_indices, + scale, + eps, + lb_log2e, + precomputed, + H, + K, + V, + tile_v, + num_v_tiles, + ilp_rows, + n_stages, + chunk_rows, + cp_l2_hint, + heads_per_cta, + bulk_store, + tma_read, + l2_policy_mode, + private_ring, + use_packed_fma, + ).launch( + grid=[grid_size, 1, 1], + block=[_NUM_THREADS, 1, 1], + smem=smem_bytes, + stream=stream, + # Full carveout: the ring buffer is the only smem client and the + # default split caps residency below the register limit. + preferred_smem_carveout=100, + # __launch_bounds__ min-blocks: bounds the register allocation so the + # requested CTA residency is achievable (0 = compiler's choice). + min_blocks_per_mp=min_blocks, + ) + + +@cute.jit +def _kda_packed_t1_launch( + state: cute.Tensor, + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + g: cute.Tensor, + beta: cute.Tensor, + A_log: cute.Tensor, + dt_bias: cute.Tensor, + o: cute.Tensor, + state_indices: cute.Tensor, + stream: cuda.CUstream, + scale: cutlass.Constexpr[float], + eps: cutlass.Constexpr[float], + lb_log2e: cutlass.Constexpr[float], + precomputed: cutlass.Constexpr[bool], + H: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + tile_v: cutlass.Constexpr[int], + ilp_rows: cutlass.Constexpr[int], + num_groups: cutlass.Constexpr[int], + use_packed_fma: cutlass.Constexpr[bool], + evict_first_state: cutlass.Constexpr[bool], +): + num_v_tiles: cutlass.Constexpr[int] = V // tile_v + B = cute.size(q.shape[0]) + grid_size = B * H * num_v_tiles + _kda_packed_t1_kernel( + state, + q, + k, + v, + g, + beta, + A_log, + dt_bias, + o, + state_indices, + scale, + eps, + lb_log2e, + precomputed, + H, + K, + V, + tile_v, + num_v_tiles, + ilp_rows, + num_groups, + use_packed_fma, + evict_first_state, + ).launch( + grid=[grid_size, 1, 1], + block=[num_groups * _LANES_PER_ROW, 1, 1], + stream=stream, + # No shared memory: leave the unified cache to L1 (q/k/gate rows are + # re-read by all warps and by every V-tile CTA of a head). + preferred_smem_carveout=0, + ) + + +def _make_compile_inputs( + qkv_div: int, gate_div: int, pool_div: int, aux_aligned: bool = True +): + """Build symbolic compile-time tensor specs. + + ``*_div`` are element divisibility guarantees for the dynamic outer + strides; 8 elements = 16 bytes lets autovec emit LDG.128/STG.128 on bf16, + 1 disables vectorisation (compatibility path for oddly padded callers). + """ + B = cute.sym_int() + N = cute.sym_int() + H, K, V = _HEADS, _HEAD_DIM, _HEAD_DIM + + def align_for(div: int) -> int: + return 16 if div % 8 == 0 else 2 + + def mixed_view(div: int): + return cute.runtime.make_fake_tensor( + cute.BFloat16, + shape=(B, H, K), + stride=(cute.sym_int64(divisibility=div), K, 1), + assumed_align=align_for(div), + ) + + state_fake = cute.runtime.make_fake_tensor( + cute.BFloat16, + shape=(N, H, V, K), + stride=(cute.sym_int64(divisibility=pool_div), V * K, K, 1), + assumed_align=align_for(pool_div), + ) + gate_fake = cute.runtime.make_fake_tensor( + cute.BFloat16, + shape=(B, H, K), + stride=(cute.sym_int64(divisibility=gate_div), K, 1), + assumed_align=align_for(gate_div), + ) + beta_fake = cute.runtime.make_fake_tensor( + cute.BFloat16, + shape=(B, H), + stride=(cute.sym_int64(divisibility=1), 1), + assumed_align=2, + ) + + def make_compact(shape, dtype=cute.BFloat16): + # dt_bias / output may sit at any element-aligned address (e.g. + # views shifted off a 16 B boundary); only claim 16 B when the + # caller verified it. + if aux_aligned: + align = 16 + else: + align = 4 if dtype is cute.Float32 else 2 + return cute.runtime.make_fake_compact_tensor( + dtype, + shape, + assumed_align=align, + stride_order=tuple(reversed(range(len(shape)))), + ) + + a_log_fake = cute.runtime.make_fake_compact_tensor( + cute.Float32, (H,), assumed_align=4, stride_order=(0,) + ) + idx_fake = cute.runtime.make_fake_compact_tensor( + cute.Int32, (cute.sym_int(),), assumed_align=4, stride_order=(0,) + ) + + return ( + state_fake, + mixed_view(qkv_div), + mixed_view(qkv_div), + mixed_view(qkv_div), + gate_fake, + beta_fake, + a_log_fake, + make_compact((H, K), dtype=cute.Float32), + make_compact((B, H, V)), + idx_fake, + make_fake_stream(use_tvm_ffi_env_stream=True), + ) + + +@functools.cache +def _get_compiled( + tile_v: int, + ilp_rows: int, + qkv_div: int, + gate_div: int, + pool_div: int, + aux_aligned: bool, + use_packed_fma: bool, + evict_first_state: bool, + maxrreg: int = 0, + n_stages: int = 0, + num_groups: int = _NUM_GROUPS, + cp_l2_hint: int = 0, + chunk_rows: int = 0, + bulk_store: bool = False, + heads_per_cta: int = 1, + tma_read: bool = False, + persistent: bool = False, + l2_policy_mode: int = 0, + private_ring: bool = False, + min_blocks: int = 0, + precomputed: bool = False, +): + options = "--enable-tvm-ffi --generate-line-info --opt-level 3" + if maxrreg: + options += f" --ptxas-options='-maxrregcount={maxrreg}'" + if persistent: + return cute.compile( + _kda_packed_t1_persist_launch, + *_make_compile_inputs(qkv_div, gate_div, pool_div, aux_aligned)[:-1], + cutlass.Int32(0), + make_fake_stream(use_tvm_ffi_env_stream=True), + _SCALE, + _EPS, + _LOWER_BOUND * _LOG2_E, + precomputed, + _HEADS, + _HEAD_DIM, + _HEAD_DIM, + n_stages, + cp_l2_hint, + use_packed_fma, + options=options, + ) + if n_stages: + # cp.async pipelined variant (16-byte-aligned pools, tile_v >= 64). + return cute.compile( + _kda_packed_t1_smem_launch, + *_make_compile_inputs(qkv_div, gate_div, pool_div, aux_aligned), + _SCALE, + _EPS, + _LOWER_BOUND * _LOG2_E, + precomputed, + _HEADS, + _HEAD_DIM, + _HEAD_DIM, + tile_v, + ilp_rows, + n_stages, + chunk_rows or _NUM_GROUPS * ilp_rows, + cp_l2_hint, + heads_per_cta, + bulk_store, + tma_read, + l2_policy_mode, + private_ring, + min_blocks, + use_packed_fma, + options=options, + ) + return cute.compile( + _kda_packed_t1_launch, + *_make_compile_inputs(qkv_div, gate_div, pool_div, aux_aligned), + _SCALE, + _EPS, + _LOWER_BOUND * _LOG2_E, + precomputed, + _HEADS, + _HEAD_DIM, + _HEAD_DIM, + tile_v, + ilp_rows, + num_groups, + use_packed_fma, + evict_first_state, + options=options, + ) + + +@functools.cache +def _use_packed_fma() -> bool: + return torch.cuda.get_device_capability(0)[0] >= 10 + + +# tile_v thresholds on work_units = B * H: bigger tiles amortise the q/k/gate +# staging and CTA overheads, smaller tiles add CTAs so small batches still +# fill the SMs. Tuned on B200 (see benchmarks/bench_packed_kda_decode.py). + + +# Benchmark-override schedules: the best-known (ilp, groups, stages, evict) +# for each forced tile width, so ``tile_v=`` forcing exercises every kernel +# shape without hitting untuned/invalid combinations. +_FORCED_TILE_CONFIGS = { + 8: (8, 4, 2, 0, True), + 16: (16, 2, 8, 0, True), + 32: (32, 2, 8, 0, True), + 64: (64, 2, 8, 4, False), + 128: (128, 2, 8, 5, False), +} +_SUPPORTED_TILE_V = tuple(sorted(_FORCED_TILE_CONFIGS)) + + +def _select_config( + batch: int, forced_tile_v: Optional[int] = None +) -> tuple[int, int, int, int, bool]: + """Per-batch schedule: ``(tile_v, ilp_rows, num_groups, stages, evict)``. + + ``stages > 0`` selects the cp.async pipelined kernel (128 threads); + ``stages == 0`` selects the register-prefetch kernel with + ``num_groups * 16`` threads. Tuned on B200 against the Cake tile8/tile16 + schedules (see PR #4378 benchmark protocol); every row was the best of a + tile/ilp/stages/CTA-width sweep at that batch size. + """ + if forced_tile_v is not None: + return _FORCED_TILE_CONFIGS[forced_tile_v] + if batch <= 11: + cfg = (16, 2, 8, 0, True) # latency floor: 96 CTAs/head-tile at B=1 + elif batch <= 23: + cfg = (8, 4, 2, 0, True) # 32-thread CTAs: finest work granularity + elif batch <= 37: + # Half-head tiles keep 24*B CTAs in flight through the 148-SM + # occupancy valley; the 4-deep ring covers the shorter pipeline. + cfg = (64, 2, 8, 4, False) + else: + # Whole-head CTAs, fine chunks, 5-deep ring: with the fused output + # form, LDS double-buffering and the barrier-free private ring this + # shape wins everywhere from B=38 up (1.17-1.26x vs Cake). + cfg = (128, 2, 8, 5, False) + + tile_v, ilp_rows, num_groups, stages, evict = cfg + # Looked up per call, like every other tuning override below. + tile_v_env = os.environ.get("FLASHINFER_PACKED_KDA_TILE_V") + if tile_v_env: + tile_v = int(tile_v_env) + ilp_rows = max(1, min(4, tile_v // _NUM_GROUPS)) + num_groups = _NUM_GROUPS + ilp_env = os.environ.get("FLASHINFER_PACKED_KDA_ILP") + if ilp_env: + ilp_rows = int(ilp_env) + threads_env = os.environ.get("FLASHINFER_PACKED_KDA_THREADS") + if threads_env: + num_groups = int(threads_env) // _LANES_PER_ROW + if os.environ.get("FLASHINFER_PACKED_KDA_EVICT_FIRST"): + evict = os.environ["FLASHINFER_PACKED_KDA_EVICT_FIRST"] == "1" + + # Reject schedules the kernels cannot trace correctly: the grid covers V + # with tile_v-wide tiles, every group must own a whole number of ilp-row + # blocks, and the DSL needs vector lengths of at least two. + if tile_v not in _FORCED_TILE_CONFIGS: + raise ValueError(f"tile_v must be one of {_SUPPORTED_TILE_V}, got {tile_v}") + if num_groups <= 0 or _NUM_GROUPS % num_groups != 0: + raise ValueError( + f"num_groups (THREADS/16) must be a divisor of {_NUM_GROUPS}, " + f"got {num_groups}" + ) + rows_per_group = tile_v // num_groups + if ilp_rows < 2 or tile_v % num_groups != 0 or rows_per_group % ilp_rows != 0: + raise ValueError( + f"invalid schedule tile_v={tile_v}, ilp_rows={ilp_rows}, " + f"num_groups={num_groups}: each of the {num_groups} groups must " + f"cover a whole number of ilp_rows>=2 blocks" + ) + return tile_v, ilp_rows, num_groups, stages, evict + + +def _select_tile_v(batch: int) -> int: + """Tile width the per-batch policy would pick (benchmark display).""" + return _select_config(batch)[0] + + +def _div_class(stride_elems: int, ptr: int) -> int: + """Largest supported element divisibility for a dynamic outer stride.""" + if stride_elems % 8 == 0 and ptr % 16 == 0: + return 8 + return 1 + + +def launch_packed_kda_decode_cute( + mixed_qkv: torch.Tensor, + raw_gate: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + state: torch.Tensor, + state_indices: torch.Tensor, + output_view: torch.Tensor, + forced_tile_v: Optional[int] = None, +) -> None: + """Launch the CuTe-DSL packed KDA T=1 kernel on the current stream. + + ``output_view`` must be the ``[B, H, V]`` contiguous view of the caller's + ``[B, 1, H, V]`` output tensor; validation happens in the public facade. + """ + batch = mixed_qkv.shape[0] + hk = _HEADS * _HEAD_DIM + + mixed_stride = mixed_qkv.stride(0) + base = mixed_qkv.storage_offset() + q = mixed_qkv.as_strided( + (batch, _HEADS, _HEAD_DIM), (mixed_stride, _HEAD_DIM, 1), base + ) + k = mixed_qkv.as_strided( + (batch, _HEADS, _HEAD_DIM), (mixed_stride, _HEAD_DIM, 1), base + hk + ) + v = mixed_qkv.as_strided( + (batch, _HEADS, _HEAD_DIM), (mixed_stride, _HEAD_DIM, 1), base + 2 * hk + ) + g = raw_gate.as_strided( + (batch, _HEADS, _HEAD_DIM), + (raw_gate.stride(0), _HEAD_DIM, 1), + raw_gate.storage_offset(), + ) + dtb = dt_bias.view(_HEADS, _HEAD_DIM) + + qkv_div = _div_class(mixed_stride, mixed_qkv.data_ptr()) + gate_div = _div_class(raw_gate.stride(0), raw_gate.data_ptr()) + _launch_from_views( + q, + k, + v, + g, + raw_beta, + A_log, + dtb, + state, + state_indices, + output_view, + forced_tile_v, + qkv_div, + gate_div, + ) + + +def launch_unpacked_kda_decode_cute( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + state: torch.Tensor, + state_indices: torch.Tensor, + output_view: torch.Tensor, + forced_tile_v: Optional[int] = None, + precomputed_gate: bool = False, +) -> None: + """Launch the same T=1 kernel on separately allocated q/k/v/g tensors. + + ``precomputed_gate=True`` selects the pre-computed convention: ``g`` is + the log-space decay (``d = exp(g)``) and ``beta`` is already sigmoided; + ``A_log``/``dt_bias`` are ignored (pass any valid tensors). + + ``q``/``k``/``v``/``g`` are ``[B, H, K]`` bf16 views with contiguous + inner ``[H, K]`` (any row stride); ``beta`` is ``[B, H]`` raw logits. + The compiled kernel takes the five tensors independently, so this is the + identical cubin the packed entry point launches -- the packed layout only + ever existed in the view construction above. + """ + qkv_div = min(_div_class(t.stride(0), t.data_ptr()) for t in (q, k, v)) + gate_div = _div_class(g.stride(0), g.data_ptr()) + _launch_from_views( + q, + k, + v, + g, + beta, + A_log, + dt_bias.view(_HEADS, _HEAD_DIM), + state, + state_indices, + output_view, + forced_tile_v, + qkv_div, + gate_div, + precomputed_gate, + ) + + +def _launch_from_views( + q, + k, + v, + g, + raw_beta, + A_log, + dtb, + state, + state_indices, + output_view, + forced_tile_v, + qkv_div, + gate_div, + precomputed_gate=False, +): + batch = q.shape[0] + tile_v, ilp_rows, num_groups, n_stages, evict = _select_config(batch, forced_tile_v) + pool_div = _div_class(state.stride(0), state.data_ptr()) + + # The cp.async pipeline requires a 16 B-aligned pool and at least two + # chunks per CTA; otherwise fall back to the register-prefetch kernel + # (which handles any alignment via scalar copies). + chunk_rows = ( + int(os.environ.get("FLASHINFER_PACKED_KDA_CHUNKR", "0")) + or _NUM_GROUPS * ilp_rows + ) + # A chunk must cover at least one compute sub-iteration. + chunk_rows = max(chunk_rows, _NUM_GROUPS * ilp_rows) + iters = tile_v // chunk_rows + if pool_div != 8 or iters < 2: + n_stages = 0 + if n_stages and ( + tile_v % chunk_rows != 0 or chunk_rows % (_NUM_GROUPS * ilp_rows) != 0 + ): + raise ValueError( + f"FLASHINFER_PACKED_KDA_CHUNKR={chunk_rows} must divide " + f"tile_v={tile_v} and be a multiple of " + f"{_NUM_GROUPS * ilp_rows} rows" + ) + if n_stages: + # iters + 1 stages lets the prologue put ALL chunks of a short + # pipeline in flight before the q/k/gate staging math. + stages_env = int(os.environ.get("FLASHINFER_PACKED_KDA_STAGES", str(n_stages))) + if stages_env == 1: + raise ValueError( + "FLASHINFER_PACKED_KDA_STAGES=1 is invalid: cp.async refill " + "requires a slot to stay cold for two iterations (use 0 for " + "the register-prefetch kernel, or >=2)" + ) + n_stages = 0 if stages_env <= 0 else min(stages_env, iters + 1) + num_groups = _NUM_GROUPS # smem kernel is fixed at 128 threads + if os.environ.get("FLASHINFER_PACKED_KDA_NO_CPASYNC", "0") == "1": + n_stages = 0 + + persistent = ( + os.environ.get("FLASHINFER_PACKED_KDA_PERSIST", "0") == "1" + and pool_div == 8 + and n_stages >= 2 + ) + if persistent: + # Slot/parity arithmetic requires n_stages to divide the 8 chunks + # per item. + n_stages = 8 if n_stages >= 8 else 4 + # The pipelined kernel has three mutually exclusive staging/wait modes + # (private ring / TMA reads / barriered ring with optional bulk stores); + # normalize the experimental env toggles here so the kernel never sees an + # inconsistent combination (e.g. TMA prologue + private-ring waits). + tma_read = os.environ.get("FLASHINFER_PACKED_KDA_TMA", "0") == "1" and n_stages > 0 + bulk_store = ( + os.environ.get("FLASHINFER_PACKED_KDA_BULK", "0") == "1" + and n_stages > 0 + and chunk_rows == _NUM_GROUPS * ilp_rows + ) + if tma_read and bulk_store: + raise ValueError( + "FLASHINFER_PACKED_KDA_TMA and FLASHINFER_PACKED_KDA_BULK are " + "mutually exclusive" + ) + private_ring = ( + os.environ.get("FLASHINFER_PACKED_KDA_PRIVRING", "1") == "1" + and not tma_read + and not bulk_store + ) + aux_aligned = dtb.data_ptr() % 16 == 0 and output_view.data_ptr() % 16 == 0 + compiled = _get_compiled( + tile_v, + ilp_rows, + qkv_div, + gate_div, + pool_div, + aux_aligned, + _use_packed_fma(), + evict, + int(os.environ.get("FLASHINFER_PACKED_KDA_MAXRREG", "0")), + n_stages, + num_groups, + int(os.environ.get("FLASHINFER_PACKED_KDA_L2HINT", "256")), + chunk_rows, + bulk_store, + ( + int(os.environ.get("FLASHINFER_PACKED_KDA_HPC", "1")) + if n_stages > 0 + and tile_v == _HEAD_DIM + and _HEADS % max(1, int(os.environ.get("FLASHINFER_PACKED_KDA_HPC", "1"))) + == 0 + else 1 + ), + tma_read, + persistent, + int(os.environ.get("FLASHINFER_PACKED_KDA_L2POL", "0")), + private_ring, + int(os.environ.get("FLASHINFER_PACKED_KDA_MINBLOCKS", "0")), + precomputed_gate, + ) + if persistent: + sms = torch.cuda.get_device_properties(q.device).multi_processor_count + # Balanced grid: the largest divisor of the item count that fits the + # resident capacity gives every CTA the same item count (no tail). + items = batch * _HEADS + cap = sms * 6 # 6 blocks/SM at the champion's register footprint + grid_ctas = min(items, cap) + for div in range(cap, cap // 2, -1): + if items % div == 0: + grid_ctas = div + break + compiled( + state, + q, + k, + v, + g, + raw_beta, + A_log, + dtb, + output_view, + state_indices, + grid_ctas, + ) + return + compiled( + state, + q, + k, + v, + g, + raw_beta, + A_log, + dtb, + output_view, + state_indices, + ) + + +def _check_cuda_tensor(name, tensor, dtype): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if 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}") + + +def _check_b200(device: torch.device) -> None: + capability = get_compute_capability(device) + if capability != (10, 0): + raise RuntimeError( + "experimental CuTe packed KDA requires exact compute capability " + f"10.0 (B200), got {capability[0]}.{capability[1]}" + ) + if not is_cuda_version_at_least("12.8"): + raise RuntimeError( + "experimental CuTe packed KDA on compute capability 10.0 requires " + "CUDA 12.8 or newer" + ) + + +@torch.no_grad() +def run_packed_kda_decode_cute( + mixed_qkv: torch.Tensor, + raw_gate: torch.Tensor, + raw_beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + state: torch.Tensor, + state_indices: torch.Tensor, + output: Optional[torch.Tensor] = None, + *, + tile_v: Optional[int] = None, +) -> torch.Tensor: + """Run the B200 CuTe packed KDA T=1 kernel. + + Packed bf16 layouts, fp32 internal math, out-of-pool ``state_indices`` rows are inactive + (zero output, untouched state). ``tile_v`` is a benchmark override that + forces one of the tuned per-tile schedules; production-style calls + should leave it as ``None`` so the per-batch policy picks the kernel. + """ + _check_cuda_tensor("mixed_qkv", mixed_qkv, 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", state, torch.bfloat16) + _check_cuda_tensor("state_indices", state_indices, torch.int32) + + for name, tensor in ( + ("raw_gate", raw_gate), + ("raw_beta", raw_beta), + ("A_log", A_log), + ("dt_bias", dt_bias), + ("state", state), + ("state_indices", state_indices), + ): + if tensor.device != mixed_qkv.device: + raise ValueError(f"{name} must be on the same device as mixed_qkv") + + if mixed_qkv.ndim != 2 or mixed_qkv.shape[1] != _MIXED_WIDTH: + raise ValueError(f"mixed_qkv must have shape [B, {_MIXED_WIDTH}]") + batch = int(mixed_qkv.shape[0]) + if batch <= 0: + raise ValueError(f"packed KDA T=1 batch must be positive, got {batch}") + if mixed_qkv.stride(1) != 1 or mixed_qkv.stride(0) < _MIXED_WIDTH: + raise ValueError("mixed_qkv must have contiguous, non-overlapping rows") + gate_width = _HEADS * _HEAD_DIM + if raw_gate.shape != (batch, gate_width) or raw_gate.stride(1) != 1: + raise ValueError(f"raw_gate must have shape [B, {gate_width}] with stride 1") + if raw_gate.stride(0) < gate_width: + raise ValueError("raw_gate rows must not overlap") + if raw_beta.shape != (batch, _HEADS) or raw_beta.stride(1) != 1: + raise ValueError(f"raw_beta must have shape [B, {_HEADS}] with stride 1") + if raw_beta.stride(0) < _HEADS: + raise ValueError("raw_beta rows must not overlap") + if A_log.shape != (_HEADS,) or not A_log.is_contiguous(): + raise ValueError(f"A_log must be contiguous with shape [{_HEADS}]") + if dt_bias.shape != (gate_width,) or not dt_bias.is_contiguous(): + raise ValueError(f"dt_bias must be contiguous with shape [{gate_width}]") + if state_indices.shape != (batch,) or not state_indices.is_contiguous(): + raise ValueError("state_indices must be contiguous with shape [B]") + if ( + state.ndim != 4 + or state.shape[1:] != (_HEADS, _HEAD_DIM, _HEAD_DIM) + or state.stride(0) < _HEADS * _HEAD_DIM * _HEAD_DIM + or tuple(state.stride()[1:]) != (_HEAD_DIM * _HEAD_DIM, _HEAD_DIM, 1) + ): + raise ValueError( + "state must have shape [N,12,128,128] with compact inner dimensions" + ) + + expected_output_shape = (batch, 1, _HEADS, _HEAD_DIM) + if output is None: + output = mixed_qkv.new_empty(expected_output_shape) + else: + _check_cuda_tensor("output", output, torch.bfloat16) + if output.device != mixed_qkv.device: + raise ValueError("output must be on the same device as mixed_qkv") + if output.shape != expected_output_shape or not output.is_contiguous(): + raise ValueError("output must be contiguous with shape [B,1,12,128]") + + _check_b200(mixed_qkv.device) + if tile_v is not None and tile_v not in _SUPPORTED_TILE_V: + raise ValueError(f"tile_v must be one of {_SUPPORTED_TILE_V}") + + launch_packed_kda_decode_cute( + mixed_qkv, + raw_gate, + raw_beta, + A_log, + dt_bias, + state, + state_indices, + output.view(batch, _HEADS, _HEAD_DIM), + forced_tile_v=tile_v, + ) + return output + + +__all__ = [ + "launch_packed_kda_decode_cute", + "launch_unpacked_kda_decode_cute", + "run_packed_kda_decode_cute", +] diff --git a/flashinfer/kda_kernels/recurrent_kda.py b/flashinfer/kda_kernels/recurrent_kda.py index 7d30a80d93c..067709c84ef 100644 --- a/flashinfer/kda_kernels/recurrent_kda.py +++ b/flashinfer/kda_kernels/recurrent_kda.py @@ -36,6 +36,7 @@ import functools import math +import os from typing import Callable, Literal, Optional, cast import cutlass @@ -55,6 +56,7 @@ get_flash_kda_decode_module, ) from ..jit.cpp_ext import is_cuda_version_at_least +from .packed_kda_decode_cute import launch_unpacked_kda_decode_cute from ..utils import get_compute_capability FlashKDADecodeDeviceArch = Literal["sm100a", "sm103a"] @@ -79,6 +81,132 @@ # one-warp grid parallelism without a head-dimension or benchmark-row table. ONE_WARP_MIN_SEQUENCE_HEADS = 128 +# T=1 fast path: eligible decode calls are routed to the pipelined packed-KDA +# kernel (flashinfer.kda_kernels.packed_kda_decode_cute), which indexes the +# state pool in-kernel instead of gathering/scattering slots on the host. +# Toggle with FLASHINFER_KDA_T1_FAST_PATH=0 (default: enabled). +_T1_FAST_PATH_HEADS = 12 +_T1_FAST_PATH_HEAD_DIM = 128 +_T1_FAST_PATH_LOWER_BOUND = -5.0 + + +@functools.cache +def _t1_fast_path_dummy_params(device_index: int): + """Placeholder A_log/dt_bias for the pre-computed gate convention (the + kernel signature keeps them; the precomputed variant never reads them).""" + device = torch.device("cuda", device_index) + return ( + torch.zeros(_T1_FAST_PATH_HEADS, dtype=torch.float32, device=device), + torch.zeros( + _T1_FAST_PATH_HEADS * _T1_FAST_PATH_HEAD_DIM, + dtype=torch.float32, + device=device, + ), + ) + + +def _t1_fast_path_mode( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + use_gate_in_kernel, + lower_bound, + cu_seqlens, + ssm_state_indices, + num_spec_tokens, + num_accepted_tokens, + output, + initial_state_source, + beta_is_logit, +): + """Return "raw", "precomputed", or None when not eligible.""" + H = _T1_FAST_PATH_HEADS + K = _T1_FAST_PATH_HEAD_DIM + if os.environ.get("FLASHINFER_KDA_T1_FAST_PATH", "1") != "1": + return None + if not use_qk_l2norm_in_kernel or output_final_state: + return None + if ( + use_gate_in_kernel + and beta_is_logit + and lower_bound == _T1_FAST_PATH_LOWER_BOUND + and A_log is not None + and dt_bias is not None + ): + mode = "raw" + elif ( + not use_gate_in_kernel + and not beta_is_logit + and lower_bound is None + and dt_bias is None + ): + # g is the pre-computed log-space decay; beta is pre-sigmoided. + mode = "precomputed" + else: + return None + if ( + cu_seqlens is not None + or num_spec_tokens is not None + or num_accepted_tokens is not None + or initial_state_source is not None + ): + return None + if scale is not None and abs(scale - K**-0.5) > 1e-9: + return None + if q.shape[1] != 1 or q.shape[2] != H or q.shape[3] != K: + return None + if v.shape[2] != H or v.shape[3] != K: + return None + B = q.shape[0] + # The kernel reads each row as one contiguous [H, K] block (any row + # stride), and needs the pool + a 1-D slot map it can index in-kernel. + for t in (q, k, v, g): + if t.stride(3) != 1 or t.stride(2) != K: + return None + if beta.stride(2) != 1: + return None + if ( + initial_state is None + or initial_state.dtype != torch.bfloat16 + or initial_state.ndim != 4 + or initial_state.shape[1:] != (H, K, K) + or tuple(initial_state.stride()[1:]) != (K * K, K, 1) + or initial_state.stride(0) < H * K * K + ): + return None + if ( + ssm_state_indices is None + or ssm_state_indices.ndim != 1 + or ssm_state_indices.dtype != torch.int32 + or ssm_state_indices.shape[0] != B + or not ssm_state_indices.is_contiguous() + ): + return None + if mode == "raw": + if A_log.shape != (H,) or not A_log.is_contiguous(): + return None + if dt_bias.numel() != H * K or not dt_bias.is_contiguous(): + return None + if output is not None and ( + output.dtype != torch.bfloat16 + or output.shape != (B, 1, H, K) + or not output.is_contiguous() + ): + return None + if get_compute_capability(q.device) != (10, 0): + return None + if not is_cuda_version_at_least("12.8"): + return None + return mode + # ============================================================================== # SHARED HELPER FUNCTIONS @@ -1689,6 +1817,58 @@ def run_recurrent_kda( if lower_bound >= 0.0: raise ValueError("lower_bound must be negative") + if ( + T == 1 + and H == HV == _T1_FAST_PATH_HEADS + and K == _T1_FAST_PATH_HEAD_DIM + and backend == "cute-dsl" + and initial_state_indices is None + and ( + fast_path_mode := _t1_fast_path_mode( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + scale, + initial_state, + output_final_state, + use_qk_l2norm_in_kernel, + use_gate_in_kernel, + lower_bound, + cu_seqlens, + ssm_state_indices, + num_spec_tokens, + num_accepted_tokens, + output, + initial_state_source, + beta_is_logit, + ) + ) + is not None + ): + o = output if output is not None else q.new_empty((B, 1, H, V)) + if fast_path_mode == "raw": + fp_a_log, fp_dt_bias = A_log, dt_bias.view(-1) + else: + fp_a_log, fp_dt_bias = _t1_fast_path_dummy_params(q.device.index) + launch_unpacked_kda_decode_cute( + q[:, 0], + k[:, 0], + v[:, 0], + g[:, 0], + beta[:, 0], + fp_a_log, + fp_dt_bias, + initial_state, + ssm_state_indices, + o.view(B, H, V), + precomputed_gate=(fast_path_mode == "precomputed"), + ) + return o, None + if (initial_state_source is None) != (initial_state_indices is None): raise ValueError( "initial_state_source and initial_state_indices must be provided together" diff --git a/tests/kda/test_packed_kda_decode_cute.py b/tests/kda/test_packed_kda_decode_cute.py new file mode 100644 index 00000000000..8283f4b44c4 --- /dev/null +++ b/tests/kda/test_packed_kda_decode_cute.py @@ -0,0 +1,926 @@ +# 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. + +"""Correctness tests for packed-input CuTe KDA T=1 decode.""" + +import pytest +import torch + +pytest.importorskip("cutlass") + +from flashinfer.kda_kernels.packed_kda_decode_cute import _select_tile_v +from flashinfer.kda_kernels.packed_kda_decode_cute import ( + run_packed_kda_decode_cute, +) + + +_HEADS = 12 +_HEAD_DIM = 128 +_MIXED_WIDTH = 3 * _HEADS * _HEAD_DIM +_GATE_WIDTH = _HEADS * _HEAD_DIM +_LOGICAL_STATE_SLOT = _HEADS * _HEAD_DIM * _HEAD_DIM +_PRODUCTION_MIXED_STRIDE = 6144 +_PRODUCTION_STATE_PADDING = 256 +_ATOL = 1.0e-2 +_RTOL = 1.0e-2 +_SCALE = _HEAD_DIM**-0.5 +_EPSILON = 1.0e-6 +_LOWER_BOUND = -5.0 + + +@pytest.fixture +def packed_kda_cute_device(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + device = torch.device("cuda") + if torch.cuda.get_device_capability(device) != (10, 0): + pytest.skip("packed-input CuTe KDA requires exact CC 10.0") + return device + + +def _state_view(storage, slots, slot_stride): + return storage.as_strided( + (slots, _HEADS, _HEAD_DIM, _HEAD_DIM), + (slot_stride, _HEAD_DIM * _HEAD_DIM, _HEAD_DIM, 1), + ) + + +def _make_case( + batch, + device, + *, + seed, + inactive=True, + state_padding=_PRODUCTION_STATE_PADDING, + gate_padding=0, + beta_padding=0, +): + generator = torch.Generator(device=device).manual_seed(seed) + mixed_storage = torch.randn( + (batch, _PRODUCTION_MIXED_STRIDE), + dtype=torch.bfloat16, + device=device, + generator=generator, + ).mul_(0.5) + gate_storage = torch.randn( + (batch, _GATE_WIDTH + gate_padding), + dtype=torch.bfloat16, + device=device, + generator=generator, + ).mul_(0.5) + beta_storage = torch.randn( + (batch, _HEADS + beta_padding), + dtype=torch.bfloat16, + device=device, + generator=generator, + ).mul_(0.5) + + slots = batch + 9 + state_slot_stride = _LOGICAL_STATE_SLOT + state_padding + state_storage = torch.randn( + slots * state_slot_stride, + dtype=torch.bfloat16, + device=device, + generator=generator, + ).mul_(0.05) + state = _state_view(state_storage, slots, state_slot_stride) + + # Shifted slots prove that state_indices, rather than the batch row, owns + # the recurrent state. The final row is graph padding for B > 1. + indices_host = [batch + 2 - row for row in range(batch)] + if inactive and batch > 1: + indices_host[-1] = -1 + state_indices = torch.tensor(indices_host, dtype=torch.int32, device=device) + + return { + "mixed_storage": mixed_storage, + "mixed_qkv": mixed_storage[:, :_MIXED_WIDTH], + "gate_storage": gate_storage, + "raw_gate": gate_storage[:, :_GATE_WIDTH], + "beta_storage": beta_storage, + "raw_beta": beta_storage[:, :_HEADS], + "A_log": torch.randn( + _HEADS, + dtype=torch.float32, + device=device, + generator=generator, + ) + .mul_(0.2) + .sub_(2.0), + "dt_bias": torch.randn( + _GATE_WIDTH, + dtype=torch.float32, + device=device, + generator=generator, + ).mul_(0.25), + "state_storage": state_storage, + "state": state, + "state_slot_stride": state_slot_stride, + "slots": slots, + "indices_host": indices_host, + "state_indices": state_indices, + "output": torch.full( + (batch, 1, _HEADS, _HEAD_DIM), + 123.0, + dtype=torch.bfloat16, + device=device, + ), + } + + +def _call_cute(case, *, state=None, output=None, tile_v=None): + return run_packed_kda_decode_cute( + case["mixed_qkv"], + case["raw_gate"], + case["raw_beta"], + case["A_log"], + case["dt_bias"], + case["state"] if state is None else state, + case["state_indices"], + output=case["output"] if output is None else output, + tile_v=tile_v, + ) + + +def _reference_step( + mixed_qkv, + raw_gate, + raw_beta, + A_log, + dt_bias, + state, + state_indices, + *, + work_dtype=torch.float32, + output_dtype=torch.bfloat16, +): + batch = mixed_qkv.shape[0] + packed = mixed_qkv.to(work_dtype).reshape(batch, 3, _HEADS, _HEAD_DIM) + q_raw = packed[:, 0] + k_raw = packed[:, 1] + q = ( + q_raw + * torch.rsqrt(torch.sum(q_raw * q_raw, dim=-1, keepdim=True) + _EPSILON) + * _SCALE + ) + k = k_raw * torch.rsqrt(torch.sum(k_raw * k_raw, dim=-1, keepdim=True) + _EPSILON) + value = packed[:, 2] + gate_x = raw_gate.to(work_dtype).reshape(batch, _HEADS, _HEAD_DIM) + gate_x = gate_x + dt_bias.to(work_dtype).reshape(_HEADS, _HEAD_DIM) + decay = torch.exp( + _LOWER_BOUND + * torch.sigmoid(torch.exp(A_log.to(work_dtype))[None, :, None] * gate_x) + ) + beta = torch.sigmoid(raw_beta.to(work_dtype)) + + active = (state_indices >= 0) & (state_indices < state.shape[0]) + safe_indices = state_indices.clamp(0, state.shape[0] - 1).to(torch.long) + selected = state.index_select(0, safe_indices).to(work_dtype) + decayed = selected * decay[:, :, None, :] + prediction = torch.einsum("bhvk,bhk->bhv", decayed, k) + delta = (value - prediction) * beta[:, :, None] + updated = decayed + delta[:, :, :, None] * k[:, :, None, :] + projected = torch.einsum("bhvk,bhk->bhv", updated, q) + + active_slots = state_indices[active].to(torch.long) + state.index_copy_(0, active_slots, updated[active].to(state.dtype)) + output = torch.where(active[:, None, None], projected, 0.0).to(output_dtype) + return output.unsqueeze(1) + + +def _clone_padded_state(case): + storage = case["state_storage"].clone() + return storage, _state_view(storage, case["slots"], case["state_slot_stride"]) + + +def _assert_close(actual, expected): + torch.testing.assert_close( + actual, + expected, + atol=_ATOL, + rtol=_RTOL, + check_dtype=False, + ) + + +def _assert_mutation_contract(case, before_storage): + selected = {slot for slot in case["indices_host"] if 0 <= slot < case["slots"]} + before_state = _state_view(before_storage, case["slots"], case["state_slot_stride"]) + state_bits = case["state"].contiguous().view(torch.int16) + before_bits = before_state.contiguous().view(torch.int16) + changed_by_slot = (state_bits != before_bits).reshape(case["slots"], -1).any(dim=1) + for slot, changed in enumerate(changed_by_slot.cpu().tolist()): + unchanged = not changed + if slot in selected: + assert not unchanged, f"selected state slot {slot} was not updated" + else: + assert unchanged, f"unselected state slot {slot} changed" + + storage_rows = case["state_storage"].as_strided( + (case["slots"], case["state_slot_stride"]), + (case["state_slot_stride"], 1), + ) + before_rows = before_storage.as_strided( + (case["slots"], case["state_slot_stride"]), + (case["state_slot_stride"], 1), + ) + assert torch.equal( + storage_rows[:, _LOGICAL_STATE_SLOT:].contiguous().view(torch.int16), + before_rows[:, _LOGICAL_STATE_SLOT:].contiguous().view(torch.int16), + ) + + inactive_rows = [ + row + for row, slot in enumerate(case["indices_host"]) + if slot < 0 or slot >= case["slots"] + ] + if inactive_rows: + inactive = case["output"][inactive_rows] + assert torch.equal( + inactive.contiguous().view(torch.int16), + torch.zeros_like(inactive).view(torch.int16), + ) + + +_BATCH_CASES = [ + 1, + 8, + 16, + 31, + 32, + 64, + 128, + pytest.param(256, marks=pytest.mark.long_running), + pytest.param(512, marks=pytest.mark.long_running), +] + + +def test_packed_kda_cute_tile_selection_cpu(): + assert _select_tile_v(1) == 16 + assert _select_tile_v(11) == 16 + assert _select_tile_v(12) == 8 + assert _select_tile_v(23) == 8 + assert _select_tile_v(24) == 64 + assert _select_tile_v(37) == 64 + assert _select_tile_v(38) == 128 + assert _select_tile_v(512) == 128 + + +@pytest.mark.arch_blackwell +@pytest.mark.parametrize("batch", _BATCH_CASES) +def test_packed_kda_cute_matches_reference_and_preserves_pool( + packed_kda_cute_device, batch +): + case = _make_case(batch, packed_kda_cute_device, seed=20261000 + batch) + before_storage = case["state_storage"].clone() + _, reference_state = _clone_padded_state(case) + reference_output = _reference_step( + case["mixed_qkv"], + case["raw_gate"], + case["raw_beta"], + case["A_log"], + case["dt_bias"], + reference_state, + case["state_indices"], + ) + + result = _call_cute(case) + torch.cuda.synchronize(packed_kda_cute_device) + + assert result is case["output"] + _assert_close(result, reference_output) + _assert_close(case["state"], reference_state) + _assert_mutation_contract(case, before_storage) + + +@pytest.mark.arch_blackwell +@pytest.mark.parametrize("tile_v", [8, 16, 32, 64, 128]) +def test_packed_kda_cute_forced_tiles_match_reference(packed_kda_cute_device, tile_v): + case = _make_case(8, packed_kda_cute_device, seed=20261100 + tile_v) + _, reference_state = _clone_padded_state(case) + reference_output = _reference_step( + case["mixed_qkv"], + case["raw_gate"], + case["raw_beta"], + case["A_log"], + case["dt_bias"], + reference_state, + case["state_indices"], + ) + + result = _call_cute(case, tile_v=tile_v) + torch.cuda.synchronize(packed_kda_cute_device) + + _assert_close(result, reference_output) + _assert_close(case["state"], reference_state) + + +@pytest.mark.arch_blackwell +@pytest.mark.parametrize( + ("batch", "state_padding"), + [ + (8, _PRODUCTION_STATE_PADDING), + (64, _PRODUCTION_STATE_PADDING), + (8, 17), + (64, 17), + ], +) +def test_packed_kda_cute_sanitizer_schedules( + packed_kda_cute_device, batch, state_padding +): + """Named aligned and unaligned tile8/tile16 sanitizer entry points.""" + case = _make_case( + batch, + packed_kda_cute_device, + seed=20261150 + batch + state_padding, + inactive=False, + state_padding=state_padding, + ) + _call_cute(case) + torch.cuda.synchronize(packed_kda_cute_device) + + +def _shifted_contiguous(tensor): + storage = torch.empty( + tensor.numel() + 1, + dtype=tensor.dtype, + device=tensor.device, + ) + shifted = storage[1:].view(tensor.shape) + shifted.copy_(tensor) + assert shifted.is_contiguous() + assert shifted.data_ptr() % 16 != 0 + return shifted + + +@pytest.mark.arch_blackwell +@pytest.mark.parametrize( + "target", + [ + "mixed_qkv", + "raw_gate", + "raw_beta", + "A_log", + "dt_bias", + "state", + "state_indices", + "output", + ], +) +def test_packed_kda_cute_shifted_tensors_match_reference( + packed_kda_cute_device, target +): + case = _make_case(8, packed_kda_cute_device, seed=20261191, inactive=False) + case[target] = _shifted_contiguous(case[target]) + reference_state = case["state"].clone() + reference_output = _reference_step( + case["mixed_qkv"], + case["raw_gate"], + case["raw_beta"], + case["A_log"], + case["dt_bias"], + reference_state, + case["state_indices"], + ) + + result = _call_cute(case, tile_v=8) + torch.cuda.synchronize(packed_kda_cute_device) + + _assert_close(result, reference_output) + _assert_close(case["state"], reference_state) + + +@pytest.mark.arch_blackwell +@pytest.mark.parametrize("tile_v", [8, 64]) +def test_packed_kda_cute_out_of_range_slots_are_inactive( + packed_kda_cute_device, tile_v +): + case = _make_case(8, packed_kda_cute_device, seed=20261192, inactive=False) + case["indices_host"][-2:] = [case["slots"], case["slots"] + 7] + case["state_indices"].copy_( + torch.tensor( + case["indices_host"], dtype=torch.int32, device=packed_kda_cute_device + ) + ) + before_storage = case["state_storage"].clone() + _, reference_state = _clone_padded_state(case) + reference_output = _reference_step( + case["mixed_qkv"], + case["raw_gate"], + case["raw_beta"], + case["A_log"], + case["dt_bias"], + reference_state, + case["state_indices"], + ) + + result = _call_cute(case, tile_v=tile_v) + torch.cuda.synchronize(packed_kda_cute_device) + + _assert_close(result, reference_output) + _assert_close(case["state"], reference_state) + _assert_mutation_contract(case, before_storage) + + +@pytest.mark.arch_blackwell +def test_packed_kda_cute_all_inactive_is_bitwise_noop(packed_kda_cute_device): + case = _make_case(1, packed_kda_cute_device, seed=20261200, inactive=False) + case["indices_host"] = [-1] + case["state_indices"].fill_(-1) + before_storage = case["state_storage"].clone() + + result = _call_cute(case) + torch.cuda.synchronize(packed_kda_cute_device) + + assert torch.equal( + case["state_storage"].view(torch.int16), before_storage.view(torch.int16) + ) + assert torch.equal( + result.view(torch.int16), torch.zeros_like(result).view(torch.int16) + ) + + +@pytest.mark.arch_blackwell +def test_packed_kda_cute_cuda_graph_replay(packed_kda_cute_device): + case = _make_case(8, packed_kda_cute_device, seed=20261300) + initial_storage = case["state_storage"].clone() + + # Materialize the DSL kernel and its persistent cache before capture. + _call_cute(case) + torch.cuda.synchronize(packed_kda_cute_device) + case["state_storage"].copy_(initial_storage) + + graph = torch.cuda.CUDAGraph() + capture_stream = torch.cuda.Stream(device=packed_kda_cute_device) + capture_stream.wait_stream(torch.cuda.current_stream(packed_kda_cute_device)) + with torch.cuda.graph(graph, stream=capture_stream): + captured_result = _call_cute(case) + torch.cuda.synchronize(packed_kda_cute_device) + assert captured_result is case["output"] + + case["state_storage"].copy_(initial_storage) + reference_storage = initial_storage.clone() + reference_state = _state_view( + reference_storage, case["slots"], case["state_slot_stride"] + ) + reference_output = _reference_step( + case["mixed_qkv"], + case["raw_gate"], + case["raw_beta"], + case["A_log"], + case["dt_bias"], + reference_state, + case["state_indices"], + ) + + graph.replay() + torch.cuda.synchronize(packed_kda_cute_device) + _assert_close(case["output"], reference_output) + _assert_close(case["state"], reference_state) + + +@pytest.mark.arch_blackwell +@pytest.mark.parametrize("batch", [8, 64]) +def test_packed_kda_cute_uses_current_stream(packed_kda_cute_device, batch): + case = _make_case(batch, packed_kda_cute_device, seed=20261400 + batch) + _, reference_state = _clone_padded_state(case) + reference_output = _reference_step( + case["mixed_qkv"], + case["raw_gate"], + case["raw_beta"], + case["A_log"], + case["dt_bias"], + reference_state, + case["state_indices"], + ) + + stream = torch.cuda.Stream(device=packed_kda_cute_device) + stream.wait_stream(torch.cuda.current_stream(packed_kda_cute_device)) + with torch.cuda.stream(stream): + result = _call_cute(case) + torch.cuda.current_stream(packed_kda_cute_device).wait_stream(stream) + + assert result is case["output"] + _assert_close(result, reference_output) + _assert_close(case["state"], reference_state) + + +@pytest.mark.arch_blackwell +@pytest.mark.parametrize("batch", [8, 64]) +def test_packed_kda_cute_cuda_graph_replays_changed_inputs_and_indices( + packed_kda_cute_device, batch +): + case = _make_case(batch, packed_kda_cute_device, seed=20261500 + batch) + initial_storage = case["state_storage"].clone() + + # Compile and initialize all lazy runtime state before capture. + _call_cute(case) + torch.cuda.synchronize(packed_kda_cute_device) + case["state_storage"].copy_(initial_storage) + case["output"].fill_(123.0) + + graph = torch.cuda.CUDAGraph() + capture_stream = torch.cuda.Stream(device=packed_kda_cute_device) + capture_stream.wait_stream(torch.cuda.current_stream(packed_kda_cute_device)) + with torch.cuda.graph(graph, stream=capture_stream): + captured_result = _call_cute(case) + torch.cuda.synchronize(packed_kda_cute_device) + assert captured_result is case["output"] + + generator = torch.Generator(device=packed_kda_cute_device).manual_seed( + 20261600 + batch + ) + case["mixed_qkv"].copy_( + torch.randn( + case["mixed_qkv"].shape, + dtype=torch.bfloat16, + device=packed_kda_cute_device, + generator=generator, + ).mul_(0.25) + ) + case["raw_gate"].copy_( + torch.randn( + case["raw_gate"].shape, + dtype=torch.bfloat16, + device=packed_kda_cute_device, + generator=generator, + ).mul_(0.25) + ) + case["raw_beta"].copy_( + torch.randn( + case["raw_beta"].shape, + dtype=torch.bfloat16, + device=packed_kda_cute_device, + generator=generator, + ).mul_(0.25) + ) + changed_indices_host = [row + 4 for row in range(batch)] + changed_indices_host[0] = -1 + case["state_indices"].copy_( + torch.tensor( + changed_indices_host, + dtype=torch.int32, + device=packed_kda_cute_device, + ) + ) + case["indices_host"] = changed_indices_host + case["state_storage"].copy_(initial_storage) + case["output"].fill_(123.0) + + reference_storage = initial_storage.clone() + reference_state = _state_view( + reference_storage, case["slots"], case["state_slot_stride"] + ) + reference_output = _reference_step( + case["mixed_qkv"], + case["raw_gate"], + case["raw_beta"], + case["A_log"], + case["dt_bias"], + reference_state, + case["state_indices"], + ) + + graph.replay() + torch.cuda.synchronize(packed_kda_cute_device) + + _assert_close(case["output"], reference_output) + _assert_close(case["state"], reference_state) + _assert_mutation_contract(case, initial_storage) + + +def _run_packed_kda_512_step_fp64_diagnostic(device): + steps = 512 + batch = 8 + generator = torch.Generator(device=device).manual_seed(20260818) + mixed_storage = torch.randn( + (steps, batch, _PRODUCTION_MIXED_STRIDE), + dtype=torch.bfloat16, + device=device, + generator=generator, + ).mul_(0.5) + gate_storage = torch.randn( + (steps, batch, _GATE_WIDTH + 17), + dtype=torch.bfloat16, + device=device, + generator=generator, + ).mul_(0.5) + beta_storage = torch.randn( + (steps, batch, _HEADS + 5), + dtype=torch.bfloat16, + device=device, + generator=generator, + ).mul_(0.5) + slots = batch + 9 + state_slot_stride = _LOGICAL_STATE_SLOT + 17 + state_storage = torch.randn( + slots * state_slot_stride, + dtype=torch.bfloat16, + device=device, + generator=generator, + ).mul_(0.05) + state = _state_view(state_storage, slots, state_slot_stride) + oracle_state = state.to(torch.float64) + state_indices = torch.tensor( + [3, 4, 5, 6, 7, 8, 9, -1], + dtype=torch.int32, + device=device, + ) + A_log = ( + torch.randn( + _HEADS, + dtype=torch.float32, + device=device, + generator=generator, + ) + .mul_(0.2) + .sub_(2.0) + ) + dt_bias = torch.randn( + _GATE_WIDTH, + dtype=torch.float32, + device=device, + generator=generator, + ).mul_(0.25) + output = torch.empty( + (batch, 1, _HEADS, _HEAD_DIM), + dtype=torch.bfloat16, + device=device, + ) + checkpoints = {1, 2, 4, 8, 16, 32, 64, 128, 256, 512} + + for token in range(steps): + arguments = ( + mixed_storage[token, :, :_MIXED_WIDTH], + gate_storage[token, :, :_GATE_WIDTH], + beta_storage[token, :, :_HEADS], + A_log, + dt_bias, + state, + state_indices, + ) + result = run_packed_kda_decode_cute(*arguments, output=output) + oracle_output = _reference_step( + mixed_storage[token, :, :_MIXED_WIDTH], + gate_storage[token, :, :_GATE_WIDTH], + beta_storage[token, :, :_HEADS], + A_log, + dt_bias, + oracle_state, + state_indices, + work_dtype=torch.float64, + output_dtype=torch.float64, + ) + if token + 1 in checkpoints: + torch.cuda.synchronize(device) + _assert_close(result, oracle_output) + _assert_close(state, oracle_state) + + output_error = (result.to(torch.float64) - oracle_output).abs() + state_error = (state.to(torch.float64) - oracle_state).abs() + assert float(output_error.max()) <= _ATOL + assert float(state_error.max()) <= _ATOL + + +@pytest.mark.arch_blackwell +@pytest.mark.long_running +def test_packed_kda_cute_512_step_fp64_diagnostic(packed_kda_cute_device): + _run_packed_kda_512_step_fp64_diagnostic(packed_kda_cute_device) + + +# --------------------------------------------------------------------------- +# recurrent_kda T=1 fast path (FLASHINFER_KDA_T1_FAST_PATH, default on) +# --------------------------------------------------------------------------- + + +def _recurrent_kda_views(case, batch): + """Unpacked [B,1,H,K] views over the packed case tensors.""" + mixed = case["mixed_qkv"] + width = _HEADS * _HEAD_DIM + q = mixed[:, :width].view(batch, 1, _HEADS, _HEAD_DIM) + k = mixed[:, width : 2 * width].view(batch, 1, _HEADS, _HEAD_DIM) + v = mixed[:, 2 * width :].view(batch, 1, _HEADS, _HEAD_DIM) + g = case["raw_gate"].view(batch, 1, _HEADS, _HEAD_DIM) + beta = case["raw_beta"].view(batch, 1, _HEADS) + return q, k, v, g, beta + + +def _call_recurrent_kda(case, batch, contiguous=False): + from flashinfer import recurrent_kda + + q, k, v, g, beta = _recurrent_kda_views(case, batch) + if contiguous: + # The pre-existing grouped-CTA path silently misreads q/k/v/g views + # whose row stride exceeds the logical row (observed on main); feed + # it compact copies. The fast path reads the strided views directly. + q, k, v, g, beta = (x.contiguous() for x in (q, k, v, g, beta)) + out, _ = recurrent_kda( + q, + k, + v, + g, + beta, + A_log=case["A_log"], + dt_bias=case["dt_bias"], + initial_state=case["state"], + ssm_state_indices=case["state_indices"], + use_gate_in_kernel=True, + lower_bound=-5.0, + beta_is_logit=True, + use_qk_l2norm_in_kernel=True, + ) + return out + + +@pytest.mark.arch_blackwell +@pytest.mark.parametrize("batch", [8, 64]) +def test_recurrent_kda_t1_fast_path_matches_reference( + packed_kda_cute_device, monkeypatch, batch +): + """Eligible recurrent_kda decode calls route to the packed kernel.""" + # inactive=False: the generic path defines no semantics for -1 rows + # (the fast path's -1 handling is covered by the packed-kernel tests). + case = _make_case( + batch, + packed_kda_cute_device, + seed=20261700 + batch, + state_padding=0, + inactive=False, + ) + _, reference_state = _clone_padded_state(case) + reference_output = _reference_step( + case["mixed_qkv"], + case["raw_gate"], + case["raw_beta"], + case["A_log"], + case["dt_bias"], + reference_state, + case["state_indices"], + ) + + monkeypatch.setenv("FLASHINFER_KDA_T1_FAST_PATH", "1") + fast_out = _call_recurrent_kda(case, batch) + torch.cuda.synchronize(packed_kda_cute_device) + _assert_close(fast_out, reference_output) + _assert_close(case["state"], reference_state) + + # The generic path must agree on the same inputs. + case2 = _make_case( + batch, + packed_kda_cute_device, + seed=20261700 + batch, + state_padding=0, + inactive=False, + ) + monkeypatch.setenv("FLASHINFER_KDA_T1_FAST_PATH", "0") + slow_out = _call_recurrent_kda(case2, batch, contiguous=True) + torch.cuda.synchronize(packed_kda_cute_device) + _assert_close(slow_out, reference_output) + _assert_close(case2["state"], reference_state) + + +@pytest.mark.arch_blackwell +def test_recurrent_kda_t1_fast_path_toggle(packed_kda_cute_device, monkeypatch): + """The env toggle switches dispatch: the fast path accepts a padded state + pool that the generic path rejects, which proves which path ran.""" + case = _make_case(8, packed_kda_cute_device, seed=20261800, inactive=False) + + monkeypatch.setenv("FLASHINFER_KDA_T1_FAST_PATH", "1") + _call_recurrent_kda(case, 8) + torch.cuda.synchronize(packed_kda_cute_device) + + monkeypatch.setenv("FLASHINFER_KDA_T1_FAST_PATH", "0") + with pytest.raises(ValueError, match="non-contiguous initial_state"): + _call_recurrent_kda(case, 8) + + +@pytest.mark.arch_blackwell +def test_recurrent_kda_t1_ineligible_calls_fall_back( + packed_kda_cute_device, monkeypatch +): + """A pre-sigmoided-beta call is ineligible and must still work.""" + from flashinfer import recurrent_kda + + batch = 8 + case = _make_case( + batch, + packed_kda_cute_device, + seed=20261900, + state_padding=0, + inactive=False, + ) + _, reference_state = _clone_padded_state(case) + reference_output = _reference_step( + case["mixed_qkv"], + case["raw_gate"], + case["raw_beta"], + case["A_log"], + case["dt_bias"], + reference_state, + case["state_indices"], + ) + + monkeypatch.setenv("FLASHINFER_KDA_T1_FAST_PATH", "1") + q, k, v, g, beta = _recurrent_kda_views(case, batch) + q, k, v, g, beta = (x.contiguous() for x in (q, k, v, g, beta)) + out, _ = recurrent_kda( + q, + k, + v, + g, + torch.sigmoid(beta.float()).to(torch.bfloat16), + A_log=case["A_log"], + dt_bias=case["dt_bias"], + initial_state=case["state"], + ssm_state_indices=case["state_indices"], + use_gate_in_kernel=True, + lower_bound=-5.0, + beta_is_logit=False, + use_qk_l2norm_in_kernel=True, + ) + torch.cuda.synchronize(packed_kda_cute_device) + _assert_close(out, reference_output) + _assert_close(case["state"], reference_state) + + +@pytest.mark.arch_blackwell +@pytest.mark.parametrize("batch", [8, 64]) +def test_recurrent_kda_t1_fast_path_precomputed_gate( + packed_kda_cute_device, monkeypatch, batch +): + """The pre-computed convention (log-space g, sigmoided beta) is also + routed to the fast path and matches an fp32 reference and the generic + path.""" + from flashinfer import recurrent_kda + + device = packed_kda_cute_device + generator = torch.Generator(device=device).manual_seed(20262000 + batch) + + def randn(*shape, dtype=torch.bfloat16): + return torch.randn(*shape, dtype=dtype, device=device, generator=generator) + + q = randn(batch, 1, _HEADS, _HEAD_DIM).mul_(0.5) + k = randn(batch, 1, _HEADS, _HEAD_DIM).mul_(0.5) + v = randn(batch, 1, _HEADS, _HEAD_DIM).mul_(0.5) + g = torch.nn.functional.logsigmoid( + randn(batch, 1, _HEADS, _HEAD_DIM, dtype=torch.float32) + ).to(torch.bfloat16) + beta = torch.sigmoid(randn(batch, 1, _HEADS).float()).to(torch.bfloat16) + slots = batch + 3 + state = randn(slots, _HEADS, _HEAD_DIM, _HEAD_DIM).mul_(0.05) + indices = torch.arange(batch, 0, -1, dtype=torch.int32, device=device) + + # fp32 reference: decay = exp(g), beta used as-is + qf = q.float().squeeze(1) + kf = k.float().squeeze(1) + vf = v.float().squeeze(1) + decay = torch.exp(g.float().squeeze(1)) + qn = qf * torch.rsqrt((qf * qf).sum(-1, keepdim=True) + 1e-6) * _HEAD_DIM**-0.5 + kn = kf * torch.rsqrt((kf * kf).sum(-1, keepdim=True) + 1e-6) + bt = beta.float().squeeze(1) + h = state.float().index_select(0, indices.long()) + hd = h * decay[:, :, None, :] + pred = torch.einsum("bhvk,bhk->bhv", hd, kn) + delta = (vf - pred) * bt[:, :, None] + hn = hd + torch.einsum("bhv,bhk->bhvk", delta, kn) + ref_out = torch.einsum("bhvk,bhk->bhv", hn, qn).unsqueeze(1) + + state_before = state.clone() + + def call(): + out, _ = recurrent_kda( + q, + k, + v, + g, + beta, + initial_state=state, + ssm_state_indices=indices, + use_qk_l2norm_in_kernel=True, + ) + return out + + monkeypatch.setenv("FLASHINFER_KDA_T1_FAST_PATH", "1") + fast_out = call() + torch.cuda.synchronize(device) + _assert_close(fast_out.view(batch, 1, _HEADS, _HEAD_DIM), ref_out) + _assert_close(state.float().index_select(0, indices.long()), hn) + fast_state = state.clone() + + state.copy_(state_before) + monkeypatch.setenv("FLASHINFER_KDA_T1_FAST_PATH", "0") + slow_out = call() + torch.cuda.synchronize(device) + _assert_close(slow_out.view(batch, 1, _HEADS, _HEAD_DIM), ref_out) + _assert_close(state, fast_state)