From 1a2cac21278bdd97e363d17201eb10be25a93e83 Mon Sep 17 00:00:00 2001 From: Yang Zhe <108724832+zheyang0825@users.noreply.github.com> Date: Fri, 14 Aug 2026 12:18:43 +0800 Subject: [PATCH] feat(kda): integrate persistent SM90 bwd intra kernel --- BENCHMARK_KDA_BWD_INTRA_SM90.md | 72 ++ README.md | 5 + benchmarks/bench_kda_bwd_intra_sm90.py | 198 ++++++ benchmarks/ncu_profile_kda_bwd_intra.py | 57 ++ csrc/api/kda_bwd_intra.cu | 162 +++++ csrc/api/pybind.cu | 21 + csrc/kda/sm90/bwd/kda_bwd_intra_sm90.cu | 848 ++++++++++++++++++++++++ csrc/kda/sm90/bwd/kda_config.h | 26 + cula/kda/chunk_intra.py | 122 +++- cula/ops/kda/sm90/bwd_intra.py | 87 +++ setup.py | 4 + tests/test_kda_sm90_bwd_intra.py | 136 ++++ 12 files changed, 1736 insertions(+), 2 deletions(-) create mode 100644 BENCHMARK_KDA_BWD_INTRA_SM90.md create mode 100644 benchmarks/bench_kda_bwd_intra_sm90.py create mode 100644 benchmarks/ncu_profile_kda_bwd_intra.py create mode 100644 csrc/api/kda_bwd_intra.cu create mode 100644 csrc/kda/sm90/bwd/kda_bwd_intra_sm90.cu create mode 100644 csrc/kda/sm90/bwd/kda_config.h create mode 100644 cula/ops/kda/sm90/bwd_intra.py create mode 100644 tests/test_kda_sm90_bwd_intra.py diff --git a/BENCHMARK_KDA_BWD_INTRA_SM90.md b/BENCHMARK_KDA_BWD_INTRA_SM90.md new file mode 100644 index 00000000..7ff21036 --- /dev/null +++ b/BENCHMARK_KDA_BWD_INTRA_SM90.md @@ -0,0 +1,72 @@ +# Benchmark Results — KDA Backward Intra (SM90) + +> Measured on 2026-08-14. + +> **GPU:** NVIDIA L20X, compute capability 9.0 | **CUDA:** 12.9 | +> **PyTorch:** 2.9.1+cu129 | **Triton:** 3.5.1 + +This report compares cuLA's persistent CUDA C++ KDA intra-chunk backward +kernel with the Triton implementation in flash-linear-attention (FLA). The +kernel uses warp-level `mma.sync.m16n8k8`, BF16 Q/K/beta inputs, FP32 gate and +gradient inputs, head dimension 128, and chunk size 64. + +Each configuration uses `triton.testing.do_bench` with 25 ms of warmup and a +100 ms measurement window. The table reports the median latency. The Triton +cache was cleared before each FLA version was measured. Variable-length cases +contain eight quasi-balanced sequences with the indicated total token count. + +## FLA v0.5.0 + +| Configuration | cuLA CUDA (ms) | FLA Triton (ms) | Speedup | +|---|---:|---:|---:| +| H=32, uniform, T=8192, N=1 | 0.502 | 0.808 | **1.61x** | +| H=32, uniform, T=32768, N=1 | 1.850 | 3.198 | **1.73x** | +| H=32, varlen, T=8192, N=8 | 0.505 | 0.813 | **1.61x** | +| H=32, varlen, T=32768, N=8 | 1.842 | 3.190 | **1.73x** | +| H=64, uniform, T=8192, N=1 | 0.952 | 1.612 | **1.69x** | +| H=64, uniform, T=32768, N=1 | 3.631 | 6.398 | **1.76x** | +| H=64, varlen, T=8192, N=8 | 0.949 | 1.610 | **1.70x** | +| H=64, varlen, T=32768, N=8 | 3.630 | 6.413 | **1.77x** | + +The geometric-mean speedup over FLA v0.5.0 is **1.699x**. + +## FLA v0.4.2 + +FLA v0.4.2 is about 4.5–5.5% faster than v0.5.0 for these shapes. The cuLA +kernel latency remains effectively unchanged, so its geometric-mean speedup is +lower against this baseline. + +| Configuration | cuLA CUDA (ms) | FLA Triton (ms) | Speedup | +|---|---:|---:|---:| +| H=32, uniform, T=8192, N=1 | 0.503 | 0.768 | **1.53x** | +| H=32, uniform, T=32768, N=1 | 1.851 | 3.035 | **1.64x** | +| H=32, varlen, T=8192, N=8 | 0.506 | 0.776 | **1.53x** | +| H=32, varlen, T=32768, N=8 | 1.852 | 3.046 | **1.64x** | +| H=64, uniform, T=8192, N=1 | 0.952 | 1.524 | **1.60x** | +| H=64, uniform, T=32768, N=1 | 3.636 | 6.071 | **1.67x** | +| H=64, varlen, T=8192, N=8 | 0.951 | 1.534 | **1.61x** | +| H=64, varlen, T=32768, N=8 | 3.627 | 6.086 | **1.68x** | + +The geometric-mean speedup over FLA v0.4.2 is **1.612x**. + +## Correctness + +The SM90 correctness suite compares `dq`, `dk`, `db`, and `dg` against FLA for +fixed-length, ragged variable-length, and dense-batch inputs. It also checks +deterministic output, dispatcher behavior, device validation, and the +unsupported-beta fallback. All nine tests pass on SM90. + +## Reproduction + +```bash +python benchmarks/bench_kda_bwd_intra_sm90.py \ + --heads 32 64 \ + --warmup 25 \ + --rep 100 + +python -m pytest tests/test_kda_sm90_bwd_intra.py -v +``` + +The repository pins FLA v0.5.0. To reproduce the v0.4.2 comparison, check out +the `v0.4.2` tag in `third_party/flash-linear-attention` and reinstall that +submodule in editable mode before running the same command. diff --git a/README.md b/README.md index 91b25a08..a1169ce2 100644 --- a/README.md +++ b/README.md @@ -157,6 +157,7 @@ See [BENCHMARK_GB200_CUDA_130.md](BENCHMARK_GB200_CUDA_130.md) tested with CUDA **Hopper (SM90)** See [BENCHMARK_H200.md](BENCHMARK_H200.md) for CuTe DSL FlashKDA results on an H200 141GB with CUDA 12.9. +See [BENCHMARK_KDA_BWD_INTRA_SM90.md](BENCHMARK_KDA_BWD_INTRA_SM90.md) for the persistent CUDA C++ KDA intra-chunk backward benchmark. **Highlights:** - **KDA Modular Forward (Blackwell):** **avg 1.33x** speedup on fixed-length, **avg 1.35x** on variable-length (18 configs, uniform/skewed/random). @@ -164,6 +165,7 @@ See [BENCHMARK_H200.md](BENCHMARK_H200.md) for CuTe DSL FlashKDA results on an H - **Lightning Attention Varlen (Blackwell):** **avg 1.47x** speedup across 126 configs (uniform/skewed/random). - **FlashKDA Prefill (Hopper):** **avg 2.72x** speedup over FLA across 28 fixed-length and variable-length configs, up to **7.56x**. - **FlashKDA Intracard CP (Hopper):** **4.29x geo-mean** speedup over serial FlashKDA on 28 CP-engaged long-sequence configs, up to **7.83x**. +- **KDA Backward Intra (Hopper):** **1.70x geo-mean** speedup over FLA v0.5.0 across eight fixed-length and variable-length configs. To reproduce the benchmark suites directly: @@ -176,6 +178,7 @@ python benchmarks/bench_la_decode_vs_fla.py --heads 64 --head-dim 128 # Hopper (SM90) python benchmarks/bench_kda_sm90_prefill.py --mode both python benchmarks/bench_kda_sm90_cp.py +python benchmarks/bench_kda_bwd_intra_sm90.py --heads 32 64 ``` ## Tests @@ -187,6 +190,8 @@ python -m pytest tests/test_kda_sm100_chunk_vs_fla.py -v python -m pytest tests/test_kda_sm100_chunk_vs_naive.py -v # 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 the persistent SM90 KDA intra-chunk backward kernel (vs FLA) +python -m pytest tests/test_kda_sm90_bwd_intra.py -v # Tests for Lightning Attention prefill on SM100 python tests/test_lightning_sm100_prefill.py # Tests for the SM90 Lightning public dispatch, semantics, and kernel structure diff --git a/benchmarks/bench_kda_bwd_intra_sm90.py b/benchmarks/bench_kda_bwd_intra_sm90.py new file mode 100644 index 00000000..5178217d --- /dev/null +++ b/benchmarks/bench_kda_bwd_intra_sm90.py @@ -0,0 +1,198 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark and determinism stress test for the persistent CUDA SM90 bwd-intra kernel. + +The kernel uses warp-level ``mma.sync.m16n8k8`` and can also run on +SM100/SM103. This script calls the low-level kernel directly so the same +benchmark can validate every supported architecture. +""" + +import argparse +import itertools +import pathlib +import random +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +import torch # noqa: E402 +import triton # noqa: E402 +from fla.ops.kda.chunk_intra import chunk_kda_bwd_intra as fla_bwd_intra # noqa: E402 +from fla.ops.utils import prepare_chunk_indices # noqa: E402 + +from cula.ops.kda.sm90.bwd_intra import kda_bwd_intra_mma # noqa: E402 + +K = 128 +BT = 64 +DEVICE = torch.device("cuda") + + +def _balanced_lengths(total_tokens: int, num_seqs: int) -> list[int]: + base, remainder = divmod(total_tokens, num_seqs) + return [base] * (num_seqs - remainder) + [base + 1] * remainder + + +def _quasi_balanced_lengths(total_tokens: int, num_seqs: int, seed: int = 42) -> list[int]: + rng = random.Random(seed) + weights = [rng.uniform(1.0, 2.5) for _ in range(num_seqs)] + lengths = [max(BT, int(total_tokens * weight / sum(weights))) for weight in weights] + delta = total_tokens - sum(lengths) + order = sorted(range(num_seqs), key=lengths.__getitem__, reverse=delta < 0) + for idx in itertools.cycle(order): + if delta == 0: + break + if delta < 0 and lengths[idx] == BT: + continue + lengths[idx] += 1 if delta > 0 else -1 + delta += -1 if delta > 0 else 1 + return lengths + + +def _make_inputs(lengths: list[int], heads: int, beta_dtype: torch.dtype = torch.bfloat16): + torch.manual_seed(42) + total_tokens = sum(lengths) + offsets = list(itertools.accumulate(lengths, initial=0)) + cu_seqlens = torch.tensor(offsets, device=DEVICE, dtype=torch.int32) + chunk_indices = prepare_chunk_indices(cu_seqlens.to(torch.long), BT).to(torch.int32).contiguous() + + q = torch.randn(1, total_tokens, heads, K, device=DEVICE, dtype=torch.bfloat16) + k = torch.randn_like(q) + g = torch.randn(1, total_tokens, heads, K, device=DEVICE, dtype=torch.float32) / 10 + beta = torch.randn(1, total_tokens, heads, device=DEVICE, dtype=beta_dtype) + d_aq = torch.randn(1, total_tokens, heads, BT, device=DEVICE, dtype=torch.float32) + d_ak = torch.randn_like(d_aq) + dq = torch.randn(1, total_tokens, heads, K, device=DEVICE, dtype=torch.float32) + dk = torch.randn_like(dq) + db = torch.randn(1, total_tokens, heads, device=DEVICE, dtype=torch.float32) + dg = torch.randn_like(dq) + return (q, k, g, beta, d_aq, d_ak, dq, dk, db, dg, cu_seqlens, chunk_indices) + + +def _prepare_cula(inputs): + q, k, _, _, _, _, _, _, db, dg, _, _ = inputs + outputs = ( + torch.empty_like(q), + torch.empty_like(k), + torch.empty_like(db), + torch.empty_like(dg), + ) + + def run(): + return kda_bwd_intra_mma(*inputs, *outputs, BT) + + return run, outputs + + +def _run_fla(inputs): + return fla_bwd_intra(*inputs, chunk_size=BT, safe_gate=True) + + +def _error_metrics(reference: torch.Tensor, actual: torch.Tensor) -> tuple[float, float, float]: + reference = reference.float() + actual = actual.float() + diff = reference - actual + rmse = diff.square().mean().sqrt().item() + relative_rmse = rmse / (reference.square().mean().sqrt().item() + 1e-8) + relative_max = diff.abs().max().item() / (reference.abs().max().item() + 1e-8) + return rmse, relative_rmse, relative_max + + +def check_accuracy(lengths: list[int], heads: int) -> tuple[float, float]: + inputs = _make_inputs(lengths, heads) + run_cula, outputs = _prepare_cula(inputs) + run_cula() + reference = _run_fla(inputs) + torch.cuda.synchronize() + + names = ("dq", "dk", "db", "dg") + metrics = [_error_metrics(ref, out) for ref, out in zip(reference, outputs)] + for name, (rmse, relative_rmse, relative_max) in zip(names, metrics): + print(f" {name}: RMSE={rmse:.6e} rRMSE={relative_rmse:.6e} rMAX={relative_max:.6e}") + return max(value[1] for value in metrics), max(value[2] for value in metrics) + + +def check_determinism(iters: int, heads: int = 4, total_tokens: int = 512, num_seqs: int = 4) -> None: + lengths = _balanced_lengths(total_tokens, num_seqs) + inputs = _make_inputs(lengths, heads) + run, outputs = _prepare_cula(inputs) + run() + torch.cuda.synchronize() + reference = tuple(output.clone() for output in outputs) + + for iteration in range(iters): + run() + for name, output, expected in zip(("dq", "dk", "db", "dg"), outputs, reference): + if not torch.equal(output, expected): + max_diff = (output.float() - expected.float()).abs().max().item() + raise AssertionError(f"{name} is non-deterministic at iteration {iteration}: max_diff={max_diff}") + if (iteration + 1) % 1000 == 0 or iteration + 1 == iters: + print(f" determinism: {iteration + 1}/{iters}", flush=True) + + +def _do_bench(fn, warmup: int, rep: int) -> tuple[float, float, float]: + result = triton.testing.do_bench(fn, warmup=warmup, rep=rep, quantiles=[0.5, 0.2, 0.8]) + return tuple(float(value) for value in result) + + +def run_benchmarks(heads_list: list[int], warmup: int, rep: int) -> None: + configs = ( + ("uniform", [8192]), + ("uniform", [32768]), + ("varlen", _quasi_balanced_lengths(8192, 8)), + ("varlen", _quasi_balanced_lengths(32768, 8)), + ) + print(f"{'Config':<32} {'cuLA p50':>11} {'FLA p50':>11} {'speedup':>9} {'cuLA p20-p80':>22} {'FLA p20-p80':>22}") + print("-" * 114) + speedups = [] + for heads in heads_list: + for kind, lengths in configs: + inputs = _make_inputs(lengths, heads) + run_cula, _ = _prepare_cula(inputs) + run_cula() + _run_fla(inputs) + torch.cuda.synchronize() + + cula_p50, cula_p20, cula_p80 = _do_bench(run_cula, warmup, rep) + fla_p50, fla_p20, fla_p80 = _do_bench(lambda: _run_fla(inputs), warmup, rep) + speedup = fla_p50 / cula_p50 + speedups.append(speedup) + label = f"H={heads} {kind} T={sum(lengths)} N={len(lengths)}" + print( + f"{label:<32} {cula_p50:>9.3f}ms {fla_p50:>9.3f}ms {speedup:>8.2f}x " + f"{cula_p20:>8.3f}-{cula_p80:<8.3f} {fla_p20:>8.3f}-{fla_p80:<8.3f}" + ) + relative_rmse, relative_max = check_accuracy(lengths, heads) + print(f" worst: rRMSE={relative_rmse:.6e} rMAX={relative_max:.6e}") + torch.cuda.empty_cache() + geometric_mean = torch.tensor(speedups, dtype=torch.float64).log().mean().exp().item() + print(f"geomean speedup over FLA: {geometric_mean:.3f}x") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--heads", type=int, nargs="+", default=[32, 64]) + parser.add_argument("--warmup", type=int, default=25) + parser.add_argument("--rep", type=int, default=100) + parser.add_argument("--determinism-iters", type=int, default=0) + parser.add_argument("--determinism-only", action="store_true") + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA GPU is required") + capability = torch.cuda.get_device_capability() + print(f"device: {torch.cuda.get_device_name()} SM{capability[0]}{capability[1]}") + print(f"torch: {torch.__version__} cuda: {torch.version.cuda}") + if capability not in ((9, 0), (10, 0), (10, 3)): + raise RuntimeError(f"requires SM90, SM100, or SM103; got SM{capability[0]}{capability[1]}") + + if args.determinism_iters: + check_determinism(args.determinism_iters) + print(f"determinism PASS ({args.determinism_iters} iterations)") + if not args.determinism_only: + run_benchmarks(args.heads, args.warmup, args.rep) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/ncu_profile_kda_bwd_intra.py b/benchmarks/ncu_profile_kda_bwd_intra.py new file mode 100644 index 00000000..691d937f --- /dev/null +++ b/benchmarks/ncu_profile_kda_bwd_intra.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Profile one warmed-up portable KDA bwd-intra kernel launch with NCU.""" + +import argparse +import pathlib +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +import torch # noqa: E402 + +from benchmarks.bench_kda_bwd_intra_sm90 import ( # noqa: E402 + _make_inputs, + _prepare_cula, + _quasi_balanced_lengths, +) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--heads", type=int, default=64) + parser.add_argument("--total-tokens", type=int, default=32768) + parser.add_argument("--num-seqs", type=int, default=8) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA GPU is required") + + lengths = _quasi_balanced_lengths(args.total_tokens, args.num_seqs) + inputs = _make_inputs(lengths, args.heads) + run, outputs = _prepare_cula(inputs) + + # Compile and warm up outside the profiler range. The range below then + # contains exactly one launch of the CUDA/CuTe bwd-intra kernel. + run() + torch.cuda.synchronize() + + cudart = torch.cuda.cudart() + cudart.cudaProfilerStart() + try: + run() + torch.cuda.synchronize() + finally: + cudart.cudaProfilerStop() + + checksum = sum(output.float().sum().item() for output in outputs) + capability = torch.cuda.get_device_capability() + print( + f"profiled H={args.heads} T={sum(lengths)} N={len(lengths)} SM{capability[0]}{capability[1]} checksum={checksum:.6e}" + ) + + +if __name__ == "__main__": + main() diff --git a/csrc/api/kda_bwd_intra.cu b/csrc/api/kda_bwd_intra.cu new file mode 100644 index 00000000..22950289 --- /dev/null +++ b/csrc/api/kda_bwd_intra.cu @@ -0,0 +1,162 @@ +// Copyright 2025-2026 Ant Group Co., Ltd. +// SPDX-License-Identifier: Apache-2.0 + +#include + +#include +#include +#include +#include + +#include "kda/sm90/bwd/kda_config.h" + +namespace sm90 { +void +run_kda_bwd_intra_sm90(KDA_bwd_intra_params& params, cudaStream_t stream); +} + +namespace { + +void +check_cuda_contiguous(at::Tensor const& tensor, char const* name, at::Device const& device) { + TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(tensor.device() == device, name, " must be on ", device, ", got ", tensor.device()); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); +} + +void +check_same_shape(at::Tensor const& tensor, at::Tensor const& expected, char const* name, char const* expected_name) { + TORCH_CHECK( + tensor.sizes() == expected.sizes(), + name, + " must have the same shape as ", + expected_name, + ", got ", + tensor.sizes(), + " vs ", + expected.sizes()); +} + +} // namespace + +void +ChunkKDABwdIntra( + at::Tensor q, + at::Tensor k, + at::Tensor g, + at::Tensor beta, + at::Tensor dAqk, + at::Tensor dAkk, + at::Tensor dq, + at::Tensor dk, + at::Tensor db, + at::Tensor dg, + at::Tensor cu_seqlens, + at::Tensor chunk_indices, + at::Tensor dq_out, + at::Tensor dk_out, + at::Tensor db_out, + at::Tensor dg_out, + int64_t chunk_size) { + auto const device = q.device(); + check_cuda_contiguous(q, "q", device); + check_cuda_contiguous(k, "k", device); + check_cuda_contiguous(g, "g", device); + check_cuda_contiguous(beta, "beta", device); + check_cuda_contiguous(dAqk, "dAqk", device); + check_cuda_contiguous(dAkk, "dAkk", device); + check_cuda_contiguous(dq, "dq", device); + check_cuda_contiguous(dk, "dk", device); + check_cuda_contiguous(db, "db", device); + check_cuda_contiguous(dg, "dg", device); + check_cuda_contiguous(cu_seqlens, "cu_seqlens", device); + check_cuda_contiguous(chunk_indices, "chunk_indices", device); + check_cuda_contiguous(dq_out, "dq_out", device); + check_cuda_contiguous(dk_out, "dk_out", device); + check_cuda_contiguous(db_out, "db_out", device); + check_cuda_contiguous(dg_out, "dg_out", device); + + TORCH_CHECK(q.scalar_type() == at::kBFloat16, "q must be bfloat16"); + TORCH_CHECK(k.scalar_type() == at::kBFloat16, "k must be bfloat16"); + TORCH_CHECK(beta.scalar_type() == at::kFloat, "beta must be float32"); + TORCH_CHECK(dq_out.scalar_type() == at::kBFloat16, "dq_out must be bfloat16"); + TORCH_CHECK(dk_out.scalar_type() == at::kBFloat16, "dk_out must be bfloat16"); + for (auto const& item : { + std::pair{&g, "g"}, + {&dAqk, "dAqk"}, + {&dAkk, "dAkk"}, + {&dq, "dq"}, + {&dk, "dk"}, + {&db, "db"}, + {&dg, "dg"}, + {&db_out, "db_out"}, + {&dg_out, "dg_out"}, + }) { + TORCH_CHECK(item.first->scalar_type() == at::kFloat, item.second, " must be float32"); + } + TORCH_CHECK(cu_seqlens.scalar_type() == at::kInt, "cu_seqlens must be int32"); + TORCH_CHECK(chunk_indices.scalar_type() == at::kInt, "chunk_indices must be int32"); + + TORCH_CHECK(q.dim() == 4, "q must have shape [B, T, H, K]"); + check_same_shape(k, q, "k", "q"); + check_same_shape(g, q, "g", "q"); + check_same_shape(dq, q, "dq", "q"); + check_same_shape(dk, q, "dk", "q"); + check_same_shape(dg, q, "dg", "q"); + check_same_shape(dq_out, q, "dq_out", "q"); + check_same_shape(dk_out, q, "dk_out", "q"); + check_same_shape(dg_out, q, "dg_out", "q"); + TORCH_CHECK(beta.dim() == 3, "beta must have shape [B, T, H]"); + TORCH_CHECK(db.sizes() == beta.sizes(), "db must have the same shape as beta"); + TORCH_CHECK( + beta.size(0) == q.size(0) && beta.size(1) == q.size(1) && beta.size(2) == q.size(2), + "beta shape must match q[:3]"); + TORCH_CHECK( + dAqk.dim() == 4 && dAqk.size(0) == q.size(0) && dAqk.size(1) == q.size(1) && dAqk.size(2) == q.size(2) && + dAqk.size(3) == chunk_size, + "dAqk must have shape [B, T, H, chunk_size]"); + TORCH_CHECK(dAkk.sizes() == dAqk.sizes(), "dAkk must have the same shape as dAqk"); + TORCH_CHECK(cu_seqlens.dim() == 1 && cu_seqlens.numel() >= 2, "cu_seqlens must have shape [num_sequences + 1]"); + TORCH_CHECK( + chunk_indices.dim() == 2 && chunk_indices.size(1) == 2, "chunk_indices must have shape [num_chunks, 2]"); + TORCH_CHECK(db_out.sizes() == beta.sizes(), "db_out must have the same shape as beta"); + + TORCH_CHECK(chunk_size == 64, "chunk_kda_bwd_intra_cuda supports only chunk_size=64, got ", chunk_size); + TORCH_CHECK(q.size(3) == 128, "chunk_kda_bwd_intra_cuda supports only K=128, got ", q.size(3)); + TORCH_CHECK(q.numel() > 0, "q must be non-empty"); + auto const total_q_len = q.size(0) * q.size(1); + TORCH_CHECK(total_q_len <= std::numeric_limits::max(), "B*T exceeds int32 range"); + TORCH_CHECK(q.size(2) <= std::numeric_limits::max(), "H exceeds int32 range"); + TORCH_CHECK(chunk_indices.size(0) > 0, "chunk_indices must contain at least one chunk"); + TORCH_CHECK(chunk_indices.size(0) <= std::numeric_limits::max(), "number of chunks exceeds int32 range"); + TORCH_CHECK( + chunk_indices.size(0) <= std::numeric_limits::max() / q.size(2), "num_chunks * H exceeds int32 range"); + + c10::cuda::CUDAGuard device_guard(device); + KDA_bwd_intra_params params{}; + params.total_q_len = static_cast(total_q_len); + params.h = static_cast(q.size(2)); + params.d = static_cast(q.size(3)); + params.q_ptr = q.data_ptr(); + params.k_ptr = k.data_ptr(); + params.g_ptr = g.data_ptr(); + params.beta_ptr = beta.data_ptr(); + params.dAqk_ptr = dAqk.data_ptr(); + params.dAkk_ptr = dAkk.data_ptr(); + params.dq_ptr = dq.data_ptr(); + params.dk_ptr = dk.data_ptr(); + params.dg_ptr = dg.data_ptr(); + params.cu_seqlens_ptr = cu_seqlens.data_ptr(); + params.chunk_indices_ptr = chunk_indices.data_ptr(); + params.dq_out_ptr = dq_out.data_ptr(); + params.dk_out_ptr = dk_out.data_ptr(); + params.dg_out_ptr = dg_out.data_ptr(); + params.num_chunks = static_cast(chunk_indices.size(0)); + + auto const num_k_tiles = q.size(3) / 32; + auto db_partials = at::zeros({num_k_tiles, beta.size(0), beta.size(1), beta.size(2)}, db.options()); + params.db2_ptr = db_partials.data_ptr(); + + sm90::run_kda_bwd_intra_sm90(params, at::cuda::getCurrentCUDAStream()); + db_out.copy_(db_partials.sum(0).add_(db)); +} diff --git a/csrc/api/pybind.cu b/csrc/api/pybind.cu index 5a0f6299..e179ce4b 100644 --- a/csrc/api/pybind.cu +++ b/csrc/api/pybind.cu @@ -17,6 +17,26 @@ #include #include +void +ChunkKDABwdIntra( + at::Tensor q, + at::Tensor k, + at::Tensor g, + at::Tensor beta, + at::Tensor dAqk, + at::Tensor dAkk, + at::Tensor dq, + at::Tensor dk, + at::Tensor db, + at::Tensor dg, + at::Tensor cu_seqlens, + at::Tensor chunk_indices, + at::Tensor dq_out, + at::Tensor dk_out, + at::Tensor db_out, + at::Tensor dg_out, + int64_t chunk_size); + #if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) void ChunkKDAFwdIntra( @@ -72,6 +92,7 @@ kda_fwd_prefill( PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "cuLA"; + m.def("chunk_kda_bwd_intra_cuda", &ChunkKDABwdIntra); #if defined(CULA_SM100_ENABLED) || defined(CULA_SM103_ENABLED) m.def("chunk_kda_fwd_intra_cuda", &ChunkKDAFwdIntra); m.def("recompute_w_u_cuda", &ChunkKDAFwdRecompWU); diff --git a/csrc/kda/sm90/bwd/kda_bwd_intra_sm90.cu b/csrc/kda/sm90/bwd/kda_bwd_intra_sm90.cu new file mode 100644 index 00000000..d666c8fd --- /dev/null +++ b/csrc/kda/sm90/bwd/kda_bwd_intra_sm90.cu @@ -0,0 +1,848 @@ +// KDA backward intra-chunk kernel for SM90 (Hopper) - v16 +// 128 threads (4 warps), each warp handles one sub-chunk +// dA matrices cached in shared memory across K-iterations +// Persistent kernel: eliminates wave quantization overhead + +#include + +#include +#include + +#include + +#include "kda/sm90/bwd/kda_config.h" + +namespace sm90 { + +constexpr int BT = 64; +constexpr int BC = 16; +constexpr int BK = 32; +constexpr int K_SIZE = 128; +constexpr int NC = BT / BC; +constexpr int NK = K_SIZE / BK; +constexpr int WARP_SIZE = 32; +constexpr int BLOCK_THREADS = NC * WARP_SIZE; +constexpr int NT = BK / 8; +constexpr uint32_t TF32_MASK = 0xFFFFE000u; +constexpr int BK_S = BK + 4; +constexpr int BT_S = BT + 4; + +struct WarpWork { + float B_a[BC][BK_S]; + float B_b[BC][BK_S]; +}; + +struct SmemLayout { + __nv_bfloat16 q_s[BT][BK]; + __nv_bfloat16 k_s[BT][BK]; + float g_s[BT][BK_S]; + float beta_s[BT]; + float dAqk_cache[BT][BT_S]; + float dAkk_cache[BT][BT_S]; + WarpWork ww[NC]; + int s_tile_id; +}; + +__device__ __forceinline__ void +cp_async_16(void* smem, const void* global) { + uint32_t sa = __cvta_generic_to_shared(smem); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(sa), "l"(global)); +} +__device__ __forceinline__ void +cp_async_8(void* smem, const void* global) { + uint32_t sa = __cvta_generic_to_shared(smem); + asm volatile("cp.async.ca.shared.global [%0], [%1], 8;\n" ::"r"(sa), "l"(global)); +} +__device__ __forceinline__ float +bf2f(__nv_bfloat16 x) { + return __bfloat162float(x); +} +__device__ __forceinline__ float4 +load_bf16x4(const __nv_bfloat16* p) { + __nv_bfloat16 tmp[4]; + *reinterpret_cast(tmp) = *reinterpret_cast(p); + return {__bfloat162float(tmp[0]), __bfloat162float(tmp[1]), __bfloat162float(tmp[2]), __bfloat162float(tmp[3])}; +} +__device__ __forceinline__ void +cp_async_commit() { + asm volatile("cp.async.commit_group;\n"); +} +__device__ __forceinline__ void +cp_async_wait_all() { + asm volatile("cp.async.wait_group 0;\n"); +} +__device__ __forceinline__ void +st_global_cg_u32(void* addr, uint32_t val) { + asm volatile("st.global.cg.u32 [%0], %1;\n" ::"l"(addr), "r"(val)); +} +__device__ __forceinline__ void +st_global_cg_f32(void* addr, float val) { + asm volatile("st.global.cg.f32 [%0], %1;\n" ::"l"(addr), "f"(val)); +} +__device__ __forceinline__ void +st_global_cg_f32x2(void* addr, float v0, float v1) { + asm volatile("st.global.cg.v2.f32 [%0], {%1, %2};\n" ::"l"(addr), "f"(v0), "f"(v1)); +} + +__device__ __forceinline__ void +mma_m16n8k8_acc( + float& c0, + float& c1, + float& c2, + float& c3, + uint32_t a0, + uint32_t a1, + uint32_t a2, + uint32_t a3, + uint32_t b0, + uint32_t b1) { + asm volatile( + "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(c0), "+f"(c1), "+f"(c2), "+f"(c3) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1)); +} + +__device__ __forceinline__ void +matmul_1warp_2A_from_cache( + float acc1[], + float acc2[], + const float cacheA1[][BT_S], + const float cacheA2[][BT_S], + int row_off, + int col_off, + const float B[][BK_S], + int gid, + int tid_in_grp, + bool apply_mask, + int sub_seq_len) { + float a1[8], a2[8]; +#pragma unroll + for (int kk = 0; kk < 2; kk++) { + const int a_col = kk * 8 + 2 * tid_in_grp; + a1[kk * 4 + 0] = cacheA1[row_off + gid][col_off + a_col]; + a1[kk * 4 + 1] = cacheA1[row_off + gid + 8][col_off + a_col]; + a1[kk * 4 + 2] = cacheA1[row_off + gid][col_off + a_col + 1]; + a1[kk * 4 + 3] = cacheA1[row_off + gid + 8][col_off + a_col + 1]; + a2[kk * 4 + 0] = cacheA2[row_off + gid][col_off + a_col]; + a2[kk * 4 + 1] = cacheA2[row_off + gid + 8][col_off + a_col]; + a2[kk * 4 + 2] = cacheA2[row_off + gid][col_off + a_col + 1]; + a2[kk * 4 + 3] = cacheA2[row_off + gid + 8][col_off + a_col + 1]; + if (apply_mask) { + int col0 = a_col, col1 = a_col + 1; + if (!(col0 <= gid && gid < sub_seq_len && col0 < sub_seq_len)) { + a1[kk * 4 + 0] = 0.0f; + a2[kk * 4 + 0] = 0.0f; + } + if (!(col0 <= gid + 8 && gid + 8 < sub_seq_len && col0 < sub_seq_len)) { + a1[kk * 4 + 1] = 0.0f; + a2[kk * 4 + 1] = 0.0f; + } + if (!(col1 <= gid && gid < sub_seq_len && col1 < sub_seq_len)) { + a1[kk * 4 + 2] = 0.0f; + a2[kk * 4 + 2] = 0.0f; + } + if (!(col1 <= gid + 8 && gid + 8 < sub_seq_len && col1 < sub_seq_len)) { + a1[kk * 4 + 3] = 0.0f; + a2[kk * 4 + 3] = 0.0f; + } + } + a1[kk * 4 + 0] = __uint_as_float(__float_as_uint(a1[kk * 4 + 0]) & TF32_MASK); + a1[kk * 4 + 1] = __uint_as_float(__float_as_uint(a1[kk * 4 + 1]) & TF32_MASK); + a1[kk * 4 + 2] = __uint_as_float(__float_as_uint(a1[kk * 4 + 2]) & TF32_MASK); + a1[kk * 4 + 3] = __uint_as_float(__float_as_uint(a1[kk * 4 + 3]) & TF32_MASK); + a2[kk * 4 + 0] = __uint_as_float(__float_as_uint(a2[kk * 4 + 0]) & TF32_MASK); + a2[kk * 4 + 1] = __uint_as_float(__float_as_uint(a2[kk * 4 + 1]) & TF32_MASK); + a2[kk * 4 + 2] = __uint_as_float(__float_as_uint(a2[kk * 4 + 2]) & TF32_MASK); + a2[kk * 4 + 3] = __uint_as_float(__float_as_uint(a2[kk * 4 + 3]) & TF32_MASK); + } +#pragma unroll + for (int nt = 0; nt < NT; nt++) { + const int n_base = nt << 3; +#pragma unroll + for (int kk = 0; kk < 2; kk++) { + const int b_k = kk * 8 + 2 * tid_in_grp; + uint32_t ub0 = __float_as_uint(B[b_k][n_base + gid]) & TF32_MASK; + uint32_t ub1 = __float_as_uint(B[b_k + 1][n_base + gid]) & TF32_MASK; + mma_m16n8k8_acc( + acc1[nt * 4], + acc1[nt * 4 + 1], + acc1[nt * 4 + 2], + acc1[nt * 4 + 3], + __float_as_uint(a1[kk * 4 + 0]), + __float_as_uint(a1[kk * 4 + 1]), + __float_as_uint(a1[kk * 4 + 2]), + __float_as_uint(a1[kk * 4 + 3]), + ub0, + ub1); + mma_m16n8k8_acc( + acc2[nt * 4], + acc2[nt * 4 + 1], + acc2[nt * 4 + 2], + acc2[nt * 4 + 3], + __float_as_uint(a2[kk * 4 + 0]), + __float_as_uint(a2[kk * 4 + 1]), + __float_as_uint(a2[kk * 4 + 2]), + __float_as_uint(a2[kk * 4 + 3]), + ub0, + ub1); + } + } +} + +__device__ __forceinline__ void +matmul_1warp_2B_transA_from_cache( + float acc[], + const float cacheA1[][BT_S], + const float cacheA2[][BT_S], + int row_off, + int col_off, + const float B_x[][BK_S], + const float B_y[][BK_S], + int gid, + int tid_in_grp, + bool apply_mask, + int sub_seq_len) { + { + float a[8]; +#pragma unroll + for (int kk = 0; kk < 2; kk++) { + const int a_col = kk * 8 + 2 * tid_in_grp; + a[kk * 4 + 0] = cacheA1[row_off + a_col][col_off + gid]; + a[kk * 4 + 1] = cacheA1[row_off + a_col][col_off + gid + 8]; + a[kk * 4 + 2] = cacheA1[row_off + a_col + 1][col_off + gid]; + a[kk * 4 + 3] = cacheA1[row_off + a_col + 1][col_off + gid + 8]; + if (apply_mask) { + int row_A0 = a_col, row_A1 = a_col + 1; + if (!(gid <= row_A0 && row_A0 < sub_seq_len && gid < sub_seq_len)) + a[kk * 4 + 0] = 0.0f; + if (!(gid + 8 <= row_A0 && row_A0 < sub_seq_len && gid + 8 < sub_seq_len)) + a[kk * 4 + 1] = 0.0f; + if (!(gid <= row_A1 && row_A1 < sub_seq_len && gid < sub_seq_len)) + a[kk * 4 + 2] = 0.0f; + if (!(gid + 8 <= row_A1 && row_A1 < sub_seq_len && gid + 8 < sub_seq_len)) + a[kk * 4 + 3] = 0.0f; + } + a[kk * 4 + 0] = __uint_as_float(__float_as_uint(a[kk * 4 + 0]) & TF32_MASK); + a[kk * 4 + 1] = __uint_as_float(__float_as_uint(a[kk * 4 + 1]) & TF32_MASK); + a[kk * 4 + 2] = __uint_as_float(__float_as_uint(a[kk * 4 + 2]) & TF32_MASK); + a[kk * 4 + 3] = __uint_as_float(__float_as_uint(a[kk * 4 + 3]) & TF32_MASK); + } +#pragma unroll + for (int nt = 0; nt < NT; nt++) { + const int n_base = nt << 3; +#pragma unroll + for (int kk = 0; kk < 2; kk++) { + const int b_k = kk * 8 + 2 * tid_in_grp; + uint32_t ub0 = __float_as_uint(B_x[b_k][n_base + gid]) & TF32_MASK; + uint32_t ub1 = __float_as_uint(B_x[b_k + 1][n_base + gid]) & TF32_MASK; + mma_m16n8k8_acc( + acc[nt * 4], + acc[nt * 4 + 1], + acc[nt * 4 + 2], + acc[nt * 4 + 3], + __float_as_uint(a[kk * 4 + 0]), + __float_as_uint(a[kk * 4 + 1]), + __float_as_uint(a[kk * 4 + 2]), + __float_as_uint(a[kk * 4 + 3]), + ub0, + ub1); + } + } + } + { + float a[8]; +#pragma unroll + for (int kk = 0; kk < 2; kk++) { + const int a_col = kk * 8 + 2 * tid_in_grp; + a[kk * 4 + 0] = cacheA2[row_off + a_col][col_off + gid]; + a[kk * 4 + 1] = cacheA2[row_off + a_col][col_off + gid + 8]; + a[kk * 4 + 2] = cacheA2[row_off + a_col + 1][col_off + gid]; + a[kk * 4 + 3] = cacheA2[row_off + a_col + 1][col_off + gid + 8]; + if (apply_mask) { + int row_A0 = a_col, row_A1 = a_col + 1; + if (!(gid <= row_A0 && row_A0 < sub_seq_len && gid < sub_seq_len)) + a[kk * 4 + 0] = 0.0f; + if (!(gid + 8 <= row_A0 && row_A0 < sub_seq_len && gid + 8 < sub_seq_len)) + a[kk * 4 + 1] = 0.0f; + if (!(gid <= row_A1 && row_A1 < sub_seq_len && gid < sub_seq_len)) + a[kk * 4 + 2] = 0.0f; + if (!(gid + 8 <= row_A1 && row_A1 < sub_seq_len && gid + 8 < sub_seq_len)) + a[kk * 4 + 3] = 0.0f; + } + a[kk * 4 + 0] = __uint_as_float(__float_as_uint(a[kk * 4 + 0]) & TF32_MASK); + a[kk * 4 + 1] = __uint_as_float(__float_as_uint(a[kk * 4 + 1]) & TF32_MASK); + a[kk * 4 + 2] = __uint_as_float(__float_as_uint(a[kk * 4 + 2]) & TF32_MASK); + a[kk * 4 + 3] = __uint_as_float(__float_as_uint(a[kk * 4 + 3]) & TF32_MASK); + } +#pragma unroll + for (int nt = 0; nt < NT; nt++) { + const int n_base = nt << 3; +#pragma unroll + for (int kk = 0; kk < 2; kk++) { + const int b_k = kk * 8 + 2 * tid_in_grp; + uint32_t ub0 = __float_as_uint(B_y[b_k][n_base + gid]) & TF32_MASK; + uint32_t ub1 = __float_as_uint(B_y[b_k + 1][n_base + gid]) & TF32_MASK; + mma_m16n8k8_acc( + acc[nt * 4], + acc[nt * 4 + 1], + acc[nt * 4 + 2], + acc[nt * 4 + 3], + __float_as_uint(a[kk * 4 + 0]), + __float_as_uint(a[kk * 4 + 1]), + __float_as_uint(a[kk * 4 + 2]), + __float_as_uint(a[kk * 4 + 3]), + ub0, + ub1); + } + } + } +} + +__device__ __forceinline__ void +load_block_cp_async(float dst[][BK_S], const float* src, int row_base, int stride, int tile_seq_len, int tid) { +#pragma unroll + for (int pass = 0; pass < 4; pass++) { + int f4_r = pass * 4 + (tid >> 3); + int f4_c = (tid & 7) << 2; + int r = row_base + f4_r; + if (r < tile_seq_len) { + cp_async_16(&dst[f4_r][f4_c], &src[r * stride + f4_c]); + } else { + float4 zero = {0, 0, 0, 0}; + *(float4*)(&dst[f4_r][f4_c]) = zero; + } + } +} + +__global__ void +__launch_bounds__(BLOCK_THREADS, 3) kda_bwd_intra_sm90_kernel(const KDA_bwd_intra_params params) { + extern __shared__ char shared_buf[]; + SmemLayout* smem = reinterpret_cast(shared_buf); + + const int warp_id = threadIdx.x / WARP_SIZE; + const int tid = threadIdx.x % WARP_SIZE; + const int i_i = warp_id; + + const int gid = tid >> 2; + const int tid_in_grp = tid & 3; + + const int* chunk_indices = (const int*)params.chunk_indices_ptr; + const int* cu_seqlens = (const int*)params.cu_seqlens_ptr; + int* tile_counter = (int*)params.tile_counter_ptr; + + const int H = params.h; + const int K = params.d; + const int stride_qk = H * K; + const int stride_dA = H * BT; + const int stride_b = H; + const int total_tiles = params.num_chunks * H; + + const __nv_bfloat16(*my_q)[BK] = &smem->q_s[i_i * BC]; + const __nv_bfloat16(*my_k)[BK] = &smem->k_s[i_i * BC]; + const float(*my_g)[BK_S] = (const float(*)[BK_S]) & smem->g_s[i_i * BC]; + WarpWork& ww = smem->ww[i_i]; + const int row0 = gid; + const int row1 = gid + 8; + + while (true) { + // ==================== PERSISTENT TILE DISPATCH ==================== + if (threadIdx.x == 0) { + smem->s_tile_id = atomicAdd(tile_counter, 1); + } + __syncthreads(); + const int tile_id = smem->s_tile_id; + if (tile_id >= total_tiles) + return; + + const int i_t = tile_id / H; + const int i_h = tile_id % H; + + const int batch_idx = chunk_indices[i_t * 2]; + const int seq_idx = chunk_indices[i_t * 2 + 1]; + const int start_offset = cu_seqlens[batch_idx]; + const int seq_len = cu_seqlens[batch_idx + 1] - start_offset; + + const int tile_seq_len = min(BT, seq_len - seq_idx * BT); + + const int tile_offset = start_offset + seq_idx * BT; + + const float* beta_base = (const float*)params.beta_ptr + tile_offset * stride_b + i_h; + const float* dAqk_base = (const float*)params.dAqk_ptr + tile_offset * stride_dA + i_h * BT; + const float* dAkk_base = (const float*)params.dAkk_ptr + tile_offset * stride_dA + i_h * BT; + + const int sub_seq_len = min(BC, tile_seq_len - i_i * BC); + const bool warp_active = (sub_seq_len > 0); + const int NC_actual = min(NC, (tile_seq_len + BC - 1) / BC); + + // Load beta (k-independent) + if (threadIdx.x < BT) { + smem->beta_s[threadIdx.x] = (threadIdx.x < tile_seq_len) ? beta_base[threadIdx.x * stride_b] : 0.0f; + } + +// ==================== dA CACHE LOAD ==================== +#pragma unroll 1 + for (int elem4 = threadIdx.x; elem4 < BT * (BT >> 2); elem4 += BLOCK_THREADS) { + int r = elem4 >> 4; + int c = (elem4 & 15) << 2; + if (r < tile_seq_len) { + cp_async_16(&smem->dAqk_cache[r][c], &dAqk_base[r * stride_dA + c]); + cp_async_16(&smem->dAkk_cache[r][c], &dAkk_base[r * stride_dA + c]); + } else { + float4 zero = {0, 0, 0, 0}; + *reinterpret_cast(&smem->dAqk_cache[r][c]) = zero; + *reinterpret_cast(&smem->dAkk_cache[r][c]) = zero; + } + } + cp_async_commit(); + cp_async_wait_all(); + __syncthreads(); + + // ==================== K-SLICE LOOP ==================== + for (int i_k = 0; i_k < NK; i_k++) { + const int k_off = i_k * BK; + + // Load Q, K, G for this K-iteration + { + const __nv_bfloat16* q_base = + (const __nv_bfloat16*)params.q_ptr + tile_offset * stride_qk + i_h * K + k_off; + const __nv_bfloat16* k_base = + (const __nv_bfloat16*)params.k_ptr + tile_offset * stride_qk + i_h * K + k_off; + const float* g_base = (const float*)params.g_ptr + tile_offset * stride_qk + i_h * K + k_off; + for (int elem = threadIdx.x; elem < BT * (BK / 8); elem += BLOCK_THREADS) { + int r = elem >> 2; + int c = (elem & 3) << 3; + if (r < tile_seq_len) { + int goff = r * stride_qk + c; + cp_async_16(&smem->q_s[r][c], &q_base[goff]); + cp_async_16(&smem->k_s[r][c], &k_base[goff]); + } else { + uint4 zero4 = {0, 0, 0, 0}; + *reinterpret_cast(&smem->q_s[r][c]) = zero4; + *reinterpret_cast(&smem->k_s[r][c]) = zero4; + } + } + for (int elem = threadIdx.x; elem < BT * (BK / 4); elem += BLOCK_THREADS) { + int r = elem >> 3; + int c = (elem & 7) << 2; + if (r < tile_seq_len) { + cp_async_16(&smem->g_s[r][c], &g_base[r * stride_qk + c]); + } else { + float4 zero = {0, 0, 0, 0}; + *reinterpret_cast(&smem->g_s[r][c]) = zero; + } + } + } + cp_async_commit(); + cp_async_wait_all(); + __syncthreads(); + + const float* dq_base = (const float*)params.dq_ptr + tile_offset * stride_qk + i_h * K + k_off; + const float* dk_base = (const float*)params.dk_ptr + tile_offset * stride_qk + i_h * K + k_off; + const float* dg_base = (const float*)params.dg_ptr + tile_offset * stride_qk + i_h * K + k_off; + __nv_bfloat16* dq_out = (__nv_bfloat16*)params.dq_out_ptr + tile_offset * stride_qk + i_h * K + k_off; + __nv_bfloat16* dk_out = (__nv_bfloat16*)params.dk_out_ptr + tile_offset * stride_qk + i_h * K + k_off; + float* db2_base = (float*)params.db2_ptr + i_k * params.total_q_len * H + tile_offset * stride_b + i_h; + float* dg_out = (float*)params.dg_out_ptr + tile_offset * stride_qk + i_h * K + k_off; + + if (warp_active) { + // ==================== FORWARD OFF-DIAGONAL ==================== + float dq2[16] = {0}; + float dk2[16] = {0}; + + if (i_i > 0) { +#pragma unroll 1 + for (int j = 0; j < i_i; j++) { +#pragma unroll + for (int pass = 0; pass < 4; pass++) { + int f4_r = pass * 4 + (tid >> 3); + int f4_c = (tid & 7) << 2; + float4 ks = load_bf16x4(&smem->k_s[j * BC + f4_r][f4_c]); + float4 gs_j = *reinterpret_cast(&smem->g_s[j * BC + f4_r][f4_c]); + float4 gs_gn = *reinterpret_cast(&my_g[0][f4_c]); + float4 ba; + ba.x = ks.x * exp2f(gs_gn.x - gs_j.x); + ba.y = ks.y * exp2f(gs_gn.y - gs_j.y); + ba.z = ks.z * exp2f(gs_gn.z - gs_j.z); + ba.w = ks.w * exp2f(gs_gn.w - gs_j.w); + *reinterpret_cast(&ww.B_a[f4_r][f4_c]) = ba; + } + __syncwarp(); + matmul_1warp_2A_from_cache( + dq2, + dk2, + smem->dAqk_cache, + smem->dAkk_cache, + i_i * BC, + j * BC, + ww.B_a, + gid, + tid_in_grp, + false, + BC); + } +#pragma unroll + for (int nt = 0; nt < NT; nt++) { + int col0 = nt * 8 + tid_in_grp * 2; + int col1 = col0 + 1; + float gn0 = my_g[0][col0], gn1 = my_g[0][col1]; + float s00 = exp2f(my_g[row0][col0] - gn0); + float s01 = exp2f(my_g[row0][col1] - gn1); + float s10 = exp2f(my_g[row1][col0] - gn0); + float s11 = exp2f(my_g[row1][col1] - gn1); + dq2[nt * 4 + 0] *= s00; + dq2[nt * 4 + 1] *= s01; + dq2[nt * 4 + 2] *= s10; + dq2[nt * 4 + 3] *= s11; + dk2[nt * 4 + 0] *= s00; + dk2[nt * 4 + 1] *= s01; + dk2[nt * 4 + 2] *= s10; + dk2[nt * 4 + 3] *= s11; + } + } + + // ==================== FORWARD DIAGONAL ==================== + { + int gn_row = min(BC / 2, sub_seq_len - 1); + +#pragma unroll + for (int pass = 0; pass < 4; pass++) { + int f4_r = pass * 4 + (tid >> 3); + int f4_c = (tid & 7) << 2; + if (f4_r < sub_seq_len) { + float4 ks = load_bf16x4(&my_k[f4_r][f4_c]); + float4 gs_gn = *reinterpret_cast(&my_g[gn_row][f4_c]); + float4 gs_r = *reinterpret_cast(&my_g[f4_r][f4_c]); + float4 ba; + ba.x = ks.x * exp2f(gs_gn.x - gs_r.x); + ba.y = ks.y * exp2f(gs_gn.y - gs_r.y); + ba.z = ks.z * exp2f(gs_gn.z - gs_r.z); + ba.w = ks.w * exp2f(gs_gn.w - gs_r.w); + *reinterpret_cast(&ww.B_a[f4_r][f4_c]) = ba; + } else { + float4 zero = {0, 0, 0, 0}; + *reinterpret_cast(&ww.B_a[f4_r][f4_c]) = zero; + } + } + __syncwarp(); + + float dqd[16] = {0}; + float dkd[16] = {0}; + matmul_1warp_2A_from_cache( + dqd, + dkd, + smem->dAqk_cache, + smem->dAkk_cache, + i_i * BC, + i_i * BC, + ww.B_a, + gid, + tid_in_grp, + true, + sub_seq_len); + + // Start dq_in load into B_a (MMA is done reading B_a) + load_block_cp_async(ww.B_a, dq_base, i_i * BC, stride_qk, tile_seq_len, tid); + cp_async_commit(); + +#pragma unroll + for (int nt = 0; nt < NT; nt++) { + int col0 = nt * 8 + tid_in_grp * 2; + int col1 = col0 + 1; + float gn0 = my_g[gn_row][col0], gn1 = my_g[gn_row][col1]; + float s00 = exp2f(my_g[row0][col0] - gn0); + float s01 = exp2f(my_g[row0][col1] - gn1); + float s10 = exp2f(my_g[row1][col0] - gn0); + float s11 = exp2f(my_g[row1][col1] - gn1); + dq2[nt * 4 + 0] += dqd[nt * 4 + 0] * s00; + dq2[nt * 4 + 1] += dqd[nt * 4 + 1] * s01; + dq2[nt * 4 + 2] += dqd[nt * 4 + 2] * s10; + dq2[nt * 4 + 3] += dqd[nt * 4 + 3] * s11; + dk2[nt * 4 + 0] += dkd[nt * 4 + 0] * s00; + dk2[nt * 4 + 1] += dkd[nt * 4 + 1] * s01; + dk2[nt * 4 + 2] += dkd[nt * 4 + 2] * s10; + dk2[nt * 4 + 3] += dkd[nt * 4 + 3] * s11; + } + } + + // Wait for dq_in load + cp_async_wait_all(); + __syncwarp(); + + // Write dq_out = dq2 + dq_in, then reuse dq2 for dg_p = q * dq2 + { + int tile_r0 = i_i * BC + row0; + int tile_r1 = i_i * BC + row1; +#pragma unroll + for (int nt = 0; nt < NT; nt++) { + int col0 = nt * 8 + tid_in_grp * 2; + if (row0 < sub_seq_len) { + __nv_bfloat162 pair = { + __float2bfloat16(dq2[nt * 4 + 0] + ww.B_a[row0][col0]), + __float2bfloat16(dq2[nt * 4 + 1] + ww.B_a[row0][col0 + 1])}; + st_global_cg_u32(&dq_out[tile_r0 * stride_qk + col0], *reinterpret_cast(&pair)); + } + if (row1 < sub_seq_len) { + __nv_bfloat162 pair = { + __float2bfloat16(dq2[nt * 4 + 2] + ww.B_a[row1][col0]), + __float2bfloat16(dq2[nt * 4 + 3] + ww.B_a[row1][col0 + 1])}; + st_global_cg_u32(&dq_out[tile_r1 * stride_qk + col0], *reinterpret_cast(&pair)); + } + dq2[nt * 4 + 0] = bf2f(my_q[row0][col0]) * dq2[nt * 4 + 0]; + dq2[nt * 4 + 1] = bf2f(my_q[row0][col0 + 1]) * dq2[nt * 4 + 1]; + dq2[nt * 4 + 2] = bf2f(my_q[row1][col0]) * dq2[nt * 4 + 2]; + dq2[nt * 4 + 3] = bf2f(my_q[row1][col0 + 1]) * dq2[nt * 4 + 3]; + } + } + + // db reduction + { + float db0 = 0, db1 = 0; +#pragma unroll + for (int nt = 0; nt < NT; nt++) { + int col0 = nt * 8 + tid_in_grp * 2; + int col1 = col0 + 1; + db0 += dk2[nt * 4 + 0] * bf2f(my_k[row0][col0]) + dk2[nt * 4 + 1] * bf2f(my_k[row0][col1]); + db1 += dk2[nt * 4 + 2] * bf2f(my_k[row1][col0]) + dk2[nt * 4 + 3] * bf2f(my_k[row1][col1]); + } + db0 += __shfl_xor_sync(0xFFFFFFFF, db0, 1); + db0 += __shfl_xor_sync(0xFFFFFFFF, db0, 2); + db1 += __shfl_xor_sync(0xFFFFFFFF, db1, 1); + db1 += __shfl_xor_sync(0xFFFFFFFF, db1, 2); + if (tid_in_grp == 0) { + if (row0 < sub_seq_len) + st_global_cg_f32(&db2_base[(i_i * BC + row0) * stride_b], db0); + if (row1 < sub_seq_len) + st_global_cg_f32(&db2_base[(i_i * BC + row1) * stride_b], db1); + } + } + + // Scale dk2 by beta + { + float beta0 = smem->beta_s[i_i * BC + row0]; + float beta1 = smem->beta_s[i_i * BC + row1]; +#pragma unroll + for (int nt = 0; nt < NT; nt++) { + dk2[nt * 4 + 0] *= beta0; + dk2[nt * 4 + 1] *= beta0; + dk2[nt * 4 + 2] *= beta1; + dk2[nt * 4 + 3] *= beta1; + } + } + + // ==================== BACKWARD OFF-DIAGONAL ==================== + float dkt[16] = {0}; + + if (i_i < NC_actual - 1) { + int gn_bwd_row = min(BC - 1, sub_seq_len - 1); + +#pragma unroll 1 + for (int j = i_i + 1; j < NC_actual; j++) { + int j_sub_seq = min(BC, tile_seq_len - j * BC); +#pragma unroll + for (int pass = 0; pass < 4; pass++) { + int f4_r = pass * 4 + (tid >> 3); + int f4_c = (tid & 7) << 2; + if (f4_r < j_sub_seq) { + float4 qs = load_bf16x4(&smem->q_s[j * BC + f4_r][f4_c]); + float4 ks = load_bf16x4(&smem->k_s[j * BC + f4_r][f4_c]); + float4 gs_j = *reinterpret_cast(&smem->g_s[j * BC + f4_r][f4_c]); + float4 gs_gn = *reinterpret_cast(&my_g[gn_bwd_row][f4_c]); + float beta_val = smem->beta_s[j * BC + f4_r]; + + float eg0 = exp2f(gs_j.x - gs_gn.x), eg1 = exp2f(gs_j.y - gs_gn.y), + eg2 = exp2f(gs_j.z - gs_gn.z), eg3 = exp2f(gs_j.w - gs_gn.w); + + float4 ba = {qs.x * eg0, qs.y * eg1, qs.z * eg2, qs.w * eg3}; + float4 bb; + bb.x = __bfloat162float(__float2bfloat16(ks.x * beta_val)) * eg0; + bb.y = __bfloat162float(__float2bfloat16(ks.y * beta_val)) * eg1; + bb.z = __bfloat162float(__float2bfloat16(ks.z * beta_val)) * eg2; + bb.w = __bfloat162float(__float2bfloat16(ks.w * beta_val)) * eg3; + *reinterpret_cast(&ww.B_a[f4_r][f4_c]) = ba; + *reinterpret_cast(&ww.B_b[f4_r][f4_c]) = bb; + } else { + float4 zero = {0, 0, 0, 0}; + *reinterpret_cast(&ww.B_a[f4_r][f4_c]) = zero; + *reinterpret_cast(&ww.B_b[f4_r][f4_c]) = zero; + } + } + __syncwarp(); + matmul_1warp_2B_transA_from_cache( + dkt, + smem->dAqk_cache, + smem->dAkk_cache, + j * BC, + i_i * BC, + ww.B_a, + ww.B_b, + gid, + tid_in_grp, + false, + BC); + } +#pragma unroll + for (int nt = 0; nt < NT; nt++) { + int col0 = nt * 8 + tid_in_grp * 2; + int col1 = col0 + 1; + float gn0 = my_g[gn_bwd_row][col0], gn1 = my_g[gn_bwd_row][col1]; + float s00 = exp2f(gn0 - my_g[row0][col0]); + float s01 = exp2f(gn1 - my_g[row0][col1]); + float s10 = exp2f(gn0 - my_g[row1][col0]); + float s11 = exp2f(gn1 - my_g[row1][col1]); + dkt[nt * 4 + 0] *= s00; + dkt[nt * 4 + 1] *= s01; + dkt[nt * 4 + 2] *= s10; + dkt[nt * 4 + 3] *= s11; + } + } + + // ==================== BACKWARD DIAGONAL ==================== + { + int gn_row = min(BC / 2, sub_seq_len - 1); + +#pragma unroll + for (int pass = 0; pass < 4; pass++) { + int f4_r = pass * 4 + (tid >> 3); + int f4_c = (tid & 7) << 2; + if (f4_r < sub_seq_len) { + float4 qs = load_bf16x4(&my_q[f4_r][f4_c]); + float4 ks = load_bf16x4(&my_k[f4_r][f4_c]); + float4 gs_r = *reinterpret_cast(&my_g[f4_r][f4_c]); + float4 gs_gn = *reinterpret_cast(&my_g[gn_row][f4_c]); + float beta_r = smem->beta_s[i_i * BC + f4_r]; + + float eg0 = exp2f(gs_r.x - gs_gn.x), eg1 = exp2f(gs_r.y - gs_gn.y), + eg2 = exp2f(gs_r.z - gs_gn.z), eg3 = exp2f(gs_r.w - gs_gn.w); + + float4 ba = {qs.x * eg0, qs.y * eg1, qs.z * eg2, qs.w * eg3}; + float4 bb; + bb.x = __bfloat162float(__float2bfloat16(ks.x * beta_r)) * eg0; + bb.y = __bfloat162float(__float2bfloat16(ks.y * beta_r)) * eg1; + bb.z = __bfloat162float(__float2bfloat16(ks.z * beta_r)) * eg2; + bb.w = __bfloat162float(__float2bfloat16(ks.w * beta_r)) * eg3; + *reinterpret_cast(&ww.B_a[f4_r][f4_c]) = ba; + *reinterpret_cast(&ww.B_b[f4_r][f4_c]) = bb; + } else { + float4 zero = {0, 0, 0, 0}; + *reinterpret_cast(&ww.B_a[f4_r][f4_c]) = zero; + *reinterpret_cast(&ww.B_b[f4_r][f4_c]) = zero; + } + } + __syncwarp(); + + float dktd[16] = {0}; + matmul_1warp_2B_transA_from_cache( + dktd, + smem->dAqk_cache, + smem->dAkk_cache, + i_i * BC, + i_i * BC, + ww.B_a, + ww.B_b, + gid, + tid_in_grp, + true, + sub_seq_len); + + // Start epilogue loads now (B_a/B_b no longer needed by MMA) + load_block_cp_async(ww.B_a, dk_base, i_i * BC, stride_qk, tile_seq_len, tid); + load_block_cp_async(ww.B_b, dg_base, i_i * BC, stride_qk, tile_seq_len, tid); + cp_async_commit(); + +#pragma unroll + for (int nt = 0; nt < NT; nt++) { + int col0 = nt * 8 + tid_in_grp * 2; + int col1 = col0 + 1; + float gn0 = my_g[gn_row][col0], gn1 = my_g[gn_row][col1]; + float s00 = exp2f(gn0 - my_g[row0][col0]); + float s01 = exp2f(gn1 - my_g[row0][col1]); + float s10 = exp2f(gn0 - my_g[row1][col0]); + float s11 = exp2f(gn1 - my_g[row1][col1]); + dkt[nt * 4 + 0] += dktd[nt * 4 + 0] * s00; + dkt[nt * 4 + 1] += dktd[nt * 4 + 1] * s01; + dkt[nt * 4 + 2] += dktd[nt * 4 + 2] * s10; + dkt[nt * 4 + 3] += dktd[nt * 4 + 3] * s11; + } + } + + // ==================== EPILOGUE ==================== + cp_async_wait_all(); + __syncwarp(); + + { + int tile_r0 = i_i * BC + row0; + int tile_r1 = i_i * BC + row1; + +#pragma unroll + for (int nt = 0; nt < NT; nt++) { + int col0 = nt * 8 + tid_in_grp * 2; + if (row0 < sub_seq_len) { + int off0 = tile_r0 * stride_qk + col0; + __nv_bfloat162 dk_pair = { + __float2bfloat16(ww.B_a[row0][col0] + dk2[nt * 4 + 0] + dkt[nt * 4 + 0]), + __float2bfloat16(ww.B_a[row0][col0 + 1] + dk2[nt * 4 + 1] + dkt[nt * 4 + 1])}; + st_global_cg_u32(&dk_out[off0], *reinterpret_cast(&dk_pair)); + st_global_cg_f32x2( + &dg_out[off0], + dq2[nt * 4 + 0] + (dk2[nt * 4 + 0] - dkt[nt * 4 + 0]) * bf2f(my_k[row0][col0]) + + ww.B_b[row0][col0], + dq2[nt * 4 + 1] + (dk2[nt * 4 + 1] - dkt[nt * 4 + 1]) * bf2f(my_k[row0][col0 + 1]) + + ww.B_b[row0][col0 + 1]); + } + if (row1 < sub_seq_len) { + int off0 = tile_r1 * stride_qk + col0; + __nv_bfloat162 dk_pair = { + __float2bfloat16(ww.B_a[row1][col0] + dk2[nt * 4 + 2] + dkt[nt * 4 + 2]), + __float2bfloat16(ww.B_a[row1][col0 + 1] + dk2[nt * 4 + 3] + dkt[nt * 4 + 3])}; + st_global_cg_u32(&dk_out[off0], *reinterpret_cast(&dk_pair)); + st_global_cg_f32x2( + &dg_out[off0], + dq2[nt * 4 + 2] + (dk2[nt * 4 + 2] - dkt[nt * 4 + 2]) * bf2f(my_k[row1][col0]) + + ww.B_b[row1][col0], + dq2[nt * 4 + 3] + (dk2[nt * 4 + 3] - dkt[nt * 4 + 3]) * bf2f(my_k[row1][col0 + 1]) + + ww.B_b[row1][col0 + 1]); + } + } + } + + } // warp_active + + __syncthreads(); + + } // k_idx loop + + } // persistent while loop +} + +void +run_kda_bwd_intra_sm90(KDA_bwd_intra_params& params, cudaStream_t stream) { + constexpr size_t smem_size = sizeof(SmemLayout); + auto kernel = &kda_bwd_intra_sm90_kernel; + C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + int num_blocks_per_sm; + C10_CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks_per_sm, kernel, BLOCK_THREADS, smem_size)); + + int device; + C10_CUDA_CHECK(cudaGetDevice(&device)); + int num_sms; + C10_CUDA_CHECK(cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, device)); + + int total_tiles = params.num_chunks * params.h; + int num_blocks = min(num_sms * num_blocks_per_sm, total_tiles); + + int* tile_counter; + C10_CUDA_CHECK(cudaMallocAsync(&tile_counter, sizeof(int), stream)); + C10_CUDA_CHECK(cudaMemsetAsync(tile_counter, 0, sizeof(int), stream)); + params.tile_counter_ptr = tile_counter; + + dim3 grid(num_blocks, 1, 1); + dim3 block(BLOCK_THREADS, 1, 1); + kernel<<>>(params); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + C10_CUDA_CHECK(cudaFreeAsync(tile_counter, stream)); +} + +} // namespace sm90 diff --git a/csrc/kda/sm90/bwd/kda_config.h b/csrc/kda/sm90/bwd/kda_config.h new file mode 100644 index 00000000..60da2e8e --- /dev/null +++ b/csrc/kda/sm90/bwd/kda_config.h @@ -0,0 +1,26 @@ +#pragma once + +struct KDA_bwd_intra_params { + int total_q_len; + int h; + int d; + + void* __restrict__ q_ptr; //[b, t, h, d] + void* __restrict__ k_ptr; //[b, t, h, d] + void* __restrict__ g_ptr; //[b, t, h, d] + void* __restrict__ beta_ptr; //[b, t, h] + void* __restrict__ dAqk_ptr; //[b, t, h, BT] + void* __restrict__ dAkk_ptr; //[b, t, h, BT] + void* __restrict__ dq_ptr; //[b, t, h, d] + void* __restrict__ dk_ptr; //[b, t, h, d] + void* __restrict__ dg_ptr; //[b, t, h, d] + void* __restrict__ dq_out_ptr; //[b, t, h, d] + void* __restrict__ dk_out_ptr; //[b, t, h, d] + void* __restrict__ db2_ptr; //[NK, total_q_len, h] - per-K-tile db partials + void* __restrict__ dg_out_ptr; //[b, t, h, d] + void* __restrict__ cu_seqlens_ptr; //[b + 1] + void* __restrict__ chunk_indices_ptr; //[num_chunks, 2] + + int num_chunks; + void* tile_counter_ptr; +}; diff --git a/cula/kda/chunk_intra.py b/cula/kda/chunk_intra.py index f864849c..6320df76 100644 --- a/cula/kda/chunk_intra.py +++ b/cula/kda/chunk_intra.py @@ -24,7 +24,39 @@ from fla.utils import IS_GATHER_SUPPORTED, autotune_cache_kwargs import cula.cudac as cula_cuda -from cula.utils import prepare_uniform_cu_seqlens +from cula.utils import get_device_sm_version, prepare_uniform_cu_seqlens + + +def _is_mma_bwd_intra_supported( + q: torch.Tensor, + k: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + dAqk: torch.Tensor, + dAkk: torch.Tensor, + dq: torch.Tensor, + dk: torch.Tensor, + db: torch.Tensor, + dg: torch.Tensor, + chunk_size: int, + safe_gate: bool, +) -> bool: + tensors = (q, k, g, beta, dAqk, dAkk, dq, dk, db, dg) + if not q.is_cuda or any(tensor.device != q.device for tensor in tensors[1:]): + return False + capability = get_device_sm_version(q.device) + return ( + capability in ((9, 0), (10, 0), (10, 3)) + and safe_gate + and chunk_size == 64 + and q.ndim == 4 + and q.shape[-1] == 128 + and q.shape == k.shape == g.shape == dq.shape == dk.shape == dg.shape + and beta.shape == db.shape == q.shape[:-1] + and dAqk.shape == dAkk.shape == (*q.shape[:-1], chunk_size) + and q.dtype == k.dtype == beta.dtype == torch.bfloat16 + and g.dtype == dAqk.dtype == dAkk.dtype == dq.dtype == dk.dtype == db.dtype == dg.dtype == torch.float32 + ) @triton.heuristics( @@ -360,7 +392,7 @@ def chunk_kda_fwd_intra( return w, u, qg, kg, Aqk, Akk -def chunk_kda_bwd_intra( +def _chunk_kda_bwd_intra_triton( q: torch.Tensor, k: torch.Tensor, g: torch.Tensor, @@ -426,3 +458,89 @@ def chunk_kda_bwd_intra( dg = dg2 return dq, dk, db, dg + + +def _chunk_kda_bwd_intra_mma( + q: torch.Tensor, + k: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + dAqk: torch.Tensor, + dAkk: torch.Tensor, + dq: torch.Tensor, + dk: torch.Tensor, + db: torch.Tensor, + dg: torch.Tensor, + cu_seqlens: torch.IntTensor | None = None, + chunk_indices: torch.IntTensor | None = None, + chunk_size: int = 64, +): + from cula.ops.kda.sm90.bwd_intra import kda_bwd_intra_mma + + B, T, _, _ = q.shape + if cu_seqlens is None: + cu_seqlens = prepare_uniform_cu_seqlens(B, T, q.device, torch.int32) + else: + cu_seqlens = cu_seqlens.to(device=q.device, dtype=torch.int32).contiguous() + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + else: + chunk_indices = chunk_indices.to(device=q.device, dtype=torch.int32).contiguous() + + inputs = (q, k, g, beta, dAqk, dAkk, dq, dk, db, dg) + return kda_bwd_intra_mma( + *(tensor.contiguous() for tensor in inputs), + cu_seqlens, + chunk_indices, + chunk_size=chunk_size, + ) + + +def chunk_kda_bwd_intra( + q: torch.Tensor, + k: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + dAqk: torch.Tensor, + dAkk: torch.Tensor, + dq: torch.Tensor, + dk: torch.Tensor, + db: torch.Tensor, + dg: torch.Tensor, + cu_seqlens: torch.IntTensor | None = None, + chunk_indices: torch.IntTensor | None = None, + chunk_size: int = 64, + safe_gate: bool = False, +): + if _is_mma_bwd_intra_supported(q, k, g, beta, dAqk, dAkk, dq, dk, db, dg, chunk_size, safe_gate): + return _chunk_kda_bwd_intra_mma( + q, + k, + g, + beta, + dAqk, + dAkk, + dq, + dk, + db, + dg, + cu_seqlens, + chunk_indices, + chunk_size, + ) + return _chunk_kda_bwd_intra_triton( + q, + k, + g, + beta, + dAqk, + dAkk, + dq, + dk, + db, + dg, + cu_seqlens, + chunk_indices, + chunk_size, + safe_gate, + ) diff --git a/cula/ops/kda/sm90/bwd_intra.py b/cula/ops/kda/sm90/bwd_intra.py new file mode 100644 index 00000000..b93f3b6c --- /dev/null +++ b/cula/ops/kda/sm90/bwd_intra.py @@ -0,0 +1,87 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Persistent CUDA ``mma.sync`` backend for KDA intra-chunk backward. + +The source remains under the SM90 directory to match the current repository +layout, but the low-level kernel is also compiled for and supported on SM100 +and SM103. +""" + +import torch + +import cula.cudac as cula_cuda +from cula.utils import get_device_sm_version + +_SUPPORTED_CAPABILITIES = {(9, 0), (10, 0), (10, 3)} +_CHUNK_SIZE = 64 +_HEAD_DIM = 128 + + +def kda_bwd_intra_mma( + q: torch.Tensor, + k: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + d_aq: torch.Tensor, + d_ak: torch.Tensor, + dq: torch.Tensor, + dk: torch.Tensor, + db: torch.Tensor, + dg: torch.Tensor, + cu_seqlens: torch.Tensor, + chunk_indices: torch.Tensor, + dq_out: torch.Tensor | None = None, + dk_out: torch.Tensor | None = None, + db_out: torch.Tensor | None = None, + dg_out: torch.Tensor | None = None, + chunk_size: int = _CHUNK_SIZE, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Run the portable SM90-style MMA kernel directly. + + This is the low-level validation/benchmark entry point. Unsupported + production cases are handled by the higher-level Triton fallback in + :func:`cula.kda.chunk_intra.chunk_kda_bwd_intra`. + """ + + if not q.is_cuda: + raise ValueError("kda_bwd_intra_mma requires CUDA tensors") + capability = get_device_sm_version(q.device) + if capability not in _SUPPORTED_CAPABILITIES: + raise RuntimeError(f"kda_bwd_intra_mma requires SM90, SM100, or SM103, got SM{capability[0]}{capability[1]}") + if chunk_size != _CHUNK_SIZE: + raise ValueError(f"kda_bwd_intra_mma supports only chunk_size={_CHUNK_SIZE}, got {chunk_size}") + if q.shape[-1] != _HEAD_DIM: + raise ValueError(f"kda_bwd_intra_mma supports only head dimension {_HEAD_DIM}, got {q.shape[-1]}") + + cu_seqlens = cu_seqlens.to(device=q.device, dtype=torch.int32).contiguous() + chunk_indices = chunk_indices.to(device=q.device, dtype=torch.int32).contiguous() + beta_fp32 = beta.float().contiguous() + dq_out = torch.empty_like(q) if dq_out is None else dq_out + dk_out = torch.empty_like(k) if dk_out is None else dk_out + db_out = torch.empty_like(db, dtype=torch.float32) if db_out is None else db_out + dg_out = torch.empty_like(dg, dtype=torch.float32) if dg_out is None else dg_out + + cula_cuda.chunk_kda_bwd_intra_cuda( + q, + k, + g, + beta_fp32, + d_aq, + d_ak, + dq, + dk, + db, + dg, + cu_seqlens, + chunk_indices, + dq_out, + dk_out, + db_out, + dg_out, + chunk_size, + ) + return dq_out, dk_out, db_out, dg_out + + +__all__ = ["kda_bwd_intra_mma"] diff --git a/setup.py b/setup.py index a187764a..073df5de 100644 --- a/setup.py +++ b/setup.py @@ -172,8 +172,10 @@ def get_nvcc_thread_args(): CUDAExtension( name="cula._cudac_sm100", sources=[ + "csrc/api/kda_bwd_intra.cu", "csrc/api/kda_sm100.cu", "csrc/api/pybind_sm100.cu", + "csrc/kda/sm90/bwd/kda_bwd_intra_sm90.cu", "csrc/kda/sm100/kda_fwd_sm100.cu", ], extra_compile_args={ @@ -195,8 +197,10 @@ def get_nvcc_thread_args(): CUDAExtension( name="cula._cudac_sm90", sources=[ + "csrc/api/kda_bwd_intra.cu", "csrc/api/kda_sm90.cu", "csrc/api/pybind_sm90.cu", + "csrc/kda/sm90/bwd/kda_bwd_intra_sm90.cu", "csrc/kda/sm90/kda_fwd_sm90.cu", "csrc/kda/sm90/kda_fwd_sm90_safe_gate.cu", ], diff --git a/tests/test_kda_sm90_bwd_intra.py b/tests/test_kda_sm90_bwd_intra.py new file mode 100644 index 00000000..f8a6a8b5 --- /dev/null +++ b/tests/test_kda_sm90_bwd_intra.py @@ -0,0 +1,136 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from fla.ops.kda.chunk_intra import chunk_kda_bwd_intra as fla_bwd_intra +from fla.ops.utils import prepare_chunk_indices + +import cula.kda.chunk_intra as chunk_intra_module +from cula.kda.chunk_intra import _is_mma_bwd_intra_supported +from cula.kda.chunk_intra import chunk_kda_bwd_intra as cula_bwd_intra +from cula.ops.kda.sm90.bwd_intra import kda_bwd_intra_mma + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +_SUPPORTED_CAPABILITIES = {(9, 0), (10, 0), (10, 3)} +_LIMITS = (8.0e-3, 8.0e-3, 2.0e-2, 2.0e-2) + + +def _require_supported_device() -> None: + capability = torch.cuda.get_device_capability() + if capability not in _SUPPORTED_CAPABILITIES: + pytest.skip(f"mma.sync kernel does not support SM{capability[0]}{capability[1]}") + + +def _relative_rmse(actual: torch.Tensor, expected: torch.Tensor) -> float: + actual = actual.float() + expected = expected.float() + rmse = (actual - expected).square().mean().sqrt() + return (rmse / (expected.square().mean().sqrt() + 1e-8)).item() + + +def _make_inputs(lengths: list[int], heads: int = 4): + torch.manual_seed(42) + total = sum(lengths) + dim, chunk_size = 128, 64 + device = torch.device("cuda") + offsets = [0] + for length in lengths: + offsets.append(offsets[-1] + length) + cu_seqlens = torch.tensor(offsets, device=device, dtype=torch.int32) + chunk_indices = prepare_chunk_indices(cu_seqlens.to(torch.long), chunk_size).to(torch.int32).contiguous() + + q = torch.randn(1, total, heads, dim, device=device, dtype=torch.bfloat16) + k = torch.randn_like(q) + g = torch.randn(1, total, heads, dim, device=device, dtype=torch.float32) / 10 + beta = torch.randn(1, total, heads, device=device, dtype=torch.bfloat16) + d_aq = torch.randn(1, total, heads, chunk_size, device=device, dtype=torch.float32) + d_ak = torch.randn_like(d_aq) + dq = torch.randn(1, total, heads, dim, device=device, dtype=torch.float32) + dk = torch.randn_like(dq) + db = torch.randn(1, total, heads, device=device, dtype=torch.float32) + dg = torch.randn_like(dq) + return q, k, g, beta, d_aq, d_ak, dq, dk, db, dg, cu_seqlens, chunk_indices + + +def _assert_matches_fla(inputs) -> None: + reference = fla_bwd_intra(*inputs, chunk_size=64, safe_gate=True) + actual = kda_bwd_intra_mma(*inputs, chunk_size=64) + torch.cuda.synchronize() + errors = tuple(_relative_rmse(got, ref) for got, ref in zip(actual, reference)) + assert all(error < limit for error, limit in zip(errors, _LIMITS)), (errors, _LIMITS) + assert tuple(value.dtype for value in actual) == ( + torch.bfloat16, + torch.bfloat16, + torch.float32, + torch.float32, + ) + assert all(torch.isfinite(value).all() for value in actual) + + +@pytest.mark.kda_fast +@pytest.mark.parametrize("lengths", [[64], [65], [64, 127, 129]]) +def test_kda_bwd_intra_mma_matches_fla(lengths: list[int]): + _require_supported_device() + _assert_matches_fla(_make_inputs(lengths)) + + +@pytest.mark.kda_fast +def test_kda_bwd_intra_mma_is_deterministic(): + _require_supported_device() + inputs = _make_inputs([64, 65], heads=2) + expected = tuple(value.clone() for value in kda_bwd_intra_mma(*inputs, chunk_size=64)) + for _ in range(20): + actual = kda_bwd_intra_mma(*inputs, chunk_size=64) + assert all(torch.equal(got, ref) for got, ref in zip(actual, expected)) + + +@pytest.mark.kda_fast +def test_kda_bwd_intra_dispatch_matches_direct_kernel(): + _require_supported_device() + inputs = _make_inputs([64, 65], heads=2) + direct = kda_bwd_intra_mma(*inputs, chunk_size=64) + dispatched = cula_bwd_intra(*inputs, chunk_size=64, safe_gate=True) + torch.cuda.synchronize() + assert all(torch.equal(got, ref) for got, ref in zip(dispatched, direct)) + + +@pytest.mark.kda_fast +def test_kda_bwd_intra_dispatch_handles_dense_batches(): + _require_supported_device() + flat_inputs = _make_inputs([70, 70], heads=2) + reference = fla_bwd_intra(*flat_inputs, chunk_size=64, safe_gate=True) + batched = tuple(value.reshape(2, 70, *value.shape[2:]) for value in flat_inputs[:10]) + actual = cula_bwd_intra(*batched, chunk_size=64, safe_gate=True) + torch.cuda.synchronize() + errors = tuple(_relative_rmse(got.reshape_as(ref), ref) for got, ref in zip(actual, reference)) + assert all(error < limit for error, limit in zip(errors, _LIMITS)), errors + + +def test_kda_bwd_intra_float_beta_falls_back_to_triton(monkeypatch: pytest.MonkeyPatch): + _require_supported_device() + inputs = list(_make_inputs([64], heads=1)) + inputs[3] = inputs[3].float() + sentinel = object() + + def fake_triton(*args, **kwargs): + return sentinel + + monkeypatch.setattr(chunk_intra_module, "_chunk_kda_bwd_intra_triton", fake_triton) + assert cula_bwd_intra(*inputs, chunk_size=64, safe_gate=True) is sentinel + + +def test_kda_bwd_intra_support_predicate_rejects_cpu_inputs(): + tensors = _make_inputs([64], heads=1)[:10] + cpu_tensors = tuple(tensor.cpu() for tensor in tensors) + assert not _is_mma_bwd_intra_supported(*cpu_tensors, chunk_size=64, safe_gate=True) + + +@pytest.mark.kda_fast +def test_kda_bwd_intra_mma_rejects_empty_chunk_indices(): + _require_supported_device() + inputs = list(_make_inputs([64], heads=1)) + inputs[-1] = torch.empty((0, 2), device="cuda", dtype=torch.int32) + with pytest.raises(RuntimeError, match="chunk_indices must contain at least one chunk"): + kda_bwd_intra_mma(*inputs, chunk_size=64)