diff --git a/BENCHMARK_H200.md b/BENCHMARK_H200.md index ce0d46f7..5c82b664 100644 --- a/BENCHMARK_H200.md +++ b/BENCHMARK_H200.md @@ -55,5 +55,5 @@ Summary (28 configs): **avg=1.58x**, min=1.02x, max=2.51x. To reproduce: ```bash -python benchmarks/bench_kda_fused_fwd.py --mode both +python benchmarks/bench_kda_sm90_prefill.py --mode both ``` diff --git a/README.md b/README.md index 09418811..c5e145f8 100644 --- a/README.md +++ b/README.md @@ -149,23 +149,25 @@ python benchmarks/generate_benchmark_hopper_md.py ```bash # Tests for modular KDA forward against FLA Triton implementation -python -m pytest tests/test_kda_compare_fla.py -v +python -m pytest tests/test_kda_sm100_chunk_vs_fla.py -v # Tests for modular KDA forward against naive KDA reference -python -m pytest tests/test_kda.py -v -# Tests for KDA fused forward +python -m pytest tests/test_kda_sm100_chunk_vs_naive.py -v +# Tests for the KDA fused forward (SM90 CUDA C++) python -m pytest tests/test_kda_fused_fwd.py -v -# Tests for Lightning Attention fused forward -python tests/test_lightning_attn.py +# Tests for the SM90 CuTeDSL two-kernel prefill + intracard CP (vs FLA) +python -m pytest tests/test_kda_sm90_prefill_vs_fla.py tests/test_kda_sm90_intracard_cp.py -v +# Tests for Lightning Attention prefill +python tests/test_lightning_sm100_prefill.py # Tests for Lightning Attention decode -python -m pytest tests/test_la_decode.py -v +python -m pytest tests/test_lightning_decode.py -v -# test_kda.py and test_kda_compare_fla.py support a fast/slow split. +# test_kda_sm100_chunk_vs_naive.py and test_kda_sm100_chunk_vs_fla.py support a fast/slow split. # Fast (default) — representative correctness paths for default CI and local iteration -python -m pytest tests/test_kda.py tests/test_kda_compare_fla.py -v +python -m pytest tests/test_kda_sm100_chunk_vs_naive.py tests/test_kda_sm100_chunk_vs_fla.py -v # Slow — broader stress coverage for nightly or manual runs -python -m pytest -m kda_slow tests/test_kda.py tests/test_kda_compare_fla.py -v +python -m pytest -m kda_slow tests/test_kda_sm100_chunk_vs_naive.py tests/test_kda_sm100_chunk_vs_fla.py -v # Full sweep (fast + slow) — run before submitting a PR -python -m pytest -m kda_full tests/test_kda.py tests/test_kda_compare_fla.py -v +python -m pytest -m kda_full tests/test_kda_sm100_chunk_vs_naive.py tests/test_kda_sm100_chunk_vs_fla.py -v ```
diff --git a/REPO_LAYOUT.md b/REPO_LAYOUT.md index 81358296..f66af2af 100644 --- a/REPO_LAYOUT.md +++ b/REPO_LAYOUT.md @@ -7,7 +7,6 @@ cuLA/ ├── cula/ # Python package (pip install -e .) │ ├── __init__.py │ ├── utils.py # arch asserts, get_pre_scan, cu_seqlens helpers, ... -│ ├── cudac.py # re-export shim over the compiled C++ extension(s) │ │ │ ├── kda/ # KDA PUBLIC API + autograd + dispatch (NO kernels) │ │ ├── __init__.py # lazy PUBLIC API: chunk_kda, kda_prefill_hopper, @@ -16,7 +15,7 @@ cuLA/ │ │ ├── chunk_fwd.py # chunk_kda_fwd — fwd orchestration (lazy-imports kernels) │ │ ├── chunk_intra.py # fwd intra (C++ ext) + bwd intra (Triton) │ │ ├── chunk_bwd.py # chunk_kda_bwd — Triton + FLA + CuTeDSL + C++ mix -│ │ └── hopper_fused_fwd.py # cula_kda_prefill (=kda_prefill_hopper) — SM90 prefill via the C++ kernel (cula.cudac) +│ │ └── flashkda.py # cula_kda_prefill (=flashkda_prefill) — SM90 two-kernel K1+K2 prefill, fwd-only │ │ │ ├── lightning/ # [non-KDA] Lightning Attention operator (LinearAttentionChunkwiseDecay, lightning_attn_fwd, linear_attention_decode) │ │ └── __init__.py @@ -25,15 +24,19 @@ cuLA/ │ ├── __init__.py # exports kda_decode, fused_sigmoid_..., linear_attention_decode │ ├── inv.py / ptx.py # shared low-level helpers │ ├── sm100/ # SM100 shared helper only +│ │ ├── __init__.py │ │ └── ptx.py # shared PTX helpers (used by KDA + lightning kernels) │ │ -│ ├── kda/ # ★ KDA Python backends — by arch (sm100 today) -│ │ ├── policy.py # SM100 CP dispatch policy: use_intracard_cp:"auto"|bool +│ ├── kda/ # ★ ALL KDA backend kernels — by arch (sm90 / sm100) +│ │ ├── policy.py # CP dispatch policy: sm90/sm100 decisions, use_intracard_cp:"auto"|bool │ │ ├── sm100/ # SM100 (Blackwell) modular-chunk kernels │ │ │ ├── delta_h.py # recurrence (chunk_gated_delta_rule_fwd_h) │ │ │ ├── fwd_o.py # output (chunk_gla_fwd_o) │ │ │ ├── bwd_wy_dqkg.py# backward wy/dqkg fused (used by chunk_bwd) │ │ │ └── cp/ # SM100 intracard-CP: chunk_delta_h, pre_scan, merge +│ │ ├── sm90/ # SM90 (Hopper) two-kernel FlashKDA prefill, fwd-only +│ │ │ ├── fwd.py k1.py k2.py # flash_kda_fwd → launch_k1 (prepare) + launch_k2 (recurrence) +│ │ │ └── cp/ # SM90 intracard-CP: flashkda, pre_scan, merge, plan │ │ ├── decode/ # single-token + MTP decode │ │ │ ├── cute.py # kda_decode / fused_sigmoid_gating_delta_rule_update (CuTe DSL) │ │ │ ├── mtp.py # kda_decode_mtp recurrent / recurrent_ws MTP verify (CuTe DSL) @@ -49,10 +52,9 @@ cuLA/ │ └── experimental/ │ └── linear_attn_prototype.py # [non-KDA] unwired normalized-linear-attn prototype │ -├── csrc/ # CUDA C++ / CUTLASS -│ ├── api/{kda_sm90.cu, kda_sm100.cu} # PyBind11 (cula.cudac): SM90 prefill + SM100 chunk intra/recompute_w_u -│ ├── kda/sm90/ # SM90 (Hopper) KDA C++ kernels (CUTLASS 3.x, TMA/wgmma) -│ ├── kda/sm100/ # Blackwell KDA C++ kernels (CUTLASS 3.x + UMMA) +├── csrc/ # CUDA C++ / CUTLASS — SM100 ONLY +│ ├── api/{pybind.cu, kda_sm100.cu} # PyBind11 (module cula.cudac): chunk intra + recompute_w_u +│ ├── kda/sm100/ # Blackwell KDA C++ kernels (CUTLASS 3.x + UMMA); kda_fwd_sm100.cu is the only .cu │ └── kerutils/include/ # shared C++ headers (generic device helpers sm80/sm90/sm100, host) │ ├── benchmarks/ tests/ docs/ @@ -66,8 +68,8 @@ cuLA/ | Directory | Language | Description | |-----------|----------|-------------| -| `cula/kda/` | Python | KDA **public API only** — autograd + dispatch, no kernels. Two prefill entries: modular chunk `chunk_kda` (SM100) and `kda_prefill_hopper` (SM90, driving the C++ kernel). | -| `cula/ops/kda/` | Python (CuTe DSL) | **KDA Python backends**, by arch: `sm100/` (+cp), `decode/`, `experimental/`, plus `policy.py` (CP dispatch). | +| `cula/kda/` | Python | KDA **public API only** — autograd + dispatch, no kernels. Two prefill entries: modular chunk `chunk_kda` (SM100) and two-kernel K1+K2 `flashkda_prefill` (SM90; the CUDA C++ fused path stays on `kda_prefill_hopper`). See [`cula/kda/README.md`](cula/kda/README.md). | +| `cula/ops/kda/` | Python (CuTe DSL) | **All KDA backends**, by arch: `sm100/` (+cp), `sm90/` (+cp), `decode/`, `experimental/`, plus `policy.py` (CP dispatch). Both prefill backends are chunked forward; arch is the discriminator (1 impl each). | | `cula/ops/lightning/` · `cula/ops/experimental/` | Python (CuTe DSL) | `[non-KDA]` Lightning/linear attention kernels. | | `cula/ops/{inv,ptx}.py`, `cula/ops/sm100/ptx.py` | Python | Shared low-level helpers (kept in place; not KDA-specific). | | `csrc/kda/{sm90,sm100}/` · `csrc/api/` | CUDA C++ | Hopper SM90 prefill + Blackwell SM100 (chunk intra + recompute_w_u), exposed as `cula.cudac`. | diff --git a/USAGE.md b/USAGE.md index 6a4c1d8e..9c9a3664 100644 --- a/USAGE.md +++ b/USAGE.md @@ -12,6 +12,7 @@ cuLA provides two KDA kernel implementations targeting different GPU architectur |---|---|---| | Modular Forward | Blackwell (SM100) | `from cula.kda import chunk_kda` | | Fused Forward | Hopper (SM90) | `from cula.kda import kda_prefill_hopper` | +| Two-Kernel Prefill | Hopper (SM90) | `from cula.kda import flashkda_prefill` | Both are drop-in replacements for [FLA](https://github.com/fla-org/flash-linear-attention)'s `chunk_kda` — just change the import. @@ -115,9 +116,77 @@ print(f'Final state shape: {final_state.shape}') # [2, 32, 128, 128] --- -## Intra-Card Context Parallel (chunk_delta_h) +### Two-Kernel Prefill (SM90 — Hopper) -cuLA includes an intra-card context parallel (CP) path for `chunk_gated_delta_rule_fwd_h`. Long sequences are split into sub-sequences, processed independently in parallel, then merged via a prefix-scan step — unlocking sequence-dimension parallelism on a single GPU. +The SM90 prefill is a **two-kernel pipeline** — K1 (Prepare) → K2 (Recurrence) — *not* a single fused kernel. Intra-chunk attention, inter-chunk state propagation, and output are computed across the two kernels. **Forward-only; backward is not yet implemented.** + +#### Example + +```python +import torch +from cula.kda import flashkda_prefill + +B, T, H, K, V = 2, 2048, 32, 128, 128 +device = 'cuda' + +q = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) +k = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) +v = torch.randn(B, T, H, V, device=device, dtype=torch.bfloat16) +g = torch.randn(B, T, H, K, device=device, dtype=torch.bfloat16) * 0.1 +beta = torch.randn(B, T, H, device=device, dtype=torch.bfloat16).sigmoid() +A_log = torch.randn(H, device=device, dtype=torch.float32) * 0.01 +dt_bias = torch.zeros(H * K, device=device, dtype=torch.float32) +init_state = torch.zeros(B, H, K, V, device=device, dtype=torch.float32) + +o, final_state = flashkda_prefill( + q=q, k=k, v=v, g=g, beta=beta, + A_log=A_log, dt_bias=dt_bias, + initial_state=init_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=-5.0, +) + +print(f'Output shape: {o.shape}') # [2, 2048, 32, 128] +print(f'Final state shape: {final_state.shape}') # [2, 32, 128, 128] +``` + +**Notes** + +- Mainly **suitable for large-batch inference**. When both batch size and head count are small, throughput on long sequences is recovered by enabling **intra-card CP** (`use_intracard_cp`), which parallelizes the sequence-length dimension on a single GPU — see [Intra-Card Context Parallel](#intra-card-context-parallel). +- **Matrix inversion uses fp16 precision**, which is faster and occupies less shared memory but introduces minor numerical differences compared to tf32 inversion. +- **Intra-subchunk attention uses g-first as anchor**, which causes some numerical differences compared with the FLA Triton implementation (FLA uses g-half as anchor in the diagonal). + +--- + +## Intra-Card Context Parallel + +Long sequences can be split into sub-sequences, processed in parallel on one GPU, and merged via a prefix scan — unlocking sequence-dimension parallelism. **Default off; inference-only.** Two surfaces: + +### SM90 — via `flashkda_prefill(use_intracard_cp=...)` + +Pass `use_intracard_cp` (alias `use_cp`) to the Hopper prefill: + +- **`"auto"`** — the auto-router enables CP only when it predicts a speedup, otherwise runs serial (no error, no regression). +- **`True`** — force CP; raises if the shape cannot be meaningfully split. +- **`False`** (default) — CP off. + +Works with **any sequence length** (non-CHUNK-aligned is handled internally) and **dense or varlen** input. The auto decision uses the device SM count plus two tunables in `cula/ops/kda/sm90/cp/plan.py` (`CULA_KDA_CP_RERUN_RATIO`, `CULA_KDA_CP_AUTO_MIN_SEG_TILES`). + +```python +o, final_state = flashkda_prefill( + q=q, k=k, v=v, g=g, beta=beta, A_log=A_log, dt_bias=dt_bias, + cu_seqlens=cu_seqlens, # varlen packed (int32) + output_final_state=True, use_gate_in_kernel=True, safe_gate=True, lower_bound=-5.0, + use_intracard_cp="auto", # "auto" | True | False (default False) +) +``` + +### SM100 — low-level `chunk_gated_delta_rule_fwd_h` + +cuLA also exposes intra-card CP directly on the SM100 recurrence `chunk_gated_delta_rule_fwd_h`, gated by an environment variable. **Requirements** diff --git a/benchmarks/bench_kda_sm90_cp.py b/benchmarks/bench_kda_sm90_cp.py new file mode 100644 index 00000000..284e1894 --- /dev/null +++ b/benchmarks/bench_kda_sm90_cp.py @@ -0,0 +1,194 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +""" +bench_kda_sm90_cp.py — Benchmark: SM90 intracard CP speedup (CP-on vs CP-off) + +Measures the speedup of the SM90 intracard context-parallel path against +the serial K1+K2 baseline across varlen configurations. + +Usage: + python bench_kda_sm90_cp.py [--ncu] [--sanitizer] + +With --ncu, warmup=1 and iters=1 for ncu profiling: + ncu --set full -o report python bench_kda_sm90_cp.py --ncu +""" + +import argparse +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +import torch + +from benchmarks.utils import SEED, exclusive_cumsum, prepare_safe_gate_inputs, set_seed +from cula.kda import flashkda_prefill as cula_kda_prefill +from cula.utils import assert_hopper, get_device_sm_count + +D = 128 +H_VALUES = [4, 8] +WARMUP = 10 +N_ITERS = 100 +NCU_MODE = False +SANITIZER_MODE = False + +# (tag, seq_lens) — each entry is tested at every H in H_VALUES. +# SM90 CHUNK=16, so sequences need to be long enough (>= ~8K tiles) for CP to pay off. +CONFIGS = [ + ("T=4K", [4096]), + ("T=8K", [8192]), + ("T=16K", [16384]), + ("T=32K", [32768]), + ("T=64K", [65536]), + ("2x16K", [16384, 16384]), + ("32K+4K", [32768, 4096]), + ("32K+1K", [32768, 1024]), + ("64K+1K", [65536, 1024]), + ("64K+2x1K", [65536, 1024, 1024]), + ("64K+5x1K", [65536] + [1024] * 5), +] + + +def time_kernel(fn, warmup=None, n_iters=None): + if warmup is None: + warmup = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP + if n_iters is None: + n_iters = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start_evt = torch.cuda.Event(enable_timing=True) + end_evt = torch.cuda.Event(enable_timing=True) + start_evt.record() + for _ in range(n_iters): + fn() + end_evt.record() + torch.cuda.synchronize() + return start_evt.elapsed_time(end_evt) / n_iters + + +def run_kernel(q, k, v, g, beta, scale, A_log, dt_bias, cu_seqlens, lower_bound, *, use_cp): + cula_kda_prefill( + q, + k, + v, + g, + beta, + scale=scale, + A_log=A_log, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + output_final_state=False, + safe_gate=True, + lower_bound=lower_bound, + use_intracard_cp="auto" if use_cp else False, + ) + + +def bench_cp(h_values, configs): + print("\n" + "=" * 100) + print(" SM90 Intracard CP Benchmark: CP-on vs CP-off") + print("=" * 100) + + device = torch.device("cuda") + assert_hopper(device) + num_sms = get_device_sm_count(device) + print(f" [Device] {torch.cuda.get_device_name(0)} SMs={num_sms}") + results = [] + + for H in h_values: + for tag, seq_lens in configs: + set_seed(SEED) + torch.cuda.empty_cache() + + total_T = sum(seq_lens) + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) + inputs = prepare_safe_gate_inputs(1, total_T, H, D, device, cu_seqlens=cu_seqlens, seed=SEED) + q, k, v, g, beta = inputs["q"], inputs["k"], inputs["v"], inputs["g"], inputs["beta"] + A_log, dt_bias = inputs["A_log"], inputs["dt_bias"] + scale, lower_bound = inputs["scale"], inputs["lower_bound"] + + common = dict( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + A_log=A_log, + dt_bias=dt_bias, + cu_seqlens=cu_seqlens, + lower_bound=lower_bound, + ) + + ms_off = time_kernel(lambda: run_kernel(**common, use_cp=False)) + ms_on = time_kernel(lambda: run_kernel(**common, use_cp=True)) + + speedup = ms_off / ms_on if ms_on > 0 else float("inf") + r = dict(tag=tag, H=H, total_T=total_T, ms_off=ms_off, ms_on=ms_on, speedup=speedup) + results.append(r) + + del q, k, v, g, beta, A_log, dt_bias, inputs + torch.cuda.empty_cache() + + return results + + +def print_report(results, h_values): + sep = "=" * 95 + print(f"\n\n{sep}") + print(" BENCHMARK REPORT: SM90 Intracard CP") + print(" CP-on vs CP-off (intracard_prefill vs flash_kda_fwd)") + print(f" D={D} dtype=bf16 safe_gate=True") + wu = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP + ni = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS + mode_tag = " [NCU mode]" if NCU_MODE else (" [Sanitizer mode]" if SANITIZER_MODE else "") + print(f" Warmup={wu} Iters={ni}{mode_tag}") + print(sep) + + for H_val in h_values: + h_results = [r for r in results if r["H"] == H_val] + if not h_results: + continue + print(f"\n [H={H_val}]") + print(f" {'─' * 80}") + print(f" {'config':<20s} {'T':>7s} │ {'CP_off(ms)':>10s} {'CP_on(ms)':>10s} {'Speedup':>8s}") + print(f" {'─' * 80}") + for r in h_results: + print(f" {r['tag']:<20s} {r['total_T']:>7d} │ {r['ms_off']:>10.4f} {r['ms_on']:>10.4f} {r['speedup']:>7.2f}x") + print(f" {'─' * 80}") + + speedups = [r["speedup"] for r in results] + if speedups: + geo = 1.0 + for s in speedups: + geo *= s + geo = geo ** (1 / len(speedups)) + print(f"\n All configs: geo-mean={geo:.2f}x best={max(speedups):.2f}x worst={min(speedups):.2f}x") + + print(f"\n{sep}\n") + + +def main(): + parser = argparse.ArgumentParser(description="bench_kda_sm90_cp: SM90 intracard CP speedup") + parser.add_argument("--ncu", action="store_true", help="NCU profiling mode: warmup=1, iters=1") + parser.add_argument("--sanitizer", action="store_true", help="Sanitizer mode: warmup=1, iters=1") + args = parser.parse_args() + + global NCU_MODE, SANITIZER_MODE + if args.ncu: + NCU_MODE = True + print("[NCU mode] warmup=1, iters=1") + if args.sanitizer: + SANITIZER_MODE = True + print("[Sanitizer mode] warmup=1, iters=1") + + results = bench_cp(H_VALUES, CONFIGS) + print_report(results, H_VALUES) + return results + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_kda_sm90_prefill.py b/benchmarks/bench_kda_sm90_prefill.py new file mode 100644 index 00000000..fcb9add5 --- /dev/null +++ b/benchmarks/bench_kda_sm90_prefill.py @@ -0,0 +1,324 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +bench_kda_sm90_prefill.py — Benchmark: SM90 CuTeDSL KDA prefill vs FLA Triton baseline + +Hopper-only. Calls cula.kda.flashkda (cula_kda_prefill, the K1+K2 +two-kernel prefill) directly -- no get_kda_fused_fwd dispatcher. + +Compares: + - Accuracy: relative_rms_error, rel_max, mean_diff of the output o vs FLA Triton + - Performance: kernel execution time (ms) via CUDA events + +Modes: + - Fixed-length: various (B, T) configs + - Varlen: packed sequences with 2-3x length variation + +The SM90 prefill is MHA only (HV == H); grouped-value attention is not yet +supported on this backend, so there is no --hv option. + +Usage: + python bench_kda_sm90_prefill.py [--mode fixed|varlen|both] [--heads H] [--init_state] [--ncu] + +With --ncu, warmup=1 and iters=1 for ncu profiling: + ncu --set full -o report python bench_kda_sm90_prefill.py --mode varlen --ncu +""" + +import argparse +import os +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) +os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) # fair comparison with FLA + +import torch +from fla.ops.kda import chunk_kda as fla_chunk_kda + +from benchmarks.utils import ( + SEED, + benchmark_cuda_mode_fn, + build_varlen_configs, + exclusive_cumsum, + prepare_safe_gate_inputs, + relative_rms_error_rel_max_mean_abs, + set_seed, +) +from cula.kda import flashkda_prefill as cula_kda_prefill +from cula.utils import assert_hopper, get_device_sm_version + +_device = torch.device("cuda") +assert_hopper(_device) +_major, _minor = get_device_sm_version(_device) +_SM_TAG = f"sm{_major}{_minor}" + +H, D = 64, 128 +WARMUP = 25 +N_ITERS = 100 +NCU_MODE = False +SANITIZER_MODE = False +HAS_INIT_STATE = False + + +def run_fla(q, k, v, g, beta, scale, A_log, dt_bias, init_state, cu_seqlens, lower_bound): + return fla_chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + A_log=A_log, + dt_bias=dt_bias, + initial_state=init_state, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu_seqlens, + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=lower_bound, + transpose_state_layout=True, # match the SM90 VK-transposed state layout + ) + + +def run_cula(q, k, v, g, beta, scale, A_log, dt_bias, init_state, cu_seqlens, lower_bound): + return cula_kda_prefill( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + A_log=A_log, + dt_bias=dt_bias, + initial_state=init_state, + output_final_state=True, + cu_seqlens=cu_seqlens, + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=lower_bound, + ) + + +def _bench_one(common): + """Accuracy (on o) + perf (ms) for a single config.""" + o_fla, _ = run_fla(**common) + o_cula, _ = run_cula(**common) + torch.cuda.synchronize() + rel_rmse, rel_max, mean_diff = relative_rms_error_rel_max_mean_abs(o_fla, o_cula) + # iqr_mean: a stray transient iteration must not poison the reported mean. + ms_fla = benchmark_cuda_mode_fn( + lambda: run_fla(**common), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + sanitizer_mode=SANITIZER_MODE, + aggregate="iqr_mean", + ) + ms_cula = benchmark_cuda_mode_fn( + lambda: run_cula(**common), + default_warmup=WARMUP, + default_rep=N_ITERS, + ncu_mode=NCU_MODE, + sanitizer_mode=SANITIZER_MODE, + aggregate="iqr_mean", + ) + speedup = ms_fla / ms_cula if ms_cula > 0 else float("inf") + del o_fla, o_cula + return rel_rmse, rel_max, mean_diff, ms_fla, ms_cula, speedup + + +def _make_common(B, T, cu_seqlens, device): + inputs = prepare_safe_gate_inputs(B, T, H, D, device, cu_seqlens=cu_seqlens, has_init_state=HAS_INIT_STATE, num_v_heads=H) + return dict( + q=inputs["q"], + k=inputs["k"], + v=inputs["v"], + g=inputs["g"], + beta=inputs["beta"], + scale=inputs["scale"], + A_log=inputs["A_log"], + dt_bias=inputs["dt_bias"], + init_state=inputs["init_state"], + cu_seqlens=cu_seqlens, + lower_bound=inputs["lower_bound"], + ) + + +def bench_fixed(configs): + print("\n" + "=" * 100) + print(f" Fixed-Length Benchmark: SM90 cula_kda_prefill ({_SM_TAG}) vs FLA Triton") + print("=" * 100) + results = [] + for B, T in configs: + set_seed(SEED) + device = torch.device("cuda") + torch.cuda.empty_cache() + cu_seqlens = torch.tensor(exclusive_cumsum([T] * B), dtype=torch.int32, device=device) + common = _make_common(B, T, cu_seqlens, device) + rel_rmse, rel_max, mean_diff, ms_fla, ms_cula, speedup = _bench_one(common) + results.append( + { + "B": B, + "T": T, + "relative_rms_error": rel_rmse, + "rel_max": rel_max, + "mean_diff": mean_diff, + "ms_fla": ms_fla, + "ms_cula": ms_cula, + "speedup": speedup, + } + ) + del common + torch.cuda.empty_cache() + return results + + +def bench_varlen(configs): + print("\n" + "=" * 100) + print(f" Varlen Benchmark: SM90 cula_kda_prefill ({_SM_TAG}) vs FLA Triton") + print("=" * 100) + results = [] + for seq_lens, total_len, dist in configs: + set_seed(SEED) + device = torch.device("cuda") + torch.cuda.empty_cache() + T = total_len + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) + common = _make_common(1, T, cu_seqlens, device) + rel_rmse, rel_max, mean_diff, ms_fla, ms_cula, speedup = _bench_one(common) + n_seqs = len(seq_lens) + tag = f"{dist:>7s} {n_seqs:>2d}seqs T={T} [{min(seq_lens)}..{max(seq_lens)}] avg={T // n_seqs}" + results.append( + { + "tag": tag, + "relative_rms_error": rel_rmse, + "rel_max": rel_max, + "mean_diff": mean_diff, + "ms_fla": ms_fla, + "ms_cula": ms_cula, + "speedup": speedup, + } + ) + del common + torch.cuda.empty_cache() + return results + + +def print_report(fixed_results, varlen_results): + sep = "=" * 110 + print(f"\n\n{sep}") + print(" BENCHMARK REPORT: SM90 cula_kda_prefill (K1+K2 two-kernel)") + print(f" cuLA {_SM_TAG} vs FLA Triton") + print(f" H={H} D={D} dtype=bf16 safe_gate=True has_init_state={HAS_INIT_STATE}") + wu = 1 if (NCU_MODE or SANITIZER_MODE) else WARMUP + ni = 1 if (NCU_MODE or SANITIZER_MODE) else N_ITERS + mode_tag = " [NCU mode]" if NCU_MODE else (" [Sanitizer mode]" if SANITIZER_MODE else "") + print(f" Warmup={wu} Iters={ni}{mode_tag}") + print(sep) + + if fixed_results: + print("\n [Fixed-Length]") + print(f" {'─' * 96}") + print( + f" {'B':>3s} {'T':>6s} │ {'rel_rmse':>10s} {'rel_max':>10s} {'mean_diff':>10s} │ " + f"{'FLA(ms)':>9s} {'cuLA(ms)':>10s} {'Speedup':>8s}" + ) + print(f" {'─' * 96}") + for r in fixed_results: + print( + f" {r['B']:3d} {r['T']:6d} │ " + f"{r['relative_rms_error']:10.6f} {r['rel_max']:10.6f} {r['mean_diff']:10.6f} │ " + f"{r['ms_fla']:9.4f} {r['ms_cula']:10.4f} {r['speedup']:7.2f}x" + ) + print(f" {'─' * 96}") + + if varlen_results: + print("\n [Varlen]") + print(f" {'─' * 106}") + print( + f" {'Config':>45s} │ {'rel_rmse':>10s} {'rel_max':>10s} {'mean_diff':>10s} │ " + f"{'FLA(ms)':>9s} {'cuLA(ms)':>10s} {'Speedup':>8s}" + ) + print(f" {'─' * 106}") + for r in varlen_results: + print( + f" {r['tag']:>45s} │ " + f"{r['relative_rms_error']:10.6f} {r['rel_max']:10.6f} {r['mean_diff']:10.6f} │ " + f"{r['ms_fla']:9.4f} {r['ms_cula']:10.4f} {r['speedup']:7.2f}x" + ) + print(f" {'─' * 106}") + + print(f"\n{sep}\n") + + +def main(): + parser = argparse.ArgumentParser(description="bench_kda_sm90_prefill: SM90 cula_kda_prefill vs FLA Triton") + parser.add_argument( + "--mode", + type=str, + default="both", + choices=["fixed", "varlen", "both"], + help="Which benchmark mode to run (default: both).", + ) + parser.add_argument("--ncu", action="store_true", help="NCU profiling mode: warmup=1, iters=1") + parser.add_argument("--sanitizer", action="store_true", help="Sanitizer mode: warmup=1, iters=1") + parser.add_argument("--init_state", action="store_true", help="Use non-zero initial state (default: False)") + global H + parser.add_argument("--heads", type=int, default=H, help=f"Number of heads (H == HV, MHA). Default: {H}") + args = parser.parse_args() + + global NCU_MODE, SANITIZER_MODE, HAS_INIT_STATE + H = args.heads + if args.ncu: + NCU_MODE = True + print("[NCU mode] warmup=1, iters=1") + if args.sanitizer: + SANITIZER_MODE = True + print("[Sanitizer mode] warmup=1, iters=1") + if args.init_state: + HAS_INIT_STATE = True + print("[init_state] using non-zero initial state") + + print(f"[Device] {torch.cuda.get_device_name(0)} compute capability {_SM_TAG} → cula_kda_prefill (SM90 K1+K2)") + + fixed_configs = [ + (1, 512), + (1, 1024), + (1, 4096), + (1, 8192), + (1, 16384), + (2, 512), + (2, 1024), + (2, 4096), + (2, 8192), + (2, 16384), + ] + varlen_configs = build_varlen_configs( + num_seqs_list=(10, 20), total_lens=(4096, 8192, 16384), dists=("uniform", "random", "skewed") + ) + + fixed_res, varlen_res = [], [] + if args.mode in ("fixed", "both"): + fixed_res = bench_fixed(fixed_configs) + if args.mode in ("varlen", "both"): + varlen_res = bench_varlen(varlen_configs) + print_report(fixed_res, varlen_res) + return fixed_res, varlen_res + + +if __name__ == "__main__": + main() diff --git a/benchmarks/generate_benchmark_hopper_md.py b/benchmarks/generate_benchmark_hopper_md.py index d05dc0df..cb9e8e08 100644 --- a/benchmarks/generate_benchmark_hopper_md.py +++ b/benchmarks/generate_benchmark_hopper_md.py @@ -3,10 +3,9 @@ generate_benchmark_hopper_md.py — Run Hopper (SM90) benchmarks and generate BENCHMARK_hopper.md Currently supported Hopper benchmarks: - - KDA Fused Forward (cula.kda.hopper_fused_fwd) + - KDA prefill (cula.kda.flashkda — SM90 K1+K2 two-kernel) -Reuses bench_kda_fused_fwd.py which auto-selects the correct cuLA implementation -based on GPU architecture (SM90 → hopper_fused_fwd). +Reuses bench_kda_sm90_prefill.py (calls cula_kda_prefill directly). Usage: python benchmarks/generate_benchmark_hopper_md.py @@ -33,17 +32,17 @@ os.environ.setdefault("FLA_USE_FAST_OPS", os.getenv("CULA_USE_FAST_MATH", "1")) # Enable fast ops in FLA for fair comparison -from benchmarks.bench_kda_fused_fwd import ( # noqa: E402 +from benchmarks.bench_kda_sm90_prefill import ( # noqa: E402 _SM_TAG, ) -from benchmarks.bench_kda_fused_fwd import ( # noqa: E402 - D as KDA_FUSED_D, +from benchmarks.bench_kda_sm90_prefill import ( # noqa: E402 + D as KDA_D, ) -from benchmarks.bench_kda_fused_fwd import ( # noqa: E402 - H as KDA_FUSED_H, +from benchmarks.bench_kda_sm90_prefill import ( # noqa: E402 + H as KDA_H, ) -from benchmarks.bench_kda_fused_fwd import ( # noqa: E402 - main as kda_fused_fwd_main, +from benchmarks.bench_kda_sm90_prefill import ( # noqa: E402 + main as kda_prefill_main, ) from benchmarks.utils import get_env_info # noqa: E402 @@ -55,20 +54,18 @@ # ============================================================ -def run_kda_fused_fwd_benchmarks(has_init_state: bool = False, heads=None, hv=None): - """Run bench_kda_fused_fwd.main() with programmatic args and return (fixed, varlen) results.""" - print("\n>>> Running KDA Fused Forward benchmarks (via bench_kda_fused_fwd.main)...") +def run_kda_prefill_benchmarks(has_init_state: bool = False, heads=None): + """Run bench_kda_sm90_prefill.main() with programmatic args and return (fixed, varlen) results.""" + print("\n>>> Running SM90 KDA prefill benchmarks (via bench_kda_sm90_prefill.main)...") orig_argv = sys.argv - argv = ["bench_kda_fused_fwd.py", "--mode", "both"] + argv = ["bench_kda_sm90_prefill.py", "--mode", "both"] if has_init_state: argv.append("--init_state") if heads is not None: argv += ["--heads", str(heads)] - if hv is not None: - argv += ["--hv", str(hv)] sys.argv = argv try: - fixed_res, varlen_res = kda_fused_fwd_main() + fixed_res, varlen_res = kda_prefill_main() finally: sys.argv = orig_argv return fixed_res, varlen_res @@ -79,7 +76,7 @@ def run_kda_fused_fwd_benchmarks(has_init_state: bool = False, heads=None, hv=No # ============================================================ -def format_benchmark_md(env, kda_fused_fixed, kda_fused_varlen, has_init_state: bool = False): +def format_benchmark_md(env, kda_fixed, kda_varlen, has_init_state: bool = False): lines = [] w = lines.append @@ -92,32 +89,32 @@ def format_benchmark_md(env, kda_fused_fixed, kda_fused_varlen, has_init_state: w("") # ------------------------------------------------------------------- - # KDA Fused Forward + # KDA prefill # ------------------------------------------------------------------- - w("\n## KDA Fused Forward (Kimi Delta Attention)\n") - w(f"Fully-fused KDA forward prefill kernel ({_SM_TAG}).\n") + w("\n## KDA Prefill (Kimi Delta Attention)\n") + w(f"SM90 K1+K2 two-kernel prefill ({_SM_TAG}).\n") # Fixed-length - if kda_fused_fixed: - w(f"### Fixed-Length (H={KDA_FUSED_H}, D={KDA_FUSED_D}, bf16)\n") - w("| B | T | FLA Triton (ms) | cuLA Fused (ms) | Speedup |") - w("|---|---|-----------------|-----------------|---------|") - for r in kda_fused_fixed: + if kda_fixed: + w(f"### Fixed-Length (H={KDA_H}, D={KDA_D}, bf16)\n") + w("| B | T | FLA Triton (ms) | cuLA (ms) | Speedup |") + w("|---|---|-----------------|-----------|---------|") + for r in kda_fixed: sp = f"**{r['speedup']:.2f}x**" w(f"| {r['B']} | {r['T']} | {r['ms_fla']:.3f} | {r['ms_cula']:.3f} | {sp} |") # Varlen - if kda_fused_varlen: - w(f"\n### Variable-Length (H={KDA_FUSED_H}, D={KDA_FUSED_D}, bf16)\n") - w("| Config | FLA Triton (ms) | cuLA Fused (ms) | Speedup |") - w("|--------|-----------------|-----------------|---------|") - for r in kda_fused_varlen: + if kda_varlen: + w(f"\n### Variable-Length (H={KDA_H}, D={KDA_D}, bf16)\n") + w("| Config | FLA Triton (ms) | cuLA (ms) | Speedup |") + w("|--------|-----------------|-----------|---------|") + for r in kda_varlen: tag = r.get("tag", "unknown") sp = f"**{r['speedup']:.2f}x**" w(f"| {tag} | {r['ms_fla']:.3f} | {r['ms_cula']:.3f} | {sp} |") # Summary stats - all_results = (kda_fused_fixed or []) + (kda_fused_varlen or []) + all_results = (kda_fixed or []) + (kda_varlen or []) all_sp = [r["speedup"] for r in all_results if r.get("speedup", 0) > 0] if all_sp: w( @@ -128,7 +125,7 @@ def format_benchmark_md(env, kda_fused_fixed, kda_fused_varlen, has_init_state: w("To reproduce:\n") w("```bash") init_state_flag = " --init_state" if has_init_state else "" - w(f"python benchmarks/bench_kda_fused_fwd.py --mode both{init_state_flag}") + w(f"python benchmarks/bench_kda_sm90_prefill.py --mode both{init_state_flag}") w("```\n") return "\n".join(lines) @@ -160,13 +157,7 @@ def main(): "--heads", type=int, default=None, - help="Number of Q/K heads (H) for KDA benchmarks. Default: use bench_kda_fused_fwd default.", - ) - parser.add_argument( - "--hv", - type=int, - default=None, - help="Number of V heads (HV) for KDA benchmarks. For GVA, set HV > H with HV %% H == 0.", + help="Number of heads (H == HV, MHA) for KDA benchmarks. Default: use bench_kda_sm90_prefill default.", ) args = parser.parse_args() @@ -176,12 +167,10 @@ def main(): print(f"Loading cached results from {args.cache}") with open(args.cache) as f: data = json.load(f) - kda_fused_fixed = data["kda_fused_fixed"] - kda_fused_varlen = data["kda_fused_varlen"] + kda_fixed = data["kda_fixed"] + kda_varlen = data["kda_varlen"] else: - kda_fused_fixed, kda_fused_varlen = run_kda_fused_fwd_benchmarks( - has_init_state=args.init_state, heads=args.heads, hv=args.hv - ) + kda_fixed, kda_varlen = run_kda_prefill_benchmarks(has_init_state=args.init_state, heads=args.heads) if args.save_cache: cache_path = Path(args.save_cache) @@ -204,15 +193,15 @@ def sanitize(results): with open(cache_path, "w") as f: json.dump( { - "kda_fused_fixed": sanitize(kda_fused_fixed), - "kda_fused_varlen": sanitize(kda_fused_varlen), + "kda_fixed": sanitize(kda_fixed), + "kda_varlen": sanitize(kda_varlen), }, f, indent=2, ) print(f"Cached results to {cache_path}") - md = format_benchmark_md(env, kda_fused_fixed, kda_fused_varlen, has_init_state=args.init_state) + md = format_benchmark_md(env, kda_fixed, kda_varlen, has_init_state=args.init_state) output_path = ROOT / (args.output if args.output else BENCHMARK_MD_DEFAULT) output_path.write_text(md) diff --git a/benchmarks/utils.py b/benchmarks/utils.py index 75786e1f..164d3a30 100644 --- a/benchmarks/utils.py +++ b/benchmarks/utils.py @@ -115,15 +115,21 @@ def benchmark_cuda_mode_fn( ncu_mode=False, sanitizer_mode=False, setup_fn=None, + aggregate="mean", ): - """Benchmark a CUDA callable using standard repo warmup/repeat mode rules.""" + """Benchmark a CUDA callable using standard repo warmup/repeat mode rules. + + ``aggregate`` controls how per-iteration times are reduced: ``"mean"`` (default) + or ``"iqr_mean"`` (interquartile mean — robust to a stray transient iteration + that would otherwise poison the mean). + """ warmup, rep = resolve_benchmark_repeats( default_warmup, default_rep, ncu_mode=ncu_mode, sanitizer_mode=sanitizer_mode, ) - return benchmark_cuda_fn(fn, setup_fn=setup_fn, warmup=warmup, rep=rep, aggregate="mean") + return benchmark_cuda_fn(fn, setup_fn=setup_fn, warmup=warmup, rep=rep, aggregate=aggregate) def triton_bench_fn(fn, **kwargs): diff --git a/cula/kda/README.md b/cula/kda/README.md new file mode 100644 index 00000000..4e7def69 --- /dev/null +++ b/cula/kda/README.md @@ -0,0 +1,91 @@ +# `cula.kda` — KDA (Kimi Delta Attention) operators + +## API + +| Symbol | Description | Arch | Direction | +|--------|-------------|------|-----------| +| `chunk_kda` | Chunked KDA prefill | SM100 (Blackwell) | Forward + backward | +| `kda_prefill_hopper` / `_opt` / `_auto` | KDA prefill (CUDA C++ fused) | SM90 (Hopper) | Forward only | +| `flashkda_prefill` | KDA prefill (CuTeDSL K1+K2, intracard-CP capable) | SM90 (Hopper) | Forward only | +| `kda_decode` | Single-token decode | SM100 | Forward | +| `fused_sigmoid_gating_delta_rule_update` | Decode state update | SM100 | Forward | + +## Quick start + +### Prefill (Hopper) + +```python +import torch +from cula.kda import flashkda_prefill + +B, T, H, D = 1, 1024, 4, 128 +q = torch.randn(B, T, H, D, dtype=torch.bfloat16, device="cuda") +k = torch.randn(B, T, H, D, dtype=torch.bfloat16, device="cuda") +v = torch.randn(B, T, H, D, dtype=torch.bfloat16, device="cuda") +g = torch.randn(B, T, H, D, dtype=torch.bfloat16, device="cuda") +beta = torch.randn(B, T, H, dtype=torch.bfloat16, device="cuda") +A_log = torch.randn(H, dtype=torch.float32, device="cuda") +dt_bias = torch.randn(H * D, dtype=torch.float32, device="cuda") + +o, ht = flashkda_prefill( + q, k, v, g, beta, + A_log=A_log, dt_bias=dt_bias, + scale=D**-0.5, lower_bound=-5.0, + safe_gate=True, use_gate_in_kernel=True, + output_final_state=True, +) +``` + +### Prefill (Blackwell, with backward) + +```python +from cula.kda import chunk_kda + +o, ht = chunk_kda( + q, k, v, g, beta, + A_log=A_log, dt_bias=dt_bias, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + safe_gate=True, lower_bound=-5.0, + output_final_state=True, +) +``` + +### Variable-length (packed) + +```python +cu_seqlens = torch.tensor([0, 256, 500, 1000], dtype=torch.int32, device="cuda") +q = torch.randn(1, 1000, H, D, dtype=torch.bfloat16, device="cuda") +# ... k, v, g, beta shaped [1, 1000, H, D] / [1, 1000, H] + +o, ht = flashkda_prefill( + q, k, v, g, beta, + A_log=A_log, dt_bias=dt_bias, + scale=D**-0.5, lower_bound=-5.0, + safe_gate=True, use_gate_in_kernel=True, + cu_seqlens=cu_seqlens, + output_final_state=True, +) +``` + +### Intracard context-parallel + +```python +# "auto": use CP only when beneficial for the given sequence lengths +o, ht = flashkda_prefill( + q, k, v, g, beta, + A_log=A_log, dt_bias=dt_bias, + scale=D**-0.5, lower_bound=-5.0, + safe_gate=True, use_gate_in_kernel=True, + cu_seqlens=cu_seqlens, + use_intracard_cp="auto", +) +``` + +## Requirements + +- **D = 128** (head dimension, currently the only supported value) +- **bf16** for q/k/v/g/beta, **fp32** for A_log/dt_bias +- All tensors must be CUDA and contiguous +- `safe_gate=True` + `lower_bound` in `[-5, 0)` required for Hopper prefill +- `use_gate_in_kernel=True` required for Hopper prefill diff --git a/cula/kda/__init__.py b/cula/kda/__init__.py index 73f33219..61ed1cb8 100644 --- a/cula/kda/__init__.py +++ b/cula/kda/__init__.py @@ -25,6 +25,7 @@ "kda_prefill_hopper", "kda_prefill_hopper_opt", "kda_prefill_hopper_auto", + "flashkda_prefill", ] _LAZY = { @@ -32,6 +33,7 @@ "kda_prefill_hopper": ("cula.kda.hopper_fused_fwd", "cula_kda_prefill"), "kda_prefill_hopper_opt": ("cula.kda.hopper_fused_fwd_opt", "cula_kda_prefill_opt"), "kda_prefill_hopper_auto": ("cula.kda.auto_route", "cula_kda_prefill_auto"), + "flashkda_prefill": ("cula.kda.flashkda", "cula_kda_prefill"), "kda_decode": ("cula.ops.kda.decode.cute", "kda_decode"), "kda_decode_mtp": ("cula.ops.kda.decode.mtp", "kda_decode_mtp"), "kda_decode_mtp_recurrent": ("cula.ops.kda.decode.mtp", "kda_decode_mtp_recurrent"), diff --git a/cula/kda/chunk.py b/cula/kda/chunk.py index 2c9daad6..d07ba755 100644 --- a/cula/kda/chunk.py +++ b/cula/kda/chunk.py @@ -27,7 +27,7 @@ from cula.kda.chunk_bwd import chunk_kda_bwd from cula.kda.chunk_fwd import chunk_kda_fwd -from cula.ops.kda.policy import IntracardCPMode, resolve_intracard_cp_mode +from cula.ops.kda.cp_mode import CPMode class ChunkKDAFunction(torch.autograd.Function): @@ -55,7 +55,7 @@ def forward( disable_recompute: bool = False, return_intermediate_states: bool = False, cp_context: FLACPContext | None = None, - use_intracard_cp: IntracardCPMode | None = None, + use_intracard_cp: CPMode | None = None, ): chunk_size = 64 @@ -357,12 +357,12 @@ def chunk_kda( # just for backward compatibility, resolve the deprecated `use_cp` argument # TODO: maybe we can remove this in the future use_cp_alias = kwargs.pop("use_cp", None) - use_intracard_cp = resolve_intracard_cp_mode(use_intracard_cp, use_cp_alias) + use_intracard_cp = CPMode.parse(use_intracard_cp, use_cp_alias) if cp_context is not None: - if use_intracard_cp is True: + if use_intracard_cp is CPMode.FORCE: raise ValueError("use_intracard_cp=True cannot be combined with FLA cp_context.") - use_intracard_cp = False + use_intracard_cp = CPMode.OFF assert initial_state is None, "Initial state is not supported for CP" assert output_final_state is False, "Output final state is not supported for CP" assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP" @@ -372,9 +372,9 @@ def chunk_kda( cu_seqlens_cpu = cp_context.cu_seqlens_cpu if return_intermediate_states: - if use_intracard_cp is True: + if use_intracard_cp is CPMode.FORCE: raise ValueError("use_intracard_cp=True is not supported with return_intermediate_states=True.") - use_intracard_cp = False + use_intracard_cp = CPMode.OFF if cu_seqlens is not None: if q.shape[0] != 1: diff --git a/cula/kda/flashkda.py b/cula/kda/flashkda.py new file mode 100644 index 00000000..b644f312 --- /dev/null +++ b/cula/kda/flashkda.py @@ -0,0 +1,391 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""SM90 KDA prefill wrapper for the two-kernel K1+K2 CuTeDSL path""" + +import contextlib +from typing import Literal + +import torch +from torch.amp import custom_bwd, custom_fwd + +from cula.ops.kda.cp_mode import CPMode +from cula.ops.kda.sm90.cp.plan import plan_prefill +from cula.ops.kda.sm90.fwd import flash_kda_fwd +from cula.utils import assert_hopper + + +def _cast_g_bf16(g: torch.Tensor) -> torch.Tensor: + return g if g.dtype == torch.bfloat16 else g.to(torch.bfloat16) + + +def _beta_logits_bf16(beta: torch.Tensor) -> torch.Tensor: + return torch.logit(beta.float(), eps=1e-6).to(torch.bfloat16) + + +def _cast_beta_bf16(beta: torch.Tensor) -> torch.Tensor: + return beta if beta.dtype == torch.bfloat16 else beta.to(torch.bfloat16) + + +def _contig(t: torch.Tensor | None) -> torch.Tensor | None: + if t is None or t.is_contiguous(): + return t + return t.contiguous() + + +def _guarded_forward( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + scale: float, + initial_state: torch.Tensor, + output_final_state: bool = False, + lower_bound: float | None = None, + cu_seqlens: torch.IntTensor | None = None, + cu_seqlens_cpu: torch.IntTensor | None = None, + use_intracard_cp: CPMode | None = None, + beta_is_logits: bool = False, + out: torch.Tensor | None = None, + final_state: torch.Tensor | None = None, +): + q, k, v, g, beta = _contig(q), _contig(k), _contig(v), _contig(g), _contig(beta) + A_log, dt_bias = _contig(A_log), _contig(dt_bias) + initial_state, cu_seqlens = _contig(initial_state), _contig(cu_seqlens) + device_ctx = ( + torch.cuda.device(q.device.index) if q.device.index != torch.cuda.current_device() else contextlib.nullcontext() + ) + with device_ctx: + return HopperChunkKDAFunction._forward_impl( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + scale, + initial_state, + output_final_state, + lower_bound, + cu_seqlens, + cu_seqlens_cpu, + use_intracard_cp, + beta_is_logits, + out, + final_state, + ) + + +class HopperChunkKDAFunction(torch.autograd.Function): + @staticmethod + @custom_fwd(device_type="cuda") + def forward(ctx, *args): + return _guarded_forward(*args) + + @staticmethod + def _forward_impl( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + scale, + initial_state, + output_final_state, + lower_bound, + cu_seqlens, + cu_seqlens_cpu, + use_intracard_cp, + beta_is_logits, + out, + final_state, + ): + batch_size, seq_len, num_heads, head_dim = q.shape + + if out is None: + out = torch.empty_like(v) + + n_seqs = batch_size + if cu_seqlens is not None: + n_seqs = cu_seqlens.numel() - 1 + + if not output_final_state: + final_state = None + elif final_state is None: + final_state = torch.empty( + n_seqs, + num_heads, + head_dim, + head_dim, + dtype=torch.float32, + device=q.device, + ) + + g = _cast_g_bf16(g) + if beta_is_logits: + beta = _cast_beta_bf16(beta) + else: + beta = _beta_logits_bf16(beta) + + plan = plan_prefill(q, cu_seqlens, cu_seqlens_cpu, use_intracard_cp) + if not plan.trivial: + from cula.ops.kda.sm90.cp.driver import run_cp + + run_cp( + plan, + q, + k, + v, + g, + beta, + scale=scale, + out=out, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=initial_state, + final_state=final_state, + cu_seqlens=cu_seqlens, + state_transposed=False, + ) + else: + flash_kda_fwd( + q, + k, + v, + g, + beta, + scale=scale, + out=out, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=initial_state, + final_state=final_state, + cu_seqlens=cu_seqlens, + cu_seqlens_cpu=cu_seqlens_cpu, + state_transposed=False, + use_gate_in_kernel=True, + ) + + return out.to(q.dtype), final_state + + @staticmethod + @custom_bwd(device_type="cuda") + def backward(ctx, do, dht): + raise NotImplementedError("Backward pass is not implemented yet.") + + +@torch.compiler.disable +def cula_kda_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float = None, + initial_state: torch.Tensor = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + use_beta_sigmoid_in_kernel: bool = False, + use_gate_in_kernel: bool = True, + safe_gate: bool = False, + lower_bound: float | None = None, + cu_seqlens: torch.IntTensor | None = None, + chunk_indices: torch.IntTensor | None = None, + use_intracard_cp: Literal["auto"] | bool | None = None, + out: torch.Tensor | None = None, + final_state: torch.Tensor | None = None, + **kwargs, +): + r""" + Hopper (SM90) KDA forward prefill using CuTeDSL two-kernel pipeline. + + Gate preprocessing (A_log, dt_bias, lower_bound) and L2-norm are handled + internally by the K1 kernel. This SM90 CuTeDSL path supports only the safe + in-kernel gate mode: ``use_gate_in_kernel=True`` and ``safe_gate=True``. + ``use_qk_l2norm_in_kernel`` is accepted for API compatibility; CuTeDSL + always applies L2-norm internally. + + Args: + q (torch.Tensor): + queries of shape `[B, T, H, K]`. + k (torch.Tensor): + keys of shape `[B, T, H, K]`. + v (torch.Tensor): + values of shape `[B, T, H, V]`. + g (torch.Tensor): + (forget) gating tensor (in log space!) of shape `[B, T, H, K]`. + beta (torch.Tensor): + betas of shape `[B, T, H]`. + scale (Optional[float]): + Scale factor for the KDA attention scores. + If not provided, it will default to `1 / sqrt(K)`. Default: `None`. + initial_state (Optional[torch.Tensor]): + Initial state of shape `[N, H, K, V]` for `N` input sequences. + Default: `None`. + output_final_state (Optional[bool]): + Whether to output the final state of shape `[N, H, K, V]`. Default: `False`. + use_qk_l2norm_in_kernel (bool): + Accepted for API compatibility; CuTeDSL always applies L2-norm + internally. Default: `False`. + use_beta_sigmoid_in_kernel (bool): + When `True`, `beta` is pre-sigmoid logits and is passed straight + to the kernel (which always applies sigmoid internally) — no + host-side logit round-trip. When `False` (FLA convention), + `beta` is post-sigmoid and is converted back to logits first. + Default: `False`. + use_gate_in_kernel (bool): + Must be `True`; CuTeDSL computes the gate internally. Default: `True`. + safe_gate (bool): + Must be `True`; unsupported unsafe gate inputs are rejected. Default: `False`. + lower_bound (Optional[float]): + Lower bound for the forget gate activation function. Required when + `safe_gate=True`; must be in `[-5, 0)`. Default: `None`. + cu_seqlens (torch.IntTensor): + Cumulative sequence lengths of shape `[N+1]`, int32. + cu_seqlens_cpu (Optional[torch.IntTensor]): + Optional CPU copy of `cu_seqlens` (same values), passed via kwargs, to + skip the GPU->host sync when first building varlen metadata. Trusted, + not verified (FLA convention). Default: `None`. + chunk_indices (torch.IntTensor): + Accepted for API compatibility; unused by CuTeDSL. + use_intracard_cp (Literal["auto"] | bool): + Whether to use the SM90 intracard-CP path when profitable. ``True`` + requires CP support and raises on rejection, ``"auto"`` falls back + to the serial K1+K2 path, and ``False`` disables CP. ``use_cp`` is + accepted as a compatibility alias. + out (Optional[torch.Tensor]): + Preallocated output buffer, bf16, same shape as ``v``, written in + place (also returned). ``None`` allocates per call. Default: `None`. + final_state (Optional[torch.Tensor]): + Preallocated final-state buffer, fp32, shape `[N, H, K, V]`, + written in place (also returned). Requires + ``output_final_state=True``. ``None`` allocates per call. + Default: `None`. + + Returns: + o (torch.Tensor): + Outputs of shape `[B, T, H, V]`. + final_state (torch.Tensor): + Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`. + """ + assert_hopper() + if not use_gate_in_kernel: + raise NotImplementedError( + "SM90 CuTeDSL KDA prefill only supports use_gate_in_kernel=True. " + "Passing preprocessed gates would otherwise fall back to the slow reference path." + ) + if not safe_gate: + raise NotImplementedError("SM90 CuTeDSL KDA prefill only supports safe_gate=True.") + if lower_bound is None: + raise ValueError("`lower_bound` must be specified when `safe_gate=True` and `use_gate_in_kernel=True`.") + if not (-5 <= lower_bound < 0): + raise ValueError(f"`lower_bound` must be in the safe range [-5, 0), got {lower_bound}.") + if cu_seqlens is not None: + if q.shape[0] != 1: + raise ValueError( + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." + f"Please flatten variable-length inputs before processing.", + ) + if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.", + ) + if initial_state is not None: + if initial_state.dtype != torch.float32: + raise TypeError("initial_state must be in float32.") + + num_qk_heads, head_dim = q.shape[2], q.shape[3] + num_kv_heads = v.shape[2] + A_log = kwargs.pop("A_log", None) + dt_bias = kwargs.pop("dt_bias", None) + cu_seqlens_cpu = kwargs.pop("cu_seqlens_cpu", None) + use_cp_alias = kwargs.pop("use_cp", None) + use_intracard_cp = CPMode.parse(use_intracard_cp, use_cp_alias) + if kwargs: + raise TypeError(f"cula_kda_prefill got unexpected keyword arguments: {set(kwargs)}") + if A_log is None: + raise ValueError("A_log must be provided when use_gate_in_kernel=True.") + if dt_bias is None: + raise ValueError("dt_bias must be provided when use_gate_in_kernel=True.") + elif dt_bias.ndim == 1: + dt_bias = dt_bias.view(num_kv_heads, head_dim) + + if q.shape != k.shape: + raise ValueError(f"q and k must have the same shape, got q={tuple(q.shape)}, k={tuple(k.shape)}") + if g.shape != v.shape: + raise ValueError(f"g and v must have the same shape, got g={tuple(g.shape)}, v={tuple(v.shape)}") + if beta.shape != v.shape[:3]: + raise ValueError(f"beta must have shape {tuple(v.shape[:3])}, got {tuple(beta.shape)}") + if q.dtype != torch.bfloat16 or k.dtype != torch.bfloat16 or v.dtype != torch.bfloat16: + raise TypeError("q, k, v must be in bfloat16.") + if beta.dtype not in (torch.bfloat16, torch.float32): + raise TypeError("beta must be in bfloat16 or float32.") + if q.shape[-1] != 128 or k.shape[-1] != 128 or v.shape[-1] != 128: + raise ValueError("Currently we only support head dim of 128 for KDA.") + if num_kv_heads != num_qk_heads: + raise NotImplementedError( + "SM90 CuTeDSL KDA prefill does not support grouped-value attention yet " + f"(num_kv_heads={num_kv_heads} != num_qk_heads={num_qk_heads}); native GVA is a follow-up change." + ) + + if out is not None: + if out.shape != v.shape or out.dtype != torch.bfloat16 or not out.is_cuda or not out.is_contiguous(): + raise ValueError( + f"out must be a contiguous CUDA bfloat16 tensor of shape {tuple(v.shape)}, " + f"got dtype={out.dtype}, shape={tuple(out.shape)}" + ) + if final_state is not None: + if not output_final_state: + raise ValueError("final_state buffer requires output_final_state=True.") + n_seqs = cu_seqlens.numel() - 1 if cu_seqlens is not None else q.shape[0] + expected = (n_seqs, num_kv_heads, head_dim, head_dim) + if ( + final_state.shape != expected + or final_state.dtype != torch.float32 + or not final_state.is_cuda + or not final_state.is_contiguous() + ): + raise ValueError( + f"final_state must be a contiguous CUDA float32 tensor of shape {expected}, " + f"got dtype={final_state.dtype}, shape={tuple(final_state.shape)}" + ) + + if scale is None: + scale = k.shape[-1] ** -0.5 + fwd_args = ( + q, + k, + v, + g, + beta, + A_log, + dt_bias, + scale, + initial_state, + output_final_state, + lower_bound, + cu_seqlens, + cu_seqlens_cpu, + use_intracard_cp, + use_beta_sigmoid_in_kernel, + out, + final_state, + ) + # Forward-only op: skip autograd.Function unless a graph could be recorded. + needs_grad = torch.is_grad_enabled() and any( + t is not None and t.requires_grad for t in (q, k, v, g, beta, A_log, dt_bias, initial_state) + ) + if needs_grad: + o, final_state = HopperChunkKDAFunction.apply(*fwd_args) + else: + o, final_state = _guarded_forward(*fwd_args) + return o, final_state diff --git a/cula/ops/experimental/linear_attn_prototype.py b/cula/ops/experimental/linear_attn_prototype.py index 5cfa0248..4f4d300b 100644 --- a/cula/ops/experimental/linear_attn_prototype.py +++ b/cula/ops/experimental/linear_attn_prototype.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 ANTGROUP. All rights reserved. +# Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 # Redistribution and use in source and binary forms, with or without diff --git a/cula/ops/kda/__init__.py b/cula/ops/kda/__init__.py index a7d2cbd2..d04d8176 100644 --- a/cula/ops/kda/__init__.py +++ b/cula/ops/kda/__init__.py @@ -1,14 +1,14 @@ # Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 -"""KDA backend kernels migrated to the arch-first layout. +"""KDA backend kernels, organized by arch (sm90 / sm100). -sm100/ SM100 (Blackwell) modular-chunk recurrence/output/bwd kernels -decode/ single-token decode (CuTe DSL + FLA reference) -experimental/ unwired fully-fused WIP -policy.py CP dispatch policy (use_cp / use_intracard_cp) + sm90/ SM90 (Hopper) two-kernel (K1+K2) FlashKDA prefill, fwd-only (+ cp/) + sm100/ SM100 (Blackwell) modular-chunk recurrence/output/bwd kernels (+ cp/) + decode/ single-token decode (CuTe DSL + FLA reference) + experimental/ unwired fully-fused WIP + policy.py CP dispatch policy (use_cp / use_intracard_cp) +Both prefill backends are chunked forward computations; arch is the discriminator +(one implementation per arch), so there is no descriptive family layer. """ - -# TODO: The SM90 (Hopper) prefill is still the C++ kernel under csrc/kda/sm90 (CuTeDSL -# port pending); it is not yet part of this package. diff --git a/cula/ops/kda/cp_mode.py b/cula/ops/kda/cp_mode.py new file mode 100644 index 00000000..71106f9c --- /dev/null +++ b/cula/ops/kda/cp_mode.py @@ -0,0 +1,40 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Intracard context-parallel dispatch vocabulary shared by the KDA backends.""" + +from __future__ import annotations + +from enum import Enum + + +class NotSplittableError(ValueError): + """Intracard CP cannot split the given shape. Subclasses ValueError so + existing ``except ValueError`` callers keep working.""" + + +class CPMode(Enum): + """OFF = serial path; AUTO = engage when profitable; FORCE = engage + whenever splittable, else NotSplittableError.""" + + OFF = "off" + AUTO = "auto" + FORCE = "force" + + @classmethod + def parse(cls, use_intracard_cp, use_cp=None) -> CPMode | None: + """Public use_intracard_cp value ("auto"/True/False; use_cp is a + deprecated alias) -> mode. None stays None (backend default).""" + if use_intracard_cp is not None and use_cp is not None: + raise TypeError("Pass only one of use_intracard_cp or use_cp.") + value = use_intracard_cp if use_intracard_cp is not None else use_cp + if value is None or isinstance(value, cls): + return value + # Identity checks (not ==): `1 == True` would match stray ints. + if value is True: + return cls.FORCE + if value is False: + return cls.OFF + if value == "auto": + return cls.AUTO + raise ValueError(f'use_intracard_cp must be "auto", True, or False, got {value!r}') diff --git a/cula/ops/kda/experimental/sm100_fused/kda_fully_fused_wip.py b/cula/ops/kda/experimental/sm100_fused/kda_fully_fused_wip.py index ab09f0b2..8e5ad3d4 100644 --- a/cula/ops/kda/experimental/sm100_fused/kda_fully_fused_wip.py +++ b/cula/ops/kda/experimental/sm100_fused/kda_fully_fused_wip.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 ANTGROUP. All rights reserved. +# Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 # Redistribution and use in source and binary forms, with or without diff --git a/cula/ops/kda/experimental/sm100_fused/wrapper.py b/cula/ops/kda/experimental/sm100_fused/wrapper.py index deec324a..1896d2c4 100644 --- a/cula/ops/kda/experimental/sm100_fused/wrapper.py +++ b/cula/ops/kda/experimental/sm100_fused/wrapper.py @@ -269,7 +269,6 @@ def flash_kda_prefill( **kwargs, ): assert_blackwell() - assert cu_seqlens is None or q.shape[0] == 1, "For varlen, batch size must be 1. Flatten sequences first." if cu_seqlens is not None: if q.shape[0] != 1: raise ValueError( diff --git a/cula/ops/kda/policy.py b/cula/ops/kda/policy.py deleted file mode 100644 index bbd6e9b2..00000000 --- a/cula/ops/kda/policy.py +++ /dev/null @@ -1,98 +0,0 @@ -# Copyright 2025-2026 Ant Group Co., Ltd. -# SPDX-License-Identifier: Apache-2.0 - -"""Context-parallel dispatch policy for KDA wrappers.""" - -from __future__ import annotations - -import os -from collections.abc import Callable -from dataclasses import dataclass -from typing import Literal - -import torch - -IntracardCPMode = Literal["auto"] | bool - - -@dataclass(frozen=True) -class IntracardCPDecision: - enabled: bool - reason: str | None = None - force: bool = False - - -class NotSplittableError(ValueError): - """Raised when intracard CP cannot meaningfully split the given shape. - - Subclasses ValueError so existing ``except ValueError`` callers keep working, - while new code can catch it narrowly and fall back to the serial path. - """ - - -def normalize_intracard_cp_mode(mode: IntracardCPMode) -> IntracardCPMode: - # Identity checks (not `in`): `1 == True` / `0 == False` would match stray ints. - if mode != "auto" and mode is not True and mode is not False: - raise ValueError(f'use_intracard_cp must be "auto", True, or False, got {mode!r}') - return mode - - -def resolve_intracard_cp_mode( - use_intracard_cp: IntracardCPMode | None, - use_cp_alias: IntracardCPMode | None, -) -> IntracardCPMode | None: - if use_intracard_cp is not None and use_cp_alias is not None: - raise TypeError("Pass only one of use_intracard_cp or use_cp.") - mode = use_intracard_cp if use_intracard_cp is not None else use_cp_alias - if mode is None: - return None - return normalize_intracard_cp_mode(mode) - - -def _reject_or_disable(mode: IntracardCPMode, reason: str) -> IntracardCPDecision: - if mode is True: - raise ValueError(reason) - return IntracardCPDecision(False, reason) - - -def _sm100_env_cp_enabled() -> bool: - return os.environ.get("CULA_INTRACARD_CP", "0") != "0" - - -def sm100_intracard_cp_decision( - *, - mode: IntracardCPMode | None, - cu_seqlens: torch.Tensor | None, - cu_seqlens_cpu: torch.Tensor | None, - g: torch.Tensor | None, - num_qk_heads: int, - chunk_size: int, - is_inference: bool, - sm_count_provider: Callable[[], int], - no_cp: bool = False, -) -> IntracardCPDecision: - if mode is None: - mode = "auto" if _sm100_env_cp_enabled() else False - mode = normalize_intracard_cp_mode(mode) - # no_cp is the recursion guard: intracard_fwd_h re-invokes fwd_h with _no_cp=True - # so sub-sequences do not recursively re-trigger CP. - if mode is False or no_cp: - return IntracardCPDecision(False, "disabled") - - if cu_seqlens is None: - return _reject_or_disable(mode, "SM100 intracard CP requires varlen cu_seqlens.") - if g is not None: - return _reject_or_disable(mode, "SM100 intracard CP requires g is None; pass gate through gk.") - if not is_inference: - return _reject_or_disable(mode, "SM100 intracard CP is inference-only.") - - if mode is True: - return IntracardCPDecision(True, force=True) - - # auto: consult the CPU-only perf heuristic. - from cula.ops.kda.sm100.cp.chunk_delta_h import should_use_intracard_cp - - cpu = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens.cpu() - if should_use_intracard_cp(cpu, sm_count_provider(), num_qk_heads, chunk_size): - return IntracardCPDecision(True) - return IntracardCPDecision(False, "SM100 intracard CP heuristic declined for this shape.") diff --git a/cula/ops/kda/sm100/cp/chunk_delta_h.py b/cula/ops/kda/sm100/cp/chunk_delta_h.py index 74937cd0..23e355c3 100644 --- a/cula/ops/kda/sm100/cp/chunk_delta_h.py +++ b/cula/ops/kda/sm100/cp/chunk_delta_h.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 ANTGROUP. All rights reserved. +# Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 """ @@ -460,7 +460,7 @@ def intracard_fwd_h( split_info = False if not split_info: - from cula.ops.kda.policy import NotSplittableError + from cula.ops.kda.cp_mode import NotSplittableError raise NotSplittableError("SM100 intracard CP is not meaningfully splittable for this shape.") @@ -560,7 +560,7 @@ def intracard_fwd_h( save_new_value=save_new_value, cu_seqlens=cu_seqlens_subseq_gpu, chunk_indices=chunk_indices_subseq, - _no_cp=True, + use_intracard_cp=False, ) final_state = _gather_final_states( diff --git a/cula/ops/kda/sm100/cp/merge.py b/cula/ops/kda/sm100/cp/merge.py index 75d3fd12..e03fa7b6 100644 --- a/cula/ops/kda/sm100/cp/merge.py +++ b/cula/ops/kda/sm100/cp/merge.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 ANTGROUP. All rights reserved. +# Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 """ @@ -14,8 +14,6 @@ Output: h [num_non_first, H, K, V] fp32 """ -from __future__ import annotations - import functools import cuda.bindings.driver as cuda diff --git a/cula/ops/kda/sm100/cp/pre_scan.py b/cula/ops/kda/sm100/cp/pre_scan.py index befef74d..709086fa 100644 --- a/cula/ops/kda/sm100/cp/pre_scan.py +++ b/cula/ops/kda/sm100/cp/pre_scan.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 ANTGROUP. All rights reserved. +# Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 """ diff --git a/cula/ops/kda/sm100/delta_h.py b/cula/ops/kda/sm100/delta_h.py index c341bae9..dd0ca912 100644 --- a/cula/ops/kda/sm100/delta_h.py +++ b/cula/ops/kda/sm100/delta_h.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 ANTGROUP. All rights reserved. +# Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 """ @@ -35,7 +35,7 @@ from fla.ops.utils import prepare_chunk_indices, prepare_lens from fla.utils import tensor_cache -from cula.ops.kda.policy import sm100_intracard_cp_decision +from cula.ops.kda.sm100.policy import sm100_intracard_cp_decision from cula.utils import USE_FAST_MATH, assert_blackwell, get_device_sm_count COMPILE_OPTIONS = "--enable-tvm-ffi --generate-line-info --ptxas-options '--verbose'" @@ -2016,7 +2016,6 @@ def chunk_gated_delta_rule_fwd_h( cu_seqlens: torch.Tensor | None = None, chunk_indices: torch.Tensor | None = None, persistent: bool = True, - _no_cp: bool = False, cu_seqlens_cpu: torch.Tensor | None = None, use_intracard_cp=None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: @@ -2050,7 +2049,7 @@ def chunk_gated_delta_rule_fwd_h( v_new: [B, T, HV, V] bf16 (or None if save_new_value=False) final_state: [N, HV, K, V] fp32 (or None if output_final_state=False) """ - # --- Intracard CP dispatch (policy: cula.ops.kda.policy) --- + # --- Intracard CP dispatch (policy: cula.ops.kda.sm100.policy) --- cp_decision = sm100_intracard_cp_decision( mode=use_intracard_cp, cu_seqlens=cu_seqlens, @@ -2060,10 +2059,11 @@ def chunk_gated_delta_rule_fwd_h( chunk_size=chunk_size, is_inference=torch.is_inference_mode_enabled(), sm_count_provider=lambda: get_device_sm_count(k.device), - no_cp=_no_cp, ) + if cp_decision.cu_seqlens_cpu is not None: + cu_seqlens_cpu = cp_decision.cu_seqlens_cpu if cp_decision.enabled: - from cula.ops.kda.policy import NotSplittableError + from cula.ops.kda.cp_mode import NotSplittableError from cula.ops.kda.sm100.cp.chunk_delta_h import intracard_fwd_h try: diff --git a/cula/ops/kda/sm100/fwd_o.py b/cula/ops/kda/sm100/fwd_o.py index 9d4bcb46..0fd4bbd8 100644 --- a/cula/ops/kda/sm100/fwd_o.py +++ b/cula/ops/kda/sm100/fwd_o.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 ANTGROUP. All rights reserved. +# Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 """ diff --git a/cula/ops/kda/sm100/policy.py b/cula/ops/kda/sm100/policy.py new file mode 100644 index 00000000..3de2ea5f --- /dev/null +++ b/cula/ops/kda/sm100/policy.py @@ -0,0 +1,68 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Intracard-CP dispatch for the SM100 KDA backend.""" + +from __future__ import annotations + +import os +from collections.abc import Callable +from dataclasses import dataclass + +import torch + +from cula.ops.kda.cp_mode import CPMode, NotSplittableError + + +@dataclass(frozen=True) +class CPDecision: + enabled: bool + reason: str = "" + force: bool = False + # cu_seqlens materialized on host by the auto heuristic, returned so the + # caller reuses it instead of paying a second D2H sync. + cu_seqlens_cpu: torch.Tensor | None = None + + +def _env_default() -> CPMode: + return CPMode.AUTO if os.environ.get("CULA_INTRACARD_CP", "0") != "0" else CPMode.OFF + + +def sm100_intracard_cp_decision( + *, + mode: CPMode | bool | str | None, + cu_seqlens: torch.Tensor | None, + cu_seqlens_cpu: torch.Tensor | None, + g: torch.Tensor | None, + num_qk_heads: int, + chunk_size: int, + is_inference: bool, + sm_count_provider: Callable[[], int], +) -> CPDecision: + mode = CPMode.parse(mode) + if mode is None: + mode = _env_default() + if mode is CPMode.OFF: + return CPDecision(False, "disabled") + + def declined(reason: str) -> CPDecision: + if mode is CPMode.FORCE: + raise NotSplittableError(reason) + return CPDecision(False, reason) + + if cu_seqlens is None: + return declined("SM100 intracard CP requires varlen cu_seqlens.") + if g is not None: + return declined("SM100 intracard CP requires g is None; pass gate through gk.") + if not is_inference: + return declined("SM100 intracard CP is inference-only.") + + if mode is CPMode.FORCE: + return CPDecision(True, force=True) + + from cula.ops.kda.sm100.cp.chunk_delta_h import should_use_intracard_cp + + cpu = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens.cpu() + if should_use_intracard_cp(cpu, sm_count_provider(), num_qk_heads, chunk_size): + return CPDecision(True, cu_seqlens_cpu=cpu) + return CPDecision(False, "SM100 intracard CP heuristic declined for this shape.", cu_seqlens_cpu=cpu) diff --git a/cula/ops/kda/sm90/__init__.py b/cula/ops/kda/sm90/__init__.py new file mode 100644 index 00000000..8d2c8953 --- /dev/null +++ b/cula/ops/kda/sm90/__init__.py @@ -0,0 +1,4 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""SM90 (Hopper) two-kernel (K1+K2) FlashKDA prefill backend.""" diff --git a/cula/ops/kda/sm90/_common.py b/cula/ops/kda/sm90/_common.py new file mode 100644 index 00000000..f813deb9 --- /dev/null +++ b/cula/ops/kda/sm90/_common.py @@ -0,0 +1,77 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared low-level helpers for the SM90 FlashKDA kernels.""" + +import weakref + +import cutlass +import torch +from cutlass import Int32 +from cutlass._mlir.dialects import llvm as _llvm +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T as _T + + +@cutlass.dsl_user_op +def movm_t_b16(src_u32: Int32, *, loc=None, ip=None) -> Int32: + """``movmatrix.sync.aligned.m8n8.trans.b16`` -- register-file 8x8 b16 transpose.""" + result = _llvm.inline_asm( + _T.i32(), + [Int32(src_u32).ir_value(loc=loc, ip=ip)], + "movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Int32(result) + + +@cutlass.dsl_user_op +def add_f16x2_u32(a_u32: Int32, b_u32: Int32, *, loc=None, ip=None) -> Int32: + """Packed ``add.f16x2`` on two u32 registers.""" + result = _llvm.inline_asm( + _T.i32(), + [ + Int32(a_u32).ir_value(loc=loc, ip=ip), + Int32(b_u32).ir_value(loc=loc, ip=ip), + ], + "add.f16x2 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return Int32(result) + + +_WRAP_CACHE: dict = {} +_WRAP_CACHE_MAXSIZE = 512 + + +def _wrap_input(t: torch.Tensor, align: int, *, view_shape=None, cache: bool = False): + """Wrap a tensor as a CuTe tensor via from_dlpack. + + ``cache=True``: reuse across launches, keyed by (id, _version, align, view_shape) + and verified by weakref. + + Use ``cache=False`` for per-call buffers (workspaces, states). + """ + if not cache: + src = t if view_shape is None else t.view(view_shape) + return from_dlpack(src.detach(), assumed_align=align) + ckey = (id(t), t._version, align, view_shape) + entry = _WRAP_CACHE.get(ckey) + if entry is not None and entry[0]() is t: + return entry[1] + src = t if view_shape is None else t.view(view_shape) + w = from_dlpack(src.detach(), assumed_align=align) + if len(_WRAP_CACHE) >= _WRAP_CACHE_MAXSIZE: + _WRAP_CACHE.pop(next(iter(_WRAP_CACHE))) + _WRAP_CACHE[ckey] = (weakref.ref(t), w) + return w diff --git a/cula/ops/kda/sm90/cp/__init__.py b/cula/ops/kda/sm90/cp/__init__.py new file mode 100644 index 00000000..40dd6eef --- /dev/null +++ b/cula/ops/kda/sm90/cp/__init__.py @@ -0,0 +1,14 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""SM90 intracard context-parallel prefill backend.""" + +__all__ = ["intracard_prefill"] + + +def __getattr__(name): + if name == "intracard_prefill": + from cula.ops.kda.sm90.cp.driver import intracard_prefill + + return intracard_prefill + raise AttributeError(name) diff --git a/cula/ops/kda/sm90/cp/driver.py b/cula/ops/kda/sm90/cp/driver.py new file mode 100644 index 00000000..88f17594 --- /dev/null +++ b/cula/ops/kda/sm90/cp/driver.py @@ -0,0 +1,377 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Intracard-CP prefill executor: K1 once → pre_scan → merge → segment-K2. + +Only *runs* a given CPPlan; planning lives in cp.plan. The wrapper calls +run_cp with a plan from plan_prefill; intracard_prefill is the plan-and-run +entry for direct callers. +""" + +from __future__ import annotations + +import torch + +from cula.ops.kda.cp_mode import NotSplittableError +from cula.ops.kda.sm90.cp.merge import launch_merge +from cula.ops.kda.sm90.cp.plan import CPPlan, plan_auto, plan_manual, split_balanced +from cula.ops.kda.sm90.cp.pre_scan import launch_pre_scan +from cula.ops.kda.sm90.fwd import ( + _cute_arch_for_device, + _get_or_alloc_workspaces, + _get_or_build_varlen_metadata, + flash_kda_fwd, +) +from cula.ops.kda.sm90.k1 import launch_k1 +from cula.ops.kda.sm90.k2 import CHUNK, D, launch_k2 +from cula.utils import get_device_sm_count + +_SCRATCH_CACHE: dict = {} +_PLAN_TENSOR_CACHE: dict = {} +_SCRATCH_CACHE_MAXSIZE = 8 +_PLAN_TENSOR_CACHE_MAXSIZE = 64 + + +def _get_plan_tensor(values: tuple, dtype, device: torch.device) -> torch.Tensor: + key = (values, dtype, str(device)) + cached = _PLAN_TENSOR_CACHE.get(key) + if cached is None: + if len(_PLAN_TENSOR_CACHE) >= _PLAN_TENSOR_CACHE_MAXSIZE: + _PLAN_TENSOR_CACHE.pop(next(iter(_PLAN_TENSOR_CACHE))) + cached = torch.tensor(values, dtype=dtype, device=device) + _PLAN_TENSOR_CACHE[key] = cached + return cached + + +def _get_scratch(key_name: str, shape: tuple, dtype, device) -> torch.Tensor: + key = (key_name, shape, dtype, str(device)) + cached = _SCRATCH_CACHE.get(key) + if cached is None: + if len(_SCRATCH_CACHE) >= _SCRATCH_CACHE_MAXSIZE: + _SCRATCH_CACHE.pop(next(iter(_SCRATCH_CACHE))) + cached = torch.empty(shape, dtype=dtype, device=device) + _SCRATCH_CACHE[key] = cached + return cached + + +def _seq_tiles_of(q: torch.Tensor, cu_seqlens: torch.Tensor | None) -> list[int]: + B, T, _H, _K = q.shape + if cu_seqlens is None: + return [(T + CHUNK - 1) // CHUNK] * B + assert B == 1 and cu_seqlens.dtype == torch.int32 + meta = _get_or_build_varlen_metadata(cu_seqlens) + return [(sl + CHUNK - 1) // CHUNK for sl in meta.seq_lens] + + +def run_cp( + plan: CPPlan, + q, + k, + v, + g, + beta, + *, + scale, + out, + A_log, + dt_bias, + lower_bound, + initial_state=None, + final_state=None, + cu_seqlens=None, + state_transposed=False, +) -> None: + """Execute a non-trivial CPPlan verbatim -- no re-planning, no fallback.""" + assert not plan.trivial + assert q.is_cuda and q.dtype == torch.bfloat16 + with _cute_arch_for_device(q.device): + if cu_seqlens is None and q.shape[1] % CHUNK != 0: + _run_padded_dense( + plan, + q, + k, + v, + g, + beta, + scale, + out, + A_log, + dt_bias, + lower_bound, + initial_state, + final_state, + state_transposed, + ) + else: + _run_pipeline( + plan, + q, + k, + v, + g, + beta, + scale, + out, + A_log, + dt_bias, + lower_bound, + initial_state, + final_state, + cu_seqlens, + state_transposed, + ) + + +def intracard_prefill( + q, + k, + v, + g, + beta, + scale, + out, + A_log, + dt_bias, + lower_bound, + initial_state=None, + final_state=None, + cu_seqlens=None, + state_transposed=False, + s_split=None, + allow_fallback=True, +) -> None: + """Plan-and-run for direct callers. s_split forces a manual split; + allow_fallback=False means forced CP (structural split only, raise on a + trivial plan); otherwise a trivial plan falls back to the serial kernel.""" + seq_tiles = _seq_tiles_of(q, cu_seqlens) + H = q.shape[2] + if s_split is not None: + plan = plan_manual(seq_tiles, s_split) + elif allow_fallback: + plan = plan_auto(seq_tiles, H, get_device_sm_count(q.device)) + else: + plan = split_balanced(seq_tiles, H, get_device_sm_count(q.device)) + if plan.trivial: + if not allow_fallback: + raise NotSplittableError("SM90 intracard CP cannot split this shape.") + flash_kda_fwd( + q, + k, + v, + g, + beta, + scale=scale, + out=out, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=initial_state, + final_state=final_state, + cu_seqlens=cu_seqlens, + state_transposed=state_transposed, + ) + return + run_cp( + plan, + q, + k, + v, + g, + beta, + scale=scale, + out=out, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=initial_state, + final_state=final_state, + cu_seqlens=cu_seqlens, + state_transposed=state_transposed, + ) + + +# Dense partial-tile support — pad to CHUNK multiple, run aligned CP, strip back. +# (Varlen partial-tile is handled natively via ceil tiles + tile_starts mask.) +# The plan is built on ceil tile counts, so it is valid for the padded tensors. +def _pad_cp_inputs(pad, q, k, v, g, beta): + """No-op sentinels: q/k/v=0, g=-1e6 (decay~0), beta=-80 (sigmoid~0).""" + return pad(q, 0.0), pad(k, 0.0), pad(v, 0.0), pad(g, -1e6), pad(beta, -80.0) + + +def _run_padded_dense( + plan, + q, + k, + v, + g, + beta, + scale, + out, + A_log, + dt_bias, + lower_bound, + initial_state, + final_state, + state_transposed, +) -> None: + B, T, H, _ = q.shape + pad_t = ((T + CHUNK - 1) // CHUNK) * CHUNK - T + + def _pad(x: torch.Tensor, fill: float) -> torch.Tensor: + spec = (0, 0, 0, pad_t) if x.ndim == 3 else (0, 0, 0, 0, 0, pad_t) + return torch.nn.functional.pad(x, spec, value=fill) + + pq, pk, pv, pg, pbeta = _pad_cp_inputs(_pad, q, k, v, g, beta) + pout = out.new_empty((B, T + pad_t, H, D)) + _run_pipeline( + plan, + pq, + pk, + pv, + pg, + pbeta, + scale, + pout, + A_log, + dt_bias, + lower_bound, + initial_state, + final_state, + None, + state_transposed, + ) + out.copy_(pout[:, :T]) + + +def _run_pipeline( + plan: CPPlan, + q, + k, + v, + g, + beta, + scale, + out, + A_log, + dt_bias, + lower_bound, + initial_state, + final_state, + cu_seqlens, + state_transposed, +) -> None: + B, T, H, K = q.shape + device = q.device + + if cu_seqlens is None: + T_total = B * T + tile_starts = tile_actual_lens = None + is_varlen_padded = False + else: + assert B == 1 and cu_seqlens.dtype == torch.int32 + varlen_meta = _get_or_build_varlen_metadata(cu_seqlens) + T_total = T + is_varlen_padded = varlen_meta.needs_padding + tile_starts = varlen_meta.tile_starts if is_varlen_padded else None + tile_actual_lens = varlen_meta.tile_actual_lens if is_varlen_padded else None + + n_seg = plan.n_seg_total + seg_cu_tiles = _get_plan_tensor(plan.seg_cu, torch.int32, device) + + n_qk = plan.total_tiles * H * CHUNK * D + n_cc = plan.total_tiles * H * CHUNK * CHUNK + # ws_beta uses tile layout (total_tiles*CHUNK*H), not packed token layout (T_total*H). + ws_qd, ws_kd, ws_kr, ws_gt, ws_inv, ws_mqk, ws_beta = _get_or_alloc_workspaces( + n_qk, n_cc, plan.total_tiles * H * D, plan.total_tiles * CHUNK * H, device, beta.dtype + ) + launch_k1( + q, + k, + g, + A_log, + dt_bias, + beta.reshape(-1), + scale, + lower_bound, + ws_qd, + ws_kd, + ws_kr, + ws_gt, + ws_inv, + ws_mqk, + ws_beta, + tile_starts=tile_starts, + tile_actual_lens=tile_actual_lens, + total_tiles=plan.total_tiles if is_varlen_padded else None, + is_varlen=is_varlen_padded, + ) + + b_seg = _get_scratch("b_seg", (n_seg, H, D, D), torch.float32, device) + m_seg = _get_scratch("m_seg", (n_seg, H, D, D), torch.float32, device) + v_flat = v.view(1, T_total, H, D) if B > 1 else v + # Longest-first launch order: stable sort -> identity for uniform splits. + seg_tiles = plan.seg_tiles + seg_order = _get_plan_tensor(tuple(sorted(range(n_seg), key=lambda i: -seg_tiles[i])), torch.int32, device) + launch_pre_scan( + v_flat, + ws_beta, + ws_kd, + ws_kr, + ws_gt, + ws_inv, + b_seg, + m_seg, + seg_cu_tiles, + v_tile_starts=tile_starts, + v_tile_actual_lens=tile_actual_lens, + total_tiles=plan.total_tiles, + seg_order=seg_order, + ) + + init_bhvk = None + if initial_state is not None: + assert initial_state.shape == (plan.n_seqs, H, D, D) + init_bhvk = initial_state.to(torch.float32) + if state_transposed: + init_bhvk = init_bhvk.transpose(-1, -2) + init_bhvk = init_bhvk.contiguous() + carries = _get_scratch("carries", (n_seg, H, D, D), torch.float32, device) + launch_merge(carries, m_seg, b_seg, plan.per_seq, init_bhvk) + + out_flat = out.view(1, T_total, H, D) if B > 1 else out + seg_final = None + if final_state is not None: + seg_final = _get_scratch("seg_final", (n_seg, H, D, D), torch.float32, device) + launch_k2( + v_flat, + ws_beta, + ws_qd, + ws_kd, + ws_kr, + ws_gt, + ws_inv, + ws_mqk, + out_flat, + seg_cu_tiles, + initial_state=carries, + final_state=seg_final, + state_transposed=False, + v_tile_starts=tile_starts, + v_tile_actual_lens=tile_actual_lens, + seq_order=seg_order, + ) + + if final_state is not None: + last_idx = _get_plan_tensor(tuple(first + n - 1 for first, n in plan.per_seq), torch.long, device) + if ( + not state_transposed + and final_state.dtype == torch.float32 + and final_state.is_contiguous() + and final_state.shape == (plan.n_seqs, H, D, D) + ): + torch.index_select(seg_final, 0, last_idx, out=final_state) + else: + fin = seg_final.index_select(0, last_idx) + if state_transposed: + fin = fin.transpose(-1, -2) + final_state.copy_(fin.to(final_state.dtype)) diff --git a/cula/ops/kda/sm90/cp/merge.py b/cula/ops/kda/sm90/cp/merge.py new file mode 100644 index 00000000..7982d339 --- /dev/null +++ b/cula/ops/kda/sm90/cp/merge.py @@ -0,0 +1,433 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""CuTeDSL merge kernel for intracard-CP (SM90). + +Carries kept in SMEM across segment iterations — eliminates gmem round-trips. +Uses SM80 TF32 MMA (mma.sync.m16n8k8), available on SM80+. + +Layout: SM90 uses separate b_seg/m_seg [S,H,D,D] fp32 in bhvk order. +Matmul: carries[i+1] = carries[i] @ M_seg[i] + B_seg[i] (carry on left). + +Compared to SM100 merge: + - A=carry (sH), B=transition (sM) — SM100 has A=transition, B=state + - Row-tiles carry with BR=64 (SM100 col-tiles state with BV=64) + - 4 warps × 16 rows (SM100: 4 warps × 32 rows) +""" + +import functools + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass.cute.nvgpu import cpasync +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream + +from cula.ops.kda.sm90.k2 import D, _get_current_custream +from cula.ops.ptx import cvt_f32_to_tf32, mma_m16n8k8_tf32 + +_BR = 64 +_BN = 64 +_M_THR = 8 +_N_THR = 16 +_NUM_THREADS = _M_THR * _N_THR # 128 +_VEC = 4 +_PAD = 8 + + +class Merge: + """SMEM-resident carry merge. + + Grid: (D // BR, n_seqs, H). Each CTA handles BR=64 rows of carry. + 4 warps × 16 rows/warp = 64 rows. MMA: 1 M-tile × 16 N-tiles × 16 K-iters. + """ + + def __init__(self, H: int, has_init: int = 0): + self.H = H + self.has_init = int(has_init) + self.num_row_tiles = D // _BR # 2 + self.num_col_tiles = D // _BN # 2 + + @cute.jit + def __call__( + self, + b_seg: cute.Tensor, + m_seg: cute.Tensor, + carries: cute.Tensor, + init: cute.Tensor, + first_ptr: cute.Tensor, + nseg_ptr: cute.Tensor, + num_seqs: cutlass.Int32, + stream: cuda.CUstream, + ): + sH_layout = cute.make_layout((_BR, D), stride=(D + _PAD, 1)) + sM_layout = cute.make_layout((D, D), stride=(D + _PAD, 1)) + sB_layout = cute.make_layout((_BR, D), stride=(D + _PAD, 1)) + + @cute.struct + class SharedStorage: + sH: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, cute.cosize(sH_layout)], + 128, + ] + sM: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, cute.cosize(sM_layout)], + 128, + ] + sB: cute.struct.Align[ + cute.struct.MemRange[cutlass.Float32, cute.cosize(sB_layout)], + 128, + ] + + self.shared_storage_ty = SharedStorage + + copy_atom = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + cutlass.Float32, + num_bits_per_copy=_VEC * 32, + ) + thr_layout = cute.make_layout((_M_THR, _N_THR), stride=(_N_THR, 1)) + val_layout = cute.make_layout((1, _VEC)) + tiled_copy = cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout) + + store_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + cutlass.Float32, + num_bits_per_copy=_VEC * 32, + ) + tiled_store = cute.make_tiled_copy_tv(store_atom, thr_layout, val_layout) + + self.kernel( + b_seg, + m_seg, + carries, + init, + first_ptr, + nseg_ptr, + sH_layout, + sM_layout, + sB_layout, + tiled_copy, + tiled_store, + ).launch( + grid=(self.num_row_tiles, num_seqs, self.H), + block=(_NUM_THREADS, 1, 1), + stream=stream, + ) + + @cute.kernel + def kernel( + self, + b_seg: cute.Tensor, + m_seg: cute.Tensor, + carries: cute.Tensor, + init: cute.Tensor, + first_ptr: cute.Tensor, + nseg_ptr: cute.Tensor, + sH_layout: cute.Layout, + sM_layout: cute.Layout, + sB_layout: cute.Layout, + tiled_copy: cute.TiledCopy, + tiled_store: cute.TiledCopy, + ): + tidx, _, _ = cute.arch.thread_idx() + i_r, i_seq, i_h = cute.arch.block_idx() + + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage_ty) + sH = cute.make_tensor(storage.sH.data_ptr(), sH_layout) + sM = cute.make_tensor(storage.sM.data_ptr(), sM_layout) + sB = cute.make_tensor(storage.sB.data_ptr(), sB_layout) + + thr_copy = tiled_copy.get_slice(tidx) + thr_store = tiled_store.get_slice(tidx) + + first = first_ptr[i_seq] + n_seg = nseg_ptr[i_seq] + + t_m = tidx // _N_THR + t_n = tidx % _N_THR + + if cutlass.const_expr(self.has_init): + g_init = init[i_seq, i_h, None, None] + for j in cutlass.range_constexpr(self.num_col_tiles): + gI = cute.local_tile(g_init, tiler=(_BR, _BN), coord=(i_r, j)) + sHj = cute.local_tile(sH, tiler=(_BR, _BN), coord=(0, j)) + cute.copy( + tiled_copy, + thr_copy.partition_S(gI), + thr_copy.partition_D(sHj), + ) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.barrier() + else: + _rows_per_thr: cutlass.Constexpr[int] = _BR // _M_THR + _col_stride: cutlass.Constexpr[int] = _N_THR * _VEC + _n_col_groups: cutlass.Constexpr[int] = D // _col_stride + for ri in cutlass.range_constexpr(_rows_per_thr): + r = t_m + _M_THR * ri + for g in cutlass.range_constexpr(_n_col_groups): + for c in cutlass.range_constexpr(_VEC): + sH[r, g * _col_stride + t_n * _VEC + c] = cutlass.Float32(0.0) + cute.arch.barrier() + + g_c0 = carries[first, i_h, None, None] + for j in cutlass.range_constexpr(self.num_col_tiles): + gC = cute.local_tile(g_c0, tiler=(_BR, _BN), coord=(i_r, j)) + sHj = cute.local_tile(sH, tiler=(_BR, _BN), coord=(0, j)) + cute.copy( + tiled_store, + thr_store.partition_S(sHj), + thr_store.partition_D(gC), + ) + + seg_idx = cutlass.Int32(0) + idx = cutlass.Int32(0) + + for idx in cutlass.range(0, n_seg - 1, unroll=0): + seg_idx = first + idx + + g_b = b_seg[seg_idx, i_h, None, None] + for j in cutlass.range_constexpr(self.num_col_tiles): + gB = cute.local_tile(g_b, tiler=(_BR, _BN), coord=(i_r, j)) + sBj = cute.local_tile(sB, tiler=(_BR, _BN), coord=(0, j)) + cute.copy( + tiled_copy, + thr_copy.partition_S(gB), + thr_copy.partition_D(sBj), + ) + + g_m = m_seg[seg_idx, i_h, None, None] + for j in cutlass.range_constexpr(self.num_col_tiles): + gM = cute.local_tile(g_m, tiler=(D, _BN), coord=(0, j)) + sMj = cute.local_tile(sM, tiler=(D, _BN), coord=(0, j)) + cute.copy( + tiled_copy, + thr_copy.partition_S(gM), + thr_copy.partition_D(sMj), + ) + + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.barrier() + + # A = sH (carry rows), B = sM (transition cols) + warp_id = tidx // 32 + lane = tidx % 32 + q = lane // 4 + rp = lane % 4 + + M_TILES: cutlass.Constexpr[int] = 1 + N_TILES: cutlass.Constexpr[int] = D // 8 # 16 + K_TILES: cutlass.Constexpr[int] = D // 8 # 16 + + acc = cute.make_rmem_tensor( + cute.make_layout((M_TILES, N_TILES, 4)), + cutlass.Float32, + ) + + for mi in cutlass.range_constexpr(M_TILES): + row_a = warp_id * 16 + mi * 16 + q + row_b = row_a + 8 + for nj in cutlass.range_constexpr(N_TILES): + col_a = nj * 8 + rp * 2 + acc[mi, nj, 0] = sB[row_a, col_a] + acc[mi, nj, 1] = sB[row_a, col_a + 1] + acc[mi, nj, 2] = sB[row_b, col_a] + acc[mi, nj, 3] = sB[row_b, col_a + 1] + + a_frag = cute.make_rmem_tensor( + cute.make_layout((M_TILES, 4)), + cutlass.Int32, + ) + b_frag = cute.make_rmem_tensor( + cute.make_layout((N_TILES, 2)), + cutlass.Int32, + ) + + for ki in cutlass.range_constexpr(K_TILES): + k_base = ki * 8 + for mi in cutlass.range_constexpr(M_TILES): + row_a = warp_id * 16 + mi * 16 + q + row_b = row_a + 8 + a_frag[mi, 0] = cvt_f32_to_tf32(sH[row_a, k_base + rp]) + a_frag[mi, 1] = cvt_f32_to_tf32(sH[row_b, k_base + rp]) + a_frag[mi, 2] = cvt_f32_to_tf32(sH[row_a, k_base + rp + 4]) + a_frag[mi, 3] = cvt_f32_to_tf32(sH[row_b, k_base + rp + 4]) + for nj in cutlass.range_constexpr(N_TILES): + col_b = nj * 8 + q + b_frag[nj, 0] = cvt_f32_to_tf32(sM[k_base + rp, col_b]) + b_frag[nj, 1] = cvt_f32_to_tf32(sM[k_base + rp + 4, col_b]) + for mi in cutlass.range_constexpr(M_TILES): + for nj in cutlass.range_constexpr(N_TILES): + d0, d1, d2, d3 = mma_m16n8k8_tf32( + a_frag[mi, 0], + a_frag[mi, 1], + a_frag[mi, 2], + a_frag[mi, 3], + b_frag[nj, 0], + b_frag[nj, 1], + acc[mi, nj, 0], + acc[mi, nj, 1], + acc[mi, nj, 2], + acc[mi, nj, 3], + ) + acc[mi, nj, 0] = d0 + acc[mi, nj, 1] = d1 + acc[mi, nj, 2] = d2 + acc[mi, nj, 3] = d3 + + cute.arch.barrier() + for mi in cutlass.range_constexpr(M_TILES): + row_a = warp_id * 16 + mi * 16 + q + row_b = row_a + 8 + for nj in cutlass.range_constexpr(N_TILES): + col_a = nj * 8 + rp * 2 + sH[row_a, col_a] = acc[mi, nj, 0] + sH[row_a, col_a + 1] = acc[mi, nj, 1] + sH[row_b, col_a] = acc[mi, nj, 2] + sH[row_b, col_a + 1] = acc[mi, nj, 3] + + cute.arch.barrier() + g_c_out = carries[seg_idx + 1, i_h, None, None] + for j in cutlass.range_constexpr(self.num_col_tiles): + gC = cute.local_tile(g_c_out, tiler=(_BR, _BN), coord=(i_r, j)) + sHj = cute.local_tile(sH, tiler=(_BR, _BN), coord=(0, j)) + cute.copy( + tiled_store, + thr_store.partition_S(sHj), + thr_store.partition_D(gC), + ) + + cute.arch.barrier() + + +def _compile_merge(H: int, has_init: int): + kernel_obj = Merge(H=H, has_init=has_init) + + sym_s = cute.sym_int() + sym_n = cute.sym_int() + sym_nseq = cute.sym_int() + + b_seg_fake = make_fake_compact_tensor( + cutlass.Float32, + (sym_s, H, D, D), + stride_order=(3, 2, 1, 0), + assumed_align=128, + ) + m_seg_fake = make_fake_compact_tensor( + cutlass.Float32, + (sym_s, H, D, D), + stride_order=(3, 2, 1, 0), + assumed_align=128, + ) + carries_fake = make_fake_compact_tensor( + cutlass.Float32, + (sym_s, H, D, D), + stride_order=(3, 2, 1, 0), + assumed_align=128, + ) + init_fake = make_fake_compact_tensor( + cutlass.Float32, + (sym_n, H, D, D), + stride_order=(3, 2, 1, 0), + assumed_align=128, + ) + first_fake = make_fake_compact_tensor(cutlass.Int32, (sym_nseq,), assumed_align=16) + nseg_fake = make_fake_compact_tensor(cutlass.Int32, (sym_nseq,), assumed_align=16) + stream_fake = make_fake_stream() + + return cute.compile( + kernel_obj, + b_seg_fake, + m_seg_fake, + carries_fake, + init_fake, + first_fake, + nseg_fake, + cutlass.Int32(1), + stream_fake, + options="--enable-tvm-ffi", + ) + + +@functools.lru_cache(maxsize=32) +def _get_compiled_merge(H: int, has_init: int): + return _compile_merge(H, has_init) + + +# Cached per_seq-derived tensors and dummy init: building them per call +# costs two synchronous pageable H2D copies plus a memset. +_PER_SEQ_TENSOR_CACHE: dict[tuple, tuple[torch.Tensor, torch.Tensor]] = {} +_PER_SEQ_TENSOR_CACHE_MAXSIZE = 64 +_DUMMY_INIT_CACHE: dict[tuple, torch.Tensor] = {} + + +def _get_per_seq_tensors(per_seq: tuple, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + key = (per_seq, str(device)) + cached = _PER_SEQ_TENSOR_CACHE.get(key) + if cached is not None: + return cached + if len(_PER_SEQ_TENSOR_CACHE) >= _PER_SEQ_TENSOR_CACHE_MAXSIZE: + _PER_SEQ_TENSOR_CACHE.pop(next(iter(_PER_SEQ_TENSOR_CACHE))) + firsts = torch.tensor([f for f, _ in per_seq], dtype=torch.int32, device=device) + nsegs = torch.tensor([n for _, n in per_seq], dtype=torch.int32, device=device) + _PER_SEQ_TENSOR_CACHE[key] = (firsts, nsegs) + return firsts, nsegs + + +def _get_dummy_init(H: int, device: torch.device) -> torch.Tensor: + key = (H, str(device)) + cached = _DUMMY_INIT_CACHE.get(key) + if cached is None: + cached = torch.zeros(1, H, D, D, dtype=torch.float32, device=device) + _DUMMY_INIT_CACHE[key] = cached + return cached + + +def launch_merge( + carries: torch.Tensor, + m_seg: torch.Tensor, + b_seg: torch.Tensor, + per_seq: list[tuple[int, int]], + init_bhvk: torch.Tensor | None, +) -> torch.Tensor: + """CuTeDSL merge: SMEM-resident carry, TF32 MMA.""" + n_seqs = len(per_seq) + H = carries.shape[1] + device = carries.device + + firsts, nsegs = _get_per_seq_tensors(tuple(per_seq), device) + has_init = 1 if init_bhvk is not None else 0 + init_arg = init_bhvk if init_bhvk is not None else _get_dummy_init(H, device) + + compiled_fn = _get_compiled_merge(H, has_init) + stream = _get_current_custream() + + # tvm-ffi launch: torch tensors pass straight through, positional args unvalidated. + compiled_fn( + b_seg, + m_seg, + carries, + init_arg, + firsts, + nsegs, + n_seqs, + stream, + ) + return carries diff --git a/cula/ops/kda/sm90/cp/plan.py b/cula/ops/kda/sm90/cp/plan.py new file mode 100644 index 00000000..0eaea5e6 --- /dev/null +++ b/cula/ops/kda/sm90/cp/plan.py @@ -0,0 +1,205 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Segment planning for SM90 KDA intracard context-parallel (CP) prefill. + +CP re-runs the recurrence (pre_scan + segment-K2) to cut the serial-K2 +critical path, so a split must shrink the longest chain by more than the +re-run cost. plan_* return the plan the executor runs verbatim; a trivial +plan (nothing split) means "serial path". + +Decisions use whole tiles + SM count. The two constants below are chain-cost +ratios fitted on H100; ratios of same-species chains cancel absolute speed, +so no per-SKU recalibration. +""" + +from __future__ import annotations + +import os +import weakref +from dataclasses import dataclass + +import torch + +from cula.ops.kda.cp_mode import CPMode, NotSplittableError +from cula.utils import get_device_sm_count + +CHUNK = 16 + +# Manual-plan floor (s_split given): keeps hand-crafted splits non-degenerate. +MIN_SEG_TILES = 4 + +# Floor on auto segment length, tuned on H100. +AUTO_MIN_SEG_TILES = int(os.environ.get("CULA_KDA_CP_AUTO_MIN_SEG_TILES", "32")) + +# Per-tile cost of the CP re-run (pre_scan + segment-K2) relative to the +# serial K2 chain, tuned on H100: splitting engages only when the critical +# path shrinks by more than this. +RERUN_RATIO = float(os.environ.get("CULA_KDA_CP_RERUN_RATIO", "3")) + + +def _ceil_div(a: int, b: int) -> int: + return -(-a // b) + + +@dataclass(frozen=True) +class CPPlan: + """Segmentation in whole tiles: segment i covers packed tile range + [seg_cu[i], seg_cu[i+1]); per_seq[s] = (first_segment, n_segments). + Trivial (nothing split) means "serial path"; reason says why.""" + + seq_tiles: tuple[int, ...] + seg_cu: tuple[int, ...] + per_seq: tuple[tuple[int, int], ...] + reason: str = "" + + @classmethod + def serial(cls, seq_tiles=(), reason: str = "") -> CPPlan: + return cls(tuple(seq_tiles), (0,), (), reason) + + @property + def trivial(self) -> bool: + return self.n_seg_total <= self.n_seqs + + @property + def n_seqs(self) -> int: + return len(self.seq_tiles) + + @property + def n_seg_total(self) -> int: + return len(self.seg_cu) - 1 + + @property + def total_tiles(self) -> int: + return self.seg_cu[-1] + + @property + def seg_tiles(self) -> tuple[int, ...]: + return tuple(self.seg_cu[i + 1] - self.seg_cu[i] for i in range(self.n_seg_total)) + + @property + def max_seg_tiles(self) -> int: + return max(self.seg_tiles, default=0) + + +def _materialize(seq_tiles: list[int], n_segs: list[int]) -> CPPlan: + """Split each sequence into its requested number of near-equal segments.""" + seg_cu = [0] + per_seq = [] + for tiles, n in zip(seq_tiles, n_segs): + n = min(n, tiles) # zero-tile sequences get zero segments + first = len(seg_cu) - 1 + base, rem = divmod(tiles, max(n, 1)) + for i in range(n): + seg_cu.append(seg_cu[-1] + base + (1 if i < rem else 0)) + per_seq.append((first, n)) + return CPPlan(tuple(seq_tiles), tuple(seg_cu), tuple(per_seq)) + + +def split_balanced(seq_tiles: list[int], H: int, sm_count: int) -> CPPlan: + """One-wave split: target segment length spreads the total load over the + chain slots (sm_count // H); each sequence gets ceil(tiles/target) + segments. Equal lengths (not counts) minimize the critical path; a full + machine pushes the target past every sequence, so nothing splits.""" + parallel_segs = max(1, sm_count // H) # segments the card can run at once + target_seg_tiles = max(AUTO_MIN_SEG_TILES, _ceil_div(sum(seq_tiles), parallel_segs)) + n_segs = [max(1, min(_ceil_div(tiles, target_seg_tiles), tiles // AUTO_MIN_SEG_TILES)) for tiles in seq_tiles] + return _materialize(seq_tiles, n_segs) + + +def plan_auto(seq_tiles: list[int], H: int, sm_count: int) -> CPPlan: + """split_balanced + profitability: wall = max(critical chain, per-slot + load) on both sides, with CP tiles costing RERUN_RATIO more; engage only + when the serial wall exceeds the CP wall. + + Example: seq_tiles = [896] + [8] * 16, H=16, 132 SMs: + parallel_segs = 132 // 16 = 8 segments running at once + target_seg_tiles = max(32, ceil(1024 / 8)) = 128 + split = 896 -> 7 segments of 128; the 8-tile sequences stay whole + engage = serial wall max(896, 128) >= CP wall 3 * max(128, 128) -> CP + + Counter-example: seq_tiles = [64], H=16 -> 2 segments of 32, but + serial wall 64 < CP wall 3 * 32 = 96 -> trivial plan (serial path). + """ + if not seq_tiles: + return CPPlan.serial(seq_tiles, "empty input") + plan = split_balanced(seq_tiles, H, sm_count) + if plan.trivial: + return CPPlan.serial(seq_tiles, "machine already full or sequences too short to split") + parallel_segs = max(1, sm_count // H) + # No schedule finishes before the total work spread over the parallel + # slots, nor before its longest chain: wall = max of the two bounds. + load_bound = _ceil_div(sum(seq_tiles), parallel_segs) + serial_wall = max(max(seq_tiles), load_bound) + cp_wall = RERUN_RATIO * max(plan.max_seg_tiles, load_bound) + if serial_wall < cp_wall: + return CPPlan.serial( + seq_tiles, + f"serial wall {serial_wall} tiles vs CP wall ~{cp_wall:g}: " + f"splitting does not pay the {RERUN_RATIO:g}x re-run cost", + ) + return plan + + +def plan_manual(seq_tiles: list[int], s_split: int) -> CPPlan: + """Up to s_split near-equal segments per sequence; no profitability + judgment (tests, experiments).""" + n_segs = [max(1, min(s_split, tiles // max(1, MIN_SEG_TILES))) for tiles in seq_tiles] + return _materialize(seq_tiles, n_segs) + + +_SEQ_LENS_CACHE: dict = {} + + +def _seq_lens_from_cu(cu_seqlens: torch.Tensor, cu_seqlens_cpu: torch.Tensor | None) -> list[int]: + """Per-sequence lengths from cu_seqlens, cached by tensor identity so the + planner pays a GPU->host sync only on a cache miss, not on every call. + Passing cu_seqlens_cpu avoids the sync entirely even on a miss.""" + key = id(cu_seqlens) + stamp = (cu_seqlens.data_ptr(), int(cu_seqlens._version), cu_seqlens.numel()) + cached = _SEQ_LENS_CACHE.get(key) + if cached is not None: + ref, cstamp, seq_lens = cached + if ref() is cu_seqlens and cstamp == stamp: + return seq_lens + _SEQ_LENS_CACHE.pop(key, None) + src = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens.cpu() + cu_list = [int(x) for x in src.tolist()] + seq_lens = [cu_list[i + 1] - cu_list[i] for i in range(len(cu_list) - 1)] + if len(_SEQ_LENS_CACHE) >= 32: + _SEQ_LENS_CACHE.pop(next(iter(_SEQ_LENS_CACHE))) + _SEQ_LENS_CACHE[key] = (weakref.ref(cu_seqlens), stamp, seq_lens) + return seq_lens + + +def plan_prefill( + q: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + cu_seqlens_cpu: torch.Tensor | None = None, + mode: CPMode | None = None, +) -> CPPlan: + """Planning entry for the prefill wrapper. Trivial plan -> serial path. + AUTO declines unprofitable splits; FORCE splits whenever possible and + raises NotSplittableError otherwise.""" + if mode is None or mode is CPMode.OFF: + return CPPlan.serial((), "disabled") + B, T, H, _K = q.shape + if cu_seqlens is None: + # Non-CHUNK-aligned lengths run CP via pad-to-tile, hence ceil tiles. + seq_tiles = [_ceil_div(T, CHUNK)] * B + elif B != 1: + if mode is CPMode.FORCE: + raise NotSplittableError("SM90 intracard CP varlen mode requires packed B=1.") + return CPPlan.serial((), "varlen requires packed B=1") + else: + seq_lens = _seq_lens_from_cu(cu_seqlens, cu_seqlens_cpu) + seq_tiles = [_ceil_div(sl, CHUNK) for sl in seq_lens] + + if mode is CPMode.FORCE: + if not seq_tiles: + raise NotSplittableError("SM90 intracard CP requires at least one sequence.") + plan = split_balanced(seq_tiles, H, get_device_sm_count(q.device)) + if plan.trivial: + raise NotSplittableError("SM90 intracard CP cannot split this shape.") + return plan + return plan_auto(seq_tiles, H, get_device_sm_count(q.device)) diff --git a/cula/ops/kda/sm90/cp/pre_scan.py b/cula/ops/kda/sm90/cp/pre_scan.py new file mode 100644 index 00000000..4d011655 --- /dev/null +++ b/cula/ops/kda/sm90/cp/pre_scan.py @@ -0,0 +1,699 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""Intracard-CP pre-scan kernel (SM90) — stage 1 (fused S+M chains). + +Computes B_seg (S-chain, S0=0) and M_seg (M-chain, M0=I) per segment in one +pass; the per-tile recurrence reuses the FlashKDA-derived K2 math it imports. + +160 threads = 4 MMA warps + 1 LOAD warp; outputs fp32 bhvk layout. +""" + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.cpasync as cpasync +import torch +from cutlass.cute.nvgpu import warp +from cutlass.cute.nvgpu.warpgroup import SmemLayoutAtomKind, make_smem_layout_atom +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream + +from cula.ops.kda.sm90._common import movm_t_b16 +from cula.ops.kda.sm90.k2 import ( + CHUNK, + D, + _get_current_custream, + _get_dummy_int32, + _get_identity_order, + _make_out_kinter_one_stage, + _make_state_smem_layout, +) + +THREADS_PER_CTA = 160 # 4 MMA warps + 1 LOAD warp +LOAD_WARP_IDX = 4 + + +@cute.kernel +def pre_scan_kernel( + tma_atom_v: cute.CopyAtom, + tma_tensor_v: cute.Tensor, + tma_atom_kd: cute.CopyAtom, + tma_tensor_kd: cute.Tensor, + tma_atom_kr: cute.CopyAtom, + tma_tensor_kr: cute.Tensor, + tma_atom_inv: cute.CopyAtom, + tma_tensor_inv: cute.Tensor, + tma_atom_gt: cute.CopyAtom, + tma_tensor_gt: cute.Tensor, + tma_atom_beta: cute.CopyAtom, + tma_tensor_beta: cute.Tensor, + H: cutlass.Constexpr[int], + total_tiles: cutlass.Int32, + T_total: cutlass.Int32, + seg_cu_tiles: cute.Tensor, + seg_order: cute.Tensor, # int32 [S]: launch-slot -> segment index + b_state_g: cute.Tensor, # flat fp32 [S*H*D*D] gmem, bhvk + m_state_g: cute.Tensor, # flat fp32 [S*H*D*D] gmem, bhvk + v_tile_starts: cute.Tensor, + v_tile_actual_lens: cute.Tensor, + v_is_varlen: cutlass.Constexpr[bool], +): + # Longest-first launch order; pure reordering. + seg_slot, head_idx, _ = cute.arch.block_idx() + seg_idx = cutlass.Int32(seg_order[seg_slot]) + tidx, _, _ = cute.arch.thread_idx() + + smem = cutlass.utils.SmemAllocator() + + state_layout = _make_state_smem_layout() + + STAGES: cutlass.Constexpr[int] = 2 + cc_stage_layout = cute.make_layout((CHUNK, CHUNK, STAGES), stride=(CHUNK, 1, CHUNK * CHUNK)) + v_kinter_atom = make_smem_layout_atom(SmemLayoutAtomKind.K_INTER, cutlass.BFloat16) + v_stage_layout = cute.tile_to_shape(v_kinter_atom, (CHUNK, D, STAGES), order=(0, 1, 2)) + kd_stage_layout = cute.tile_to_shape(v_kinter_atom, (CHUNK, D, STAGES), order=(0, 1, 2)) + kr_stage_layout = cute.tile_to_shape(v_kinter_atom, (CHUNK, D, STAGES), order=(0, 1, 2)) + kr_mninter_atom = make_smem_layout_atom(SmemLayoutAtomKind.MN_INTER, cutlass.BFloat16) + kr_t_stage_layout = cute.tile_to_shape(kr_mninter_atom, (D, CHUNK, STAGES), order=(1, 0, 2)) + + sV = smem.allocate_tensor(cutlass.BFloat16, v_stage_layout, 128) + sKd = smem.allocate_tensor(cutlass.BFloat16, kd_stage_layout, 128) + sKr = smem.allocate_tensor(cutlass.BFloat16, kr_stage_layout, 128) + sINV = smem.allocate_tensor(cutlass.BFloat16, cc_stage_layout, 128) + sState = smem.allocate_tensor(cutlass.BFloat16, state_layout, 128) + sM = smem.allocate_tensor(cutlass.BFloat16, state_layout, 128) + sGt = smem.allocate_tensor(cutlass.Float32, cute.make_layout((D, 1, STAGES), stride=(1, D, D)), 128) + sBeta = smem.allocate_tensor(cutlass.BFloat16, cute.make_layout((CHUNK, 1, STAGES), stride=(1, 64, 64)), 128) + sMbar = smem.allocate_tensor(cutlass.Int64, cute.make_layout((STAGES,)), 16) + sMbar_ptr = sMbar.iterator + sMbarE = smem.allocate_tensor(cutlass.Int64, cute.make_layout((STAGES,)), 16) + sMbarE_ptr = sMbarE.iterator + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + if warp_idx == 0: + with cute.arch.elect_one(): + for s in cutlass.range_constexpr(STAGES): + cute.arch.mbarrier_init(sMbar_ptr + cutlass.Int32(s), cutlass.Int32(1)) + cute.arch.mbarrier_init(sMbarE_ptr + cutlass.Int32(s), cutlass.Int32(1)) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + + gSrc_v = cute.local_tile(tma_tensor_v, (CHUNK, D), (None, None, None)) + tVs, tVg = cpasync.tma_partition( + tma_atom_v, + 0, + cute.make_layout(1), + cute.group_modes(sV, 0, 2), + cute.group_modes(gSrc_v, 0, 2), + ) + gSrc_kd = cute.local_tile(tma_tensor_kd, (CHUNK, D), (None, None, None)) + tKDs, tKDg = cpasync.tma_partition( + tma_atom_kd, + 0, + cute.make_layout(1), + cute.group_modes(sKd, 0, 2), + cute.group_modes(gSrc_kd, 0, 2), + ) + gSrc_kr = cute.local_tile(tma_tensor_kr, (CHUNK, D), (None, None, None)) + tKRs, tKRg = cpasync.tma_partition( + tma_atom_kr, + 0, + cute.make_layout(1), + cute.group_modes(sKr, 0, 2), + cute.group_modes(gSrc_kr, 0, 2), + ) + gSrc_inv = cute.local_tile(tma_tensor_inv, (CHUNK, CHUNK), (None, None, None)) + tIs, tIg = cpasync.tma_partition( + tma_atom_inv, + 0, + cute.make_layout(1), + cute.group_modes(sINV, 0, 2), + cute.group_modes(gSrc_inv, 0, 2), + ) + gSrc_gt = cute.local_tile(tma_tensor_gt, (D, 1), (None, None, None)) + tGTs, tGTg = cpasync.tma_partition( + tma_atom_gt, + 0, + cute.make_layout(1), + cute.group_modes(sGt, 0, 2), + cute.group_modes(gSrc_gt, 0, 2), + ) + gSrc_beta = cute.local_tile(tma_tensor_beta, (CHUNK, 1), (None, None, None)) + tBs, tBg = cpasync.tma_partition( + tma_atom_beta, + 0, + cute.make_layout(1), + cute.group_modes(sBeta, 0, 2), + cute.group_modes(gSrc_beta, 0, 2), + ) + + # sState=0, sM=I + if tidx < D: + for e in cutlass.range_constexpr(D): + sState[tidx, e] = cutlass.BFloat16(0.0) + sM[tidx, e] = cutlass.BFloat16(0.0) + sM[tidx, tidx] = cutlass.BFloat16(1.0) + cute.arch.barrier() + + mma_atom = warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16)) + tiled_mma = cute.make_tiled_mma( + mma_atom, + atom_layout_mnk=(1, 4, 1), + permutation_mnk=(16, 32, 16), + ) + thr_mma = tiled_mma.get_slice(tidx) + + copy_atom_AB = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), + cutlass.BFloat16, + ) + smem_tiled_copy_A = cute.make_tiled_copy_A(copy_atom_AB, tiled_mma) + smem_thr_copy_A = smem_tiled_copy_A.get_slice(tidx) + copy_atom_B_T = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4), + cutlass.BFloat16, + ) + smem_tiled_copy_B_T = cute.make_tiled_copy_B(copy_atom_B_T, tiled_mma) + smem_thr_copy_B_T = smem_tiled_copy_B_T.get_slice(tidx) + + tiled_mma_state = cute.make_tiled_mma( + mma_atom, + atom_layout_mnk=(1, 4, 1), + permutation_mnk=(16, 32, 16), + ) + thr_mma_state = tiled_mma_state.get_slice(tidx) + smem_tiled_copy_A_state = cute.make_tiled_copy_A(copy_atom_B_T, tiled_mma_state) + smem_thr_copy_A_state = smem_tiled_copy_A_state.get_slice(tidx) + + sKd_s0 = sKd[(None, None, 0)] + sKd_tile0 = cute.flat_divide(sKd_s0, (CHUNK, 16)) + # state[K_in, D_out] transposed B-view for Phase 1 + sState_B_view = cute.make_tensor(sState.iterator, layout=cute.select(sState.layout, mode=[1, 0])) + sState_tile = cute.flat_divide(sState_B_view, (D, 16)) + sM_B_view = cute.make_tensor(sM.iterator, layout=cute.select(sM.layout, mode=[1, 0])) + sM_tile = cute.flat_divide(sM_B_view, (D, 16)) + + sKd_ref = sKd_tile0[None, None, 0, 0] + sState_ref = sState_tile[None, None, 0, 0] + + tCrKd = thr_mma.make_fragment_A(thr_mma.partition_A(sKd_ref)) + tCrState = thr_mma.make_fragment_B(thr_mma.partition_B(sState_ref)) + tCrU = thr_mma.make_fragment_C(tiled_mma.partition_shape_C((CHUNK, D))) + # Separate accumulator for the M-chain's kd@M, hoisted ahead of the + # S-chain's dependent stages instead of serializing behind them. + tCrU_m = thr_mma.make_fragment_C(tiled_mma.partition_shape_C((CHUNK, D))) + + tCrKd_cv = smem_thr_copy_A.retile(tCrKd) + tCrState_cv = smem_thr_copy_B_T.retile(tCrState) + + # sKr transposed view (MN_INTER swizzle, aliased storage) + sKr_T_view = cute.make_tensor(sKr.iterator, kr_t_stage_layout) + sKr_T_view_s0 = sKr_T_view[(None, None, 0)] + sKr_T_ref = cute.flat_divide(sKr_T_view_s0, (D, CHUNK))[None, None, 0, 0] + + sKr_T_blk_for_frag = cute.flat_divide(sKr_T_view_s0, (CHUNK, CHUNK))[None, None, 0, 0] + tCrKrA_state_blk = thr_mma_state.make_fragment_A(thr_mma_state.partition_A(sKr_T_blk_for_frag)) + tCrKrA_state_blk_cv = smem_thr_copy_A_state.retile(tCrKrA_state_blk) + tCrUpd_blk = thr_mma_state.make_fragment_C(tiled_mma_state.partition_shape_C((CHUNK, D))) + sState_blk_tile = cute.flat_divide(sState, (CHUNK, D)) + sM_blk_tile = cute.flat_divide(sM, (CHUNK, D)) + coord_state_blk = cute.make_identity_tensor((CHUNK, D)) + tCcState_blk = thr_mma_state.partition_C(coord_state_blk) + + tCrU_T = thr_mma.make_fragment_B(thr_mma.partition_B(sKr_T_ref)) + + sINV_s0 = sINV[(None, None, 0)] + sINV_tile0 = cute.flat_divide(sINV_s0, (CHUNK, CHUNK)) + sINV_ref = sINV_tile0[None, None, 0, 0] + tCrInv = thr_mma.make_fragment_A(thr_mma.partition_A(sINV_ref)) + tCrInv_cv = smem_thr_copy_A.retile(tCrInv) + tCrU_post = thr_mma.make_fragment_C(tiled_mma.partition_shape_C((CHUNK, D))) + tCrU_T_post = cute.make_fragment_like(tCrU_T) + tCrU_post_bf16 = cute.make_fragment_like(tCrU_post, cutlass.BFloat16) + tCrU_pre_bf16 = cute.make_fragment_like(tCrU, cutlass.BFloat16) + + tile_base = seg_cu_tiles[seg_idx] + t_tiles = seg_cu_tiles[seg_idx + 1] - tile_base + TMA_BYTES: cutlass.Constexpr[int] = 3 * CHUNK * D * 2 + CHUNK * CHUNK * 2 + D * 4 + CHUNK * 2 + + if warp_idx == LOAD_WARP_IDX: + s_dyn_l = cutlass.Int32(0) + phase_emp = cutlass.Int32(1) + if cutlass.const_expr(v_is_varlen): + seq_v_start = v_tile_starts[tile_base] + tma_tensor_v_seq = cute.domain_offset((seq_v_start, 0, 0), tma_tensor_v) + gSrc_v_seq = cute.local_tile(tma_tensor_v_seq, (CHUNK, D), (None, None, None)) + tVs_seq, tVg_seq = cpasync.tma_partition( + tma_atom_v, + 0, + cute.make_layout(1), + cute.group_modes(sV, 0, 2), + cute.group_modes(gSrc_v_seq, 0, 2), + ) + for t in cutlass.range(t_tiles, unroll=1): + cute.arch.mbarrier_wait(sMbarE_ptr + s_dyn_l, phase_emp) + tg_l = tile_base + t + wt_l = head_idx * total_tiles + tg_l + bar_l = sMbar_ptr + s_dyn_l + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(bar_l, cutlass.Int32(TMA_BYTES)) + if cutlass.const_expr(v_is_varlen): + cute.copy(tma_atom_v, tVg_seq[(None, t, 0, head_idx)], tVs_seq[(None, s_dyn_l)], tma_bar_ptr=bar_l) + else: + cute.copy(tma_atom_v, tVg[(None, tg_l, 0, head_idx)], tVs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(tma_atom_kd, tKDg[(None, 0, 0, wt_l)], tKDs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(tma_atom_kr, tKRg[(None, 0, 0, wt_l)], tKRs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(tma_atom_inv, tIg[(None, 0, 0, wt_l)], tIs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(tma_atom_gt, tGTg[(None, 0, 0, wt_l)], tGTs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(tma_atom_beta, tBg[(None, 0, 0, wt_l)], tBs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + s_dyn_l = s_dyn_l + cutlass.Int32(1) + if s_dyn_l == cutlass.Int32(STAGES): + s_dyn_l = cutlass.Int32(0) + phase_emp = phase_emp ^ cutlass.Int32(1) + else: + phase_full = cutlass.Int32(0) + s_dyn = cutlass.Int32(0) + + for t in cutlass.range(t_tiles, unroll=1): + sV_s = sV[(None, None, s_dyn)] + sKd_tile = cute.flat_divide(sKd[(None, None, s_dyn)], (CHUNK, 16)) + sINV_ref_s = cute.flat_divide(sINV[(None, None, s_dyn)], (CHUNK, CHUNK))[None, None, 0, 0] + sGt_s = sGt[(None, 0, s_dyn)] + sBeta_s = sBeta[(None, 0, s_dyn)] + + cute.arch.mbarrier_wait(sMbar_ptr + s_dyn, phase_full) + + if cutlass.const_expr(v_is_varlen): + actual_len = v_tile_actual_lens[tile_base + t] + if actual_len < cutlass.Int32(CHUNK): + if tidx < D: + col_v_tail = tidx + for r in cutlass.range_constexpr(CHUNK): + if actual_len <= cutlass.Int32(r): + sV_s[r, col_v_tail] = cutlass.BFloat16(0.0) + cute.arch.barrier(barrier_id=1, number_of_threads=128) + + sKr_T_s = sKr_T_view[(None, None, s_dyn)] + sKr_T_blk_tile_s = cute.flat_divide(sKr_T_s, (CHUNK, CHUNK)) + + # kd @ state + tCrU.fill(0.0) + for k in cutlass.range_constexpr(D // 16): + sKd_k = sKd_tile[None, None, 0, k] + sState_k = sState_tile[None, None, 0, k] + cute.copy(smem_tiled_copy_B_T, smem_thr_copy_B_T.partition_S(sState_k), tCrState_cv) + cute.copy(smem_tiled_copy_A, smem_thr_copy_A.partition_S(sKd_k), tCrKd_cv) + cute.gemm(tiled_mma, tCrU, tCrKd, tCrState, tCrU) + + # M-chain front half (kd @ M), hoisted to overlap the S-chain's + # dependent sigmoid/movmatrix/INV stages. + tCrU_m.fill(0.0) + for k in cutlass.range_constexpr(D // 16): + sKd_k = sKd_tile[None, None, 0, k] + sM_k = sM_tile[None, None, 0, k] + cute.copy(smem_tiled_copy_B_T, smem_thr_copy_B_T.partition_S(sM_k), tCrState_cv) + cute.copy(smem_tiled_copy_A, smem_thr_copy_A.partition_S(sKd_k), tCrKd_cv) + cute.gemm(tiled_mma, tCrU_m, tCrKd, tCrState, tCrU_m) + + # sigmoid(beta) * (v - u_pre) + lane_in_warp = tidx % 32 + Rrow0 = lane_in_warp // 4 + Rrow1 = Rrow0 + 8 + b0 = cutlass.Float32(sBeta_s[Rrow0]) + b1 = cutlass.Float32(sBeta_s[Rrow1]) + sig0 = cutlass.Float32(0.5) * (cute.tanh(b0 * cutlass.Float32(0.5), fastmath=True) + cutlass.Float32(1.0)) + sig1 = cutlass.Float32(0.5) * (cute.tanh(b1 * cutlass.Float32(0.5), fastmath=True) + cutlass.Float32(1.0)) + tCsV = thr_mma.partition_C(sV_s) + for i in cutlass.range_constexpr(cute.size(tCrU)): + ii: cutlass.Constexpr[int] = i + sub_i: cutlass.Constexpr[int] = (ii % 4) // 2 + sig = sig0 if sub_i == 0 else sig1 + diff = cutlass.Float32(tCsV[ii]) - tCrU[ii] + tCrU_pre_bf16[ii] = cutlass.BFloat16(diff * sig) + tCrU_pre_u32 = cute.recast_tensor(tCrU_pre_bf16, dtype=cutlass.Int32) + tCrU_T_u32 = cute.recast_tensor(tCrU_T, dtype=cutlass.Int32) + for i in cutlass.range_constexpr(cute.size(tCrU_pre_u32)): + ii: cutlass.Constexpr[int] = i + tCrU_T_u32[ii] = movm_t_b16(cutlass.Int32(tCrU_pre_u32[ii])) + + # U_post = INV @ U_pre + cute.copy(smem_tiled_copy_A, smem_thr_copy_A.partition_S(sINV_ref_s), tCrInv_cv) + tCrU_post.fill(0.0) + cute.gemm(tiled_mma, tCrU_post, tCrInv, tCrU_T, tCrU_post) + for i in cutlass.range_constexpr(cute.size(tCrU_post)): + ii: cutlass.Constexpr[int] = i + tCrU_post_bf16[ii] = cutlass.BFloat16(tCrU_post[ii]) + tCrU_post_u32 = cute.recast_tensor(tCrU_post_bf16, dtype=cutlass.Int32) + tCrU_T_post_u32 = cute.recast_tensor(tCrU_T_post, dtype=cutlass.Int32) + for i in cutlass.range_constexpr(cute.size(tCrU_post_u32)): + ii: cutlass.Constexpr[int] = i + tCrU_T_post_u32[ii] = movm_t_b16(cutlass.Int32(tCrU_post_u32[ii])) + + # State update: state = state*gt + kr^T @ U + M_BLOCKS: cutlass.Constexpr[int] = D // CHUNK + for mi in cutlass.range_constexpr(M_BLOCKS): + sKr_T_blk_s = sKr_T_blk_tile_s[None, None, mi, 0] + cute.copy( + smem_tiled_copy_A_state, + smem_thr_copy_A_state.partition_S(sKr_T_blk_s), + tCrKrA_state_blk_cv, + ) + tCrUpd_blk.fill(0.0) + cute.gemm(tiled_mma_state, tCrUpd_blk, tCrKrA_state_blk, tCrU_T_post, tCrUpd_blk) + + sState_blk = sState_blk_tile[None, None, mi, 0] + tCsState_blk = thr_mma_state.partition_C(sState_blk) + state_frag_blk = cute.make_fragment_like(tCsState_blk, cutlass.BFloat16) + gt_frag_blk = cute.make_fragment_like(tCsState_blk, cutlass.Float32) + m_off: cutlass.Constexpr[int] = mi * CHUNK + for i in cutlass.range_constexpr(cute.size(state_frag_blk)): + ii: cutlass.Constexpr[int] = i + state_frag_blk[ii] = tCsState_blk[ii] + gt_frag_blk[ii] = sGt_s[m_off + tCcState_blk[ii][0]] + for i in cutlass.range_constexpr(cute.size(tCrUpd_blk)): + ii: cutlass.Constexpr[int] = i + old = cutlass.Float32(state_frag_blk[ii]) * gt_frag_blk[ii] + tCsState_blk[ii] = cutlass.BFloat16(old + tCrUpd_blk[ii]) + + # M-chain: same pipeline with V := 0 (duality). + # (kd @ M already issued above.) + + # sigmoid(beta) * (0 - u_pre) + for i in cutlass.range_constexpr(cute.size(tCrU_m)): + ii: cutlass.Constexpr[int] = i + sub_i: cutlass.Constexpr[int] = (ii % 4) // 2 + sig = sig0 if sub_i == 0 else sig1 + diff = cutlass.Float32(0.0) - tCrU_m[ii] + tCrU_pre_bf16[ii] = cutlass.BFloat16(diff * sig) + for i in cutlass.range_constexpr(cute.size(tCrU_pre_u32)): + ii: cutlass.Constexpr[int] = i + tCrU_T_u32[ii] = movm_t_b16(cutlass.Int32(tCrU_pre_u32[ii])) + + # U_post = INV @ U_pre + tCrU_post.fill(0.0) + cute.gemm(tiled_mma, tCrU_post, tCrInv, tCrU_T, tCrU_post) + for i in cutlass.range_constexpr(cute.size(tCrU_post)): + ii: cutlass.Constexpr[int] = i + tCrU_post_bf16[ii] = cutlass.BFloat16(tCrU_post[ii]) + for i in cutlass.range_constexpr(cute.size(tCrU_post_u32)): + ii: cutlass.Constexpr[int] = i + tCrU_T_post_u32[ii] = movm_t_b16(cutlass.Int32(tCrU_post_u32[ii])) + + # State update': M = M*gt + kr^T @ U + for mi in cutlass.range_constexpr(M_BLOCKS): + sKr_T_blk_s = sKr_T_blk_tile_s[None, None, mi, 0] + cute.copy( + smem_tiled_copy_A_state, + smem_thr_copy_A_state.partition_S(sKr_T_blk_s), + tCrKrA_state_blk_cv, + ) + tCrUpd_blk.fill(0.0) + cute.gemm(tiled_mma_state, tCrUpd_blk, tCrKrA_state_blk, tCrU_T_post, tCrUpd_blk) + + sM_blk = sM_blk_tile[None, None, mi, 0] + tCsM_blk = thr_mma_state.partition_C(sM_blk) + m_frag_blk = cute.make_fragment_like(tCsM_blk, cutlass.BFloat16) + gt_frag_blk_m = cute.make_fragment_like(tCsM_blk, cutlass.Float32) + m_off2: cutlass.Constexpr[int] = mi * CHUNK + for i in cutlass.range_constexpr(cute.size(m_frag_blk)): + ii: cutlass.Constexpr[int] = i + m_frag_blk[ii] = tCsM_blk[ii] + gt_frag_blk_m[ii] = sGt_s[m_off2 + tCcState_blk[ii][0]] + for i in cutlass.range_constexpr(cute.size(tCrUpd_blk)): + ii: cutlass.Constexpr[int] = i + old = cutlass.Float32(m_frag_blk[ii]) * gt_frag_blk_m[ii] + tCsM_blk[ii] = cutlass.BFloat16(old + tCrUpd_blk[ii]) + + cute.arch.barrier(barrier_id=1, number_of_threads=128) + cute.arch.fence_view_async_shared() + if warp_idx == 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(sMbarE_ptr + s_dyn) + s_dyn = s_dyn + cutlass.Int32(1) + if s_dyn == cutlass.Int32(STAGES): + s_dyn = cutlass.Int32(0) + phase_full = phase_full ^ cutlass.Int32(1) + cute.arch.barrier() + # Epilogue: write both states fp32. The index swap below (gmem[d_out, tidx] = + # smem[tidx, d_out]) stores the TRANSPOSE S^T / M^T on purpose — merge consumes + # this transposed convention (carry @ M^T), so do not "fix" the orientation here. + state_base_f = cutlass.Int32(seg_idx) * cutlass.Int32(H * D * D) + cutlass.Int32(head_idx) * cutlass.Int32(D * D) + if tidx < D: + for d_out in cutlass.range_constexpr(D): + b_state_g[state_base_f + cutlass.Int32(d_out * D) + cutlass.Int32(tidx)] = cutlass.Float32(sState[tidx, d_out]) + m_state_g[state_base_f + cutlass.Int32(d_out * D) + cutlass.Int32(tidx)] = cutlass.Float32(sM[tidx, d_out]) + + +@cute.jit +def run_pre_scan( + v: cute.Tensor, + beta: cute.Tensor, + ws_kd: cute.Tensor, + ws_kr: cute.Tensor, + ws_gt: cute.Tensor, + ws_inv: cute.Tensor, + seg_cu_tiles: cute.Tensor, + seg_order: cute.Tensor, + b_state_g: cute.Tensor, # flat fp32 [S*H*D*D] + m_state_g: cute.Tensor, # flat fp32 [S*H*D*D] + v_tile_starts: cute.Tensor, + v_tile_actual_lens: cute.Tensor, + H: cutlass.Constexpr[int], + total_tiles: cutlass.Int32, + T_total: cutlass.Int32, + v_is_varlen: cutlass.Constexpr[bool], + S: cutlass.Int32, + stream: cuda_drv.CUstream, +): + cc_smem = cute.make_layout((CHUNK, CHUNK), stride=(CHUNK, 1)) + kinter_smem = _make_out_kinter_one_stage() + + def make_thd_atom(t, op): + view = cute.make_tensor( + t.iterator, + cute.make_layout((T_total, D, H), stride=(H * D, 1, D)), + ) + return cpasync.make_tiled_tma_atom(op, view, kinter_smem, (CHUNK, D)) + + def make_ws_qkd_atom(t): + view = cute.make_tensor( + t.iterator, + cute.make_layout((CHUNK, D, total_tiles * H), stride=(D, 1, CHUNK * D)), + ) + return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, kinter_smem, (CHUNK, D)) + + def make_ws_cc_atom(t): + view = cute.make_tensor( + t.iterator, + cute.make_layout((CHUNK, CHUNK, total_tiles * H), stride=(CHUNK, 1, CHUNK * CHUNK)), + ) + return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, cc_smem, (CHUNK, CHUNK)) + + tma_atom_v, tma_tensor_v = make_thd_atom(v, cpasync.CopyBulkTensorTileG2SOp()) + tma_atom_kd, tma_tensor_kd = make_ws_qkd_atom(ws_kd) + tma_atom_kr, tma_tensor_kr = make_ws_qkd_atom(ws_kr) + tma_atom_inv, tma_tensor_inv = make_ws_cc_atom(ws_inv) + + gt_smem = cute.make_layout((D, 1), stride=(1, D)) + beta_smem = cute.make_layout((CHUNK, 1), stride=(1, 64)) + + def make_gt_atom(t): + view = cute.make_tensor( + t.iterator, + cute.make_layout((D, 1, total_tiles * H), stride=(1, D, D)), + ) + return cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + view, + gt_smem, + (D, 1), + ) + + def make_beta_atom(t): + view = cute.make_tensor( + t.iterator, + cute.make_layout((CHUNK, 1, total_tiles * H), stride=(1, CHUNK, CHUNK)), + ) + return cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + view, + beta_smem, + (CHUNK, 1), + ) + + tma_atom_gt, tma_tensor_gt = make_gt_atom(ws_gt) + tma_atom_beta, tma_tensor_beta = make_beta_atom(beta) + + STAGES_LOCAL = 2 + smem_bytes = ( + 2 * D * D * 2 + + STAGES_LOCAL * 3 * (CHUNK * D * 2) + + STAGES_LOCAL * (CHUNK * CHUNK * 2) + + STAGES_LOCAL * (D * 4) + + STAGES_LOCAL * (64 * 2) + + STAGES_LOCAL * 8 + + STAGES_LOCAL * 8 + + 2048 + ) + + pre_scan_kernel( + tma_atom_v, + tma_tensor_v, + tma_atom_kd, + tma_tensor_kd, + tma_atom_kr, + tma_tensor_kr, + tma_atom_inv, + tma_tensor_inv, + tma_atom_gt, + tma_tensor_gt, + tma_atom_beta, + tma_tensor_beta, + H, + total_tiles, + T_total, + seg_cu_tiles, + seg_order, + b_state_g, + m_state_g, + v_tile_starts, + v_tile_actual_lens, + v_is_varlen, + ).launch( + grid=(S, H, 1), + block=[THREADS_PER_CTA, 1, 1], + smem=smem_bytes, + stream=stream, + ) + + +# Compile cache keyed on CONFIG ONLY — total_tiles/T_total/S are dynamic +# cutlass.Int32, so one compiled kernel serves every batch shape. +_prescan_kernel_cache: dict = {} + + +def _compile_pre_scan(H, v_is_varlen): + sym_t = cute.sym_int() # T_total (v token extent) + sym_bt = cute.sym_int() # beta/ws_beta (total_tiles*CHUNK*H) + sym_qk = cute.sym_int() # ws_kd/kr (total_tiles*H*CHUNK*D) + sym_gt = cute.sym_int() # ws_gt (total_tiles*H*D) + sym_cc = cute.sym_int() # ws_inv (total_tiles*H*CHUNK*CHUNK) + sym_seg = cute.sym_int() # seg_cu_tiles (S+1) + sym_so = cute.sym_int() # seg_order (S) — own sym: length differs from seg_cu_tiles + sym_st = cute.sym_int() # b_flat / m_flat (S*H*D*D) + sym_vs = cute.sym_int() # v_tile_starts + sym_va = cute.sym_int() # v_tile_actual_lens + + v_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_t, H, D), stride_order=(2, 1, 0), assumed_align=16) + beta_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_bt,), assumed_align=16) + ws_kd_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_qk,), assumed_align=16) + ws_kr_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_qk,), assumed_align=16) + ws_gt_fake = make_fake_compact_tensor(cutlass.Float32, (sym_gt,), assumed_align=16) + ws_inv_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_cc,), assumed_align=16) + seg_fake = make_fake_compact_tensor(cutlass.Int32, (sym_seg,), assumed_align=4) + so_fake = make_fake_compact_tensor(cutlass.Int32, (sym_so,), assumed_align=4) + b_fake = make_fake_compact_tensor(cutlass.Float32, (sym_st,), assumed_align=16) + m_fake = make_fake_compact_tensor(cutlass.Float32, (sym_st,), assumed_align=16) + vts_fake = make_fake_compact_tensor(cutlass.Int32, (sym_vs,), assumed_align=4) + vtal_fake = make_fake_compact_tensor(cutlass.Int32, (sym_va,), assumed_align=4) + stream_fake = make_fake_stream() + + return cute.compile( + run_pre_scan, + v_fake, + beta_fake, + ws_kd_fake, + ws_kr_fake, + ws_gt_fake, + ws_inv_fake, + seg_fake, + so_fake, + b_fake, + m_fake, + vts_fake, + vtal_fake, + H, # Constexpr -> baked + cutlass.Int32(1), # total_tiles -> runtime (placeholder) + cutlass.Int32(1), # T_total -> runtime + v_is_varlen, # Constexpr + cutlass.Int32(1), # S -> runtime + stream_fake, + options="--enable-tvm-ffi", + ) + + +def _get_compiled_pre_scan(H, v_is_varlen): + key = (H, v_is_varlen) + cached = _prescan_kernel_cache.get(key) + if cached is None: + cached = _compile_pre_scan(H, v_is_varlen) + _prescan_kernel_cache[key] = cached + return cached + + +def launch_pre_scan( + v: torch.Tensor, + beta: torch.Tensor, + ws_kd: torch.Tensor, + ws_kr: torch.Tensor, + ws_gt: torch.Tensor, + ws_inv: torch.Tensor, + b_state: torch.Tensor, # [S, H, D, D] fp32, written (bhvk) + m_state: torch.Tensor, # [S, H, D, D] fp32, written (bhvk) + seg_cu_tiles: torch.Tensor, # int32 [S+1], global tile prefix sum + v_tile_starts: torch.Tensor | None = None, # per-tile packed offset (native varlen) + v_tile_actual_lens: torch.Tensor | None = None, # per-tile valid rows (partial-tile mask) + total_tiles: int | None = None, # ceil tile sum (varlen); None -> T_total // CHUNK + seg_order: torch.Tensor | None = None, # int32 [S] launch order; None = identity +) -> None: + """Launch fused pre_scan: B_seg and M_seg for all segments.""" + assert v.is_cuda and v.dtype == torch.bfloat16 and v.is_contiguous() + B, T, H, K = v.shape + assert K == D + T_total = B * T + v_is_varlen = v_tile_starts is not None + if not v_is_varlen: + dummy = _get_dummy_int32(v.device) + v_tile_starts = dummy + v_tile_actual_lens = dummy + if total_tiles is None: + total_tiles = T_total // CHUNK + S_total = seg_cu_tiles.numel() - 1 + if seg_order is None: + seg_order = _get_identity_order(S_total, v.device) + + compiled_fn = _get_compiled_pre_scan(H, v_is_varlen) + stream = _get_current_custream() + # TMA descriptors are (re)built inside run_pre_scan from the dynamic Int32 + # dims every launch, so one compiled kernel serves all shapes. + compiled_fn( + v.view(T_total, H, D), + beta, + ws_kd, + ws_kr, + ws_gt, + ws_inv, + seg_cu_tiles, + seg_order, + b_state.reshape(-1), + m_state.reshape(-1), + v_tile_starts, + v_tile_actual_lens, + total_tiles, + T_total, + S_total, + stream, + ) diff --git a/cula/ops/kda/sm90/fwd.py b/cula/ops/kda/sm90/fwd.py new file mode 100644 index 00000000..8d958ff8 --- /dev/null +++ b/cula/ops/kda/sm90/fwd.py @@ -0,0 +1,599 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright (c) 2026 MoonshotAI +# Licensed under the MIT License. +# Based on MoonshotAI/FlashKDA (https://github.com/MoonshotAI/FlashKDA) + +""" +FlashKDA Prefill + +two-kernel (K1 Prepare + K2 Recurrence), CHUNK=16, D=128. +""" + +from __future__ import annotations + +import os +import weakref +from contextlib import contextmanager +from dataclasses import dataclass + +import torch + +CHUNK: int = 16 +D: int = 128 # only 128 supported + +_CUTE_ARCH_BY_CC = {(9, 0): "sm_90a", (10, 0): "sm_100a", (10, 3): "sm_103a"} +_VARLEN_LAYOUT_CACHE_MAXSIZE = 64 + + +def _compute_total_tiles(seq_lens: list[int] | tuple[int, ...]) -> int: + return sum((sl + CHUNK - 1) // CHUNK for sl in seq_lens) + + +@dataclass +class _VarlenMetadata: + cu_values: tuple[int, ...] + seq_lens: tuple[int, ...] + total_tiles: int + needs_padding: bool + total_aligned: int + cu_tiles: torch.Tensor | None + tile_starts: torch.Tensor + tile_actual_lens: torch.Tensor + + +@dataclass +class _PrefillProblem: + B: int + T: int + H: int + N: int # number of sequences (=B for fixed-len, =len(cu_seqlens)-1 for varlen) + total_tiles: int + is_varlen: bool + has_state_in: bool + has_state_out: bool + varlen_meta: _VarlenMetadata | None = None + + +def _validate_inputs( + q, k, v, g, beta, A_log, dt_bias, initial_state, final_state, cu_seqlens, cu_seqlens_cpu=None +) -> _PrefillProblem: + if q.ndim != 4: + raise ValueError(f"q must have shape [B, T, H, D], got {tuple(q.shape)}") + if not q.is_cuda or q.dtype != torch.bfloat16: + raise TypeError(f"q must be a CUDA bfloat16 tensor, got dtype={q.dtype}, device={q.device}") + for name, tensor in (("k", k), ("v", v), ("g", g), ("beta", beta)): + if not tensor.is_cuda or tensor.dtype != torch.bfloat16: + raise TypeError(f"{name} must be a CUDA bfloat16 tensor, got dtype={tensor.dtype}, device={tensor.device}") + if q.shape != k.shape or q.shape != g.shape: + raise ValueError(f"q/k/g shapes must match, got q={tuple(q.shape)}, k={tuple(k.shape)}, g={tuple(g.shape)}") + if v.shape != q.shape: + raise ValueError(f"v shape {tuple(v.shape)} must match q shape {tuple(q.shape)}") + + B, T, H, K = q.shape + if B <= 0 or T <= 0 or H <= 0: + raise ValueError(f"B, T and H must be positive, got B={B}, T={T}, H={H}") + if K != D or v.shape[-1] != D: + raise ValueError(f"only K=V={D} supported, got K={K} V={v.shape[-1]}") + if beta.shape != (B, T, H): + raise ValueError(f"beta shape mismatch: {tuple(beta.shape)} vs ({B},{T},{H})") + if A_log is None or not A_log.is_cuda or not A_log.is_contiguous() or A_log.shape != (H,) or A_log.dtype != torch.float32: + raise ValueError( + f"A_log must be float32 with shape ({H},), got {None if A_log is None else (A_log.dtype, tuple(A_log.shape))}" + ) + if ( + dt_bias is None + or not dt_bias.is_cuda + or not dt_bias.is_contiguous() + or dt_bias.shape != (H, K) + or dt_bias.dtype != torch.float32 + ): + raise ValueError( + f"dt_bias must be float32 with shape ({H}, {K}), " + f"got {None if dt_bias is None else (dt_bias.dtype, tuple(dt_bias.shape))}" + ) + + is_varlen = cu_seqlens is not None + if is_varlen: + if B != 1: + raise ValueError(f"varlen requires B=1, got B={B}") + if not cu_seqlens.is_cuda or cu_seqlens.ndim != 1: + raise ValueError("cu_seqlens must be a 1D CUDA tensor") + if cu_seqlens.dtype != torch.int32: + raise TypeError(f"cu_seqlens must be int32, got {cu_seqlens.dtype}") + if cu_seqlens.numel() < 2: + raise ValueError("cu_seqlens must contain at least two entries") + if cu_seqlens_cpu is not None and ( + cu_seqlens_cpu.device.type != "cpu" or cu_seqlens_cpu.ndim != 1 or cu_seqlens_cpu.numel() != cu_seqlens.numel() + ): + raise ValueError( + "cu_seqlens_cpu must be a 1D CPU tensor with the same numel as " + f"cu_seqlens ({cu_seqlens.numel()}), got device={cu_seqlens_cpu.device}, " + f"shape={tuple(cu_seqlens_cpu.shape)}" + ) + varlen_meta = _get_or_build_varlen_metadata(cu_seqlens, cu_seqlens_cpu) + N = len(varlen_meta.seq_lens) + if varlen_meta.cu_values[0] != 0: + raise ValueError("cu_seqlens must start at 0") + if varlen_meta.cu_values[-1] != T: + raise ValueError(f"cu_seqlens[-1] must equal packed T={T}, got {varlen_meta.cu_values[-1]}") + seq_lens = varlen_meta.seq_lens + if any(sl <= 0 for sl in seq_lens): + raise ValueError(f"all variable-length sequences must be non-empty, got seq_lens={seq_lens}") + total_tiles = varlen_meta.total_tiles + else: + N = B + total_tiles = B * ((T + CHUNK - 1) // CHUNK) + varlen_meta = None + + has_state_in = initial_state is not None + has_state_out = final_state is not None + if has_state_in: + if initial_state.shape != (N, H, D, D): + raise ValueError(f"initial_state shape must be ({N}, {H}, {D}, {D}), got {tuple(initial_state.shape)}") + if not initial_state.is_cuda or initial_state.dtype != torch.float32 or not initial_state.is_contiguous(): + raise TypeError("initial_state must be a contiguous CUDA float32 tensor") + if has_state_out: + if final_state.shape != (N, H, D, D): + raise ValueError(f"final_state shape must be ({N}, {H}, {D}, {D}), got {tuple(final_state.shape)}") + if not final_state.is_cuda or final_state.dtype != torch.float32 or not final_state.is_contiguous(): + raise TypeError("final_state must be a contiguous CUDA float32 tensor") + + return _PrefillProblem( + B=B, + T=T, + H=H, + N=N, + total_tiles=total_tiles, + is_varlen=is_varlen, + has_state_in=has_state_in, + has_state_out=has_state_out, + varlen_meta=varlen_meta, + ) + + +_DEVICE_ARCH_CACHE: dict[int, str] = {} + + +@contextmanager +def _cute_arch_for_device(device: torch.device): + """Ensure CUTE_DSL_ARCH matches the device before any lazy cute.compile. + + Cached + check-and-set (no pop): compiles only happen on the first call + per kernel config, so re-writing the env var on every dispatch was pure + overhead. Leaving it set is safe — other arches' dispatch paths perform + the same check-and-set before their own compiles. + """ + idx = device.index if device.index is not None else torch.cuda.current_device() + arch = _DEVICE_ARCH_CACHE.get(idx) + if arch is None: + major, minor = torch.cuda.get_device_capability(device) + arch = _CUTE_ARCH_BY_CC.get((major, minor)) + if arch is None: + raise RuntimeError(f"unsupported compute capability sm_{major}{minor}") + _DEVICE_ARCH_CACHE[idx] = arch + if os.environ.get("CUTE_DSL_ARCH") != arch: + os.environ["CUTE_DSL_ARCH"] = arch + yield + + +_VARLEN_LAYOUT_CACHE: dict = {} +_VARLEN_METADATA_CACHE: dict[int, tuple[weakref.ReferenceType[torch.Tensor], tuple, _VarlenMetadata]] = {} +_K1_SYMBOLS = None +_K2_LAUNCHER = None + + +_WS_ARENA_ALIGN = 256 +_WS_ARENA: dict = {} # (device, stream_ptr) -> [arena uint8 tensor, {sizes_key: views}] +_WS_ARENA_MAXSIZE = 8 +_WS_VIEWS_MAXSIZE = 32 + + +def _get_or_alloc_workspaces(n_qk: int, n_cc: int, n_gt: int, n_beta: int, device, dtype): + """Carve K1/K2 scratch (ws_qd/kd/kr/gt/inv/mqk, ws_beta) out of a grow-only + per-(device, stream) arena instead of allocating per call. + + Reusing the arena is safe because every producer/consumer runs on the + keyed stream: the next call's K1 cannot overwrite a workspace before this + call's K2 finished reading it. + """ + stream_ptr = int(torch.cuda.current_stream(device).cuda_stream) + arena_key = (str(device), stream_ptr) + sizes_key = (n_qk, n_cc, n_gt, n_beta, dtype) + entry = _WS_ARENA.get(arena_key) + if entry is not None: + views = entry[1].get(sizes_key) + if views is not None: + return views + + nbytes_list = ( + n_qk * 2, # ws_qd bf16 + n_qk * 2, # ws_kd + n_qk * 2, # ws_kr + n_gt * 4, # ws_gt fp32 + n_cc * 2, # ws_inv bf16 + n_cc * 2, # ws_mqk + n_beta * dtype.itemsize, # ws_beta + ) + offsets = [] + total = 0 + for nbytes in nbytes_list: + offsets.append(total) + total += -(-nbytes // _WS_ARENA_ALIGN) * _WS_ARENA_ALIGN + + if entry is None or entry[0].numel() < total: + if entry is None and len(_WS_ARENA) >= _WS_ARENA_MAXSIZE: + _WS_ARENA.pop(next(iter(_WS_ARENA))) + # Growing replaces the arena; stale views die with the old entry. + entry = [torch.empty(total, dtype=torch.uint8, device=device), {}] + _WS_ARENA[arena_key] = entry + arena = entry[0] + + def carve(idx: int, numel: int, view_dtype: torch.dtype): + return arena.narrow(0, offsets[idx], numel * view_dtype.itemsize).view(view_dtype) + + views = ( + carve(0, n_qk, torch.bfloat16), + carve(1, n_qk, torch.bfloat16), + carve(2, n_qk, torch.bfloat16), + carve(3, n_gt, torch.float32), + carve(4, n_cc, torch.bfloat16), + carve(5, n_cc, torch.bfloat16), + carve(6, n_beta, dtype), + ) + if len(entry[1]) >= _WS_VIEWS_MAXSIZE: + entry[1].pop(next(iter(entry[1]))) + entry[1][sizes_key] = views + return views + + +def _get_or_build_varlen_layout(seq_lens: tuple[int, ...], device, cu_dtype): + """CHUNK-aligned cumulative token offsets and tile counts for non-aligned varlen.""" + key = (seq_lens, str(device), cu_dtype) + cached = _VARLEN_LAYOUT_CACHE.get(key) + if cached is not None: + return cached + + out_offsets = [0] + for sl in seq_lens: + aligned = ((sl + CHUNK - 1) // CHUNK) * CHUNK + out_offsets.append(out_offsets[-1] + aligned) + + cu_pad = torch.tensor(out_offsets, dtype=cu_dtype, device=device) + cu_tiles = torch.tensor([off // CHUNK for off in out_offsets], dtype=torch.int32, device=device) + cached = (cu_pad, cu_tiles) + if len(_VARLEN_LAYOUT_CACHE) >= _VARLEN_LAYOUT_CACHE_MAXSIZE: + _VARLEN_LAYOUT_CACHE.pop(next(iter(_VARLEN_LAYOUT_CACHE))) + _VARLEN_LAYOUT_CACHE[key] = cached + return cached + + +def _get_or_build_varlen_metadata(cu_seqlens: torch.Tensor, cu_seqlens_cpu: torch.Tensor | None = None) -> _VarlenMetadata: + """Cache varlen metadata (seq_lens, tile offsets, padding flags) for cu_seqlens.""" + cache_key = id(cu_seqlens) + attrs = ( + cu_seqlens.data_ptr(), + tuple(cu_seqlens.shape), + str(cu_seqlens.device), + cu_seqlens.dtype, + int(cu_seqlens._version), + ) + cached = _VARLEN_METADATA_CACHE.get(cache_key) + if cached is not None: + tensor_ref, cached_attrs, meta = cached + if tensor_ref() is cu_seqlens and cached_attrs == attrs: + return meta + _VARLEN_METADATA_CACHE.pop(cache_key, None) + + src_cpu = cu_seqlens_cpu if cu_seqlens_cpu is not None else cu_seqlens.detach().to("cpu") + cu_values = tuple(int(v) for v in src_cpu.tolist()) + seq_lens = tuple(cu_values[i + 1] - cu_values[i] for i in range(len(cu_values) - 1)) + total_tiles = _compute_total_tiles(seq_lens) + needs_padding = any((sl % CHUNK) != 0 for sl in seq_lens) + aligned_lens = tuple(((sl + CHUNK - 1) // CHUNK) * CHUNK for sl in seq_lens) + total_aligned = sum(aligned_lens) + tile_starts_list: list[int] = [] + tile_actual_lens_list: list[int] = [] + for bos, sl in zip(cu_values[:-1], seq_lens): + for offset in range(0, sl, CHUNK): + tile_starts_list.append(bos + offset) + tile_actual_lens_list.append(min(CHUNK, sl - offset)) + tile_starts = torch.tensor(tile_starts_list, dtype=torch.int32, device=cu_seqlens.device) + tile_actual_lens = torch.tensor(tile_actual_lens_list, dtype=torch.int32, device=cu_seqlens.device) + cu_tiles = None + if not needs_padding: + cu_tiles = torch.tensor( + [v // CHUNK for v in cu_values], + dtype=torch.int32, + device=cu_seqlens.device, + ) + + meta = _VarlenMetadata( + cu_values=cu_values, + seq_lens=seq_lens, + total_tiles=total_tiles, + needs_padding=needs_padding, + total_aligned=total_aligned, + cu_tiles=cu_tiles, + tile_starts=tile_starts, + tile_actual_lens=tile_actual_lens, + ) + if len(_VARLEN_METADATA_CACHE) >= _VARLEN_LAYOUT_CACHE_MAXSIZE: + for k, (ref, _a, _m) in list(_VARLEN_METADATA_CACHE.items()): + if ref() is None: + _VARLEN_METADATA_CACHE.pop(k, None) + if len(_VARLEN_METADATA_CACHE) >= _VARLEN_LAYOUT_CACHE_MAXSIZE: + _VARLEN_METADATA_CACHE.pop(next(iter(_VARLEN_METADATA_CACHE))) + _VARLEN_METADATA_CACHE[cache_key] = (weakref.ref(cu_seqlens), attrs, meta) + return meta + + +def _get_k1_symbols(): + global _K1_SYMBOLS + if _K1_SYMBOLS is None: + from cula.ops.kda.sm90.k1 import CHUNK as k1_chunk + from cula.ops.kda.sm90.k1 import D as k1_d + from cula.ops.kda.sm90.k1 import launch_k1 as k1_launch + + _K1_SYMBOLS = (k1_chunk, k1_d, k1_launch) + return _K1_SYMBOLS + + +def _get_k2_launcher(): + global _K2_LAUNCHER + if _K2_LAUNCHER is not None: + return _K2_LAUNCHER + from cula.ops.kda.sm90.k2 import launch_k2 + + _K2_LAUNCHER = launch_k2 + return launch_k2 + + +def flash_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + scale: float, + out: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + lower_bound: float, + initial_state: torch.Tensor | None = None, + final_state: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + cu_seqlens_cpu: torch.Tensor | None = None, + state_transposed: bool = False, + use_gate_in_kernel: bool = True, +) -> None: + """FlashKDA fwd. ``out`` and ``final_state`` are written in-place. + + Args: + q, k, v, g: [B, T, H, D] bf16. + beta: [B, T, H] bf16 (pre-sigmoid). + scale: attention scale. + out: [B, T, H, D] bf16 output (written in-place). + A_log: [H] fp32. + dt_bias: [H, D] fp32. + lower_bound: gate floor (negative). + initial_state: [N, H, D, D] bf16/fp32 or None. + final_state: [N, H, D, D] bf16/fp32 or None (written in-place). + cu_seqlens: [N+1] int32/int64 for variable-length, or None. + cu_seqlens_cpu: optional CPU copy of cu_seqlens (same values) to skip the + GPU->host sync when first building varlen metadata. + state_transposed: False -> [N,H,V,K] (default), True -> [N,H,K,V]. + """ + problem = _validate_inputs(q, k, v, g, beta, A_log, dt_bias, initial_state, final_state, cu_seqlens, cu_seqlens_cpu) + if out.shape != q.shape or not out.is_cuda or out.dtype != torch.bfloat16: + raise ValueError( + f"out must be CUDA bfloat16 with shape {tuple(q.shape)}, got dtype={out.dtype}, shape={tuple(out.shape)}" + ) + if not use_gate_in_kernel: + raise NotImplementedError( + "CuTeDSL FlashKDA prefill only supports use_gate_in_kernel=True. " + "Pre-gated inputs would require the torch reference, which is test-only." + ) + if lower_bound is None: + raise ValueError("lower_bound must be specified.") + if not (-5 <= lower_bound < 0): + raise ValueError(f"lower_bound must be in the safe range [-5, 0), got {lower_bound}.") + + with _cute_arch_for_device(q.device): + _dispatch_cute( + q, + k, + v, + g, + beta, + scale, + out, + A_log, + dt_bias, + lower_bound, + initial_state, + final_state, + cu_seqlens, + problem, + state_transposed=state_transposed, + ) + + +def _dispatch_cute( + q, + k, + v, + g, + beta, + scale, + out, + A_log, + dt_bias, + lower_bound, + initial_state, + final_state, + cu_seqlens, + problem: _PrefillProblem, + *, + state_transposed: bool = False, +): + """Launch K1 + K2.""" + K1_CHUNK, K1_D, launch_k1 = _get_k1_symbols() + + T_orig = problem.T + need_t_pad = (not problem.is_varlen) and (T_orig % K1_CHUNK != 0) + if need_t_pad: + T_pad = ((T_orig + K1_CHUNK - 1) // K1_CHUNK) * K1_CHUNK + B, H = problem.B, problem.H + pad_len = T_pad - T_orig + q = torch.nn.functional.pad(q, (0, 0, 0, 0, 0, pad_len)) + k = torch.nn.functional.pad(k, (0, 0, 0, 0, 0, pad_len)) + v = torch.nn.functional.pad(v, (0, 0, 0, 0, 0, pad_len)) + g = torch.nn.functional.pad(g, (0, 0, 0, 0, 0, pad_len), value=-1e6) + beta = torch.nn.functional.pad(beta, (0, 0, 0, pad_len), value=-80.0) + out_orig = out + out = torch.empty_like(q) + problem = _PrefillProblem( + B=B, + T=T_pad, + H=H, + N=problem.N, + total_tiles=B * (T_pad // K1_CHUNK), + is_varlen=False, + has_state_in=problem.has_state_in, + has_state_out=problem.has_state_out, + ) + + k1_q, k1_k, k1_g, k1_beta = q, k, g, beta + k1_total_tiles = problem.total_tiles + k1_tile_starts = None + k1_tile_actual_lens = None + k1_is_varlen = False + + # Varlen: K1/K2 read original q/k/g/v; beta remains padded for the + # existing compact workspace layout. + k2_cu_seqlens_tiles_cached = None + k2_v_tile_starts = None + k2_v_tile_actual_lens = None + if problem.is_varlen: + varlen_meta = problem.varlen_meta + seq_lens_list = varlen_meta.seq_lens + if varlen_meta.needs_padding: + total_aligned = varlen_meta.total_aligned + + k1_q = q.contiguous() + k1_k = k.contiguous() + k1_g = g.contiguous() + k1_beta = beta.contiguous() + k1_total_tiles = varlen_meta.total_tiles + k1_tile_starts = varlen_meta.tile_starts + k1_tile_actual_lens = varlen_meta.tile_actual_lens + k1_is_varlen = True + k2_v_tile_starts = varlen_meta.tile_starts + k2_v_tile_actual_lens = varlen_meta.tile_actual_lens + + # Padded tile boundaries for K2's per-sequence recurrence. K1 emits ws_beta + # directly, so varlen needs no host-side beta padding/gather. + cu_pad, k2_cu_seqlens_tiles_cached = _get_or_build_varlen_layout( + tuple(seq_lens_list), + q.device, + cu_seqlens.dtype, + ) + + problem_pad = _PrefillProblem( + B=1, + T=total_aligned, + H=problem.H, + N=problem.N, + total_tiles=total_aligned // K1_CHUNK, + is_varlen=True, + has_state_in=problem.has_state_in, + has_state_out=problem.has_state_out, + ) + cu_seqlens, problem = cu_pad, problem_pad + else: + k2_cu_seqlens_tiles_cached = varlen_meta.cu_tiles + + _launch_k2 = _get_k2_launcher() + + B, T, H = problem.B, problem.T, problem.H + + if problem.is_varlen: + T_total = T + k2_cu_seqlens_tiles = k2_cu_seqlens_tiles_cached + else: + T_total = B * T + k2_cu_seqlens_tiles = None + + total_tiles = T_total // K1_CHUNK + + n_qk = total_tiles * H * K1_CHUNK * K1_D + n_cc = total_tiles * H * K1_CHUNK * K1_CHUNK + ws_qd, ws_kd, ws_kr, ws_gt, ws_inv, ws_mqk, ws_beta = _get_or_alloc_workspaces( + n_qk, n_cc, total_tiles * H * K1_D, T_total * H, q.device, beta.dtype + ) + + k2_initial_state = None + if problem.has_state_in: + k2_initial_state = initial_state.contiguous() + + k2_final_state = None + if problem.has_state_out: + k2_final_state = final_state + + # K1 reads beta from its original packed [T, H] layout and emits raw beta + # into ws_beta (tail rows = -80); K2 reads ws_beta directly, so no + # host-side transpose/padding/gather of beta is needed. + launch_k1( + k1_q, + k1_k, + k1_g, + A_log, + dt_bias, + k1_beta.reshape(-1), + scale, + lower_bound, + ws_qd, + ws_kd, + ws_kr, + ws_gt, + ws_inv, + ws_mqk, + ws_beta, + tile_starts=k1_tile_starts, + tile_actual_lens=k1_tile_actual_lens, + total_tiles=k1_total_tiles, + is_varlen=k1_is_varlen, + ) + _launch_k2( + v, + ws_beta, + ws_qd, + ws_kd, + ws_kr, + ws_gt, + ws_inv, + ws_mqk, + out, + k2_cu_seqlens_tiles, + initial_state=k2_initial_state, + final_state=k2_final_state, + state_transposed=state_transposed, + v_tile_starts=k2_v_tile_starts, + v_tile_actual_lens=k2_v_tile_actual_lens, + ) + + if need_t_pad: + out_orig.copy_(out[:, :T_orig]) diff --git a/cula/ops/kda/sm90/k1.py b/cula/ops/kda/sm90/k1.py new file mode 100644 index 00000000..63d8a41c --- /dev/null +++ b/cula/ops/kda/sm90/k1.py @@ -0,0 +1,824 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright (c) 2026 MoonshotAI +# Licensed under the MIT License. +# Based on MoonshotAI/FlashKDA (https://github.com/MoonshotAI/FlashKDA) + +""" +CuteDSL port of FlashKDA K1. + +Produces 6 workspace tensors consumed by K2: ws_qd, ws_kd, ws_kr, ws_gt, ws_inv, ws_mqk. +""" + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import torch +from cutlass.cute.nvgpu import cpasync, warp +from cutlass.cute.nvgpu.warpgroup import SmemLayoutAtomKind, make_smem_layout_atom +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream + +from cula.ops.kda.sm90._common import add_f16x2_u32, movm_t_b16 + +CHUNK: int = 16 +D: int = 128 +THREADS_PER_CTA: int = 256 + + +@cute.kernel +def k1_kernel( + tma_atom_q: cute.CopyAtom, + tma_tensor_q: cute.Tensor, + tma_atom_k: cute.CopyAtom, + tma_tensor_k: cute.Tensor, + tma_atom_g: cute.CopyAtom, + tma_tensor_g: cute.Tensor, + tma_atom_ws_qd: cute.CopyAtom, + tma_tensor_ws_qd: cute.Tensor, + tma_atom_ws_kd: cute.CopyAtom, + tma_tensor_ws_kd: cute.Tensor, + tma_atom_ws_kr: cute.CopyAtom, + tma_tensor_ws_kr: cute.Tensor, + tma_atom_ws_inv: cute.CopyAtom, + tma_tensor_ws_inv: cute.Tensor, + tma_atom_ws_mqk: cute.CopyAtom, + tma_tensor_ws_mqk: cute.Tensor, + a_log: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + ws_gt: cute.Tensor, + ws_inv: cute.Tensor, + ws_mqk: cute.Tensor, + ws_beta: cute.Tensor, + tile_starts: cute.Tensor, + tile_actual_lens: cute.Tensor, + H: cutlass.Constexpr[int], + total_tiles: cutlass.Int32, + T_total: cutlass.Int32, + scale: cutlass.Constexpr[float], + gate_scale: cutlass.Constexpr[float], + is_varlen: cutlass.Constexpr[bool], +): + tile_idx, head_idx, _ = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + + smem = cutlass.utils.SmemAllocator() + qk_layout = cute.make_layout((CHUNK, D), stride=(D, 1)) + kinter_atom_qk = make_smem_layout_atom(SmemLayoutAtomKind.K_INTER, cutlass.BFloat16) + kinter_qk_layout = cute.tile_to_shape(kinter_atom_qk, (CHUNK, D), order=(0, 1)) + cc_layout = cute.make_layout((CHUNK, CHUNK), stride=(CHUNK, 1)) + + sQ = smem.allocate_tensor(cutlass.BFloat16, qk_layout, 128) + sK = smem.allocate_tensor(cutlass.BFloat16, qk_layout, 128) + sG_raw = smem.allocate_tensor(cutlass.BFloat16, qk_layout, 128) + sG_cumsum = smem.allocate_tensor(cutlass.Float32, qk_layout, 128) + s_g_total = smem.allocate_tensor(cutlass.Float32, cute.make_layout((D,)), 128) + # Outputs (K_INTER swizzled, alias input bytes after barrier) + s_k_inv = cute.make_tensor(cute.recast_ptr(sG_cumsum.iterator, dtype=cutlass.BFloat16), kinter_qk_layout) + s_q_decayed = cute.make_tensor(sQ.iterator, kinter_qk_layout) + s_k_decayed = cute.make_tensor(sK.iterator, kinter_qk_layout) + s_k_restored = cute.make_tensor(sG_raw.iterator, kinter_qk_layout) + sL_bf16 = smem.allocate_tensor(cutlass.BFloat16, cc_layout, 128) + sMqk_bf16 = smem.allocate_tensor(cutlass.BFloat16, cc_layout, 128) + sBetaSig = smem.allocate_tensor(cutlass.Float32, cute.make_layout((CHUNK,)), 128) + # Neumann buffers (alias sG_cumsum after decay_apply barrier) + sG_cumsum_fp16_ptr = cute.recast_ptr(sG_cumsum.iterator, dtype=cutlass.Float16) + sG_cumsum_bf16_ptr = cute.recast_ptr(sG_cumsum.iterator, dtype=cutlass.BFloat16) + sL_fp16 = cute.make_tensor(sG_cumsum_fp16_ptr, cc_layout) + sINV_fp16 = cute.make_tensor(sG_cumsum_fp16_ptr + (CHUNK * CHUNK), cc_layout) + sINV_bf16 = cute.make_tensor(sG_cumsum_bf16_ptr + (2 * CHUNK * CHUNK), cc_layout) + sMbar = smem.allocate_tensor(cutlass.Int64, cute.make_layout((1,)), 8) + sMbar_ptr = sMbar.iterator + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_init(sMbar_ptr, cutlass.Int32(1)) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + + token_start = tile_idx * CHUNK + actual_len = cutlass.Int32(CHUNK) + load_tile_idx = tile_idx + tma_tensor_q_use = tma_tensor_q + tma_tensor_k_use = tma_tensor_k + tma_tensor_g_use = tma_tensor_g + if cutlass.const_expr(is_varlen): + token_start = tile_starts[tile_idx] + actual_len = tile_actual_lens[tile_idx] + load_tile_idx = cutlass.Int32(0) + tma_tensor_q_use = cute.domain_offset((token_start, 0, 0), tma_tensor_q) + tma_tensor_k_use = cute.domain_offset((token_start, 0, 0), tma_tensor_k) + tma_tensor_g_use = cute.domain_offset((token_start, 0, 0), tma_tensor_g) + + gSrc_q = cute.local_tile(tma_tensor_q_use, (CHUNK, D), (None, None, None)) + gSrc_k = cute.local_tile(tma_tensor_k_use, (CHUNK, D), (None, None, None)) + gSrc_g = cute.local_tile(tma_tensor_g_use, (CHUNK, D), (None, None, None)) + tQs, tQg = cpasync.tma_partition( + tma_atom_q, + 0, + cute.make_layout(1), + cute.group_modes(sQ, 0, 2), + cute.group_modes(gSrc_q, 0, 2), + ) + tKs, tKg = cpasync.tma_partition( + tma_atom_k, + 0, + cute.make_layout(1), + cute.group_modes(sK, 0, 2), + cute.group_modes(gSrc_k, 0, 2), + ) + tGs, tGg = cpasync.tma_partition( + tma_atom_g, + 0, + cute.make_layout(1), + cute.group_modes(sG_raw, 0, 2), + cute.group_modes(gSrc_g, 0, 2), + ) + + gDst_qd = cute.local_tile(tma_tensor_ws_qd, (CHUNK, D), (None, None, None)) + tQDws_s, tQDws_g = cpasync.tma_partition( + tma_atom_ws_qd, + 0, + cute.make_layout(1), + cute.group_modes(s_q_decayed, 0, 2), + cute.group_modes(gDst_qd, 0, 2), + ) + gDst_kd = cute.local_tile(tma_tensor_ws_kd, (CHUNK, D), (None, None, None)) + tKDws_s, tKDws_g = cpasync.tma_partition( + tma_atom_ws_kd, + 0, + cute.make_layout(1), + cute.group_modes(s_k_decayed, 0, 2), + cute.group_modes(gDst_kd, 0, 2), + ) + gDst_kr = cute.local_tile(tma_tensor_ws_kr, (CHUNK, D), (None, None, None)) + tKRws_s, tKRws_g = cpasync.tma_partition( + tma_atom_ws_kr, + 0, + cute.make_layout(1), + cute.group_modes(s_k_restored, 0, 2), + cute.group_modes(gDst_kr, 0, 2), + ) + gDst_inv = cute.local_tile(tma_tensor_ws_inv, (CHUNK, CHUNK), (None, None, None)) + tINVws_s, tINVws_g = cpasync.tma_partition( + tma_atom_ws_inv, + 0, + cute.make_layout(1), + cute.group_modes(sINV_bf16, 0, 2), + cute.group_modes(gDst_inv, 0, 2), + ) + gDst_mqk = cute.local_tile(tma_tensor_ws_mqk, (CHUNK, CHUNK), (None, None, None)) + tMQKws_s, tMQKws_g = cpasync.tma_partition( + tma_atom_ws_mqk, + 0, + cute.make_layout(1), + cute.group_modes(sMqk_bf16, 0, 2), + cute.group_modes(gDst_mqk, 0, 2), + ) + ws_slot = head_idx * total_tiles + tile_idx + + if warp_idx == 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(sMbar_ptr, cutlass.Int32(3 * CHUNK * D * 2)) + cute.copy(tma_atom_q, tQg[(None, load_tile_idx, 0, head_idx)], tQs[(None,)], tma_bar_ptr=sMbar_ptr) + cute.copy(tma_atom_k, tKg[(None, load_tile_idx, 0, head_idx)], tKs[(None,)], tma_bar_ptr=sMbar_ptr) + cute.copy(tma_atom_g, tGg[(None, load_tile_idx, 0, head_idx)], tGs[(None,)], tma_bar_ptr=sMbar_ptr) + + cute.arch.mbarrier_wait(sMbar_ptr, cutlass.Int32(0)) + + if tidx < 128: + col_tail = tidx + for r in cutlass.range_constexpr(CHUNK): + if actual_len <= cutlass.Int32(r): + sQ[r, col_tail] = cutlass.BFloat16(0.0) + sK[r, col_tail] = cutlass.BFloat16(0.0) + sG_raw[r, col_tail] = cutlass.BFloat16(0.0) + cute.arch.barrier() + + row = tidx // 16 + sQ_tile = cute.flat_divide(sQ, (1, 8)) # ((1,8), CHUNK, D//8) + sK_tile = cute.flat_divide(sK, (1, 8)) + cb = tidx % 16 + sQ_my = sQ_tile[(None, None, row, cb)] + sK_my = sK_tile[(None, None, row, cb)] + r_q_bf = cute.make_rmem_tensor(cute.make_layout((1, 8)), cutlass.BFloat16) + r_k_bf = cute.make_rmem_tensor(cute.make_layout((1, 8)), cutlass.BFloat16) + cute.autovec_copy(sQ_my, r_q_bf) + cute.autovec_copy(sK_my, r_k_bf) + q_sq = cutlass.Float32(0.0) + k_sq = cutlass.Float32(0.0) + q_vals = cute.make_rmem_tensor(cute.make_layout((8,), stride=(1,)), cutlass.Float32) + k_vals = cute.make_rmem_tensor(cute.make_layout((8,), stride=(1,)), cutlass.Float32) + for j in cutlass.range_constexpr(8): + qv = cutlass.Float32(r_q_bf[0, j]) + kv = cutlass.Float32(r_k_bf[0, j]) + q_vals[j] = qv + k_vals[j] = kv + q_sq = q_sq + qv * qv + k_sq = k_sq + kv * kv + q_sq = cute.arch.warp_reduction(q_sq, lambda a, b: a + b, threads_in_group=16) + k_sq = cute.arch.warp_reduction(k_sq, lambda a, b: a + b, threads_in_group=16) + q_inv = cute.rsqrt(q_sq + cutlass.Float32(1.0e-6), fastmath=True) + k_inv = cute.rsqrt(k_sq + cutlass.Float32(1.0e-6), fastmath=True) + for j in cutlass.range_constexpr(8): + r_q_bf[0, j] = cutlass.BFloat16(q_vals[j] * q_inv) + r_k_bf[0, j] = cutlass.BFloat16(k_vals[j] * k_inv) + cute.autovec_copy(r_q_bf, sQ_my) + cute.autovec_copy(r_k_bf, sK_my) + a_log_exp = cute.exp(cutlass.Float32(a_log[head_idx]), fastmath=True) + if tidx < 128: + col_c = tidx + dt = cutlass.Float32(dt_bias[head_idx, col_c]) + s = cutlass.Float32(0.0) + for r in cutlass.range_constexpr(CHUNK): + if actual_len > cutlass.Int32(r): + x = cutlass.Float32(sG_raw[r, col_c]) + dt + x = a_log_exp * x + sig = cutlass.Float32(0.5) * (cute.tanh(x * cutlass.Float32(0.5), fastmath=True) + cutlass.Float32(1.0)) + s = s + cutlass.Float32(gate_scale) * sig + sG_cumsum[r, col_c] = s + s_g_total[col_c] = cute.exp(s, fastmath=True) + cute.arch.barrier() + + # Pre-compute per-row sigmoid(beta). beta is the original packed [T, H] + # layout, so no host-side transpose is needed. + if tidx < CHUNK: + bv = cutlass.Float32(-80.0) + if actual_len > tidx: + bv = cutlass.Float32(beta[(token_start + tidx) * H + head_idx]) + sBetaSig[tidx] = cutlass.Float32(0.5) * (cute.tanh(bv * cutlass.Float32(0.5), fastmath=True) + cutlass.Float32(1.0)) + # Emit raw beta into the compact wt_l workspace (tail rows = -80 -> sigmoid~0). + ws_beta[(head_idx * total_tiles + tile_idx) * CHUNK + tidx] = cutlass.BFloat16(bv) + + lane = tidx % 32 + warp_id = tidx // 32 + group = lane // 4 + t_in_group = lane % 4 + N_M: cutlass.Constexpr[int] = CHUNK // 8 # = 2 + N_N: cutlass.Constexpr[int] = D // 64 # = 2 + N_TILES: cutlass.Constexpr[int] = N_M * N_N # = 4 + + reg_g = cute.make_rmem_tensor(cute.make_layout((N_TILES, 2)), cutlass.Float32) + reg_q = cute.make_rmem_tensor(cute.make_layout((N_TILES, 2)), cutlass.BFloat16) + reg_k = cute.make_rmem_tensor(cute.make_layout((N_TILES, 2)), cutlass.BFloat16) + reg_gt = cute.make_rmem_tensor(cute.make_layout((N_TILES, 2)), cutlass.Float32) + sG_cumsum_zipped = cute.zipped_divide(sG_cumsum, (1, 2)) + sQ_zipped = cute.zipped_divide(sQ, (1, 2)) + sK_zipped = cute.zipped_divide(sK, (1, 2)) + s_g_total_zipped = cute.zipped_divide(s_g_total, (2,)) + for m_blk in cutlass.range_constexpr(0, CHUNK, 8): + for n_blk in cutlass.range_constexpr(0, D, 64): + tile_idx_d: cutlass.Constexpr[int] = (m_blk // 8) * N_N + (n_blk // 64) + row = m_blk + ((warp_id + group) % 8) + col = n_blk + group * 8 + t_in_group * 2 + cute.autovec_copy(sG_cumsum_zipped[None, (row, col // 2)], reg_g[tile_idx_d, None]) + cute.autovec_copy(sQ_zipped[None, (row, col // 2)], reg_q[tile_idx_d, None]) + cute.autovec_copy(sK_zipped[None, (row, col // 2)], reg_k[tile_idx_d, None]) + cute.autovec_copy(s_g_total_zipped[None, col // 2], reg_gt[tile_idx_d, None]) + + cute.arch.barrier() + + r_qd_pack = cute.make_rmem_tensor(cute.make_layout((1, 2)), cutlass.BFloat16) + r_kd_pack = cute.make_rmem_tensor(cute.make_layout((1, 2)), cutlass.BFloat16) + r_ki_pack = cute.make_rmem_tensor(cute.make_layout((1, 2)), cutlass.BFloat16) + r_kr_pack = cute.make_rmem_tensor(cute.make_layout((1, 2)), cutlass.BFloat16) + s_q_decayed_zipped = cute.zipped_divide(s_q_decayed, (1, 2)) + s_k_decayed_zipped = cute.zipped_divide(s_k_decayed, (1, 2)) + s_k_inv_zipped = cute.zipped_divide(s_k_inv, (1, 2)) + s_k_restored_zipped = cute.zipped_divide(s_k_restored, (1, 2)) + for m_blk in cutlass.range_constexpr(0, CHUNK, 8): + for n_blk in cutlass.range_constexpr(0, D, 64): + tile_idx_d: cutlass.Constexpr[int] = (m_blk // 8) * N_N + (n_blk // 64) + row = m_blk + ((warp_id + group) % 8) + col = n_blk + group * 8 + t_in_group * 2 + for v in cutlass.range_constexpr(2): + vv: cutlass.Constexpr[int] = v + gv = reg_g[tile_idx_d, vv] + qv = cutlass.Float32(reg_q[tile_idx_d, vv]) + kv = cutlass.Float32(reg_k[tile_idx_d, vv]) + gtv = reg_gt[tile_idx_d, vv] + exp_pos = cute.exp(gv, fastmath=True) + inv_pos = cutlass.Float32(1.0) / exp_pos + r_qd_pack[0, vv] = cutlass.BFloat16(qv * exp_pos * cutlass.Float32(scale)) + r_kd_pack[0, vv] = cutlass.BFloat16(kv * exp_pos) + r_ki_pack[0, vv] = cutlass.BFloat16(kv * inv_pos) + r_kr_pack[0, vv] = cutlass.BFloat16(kv * gtv * inv_pos) + cute.autovec_copy(r_qd_pack, s_q_decayed_zipped[None, (row, col // 2)]) + cute.autovec_copy(r_kd_pack, s_k_decayed_zipped[None, (row, col // 2)]) + cute.autovec_copy(r_ki_pack, s_k_inv_zipped[None, (row, col // 2)]) + cute.autovec_copy(r_kr_pack, s_k_restored_zipped[None, (row, col // 2)]) + + if tidx < 128: + gt_base = (head_idx * total_tiles + tile_idx) * D + ws_gt[gt_base + tidx] = s_g_total[tidx] + cute.arch.barrier() + + # All 5 TMA stores must come from one thread (cp.async.bulk groups are per-thread). + + mma_atom_mask_mma = warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16)) + tiled_mma_mask_mma = cute.make_tiled_mma( + mma_atom_mask_mma, + atom_layout_mnk=(1, 1, 1), + permutation_mnk=(16, 16, 16), + ) + warp_lane = tidx % 32 + thr_mma_mask_mma = tiled_mma_mask_mma.get_slice(warp_lane) + + copy_atom_AB_mask_mma = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), + cutlass.BFloat16, + ) + smem_tiled_copy_A_mask_mma = cute.make_tiled_copy_A(copy_atom_AB_mask_mma, tiled_mma_mask_mma) + smem_tiled_copy_B_mask_mma = cute.make_tiled_copy_B(copy_atom_AB_mask_mma, tiled_mma_mask_mma) + smem_thr_copy_A_mask_mma = smem_tiled_copy_A_mask_mma.get_slice(warp_lane) + smem_thr_copy_B_mask_mma = smem_tiled_copy_B_mask_mma.get_slice(warp_lane) + + copy_atom_stsm_mask_mma = cute.make_copy_atom( + warp.StMatrix8x8x16bOp(transpose=False, num_matrices=2), + cutlass.BFloat16, + ) + smem_tiled_store_mask_mma = cute.make_tiled_copy_C_atom(copy_atom_stsm_mask_mma, tiled_mma_mask_mma) + smem_thr_store_mask_mma = smem_tiled_store_mask_mma.get_slice(warp_lane) + + sB_tile = cute.flat_divide(s_k_inv, (CHUNK, 16)) # ((16,16), 1, D//16) + sB_ref = sB_tile[None, None, 0, 0] + + if warp_idx == 0: + sA_tile_l = cute.flat_divide(s_k_decayed, (CHUNK, 16)) + sA_ref_l = sA_tile_l[None, None, 0, 0] + tCrA_l = thr_mma_mask_mma.make_fragment_A(thr_mma_mask_mma.partition_A(sA_ref_l)) + tCrB_l = thr_mma_mask_mma.make_fragment_B(thr_mma_mask_mma.partition_B(sB_ref)) + tCrC_l = thr_mma_mask_mma.make_fragment_C(tiled_mma_mask_mma.partition_shape_C((CHUNK, CHUNK))) + tCrA_l_cv = smem_thr_copy_A_mask_mma.retile(tCrA_l) + tCrB_l_cv = smem_thr_copy_B_mask_mma.retile(tCrB_l) + + tCrC_l.fill(0.0) + for k_blk in cutlass.range_constexpr(D // 16): + sA_k = sA_tile_l[None, None, 0, k_blk] + sB_k = sB_tile[None, None, 0, k_blk] + cute.copy(smem_tiled_copy_A_mask_mma, smem_thr_copy_A_mask_mma.partition_S(sA_k), tCrA_l_cv) + cute.copy(smem_tiled_copy_B_mask_mma, smem_thr_copy_B_mask_mma.partition_S(sB_k), tCrB_l_cv) + cute.gemm(tiled_mma_mask_mma, tCrC_l, tCrA_l, tCrB_l, tCrC_l) + + # L mask fold: factor = float(m > n) * sigmoid(beta[m]). + coord_Cl = cute.make_identity_tensor((CHUNK, CHUNK)) + tCcC_l = thr_mma_mask_mma.partition_C(coord_Cl) + for ii in cutlass.range_constexpr(cute.size(tCrC_l)): + crd = tCcC_l[ii] + m = crd[0] + n = crd[1] + keep = cutlass.Float32(1.0) if m > n else cutlass.Float32(0.0) + tCrC_l[ii] = tCrC_l[ii] * keep * sBetaSig[m] + + tCrC_l_bf16 = cute.make_fragment_like(tCrC_l, cutlass.BFloat16) + for ii in cutlass.range_constexpr(cute.size(tCrC_l)): + tCrC_l_bf16[ii] = cutlass.BFloat16(tCrC_l[ii]) + cute.copy( + smem_tiled_store_mask_mma, + smem_thr_store_mask_mma.retile(tCrC_l_bf16), + smem_thr_store_mask_mma.partition_D(sL_bf16), + ) + elif warp_idx == 1: + sA_tile_m = cute.flat_divide(s_q_decayed, (CHUNK, 16)) + sA_ref_m = sA_tile_m[None, None, 0, 0] + tCrA_m = thr_mma_mask_mma.make_fragment_A(thr_mma_mask_mma.partition_A(sA_ref_m)) + tCrB_m = thr_mma_mask_mma.make_fragment_B(thr_mma_mask_mma.partition_B(sB_ref)) + tCrC_m = thr_mma_mask_mma.make_fragment_C(tiled_mma_mask_mma.partition_shape_C((CHUNK, CHUNK))) + tCrA_m_cv = smem_thr_copy_A_mask_mma.retile(tCrA_m) + tCrB_m_cv = smem_thr_copy_B_mask_mma.retile(tCrB_m) + + tCrC_m.fill(0.0) + for k_blk in cutlass.range_constexpr(D // 16): + sA_k = sA_tile_m[None, None, 0, k_blk] + sB_k = sB_tile[None, None, 0, k_blk] + cute.copy(smem_tiled_copy_A_mask_mma, smem_thr_copy_A_mask_mma.partition_S(sA_k), tCrA_m_cv) + cute.copy(smem_tiled_copy_B_mask_mma, smem_thr_copy_B_mask_mma.partition_S(sB_k), tCrB_m_cv) + cute.gemm(tiled_mma_mask_mma, tCrC_m, tCrA_m, tCrB_m, tCrC_m) + + # Mqk mask fold: factor = float(m >= n). + coord_Cm = cute.make_identity_tensor((CHUNK, CHUNK)) + tCcC_m = thr_mma_mask_mma.partition_C(coord_Cm) + for ii in cutlass.range_constexpr(cute.size(tCrC_m)): + crd = tCcC_m[ii] + m = crd[0] + n = crd[1] + keep = cutlass.Float32(1.0) if m >= n else cutlass.Float32(0.0) + tCrC_m[ii] = tCrC_m[ii] * keep + + tCrC_m_bf16 = cute.make_fragment_like(tCrC_m, cutlass.BFloat16) + for ii in cutlass.range_constexpr(cute.size(tCrC_m)): + tCrC_m_bf16[ii] = cutlass.BFloat16(tCrC_m[ii]) + cute.copy( + smem_tiled_store_mask_mma, + smem_thr_store_mask_mma.retile(tCrC_m_bf16), + smem_thr_store_mask_mma.partition_D(sMqk_bf16), + ) + cute.arch.barrier() + + i = tidx // CHUNK + col = tidx % CHUNK + l_bf = cutlass.Float32(sL_bf16[i, col]) + sL_fp16[i, col] = cutlass.Float16(l_bf) + inv_init = cutlass.Float32(1.0 if i == col else 0.0) - l_bf + sINV_fp16[i, col] = cutlass.Float16(inv_init) + cute.arch.barrier() + + if warp_idx == 0: + mma_atom_neumann = warp.MmaF16BF16Op(cutlass.Float16, cutlass.Float16, (16, 8, 16)) + tiled_mma_neumann = cute.make_tiled_mma( + mma_atom_neumann, + atom_layout_mnk=(1, 1, 1), + permutation_mnk=(16, 16, 16), + ) + thr_mma_neumann = tiled_mma_neumann.get_slice(tidx) + + copy_atom_A_neumann = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), + cutlass.Float16, + ) + smem_tiled_copy_A_neumann = cute.make_tiled_copy_A(copy_atom_A_neumann, tiled_mma_neumann) + smem_thr_copy_A_neumann = smem_tiled_copy_A_neumann.get_slice(tidx) + + tCrL = thr_mma_neumann.make_fragment_A(thr_mma_neumann.partition_A(sL_fp16)) + tCrL_cv = smem_thr_copy_A_neumann.retile(tCrL) + cute.copy(smem_tiled_copy_A_neumann, smem_thr_copy_A_neumann.partition_S(sL_fp16), tCrL_cv) + + tCrInv = thr_mma_neumann.make_fragment_A(thr_mma_neumann.partition_A(sINV_fp16)) + tCrInv_cv = smem_thr_copy_A_neumann.retile(tCrInv) + cute.copy(smem_tiled_copy_A_neumann, smem_thr_copy_A_neumann.partition_S(sINV_fp16), tCrInv_cv) + + tCrLpowB = thr_mma_neumann.make_fragment_B(thr_mma_neumann.partition_B(sL_fp16)) + tCrLpow = thr_mma_neumann.make_fragment_C(tiled_mma_neumann.partition_shape_C((CHUNK, CHUNK))) + tCrDelta = thr_mma_neumann.make_fragment_C(tiled_mma_neumann.partition_shape_C((CHUNK, CHUNK))) + tCrLpowA = thr_mma_neumann.make_fragment_A(thr_mma_neumann.partition_A(sL_fp16)) + + tCrL_u32 = cute.recast_tensor(tCrL, dtype=cutlass.Int32) + tCrInv_u32 = cute.recast_tensor(tCrInv, dtype=cutlass.Int32) + tCrLpowB_u32 = cute.recast_tensor(tCrLpowB, dtype=cutlass.Int32) + tCrLpow_u32 = cute.recast_tensor(tCrLpow, dtype=cutlass.Int32) + tCrDelta_u32 = cute.recast_tensor(tCrDelta, dtype=cutlass.Int32) + tCrLpowA_u32 = cute.recast_tensor(tCrLpowA, dtype=cutlass.Int32) + + N_REGS_U32: cutlass.Constexpr[int] = 4 # 8 fp16 / thread = 4 u32 + + for ii in cutlass.range_constexpr(N_REGS_U32): + tCrLpowB_u32[ii] = movm_t_b16(cutlass.Int32(tCrL_u32[ii])) + tCrLpow.fill(0.0) + cute.gemm(tiled_mma_neumann, tCrLpow, tCrL, tCrLpowB, tCrLpow) + + for ii in cutlass.range_constexpr(N_REGS_U32): + tCrLpowB_u32[ii] = movm_t_b16(cutlass.Int32(tCrLpow_u32[ii])) + tCrDelta.fill(0.0) + cute.gemm(tiled_mma_neumann, tCrDelta, tCrInv, tCrLpowB, tCrDelta) + for ii in cutlass.range_constexpr(N_REGS_U32): + tCrInv_u32[ii] = add_f16x2_u32(cutlass.Int32(tCrInv_u32[ii]), cutlass.Int32(tCrDelta_u32[ii])) + + for ii in cutlass.range_constexpr(N_REGS_U32): + tCrLpowA_u32[ii] = tCrLpow_u32[ii] + tCrLpow.fill(0.0) + cute.gemm(tiled_mma_neumann, tCrLpow, tCrLpowA, tCrLpowB, tCrLpow) + + for ii in cutlass.range_constexpr(N_REGS_U32): + tCrLpowB_u32[ii] = movm_t_b16(cutlass.Int32(tCrLpow_u32[ii])) + tCrDelta.fill(0.0) + cute.gemm(tiled_mma_neumann, tCrDelta, tCrInv, tCrLpowB, tCrDelta) + for ii in cutlass.range_constexpr(N_REGS_U32): + tCrInv_u32[ii] = add_f16x2_u32(cutlass.Int32(tCrInv_u32[ii]), cutlass.Int32(tCrDelta_u32[ii])) + + for ii in cutlass.range_constexpr(N_REGS_U32): + tCrLpowA_u32[ii] = tCrLpow_u32[ii] + tCrLpow.fill(0.0) + cute.gemm(tiled_mma_neumann, tCrLpow, tCrLpowA, tCrLpowB, tCrLpow) + + for ii in cutlass.range_constexpr(N_REGS_U32): + tCrLpowB_u32[ii] = movm_t_b16(cutlass.Int32(tCrLpow_u32[ii])) + tCrDelta.fill(0.0) + cute.gemm(tiled_mma_neumann, tCrDelta, tCrInv, tCrLpowB, tCrDelta) + for ii in cutlass.range_constexpr(N_REGS_U32): + tCrInv_u32[ii] = add_f16x2_u32(cutlass.Int32(tCrInv_u32[ii]), cutlass.Int32(tCrDelta_u32[ii])) + + tCrInvC = thr_mma_neumann.make_fragment_C(tiled_mma_neumann.partition_shape_C((CHUNK, CHUNK))) + tCrInvC_u32 = cute.recast_tensor(tCrInvC, dtype=cutlass.Int32) + for ii in cutlass.range_constexpr(N_REGS_U32): + tCrInvC_u32[ii] = tCrInv_u32[ii] + tCrInvC_bf16 = cute.make_fragment_like(tCrInvC, cutlass.BFloat16) + for ii in cutlass.range_constexpr(cute.size(tCrInvC)): + tCrInvC_bf16[ii] = cutlass.BFloat16(cutlass.Float32(tCrInvC[ii])) + + copy_atom_stsm = cute.make_copy_atom( + warp.StMatrix8x8x16bOp(transpose=False, num_matrices=2), + cutlass.BFloat16, + ) + smem_tiled_store_C = cute.make_tiled_copy_C_atom(copy_atom_stsm, tiled_mma_neumann) + smem_thr_store_C = smem_tiled_store_C.get_slice(tidx) + cute.copy( + smem_tiled_store_C, + smem_thr_store_C.retile(tCrInvC_bf16), + smem_thr_store_C.partition_D(sINV_bf16), + ) + cute.arch.barrier() + + # TMA bulk store all 5 workspace tensors (one elect_one, one thread). + if warp_idx == 0: + with cute.arch.elect_one(): + cute.copy(tma_atom_ws_qd, tQDws_s[(None,)], tQDws_g[(None, 0, 0, ws_slot)]) + cute.copy(tma_atom_ws_kd, tKDws_s[(None,)], tKDws_g[(None, 0, 0, ws_slot)]) + cute.copy(tma_atom_ws_kr, tKRws_s[(None,)], tKRws_g[(None, 0, 0, ws_slot)]) + cute.copy(tma_atom_ws_inv, tINVws_s[(None,)], tINVws_g[(None, 0, 0, ws_slot)]) + cute.copy(tma_atom_ws_mqk, tMQKws_s[(None,)], tMQKws_g[(None, 0, 0, ws_slot)]) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + + +@cute.jit +def run_k1( + q: cute.Tensor, + k: cute.Tensor, + g: cute.Tensor, + a_log: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + ws_qd: cute.Tensor, + ws_kd: cute.Tensor, + ws_kr: cute.Tensor, + ws_gt: cute.Tensor, + ws_inv: cute.Tensor, + ws_mqk: cute.Tensor, + ws_beta: cute.Tensor, + tile_starts: cute.Tensor, + tile_actual_lens: cute.Tensor, + H: cutlass.Constexpr[int], + total_tiles: cutlass.Int32, + T_total: cutlass.Int32, + scale: cutlass.Constexpr[float], + gate_scale: cutlass.Constexpr[float], + is_varlen: cutlass.Constexpr[bool], + stream: cuda_drv.CUstream, +): + smem_layout_qk = cute.make_layout((CHUNK, D), stride=(D, 1)) + # K_INTER swizzled layout — must match kernel SMEM layout for TMA stores. + kinter_atom = make_smem_layout_atom(SmemLayoutAtomKind.K_INTER, cutlass.BFloat16) + smem_layout_qk_kinter = cute.tile_to_shape(kinter_atom, (CHUNK, D), order=(0, 1)) + + def make_atom(t): + view = cute.make_tensor( + t.iterator, + cute.make_layout((T_total, D, H), stride=(H * D, 1, D)), + ) + return cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + view, + smem_layout_qk, + (CHUNK, D), + ) + + def make_ws_store_atom(t): + view = cute.make_tensor( + t.iterator, + cute.make_layout((CHUNK, D, total_tiles * H), stride=(D, 1, CHUNK * D)), + ) + return cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + view, + smem_layout_qk_kinter, + (CHUNK, D), + ) + + # (CHUNK, CHUNK) bf16 plain layout for ws_inv / ws_mqk TMA bulk store. + smem_layout_cc = cute.make_layout((CHUNK, CHUNK), stride=(CHUNK, 1)) + + def make_ws_cc_store_atom(t): + view = cute.make_tensor( + t.iterator, + cute.make_layout( + (CHUNK, CHUNK, total_tiles * H), + stride=(CHUNK, 1, CHUNK * CHUNK), + ), + ) + return cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + view, + smem_layout_cc, + (CHUNK, CHUNK), + ) + + tma_atom_q, tma_tensor_q = make_atom(q) + tma_atom_k, tma_tensor_k = make_atom(k) + tma_atom_g, tma_tensor_g = make_atom(g) + tma_atom_ws_qd, tma_tensor_ws_qd = make_ws_store_atom(ws_qd) + tma_atom_ws_kd, tma_tensor_ws_kd = make_ws_store_atom(ws_kd) + tma_atom_ws_kr, tma_tensor_ws_kr = make_ws_store_atom(ws_kr) + tma_atom_ws_inv, tma_tensor_ws_inv = make_ws_cc_store_atom(ws_inv) + tma_atom_ws_mqk, tma_tensor_ws_mqk = make_ws_cc_store_atom(ws_mqk) + + smem_bytes = 24 * 1024 + + k1_kernel( + tma_atom_q, + tma_tensor_q, + tma_atom_k, + tma_tensor_k, + tma_atom_g, + tma_tensor_g, + tma_atom_ws_qd, + tma_tensor_ws_qd, + tma_atom_ws_kd, + tma_tensor_ws_kd, + tma_atom_ws_kr, + tma_tensor_ws_kr, + tma_atom_ws_inv, + tma_tensor_ws_inv, + tma_atom_ws_mqk, + tma_tensor_ws_mqk, + a_log, + dt_bias, + beta, + ws_gt, + ws_inv, + ws_mqk, + ws_beta, + tile_starts, + tile_actual_lens, + H, + total_tiles, + T_total, + scale, + gate_scale, + is_varlen, + ).launch( + grid=(total_tiles, H, 1), + block=[THREADS_PER_CTA, 1, 1], + smem=smem_bytes, + stream=stream, + min_blocks_per_mp=8, + ) + + +# Compile cache keyed on CONFIG ONLY — total_tiles/T_total are dynamic +# cutlass.Int32, so one compiled kernel serves every batch shape. +_k1_kernel_cache: dict = {} +_CU_STREAM_CACHE: dict[int, object] = {} +_DUMMY_INT32_CACHE: dict[str, torch.Tensor] = {} +_CU_STREAM_CACHE_MAXSIZE = 64 + + +def _compile_k1(H, scale, gate_scale, is_varlen): + sym_t = cute.sym_int() # T_total (q/k/g token extent) + sym_al = cute.sym_int() # a_log + sym_bt = cute.sym_int() # beta, original packed [T_total, H] flattened (T_total*H) + sym_qk = cute.sym_int() # ws_qd/kd/kr (total_tiles*H*CHUNK*D) + sym_gt = cute.sym_int() # ws_gt (total_tiles*H*D) + sym_cc = cute.sym_int() # ws_inv/mqk (total_tiles*H*CHUNK*CHUNK) + sym_wb = cute.sym_int() # ws_beta (total_tiles*CHUNK*H) + sym_ts = cute.sym_int() # tile_starts + sym_tal = cute.sym_int() # tile_actual_lens + + q_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_t, H, D), stride_order=(2, 1, 0), assumed_align=16) + k_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_t, H, D), stride_order=(2, 1, 0), assumed_align=16) + g_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_t, H, D), stride_order=(2, 1, 0), assumed_align=16) + alog_fake = make_fake_compact_tensor(cutlass.Float32, (sym_al,), assumed_align=16) + dt_fake = make_fake_compact_tensor(cutlass.Float32, (H, D), stride_order=(1, 0), assumed_align=16) + beta_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_bt,), assumed_align=16) + ws_qd_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_qk,), assumed_align=16) + ws_kd_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_qk,), assumed_align=16) + ws_kr_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_qk,), assumed_align=16) + ws_gt_fake = make_fake_compact_tensor(cutlass.Float32, (sym_gt,), assumed_align=16) + ws_inv_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_cc,), assumed_align=16) + ws_mqk_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_cc,), assumed_align=16) + ws_beta_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_wb,), assumed_align=16) + ts_fake = make_fake_compact_tensor(cutlass.Int32, (sym_ts,), assumed_align=4) + tal_fake = make_fake_compact_tensor(cutlass.Int32, (sym_tal,), assumed_align=4) + stream_fake = make_fake_stream() + + return cute.compile( + run_k1, + q_fake, + k_fake, + g_fake, + alog_fake, + dt_fake, + beta_fake, + ws_qd_fake, + ws_kd_fake, + ws_kr_fake, + ws_gt_fake, + ws_inv_fake, + ws_mqk_fake, + ws_beta_fake, + ts_fake, + tal_fake, + H, # Constexpr -> baked + cutlass.Int32(1), # total_tiles -> runtime (placeholder) + cutlass.Int32(1), # T_total -> runtime + scale, # Constexpr + gate_scale, # Constexpr + is_varlen, # Constexpr + stream_fake, + options="--enable-tvm-ffi --opt-level=3", + ) + + +def _get_compiled_k1(H, scale, gate_scale, is_varlen): + key = (H, scale, gate_scale, is_varlen) + cached = _k1_kernel_cache.get(key) + if cached is None: + cached = _compile_k1(H, scale, gate_scale, is_varlen) + _k1_kernel_cache[key] = cached + return cached + + +def _get_current_custream(): + stream_ptr = int(torch.cuda.current_stream().cuda_stream) + cached = _CU_STREAM_CACHE.get(stream_ptr) + if cached is not None: + return cached + if len(_CU_STREAM_CACHE) >= _CU_STREAM_CACHE_MAXSIZE: + _CU_STREAM_CACHE.pop(next(iter(_CU_STREAM_CACHE))) + cached = cuda_drv.CUstream(stream_ptr) + _CU_STREAM_CACHE[stream_ptr] = cached + return cached + + +def _get_dummy_int32(device: torch.device) -> torch.Tensor: + key = str(device) + cached = _DUMMY_INT32_CACHE.get(key) + if cached is not None: + return cached + cached = torch.zeros(1, dtype=torch.int32, device=device) + _DUMMY_INT32_CACHE[key] = cached + return cached + + +def launch_k1( + q: torch.Tensor, + k: torch.Tensor, + g: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + beta: torch.Tensor, + scale: float, + gate_scale: float, + ws_qd: torch.Tensor, + ws_kd: torch.Tensor, + ws_kr: torch.Tensor, + ws_gt: torch.Tensor, + ws_inv: torch.Tensor, + ws_mqk: torch.Tensor, + ws_beta: torch.Tensor, + *, + tile_starts: torch.Tensor | None = None, + tile_actual_lens: torch.Tensor | None = None, + total_tiles: int | None = None, + is_varlen: bool = False, +) -> None: + """Run K1 pipeline; produces all 6 K2-ready workspace tensors.""" + for t in (q, k, g, beta): + assert t.dtype == torch.bfloat16 and t.is_cuda and t.is_contiguous() + assert A_log.dtype == torch.float32 and A_log.is_contiguous() + assert dt_bias.dtype == torch.float32 and dt_bias.is_contiguous() + B, T, H, K = q.shape + assert K == D + T_total = B * T + if not is_varlen: + total_tiles = (B * T) // CHUNK + dummy = _get_dummy_int32(q.device) + tile_starts = dummy + tile_actual_lens = dummy + + compiled_fn = _get_compiled_k1(H, scale, gate_scale, is_varlen) + stream = _get_current_custream() + # TMA descriptors are (re)built inside run_k1 from the dynamic Int32 dims + # every launch, so one compiled kernel serves all shapes. + compiled_fn( + q.view(T_total, H, D), + k.view(T_total, H, D), + g.view(T_total, H, D), + A_log, + dt_bias, + beta, + ws_qd, + ws_kd, + ws_kr, + ws_gt, + ws_inv, + ws_mqk, + ws_beta, + tile_starts, + tile_actual_lens, + total_tiles, + T_total, + stream, + ) diff --git a/cula/ops/kda/sm90/k2.py b/cula/ops/kda/sm90/k2.py new file mode 100644 index 00000000..2fbc15b6 --- /dev/null +++ b/cula/ops/kda/sm90/k2.py @@ -0,0 +1,995 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright (c) 2026 MoonshotAI +# Licensed under the MIT License. +# Based on MoonshotAI/FlashKDA (https://github.com/MoonshotAI/FlashKDA) + +""" +CuteDSL port of FlashKDA K2 (Recurrence). + +Produces workspace tensors consumed by subsequent stages. +""" + +import cuda.bindings.driver as cuda_drv +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.cpasync as cpasync +import torch +from cutlass.cute.nvgpu import warp +from cutlass.cute.nvgpu.warpgroup import SmemLayoutAtomKind, make_smem_layout_atom +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream + +CHUNK: int = 16 +D: int = 128 + + +def _make_state_smem_layout(): + atom = cute.make_composed_layout( + cute.make_swizzle(3, 3, 3), + 0, + cute.make_layout((8, 64), stride=(64, 1)), + ) + return cute.tile_to_shape(atom, (D, D), (0, 1)) + + +from cula.ops.kda.sm90._common import movm_t_b16 # noqa: E402 + + +def _make_out_kinter_one_stage(): + """K_INTER swizzled (CHUNK, D) bf16 SMEM layout.""" + atom = make_smem_layout_atom(SmemLayoutAtomKind.K_INTER, cutlass.BFloat16) + return cute.tile_to_shape(atom, (CHUNK, D), order=(0, 1)) + + +THREADS_PER_CTA = 192 # 128 compute (4 MMA warps) + 32 load + 32 store +N_WARPS = 4 +LOAD_WARP_IDX = 4 +STORE_WARP_IDX = 5 + + +@cute.kernel +def k2_kernel( + tma_atom_v: cute.CopyAtom, + tma_tensor_v: cute.Tensor, + tma_atom_kd: cute.CopyAtom, + tma_tensor_kd: cute.Tensor, + tma_atom_qd: cute.CopyAtom, + tma_tensor_qd: cute.Tensor, + tma_atom_kr: cute.CopyAtom, + tma_tensor_kr: cute.Tensor, + tma_atom_inv: cute.CopyAtom, + tma_tensor_inv: cute.Tensor, + tma_atom_mqk: cute.CopyAtom, + tma_tensor_mqk: cute.Tensor, + tma_atom_out: cute.CopyAtom, + tma_tensor_out: cute.Tensor, + out_gmem: cute.Tensor, + tma_atom_gt: cute.CopyAtom, + tma_tensor_gt: cute.Tensor, + tma_atom_beta: cute.CopyAtom, + tma_tensor_beta: cute.Tensor, + H: cutlass.Constexpr[int], + total_tiles: cutlass.Int32, + cu_seqlens_tiles: cute.Tensor, + seq_order: cute.Tensor, # int32 [N]: launch-slot -> sequence/segment index + v_tile_starts: cute.Tensor, + v_tile_actual_lens: cute.Tensor, + initial_state_g: cute.Tensor, # flat fp32 [N*H*D*D] gmem (layout per state_transposed) + final_state_g: cute.Tensor, # flat fp32 [N*H*D*D] gmem (layout per state_transposed) + has_initial_state: cutlass.Constexpr[bool], + has_final_state: cutlass.Constexpr[bool], + state_transposed: cutlass.Constexpr[bool], + v_is_varlen: cutlass.Constexpr[bool], +): + # Longest-first launch order: pure reordering — no math changes. + seq_slot, head_idx, _ = cute.arch.block_idx() + seq_idx = cutlass.Int32(seq_order[seq_slot]) + tidx, _, _ = cute.arch.thread_idx() + + smem = cutlass.utils.SmemAllocator() + + state_layout = _make_state_smem_layout() + + STAGES: cutlass.Constexpr[int] = 2 + OUT_STAGES: cutlass.Constexpr[int] = 2 + cc_stage_layout = cute.make_layout((CHUNK, CHUNK, STAGES), stride=(CHUNK, 1, CHUNK * CHUNK)) + out_kinter_atom = make_smem_layout_atom(SmemLayoutAtomKind.K_INTER, cutlass.BFloat16) + out_stage_layout = cute.tile_to_shape(out_kinter_atom, (CHUNK, D, OUT_STAGES), order=(0, 1, 2)) + v_kinter_atom = make_smem_layout_atom(SmemLayoutAtomKind.K_INTER, cutlass.BFloat16) + v_stage_layout = cute.tile_to_shape(v_kinter_atom, (CHUNK, D, STAGES), order=(0, 1, 2)) + kd_stage_layout = cute.tile_to_shape(v_kinter_atom, (CHUNK, D, STAGES), order=(0, 1, 2)) + qd_stage_layout = cute.tile_to_shape(v_kinter_atom, (CHUNK, D, STAGES), order=(0, 1, 2)) + kr_stage_layout = cute.tile_to_shape(v_kinter_atom, (CHUNK, D, STAGES), order=(0, 1, 2)) + # MN_INTER transposed view of sKr for state update (same bytes, transposed swizzle). + kr_mninter_atom = make_smem_layout_atom(SmemLayoutAtomKind.MN_INTER, cutlass.BFloat16) + kr_t_stage_layout = cute.tile_to_shape(kr_mninter_atom, (D, CHUNK, STAGES), order=(1, 0, 2)) + + sV = smem.allocate_tensor(cutlass.BFloat16, v_stage_layout, 128) + sKd = smem.allocate_tensor(cutlass.BFloat16, kd_stage_layout, 128) + sQd = smem.allocate_tensor(cutlass.BFloat16, qd_stage_layout, 128) + sKr = smem.allocate_tensor(cutlass.BFloat16, kr_stage_layout, 128) + sINV = smem.allocate_tensor(cutlass.BFloat16, cc_stage_layout, 128) + sMqk = smem.allocate_tensor(cutlass.BFloat16, cc_stage_layout, 128) + sOut = smem.allocate_tensor(cutlass.BFloat16, out_stage_layout, 128) + sState = smem.allocate_tensor(cutlass.BFloat16, state_layout, 128) + sGt = smem.allocate_tensor(cutlass.Float32, cute.make_layout((D, 1, STAGES), stride=(1, D, D)), 128) + sBeta = smem.allocate_tensor(cutlass.BFloat16, cute.make_layout((CHUNK, 1, STAGES), stride=(1, 64, 64)), 128) + sMbar = smem.allocate_tensor(cutlass.Int64, cute.make_layout((STAGES,)), 16) + sMbar_ptr = sMbar.iterator + sMbarE = smem.allocate_tensor(cutlass.Int64, cute.make_layout((STAGES,)), 16) + sMbarE_ptr = sMbarE.iterator + sMbarSF = smem.allocate_tensor(cutlass.Int64, cute.make_layout((OUT_STAGES,)), 16) + sMbarSF_ptr = sMbarSF.iterator + sMbarSE = smem.allocate_tensor(cutlass.Int64, cute.make_layout((OUT_STAGES,)), 16) + sMbarSE_ptr = sMbarSE.iterator + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + if warp_idx == 0: + with cute.arch.elect_one(): + for s in cutlass.range_constexpr(STAGES): + cute.arch.mbarrier_init(sMbar_ptr + cutlass.Int32(s), cutlass.Int32(1)) + cute.arch.mbarrier_init(sMbarE_ptr + cutlass.Int32(s), cutlass.Int32(1)) + for s in cutlass.range_constexpr(OUT_STAGES): + cute.arch.mbarrier_init(sMbarSF_ptr + cutlass.Int32(s), cutlass.Int32(1)) + cute.arch.mbarrier_init(sMbarSE_ptr + cutlass.Int32(s), cutlass.Int32(1)) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + + gSrc_v = cute.local_tile(tma_tensor_v, (CHUNK, D), (None, None, None)) + tVs, tVg = cpasync.tma_partition( + tma_atom_v, + 0, + cute.make_layout(1), + cute.group_modes(sV, 0, 2), + cute.group_modes(gSrc_v, 0, 2), + ) + gSrc_kd = cute.local_tile(tma_tensor_kd, (CHUNK, D), (None, None, None)) + tKDs, tKDg = cpasync.tma_partition( + tma_atom_kd, + 0, + cute.make_layout(1), + cute.group_modes(sKd, 0, 2), + cute.group_modes(gSrc_kd, 0, 2), + ) + gSrc_qd = cute.local_tile(tma_tensor_qd, (CHUNK, D), (None, None, None)) + tQDs, tQDg = cpasync.tma_partition( + tma_atom_qd, + 0, + cute.make_layout(1), + cute.group_modes(sQd, 0, 2), + cute.group_modes(gSrc_qd, 0, 2), + ) + gSrc_kr = cute.local_tile(tma_tensor_kr, (CHUNK, D), (None, None, None)) + tKRs, tKRg = cpasync.tma_partition( + tma_atom_kr, + 0, + cute.make_layout(1), + cute.group_modes(sKr, 0, 2), + cute.group_modes(gSrc_kr, 0, 2), + ) + gSrc_inv = cute.local_tile(tma_tensor_inv, (CHUNK, CHUNK), (None, None, None)) + tIs, tIg = cpasync.tma_partition( + tma_atom_inv, + 0, + cute.make_layout(1), + cute.group_modes(sINV, 0, 2), + cute.group_modes(gSrc_inv, 0, 2), + ) + gSrc_mqk = cute.local_tile(tma_tensor_mqk, (CHUNK, CHUNK), (None, None, None)) + tMs, tMg = cpasync.tma_partition( + tma_atom_mqk, + 0, + cute.make_layout(1), + cute.group_modes(sMqk, 0, 2), + cute.group_modes(gSrc_mqk, 0, 2), + ) + gDst_o = cute.local_tile(tma_tensor_out, (CHUNK, D), (None, None, None)) + tOs, tOg = cpasync.tma_partition( + tma_atom_out, + 0, + cute.make_layout(1), + cute.group_modes(sOut, 0, 2), + cute.group_modes(gDst_o, 0, 2), + ) + gSrc_gt = cute.local_tile(tma_tensor_gt, (D, 1), (None, None, None)) + tGTs, tGTg = cpasync.tma_partition( + tma_atom_gt, + 0, + cute.make_layout(1), + cute.group_modes(sGt, 0, 2), + cute.group_modes(gSrc_gt, 0, 2), + ) + gSrc_beta = cute.local_tile(tma_tensor_beta, (CHUNK, 1), (None, None, None)) + tBs, tBg = cpasync.tma_partition( + tma_atom_beta, + 0, + cute.make_layout(1), + cute.group_modes(sBeta, 0, 2), + cute.group_modes(gSrc_beta, 0, 2), + ) + + if tidx < D: + for e in cutlass.range_constexpr(D): + sState[tidx, e] = cutlass.BFloat16(0.0) + # Load initial_state -> sState[K_in, D_out]. + if has_initial_state: + state_base = cutlass.Int32(seq_idx) * cutlass.Int32(H * D * D) + cutlass.Int32(head_idx) * cutlass.Int32(D * D) + if tidx < D: + if state_transposed: + for k_in in cutlass.range_constexpr(D): + sState[k_in, tidx] = cutlass.BFloat16( + initial_state_g[state_base + cutlass.Int32(k_in * D) + cutlass.Int32(tidx)] + ) + else: + for d_out in cutlass.range_constexpr(D): + sState[tidx, d_out] = cutlass.BFloat16( + initial_state_g[state_base + cutlass.Int32(d_out * D) + cutlass.Int32(tidx)] + ) + cute.arch.barrier() + + mma_atom = warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16)) + tiled_mma = cute.make_tiled_mma( + mma_atom, + atom_layout_mnk=(1, 4, 1), + permutation_mnk=(16, 32, 16), + ) + thr_mma = tiled_mma.get_slice(tidx) + + copy_atom_AB = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), + cutlass.BFloat16, + ) + smem_tiled_copy_A = cute.make_tiled_copy_A(copy_atom_AB, tiled_mma) + smem_thr_copy_A = smem_tiled_copy_A.get_slice(tidx) + copy_atom_B_T = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4), + cutlass.BFloat16, + ) + smem_tiled_copy_B_T = cute.make_tiled_copy_B(copy_atom_B_T, tiled_mma) + smem_thr_copy_B_T = smem_tiled_copy_B_T.get_slice(tidx) + + copy_atom_stsm = cute.make_copy_atom( + warp.StMatrix8x8x16bOp(transpose=False, num_matrices=2), + cutlass.BFloat16, + ) + smem_tiled_store_C = cute.make_tiled_copy_C_atom(copy_atom_stsm, tiled_mma) + smem_thr_store_C = smem_tiled_store_C.get_slice(tidx) + + tiled_mma_state = cute.make_tiled_mma( + mma_atom, + atom_layout_mnk=(1, 4, 1), + permutation_mnk=(16, 32, 16), + ) + thr_mma_state = tiled_mma_state.get_slice(tidx) + smem_tiled_copy_A_state = cute.make_tiled_copy_A(copy_atom_B_T, tiled_mma_state) + smem_thr_copy_A_state = smem_tiled_copy_A_state.get_slice(tidx) + + sKd_s0 = sKd[(None, None, 0)] + sQd_s0 = sQd[(None, None, 0)] + sKd_tile0 = cute.flat_divide(sKd_s0, (CHUNK, 16)) + sQd_tile0 = cute.flat_divide(sQd_s0, (CHUNK, 16)) + # sState[K_in, D_out] transposed view for B-operand in Phase 1/4. + sState_B_view = cute.make_tensor(sState.iterator, layout=cute.select(sState.layout, mode=[1, 0])) + sState_tile = cute.flat_divide(sState_B_view, (D, 16)) + + sKd_ref = sKd_tile0[None, None, 0, 0] + sQd_ref = sQd_tile0[None, None, 0, 0] + sState_ref = sState_tile[None, None, 0, 0] + + tCrKd = thr_mma.make_fragment_A(thr_mma.partition_A(sKd_ref)) + tCrQd = thr_mma.make_fragment_A(thr_mma.partition_A(sQd_ref)) + tCrState = thr_mma.make_fragment_B(thr_mma.partition_B(sState_ref)) + tCrU = thr_mma.make_fragment_C(tiled_mma.partition_shape_C((CHUNK, D))) + tCrOut = thr_mma.make_fragment_C(tiled_mma.partition_shape_C((CHUNK, D))) + + tCrKd_cv = smem_thr_copy_A.retile(tCrKd) + tCrQd_cv = smem_thr_copy_A.retile(tCrQd) + tCrState_cv = smem_thr_copy_B_T.retile(tCrState) + + sMqk_s0 = sMqk[(None, None, 0)] + sMqk_tile0 = cute.flat_divide(sMqk_s0, (CHUNK, CHUNK)) + sMqk_ref = sMqk_tile0[None, None, 0, 0] + tCrMqk = thr_mma.make_fragment_A(thr_mma.partition_A(sMqk_ref)) + tCrMqk_cv = smem_thr_copy_A.retile(tCrMqk) + + # State update: MN_INTER transposed view of sKr. + sKr_T_view = cute.make_tensor(sKr.iterator, kr_t_stage_layout) + sKr_T_view_s0 = sKr_T_view[(None, None, 0)] + sKr_T_ref = cute.flat_divide(sKr_T_view_s0, (D, CHUNK))[None, None, 0, 0] + + sKr_T_blk_for_frag = cute.flat_divide(sKr_T_view_s0, (CHUNK, CHUNK))[None, None, 0, 0] + tCrKrA_state_blk = thr_mma_state.make_fragment_A(thr_mma_state.partition_A(sKr_T_blk_for_frag)) + tCrKrA_state_blk_cv = smem_thr_copy_A_state.retile(tCrKrA_state_blk) + tCrUpd_blk = thr_mma_state.make_fragment_C(tiled_mma_state.partition_shape_C((CHUNK, D))) + sState_blk_tile = cute.flat_divide(sState, (CHUNK, D)) + coord_state_blk = cute.make_identity_tensor((CHUNK, D)) + tCcState_blk = thr_mma_state.partition_C(coord_state_blk) + + tCrU_T = thr_mma.make_fragment_B(thr_mma.partition_B(sKr_T_ref)) + + sINV_s0 = sINV[(None, None, 0)] + sINV_tile0 = cute.flat_divide(sINV_s0, (CHUNK, CHUNK)) + sINV_ref = sINV_tile0[None, None, 0, 0] + tCrInv = thr_mma.make_fragment_A(thr_mma.partition_A(sINV_ref)) + tCrInv_cv = smem_thr_copy_A.retile(tCrInv) + tCrU_post = thr_mma.make_fragment_C(tiled_mma.partition_shape_C((CHUNK, D))) + tCrU_T_post = cute.make_fragment_like(tCrU_T) + tCrU_post_bf16 = cute.make_fragment_like(tCrU_post, cutlass.BFloat16) + tCrU_pre_bf16 = cute.make_fragment_like(tCrU, cutlass.BFloat16) + + tile_base = cu_seqlens_tiles[seq_idx] + t_tiles = cu_seqlens_tiles[seq_idx + 1] - tile_base # dynamic tile count for this sequence + TMA_BYTES: cutlass.Constexpr[int] = 4 * CHUNK * D * 2 + 2 * CHUNK * CHUNK * 2 + D * 4 + CHUNK * 2 + + if warp_idx == LOAD_WARP_IDX: + s_dyn_l = cutlass.Int32(0) + phase_emp = cutlass.Int32(1) + if cutlass.const_expr(v_is_varlen): + seq_v_start = v_tile_starts[tile_base] + tma_tensor_v_seq = cute.domain_offset((seq_v_start, 0, 0), tma_tensor_v) + gSrc_v_seq = cute.local_tile(tma_tensor_v_seq, (CHUNK, D), (None, None, None)) + tVs_seq, tVg_seq = cpasync.tma_partition( + tma_atom_v, + 0, + cute.make_layout(1), + cute.group_modes(sV, 0, 2), + cute.group_modes(gSrc_v_seq, 0, 2), + ) + for t in cutlass.range(t_tiles, unroll=1): + cute.arch.mbarrier_wait(sMbarE_ptr + s_dyn_l, phase_emp) + tg_l = tile_base + t + wt_l = head_idx * total_tiles + tg_l + bar_l = sMbar_ptr + s_dyn_l + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(bar_l, cutlass.Int32(TMA_BYTES)) + if cutlass.const_expr(v_is_varlen): + cute.copy(tma_atom_v, tVg_seq[(None, t, 0, head_idx)], tVs_seq[(None, s_dyn_l)], tma_bar_ptr=bar_l) + else: + cute.copy(tma_atom_v, tVg[(None, tg_l, 0, head_idx)], tVs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(tma_atom_kd, tKDg[(None, 0, 0, wt_l)], tKDs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(tma_atom_qd, tQDg[(None, 0, 0, wt_l)], tQDs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(tma_atom_kr, tKRg[(None, 0, 0, wt_l)], tKRs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(tma_atom_inv, tIg[(None, 0, 0, wt_l)], tIs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(tma_atom_mqk, tMg[(None, 0, 0, wt_l)], tMs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(tma_atom_gt, tGTg[(None, 0, 0, wt_l)], tGTs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + cute.copy(tma_atom_beta, tBg[(None, 0, 0, wt_l)], tBs[(None, s_dyn_l)], tma_bar_ptr=bar_l) + s_dyn_l = s_dyn_l + cutlass.Int32(1) + if s_dyn_l == cutlass.Int32(STAGES): + s_dyn_l = cutlass.Int32(0) + phase_emp = phase_emp ^ cutlass.Int32(1) + elif warp_idx == STORE_WARP_IDX: + if cutlass.const_expr(v_is_varlen): + universal_copy_bits: cutlass.Constexpr[int] = 128 + async_copy_elems: cutlass.Constexpr[int] = universal_copy_bits // cutlass.BFloat16.width + atom_universal_copy_o = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + cutlass.BFloat16, + num_bits_per_copy=universal_copy_bits, + ) + o_thr_layout = cute.make_ordered_layout( + (THREADS_PER_CTA // 6 // (D // async_copy_elems), D // async_copy_elems), + order=(1, 0), + ) + o_val_layout = cute.make_layout((1, async_copy_elems)) + gmem_tiled_copy_o = cute.make_tiled_copy_tv( + atom_universal_copy_o, + o_thr_layout, + o_val_layout, + ) + gmem_thr_copy_o = gmem_tiled_copy_o.get_slice(tidx % 32) + seq_out_start = v_tile_starts[tile_base] + tma_tensor_out_seq = cute.domain_offset((seq_out_start, 0, 0), tma_tensor_out) + gDst_o_seq = cute.local_tile(tma_tensor_out_seq, (CHUNK, D), (None, None, None)) + tOs_seq, tOg_seq = cpasync.tma_partition( + tma_atom_out, + 0, + cute.make_layout(1), + cute.group_modes(sOut, 0, 2), + cute.group_modes(gDst_o_seq, 0, 2), + ) + + s_out_s = cutlass.Int32(0) + phase_sf = cutlass.Int32(0) + for t in cutlass.range(t_tiles, unroll=1): + cute.arch.mbarrier_wait(sMbarSF_ptr + s_out_s, phase_sf) + t_g_s = tile_base + t + if cutlass.const_expr(v_is_varlen): + actual_len_o = v_tile_actual_lens[t_g_s] + if actual_len_o == cutlass.Int32(CHUNK): + cute.copy(tma_atom_out, tOs_seq[(None, s_out_s)], tOg_seq[(None, t, 0, head_idx)]) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + else: + out_start = v_tile_starts[t_g_s] + sOut_stage = sOut[(None, None, s_out_s)] + tOsO = gmem_thr_copy_o.partition_S(sOut_stage) + cO = cute.make_identity_tensor((CHUNK, D)) + tOcO = gmem_thr_copy_o.partition_S(cO) + tOrO = cute.make_fragment_like(tOsO, cutlass.BFloat16) + cute.autovec_copy(tOsO, tOrO) + + out_chunk_raw = out_gmem.iterator + out_start * H * D + head_idx * D + out_chunk_ptr = cute.make_ptr( + cutlass.BFloat16, + out_chunk_raw.toint(), + cute.AddressSpace.gmem, + assumed_align=16, + ) + out_stride_t: cutlass.Constexpr[int] = H * D + gOut_chunk = cute.make_tensor( + out_chunk_ptr, + cute.make_layout((CHUNK, D), stride=(out_stride_t, 1)), + ) + tOgO = gmem_thr_copy_o.partition_D(gOut_chunk) + for m1 in cutlass.range_constexpr(cute.size(tOsO.shape[1])): + row_coord = tOcO[(0, 0), m1, 0][0] + if row_coord < actual_len_o: + cute.autovec_copy(tOrO[(None, m1, None)], tOgO[(None, m1, None)]) + else: + cute.copy(tma_atom_out, tOs[(None, s_out_s)], tOg[(None, t_g_s, 0, head_idx)]) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(sMbarSE_ptr + s_out_s) + s_out_s = s_out_s + cutlass.Int32(1) + if s_out_s == cutlass.Int32(OUT_STAGES): + s_out_s = cutlass.Int32(0) + phase_sf = phase_sf ^ cutlass.Int32(1) + else: + phase_full = cutlass.Int32(0) + s_dyn = cutlass.Int32(0) + s_out = cutlass.Int32(0) + phase_se = cutlass.Int32(1) + + for t in cutlass.range(t_tiles, unroll=1): + sV_s = sV[(None, None, s_dyn)] + sKd_tile = cute.flat_divide(sKd[(None, None, s_dyn)], (CHUNK, 16)) + sQd_tile = cute.flat_divide(sQd[(None, None, s_dyn)], (CHUNK, 16)) + sINV_ref_s = cute.flat_divide(sINV[(None, None, s_dyn)], (CHUNK, CHUNK))[None, None, 0, 0] + sMqk_ref_s = cute.flat_divide(sMqk[(None, None, s_dyn)], (CHUNK, CHUNK))[None, None, 0, 0] + sGt_s = sGt[(None, 0, s_dyn)] + sBeta_s = sBeta[(None, 0, s_dyn)] + + cute.arch.mbarrier_wait(sMbar_ptr + s_dyn, phase_full) + if cutlass.const_expr(v_is_varlen): + actual_len = v_tile_actual_lens[tile_base + t] + if actual_len < cutlass.Int32(CHUNK): + if tidx < D: + col_v_tail = tidx + for r in cutlass.range_constexpr(CHUNK): + if actual_len <= cutlass.Int32(r): + sV_s[r, col_v_tail] = cutlass.BFloat16(0.0) + cute.arch.barrier(barrier_id=1, number_of_threads=128) + + sKr_T_s = sKr_T_view[(None, None, s_dyn)] + sKr_T_blk_tile_s = cute.flat_divide(sKr_T_s, (CHUNK, CHUNK)) + + # kd @ state and qd @ state share the same state tile load. + tCrU.fill(0.0) + tCrOut.fill(0.0) + for k in cutlass.range_constexpr(D // 16): + sKd_k = sKd_tile[None, None, 0, k] + sQd_k = sQd_tile[None, None, 0, k] + sState_k = sState_tile[None, None, 0, k] + cute.copy(smem_tiled_copy_B_T, smem_thr_copy_B_T.partition_S(sState_k), tCrState_cv) + cute.copy(smem_tiled_copy_A, smem_thr_copy_A.partition_S(sKd_k), tCrKd_cv) + cute.copy(smem_tiled_copy_A, smem_thr_copy_A.partition_S(sQd_k), tCrQd_cv) + cute.gemm(tiled_mma, tCrU, tCrKd, tCrState, tCrU) + cute.gemm(tiled_mma, tCrOut, tCrQd, tCrState, tCrOut) + + # sigmoid(beta) * (v - u_pre) + lane_in_warp = tidx % 32 + Rrow0 = lane_in_warp // 4 + Rrow1 = Rrow0 + 8 + b0 = cutlass.Float32(sBeta_s[Rrow0]) + b1 = cutlass.Float32(sBeta_s[Rrow1]) + sig0 = cutlass.Float32(0.5) * (cute.tanh(b0 * cutlass.Float32(0.5), fastmath=True) + cutlass.Float32(1.0)) + sig1 = cutlass.Float32(0.5) * (cute.tanh(b1 * cutlass.Float32(0.5), fastmath=True) + cutlass.Float32(1.0)) + tCsV = thr_mma.partition_C(sV_s) + for i in cutlass.range_constexpr(cute.size(tCrU)): + ii: cutlass.Constexpr[int] = i + sub_i: cutlass.Constexpr[int] = (ii % 4) // 2 + sig = sig0 if sub_i == 0 else sig1 + diff = cutlass.Float32(tCsV[ii]) - tCrU[ii] + tCrU_pre_bf16[ii] = cutlass.BFloat16(diff * sig) + tCrU_pre_u32 = cute.recast_tensor(tCrU_pre_bf16, dtype=cutlass.Int32) + tCrU_T_u32 = cute.recast_tensor(tCrU_T, dtype=cutlass.Int32) + for i in cutlass.range_constexpr(cute.size(tCrU_pre_u32)): + ii: cutlass.Constexpr[int] = i + tCrU_T_u32[ii] = movm_t_b16(cutlass.Int32(tCrU_pre_u32[ii])) + + # U_post = INV @ U_pre + cute.copy(smem_tiled_copy_A, smem_thr_copy_A.partition_S(sINV_ref_s), tCrInv_cv) + tCrU_post.fill(0.0) + cute.gemm(tiled_mma, tCrU_post, tCrInv, tCrU_T, tCrU_post) + for i in cutlass.range_constexpr(cute.size(tCrU_post)): + ii: cutlass.Constexpr[int] = i + tCrU_post_bf16[ii] = cutlass.BFloat16(tCrU_post[ii]) + tCrU_post_u32 = cute.recast_tensor(tCrU_post_bf16, dtype=cutlass.Int32) + tCrU_T_post_u32 = cute.recast_tensor(tCrU_T_post, dtype=cutlass.Int32) + for i in cutlass.range_constexpr(cute.size(tCrU_post_u32)): + ii: cutlass.Constexpr[int] = i + tCrU_T_post_u32[ii] = movm_t_b16(cutlass.Int32(tCrU_post_u32[ii])) + + # out += Mqk @ U_post + cute.copy(smem_tiled_copy_A, smem_thr_copy_A.partition_S(sMqk_ref_s), tCrMqk_cv) + cute.gemm(tiled_mma, tCrOut, tCrMqk, tCrU_T_post, tCrOut) + + cute.arch.mbarrier_wait(sMbarSE_ptr + s_out, phase_se) + sOut_s = sOut[(None, None, s_out)] + tCrOut_bf16 = cute.make_fragment_like(tCrOut, cutlass.BFloat16) + for i in cutlass.range_constexpr(cute.size(tCrOut)): + tCrOut_bf16[i] = cutlass.BFloat16(tCrOut[i]) + cute.copy( + smem_tiled_store_C, + smem_thr_store_C.retile(tCrOut_bf16), + smem_thr_store_C.partition_D(sOut_s), + ) + + # State update: state = state*gt + kr^T @ U (blocked M-loop) + M_BLOCKS: cutlass.Constexpr[int] = D // CHUNK + for mi in cutlass.range_constexpr(M_BLOCKS): + sKr_T_blk_s = sKr_T_blk_tile_s[None, None, mi, 0] + cute.copy( + smem_tiled_copy_A_state, + smem_thr_copy_A_state.partition_S(sKr_T_blk_s), + tCrKrA_state_blk_cv, + ) + tCrUpd_blk.fill(0.0) + cute.gemm(tiled_mma_state, tCrUpd_blk, tCrKrA_state_blk, tCrU_T_post, tCrUpd_blk) + + sState_blk = sState_blk_tile[None, None, mi, 0] + tCsState_blk = thr_mma_state.partition_C(sState_blk) + state_frag_blk = cute.make_fragment_like(tCsState_blk, cutlass.BFloat16) + gt_frag_blk = cute.make_fragment_like(tCsState_blk, cutlass.Float32) + m_off: cutlass.Constexpr[int] = mi * CHUNK + for i in cutlass.range_constexpr(cute.size(state_frag_blk)): + ii: cutlass.Constexpr[int] = i + state_frag_blk[ii] = tCsState_blk[ii] + gt_frag_blk[ii] = sGt_s[m_off + tCcState_blk[ii][0]] + for i in cutlass.range_constexpr(cute.size(tCrUpd_blk)): + ii: cutlass.Constexpr[int] = i + old = cutlass.Float32(state_frag_blk[ii]) * gt_frag_blk[ii] + tCsState_blk[ii] = cutlass.BFloat16(old + tCrUpd_blk[ii]) + cute.arch.barrier(barrier_id=1, number_of_threads=128) + cute.arch.fence_view_async_shared() + if warp_idx == 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(sMbarSF_ptr + s_out) + cute.arch.mbarrier_arrive(sMbarE_ptr + s_dyn) + s_dyn = s_dyn + cutlass.Int32(1) + if s_dyn == cutlass.Int32(STAGES): + s_dyn = cutlass.Int32(0) + phase_full = phase_full ^ cutlass.Int32(1) + s_out = s_out + cutlass.Int32(1) + if s_out == cutlass.Int32(OUT_STAGES): + s_out = cutlass.Int32(0) + phase_se = phase_se ^ cutlass.Int32(1) + cute.arch.barrier() + if has_final_state: + state_base_f = cutlass.Int32(seq_idx) * cutlass.Int32(H * D * D) + cutlass.Int32(head_idx) * cutlass.Int32(D * D) + if tidx < D: + if state_transposed: + for k_in in cutlass.range_constexpr(D): + final_state_g[state_base_f + cutlass.Int32(k_in * D) + cutlass.Int32(tidx)] = cutlass.Float32( + sState[k_in, tidx] + ) + else: + for d_out in cutlass.range_constexpr(D): + final_state_g[state_base_f + cutlass.Int32(d_out * D) + cutlass.Int32(tidx)] = cutlass.Float32( + sState[tidx, d_out] + ) + + +@cute.jit +def run_k2( + v: cute.Tensor, + beta: cute.Tensor, + ws_qd: cute.Tensor, + ws_kd: cute.Tensor, + ws_kr: cute.Tensor, + ws_gt: cute.Tensor, + ws_inv: cute.Tensor, + ws_mqk: cute.Tensor, + out: cute.Tensor, + cu_seqlens_tiles: cute.Tensor, + seq_order: cute.Tensor, + initial_state_g: cute.Tensor, # flat fp32 [N*H*D*D] or dummy [1] + final_state_g: cute.Tensor, # flat fp32 [N*H*D*D] or dummy [1] + v_tile_starts: cute.Tensor, + v_tile_actual_lens: cute.Tensor, + H: cutlass.Constexpr[int], + total_tiles: cutlass.Int32, + O_T_total: cutlass.Int32, + V_T_total: cutlass.Int32, + N: cutlass.Int32, + has_initial_state: cutlass.Constexpr[bool], + has_final_state: cutlass.Constexpr[bool], + state_transposed: cutlass.Constexpr[bool], + v_is_varlen: cutlass.Constexpr[bool], + stream: cuda_drv.CUstream, +): + cc_smem = cute.make_layout((CHUNK, CHUNK), stride=(CHUNK, 1)) + kinter_smem = _make_out_kinter_one_stage() + + def make_thd_atom(t, op, t_total: cutlass.Int32): + view = cute.make_tensor( + t.iterator, + cute.make_layout((t_total, D, H), stride=(H * D, 1, D)), + ) + return cpasync.make_tiled_tma_atom(op, view, kinter_smem, (CHUNK, D)) + + def make_ws_qkd_atom(t): + view = cute.make_tensor( + t.iterator, + cute.make_layout((CHUNK, D, total_tiles * H), stride=(D, 1, CHUNK * D)), + ) + return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, kinter_smem, (CHUNK, D)) + + def make_ws_cc_atom(t): + view = cute.make_tensor( + t.iterator, + cute.make_layout((CHUNK, CHUNK, total_tiles * H), stride=(CHUNK, 1, CHUNK * CHUNK)), + ) + return cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), view, cc_smem, (CHUNK, CHUNK)) + + tma_atom_v, tma_tensor_v = make_thd_atom(v, cpasync.CopyBulkTensorTileG2SOp(), V_T_total) + tma_atom_out, tma_tensor_out = make_thd_atom(out, cpasync.CopyBulkTensorTileS2GOp(), O_T_total) + tma_atom_kd, tma_tensor_kd = make_ws_qkd_atom(ws_kd) + tma_atom_qd, tma_tensor_qd = make_ws_qkd_atom(ws_qd) + tma_atom_kr, tma_tensor_kr = make_ws_qkd_atom(ws_kr) + tma_atom_inv, tma_tensor_inv = make_ws_cc_atom(ws_inv) + tma_atom_mqk, tma_tensor_mqk = make_ws_cc_atom(ws_mqk) + + gt_smem = cute.make_layout((D, 1), stride=(1, D)) + beta_smem = cute.make_layout((CHUNK, 1), stride=(1, 64)) + + def make_gt_atom(t): + view = cute.make_tensor( + t.iterator, + cute.make_layout((D, 1, total_tiles * H), stride=(1, D, D)), + ) + return cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + view, + gt_smem, + (D, 1), + ) + + def make_beta_atom(t): + view = cute.make_tensor( + t.iterator, + cute.make_layout((CHUNK, 1, total_tiles * H), stride=(1, CHUNK, CHUNK)), + ) + return cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + view, + beta_smem, + (CHUNK, 1), + ) + + tma_atom_gt, tma_tensor_gt = make_gt_atom(ws_gt) + tma_atom_beta, tma_tensor_beta = make_beta_atom(beta) + + STAGES_LOCAL = 2 + OUT_STAGES_LOCAL = 2 + smem_bytes = ( + D * D * 2 + + STAGES_LOCAL * 4 * (CHUNK * D * 2) + + STAGES_LOCAL * 2 * (CHUNK * CHUNK * 2) + + OUT_STAGES_LOCAL * (CHUNK * D * 2) + + STAGES_LOCAL * (D * 4) + + STAGES_LOCAL * (64 * 2) + + (STAGES_LOCAL + OUT_STAGES_LOCAL) * 2 * 8 + + 2048 + ) + + k2_kernel( + tma_atom_v, + tma_tensor_v, + tma_atom_kd, + tma_tensor_kd, + tma_atom_qd, + tma_tensor_qd, + tma_atom_kr, + tma_tensor_kr, + tma_atom_inv, + tma_tensor_inv, + tma_atom_mqk, + tma_tensor_mqk, + tma_atom_out, + tma_tensor_out, + out, + tma_atom_gt, + tma_tensor_gt, + tma_atom_beta, + tma_tensor_beta, + H, + total_tiles, + cu_seqlens_tiles, + seq_order, + v_tile_starts, + v_tile_actual_lens, + initial_state_g, + final_state_g, + has_initial_state, + has_final_state, + state_transposed, + v_is_varlen, + ).launch( + grid=(N, H, 1), + block=[THREADS_PER_CTA, 1, 1], + smem=smem_bytes, + stream=stream, + ) + + +# Compile cache keyed on CONFIG ONLY — the per-batch shape dims +# (total_tiles, O_T_total, V_T_total, N) are dynamic cutlass.Int32, so one +# compiled kernel serves every batch shape. +_k2_kernel_cache: dict = {} +_DUMMY_FP32_CACHE: dict[str, torch.Tensor] = {} +_DUMMY_INT32_CACHE: dict[str, torch.Tensor] = {} +_CU_STREAM_CACHE: dict[int, object] = {} +_CU_STREAM_CACHE_MAXSIZE = 64 +_FIXED_CU_TILES_CACHE: dict[tuple, torch.Tensor] = {} +_FIXED_CU_TILES_CACHE_MAXSIZE = 64 +_IDENTITY_ORDER_CACHE: dict[tuple, torch.Tensor] = {} +_IDENTITY_ORDER_CACHE_MAXSIZE = 64 + + +def _compile_k2(H, has_initial_state, has_final_state, state_transposed, v_is_varlen): + # One sym_int per dynamic tensor extent; shape scalars are runtime Int32. + sym_vt = cute.sym_int() # V_T_total (v token extent) + sym_ot = cute.sym_int() # O_T_total (out token extent; differs from V in varlen) + sym_qk = cute.sym_int() # ws_qd/kd/kr (total_tiles*H*CHUNK*D) + sym_cc = cute.sym_int() # ws_inv/mqk (total_tiles*H*CHUNK*CHUNK) + sym_gt = cute.sym_int() # ws_gt (total_tiles*H*D) + sym_bt = cute.sym_int() # beta/ws_beta (total_tiles*CHUNK*H) + sym_cu = cute.sym_int() # cu_seqlens_tiles (N+1) + sym_so = cute.sym_int() # seq_order (N) — own sym: length differs from cu_seqlens_tiles + # init/final need separate syms: their runtime sizes differ when one is + # the 1-element dummy, and tvm-ffi binds each sym to one value. + sym_sti = cute.sym_int() # initial state flat (N*H*D*D or 1) + sym_stf = cute.sym_int() # final state flat (N*H*D*D or 1) + sym_vs = cute.sym_int() # v_tile_starts (total_tiles or 1) + sym_va = cute.sym_int() # v_tile_actual_lens (total_tiles or 1) + + v_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_vt, H, D), stride_order=(2, 1, 0), assumed_align=16) + beta_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_bt,), assumed_align=16) + ws_qd_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_qk,), assumed_align=16) + ws_kd_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_qk,), assumed_align=16) + ws_kr_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_qk,), assumed_align=16) + ws_gt_fake = make_fake_compact_tensor(cutlass.Float32, (sym_gt,), assumed_align=16) + ws_inv_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_cc,), assumed_align=16) + ws_mqk_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_cc,), assumed_align=16) + out_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_ot, H, D), stride_order=(2, 1, 0), assumed_align=16) + cu_fake = make_fake_compact_tensor(cutlass.Int32, (sym_cu,), assumed_align=4) + so_fake = make_fake_compact_tensor(cutlass.Int32, (sym_so,), assumed_align=4) + init_fake = make_fake_compact_tensor(cutlass.Float32, (sym_sti,), assumed_align=16) + final_fake = make_fake_compact_tensor(cutlass.Float32, (sym_stf,), assumed_align=16) + vts_fake = make_fake_compact_tensor(cutlass.Int32, (sym_vs,), assumed_align=4) + vtal_fake = make_fake_compact_tensor(cutlass.Int32, (sym_va,), assumed_align=4) + stream_fake = make_fake_stream() + + return cute.compile( + run_k2, + v_fake, + beta_fake, + ws_qd_fake, + ws_kd_fake, + ws_kr_fake, + ws_gt_fake, + ws_inv_fake, + ws_mqk_fake, + out_fake, + cu_fake, + so_fake, + init_fake, + final_fake, + vts_fake, + vtal_fake, + H, # Constexpr -> baked + cutlass.Int32(1), # total_tiles -> runtime (placeholder) + cutlass.Int32(1), # O_T_total + cutlass.Int32(1), # V_T_total + cutlass.Int32(1), # N + has_initial_state, # Constexpr + has_final_state, # Constexpr + state_transposed, # Constexpr + v_is_varlen, # Constexpr + stream_fake, + options="--enable-tvm-ffi", + ) + + +def _get_compiled_k2(H, has_initial_state, has_final_state, state_transposed, v_is_varlen): + key = (H, has_initial_state, has_final_state, state_transposed, v_is_varlen) + cached = _k2_kernel_cache.get(key) + if cached is None: + cached = _compile_k2(H, has_initial_state, has_final_state, state_transposed, v_is_varlen) + _k2_kernel_cache[key] = cached + return cached + + +def _get_current_custream(): + stream_ptr = int(torch.cuda.current_stream().cuda_stream) + cached = _CU_STREAM_CACHE.get(stream_ptr) + if cached is not None: + return cached + if len(_CU_STREAM_CACHE) >= _CU_STREAM_CACHE_MAXSIZE: + _CU_STREAM_CACHE.pop(next(iter(_CU_STREAM_CACHE))) + cached = cuda_drv.CUstream(stream_ptr) + _CU_STREAM_CACHE[stream_ptr] = cached + return cached + + +def _get_dummy_fp32(device: torch.device) -> torch.Tensor: + key = str(device) + cached = _DUMMY_FP32_CACHE.get(key) + if cached is not None: + return cached + cached = torch.zeros(1, dtype=torch.float32, device=device) + _DUMMY_FP32_CACHE[key] = cached + return cached + + +def _get_dummy_int32(device: torch.device) -> torch.Tensor: + key = str(device) + cached = _DUMMY_INT32_CACHE.get(key) + if cached is not None: + return cached + cached = torch.zeros(1, dtype=torch.int32, device=device) + _DUMMY_INT32_CACHE[key] = cached + return cached + + +def _get_identity_order(n: int, device: torch.device) -> torch.Tensor: + """Cached identity launch order [0..n) for serial/uniform batches.""" + key = (n, str(device)) + cached = _IDENTITY_ORDER_CACHE.get(key) + if cached is not None: + return cached + if len(_IDENTITY_ORDER_CACHE) >= _IDENTITY_ORDER_CACHE_MAXSIZE: + _IDENTITY_ORDER_CACHE.pop(next(iter(_IDENTITY_ORDER_CACHE))) + cached = torch.arange(n, dtype=torch.int32, device=device) + _IDENTITY_ORDER_CACHE[key] = cached + return cached + + +def _get_fixed_cu_seqlens_tiles(B: int, t_tiles_per_seq: int, device: torch.device) -> torch.Tensor: + """Cached [0, t, 2t, ..., B*t] tile offsets for the fixed-length path.""" + key = (B, t_tiles_per_seq, str(device)) + cached = _FIXED_CU_TILES_CACHE.get(key) + if cached is not None: + return cached + if len(_FIXED_CU_TILES_CACHE) >= _FIXED_CU_TILES_CACHE_MAXSIZE: + _FIXED_CU_TILES_CACHE.pop(next(iter(_FIXED_CU_TILES_CACHE))) + cached = torch.arange( + 0, + (B + 1) * t_tiles_per_seq, + t_tiles_per_seq, + dtype=torch.int32, + device=device, + ) + _FIXED_CU_TILES_CACHE[key] = cached + return cached + + +def launch_k2( + v: torch.Tensor, + beta: torch.Tensor, + ws_qd: torch.Tensor, + ws_kd: torch.Tensor, + ws_kr: torch.Tensor, + ws_gt: torch.Tensor, + ws_inv: torch.Tensor, + ws_mqk: torch.Tensor, + out: torch.Tensor, + cu_seqlens_tiles: torch.Tensor | None = None, + initial_state: torch.Tensor | None = None, # [N, H, V, K] fp32 bhvk + final_state: torch.Tensor | None = None, # [N, H, V, K] fp32 bhvk (written in-place) + state_transposed: bool = False, + v_tile_starts: torch.Tensor | None = None, + v_tile_actual_lens: torch.Tensor | None = None, + seq_order: torch.Tensor | None = None, +) -> None: + """Run K2 recurrence. Supports fixed-len and varlen inputs. + + state_transposed: False=[N,H,V,K] (default), True=[N,H,K,V]. + seq_order: optional int32 [N] launch-slot -> sequence index permutation + (longest-first for ragged CP batches); None = identity. + """ + assert v.is_cuda and v.dtype == torch.bfloat16 and v.is_contiguous() + assert out.is_cuda and out.dtype == torch.bfloat16 and out.is_contiguous() + B, T, H, K = v.shape + assert K == D + assert out.ndim == 4 and out.shape[0] == B and out.shape[2] == H and out.shape[3] == D + V_T_total = B * T + O_T_total = out.shape[0] * out.shape[1] + v_is_varlen = v_tile_starts is not None + if v_is_varlen: + total_tiles = v_tile_starts.numel() + else: + O_T_total = V_T_total + total_tiles = V_T_total // CHUNK + dummy = _get_dummy_int32(v.device) + v_tile_starts = dummy + v_tile_actual_lens = dummy + + if cu_seqlens_tiles is None: + cu_seqlens_tiles = _get_fixed_cu_seqlens_tiles(B, T // CHUNK, v.device) + N_seqs = B + else: + N_seqs = cu_seqlens_tiles.numel() - 1 + + has_initial_state_flag = initial_state is not None + has_final_state_flag = final_state is not None + + _dummy = _get_dummy_fp32(v.device) + if has_initial_state_flag: + assert initial_state.shape == (N_seqs, H, D, D), ( + f"initial_state shape must be ({N_seqs}, {H}, {D}, {D}), got {initial_state.shape}" + ) + initial_state_fp32 = initial_state.contiguous().reshape(-1) + else: + initial_state_fp32 = _dummy + if has_final_state_flag: + assert final_state.shape == (N_seqs, H, D, D), ( + f"final_state shape must be ({N_seqs}, {H}, {D}, {D}), got {final_state.shape}" + ) + assert final_state.dtype == torch.float32, f"final_state must be fp32, got {final_state.dtype}" + final_state_fp32 = final_state.reshape(-1) + else: + final_state_fp32 = _dummy + + if seq_order is None: + seq_order = _get_identity_order(N_seqs, v.device) + + compiled_fn = _get_compiled_k2( + H, + has_initial_state_flag, + has_final_state_flag, + state_transposed, + v_is_varlen, + ) + stream = _get_current_custream() + # tvm-ffi launch: torch tensors pass straight through with NO shape/dtype + # validation of the positional args. TMA descriptors are (re)built inside + # run_k2 from the dynamic Int32 dims every launch, so one compiled kernel + # serves all shapes. + compiled_fn( + v.view(V_T_total, H, D), + beta, + ws_qd, + ws_kd, + ws_kr, + ws_gt, + ws_inv, + ws_mqk, + out.view(O_T_total, H, D), + cu_seqlens_tiles, + seq_order, + initial_state_fp32, + final_state_fp32, + v_tile_starts, + v_tile_actual_lens, + total_tiles, + O_T_total, + V_T_total, + N_seqs, + stream, + ) diff --git a/cula/ops/lightning/prefill_sm100.py b/cula/ops/lightning/prefill_sm100.py index 3e516f90..93f81f39 100644 --- a/cula/ops/lightning/prefill_sm100.py +++ b/cula/ops/lightning/prefill_sm100.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 ANTGROUP. All rights reserved. +# Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 # Redistribution and use in source and binary forms, with or without diff --git a/cula/utils.py b/cula/utils.py index 99d19cd4..186d9a5a 100644 --- a/cula/utils.py +++ b/cula/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 ANTGROUP. All rights reserved. +# Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 """ diff --git a/tests/test_kda_compare_fla.py b/tests/test_kda_sm100_chunk_vs_fla.py similarity index 100% rename from tests/test_kda_compare_fla.py rename to tests/test_kda_sm100_chunk_vs_fla.py diff --git a/tests/test_kda.py b/tests/test_kda_sm100_chunk_vs_naive.py similarity index 100% rename from tests/test_kda.py rename to tests/test_kda_sm100_chunk_vs_naive.py diff --git a/tests/test_intracard_cp.py b/tests/test_kda_sm100_intracard_cp.py similarity index 97% rename from tests/test_intracard_cp.py rename to tests/test_kda_sm100_intracard_cp.py index 85d077a6..bd44c322 100644 --- a/tests/test_intracard_cp.py +++ b/tests/test_kda_sm100_intracard_cp.py @@ -46,9 +46,9 @@ DEVICE = "cuda" # Tolerances aligned with existing cuLA tests: # * Same-kernel (CP scheduling only): torch.testing.assert_close(atol=1e-2, rtol=1e-2) -# — matches tests/test_chunk_delta_h.py CP block +# — matches tests/test_kda_sm100_recurrence.py CP block # * Cross-impl / vs ref: fla.utils.assert_close(ratio=...) -# — matches tests/test_kda_compare_fla.py +# — matches tests/test_kda_sm100_chunk_vs_fla.py ATOL_SAME_KERNEL = 1e-2 RTOL_SAME_KERNEL = 1e-2 RATIO_VS_REF = 0.005 # RMSE / RMS(ref) — matches FLA test_gated_delta.py fwd @@ -134,7 +134,7 @@ def run_intracard_direct(k, w, u, gk, h0, cu, *, output_final_state=True, save_n post-split occupancy guard rejects; mirror the production caller's graceful fallback to the serial path so configs that don't engage CP still return. """ - from cula.ops.kda.policy import NotSplittableError + from cula.ops.kda.cp_mode import NotSplittableError try: return intracard_fwd_h( @@ -223,7 +223,7 @@ def pytorch_ref(k, w, u, *, gk=None, initial_state=None, cu_seqlens, save_new_va def _assert_same_kernel(name, actual, ref): - """torch.testing.assert_close — matches tests/test_chunk_delta_h.py.""" + """torch.testing.assert_close — matches tests/test_kda_sm100_recurrence.py.""" if actual is None or ref is None: assert actual is ref, f"{name}: one is None and other isn't" return @@ -257,7 +257,7 @@ def assert_cp_splits(cu, H, total_T): def test_forced_cp_not_splittable_raises(): """use_intracard_cp=True on an unsplittable shape must raise NotSplittableError.""" - from cula.ops.kda.policy import NotSplittableError + from cula.ops.kda.cp_mode import NotSplittableError # A single one-chunk sequence cannot be meaningfully split. cu = torch.tensor([0, BT], dtype=torch.int32, device=DEVICE) @@ -289,7 +289,7 @@ def test_cp_autodispatch_matches_baseline(seq_lens, H, use_gk): """CP auto-dispatch output equals no-CP baseline (same kernel). Tolerance: `torch.testing.assert_close(atol=1e-2, rtol=1e-2)` — - matches the CP block in tests/test_chunk_delta_h.py. + matches the CP block in tests/test_kda_sm100_recurrence.py. """ k, w, u, gk, _, cu = make_varlen_inputs(seq_lens, H, use_gk=use_gk) h_base, v_base, _ = run_cula_no_cp(k, w, u, gk, None, cu) diff --git a/tests/test_fwd_o.py b/tests/test_kda_sm100_output.py similarity index 99% rename from tests/test_fwd_o.py rename to tests/test_kda_sm100_output.py index 9ef7b213..d4e2121d 100644 --- a/tests/test_fwd_o.py +++ b/tests/test_kda_sm100_output.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -# Copyright (c) 2025 ANTGROUP. All rights reserved. +# Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 """ diff --git a/tests/test_chunk_delta_h.py b/tests/test_kda_sm100_recurrence.py similarity index 99% rename from tests/test_chunk_delta_h.py rename to tests/test_kda_sm100_recurrence.py index 75223a15..195f1217 100644 --- a/tests/test_chunk_delta_h.py +++ b/tests/test_kda_sm100_recurrence.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -# Copyright (c) 2025 ANTGROUP. All rights reserved. +# Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 """ diff --git a/tests/test_kda_sm90_intracard_cp.py b/tests/test_kda_sm90_intracard_cp.py new file mode 100644 index 00000000..eb567400 --- /dev/null +++ b/tests/test_kda_sm90_intracard_cp.py @@ -0,0 +1,314 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Intracard-CP end-to-end correctness: CP vs serial, CP vs FLA, determinism.""" + +import pytest +import torch +from fla.ops.kda.chunk import chunk_kda as fla_chunk_kda +from fla.utils import assert_close + +from cula.kda import flashkda_prefill as cula_kda_prefill +from cula.ops.kda.sm90.cp import intracard_prefill +from cula.ops.kda.sm90.fwd import D, flash_kda_fwd + +H = 8 +SCALE = D**-0.5 +LB = -5.0 +TOL_MAX = 1e-2 +TOL_RMSE = 4e-3 +TOL_FLA = 5e-3 + +needs_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _make_inputs(B, T, seed=0): + torch.manual_seed(seed) + dev = torch.device("cuda") + q = torch.randn(B, T, H, D, dtype=torch.bfloat16, device=dev) + k = torch.randn(B, T, H, D, dtype=torch.bfloat16, device=dev) + v = torch.randn(B, T, H, D, dtype=torch.bfloat16, device=dev) + g = torch.randn(B, T, H, D, dtype=torch.bfloat16, device=dev) + beta = torch.randn(B, T, H, dtype=torch.bfloat16, device=dev) + A_log = torch.randn(H, dtype=torch.float32, device=dev) + dt_bias = torch.randn(H, D, dtype=torch.float32, device=dev) + return q, k, v, g, beta, A_log, dt_bias + + +def _make_fla_inputs(T, *, with_state, n_state, seed): + torch.manual_seed(seed) + dev = torch.device("cuda") + q = torch.rand(1, T, H, D, dtype=torch.bfloat16, device=dev) + k = torch.rand(1, T, H, D, dtype=torch.bfloat16, device=dev) + v = torch.rand(1, T, H, D, dtype=torch.bfloat16, device=dev) + g = torch.randn(1, T, H, D, dtype=torch.bfloat16, device=dev) + A_log = torch.randn(H, dtype=torch.float32, device=dev) + dt_bias = torch.randn(H * D, dtype=torch.float32, device=dev) + beta = torch.randn(1, T, H, dtype=torch.float32, device=dev).sigmoid().to(torch.bfloat16) + h0 = torch.randn(n_state, H, D, D, dtype=torch.float32, device=dev) if with_state else None + return q, k, v, g, beta, A_log, dt_bias, h0 + + +def _rel_max(a, b): + return (a.float() - b.float()).abs().max().item() / max(b.float().abs().max().item(), 1e-6) + + +def _rel_rmse(a, b): + a, b = a.float(), b.float() + return (a - b).pow(2).mean().sqrt().item() / max(b.pow(2).mean().sqrt().item(), 1e-6) + + +def _assert_cp_matches(actual, ref, name): + rrmse = _rel_rmse(actual, ref) + assert rrmse < TOL_RMSE, f"{name}: rel_rmse {rrmse:.2e} >= {TOL_RMSE}" + rmax = _rel_max(actual, ref) + assert rmax < TOL_MAX, f"{name}: rel_max {rmax:.2e} >= {TOL_MAX}" + + +def _alloc_final(n_seqs, device): + return torch.empty(n_seqs, H, D, D, dtype=torch.float32, device=device) + + +def _run_serial(q, k, v, g, beta, A_log, dt_bias, init=None, want_final=True, cu=None, transposed=False): + n = (cu.numel() - 1 if cu is not None else q.shape[0]) if want_final else 0 + out = torch.empty_like(v) + fin = _alloc_final(n, q.device) if want_final else None + flash_kda_fwd( + q, + k, + v, + g, + beta, + scale=SCALE, + out=out, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=LB, + initial_state=init, + final_state=fin, + cu_seqlens=cu, + state_transposed=transposed, + ) + return out, fin + + +def _run_cp(q, k, v, g, beta, A_log, dt_bias, init=None, want_final=True, cu=None, transposed=False, s_split=4): + n = (cu.numel() - 1 if cu is not None else q.shape[0]) if want_final else 0 + out = torch.empty_like(v) + fin = _alloc_final(n, q.device) if want_final else None + intracard_prefill( + q, + k, + v, + g, + beta, + scale=SCALE, + out=out, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=LB, + initial_state=init, + final_state=fin, + cu_seqlens=cu, + state_transposed=transposed, + s_split=s_split, + ) + return out, fin + + +@needs_cuda +@pytest.mark.parametrize("s_split", [1, 2, 4, 7]) +def test_cp_matches_serial_fixed(s_split): + q, k, v, g, beta, A_log, dt_bias = _make_inputs(1, 2048) + out_ref, fin_ref = _run_serial(q, k, v, g, beta, A_log, dt_bias, None, True) + out_cp, fin_cp = _run_cp(q, k, v, g, beta, A_log, dt_bias, None, True, s_split=s_split) + _assert_cp_matches(out_cp, out_ref, "o") + _assert_cp_matches(fin_cp, fin_ref, "ht") + + +@needs_cuda +def test_cp_matches_serial_fixed_b2_with_state(): + q, k, v, g, beta, A_log, dt_bias = _make_inputs(2, 1024, seed=3) + init = torch.randn(2, H, D, D, dtype=torch.float32, device="cuda") + out_ref, fin_ref = _run_serial(q, k, v, g, beta, A_log, dt_bias, init, True) + out_cp, fin_cp = _run_cp(q, k, v, g, beta, A_log, dt_bias, init, True, s_split=4) + _assert_cp_matches(out_cp, out_ref, "o") + _assert_cp_matches(fin_cp, fin_ref, "ht") + + +@needs_cuda +def test_cp_matches_serial_varlen(): + lens = [1024, 512, 2048, 256] + T = sum(lens) + q, k, v, g, beta, A_log, dt_bias = _make_inputs(1, T, seed=7) + cu = torch.tensor([0] + list(torch.tensor(lens).cumsum(0)), dtype=torch.int32, device="cuda") + init = torch.randn(len(lens), H, D, D, dtype=torch.float32, device="cuda") + out_ref, fin_ref = _run_serial(q, k, v, g, beta, A_log, dt_bias, init, True, cu=cu) + out_cp, fin_cp = _run_cp(q, k, v, g, beta, A_log, dt_bias, init, True, cu=cu, s_split=4) + _assert_cp_matches(out_cp, out_ref, "o") + _assert_cp_matches(fin_cp, fin_ref, "ht") + + +@needs_cuda +def test_cp_matches_serial_state_transposed(): + q, k, v, g, beta, A_log, dt_bias = _make_inputs(1, 1024, seed=11) + init = torch.randn(1, H, D, D, dtype=torch.float32, device="cuda") + out_ref, fin_ref = _run_serial(q, k, v, g, beta, A_log, dt_bias, init, True, transposed=True) + out_cp, fin_cp = _run_cp(q, k, v, g, beta, A_log, dt_bias, init, True, transposed=True, s_split=4) + _assert_cp_matches(out_cp, out_ref, "o") + _assert_cp_matches(fin_cp, fin_ref, "ht") + + +@needs_cuda +def test_cp_no_final_state(): + q, k, v, g, beta, A_log, dt_bias = _make_inputs(1, 512, seed=13) + out_ref, _ = _run_serial(q, k, v, g, beta, A_log, dt_bias, None, False) + out_cp, _ = _run_cp(q, k, v, g, beta, A_log, dt_bias, None, False, s_split=2) + _assert_cp_matches(out_cp, out_ref, "o") + + +@needs_cuda +@pytest.mark.parametrize( + "lens", + [ + [1024, 1, 63, 65, 129], + [28679, 4096], + [40007], + [32768, 100], + ], +) +def test_cp_matches_serial_varlen_nonaligned(lens): + T = sum(lens) + q, k, v, g, beta, A_log, dt_bias = _make_inputs(1, T, seed=5) + cu = torch.tensor([0] + list(torch.tensor(lens).cumsum(0)), dtype=torch.int32, device="cuda") + out_ref, fin_ref = _run_serial(q, k, v, g, beta, A_log, dt_bias, None, True, cu=cu) + out_cp, fin_cp = _run_cp(q, k, v, g, beta, A_log, dt_bias, None, True, cu=cu, s_split=8) + _assert_cp_matches(out_cp, out_ref, "o") + _assert_cp_matches(fin_cp, fin_ref, "ht") + + +@needs_cuda +@pytest.mark.parametrize("T", [100, 4100, 8197]) +def test_cp_matches_serial_dense_nonaligned(T): + q, k, v, g, beta, A_log, dt_bias = _make_inputs(1, T, seed=9) + out_ref, fin_ref = _run_serial(q, k, v, g, beta, A_log, dt_bias, None, True) + out_cp, fin_cp = _run_cp(q, k, v, g, beta, A_log, dt_bias, None, True, s_split=8) + _assert_cp_matches(out_cp, out_ref, "o") + _assert_cp_matches(fin_cp, fin_ref, "ht") + + +def _check_cp_vs_fla(T, *, with_state, cu, seed): + n_state = (cu.numel() - 1) if cu is not None else 1 + q, k, v, g, beta, A_log, dt_bias, h0 = _make_fla_inputs(T, with_state=with_state, n_state=n_state, seed=seed) + cu_cpu = cu.cpu() if cu is not None else None + with torch.no_grad(): + ref_o, ref_ht = fla_chunk_kda( + q, + k, + v, + g, + beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=h0, + cu_seqlens=cu, + cu_seqlens_cpu=cu_cpu, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=LB, + ) + h0_vk = h0.transpose(-2, -1).contiguous() if h0 is not None else None + cp_o, cp_ht_vk = cula_kda_prefill( + q, + k, + v, + g, + beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=h0_vk, + cu_seqlens=cu, + output_final_state=True, + safe_gate=True, + lower_bound=LB, + use_intracard_cp=True, + ) + assert_close("o", ref_o, cp_o, TOL_FLA) + assert_close("ht", ref_ht, cp_ht_vk.transpose(-2, -1), TOL_FLA) + + +@needs_cuda +@pytest.mark.parametrize("T", [8192, 16384]) +def test_cp_vs_fla_dense(T): + _check_cp_vs_fla(T, with_state=False, cu=None, seed=42) + + +@needs_cuda +def test_cp_vs_fla_dense_with_state(): + _check_cp_vs_fla(16384, with_state=True, cu=None, seed=42) + + +@needs_cuda +def test_cp_vs_fla_varlen_with_state(): + lens = [16384, 8192, 4096] + cu = torch.tensor([0] + list(torch.tensor(lens).cumsum(0)), dtype=torch.int32, device="cuda") + _check_cp_vs_fla(sum(lens), with_state=True, cu=cu, seed=7) + + +_DETERMINISM_ITERS = 10000 + + +@needs_cuda +def test_cp_determinism(): + q, k, v, g, beta, A_log, dt_bias = _make_inputs(1, 4096, seed=17) + out0, fin0 = _run_cp(q, k, v, g, beta, A_log, dt_bias, None, True, s_split=8) + for i in range(_DETERMINISM_ITERS): + out, fin = _run_cp(q, k, v, g, beta, A_log, dt_bias, None, True, s_split=8) + assert torch.equal(out, out0), f"non-deterministic out at iter {i}" + assert torch.equal(fin, fin0), f"non-deterministic ht at iter {i}" + + +@needs_cuda +def test_cp_determinism_varlen(): + lens = [1024, 1, 63, 65, 129] + cu = torch.tensor([0] + list(torch.tensor(lens).cumsum(0)), dtype=torch.int32, device="cuda") + q, k, v, g, beta, A_log, dt_bias = _make_inputs(1, sum(lens), seed=5) + out0, fin0 = _run_cp(q, k, v, g, beta, A_log, dt_bias, None, True, cu=cu, s_split=8) + for i in range(_DETERMINISM_ITERS): + out, fin = _run_cp(q, k, v, g, beta, A_log, dt_bias, None, True, cu=cu, s_split=8) + assert torch.equal(out, out0), f"non-deterministic varlen out at iter {i}" + assert torch.equal(fin, fin0), f"non-deterministic varlen ht at iter {i}" + + +@needs_cuda +@pytest.mark.parametrize( + "lens", + [ + pytest.param([14336] + [128] * 16, id="skewed-1long-16short"), + pytest.param([8192, 2048], id="mixed-lengths"), + ], +) +def test_cp_matches_serial_ragged_auto_plan(lens): + """Ragged varlen batches planned by the duration-balanced auto planner + (s_split=None). Splitting must not change any accumulation order, so CP + output is required to be BIT-IDENTICAL to the serial path.""" + from cula.ops.kda.sm90.cp.plan import plan_auto + from cula.utils import get_device_sm_count + + seq_tiles = [(sl + 15) // 16 for sl in lens] + plan = plan_auto(seq_tiles, H, get_device_sm_count(torch.device("cuda"))) + assert not plan.trivial, "precondition: auto planner must split this batch" + + total = sum(lens) + q, k, v, g, beta, A_log, dt_bias = _make_inputs(1, total, seed=7) + cu = [0] + for sl in lens: + cu.append(cu[-1] + sl) + cu = torch.tensor(cu, dtype=torch.int32, device=q.device) + + ref_o, ref_f = _run_serial(q, k, v, g, beta, A_log, dt_bias, cu=cu) + out, fin = _run_cp(q, k, v, g, beta, A_log, dt_bias, cu=cu, s_split=None) + assert torch.equal(out, ref_o) + assert torch.equal(fin, ref_f) diff --git a/tests/test_kda_sm90_prefill_vs_fla.py b/tests/test_kda_sm90_prefill_vs_fla.py new file mode 100644 index 00000000..e1412b0c --- /dev/null +++ b/tests/test_kda_sm90_prefill_vs_fla.py @@ -0,0 +1,196 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang + +"""Precision tests for SM90 CuTeDSL prefill vs FLA baseline. + +Forward-only: the SM90 path has no backward, so only ``o`` and ``ht`` are compared. + +NOTE: the SM90 kernel keeps its recurrence state in **VK-transposed** layout +(``[B, H, V, K]``) while FLA uses KV layout (``[B, H, K, V]``). +``initial_state`` is transposed when passed to cuLA, and cuLA's ``ht`` is +transposed back to KV layout before comparison. +""" + +import pytest +import torch +from fla.ops.kda.chunk import chunk_kda as fla_chunk_kda +from fla.utils import assert_close, device + +from cula.kda import flashkda_prefill as cula_kda_prefill + +pytestmark = [ + pytest.mark.sm90_only, + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA"), +] + +D = 128 + +# (B, T, H, with_state) — non-aligned T exercises tail-chunk handling (CHUNK=16). +_DENSE = [ + pytest.param(1, 63, 1, False, marks=pytest.mark.kda_fast, id="B1-T63-H1-tail"), + pytest.param(1, 256, 2, False, marks=pytest.mark.kda_fast, id="B1-T256-H2-aligned"), + pytest.param(1, 500, 2, False, marks=pytest.mark.kda_fast, id="B1-T500-H2-tail"), + pytest.param(1, 512, 2, True, marks=pytest.mark.kda_fast, id="B1-T512-H2-init_state"), + pytest.param(2, 512, 2, False, marks=pytest.mark.kda_fast, id="B2-T512-H2-stride"), + pytest.param(1, 1024, 4, True, marks=pytest.mark.kda_fast, id="B1-T1024-H4-init_state"), + pytest.param(2, 1024, 4, False, marks=pytest.mark.kda_slow, id="B2-T1024-H4"), +] + +_VARLEN = [ + pytest.param([0, 15], 2, False, marks=pytest.mark.kda_fast, id="cu[0,15]-H2"), + pytest.param([0, 256, 500, 1000], 4, False, marks=pytest.mark.kda_slow, id="cu[0,256,500,1000]-H4"), + pytest.param([0, 15, 100, 300, 1200], 4, True, marks=pytest.mark.kda_slow, id="cu[0,15,100,300,1200]-H4-init"), +] + + +def _make_inputs(B, T, H, *, with_state, n_state=None, seed=42): + torch.manual_seed(seed) + q = torch.rand(B, T, H, D, dtype=torch.bfloat16, device=device) + k = torch.rand(B, T, H, D, dtype=torch.bfloat16, device=device) + v = torch.rand(B, T, H, D, dtype=torch.bfloat16, device=device) + g = torch.randn(B, T, H, D, dtype=torch.bfloat16, device=device) + A_log = torch.randn(H, dtype=torch.float32, device=device) + dt_bias = torch.randn(H * D, dtype=torch.float32, device=device) + beta = torch.randn(B, T, H, dtype=torch.float32, device=device).sigmoid().to(torch.bfloat16) + n_state = B if n_state is None else n_state + h0 = torch.randn(n_state, H, D, D, dtype=torch.float32, device=device) if with_state else None + return q, k, v, g, beta, A_log, dt_bias, h0 + + +@pytest.mark.parametrize(("B", "T", "H", "with_state"), _DENSE) +def test_prefill_dense_matches_fla(B, T, H, with_state): + q, k, v, g, beta, A_log, dt_bias, h0 = _make_inputs(B, T, H, with_state=with_state) + + with torch.no_grad(): + ref_o, ref_ht = fla_chunk_kda( + q, + k, + v, + g, + beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=h0, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=-5.0, + ) + + h0_vk = h0.transpose(-2, -1).contiguous() if h0 is not None else None + tri_o, tri_ht_vk = cula_kda_prefill( + q, + k, + v, + g, + beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=h0_vk, + output_final_state=True, + safe_gate=True, + lower_bound=-5.0, + ) + tri_ht = tri_ht_vk.transpose(-2, -1) + + assert_close("o", ref_o, tri_o, 0.005) + assert_close("ht", ref_ht, tri_ht, 0.005) + + +@pytest.mark.parametrize(("cu_seqlens", "H", "with_state"), _VARLEN) +def test_prefill_varlen_matches_fla(cu_seqlens, H, with_state): + total_t = cu_seqlens[-1] + q, k, v, g, beta, A_log, dt_bias, h0 = _make_inputs(1, total_t, H, with_state=with_state, n_state=len(cu_seqlens) - 1) + cu = torch.tensor(cu_seqlens, dtype=torch.int32, device=device) + cu_cpu = cu.cpu() + + with torch.no_grad(): + ref_o, ref_ht = fla_chunk_kda( + q, + k, + v, + g, + beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=h0, + cu_seqlens=cu, + cu_seqlens_cpu=cu_cpu, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=-5.0, + ) + + h0_vk = h0.transpose(-2, -1).contiguous() if h0 is not None else None + tri_o, tri_ht_vk = cula_kda_prefill( + q, + k, + v, + g, + beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=h0_vk, + cu_seqlens=cu, + cu_seqlens_cpu=cu_cpu, + output_final_state=True, + safe_gate=True, + lower_bound=-5.0, + ) + tri_ht = tri_ht_vk.transpose(-2, -1) + + assert_close("o", ref_o, tri_o, 0.005) + assert_close("ht", ref_ht, tri_ht, 0.005) + + +def test_prefill_workspace_reuse_across_shapes(): + """Back-to-back calls with different shapes share the grow-only workspace + arena; earlier shapes must keep matching FLA after the arena has been + re-carved for larger and smaller shapes (including a tail-chunk one).""" + shapes = [(1, 512, 2), (1, 1024, 2), (2, 512, 2), (1, 500, 2), (1, 512, 2)] + for B, T, H in shapes: + q, k, v, g, beta, A_log, dt_bias, _ = _make_inputs(B, T, H, with_state=False) + with torch.no_grad(): + ref_o, _ = fla_chunk_kda( + q, + k, + v, + g, + beta, + A_log=A_log, + dt_bias=dt_bias, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=-5.0, + ) + tri_o, _ = cula_kda_prefill( + q, + k, + v, + g, + beta, + A_log=A_log, + dt_bias=dt_bias, + output_final_state=True, + safe_gate=True, + lower_bound=-5.0, + ) + assert_close("o", ref_o, tri_o, 0.005) diff --git a/tests/test_la_decode.py b/tests/test_lightning_decode.py similarity index 100% rename from tests/test_la_decode.py rename to tests/test_lightning_decode.py diff --git a/tests/test_la_decode_pool.py b/tests/test_lightning_decode_pool.py similarity index 100% rename from tests/test_la_decode_pool.py rename to tests/test_lightning_decode_pool.py diff --git a/tests/test_lightning_attn.py b/tests/test_lightning_sm100_prefill.py similarity index 100% rename from tests/test_lightning_attn.py rename to tests/test_lightning_sm100_prefill.py diff --git a/tests/test_ptx_umma_masked.py b/tests/test_ptx_umma_masked.py index 3d6859b3..38d6389a 100644 --- a/tests/test_ptx_umma_masked.py +++ b/tests/test_ptx_umma_masked.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 ANTGROUP. All rights reserved. +# Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 """ diff --git a/tests/test_ptx_umma_ws.py b/tests/test_ptx_umma_ws.py index 210028b4..46a55709 100644 --- a/tests/test_ptx_umma_ws.py +++ b/tests/test_ptx_umma_ws.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 ANTGROUP. All rights reserved. +# Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 """