From dc2171f82fdb95dd6f1454d8aac93b802d9376b8 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 29 Jul 2026 15:28:25 -0700 Subject: [PATCH 1/8] feat(kda): add optimized B200 recurrent prefill backend Dispatch the exact BF16 multi-token recurrent KDA contract to exported SM100a M64/M128 kernels while retaining the existing CuTe backend for decode, speculative decode, GQA, and unsupported shapes. Add graph-safe caller workspaces, JIT/AOT integration, frozen-source integrity checks, correctness coverage, trace updates, documentation, and the six-shape CUPTI benchmark. Signed-off-by: Yingyi Huang --- benchmarks/bench_recurrent_kda_prefill.py | 187 ++ csrc/kda/flashkda_bf16_fused_m128.cu | 2553 ++++++++++++++++ csrc/kda/flashkda_bf16_fused_m128_binding.cu | 93 + csrc/kda/flashkda_bf16_fused_m64.cu | 2564 +++++++++++++++++ csrc/kda/flashkda_bf16_fused_m64_binding.cu | 100 + csrc/kda/flashkda_binding_common.cuh | 442 +++ docs/api/kda_decode.rst | 95 +- flashinfer/__init__.py | 3 + flashinfer/aot.py | 17 + flashinfer/jit/__init__.py | 15 + flashinfer/jit/flash_kda.py | 148 + flashinfer/kda_decode.py | 819 +++++- flashinfer/trace/templates/kda.py | 51 +- tests/jit/test_flash_kda_jit.py | 164 ++ tests/kda/test_recurrent_kda_prefill.py | 864 ++++++ .../recurrent_kda_q8_v16_d128.json | 51 +- 16 files changed, 8139 insertions(+), 27 deletions(-) create mode 100644 benchmarks/bench_recurrent_kda_prefill.py create mode 100644 csrc/kda/flashkda_bf16_fused_m128.cu create mode 100644 csrc/kda/flashkda_bf16_fused_m128_binding.cu create mode 100644 csrc/kda/flashkda_bf16_fused_m64.cu create mode 100644 csrc/kda/flashkda_bf16_fused_m64_binding.cu create mode 100644 csrc/kda/flashkda_binding_common.cuh create mode 100644 flashinfer/jit/flash_kda.py create mode 100644 tests/jit/test_flash_kda_jit.py create mode 100644 tests/kda/test_recurrent_kda_prefill.py diff --git a/benchmarks/bench_recurrent_kda_prefill.py b/benchmarks/bench_recurrent_kda_prefill.py new file mode 100644 index 00000000000..8d0b66e712e --- /dev/null +++ b/benchmarks/bench_recurrent_kda_prefill.py @@ -0,0 +1,187 @@ +# 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. + +"""CUPTI benchmark for the six frozen recurrent-KDA prefill contract shapes. + +The reported public end-to-end time includes the required same-stream final +state copy-back into ``initial_state``. Input allocation, packed metadata, +sequence ordering, output/state allocation, and JIT/cache warmup are outside +the measured region. +""" + +import argparse +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Optional + +import numpy as np +import torch + +from flashinfer.kda_decode import ( + RecurrentKDAPrefillWorkspace, + recurrent_kda, +) +from flashinfer.testing import bench_gpu_time +from flashinfer.utils import get_compute_capability + + +@dataclass(frozen=True) +class Case: + name: str + num_heads: int + seq_lens: tuple[int, ...] + packed: bool + + +CASES = tuple( + Case( + name=f"h{num_heads}_{name}", + num_heads=num_heads, + seq_lens=seq_lens, + packed=packed, + ) + for num_heads in (96, 64) + for name, seq_lens, packed in ( + ("fixed8192", (8192,), False), + ("mixed", (1300, 547, 2048, 963, 271, 3063), True), + ("uniform", (1024,) * 8, True), + ) +) + + +def _make_case( + case: Case, +) -> tuple[ + Callable[[], tuple[torch.Tensor, Optional[torch.Tensor]]], + dict, +]: + total_tokens = sum(case.seq_lens) + shape = (1, total_tokens, case.num_heads, 128) + q = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + k = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + v = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + g = (0.1 * torch.randn(shape, dtype=torch.float32, device="cuda")).to( + torch.bfloat16 + ) + beta = torch.randn( + (1, total_tokens, case.num_heads), + dtype=torch.bfloat16, + device="cuda", + ) + A_log = 0.1 * torch.randn(case.num_heads, dtype=torch.float32, device="cuda") + dt_bias = 0.1 * torch.randn( + (case.num_heads, 128), dtype=torch.float32, device="cuda" + ) + state = torch.zeros( + (len(case.seq_lens), case.num_heads, 128, 128), + dtype=torch.bfloat16, + device="cuda", + ) + output = torch.empty_like(q) + workspace = RecurrentKDAPrefillWorkspace(q.device) + + cu_seqlens = None + seq_order = None + if case.packed: + offsets = [0] + for seq_len in case.seq_lens: + offsets.append(offsets[-1] + seq_len) + cu_seqlens = torch.tensor(offsets, dtype=torch.int64, device="cuda") + seq_order = torch.tensor( + sorted( + range(len(case.seq_lens)), + key=case.seq_lens.__getitem__, + reverse=True, + ), + dtype=torch.int32, + device="cuda", + ) + + def run(): + return recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=state, + output=output, + output_final_state=False, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + lower_bound=-5.0, + cu_seqlens=cu_seqlens, + beta_is_logit=True, + seq_order=seq_order, + prefill_workspace=workspace, + ) + + metadata = { + "name": case.name, + "num_heads": case.num_heads, + "seq_lens": list(case.seq_lens), + "total_tokens": total_tokens, + "layout": "packed" if case.packed else "fixed", + "variant": "m64" if case.name == "h64_fixed8192" else "m128", + } + return run, metadata + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iters", type=int, default=100) + parser.add_argument( + "--json", + type=Path, + help="Optionally write the result list as JSON.", + ) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if get_compute_capability(torch.device("cuda")) != (10, 0): + raise RuntimeError("frozen recurrent-KDA prefill requires B200 (cc 10.0)") + + results = [] + for case in CASES: + run, result = _make_case(case) + run() + torch.cuda.synchronize() + measurements = bench_gpu_time( + run, + enable_cupti=True, + cold_l2_cache=True, + use_cuda_graph=False, + dry_run_iters=args.warmup, + repeat_iters=args.iters, + ) + result["median_ms"] = float(np.median(measurements)) + result["median_us"] = result["median_ms"] * 1000.0 + result["timing_scope"] = "public_end_to_end_with_state_copy_back" + results.append(result) + print( + f"{result['name']:<18} {result['variant']:<4} " + f"{result['median_us']:10.3f} us" + ) + + if args.json is not None: + args.json.write_text(json.dumps(results, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/csrc/kda/flashkda_bf16_fused_m128.cu b/csrc/kda/flashkda_bf16_fused_m128.cu new file mode 100644 index 00000000000..bf2a8618986 --- /dev/null +++ b/csrc/kda/flashkda_bf16_fused_m128.cu @@ -0,0 +1,2553 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +// clang-format off +// Generated by tools/export-generated-programs (device kernel). +// Provenance: loom @ 8437e0515b212e7973b196c4ab680d3d90c1209c. schedule symbol 'flashkda_bf16_fused_m128'; embedded in the host TU as flashkda_bf16_fused_m128_f0217be48b. +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_TMEM_STATE_OFFSET 64 +#define TMEM_TMEM_STATE_INP_OFFSET 0 +#define TMEM_TMEM_U_ACC_OFFSET 224 +#define TMEM_TMEM_U2_INP_OFFSET 224 +#define TMEM_TMEM_U2_ACC_OFFSET 0 +#define TMEM_TMEM_OUT_OFFSET 192 +#define TMEM_TMEM_STATE_OUT_OFFSET 64 +#define NUM_CHUNK_PIPE_STAGES 5 +#define SMEM_SMEM_QD_OFF 1024 +#define SMEM_SMEM_QD_STAGE_BYTES 8192 +#define SMEM_SMEM_QD_STRIDE 41984 +#define SMEM_SMEM_G_RAW_OFF 1024 +#define SMEM_SMEM_G_RAW_STAGE_BYTES 8192 +#define SMEM_SMEM_G_RAW_STRIDE 41984 +#define SMEM_SMEM_G_RAW_ALL_OFF 1024 +#define SMEM_SMEM_G_RAW_ALL_STAGE_BYTES 176128 +#define SMEM_SMEM_G_RAW_ALL_STRIDE 176128 +#define SMEM_SMEM_KD_OFF 9216 +#define SMEM_SMEM_KD_STAGE_BYTES 8192 +#define SMEM_SMEM_KD_STRIDE 41984 +#define SMEM_SMEM_Q_RAW_PREFETCH_OFF 17408 +#define SMEM_SMEM_Q_RAW_PREFETCH_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_RAW_PREFETCH_STRIDE 41984 +#define SMEM_SMEM_FINAL_TRANS_OFF 17408 +#define SMEM_SMEM_FINAL_TRANS_STAGE_BYTES 12288 +#define SMEM_SMEM_FINAL_TRANS_STRIDE 41984 +#define SMEM_SMEM_KR_TRANS_OFF 17408 +#define SMEM_SMEM_KR_TRANS_STAGE_BYTES 8192 +#define SMEM_SMEM_KR_TRANS_STRIDE 41984 +#define SMEM_SMEM_MQK_TRANS_OFF 25600 +#define SMEM_SMEM_MQK_TRANS_STAGE_BYTES 2048 +#define SMEM_SMEM_MQK_TRANS_STRIDE 41984 +#define SMEM_SMEM_INV_OFF 29696 +#define SMEM_SMEM_INV_STAGE_BYTES 2048 +#define SMEM_SMEM_INV_STRIDE 41984 +#define SMEM_SMEM_V_OFF 32384 +#define SMEM_SMEM_V_STAGE_BYTES 8192 +#define SMEM_SMEM_V_STRIDE 41984 +#define SMEM_SMEM_KI_OFF 17408 +#define SMEM_SMEM_KI_STAGE_BYTES 8192 +#define SMEM_SMEM_KI_STRIDE 41984 +#define SMEM_SMEM_GATE_OFF 25600 +#define SMEM_SMEM_GATE_STAGE_BYTES 16384 +#define SMEM_SMEM_GATE_STRIDE 41984 +#define SMEM_SMEM_BETA_RAW_OFF 41984 +#define SMEM_SMEM_BETA_RAW_STAGE_BYTES 512 +#define SMEM_SMEM_BETA_RAW_STRIDE 41984 +#define SMEM_SMEM_INV_WORK_OFF 32384 +#define SMEM_SMEM_INV_WORK_STAGE_BYTES 4096 +#define SMEM_SMEM_INV_WORK_STRIDE 41984 +#define SMEM_SMEM_OUT_OFF 210944 +#define SMEM_SMEM_OUT_STAGE_BYTES 8192 +#define SMEM_SMEM_OUT_STRIDE 8192 +#define SMEM_SMEM_RESTORE_FACTOR_ALL_OFF 41984 +#define SMEM_SMEM_RESTORE_FACTOR_ALL_STAGE_BYTES 168452 +#define SMEM_SMEM_RESTORE_FACTOR_ALL_STRIDE 168452 +#define SMEM_SMEM_GT_PREFIX_ALL_OFF 41472 +#define SMEM_SMEM_GT_PREFIX_ALL_STAGE_BYTES 168448 +#define SMEM_SMEM_GT_PREFIX_ALL_STRIDE 168448 +#define SMEM_SMEM_GT_ALL_OFF 31744 +#define SMEM_SMEM_GT_ALL_STAGE_BYTES 168448 +#define SMEM_SMEM_GT_ALL_STRIDE 168448 +#define SMEM_SMEM_PREP_BETA_ALL_OFF 42500 +#define SMEM_SMEM_PREP_BETA_ALL_STAGE_BYTES 168064 +#define SMEM_SMEM_PREP_BETA_ALL_STRIDE 168064 +#define SMEM_SMEM_GATE_RATE_ALL_OFF 42628 +#define SMEM_SMEM_GATE_RATE_ALL_STAGE_BYTES 167940 +#define SMEM_SMEM_GATE_RATE_ALL_STRIDE 167940 +#define SMEM_SMEM_V_ALL_OFF 32384 +#define SMEM_SMEM_V_ALL_STAGE_BYTES 176128 +#define SMEM_SMEM_V_ALL_STRIDE 176128 +#define SMEM_SMEM_GATE_ALL_OFF 25600 +#define SMEM_SMEM_GATE_ALL_STAGE_BYTES 184320 +#define SMEM_SMEM_GATE_ALL_STRIDE 184320 +#define SMEM_TOTAL 227328 +#define THREADS 1024 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ts_step( + int taddr_out, int taddr_a, int b_lo, uint32_t b_dhi, + uint32_t i_desc, int enable_d) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 dhi;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %5, 0;\n\t" + "mov.b32 dhi, %3;\n\t" + "mov.b64 db, {%2, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%1], db, %4, p;\n\t" + "}\n" + :: "r"(taddr_out), "r"(taddr_a), "r"(b_lo), "r"(b_dhi), + "r"(i_desc), "r"(enable_d)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x32_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x32.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16," + " %17, %18, %19, %20, %21, %22, %23, %24," + " %25, %26, %27, %28, %29, %30, %31, %32};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15]), + "f"(src[16]), "f"(src[17]), "f"(src[18]), "f"(src[19]), + "f"(src[20]), "f"(src[21]), "f"(src[22]), "f"(src[23]), + "f"(src[24]), "f"(src[25]), "f"(src[26]), "f"(src[27]), + "f"(src[28]), "f"(src[29]), "f"(src[30]), "f"(src[31])); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + + +__device__ __forceinline__ void elect_commit2(int mbar_addr0, int mbar_addr1) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%1];\n\t" + "}\n" + :: "r"(mbar_addr0), "r"(mbar_addr1) : "memory"); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_store_4d( + const void *tmap, int x, int y, int z, int w, unsigned smem_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%1, %2, %3, %4}], [%5];" + :: "l"(tmap), "r"(x), "r"(y), "r"(z), "r"(w), "r"(smem_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" + :: "r"(addr), + "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), + "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(1024) void +kernel_flashkda_bf16_fused_m128(__nv_bfloat16* __restrict__ q, const void* __restrict__ q_tma, __nv_bfloat16* __restrict__ k, const void* __restrict__ k_tma, __nv_bfloat16* __restrict__ v, const void* __restrict__ v_tma, __nv_bfloat16* __restrict__ g, const void* __restrict__ g_tma, __nv_bfloat16* __restrict__ beta, const void* __restrict__ beta_tma, float* __restrict__ A_log, float* __restrict__ dt_bias, long long* __restrict__ cu_seqlens, int* __restrict__ seq_order, __nv_bfloat16* __restrict__ initial_state, __nv_bfloat16* __restrict__ out, const void* __restrict__ out_tma, __nv_bfloat16* __restrict__ final_state, int num_heads, int use_initial_state, int store_final_state, float scale, float lower_bound) +{ + // FLASHINFER INTEGRATION BEGIN: acquire global tensor maps + // CUDA kernel-start ordering does not acquire the tensor-map proxy. + // One thread acquires each 128-byte map; the CTA barrier publishes those + // acquires to every thread before any TMA instruction can use a map. + if (threadIdx.x == 0) { + asm volatile( + "fence.proxy.tensormap::generic.acquire.gpu [%0], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%1], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%2], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%3], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%4], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%5], 128;\n" + :: "l"(q_tma), "l"(k_tma), "l"(v_tma), "l"(g_tma), + "l"(beta_tma), "l"(out_tma) + : "memory"); + } + __syncthreads(); + // FLASHINFER INTEGRATION END: acquire global tensor maps + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_qd = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_qd_addr = smem + 1024; + __nv_bfloat16* smem_g_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_g_raw_addr = smem + 1024; + __nv_bfloat16* smem_g_raw_all = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_g_raw_all_addr = smem + 1024; + __nv_bfloat16* smem_kd = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_kd_addr = smem + 9216; + __nv_bfloat16* smem_q_raw_prefetch = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_q_raw_prefetch_addr = smem + 17408; + __nv_bfloat16* smem_final_trans = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_final_trans_addr = smem + 17408; + __nv_bfloat16* smem_kr_trans = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_kr_trans_addr = smem + 17408; + __nv_bfloat16* smem_mqk_trans = reinterpret_cast<__nv_bfloat16*>(smem_raw + 25600); + const int smem_mqk_trans_addr = smem + 25600; + __nv_bfloat16* smem_inv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 29696); + const int smem_inv_addr = smem + 29696; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 32384); + const int smem_v_addr = smem + 32384; + __nv_bfloat16* smem_ki = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_ki_addr = smem + 17408; + float* smem_gate = reinterpret_cast(smem_raw + 25600); + const int smem_gate_addr = smem + 25600; + __nv_bfloat16* smem_beta_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 41984); + const int smem_beta_raw_addr = smem + 41984; + __nv_bfloat16* smem_inv_work = reinterpret_cast<__nv_bfloat16*>(smem_raw + 32384); + const int smem_inv_work_addr = smem + 32384; + __nv_bfloat16* smem_out = reinterpret_cast<__nv_bfloat16*>(smem_raw + 210944); + const int smem_out_addr = smem + 210944; + float* smem_restore_factor_all = reinterpret_cast(smem_raw + 41984); + const int smem_restore_factor_all_addr = smem + 41984; + float* smem_gt_prefix_all = reinterpret_cast(smem_raw + 41472); + const int smem_gt_prefix_all_addr = smem + 41472; + float* smem_gt_all = reinterpret_cast(smem_raw + 31744); + const int smem_gt_all_addr = smem + 31744; + float* smem_prep_beta_all = reinterpret_cast(smem_raw + 42500); + const int smem_prep_beta_all_addr = smem + 42500; + float* smem_gate_rate_all = reinterpret_cast(smem_raw + 42628); + const int smem_gate_rate_all_addr = smem + 42628; + __nv_bfloat16* smem_v_all = reinterpret_cast<__nv_bfloat16*>(smem_raw + 32384); + const int smem_v_all_addr = smem + 32384; + float* smem_gate_all = reinterpret_cast(smem_raw + 25600); + const int smem_gate_all_addr = smem + 25600; + + // Mbarrier init (17 groups, 77 barriers) + // Mbarriers at smem_raw[0..616) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // --- pipeline 'chunk_pipe' --- + // qk_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + mbarrier_init_pred(smem + 8, 1, leader); + mbarrier_init_pred(smem + 16, 1, leader); + mbarrier_init_pred(smem + 24, 1, leader); + mbarrier_init_pred(smem + 32, 1, leader); + // gate_raw_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + mbarrier_init_pred(smem + 48, 1, leader); + mbarrier_init_pred(smem + 56, 1, leader); + mbarrier_init_pred(smem + 64, 1, leader); + mbarrier_init_pred(smem + 72, 1, leader); + // qk_raw_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 80, 1, leader); + mbarrier_init_pred(smem + 88, 1, leader); + mbarrier_init_pred(smem + 96, 1, leader); + mbarrier_init_pred(smem + 104, 1, leader); + mbarrier_init_pred(smem + 112, 1, leader); + // v_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 120, 1, leader); + mbarrier_init_pred(smem + 128, 1, leader); + mbarrier_init_pred(smem + 136, 1, leader); + mbarrier_init_pred(smem + 144, 1, leader); + mbarrier_init_pred(smem + 152, 1, leader); + // v_free: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 160, 4, leader); + mbarrier_init_pred(smem + 168, 4, leader); + mbarrier_init_pred(smem + 176, 4, leader); + mbarrier_init_pred(smem + 184, 4, leader); + mbarrier_init_pred(smem + 192, 4, leader); + // smem_free: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 200, 1, leader); + mbarrier_init_pred(smem + 208, 1, leader); + mbarrier_init_pred(smem + 216, 1, leader); + mbarrier_init_pred(smem + 224, 1, leader); + mbarrier_init_pred(smem + 232, 1, leader); + // raw_inputs_free: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 240, 1, leader); + mbarrier_init_pred(smem + 248, 1, leader); + mbarrier_init_pred(smem + 256, 1, leader); + mbarrier_init_pred(smem + 264, 1, leader); + mbarrier_init_pred(smem + 272, 1, leader); + // state_inp_ready: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 280, 4, leader); + mbarrier_init_pred(smem + 288, 4, leader); + mbarrier_init_pred(smem + 296, 4, leader); + mbarrier_init_pred(smem + 304, 4, leader); + mbarrier_init_pred(smem + 312, 4, leader); + // old_out_ready: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 320, 1, leader); + mbarrier_init_pred(smem + 328, 1, leader); + mbarrier_init_pred(smem + 336, 1, leader); + mbarrier_init_pred(smem + 344, 1, leader); + mbarrier_init_pred(smem + 352, 1, leader); + // u_inp_ready: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 360, 4, leader); + mbarrier_init_pred(smem + 368, 4, leader); + mbarrier_init_pred(smem + 376, 4, leader); + mbarrier_init_pred(smem + 384, 4, leader); + mbarrier_init_pred(smem + 392, 4, leader); + // u2_acc_ready: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 400, 1, leader); + mbarrier_init_pred(smem + 408, 1, leader); + mbarrier_init_pred(smem + 416, 1, leader); + mbarrier_init_pred(smem + 424, 1, leader); + mbarrier_init_pred(smem + 432, 1, leader); + // u2_inp_ready: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 440, 4, leader); + mbarrier_init_pred(smem + 448, 4, leader); + mbarrier_init_pred(smem + 456, 4, leader); + mbarrier_init_pred(smem + 464, 4, leader); + mbarrier_init_pred(smem + 472, 4, leader); + // final_ready: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 480, 1, leader); + mbarrier_init_pred(smem + 488, 1, leader); + mbarrier_init_pred(smem + 496, 1, leader); + mbarrier_init_pred(smem + 504, 1, leader); + mbarrier_init_pred(smem + 512, 1, leader); + // out_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 520, 1, leader); + // tmem_dealloc_ready: 1 barriers, init_count=2 + mbarrier_init_pred(smem + 528, 2, leader); + // prep_diag_ready: 5 barriers, init_count=2 + mbarrier_init_pred(smem + 536, 2, leader); + mbarrier_init_pred(smem + 544, 2, leader); + mbarrier_init_pred(smem + 552, 2, leader); + mbarrier_init_pred(smem + 560, 2, leader); + mbarrier_init_pred(smem + 568, 2, leader); + // prep_inv16_ready: 5 barriers, init_count=2 + mbarrier_init_pred(smem + 576, 2, leader); + mbarrier_init_pred(smem + 584, 2, leader); + mbarrier_init_pred(smem + 592, 2, leader); + mbarrier_init_pred(smem + 600, 2, leader); + mbarrier_init_pred(smem + 608, 2, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 616); + if (warp == 0) { + int _tmem_hold = smem + 616; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define qk_full_addr (mbar_base + 0) + #define gate_raw_full_addr (mbar_base + 40) + #define qk_raw_full_addr (mbar_base + 80) + #define v_full_addr (mbar_base + 120) + #define v_free_addr (mbar_base + 160) + #define smem_free_addr (mbar_base + 200) + #define raw_inputs_free_addr (mbar_base + 240) + #define state_inp_ready_addr (mbar_base + 280) + #define old_out_ready_addr (mbar_base + 320) + #define u_inp_ready_addr (mbar_base + 360) + #define u2_acc_ready_addr (mbar_base + 400) + #define u2_inp_ready_addr (mbar_base + 440) + #define final_ready_addr (mbar_base + 480) + #define out_empty_addr (mbar_base + 520) + #define tmem_dealloc_ready_addr (mbar_base + 528) + #define prep_diag_ready_addr (mbar_base + 536) + #define prep_inv16_ready_addr (mbar_base + 576) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_state = taddr + 64; + const int tmem_tmem_state_inp = taddr; + const int tmem_tmem_u_acc = taddr + 224; + const int tmem_tmem_u2_inp = taddr + 224; + const int tmem_tmem_u2_acc = taddr; + const int tmem_tmem_out = taddr + 192; + const int tmem_tmem_state_out = taddr + 64; + + // ---- Register redistribution for WGs split across roles ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); + } + + // ---- Role: compute ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 168;"); + { // compute_main + int task_idx = blockIdx.x; + int seq_idx = seq_order[task_idx / num_heads]; + int head_idx = task_idx % num_heads; + long long bos = cu_seqlens[seq_idx]; + long long eos = cu_seqlens[seq_idx + 1]; + int seq_len = (int)(eos - bos); + int num_chunks = (seq_len + 32 - 1) / 32; + int warp_in_wg = warp % 4; + const int tmem_row_base = warp_in_wg * 32 << 16; + int state_row = warp_in_wg * 32 + lane; + int warp_id_in_role = (warp - 0); + int compute_local_warp = warp_id_in_role; + long long state_base = (((long long)seq_idx * (long long)num_heads + (long long)head_idx) * 128 + (long long)state_row) * 128; + #pragma unroll + for (int state_col_block = 0; state_col_block < 4; state_col_block++) { + float state_frag[32]; + state_frag[0] = 0.0f; + state_frag[1] = 0.0f; + state_frag[2] = 0.0f; + state_frag[3] = 0.0f; + state_frag[4] = 0.0f; + state_frag[5] = 0.0f; + state_frag[6] = 0.0f; + state_frag[7] = 0.0f; + state_frag[8] = 0.0f; + state_frag[9] = 0.0f; + state_frag[10] = 0.0f; + state_frag[11] = 0.0f; + state_frag[12] = 0.0f; + state_frag[13] = 0.0f; + state_frag[14] = 0.0f; + state_frag[15] = 0.0f; + state_frag[16] = 0.0f; + state_frag[17] = 0.0f; + state_frag[18] = 0.0f; + state_frag[19] = 0.0f; + state_frag[20] = 0.0f; + state_frag[21] = 0.0f; + state_frag[22] = 0.0f; + state_frag[23] = 0.0f; + state_frag[24] = 0.0f; + state_frag[25] = 0.0f; + state_frag[26] = 0.0f; + state_frag[27] = 0.0f; + state_frag[28] = 0.0f; + state_frag[29] = 0.0f; + state_frag[30] = 0.0f; + state_frag[31] = 0.0f; + if (use_initial_state != 0) { + { + const uint4* _vptr_0 = reinterpret_cast(initial_state + state_base + (long long)(state_col_block * 32)); + uint4 _vld_0[2]; + #pragma unroll + for (int _blk = 0; _blk < 2; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_frag[0 + _blk * 8 + _pair * 2])[0]), "=f"((&state_frag[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + } + { + const uint4* _vptr_1 = reinterpret_cast(initial_state + state_base + (long long)(state_col_block * 32) + 16); + uint4 _vld_1[2]; + #pragma unroll + for (int _blk = 0; _blk < 2; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_frag[16 + _blk * 8 + _pair * 2])[0]), "=f"((&state_frag[16 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + } + } + tmem_st_x32_f32(taddr + 64 + (unsigned int)tmem_row_base + (unsigned int)(state_col_block * 32), state_frag); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + unsigned int compute_stage = 0; + unsigned int _phase_qk_full = 0; + unsigned int _phase_v_full = 0; + unsigned int _phase_old_out_ready = 0; + unsigned int _phase_u2_acc_ready = 0; + unsigned int _phase_final_ready = 0; + #pragma unroll 1 + for (int chunk_idx = 0; chunk_idx < num_chunks; chunk_idx++) { + mbarrier_wait(qk_full_addr + (compute_stage) * 8, _phase_qk_full); + #pragma unroll 1 + for (int state_col_block_1 = 0; state_col_block_1 < 4; state_col_block_1++) { + int state_addr = taddr + 64 + (unsigned int)tmem_row_base + (unsigned int)(state_col_block_1 * 32); + float _tmem_load_0[32]; + tmem_ld_x32(&_tmem_load_0[0], state_addr); + uint32_t _tmem_load_0_bf16[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(taddr + (unsigned int)tmem_row_base + (unsigned int)(state_col_block_1 * 16)), "r"(*reinterpret_cast(&_tmem_load_0_bf16[0])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[1])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[2])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[3])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[4])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[5])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[6])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[7])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[8])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[9])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[10])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[11])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[12])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[13])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[14])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[15])) + : "memory"); + float state_scale[16]; + #pragma unroll + for (int state_half = 0; state_half < 2; state_half++) { + #pragma unroll + for (int state_col = 0; state_col < 16; state_col++) { + state_scale[state_col] = smem_gt_all[compute_stage * 10496 + (unsigned int)(state_col_block_1 * 32) + (unsigned int)(state_half * 16) + (unsigned int)state_col]; + } + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast((_tmem_load_0 + state_half * 16))[_ls], reinterpret_cast(state_scale)[_ls]); + } + tmem_st_x32_f32(state_addr, _tmem_load_0); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(state_inp_ready_addr + (compute_stage) * 8); + } + mbarrier_wait(v_full_addr + (compute_stage) * 8, _phase_v_full); + mbarrier_wait(old_out_ready_addr + (compute_stage) * 8, _phase_old_out_ready); + float _tmem_load_1[32]; + tmem_ld_x32(&_tmem_load_1[0], taddr + 224 + (unsigned int)tmem_row_base); + #pragma unroll + for (int residual_half = 0; residual_half < 2; residual_half++) { + float residual_v[16]; + float residual_beta[16]; + #pragma unroll + for (int residual_col = 0; residual_col < 16; residual_col++) { + int token_col = residual_half * 16 + residual_col; + __nv_bfloat16 v_value = smem_v_all[compute_stage * 20992 + (unsigned int)(token_col * 128) + (unsigned int)state_row]; + float _cvt_f32_2 = __bfloat162float(v_value); + residual_v[residual_col] = _cvt_f32_2; + residual_beta[residual_col] = smem_prep_beta_all[compute_stage * 10496 + (unsigned int)token_col]; + } + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + sub_f32x2_inplace(&reinterpret_cast(residual_v)[_ls], reinterpret_cast((_tmem_load_1 + residual_half * 16))[_ls]); + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast(residual_v)[_ls], reinterpret_cast(residual_beta)[_ls]); + uint32_t residual_v_bf16[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(residual_v[_lp*2 + 0], residual_v[_lp*2+1 + 0])); + residual_v_bf16[_lp] = *(uint32_t*)&_bf2; + } + tmem_st_x8_u32(taddr + 224 + (unsigned int)tmem_row_base + (unsigned int)(residual_half * 8), (const uint32_t*)residual_v_bf16); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(v_free_addr + (compute_stage) * 8); + mbarrier_arrive(u_inp_ready_addr + (compute_stage) * 8); + } + mbarrier_wait(u2_acc_ready_addr + (compute_stage) * 8, _phase_u2_acc_ready); + float _tmem_load_2[32]; + tmem_ld_x32(&_tmem_load_2[0], taddr + (unsigned int)tmem_row_base); + uint32_t _tmem_load_2_bf16[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_2[_lp*2 + 0], _tmem_load_2[_lp*2+1 + 0])); + _tmem_load_2_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(taddr + 224 + (unsigned int)tmem_row_base), "r"(*reinterpret_cast(&_tmem_load_2_bf16[0])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[1])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[2])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[3])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[4])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[5])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[6])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[7])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[8])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[9])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[10])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[11])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[12])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[13])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[14])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[15])) + : "memory"); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(u2_inp_ready_addr + (compute_stage) * 8); + } + mbarrier_wait(final_ready_addr + (compute_stage) * 8, _phase_final_ready); + compute_stage += 1; + if (compute_stage == 5) { compute_stage = 0; _phase_qk_full ^= 1; _phase_v_full ^= 1; _phase_old_out_ready ^= 1; _phase_u2_acc_ready ^= 1; _phase_final_ready ^= 1; } + } + if (store_final_state != 0) { + #pragma unroll + for (int state_col_block_2 = 0; state_col_block_2 < 4; state_col_block_2++) { + float _tmem_load_3[32]; + tmem_ld_x32(&_tmem_load_3[0], taddr + 64 + (unsigned int)tmem_row_base + (unsigned int)(state_col_block_2 * 32)); + { + __nv_bfloat162 _pk[8]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_3[0 + 0], _tmem_load_3[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_3[0 + 2], _tmem_load_3[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_3[0 + 4], _tmem_load_3[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_3[0 + 6], _tmem_load_3[0 + 7]); + _pk[4] = __floats2bfloat162_rn(_tmem_load_3[0 + 8], _tmem_load_3[0 + 9]); + _pk[5] = __floats2bfloat162_rn(_tmem_load_3[0 + 10], _tmem_load_3[0 + 11]); + _pk[6] = __floats2bfloat162_rn(_tmem_load_3[0 + 12], _tmem_load_3[0 + 13]); + _pk[7] = __floats2bfloat162_rn(_tmem_load_3[0 + 14], _tmem_load_3[0 + 15]); + *reinterpret_cast(&((__nv_bfloat16*)(final_state + (state_base + (long long)(state_col_block_2 * 32))))[0]) = *reinterpret_cast(&_pk[0]); + *reinterpret_cast(&((__nv_bfloat16*)(final_state + (state_base + (long long)(state_col_block_2 * 32))))[8]) = *reinterpret_cast(&_pk[4]); + } + { + __nv_bfloat162 _pk[8]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_3[16 + 0], _tmem_load_3[16 + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_3[16 + 2], _tmem_load_3[16 + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_3[16 + 4], _tmem_load_3[16 + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_3[16 + 6], _tmem_load_3[16 + 7]); + _pk[4] = __floats2bfloat162_rn(_tmem_load_3[16 + 8], _tmem_load_3[16 + 9]); + _pk[5] = __floats2bfloat162_rn(_tmem_load_3[16 + 10], _tmem_load_3[16 + 11]); + _pk[6] = __floats2bfloat162_rn(_tmem_load_3[16 + 12], _tmem_load_3[16 + 13]); + _pk[7] = __floats2bfloat162_rn(_tmem_load_3[16 + 14], _tmem_load_3[16 + 15]); + *reinterpret_cast(&((__nv_bfloat16*)(final_state + (state_base + (long long)(state_col_block_2 * 32) + 16)))[0]) = *reinterpret_cast(&_pk[0]); + *reinterpret_cast(&((__nv_bfloat16*)(final_state + (state_base + (long long)(state_col_block_2 * 32) + 16)))[8]) = *reinterpret_cast(&_pk[4]); + } + } + } + asm volatile("barrier.sync 10, 128;" ::: "memory"); + if (compute_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(tmem_dealloc_ready_addr); + } + } + } + // ---- Role: epilogue ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); + { // epilogue_main + int task_idx_1 = blockIdx.x; + int seq_idx_1 = seq_order[task_idx_1 / num_heads]; + int head_idx_1 = task_idx_1 % num_heads; + long long bos_1 = cu_seqlens[seq_idx_1]; + long long eos_1 = cu_seqlens[seq_idx_1 + 1]; + int seq_len_1 = (int)(eos_1 - bos_1); + int num_chunks_1 = (seq_len_1 + 32 - 1) / 32; + int warp_id_in_role_1 = (warp - 4); + int epilogue_local_warp = warp_id_in_role_1; + int warp_in_wg_1 = warp % 4; + const int tmem_row_base_1 = warp_in_wg_1 * 32 << 16; + int state_row_1 = warp_in_wg_1 * 32 + lane; + unsigned int epilogue_stage = 0; + unsigned int output_stage = 0; + unsigned int _phase_final_ready_1 = 0; + #pragma unroll 1 + for (int chunk_idx_1 = 0; chunk_idx_1 < num_chunks_1; chunk_idx_1++) { + mbarrier_wait(final_ready_addr + (epilogue_stage) * 8, _phase_final_ready_1); + int chunk_is_full = ((seq_len_1 >= (chunk_idx_1 + 1) * 32) ? 1 : 0); + if (chunk_is_full != 0) { + float _tmem_load_4[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x4.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])), "=r"(*reinterpret_cast(&_tmem_load_4[8])), "=r"(*reinterpret_cast(&_tmem_load_4[9])), "=r"(*reinterpret_cast(&_tmem_load_4[10])), "=r"(*reinterpret_cast(&_tmem_load_4[11])), "=r"(*reinterpret_cast(&_tmem_load_4[12])), "=r"(*reinterpret_cast(&_tmem_load_4[13])), "=r"(*reinterpret_cast(&_tmem_load_4[14])), "=r"(*reinterpret_cast(&_tmem_load_4[15])) + : "r"(taddr + 192 + (unsigned int)tmem_row_base_1) + : "memory"); + float _tmem_load_5[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x4.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=r"(*reinterpret_cast(&_tmem_load_5[0])), "=r"(*reinterpret_cast(&_tmem_load_5[1])), "=r"(*reinterpret_cast(&_tmem_load_5[2])), "=r"(*reinterpret_cast(&_tmem_load_5[3])), "=r"(*reinterpret_cast(&_tmem_load_5[4])), "=r"(*reinterpret_cast(&_tmem_load_5[5])), "=r"(*reinterpret_cast(&_tmem_load_5[6])), "=r"(*reinterpret_cast(&_tmem_load_5[7])), "=r"(*reinterpret_cast(&_tmem_load_5[8])), "=r"(*reinterpret_cast(&_tmem_load_5[9])), "=r"(*reinterpret_cast(&_tmem_load_5[10])), "=r"(*reinterpret_cast(&_tmem_load_5[11])), "=r"(*reinterpret_cast(&_tmem_load_5[12])), "=r"(*reinterpret_cast(&_tmem_load_5[13])), "=r"(*reinterpret_cast(&_tmem_load_5[14])), "=r"(*reinterpret_cast(&_tmem_load_5[15])) + : "r"(taddr + 192 + (unsigned int)tmem_row_base_1 + 1048576) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("barrier.sync 9, 128;" ::: "memory"); + if (epilogue_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(out_empty_addr); + } + } + if (epilogue_local_warp == 0) { + if (chunk_idx_1 >= 2) { + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 9, 128;" ::: "memory"); + int out_stage_addr = smem_out_addr + output_stage * 8192; + #pragma unroll + for (int dim_half = 0; dim_half < 2; dim_half++) { + unsigned int out_packed[8]; + if (dim_half == 0) { + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_4[_lp*2 + 0], _tmem_load_4[_lp*2+1 + 0])); + out_packed[_lp] = *(uint32_t*)&_bf2; + } + } else { + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_5[_lp*2 + 0], _tmem_load_5[_lp*2+1 + 0])); + out_packed[_lp] = *(uint32_t*)&_bf2; + } + } + #pragma unroll + for (int token_group = 0; token_group < 2; token_group++) { + int mtx_idx = lane / 8; + int row_addr = lane & 7; + int dim_base = epilogue_local_warp * 32 + dim_half * 16 + (mtx_idx & 1) * 8; + int token_base = token_group * 16 + mtx_idx / 2 * 8; + int token_addr = token_base + row_addr; + int token_pair = token_addr / 2; + int token_parity = token_addr & 1; + int raw_row = token_pair + dim_base / 64 * 16; + int raw_col = (dim_base & 63 ^ (token_pair & 3) << 4 ^ token_parity << 3) + token_parity * 64; + int stsm_offset = (raw_row * 128 + raw_col) * 2; + const int pack_base = token_group * 4; + uint32_t _stmatrix_addr_0 = static_cast((unsigned long long)(out_stage_addr + stsm_offset)); + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_0), "r"(*reinterpret_cast(&out_packed[pack_base])), "r"(*reinterpret_cast(&out_packed[pack_base + 1])), "r"(*reinterpret_cast(&out_packed[pack_base + 2])), "r"(*reinterpret_cast(&out_packed[pack_base + 3])) + : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 9, 128;" ::: "memory"); + if (epilogue_local_warp == 0) { + if (elect_sync()) { + tma_store_4d(out_tma, 0, (int)(bos_1 + (long long)(chunk_idx_1 * 32)), head_idx_1, 0, smem_out_addr + output_stage * 8192); + } + asm volatile("cp.async.bulk.commit_group;"); + } + output_stage = output_stage ^ 1; + } else { + float _tmem_load_6[32]; + tmem_ld_x32(&_tmem_load_6[0], taddr + 192 + (unsigned int)tmem_row_base_1); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("barrier.sync 9, 128;" ::: "memory"); + if (epilogue_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(out_empty_addr); + } + } + #pragma unroll + for (int token_col_1 = 0; token_col_1 < 32; token_col_1++) { + long long out_token = bos_1 + (long long)(chunk_idx_1 * 32 + token_col_1); + if (out_token < eos_1) { + long long out_idx = (out_token * (long long)num_heads + (long long)head_idx_1) * 128 + (long long)state_row_1; + out[out_idx] = _tmem_load_6[token_col_1]; + } + } + } + epilogue_stage += 1; + if (epilogue_stage == 5) { epilogue_stage = 0; _phase_final_ready_1 ^= 1; } + } + if (epilogue_local_warp == 0) { + asm volatile("cp.async.bulk.wait_group 0;"); + } + asm volatile("barrier.sync 9, 128;" ::: "memory"); + if (epilogue_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(tmem_dealloc_ready_addr); + } + } + } + // ---- Role: mma ---- + } else if (warp == 9) { + { // mma_main + int task_idx_2 = blockIdx.x; + int seq_idx_2 = seq_order[task_idx_2 / num_heads]; + long long bos_2 = cu_seqlens[seq_idx_2]; + long long eos_2 = cu_seqlens[seq_idx_2 + 1]; + int seq_len_2 = (int)(eos_2 - bos_2); + int num_chunks_2 = (seq_len_2 + 32 - 1) / 32; + unsigned int mma_stage = 0; + unsigned int _phase_qk_full_1 = 0; + unsigned int _phase_state_inp_ready = 0; + unsigned int _phase_out_empty_0 = 1; + unsigned int _phase_u_inp_ready = 0; + unsigned int _phase_u2_inp_ready = 0; + #pragma unroll 1 + for (int _chunk_idx = 0; _chunk_idx < num_chunks_2; _chunk_idx++) { + mbarrier_wait(qk_full_addr + (mma_stage) * 8, _phase_qk_full_1); + mbarrier_wait(state_inp_ready_addr + (mma_stage) * 8, _phase_state_inp_ready); + mbarrier_wait(out_empty_addr, _phase_out_empty_0); + _phase_out_empty_0 ^= 1; + int _mma_b_addr_0 = smem_qd_addr + mma_stage * 41984; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 134743184;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 16], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 24], db, id, p1;\n\t" + "add.u32 blo, blo, 250;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 32], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 40], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 48], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 56], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_out), "r"(_mma_b_lo_0), "r"(tmem_tmem_state_inp), "r"(0)); + int _mma_b_addr_1 = smem_kd_addr + mma_stage * 41984; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 134743184;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 16], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 24], db, id, p1;\n\t" + "add.u32 blo, blo, 250;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 32], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 40], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 48], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 56], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_u_acc), "r"(_mma_b_lo_1), "r"(tmem_tmem_state_inp), "r"(0)); + elect_commit2(old_out_ready_addr + (mma_stage) * 8, raw_inputs_free_addr + (mma_stage) * 8); + mbarrier_wait(u_inp_ready_addr + (mma_stage) * 8, _phase_u_inp_ready); + int _mma_b_addr_2 = smem_inv_addr + mma_stage * 41984; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0xC0004010;\n\t" + "mov.b32 id, 134743184;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 64;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_u2_acc), "r"(_mma_b_lo_2), "r"(tmem_tmem_u2_inp), "r"(0)); + elect_commit(u2_acc_ready_addr + (mma_stage) * 8); + mbarrier_wait(u2_inp_ready_addr + (mma_stage) * 8, _phase_u2_inp_ready); + int _mma_b_addr_3 = smem_final_trans_addr + mma_stage * 41984; + int _mma_b_lo_3 = make_warp_uniform(((_mma_b_addr_3 >> 4) & 0x3FFF) | 0x1000000); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 136905872;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_state_out), "r"(_mma_b_lo_3), "r"(tmem_tmem_u2_inp), "r"(1)); + elect_commit2(final_ready_addr + (mma_stage) * 8, smem_free_addr + (mma_stage) * 8); + mma_stage += 1; + if (mma_stage == 5) { mma_stage = 0; _phase_qk_full_1 ^= 1; _phase_state_inp_ready ^= 1; _phase_u_inp_ready ^= 1; _phase_u2_inp_ready ^= 1; } + } + unsigned int _phase_tmem_dealloc_ready_0 = 0; + mbarrier_wait(tmem_dealloc_ready_addr, _phase_tmem_dealloc_ready_0); + _phase_tmem_dealloc_ready_0 ^= 1; + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + // ---- Role: load ---- + } else if (warp == 10) { + { // load_main + int task_idx_3 = blockIdx.x; + int seq_idx_3 = seq_order[task_idx_3 / num_heads]; + int head_idx_2 = task_idx_3 % num_heads; + long long bos_3 = cu_seqlens[seq_idx_3]; + long long eos_3 = cu_seqlens[seq_idx_3 + 1]; + int seq_len_3 = (int)(eos_3 - bos_3); + int num_chunks_3 = (seq_len_3 + 32 - 1) / 32; + unsigned int load_stage = 0; + unsigned int _phase_v_free = 1; + unsigned int _phase_qk_full_2 = 0; + #pragma unroll 1 + for (int chunk_idx_2 = 0; chunk_idx_2 < num_chunks_3; chunk_idx_2++) { + mbarrier_wait(v_free_addr + (load_stage) * 8, _phase_v_free); + mbarrier_wait(qk_full_addr + (load_stage) * 8, _phase_qk_full_2); + int chunk_is_full_1 = ((seq_len_3 >= (chunk_idx_2 + 1) * 32) ? 1 : 0); + if (elect_sync()) { + if (chunk_is_full_1 != 0) { + mbarrier_arrive_expect_tx(v_full_addr + (load_stage) * 8, 8192); + tma_3d_gmem2smem(smem_v_addr + load_stage * 41984, v_tma, 0, head_idx_2, (int)(bos_3 + (long long)(chunk_idx_2 * 32)), v_full_addr + (load_stage) * 8); + } + } + if (chunk_is_full_1 == 0) { + #pragma unroll + for (int v_load_iter = 0; v_load_iter < 16; v_load_iter++) { + int v_item = v_load_iter * 32 + lane; + int row = v_item / 16; + int segment = v_item % 16; + long long token = bos_3 + (long long)(chunk_idx_2 * 32 + row); + int token_valid = ((token < eos_3) ? 1 : 0); + long long v_src = (token * (long long)num_heads + (long long)head_idx_2) * 128 + (long long)(segment * 8); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16, %2;" + :: "r"(smem_v_addr + load_stage * 41984 + (unsigned int)((row * 128 + segment * 8) * 2)), "l"(v + v_src), "r"((token_valid != 0) ? 16 : 0)); + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + } + asm volatile("barrier.sync 8, 32;" ::: "memory"); + if (elect_sync()) { + if (chunk_is_full_1 == 0) { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(v_full_addr + (load_stage) * 8); + } + } + load_stage += 1; + if (load_stage == 5) { load_stage = 0; _phase_v_free ^= 1; _phase_qk_full_2 ^= 1; } + } + } + // ---- Role: prep ---- + } else if (warp >= 12 && warp <= 31) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); + { // prep_main + int task_idx_4 = blockIdx.x; + int seq_idx_4 = seq_order[task_idx_4 / num_heads]; + int head_idx_3 = task_idx_4 % num_heads; + long long bos_4 = cu_seqlens[seq_idx_4]; + long long eos_4 = cu_seqlens[seq_idx_4 + 1]; + int seq_len_4 = (int)(eos_4 - bos_4); + int num_chunks_4 = (seq_len_4 + 32 - 1) / 32; + int instance_id = (warp - 12) / 4; + int prep_instance = instance_id; + int warp_id_in_role_2 = (warp - 12); + int prep_local_warp = warp_id_in_role_2 - prep_instance * 4; + int prep_tid = prep_local_warp * 32 + lane; + int num_prep_iters = (num_chunks_4 + 4 - prep_instance) / 5; + unsigned int prep_stage = (unsigned int)prep_instance; + int gate_rate_stage_f32 = prep_instance * 10496; + if (prep_tid == 0) { + float _expf_0 = __expf(A_log[head_idx_3]); + smem_gate_rate_all[gate_rate_stage_f32] = _expf_0; + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + unsigned int _phase_raw_inputs_free = 1; + unsigned int _phase_gate_raw_full = 0; + unsigned int _phase_smem_free = 1; + unsigned int _phase_qk_raw_full = 0; + unsigned int _phase_prep_diag_ready = 0; + unsigned int _phase_prep_inv16_ready = 0; + #pragma unroll 1 + for (int prep_iter = 0; prep_iter < num_prep_iters; prep_iter++) { + int chunk_idx_3 = prep_iter * 5 + prep_instance; + int stage_f32 = prep_stage * 10496; + int stage_bf16 = prep_stage * 20992; + int chunk_is_full_2 = ((seq_len_4 >= (chunk_idx_3 + 1) * 32) ? 1 : 0); + float early_beta_value = 0.0f; + float early_gate0 = 0.0f; + if (chunk_is_full_2 != 0) { + mbarrier_wait(raw_inputs_free_addr + (prep_stage) * 8, _phase_raw_inputs_free); + if (prep_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(gate_raw_full_addr + (prep_stage) * 8, 8704); + tma_3d_gmem2smem(smem_g_raw_addr + prep_stage * 41984, g_tma, 0, head_idx_3, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), gate_raw_full_addr + (prep_stage) * 8); + tma_2d_gmem2smem(smem_beta_raw_addr + prep_stage * 41984, beta_tma, head_idx_3 / 8 * 8, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), gate_raw_full_addr + (prep_stage) * 8); + mbarrier_arrive_expect_tx(qk_raw_full_addr + (prep_stage) * 8, 16384); + tma_4d_gmem2smem(smem_kd_addr + prep_stage * 41984, k_tma, 0, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), head_idx_3, 0, qk_raw_full_addr + (prep_stage) * 8); + } + } + mbarrier_wait(gate_raw_full_addr + (prep_stage) * 8, _phase_gate_raw_full); + if (prep_local_warp == 2 && lane < 32) { + unsigned int beta_raw_pair[1]; + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&beta_raw_pair[0])) : "r"(smem_beta_raw_addr + prep_stage * 41984 + (unsigned int)(lane * 16) + (unsigned int)(head_idx_3 % 8 / 2 * 4))); + float beta_raw_pair_fp32[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&beta_raw_pair_fp32[_pair * 2])[0]), "=f"((&beta_raw_pair_fp32[_pair * 2])[1]) + : "r"(beta_raw_pair[_pair + 0])); + } + float beta_logit = beta_raw_pair_fp32[0]; + if (head_idx_3 % 2 != 0) { + beta_logit = beta_raw_pair_fp32[1]; + } + float _tanh_approx_0; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_0) : "f"(beta_logit * 0.5f)); + early_beta_value = _tanh_approx_0 * 0.5f + 0.5f; + } + if (prep_tid < 128) { + float early_gate_rate = smem_gate_rate_all[stage_f32]; + float early_gate_bias = dt_bias[head_idx_3 * 128 + prep_tid]; + __nv_bfloat16 early_gate_raw = smem_g_raw_all[stage_bf16 + prep_tid]; + float _cvt_f32_0 = __bfloat162float(early_gate_raw); + float early_gate_arg = early_gate_rate * (_cvt_f32_0 + early_gate_bias); + float _tanh_approx_1; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_1) : "f"(early_gate_arg * 0.5f)); + float early_gate_sigmoid = _tanh_approx_1 * 0.5f + 0.5f; + early_gate0 = lower_bound * 1.4426950408889634f * early_gate_sigmoid; + } + } + mbarrier_wait(smem_free_addr + (prep_stage) * 8, _phase_smem_free); + if (chunk_is_full_2 != 0) { + if (prep_local_warp == 0) { + if (elect_sync()) { + tma_4d_gmem2smem(smem_q_raw_prefetch_addr + prep_stage * 41984, q_tma, 0, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), head_idx_3, 0, qk_raw_full_addr + (prep_stage) * 8); + } + } + } + if (chunk_is_full_2 == 0) { + #pragma unroll + for (int gate_load_pass = 0; gate_load_pass < 4; gate_load_pass++) { + int gate_load_item = gate_load_pass * 128 + prep_tid; + int gate_load_row = gate_load_item / 16; + int gate_load_segment = gate_load_item % 16; + long long gate_load_token = bos_4 + (long long)(chunk_idx_3 * 32 + gate_load_row); + long long gate_load_base = (gate_load_token * (long long)num_heads + (long long)head_idx_3) * 128 + (long long)(gate_load_segment * 8); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16, %2;" + :: "r"(smem_g_raw_addr + prep_stage * 41984 + (unsigned int)(gate_load_item * 16)), "l"(g + gate_load_base), "r"((gate_load_token < eos_4) ? 16 : 0)); + } + } + if (chunk_is_full_2 == 0) { + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + } + if (prep_local_warp == 2 && lane < 32) { + float beta_value = early_beta_value; + if (chunk_is_full_2 == 0) { + long long beta_token = bos_4 + (long long)(chunk_idx_3 * 32 + lane); + if (beta_token < eos_4) { + float beta_logit_1 = (float)beta[beta_token * (long long)num_heads + (long long)head_idx_3]; + float _tanh_approx_2; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_2) : "f"(beta_logit_1 * 0.5f)); + beta_value = _tanh_approx_2 * 0.5f + 0.5f; + } + } + smem_prep_beta_all[stage_f32 + lane] = beta_value; + } + if (prep_tid < 128) { + int gate_col = prep_tid; + float gate_rate = smem_gate_rate_all[stage_f32]; + float gate_bias = dt_bias[head_idx_3 * 128 + gate_col]; + float prefix_log2 = 0.0f; + for (int gate_row = 0; gate_row < 32; gate_row++) { + long long gate_token = bos_4 + (long long)(chunk_idx_3 * 32 + gate_row); + float gate_log2 = 0.0f; + int gate_needs_compute = 1; + if (gate_row == 0) { + if (chunk_is_full_2 != 0) { + gate_log2 = early_gate0; + gate_needs_compute = 0; + } + } + if (gate_needs_compute != 0) { + if (gate_token < eos_4) { + __nv_bfloat16 gate_raw = smem_g_raw_all[stage_bf16 + gate_row * 128 + gate_col]; + float _cvt_f32_1 = __bfloat162float(gate_raw); + float gate_arg = gate_rate * (_cvt_f32_1 + gate_bias); + float _tanh_approx_3; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_3) : "f"(gate_arg * 0.5f)); + float gate_sigmoid = _tanh_approx_3 * 0.5f + 0.5f; + gate_log2 = lower_bound * 1.4426950408889634f * gate_sigmoid; + } + } + prefix_log2 += gate_log2; + smem_gate_all[stage_f32 + gate_row * 128 + gate_col] = prefix_log2; + } + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + if (chunk_is_full_2 != 0) { + mbarrier_wait(qk_raw_full_addr + (prep_stage) * 8, _phase_qk_raw_full); + } + if (prep_tid < 128) { + float total_log2 = smem_gt_prefix_all[stage_f32 + prep_tid]; + float _exp2_0 = approx_exp2(total_log2 - lower_bound * 1.4426950408889634f * 16.0f); + smem_restore_factor_all[stage_f32 + prep_tid] = _exp2_0; + } + if (prep_tid == 0) { + float _exp2_1 = approx_exp2(lower_bound * 1.4426950408889634f * 16.0f); + smem_restore_factor_all[stage_f32 + 128] = _exp2_1; + } + #pragma unroll 1 + for (int work_pass = 0; work_pass < 4; work_pass++) { + int work_item = work_pass * 128 + prep_tid; + int row_1 = work_item / 16; + int segment_1 = work_item % 16; + long long token_1 = bos_4 + (long long)(chunk_idx_3 * 32 + row_1); + int token_valid_1 = ((token_1 < eos_4) ? 1 : 0); + long long gmem_base = (token_1 * (long long)num_heads + (long long)head_idx_3) * 128 + (long long)(segment_1 * 8); + float q_raw_vec[8]; + float k_raw_vec[8]; + q_raw_vec[0] = 0.0f; + q_raw_vec[1] = 0.0f; + q_raw_vec[2] = 0.0f; + q_raw_vec[3] = 0.0f; + q_raw_vec[4] = 0.0f; + q_raw_vec[5] = 0.0f; + q_raw_vec[6] = 0.0f; + q_raw_vec[7] = 0.0f; + k_raw_vec[0] = 0.0f; + k_raw_vec[1] = 0.0f; + k_raw_vec[2] = 0.0f; + k_raw_vec[3] = 0.0f; + k_raw_vec[4] = 0.0f; + k_raw_vec[5] = 0.0f; + k_raw_vec[6] = 0.0f; + k_raw_vec[7] = 0.0f; + if (chunk_is_full_2 != 0) { + unsigned int packed[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed[0])), "=r"(*reinterpret_cast(&packed[(0) + 1])), "=r"(*reinterpret_cast(&packed[(0) + 2])), "=r"(*reinterpret_cast(&packed[(0) + 3])) + : "r"((smem_q_raw_prefetch_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_fp32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32[_pair * 2])[0]), "=f"((&packed_fp32[_pair * 2])[1]) + : "r"(packed[_pair + 0])); + } + #pragma unroll + for (int value_idx = 0; value_idx < 8; value_idx++) { + q_raw_vec[value_idx] = packed_fp32[value_idx]; + } + unsigned int packed_0[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_0[0])), "=r"(*reinterpret_cast(&packed_0[(0) + 1])), "=r"(*reinterpret_cast(&packed_0[(0) + 2])), "=r"(*reinterpret_cast(&packed_0[(0) + 3])) + : "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_0_fp32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_0_fp32[_pair * 2])[0]), "=f"((&packed_0_fp32[_pair * 2])[1]) + : "r"(packed_0[_pair + 0])); + } + #pragma unroll + for (int value_idx_1 = 0; value_idx_1 < 8; value_idx_1++) { + k_raw_vec[value_idx_1] = packed_0_fp32[value_idx_1]; + } + } else if (token_valid_1 != 0) { + { + const uint4* _vptr_0 = reinterpret_cast(q + gmem_base); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&q_raw_vec[0 + _blk * 8 + _pair * 2])[0]), "=f"((&q_raw_vec[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + } + { + const uint4* _vptr_1 = reinterpret_cast(k + gmem_base); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&k_raw_vec[0 + _blk * 8 + _pair * 2])[0]), "=f"((&k_raw_vec[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + } + } + float q_sum = 0.0f; + float k_sum = 0.0f; + for (int elem_in_segment = 0; elem_in_segment < 8; elem_in_segment++) { + float q_raw = q_raw_vec[elem_in_segment]; + float k_raw = k_raw_vec[elem_in_segment]; + float _fma_0 = __fmaf_rn(q_raw, q_raw, q_sum); + q_sum = _fma_0; + float _fma_1 = __fmaf_rn(k_raw, k_raw, k_sum); + k_sum = _fma_1; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 8); + q_sum += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 8); + k_sum += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 4); + q_sum += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 4); + k_sum += _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 2); + q_sum += _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 2); + k_sum += _shfl_xor_5; + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 1); + q_sum += _shfl_xor_6; + float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 1); + k_sum += _shfl_xor_7; + float _rsqrt_0 = rsqrtf(q_sum + 1e-06f); + float q_inv = _rsqrt_0; + float _rsqrt_1 = rsqrtf(k_sum + 1e-06f); + float k_inv = _rsqrt_1; + const float2 _scale2_2 = {q_inv, q_inv}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(q_raw_vec)[_ls], _scale2_2); + const float2 _scale2_3 = {k_inv, k_inv}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(k_raw_vec)[_ls], _scale2_3); + float qd_vec[8]; + float kd_vec[8]; + float ki_vec[8]; + for (int elem_in_segment_1 = 0; elem_in_segment_1 < 8; elem_in_segment_1++) { + int col = segment_1 * 8 + elem_in_segment_1; + float prefix = smem_gate_all[stage_f32 + row_1 * 128 + col]; + float common_log2 = lower_bound * 1.4426950408889634f * 16.0f; + float _exp2_2 = approx_exp2(prefix - common_log2); + float decay = _exp2_2; + qd_vec[elem_in_segment_1] = decay; + kd_vec[elem_in_segment_1] = decay; + ki_vec[elem_in_segment_1] = k_raw_vec[elem_in_segment_1] / decay; + } + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(qd_vec)[_ls], reinterpret_cast(q_raw_vec)[_ls]); + const float2 _scale2_4 = {scale, scale}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(qd_vec)[_ls], _scale2_4); + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(kd_vec)[_ls], reinterpret_cast(k_raw_vec)[_ls]); + unsigned int packed_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(qd_vec[_lp*2 + 0], qd_vec[_lp*2+1 + 0])); + packed_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word = 0; word < 4; word++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word * 4)), "r"(packed_1[word])); + } + unsigned int packed_0_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(kd_vec[_lp*2 + 0], kd_vec[_lp*2+1 + 0])); + packed_0_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_1 = 0; word_1 < 4; word_1++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_1 * 4)), "r"(packed_0_1[word_1])); + } + unsigned int packed_1_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ki_vec[_lp*2 + 0], ki_vec[_lp*2+1 + 0])); + packed_1_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_2 = 0; word_2 < 4; word_2++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_ki_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_2 * 4)), "r"(packed_1_1[word_2])); + } + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + int pair_row_base = prep_local_warp / 2 * 16; + int pair_col_base = prep_local_warp % 2 * 16; + unsigned int a_frag[4]; + unsigned int b_frag[4]; + float acc[8]; + if (pair_row_base >= pair_col_base) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[4]), "=f"(acc[(4) + 1]), "=f"(acc[(4) + 2]), "=f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + int row0 = pair_row_base + lane / 4; + int row1 = row0 + 8; + int col0 = pair_col_base + lane % 4 * 2; + float beta0 = smem_prep_beta_all[stage_f32 + row0]; + float beta1 = smem_prep_beta_all[stage_f32 + row1]; + float seed[8]; + seed[0] = 0.0f; + seed[1] = 0.0f; + seed[2] = 0.0f; + seed[3] = 0.0f; + seed[4] = 0.0f; + seed[5] = 0.0f; + seed[6] = 0.0f; + seed[7] = 0.0f; + if (row0 > col0) { + seed[0] = acc[0] * beta0; + } + if (row0 > col0 + 1) { + seed[1] = acc[1] * beta0; + } + if (row1 > col0) { + seed[2] = acc[2] * beta1; + } + if (row1 > col0 + 1) { + seed[3] = acc[3] * beta1; + } + if (row0 > col0 + 8) { + seed[4] = acc[4] * beta0; + } + if (row0 > col0 + 9) { + seed[5] = acc[5] * beta0; + } + if (row1 > col0 + 8) { + seed[6] = acc[6] * beta1; + } + if (row1 > col0 + 9) { + seed[7] = acc[7] * beta1; + } + unsigned int seed_packed[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(seed[_lp*2 + 0], seed[_lp*2+1 + 0])); + seed_packed[_lp] = *(uint32_t*)&_bf2; + } + int seed_lane_row = lane % 16; + int seed_lane_col = lane / 16 * 8; + int byte_off = (pair_row_base + seed_lane_row) * 128 + (pair_col_base + seed_lane_col) * 2; + int swizzled_off = byte_off ^ (byte_off >> 7 & 7) << 4; + int seed_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off; + uint32_t _stmatrix_addr_5 = static_cast((unsigned long long)seed_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_5), "r"(*reinterpret_cast(&seed_packed[0])), "r"(*reinterpret_cast(&seed_packed[1])), "r"(*reinterpret_cast(&seed_packed[2])), "r"(*reinterpret_cast(&seed_packed[3])) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[4]), "=f"(acc[(4) + 1]), "=f"(acc[(4) + 2]), "=f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + } else { + acc[0] = 0.0f; + acc[1] = 0.0f; + acc[2] = 0.0f; + acc[3] = 0.0f; + acc[4] = 0.0f; + acc[5] = 0.0f; + acc[6] = 0.0f; + acc[7] = 0.0f; + } + int row0_1 = pair_row_base + lane / 4; + int row1_1 = row0_1 + 8; + int col0_1 = pair_col_base + lane % 4 * 2; + float mqk[8]; + mqk[0] = 0.0f; + mqk[1] = 0.0f; + mqk[2] = 0.0f; + mqk[3] = 0.0f; + mqk[4] = 0.0f; + mqk[5] = 0.0f; + mqk[6] = 0.0f; + mqk[7] = 0.0f; + if (row0_1 >= col0_1) { + mqk[0] = acc[0]; + } + if (row0_1 >= col0_1 + 1) { + mqk[1] = acc[1]; + } + if (row1_1 >= col0_1) { + mqk[2] = acc[2]; + } + if (row1_1 >= col0_1 + 1) { + mqk[3] = acc[3]; + } + if (row0_1 >= col0_1 + 8) { + mqk[4] = acc[4]; + } + if (row0_1 >= col0_1 + 9) { + mqk[5] = acc[5]; + } + if (row1_1 >= col0_1 + 8) { + mqk[6] = acc[6]; + } + if (row1_1 >= col0_1 + 9) { + mqk[7] = acc[7]; + } + unsigned int mqk_packed[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(mqk[_lp*2 + 0], mqk[_lp*2+1 + 0])); + mqk_packed[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int publish_pair = 0; publish_pair < 2; publish_pair++) { + int publish_row = pair_col_base + publish_pair * 8 + (lane & 7); + int publish_col = 128 + pair_row_base + lane / 8 * 8; + uint32_t _stmatrix_addr_6 = static_cast((unsigned long long)(smem_final_trans_addr + prep_stage * 41984 + (unsigned int)(publish_col / 64 * 4096 + publish_row * 128 + publish_col % 64 * 2 ^ (publish_col / 64 * 4096 + publish_row * 128 + publish_col % 64 * 2 >> 7 & 7) << 4))); + asm volatile("stmatrix.sync.aligned.m8n8.x2.trans.shared.b16 [%0], {%1, %2};\n" + :: "r"(_stmatrix_addr_6), "r"(*reinterpret_cast(&mqk_packed[publish_pair * 2])), "r"(*reinterpret_cast(&mqk_packed[publish_pair * 2 + 1])) + : "memory"); + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + if (prep_tid < 128) { + float total_log2_1 = smem_gt_prefix_all[stage_f32 + prep_tid]; + float _exp2_3 = approx_exp2(total_log2_1); + smem_gt_all[stage_f32 + prep_tid] = _exp2_3; + } + if (prep_local_warp >= 2) { + int stage_f32_0 = prep_stage * 10496; + float restore_scale = smem_restore_factor_all[stage_f32_0 + 128]; + float restore_factor[8]; + int restore_segment = lane & 15; + #pragma unroll + for (int restore_elem = 0; restore_elem < 8; restore_elem++) { + int restore_col = restore_segment * 8 + restore_elem; + restore_factor[restore_elem] = smem_restore_factor_all[stage_f32_0 + restore_col]; + } + #pragma unroll 1 + for (int restore_pass = 0; restore_pass < 6; restore_pass++) { + int restore_row = 8 + (prep_local_warp - 2) * 12 + restore_pass * 2 + (lane >> 4); + float restore_qd_values[8]; + float restore_kd_values[8]; + float restore_ki_values[8]; + unsigned int packed_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_2[0])), "=r"(*reinterpret_cast(&packed_2[(0) + 1])), "=r"(*reinterpret_cast(&packed_2[(0) + 2])), "=r"(*reinterpret_cast(&packed_2[(0) + 3])) + : "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_fp32_1[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32_1[_pair * 2])[0]), "=f"((&packed_fp32_1[_pair * 2])[1]) + : "r"(packed_2[_pair + 0])); + } + #pragma unroll + for (int value_idx_2 = 0; value_idx_2 < 8; value_idx_2++) { + restore_qd_values[value_idx_2] = packed_fp32_1[value_idx_2]; + } + unsigned int packed_0_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_0_2[0])), "=r"(*reinterpret_cast(&packed_0_2[(0) + 1])), "=r"(*reinterpret_cast(&packed_0_2[(0) + 2])), "=r"(*reinterpret_cast(&packed_0_2[(0) + 3])) + : "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_0_fp32_1[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_0_fp32_1[_pair * 2])[0]), "=f"((&packed_0_fp32_1[_pair * 2])[1]) + : "r"(packed_0_2[_pair + 0])); + } + #pragma unroll + for (int value_idx_3 = 0; value_idx_3 < 8; value_idx_3++) { + restore_kd_values[value_idx_3] = packed_0_fp32_1[value_idx_3]; + } + unsigned int packed_1_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_1_2[0])), "=r"(*reinterpret_cast(&packed_1_2[(0) + 1])), "=r"(*reinterpret_cast(&packed_1_2[(0) + 2])), "=r"(*reinterpret_cast(&packed_1_2[(0) + 3])) + : "r"((smem_ki_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_1_fp32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_1_fp32[_pair * 2])[0]), "=f"((&packed_1_fp32[_pair * 2])[1]) + : "r"(packed_1_2[_pair + 0])); + } + #pragma unroll + for (int value_idx_4 = 0; value_idx_4 < 8; value_idx_4++) { + restore_ki_values[value_idx_4] = packed_1_fp32[value_idx_4]; + } + float restore_kr_values[8]; + #pragma unroll + for (int restore_elem_1 = 0; restore_elem_1 < 8; restore_elem_1++) { + restore_kr_values[restore_elem_1] = restore_ki_values[restore_elem_1] * restore_factor[restore_elem_1]; + } + const float2 _scale2_7 = {restore_scale, restore_scale}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_qd_values)[_ls], _scale2_7); + const float2 _scale2_8 = {restore_scale, restore_scale}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_kd_values)[_ls], _scale2_8); + unsigned int packed_2_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_qd_values[_lp*2 + 0], restore_qd_values[_lp*2+1 + 0])); + packed_2_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_3 = 0; word_3 < 4; word_3++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_3 * 4)), "r"(packed_2_1[word_3])); + } + unsigned int packed_3[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kd_values[_lp*2 + 0], restore_kd_values[_lp*2+1 + 0])); + packed_3[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_4 = 0; word_4 < 4; word_4++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_4 * 4)), "r"(packed_3[word_4])); + } + unsigned int packed_4[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kr_values[_lp*2 + 0], restore_kr_values[_lp*2+1 + 0])); + packed_4[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_5 = 0; word_5 < 4; word_5++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kr_trans_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_5 * 4)), "r"(packed_4[word_5])); + } + } + } + if (prep_local_warp == 0) { + int inverse_row = lane; + int diag_block = inverse_row / 8; + int lane_in_diag = lane & 7; + float inv_row[8]; + unsigned int packed_5[4]; + int byte_off_1 = inverse_row * 128 + diag_block * 8 * 2; + int swizzled_off_1 = byte_off_1 ^ (byte_off_1 >> 7 & 7) << 4; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_5[0])), "=r"(*reinterpret_cast(&packed_5[(0) + 1])), "=r"(*reinterpret_cast(&packed_5[(0) + 2])), "=r"(*reinterpret_cast(&packed_5[(0) + 3])) + : "r"(smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_1)); + float packed_fp32_2[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32_2[_pair * 2])[0]), "=f"((&packed_fp32_2[_pair * 2])[1]) + : "r"(packed_5[_pair + 0])); + } + #pragma unroll + for (int value_idx_5 = 0; value_idx_5 < 8; value_idx_5++) { + inv_row[value_idx_5] = packed_fp32_2[value_idx_5]; + } + #pragma unroll + for (int diag_elem = 0; diag_elem < 8; diag_elem++) { + if (lane_in_diag == diag_elem) { + inv_row[diag_elem] = 1.0f; + } + } + int diag_group_base = lane - lane_in_diag; + #pragma unroll + for (int src_row = 0; src_row < 7; src_row++) { + float row_scale = -inv_row[src_row]; + #pragma unroll + for (int prev_col = 0; prev_col < src_row; prev_col++) { + int pivot_lane = diag_group_base + src_row; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, inv_row[prev_col], pivot_lane); + float pivot = _shfl_0; + if (lane_in_diag > src_row) { + float _fma_2 = __fmaf_rn(row_scale, pivot, inv_row[prev_col]); + inv_row[prev_col] = _fma_2; + } + } + if (lane_in_diag > src_row) { + inv_row[src_row] = row_scale; + } + } + unsigned int packed_0_3[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(inv_row[_lp*2 + 0], inv_row[_lp*2+1 + 0])); + packed_0_3[_lp] = *(uint32_t*)&_bf2; + } + int byte_off_1_1 = inverse_row * 128 + diag_block * 8 * 2; + int swizzled_off_2 = byte_off_1_1 ^ (byte_off_1_1 >> 7 & 7) << 4; + #pragma unroll + for (int word_6 = 0; word_6 < 4; word_6++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"(smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_2 + (unsigned int)(word_6 * 4)), "r"(packed_0_3[word_6])); + } + } + if (prep_local_warp < 2) { + if (elect_sync()) { + mbarrier_arrive(prep_diag_ready_addr + (prep_stage) * 8); + } + mbarrier_wait(prep_diag_ready_addr + (prep_stage) * 8, _phase_prep_diag_ready); + } + if (prep_local_warp < 2) { + int lane_row = lane & 7; + int byte_off_2 = (prep_local_warp * 16 + 8 + lane_row) * 128 + (prep_local_warp * 16 + 8) * 2; + int swizzled_off_3 = byte_off_2 ^ (byte_off_2 >> 7 & 7) << 4; + int d_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_3; + int byte_off_0 = (prep_local_warp * 16 + 8 + lane_row) * 128 + prep_local_warp * 16 * 2; + int swizzled_off_1_1 = byte_off_0 ^ (byte_off_0 >> 7 & 7) << 4; + int c_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_1_1; + int byte_off_2_1 = (prep_local_warp * 16 + lane_row) * 128 + prep_local_warp * 16 * 2; + int swizzled_off_3_1 = byte_off_2_1 ^ (byte_off_2_1 >> 7 & 7) << 4; + int a_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_3_1; + unsigned int d_frag[2]; + unsigned int c_frag[1]; + float dc_acc[4]; + unsigned int dc_bf16[2]; + unsigned int inv_a_frag[1]; + float o_acc[4]; + unsigned int o_bf16[2]; + asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" + : "=r"(d_frag[0]) + : "r"(d_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" + : "=r"(d_frag[1]) + : "r"(d_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x1.trans.shared.b16 {%0}, [%1];\n" + : "=r"(c_frag[0]) + : "r"(c_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(dc_acc[0]), "=f"(dc_acc[1]), "=f"(dc_acc[2]), "=f"(dc_acc[3]) + : "r"(d_frag[0]), "r"(d_frag[1]), "r"(c_frag[0])); + const float2 _scale2_9 = {-1.0f, -1.0f}; + #pragma unroll + for (int _ls = 0; _ls < 2; _ls++) + mul_f32x2_inplace(&reinterpret_cast(dc_acc)[_ls], _scale2_9); + #pragma unroll + for (int _lp = 0; _lp < 2; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dc_acc[_lp*2 + 0], dc_acc[_lp*2+1 + 0])); + dc_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("ldmatrix.sync.aligned.m8n8.x1.trans.shared.b16 {%0}, [%1];\n" + : "=r"(inv_a_frag[0]) + : "r"(a_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(o_acc[0]), "=f"(o_acc[1]), "=f"(o_acc[2]), "=f"(o_acc[3]) + : "r"(dc_bf16[0]), "r"(dc_bf16[1]), "r"(inv_a_frag[0])); + #pragma unroll + for (int _lp = 0; _lp < 2; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(o_acc[_lp*2 + 0], o_acc[_lp*2+1 + 0])); + o_bf16[_lp] = *(uint32_t*)&_bf2; + } + int byte_off_4 = (prep_local_warp * 16 + 8 + lane_row) * 128 + prep_local_warp * 16 * 2; + int swizzled_off_5 = byte_off_4 ^ (byte_off_4 >> 7 & 7) << 4; + int o_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_5; + uint32_t _stmatrix_addr_10 = static_cast((unsigned long long)o_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x1.shared.b16 [%0], {%1};\n" + :: "r"(_stmatrix_addr_10), "r"(*reinterpret_cast(&o_bf16[0])) + : "memory"); + if (elect_sync()) { + mbarrier_arrive(prep_inv16_ready_addr + (prep_stage) * 8); + } + mbarrier_wait(prep_inv16_ready_addr + (prep_stage) * 8, _phase_prep_inv16_ready); + } + if (prep_local_warp == 0) { + int lane_row_1 = lane % 16; + int lane_col = lane / 16 * 8; + int byte_off_3 = (16 + lane_row_1) * 128 + (16 + lane_col) * 2; + int swizzled_off_4 = byte_off_3 ^ (byte_off_3 >> 7 & 7) << 4; + int d_addr_1 = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_4; + int byte_off_0_1 = (16 + lane_row_1) * 128 + lane_col * 2; + int swizzled_off_1_2 = byte_off_0_1 ^ (byte_off_0_1 >> 7 & 7) << 4; + int c_addr_1 = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_1_2; + int byte_off_2_2 = lane_row_1 * 128 + lane_col * 2; + int swizzled_off_3_2 = byte_off_2_2 ^ (byte_off_2_2 >> 7 & 7) << 4; + int a_addr_1 = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_3_2; + unsigned int d32_frag[4]; + unsigned int c32_frag[4]; + float dc32_acc[8]; + unsigned int dc32_bf16[4]; + unsigned int a32_frag[4]; + float o32_acc[8]; + unsigned int o32_bf16[4]; + unsigned int zero32_bf16[4]; + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(d32_frag[0]), "=r"(d32_frag[1]), "=r"(d32_frag[2]), "=r"(d32_frag[3]) + : "r"(d_addr_1) + : "memory"); + int d_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)((16 + lane_col) / 16 * 1024 + (16 + lane_row_1) * 32 + (16 + lane_col) % 16 * 2 ^ ((16 + lane_col) / 16 * 1024 + (16 + lane_row_1) * 32 + (16 + lane_col) % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_11 = static_cast((unsigned long long)d_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_11), "r"(*reinterpret_cast(&d32_frag[0])), "r"(*reinterpret_cast(&d32_frag[1])), "r"(*reinterpret_cast(&d32_frag[2])), "r"(*reinterpret_cast(&d32_frag[3])) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(c32_frag[0]), "=r"(c32_frag[1]), "=r"(c32_frag[2]), "=r"(c32_frag[3]) + : "r"(c_addr_1) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(dc32_acc[0]), "=f"(dc32_acc[1]), "=f"(dc32_acc[2]), "=f"(dc32_acc[3]) + : "r"(d32_frag[0]), "r"(d32_frag[1]), "r"(d32_frag[2]), "r"(d32_frag[3]), "r"(c32_frag[0]), "r"(c32_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(dc32_acc[4]), "=f"(dc32_acc[(4) + 1]), "=f"(dc32_acc[(4) + 2]), "=f"(dc32_acc[(4) + 3]) + : "r"(d32_frag[0]), "r"(d32_frag[1]), "r"(d32_frag[2]), "r"(d32_frag[3]), "r"(c32_frag[2]), "r"(c32_frag[(2) + 1])); + const float2 _scale2_12 = {-1.0f, -1.0f}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(dc32_acc)[_ls], _scale2_12); + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dc32_acc[_lp*2 + 0], dc32_acc[_lp*2+1 + 0])); + dc32_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a32_frag[0]), "=r"(a32_frag[1]), "=r"(a32_frag[2]), "=r"(a32_frag[3]) + : "r"(a_addr_1) + : "memory"); + int a_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)(lane_col / 16 * 1024 + lane_row_1 * 32 + lane_col % 16 * 2 ^ (lane_col / 16 * 1024 + lane_row_1 * 32 + lane_col % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_13 = static_cast((unsigned long long)a_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_13), "r"(*reinterpret_cast(&a32_frag[0])), "r"(*reinterpret_cast(&a32_frag[1])), "r"(*reinterpret_cast(&a32_frag[2])), "r"(*reinterpret_cast(&a32_frag[3])) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(o32_acc[0]), "=f"(o32_acc[1]), "=f"(o32_acc[2]), "=f"(o32_acc[3]) + : "r"(dc32_bf16[0]), "r"(dc32_bf16[1]), "r"(dc32_bf16[2]), "r"(dc32_bf16[3]), "r"(a32_frag[0]), "r"(a32_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(o32_acc[4]), "=f"(o32_acc[(4) + 1]), "=f"(o32_acc[(4) + 2]), "=f"(o32_acc[(4) + 3]) + : "r"(dc32_bf16[0]), "r"(dc32_bf16[1]), "r"(dc32_bf16[2]), "r"(dc32_bf16[3]), "r"(a32_frag[2]), "r"(a32_frag[(2) + 1])); + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(o32_acc[_lp*2 + 0], o32_acc[_lp*2+1 + 0])); + o32_bf16[_lp] = *(uint32_t*)&_bf2; + } + int o_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)(lane_col / 16 * 1024 + (16 + lane_row_1) * 32 + lane_col % 16 * 2 ^ (lane_col / 16 * 1024 + (16 + lane_row_1) * 32 + lane_col % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_14 = static_cast((unsigned long long)o_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_14), "r"(*reinterpret_cast(&o32_bf16[0])), "r"(*reinterpret_cast(&o32_bf16[1])), "r"(*reinterpret_cast(&o32_bf16[2])), "r"(*reinterpret_cast(&o32_bf16[3])) + : "memory"); + #pragma unroll + for (int zero_word = 0; zero_word < 4; zero_word++) { + zero32_bf16[zero_word] = 0; + } + int zero_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)((16 + lane_col) / 16 * 1024 + lane_row_1 * 32 + (16 + lane_col) % 16 * 2 ^ ((16 + lane_col) / 16 * 1024 + lane_row_1 * 32 + (16 + lane_col) % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_15 = static_cast((unsigned long long)zero_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_15), "r"(*reinterpret_cast(&zero32_bf16[0])), "r"(*reinterpret_cast(&zero32_bf16[1])), "r"(*reinterpret_cast(&zero32_bf16[2])), "r"(*reinterpret_cast(&zero32_bf16[3])) + : "memory"); + } else if (prep_local_warp == 1) { + int stage_f32_0_1 = prep_stage * 10496; + float restore_scale_1 = smem_restore_factor_all[stage_f32_0_1 + 128]; + float restore_factor_1[8]; + int restore_segment_1 = lane & 15; + #pragma unroll + for (int restore_elem_2 = 0; restore_elem_2 < 8; restore_elem_2++) { + int restore_col_1 = restore_segment_1 * 8 + restore_elem_2; + restore_factor_1[restore_elem_2] = smem_restore_factor_all[stage_f32_0_1 + restore_col_1]; + } + #pragma unroll 1 + for (int restore_pass_1 = 0; restore_pass_1 < 4; restore_pass_1++) { + int restore_row_1 = restore_pass_1 * 2 + (lane >> 4); + float restore_qd_values_1[8]; + float restore_kd_values_1[8]; + float restore_ki_values_1[8]; + unsigned int packed_6[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_6[0])), "=r"(*reinterpret_cast(&packed_6[(0) + 1])), "=r"(*reinterpret_cast(&packed_6[(0) + 2])), "=r"(*reinterpret_cast(&packed_6[(0) + 3])) + : "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_fp32_3[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32_3[_pair * 2])[0]), "=f"((&packed_fp32_3[_pair * 2])[1]) + : "r"(packed_6[_pair + 0])); + } + #pragma unroll + for (int value_idx_6 = 0; value_idx_6 < 8; value_idx_6++) { + restore_qd_values_1[value_idx_6] = packed_fp32_3[value_idx_6]; + } + unsigned int packed_0_4[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_0_4[0])), "=r"(*reinterpret_cast(&packed_0_4[(0) + 1])), "=r"(*reinterpret_cast(&packed_0_4[(0) + 2])), "=r"(*reinterpret_cast(&packed_0_4[(0) + 3])) + : "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_0_fp32_2[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_0_fp32_2[_pair * 2])[0]), "=f"((&packed_0_fp32_2[_pair * 2])[1]) + : "r"(packed_0_4[_pair + 0])); + } + #pragma unroll + for (int value_idx_7 = 0; value_idx_7 < 8; value_idx_7++) { + restore_kd_values_1[value_idx_7] = packed_0_fp32_2[value_idx_7]; + } + unsigned int packed_1_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_1_3[0])), "=r"(*reinterpret_cast(&packed_1_3[(0) + 1])), "=r"(*reinterpret_cast(&packed_1_3[(0) + 2])), "=r"(*reinterpret_cast(&packed_1_3[(0) + 3])) + : "r"((smem_ki_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_1_fp32_1[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_1_fp32_1[_pair * 2])[0]), "=f"((&packed_1_fp32_1[_pair * 2])[1]) + : "r"(packed_1_3[_pair + 0])); + } + #pragma unroll + for (int value_idx_8 = 0; value_idx_8 < 8; value_idx_8++) { + restore_ki_values_1[value_idx_8] = packed_1_fp32_1[value_idx_8]; + } + float restore_kr_values_1[8]; + #pragma unroll + for (int restore_elem_3 = 0; restore_elem_3 < 8; restore_elem_3++) { + restore_kr_values_1[restore_elem_3] = restore_ki_values_1[restore_elem_3] * restore_factor_1[restore_elem_3]; + } + const float2 _scale2_16 = {restore_scale_1, restore_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_qd_values_1)[_ls], _scale2_16); + const float2 _scale2_17 = {restore_scale_1, restore_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_kd_values_1)[_ls], _scale2_17); + unsigned int packed_2_2[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_qd_values_1[_lp*2 + 0], restore_qd_values_1[_lp*2+1 + 0])); + packed_2_2[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_7 = 0; word_7 < 4; word_7++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_7 * 4)), "r"(packed_2_2[word_7])); + } + unsigned int packed_3_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kd_values_1[_lp*2 + 0], restore_kd_values_1[_lp*2+1 + 0])); + packed_3_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_8 = 0; word_8 < 4; word_8++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_8 * 4)), "r"(packed_3_1[word_8])); + } + unsigned int packed_4_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kr_values_1[_lp*2 + 0], restore_kr_values_1[_lp*2+1 + 0])); + packed_4_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_9 = 0; word_9 < 4; word_9++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kr_trans_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_9 * 4)), "r"(packed_4_1[word_9])); + } + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + if (prep_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(qk_full_addr + (prep_stage) * 8); + } + } + for (int _advance = 0; _advance < 5; _advance++) { + prep_stage += 1; + if (prep_stage == 5) { prep_stage = 0; _phase_raw_inputs_free ^= 1; _phase_gate_raw_full ^= 1; _phase_smem_free ^= 1; _phase_qk_raw_full ^= 1; _phase_prep_diag_ready ^= 1; _phase_prep_inv16_ready ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + +// clang-format on diff --git a/csrc/kda/flashkda_bf16_fused_m128_binding.cu b/csrc/kda/flashkda_bf16_fused_m128_binding.cu new file mode 100644 index 00000000000..8fe63fe35af --- /dev/null +++ b/csrc/kda/flashkda_bf16_fused_m128_binding.cu @@ -0,0 +1,93 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +#include "flashkda_binding_common.cuh" + +// See the M64 binding for why the frozen standalone typedefs are isolated. +#define uint8_t flashkda_generated_uint8_t +#define uint16_t flashkda_generated_uint16_t +#define uint32_t flashkda_generated_uint32_t +#define uint64_t flashkda_generated_uint64_t +#define int32_t flashkda_generated_int32_t +#define int16_t flashkda_generated_int16_t +#include "flashkda_bf16_fused_m128.cu" +#undef uint8_t +#undef uint16_t +#undef uint32_t +#undef uint64_t +#undef int32_t +#undef int16_t + +namespace flashinfer { +namespace flash_kda { + +static_assert(THREADS == 1024); +static_assert(SMEM_TOTAL == 227328); + +void RunM128(TensorView q, TensorView k, TensorView v, TensorView g, TensorView beta, + TensorView beta_tma, TensorView A_log, TensorView dt_bias, TensorView cu_seqlens, + TensorView seq_order, TensorView initial_state, TensorView out, TensorView final_state, + TensorView descriptor_storage, int64_t prepare_descriptors, int64_t num_heads, + int64_t use_initial_state, int64_t store_final_state, double scale, double lower_bound, + int64_t cuda_stream) { + TVM_FFI_ICHECK(cuda_stream >= 0) << "cuda_stream must be a non-negative stream handle"; + TVM_FFI_ICHECK(q.device().device_type == kDLCUDA) << "q must be a CUDA tensor"; + const int32_t device_id = q.device().device_id; + ffi::CUDADeviceGuard device_guard(device_id); + CheckExactSm100a(device_id); + + const int64_t num_seqs = + CheckCommonInputs(q, k, v, g, beta, beta_tma, A_log, dt_bias, cu_seqlens, seq_order, + initial_state, out, final_state, descriptor_storage, prepare_descriptors, + num_heads, use_initial_state, store_final_state, scale, lower_bound); + + constexpr int32_t kSmemBytes = SMEM_TOTAL; + CheckDynamicSmemCapacity(device_id, kSmemBytes); + + CheckCuda(cudaFuncSetAttribute(kernel_flashkda_bf16_fused_m128, + cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes), + "cudaFuncSetAttribute(kernel_flashkda_bf16_fused_m128)"); + + const int64_t grid_x_i64 = num_seqs * num_heads; + TVM_FFI_ICHECK(grid_x_i64 > 0 && grid_x_i64 <= std::numeric_limits::max()) + << "M128 FlashKDA grid.x is out of range: " << grid_x_i64; + const dim3 grid(static_cast(grid_x_i64), 1, 1); + const dim3 block(THREADS, 1, 1); + const cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream)); + const TmaPointers tma = EncodeTmaPointers<128>(q, k, v, g, beta_tma, out, descriptor_storage, + prepare_descriptors, stream); + + kernel_flashkda_bf16_fused_m128<<>>( + reinterpret_cast<__nv_bfloat16*>(q.data_ptr()), tma.q, + reinterpret_cast<__nv_bfloat16*>(k.data_ptr()), tma.k, + reinterpret_cast<__nv_bfloat16*>(v.data_ptr()), tma.v, + reinterpret_cast<__nv_bfloat16*>(g.data_ptr()), tma.g, + reinterpret_cast<__nv_bfloat16*>(beta.data_ptr()), tma.beta, + reinterpret_cast(A_log.data_ptr()), reinterpret_cast(dt_bias.data_ptr()), + reinterpret_cast(cu_seqlens.data_ptr()), + reinterpret_cast(seq_order.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(initial_state.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), tma.out, + reinterpret_cast<__nv_bfloat16*>(final_state.data_ptr()), static_cast(num_heads), + static_cast(use_initial_state), static_cast(store_final_state), + static_cast(scale), static_cast(lower_bound)); + CheckCuda(cudaGetLastError(), "kernel_flashkda_bf16_fused_m128 launch"); +} + +} // namespace flash_kda +} // namespace flashinfer + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, flashinfer::flash_kda::RunM128); diff --git a/csrc/kda/flashkda_bf16_fused_m64.cu b/csrc/kda/flashkda_bf16_fused_m64.cu new file mode 100644 index 00000000000..4849c355b51 --- /dev/null +++ b/csrc/kda/flashkda_bf16_fused_m64.cu @@ -0,0 +1,2564 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +// clang-format off +// Generated by tools/export-generated-programs (device kernel). +// Provenance: loom @ 8437e0515b212e7973b196c4ab680d3d90c1209c. schedule symbol 'flashkda_bf16_fused_m64'; embedded in the host TU as flashkda_bf16_fused_m64_f0217be48b. +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_TMEM_STATE_OFFSET 64 +#define TMEM_TMEM_STATE_INP_OFFSET 0 +#define TMEM_TMEM_U_ACC_OFFSET 224 +#define TMEM_TMEM_U2_INP_OFFSET 224 +#define TMEM_TMEM_U2_ACC_OFFSET 0 +#define TMEM_TMEM_OUT_OFFSET 192 +#define TMEM_TMEM_STATE_OUT_OFFSET 64 +#define NUM_CHUNK_PIPE_STAGES 5 +#define SMEM_SMEM_QD_OFF 1024 +#define SMEM_SMEM_QD_STAGE_BYTES 8192 +#define SMEM_SMEM_QD_STRIDE 41984 +#define SMEM_SMEM_G_RAW_OFF 1024 +#define SMEM_SMEM_G_RAW_STAGE_BYTES 8192 +#define SMEM_SMEM_G_RAW_STRIDE 41984 +#define SMEM_SMEM_G_RAW_ALL_OFF 1024 +#define SMEM_SMEM_G_RAW_ALL_STAGE_BYTES 176128 +#define SMEM_SMEM_G_RAW_ALL_STRIDE 176128 +#define SMEM_SMEM_KD_OFF 9216 +#define SMEM_SMEM_KD_STAGE_BYTES 8192 +#define SMEM_SMEM_KD_STRIDE 41984 +#define SMEM_SMEM_Q_RAW_PREFETCH_OFF 17408 +#define SMEM_SMEM_Q_RAW_PREFETCH_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_RAW_PREFETCH_STRIDE 41984 +#define SMEM_SMEM_FINAL_TRANS_OFF 17408 +#define SMEM_SMEM_FINAL_TRANS_STAGE_BYTES 12288 +#define SMEM_SMEM_FINAL_TRANS_STRIDE 41984 +#define SMEM_SMEM_KR_TRANS_OFF 17408 +#define SMEM_SMEM_KR_TRANS_STAGE_BYTES 8192 +#define SMEM_SMEM_KR_TRANS_STRIDE 41984 +#define SMEM_SMEM_MQK_TRANS_OFF 25600 +#define SMEM_SMEM_MQK_TRANS_STAGE_BYTES 2048 +#define SMEM_SMEM_MQK_TRANS_STRIDE 41984 +#define SMEM_SMEM_INV_OFF 29696 +#define SMEM_SMEM_INV_STAGE_BYTES 2048 +#define SMEM_SMEM_INV_STRIDE 41984 +#define SMEM_SMEM_V_OFF 32384 +#define SMEM_SMEM_V_STAGE_BYTES 4096 +#define SMEM_SMEM_V_STRIDE 41984 +#define SMEM_SMEM_KI_OFF 17408 +#define SMEM_SMEM_KI_STAGE_BYTES 8192 +#define SMEM_SMEM_KI_STRIDE 41984 +#define SMEM_SMEM_GATE_OFF 25600 +#define SMEM_SMEM_GATE_STAGE_BYTES 16384 +#define SMEM_SMEM_GATE_STRIDE 41984 +#define SMEM_SMEM_BETA_RAW_OFF 41984 +#define SMEM_SMEM_BETA_RAW_STAGE_BYTES 512 +#define SMEM_SMEM_BETA_RAW_STRIDE 41984 +#define SMEM_SMEM_INV_WORK_OFF 32384 +#define SMEM_SMEM_INV_WORK_STAGE_BYTES 4096 +#define SMEM_SMEM_INV_WORK_STRIDE 41984 +#define SMEM_SMEM_OUT_OFF 210944 +#define SMEM_SMEM_OUT_STAGE_BYTES 4096 +#define SMEM_SMEM_OUT_STRIDE 4096 +#define SMEM_SMEM_RESTORE_FACTOR_ALL_OFF 41984 +#define SMEM_SMEM_RESTORE_FACTOR_ALL_STAGE_BYTES 168452 +#define SMEM_SMEM_RESTORE_FACTOR_ALL_STRIDE 168452 +#define SMEM_SMEM_GT_PREFIX_ALL_OFF 41472 +#define SMEM_SMEM_GT_PREFIX_ALL_STAGE_BYTES 168448 +#define SMEM_SMEM_GT_PREFIX_ALL_STRIDE 168448 +#define SMEM_SMEM_GT_ALL_OFF 31744 +#define SMEM_SMEM_GT_ALL_STAGE_BYTES 168448 +#define SMEM_SMEM_GT_ALL_STRIDE 168448 +#define SMEM_SMEM_PREP_BETA_ALL_OFF 42500 +#define SMEM_SMEM_PREP_BETA_ALL_STAGE_BYTES 168064 +#define SMEM_SMEM_PREP_BETA_ALL_STRIDE 168064 +#define SMEM_SMEM_GATE_RATE_ALL_OFF 42628 +#define SMEM_SMEM_GATE_RATE_ALL_STAGE_BYTES 167940 +#define SMEM_SMEM_GATE_RATE_ALL_STRIDE 167940 +#define SMEM_SMEM_GATE_ALL_OFF 25600 +#define SMEM_SMEM_GATE_ALL_STAGE_BYTES 184320 +#define SMEM_SMEM_GATE_ALL_STRIDE 184320 +#define SMEM_TOTAL 219136 +#define THREADS 1024 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ts_step( + int taddr_out, int taddr_a, int b_lo, uint32_t b_dhi, + uint32_t i_desc, int enable_d) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 dhi;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %5, 0;\n\t" + "mov.b32 dhi, %3;\n\t" + "mov.b64 db, {%2, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%1], db, %4, p;\n\t" + "}\n" + :: "r"(taddr_out), "r"(taddr_a), "r"(b_lo), "r"(b_dhi), + "r"(i_desc), "r"(enable_d)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + +__device__ __forceinline__ void tmem_ld_x32(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15," + " %16, %17, %18, %19, %20, %21, %22, %23," + " %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_st_x32_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x32.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16," + " %17, %18, %19, %20, %21, %22, %23, %24," + " %25, %26, %27, %28, %29, %30, %31, %32};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15]), + "f"(src[16]), "f"(src[17]), "f"(src[18]), "f"(src[19]), + "f"(src[20]), "f"(src[21]), "f"(src[22]), "f"(src[23]), + "f"(src[24]), "f"(src[25]), "f"(src[26]), "f"(src[27]), + "f"(src[28]), "f"(src[29]), "f"(src[30]), "f"(src[31])); +} + + +__device__ __forceinline__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + + +__device__ __forceinline__ void elect_commit2(int mbar_addr0, int mbar_addr1) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%1];\n\t" + "}\n" + :: "r"(mbar_addr0), "r"(mbar_addr1) : "memory"); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_store_4d( + const void *tmap, int x, int y, int z, int w, unsigned smem_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%1, %2, %3, %4}], [%5];" + :: "l"(tmap), "r"(x), "r"(y), "r"(z), "r"(w), "r"(smem_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" + :: "r"(addr), + "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), + "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(1024) void +kernel_flashkda_bf16_fused_m64(__nv_bfloat16* __restrict__ q, const void* __restrict__ q_tma, __nv_bfloat16* __restrict__ k, const void* __restrict__ k_tma, __nv_bfloat16* __restrict__ v, const void* __restrict__ v_tma, __nv_bfloat16* __restrict__ g, const void* __restrict__ g_tma, __nv_bfloat16* __restrict__ beta, const void* __restrict__ beta_tma, float* __restrict__ A_log, float* __restrict__ dt_bias, long long* __restrict__ cu_seqlens, int* __restrict__ seq_order, __nv_bfloat16* __restrict__ initial_state, __nv_bfloat16* __restrict__ out, const void* __restrict__ out_tma, __nv_bfloat16* __restrict__ final_state, int num_heads, int use_initial_state, int store_final_state, float scale, float lower_bound) +{ + // FLASHINFER INTEGRATION BEGIN: acquire global tensor maps + // CUDA kernel-start ordering does not acquire the tensor-map proxy. + // One thread acquires each 128-byte map; the CTA barrier publishes those + // acquires to every thread before any TMA instruction can use a map. + if (threadIdx.x == 0) { + asm volatile( + "fence.proxy.tensormap::generic.acquire.gpu [%0], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%1], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%2], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%3], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%4], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%5], 128;\n" + :: "l"(q_tma), "l"(k_tma), "l"(v_tma), "l"(g_tma), + "l"(beta_tma), "l"(out_tma) + : "memory"); + } + __syncthreads(); + // FLASHINFER INTEGRATION END: acquire global tensor maps + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_qd = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_qd_addr = smem + 1024; + __nv_bfloat16* smem_g_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_g_raw_addr = smem + 1024; + __nv_bfloat16* smem_g_raw_all = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_g_raw_all_addr = smem + 1024; + __nv_bfloat16* smem_kd = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_kd_addr = smem + 9216; + __nv_bfloat16* smem_q_raw_prefetch = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_q_raw_prefetch_addr = smem + 17408; + __nv_bfloat16* smem_final_trans = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_final_trans_addr = smem + 17408; + __nv_bfloat16* smem_kr_trans = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_kr_trans_addr = smem + 17408; + __nv_bfloat16* smem_mqk_trans = reinterpret_cast<__nv_bfloat16*>(smem_raw + 25600); + const int smem_mqk_trans_addr = smem + 25600; + __nv_bfloat16* smem_inv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 29696); + const int smem_inv_addr = smem + 29696; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 32384); + const int smem_v_addr = smem + 32384; + __nv_bfloat16* smem_ki = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_ki_addr = smem + 17408; + float* smem_gate = reinterpret_cast(smem_raw + 25600); + const int smem_gate_addr = smem + 25600; + __nv_bfloat16* smem_beta_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 41984); + const int smem_beta_raw_addr = smem + 41984; + __nv_bfloat16* smem_inv_work = reinterpret_cast<__nv_bfloat16*>(smem_raw + 32384); + const int smem_inv_work_addr = smem + 32384; + __nv_bfloat16* smem_out = reinterpret_cast<__nv_bfloat16*>(smem_raw + 210944); + const int smem_out_addr = smem + 210944; + float* smem_restore_factor_all = reinterpret_cast(smem_raw + 41984); + const int smem_restore_factor_all_addr = smem + 41984; + float* smem_gt_prefix_all = reinterpret_cast(smem_raw + 41472); + const int smem_gt_prefix_all_addr = smem + 41472; + float* smem_gt_all = reinterpret_cast(smem_raw + 31744); + const int smem_gt_all_addr = smem + 31744; + float* smem_prep_beta_all = reinterpret_cast(smem_raw + 42500); + const int smem_prep_beta_all_addr = smem + 42500; + float* smem_gate_rate_all = reinterpret_cast(smem_raw + 42628); + const int smem_gate_rate_all_addr = smem + 42628; + float* smem_gate_all = reinterpret_cast(smem_raw + 25600); + const int smem_gate_all_addr = smem + 25600; + + // Mbarrier init (17 groups, 77 barriers) + // Mbarriers at smem_raw[0..616) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // --- pipeline 'chunk_pipe' --- + // qk_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + mbarrier_init_pred(smem + 8, 1, leader); + mbarrier_init_pred(smem + 16, 1, leader); + mbarrier_init_pred(smem + 24, 1, leader); + mbarrier_init_pred(smem + 32, 1, leader); + // gate_raw_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + mbarrier_init_pred(smem + 48, 1, leader); + mbarrier_init_pred(smem + 56, 1, leader); + mbarrier_init_pred(smem + 64, 1, leader); + mbarrier_init_pred(smem + 72, 1, leader); + // qk_raw_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 80, 1, leader); + mbarrier_init_pred(smem + 88, 1, leader); + mbarrier_init_pred(smem + 96, 1, leader); + mbarrier_init_pred(smem + 104, 1, leader); + mbarrier_init_pred(smem + 112, 1, leader); + // v_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 120, 1, leader); + mbarrier_init_pred(smem + 128, 1, leader); + mbarrier_init_pred(smem + 136, 1, leader); + mbarrier_init_pred(smem + 144, 1, leader); + mbarrier_init_pred(smem + 152, 1, leader); + // v_free: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 160, 4, leader); + mbarrier_init_pred(smem + 168, 4, leader); + mbarrier_init_pred(smem + 176, 4, leader); + mbarrier_init_pred(smem + 184, 4, leader); + mbarrier_init_pred(smem + 192, 4, leader); + // smem_free: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 200, 1, leader); + mbarrier_init_pred(smem + 208, 1, leader); + mbarrier_init_pred(smem + 216, 1, leader); + mbarrier_init_pred(smem + 224, 1, leader); + mbarrier_init_pred(smem + 232, 1, leader); + // raw_inputs_free: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 240, 1, leader); + mbarrier_init_pred(smem + 248, 1, leader); + mbarrier_init_pred(smem + 256, 1, leader); + mbarrier_init_pred(smem + 264, 1, leader); + mbarrier_init_pred(smem + 272, 1, leader); + // state_inp_ready: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 280, 4, leader); + mbarrier_init_pred(smem + 288, 4, leader); + mbarrier_init_pred(smem + 296, 4, leader); + mbarrier_init_pred(smem + 304, 4, leader); + mbarrier_init_pred(smem + 312, 4, leader); + // old_out_ready: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 320, 1, leader); + mbarrier_init_pred(smem + 328, 1, leader); + mbarrier_init_pred(smem + 336, 1, leader); + mbarrier_init_pred(smem + 344, 1, leader); + mbarrier_init_pred(smem + 352, 1, leader); + // u_inp_ready: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 360, 4, leader); + mbarrier_init_pred(smem + 368, 4, leader); + mbarrier_init_pred(smem + 376, 4, leader); + mbarrier_init_pred(smem + 384, 4, leader); + mbarrier_init_pred(smem + 392, 4, leader); + // u2_acc_ready: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 400, 1, leader); + mbarrier_init_pred(smem + 408, 1, leader); + mbarrier_init_pred(smem + 416, 1, leader); + mbarrier_init_pred(smem + 424, 1, leader); + mbarrier_init_pred(smem + 432, 1, leader); + // u2_inp_ready: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 440, 4, leader); + mbarrier_init_pred(smem + 448, 4, leader); + mbarrier_init_pred(smem + 456, 4, leader); + mbarrier_init_pred(smem + 464, 4, leader); + mbarrier_init_pred(smem + 472, 4, leader); + // final_ready: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 480, 1, leader); + mbarrier_init_pred(smem + 488, 1, leader); + mbarrier_init_pred(smem + 496, 1, leader); + mbarrier_init_pred(smem + 504, 1, leader); + mbarrier_init_pred(smem + 512, 1, leader); + // out_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 520, 4, leader); + // tmem_dealloc_ready: 1 barriers, init_count=2 + mbarrier_init_pred(smem + 528, 2, leader); + // prep_diag_ready: 5 barriers, init_count=2 + mbarrier_init_pred(smem + 536, 2, leader); + mbarrier_init_pred(smem + 544, 2, leader); + mbarrier_init_pred(smem + 552, 2, leader); + mbarrier_init_pred(smem + 560, 2, leader); + mbarrier_init_pred(smem + 568, 2, leader); + // prep_inv16_ready: 5 barriers, init_count=2 + mbarrier_init_pred(smem + 576, 2, leader); + mbarrier_init_pred(smem + 584, 2, leader); + mbarrier_init_pred(smem + 592, 2, leader); + mbarrier_init_pred(smem + 600, 2, leader); + mbarrier_init_pred(smem + 608, 2, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 616); + if (warp == 0) { + int _tmem_hold = smem + 616; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define qk_full_addr (mbar_base + 0) + #define gate_raw_full_addr (mbar_base + 40) + #define qk_raw_full_addr (mbar_base + 80) + #define v_full_addr (mbar_base + 120) + #define v_free_addr (mbar_base + 160) + #define smem_free_addr (mbar_base + 200) + #define raw_inputs_free_addr (mbar_base + 240) + #define state_inp_ready_addr (mbar_base + 280) + #define old_out_ready_addr (mbar_base + 320) + #define u_inp_ready_addr (mbar_base + 360) + #define u2_acc_ready_addr (mbar_base + 400) + #define u2_inp_ready_addr (mbar_base + 440) + #define final_ready_addr (mbar_base + 480) + #define out_empty_addr (mbar_base + 520) + #define tmem_dealloc_ready_addr (mbar_base + 528) + #define prep_diag_ready_addr (mbar_base + 536) + #define prep_inv16_ready_addr (mbar_base + 576) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_state = taddr + 64; + const int tmem_tmem_state_inp = taddr; + const int tmem_tmem_u_acc = taddr + 224; + const int tmem_tmem_u2_inp = taddr + 224; + const int tmem_tmem_u2_acc = taddr; + const int tmem_tmem_out = taddr + 192; + const int tmem_tmem_state_out = taddr + 64; + + // ---- Register redistribution for WGs split across roles ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); + } + + // ---- Role: compute ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 168;"); + { // compute_main + int split_task_idx = blockIdx.x; + int task_idx = split_task_idx / 2; + int value_split_idx = split_task_idx % 2; + int value_row_offset = value_split_idx * 64; + int seq_idx = seq_order[task_idx / num_heads]; + int head_idx = task_idx % num_heads; + long long bos = cu_seqlens[seq_idx]; + long long eos = cu_seqlens[seq_idx + 1]; + int seq_len = (int)(eos - bos); + int num_chunks = (seq_len + 32 - 1) / 32; + int warp_in_wg = warp % 4; + const int tmem_row_base = warp_in_wg * 32 << 16; + int lane_quad = lane & 3; + int local_row_top = warp_in_wg * 16 + lane / 4; + int local_row_bot = local_row_top + 8; + int state_row_top = value_row_offset + local_row_top; + int state_row_bot = value_row_offset + local_row_bot; + int warp_id_in_role = (warp - 0); + int compute_local_warp = warp_id_in_role; + long long state_head_base = ((long long)seq_idx * (long long)num_heads + (long long)head_idx) * 128 * 128; + long long state_base_top = state_head_base + (long long)state_row_top * 128; + long long state_base_bot = state_head_base + (long long)state_row_bot * 128; + #pragma unroll + for (int state_col_half = 0; state_col_half < 2; state_col_half++) { + float state_init[32]; + state_init[0] = 0.0f; + state_init[1] = 0.0f; + state_init[2] = 0.0f; + state_init[3] = 0.0f; + state_init[4] = 0.0f; + state_init[5] = 0.0f; + state_init[6] = 0.0f; + state_init[7] = 0.0f; + state_init[8] = 0.0f; + state_init[9] = 0.0f; + state_init[10] = 0.0f; + state_init[11] = 0.0f; + state_init[12] = 0.0f; + state_init[13] = 0.0f; + state_init[14] = 0.0f; + state_init[15] = 0.0f; + state_init[16] = 0.0f; + state_init[17] = 0.0f; + state_init[18] = 0.0f; + state_init[19] = 0.0f; + state_init[20] = 0.0f; + state_init[21] = 0.0f; + state_init[22] = 0.0f; + state_init[23] = 0.0f; + state_init[24] = 0.0f; + state_init[25] = 0.0f; + state_init[26] = 0.0f; + state_init[27] = 0.0f; + state_init[28] = 0.0f; + state_init[29] = 0.0f; + state_init[30] = 0.0f; + state_init[31] = 0.0f; + if (use_initial_state != 0) { + #pragma unroll + for (int state_col_group = 0; state_col_group < 8; state_col_group++) { + int state_col_pair = state_col_half * 64 + state_col_group * 8 + lane_quad * 2; + const int state_reg_base = state_col_group * 4; + state_init[state_reg_base] = (float)initial_state[state_base_top + (long long)state_col_pair]; + state_init[state_reg_base + 1] = (float)initial_state[state_base_top + (long long)state_col_pair + 1]; + state_init[state_reg_base + 2] = (float)initial_state[state_base_bot + (long long)state_col_pair]; + state_init[state_reg_base + 3] = (float)initial_state[state_base_bot + (long long)state_col_pair + 1]; + } + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x8.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32};" + :: "r"(taddr + 64 + (unsigned int)tmem_row_base + (unsigned int)(state_col_half * 64)), "r"(*reinterpret_cast(&state_init[0])), "r"(*reinterpret_cast(&state_init[1])), "r"(*reinterpret_cast(&state_init[2])), "r"(*reinterpret_cast(&state_init[3])), "r"(*reinterpret_cast(&state_init[4])), "r"(*reinterpret_cast(&state_init[5])), "r"(*reinterpret_cast(&state_init[6])), "r"(*reinterpret_cast(&state_init[7])), "r"(*reinterpret_cast(&state_init[8])), "r"(*reinterpret_cast(&state_init[9])), "r"(*reinterpret_cast(&state_init[10])), "r"(*reinterpret_cast(&state_init[11])), "r"(*reinterpret_cast(&state_init[12])), "r"(*reinterpret_cast(&state_init[13])), "r"(*reinterpret_cast(&state_init[14])), "r"(*reinterpret_cast(&state_init[15])), "r"(*reinterpret_cast(&state_init[16])), "r"(*reinterpret_cast(&state_init[17])), "r"(*reinterpret_cast(&state_init[18])), "r"(*reinterpret_cast(&state_init[19])), "r"(*reinterpret_cast(&state_init[20])), "r"(*reinterpret_cast(&state_init[21])), "r"(*reinterpret_cast(&state_init[22])), "r"(*reinterpret_cast(&state_init[23])), "r"(*reinterpret_cast(&state_init[24])), "r"(*reinterpret_cast(&state_init[25])), "r"(*reinterpret_cast(&state_init[26])), "r"(*reinterpret_cast(&state_init[27])), "r"(*reinterpret_cast(&state_init[28])), "r"(*reinterpret_cast(&state_init[29])), "r"(*reinterpret_cast(&state_init[30])), "r"(*reinterpret_cast(&state_init[31])) + : "memory"); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + unsigned int compute_stage = 0; + unsigned int _phase_qk_full = 0; + unsigned int _phase_v_full = 0; + unsigned int _phase_old_out_ready = 0; + unsigned int _phase_u2_acc_ready = 0; + unsigned int _phase_final_ready = 0; + #pragma unroll 1 + for (int chunk_idx = 0; chunk_idx < num_chunks; chunk_idx++) { + mbarrier_wait(qk_full_addr + (compute_stage) * 8, _phase_qk_full); + #pragma unroll + for (int state_col_half_1 = 0; state_col_half_1 < 2; state_col_half_1++) { + int state_addr = taddr + 64 + (unsigned int)tmem_row_base + (unsigned int)(state_col_half_1 * 64); + float _tmem_load_0[32]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x8.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=r"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) + : "r"(state_addr) + : "memory"); + uint32_t _tmem_load_0_bf16[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile( + "tcgen05.st.sync.aligned.16x128b.x8.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(taddr + (unsigned int)tmem_row_base + (unsigned int)(state_col_half_1 * 32)), "r"(*reinterpret_cast(&_tmem_load_0_bf16[0])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[1])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[2])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[3])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[4])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[5])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[6])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[7])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[8])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[9])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[10])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[11])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[12])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[13])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[14])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[15])) + : "memory"); + #pragma unroll + for (int state_col_group_1 = 0; state_col_group_1 < 8; state_col_group_1++) { + int state_col_pair_1 = state_col_half_1 * 64 + state_col_group_1 * 8 + lane_quad * 2; + const int state_reg_base_1 = state_col_group_1 * 4; + float state_scale_0 = smem_gt_all[compute_stage * 10496 + (unsigned int)state_col_pair_1]; + float state_scale_1 = smem_gt_all[compute_stage * 10496 + (unsigned int)state_col_pair_1 + 1]; + _tmem_load_0[state_reg_base_1] = _tmem_load_0[state_reg_base_1] * state_scale_0; + _tmem_load_0[state_reg_base_1 + 1] = _tmem_load_0[state_reg_base_1 + 1] * state_scale_1; + _tmem_load_0[state_reg_base_1 + 2] = _tmem_load_0[state_reg_base_1 + 2] * state_scale_0; + _tmem_load_0[state_reg_base_1 + 3] = _tmem_load_0[state_reg_base_1 + 3] * state_scale_1; + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x8.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32};" + :: "r"(state_addr), "r"(*reinterpret_cast(&_tmem_load_0[0])), "r"(*reinterpret_cast(&_tmem_load_0[1])), "r"(*reinterpret_cast(&_tmem_load_0[2])), "r"(*reinterpret_cast(&_tmem_load_0[3])), "r"(*reinterpret_cast(&_tmem_load_0[4])), "r"(*reinterpret_cast(&_tmem_load_0[5])), "r"(*reinterpret_cast(&_tmem_load_0[6])), "r"(*reinterpret_cast(&_tmem_load_0[7])), "r"(*reinterpret_cast(&_tmem_load_0[8])), "r"(*reinterpret_cast(&_tmem_load_0[9])), "r"(*reinterpret_cast(&_tmem_load_0[10])), "r"(*reinterpret_cast(&_tmem_load_0[11])), "r"(*reinterpret_cast(&_tmem_load_0[12])), "r"(*reinterpret_cast(&_tmem_load_0[13])), "r"(*reinterpret_cast(&_tmem_load_0[14])), "r"(*reinterpret_cast(&_tmem_load_0[15])), "r"(*reinterpret_cast(&_tmem_load_0[16])), "r"(*reinterpret_cast(&_tmem_load_0[17])), "r"(*reinterpret_cast(&_tmem_load_0[18])), "r"(*reinterpret_cast(&_tmem_load_0[19])), "r"(*reinterpret_cast(&_tmem_load_0[20])), "r"(*reinterpret_cast(&_tmem_load_0[21])), "r"(*reinterpret_cast(&_tmem_load_0[22])), "r"(*reinterpret_cast(&_tmem_load_0[23])), "r"(*reinterpret_cast(&_tmem_load_0[24])), "r"(*reinterpret_cast(&_tmem_load_0[25])), "r"(*reinterpret_cast(&_tmem_load_0[26])), "r"(*reinterpret_cast(&_tmem_load_0[27])), "r"(*reinterpret_cast(&_tmem_load_0[28])), "r"(*reinterpret_cast(&_tmem_load_0[29])), "r"(*reinterpret_cast(&_tmem_load_0[30])), "r"(*reinterpret_cast(&_tmem_load_0[31])) + : "memory"); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(state_inp_ready_addr + (compute_stage) * 8); + } + mbarrier_wait(v_full_addr + (compute_stage) * 8, _phase_v_full); + mbarrier_wait(old_out_ready_addr + (compute_stage) * 8, _phase_old_out_ready); + float _tmem_load_1[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x4.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])) + : "r"(taddr + 224 + (unsigned int)tmem_row_base) + : "memory"); + float residual_values[16]; + int v_stage_addr = smem_v_addr + compute_stage * 41984; + unsigned int v_ld_bits[2]; + #pragma unroll + for (int token_group = 0; token_group < 4; token_group++) { + int token_pair = token_group * 8 + lane_quad * 2; + const int residual_reg_base = token_group * 4; + float beta_0 = smem_prep_beta_all[compute_stage * 10496 + (unsigned int)token_pair]; + float beta_1 = smem_prep_beta_all[compute_stage * 10496 + (unsigned int)token_pair + 1]; + int v_ld_matrix = lane / 8 & 1; + int v_ld_token = token_group * 8 + (lane & 7); + int v_ld_row = warp_in_wg * 16 + v_ld_matrix * 8; + int v_ld_row_addr = v_stage_addr + v_ld_token * 64 * 2; + int v_ld_addr = (v_ld_row_addr + (v_ld_row * 2 ^ (v_ld_row_addr >> 7 & 7) << 4)); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0, %1}, [%2];\n" + : "=r"(v_ld_bits[0]), "=r"(v_ld_bits[1]) + : "r"(v_ld_addr) + : "memory"); + float v_ld_bits_fp32[4]; + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&v_ld_bits_fp32[_pair * 2])[0]), "=f"((&v_ld_bits_fp32[_pair * 2])[1]) + : "r"(v_ld_bits[_pair + 0])); + } + residual_values[residual_reg_base] = (v_ld_bits_fp32[0] - _tmem_load_1[residual_reg_base]) * beta_0; + residual_values[residual_reg_base + 1] = (v_ld_bits_fp32[1] - _tmem_load_1[residual_reg_base + 1]) * beta_1; + residual_values[residual_reg_base + 2] = (v_ld_bits_fp32[2] - _tmem_load_1[residual_reg_base + 2]) * beta_0; + residual_values[residual_reg_base + 3] = (v_ld_bits_fp32[3] - _tmem_load_1[residual_reg_base + 3]) * beta_1; + } + uint32_t residual_values_bf16[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(residual_values[_lp*2 + 0], residual_values[_lp*2+1 + 0])); + residual_values_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile( + "tcgen05.st.sync.aligned.16x128b.x4.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "r"(taddr + 224 + (unsigned int)tmem_row_base), "r"(*reinterpret_cast(&residual_values_bf16[0])), "r"(*reinterpret_cast(&residual_values_bf16[1])), "r"(*reinterpret_cast(&residual_values_bf16[2])), "r"(*reinterpret_cast(&residual_values_bf16[3])), "r"(*reinterpret_cast(&residual_values_bf16[4])), "r"(*reinterpret_cast(&residual_values_bf16[5])), "r"(*reinterpret_cast(&residual_values_bf16[6])), "r"(*reinterpret_cast(&residual_values_bf16[7])) + : "memory"); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(v_free_addr + (compute_stage) * 8); + mbarrier_arrive(u_inp_ready_addr + (compute_stage) * 8); + } + mbarrier_wait(u2_acc_ready_addr + (compute_stage) * 8, _phase_u2_acc_ready); + float _tmem_load_2[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x4.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=r"(*reinterpret_cast(&_tmem_load_2[0])), "=r"(*reinterpret_cast(&_tmem_load_2[1])), "=r"(*reinterpret_cast(&_tmem_load_2[2])), "=r"(*reinterpret_cast(&_tmem_load_2[3])), "=r"(*reinterpret_cast(&_tmem_load_2[4])), "=r"(*reinterpret_cast(&_tmem_load_2[5])), "=r"(*reinterpret_cast(&_tmem_load_2[6])), "=r"(*reinterpret_cast(&_tmem_load_2[7])), "=r"(*reinterpret_cast(&_tmem_load_2[8])), "=r"(*reinterpret_cast(&_tmem_load_2[9])), "=r"(*reinterpret_cast(&_tmem_load_2[10])), "=r"(*reinterpret_cast(&_tmem_load_2[11])), "=r"(*reinterpret_cast(&_tmem_load_2[12])), "=r"(*reinterpret_cast(&_tmem_load_2[13])), "=r"(*reinterpret_cast(&_tmem_load_2[14])), "=r"(*reinterpret_cast(&_tmem_load_2[15])) + : "r"(taddr + (unsigned int)tmem_row_base) + : "memory"); + uint32_t _tmem_load_2_bf16[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_2[_lp*2 + 0], _tmem_load_2[_lp*2+1 + 0])); + _tmem_load_2_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile( + "tcgen05.st.sync.aligned.16x128b.x4.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "r"(taddr + 224 + (unsigned int)tmem_row_base), "r"(*reinterpret_cast(&_tmem_load_2_bf16[0])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[1])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[2])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[3])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[4])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[5])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[6])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[7])) + : "memory"); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(u2_inp_ready_addr + (compute_stage) * 8); + } + mbarrier_wait(final_ready_addr + (compute_stage) * 8, _phase_final_ready); + compute_stage += 1; + if (compute_stage == 5) { compute_stage = 0; _phase_qk_full ^= 1; _phase_v_full ^= 1; _phase_old_out_ready ^= 1; _phase_u2_acc_ready ^= 1; _phase_final_ready ^= 1; } + } + if (store_final_state != 0) { + #pragma unroll + for (int state_col_half_2 = 0; state_col_half_2 < 2; state_col_half_2++) { + float _tmem_load_3[32]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x8.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=r"(*reinterpret_cast(&_tmem_load_3[0])), "=r"(*reinterpret_cast(&_tmem_load_3[1])), "=r"(*reinterpret_cast(&_tmem_load_3[2])), "=r"(*reinterpret_cast(&_tmem_load_3[3])), "=r"(*reinterpret_cast(&_tmem_load_3[4])), "=r"(*reinterpret_cast(&_tmem_load_3[5])), "=r"(*reinterpret_cast(&_tmem_load_3[6])), "=r"(*reinterpret_cast(&_tmem_load_3[7])), "=r"(*reinterpret_cast(&_tmem_load_3[8])), "=r"(*reinterpret_cast(&_tmem_load_3[9])), "=r"(*reinterpret_cast(&_tmem_load_3[10])), "=r"(*reinterpret_cast(&_tmem_load_3[11])), "=r"(*reinterpret_cast(&_tmem_load_3[12])), "=r"(*reinterpret_cast(&_tmem_load_3[13])), "=r"(*reinterpret_cast(&_tmem_load_3[14])), "=r"(*reinterpret_cast(&_tmem_load_3[15])), "=r"(*reinterpret_cast(&_tmem_load_3[16])), "=r"(*reinterpret_cast(&_tmem_load_3[17])), "=r"(*reinterpret_cast(&_tmem_load_3[18])), "=r"(*reinterpret_cast(&_tmem_load_3[19])), "=r"(*reinterpret_cast(&_tmem_load_3[20])), "=r"(*reinterpret_cast(&_tmem_load_3[21])), "=r"(*reinterpret_cast(&_tmem_load_3[22])), "=r"(*reinterpret_cast(&_tmem_load_3[23])), "=r"(*reinterpret_cast(&_tmem_load_3[24])), "=r"(*reinterpret_cast(&_tmem_load_3[25])), "=r"(*reinterpret_cast(&_tmem_load_3[26])), "=r"(*reinterpret_cast(&_tmem_load_3[27])), "=r"(*reinterpret_cast(&_tmem_load_3[28])), "=r"(*reinterpret_cast(&_tmem_load_3[29])), "=r"(*reinterpret_cast(&_tmem_load_3[30])), "=r"(*reinterpret_cast(&_tmem_load_3[31])) + : "r"(taddr + 64 + (unsigned int)tmem_row_base + (unsigned int)(state_col_half_2 * 64)) + : "memory"); + #pragma unroll + for (int state_col_group_2 = 0; state_col_group_2 < 8; state_col_group_2++) { + int state_col_pair_2 = state_col_half_2 * 64 + state_col_group_2 * 8 + lane_quad * 2; + const int state_reg_base_2 = state_col_group_2 * 4; + final_state[state_base_top + (long long)state_col_pair_2] = _tmem_load_3[state_reg_base_2]; + final_state[state_base_top + (long long)state_col_pair_2 + 1] = _tmem_load_3[state_reg_base_2 + 1]; + final_state[state_base_bot + (long long)state_col_pair_2] = _tmem_load_3[state_reg_base_2 + 2]; + final_state[state_base_bot + (long long)state_col_pair_2 + 1] = _tmem_load_3[state_reg_base_2 + 3]; + } + } + } + asm volatile("barrier.sync 10, 128;" ::: "memory"); + if (compute_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(tmem_dealloc_ready_addr); + } + } + } + // ---- Role: epilogue ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); + { // epilogue_main + int split_task_idx_1 = blockIdx.x; + int task_idx_1 = split_task_idx_1 / 2; + int value_split_idx_1 = split_task_idx_1 % 2; + int value_row_offset_1 = value_split_idx_1 * 64; + int seq_idx_1 = seq_order[task_idx_1 / num_heads]; + int head_idx_1 = task_idx_1 % num_heads; + long long bos_1 = cu_seqlens[seq_idx_1]; + long long eos_1 = cu_seqlens[seq_idx_1 + 1]; + int seq_len_1 = (int)(eos_1 - bos_1); + int num_chunks_1 = (seq_len_1 + 32 - 1) / 32; + int warp_id_in_role_1 = (warp - 4); + int epilogue_local_warp = warp_id_in_role_1; + int warp_in_wg_1 = warp % 4; + const int tmem_row_base_1 = warp_in_wg_1 * 32 << 16; + int lane_quad_1 = lane & 3; + int local_row_top_1 = warp_in_wg_1 * 16 + lane / 4; + int local_row_bot_1 = local_row_top_1 + 8; + int state_row_top_1 = value_row_offset_1 + local_row_top_1; + int state_row_bot_1 = value_row_offset_1 + local_row_bot_1; + unsigned int epilogue_stage = 0; + unsigned int output_stage = 0; + unsigned int _phase_final_ready_1 = 0; + #pragma unroll 1 + for (int chunk_idx_1 = 0; chunk_idx_1 < num_chunks_1; chunk_idx_1++) { + mbarrier_wait(final_ready_addr + (epilogue_stage) * 8, _phase_final_ready_1); + int chunk_is_full = ((seq_len_1 >= (chunk_idx_1 + 1) * 32) ? 1 : 0); + if (chunk_is_full != 0) { + float _tmem_load_4[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x4.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])), "=r"(*reinterpret_cast(&_tmem_load_4[8])), "=r"(*reinterpret_cast(&_tmem_load_4[9])), "=r"(*reinterpret_cast(&_tmem_load_4[10])), "=r"(*reinterpret_cast(&_tmem_load_4[11])), "=r"(*reinterpret_cast(&_tmem_load_4[12])), "=r"(*reinterpret_cast(&_tmem_load_4[13])), "=r"(*reinterpret_cast(&_tmem_load_4[14])), "=r"(*reinterpret_cast(&_tmem_load_4[15])) + : "r"(taddr + 192 + (unsigned int)tmem_row_base_1) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(out_empty_addr); + } + if (epilogue_local_warp == 0) { + if (chunk_idx_1 >= 2) { + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 9, 128;" ::: "memory"); + int out_stage_addr = smem_out_addr + output_stage * 4096; + unsigned int out_packed[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_4[_lp*2 + 0], _tmem_load_4[_lp*2+1 + 0])); + out_packed[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int token_group_1 = 0; token_group_1 < 2; token_group_1++) { + int mtx_idx = lane / 8; + int row_addr = lane & 7; + int dim_base = epilogue_local_warp * 16 + (mtx_idx & 1) * 8; + int token_base = token_group_1 * 16 + mtx_idx / 2 * 8; + int token_addr = token_base + row_addr; + int token_pair_1 = token_addr / 2; + int token_parity = token_addr & 1; + int raw_row = token_pair_1; + int raw_col = (dim_base & 63 ^ (token_pair_1 & 3) << 4 ^ token_parity << 3) + token_parity * 64; + int stsm_offset = (raw_row * 128 + raw_col) * 2; + const int pack_base = token_group_1 * 4; + uint32_t _stmatrix_addr_0 = static_cast((unsigned long long)(out_stage_addr + stsm_offset)); + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_0), "r"(*reinterpret_cast(&out_packed[pack_base])), "r"(*reinterpret_cast(&out_packed[pack_base + 1])), "r"(*reinterpret_cast(&out_packed[pack_base + 2])), "r"(*reinterpret_cast(&out_packed[pack_base + 3])) + : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 9, 128;" ::: "memory"); + if (epilogue_local_warp == 0) { + if (elect_sync()) { + tma_store_4d(out_tma, 0, (int)(bos_1 + (long long)(chunk_idx_1 * 32)), head_idx_1, value_split_idx_1, smem_out_addr + output_stage * 4096); + } + asm volatile("cp.async.bulk.commit_group;"); + } + output_stage = output_stage ^ 1; + } else { + float _tmem_load_5[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x4.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=r"(*reinterpret_cast(&_tmem_load_5[0])), "=r"(*reinterpret_cast(&_tmem_load_5[1])), "=r"(*reinterpret_cast(&_tmem_load_5[2])), "=r"(*reinterpret_cast(&_tmem_load_5[3])), "=r"(*reinterpret_cast(&_tmem_load_5[4])), "=r"(*reinterpret_cast(&_tmem_load_5[5])), "=r"(*reinterpret_cast(&_tmem_load_5[6])), "=r"(*reinterpret_cast(&_tmem_load_5[7])), "=r"(*reinterpret_cast(&_tmem_load_5[8])), "=r"(*reinterpret_cast(&_tmem_load_5[9])), "=r"(*reinterpret_cast(&_tmem_load_5[10])), "=r"(*reinterpret_cast(&_tmem_load_5[11])), "=r"(*reinterpret_cast(&_tmem_load_5[12])), "=r"(*reinterpret_cast(&_tmem_load_5[13])), "=r"(*reinterpret_cast(&_tmem_load_5[14])), "=r"(*reinterpret_cast(&_tmem_load_5[15])) + : "r"(taddr + 192 + (unsigned int)tmem_row_base_1) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(out_empty_addr); + } + #pragma unroll + for (int token_group_2 = 0; token_group_2 < 4; token_group_2++) { + int token_pair_2 = token_group_2 * 8 + lane_quad_1 * 2; + const int out_reg_base = token_group_2 * 4; + long long out_token_0 = bos_1 + (long long)(chunk_idx_1 * 32 + token_pair_2); + long long out_token_1 = out_token_0 + 1; + if (out_token_0 < eos_1) { + long long out_idx_top_0 = (out_token_0 * (long long)num_heads + (long long)head_idx_1) * 128 + (long long)state_row_top_1; + long long out_idx_bot_0 = (out_token_0 * (long long)num_heads + (long long)head_idx_1) * 128 + (long long)state_row_bot_1; + out[out_idx_top_0] = _tmem_load_5[out_reg_base]; + out[out_idx_bot_0] = _tmem_load_5[out_reg_base + 2]; + } + if (out_token_1 < eos_1) { + long long out_idx_top_1 = (out_token_1 * (long long)num_heads + (long long)head_idx_1) * 128 + (long long)state_row_top_1; + long long out_idx_bot_1 = (out_token_1 * (long long)num_heads + (long long)head_idx_1) * 128 + (long long)state_row_bot_1; + out[out_idx_top_1] = _tmem_load_5[out_reg_base + 1]; + out[out_idx_bot_1] = _tmem_load_5[out_reg_base + 3]; + } + } + } + epilogue_stage += 1; + if (epilogue_stage == 5) { epilogue_stage = 0; _phase_final_ready_1 ^= 1; } + } + if (epilogue_local_warp == 0) { + asm volatile("cp.async.bulk.wait_group 0;"); + } + asm volatile("barrier.sync 9, 128;" ::: "memory"); + if (epilogue_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(tmem_dealloc_ready_addr); + } + } + } + // ---- Role: mma ---- + } else if (warp == 9) { + { // mma_main + int split_task_idx_2 = blockIdx.x; + int task_idx_2 = split_task_idx_2 / 2; + int seq_idx_2 = seq_order[task_idx_2 / num_heads]; + long long bos_2 = cu_seqlens[seq_idx_2]; + long long eos_2 = cu_seqlens[seq_idx_2 + 1]; + int seq_len_2 = (int)(eos_2 - bos_2); + int num_chunks_2 = (seq_len_2 + 32 - 1) / 32; + unsigned int mma_stage = 0; + unsigned int _phase_qk_full_1 = 0; + unsigned int _phase_state_inp_ready = 0; + unsigned int _phase_out_empty_0 = 1; + unsigned int _phase_u_inp_ready = 0; + unsigned int _phase_u2_inp_ready = 0; + #pragma unroll 1 + for (int _chunk_idx = 0; _chunk_idx < num_chunks_2; _chunk_idx++) { + mbarrier_wait(qk_full_addr + (mma_stage) * 8, _phase_qk_full_1); + mbarrier_wait(state_inp_ready_addr + (mma_stage) * 8, _phase_state_inp_ready); + mbarrier_wait(out_empty_addr, _phase_out_empty_0); + _phase_out_empty_0 ^= 1; + int _mma_b_addr_0 = smem_qd_addr + mma_stage * 41984; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 67634320;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 16], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 24], db, id, p1;\n\t" + "add.u32 blo, blo, 250;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 32], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 40], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 48], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 56], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_out), "r"(_mma_b_lo_0), "r"(tmem_tmem_state_inp), "r"(0)); + int _mma_b_addr_1 = smem_kd_addr + mma_stage * 41984; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 67634320;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 16], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 24], db, id, p1;\n\t" + "add.u32 blo, blo, 250;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 32], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 40], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 48], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 56], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_u_acc), "r"(_mma_b_lo_1), "r"(tmem_tmem_state_inp), "r"(0)); + elect_commit2(old_out_ready_addr + (mma_stage) * 8, raw_inputs_free_addr + (mma_stage) * 8); + mbarrier_wait(u_inp_ready_addr + (mma_stage) * 8, _phase_u_inp_ready); + int _mma_b_addr_2 = smem_inv_addr + mma_stage * 41984; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0xC0004010;\n\t" + "mov.b32 id, 67634320;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 64;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_u2_acc), "r"(_mma_b_lo_2), "r"(tmem_tmem_u2_inp), "r"(0)); + elect_commit(u2_acc_ready_addr + (mma_stage) * 8); + mbarrier_wait(u2_inp_ready_addr + (mma_stage) * 8, _phase_u2_inp_ready); + int _mma_b_addr_3 = smem_final_trans_addr + mma_stage * 41984; + int _mma_b_lo_3 = make_warp_uniform(((_mma_b_addr_3 >> 4) & 0x3FFF) | 0x1000000); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 69797008;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_state_out), "r"(_mma_b_lo_3), "r"(tmem_tmem_u2_inp), "r"(1)); + elect_commit2(final_ready_addr + (mma_stage) * 8, smem_free_addr + (mma_stage) * 8); + mma_stage += 1; + if (mma_stage == 5) { mma_stage = 0; _phase_qk_full_1 ^= 1; _phase_state_inp_ready ^= 1; _phase_u_inp_ready ^= 1; _phase_u2_inp_ready ^= 1; } + } + unsigned int _phase_tmem_dealloc_ready_0 = 0; + mbarrier_wait(tmem_dealloc_ready_addr, _phase_tmem_dealloc_ready_0); + _phase_tmem_dealloc_ready_0 ^= 1; + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + // ---- Role: load ---- + } else if (warp == 10) { + { // load_main + int split_task_idx_3 = blockIdx.x; + int task_idx_3 = split_task_idx_3 / 2; + int value_split_idx_2 = split_task_idx_3 % 2; + int value_row_offset_2 = value_split_idx_2 * 64; + int seq_idx_3 = seq_order[task_idx_3 / num_heads]; + int head_idx_2 = task_idx_3 % num_heads; + long long bos_3 = cu_seqlens[seq_idx_3]; + long long eos_3 = cu_seqlens[seq_idx_3 + 1]; + int seq_len_3 = (int)(eos_3 - bos_3); + int num_chunks_3 = (seq_len_3 + 32 - 1) / 32; + unsigned int load_stage = 0; + unsigned int _phase_v_free = 1; + unsigned int _phase_qk_full_2 = 0; + #pragma unroll 1 + for (int chunk_idx_2 = 0; chunk_idx_2 < num_chunks_3; chunk_idx_2++) { + mbarrier_wait(v_free_addr + (load_stage) * 8, _phase_v_free); + mbarrier_wait(qk_full_addr + (load_stage) * 8, _phase_qk_full_2); + int chunk_is_full_1 = ((seq_len_3 >= (chunk_idx_2 + 1) * 32) ? 1 : 0); + if (elect_sync()) { + if (chunk_is_full_1 != 0) { + mbarrier_arrive_expect_tx(v_full_addr + (load_stage) * 8, 4096); + tma_3d_gmem2smem(smem_v_addr + load_stage * 41984, v_tma, value_row_offset_2, head_idx_2, (int)(bos_3 + (long long)(chunk_idx_2 * 32)), v_full_addr + (load_stage) * 8); + } + } + if (chunk_is_full_1 == 0) { + #pragma unroll + for (int v_load_iter = 0; v_load_iter < 8; v_load_iter++) { + int v_item = v_load_iter * 32 + lane; + int row = v_item / 8; + int segment = v_item % 8; + long long token = bos_3 + (long long)(chunk_idx_2 * 32 + row); + int token_valid = ((token < eos_3) ? 1 : 0); + long long v_src = (token * (long long)num_heads + (long long)head_idx_2) * 128 + (long long)value_row_offset_2 + (long long)(segment * 8); + int v_dst_row_addr = smem_v_addr + load_stage * 41984 + (unsigned int)(row * 64 * 2); + int v_dst_addr = (v_dst_row_addr + (segment * 8 * 2 ^ (v_dst_row_addr >> 7 & 7) << 4)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16, %2;" + :: "r"(v_dst_addr), "l"(v + v_src), "r"((token_valid != 0) ? 16 : 0)); + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + } + asm volatile("barrier.sync 8, 32;" ::: "memory"); + if (elect_sync()) { + if (chunk_is_full_1 == 0) { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(v_full_addr + (load_stage) * 8); + } + } + load_stage += 1; + if (load_stage == 5) { load_stage = 0; _phase_v_free ^= 1; _phase_qk_full_2 ^= 1; } + } + } + // ---- Role: prep ---- + } else if (warp >= 12 && warp <= 31) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); + { // prep_main + int split_task_idx_4 = blockIdx.x; + int task_idx_4 = split_task_idx_4 / 2; + int seq_idx_4 = seq_order[task_idx_4 / num_heads]; + int head_idx_3 = task_idx_4 % num_heads; + long long bos_4 = cu_seqlens[seq_idx_4]; + long long eos_4 = cu_seqlens[seq_idx_4 + 1]; + int seq_len_4 = (int)(eos_4 - bos_4); + int num_chunks_4 = (seq_len_4 + 32 - 1) / 32; + int instance_id = (warp - 12) / 4; + int prep_instance = instance_id; + int warp_id_in_role_2 = (warp - 12); + int prep_local_warp = warp_id_in_role_2 - prep_instance * 4; + int prep_tid = prep_local_warp * 32 + lane; + int num_prep_iters = (num_chunks_4 + 4 - prep_instance) / 5; + unsigned int prep_stage = (unsigned int)prep_instance; + int gate_rate_stage_f32 = prep_instance * 10496; + if (prep_tid == 0) { + float _expf_0 = __expf(A_log[head_idx_3]); + smem_gate_rate_all[gate_rate_stage_f32] = _expf_0; + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + unsigned int _phase_raw_inputs_free = 1; + unsigned int _phase_gate_raw_full = 0; + unsigned int _phase_smem_free = 1; + unsigned int _phase_qk_raw_full = 0; + unsigned int _phase_prep_diag_ready = 0; + unsigned int _phase_prep_inv16_ready = 0; + #pragma unroll 1 + for (int prep_iter = 0; prep_iter < num_prep_iters; prep_iter++) { + int chunk_idx_3 = prep_iter * 5 + prep_instance; + int stage_f32 = prep_stage * 10496; + int stage_bf16 = prep_stage * 20992; + int chunk_is_full_2 = ((seq_len_4 >= (chunk_idx_3 + 1) * 32) ? 1 : 0); + float early_beta_value = 0.0f; + float early_gate0 = 0.0f; + if (chunk_is_full_2 != 0) { + mbarrier_wait(raw_inputs_free_addr + (prep_stage) * 8, _phase_raw_inputs_free); + if (prep_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(gate_raw_full_addr + (prep_stage) * 8, 8704); + tma_3d_gmem2smem(smem_g_raw_addr + prep_stage * 41984, g_tma, 0, head_idx_3, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), gate_raw_full_addr + (prep_stage) * 8); + tma_2d_gmem2smem(smem_beta_raw_addr + prep_stage * 41984, beta_tma, head_idx_3 / 8 * 8, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), gate_raw_full_addr + (prep_stage) * 8); + mbarrier_arrive_expect_tx(qk_raw_full_addr + (prep_stage) * 8, 16384); + tma_4d_gmem2smem(smem_kd_addr + prep_stage * 41984, k_tma, 0, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), head_idx_3, 0, qk_raw_full_addr + (prep_stage) * 8); + } + } + mbarrier_wait(gate_raw_full_addr + (prep_stage) * 8, _phase_gate_raw_full); + if (prep_local_warp == 2 && lane < 32) { + unsigned int beta_raw_pair[1]; + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&beta_raw_pair[0])) : "r"(smem_beta_raw_addr + prep_stage * 41984 + (unsigned int)(lane * 16) + (unsigned int)(head_idx_3 % 8 / 2 * 4))); + float beta_raw_pair_fp32[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&beta_raw_pair_fp32[_pair * 2])[0]), "=f"((&beta_raw_pair_fp32[_pair * 2])[1]) + : "r"(beta_raw_pair[_pair + 0])); + } + float beta_logit = beta_raw_pair_fp32[0]; + if (head_idx_3 % 2 != 0) { + beta_logit = beta_raw_pair_fp32[1]; + } + float _tanh_approx_0; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_0) : "f"(beta_logit * 0.5f)); + early_beta_value = _tanh_approx_0 * 0.5f + 0.5f; + } + if (prep_tid < 128) { + float early_gate_rate = smem_gate_rate_all[stage_f32]; + float early_gate_bias = dt_bias[head_idx_3 * 128 + prep_tid]; + __nv_bfloat16 early_gate_raw = smem_g_raw_all[stage_bf16 + prep_tid]; + float _cvt_f32_0 = __bfloat162float(early_gate_raw); + float early_gate_arg = early_gate_rate * (_cvt_f32_0 + early_gate_bias); + float _tanh_approx_1; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_1) : "f"(early_gate_arg * 0.5f)); + float early_gate_sigmoid = _tanh_approx_1 * 0.5f + 0.5f; + early_gate0 = lower_bound * 1.4426950408889634f * early_gate_sigmoid; + } + } + mbarrier_wait(smem_free_addr + (prep_stage) * 8, _phase_smem_free); + if (chunk_is_full_2 != 0) { + if (prep_local_warp == 0) { + if (elect_sync()) { + tma_4d_gmem2smem(smem_q_raw_prefetch_addr + prep_stage * 41984, q_tma, 0, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), head_idx_3, 0, qk_raw_full_addr + (prep_stage) * 8); + } + } + } + if (chunk_is_full_2 == 0) { + #pragma unroll + for (int gate_load_pass = 0; gate_load_pass < 4; gate_load_pass++) { + int gate_load_item = gate_load_pass * 128 + prep_tid; + int gate_load_row = gate_load_item / 16; + int gate_load_segment = gate_load_item % 16; + long long gate_load_token = bos_4 + (long long)(chunk_idx_3 * 32 + gate_load_row); + long long gate_load_base = (gate_load_token * (long long)num_heads + (long long)head_idx_3) * 128 + (long long)(gate_load_segment * 8); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16, %2;" + :: "r"(smem_g_raw_addr + prep_stage * 41984 + (unsigned int)(gate_load_item * 16)), "l"(g + gate_load_base), "r"((gate_load_token < eos_4) ? 16 : 0)); + } + } + if (chunk_is_full_2 == 0) { + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + } + if (prep_local_warp == 2 && lane < 32) { + float beta_value = early_beta_value; + if (chunk_is_full_2 == 0) { + long long beta_token = bos_4 + (long long)(chunk_idx_3 * 32 + lane); + if (beta_token < eos_4) { + float beta_logit_1 = (float)beta[beta_token * (long long)num_heads + (long long)head_idx_3]; + float _tanh_approx_2; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_2) : "f"(beta_logit_1 * 0.5f)); + beta_value = _tanh_approx_2 * 0.5f + 0.5f; + } + } + smem_prep_beta_all[stage_f32 + lane] = beta_value; + } + if (prep_tid < 128) { + int gate_col = prep_tid; + float gate_rate = smem_gate_rate_all[stage_f32]; + float gate_bias = dt_bias[head_idx_3 * 128 + gate_col]; + float prefix_log2 = 0.0f; + for (int gate_row = 0; gate_row < 32; gate_row++) { + long long gate_token = bos_4 + (long long)(chunk_idx_3 * 32 + gate_row); + float gate_log2 = 0.0f; + int gate_needs_compute = 1; + if (gate_row == 0) { + if (chunk_is_full_2 != 0) { + gate_log2 = early_gate0; + gate_needs_compute = 0; + } + } + if (gate_needs_compute != 0) { + if (gate_token < eos_4) { + __nv_bfloat16 gate_raw = smem_g_raw_all[stage_bf16 + gate_row * 128 + gate_col]; + float _cvt_f32_1 = __bfloat162float(gate_raw); + float gate_arg = gate_rate * (_cvt_f32_1 + gate_bias); + float _tanh_approx_3; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_3) : "f"(gate_arg * 0.5f)); + float gate_sigmoid = _tanh_approx_3 * 0.5f + 0.5f; + gate_log2 = lower_bound * 1.4426950408889634f * gate_sigmoid; + } + } + prefix_log2 += gate_log2; + smem_gate_all[stage_f32 + gate_row * 128 + gate_col] = prefix_log2; + } + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + if (chunk_is_full_2 != 0) { + mbarrier_wait(qk_raw_full_addr + (prep_stage) * 8, _phase_qk_raw_full); + } + if (prep_tid < 128) { + float total_log2 = smem_gt_prefix_all[stage_f32 + prep_tid]; + float _exp2_0 = approx_exp2(total_log2 - lower_bound * 1.4426950408889634f * 16.0f); + smem_restore_factor_all[stage_f32 + prep_tid] = _exp2_0; + } + if (prep_tid == 0) { + float _exp2_1 = approx_exp2(lower_bound * 1.4426950408889634f * 16.0f); + smem_restore_factor_all[stage_f32 + 128] = _exp2_1; + } + #pragma unroll 1 + for (int work_pass = 0; work_pass < 4; work_pass++) { + int work_item = work_pass * 128 + prep_tid; + int row_1 = work_item / 16; + int segment_1 = work_item % 16; + long long token_1 = bos_4 + (long long)(chunk_idx_3 * 32 + row_1); + int token_valid_1 = ((token_1 < eos_4) ? 1 : 0); + long long gmem_base = (token_1 * (long long)num_heads + (long long)head_idx_3) * 128 + (long long)(segment_1 * 8); + float q_raw_vec[8]; + float k_raw_vec[8]; + q_raw_vec[0] = 0.0f; + q_raw_vec[1] = 0.0f; + q_raw_vec[2] = 0.0f; + q_raw_vec[3] = 0.0f; + q_raw_vec[4] = 0.0f; + q_raw_vec[5] = 0.0f; + q_raw_vec[6] = 0.0f; + q_raw_vec[7] = 0.0f; + k_raw_vec[0] = 0.0f; + k_raw_vec[1] = 0.0f; + k_raw_vec[2] = 0.0f; + k_raw_vec[3] = 0.0f; + k_raw_vec[4] = 0.0f; + k_raw_vec[5] = 0.0f; + k_raw_vec[6] = 0.0f; + k_raw_vec[7] = 0.0f; + if (chunk_is_full_2 != 0) { + unsigned int packed[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed[0])), "=r"(*reinterpret_cast(&packed[(0) + 1])), "=r"(*reinterpret_cast(&packed[(0) + 2])), "=r"(*reinterpret_cast(&packed[(0) + 3])) + : "r"((smem_q_raw_prefetch_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_fp32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32[_pair * 2])[0]), "=f"((&packed_fp32[_pair * 2])[1]) + : "r"(packed[_pair + 0])); + } + #pragma unroll + for (int value_idx = 0; value_idx < 8; value_idx++) { + q_raw_vec[value_idx] = packed_fp32[value_idx]; + } + unsigned int packed_0[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_0[0])), "=r"(*reinterpret_cast(&packed_0[(0) + 1])), "=r"(*reinterpret_cast(&packed_0[(0) + 2])), "=r"(*reinterpret_cast(&packed_0[(0) + 3])) + : "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_0_fp32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_0_fp32[_pair * 2])[0]), "=f"((&packed_0_fp32[_pair * 2])[1]) + : "r"(packed_0[_pair + 0])); + } + #pragma unroll + for (int value_idx_1 = 0; value_idx_1 < 8; value_idx_1++) { + k_raw_vec[value_idx_1] = packed_0_fp32[value_idx_1]; + } + } else if (token_valid_1 != 0) { + { + const uint4* _vptr_0 = reinterpret_cast(q + gmem_base); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&q_raw_vec[0 + _blk * 8 + _pair * 2])[0]), "=f"((&q_raw_vec[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + } + { + const uint4* _vptr_1 = reinterpret_cast(k + gmem_base); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&k_raw_vec[0 + _blk * 8 + _pair * 2])[0]), "=f"((&k_raw_vec[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + } + } + float q_sum = 0.0f; + float k_sum = 0.0f; + for (int elem_in_segment = 0; elem_in_segment < 8; elem_in_segment++) { + float q_raw = q_raw_vec[elem_in_segment]; + float k_raw = k_raw_vec[elem_in_segment]; + float _fma_0 = __fmaf_rn(q_raw, q_raw, q_sum); + q_sum = _fma_0; + float _fma_1 = __fmaf_rn(k_raw, k_raw, k_sum); + k_sum = _fma_1; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 8); + q_sum += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 8); + k_sum += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 4); + q_sum += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 4); + k_sum += _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 2); + q_sum += _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 2); + k_sum += _shfl_xor_5; + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 1); + q_sum += _shfl_xor_6; + float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 1); + k_sum += _shfl_xor_7; + float _rsqrt_0 = rsqrtf(q_sum + 1e-06f); + float q_inv = _rsqrt_0; + float _rsqrt_1 = rsqrtf(k_sum + 1e-06f); + float k_inv = _rsqrt_1; + const float2 _scale2_2 = {q_inv, q_inv}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(q_raw_vec)[_ls], _scale2_2); + const float2 _scale2_3 = {k_inv, k_inv}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(k_raw_vec)[_ls], _scale2_3); + float qd_vec[8]; + float kd_vec[8]; + float ki_vec[8]; + for (int elem_in_segment_1 = 0; elem_in_segment_1 < 8; elem_in_segment_1++) { + int col = segment_1 * 8 + elem_in_segment_1; + float prefix = smem_gate_all[stage_f32 + row_1 * 128 + col]; + float common_log2 = lower_bound * 1.4426950408889634f * 16.0f; + float _exp2_2 = approx_exp2(prefix - common_log2); + float decay = _exp2_2; + qd_vec[elem_in_segment_1] = decay; + kd_vec[elem_in_segment_1] = decay; + ki_vec[elem_in_segment_1] = k_raw_vec[elem_in_segment_1] / decay; + } + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(qd_vec)[_ls], reinterpret_cast(q_raw_vec)[_ls]); + const float2 _scale2_4 = {scale, scale}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(qd_vec)[_ls], _scale2_4); + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(kd_vec)[_ls], reinterpret_cast(k_raw_vec)[_ls]); + unsigned int packed_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(qd_vec[_lp*2 + 0], qd_vec[_lp*2+1 + 0])); + packed_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word = 0; word < 4; word++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word * 4)), "r"(packed_1[word])); + } + unsigned int packed_0_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(kd_vec[_lp*2 + 0], kd_vec[_lp*2+1 + 0])); + packed_0_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_1 = 0; word_1 < 4; word_1++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_1 * 4)), "r"(packed_0_1[word_1])); + } + unsigned int packed_1_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ki_vec[_lp*2 + 0], ki_vec[_lp*2+1 + 0])); + packed_1_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_2 = 0; word_2 < 4; word_2++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_ki_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_2 * 4)), "r"(packed_1_1[word_2])); + } + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + int pair_row_base = prep_local_warp / 2 * 16; + int pair_col_base = prep_local_warp % 2 * 16; + unsigned int a_frag[4]; + unsigned int b_frag[4]; + float acc[8]; + if (pair_row_base >= pair_col_base) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[4]), "=f"(acc[(4) + 1]), "=f"(acc[(4) + 2]), "=f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + int row0 = pair_row_base + lane / 4; + int row1 = row0 + 8; + int col0 = pair_col_base + lane % 4 * 2; + float beta0 = smem_prep_beta_all[stage_f32 + row0]; + float beta1 = smem_prep_beta_all[stage_f32 + row1]; + float seed[8]; + seed[0] = 0.0f; + seed[1] = 0.0f; + seed[2] = 0.0f; + seed[3] = 0.0f; + seed[4] = 0.0f; + seed[5] = 0.0f; + seed[6] = 0.0f; + seed[7] = 0.0f; + if (row0 > col0) { + seed[0] = acc[0] * beta0; + } + if (row0 > col0 + 1) { + seed[1] = acc[1] * beta0; + } + if (row1 > col0) { + seed[2] = acc[2] * beta1; + } + if (row1 > col0 + 1) { + seed[3] = acc[3] * beta1; + } + if (row0 > col0 + 8) { + seed[4] = acc[4] * beta0; + } + if (row0 > col0 + 9) { + seed[5] = acc[5] * beta0; + } + if (row1 > col0 + 8) { + seed[6] = acc[6] * beta1; + } + if (row1 > col0 + 9) { + seed[7] = acc[7] * beta1; + } + unsigned int seed_packed[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(seed[_lp*2 + 0], seed[_lp*2+1 + 0])); + seed_packed[_lp] = *(uint32_t*)&_bf2; + } + int seed_lane_row = lane % 16; + int seed_lane_col = lane / 16 * 8; + int byte_off = (pair_row_base + seed_lane_row) * 128 + (pair_col_base + seed_lane_col) * 2; + int swizzled_off = byte_off ^ (byte_off >> 7 & 7) << 4; + int seed_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off; + uint32_t _stmatrix_addr_5 = static_cast((unsigned long long)seed_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_5), "r"(*reinterpret_cast(&seed_packed[0])), "r"(*reinterpret_cast(&seed_packed[1])), "r"(*reinterpret_cast(&seed_packed[2])), "r"(*reinterpret_cast(&seed_packed[3])) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[4]), "=f"(acc[(4) + 1]), "=f"(acc[(4) + 2]), "=f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + } else { + acc[0] = 0.0f; + acc[1] = 0.0f; + acc[2] = 0.0f; + acc[3] = 0.0f; + acc[4] = 0.0f; + acc[5] = 0.0f; + acc[6] = 0.0f; + acc[7] = 0.0f; + } + int row0_1 = pair_row_base + lane / 4; + int row1_1 = row0_1 + 8; + int col0_1 = pair_col_base + lane % 4 * 2; + float mqk[8]; + mqk[0] = 0.0f; + mqk[1] = 0.0f; + mqk[2] = 0.0f; + mqk[3] = 0.0f; + mqk[4] = 0.0f; + mqk[5] = 0.0f; + mqk[6] = 0.0f; + mqk[7] = 0.0f; + if (row0_1 >= col0_1) { + mqk[0] = acc[0]; + } + if (row0_1 >= col0_1 + 1) { + mqk[1] = acc[1]; + } + if (row1_1 >= col0_1) { + mqk[2] = acc[2]; + } + if (row1_1 >= col0_1 + 1) { + mqk[3] = acc[3]; + } + if (row0_1 >= col0_1 + 8) { + mqk[4] = acc[4]; + } + if (row0_1 >= col0_1 + 9) { + mqk[5] = acc[5]; + } + if (row1_1 >= col0_1 + 8) { + mqk[6] = acc[6]; + } + if (row1_1 >= col0_1 + 9) { + mqk[7] = acc[7]; + } + unsigned int mqk_packed[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(mqk[_lp*2 + 0], mqk[_lp*2+1 + 0])); + mqk_packed[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int publish_pair = 0; publish_pair < 2; publish_pair++) { + int publish_row = pair_col_base + publish_pair * 8 + (lane & 7); + int publish_col = 128 + pair_row_base + lane / 8 * 8; + uint32_t _stmatrix_addr_6 = static_cast((unsigned long long)(smem_final_trans_addr + prep_stage * 41984 + (unsigned int)(publish_col / 64 * 4096 + publish_row * 128 + publish_col % 64 * 2 ^ (publish_col / 64 * 4096 + publish_row * 128 + publish_col % 64 * 2 >> 7 & 7) << 4))); + asm volatile("stmatrix.sync.aligned.m8n8.x2.trans.shared.b16 [%0], {%1, %2};\n" + :: "r"(_stmatrix_addr_6), "r"(*reinterpret_cast(&mqk_packed[publish_pair * 2])), "r"(*reinterpret_cast(&mqk_packed[publish_pair * 2 + 1])) + : "memory"); + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + if (prep_tid < 128) { + float total_log2_1 = smem_gt_prefix_all[stage_f32 + prep_tid]; + float _exp2_3 = approx_exp2(total_log2_1); + smem_gt_all[stage_f32 + prep_tid] = _exp2_3; + } + if (prep_local_warp >= 2) { + int stage_f32_0 = prep_stage * 10496; + float restore_scale = smem_restore_factor_all[stage_f32_0 + 128]; + float restore_factor[8]; + int restore_segment = lane & 15; + #pragma unroll + for (int restore_elem = 0; restore_elem < 8; restore_elem++) { + int restore_col = restore_segment * 8 + restore_elem; + restore_factor[restore_elem] = smem_restore_factor_all[stage_f32_0 + restore_col]; + } + #pragma unroll 1 + for (int restore_pass = 0; restore_pass < 6; restore_pass++) { + int restore_row = 8 + (prep_local_warp - 2) * 12 + restore_pass * 2 + (lane >> 4); + float restore_qd_values[8]; + float restore_kd_values[8]; + float restore_ki_values[8]; + unsigned int packed_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_2[0])), "=r"(*reinterpret_cast(&packed_2[(0) + 1])), "=r"(*reinterpret_cast(&packed_2[(0) + 2])), "=r"(*reinterpret_cast(&packed_2[(0) + 3])) + : "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_fp32_1[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32_1[_pair * 2])[0]), "=f"((&packed_fp32_1[_pair * 2])[1]) + : "r"(packed_2[_pair + 0])); + } + #pragma unroll + for (int value_idx_2 = 0; value_idx_2 < 8; value_idx_2++) { + restore_qd_values[value_idx_2] = packed_fp32_1[value_idx_2]; + } + unsigned int packed_0_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_0_2[0])), "=r"(*reinterpret_cast(&packed_0_2[(0) + 1])), "=r"(*reinterpret_cast(&packed_0_2[(0) + 2])), "=r"(*reinterpret_cast(&packed_0_2[(0) + 3])) + : "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_0_fp32_1[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_0_fp32_1[_pair * 2])[0]), "=f"((&packed_0_fp32_1[_pair * 2])[1]) + : "r"(packed_0_2[_pair + 0])); + } + #pragma unroll + for (int value_idx_3 = 0; value_idx_3 < 8; value_idx_3++) { + restore_kd_values[value_idx_3] = packed_0_fp32_1[value_idx_3]; + } + unsigned int packed_1_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_1_2[0])), "=r"(*reinterpret_cast(&packed_1_2[(0) + 1])), "=r"(*reinterpret_cast(&packed_1_2[(0) + 2])), "=r"(*reinterpret_cast(&packed_1_2[(0) + 3])) + : "r"((smem_ki_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_1_fp32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_1_fp32[_pair * 2])[0]), "=f"((&packed_1_fp32[_pair * 2])[1]) + : "r"(packed_1_2[_pair + 0])); + } + #pragma unroll + for (int value_idx_4 = 0; value_idx_4 < 8; value_idx_4++) { + restore_ki_values[value_idx_4] = packed_1_fp32[value_idx_4]; + } + float restore_kr_values[8]; + #pragma unroll + for (int restore_elem_1 = 0; restore_elem_1 < 8; restore_elem_1++) { + restore_kr_values[restore_elem_1] = restore_ki_values[restore_elem_1] * restore_factor[restore_elem_1]; + } + const float2 _scale2_7 = {restore_scale, restore_scale}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_qd_values)[_ls], _scale2_7); + const float2 _scale2_8 = {restore_scale, restore_scale}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_kd_values)[_ls], _scale2_8); + unsigned int packed_2_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_qd_values[_lp*2 + 0], restore_qd_values[_lp*2+1 + 0])); + packed_2_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_3 = 0; word_3 < 4; word_3++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_3 * 4)), "r"(packed_2_1[word_3])); + } + unsigned int packed_3[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kd_values[_lp*2 + 0], restore_kd_values[_lp*2+1 + 0])); + packed_3[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_4 = 0; word_4 < 4; word_4++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_4 * 4)), "r"(packed_3[word_4])); + } + unsigned int packed_4[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kr_values[_lp*2 + 0], restore_kr_values[_lp*2+1 + 0])); + packed_4[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_5 = 0; word_5 < 4; word_5++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kr_trans_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_5 * 4)), "r"(packed_4[word_5])); + } + } + } + if (prep_local_warp == 0) { + int inverse_row = lane; + int diag_block = inverse_row / 8; + int lane_in_diag = lane & 7; + float inv_row[8]; + unsigned int packed_5[4]; + int byte_off_1 = inverse_row * 128 + diag_block * 8 * 2; + int swizzled_off_1 = byte_off_1 ^ (byte_off_1 >> 7 & 7) << 4; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_5[0])), "=r"(*reinterpret_cast(&packed_5[(0) + 1])), "=r"(*reinterpret_cast(&packed_5[(0) + 2])), "=r"(*reinterpret_cast(&packed_5[(0) + 3])) + : "r"(smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_1)); + float packed_fp32_2[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32_2[_pair * 2])[0]), "=f"((&packed_fp32_2[_pair * 2])[1]) + : "r"(packed_5[_pair + 0])); + } + #pragma unroll + for (int value_idx_5 = 0; value_idx_5 < 8; value_idx_5++) { + inv_row[value_idx_5] = packed_fp32_2[value_idx_5]; + } + #pragma unroll + for (int diag_elem = 0; diag_elem < 8; diag_elem++) { + if (lane_in_diag == diag_elem) { + inv_row[diag_elem] = 1.0f; + } + } + int diag_group_base = lane - lane_in_diag; + #pragma unroll + for (int src_row = 0; src_row < 7; src_row++) { + float row_scale = -inv_row[src_row]; + #pragma unroll + for (int prev_col = 0; prev_col < src_row; prev_col++) { + int pivot_lane = diag_group_base + src_row; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, inv_row[prev_col], pivot_lane); + float pivot = _shfl_0; + if (lane_in_diag > src_row) { + float _fma_2 = __fmaf_rn(row_scale, pivot, inv_row[prev_col]); + inv_row[prev_col] = _fma_2; + } + } + if (lane_in_diag > src_row) { + inv_row[src_row] = row_scale; + } + } + unsigned int packed_0_3[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(inv_row[_lp*2 + 0], inv_row[_lp*2+1 + 0])); + packed_0_3[_lp] = *(uint32_t*)&_bf2; + } + int byte_off_1_1 = inverse_row * 128 + diag_block * 8 * 2; + int swizzled_off_2 = byte_off_1_1 ^ (byte_off_1_1 >> 7 & 7) << 4; + #pragma unroll + for (int word_6 = 0; word_6 < 4; word_6++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"(smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_2 + (unsigned int)(word_6 * 4)), "r"(packed_0_3[word_6])); + } + } + if (prep_local_warp < 2) { + if (elect_sync()) { + mbarrier_arrive(prep_diag_ready_addr + (prep_stage) * 8); + } + mbarrier_wait(prep_diag_ready_addr + (prep_stage) * 8, _phase_prep_diag_ready); + } + if (prep_local_warp < 2) { + int lane_row = lane & 7; + int byte_off_2 = (prep_local_warp * 16 + 8 + lane_row) * 128 + (prep_local_warp * 16 + 8) * 2; + int swizzled_off_3 = byte_off_2 ^ (byte_off_2 >> 7 & 7) << 4; + int d_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_3; + int byte_off_0 = (prep_local_warp * 16 + 8 + lane_row) * 128 + prep_local_warp * 16 * 2; + int swizzled_off_1_1 = byte_off_0 ^ (byte_off_0 >> 7 & 7) << 4; + int c_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_1_1; + int byte_off_2_1 = (prep_local_warp * 16 + lane_row) * 128 + prep_local_warp * 16 * 2; + int swizzled_off_3_1 = byte_off_2_1 ^ (byte_off_2_1 >> 7 & 7) << 4; + int a_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_3_1; + unsigned int d_frag[2]; + unsigned int c_frag[1]; + float dc_acc[4]; + unsigned int dc_bf16[2]; + unsigned int inv_a_frag[1]; + float o_acc[4]; + unsigned int o_bf16[2]; + asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" + : "=r"(d_frag[0]) + : "r"(d_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" + : "=r"(d_frag[1]) + : "r"(d_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x1.trans.shared.b16 {%0}, [%1];\n" + : "=r"(c_frag[0]) + : "r"(c_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(dc_acc[0]), "=f"(dc_acc[1]), "=f"(dc_acc[2]), "=f"(dc_acc[3]) + : "r"(d_frag[0]), "r"(d_frag[1]), "r"(c_frag[0])); + const float2 _scale2_9 = {-1.0f, -1.0f}; + #pragma unroll + for (int _ls = 0; _ls < 2; _ls++) + mul_f32x2_inplace(&reinterpret_cast(dc_acc)[_ls], _scale2_9); + #pragma unroll + for (int _lp = 0; _lp < 2; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dc_acc[_lp*2 + 0], dc_acc[_lp*2+1 + 0])); + dc_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("ldmatrix.sync.aligned.m8n8.x1.trans.shared.b16 {%0}, [%1];\n" + : "=r"(inv_a_frag[0]) + : "r"(a_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(o_acc[0]), "=f"(o_acc[1]), "=f"(o_acc[2]), "=f"(o_acc[3]) + : "r"(dc_bf16[0]), "r"(dc_bf16[1]), "r"(inv_a_frag[0])); + #pragma unroll + for (int _lp = 0; _lp < 2; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(o_acc[_lp*2 + 0], o_acc[_lp*2+1 + 0])); + o_bf16[_lp] = *(uint32_t*)&_bf2; + } + int byte_off_4 = (prep_local_warp * 16 + 8 + lane_row) * 128 + prep_local_warp * 16 * 2; + int swizzled_off_5 = byte_off_4 ^ (byte_off_4 >> 7 & 7) << 4; + int o_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_5; + uint32_t _stmatrix_addr_10 = static_cast((unsigned long long)o_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x1.shared.b16 [%0], {%1};\n" + :: "r"(_stmatrix_addr_10), "r"(*reinterpret_cast(&o_bf16[0])) + : "memory"); + if (elect_sync()) { + mbarrier_arrive(prep_inv16_ready_addr + (prep_stage) * 8); + } + mbarrier_wait(prep_inv16_ready_addr + (prep_stage) * 8, _phase_prep_inv16_ready); + } + if (prep_local_warp == 0) { + int lane_row_1 = lane % 16; + int lane_col = lane / 16 * 8; + int byte_off_3 = (16 + lane_row_1) * 128 + (16 + lane_col) * 2; + int swizzled_off_4 = byte_off_3 ^ (byte_off_3 >> 7 & 7) << 4; + int d_addr_1 = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_4; + int byte_off_0_1 = (16 + lane_row_1) * 128 + lane_col * 2; + int swizzled_off_1_2 = byte_off_0_1 ^ (byte_off_0_1 >> 7 & 7) << 4; + int c_addr_1 = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_1_2; + int byte_off_2_2 = lane_row_1 * 128 + lane_col * 2; + int swizzled_off_3_2 = byte_off_2_2 ^ (byte_off_2_2 >> 7 & 7) << 4; + int a_addr_1 = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_3_2; + unsigned int d32_frag[4]; + unsigned int c32_frag[4]; + float dc32_acc[8]; + unsigned int dc32_bf16[4]; + unsigned int a32_frag[4]; + float o32_acc[8]; + unsigned int o32_bf16[4]; + unsigned int zero32_bf16[4]; + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(d32_frag[0]), "=r"(d32_frag[1]), "=r"(d32_frag[2]), "=r"(d32_frag[3]) + : "r"(d_addr_1) + : "memory"); + int d_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)((16 + lane_col) / 16 * 1024 + (16 + lane_row_1) * 32 + (16 + lane_col) % 16 * 2 ^ ((16 + lane_col) / 16 * 1024 + (16 + lane_row_1) * 32 + (16 + lane_col) % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_11 = static_cast((unsigned long long)d_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_11), "r"(*reinterpret_cast(&d32_frag[0])), "r"(*reinterpret_cast(&d32_frag[1])), "r"(*reinterpret_cast(&d32_frag[2])), "r"(*reinterpret_cast(&d32_frag[3])) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(c32_frag[0]), "=r"(c32_frag[1]), "=r"(c32_frag[2]), "=r"(c32_frag[3]) + : "r"(c_addr_1) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(dc32_acc[0]), "=f"(dc32_acc[1]), "=f"(dc32_acc[2]), "=f"(dc32_acc[3]) + : "r"(d32_frag[0]), "r"(d32_frag[1]), "r"(d32_frag[2]), "r"(d32_frag[3]), "r"(c32_frag[0]), "r"(c32_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(dc32_acc[4]), "=f"(dc32_acc[(4) + 1]), "=f"(dc32_acc[(4) + 2]), "=f"(dc32_acc[(4) + 3]) + : "r"(d32_frag[0]), "r"(d32_frag[1]), "r"(d32_frag[2]), "r"(d32_frag[3]), "r"(c32_frag[2]), "r"(c32_frag[(2) + 1])); + const float2 _scale2_12 = {-1.0f, -1.0f}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(dc32_acc)[_ls], _scale2_12); + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dc32_acc[_lp*2 + 0], dc32_acc[_lp*2+1 + 0])); + dc32_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a32_frag[0]), "=r"(a32_frag[1]), "=r"(a32_frag[2]), "=r"(a32_frag[3]) + : "r"(a_addr_1) + : "memory"); + int a_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)(lane_col / 16 * 1024 + lane_row_1 * 32 + lane_col % 16 * 2 ^ (lane_col / 16 * 1024 + lane_row_1 * 32 + lane_col % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_13 = static_cast((unsigned long long)a_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_13), "r"(*reinterpret_cast(&a32_frag[0])), "r"(*reinterpret_cast(&a32_frag[1])), "r"(*reinterpret_cast(&a32_frag[2])), "r"(*reinterpret_cast(&a32_frag[3])) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(o32_acc[0]), "=f"(o32_acc[1]), "=f"(o32_acc[2]), "=f"(o32_acc[3]) + : "r"(dc32_bf16[0]), "r"(dc32_bf16[1]), "r"(dc32_bf16[2]), "r"(dc32_bf16[3]), "r"(a32_frag[0]), "r"(a32_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(o32_acc[4]), "=f"(o32_acc[(4) + 1]), "=f"(o32_acc[(4) + 2]), "=f"(o32_acc[(4) + 3]) + : "r"(dc32_bf16[0]), "r"(dc32_bf16[1]), "r"(dc32_bf16[2]), "r"(dc32_bf16[3]), "r"(a32_frag[2]), "r"(a32_frag[(2) + 1])); + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(o32_acc[_lp*2 + 0], o32_acc[_lp*2+1 + 0])); + o32_bf16[_lp] = *(uint32_t*)&_bf2; + } + int o_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)(lane_col / 16 * 1024 + (16 + lane_row_1) * 32 + lane_col % 16 * 2 ^ (lane_col / 16 * 1024 + (16 + lane_row_1) * 32 + lane_col % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_14 = static_cast((unsigned long long)o_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_14), "r"(*reinterpret_cast(&o32_bf16[0])), "r"(*reinterpret_cast(&o32_bf16[1])), "r"(*reinterpret_cast(&o32_bf16[2])), "r"(*reinterpret_cast(&o32_bf16[3])) + : "memory"); + #pragma unroll + for (int zero_word = 0; zero_word < 4; zero_word++) { + zero32_bf16[zero_word] = 0; + } + int zero_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)((16 + lane_col) / 16 * 1024 + lane_row_1 * 32 + (16 + lane_col) % 16 * 2 ^ ((16 + lane_col) / 16 * 1024 + lane_row_1 * 32 + (16 + lane_col) % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_15 = static_cast((unsigned long long)zero_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_15), "r"(*reinterpret_cast(&zero32_bf16[0])), "r"(*reinterpret_cast(&zero32_bf16[1])), "r"(*reinterpret_cast(&zero32_bf16[2])), "r"(*reinterpret_cast(&zero32_bf16[3])) + : "memory"); + } else if (prep_local_warp == 1) { + int stage_f32_0_1 = prep_stage * 10496; + float restore_scale_1 = smem_restore_factor_all[stage_f32_0_1 + 128]; + float restore_factor_1[8]; + int restore_segment_1 = lane & 15; + #pragma unroll + for (int restore_elem_2 = 0; restore_elem_2 < 8; restore_elem_2++) { + int restore_col_1 = restore_segment_1 * 8 + restore_elem_2; + restore_factor_1[restore_elem_2] = smem_restore_factor_all[stage_f32_0_1 + restore_col_1]; + } + #pragma unroll 1 + for (int restore_pass_1 = 0; restore_pass_1 < 4; restore_pass_1++) { + int restore_row_1 = restore_pass_1 * 2 + (lane >> 4); + float restore_qd_values_1[8]; + float restore_kd_values_1[8]; + float restore_ki_values_1[8]; + unsigned int packed_6[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_6[0])), "=r"(*reinterpret_cast(&packed_6[(0) + 1])), "=r"(*reinterpret_cast(&packed_6[(0) + 2])), "=r"(*reinterpret_cast(&packed_6[(0) + 3])) + : "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_fp32_3[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32_3[_pair * 2])[0]), "=f"((&packed_fp32_3[_pair * 2])[1]) + : "r"(packed_6[_pair + 0])); + } + #pragma unroll + for (int value_idx_6 = 0; value_idx_6 < 8; value_idx_6++) { + restore_qd_values_1[value_idx_6] = packed_fp32_3[value_idx_6]; + } + unsigned int packed_0_4[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_0_4[0])), "=r"(*reinterpret_cast(&packed_0_4[(0) + 1])), "=r"(*reinterpret_cast(&packed_0_4[(0) + 2])), "=r"(*reinterpret_cast(&packed_0_4[(0) + 3])) + : "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_0_fp32_2[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_0_fp32_2[_pair * 2])[0]), "=f"((&packed_0_fp32_2[_pair * 2])[1]) + : "r"(packed_0_4[_pair + 0])); + } + #pragma unroll + for (int value_idx_7 = 0; value_idx_7 < 8; value_idx_7++) { + restore_kd_values_1[value_idx_7] = packed_0_fp32_2[value_idx_7]; + } + unsigned int packed_1_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_1_3[0])), "=r"(*reinterpret_cast(&packed_1_3[(0) + 1])), "=r"(*reinterpret_cast(&packed_1_3[(0) + 2])), "=r"(*reinterpret_cast(&packed_1_3[(0) + 3])) + : "r"((smem_ki_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_1_fp32_1[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_1_fp32_1[_pair * 2])[0]), "=f"((&packed_1_fp32_1[_pair * 2])[1]) + : "r"(packed_1_3[_pair + 0])); + } + #pragma unroll + for (int value_idx_8 = 0; value_idx_8 < 8; value_idx_8++) { + restore_ki_values_1[value_idx_8] = packed_1_fp32_1[value_idx_8]; + } + float restore_kr_values_1[8]; + #pragma unroll + for (int restore_elem_3 = 0; restore_elem_3 < 8; restore_elem_3++) { + restore_kr_values_1[restore_elem_3] = restore_ki_values_1[restore_elem_3] * restore_factor_1[restore_elem_3]; + } + const float2 _scale2_16 = {restore_scale_1, restore_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_qd_values_1)[_ls], _scale2_16); + const float2 _scale2_17 = {restore_scale_1, restore_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_kd_values_1)[_ls], _scale2_17); + unsigned int packed_2_2[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_qd_values_1[_lp*2 + 0], restore_qd_values_1[_lp*2+1 + 0])); + packed_2_2[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_7 = 0; word_7 < 4; word_7++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_7 * 4)), "r"(packed_2_2[word_7])); + } + unsigned int packed_3_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kd_values_1[_lp*2 + 0], restore_kd_values_1[_lp*2+1 + 0])); + packed_3_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_8 = 0; word_8 < 4; word_8++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_8 * 4)), "r"(packed_3_1[word_8])); + } + unsigned int packed_4_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kr_values_1[_lp*2 + 0], restore_kr_values_1[_lp*2+1 + 0])); + packed_4_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_9 = 0; word_9 < 4; word_9++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kr_trans_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_9 * 4)), "r"(packed_4_1[word_9])); + } + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + if (prep_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(qk_full_addr + (prep_stage) * 8); + } + } + for (int _advance = 0; _advance < 5; _advance++) { + prep_stage += 1; + if (prep_stage == 5) { prep_stage = 0; _phase_raw_inputs_free ^= 1; _phase_gate_raw_full ^= 1; _phase_smem_free ^= 1; _phase_qk_raw_full ^= 1; _phase_prep_diag_ready ^= 1; _phase_prep_inv16_ready ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + +// clang-format on diff --git a/csrc/kda/flashkda_bf16_fused_m64_binding.cu b/csrc/kda/flashkda_bf16_fused_m64_binding.cu new file mode 100644 index 00000000000..e3ef52f3445 --- /dev/null +++ b/csrc/kda/flashkda_bf16_fused_m64_binding.cu @@ -0,0 +1,100 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +#include "flashkda_binding_common.cuh" + +// The standalone exporter deliberately emits its own fixed-width typedefs so +// the frozen device TU has no host-library dependency. Isolate those names +// when the source is included into FlashInfer's C++ binding, whose CUDA/TVM +// headers already define the standard typedefs. +#define uint8_t flashkda_generated_uint8_t +#define uint16_t flashkda_generated_uint16_t +#define uint32_t flashkda_generated_uint32_t +#define uint64_t flashkda_generated_uint64_t +#define int32_t flashkda_generated_int32_t +#define int16_t flashkda_generated_int16_t +#include "flashkda_bf16_fused_m64.cu" +#undef uint8_t +#undef uint16_t +#undef uint32_t +#undef uint64_t +#undef int32_t +#undef int16_t + +namespace flashinfer { +namespace flash_kda { + +static_assert(THREADS == 1024); +static_assert(SMEM_TOTAL == 219136); + +void RunM64(TensorView q, TensorView k, TensorView v, TensorView g, TensorView beta, + TensorView beta_tma, TensorView A_log, TensorView dt_bias, TensorView cu_seqlens, + TensorView seq_order, TensorView initial_state, TensorView out, TensorView final_state, + TensorView descriptor_storage, int64_t prepare_descriptors, int64_t num_heads, + int64_t use_initial_state, int64_t store_final_state, double scale, double lower_bound, + int64_t cuda_stream) { + TVM_FFI_ICHECK(cuda_stream >= 0) << "cuda_stream must be a non-negative stream handle"; + TVM_FFI_ICHECK(q.device().device_type == kDLCUDA) << "q must be a CUDA tensor"; + const int32_t device_id = q.device().device_id; + ffi::CUDADeviceGuard device_guard(device_id); + CheckExactSm100a(device_id); + + const int64_t num_seqs = + CheckCommonInputs(q, k, v, g, beta, beta_tma, A_log, dt_bias, cu_seqlens, seq_order, + initial_state, out, final_state, descriptor_storage, prepare_descriptors, + num_heads, use_initial_state, store_final_state, scale, lower_bound); + TVM_FFI_ICHECK(num_seqs == 1 && num_heads == 64) + << "the M64 FlashKDA variant is specialized for fixed N=1, H=64; got " + "N=" + << num_seqs << ", H=" << num_heads; + + constexpr int32_t kSmemBytes = SMEM_TOTAL; + CheckDynamicSmemCapacity(device_id, kSmemBytes); + + CheckCuda(cudaFuncSetAttribute(kernel_flashkda_bf16_fused_m64, + cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes), + "cudaFuncSetAttribute(kernel_flashkda_bf16_fused_m64)"); + + const int64_t grid_x_i64 = 2 * num_seqs * num_heads; + TVM_FFI_ICHECK(grid_x_i64 > 0 && grid_x_i64 <= std::numeric_limits::max()) + << "M64 FlashKDA grid.x is out of range: " << grid_x_i64; + const dim3 grid(static_cast(grid_x_i64), 1, 1); + const dim3 block(THREADS, 1, 1); + const cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream)); + const TmaPointers tma = EncodeTmaPointers<64>(q, k, v, g, beta_tma, out, descriptor_storage, + prepare_descriptors, stream); + + kernel_flashkda_bf16_fused_m64<<>>( + reinterpret_cast<__nv_bfloat16*>(q.data_ptr()), tma.q, + reinterpret_cast<__nv_bfloat16*>(k.data_ptr()), tma.k, + reinterpret_cast<__nv_bfloat16*>(v.data_ptr()), tma.v, + reinterpret_cast<__nv_bfloat16*>(g.data_ptr()), tma.g, + reinterpret_cast<__nv_bfloat16*>(beta.data_ptr()), tma.beta, + reinterpret_cast(A_log.data_ptr()), reinterpret_cast(dt_bias.data_ptr()), + reinterpret_cast(cu_seqlens.data_ptr()), + reinterpret_cast(seq_order.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(initial_state.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), tma.out, + reinterpret_cast<__nv_bfloat16*>(final_state.data_ptr()), static_cast(num_heads), + static_cast(use_initial_state), static_cast(store_final_state), + static_cast(scale), static_cast(lower_bound)); + CheckCuda(cudaGetLastError(), "kernel_flashkda_bf16_fused_m64 launch"); +} + +} // namespace flash_kda +} // namespace flashinfer + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, flashinfer::flash_kda::RunM64); diff --git a/csrc/kda/flashkda_binding_common.cuh b/csrc/kda/flashkda_binding_common.cuh new file mode 100644 index 00000000000..c8a5b3afcd6 --- /dev/null +++ b/csrc/kda/flashkda_binding_common.cuh @@ -0,0 +1,442 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +#pragma once + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +namespace flashinfer { +namespace flash_kda { + +constexpr int64_t kHeadDim = 128; +constexpr size_t kTensorMapCount = 6; +constexpr size_t kTensorMapAlignment = 64; +static_assert(sizeof(CUtensorMap) == 128); +constexpr size_t kDescriptorStorageBytes = kTensorMapCount * sizeof(CUtensorMap); + +inline void CheckCuda(cudaError_t status, const char* operation) { + TVM_FFI_ICHECK(status == cudaSuccess) << operation << " failed: " << cudaGetErrorString(status); +} + +inline void CheckCudaTensor(const TensorView& tensor, const char* name, int32_t device_id) { + TVM_FFI_ICHECK(tensor.device().device_type == kDLCUDA) << name << " must be a CUDA tensor"; + TVM_FFI_ICHECK(tensor.device().device_id == device_id) + << name << " must be on CUDA device " << device_id << ", got " << tensor.device().device_id; + TVM_FFI_ICHECK(tensor.IsContiguous()) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& tensor, const char* name, DLDataType expected) { + const DLDataType actual = tensor.dtype(); + TVM_FFI_ICHECK(actual.code == expected.code && actual.bits == expected.bits && + actual.lanes == expected.lanes) + << name << " has wrong dtype: expected (code=" << int(expected.code) + << ", bits=" << int(expected.bits) << ", lanes=" << int(expected.lanes) + << "), got (code=" << int(actual.code) << ", bits=" << int(actual.bits) + << ", lanes=" << int(actual.lanes) << ")"; +} + +struct TensorByteRange { + uintptr_t begin; + uintptr_t end; +}; + +inline TensorByteRange GetTensorByteRange(const TensorView& tensor, const char* name) { + const DLDataType dtype = tensor.dtype(); + const uint64_t bits = static_cast(dtype.bits) * static_cast(dtype.lanes); + TVM_FFI_ICHECK(bits > 0 && bits % 8 == 0) << name << " has a non-byte-addressable dtype"; + const uint64_t bytes = static_cast(tensor.numel()) * (bits / 8); + const uintptr_t begin = reinterpret_cast(tensor.data_ptr()); + TVM_FFI_ICHECK(bytes <= std::numeric_limits::max() - begin) + << name << " byte range overflows uintptr_t"; + return {begin, begin + static_cast(bytes)}; +} + +inline void CheckNoOverlap(const TensorView& lhs, const char* lhs_name, const TensorView& rhs, + const char* rhs_name) { + const TensorByteRange lhs_range = GetTensorByteRange(lhs, lhs_name); + const TensorByteRange rhs_range = GetTensorByteRange(rhs, rhs_name); + const bool overlaps = lhs_range.begin < rhs_range.end && rhs_range.begin < lhs_range.end; + TVM_FFI_ICHECK(!overlaps) << lhs_name << " must not overlap " << rhs_name + << ": the frozen kernel uses __restrict__ pointers"; +} + +inline void CheckExactSm100a(int32_t device_id) { + int major = 0; + int minor = 0; + CheckCuda(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device_id), + "cudaDeviceGetAttribute(major)"); + CheckCuda(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device_id), + "cudaDeviceGetAttribute(minor)"); + TVM_FFI_ICHECK(major == 10 && minor == 0) + << "FlashKDA frozen kernels require exact compute capability 10.0 " + "(sm_100a), got " + << major << "." << minor; +} + +inline int64_t CheckCommonInputs(const TensorView& q, const TensorView& k, const TensorView& v, + const TensorView& g, const TensorView& beta, + const TensorView& beta_tma, const TensorView& A_log, + const TensorView& dt_bias, const TensorView& cu_seqlens, + const TensorView& seq_order, const TensorView& initial_state, + const TensorView& out, const TensorView& final_state, + const TensorView& descriptor_storage, int64_t prepare_descriptors, + int64_t num_heads, int64_t use_initial_state, + int64_t store_final_state, double scale, double lower_bound) { + TVM_FFI_ICHECK(prepare_descriptors == 0 || prepare_descriptors == 1) + << "prepare_descriptors must be 0 or 1, got " << prepare_descriptors; + TVM_FFI_ICHECK(num_heads > 0 && num_heads <= std::numeric_limits::max()) + << "num_heads must be in the positive int32 range, got " << num_heads; + TVM_FFI_ICHECK(use_initial_state == 0 || use_initial_state == 1) + << "use_initial_state must be 0 or 1, got " << use_initial_state; + TVM_FFI_ICHECK(store_final_state == 0 || store_final_state == 1) + << "store_final_state must be 0 or 1, got " << store_final_state; + TVM_FFI_ICHECK(std::isfinite(scale) && std::isfinite(static_cast(scale))) + << "scale must be finite and representable as float32, got " << scale; + TVM_FFI_ICHECK(std::isfinite(lower_bound) && lower_bound < 0.0 && + std::isfinite(static_cast(lower_bound))) + << "lower_bound must be finite, negative, and representable as " + "float32, got " + << lower_bound; + + const int32_t device_id = q.device().device_id; + CheckCudaTensor(q, "q", device_id); + CheckCudaTensor(k, "k", device_id); + CheckCudaTensor(v, "v", device_id); + CheckCudaTensor(g, "g", device_id); + CheckCudaTensor(beta, "beta", device_id); + CheckCudaTensor(beta_tma, "beta_tma", device_id); + CheckCudaTensor(A_log, "A_log", device_id); + CheckCudaTensor(dt_bias, "dt_bias", device_id); + CheckCudaTensor(cu_seqlens, "cu_seqlens", device_id); + CheckCudaTensor(seq_order, "seq_order", device_id); + CheckCudaTensor(initial_state, "initial_state", device_id); + CheckCudaTensor(out, "out", device_id); + CheckCudaTensor(final_state, "final_state", device_id); + CheckCudaTensor(descriptor_storage, "descriptor_storage", device_id); + + CheckDtype(q, "q", dl_bfloat16); + CheckDtype(k, "k", dl_bfloat16); + CheckDtype(v, "v", dl_bfloat16); + CheckDtype(g, "g", dl_bfloat16); + CheckDtype(beta, "beta", dl_bfloat16); + CheckDtype(beta_tma, "beta_tma", dl_bfloat16); + CheckDtype(A_log, "A_log", dl_float32); + CheckDtype(dt_bias, "dt_bias", dl_float32); + CheckDtype(cu_seqlens, "cu_seqlens", dl_int64); + CheckDtype(seq_order, "seq_order", dl_int32); + CheckDtype(initial_state, "initial_state", dl_bfloat16); + CheckDtype(out, "out", dl_bfloat16); + CheckDtype(final_state, "final_state", dl_bfloat16); + CheckDtype(descriptor_storage, "descriptor_storage", dl_uint8); + + TVM_FFI_ICHECK(descriptor_storage.numel() >= static_cast(kDescriptorStorageBytes)) + << "descriptor_storage must contain at least " << kDescriptorStorageBytes << " bytes"; + TVM_FFI_ICHECK(reinterpret_cast(descriptor_storage.data_ptr()) % kTensorMapAlignment == + 0) + << "descriptor_storage must be aligned to " << kTensorMapAlignment << " bytes"; + + TVM_FFI_ICHECK(q.ndim() >= 3) << "q must have trailing [H, 128] dimensions"; + TVM_FFI_ICHECK(q.size(q.ndim() - 1) == kHeadDim && q.size(q.ndim() - 2) == num_heads) + << "q must have trailing shape [" << num_heads << ", 128]"; + const int64_t token_count = q.numel() / (num_heads * kHeadDim); + TVM_FFI_ICHECK(token_count > 0) << "q must contain at least one token"; + + for (const auto& named : {std::pair(&k, "k"), + std::pair(&v, "v"), + std::pair(&g, "g"), + std::pair(&out, "out")}) { + const TensorView& tensor = *named.first; + TVM_FFI_ICHECK(tensor.ndim() >= 3 && tensor.size(tensor.ndim() - 1) == kHeadDim && + tensor.size(tensor.ndim() - 2) == num_heads && tensor.numel() == q.numel()) + << named.second << " must match q's flattened [tokens, H, 128] shape"; + } + + TVM_FFI_ICHECK(beta.ndim() >= 2 && beta.size(beta.ndim() - 1) == num_heads && + beta.numel() == token_count * num_heads) + << "beta must match flattened [tokens, H]"; + const int64_t beta_tma_heads = std::max(num_heads, 8); + TVM_FFI_ICHECK(beta_tma.ndim() >= 2 && beta_tma.size(beta_tma.ndim() - 1) == beta_tma_heads && + beta_tma.numel() % beta_tma_heads == 0 && + beta_tma.numel() / beta_tma_heads >= std::max(token_count, 32)) + << "beta_tma must have at least [max(tokens, 32), max(H, 8)] " + "storage"; + TVM_FFI_ICHECK(A_log.numel() == num_heads) << "A_log must contain H elements"; + TVM_FFI_ICHECK(dt_bias.numel() == num_heads * kHeadDim) + << "dt_bias must contain H * 128 elements"; + TVM_FFI_ICHECK(cu_seqlens.ndim() == 1 && cu_seqlens.numel() >= 2) + << "cu_seqlens must be a one-dimensional tensor with at least two " + "elements"; + + const int64_t num_seqs = cu_seqlens.numel() - 1; + TVM_FFI_ICHECK(seq_order.ndim() == 1 && seq_order.numel() == num_seqs) + << "seq_order must contain one int32 entry per sequence"; + const int64_t state_numel = num_seqs * num_heads * kHeadDim * kHeadDim; + if (use_initial_state != 0) { + TVM_FFI_ICHECK(initial_state.numel() == state_numel) + << "initial_state must have flattened [N, H, 128, 128] size"; + } + if (store_final_state != 0) { + TVM_FFI_ICHECK(final_state.numel() == state_numel) + << "final_state must have flattened [N, H, 128, 128] size"; + } + CheckNoOverlap(out, "out", q, "q"); + CheckNoOverlap(out, "out", k, "k"); + CheckNoOverlap(out, "out", v, "v"); + CheckNoOverlap(out, "out", g, "g"); + CheckNoOverlap(out, "out", beta, "beta"); + CheckNoOverlap(out, "out", beta_tma, "beta_tma"); + CheckNoOverlap(out, "out", A_log, "A_log"); + CheckNoOverlap(out, "out", dt_bias, "dt_bias"); + CheckNoOverlap(out, "out", cu_seqlens, "cu_seqlens"); + CheckNoOverlap(out, "out", seq_order, "seq_order"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", q, "q"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", k, "k"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", v, "v"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", g, "g"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", beta, "beta"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", beta_tma, "beta_tma"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", A_log, "A_log"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", dt_bias, "dt_bias"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", cu_seqlens, "cu_seqlens"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", seq_order, "seq_order"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", out, "out"); + if (use_initial_state != 0) { + CheckNoOverlap(out, "out", initial_state, "initial_state"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", initial_state, "initial_state"); + } + if (store_final_state != 0) { + CheckNoOverlap(out, "out", final_state, "final_state"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", final_state, "final_state"); + CheckNoOverlap(final_state, "final_state", q, "q"); + CheckNoOverlap(final_state, "final_state", k, "k"); + CheckNoOverlap(final_state, "final_state", v, "v"); + CheckNoOverlap(final_state, "final_state", g, "g"); + CheckNoOverlap(final_state, "final_state", beta, "beta"); + CheckNoOverlap(final_state, "final_state", beta_tma, "beta_tma"); + CheckNoOverlap(final_state, "final_state", A_log, "A_log"); + CheckNoOverlap(final_state, "final_state", dt_bias, "dt_bias"); + CheckNoOverlap(final_state, "final_state", cu_seqlens, "cu_seqlens"); + CheckNoOverlap(final_state, "final_state", seq_order, "seq_order"); + if (use_initial_state != 0) { + CheckNoOverlap(initial_state, "initial_state", final_state, "final_state"); + } + } + return num_seqs; +} + +inline CUtensorMap EncodeQkTma(const TensorView& tensor, const char* name) { + TVM_FFI_ICHECK(tensor.ndim() >= 2) << name << " must have at least two dimensions"; + const int64_t d1 = tensor.size(tensor.ndim() - 1); + const int64_t d2 = tensor.size(tensor.ndim() - 2); + TVM_FFI_ICHECK(d1 > 0 && d2 > 0 && d1 % 64 == 0) << name << " has invalid trailing dimensions"; + const int64_t outer2 = tensor.numel() / (d1 * d2); + uint64_t global_dim[4] = {64, static_cast(outer2), static_cast(d2), + static_cast(d1 / 64)}; + TVM_FFI_ICHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] >= 1 && global_dim[3] >= 2) + << name << " cannot encode the (64, 32, 1, 2) TMA box"; + uint64_t global_strides[3] = {static_cast(d2 * d1 * sizeof(__nv_bfloat16)), + static_cast(d1 * sizeof(__nv_bfloat16)), + static_cast(64 * sizeof(__nv_bfloat16))}; + uint32_t box_dim[4] = {64, 32, 1, 2}; + uint32_t elem_strides[4] = {1, 1, 1, 1}; + CUtensorMap tensor_map{}; + const CUresult result = + cuTensorMapEncodeTiled(&tensor_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, tensor.data_ptr(), + global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_ICHECK(result == CUDA_SUCCESS) + << "cuTensorMapEncodeTiled failed for " << name << " with CUresult=" << int(result); + return tensor_map; +} + +template +inline CUtensorMap EncodeValueTma(const TensorView& tensor) { + static_assert(ValueRows == 64 || ValueRows == 128); + const int64_t d1 = tensor.size(tensor.ndim() - 1); + const int64_t d2 = tensor.size(tensor.ndim() - 2); + const int64_t outer2 = tensor.numel() / (d1 * d2); + uint64_t global_dim[3] = {static_cast(d1), static_cast(d2), + static_cast(outer2)}; + TVM_FFI_ICHECK(global_dim[0] >= ValueRows && global_dim[1] >= 1 && global_dim[2] > 0) + << "v cannot encode the (" << ValueRows << ", 1, 32) TMA box"; + uint64_t global_strides[2] = {static_cast(d1 * sizeof(__nv_bfloat16)), + static_cast(d1 * d2 * sizeof(__nv_bfloat16))}; + uint32_t box_dim[3] = {ValueRows, 1, 32}; + uint32_t elem_strides[3] = {1, 1, 1}; + CUtensorMap tensor_map{}; + constexpr CUtensorMapSwizzle swizzle = + ValueRows == 64 ? CU_TENSOR_MAP_SWIZZLE_128B : CU_TENSOR_MAP_SWIZZLE_NONE; + const CUresult result = cuTensorMapEncodeTiled( + &tensor_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, tensor.data_ptr(), global_dim, + global_strides, box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_ICHECK(result == CUDA_SUCCESS) + << "cuTensorMapEncodeTiled failed for v with CUresult=" << int(result); + return tensor_map; +} + +inline CUtensorMap EncodeGateTma(const TensorView& tensor) { + const int64_t d1 = tensor.size(tensor.ndim() - 1); + const int64_t d2 = tensor.size(tensor.ndim() - 2); + const int64_t outer2 = tensor.numel() / (d1 * d2); + uint64_t global_dim[3] = {static_cast(d1), static_cast(d2), + static_cast(outer2)}; + TVM_FFI_ICHECK(global_dim[0] >= 128 && global_dim[1] >= 1 && global_dim[2] > 0) + << "g cannot encode the (128, 1, 32) TMA box"; + uint64_t global_strides[2] = {static_cast(d1 * sizeof(__nv_bfloat16)), + static_cast(d1 * d2 * sizeof(__nv_bfloat16))}; + uint32_t box_dim[3] = {128, 1, 32}; + uint32_t elem_strides[3] = {1, 1, 1}; + CUtensorMap tensor_map{}; + const CUresult result = + cuTensorMapEncodeTiled(&tensor_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, tensor.data_ptr(), + global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_ICHECK(result == CUDA_SUCCESS) + << "cuTensorMapEncodeTiled failed for g with CUresult=" << int(result); + return tensor_map; +} + +inline CUtensorMap EncodeBetaTma(const TensorView& tensor) { + const int64_t d1 = tensor.size(tensor.ndim() - 1); + const int64_t outer1 = tensor.numel() / d1; + uint64_t global_dim[2] = {static_cast(d1), static_cast(outer1)}; + TVM_FFI_ICHECK(global_dim[0] >= 8 && global_dim[1] >= 32) + << "beta_tma cannot encode the (8, 32) TMA box"; + uint64_t global_strides[1] = {static_cast(d1 * sizeof(__nv_bfloat16))}; + uint32_t box_dim[2] = {8, 32}; + uint32_t elem_strides[2] = {1, 1}; + CUtensorMap tensor_map{}; + const CUresult result = + cuTensorMapEncodeTiled(&tensor_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, tensor.data_ptr(), + global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_ICHECK(result == CUDA_SUCCESS) + << "cuTensorMapEncodeTiled failed for beta_tma with CUresult=" << int(result); + return tensor_map; +} + +template +inline CUtensorMap EncodeOutputTma(const TensorView& tensor) { + static_assert(ValueRows == 64 || ValueRows == 128); + const int64_t d1 = tensor.size(tensor.ndim() - 1); + const int64_t d2 = tensor.size(tensor.ndim() - 2); + const int64_t outer2 = tensor.numel() / (d1 * d2); + TVM_FFI_ICHECK(d1 > 0 && d2 > 0 && d1 % 64 == 0) << "out has invalid trailing dimensions"; + uint64_t global_dim[4] = {64, static_cast(outer2), static_cast(d2), + static_cast(d1 / 64)}; + constexpr uint32_t value_splits = ValueRows / 64; + TVM_FFI_ICHECK(global_dim[0] >= 64 && global_dim[1] > 0 && global_dim[2] >= 1 && + global_dim[3] >= value_splits) + << "out cannot encode the (64, 32, 1, " << value_splits << ") TMA box"; + uint64_t global_strides[3] = {static_cast(d2 * d1 * sizeof(__nv_bfloat16)), + static_cast(d1 * sizeof(__nv_bfloat16)), + static_cast(64 * sizeof(__nv_bfloat16))}; + uint32_t box_dim[4] = {64, 32, 1, value_splits}; + uint32_t elem_strides[4] = {1, 1, 1, 1}; + CUtensorMap tensor_map{}; + const CUresult result = + cuTensorMapEncodeTiled(&tensor_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, tensor.data_ptr(), + global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_ICHECK(result == CUDA_SUCCESS) + << "cuTensorMapEncodeTiled failed for out with CUresult=" << int(result); + return tensor_map; +} + +struct TmaPointers { + void* q; + void* k; + void* v; + void* g; + void* beta; + void* out; +}; + +struct TensorMapWords { + static constexpr size_t kWordCount = kDescriptorStorageBytes / sizeof(uint64_t); + uint64_t words[kWordCount]; +}; + +static __global__ void PublishTensorMaps(uint64_t* destination, TensorMapWords source) { + const uint32_t index = threadIdx.x; + if (index < TensorMapWords::kWordCount) { + destination[index] = source.words[index]; + } +} + +template +inline TmaPointers EncodeTmaPointers(const TensorView& q, const TensorView& k, const TensorView& v, + const TensorView& g, const TensorView& beta_tma, + const TensorView& out, const TensorView& descriptor_storage, + int64_t prepare_descriptors, cudaStream_t stream) { + if (prepare_descriptors != 0) { + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + CheckCuda(cudaStreamIsCapturing(stream, &capture_status), + "cudaStreamIsCapturing(TMA descriptor preparation)"); + TVM_FFI_ICHECK(capture_status == cudaStreamCaptureStatusNone) + << "prepare_descriptors must be 0 during CUDA graph capture; warm " + "this exact workspace and tensor signature before capture"; + + const std::array host_maps = { + EncodeQkTma(q, "q"), EncodeQkTma(k, "k"), EncodeValueTma(v), + EncodeGateTma(g), EncodeBetaTma(beta_tma), EncodeOutputTma(out), + }; + static_assert(sizeof(host_maps) == kDescriptorStorageBytes); + TensorMapWords words{}; + std::memcpy(words.words, host_maps.data(), sizeof(host_maps)); + PublishTensorMaps<<<1, 128, 0, stream>>>( + reinterpret_cast(descriptor_storage.data_ptr()), words); + CheckCuda(cudaGetLastError(), "PublishTensorMaps launch"); + } + + auto* bytes = static_cast(descriptor_storage.data_ptr()); + constexpr size_t stride = sizeof(CUtensorMap); + return { + bytes + 0 * stride, bytes + 1 * stride, bytes + 2 * stride, + bytes + 3 * stride, bytes + 4 * stride, bytes + 5 * stride, + }; +} + +inline void CheckDynamicSmemCapacity(int32_t device_id, int32_t smem_bytes) { + int max_optin = 0; + CheckCuda(cudaDeviceGetAttribute(&max_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device_id), + "cudaDeviceGetAttribute(max dynamic shared memory)"); + TVM_FFI_ICHECK(max_optin >= smem_bytes) + << "device exposes only " << max_optin << " bytes of opt-in shared memory; FlashKDA requires " + << smem_bytes; +} + +} // namespace flash_kda +} // namespace flashinfer diff --git a/docs/api/kda_decode.rst b/docs/api/kda_decode.rst index e678b6963ef..ec16c55cd54 100644 --- a/docs/api/kda_decode.rst +++ b/docs/api/kda_decode.rst @@ -3,12 +3,103 @@ flashinfer.kda_decode ===================== -Key-Driven Attention (KDA) decode API. The CuTe-DSL kernel lives under -``flashinfer.kda_kernels``; this module is the public entry point. +Recurrent Key-Driven Attention (KDA) API. Single-token decode, speculative +decode, GQA, state-pool indexing, and checkpoint modes use the CuTe-DSL +backend under ``flashinfer.kda_kernels``. A strict ordinary multi-token +prefill subset is dispatched to frozen FlashKDA-compatible SM100a kernels. .. currentmodule:: flashinfer.kda_decode .. autosummary:: :toctree: ../generated + RecurrentKDAPrefillWorkspace recurrent_kda + +Optimized B200 prefill subset +----------------------------- + +``recurrent_kda`` uses the frozen prefill backend only when every condition +below holds: + +* the device has compute capability 10.0; +* input is ordinary multi-token prefill: fixed ``T > 1``, or packed input + whose total token count is greater than its number of sequences; +* Q, K, V, and G are contiguous BF16 ``[B,T,H,128]`` tensors with one shared + head count, and beta is contiguous BF16 ``[B,T,H]``; +* ``A_log`` is contiguous FP32 ``[H]`` and ``dt_bias`` is contiguous FP32 + ``[H,128]`` or flattened ``[H*128]``; +* ``use_qk_l2norm_in_kernel=True``, ``use_gate_in_kernel=True``, + ``beta_is_logit=True``, and ``lower_bound`` is a finite negative value; +* speculative decode, GQA, state indices, committed-state sources, and + accepted-token/checkpoint features are not enabled. + +Calls outside that subset retain the existing CuTe-DSL path. In particular, +T=1 decode and speculative decode are not rerouted. + +Fixed input omits ``cu_seqlens``. Packed input has ``B=1`` and accepts a +contiguous CUDA int32 or int64 ``cu_seqlens``. The frozen binding consumes +int64 offsets; pass int64 directly for CUDA graph capture to avoid an +in-capture conversion allocation. Offset values are a caller contract: +``cu_seqlens[0] == 0``, entries are strictly increasing (every sequence is +non-empty), and ``cu_seqlens[-1] == total_tokens``. FlashInfer does not +synchronize the device to inspect these values; invalid offsets may cause +out-of-bounds device access. + +Packed scheduling +----------------- + +Packed prefill optionally accepts ``seq_order``, a contiguous CUDA int32 +tensor with one entry per sequence. It is a caller contract that this tensor +is a permutation of ``[0, N)``. Ordering sequences by decreasing length +reduces the final partial wave. FlashInfer validates dtype, device, rank, and +size without synchronizing the device to inspect permutation values. + +When ``seq_order=None``, a cached identity order is used. Fixed ``B=1,H=64`` +selects the two-CTA M64 value-split kernel; every packed input and every other +head count selects M128. + +State and graph semantics +------------------------- + +The BF16 state layout remains ``[N,H,V,K]`` and an explicitly supplied +``initial_state`` is still updated in place, even when +``output_final_state=False``. The frozen CUDA ABI marks initial and final +state pointers ``__restrict__``; FlashInfer therefore launches into separate +per-stream scratch storage and enqueues a same-stream copy back into +``initial_state``. If no initial state is supplied, a final state is allocated +only when ``output_final_state=True``. + +The frozen kernel uses restricted output storage. A preallocated ``output`` +must not overlap Q, K, V, G, beta, or ``initial_state``. + +Eager calls without ``prefill_workspace`` use an internal serialized workspace +for the current CUDA stream. This default workspace is eager-only and cannot +be used during CUDA graph capture. + +CUDA graph capture requires a caller-owned +``RecurrentKDAPrefillWorkspace(device)`` and a preallocated ``output``. The +workspace owns stable state scratch, beta padding, and separate 768-byte M64 +and M128 TMA descriptor blocks. It binds to the device and CUDA stream of its +first ``recurrent_kda`` call. Warm it eagerly on the intended capture stream +with the exact Q, K, V, G, beta, and output tensors, then synchronize that +stream before capture. Packed graphs must also pass preallocated int64 +``cu_seqlens`` and int32 ``seq_order``. The warm call prepares descriptors; +capture accepts only the exact warmed pointer, shape, stride, and dtype +signature and performs no descriptor preparation. + +The workspace must outlive its graph and every replay. Use one distinct +workspace for each captured ``recurrent_kda`` invocation, including two KDA +invocations in the same graph. Once a workspace participates in capture, any +later Python use through ``recurrent_kda``—eager or another capture—is +rejected. ``graph.replay()`` does not re-enter Python and remains valid. +Sequential replay launches may be issued while a different PyTorch stream is +current, but the caller must provide normal stream ordering. The Python stream +binding applies to eager warmup and capture calls, which must use the same +stream. + +When an explicit workspace is used with ``initial_state=None`` and +``output_final_state=True``, the returned final state is workspace-owned +stable scratch. Otherwise an explicitly supplied ``initial_state`` keeps its +usual in-place copy-back semantics. The small-head ``H < 8`` path captures +the beta copy into workspace-owned padded storage before the frozen launch. diff --git a/flashinfer/__init__.py b/flashinfer/__init__.py index 41cbed50b4b..3fec69a5495 100644 --- a/flashinfer/__init__.py +++ b/flashinfer/__init__.py @@ -124,6 +124,9 @@ from .grouped_mm import grouped_mm_fp8 as grouped_mm_fp8 from .grouped_mm import grouped_mm_mxfp8 as grouped_mm_mxfp8 from .grouped_mm import grouped_mm_fp4 as grouped_mm_fp4 +from .kda_decode import ( + RecurrentKDAPrefillWorkspace as RecurrentKDAPrefillWorkspace, +) from .kda_decode import recurrent_kda as recurrent_kda from .mla import BatchMLAPagedAttentionWrapper as BatchMLAPagedAttentionWrapper from . import mhc as mhc diff --git a/flashinfer/aot.py b/flashinfer/aot.py index 04e71dd11af..0349646288a 100644 --- a/flashinfer/aot.py +++ b/flashinfer/aot.py @@ -60,6 +60,10 @@ ) from .jit.fp4_kv_dequantization import gen_fp4_kv_dequantization_module from .jit.fp4_kv_quantization import gen_fp4_kv_quantization_module +from .jit.flash_kda import ( + gen_flash_kda_m64_module, + gen_flash_kda_m128_module, +) from .jit.nvfp4_attention_sm120 import gen_nvfp4_attention_sm120_module from .jit.fp8_quantization import gen_mxfp8_quantization_sm100_module from .jit.fused_moe import ( @@ -498,6 +502,7 @@ def gen_all_modules( has_sm80 = sm_capabilities.get("sm80", False) has_sm90 = sm_capabilities.get("sm90", False) has_sm100 = sm_capabilities.get("sm100", False) + has_sm100a_exact = sm_capabilities.get("sm100a_exact", False) has_sm100f = sm_capabilities.get("sm100f", False) has_sm103 = sm_capabilities.get("sm103", False) has_sm110 = sm_capabilities.get("sm110", False) @@ -521,6 +526,16 @@ def gen_all_modules( ) if has_sm120 or has_sm121: jit_specs.append(gen_nvfp4_attention_sm120_module()) + if has_sm100a_exact: + # Frozen FlashKDA sources use the SM100a-only tcgen05/TMEM surface. + # Do not package them for SM100f, SM103, or later architectures until + # those exact cubins have been independently validated. + jit_specs.extend( + [ + gen_flash_kda_m64_module(), + gen_flash_kda_m128_module(), + ] + ) if add_act: for act_name in act_func_def_str: @@ -962,6 +977,8 @@ def has_sm(compute: str, version: str) -> bool: "sm80": has_any_sm8x and get_cuda_version() >= Version("11.0"), "sm90": has_sm("compute_90", "12.3"), "sm100": has_sm("compute_100", "12.8"), + "sm100a_exact": (10, "0a") in compilation_context.TARGET_CUDA_ARCHS + and get_cuda_version() >= Version("12.8"), "sm100f": has_sm("compute_100", "12.9"), "sm103": has_sm("compute_103", "12.9"), "sm110": has_sm("compute_110", "13.0"), diff --git a/flashinfer/jit/__init__.py b/flashinfer/jit/__init__.py index 75192911c78..cbfb10f3945 100644 --- a/flashinfer/jit/__init__.py +++ b/flashinfer/jit/__init__.py @@ -100,6 +100,21 @@ from .fp4_kv_quantization import ( gen_fp4_kv_quantization_module as gen_fp4_kv_quantization_module, ) +from .flash_kda import ( + gen_flash_kda_m64_module as gen_flash_kda_m64_module, +) +from .flash_kda import ( + gen_flash_kda_m128_module as gen_flash_kda_m128_module, +) +from .flash_kda import ( + get_flash_kda_prefill_module as get_flash_kda_prefill_module, +) +from .flash_kda import ( + load_flash_kda_m64_module as load_flash_kda_m64_module, +) +from .flash_kda import ( + load_flash_kda_m128_module as load_flash_kda_m128_module, +) from .nvfp4_attention_sm120 import ( gen_nvfp4_attention_sm120_module as gen_nvfp4_attention_sm120_module, ) diff --git a/flashinfer/jit/flash_kda.py b/flashinfer/jit/flash_kda.py new file mode 100644 index 00000000000..13d9d53a21b --- /dev/null +++ b/flashinfer/jit/flash_kda.py @@ -0,0 +1,148 @@ +""" +Copyright (c) 2026 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import functools +from pathlib import Path +from typing import Literal + +from . import env as jit_env +from .core import JitSpec, gen_jit_spec, logger, sm100a_nvcc_flags + +FlashKDAVariant = Literal["m64", "m128"] + + +def _get_flash_kda_csrc_dir() -> Path: + """Locate frozen FlashKDA sources in installed and source checkouts.""" + + installed = jit_env.FLASHINFER_CSRC_DIR / "kda" + if installed.exists(): + return installed + + checkout = Path(__file__).resolve().parents[2] / "csrc" / "kda" + if checkout.exists(): + return checkout + + raise FileNotFoundError( + "FlashKDA CUDA sources were not found. Checked:\n" + f" - {installed}\n" + f" - {checkout}" + ) + + +def _get_flash_kda_include_dir() -> Path: + """Locate FlashInfer headers in installed and source checkouts.""" + + if jit_env.FLASHINFER_INCLUDE_DIR.exists(): + return jit_env.FLASHINFER_INCLUDE_DIR + checkout = Path(__file__).resolve().parents[2] / "include" + if checkout.exists(): + return checkout + raise FileNotFoundError( + "FlashInfer headers were not found. Checked:\n" + f" - {jit_env.FLASHINFER_INCLUDE_DIR}\n" + f" - {checkout}" + ) + + +def get_flash_kda_uri(variant: FlashKDAVariant) -> str: + """Return the stable JIT/AOT cache key for one physical schedule.""" + + if variant not in ("m64", "m128"): + raise ValueError(f"unsupported FlashKDA variant: {variant}") + return f"flash_kda_bf16_fused_{variant}_sm100a" + + +@functools.cache +def gen_flash_kda_module(variant: FlashKDAVariant) -> JitSpec: + """Generate one exact-sm_100a FlashKDA JIT module. + + Each physical schedule is compiled in its own translation unit because the + checked-in frozen sources intentionally retain generated helper names and + macros. ``gen_jit_spec`` supplies FlashInfer's standard ``-use_fast_math`` + flag; the explicit architecture flags below emit only an sm_100a cubin. + """ + + csrc_dir = _get_flash_kda_csrc_dir() + include_dir = _get_flash_kda_include_dir() + uri = get_flash_kda_uri(variant) + binding = csrc_dir / f"flashkda_bf16_fused_{variant}_binding.cu" + if not binding.exists(): + raise FileNotFoundError(f"FlashKDA binding source not found: {binding}") + + spec = gen_jit_spec( + name=uri, + sources=[binding], + extra_cuda_cflags=sm100a_nvcc_flags, + extra_include_paths=[ + csrc_dir, + csrc_dir.parent, + include_dir, + ], + ) + logger.info(f"Generated FlashKDA {variant} JIT spec: {spec.name}") + return spec + + +def gen_flash_kda_m64_module() -> JitSpec: + """Generate the fixed N=1, H=64 two-CTA M64 module.""" + + return gen_flash_kda_module("m64") + + +def gen_flash_kda_m128_module() -> JitSpec: + """Generate the general packed/fixed M128 module.""" + + return gen_flash_kda_module("m128") + + +@functools.cache +def load_flash_kda_module(variant: FlashKDAVariant): + """Build or load one physical FlashKDA module.""" + + module = gen_flash_kda_module(variant).build_and_load() + logger.info(f"Loaded FlashKDA {variant} module") + return module + + +def load_flash_kda_m64_module(): + """Load the fixed N=1, H=64 two-CTA M64 module.""" + + return load_flash_kda_module("m64") + + +def load_flash_kda_m128_module(): + """Load the general packed/fixed M128 module.""" + + return load_flash_kda_module("m128") + + +def get_flash_kda_prefill_module(variant: FlashKDAVariant): + """Return the loaded module used by the recurrent-KDA prefill dispatcher.""" + + return load_flash_kda_module(variant) + + +__all__ = [ + "FlashKDAVariant", + "gen_flash_kda_m64_module", + "gen_flash_kda_m128_module", + "gen_flash_kda_module", + "get_flash_kda_prefill_module", + "get_flash_kda_uri", + "load_flash_kda_m64_module", + "load_flash_kda_m128_module", + "load_flash_kda_module", +] diff --git a/flashinfer/kda_decode.py b/flashinfer/kda_decode.py index 0c9ef48fb44..c283f38e24a 100644 --- a/flashinfer/kda_decode.py +++ b/flashinfer/kda_decode.py @@ -15,19 +15,22 @@ """ """ -Kimi Delta Attention Decode - API Layer -======================================= +Kimi Delta Attention - API Layer +================================ -This file provides the public API for recurrent KDA decode operations. -Kernel implementations are in flashinfer/kda_kernels/. +This file provides the public API and dispatch for recurrent KDA decode and +prefill operations. """ +import math +import threading from typing import Optional import torch from .api_logging import flashinfer_api from .trace.templates.kda import recurrent_kda_trace +from .utils import get_compute_capability try: from .kda_kernels.recurrent_kda import run_recurrent_kda as _run_recurrent_kda @@ -38,6 +41,694 @@ _RECURRENT_KDA_AVAILABLE = False +_FLASH_KDA_HEAD_DIM = 128 +_FLASH_KDA_BETA_TMA_MIN_HEADS = 8 +_FLASH_KDA_B200_COMPUTE_CAPABILITY = (10, 0) +_FLASH_KDA_DESCRIPTOR_STORAGE_BYTES = 6 * 128 +_flash_kda_tensor_cache: dict[tuple, torch.Tensor] = {} +_flash_kda_tensor_cache_lock = threading.Lock() + + +class _RecurrentKDAPrefillWorkspaceBase: + def __init__(self, device: torch.device | str) -> None: + normalized_device = torch.device(device) + if normalized_device.type != "cuda": + raise ValueError("RecurrentKDAPrefillWorkspace requires a CUDA device") + if normalized_device.index is None: + normalized_device = torch.device("cuda", torch.cuda.current_device()) + self.device = normalized_device + self._lock = threading.Lock() + self._state_scratch: Optional[torch.Tensor] = None + self._beta_padding: Optional[torch.Tensor] = None + self._descriptor_storages = { + variant: torch.empty( + _FLASH_KDA_DESCRIPTOR_STORAGE_BYTES, + dtype=torch.uint8, + device=self.device, + ) + for variant in ("m64", "m128") + } + self._descriptor_signatures: dict[str, tuple] = {} + self._bound_stream_ptr: Optional[int] = None + self._captured = False + + +class RecurrentKDAPrefillWorkspace(_RecurrentKDAPrefillWorkspaceBase): + """Caller-owned storage required for recurrent-KDA CUDA graph capture. + + Construct one workspace per captured :func:`recurrent_kda` invocation on + the graph's CUDA device. Warm it by invoking :func:`recurrent_kda` eagerly + with the exact tensors and capture stream, then synchronize that stream + before capture. The workspace owns stable state scratch, beta padding, and + M64/M128 TMA descriptor storage for the lifetime of the graph. + + A workspace binds to its first stream. Once it participates in capture it + cannot be passed to Python again, either eagerly or in another capture. + Graph replay does not invoke Python and remains valid for the lifetime of + the workspace. + """ + + +class _FlashKDAStreamWorkspace(_RecurrentKDAPrefillWorkspaceBase): + """Internal eager-only workspace for one CUDA stream.""" + + +_flash_kda_stream_workspaces: dict[tuple[int, int], _FlashKDAStreamWorkspace] = {} +_flash_kda_stream_workspaces_lock = threading.Lock() + + +def _is_plain_multi_token_prefill( + q: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], + num_spec_tokens: Optional[int], +) -> bool: + if num_spec_tokens is not None or not isinstance(q, torch.Tensor) or q.ndim != 4: + return False + if cu_seqlens is None: + return q.shape[1] > 1 + if not isinstance(cu_seqlens, torch.Tensor) or cu_seqlens.ndim != 1: + return False + num_sequences = cu_seqlens.numel() - 1 + return num_sequences > 0 and q.shape[1] > num_sequences + + +def _is_contiguous_cuda_tensor( + tensor: Optional[torch.Tensor], + *, + dtype: torch.dtype, + device: torch.device, +) -> bool: + return ( + isinstance(tensor, torch.Tensor) + and tensor.is_cuda + and tensor.device == device + and tensor.dtype == dtype + and tensor.is_contiguous() + ) + + +def _flash_kda_prefill_is_eligible( + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: Optional[torch.Tensor], + dt_bias: Optional[torch.Tensor], + initial_state: Optional[torch.Tensor], + use_qk_l2norm_in_kernel: bool, + use_gate_in_kernel: bool, + lower_bound: Optional[float], + cu_seqlens: Optional[torch.Tensor], + ssm_state_indices: Optional[torch.Tensor], + num_spec_tokens: Optional[int], + num_accepted_tokens: Optional[torch.Tensor], + output: Optional[torch.Tensor], + initial_state_source: Optional[torch.Tensor], + initial_state_indices: Optional[torch.Tensor], + beta_is_logit: bool, +) -> bool: + """Return whether the call exactly matches the frozen FlashKDA contract.""" + + if not _is_plain_multi_token_prefill(q, cu_seqlens, num_spec_tokens): + return False + if ( + ssm_state_indices is not None + or num_accepted_tokens is not None + or initial_state_source is not None + or initial_state_indices is not None + ): + return False + if not ( + use_qk_l2norm_in_kernel + and use_gate_in_kernel + and beta_is_logit + and lower_bound is not None + and math.isfinite(float(lower_bound)) + and float(lower_bound) < 0.0 + ): + return False + if ( + not q.is_cuda + or get_compute_capability(q.device) != _FLASH_KDA_B200_COMPUTE_CAPABILITY + ): + return False + if not _is_contiguous_cuda_tensor(q, dtype=torch.bfloat16, device=q.device): + return False + if q.ndim != 4: + return False + batch_size, total_or_fixed_tokens, num_heads, head_dim = q.shape + if ( + batch_size <= 0 + or total_or_fixed_tokens <= 1 + or num_heads <= 0 + or head_dim != _FLASH_KDA_HEAD_DIM + ): + return False + for tensor in (k, v, g): + if ( + not _is_contiguous_cuda_tensor( + tensor, dtype=torch.bfloat16, device=q.device + ) + or tensor.shape != q.shape + ): + return False + if not _is_contiguous_cuda_tensor( + beta, dtype=torch.bfloat16, device=q.device + ) or beta.shape != (batch_size, total_or_fixed_tokens, num_heads): + return False + if not _is_contiguous_cuda_tensor( + A_log, dtype=torch.float32, device=q.device + ) or A_log.shape != (num_heads,): + return False + if not _is_contiguous_cuda_tensor(dt_bias, dtype=torch.float32, device=q.device): + return False + if dt_bias.numel() != num_heads * _FLASH_KDA_HEAD_DIM or dt_bias.ndim not in (1, 2): + return False + if dt_bias.ndim == 2 and dt_bias.shape != (num_heads, _FLASH_KDA_HEAD_DIM): + return False + + if cu_seqlens is None: + num_sequences = batch_size + else: + if ( + batch_size != 1 + or not cu_seqlens.is_cuda + or cu_seqlens.device != q.device + or cu_seqlens.dtype not in (torch.int32, torch.int64) + or cu_seqlens.ndim != 1 + or not cu_seqlens.is_contiguous() + ): + return False + num_sequences = cu_seqlens.numel() - 1 + if num_sequences <= 0 or total_or_fixed_tokens <= num_sequences: + return False + + if initial_state is not None: + if not _is_contiguous_cuda_tensor( + initial_state, dtype=torch.bfloat16, device=q.device + ) or initial_state.shape != ( + num_sequences, + num_heads, + _FLASH_KDA_HEAD_DIM, + _FLASH_KDA_HEAD_DIM, + ): + return False + if output is not None: + if ( + not _is_contiguous_cuda_tensor( + output, dtype=torch.bfloat16, device=q.device + ) + or output.shape != q.shape + ): + return False + return True + + +def _select_flash_kda_prefill_variant( + *, fixed_layout: bool, num_sequences: int, num_heads: int +) -> str: + if fixed_layout and num_sequences == 1 and num_heads == 64: + return "m64" + return "m128" + + +def _cached_tensor( + key: tuple, + factory, + *, + capture_error: str, +) -> torch.Tensor: + with _flash_kda_tensor_cache_lock: + tensor = _flash_kda_tensor_cache.get(key) + if tensor is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError(capture_error) + tensor = factory() + _flash_kda_tensor_cache[key] = tensor + return tensor + + +def _fixed_cu_seqlens( + *, + device: torch.device, + batch_size: int, + seq_len: int, +) -> torch.Tensor: + key = ("fixed_cu", *_stream_cache_key(device), batch_size, seq_len) + return _cached_tensor( + key, + lambda: torch.arange( + 0, + batch_size * seq_len + 1, + seq_len, + dtype=torch.int64, + device=device, + ), + capture_error=( + "fixed-layout recurrent_kda prefill metadata is not warmed for " + "CUDA graph capture; invoke the same shape once before capture" + ), + ) + + +def _identity_seq_order( + *, + device: torch.device, + num_sequences: int, +) -> torch.Tensor: + key = ("seq_order", *_stream_cache_key(device), num_sequences) + return _cached_tensor( + key, + lambda: torch.arange(num_sequences, dtype=torch.int32, device=device), + capture_error=( + "recurrent_kda prefill seq_order is not warmed for CUDA graph " + "capture; pass a preallocated seq_order or warm the shape first" + ), + ) + + +def _dummy_bf16(device: torch.device) -> torch.Tensor: + key = ("dummy_bf16", *_stream_cache_key(device)) + return _cached_tensor( + key, + lambda: torch.empty(1, dtype=torch.bfloat16, device=device), + capture_error=( + "recurrent_kda prefill dummy state is not warmed for CUDA graph " + "capture; invoke the same device once before capture" + ), + ) + + +def _stream_cache_key(device: torch.device) -> tuple[int, int]: + stream = torch.cuda.current_stream(device) + device_index = ( + device.index if device.index is not None else torch.cuda.current_device() + ) + return device_index, int(stream.cuda_stream) + + +def _get_stream_workspace(device: torch.device) -> _FlashKDAStreamWorkspace: + key = _stream_cache_key(device) + with _flash_kda_stream_workspaces_lock: + workspace = _flash_kda_stream_workspaces.get(key) + if workspace is None: + workspace = _FlashKDAStreamWorkspace(device) + _flash_kda_stream_workspaces[key] = workspace + return workspace + + +def _workspace_buffer( + *, + workspace: _RecurrentKDAPrefillWorkspaceBase, + attribute: str, + device: torch.device, + numel: int, + capture_error: str, +) -> torch.Tensor: + buffer = getattr(workspace, attribute) + capturing = torch.cuda.is_current_stream_capturing() + if buffer is None or buffer.numel() < numel: + if capturing: + raise RuntimeError(capture_error) + buffer = torch.empty(numel, dtype=torch.bfloat16, device=device) + setattr(workspace, attribute, buffer) + return buffer[:numel] + + +def _state_scratch( + *, + workspace: _RecurrentKDAPrefillWorkspaceBase, + device: torch.device, + shape: tuple[int, int, int, int], +) -> torch.Tensor: + numel = math.prod(shape) + return _workspace_buffer( + workspace=workspace, + attribute="_state_scratch", + device=device, + numel=numel, + capture_error=( + "recurrent_kda prefill final-state workspace is not large enough " + "for CUDA graph capture; warm the largest shape on this stream " + "before capture" + ), + ).view(shape) + + +def _beta_tma_source( + beta: torch.Tensor, + workspace: _RecurrentKDAPrefillWorkspaceBase, +) -> torch.Tensor: + batch_size, seq_len, num_heads = beta.shape + total_tokens = batch_size * seq_len + beta_flat = beta.reshape(total_tokens, num_heads) + padded_tokens = max(total_tokens, 32) + padded_heads = max(num_heads, _FLASH_KDA_BETA_TMA_MIN_HEADS) + if padded_tokens == total_tokens and padded_heads == num_heads: + return beta_flat + shape = (padded_tokens, padded_heads) + padded = _workspace_buffer( + workspace=workspace, + attribute="_beta_padding", + device=beta.device, + numel=math.prod(shape), + capture_error=( + "recurrent_kda prefill beta TMA workspace is not large enough for " + "CUDA graph capture; warm the largest padded token/head shape on " + "this stream before capture" + ), + ).view(shape) + # The descriptor may fetch a full (32 token, 8 head) box at the tail. + # Clear reused padding so every legal OOB-adjacent fetch has defined data. + padded.zero_() + padded[:total_tokens, :num_heads].copy_(beta_flat) + return padded + + +def _tensor_descriptor_signature(tensor: torch.Tensor) -> tuple: + return ( + tensor.data_ptr(), + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + ) + + +def _descriptor_signature( + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta_tma: torch.Tensor, + out: torch.Tensor, +) -> tuple: + return tuple( + _tensor_descriptor_signature(tensor) for tensor in (q, k, v, g, beta_tma, out) + ) + + +def _bind_workspace( + workspace: _RecurrentKDAPrefillWorkspaceBase, + *, + device: torch.device, + stream_ptr: int, + capturing: bool, + explicit: bool, +) -> None: + if workspace.device != device: + raise ValueError( + "RecurrentKDAPrefillWorkspace is bound to " + f"{workspace.device}, but recurrent_kda inputs are on {device}" + ) + if workspace._bound_stream_ptr is None: + workspace._bound_stream_ptr = stream_ptr + elif workspace._bound_stream_ptr != stream_ptr: + raise RuntimeError( + "RecurrentKDAPrefillWorkspace is bound to a different CUDA " + "stream; warm and capture it on one stream" + ) + if explicit and workspace._captured: + reuse_kind = "captured by another CUDA graph" if capturing else "reused eagerly" + raise RuntimeError( + "RecurrentKDAPrefillWorkspace has participated in CUDA graph " + f"capture and cannot be {reuse_kind} or mutated" + ) + + +def _storage_ranges_overlap( + left: torch.Tensor, + right: torch.Tensor, +) -> bool: + if left.device != right.device or left.numel() == 0 or right.numel() == 0: + return False + left_start = left.data_ptr() + right_start = right.data_ptr() + left_end = left_start + left.numel() * left.element_size() + right_end = right_start + right.numel() * right.element_size() + return left_start < right_end and right_start < left_end + + +def _check_output_does_not_overlap_inputs( + output: torch.Tensor, + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: Optional[torch.Tensor], +) -> None: + for name, tensor in ( + ("q", q), + ("k", k), + ("v", v), + ("g", g), + ("beta", beta), + ("initial_state", initial_state), + ): + if tensor is not None and _storage_ranges_overlap(output, tensor): + raise ValueError( + f"output must not overlap {name} for frozen recurrent_kda prefill" + ) + + +def _validate_prefill_seq_order( + seq_order: Optional[torch.Tensor], + *, + fixed_layout: bool, + num_sequences: int, + device: torch.device, +) -> torch.Tensor: + if seq_order is None: + return _identity_seq_order(device=device, num_sequences=num_sequences) + if fixed_layout: + raise ValueError("seq_order is only supported for packed recurrent_kda prefill") + if not isinstance(seq_order, torch.Tensor): + raise TypeError("seq_order must be a torch.Tensor") + if ( + not seq_order.is_cuda + or seq_order.device != device + or seq_order.dtype != torch.int32 + or seq_order.ndim != 1 + or not seq_order.is_contiguous() + or seq_order.numel() != num_sequences + ): + raise ValueError( + "seq_order must be a contiguous CUDA int32 tensor with one " + f"entry per sequence ({num_sequences})" + ) + return seq_order + + +def _get_flash_kda_prefill_module(variant: str): + from .jit.flash_kda import get_flash_kda_prefill_module + + return get_flash_kda_prefill_module(variant) + + +def _run_flash_kda_prefill( + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + scale: Optional[float], + initial_state: Optional[torch.Tensor], + output_final_state: bool, + lower_bound: float, + cu_seqlens: Optional[torch.Tensor], + output: Optional[torch.Tensor], + seq_order: Optional[torch.Tensor], + prefill_workspace: Optional[RecurrentKDAPrefillWorkspace], +) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + capturing = torch.cuda.is_current_stream_capturing() + if capturing and prefill_workspace is None: + raise RuntimeError( + "CUDA graph capture of recurrent_kda prefill requires an explicit " + "RecurrentKDAPrefillWorkspace warmed with the exact tensors on " + "the capture stream" + ) + batch_size, seq_len, num_heads, _ = q.shape + fixed_layout = cu_seqlens is None + num_sequences = batch_size if fixed_layout else cu_seqlens.numel() - 1 + if fixed_layout: + cu_seqlens_i64 = _fixed_cu_seqlens( + device=q.device, batch_size=batch_size, seq_len=seq_len + ) + else: + assert cu_seqlens is not None + if cu_seqlens.dtype == torch.int32 and capturing: + raise RuntimeError( + "packed recurrent_kda prefill requires int64 cu_seqlens " + "during CUDA graph capture; convert it before capture" + ) + cu_seqlens_i64 = ( + cu_seqlens + if cu_seqlens.dtype == torch.int64 + else cu_seqlens.to(torch.int64) + ) + seq_order_i32 = _validate_prefill_seq_order( + seq_order, + fixed_layout=fixed_layout, + num_sequences=num_sequences, + device=q.device, + ) + dummy_state = _dummy_bf16(q.device) + + if output is None: + if capturing: + raise RuntimeError( + "CUDA graph capture requires a preallocated output tensor for " + "recurrent_kda prefill" + ) + out_buf = torch.empty_like(q) + else: + out_buf = output + _check_output_does_not_overlap_inputs( + out_buf, + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + ) + + state_shape = ( + num_sequences, + num_heads, + _FLASH_KDA_HEAD_DIM, + _FLASH_KDA_HEAD_DIM, + ) + use_initial_state = initial_state is not None + if initial_state is not None: + initial_state_arg = initial_state + store_final_state = True + returned_state = initial_state + elif output_final_state: + initial_state_arg = dummy_state + if prefill_workspace is None: + final_state_arg = torch.empty( + state_shape, dtype=torch.bfloat16, device=q.device + ) + returned_state = final_state_arg + else: + # Assigned to caller-owned stable state scratch under its lock. + returned_state = None + store_final_state = True + else: + initial_state_arg = dummy_state + final_state_arg = dummy_state + store_final_state = False + returned_state = None + + scale_value = ( + 1.0 / math.sqrt(_FLASH_KDA_HEAD_DIM) if scale is None else float(scale) + ) + if not math.isfinite(scale_value): + raise ValueError(f"scale must be finite, got {scale_value}") + variant = _select_flash_kda_prefill_variant( + fixed_layout=fixed_layout, + num_sequences=num_sequences, + num_heads=num_heads, + ) + stream_ptr = int(torch.cuda.current_stream(q.device).cuda_stream) + explicit_workspace = prefill_workspace is not None + workspace: _RecurrentKDAPrefillWorkspaceBase + if prefill_workspace is None: + workspace = _get_stream_workspace(q.device) + else: + workspace = prefill_workspace + # TVM FFI may release the GIL. Serialize the complete shared-workspace + # enqueue sequence so two host threads cannot interleave preparation, + # launch, or state copy-back on the same CUDA stream. + with workspace._lock: + _bind_workspace( + workspace, + device=q.device, + stream_ptr=stream_ptr, + capturing=capturing, + explicit=explicit_workspace, + ) + beta_tma = _beta_tma_source(beta, workspace) + if initial_state is not None or (output_final_state and explicit_workspace): + # The frozen ABI marks both state pointers __restrict__. Always + # launch into distinct storage, then preserve recurrent_kda's + # in-place mutation with a same-stream copy-back. + final_state_arg = _state_scratch( + workspace=workspace, + device=q.device, + shape=state_shape, + ) + if initial_state is None: + returned_state = final_state_arg + signature = _descriptor_signature( + q=q, + k=k, + v=v, + g=g, + beta_tma=beta_tma, + out=out_buf, + ) + warmed_signature = workspace._descriptor_signatures.get(variant) + if capturing: + if warmed_signature != signature: + raise RuntimeError( + "RecurrentKDAPrefillWorkspace is not warmed for the exact " + f"{variant} descriptor signature; eagerly invoke the same " + "call on this stream before capture" + ) + prepare_descriptors = 0 + else: + prepare_descriptors = int(warmed_signature != signature) + descriptor_storage = workspace._descriptor_storages[variant] + module = _get_flash_kda_prefill_module(variant) + try: + module.run( + q, + k, + v, + g, + beta, + beta_tma, + A_log, + dt_bias, + cu_seqlens_i64, + seq_order_i32, + initial_state_arg, + out_buf, + final_state_arg, + descriptor_storage, + prepare_descriptors, + num_heads, + int(use_initial_state), + int(store_final_state), + scale_value, + float(lower_bound), + stream_ptr, + ) + except Exception: + if prepare_descriptors: + workspace._descriptor_signatures.pop(variant, None) + raise + if prepare_descriptors: + workspace._descriptor_signatures[variant] = signature + if initial_state is not None: + initial_state.copy_(final_state_arg) + if capturing and explicit_workspace: + workspace._captured = True + return ( + out_buf, + returned_state if output_final_state else None, + ) + + @flashinfer_api(trace=recurrent_kda_trace) def recurrent_kda( q: torch.Tensor, @@ -61,34 +752,46 @@ def recurrent_kda( initial_state_source: Optional[torch.Tensor] = None, initial_state_indices: Optional[torch.Tensor] = None, beta_is_logit: bool = False, + seq_order: Optional[torch.Tensor] = None, + prefill_workspace: Optional[RecurrentKDAPrefillWorkspace] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - r"""Recurrent KDA (Kimi Delta Attention) decode kernel. + r"""Recurrent KDA (Kimi Delta Attention) decode and prefill kernel. This is the public API layer for the CuTe DSL implementation in ``flashinfer.kda_kernels.recurrent_kda``. It supports single-token decode, fused speculative decode, GQA, optional cu_seqlens packing, and the same - gate modes as the backend implementation. + gate modes as the backend implementation. On NVIDIA B200, the exact + FlashKDA-compatible subset of ordinary multi-token prefill is dispatched to + frozen SM100a kernels. All existing decode and speculative-decode calls + retain the CuTe DSL backend. Args: q (torch.Tensor): - Current query of shape ``[B, 1, H, K]``, or ``[1, total_tokens, H, K]`` - when using ``cu_seqlens``. Must be bfloat16. + Query of shape ``[B, T, H, K]``, or + ``[1, total_tokens, H, K]`` when using ``cu_seqlens``. Must be + bfloat16. ``T=1`` selects decode; eligible ``T>1`` calls may select + the frozen prefill backend. k (torch.Tensor): - Current key of shape ``[B, 1, H, K]``. Must be bfloat16. + Key with the same shape as ``q``. Must be bfloat16. v (torch.Tensor): - Current value of shape ``[B, 1, HV, V]``. Must be bfloat16. - GQA is applied when ``HV != H``. + Value of shape ``[B, T, HV, V]``, or + ``[1, total_tokens, HV, V]`` when packed. Must be bfloat16. GQA is + applied when ``HV != H``. g (torch.Tensor): - Per-K-dimension gate of shape ``[B, 1, HV, K]``. Must be bfloat16. - Log-space if pre-computed, raw input if ``use_gate_in_kernel=True``. + Per-K-dimension gate of shape ``[B, T, HV, K]``, or + ``[1, total_tokens, HV, K]`` when packed. Must be bfloat16. + Log-space if pre-computed, raw input if + ``use_gate_in_kernel=True``. beta (torch.Tensor): - Delta-rule learning rate of shape ``[B, 1, HV]``. Must be bfloat16. + Delta-rule learning rate of shape ``[B, T, HV]``, or + ``[1, total_tokens, HV]`` when packed. Must be bfloat16. Pre-sigmoided unless ``beta_is_logit=True``. A_log (Optional[torch.Tensor]): Log decay parameter of shape ``[H]``. Must be float32. Required when ``use_gate_in_kernel=True``. dt_bias (Optional[torch.Tensor]): - Per-head-K decay bias of shape ``[H*K]``. Must be float32. + Per-head-K decay bias of shape ``[H*K]`` or ``[H, K]``. Must be + float32. scale (Optional[float]): Scale factor for queries. If ``None``, defaults to ``1 / sqrt(K)``. initial_state (Optional[torch.Tensor]): @@ -108,7 +811,12 @@ def recurrent_kda( If set, uses ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))`` gate formula instead of softplus. Must be negative. cu_seqlens (Optional[torch.Tensor]): - Cumulative sequence lengths of shape ``[N+1]``. Must be int32. + Contiguous CUDA cumulative sequence lengths of shape ``[N+1]``. + May be int32 or int64. Frozen prefill converts int32 offsets to + int64 outside graph capture; graph capture requires caller-provided + int64 offsets. For frozen prefill, values must start at zero, be + strictly increasing, and end at the total token count. This value + contract is not host-validated to avoid a device synchronization. ssm_state_indices (Optional[torch.Tensor]): State cache indices. Shape ``[N]`` int32 for standard decode, or ``[N, 1+S]`` int32 for spec decode (``num_spec_tokens`` must also @@ -122,9 +830,11 @@ def recurrent_kda( from ``ssm_state_indices[n, 0]``. Values above ``1+S`` are clamped to the final checkpoint slot. output (Optional[torch.Tensor]): - Pre-allocated output tensor. Shape ``[B, 1, HV, V]`` for standard - decode, ``[1, N*(1+S), HV, V]`` for spec decode with - ``cu_seqlens``. If ``None``, a new tensor is allocated. + Pre-allocated output tensor. Shape ``[B, T, HV, V]`` for fixed + layout, or the corresponding packed/speculative shape when using + ``cu_seqlens``. If ``None``, a new tensor is allocated. Frozen + prefill requires storage disjoint from Q, K, V, G, beta, and + ``initial_state``. initial_state_source (Optional[torch.Tensor]): Optional read-only committed state pool ``[N0, HV, V, K]``. When provided, token 0 is loaded from this pool instead of @@ -134,6 +844,19 @@ def recurrent_kda( with ``initial_state_source``. beta_is_logit (bool): If ``True``, apply sigmoid to ``beta`` inside the recurrent kernel. + seq_order (Optional[torch.Tensor]): + Optional packed-prefill sequence order, as a contiguous CUDA int32 + permutation of shape ``[N]``. Sorting by descending sequence length + improves tail utilization. It is only consumed by the frozen + FlashKDA prefill backend; prepare it before CUDA graph capture or + timed launches. Fixed-layout prefill and decode calls must leave it + as ``None``. + prefill_workspace (Optional[RecurrentKDAPrefillWorkspace]): + Caller-owned workspace for the frozen B200 prefill backend. It is + optional for eager execution and required for CUDA graph capture. + Warm it eagerly with the exact tensors on the capture stream before + capture. Use one workspace per captured ``recurrent_kda`` + invocation. Returns: Tuple of ``(output, final_state)`` where ``final_state`` is ``None`` @@ -141,6 +864,64 @@ def recurrent_kda( :func:`flashinfer.kda_kernels.recurrent_kda.run_recurrent_kda` for the backend implementation. """ + if prefill_workspace is not None and not isinstance( + prefill_workspace, RecurrentKDAPrefillWorkspace + ): + raise TypeError("prefill_workspace must be a RecurrentKDAPrefillWorkspace") + + use_flash_kda_prefill = _flash_kda_prefill_is_eligible( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=initial_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + use_gate_in_kernel=use_gate_in_kernel, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_spec_tokens=num_spec_tokens, + num_accepted_tokens=num_accepted_tokens, + output=output, + initial_state_source=initial_state_source, + initial_state_indices=initial_state_indices, + beta_is_logit=beta_is_logit, + ) + if use_flash_kda_prefill: + assert A_log is not None + assert dt_bias is not None + assert lower_bound is not None + return _run_flash_kda_prefill( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + output=output, + seq_order=seq_order, + prefill_workspace=prefill_workspace, + ) + + if prefill_workspace is not None: + raise ValueError( + "prefill_workspace is only supported by eligible ordinary " + "prefill on the frozen B200 FlashKDA backend" + ) + if seq_order is not None: + raise ValueError( + "seq_order is only supported by eligible packed ordinary prefill " + "on the frozen B200 FlashKDA backend" + ) if _run_recurrent_kda is None: raise NotImplementedError("recurrent KDA backend is unavailable") diff --git a/flashinfer/trace/templates/kda.py b/flashinfer/trace/templates/kda.py index e27552501b2..935be32acd7 100644 --- a/flashinfer/trace/templates/kda.py +++ b/flashinfer/trace/templates/kda.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""TraceTemplate for recurrent Key-Driven Attention (KDA) decode.""" +"""TraceTemplate for recurrent Key-Driven Attention (KDA).""" from ..template import Const, Scalar, Tensor, TraceTemplate, Var @@ -21,8 +21,9 @@ op_type="kda", name_prefix="recurrent_kda", description=( - "Recurrent Key-Driven Attention decode with per-key-dimension gating " - "and an optional read-only committed-state source." + "Recurrent Key-Driven Attention decode/spec-decode plus exact " + "FlashKDA-compatible ordinary prefill on B200, with " + "per-key-dimension gating and recurrent BF16 state." ), axes={ "batch_size": Var(description="Number of input batch rows."), @@ -35,6 +36,9 @@ "state_pool_size": Var(description="Number of writable state slots."), "source_pool_size": Var(description="Number of committed-state slots."), "num_sequences": Var(description="Number of state-source indices."), + "num_sequences_plus_one": Var( + description="Number of packed cumulative-length entries." + ), }, inputs={ "q": Tensor(["batch_size", "seq_len", "num_q_heads", "head_dim"]), @@ -42,20 +46,58 @@ "v": Tensor(["batch_size", "seq_len", "num_v_heads", "head_dim"]), "g": Tensor(["batch_size", "seq_len", "num_v_heads", "head_dim"]), "beta": Tensor(["batch_size", "seq_len", "num_v_heads"]), + "A_log": Tensor( + ["num_q_heads"], + dtype="float32", + optional=True, + description="FP32 per-query-head log decay rate.", + ), + "dt_bias": Tensor( + ["num_q_heads", "head_dim"], + dtype="float32", + optional=True, + description=( + "FP32 per-head/key decay bias; flattened [H*D] or [H,D] " + "storage is accepted by the API." + ), + ), "initial_state": Tensor( ["state_pool_size", "num_v_heads", "head_dim", "head_dim"], + dtype="bfloat16", optional=True, ), "initial_state_source": Tensor( ["source_pool_size", "num_v_heads", "head_dim", "head_dim"], + dtype="bfloat16", optional=True, description="Read-only committed-state pool.", ), "initial_state_indices": Tensor( ["num_sequences"], + dtype="int32", optional=True, description="Committed-state slot selected for each sequence.", ), + "cu_seqlens": Tensor( + ["num_sequences_plus_one"], + optional=True, + description="Packed cumulative sequence lengths (int32 or int64).", + ), + "num_accepted_tokens": Tensor( + ["num_sequences"], + dtype="int32", + optional=True, + description="Accepted-token counts for speculative decode.", + ), + "seq_order": Tensor( + ["num_sequences"], + dtype="int32", + optional=True, + description=( + "Packed-prefill sequence permutation, normally descending " + "by sequence length." + ), + ), "scale": Scalar("float32", optional=True), "output_final_state": Scalar("int32", optional=True), "use_qk_l2norm_in_kernel": Scalar("int32", optional=True), @@ -68,6 +110,7 @@ "output": Tensor( ["batch_size", "seq_len", "num_v_heads", "head_dim"], dtype_from="q", + param="output", ), "final_state": Tensor( ["state_pool_size", "num_v_heads", "head_dim", "head_dim"], @@ -79,5 +122,5 @@ "num_v_heads % num_q_heads == 0", "head_dim in (64, 128)", ], - tags=["stage:decode", "status:verified"], + tags=["stage:decode", "stage:prefill", "status:verified"], ) diff --git a/tests/jit/test_flash_kda_jit.py b/tests/jit/test_flash_kda_jit.py new file mode 100644 index 00000000000..f627d151be9 --- /dev/null +++ b/tests/jit/test_flash_kda_jit.py @@ -0,0 +1,164 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import hashlib + +import pytest +from packaging.version import Version + +from flashinfer.jit import core as jit_core +from flashinfer.jit import flash_kda + + +@pytest.mark.parametrize( + ("variant", "smem_bytes", "generated_sha256"), + [ + ( + "m64", + 219136, + "468780f04768c949b22c688c1b85d235b6ffe050fd3327394fcdda7ea5112203", + ), + ( + "m128", + 227328, + "a2378074fde64fa454edb972dc51e188294dcb89369ee5ea1153f6c67200f1ab", + ), + ], +) +def test_flash_kda_uri_and_jit_spec(monkeypatch, variant, smem_bytes, generated_sha256): + monkeypatch.setattr( + jit_core.current_compilation_context, + "TARGET_CUDA_ARCHS", + {(10, "0a")}, + ) + flash_kda.gen_flash_kda_module.cache_clear() + + uri = flash_kda.get_flash_kda_uri(variant) + spec = flash_kda.gen_flash_kda_module(variant) + + assert uri == f"flash_kda_bf16_fused_{variant}_sm100a" + assert spec.name == uri + assert len(spec.sources) == 1 + assert spec.sources[0].name == f"flashkda_bf16_fused_{variant}_binding.cu" + assert spec.sources[0].is_file() + assert "-gencode=arch=compute_100a,code=sm_100a" in spec.extra_cuda_cflags + assert "-use_fast_math" in spec.extra_cuda_cflags + assert not any( + "compute_103" in flag or "compute_120" in flag + for flag in spec.extra_cuda_cflags + ) + frozen_source = spec.sources[0].parent / f"flashkda_bf16_fused_{variant}.cu" + frozen_text = frozen_source.read_text() + assert "Provenance: loom @ 8437e0515b212e7973b196c4ab680d3d90c1209c" in frozen_text + assert f"#define SMEM_TOTAL {smem_bytes}" in frozen_text + assert frozen_text.count("// clang-format off") == 1 + assert frozen_text.rstrip().endswith("// clang-format on") + generated_body = frozen_text.partition("// clang-format off\n")[2].rpartition( + "// clang-format on" + )[0] + integration_begin = ( + " // FLASHINFER INTEGRATION BEGIN: acquire global tensor maps\n" + ) + integration_end = " // FLASHINFER INTEGRATION END: acquire global tensor maps\n" + generated_prefix, begin_marker, integration_tail = generated_body.partition( + integration_begin + ) + integration_prologue, end_marker, generated_suffix = integration_tail.partition( + integration_end + ) + assert begin_marker == integration_begin + assert end_marker == integration_end + assert integration_prologue.count("fence.proxy.tensormap::generic.acquire.gpu") == 6 + assert integration_prologue.count("], 128;") == 6 + assert integration_prologue.count("__syncthreads();") == 1 + for tensor_map in ("q_tma", "k_tma", "v_tma", "g_tma", "beta_tma", "out_tma"): + assert f'"l"({tensor_map})' in integration_prologue + # Keep the exporter output immutable below the narrowly marked FlashInfer + # integration prologue. + normalized_generated_body = generated_prefix + generated_suffix + assert ( + hashlib.sha256(normalized_generated_body.encode()).hexdigest() + == generated_sha256 + ) + assert [ + line for line in frozen_text.splitlines() if line.startswith("#include") + ] == [ + "#include ", + "#include ", + ] + + binding_text = spec.sources[0].read_text() + assert "#define uint64_t flashkda_generated_uint64_t" in binding_text + assert "TensorView descriptor_storage, int64_t prepare_descriptors" in binding_text + + +def test_flash_kda_descriptor_workspace_contract(): + common_source = flash_kda._get_flash_kda_csrc_dir() / ( + "flashkda_binding_common.cuh" + ) + common_text = common_source.read_text() + + assert "static_assert(sizeof(CUtensorMap) == 128);" in common_text + assert "kTensorMapAlignment = 64" in common_text + assert 'CheckDtype(descriptor_storage, "descriptor_storage", dl_uint8)' in ( + common_text + ) + assert "PublishTensorMaps<<<1, 128, 0, stream>>>" in common_text + assert "prepare_descriptors must be 0 during CUDA graph capture" in common_text + assert "cudaMemcpyAsync(TMA descriptors)" not in common_text + + +def test_flash_kda_variant_validation_and_public_getter(monkeypatch): + with pytest.raises(ValueError, match="unsupported FlashKDA variant"): + flash_kda.get_flash_kda_uri("m32") + + sentinel = object() + monkeypatch.setattr( + flash_kda, + "load_flash_kda_module", + lambda variant: (sentinel, variant), + ) + assert flash_kda.get_flash_kda_prefill_module("m128") == ( + sentinel, + "m128", + ) + + +@pytest.mark.parametrize( + ("target_archs", "expected_exact"), + [ + ({(10, "0a")}, True), + ({(10, "0f")}, False), + ({(10, "3a")}, False), + ({(12, "0f")}, False), + ], +) +def test_aot_detects_only_exact_sm100a(monkeypatch, target_archs, expected_exact): + from flashinfer import aot + + class FakeCompilationContext: + TARGET_CUDA_ARCHS = target_archs + + def get_nvcc_flags_list(self, supported_major_versions=None): + del supported_major_versions + return [ + f"-gencode=arch=compute_{major}{minor},code=sm_{major}{minor}" + for major, minor in sorted(self.TARGET_CUDA_ARCHS) + ] + + monkeypatch.setattr(aot, "CompilationContext", FakeCompilationContext) + monkeypatch.setattr(aot, "get_cuda_version", lambda: Version("13.0")) + + capabilities = aot.detect_sm_capabilities() + assert capabilities["sm100a_exact"] is expected_exact diff --git a/tests/kda/test_recurrent_kda_prefill.py b/tests/kda/test_recurrent_kda_prefill.py new file mode 100644 index 00000000000..916832079b3 --- /dev/null +++ b/tests/kda/test_recurrent_kda_prefill.py @@ -0,0 +1,864 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib +import math + +import pytest +import torch +import torch.nn.functional as F + +from flashinfer.kda_decode import ( + RecurrentKDAPrefillWorkspace, + recurrent_kda, +) +from flashinfer.utils import get_compute_capability + +kda_api = importlib.import_module("flashinfer.kda_decode") + + +def _strict_prefill_kwargs(inputs): + return { + **inputs, + "use_qk_l2norm_in_kernel": True, + "use_gate_in_kernel": True, + "lower_bound": -5.0, + "beta_is_logit": True, + } + + +def _make_inputs( + *, + seq_lens, + num_heads: int, + packed: bool, + initial_state: bool = False, + seed: int = 0, +): + torch.manual_seed(seed) + if packed: + batch_size = 1 + seq_len = sum(seq_lens) + else: + if len(set(seq_lens)) != 1: + raise ValueError("fixed test inputs require equal sequence lengths") + batch_size = len(seq_lens) + seq_len = seq_lens[0] + shape = (batch_size, seq_len, num_heads, 128) + q = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + k = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + v = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + g = (0.1 * torch.randn(shape, dtype=torch.float32, device="cuda")).to( + torch.bfloat16 + ) + beta = torch.randn( + (batch_size, seq_len, num_heads), + dtype=torch.bfloat16, + device="cuda", + ) + A_log = 0.1 * torch.randn(num_heads, dtype=torch.float32, device="cuda") + dt_bias = 0.1 * torch.randn((num_heads, 128), dtype=torch.float32, device="cuda") + offsets = [0] + for length in seq_lens: + offsets.append(offsets[-1] + length) + state = None + if initial_state: + state = ( + 0.1 + * torch.randn( + (len(seq_lens), num_heads, 128, 128), + dtype=torch.float32, + device="cuda", + ) + ).to(torch.bfloat16) + return { + "q": q, + "k": k, + "v": v, + "g": g, + "beta": beta, + "A_log": A_log, + "dt_bias": dt_bias, + "initial_state": state, + "cu_seqlens": ( + torch.tensor(offsets, dtype=torch.int64, device="cuda") if packed else None + ), + } + + +def _reference(inputs, *, lower_bound=-5.0, scale=None): + q = inputs["q"] + batch_size, seq_len, num_heads, head_dim = q.shape + scale = head_dim**-0.5 if scale is None else scale + q_flat = F.normalize(q.float(), dim=-1).reshape(-1, num_heads, head_dim) + k_flat = F.normalize(inputs["k"].float(), dim=-1).reshape(-1, num_heads, head_dim) + v_flat = inputs["v"].float().reshape(-1, num_heads, head_dim) + g_flat = inputs["g"].float().reshape(-1, num_heads, head_dim) + beta_flat = torch.sigmoid(inputs["beta"].float().reshape(-1, num_heads)) + gate = lower_bound * torch.sigmoid( + torch.exp(inputs["A_log"]).reshape(1, num_heads, 1) + * (g_flat + inputs["dt_bias"].reshape(1, num_heads, head_dim)) + ) + decay = torch.exp(gate) + if inputs["cu_seqlens"] is None: + offsets = [index * seq_len for index in range(batch_size + 1)] + else: + offsets = [int(value) for value in inputs["cu_seqlens"].tolist()] + if inputs["initial_state"] is None: + state = torch.zeros( + (len(offsets) - 1, num_heads, head_dim, head_dim), + dtype=torch.bfloat16, + device=q.device, + ) + else: + state = inputs["initial_state"].clone() + out = torch.empty_like(q_flat) + for sequence in range(len(offsets) - 1): + for token in range(offsets[sequence], offsets[sequence + 1]): + state_f32 = state[sequence].float() + decayed = state_f32 * decay[token].unsqueeze(1) + predicted = torch.einsum("hk,hvk->hv", k_flat[token], decayed) + residual = beta_flat[token].unsqueeze(-1) * (v_flat[token] - predicted) + updated = decayed + residual.unsqueeze(-1) * k_flat[token].unsqueeze(1) + state[sequence] = updated.to(torch.bfloat16) + projected = torch.einsum( + "hk,hvk->hv", q_flat[token], state[sequence].float() + ) + out[token] = (scale * projected).to(torch.bfloat16) + return out.reshape_as(q), state + + +@pytest.fixture +def cuda_device(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + return torch.device("cuda") + + +@pytest.fixture +def b200(cuda_device): + if get_compute_capability(cuda_device) != (10, 0): + pytest.skip("frozen recurrent KDA prefill requires B200 (cc 10.0)") + return cuda_device + + +class _RecorderModule: + def __init__(self, *, final_value=None): + self.calls = [] + self.final_value = final_value + + def run(self, *args): + self.calls.append(args) + if self.final_value is not None and bool(args[17]): + args[12].fill_(self.final_value) + + +def test_decode_and_spec_stay_on_existing_backend(monkeypatch): + sentinel = (object(), object()) + calls = [] + + def old_backend(**kwargs): + calls.append(kwargs) + return sentinel + + monkeypatch.setattr(kda_api, "_run_recurrent_kda", old_backend) + monkeypatch.setattr( + kda_api, + "_get_flash_kda_prefill_module", + lambda variant: pytest.fail(f"unexpected frozen route {variant}"), + ) + q = torch.empty((2, 1, 4, 128), dtype=torch.bfloat16) + result = recurrent_kda(q, q, q, q, torch.empty((2, 1, 4))) + assert result is sentinel + result = recurrent_kda( + q.expand(2, 2, 4, 128), + q.expand(2, 2, 4, 128), + q.expand(2, 2, 4, 128), + q.expand(2, 2, 4, 128), + torch.empty((2, 2, 4)), + num_spec_tokens=1, + ) + assert result is sentinel + assert len(calls) == 2 + + +def test_multi_token_gqa_stays_on_existing_backend(cuda_device, monkeypatch): + sentinel = (object(), object()) + monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + monkeypatch.setattr(kda_api, "_run_recurrent_kda", lambda **kwargs: sentinel) + monkeypatch.setattr( + kda_api, + "_get_flash_kda_prefill_module", + lambda variant: pytest.fail(f"unexpected frozen route {variant}"), + ) + q = torch.randn((1, 2, 2, 128), dtype=torch.bfloat16, device=cuda_device) + v = torch.randn((1, 2, 4, 128), dtype=torch.bfloat16, device=cuda_device) + result = recurrent_kda( + q, + q.clone(), + v, + v.clone(), + torch.randn((1, 2, 4), dtype=torch.bfloat16, device=cuda_device), + A_log=torch.randn(2, device=cuda_device), + dt_bias=torch.randn((2, 128), device=cuda_device), + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + lower_bound=-5.0, + beta_is_logit=True, + ) + assert result is sentinel + + +@pytest.mark.parametrize( + ("packed", "num_heads", "expected_variant"), + [(False, 64, "m64"), (True, 64, "m128"), (True, 2, "m128")], +) +def test_frozen_route_and_ffi_abi( + cuda_device, + monkeypatch, + packed, + num_heads, + expected_variant, +): + monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + monkeypatch.setattr(kda_api, "_flash_kda_stream_workspaces", {}) + modules = {} + + def get_module(variant): + modules.setdefault(variant, _RecorderModule()) + return modules[variant] + + monkeypatch.setattr(kda_api, "_get_flash_kda_prefill_module", get_module) + inputs = _make_inputs( + seq_lens=[1, 2] if packed else [2], + num_heads=num_heads, + packed=packed, + ) + if packed and num_heads == 2: + inputs["cu_seqlens"] = inputs["cu_seqlens"].to(torch.int32) + output = torch.zeros_like(inputs["q"]) + seq_order = ( + torch.tensor([1, 0], dtype=torch.int32, device="cuda") if packed else None + ) + actual, state = recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + seq_order=seq_order, + ) + assert actual.data_ptr() == output.data_ptr() + assert state is None + assert set(modules) == {expected_variant} + (args,) = modules[expected_variant].calls + assert len(args) == 21 + assert args[0].data_ptr() == inputs["q"].data_ptr() + assert args[4].data_ptr() == inputs["beta"].data_ptr() + assert args[5].shape == ( + max(inputs["q"].numel() // (num_heads * 128), 32), + max(num_heads, 8), + ) + assert args[8].dtype == torch.int64 + assert args[9].dtype == torch.int32 + assert args[10].data_ptr() == args[12].data_ptr() + assert args[13].dtype == torch.uint8 + assert args[13].shape == (768,) + assert args[14] == 1 + assert args[15] == num_heads + assert args[16] == 0 + assert args[17] == 0 + assert math.isclose(args[18], 128**-0.5) + assert args[19] == -5.0 + assert args[20] == int(torch.cuda.current_stream(cuda_device).cuda_stream) + if args[5].data_ptr() != inputs["beta"].data_ptr(): + total_tokens = inputs["q"].numel() // (num_heads * 128) + torch.testing.assert_close( + args[5][:total_tokens, :num_heads], + inputs["beta"].reshape(-1, num_heads), + ) + + +def test_frozen_route_passes_nondefault_stream(cuda_device, monkeypatch): + monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + module = _RecorderModule() + monkeypatch.setattr( + kda_api, "_get_flash_kda_prefill_module", lambda variant: module + ) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) + stream = torch.cuda.Stream(device=cuda_device) + stream.wait_stream(torch.cuda.current_stream(cuda_device)) + with torch.cuda.stream(stream): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=torch.empty_like(inputs["q"]), + ) + (args,) = module.calls + assert args[20] == int(stream.cuda_stream) + + +def test_frozen_route_rejects_output_overlap(cuda_device, monkeypatch): + monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + module = _RecorderModule() + monkeypatch.setattr( + kda_api, "_get_flash_kda_prefill_module", lambda variant: module + ) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) + with pytest.raises(ValueError, match="output must not overlap q"): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=inputs["q"].view_as(inputs["q"]), + ) + assert module.calls == [] + + +def test_initial_state_uses_distinct_final_and_copies_back(cuda_device, monkeypatch): + monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + module = _RecorderModule(final_value=0.25) + monkeypatch.setattr( + kda_api, "_get_flash_kda_prefill_module", lambda variant: module + ) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False, initial_state=True) + original_state = inputs["initial_state"] + actual, returned_state = recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=torch.empty_like(inputs["q"]), + output_final_state=True, + ) + assert actual.shape == inputs["q"].shape + assert returned_state is original_state + (args,) = module.calls + assert args[10].data_ptr() == original_state.data_ptr() + assert args[12].data_ptr() != original_state.data_ptr() + assert args[16] == 1 + assert args[17] == 1 + torch.testing.assert_close( + original_state, + torch.full_like(original_state, 0.25), + ) + + +def test_stream_workspace_retains_only_largest_state_buffer(cuda_device, monkeypatch): + monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + monkeypatch.setattr(kda_api, "_flash_kda_stream_workspaces", {}) + module = _RecorderModule(final_value=0.0) + monkeypatch.setattr( + kda_api, "_get_flash_kda_prefill_module", lambda variant: module + ) + cases = [ + _make_inputs( + seq_lens=[2], + num_heads=2, + packed=False, + initial_state=True, + ), + _make_inputs( + seq_lens=[1, 1, 2], + num_heads=2, + packed=True, + initial_state=True, + ), + _make_inputs( + seq_lens=[2, 2], + num_heads=2, + packed=False, + initial_state=True, + ), + ] + for inputs in cases: + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=torch.empty_like(inputs["q"]), + ) + + assert len(kda_api._flash_kda_stream_workspaces) == 1 + (workspace,) = kda_api._flash_kda_stream_workspaces.values() + largest_state_numel = max(inputs["initial_state"].numel() for inputs in cases) + assert workspace._state_scratch.numel() == largest_state_numel + assert workspace._beta_padding.numel() == 32 * 8 + + +@pytest.mark.parametrize( + ("dtype", "size_delta"), + [(torch.int64, 0), (torch.int32, 1)], +) +def test_packed_seq_order_validation(cuda_device, monkeypatch, dtype, size_delta): + monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + monkeypatch.setattr( + kda_api, + "_get_flash_kda_prefill_module", + lambda variant: _RecorderModule(), + ) + inputs = _make_inputs(seq_lens=[1, 2], num_heads=2, packed=True) + seq_order = torch.arange(2 + size_delta, dtype=dtype, device="cuda") + with pytest.raises(ValueError, match="seq_order"): + recurrent_kda(**_strict_prefill_kwargs(inputs), seq_order=seq_order) + + +def test_fixed_prefill_rejects_seq_order(cuda_device, monkeypatch): + monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) + with pytest.raises(ValueError, match="only supported for packed"): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + seq_order=torch.zeros(1, dtype=torch.int32, device=cuda_device), + ) + + +def test_graph_capture_requires_packed_int64_offsets(cuda_device, monkeypatch): + monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + inputs = _make_inputs(seq_lens=[1, 2], num_heads=2, packed=True) + inputs["cu_seqlens"] = inputs["cu_seqlens"].to(torch.int32) + workspace = RecurrentKDAPrefillWorkspace(cuda_device) + with pytest.raises(RuntimeError, match="requires int64 cu_seqlens"): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + prefill_workspace=workspace, + ) + + +def test_graph_capture_requires_explicit_workspace(cuda_device, monkeypatch): + monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) + with pytest.raises( + RuntimeError, match="requires an explicit RecurrentKDAPrefillWorkspace" + ): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=torch.empty_like(inputs["q"]), + ) + + +def test_explicit_workspace_descriptor_prepare_and_reuse(cuda_device, monkeypatch): + monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + module = _RecorderModule() + monkeypatch.setattr( + kda_api, "_get_flash_kda_prefill_module", lambda variant: module + ) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) + output = torch.empty_like(inputs["q"]) + workspace = RecurrentKDAPrefillWorkspace(cuda_device) + + for _ in range(2): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + prefill_workspace=workspace, + ) + assert [args[14] for args in module.calls] == [1, 0] + assert ( + module.calls[0][13].data_ptr() + == module.calls[1][13].data_ptr() + == workspace._descriptor_storages["m128"].data_ptr() + ) + + changed_output = torch.empty_like(output) + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=changed_output, + prefill_workspace=workspace, + ) + assert module.calls[-1][14] == 1 + + +def test_captured_workspace_rejects_eager_reuse_and_capture_mismatch( + cuda_device, monkeypatch +): + monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + module = _RecorderModule() + monkeypatch.setattr( + kda_api, "_get_flash_kda_prefill_module", lambda variant: module + ) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) + output = torch.empty_like(inputs["q"]) + workspace = RecurrentKDAPrefillWorkspace(cuda_device) + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + prefill_workspace=workspace, + ) + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError, match="not warmed for the exact"): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=torch.empty_like(output), + prefill_workspace=workspace, + ) + + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + prefill_workspace=workspace, + ) + assert module.calls[-1][14] == 0 + assert workspace._captured + + with pytest.raises(RuntimeError, match="captured by another CUDA graph"): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + prefill_workspace=workspace, + ) + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + with pytest.raises(RuntimeError, match="cannot be reused eagerly"): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + prefill_workspace=workspace, + ) + + +def test_workspace_rejects_a_different_stream(cuda_device, monkeypatch): + monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + module = _RecorderModule() + monkeypatch.setattr( + kda_api, "_get_flash_kda_prefill_module", lambda variant: module + ) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) + output = torch.empty_like(inputs["q"]) + workspace = RecurrentKDAPrefillWorkspace(cuda_device) + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + prefill_workspace=workspace, + ) + + other_stream = torch.cuda.Stream(device=cuda_device) + with ( + torch.cuda.stream(other_stream), + pytest.raises(RuntimeError, match="different CUDA stream"), + ): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + prefill_workspace=workspace, + ) + + +def test_recurrent_kda_prefill_trace_has_semantic_inputs(): + q = torch.empty((1, 8, 2, 128), dtype=torch.bfloat16) + trace = recurrent_kda.fi_trace( + q=q, + k=q, + v=q, + g=q, + beta=torch.empty((1, 8, 2), dtype=torch.bfloat16), + A_log=torch.empty(2), + dt_bias=torch.empty((2, 128)), + cu_seqlens=torch.tensor([0, 3, 8], dtype=torch.int64), + seq_order=torch.tensor([1, 0], dtype=torch.int32), + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + lower_bound=-5.0, + beta_is_logit=True, + ) + assert trace["op_type"] == "kda" + assert "stage:prefill" in trace["tags"] + for name in ( + "A_log", + "dt_bias", + "cu_seqlens", + "seq_order", + ): + assert name in trace["inputs"] + + +def test_flash_kda_jit_getter_is_importable(): + import flashinfer + from flashinfer.jit.flash_kda import get_flash_kda_prefill_module + + assert callable(get_flash_kda_prefill_module) + assert flashinfer.RecurrentKDAPrefillWorkspace is RecurrentKDAPrefillWorkspace + + +@pytest.mark.parametrize("packed", [False, True]) +@pytest.mark.parametrize("non_default_stream", [False, True]) +def test_frozen_prefill_matches_reference(b200, packed, non_default_stream): + inputs = _make_inputs( + seq_lens=[3, 5] if packed else [4, 4], + num_heads=2, + packed=packed, + initial_state=True, + seed=2026, + ) + reference_inputs = { + **inputs, + "initial_state": inputs["initial_state"].clone(), + } + expected_output, expected_state = _reference(reference_inputs) + output = torch.empty_like(inputs["q"]) + state_identity = inputs["initial_state"] + seq_order = torch.tensor([1, 0], dtype=torch.int32, device=b200) if packed else None + + if non_default_stream: + stream = torch.cuda.Stream(device=b200) + stream.wait_stream(torch.cuda.current_stream(b200)) + with torch.cuda.stream(stream): + actual_output, actual_state = recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + output_final_state=True, + seq_order=seq_order, + ) + stream.synchronize() + else: + actual_output, actual_state = recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + output_final_state=True, + seq_order=seq_order, + ) + + assert actual_output.data_ptr() == output.data_ptr() + assert actual_state is state_identity + torch.testing.assert_close( + actual_output.float(), + expected_output.float(), + atol=1e-2, + rtol=1e-2, + ) + torch.testing.assert_close( + actual_state.float(), + expected_state.float(), + atol=1e-2, + rtol=1e-2, + ) + + +def test_frozen_prefill_without_initial_or_final_state(b200): + inputs = _make_inputs(seq_lens=[3], num_heads=2, packed=False, initial_state=False) + expected_output, _ = _reference(inputs) + output = torch.empty_like(inputs["q"]) + actual_output, actual_state = recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + output_final_state=False, + ) + assert actual_output.data_ptr() == output.data_ptr() + assert actual_state is None + torch.testing.assert_close( + actual_output.float(), + expected_output.float(), + atol=1e-2, + rtol=1e-2, + ) + + +def test_frozen_prefill_m64_matches_reference(b200): + inputs = _make_inputs( + seq_lens=[2], + num_heads=64, + packed=False, + initial_state=True, + seed=2027, + ) + reference_inputs = { + **inputs, + "initial_state": inputs["initial_state"].clone(), + } + expected_output, expected_state = _reference(reference_inputs) + output = torch.empty_like(inputs["q"]) + state_identity = inputs["initial_state"] + + actual_output, actual_state = recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + output_final_state=True, + ) + + assert actual_output.data_ptr() == output.data_ptr() + assert actual_state is state_identity + torch.testing.assert_close( + actual_output.float(), + expected_output.float(), + atol=1e-2, + rtol=1e-2, + ) + torch.testing.assert_close( + actual_state.float(), + expected_state.float(), + atol=1e-2, + rtol=1e-2, + ) + + +@pytest.mark.parametrize( + ("packed", "num_heads", "has_initial_state"), + [(False, 64, True), (True, 2, False)], +) +def test_frozen_prefill_cuda_graph_capture_and_replay( + b200, + packed, + num_heads, + has_initial_state, +): + inputs = _make_inputs( + seq_lens=[1, 2] if packed else [2], + num_heads=num_heads, + packed=packed, + initial_state=has_initial_state, + seed=2028, + ) + initial_state_seed = ( + inputs["initial_state"].clone() if inputs["initial_state"] is not None else None + ) + expected_output, expected_state = _reference( + { + **inputs, + "initial_state": ( + initial_state_seed.clone() if initial_state_seed is not None else None + ), + } + ) + output = torch.empty_like(inputs["q"]) + seq_order = torch.tensor([1, 0], dtype=torch.int32, device=b200) if packed else None + workspace = RecurrentKDAPrefillWorkspace(b200) + capture_stream = torch.cuda.Stream(device=b200) + capture_stream.wait_stream(torch.cuda.current_stream(b200)) + + call_kwargs = { + **_strict_prefill_kwargs(inputs), + "output": output, + "output_final_state": True, + "seq_order": seq_order, + "prefill_workspace": workspace, + } + with torch.cuda.stream(capture_stream): + recurrent_kda(**call_kwargs) + if initial_state_seed is not None: + inputs["initial_state"].copy_(initial_state_seed) + output.zero_() + capture_stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + captured_output, captured_state = recurrent_kda(**call_kwargs) + + with torch.cuda.stream(capture_stream): + if initial_state_seed is not None: + inputs["initial_state"].copy_(initial_state_seed) + output.fill_(float("nan")) + capture_stream.synchronize() + graph.replay() + torch.cuda.synchronize() + + assert captured_output.data_ptr() == output.data_ptr() + if inputs["initial_state"] is None: + assert captured_state.data_ptr() == workspace._state_scratch.data_ptr() + else: + assert captured_state is inputs["initial_state"] + assert workspace._captured + torch.testing.assert_close( + captured_output.float(), + expected_output.float(), + atol=1e-2, + rtol=1e-2, + ) + torch.testing.assert_close( + captured_state.float(), + expected_state.float(), + atol=1e-2, + rtol=1e-2, + ) + + +def test_frozen_prefill_cuda_graph_workspaces_are_isolated(b200): + capture_stream = torch.cuda.Stream(device=b200) + launch_stream = torch.cuda.Stream(device=b200) + bundles = [] + + for seed in (2030, 2031): + inputs = _make_inputs( + seq_lens=[2], + num_heads=2, + packed=False, + initial_state=True, + seed=seed, + ) + state_seed = inputs["initial_state"].clone() + expected_output, expected_state = _reference( + { + **inputs, + "initial_state": state_seed.clone(), + } + ) + output = torch.empty_like(inputs["q"]) + workspace = RecurrentKDAPrefillWorkspace(b200) + call_kwargs = { + **_strict_prefill_kwargs(inputs), + "output": output, + "output_final_state": True, + "prefill_workspace": workspace, + } + capture_stream.wait_stream(torch.cuda.current_stream(b200)) + with torch.cuda.stream(capture_stream): + recurrent_kda(**call_kwargs) + inputs["initial_state"].copy_(state_seed) + output.zero_() + capture_stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + recurrent_kda(**call_kwargs) + bundles.append( + ( + graph, + workspace, + inputs, + state_seed, + output, + expected_output, + expected_state, + ) + ) + + assert ( + bundles[0][1]._state_scratch.data_ptr() + != bundles[1][1]._state_scratch.data_ptr() + ) + assert ( + bundles[0][1]._descriptor_storages["m128"].data_ptr() + != bundles[1][1]._descriptor_storages["m128"].data_ptr() + ) + + for bundle_index in (0, 1, 0, 1): + ( + graph, + _workspace, + inputs, + state_seed, + output, + expected_output, + expected_state, + ) = bundles[bundle_index] + with torch.cuda.stream(launch_stream): + inputs["initial_state"].copy_(state_seed) + output.fill_(float("nan")) + launch_stream.synchronize() + with torch.cuda.stream(launch_stream): + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + output.float(), + expected_output.float(), + atol=1e-2, + rtol=1e-2, + ) + torch.testing.assert_close( + inputs["initial_state"].float(), + expected_state.float(), + atol=1e-2, + rtol=1e-2, + ) diff --git a/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json b/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json index 09e5f6ae981..6767a4bc811 100644 --- a/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json +++ b/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json @@ -1,10 +1,11 @@ { "name": "recurrent_kda_q8_v16_d128", - "description": "Recurrent Key-Driven Attention decode with per-key-dimension gating and an optional read-only committed-state source.", + "description": "Recurrent Key-Driven Attention decode/spec-decode plus exact FlashKDA-compatible ordinary prefill on B200, with per-key-dimension gating and recurrent BF16 state.", "op_type": "kda", "tags": [ "fi_api:flashinfer.kda_decode.recurrent_kda", "stage:decode", + "stage:prefill", "status:verified" ], "axes": { @@ -42,6 +43,10 @@ "num_sequences": { "type": "var", "description": "Number of state-source indices." + }, + "num_sequences_plus_one": { + "type": "var", + "description": "Number of packed cumulative-length entries." } }, "constraints": [ @@ -93,6 +98,23 @@ ], "dtype": "bfloat16" }, + "A_log": { + "shape": [ + "num_q_heads" + ], + "dtype": "float32", + "optional": true, + "description": "FP32 per-query-head log decay rate." + }, + "dt_bias": { + "shape": [ + "num_q_heads", + "head_dim" + ], + "dtype": "float32", + "optional": true, + "description": "FP32 per-head/key decay bias; flattened [H*D] or [H,D] storage is accepted by the API." + }, "initial_state": { "shape": [ "state_pool_size", @@ -122,6 +144,30 @@ "optional": true, "description": "Committed-state slot selected for each sequence." }, + "cu_seqlens": { + "shape": [ + "num_sequences_plus_one" + ], + "dtype": "unknown", + "optional": true, + "description": "Packed cumulative sequence lengths (int32 or int64)." + }, + "num_accepted_tokens": { + "shape": [ + "num_sequences" + ], + "dtype": "int32", + "optional": true, + "description": "Accepted-token counts for speculative decode." + }, + "seq_order": { + "shape": [ + "num_sequences" + ], + "dtype": "int32", + "optional": true, + "description": "Packed-prefill sequence permutation, normally descending by sequence length." + }, "scale": { "shape": null, "dtype": "float32", @@ -166,7 +212,8 @@ "num_v_heads", "head_dim" ], - "dtype": "bfloat16" + "dtype": "bfloat16", + "param": "output" }, "final_state": { "shape": [ From 0fc9909e4beb3e71cba702274f4dbf86951c4464 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 29 Jul 2026 16:35:02 -0700 Subject: [PATCH 2/8] fix(kda): type prefill variant precisely Signed-off-by: Yingyi Huang --- flashinfer/kda_decode.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/flashinfer/kda_decode.py b/flashinfer/kda_decode.py index c283f38e24a..90830857d0f 100644 --- a/flashinfer/kda_decode.py +++ b/flashinfer/kda_decode.py @@ -24,7 +24,7 @@ import math import threading -from typing import Optional +from typing import TYPE_CHECKING, Optional import torch @@ -32,6 +32,9 @@ from .trace.templates.kda import recurrent_kda_trace from .utils import get_compute_capability +if TYPE_CHECKING: + from .jit.flash_kda import FlashKDAVariant + try: from .kda_kernels.recurrent_kda import run_recurrent_kda as _run_recurrent_kda @@ -248,7 +251,7 @@ def _flash_kda_prefill_is_eligible( def _select_flash_kda_prefill_variant( *, fixed_layout: bool, num_sequences: int, num_heads: int -) -> str: +) -> "FlashKDAVariant": if fixed_layout and num_sequences == 1 and num_heads == 64: return "m64" return "m128" @@ -523,7 +526,7 @@ def _validate_prefill_seq_order( return seq_order -def _get_flash_kda_prefill_module(variant: str): +def _get_flash_kda_prefill_module(variant: "FlashKDAVariant"): from .jit.flash_kda import get_flash_kda_prefill_module return get_flash_kda_prefill_module(variant) From f799e87ea3e1df746e56e45a72b04553cbae5777 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Wed, 29 Jul 2026 18:46:47 -0700 Subject: [PATCH 3/8] bench(kda): compare MR458 peer through public semantics --- benchmarks/bench_recurrent_kda_prefill.py | 604 +++++++++++++++++++--- 1 file changed, 531 insertions(+), 73 deletions(-) diff --git a/benchmarks/bench_recurrent_kda_prefill.py b/benchmarks/bench_recurrent_kda_prefill.py index 8d0b66e712e..c5a0bb15caa 100644 --- a/benchmarks/bench_recurrent_kda_prefill.py +++ b/benchmarks/bench_recurrent_kda_prefill.py @@ -14,15 +14,27 @@ """CUPTI benchmark for the six frozen recurrent-KDA prefill contract shapes. -The reported public end-to-end time includes the required same-stream final -state copy-back into ``initial_state``. Input allocation, packed metadata, -sequence ordering, output/state allocation, and JIT/cache warmup are outside -the measured region. +The FlashInfer candidate is always invoked through the public +``recurrent_kda`` API. With ``--flash-kda-peer``, two commit-verified +MoonshotAI/FlashKDA measurements are reported: + +* the raw ``_fwd_raw`` timing scope used as the Official peer in Cake MR !458; +* a public-semantics adapter that follows ``_fwd_raw`` with the same-stream + state copy-back required by ``recurrent_kda``. + +All paths use the exact MR !458 tensors and seeds. Preinitialized rotating +state buffers ensure every timed invocation sees the same initial state. +Allocation, metadata, sequence ordering, build/JIT, and state-pool reset are +outside the measured region. """ import argparse +import hashlib import json +import subprocess from dataclasses import dataclass +from importlib import import_module +from importlib.metadata import version from pathlib import Path from typing import Callable, Optional @@ -36,6 +48,16 @@ from flashinfer.testing import bench_gpu_time from flashinfer.utils import get_compute_capability +FLASH_KDA_PEER_COMMIT = "d2ff19a6a0c82f39f796f637ebd1c36090b1268f" +FLASH_KDA_CUTLASS_COMMIT = "5c149f52a436782210263fb2f19b354443a61c6a" +# This is an opaque identity recorded by the MR !458 contract. The contract +# does not define its derivation, so it is not asserted to be a source or +# rebuilt-binary digest. The actual extension digest is recorded separately. +FLASH_KDA_CONTRACT_IDENTITY = ( + "sha256:997c3a1d1338f8bf9dba3c1a01386b1b74448214c294d64409454cc11141c04c" +) +DEFAULT_STATE_ROTATIONS = 512 + @dataclass(frozen=True) class Case: @@ -43,63 +65,180 @@ class Case: num_heads: int seq_lens: tuple[int, ...] packed: bool + seed: int -CASES = tuple( - Case( - name=f"h{num_heads}_{name}", - num_heads=num_heads, - seq_lens=seq_lens, - packed=packed, - ) - for num_heads in (96, 64) - for name, seq_lens, packed in ( - ("fixed8192", (8192,), False), - ("mixed", (1300, 547, 2048, 963, 271, 3063), True), - ("uniform", (1024,) * 8, True), - ) +@dataclass +class PreparedCase: + candidate_run: Callable[[], tuple[torch.Tensor, Optional[torch.Tensor]]] + peer_raw_run: Optional[Callable[[], None]] + peer_adapted_run: Optional[Callable[[], None]] + reset_state_pools: Callable[[], None] + candidate_output: torch.Tensor + candidate_state_pool: torch.Tensor + peer_raw_output: Optional[torch.Tensor] + peer_raw_final_state: Optional[torch.Tensor] + peer_adapted_output: Optional[torch.Tensor] + peer_adapted_state_pool: Optional[torch.Tensor] + state_cursors: dict[str, list[int]] + metadata: dict + + +CASES = ( + Case("h96_fixed8192", 96, (8192,), False, 10000), + Case("h96_mixed", 96, (1300, 547, 2048, 963, 271, 3063), True, 10001), + Case("h96_uniform", 96, (1024,) * 8, True, 10002), + Case("h64_fixed8192", 64, (8192,), False, 10003), + Case("h64_mixed", 64, (1300, 547, 2048, 963, 271, 3063), True, 10004), + Case("h64_uniform", 64, (1024,) * 8, True, 10005), ) -def _make_case( - case: Case, -) -> tuple[ - Callable[[], tuple[torch.Tensor, Optional[torch.Tensor]]], - dict, -]: +def _require_cupti() -> None: + try: + from cupti import cupti # noqa: F401 + except ImportError as error: + raise RuntimeError("cupti-python >= 13 is required") from error + cupti_version = version("cupti-python") + if int(cupti_version.split(".", 1)[0]) < 13: + raise RuntimeError(f"cupti-python >= 13 is required, found {cupti_version}") + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _git_output(root: Path, *args: str) -> str: + try: + return subprocess.check_output( + ["git", "-C", str(root), *args], + text=True, + stderr=subprocess.STDOUT, + ).strip() + except subprocess.CalledProcessError as error: + raise RuntimeError( + f"failed to verify FlashKDA source provenance at {root}: " + f"{error.output.strip()}" + ) from error + + +def _verify_peer_provenance(flash_kda, source_dir: Path) -> dict: + source_dir = source_dir.resolve(strict=True) + package_path = Path(flash_kda.__file__).resolve(strict=True) + if not package_path.is_relative_to(source_dir): + raise RuntimeError( + "flash_kda must be imported from the verified source checkout: " + f"module={package_path}, checkout={source_dir}" + ) + + source_commit = _git_output(source_dir, "rev-parse", "HEAD") + if source_commit != FLASH_KDA_PEER_COMMIT: + raise RuntimeError( + "unexpected FlashKDA source revision: " + f"expected {FLASH_KDA_PEER_COMMIT}, got {source_commit}" + ) + cutlass_dir = source_dir / "cutlass" + cutlass_commit = _git_output(cutlass_dir, "rev-parse", "HEAD") + if cutlass_commit != FLASH_KDA_CUTLASS_COMMIT: + raise RuntimeError( + "unexpected FlashKDA CUTLASS revision: " + f"expected {FLASH_KDA_CUTLASS_COMMIT}, got {cutlass_commit}" + ) + submodule_record = _git_output(source_dir, "ls-tree", "HEAD", "cutlass").split() + if len(submodule_record) < 3 or submodule_record[2] != cutlass_commit: + raise RuntimeError( + "FlashKDA CUTLASS checkout does not match the pinned gitlink" + ) + tracked_changes = _git_output( + source_dir, + "status", + "--porcelain", + "--untracked-files=no", + ) + if tracked_changes: + raise RuntimeError( + f"verified FlashKDA checkout has tracked modifications:\n{tracked_changes}" + ) + + extension = import_module("flash_kda_C") + extension_path = Path(extension.__file__).resolve(strict=True) + if not extension_path.is_relative_to(source_dir): + raise RuntimeError( + "flash_kda_C must be loaded from the verified source checkout: " + f"extension={extension_path}, checkout={source_dir}" + ) + return { + "repository": "https://github.com/MoonshotAI/FlashKDA.git", + "source_dir": str(source_dir), + "source_commit": source_commit, + "cutlass_commit": cutlass_commit, + "package_path": str(package_path), + "package_sha256": _sha256(package_path), + "extension_path": str(extension_path), + "extension_sha256": _sha256(extension_path), + "contract_recorded_peer_identity": FLASH_KDA_CONTRACT_IDENTITY, + "contract_identity_derivation": "unspecified", + "byte_identical_to_contract_artifact": "not_claimed", + } + + +def _make_state_pool( + initial_state: torch.Tensor, + rotations: int, +) -> torch.Tensor: + return initial_state.unsqueeze(0).expand(rotations, *initial_state.shape).clone() + + +def _make_case(case: Case, *, state_rotations: int, flash_kda=None) -> PreparedCase: total_tokens = sum(case.seq_lens) shape = (1, total_tokens, case.num_heads, 128) - q = torch.randn(shape, dtype=torch.bfloat16, device="cuda") - k = torch.randn(shape, dtype=torch.bfloat16, device="cuda") - v = torch.randn(shape, dtype=torch.bfloat16, device="cuda") - g = (0.1 * torch.randn(shape, dtype=torch.float32, device="cuda")).to( - torch.bfloat16 - ) + generator = torch.Generator(device="cuda").manual_seed(case.seed) + q = torch.randn(shape, generator=generator, device="cuda").to(torch.bfloat16) + k = torch.randn(shape, generator=generator, device="cuda").to(torch.bfloat16) + v = torch.randn(shape, generator=generator, device="cuda").to(torch.bfloat16) + g = torch.randn(shape, generator=generator, device="cuda").to(torch.bfloat16) beta = torch.randn( (1, total_tokens, case.num_heads), - dtype=torch.bfloat16, + generator=generator, device="cuda", + ).to(torch.bfloat16) + A_log = torch.rand( + (case.num_heads,), + generator=generator, + device="cuda", + dtype=torch.float32, ) - A_log = 0.1 * torch.randn(case.num_heads, dtype=torch.float32, device="cuda") - dt_bias = 0.1 * torch.randn( - (case.num_heads, 128), dtype=torch.float32, device="cuda" - ) - state = torch.zeros( - (len(case.seq_lens), case.num_heads, 128, 128), - dtype=torch.bfloat16, + dt_bias = torch.rand( + (case.num_heads, 128), + generator=generator, device="cuda", + dtype=torch.float32, + ) + initial_state = ( + torch.randn( + (len(case.seq_lens), case.num_heads, 128, 128), + generator=generator, + device="cuda", + ) + * 0.25 + ).to(torch.bfloat16) + candidate_state_pool = _make_state_pool(initial_state, state_rotations) + candidate_output = torch.empty_like(q) + candidate_workspace = RecurrentKDAPrefillWorkspace(q.device) + state_cursors = {"pr": [0], "adapted": [0]} + + offsets = [0] + for seq_len in case.seq_lens: + offsets.append(offsets[-1] + seq_len) + cu_seqlens = ( + torch.tensor(offsets, dtype=torch.int64, device="cuda") if case.packed else None ) - output = torch.empty_like(q) - workspace = RecurrentKDAPrefillWorkspace(q.device) - - cu_seqlens = None - seq_order = None - if case.packed: - offsets = [0] - for seq_len in case.seq_lens: - offsets.append(offsets[-1] + seq_len) - cu_seqlens = torch.tensor(offsets, dtype=torch.int64, device="cuda") - seq_order = torch.tensor( + seq_order = ( + torch.tensor( sorted( range(len(case.seq_lens)), key=case.seq_lens.__getitem__, @@ -108,8 +247,18 @@ def _make_case( dtype=torch.int32, device="cuda", ) + if case.packed + else None + ) + scale = float(1.0 / np.sqrt(128.0)) - def run(): + def candidate_run(): + state_index = state_cursors["pr"][0] + if state_index >= state_rotations: + raise RuntimeError( + f"PR state rotations exhausted: {state_index} >= {state_rotations}" + ) + state_cursors["pr"][0] += 1 return recurrent_kda( q=q, k=k, @@ -118,8 +267,9 @@ def run(): beta=beta, A_log=A_log, dt_bias=dt_bias, - initial_state=state, - output=output, + scale=scale, + initial_state=candidate_state_pool[state_index], + output=candidate_output, output_final_state=False, use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True, @@ -127,8 +277,95 @@ def run(): cu_seqlens=cu_seqlens, beta_is_logit=True, seq_order=seq_order, - prefill_workspace=workspace, + prefill_workspace=candidate_workspace, + ) + + peer_raw_run = None + peer_adapted_run = None + peer_raw_output = None + peer_raw_final_state = None + peer_adapted_output = None + peer_adapted_state_pool = None + peer_raw_initial_state = None + if flash_kda is not None: + peer_raw_initial_state = initial_state.clone() + peer_raw_final_state = torch.empty_like(initial_state) + peer_raw_output = torch.empty_like(q) + peer_adapted_state_pool = _make_state_pool( + initial_state, + state_rotations, ) + peer_adapted_final_state = torch.empty_like(initial_state) + peer_adapted_output = torch.empty_like(q) + workspace_size = flash_kda.get_workspace_size( + total_tokens, + case.num_heads, + len(case.seq_lens), + ) + peer_raw_workspace = torch.empty( + workspace_size, + dtype=torch.uint8, + device="cuda", + ) + peer_adapted_workspace = torch.empty( + workspace_size, + dtype=torch.uint8, + device="cuda", + ) + + def peer_raw_run() -> None: + flash_kda._fwd_raw( + q, + k, + v, + g, + beta, + scale, + peer_raw_output, + peer_raw_workspace, + A_log, + dt_bias, + -5.0, + initial_state=peer_raw_initial_state, + final_state=peer_raw_final_state, + cu_seqlens=cu_seqlens, + ) + + def peer_adapted_run() -> None: + state_index = state_cursors["adapted"][0] + if state_index >= state_rotations: + raise RuntimeError( + "adapted-peer state rotations exhausted: " + f"{state_index} >= {state_rotations}" + ) + state_cursors["adapted"][0] += 1 + adapted_state = peer_adapted_state_pool[state_index] + flash_kda._fwd_raw( + q, + k, + v, + g, + beta, + scale, + peer_adapted_output, + peer_adapted_workspace, + A_log, + dt_bias, + -5.0, + initial_state=adapted_state, + final_state=peer_adapted_final_state, + cu_seqlens=cu_seqlens, + ) + adapted_state.copy_(peer_adapted_final_state) + + def reset_state_pools() -> None: + candidate_state_pool.copy_(initial_state.unsqueeze(0)) + state_cursors["pr"][0] = 0 + if peer_raw_initial_state is not None: + peer_raw_initial_state.copy_(initial_state) + if peer_adapted_state_pool is not None: + peer_adapted_state_pool.copy_(initial_state.unsqueeze(0)) + state_cursors["adapted"][0] = 0 metadata = { "name": case.name, @@ -137,14 +374,115 @@ def run(): "total_tokens": total_tokens, "layout": "packed" if case.packed else "fixed", "variant": "m64" if case.name == "h64_fixed8192" else "m128", + "seed": case.seed, + "state_rotation_capacity": state_rotations, } - return run, metadata + return PreparedCase( + candidate_run=candidate_run, + peer_raw_run=peer_raw_run, + peer_adapted_run=peer_adapted_run, + reset_state_pools=reset_state_pools, + candidate_output=candidate_output, + candidate_state_pool=candidate_state_pool, + peer_raw_output=peer_raw_output, + peer_raw_final_state=peer_raw_final_state, + peer_adapted_output=peer_adapted_output, + peer_adapted_state_pool=peer_adapted_state_pool, + state_cursors=state_cursors, + metadata=metadata, + ) + + +def _check_peer(prepared: PreparedCase) -> dict[str, float]: + assert prepared.peer_raw_run is not None + assert prepared.peer_adapted_run is not None + assert prepared.peer_raw_output is not None + assert prepared.peer_raw_final_state is not None + assert prepared.peer_adapted_output is not None + assert prepared.peer_adapted_state_pool is not None + prepared.reset_state_pools() + prepared.candidate_run() + prepared.peer_raw_run() + prepared.peer_adapted_run() + torch.cuda.synchronize() + + candidate_state = prepared.candidate_state_pool[0] + adapted_state = prepared.peer_adapted_state_pool[0] + comparisons = ( + ( + "raw_output_max_abs", + prepared.candidate_output, + prepared.peer_raw_output, + ), + ( + "raw_state_max_abs", + candidate_state, + prepared.peer_raw_final_state, + ), + ( + "adapted_output_max_abs", + prepared.candidate_output, + prepared.peer_adapted_output, + ), + ("adapted_state_max_abs", candidate_state, adapted_state), + ) + diagnostics = {} + for name, actual, expected in comparisons: + diagnostics[name] = float((actual.float() - expected.float()).abs().max()) + torch.testing.assert_close( + actual, + expected, + atol=1e-2, + rtol=1e-2, + ) + return diagnostics + + +def _measure( + run: Callable[[], object], + *, + warmup_ms: int, + bench_ms: int, +) -> tuple[float, list[float]]: + measurements = bench_gpu_time( + run, + enable_cupti=True, + cold_l2_cache=True, + use_cuda_graph=False, + dry_run_time_ms=warmup_ms, + repeat_time_ms=bench_ms, + ) + samples_ms = [float(value) for value in measurements] + return float(np.median(samples_ms)), samples_ms def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--warmup", type=int, default=10) - parser.add_argument("--iters", type=int, default=100) + parser.add_argument("--warmup-ms", type=int, default=20) + parser.add_argument("--bench-ms", type=int, default=100) + parser.add_argument( + "--state-rotations", + type=int, + default=DEFAULT_STATE_ROTATIONS, + help="Number of preinitialized same-input state slots per mutable path.", + ) + parser.add_argument( + "--flash-kda-peer", + action="store_true", + help=( + "Compare against FlashKDA commit " + f"{FLASH_KDA_PEER_COMMIT} using both MR !458 raw and " + "public-semantics-adapted scopes." + ), + ) + parser.add_argument( + "--flash-kda-source-dir", + type=Path, + help=( + "Required with --flash-kda-peer. The imported editable FlashKDA " + "package and extension must resolve inside this exact checkout." + ), + ) parser.add_argument( "--json", type=Path, @@ -152,32 +490,152 @@ def main() -> None: ) args = parser.parse_args() + if args.warmup_ms <= 0 or args.bench_ms <= 0: + parser.error("--warmup-ms and --bench-ms must be positive") + if args.state_rotations <= 0: + parser.error("--state-rotations must be positive") + if args.flash_kda_peer != (args.flash_kda_source_dir is not None): + parser.error( + "--flash-kda-peer and --flash-kda-source-dir must be provided together" + ) if not torch.cuda.is_available(): raise RuntimeError("CUDA is required") if get_compute_capability(torch.device("cuda")) != (10, 0): raise RuntimeError("frozen recurrent-KDA prefill requires B200 (cc 10.0)") + _require_cupti() + + flash_kda = None + peer_provenance = None + if args.flash_kda_peer: + try: + import flash_kda as imported_flash_kda + except ImportError as error: + raise RuntimeError( + "install MoonshotAI/FlashKDA at " + f"{FLASH_KDA_PEER_COMMIT} to run the MR !458 comparison" + ) from error + flash_kda = imported_flash_kda + assert args.flash_kda_source_dir is not None + peer_provenance = _verify_peer_provenance( + flash_kda, + args.flash_kda_source_dir, + ) results = [] for case in CASES: - run, result = _make_case(case) - run() - torch.cuda.synchronize() - measurements = bench_gpu_time( - run, - enable_cupti=True, - cold_l2_cache=True, - use_cuda_graph=False, - dry_run_iters=args.warmup, - repeat_iters=args.iters, - ) - result["median_ms"] = float(np.median(measurements)) - result["median_us"] = result["median_ms"] * 1000.0 - result["timing_scope"] = "public_end_to_end_with_state_copy_back" - results.append(result) - print( - f"{result['name']:<18} {result['variant']:<4} " - f"{result['median_us']:10.3f} us" + prepared = _make_case( + case, + state_rotations=args.state_rotations, + flash_kda=flash_kda, ) + result = dict(prepared.metadata) + if prepared.peer_raw_run is None: + prepared.reset_state_pools() + prepared.candidate_run() + torch.cuda.synchronize() + prepared.reset_state_pools() + candidate_ms, candidate_samples = _measure( + prepared.candidate_run, + warmup_ms=args.warmup_ms, + bench_ms=args.bench_ms, + ) + candidate_block_medians = [candidate_ms] + result["correctness_peer"] = "not_requested" + else: + assert prepared.peer_adapted_run is not None + correctness = _check_peer(prepared) + result.update( + { + "correctness_peer": "passed", + **correctness, + } + ) + samples = {"pr": [], "raw": [], "adapted": []} + block_medians = {"pr": [], "raw": [], "adapted": []} + state_slots_used = {"pr": [], "adapted": []} + # Symmetric ABCCBA order bounds temperature/clock drift while + # retaining two independent medians for every timing scope. + for backend, run in ( + ("pr", prepared.candidate_run), + ("raw", prepared.peer_raw_run), + ("adapted", prepared.peer_adapted_run), + ("adapted", prepared.peer_adapted_run), + ("raw", prepared.peer_raw_run), + ("pr", prepared.candidate_run), + ): + prepared.reset_state_pools() + torch.cuda.synchronize() + block_median, block_samples = _measure( + run, + warmup_ms=args.warmup_ms, + bench_ms=args.bench_ms, + ) + block_medians[backend].append(block_median) + samples[backend].extend(block_samples) + if backend in state_slots_used: + state_slots_used[backend].append(prepared.state_cursors[backend][0]) + del run + + candidate_block_medians = block_medians["pr"] + candidate_samples = samples["pr"] + candidate_ms = float(np.median(candidate_block_medians)) + raw_ms = float(np.median(block_medians["raw"])) + adapted_ms = float(np.median(block_medians["adapted"])) + result.update( + { + "flash_kda_peer_raw_ms": raw_ms, + "flash_kda_peer_raw_samples_ms": samples["raw"], + "flash_kda_peer_raw_block_medians_ms": (block_medians["raw"]), + "speedup_vs_flash_kda_peer_raw": raw_ms / candidate_ms, + "flash_kda_peer_adapted_ms": adapted_ms, + "flash_kda_peer_adapted_samples_ms": samples["adapted"], + "flash_kda_peer_adapted_block_medians_ms": ( + block_medians["adapted"] + ), + "speedup_vs_flash_kda_peer_adapted": (adapted_ms / candidate_ms), + "peer_raw_timing_scope": ("mr458_official_peer_raw_fwd"), + "peer_adapted_timing_scope": ( + "raw_fwd_plus_public_state_copy_back" + ), + "pair_order": "PR/raw/adapted/adapted/raw/PR", + "same_initial_state_per_timed_call": True, + "state_slots_used_per_block": state_slots_used, + "flash_kda_peer_provenance": peer_provenance, + } + ) + + result.update( + { + "median_ms": candidate_ms, + "median_us": candidate_ms * 1000.0, + "samples_ms": candidate_samples, + "block_medians_ms": candidate_block_medians, + "timing_backend": "cupti", + "cold_l2": True, + "cuda_graph": False, + "timing_scope": ("public_recurrent_kda_with_state_copy_back"), + "warmup_ms": args.warmup_ms, + "bench_ms": args.bench_ms, + } + ) + results.append(result) + if prepared.peer_raw_run is None: + print( + f"{result['name']:<18} {result['variant']:<4} " + f"{result['median_us']:10.3f} us" + ) + else: + print( + f"{result['name']:<18} {result['variant']:<4} " + f"PR {result['median_us']:10.3f} us " + f"raw {result['flash_kda_peer_raw_ms'] * 1000.0:10.3f} us " + f"{result['speedup_vs_flash_kda_peer_raw']:.4f}x " + f"adapted " + f"{result['flash_kda_peer_adapted_ms'] * 1000.0:10.3f} us " + f"{result['speedup_vs_flash_kda_peer_adapted']:.4f}x" + ) + del prepared + torch.cuda.empty_cache() if args.json is not None: args.json.write_text(json.dumps(results, indent=2) + "\n") From f6c3b0c486787e64f76b2be91ddb239ea3ba66ce Mon Sep 17 00:00:00 2001 From: Zihao Date: Thu, 30 Jul 2026 03:06:21 +0000 Subject: [PATCH 4/8] fix(kda): update recurrent state in place --- benchmarks/bench_recurrent_kda_prefill.py | 11 +++++---- csrc/kda/flashkda_bf16_fused_m128.cu | 4 ++- csrc/kda/flashkda_bf16_fused_m64.cu | 4 ++- csrc/kda/flashkda_binding_common.cuh | 13 +++++++++- docs/api/kda_decode.rst | 30 +++++++++++------------ flashinfer/kda_decode.py | 17 ++++++------- tests/jit/test_flash_kda_jit.py | 24 +++++++++++++++--- tests/kda/test_recurrent_kda_prefill.py | 17 ++++++------- 8 files changed, 75 insertions(+), 45 deletions(-) diff --git a/benchmarks/bench_recurrent_kda_prefill.py b/benchmarks/bench_recurrent_kda_prefill.py index c5a0bb15caa..37c523cc667 100644 --- a/benchmarks/bench_recurrent_kda_prefill.py +++ b/benchmarks/bench_recurrent_kda_prefill.py @@ -20,12 +20,13 @@ * the raw ``_fwd_raw`` timing scope used as the Official peer in Cake MR !458; * a public-semantics adapter that follows ``_fwd_raw`` with the same-stream - state copy-back required by ``recurrent_kda``. + state copy-back required to emulate ``recurrent_kda`` on FlashKDA. All paths use the exact MR !458 tensors and seeds. Preinitialized rotating -state buffers ensure every timed invocation sees the same initial state. -Allocation, metadata, sequence ordering, build/JIT, and state-pool reset are -outside the measured region. +state buffers ensure every timed invocation sees the same initial state. The +FlashInfer path updates each state slot in place inside the kernel; it has no +state scratch or copy-back. Allocation, metadata, sequence ordering, build/JIT, +and state-pool reset are outside the measured region. """ import argparse @@ -613,7 +614,7 @@ def main() -> None: "timing_backend": "cupti", "cold_l2": True, "cuda_graph": False, - "timing_scope": ("public_recurrent_kda_with_state_copy_back"), + "timing_scope": ("public_recurrent_kda_with_inplace_state_update"), "warmup_ms": args.warmup_ms, "bench_ms": args.bench_ms, } diff --git a/csrc/kda/flashkda_bf16_fused_m128.cu b/csrc/kda/flashkda_bf16_fused_m128.cu index bf2a8618986..3c1e4090cfd 100644 --- a/csrc/kda/flashkda_bf16_fused_m128.cu +++ b/csrc/kda/flashkda_bf16_fused_m128.cu @@ -497,7 +497,9 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(1024) void -kernel_flashkda_bf16_fused_m128(__nv_bfloat16* __restrict__ q, const void* __restrict__ q_tma, __nv_bfloat16* __restrict__ k, const void* __restrict__ k_tma, __nv_bfloat16* __restrict__ v, const void* __restrict__ v_tma, __nv_bfloat16* __restrict__ g, const void* __restrict__ g_tma, __nv_bfloat16* __restrict__ beta, const void* __restrict__ beta_tma, float* __restrict__ A_log, float* __restrict__ dt_bias, long long* __restrict__ cu_seqlens, int* __restrict__ seq_order, __nv_bfloat16* __restrict__ initial_state, __nv_bfloat16* __restrict__ out, const void* __restrict__ out_tma, __nv_bfloat16* __restrict__ final_state, int num_heads, int use_initial_state, int store_final_state, float scale, float lower_bound) +// FLASHINFER INTEGRATION BEGIN: allow exact state alias +kernel_flashkda_bf16_fused_m128(__nv_bfloat16* __restrict__ q, const void* __restrict__ q_tma, __nv_bfloat16* __restrict__ k, const void* __restrict__ k_tma, __nv_bfloat16* __restrict__ v, const void* __restrict__ v_tma, __nv_bfloat16* __restrict__ g, const void* __restrict__ g_tma, __nv_bfloat16* __restrict__ beta, const void* __restrict__ beta_tma, float* __restrict__ A_log, float* __restrict__ dt_bias, long long* __restrict__ cu_seqlens, int* __restrict__ seq_order, __nv_bfloat16* initial_state, __nv_bfloat16* __restrict__ out, const void* __restrict__ out_tma, __nv_bfloat16* final_state, int num_heads, int use_initial_state, int store_final_state, float scale, float lower_bound) +// FLASHINFER INTEGRATION END: allow exact state alias { // FLASHINFER INTEGRATION BEGIN: acquire global tensor maps // CUDA kernel-start ordering does not acquire the tensor-map proxy. diff --git a/csrc/kda/flashkda_bf16_fused_m64.cu b/csrc/kda/flashkda_bf16_fused_m64.cu index 4849c355b51..0be522d9126 100644 --- a/csrc/kda/flashkda_bf16_fused_m64.cu +++ b/csrc/kda/flashkda_bf16_fused_m64.cu @@ -494,7 +494,9 @@ __device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { extern "C" { __global__ __launch_bounds__(1024) void -kernel_flashkda_bf16_fused_m64(__nv_bfloat16* __restrict__ q, const void* __restrict__ q_tma, __nv_bfloat16* __restrict__ k, const void* __restrict__ k_tma, __nv_bfloat16* __restrict__ v, const void* __restrict__ v_tma, __nv_bfloat16* __restrict__ g, const void* __restrict__ g_tma, __nv_bfloat16* __restrict__ beta, const void* __restrict__ beta_tma, float* __restrict__ A_log, float* __restrict__ dt_bias, long long* __restrict__ cu_seqlens, int* __restrict__ seq_order, __nv_bfloat16* __restrict__ initial_state, __nv_bfloat16* __restrict__ out, const void* __restrict__ out_tma, __nv_bfloat16* __restrict__ final_state, int num_heads, int use_initial_state, int store_final_state, float scale, float lower_bound) +// FLASHINFER INTEGRATION BEGIN: allow exact state alias +kernel_flashkda_bf16_fused_m64(__nv_bfloat16* __restrict__ q, const void* __restrict__ q_tma, __nv_bfloat16* __restrict__ k, const void* __restrict__ k_tma, __nv_bfloat16* __restrict__ v, const void* __restrict__ v_tma, __nv_bfloat16* __restrict__ g, const void* __restrict__ g_tma, __nv_bfloat16* __restrict__ beta, const void* __restrict__ beta_tma, float* __restrict__ A_log, float* __restrict__ dt_bias, long long* __restrict__ cu_seqlens, int* __restrict__ seq_order, __nv_bfloat16* initial_state, __nv_bfloat16* __restrict__ out, const void* __restrict__ out_tma, __nv_bfloat16* final_state, int num_heads, int use_initial_state, int store_final_state, float scale, float lower_bound) +// FLASHINFER INTEGRATION END: allow exact state alias { // FLASHINFER INTEGRATION BEGIN: acquire global tensor maps // CUDA kernel-start ordering does not acquire the tensor-map proxy. diff --git a/csrc/kda/flashkda_binding_common.cuh b/csrc/kda/flashkda_binding_common.cuh index c8a5b3afcd6..0e4bc0c8152 100644 --- a/csrc/kda/flashkda_binding_common.cuh +++ b/csrc/kda/flashkda_binding_common.cuh @@ -87,6 +87,17 @@ inline void CheckNoOverlap(const TensorView& lhs, const char* lhs_name, const Te << ": the frozen kernel uses __restrict__ pointers"; } +inline void CheckNoPartialOverlapOrExactAlias(const TensorView& lhs, const char* lhs_name, + const TensorView& rhs, const char* rhs_name) { + const TensorByteRange lhs_range = GetTensorByteRange(lhs, lhs_name); + const TensorByteRange rhs_range = GetTensorByteRange(rhs, rhs_name); + const bool overlaps = lhs_range.begin < rhs_range.end && rhs_range.begin < lhs_range.end; + const bool exact_alias = lhs_range.begin == rhs_range.begin && lhs_range.end == rhs_range.end; + TVM_FFI_ICHECK(!overlaps || exact_alias) + << lhs_name << " and " << rhs_name + << " must either be disjoint or exactly alias the same state storage"; +} + inline void CheckExactSm100a(int32_t device_id) { int major = 0; int minor = 0; @@ -245,7 +256,7 @@ inline int64_t CheckCommonInputs(const TensorView& q, const TensorView& k, const CheckNoOverlap(final_state, "final_state", cu_seqlens, "cu_seqlens"); CheckNoOverlap(final_state, "final_state", seq_order, "seq_order"); if (use_initial_state != 0) { - CheckNoOverlap(initial_state, "initial_state", final_state, "final_state"); + CheckNoPartialOverlapOrExactAlias(initial_state, "initial_state", final_state, "final_state"); } } return num_seqs; diff --git a/docs/api/kda_decode.rst b/docs/api/kda_decode.rst index ec16c55cd54..507febb417b 100644 --- a/docs/api/kda_decode.rst +++ b/docs/api/kda_decode.rst @@ -64,11 +64,10 @@ State and graph semantics The BF16 state layout remains ``[N,H,V,K]`` and an explicitly supplied ``initial_state`` is still updated in place, even when -``output_final_state=False``. The frozen CUDA ABI marks initial and final -state pointers ``__restrict__``; FlashInfer therefore launches into separate -per-stream scratch storage and enqueues a same-stream copy back into -``initial_state``. If no initial state is supplied, a final state is allocated -only when ``output_final_state=True``. +``output_final_state=False``. The frozen kernels load each CTA's disjoint +state rows before writing the final rows back to the same storage, so no +separate state scratch or copy-back is required. If no initial state is +supplied, a final state is allocated only when ``output_final_state=True``. The frozen kernel uses restricted output storage. A preallocated ``output`` must not overlap Q, K, V, G, beta, or ``initial_state``. @@ -79,14 +78,15 @@ be used during CUDA graph capture. CUDA graph capture requires a caller-owned ``RecurrentKDAPrefillWorkspace(device)`` and a preallocated ``output``. The -workspace owns stable state scratch, beta padding, and separate 768-byte M64 -and M128 TMA descriptor blocks. It binds to the device and CUDA stream of its -first ``recurrent_kda`` call. Warm it eagerly on the intended capture stream -with the exact Q, K, V, G, beta, and output tensors, then synchronize that -stream before capture. Packed graphs must also pass preallocated int64 -``cu_seqlens`` and int32 ``seq_order``. The warm call prepares descriptors; -capture accepts only the exact warmed pointer, shape, stride, and dtype -signature and performs no descriptor preparation. +workspace owns optional final-state scratch for calls without an initial +state, beta padding, and separate 768-byte M64 and M128 TMA descriptor blocks. +It binds to the device and CUDA stream of its first ``recurrent_kda`` call. +Warm it eagerly on the intended capture stream with the exact Q, K, V, G, +beta, and output tensors, then synchronize that stream before capture. Packed +graphs must also pass preallocated int64 ``cu_seqlens`` and int32 +``seq_order``. The warm call prepares descriptors; capture accepts only the +exact warmed pointer, shape, stride, and dtype signature and performs no +descriptor preparation. The workspace must outlive its graph and every replay. Use one distinct workspace for each captured ``recurrent_kda`` invocation, including two KDA @@ -100,6 +100,6 @@ stream. When an explicit workspace is used with ``initial_state=None`` and ``output_final_state=True``, the returned final state is workspace-owned -stable scratch. Otherwise an explicitly supplied ``initial_state`` keeps its -usual in-place copy-back semantics. The small-head ``H < 8`` path captures +stable scratch. Otherwise an explicitly supplied ``initial_state`` is updated +directly in place by the frozen kernel. The small-head ``H < 8`` path captures the beta copy into workspace-owned padded storage before the frozen launch. diff --git a/flashinfer/kda_decode.py b/flashinfer/kda_decode.py index 90830857d0f..8cad532a02e 100644 --- a/flashinfer/kda_decode.py +++ b/flashinfer/kda_decode.py @@ -82,8 +82,9 @@ class RecurrentKDAPrefillWorkspace(_RecurrentKDAPrefillWorkspaceBase): Construct one workspace per captured :func:`recurrent_kda` invocation on the graph's CUDA device. Warm it by invoking :func:`recurrent_kda` eagerly with the exact tensors and capture stream, then synchronize that stream - before capture. The workspace owns stable state scratch, beta padding, and - M64/M128 TMA descriptor storage for the lifetime of the graph. + before capture. The workspace owns optional final-state scratch for calls + without an initial state, beta padding, and M64/M128 TMA descriptor storage + for the lifetime of the graph. A workspace binds to its first stream. Once it participates in capture it cannot be passed to Python again, either eagerly or in another capture. @@ -612,6 +613,7 @@ def _run_flash_kda_prefill( use_initial_state = initial_state is not None if initial_state is not None: initial_state_arg = initial_state + final_state_arg = initial_state store_final_state = True returned_state = initial_state elif output_final_state: @@ -649,8 +651,8 @@ def _run_flash_kda_prefill( else: workspace = prefill_workspace # TVM FFI may release the GIL. Serialize the complete shared-workspace - # enqueue sequence so two host threads cannot interleave preparation, - # launch, or state copy-back on the same CUDA stream. + # enqueue sequence so two host threads cannot interleave preparation or + # launch on the same CUDA stream. with workspace._lock: _bind_workspace( workspace, @@ -660,10 +662,7 @@ def _run_flash_kda_prefill( explicit=explicit_workspace, ) beta_tma = _beta_tma_source(beta, workspace) - if initial_state is not None or (output_final_state and explicit_workspace): - # The frozen ABI marks both state pointers __restrict__. Always - # launch into distinct storage, then preserve recurrent_kda's - # in-place mutation with a same-stream copy-back. + if initial_state is None and output_final_state and explicit_workspace: final_state_arg = _state_scratch( workspace=workspace, device=q.device, @@ -722,8 +721,6 @@ def _run_flash_kda_prefill( raise if prepare_descriptors: workspace._descriptor_signatures[variant] = signature - if initial_state is not None: - initial_state.copy_(final_state_arg) if capturing and explicit_workspace: workspace._captured = True return ( diff --git a/tests/jit/test_flash_kda_jit.py b/tests/jit/test_flash_kda_jit.py index f627d151be9..6ab0e2f420c 100644 --- a/tests/jit/test_flash_kda_jit.py +++ b/tests/jit/test_flash_kda_jit.py @@ -84,9 +84,27 @@ def test_flash_kda_uri_and_jit_spec(monkeypatch, variant, smem_bytes, generated_ assert integration_prologue.count("__syncthreads();") == 1 for tensor_map in ("q_tma", "k_tma", "v_tma", "g_tma", "beta_tma", "out_tma"): assert f'"l"({tensor_map})' in integration_prologue - # Keep the exporter output immutable below the narrowly marked FlashInfer - # integration prologue. - normalized_generated_body = generated_prefix + generated_suffix + generated_body_without_tma_integration = generated_prefix + generated_suffix + alias_begin = "// FLASHINFER INTEGRATION BEGIN: allow exact state alias\n" + alias_end = "// FLASHINFER INTEGRATION END: allow exact state alias\n" + alias_prefix, begin_marker, alias_tail = ( + generated_body_without_tma_integration.partition(alias_begin) + ) + alias_signature, end_marker, alias_suffix = alias_tail.partition(alias_end) + assert begin_marker == alias_begin + assert end_marker == alias_end + assert alias_signature.count("__nv_bfloat16* initial_state") == 1 + assert alias_signature.count("__nv_bfloat16* final_state") == 1 + restricted_alias_signature = alias_signature.replace( + "__nv_bfloat16* initial_state", + "__nv_bfloat16* __restrict__ initial_state", + ).replace( + "__nv_bfloat16* final_state", + "__nv_bfloat16* __restrict__ final_state", + ) + # Keep the exporter output immutable outside the two narrowly marked + # FlashInfer integration patches. + normalized_generated_body = alias_prefix + restricted_alias_signature + alias_suffix assert ( hashlib.sha256(normalized_generated_body.encode()).hexdigest() == generated_sha256 diff --git a/tests/kda/test_recurrent_kda_prefill.py b/tests/kda/test_recurrent_kda_prefill.py index 916832079b3..782dc95e7d5 100644 --- a/tests/kda/test_recurrent_kda_prefill.py +++ b/tests/kda/test_recurrent_kda_prefill.py @@ -320,7 +320,7 @@ def test_frozen_route_rejects_output_overlap(cuda_device, monkeypatch): assert module.calls == [] -def test_initial_state_uses_distinct_final_and_copies_back(cuda_device, monkeypatch): +def test_initial_state_is_updated_in_place(cuda_device, monkeypatch): monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) module = _RecorderModule(final_value=0.25) monkeypatch.setattr( @@ -337,7 +337,7 @@ def test_initial_state_uses_distinct_final_and_copies_back(cuda_device, monkeypa assert returned_state is original_state (args,) = module.calls assert args[10].data_ptr() == original_state.data_ptr() - assert args[12].data_ptr() != original_state.data_ptr() + assert args[12].data_ptr() == original_state.data_ptr() assert args[16] == 1 assert args[17] == 1 torch.testing.assert_close( @@ -346,7 +346,9 @@ def test_initial_state_uses_distinct_final_and_copies_back(cuda_device, monkeypa ) -def test_stream_workspace_retains_only_largest_state_buffer(cuda_device, monkeypatch): +def test_stream_workspace_does_not_allocate_state_scratch_for_inplace_update( + cuda_device, monkeypatch +): monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) monkeypatch.setattr(kda_api, "_flash_kda_stream_workspaces", {}) module = _RecorderModule(final_value=0.0) @@ -381,8 +383,7 @@ def test_stream_workspace_retains_only_largest_state_buffer(cuda_device, monkeyp assert len(kda_api._flash_kda_stream_workspaces) == 1 (workspace,) = kda_api._flash_kda_stream_workspaces.values() - largest_state_numel = max(inputs["initial_state"].numel() for inputs in cases) - assert workspace._state_scratch.numel() == largest_state_numel + assert workspace._state_scratch is None assert workspace._beta_padding.numel() == 32 * 8 @@ -824,10 +825,8 @@ def test_frozen_prefill_cuda_graph_workspaces_are_isolated(b200): ) ) - assert ( - bundles[0][1]._state_scratch.data_ptr() - != bundles[1][1]._state_scratch.data_ptr() - ) + assert bundles[0][1]._state_scratch is None + assert bundles[1][1]._state_scratch is None assert ( bundles[0][1]._descriptor_storages["m128"].data_ptr() != bundles[1][1]._descriptor_storages["m128"].data_ptr() From 845b118de980aef1ff164c2bde26b6d794cbb2d5 Mon Sep 17 00:00:00 2001 From: Zihao Date: Thu, 30 Jul 2026 03:31:54 +0000 Subject: [PATCH 5/8] docs(kda): remove internal benchmark provenance --- benchmarks/bench_recurrent_kda_prefill.py | 27 ++++++++--------------- csrc/kda/flashkda_bf16_fused_m128.cu | 2 +- csrc/kda/flashkda_bf16_fused_m64.cu | 2 +- tests/jit/test_flash_kda_jit.py | 8 ++++--- 4 files changed, 16 insertions(+), 23 deletions(-) diff --git a/benchmarks/bench_recurrent_kda_prefill.py b/benchmarks/bench_recurrent_kda_prefill.py index 37c523cc667..4f91114607b 100644 --- a/benchmarks/bench_recurrent_kda_prefill.py +++ b/benchmarks/bench_recurrent_kda_prefill.py @@ -18,15 +18,15 @@ ``recurrent_kda`` API. With ``--flash-kda-peer``, two commit-verified MoonshotAI/FlashKDA measurements are reported: -* the raw ``_fwd_raw`` timing scope used as the Official peer in Cake MR !458; +* the raw ``_fwd_raw`` kernel timing scope; * a public-semantics adapter that follows ``_fwd_raw`` with the same-stream state copy-back required to emulate ``recurrent_kda`` on FlashKDA. -All paths use the exact MR !458 tensors and seeds. Preinitialized rotating -state buffers ensure every timed invocation sees the same initial state. The -FlashInfer path updates each state slot in place inside the kernel; it has no -state scratch or copy-back. Allocation, metadata, sequence ordering, build/JIT, -and state-pool reset are outside the measured region. +All paths use the same deterministic tensors and seeds. Preinitialized +rotating state buffers ensure every timed invocation sees the same initial +state. The FlashInfer path updates each state slot in place inside the kernel; +it has no state scratch or copy-back. Allocation, metadata, sequence ordering, +build/JIT, and state-pool reset are outside the measured region. """ import argparse @@ -51,12 +51,6 @@ FLASH_KDA_PEER_COMMIT = "d2ff19a6a0c82f39f796f637ebd1c36090b1268f" FLASH_KDA_CUTLASS_COMMIT = "5c149f52a436782210263fb2f19b354443a61c6a" -# This is an opaque identity recorded by the MR !458 contract. The contract -# does not define its derivation, so it is not asserted to be a source or -# rebuilt-binary digest. The actual extension digest is recorded separately. -FLASH_KDA_CONTRACT_IDENTITY = ( - "sha256:997c3a1d1338f8bf9dba3c1a01386b1b74448214c294d64409454cc11141c04c" -) DEFAULT_STATE_ROTATIONS = 512 @@ -181,9 +175,6 @@ def _verify_peer_provenance(flash_kda, source_dir: Path) -> dict: "package_sha256": _sha256(package_path), "extension_path": str(extension_path), "extension_sha256": _sha256(extension_path), - "contract_recorded_peer_identity": FLASH_KDA_CONTRACT_IDENTITY, - "contract_identity_derivation": "unspecified", - "byte_identical_to_contract_artifact": "not_claimed", } @@ -472,7 +463,7 @@ def main() -> None: action="store_true", help=( "Compare against FlashKDA commit " - f"{FLASH_KDA_PEER_COMMIT} using both MR !458 raw and " + f"{FLASH_KDA_PEER_COMMIT} using both raw and " "public-semantics-adapted scopes." ), ) @@ -513,7 +504,7 @@ def main() -> None: except ImportError as error: raise RuntimeError( "install MoonshotAI/FlashKDA at " - f"{FLASH_KDA_PEER_COMMIT} to run the MR !458 comparison" + f"{FLASH_KDA_PEER_COMMIT} to run the peer comparison" ) from error flash_kda = imported_flash_kda assert args.flash_kda_source_dir is not None @@ -594,7 +585,7 @@ def main() -> None: block_medians["adapted"] ), "speedup_vs_flash_kda_peer_adapted": (adapted_ms / candidate_ms), - "peer_raw_timing_scope": ("mr458_official_peer_raw_fwd"), + "peer_raw_timing_scope": ("flash_kda_raw_fwd"), "peer_adapted_timing_scope": ( "raw_fwd_plus_public_state_copy_back" ), diff --git a/csrc/kda/flashkda_bf16_fused_m128.cu b/csrc/kda/flashkda_bf16_fused_m128.cu index 3c1e4090cfd..ffaada1f9f6 100644 --- a/csrc/kda/flashkda_bf16_fused_m128.cu +++ b/csrc/kda/flashkda_bf16_fused_m128.cu @@ -16,7 +16,7 @@ // clang-format off // Generated by tools/export-generated-programs (device kernel). -// Provenance: loom @ 8437e0515b212e7973b196c4ab680d3d90c1209c. schedule symbol 'flashkda_bf16_fused_m128'; embedded in the host TU as flashkda_bf16_fused_m128_f0217be48b. +// Provenance: generated Loom schedule 'flashkda_bf16_fused_m128'; embedded in the host TU as flashkda_bf16_fused_m128_f0217be48b. typedef unsigned char uint8_t; typedef unsigned short uint16_t; typedef unsigned int uint32_t; diff --git a/csrc/kda/flashkda_bf16_fused_m64.cu b/csrc/kda/flashkda_bf16_fused_m64.cu index 0be522d9126..5f8541897e4 100644 --- a/csrc/kda/flashkda_bf16_fused_m64.cu +++ b/csrc/kda/flashkda_bf16_fused_m64.cu @@ -16,7 +16,7 @@ // clang-format off // Generated by tools/export-generated-programs (device kernel). -// Provenance: loom @ 8437e0515b212e7973b196c4ab680d3d90c1209c. schedule symbol 'flashkda_bf16_fused_m64'; embedded in the host TU as flashkda_bf16_fused_m64_f0217be48b. +// Provenance: generated Loom schedule 'flashkda_bf16_fused_m64'; embedded in the host TU as flashkda_bf16_fused_m64_f0217be48b. typedef unsigned char uint8_t; typedef unsigned short uint16_t; typedef unsigned int uint32_t; diff --git a/tests/jit/test_flash_kda_jit.py b/tests/jit/test_flash_kda_jit.py index 6ab0e2f420c..507683b91e5 100644 --- a/tests/jit/test_flash_kda_jit.py +++ b/tests/jit/test_flash_kda_jit.py @@ -27,12 +27,12 @@ ( "m64", 219136, - "468780f04768c949b22c688c1b85d235b6ffe050fd3327394fcdda7ea5112203", + "c28aacd475983c72ffe84acac7321a0b2e1c495d7c6e9cdc4a80ada112d76515", ), ( "m128", 227328, - "a2378074fde64fa454edb972dc51e188294dcb89369ee5ea1153f6c67200f1ab", + "e6ea814f0f2e0e0cb33c1562458de9e47272760562dbcb2364855c5b48f0b6ce", ), ], ) @@ -60,7 +60,9 @@ def test_flash_kda_uri_and_jit_spec(monkeypatch, variant, smem_bytes, generated_ ) frozen_source = spec.sources[0].parent / f"flashkda_bf16_fused_{variant}.cu" frozen_text = frozen_source.read_text() - assert "Provenance: loom @ 8437e0515b212e7973b196c4ab680d3d90c1209c" in frozen_text + assert f"Provenance: generated Loom schedule 'flashkda_bf16_fused_{variant}'" in ( + frozen_text + ) assert f"#define SMEM_TOTAL {smem_bytes}" in frozen_text assert frozen_text.count("// clang-format off") == 1 assert frozen_text.rstrip().endswith("// clang-format on") From 9d1f8342849acf6f8af997553b3ae96c08a657a8 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Sat, 1 Aug 2026 01:51:42 -0700 Subject: [PATCH 6/8] refactor(kda): use phase-neutral API module name Rename the canonical recurrent KDA API and documentation from kda_decode to kda now that one dispatcher serves decode, speculative decode, and prefill. Preserve existing direct imports through a thin compatibility package and update trace identity and Kimi Delta Attention naming.\n\nAI-assisted: Codex --- benchmarks/bench_recurrent_kda.py | 4 ++-- benchmarks/bench_recurrent_kda_prefill.py | 2 +- docs/api/{kda_decode.rst => kda.rst} | 13 ++++++++----- docs/index.rst | 2 +- flashinfer/__init__.py | 4 ++-- flashinfer/{kda_decode.py => kda.py} | 0 flashinfer/kda_decode/__init__.py | 14 ++++++++++++++ flashinfer/kda_kernels/__init__.py | 2 +- flashinfer/trace/templates/kda.py | 4 ++-- tests/kda/test_recurrent_kda.py | 2 +- tests/kda/test_recurrent_kda_prefill.py | 11 +++++++++-- tests/trace/example.py | 4 ++-- .../fi_trace_out/recurrent_kda_q8_v16_d128.json | 4 ++-- tests/trace/test_fi_trace.py | 6 +++--- tests/trace/test_fi_trace_template_consistency.py | 2 +- 15 files changed, 49 insertions(+), 25 deletions(-) rename docs/api/{kda_decode.rst => kda.rst} (94%) rename flashinfer/{kda_decode.py => kda.py} (100%) create mode 100644 flashinfer/kda_decode/__init__.py diff --git a/benchmarks/bench_recurrent_kda.py b/benchmarks/bench_recurrent_kda.py index aef5b7ca391..b0172d2aef7 100644 --- a/benchmarks/bench_recurrent_kda.py +++ b/benchmarks/bench_recurrent_kda.py @@ -35,7 +35,7 @@ # Import the recurrent KDA kernel try: - from flashinfer.kda_decode import _RECURRENT_KDA_AVAILABLE, recurrent_kda + from flashinfer.kda import _RECURRENT_KDA_AVAILABLE, recurrent_kda RECURRENT_KDA_AVAILABLE = _RECURRENT_KDA_AVAILABLE except ImportError: @@ -233,7 +233,7 @@ def run_recurrent_kda_benchmark(args, dtype): """Run recurrent KDA decode benchmarks.""" if not RECURRENT_KDA_AVAILABLE: print("Error: recurrent KDA kernel is not available.") - print("Make sure flashinfer.kda_decode.recurrent_kda is importable.") + print("Make sure flashinfer.kda.recurrent_kda is importable.") return invalid_seq_lens = [t for t in args.seq_len if t < 1] diff --git a/benchmarks/bench_recurrent_kda_prefill.py b/benchmarks/bench_recurrent_kda_prefill.py index 4f91114607b..0d5e41a588f 100644 --- a/benchmarks/bench_recurrent_kda_prefill.py +++ b/benchmarks/bench_recurrent_kda_prefill.py @@ -42,7 +42,7 @@ import numpy as np import torch -from flashinfer.kda_decode import ( +from flashinfer.kda import ( RecurrentKDAPrefillWorkspace, recurrent_kda, ) diff --git a/docs/api/kda_decode.rst b/docs/api/kda.rst similarity index 94% rename from docs/api/kda_decode.rst rename to docs/api/kda.rst index 507febb417b..e8742e12aba 100644 --- a/docs/api/kda_decode.rst +++ b/docs/api/kda.rst @@ -1,14 +1,17 @@ -.. _apikda_decode: +.. _apikda: -flashinfer.kda_decode -===================== +flashinfer.kda +============== -Recurrent Key-Driven Attention (KDA) API. Single-token decode, speculative +Recurrent Kimi Delta Attention (KDA) API. Single-token decode, speculative decode, GQA, state-pool indexing, and checkpoint modes use the CuTe-DSL backend under ``flashinfer.kda_kernels``. A strict ordinary multi-token prefill subset is dispatched to frozen FlashKDA-compatible SM100a kernels. -.. currentmodule:: flashinfer.kda_decode +The former ``flashinfer.kda_decode`` module remains as a compatibility import; +new code should import this phase-neutral module. + +.. currentmodule:: flashinfer.kda .. autosummary:: :toctree: ../generated diff --git a/docs/index.rst b/docs/index.rst index 3e2f7d25709..433e2afb284 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -53,7 +53,7 @@ FlashInfer is a library and kernel generator for Large Language Models that prov api/activation api/gdn_decode api/gdn_prefill - api/kda_decode + api/kda api/mamba api/mhc api/quantization diff --git a/flashinfer/__init__.py b/flashinfer/__init__.py index 3fec69a5495..35fb39a1260 100644 --- a/flashinfer/__init__.py +++ b/flashinfer/__init__.py @@ -124,10 +124,10 @@ from .grouped_mm import grouped_mm_fp8 as grouped_mm_fp8 from .grouped_mm import grouped_mm_mxfp8 as grouped_mm_mxfp8 from .grouped_mm import grouped_mm_fp4 as grouped_mm_fp4 -from .kda_decode import ( +from .kda import ( RecurrentKDAPrefillWorkspace as RecurrentKDAPrefillWorkspace, ) -from .kda_decode import recurrent_kda as recurrent_kda +from .kda import recurrent_kda as recurrent_kda from .mla import BatchMLAPagedAttentionWrapper as BatchMLAPagedAttentionWrapper from . import mhc as mhc from . import msa_ops as msa_ops diff --git a/flashinfer/kda_decode.py b/flashinfer/kda.py similarity index 100% rename from flashinfer/kda_decode.py rename to flashinfer/kda.py diff --git a/flashinfer/kda_decode/__init__.py b/flashinfer/kda_decode/__init__.py new file mode 100644 index 00000000000..44ed9598023 --- /dev/null +++ b/flashinfer/kda_decode/__init__.py @@ -0,0 +1,14 @@ +"""Compatibility imports for the former recurrent-KDA module name. + +The canonical recurrent KDA API lives in :mod:`flashinfer.kda`. This module +keeps existing direct imports working while new code uses the phase-neutral +module name. +""" + +from ..kda import ( + RecurrentKDAPrefillWorkspace as RecurrentKDAPrefillWorkspace, +) +from ..kda import _RECURRENT_KDA_AVAILABLE as _RECURRENT_KDA_AVAILABLE +from ..kda import recurrent_kda as recurrent_kda + +__all__ = ["RecurrentKDAPrefillWorkspace", "recurrent_kda"] diff --git a/flashinfer/kda_kernels/__init__.py b/flashinfer/kda_kernels/__init__.py index d48d87c79b8..004c8724fc5 100644 --- a/flashinfer/kda_kernels/__init__.py +++ b/flashinfer/kda_kernels/__init__.py @@ -13,7 +13,7 @@ # limitations under the License. """ -KDA (Key-Driven Attention) Kernels - CuTe DSL Implementations +KDA (Kimi Delta Attention) Kernels - CuTe DSL Implementations ============================================================== Per-K-dimension gating variant of GDN. Gate g[B,T,HV,K] applied per-lane diff --git a/flashinfer/trace/templates/kda.py b/flashinfer/trace/templates/kda.py index 935be32acd7..5813c2d5d28 100644 --- a/flashinfer/trace/templates/kda.py +++ b/flashinfer/trace/templates/kda.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""TraceTemplate for recurrent Key-Driven Attention (KDA).""" +"""TraceTemplate for recurrent Kimi Delta Attention (KDA).""" from ..template import Const, Scalar, Tensor, TraceTemplate, Var @@ -21,7 +21,7 @@ op_type="kda", name_prefix="recurrent_kda", description=( - "Recurrent Key-Driven Attention decode/spec-decode plus exact " + "Recurrent Kimi Delta Attention decode/spec-decode plus exact " "FlashKDA-compatible ordinary prefill on B200, with " "per-key-dimension gating and recurrent BF16 state." ), diff --git a/tests/kda/test_recurrent_kda.py b/tests/kda/test_recurrent_kda.py index 3eddea1ed18..7e5dc067f16 100644 --- a/tests/kda/test_recurrent_kda.py +++ b/tests/kda/test_recurrent_kda.py @@ -10,7 +10,7 @@ from flashinfer.utils import is_sm100a_supported try: - from flashinfer.kda_decode import _RECURRENT_KDA_AVAILABLE, recurrent_kda + from flashinfer.kda import _RECURRENT_KDA_AVAILABLE, recurrent_kda _has_recurrent_kda = _RECURRENT_KDA_AVAILABLE except ImportError: diff --git a/tests/kda/test_recurrent_kda_prefill.py b/tests/kda/test_recurrent_kda_prefill.py index 782dc95e7d5..21e526428b8 100644 --- a/tests/kda/test_recurrent_kda_prefill.py +++ b/tests/kda/test_recurrent_kda_prefill.py @@ -19,13 +19,20 @@ import torch import torch.nn.functional as F -from flashinfer.kda_decode import ( +from flashinfer.kda import ( RecurrentKDAPrefillWorkspace, recurrent_kda, ) from flashinfer.utils import get_compute_capability -kda_api = importlib.import_module("flashinfer.kda_decode") +kda_api = importlib.import_module("flashinfer.kda") + + +def test_legacy_kda_decode_module_reexports_canonical_api(): + legacy_api = importlib.import_module("flashinfer.kda_decode") + + assert legacy_api.recurrent_kda is recurrent_kda + assert legacy_api.RecurrentKDAPrefillWorkspace is RecurrentKDAPrefillWorkspace def _strict_prefill_kwargs(inputs): diff --git a/tests/trace/example.py b/tests/trace/example.py index f2e0dd80965..2f9d2b87c69 100644 --- a/tests/trace/example.py +++ b/tests/trace/example.py @@ -105,7 +105,7 @@ import flashinfer.sampling import flashinfer.gemm import flashinfer.gdn_decode -import flashinfer.kda_decode +import flashinfer.kda import flashinfer.fused_moe import flashinfer.activation import flashinfer.cascade @@ -691,7 +691,7 @@ rk_B + 2, rk_HV, rk_D, rk_D, dtype=torch.bfloat16, device=device ) rk_source_indices = torch.arange(rk_B, dtype=torch.int32, device=device) -flashinfer.kda_decode.recurrent_kda( +flashinfer.kda.recurrent_kda( rk_q, rk_k, rk_v, diff --git a/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json b/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json index 6767a4bc811..04169b3c75c 100644 --- a/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json +++ b/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json @@ -1,9 +1,9 @@ { "name": "recurrent_kda_q8_v16_d128", - "description": "Recurrent Key-Driven Attention decode/spec-decode plus exact FlashKDA-compatible ordinary prefill on B200, with per-key-dimension gating and recurrent BF16 state.", + "description": "Recurrent Kimi Delta Attention decode/spec-decode plus exact FlashKDA-compatible ordinary prefill on B200, with per-key-dimension gating and recurrent BF16 state.", "op_type": "kda", "tags": [ - "fi_api:flashinfer.kda_decode.recurrent_kda", + "fi_api:flashinfer.kda.recurrent_kda", "stage:decode", "stage:prefill", "status:verified" diff --git a/tests/trace/test_fi_trace.py b/tests/trace/test_fi_trace.py index 26afdbc4697..361d49edd1d 100644 --- a/tests/trace/test_fi_trace.py +++ b/tests/trace/test_fi_trace.py @@ -175,7 +175,7 @@ def test_attention_trace_check_tolerances_match_unit_tests(): def test_recurrent_kda_fi_trace(): - import flashinfer.kda_decode + import flashinfer.kda batch_size, num_q_heads, num_v_heads, head_dim = 4, 8, 16, 128 q = torch.empty(batch_size, 1, num_q_heads, head_dim, dtype=torch.bfloat16) @@ -191,7 +191,7 @@ def test_recurrent_kda_fi_trace(): ) source_indices = torch.arange(batch_size, dtype=torch.int32) - defn = flashinfer.kda_decode.recurrent_kda.fi_trace( + defn = flashinfer.kda.recurrent_kda.fi_trace( q=q, k=k, v=v, @@ -203,7 +203,7 @@ def test_recurrent_kda_fi_trace(): beta_is_logit=True, ) - _check_defn(defn, "kda", "flashinfer.kda_decode.recurrent_kda") + _check_defn(defn, "kda", "flashinfer.kda.recurrent_kda") assert defn["inputs"]["initial_state_source"]["shape"] == [ "source_pool_size", "num_v_heads", diff --git a/tests/trace/test_fi_trace_template_consistency.py b/tests/trace/test_fi_trace_template_consistency.py index 70f2d9256c9..be678c75ebf 100644 --- a/tests/trace/test_fi_trace_template_consistency.py +++ b/tests/trace/test_fi_trace_template_consistency.py @@ -329,7 +329,7 @@ def _collect_template_func_pairs() -> List[Tuple[Callable, TraceTemplate, str]]: import flashinfer.gdn_decode # gated_delta_rule_decode, gated_delta_rule_mtp import flashinfer.gdn_prefill # chunk_gated_delta_rule import flashinfer.gemm # mm_bf16, mm_fp8, mm_mxfp8, mm_fp4 - import flashinfer.kda_decode # recurrent_kda + import flashinfer.kda # recurrent_kda import flashinfer.mla # BatchMLAPagedAttentionWrapper import flashinfer.msa_ops # msa_proxy_score, msa_sparse_attention, decode import flashinfer.norm # rmsnorm, fused_add_rmsnorm From 07e4af6165e2447ae9c38b00347431ab3b6e1ced Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Sat, 1 Aug 2026 02:10:00 -0700 Subject: [PATCH 7/8] refactor(kda): separate prefill backend from decode API --- benchmarks/bench_recurrent_kda.py | 4 +- benchmarks/bench_recurrent_kda_prefill.py | 6 +- docs/api/kda_decode.rst | 16 ++ docs/api/{kda.rst => kda_prefill.rst} | 23 +- docs/index.rst | 3 +- flashinfer/__init__.py | 4 +- flashinfer/kda_decode.py | 262 ++++++++++++++++++ flashinfer/kda_decode/__init__.py | 14 - flashinfer/{kda.py => kda_prefill.py} | 254 +---------------- tests/kda/test_recurrent_kda.py | 2 +- tests/kda/test_recurrent_kda_prefill.py | 107 ++++--- tests/trace/example.py | 4 +- .../recurrent_kda_q8_v16_d128.json | 2 +- tests/trace/test_fi_trace.py | 6 +- .../test_fi_trace_template_consistency.py | 2 +- 15 files changed, 382 insertions(+), 327 deletions(-) create mode 100644 docs/api/kda_decode.rst rename docs/api/{kda.rst => kda_prefill.rst} (88%) create mode 100644 flashinfer/kda_decode.py delete mode 100644 flashinfer/kda_decode/__init__.py rename flashinfer/{kda.py => kda_prefill.py} (67%) diff --git a/benchmarks/bench_recurrent_kda.py b/benchmarks/bench_recurrent_kda.py index b0172d2aef7..aef5b7ca391 100644 --- a/benchmarks/bench_recurrent_kda.py +++ b/benchmarks/bench_recurrent_kda.py @@ -35,7 +35,7 @@ # Import the recurrent KDA kernel try: - from flashinfer.kda import _RECURRENT_KDA_AVAILABLE, recurrent_kda + from flashinfer.kda_decode import _RECURRENT_KDA_AVAILABLE, recurrent_kda RECURRENT_KDA_AVAILABLE = _RECURRENT_KDA_AVAILABLE except ImportError: @@ -233,7 +233,7 @@ def run_recurrent_kda_benchmark(args, dtype): """Run recurrent KDA decode benchmarks.""" if not RECURRENT_KDA_AVAILABLE: print("Error: recurrent KDA kernel is not available.") - print("Make sure flashinfer.kda.recurrent_kda is importable.") + print("Make sure flashinfer.kda_decode.recurrent_kda is importable.") return invalid_seq_lens = [t for t in args.seq_len if t < 1] diff --git a/benchmarks/bench_recurrent_kda_prefill.py b/benchmarks/bench_recurrent_kda_prefill.py index 0d5e41a588f..b6e4417676a 100644 --- a/benchmarks/bench_recurrent_kda_prefill.py +++ b/benchmarks/bench_recurrent_kda_prefill.py @@ -42,10 +42,8 @@ import numpy as np import torch -from flashinfer.kda import ( - RecurrentKDAPrefillWorkspace, - recurrent_kda, -) +from flashinfer.kda_decode import recurrent_kda +from flashinfer.kda_prefill import RecurrentKDAPrefillWorkspace from flashinfer.testing import bench_gpu_time from flashinfer.utils import get_compute_capability diff --git a/docs/api/kda_decode.rst b/docs/api/kda_decode.rst new file mode 100644 index 00000000000..5e950f2727b --- /dev/null +++ b/docs/api/kda_decode.rst @@ -0,0 +1,16 @@ +.. _apikda_decode: + +flashinfer.kda_decode +===================== + +Recurrent Kimi Delta Attention (KDA) public API. Decode and speculative +decode use the CuTe-DSL backend under ``flashinfer.kda_kernels``. Eligible +ordinary multi-token prefill calls dispatch to the optimized backend described +in :ref:`apikda_prefill`. + +.. currentmodule:: flashinfer.kda_decode + +.. autosummary:: + :toctree: ../generated + + recurrent_kda diff --git a/docs/api/kda.rst b/docs/api/kda_prefill.rst similarity index 88% rename from docs/api/kda.rst rename to docs/api/kda_prefill.rst index e8742e12aba..19457552cb6 100644 --- a/docs/api/kda.rst +++ b/docs/api/kda_prefill.rst @@ -1,29 +1,24 @@ -.. _apikda: +.. _apikda_prefill: -flashinfer.kda -============== +flashinfer.kda_prefill +====================== -Recurrent Kimi Delta Attention (KDA) API. Single-token decode, speculative -decode, GQA, state-pool indexing, and checkpoint modes use the CuTe-DSL -backend under ``flashinfer.kda_kernels``. A strict ordinary multi-token -prefill subset is dispatched to frozen FlashKDA-compatible SM100a kernels. +Optimized recurrent Kimi Delta Attention (KDA) prefill support. The stable +:func:`flashinfer.kda_decode.recurrent_kda` API dispatches a strict ordinary +multi-token prefill subset to frozen FlashKDA-compatible SM100a kernels. -The former ``flashinfer.kda_decode`` module remains as a compatibility import; -new code should import this phase-neutral module. - -.. currentmodule:: flashinfer.kda +.. currentmodule:: flashinfer.kda_prefill .. autosummary:: :toctree: ../generated RecurrentKDAPrefillWorkspace - recurrent_kda Optimized B200 prefill subset ----------------------------- -``recurrent_kda`` uses the frozen prefill backend only when every condition -below holds: +``flashinfer.kda_decode.recurrent_kda`` uses the frozen prefill backend only +when every condition below holds: * the device has compute capability 10.0; * input is ordinary multi-token prefill: fixed ``T > 1``, or packed input diff --git a/docs/index.rst b/docs/index.rst index 433e2afb284..adbf5e5ec8f 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -53,7 +53,8 @@ FlashInfer is a library and kernel generator for Large Language Models that prov api/activation api/gdn_decode api/gdn_prefill - api/kda + api/kda_decode + api/kda_prefill api/mamba api/mhc api/quantization diff --git a/flashinfer/__init__.py b/flashinfer/__init__.py index 35fb39a1260..67b777c1007 100644 --- a/flashinfer/__init__.py +++ b/flashinfer/__init__.py @@ -124,10 +124,10 @@ from .grouped_mm import grouped_mm_fp8 as grouped_mm_fp8 from .grouped_mm import grouped_mm_mxfp8 as grouped_mm_mxfp8 from .grouped_mm import grouped_mm_fp4 as grouped_mm_fp4 -from .kda import ( +from .kda_prefill import ( RecurrentKDAPrefillWorkspace as RecurrentKDAPrefillWorkspace, ) -from .kda import recurrent_kda as recurrent_kda +from .kda_decode import recurrent_kda as recurrent_kda from .mla import BatchMLAPagedAttentionWrapper as BatchMLAPagedAttentionWrapper from . import mhc as mhc from . import msa_ops as msa_ops diff --git a/flashinfer/kda_decode.py b/flashinfer/kda_decode.py new file mode 100644 index 00000000000..37d13e4de9d --- /dev/null +++ b/flashinfer/kda_decode.py @@ -0,0 +1,262 @@ +""" +Copyright (c) 2025 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. +""" + +""" +Kimi Delta Attention Decode - API Layer +======================================= + +This file preserves the public recurrent KDA decode API. Eligible ordinary +multi-token prefill calls are dispatched to the frozen backend implemented in +``flashinfer.kda_prefill``; decode and speculative decode retain the +CuTe-DSL backend under ``flashinfer.kda_kernels``. +""" + +from typing import Optional + +import torch + +from . import kda_prefill as _kda_prefill +from .api_logging import flashinfer_api +from .trace.templates.kda import recurrent_kda_trace + +try: + from .kda_kernels.recurrent_kda import run_recurrent_kda as _run_recurrent_kda + + _RECURRENT_KDA_AVAILABLE = True +except (ImportError, RuntimeError): + _run_recurrent_kda = None + _RECURRENT_KDA_AVAILABLE = False + + +@flashinfer_api(trace=recurrent_kda_trace) +def recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + scale: Optional[float] = None, + initial_state: Optional[torch.Tensor] = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = True, + use_gate_in_kernel: bool = False, + lower_bound: Optional[float] = None, + cu_seqlens: Optional[torch.Tensor] = None, + ssm_state_indices: Optional[torch.Tensor] = None, + num_spec_tokens: Optional[int] = None, + num_accepted_tokens: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + initial_state_source: Optional[torch.Tensor] = None, + initial_state_indices: Optional[torch.Tensor] = None, + beta_is_logit: bool = False, + seq_order: Optional[torch.Tensor] = None, + prefill_workspace: Optional[_kda_prefill.RecurrentKDAPrefillWorkspace] = None, +) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + r"""Recurrent KDA (Kimi Delta Attention) decode and prefill kernel. + + This is the public API layer for the CuTe DSL implementation in + ``flashinfer.kda_kernels.recurrent_kda``. It supports single-token decode, + fused speculative decode, GQA, optional cu_seqlens packing, and the same + gate modes as the backend implementation. On NVIDIA B200, the exact + FlashKDA-compatible subset of ordinary multi-token prefill is dispatched to + frozen SM100a kernels. All existing decode and speculative-decode calls + retain the CuTe DSL backend. + + Args: + q (torch.Tensor): + Query of shape ``[B, T, H, K]``, or + ``[1, total_tokens, H, K]`` when using ``cu_seqlens``. Must be + bfloat16. ``T=1`` selects decode; eligible ``T>1`` calls may select + the frozen prefill backend. + k (torch.Tensor): + Key with the same shape as ``q``. Must be bfloat16. + v (torch.Tensor): + Value of shape ``[B, T, HV, V]``, or + ``[1, total_tokens, HV, V]`` when packed. Must be bfloat16. GQA is + applied when ``HV != H``. + g (torch.Tensor): + Per-K-dimension gate of shape ``[B, T, HV, K]``, or + ``[1, total_tokens, HV, K]`` when packed. Must be bfloat16. + Log-space if pre-computed, raw input if + ``use_gate_in_kernel=True``. + beta (torch.Tensor): + Delta-rule learning rate of shape ``[B, T, HV]``, or + ``[1, total_tokens, HV]`` when packed. Must be bfloat16. + Pre-sigmoided unless ``beta_is_logit=True``. + A_log (Optional[torch.Tensor]): + Log decay parameter of shape ``[H]``. Must be float32. + Required when ``use_gate_in_kernel=True``. + dt_bias (Optional[torch.Tensor]): + Per-head-K decay bias of shape ``[H*K]`` or ``[H, K]``. Must be + float32. + scale (Optional[float]): + Scale factor for queries. If ``None``, defaults to ``1 / sqrt(K)``. + initial_state (Optional[torch.Tensor]): + Initial state of shape ``[N, HV, V, K]``. Must be bfloat16. + If ``None``, zero-initialized. Updated in-place. For batched spec + decode without ``cu_seqlens``, ``N`` is the packed checkpoint-slot + count ``B * (1 + num_spec_tokens)`` when ``ssm_state_indices`` is + omitted. + output_final_state (bool): + Whether to return the final state. Default: ``False``. + use_qk_l2norm_in_kernel (bool): + Whether to apply L2 normalization to Q and K. Default: ``True``. + use_gate_in_kernel (bool): + Whether to compute the gate inside the kernel from ``A_log`` and + ``g``. Default: ``False``. + lower_bound (Optional[float]): + If set, uses ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))`` + gate formula instead of softplus. Must be negative. + cu_seqlens (Optional[torch.Tensor]): + Contiguous CUDA cumulative sequence lengths of shape ``[N+1]``. + May be int32 or int64. Frozen prefill converts int32 offsets to + int64 outside graph capture; graph capture requires caller-provided + int64 offsets. For frozen prefill, values must start at zero, be + strictly increasing, and end at the total token count. This value + contract is not host-validated to avoid a device synchronization. + ssm_state_indices (Optional[torch.Tensor]): + State cache indices. Shape ``[N]`` int32 for standard decode, or + ``[N, 1+S]`` int32 for spec decode (``num_spec_tokens`` must also + be set). + num_spec_tokens (Optional[int]): + Number of speculative tokens (S). When set, processes 1+S tokens in + a single fused kernel launch. Must be >= 1. + num_accepted_tokens (Optional[torch.Tensor]): + Per-sequence accepted token count from the previous spec decode + round. Shape ``[N]`` int32. If ``None``, initial state is loaded + from ``ssm_state_indices[n, 0]``. Values above ``1+S`` are clamped + to the final checkpoint slot. + output (Optional[torch.Tensor]): + Pre-allocated output tensor. Shape ``[B, T, HV, V]`` for fixed + layout, or the corresponding packed/speculative shape when using + ``cu_seqlens``. If ``None``, a new tensor is allocated. Frozen + prefill requires storage disjoint from Q, K, V, G, beta, and + ``initial_state``. + initial_state_source (Optional[torch.Tensor]): + Optional read-only committed state pool ``[N0, HV, V, K]``. When + provided, token 0 is loaded from this pool instead of + ``initial_state``. + initial_state_indices (Optional[torch.Tensor]): + Source slot per sequence, shape ``[N]`` int32. Required together + with ``initial_state_source``. + beta_is_logit (bool): + If ``True``, apply sigmoid to ``beta`` inside the recurrent kernel. + seq_order (Optional[torch.Tensor]): + Optional packed-prefill sequence order, as a contiguous CUDA int32 + permutation of shape ``[N]``. Sorting by descending sequence length + improves tail utilization. It is only consumed by the frozen + FlashKDA prefill backend; prepare it before CUDA graph capture or + timed launches. Fixed-layout prefill and decode calls must leave it + as ``None``. + prefill_workspace (Optional[RecurrentKDAPrefillWorkspace]): + Caller-owned workspace for the frozen B200 prefill backend. It is + optional for eager execution and required for CUDA graph capture. + Warm it eagerly with the exact tensors on the capture stream before + capture. Use one workspace per captured ``recurrent_kda`` + invocation. + + Returns: + Tuple of ``(output, final_state)`` where ``final_state`` is ``None`` + when ``output_final_state=False``. See + :func:`flashinfer.kda_kernels.recurrent_kda.run_recurrent_kda` for the + backend implementation. + """ + if prefill_workspace is not None and not isinstance( + prefill_workspace, _kda_prefill.RecurrentKDAPrefillWorkspace + ): + raise TypeError("prefill_workspace must be a RecurrentKDAPrefillWorkspace") + + use_flash_kda_prefill = _kda_prefill._flash_kda_prefill_is_eligible( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=initial_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + use_gate_in_kernel=use_gate_in_kernel, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_spec_tokens=num_spec_tokens, + num_accepted_tokens=num_accepted_tokens, + output=output, + initial_state_source=initial_state_source, + initial_state_indices=initial_state_indices, + beta_is_logit=beta_is_logit, + ) + if use_flash_kda_prefill: + assert A_log is not None + assert dt_bias is not None + assert lower_bound is not None + return _kda_prefill._run_flash_kda_prefill( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + output=output, + seq_order=seq_order, + prefill_workspace=prefill_workspace, + ) + + if prefill_workspace is not None: + raise ValueError( + "prefill_workspace is only supported by eligible ordinary " + "prefill on the frozen B200 FlashKDA backend" + ) + if seq_order is not None: + raise ValueError( + "seq_order is only supported by eligible packed ordinary prefill " + "on the frozen B200 FlashKDA backend" + ) + if _run_recurrent_kda is None: + raise NotImplementedError("recurrent KDA backend is unavailable") + + return _run_recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + use_gate_in_kernel=use_gate_in_kernel, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_spec_tokens=num_spec_tokens, + num_accepted_tokens=num_accepted_tokens, + output=output, + initial_state_source=initial_state_source, + initial_state_indices=initial_state_indices, + beta_is_logit=beta_is_logit, + ) diff --git a/flashinfer/kda_decode/__init__.py b/flashinfer/kda_decode/__init__.py deleted file mode 100644 index 44ed9598023..00000000000 --- a/flashinfer/kda_decode/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -"""Compatibility imports for the former recurrent-KDA module name. - -The canonical recurrent KDA API lives in :mod:`flashinfer.kda`. This module -keeps existing direct imports working while new code uses the phase-neutral -module name. -""" - -from ..kda import ( - RecurrentKDAPrefillWorkspace as RecurrentKDAPrefillWorkspace, -) -from ..kda import _RECURRENT_KDA_AVAILABLE as _RECURRENT_KDA_AVAILABLE -from ..kda import recurrent_kda as recurrent_kda - -__all__ = ["RecurrentKDAPrefillWorkspace", "recurrent_kda"] diff --git a/flashinfer/kda.py b/flashinfer/kda_prefill.py similarity index 67% rename from flashinfer/kda.py rename to flashinfer/kda_prefill.py index 8cad532a02e..da2a8fe8453 100644 --- a/flashinfer/kda.py +++ b/flashinfer/kda_prefill.py @@ -15,11 +15,12 @@ """ """ -Kimi Delta Attention - API Layer -================================ +Kimi Delta Attention Prefill - Backend Layer +============================================ -This file provides the public API and dispatch for recurrent KDA decode and -prefill operations. +This file provides workspace management, validation, and frozen-kernel launch +support for recurrent KDA prefill. The stable public dispatcher remains in +``flashinfer.kda_decode``. """ import math @@ -28,22 +29,11 @@ import torch -from .api_logging import flashinfer_api -from .trace.templates.kda import recurrent_kda_trace from .utils import get_compute_capability if TYPE_CHECKING: from .jit.flash_kda import FlashKDAVariant -try: - from .kda_kernels.recurrent_kda import run_recurrent_kda as _run_recurrent_kda - - _RECURRENT_KDA_AVAILABLE = True -except (ImportError, RuntimeError): - _run_recurrent_kda = None - _RECURRENT_KDA_AVAILABLE = False - - _FLASH_KDA_HEAD_DIM = 128 _FLASH_KDA_BETA_TMA_MIN_HEADS = 8 _FLASH_KDA_B200_COMPUTE_CAPABILITY = (10, 0) @@ -79,12 +69,13 @@ def __init__(self, device: torch.device | str) -> None: class RecurrentKDAPrefillWorkspace(_RecurrentKDAPrefillWorkspaceBase): """Caller-owned storage required for recurrent-KDA CUDA graph capture. - Construct one workspace per captured :func:`recurrent_kda` invocation on - the graph's CUDA device. Warm it by invoking :func:`recurrent_kda` eagerly - with the exact tensors and capture stream, then synchronize that stream - before capture. The workspace owns optional final-state scratch for calls - without an initial state, beta padding, and M64/M128 TMA descriptor storage - for the lifetime of the graph. + Construct one workspace per captured + :func:`flashinfer.kda_decode.recurrent_kda` invocation on the graph's CUDA + device. Warm it by invoking that function eagerly with the exact tensors + and capture stream, then synchronize that stream before capture. The + workspace owns optional final-state scratch for calls without an initial + state, beta padding, and M64/M128 TMA descriptor storage for the lifetime + of the graph. A workspace binds to its first stream. Once it participates in capture it cannot be passed to Python again, either eagerly or in another capture. @@ -727,224 +718,3 @@ def _run_flash_kda_prefill( out_buf, returned_state if output_final_state else None, ) - - -@flashinfer_api(trace=recurrent_kda_trace) -def recurrent_kda( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - g: torch.Tensor, - beta: torch.Tensor, - A_log: Optional[torch.Tensor] = None, - dt_bias: Optional[torch.Tensor] = None, - scale: Optional[float] = None, - initial_state: Optional[torch.Tensor] = None, - output_final_state: bool = False, - use_qk_l2norm_in_kernel: bool = True, - use_gate_in_kernel: bool = False, - lower_bound: Optional[float] = None, - cu_seqlens: Optional[torch.Tensor] = None, - ssm_state_indices: Optional[torch.Tensor] = None, - num_spec_tokens: Optional[int] = None, - num_accepted_tokens: Optional[torch.Tensor] = None, - output: Optional[torch.Tensor] = None, - initial_state_source: Optional[torch.Tensor] = None, - initial_state_indices: Optional[torch.Tensor] = None, - beta_is_logit: bool = False, - seq_order: Optional[torch.Tensor] = None, - prefill_workspace: Optional[RecurrentKDAPrefillWorkspace] = None, -) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - r"""Recurrent KDA (Kimi Delta Attention) decode and prefill kernel. - - This is the public API layer for the CuTe DSL implementation in - ``flashinfer.kda_kernels.recurrent_kda``. It supports single-token decode, - fused speculative decode, GQA, optional cu_seqlens packing, and the same - gate modes as the backend implementation. On NVIDIA B200, the exact - FlashKDA-compatible subset of ordinary multi-token prefill is dispatched to - frozen SM100a kernels. All existing decode and speculative-decode calls - retain the CuTe DSL backend. - - Args: - q (torch.Tensor): - Query of shape ``[B, T, H, K]``, or - ``[1, total_tokens, H, K]`` when using ``cu_seqlens``. Must be - bfloat16. ``T=1`` selects decode; eligible ``T>1`` calls may select - the frozen prefill backend. - k (torch.Tensor): - Key with the same shape as ``q``. Must be bfloat16. - v (torch.Tensor): - Value of shape ``[B, T, HV, V]``, or - ``[1, total_tokens, HV, V]`` when packed. Must be bfloat16. GQA is - applied when ``HV != H``. - g (torch.Tensor): - Per-K-dimension gate of shape ``[B, T, HV, K]``, or - ``[1, total_tokens, HV, K]`` when packed. Must be bfloat16. - Log-space if pre-computed, raw input if - ``use_gate_in_kernel=True``. - beta (torch.Tensor): - Delta-rule learning rate of shape ``[B, T, HV]``, or - ``[1, total_tokens, HV]`` when packed. Must be bfloat16. - Pre-sigmoided unless ``beta_is_logit=True``. - A_log (Optional[torch.Tensor]): - Log decay parameter of shape ``[H]``. Must be float32. - Required when ``use_gate_in_kernel=True``. - dt_bias (Optional[torch.Tensor]): - Per-head-K decay bias of shape ``[H*K]`` or ``[H, K]``. Must be - float32. - scale (Optional[float]): - Scale factor for queries. If ``None``, defaults to ``1 / sqrt(K)``. - initial_state (Optional[torch.Tensor]): - Initial state of shape ``[N, HV, V, K]``. Must be bfloat16. - If ``None``, zero-initialized. Updated in-place. For batched spec - decode without ``cu_seqlens``, ``N`` is the packed checkpoint-slot - count ``B * (1 + num_spec_tokens)`` when ``ssm_state_indices`` is - omitted. - output_final_state (bool): - Whether to return the final state. Default: ``False``. - use_qk_l2norm_in_kernel (bool): - Whether to apply L2 normalization to Q and K. Default: ``True``. - use_gate_in_kernel (bool): - Whether to compute the gate inside the kernel from ``A_log`` and - ``g``. Default: ``False``. - lower_bound (Optional[float]): - If set, uses ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))`` - gate formula instead of softplus. Must be negative. - cu_seqlens (Optional[torch.Tensor]): - Contiguous CUDA cumulative sequence lengths of shape ``[N+1]``. - May be int32 or int64. Frozen prefill converts int32 offsets to - int64 outside graph capture; graph capture requires caller-provided - int64 offsets. For frozen prefill, values must start at zero, be - strictly increasing, and end at the total token count. This value - contract is not host-validated to avoid a device synchronization. - ssm_state_indices (Optional[torch.Tensor]): - State cache indices. Shape ``[N]`` int32 for standard decode, or - ``[N, 1+S]`` int32 for spec decode (``num_spec_tokens`` must also - be set). - num_spec_tokens (Optional[int]): - Number of speculative tokens (S). When set, processes 1+S tokens in - a single fused kernel launch. Must be >= 1. - num_accepted_tokens (Optional[torch.Tensor]): - Per-sequence accepted token count from the previous spec decode - round. Shape ``[N]`` int32. If ``None``, initial state is loaded - from ``ssm_state_indices[n, 0]``. Values above ``1+S`` are clamped - to the final checkpoint slot. - output (Optional[torch.Tensor]): - Pre-allocated output tensor. Shape ``[B, T, HV, V]`` for fixed - layout, or the corresponding packed/speculative shape when using - ``cu_seqlens``. If ``None``, a new tensor is allocated. Frozen - prefill requires storage disjoint from Q, K, V, G, beta, and - ``initial_state``. - initial_state_source (Optional[torch.Tensor]): - Optional read-only committed state pool ``[N0, HV, V, K]``. When - provided, token 0 is loaded from this pool instead of - ``initial_state``. - initial_state_indices (Optional[torch.Tensor]): - Source slot per sequence, shape ``[N]`` int32. Required together - with ``initial_state_source``. - beta_is_logit (bool): - If ``True``, apply sigmoid to ``beta`` inside the recurrent kernel. - seq_order (Optional[torch.Tensor]): - Optional packed-prefill sequence order, as a contiguous CUDA int32 - permutation of shape ``[N]``. Sorting by descending sequence length - improves tail utilization. It is only consumed by the frozen - FlashKDA prefill backend; prepare it before CUDA graph capture or - timed launches. Fixed-layout prefill and decode calls must leave it - as ``None``. - prefill_workspace (Optional[RecurrentKDAPrefillWorkspace]): - Caller-owned workspace for the frozen B200 prefill backend. It is - optional for eager execution and required for CUDA graph capture. - Warm it eagerly with the exact tensors on the capture stream before - capture. Use one workspace per captured ``recurrent_kda`` - invocation. - - Returns: - Tuple of ``(output, final_state)`` where ``final_state`` is ``None`` - when ``output_final_state=False``. See - :func:`flashinfer.kda_kernels.recurrent_kda.run_recurrent_kda` for the - backend implementation. - """ - if prefill_workspace is not None and not isinstance( - prefill_workspace, RecurrentKDAPrefillWorkspace - ): - raise TypeError("prefill_workspace must be a RecurrentKDAPrefillWorkspace") - - use_flash_kda_prefill = _flash_kda_prefill_is_eligible( - q=q, - k=k, - v=v, - g=g, - beta=beta, - A_log=A_log, - dt_bias=dt_bias, - initial_state=initial_state, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - use_gate_in_kernel=use_gate_in_kernel, - lower_bound=lower_bound, - cu_seqlens=cu_seqlens, - ssm_state_indices=ssm_state_indices, - num_spec_tokens=num_spec_tokens, - num_accepted_tokens=num_accepted_tokens, - output=output, - initial_state_source=initial_state_source, - initial_state_indices=initial_state_indices, - beta_is_logit=beta_is_logit, - ) - if use_flash_kda_prefill: - assert A_log is not None - assert dt_bias is not None - assert lower_bound is not None - return _run_flash_kda_prefill( - q=q, - k=k, - v=v, - g=g, - beta=beta, - A_log=A_log, - dt_bias=dt_bias, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - lower_bound=lower_bound, - cu_seqlens=cu_seqlens, - output=output, - seq_order=seq_order, - prefill_workspace=prefill_workspace, - ) - - if prefill_workspace is not None: - raise ValueError( - "prefill_workspace is only supported by eligible ordinary " - "prefill on the frozen B200 FlashKDA backend" - ) - if seq_order is not None: - raise ValueError( - "seq_order is only supported by eligible packed ordinary prefill " - "on the frozen B200 FlashKDA backend" - ) - if _run_recurrent_kda is None: - raise NotImplementedError("recurrent KDA backend is unavailable") - - return _run_recurrent_kda( - q=q, - k=k, - v=v, - g=g, - beta=beta, - A_log=A_log, - dt_bias=dt_bias, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - use_gate_in_kernel=use_gate_in_kernel, - lower_bound=lower_bound, - cu_seqlens=cu_seqlens, - ssm_state_indices=ssm_state_indices, - num_spec_tokens=num_spec_tokens, - num_accepted_tokens=num_accepted_tokens, - output=output, - initial_state_source=initial_state_source, - initial_state_indices=initial_state_indices, - beta_is_logit=beta_is_logit, - ) diff --git a/tests/kda/test_recurrent_kda.py b/tests/kda/test_recurrent_kda.py index 7e5dc067f16..3eddea1ed18 100644 --- a/tests/kda/test_recurrent_kda.py +++ b/tests/kda/test_recurrent_kda.py @@ -10,7 +10,7 @@ from flashinfer.utils import is_sm100a_supported try: - from flashinfer.kda import _RECURRENT_KDA_AVAILABLE, recurrent_kda + from flashinfer.kda_decode import _RECURRENT_KDA_AVAILABLE, recurrent_kda _has_recurrent_kda = _RECURRENT_KDA_AVAILABLE except ImportError: diff --git a/tests/kda/test_recurrent_kda_prefill.py b/tests/kda/test_recurrent_kda_prefill.py index 21e526428b8..6cd9a969239 100644 --- a/tests/kda/test_recurrent_kda_prefill.py +++ b/tests/kda/test_recurrent_kda_prefill.py @@ -19,20 +19,21 @@ import torch import torch.nn.functional as F -from flashinfer.kda import ( - RecurrentKDAPrefillWorkspace, - recurrent_kda, -) +import flashinfer +from flashinfer.kda_decode import recurrent_kda +from flashinfer.kda_prefill import RecurrentKDAPrefillWorkspace from flashinfer.utils import get_compute_capability -kda_api = importlib.import_module("flashinfer.kda") - +kda_decode_api = importlib.import_module("flashinfer.kda_decode") +kda_prefill_api = importlib.import_module("flashinfer.kda_prefill") -def test_legacy_kda_decode_module_reexports_canonical_api(): - legacy_api = importlib.import_module("flashinfer.kda_decode") - assert legacy_api.recurrent_kda is recurrent_kda - assert legacy_api.RecurrentKDAPrefillWorkspace is RecurrentKDAPrefillWorkspace +def test_public_api_keeps_decode_entry_and_prefill_workspace(): + assert flashinfer.recurrent_kda is kda_decode_api.recurrent_kda + assert ( + flashinfer.RecurrentKDAPrefillWorkspace + is kda_prefill_api.RecurrentKDAPrefillWorkspace + ) def _strict_prefill_kwargs(inputs): @@ -179,9 +180,9 @@ def old_backend(**kwargs): calls.append(kwargs) return sentinel - monkeypatch.setattr(kda_api, "_run_recurrent_kda", old_backend) + monkeypatch.setattr(kda_decode_api, "_run_recurrent_kda", old_backend) monkeypatch.setattr( - kda_api, + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: pytest.fail(f"unexpected frozen route {variant}"), ) @@ -202,10 +203,12 @@ def old_backend(**kwargs): def test_multi_token_gqa_stays_on_existing_backend(cuda_device, monkeypatch): sentinel = (object(), object()) - monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) - monkeypatch.setattr(kda_api, "_run_recurrent_kda", lambda **kwargs: sentinel) monkeypatch.setattr( - kda_api, + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + monkeypatch.setattr(kda_decode_api, "_run_recurrent_kda", lambda **kwargs: sentinel) + monkeypatch.setattr( + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: pytest.fail(f"unexpected frozen route {variant}"), ) @@ -238,15 +241,17 @@ def test_frozen_route_and_ffi_abi( num_heads, expected_variant, ): - monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) - monkeypatch.setattr(kda_api, "_flash_kda_stream_workspaces", {}) + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + monkeypatch.setattr(kda_prefill_api, "_flash_kda_stream_workspaces", {}) modules = {} def get_module(variant): modules.setdefault(variant, _RecorderModule()) return modules[variant] - monkeypatch.setattr(kda_api, "_get_flash_kda_prefill_module", get_module) + monkeypatch.setattr(kda_prefill_api, "_get_flash_kda_prefill_module", get_module) inputs = _make_inputs( seq_lens=[1, 2] if packed else [2], num_heads=num_heads, @@ -295,10 +300,12 @@ def get_module(variant): def test_frozen_route_passes_nondefault_stream(cuda_device, monkeypatch): - monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) module = _RecorderModule() monkeypatch.setattr( - kda_api, "_get_flash_kda_prefill_module", lambda variant: module + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: module ) inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) stream = torch.cuda.Stream(device=cuda_device) @@ -313,10 +320,12 @@ def test_frozen_route_passes_nondefault_stream(cuda_device, monkeypatch): def test_frozen_route_rejects_output_overlap(cuda_device, monkeypatch): - monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) module = _RecorderModule() monkeypatch.setattr( - kda_api, "_get_flash_kda_prefill_module", lambda variant: module + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: module ) inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) with pytest.raises(ValueError, match="output must not overlap q"): @@ -328,10 +337,12 @@ def test_frozen_route_rejects_output_overlap(cuda_device, monkeypatch): def test_initial_state_is_updated_in_place(cuda_device, monkeypatch): - monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) module = _RecorderModule(final_value=0.25) monkeypatch.setattr( - kda_api, "_get_flash_kda_prefill_module", lambda variant: module + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: module ) inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False, initial_state=True) original_state = inputs["initial_state"] @@ -356,11 +367,13 @@ def test_initial_state_is_updated_in_place(cuda_device, monkeypatch): def test_stream_workspace_does_not_allocate_state_scratch_for_inplace_update( cuda_device, monkeypatch ): - monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) - monkeypatch.setattr(kda_api, "_flash_kda_stream_workspaces", {}) + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + monkeypatch.setattr(kda_prefill_api, "_flash_kda_stream_workspaces", {}) module = _RecorderModule(final_value=0.0) monkeypatch.setattr( - kda_api, "_get_flash_kda_prefill_module", lambda variant: module + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: module ) cases = [ _make_inputs( @@ -388,8 +401,8 @@ def test_stream_workspace_does_not_allocate_state_scratch_for_inplace_update( output=torch.empty_like(inputs["q"]), ) - assert len(kda_api._flash_kda_stream_workspaces) == 1 - (workspace,) = kda_api._flash_kda_stream_workspaces.values() + assert len(kda_prefill_api._flash_kda_stream_workspaces) == 1 + (workspace,) = kda_prefill_api._flash_kda_stream_workspaces.values() assert workspace._state_scratch is None assert workspace._beta_padding.numel() == 32 * 8 @@ -399,9 +412,11 @@ def test_stream_workspace_does_not_allocate_state_scratch_for_inplace_update( [(torch.int64, 0), (torch.int32, 1)], ) def test_packed_seq_order_validation(cuda_device, monkeypatch, dtype, size_delta): - monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) monkeypatch.setattr( - kda_api, + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + monkeypatch.setattr( + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: _RecorderModule(), ) @@ -412,7 +427,9 @@ def test_packed_seq_order_validation(cuda_device, monkeypatch, dtype, size_delta def test_fixed_prefill_rejects_seq_order(cuda_device, monkeypatch): - monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) with pytest.raises(ValueError, match="only supported for packed"): recurrent_kda( @@ -422,7 +439,9 @@ def test_fixed_prefill_rejects_seq_order(cuda_device, monkeypatch): def test_graph_capture_requires_packed_int64_offsets(cuda_device, monkeypatch): - monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) inputs = _make_inputs(seq_lens=[1, 2], num_heads=2, packed=True) inputs["cu_seqlens"] = inputs["cu_seqlens"].to(torch.int32) @@ -435,7 +454,9 @@ def test_graph_capture_requires_packed_int64_offsets(cuda_device, monkeypatch): def test_graph_capture_requires_explicit_workspace(cuda_device, monkeypatch): - monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) with pytest.raises( @@ -448,10 +469,12 @@ def test_graph_capture_requires_explicit_workspace(cuda_device, monkeypatch): def test_explicit_workspace_descriptor_prepare_and_reuse(cuda_device, monkeypatch): - monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) module = _RecorderModule() monkeypatch.setattr( - kda_api, "_get_flash_kda_prefill_module", lambda variant: module + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: module ) inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) output = torch.empty_like(inputs["q"]) @@ -482,10 +505,12 @@ def test_explicit_workspace_descriptor_prepare_and_reuse(cuda_device, monkeypatc def test_captured_workspace_rejects_eager_reuse_and_capture_mismatch( cuda_device, monkeypatch ): - monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) module = _RecorderModule() monkeypatch.setattr( - kda_api, "_get_flash_kda_prefill_module", lambda variant: module + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: module ) inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) output = torch.empty_like(inputs["q"]) @@ -529,10 +554,12 @@ def test_captured_workspace_rejects_eager_reuse_and_capture_mismatch( def test_workspace_rejects_a_different_stream(cuda_device, monkeypatch): - monkeypatch.setattr(kda_api, "get_compute_capability", lambda device: (10, 0)) + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) module = _RecorderModule() monkeypatch.setattr( - kda_api, "_get_flash_kda_prefill_module", lambda variant: module + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: module ) inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) output = torch.empty_like(inputs["q"]) diff --git a/tests/trace/example.py b/tests/trace/example.py index 2f9d2b87c69..f2e0dd80965 100644 --- a/tests/trace/example.py +++ b/tests/trace/example.py @@ -105,7 +105,7 @@ import flashinfer.sampling import flashinfer.gemm import flashinfer.gdn_decode -import flashinfer.kda +import flashinfer.kda_decode import flashinfer.fused_moe import flashinfer.activation import flashinfer.cascade @@ -691,7 +691,7 @@ rk_B + 2, rk_HV, rk_D, rk_D, dtype=torch.bfloat16, device=device ) rk_source_indices = torch.arange(rk_B, dtype=torch.int32, device=device) -flashinfer.kda.recurrent_kda( +flashinfer.kda_decode.recurrent_kda( rk_q, rk_k, rk_v, diff --git a/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json b/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json index 04169b3c75c..c45683cb341 100644 --- a/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json +++ b/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json @@ -3,7 +3,7 @@ "description": "Recurrent Kimi Delta Attention decode/spec-decode plus exact FlashKDA-compatible ordinary prefill on B200, with per-key-dimension gating and recurrent BF16 state.", "op_type": "kda", "tags": [ - "fi_api:flashinfer.kda.recurrent_kda", + "fi_api:flashinfer.kda_decode.recurrent_kda", "stage:decode", "stage:prefill", "status:verified" diff --git a/tests/trace/test_fi_trace.py b/tests/trace/test_fi_trace.py index 361d49edd1d..26afdbc4697 100644 --- a/tests/trace/test_fi_trace.py +++ b/tests/trace/test_fi_trace.py @@ -175,7 +175,7 @@ def test_attention_trace_check_tolerances_match_unit_tests(): def test_recurrent_kda_fi_trace(): - import flashinfer.kda + import flashinfer.kda_decode batch_size, num_q_heads, num_v_heads, head_dim = 4, 8, 16, 128 q = torch.empty(batch_size, 1, num_q_heads, head_dim, dtype=torch.bfloat16) @@ -191,7 +191,7 @@ def test_recurrent_kda_fi_trace(): ) source_indices = torch.arange(batch_size, dtype=torch.int32) - defn = flashinfer.kda.recurrent_kda.fi_trace( + defn = flashinfer.kda_decode.recurrent_kda.fi_trace( q=q, k=k, v=v, @@ -203,7 +203,7 @@ def test_recurrent_kda_fi_trace(): beta_is_logit=True, ) - _check_defn(defn, "kda", "flashinfer.kda.recurrent_kda") + _check_defn(defn, "kda", "flashinfer.kda_decode.recurrent_kda") assert defn["inputs"]["initial_state_source"]["shape"] == [ "source_pool_size", "num_v_heads", diff --git a/tests/trace/test_fi_trace_template_consistency.py b/tests/trace/test_fi_trace_template_consistency.py index be678c75ebf..70f2d9256c9 100644 --- a/tests/trace/test_fi_trace_template_consistency.py +++ b/tests/trace/test_fi_trace_template_consistency.py @@ -329,7 +329,7 @@ def _collect_template_func_pairs() -> List[Tuple[Callable, TraceTemplate, str]]: import flashinfer.gdn_decode # gated_delta_rule_decode, gated_delta_rule_mtp import flashinfer.gdn_prefill # chunk_gated_delta_rule import flashinfer.gemm # mm_bf16, mm_fp8, mm_mxfp8, mm_fp4 - import flashinfer.kda # recurrent_kda + import flashinfer.kda_decode # recurrent_kda import flashinfer.mla # BatchMLAPagedAttentionWrapper import flashinfer.msa_ops # msa_proxy_score, msa_sparse_attention, decode import flashinfer.norm # rmsnorm, fused_add_rmsnorm From e835e0f5565b5b9786c987e00c6b39a26bfecca5 Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Sat, 1 Aug 2026 02:29:53 -0700 Subject: [PATCH 8/8] refactor(kda): add phase-neutral recurrent facade --- benchmarks/bench_recurrent_kda_prefill.py | 2 +- docs/api/kda.rst | 16 ++ docs/api/kda_decode.rst | 6 +- docs/api/kda_prefill.rst | 8 +- docs/index.rst | 1 + flashinfer/__init__.py | 2 +- flashinfer/kda.py | 255 ++++++++++++++++++ flashinfer/kda_decode.py | 127 ++------- flashinfer/kda_prefill.py | 2 +- flashinfer/trace/templates/kda.py | 51 +--- tests/kda/test_recurrent_kda_prefill.py | 40 +-- .../recurrent_kda_q8_v16_d128.json | 51 +--- 12 files changed, 313 insertions(+), 248 deletions(-) create mode 100644 docs/api/kda.rst create mode 100644 flashinfer/kda.py diff --git a/benchmarks/bench_recurrent_kda_prefill.py b/benchmarks/bench_recurrent_kda_prefill.py index b6e4417676a..51425a775eb 100644 --- a/benchmarks/bench_recurrent_kda_prefill.py +++ b/benchmarks/bench_recurrent_kda_prefill.py @@ -42,7 +42,7 @@ import numpy as np import torch -from flashinfer.kda_decode import recurrent_kda +from flashinfer.kda import recurrent_kda from flashinfer.kda_prefill import RecurrentKDAPrefillWorkspace from flashinfer.testing import bench_gpu_time from flashinfer.utils import get_compute_capability diff --git a/docs/api/kda.rst b/docs/api/kda.rst new file mode 100644 index 00000000000..43f173f10cb --- /dev/null +++ b/docs/api/kda.rst @@ -0,0 +1,16 @@ +.. _apikda: + +flashinfer.kda +============== + +Phase-neutral recurrent Kimi Delta Attention (KDA) facade. The public +``recurrent_kda`` entry point keeps decode and speculative decode on +``flashinfer.kda_decode`` while dispatching eligible ordinary multi-token +prefill to the optimized backend described in :ref:`apikda_prefill`. + +.. currentmodule:: flashinfer.kda + +.. autosummary:: + :toctree: ../generated + + recurrent_kda diff --git a/docs/api/kda_decode.rst b/docs/api/kda_decode.rst index 5e950f2727b..e678b6963ef 100644 --- a/docs/api/kda_decode.rst +++ b/docs/api/kda_decode.rst @@ -3,10 +3,8 @@ flashinfer.kda_decode ===================== -Recurrent Kimi Delta Attention (KDA) public API. Decode and speculative -decode use the CuTe-DSL backend under ``flashinfer.kda_kernels``. Eligible -ordinary multi-token prefill calls dispatch to the optimized backend described -in :ref:`apikda_prefill`. +Key-Driven Attention (KDA) decode API. The CuTe-DSL kernel lives under +``flashinfer.kda_kernels``; this module is the public entry point. .. currentmodule:: flashinfer.kda_decode diff --git a/docs/api/kda_prefill.rst b/docs/api/kda_prefill.rst index 19457552cb6..30992059bf0 100644 --- a/docs/api/kda_prefill.rst +++ b/docs/api/kda_prefill.rst @@ -3,8 +3,8 @@ flashinfer.kda_prefill ====================== -Optimized recurrent Kimi Delta Attention (KDA) prefill support. The stable -:func:`flashinfer.kda_decode.recurrent_kda` API dispatches a strict ordinary +Optimized recurrent Kimi Delta Attention (KDA) prefill support. The +:func:`flashinfer.kda.recurrent_kda` facade dispatches a strict ordinary multi-token prefill subset to frozen FlashKDA-compatible SM100a kernels. .. currentmodule:: flashinfer.kda_prefill @@ -17,8 +17,8 @@ multi-token prefill subset to frozen FlashKDA-compatible SM100a kernels. Optimized B200 prefill subset ----------------------------- -``flashinfer.kda_decode.recurrent_kda`` uses the frozen prefill backend only -when every condition below holds: +``flashinfer.kda.recurrent_kda`` uses the frozen prefill backend only when +every condition below holds: * the device has compute capability 10.0; * input is ordinary multi-token prefill: fixed ``T > 1``, or packed input diff --git a/docs/index.rst b/docs/index.rst index adbf5e5ec8f..5789e238619 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -53,6 +53,7 @@ FlashInfer is a library and kernel generator for Large Language Models that prov api/activation api/gdn_decode api/gdn_prefill + api/kda api/kda_decode api/kda_prefill api/mamba diff --git a/flashinfer/__init__.py b/flashinfer/__init__.py index 67b777c1007..754a2405154 100644 --- a/flashinfer/__init__.py +++ b/flashinfer/__init__.py @@ -127,7 +127,7 @@ from .kda_prefill import ( RecurrentKDAPrefillWorkspace as RecurrentKDAPrefillWorkspace, ) -from .kda_decode import recurrent_kda as recurrent_kda +from .kda import recurrent_kda as recurrent_kda from .mla import BatchMLAPagedAttentionWrapper as BatchMLAPagedAttentionWrapper from . import mhc as mhc from . import msa_ops as msa_ops diff --git a/flashinfer/kda.py b/flashinfer/kda.py new file mode 100644 index 00000000000..be806f3d934 --- /dev/null +++ b/flashinfer/kda.py @@ -0,0 +1,255 @@ +""" +Copyright (c) 2025 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. +""" + +""" +Kimi Delta Attention - Public Facade +==================================== + +This phase-neutral facade preserves the top-level recurrent KDA entry point. +Eligible ordinary multi-token prefill calls use ``flashinfer.kda_prefill``; +decode and speculative decode retain the backend exposed by +``flashinfer.kda_decode``. +""" + +from typing import Optional + +import torch + +from . import kda_decode as _kda_decode +from . import kda_prefill as _kda_prefill +from .api_logging import flashinfer_api +from .trace.templates.kda import recurrent_kda_trace + + +@flashinfer_api(trace=recurrent_kda_trace) +def recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + scale: Optional[float] = None, + initial_state: Optional[torch.Tensor] = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = True, + use_gate_in_kernel: bool = False, + lower_bound: Optional[float] = None, + cu_seqlens: Optional[torch.Tensor] = None, + ssm_state_indices: Optional[torch.Tensor] = None, + num_spec_tokens: Optional[int] = None, + num_accepted_tokens: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + initial_state_source: Optional[torch.Tensor] = None, + initial_state_indices: Optional[torch.Tensor] = None, + beta_is_logit: bool = False, + seq_order: Optional[torch.Tensor] = None, + prefill_workspace: Optional[_kda_prefill.RecurrentKDAPrefillWorkspace] = None, +) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + r"""Recurrent KDA (Kimi Delta Attention) decode and prefill kernel. + + This is the public API layer for the CuTe DSL implementation in + ``flashinfer.kda_kernels.recurrent_kda``. It supports single-token decode, + fused speculative decode, GQA, optional cu_seqlens packing, and the same + gate modes as the backend implementation. On NVIDIA B200, the exact + FlashKDA-compatible subset of ordinary multi-token prefill is dispatched to + frozen SM100a kernels. All existing decode and speculative-decode calls + retain the CuTe DSL backend. + + Args: + q (torch.Tensor): + Query of shape ``[B, T, H, K]``, or + ``[1, total_tokens, H, K]`` when using ``cu_seqlens``. Must be + bfloat16. ``T=1`` selects decode; eligible ``T>1`` calls may select + the frozen prefill backend. + k (torch.Tensor): + Key with the same shape as ``q``. Must be bfloat16. + v (torch.Tensor): + Value of shape ``[B, T, HV, V]``, or + ``[1, total_tokens, HV, V]`` when packed. Must be bfloat16. GQA is + applied when ``HV != H``. + g (torch.Tensor): + Per-K-dimension gate of shape ``[B, T, HV, K]``, or + ``[1, total_tokens, HV, K]`` when packed. Must be bfloat16. + Log-space if pre-computed, raw input if + ``use_gate_in_kernel=True``. + beta (torch.Tensor): + Delta-rule learning rate of shape ``[B, T, HV]``, or + ``[1, total_tokens, HV]`` when packed. Must be bfloat16. + Pre-sigmoided unless ``beta_is_logit=True``. + A_log (Optional[torch.Tensor]): + Log decay parameter of shape ``[H]``. Must be float32. + Required when ``use_gate_in_kernel=True``. + dt_bias (Optional[torch.Tensor]): + Per-head-K decay bias of shape ``[H*K]`` or ``[H, K]``. Must be + float32. + scale (Optional[float]): + Scale factor for queries. If ``None``, defaults to ``1 / sqrt(K)``. + initial_state (Optional[torch.Tensor]): + Initial state of shape ``[N, HV, V, K]``. Must be bfloat16. + If ``None``, zero-initialized. Updated in-place. For batched spec + decode without ``cu_seqlens``, ``N`` is the packed checkpoint-slot + count ``B * (1 + num_spec_tokens)`` when ``ssm_state_indices`` is + omitted. + output_final_state (bool): + Whether to return the final state. Default: ``False``. + use_qk_l2norm_in_kernel (bool): + Whether to apply L2 normalization to Q and K. Default: ``True``. + use_gate_in_kernel (bool): + Whether to compute the gate inside the kernel from ``A_log`` and + ``g``. Default: ``False``. + lower_bound (Optional[float]): + If set, uses ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))`` + gate formula instead of softplus. Must be negative. + cu_seqlens (Optional[torch.Tensor]): + Contiguous CUDA cumulative sequence lengths of shape ``[N+1]``. + May be int32 or int64. Frozen prefill converts int32 offsets to + int64 outside graph capture; graph capture requires caller-provided + int64 offsets. For frozen prefill, values must start at zero, be + strictly increasing, and end at the total token count. This value + contract is not host-validated to avoid a device synchronization. + ssm_state_indices (Optional[torch.Tensor]): + State cache indices. Shape ``[N]`` int32 for standard decode, or + ``[N, 1+S]`` int32 for spec decode (``num_spec_tokens`` must also + be set). + num_spec_tokens (Optional[int]): + Number of speculative tokens (S). When set, processes 1+S tokens in + a single fused kernel launch. Must be >= 1. + num_accepted_tokens (Optional[torch.Tensor]): + Per-sequence accepted token count from the previous spec decode + round. Shape ``[N]`` int32. If ``None``, initial state is loaded + from ``ssm_state_indices[n, 0]``. Values above ``1+S`` are clamped + to the final checkpoint slot. + output (Optional[torch.Tensor]): + Pre-allocated output tensor. Shape ``[B, T, HV, V]`` for fixed + layout, or the corresponding packed/speculative shape when using + ``cu_seqlens``. If ``None``, a new tensor is allocated. Frozen + prefill requires storage disjoint from Q, K, V, G, beta, and + ``initial_state``. + initial_state_source (Optional[torch.Tensor]): + Optional read-only committed state pool ``[N0, HV, V, K]``. When + provided, token 0 is loaded from this pool instead of + ``initial_state``. + initial_state_indices (Optional[torch.Tensor]): + Source slot per sequence, shape ``[N]`` int32. Required together + with ``initial_state_source``. + beta_is_logit (bool): + If ``True``, apply sigmoid to ``beta`` inside the recurrent kernel. + seq_order (Optional[torch.Tensor]): + Optional packed-prefill sequence order, as a contiguous CUDA int32 + permutation of shape ``[N]``. Sorting by descending sequence length + improves tail utilization. It is only consumed by the frozen + FlashKDA prefill backend; prepare it before CUDA graph capture or + timed launches. Fixed-layout prefill and decode calls must leave it + as ``None``. + prefill_workspace (Optional[RecurrentKDAPrefillWorkspace]): + Caller-owned workspace for the frozen B200 prefill backend. It is + optional for eager execution and required for CUDA graph capture. + Warm it eagerly with the exact tensors on the capture stream before + capture. Use one workspace per captured ``recurrent_kda`` + invocation. + + Returns: + Tuple of ``(output, final_state)`` where ``final_state`` is ``None`` + when ``output_final_state=False``. See + :func:`flashinfer.kda_kernels.recurrent_kda.run_recurrent_kda` for the + backend implementation. + """ + if prefill_workspace is not None and not isinstance( + prefill_workspace, _kda_prefill.RecurrentKDAPrefillWorkspace + ): + raise TypeError("prefill_workspace must be a RecurrentKDAPrefillWorkspace") + + use_flash_kda_prefill = _kda_prefill._flash_kda_prefill_is_eligible( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=initial_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + use_gate_in_kernel=use_gate_in_kernel, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_spec_tokens=num_spec_tokens, + num_accepted_tokens=num_accepted_tokens, + output=output, + initial_state_source=initial_state_source, + initial_state_indices=initial_state_indices, + beta_is_logit=beta_is_logit, + ) + if use_flash_kda_prefill: + assert A_log is not None + assert dt_bias is not None + assert lower_bound is not None + return _kda_prefill._run_flash_kda_prefill( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + output=output, + seq_order=seq_order, + prefill_workspace=prefill_workspace, + ) + + if prefill_workspace is not None: + raise ValueError( + "prefill_workspace is only supported by eligible ordinary " + "prefill on the frozen B200 FlashKDA backend" + ) + if seq_order is not None: + raise ValueError( + "seq_order is only supported by eligible packed ordinary prefill " + "on the frozen B200 FlashKDA backend" + ) + if _kda_decode._run_recurrent_kda is None: + raise NotImplementedError("recurrent KDA backend is unavailable") + + return _kda_decode._run_recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + use_gate_in_kernel=use_gate_in_kernel, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_spec_tokens=num_spec_tokens, + num_accepted_tokens=num_accepted_tokens, + output=output, + initial_state_source=initial_state_source, + initial_state_indices=initial_state_indices, + beta_is_logit=beta_is_logit, + ) diff --git a/flashinfer/kda_decode.py b/flashinfer/kda_decode.py index 37d13e4de9d..0c9ef48fb44 100644 --- a/flashinfer/kda_decode.py +++ b/flashinfer/kda_decode.py @@ -18,17 +18,14 @@ Kimi Delta Attention Decode - API Layer ======================================= -This file preserves the public recurrent KDA decode API. Eligible ordinary -multi-token prefill calls are dispatched to the frozen backend implemented in -``flashinfer.kda_prefill``; decode and speculative decode retain the -CuTe-DSL backend under ``flashinfer.kda_kernels``. +This file provides the public API for recurrent KDA decode operations. +Kernel implementations are in flashinfer/kda_kernels/. """ from typing import Optional import torch -from . import kda_prefill as _kda_prefill from .api_logging import flashinfer_api from .trace.templates.kda import recurrent_kda_trace @@ -64,46 +61,34 @@ def recurrent_kda( initial_state_source: Optional[torch.Tensor] = None, initial_state_indices: Optional[torch.Tensor] = None, beta_is_logit: bool = False, - seq_order: Optional[torch.Tensor] = None, - prefill_workspace: Optional[_kda_prefill.RecurrentKDAPrefillWorkspace] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - r"""Recurrent KDA (Kimi Delta Attention) decode and prefill kernel. + r"""Recurrent KDA (Kimi Delta Attention) decode kernel. This is the public API layer for the CuTe DSL implementation in ``flashinfer.kda_kernels.recurrent_kda``. It supports single-token decode, fused speculative decode, GQA, optional cu_seqlens packing, and the same - gate modes as the backend implementation. On NVIDIA B200, the exact - FlashKDA-compatible subset of ordinary multi-token prefill is dispatched to - frozen SM100a kernels. All existing decode and speculative-decode calls - retain the CuTe DSL backend. + gate modes as the backend implementation. Args: q (torch.Tensor): - Query of shape ``[B, T, H, K]``, or - ``[1, total_tokens, H, K]`` when using ``cu_seqlens``. Must be - bfloat16. ``T=1`` selects decode; eligible ``T>1`` calls may select - the frozen prefill backend. + Current query of shape ``[B, 1, H, K]``, or ``[1, total_tokens, H, K]`` + when using ``cu_seqlens``. Must be bfloat16. k (torch.Tensor): - Key with the same shape as ``q``. Must be bfloat16. + Current key of shape ``[B, 1, H, K]``. Must be bfloat16. v (torch.Tensor): - Value of shape ``[B, T, HV, V]``, or - ``[1, total_tokens, HV, V]`` when packed. Must be bfloat16. GQA is - applied when ``HV != H``. + Current value of shape ``[B, 1, HV, V]``. Must be bfloat16. + GQA is applied when ``HV != H``. g (torch.Tensor): - Per-K-dimension gate of shape ``[B, T, HV, K]``, or - ``[1, total_tokens, HV, K]`` when packed. Must be bfloat16. - Log-space if pre-computed, raw input if - ``use_gate_in_kernel=True``. + Per-K-dimension gate of shape ``[B, 1, HV, K]``. Must be bfloat16. + Log-space if pre-computed, raw input if ``use_gate_in_kernel=True``. beta (torch.Tensor): - Delta-rule learning rate of shape ``[B, T, HV]``, or - ``[1, total_tokens, HV]`` when packed. Must be bfloat16. + Delta-rule learning rate of shape ``[B, 1, HV]``. Must be bfloat16. Pre-sigmoided unless ``beta_is_logit=True``. A_log (Optional[torch.Tensor]): Log decay parameter of shape ``[H]``. Must be float32. Required when ``use_gate_in_kernel=True``. dt_bias (Optional[torch.Tensor]): - Per-head-K decay bias of shape ``[H*K]`` or ``[H, K]``. Must be - float32. + Per-head-K decay bias of shape ``[H*K]``. Must be float32. scale (Optional[float]): Scale factor for queries. If ``None``, defaults to ``1 / sqrt(K)``. initial_state (Optional[torch.Tensor]): @@ -123,12 +108,7 @@ def recurrent_kda( If set, uses ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))`` gate formula instead of softplus. Must be negative. cu_seqlens (Optional[torch.Tensor]): - Contiguous CUDA cumulative sequence lengths of shape ``[N+1]``. - May be int32 or int64. Frozen prefill converts int32 offsets to - int64 outside graph capture; graph capture requires caller-provided - int64 offsets. For frozen prefill, values must start at zero, be - strictly increasing, and end at the total token count. This value - contract is not host-validated to avoid a device synchronization. + Cumulative sequence lengths of shape ``[N+1]``. Must be int32. ssm_state_indices (Optional[torch.Tensor]): State cache indices. Shape ``[N]`` int32 for standard decode, or ``[N, 1+S]`` int32 for spec decode (``num_spec_tokens`` must also @@ -142,11 +122,9 @@ def recurrent_kda( from ``ssm_state_indices[n, 0]``. Values above ``1+S`` are clamped to the final checkpoint slot. output (Optional[torch.Tensor]): - Pre-allocated output tensor. Shape ``[B, T, HV, V]`` for fixed - layout, or the corresponding packed/speculative shape when using - ``cu_seqlens``. If ``None``, a new tensor is allocated. Frozen - prefill requires storage disjoint from Q, K, V, G, beta, and - ``initial_state``. + Pre-allocated output tensor. Shape ``[B, 1, HV, V]`` for standard + decode, ``[1, N*(1+S), HV, V]`` for spec decode with + ``cu_seqlens``. If ``None``, a new tensor is allocated. initial_state_source (Optional[torch.Tensor]): Optional read-only committed state pool ``[N0, HV, V, K]``. When provided, token 0 is loaded from this pool instead of @@ -156,19 +134,6 @@ def recurrent_kda( with ``initial_state_source``. beta_is_logit (bool): If ``True``, apply sigmoid to ``beta`` inside the recurrent kernel. - seq_order (Optional[torch.Tensor]): - Optional packed-prefill sequence order, as a contiguous CUDA int32 - permutation of shape ``[N]``. Sorting by descending sequence length - improves tail utilization. It is only consumed by the frozen - FlashKDA prefill backend; prepare it before CUDA graph capture or - timed launches. Fixed-layout prefill and decode calls must leave it - as ``None``. - prefill_workspace (Optional[RecurrentKDAPrefillWorkspace]): - Caller-owned workspace for the frozen B200 prefill backend. It is - optional for eager execution and required for CUDA graph capture. - Warm it eagerly with the exact tensors on the capture stream before - capture. Use one workspace per captured ``recurrent_kda`` - invocation. Returns: Tuple of ``(output, final_state)`` where ``final_state`` is ``None`` @@ -176,64 +141,6 @@ def recurrent_kda( :func:`flashinfer.kda_kernels.recurrent_kda.run_recurrent_kda` for the backend implementation. """ - if prefill_workspace is not None and not isinstance( - prefill_workspace, _kda_prefill.RecurrentKDAPrefillWorkspace - ): - raise TypeError("prefill_workspace must be a RecurrentKDAPrefillWorkspace") - - use_flash_kda_prefill = _kda_prefill._flash_kda_prefill_is_eligible( - q=q, - k=k, - v=v, - g=g, - beta=beta, - A_log=A_log, - dt_bias=dt_bias, - initial_state=initial_state, - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, - use_gate_in_kernel=use_gate_in_kernel, - lower_bound=lower_bound, - cu_seqlens=cu_seqlens, - ssm_state_indices=ssm_state_indices, - num_spec_tokens=num_spec_tokens, - num_accepted_tokens=num_accepted_tokens, - output=output, - initial_state_source=initial_state_source, - initial_state_indices=initial_state_indices, - beta_is_logit=beta_is_logit, - ) - if use_flash_kda_prefill: - assert A_log is not None - assert dt_bias is not None - assert lower_bound is not None - return _kda_prefill._run_flash_kda_prefill( - q=q, - k=k, - v=v, - g=g, - beta=beta, - A_log=A_log, - dt_bias=dt_bias, - scale=scale, - initial_state=initial_state, - output_final_state=output_final_state, - lower_bound=lower_bound, - cu_seqlens=cu_seqlens, - output=output, - seq_order=seq_order, - prefill_workspace=prefill_workspace, - ) - - if prefill_workspace is not None: - raise ValueError( - "prefill_workspace is only supported by eligible ordinary " - "prefill on the frozen B200 FlashKDA backend" - ) - if seq_order is not None: - raise ValueError( - "seq_order is only supported by eligible packed ordinary prefill " - "on the frozen B200 FlashKDA backend" - ) if _run_recurrent_kda is None: raise NotImplementedError("recurrent KDA backend is unavailable") diff --git a/flashinfer/kda_prefill.py b/flashinfer/kda_prefill.py index da2a8fe8453..11fe6255e21 100644 --- a/flashinfer/kda_prefill.py +++ b/flashinfer/kda_prefill.py @@ -70,7 +70,7 @@ class RecurrentKDAPrefillWorkspace(_RecurrentKDAPrefillWorkspaceBase): """Caller-owned storage required for recurrent-KDA CUDA graph capture. Construct one workspace per captured - :func:`flashinfer.kda_decode.recurrent_kda` invocation on the graph's CUDA + :func:`flashinfer.kda.recurrent_kda` invocation on the graph's CUDA device. Warm it by invoking that function eagerly with the exact tensors and capture stream, then synchronize that stream before capture. The workspace owns optional final-state scratch for calls without an initial diff --git a/flashinfer/trace/templates/kda.py b/flashinfer/trace/templates/kda.py index 5813c2d5d28..e27552501b2 100644 --- a/flashinfer/trace/templates/kda.py +++ b/flashinfer/trace/templates/kda.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""TraceTemplate for recurrent Kimi Delta Attention (KDA).""" +"""TraceTemplate for recurrent Key-Driven Attention (KDA) decode.""" from ..template import Const, Scalar, Tensor, TraceTemplate, Var @@ -21,9 +21,8 @@ op_type="kda", name_prefix="recurrent_kda", description=( - "Recurrent Kimi Delta Attention decode/spec-decode plus exact " - "FlashKDA-compatible ordinary prefill on B200, with " - "per-key-dimension gating and recurrent BF16 state." + "Recurrent Key-Driven Attention decode with per-key-dimension gating " + "and an optional read-only committed-state source." ), axes={ "batch_size": Var(description="Number of input batch rows."), @@ -36,9 +35,6 @@ "state_pool_size": Var(description="Number of writable state slots."), "source_pool_size": Var(description="Number of committed-state slots."), "num_sequences": Var(description="Number of state-source indices."), - "num_sequences_plus_one": Var( - description="Number of packed cumulative-length entries." - ), }, inputs={ "q": Tensor(["batch_size", "seq_len", "num_q_heads", "head_dim"]), @@ -46,58 +42,20 @@ "v": Tensor(["batch_size", "seq_len", "num_v_heads", "head_dim"]), "g": Tensor(["batch_size", "seq_len", "num_v_heads", "head_dim"]), "beta": Tensor(["batch_size", "seq_len", "num_v_heads"]), - "A_log": Tensor( - ["num_q_heads"], - dtype="float32", - optional=True, - description="FP32 per-query-head log decay rate.", - ), - "dt_bias": Tensor( - ["num_q_heads", "head_dim"], - dtype="float32", - optional=True, - description=( - "FP32 per-head/key decay bias; flattened [H*D] or [H,D] " - "storage is accepted by the API." - ), - ), "initial_state": Tensor( ["state_pool_size", "num_v_heads", "head_dim", "head_dim"], - dtype="bfloat16", optional=True, ), "initial_state_source": Tensor( ["source_pool_size", "num_v_heads", "head_dim", "head_dim"], - dtype="bfloat16", optional=True, description="Read-only committed-state pool.", ), "initial_state_indices": Tensor( ["num_sequences"], - dtype="int32", optional=True, description="Committed-state slot selected for each sequence.", ), - "cu_seqlens": Tensor( - ["num_sequences_plus_one"], - optional=True, - description="Packed cumulative sequence lengths (int32 or int64).", - ), - "num_accepted_tokens": Tensor( - ["num_sequences"], - dtype="int32", - optional=True, - description="Accepted-token counts for speculative decode.", - ), - "seq_order": Tensor( - ["num_sequences"], - dtype="int32", - optional=True, - description=( - "Packed-prefill sequence permutation, normally descending " - "by sequence length." - ), - ), "scale": Scalar("float32", optional=True), "output_final_state": Scalar("int32", optional=True), "use_qk_l2norm_in_kernel": Scalar("int32", optional=True), @@ -110,7 +68,6 @@ "output": Tensor( ["batch_size", "seq_len", "num_v_heads", "head_dim"], dtype_from="q", - param="output", ), "final_state": Tensor( ["state_pool_size", "num_v_heads", "head_dim", "head_dim"], @@ -122,5 +79,5 @@ "num_v_heads % num_q_heads == 0", "head_dim in (64, 128)", ], - tags=["stage:decode", "stage:prefill", "status:verified"], + tags=["stage:decode", "status:verified"], ) diff --git a/tests/kda/test_recurrent_kda_prefill.py b/tests/kda/test_recurrent_kda_prefill.py index 6cd9a969239..5e2582f351b 100644 --- a/tests/kda/test_recurrent_kda_prefill.py +++ b/tests/kda/test_recurrent_kda_prefill.py @@ -20,16 +20,17 @@ import torch.nn.functional as F import flashinfer -from flashinfer.kda_decode import recurrent_kda +from flashinfer.kda import recurrent_kda from flashinfer.kda_prefill import RecurrentKDAPrefillWorkspace from flashinfer.utils import get_compute_capability kda_decode_api = importlib.import_module("flashinfer.kda_decode") +kda_api = importlib.import_module("flashinfer.kda") kda_prefill_api = importlib.import_module("flashinfer.kda_prefill") -def test_public_api_keeps_decode_entry_and_prefill_workspace(): - assert flashinfer.recurrent_kda is kda_decode_api.recurrent_kda +def test_public_api_uses_phase_neutral_facade_and_prefill_workspace(): + assert flashinfer.recurrent_kda is kda_api.recurrent_kda assert ( flashinfer.RecurrentKDAPrefillWorkspace is kda_prefill_api.RecurrentKDAPrefillWorkspace @@ -181,6 +182,11 @@ def old_backend(**kwargs): return sentinel monkeypatch.setattr(kda_decode_api, "_run_recurrent_kda", old_backend) + monkeypatch.setattr( + kda_decode_api, + "recurrent_kda", + lambda *args, **kwargs: pytest.fail("facade nested the decorated decode API"), + ) monkeypatch.setattr( kda_prefill_api, "_get_flash_kda_prefill_module", @@ -582,34 +588,6 @@ def test_workspace_rejects_a_different_stream(cuda_device, monkeypatch): ) -def test_recurrent_kda_prefill_trace_has_semantic_inputs(): - q = torch.empty((1, 8, 2, 128), dtype=torch.bfloat16) - trace = recurrent_kda.fi_trace( - q=q, - k=q, - v=q, - g=q, - beta=torch.empty((1, 8, 2), dtype=torch.bfloat16), - A_log=torch.empty(2), - dt_bias=torch.empty((2, 128)), - cu_seqlens=torch.tensor([0, 3, 8], dtype=torch.int64), - seq_order=torch.tensor([1, 0], dtype=torch.int32), - use_qk_l2norm_in_kernel=True, - use_gate_in_kernel=True, - lower_bound=-5.0, - beta_is_logit=True, - ) - assert trace["op_type"] == "kda" - assert "stage:prefill" in trace["tags"] - for name in ( - "A_log", - "dt_bias", - "cu_seqlens", - "seq_order", - ): - assert name in trace["inputs"] - - def test_flash_kda_jit_getter_is_importable(): import flashinfer from flashinfer.jit.flash_kda import get_flash_kda_prefill_module diff --git a/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json b/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json index c45683cb341..09e5f6ae981 100644 --- a/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json +++ b/tests/trace/fi_trace_out/recurrent_kda_q8_v16_d128.json @@ -1,11 +1,10 @@ { "name": "recurrent_kda_q8_v16_d128", - "description": "Recurrent Kimi Delta Attention decode/spec-decode plus exact FlashKDA-compatible ordinary prefill on B200, with per-key-dimension gating and recurrent BF16 state.", + "description": "Recurrent Key-Driven Attention decode with per-key-dimension gating and an optional read-only committed-state source.", "op_type": "kda", "tags": [ "fi_api:flashinfer.kda_decode.recurrent_kda", "stage:decode", - "stage:prefill", "status:verified" ], "axes": { @@ -43,10 +42,6 @@ "num_sequences": { "type": "var", "description": "Number of state-source indices." - }, - "num_sequences_plus_one": { - "type": "var", - "description": "Number of packed cumulative-length entries." } }, "constraints": [ @@ -98,23 +93,6 @@ ], "dtype": "bfloat16" }, - "A_log": { - "shape": [ - "num_q_heads" - ], - "dtype": "float32", - "optional": true, - "description": "FP32 per-query-head log decay rate." - }, - "dt_bias": { - "shape": [ - "num_q_heads", - "head_dim" - ], - "dtype": "float32", - "optional": true, - "description": "FP32 per-head/key decay bias; flattened [H*D] or [H,D] storage is accepted by the API." - }, "initial_state": { "shape": [ "state_pool_size", @@ -144,30 +122,6 @@ "optional": true, "description": "Committed-state slot selected for each sequence." }, - "cu_seqlens": { - "shape": [ - "num_sequences_plus_one" - ], - "dtype": "unknown", - "optional": true, - "description": "Packed cumulative sequence lengths (int32 or int64)." - }, - "num_accepted_tokens": { - "shape": [ - "num_sequences" - ], - "dtype": "int32", - "optional": true, - "description": "Accepted-token counts for speculative decode." - }, - "seq_order": { - "shape": [ - "num_sequences" - ], - "dtype": "int32", - "optional": true, - "description": "Packed-prefill sequence permutation, normally descending by sequence length." - }, "scale": { "shape": null, "dtype": "float32", @@ -212,8 +166,7 @@ "num_v_heads", "head_dim" ], - "dtype": "bfloat16", - "param": "output" + "dtype": "bfloat16" }, "final_state": { "shape": [