diff --git a/benchmarks/bench_kda_sm100_intra_fused.py b/benchmarks/bench_kda_sm100_intra_fused.py new file mode 100644 index 00000000..553523d4 --- /dev/null +++ b/benchmarks/bench_kda_sm100_intra_fused.py @@ -0,0 +1,764 @@ +#!/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. + +"""Benchmark/profiling entry for SM100 KDA K123 inverse variants. + +Examples: + python benchmarks/bench_kda_sm100_intra_fused.py --mode both + python benchmarks/bench_kda_sm100_intra_fused.py --mode varlen --seq-lens 288 + /usr/local/cuda-13/bin/ncu --profile-from-start off --set full -o ncu_reports/kda_fwd_intra_sm100_varlen \ + .venv/bin/python benchmarks/bench_kda_sm100_intra_fused.py --ncu --mode varlen +""" + +from __future__ import annotations + +import argparse +import pathlib +import sys +import time +from collections.abc import Callable +from dataclasses import dataclass + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +import torch +import torch.nn.functional as F +from fla.ops.kda.chunk_intra import chunk_kda_fwd_intra as fla_chunk_kda_fwd_intra +from fla.ops.kda.gate import kda_gate_chunk_cumsum +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.constant import RCP_LN2 + +from benchmarks.utils import exclusive_cumsum, gen_random, gen_skewed, gen_uniform, set_seed +from cula.kda.chunk_intra import chunk_kda_fwd_intra as cula_chunk_kda_fwd_intra +from cula.ops.kda.sm100.intra_fused import ( + BT, + K_DIM, + chunk_kda_fwd_intra_sm100_equal, + chunk_kda_fwd_intra_sm100_varlen, +) +from cula.ops.kda.sm100.recompute_wu import recompute_w_u_from_preprocessed + + +@dataclass +class BenchCase: + name: str + detail: str + run: Callable[[], tuple[torch.Tensor, ...]] + run_fla: Callable[[], tuple[torch.Tensor, ...]] + compare: Callable[ + [tuple[torch.Tensor, ...], tuple[torch.Tensor, ...]], list[tuple[str, tuple[float, float, float, float]]] + ] + + +def _l2norm_bf16(x: torch.Tensor) -> torch.Tensor: + return F.normalize(x.float(), p=2.0, dim=-1).to(torch.bfloat16) + + +def _make_inputs(args: argparse.Namespace, batch: int, total_t: int): + set_seed(args.seed) + device = torch.device(args.device) + q = _l2norm_bf16(torch.randn(batch, total_t, args.H, args.K, device=device, dtype=torch.bfloat16)) + k = _l2norm_bf16(torch.randn(batch, total_t, args.H, args.K, device=device, dtype=torch.bfloat16)) + g = (torch.randn(batch, total_t, args.H, args.K, device=device, dtype=torch.bfloat16) * args.g_scale).bfloat16() + beta = torch.randn(batch, total_t, args.H, device=device, dtype=torch.float32).sigmoid().bfloat16() + a_log = torch.randn(args.H, device=device, dtype=torch.float32) * args.alog_scale + dt_bias = None if args.no_bias else torch.randn(args.H * args.K, device=device, dtype=torch.float32) * args.dt_bias_scale + return q, k, g, beta, a_log, dt_bias + + +def _parse_seq_lens(text: str) -> list[int]: + seq_lens = [int(x) for x in text.replace(",", " ").split()] + if not seq_lens or any(x <= 0 for x in seq_lens): + raise ValueError(f"invalid --seq-lens: {text!r}") + return seq_lens + + +def _build_seq_lens(args: argparse.Namespace) -> list[int]: + if args.seq_lens: + return _parse_seq_lens(args.seq_lens) + if args.dist == "uniform": + return gen_uniform(args.num_seqs, args.T) + if args.dist == "skewed": + return gen_skewed(args.num_seqs, args.T) + if args.dist == "random": + return gen_random(args.num_seqs, args.T, seed=args.seed) + raise ValueError(f"unknown --dist {args.dist}") + + +def _varlen_launch_summary(seq_lens: list[int]) -> tuple[int, int, bool]: + t_launch = 0 + launch_bos = [] + all_bt_aligned = True + nt = 0 + for seq_len in seq_lens: + chunks = (seq_len + BT - 1) // BT + launch_bos.append(t_launch) + nt += chunks + t_launch += chunks * BT + all_bt_aligned = all_bt_aligned and (seq_len % BT == 0) + cu = exclusive_cumsum(seq_lens) + pure = all_bt_aligned and (nt % 4) == 0 and launch_bos == cu[:-1] + return nt, t_launch, pure + + +def accuracy_stats( + ref: torch.Tensor, out: torch.Tensor, mask: torch.Tensor | None = None +) -> tuple[float, float, float, float]: + """Return rel_rmse, rel_max, max_abs, and mean_abs.""" + if mask is not None: + ref = ref[mask] + out = out[mask] + ref_f = ref.float() + out_f = out.float() + diff = (ref_f - out_f).abs() + rmse = diff.square().mean().sqrt().item() + ref_rms = ref_f.square().mean().sqrt().item() + rel_rmse = rmse / (ref_rms + 1e-8) + max_abs = diff.max().item() + ref_max = ref_f.abs().max().item() + rel_max = max_abs / ref_max if ref_max > 0 else 0.0 + mean_abs = diff.mean().item() + return rel_rmse, rel_max, max_abs, mean_abs + + +def _valid_lower_mask(batch: int, total_t: int, heads: int, device: torch.device) -> torch.Tensor: + row = torch.arange(total_t, device=device) % BT + col = torch.arange(BT, device=device) + return (col[None, :] <= row[:, None]).view(1, total_t, 1, BT).expand(batch, total_t, heads, BT) + + +def _valid_lower_mask_varlen(seq_lens: list[int], heads: int, device: torch.device) -> torch.Tensor: + total_t = sum(seq_lens) + valid = torch.zeros(1, total_t, heads, BT, device=device, dtype=torch.bool) + col = torch.arange(BT, device=device) + offset = 0 + for seq_len in seq_lens: + row = torch.arange(seq_len, device=device) % BT + seq_valid = (col[None, :] <= row[:, None]).view(1, seq_len, 1, BT).expand(1, seq_len, heads, BT) + valid[:, offset : offset + seq_len] = seq_valid + offset += seq_len + return valid + + +def _reference_gk_equal( + g: torch.Tensor, a_log: torch.Tensor, dt_bias: torch.Tensor | None, lower_bound: float +) -> torch.Tensor: + batch, total_t, _, _ = g.shape + outs = [] + for b in range(batch): + cu_seqlens = torch.tensor([0, total_t], dtype=torch.int32, device=g.device) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + outs.append( + kda_gate_chunk_cumsum( + g=g[b : b + 1].float(), + A_log=a_log, + dt_bias=dt_bias, + scale=RCP_LN2, + chunk_size=BT, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=lower_bound, + ) + ) + return torch.cat(outs, dim=0).contiguous() + + +def _gk_last_exp_equal(gk: torch.Tensor) -> torch.Tensor: + return torch.exp2(gk[:, BT - 1 :: BT]).contiguous() + + +def _gk_last_exp_varlen(gk: torch.Tensor, seq_lens: list[int]) -> torch.Tensor: + rows = [] + offset = 0 + for seq_len in seq_lens: + for chunk_start in range(0, seq_len, BT): + last = offset + min(chunk_start + BT, seq_len) - 1 + rows.append(gk[:, last]) + offset += seq_len + return torch.stack(rows, dim=1).contiguous() + + +def _stats_map(stats: list[tuple[str, tuple[float, float, float, float]]]) -> dict[str, tuple[float, float, float, float]]: + return {name: values for name, values in stats} + + +def _format_stat(values: tuple[float, float, float, float]) -> str: + rel_rmse, rel_max, max_abs, mean_abs = values + return f"rel_rmse={rel_rmse:.3e} rel_max={rel_max:.3e} max_abs={max_abs:.3e} mean_abs={mean_abs:.3e}" + + +def make_equal_case(args: argparse.Namespace) -> BenchCase: + if args.T % (4 * BT) != 0: + raise NotImplementedError(f"equal mode requires T to be a multiple of {4 * BT}, got {args.T}.") + q, k, g, beta, a_log, dt_bias = _make_inputs(args, args.B, args.T) + scale = args.K**-0.5 + if args.cutedsl_variant == "flashinfer-k123-copy": + cutedsl_fn = chunk_kda_fwd_intra_sm100_equal + elif args.cutedsl_variant == "flashinfer-k123-copy-fp32-inv": + + def cutedsl_fn(**kwargs): + return chunk_kda_fwd_intra_sm100_equal( + **kwargs, + fp32_akk_inv=True, + ) + + elif args.cutedsl_variant == "flashinfer-k123-incta-inv": + + def cutedsl_fn(**kwargs): + return chunk_kda_fwd_intra_sm100_equal( + **kwargs, + fused_akk_inv=True, + ) + + elif args.cutedsl_variant == "flashinfer-k123-copy-fp32-inv-wbeta": + + def cutedsl_fn(**kwargs): + return chunk_kda_fwd_intra_sm100_equal( + **kwargs, + fp32_akk_inv=True, + preprocess_w_beta=True, + ) + + elif args.cutedsl_variant == "flashinfer-k123-copy-fp32-inv-fp16-k": + + def cutedsl_fn(**kwargs): + return chunk_kda_fwd_intra_sm100_equal( + **kwargs, + fp32_akk_inv=True, + kscaled_fp16=True, + ) + + else: + raise NotImplementedError(f"Unsupported CuTeDSL variant: {args.cutedsl_variant}") + + def run(): + out = cutedsl_fn( + q=q, + k=k, + g=g, + beta=beta, + A_log=a_log, + dt_bias=dt_bias, + scale=scale, + safe_gate=True, + lower_bound=args.lower_bound, + ) + if args.with_recompute_wu: + w, u = recompute_w_u_from_preprocessed( + out[0], + k, + beta, + out[5], + k_includes_beta=args.cutedsl_variant == "flashinfer-k123-copy-fp32-inv-wbeta", + ) + return (*out, w, u) + return out + + def run_fla_with_gk(): + gk = _reference_gk_equal(g, a_log, dt_bias, args.lower_bound) + cu_seqlens = None + chunk_indices = None + if args.B == 1: + cu_seqlens = torch.tensor([0, args.T], dtype=torch.int32, device=q.device) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + fla = fla_chunk_kda_fwd_intra( + q=q, + k=k, + v=k, + gk=gk, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=BT, + safe_gate=True, + disable_recompute=True, + ) + return (gk, *fla) + + def run_csrc_with_gk(): + gk = kda_gate_chunk_cumsum( + g=g, + A_log=a_log, + dt_bias=dt_bias, + scale=RCP_LN2, + chunk_size=BT, + lower_bound=args.lower_bound, + ) + csrc = cula_chunk_kda_fwd_intra( + q=q, + k=k, + v=k, + gk=gk, + beta=beta, + scale=scale, + chunk_size=BT, + safe_gate=True, + ) + return (gk, *csrc) + + valid_mask = _valid_lower_mask(args.B, args.T, args.H, q.device) + + def compare(cutedsl: tuple[torch.Tensor, ...], fla_with_gk: tuple[torch.Tensor, ...]): + gk, _, _, _, kg_fla, aqk_fla, akk_fla = fla_with_gk + exp_gk = torch.exp2(gk) + if args.cutedsl_variant == "flashinfer-k123-copy-fp32-inv-wbeta": + k_scaled_ref = ((k.float() * beta.unsqueeze(-1).float()) * exp_gk).to(torch.bfloat16) + elif args.cutedsl_variant == "flashinfer-k123-copy-fp32-inv-fp16-k": + k_scaled_ref = (k.float() * exp_gk).to(torch.float16) + else: + k_scaled_ref = (k.float() * exp_gk).to(torch.bfloat16) + q_scaled_ref = (q.float() * exp_gk).to(torch.bfloat16) + stats = [ + ("k_scaled", accuracy_stats(k_scaled_ref, cutedsl[0])), + ("kg", accuracy_stats(kg_fla, cutedsl[1])), + ("q_scaled", accuracy_stats(q_scaled_ref, cutedsl[2])), + ("Aqk_valid", accuracy_stats(aqk_fla, cutedsl[4], valid_mask)), + ] + if args.cutedsl_variant in ( + "flashinfer-k123-copy-fp32-inv", + "flashinfer-k123-copy-fp32-inv-fp16-k", + "flashinfer-k123-copy-fp32-inv-wbeta", + "flashinfer-k123-incta-inv", + ): + stats.append(("Akk_valid", accuracy_stats(akk_fla, cutedsl[5], valid_mask))) + if args.with_recompute_wu: + stats.extend( + [ + ("w", accuracy_stats(fla_with_gk[1], cutedsl[6])), + ("u", accuracy_stats(fla_with_gk[2], cutedsl[7])), + ] + ) + return stats + + nt = args.B * (args.T // BT) + detail = f"B={args.B} T={args.T} H={args.H} K={args.K} NT={nt} bias={not args.no_bias} variant={args.cutedsl_variant}" + run_ref = run_csrc_with_gk if args.baseline == "csrc" else run_fla_with_gk + return BenchCase("equal", detail, run, run_ref, compare) + + +def make_varlen_case(args: argparse.Namespace) -> BenchCase: + if args.cutedsl_variant not in ( + "flashinfer-k123-copy", + "flashinfer-k123-copy-fp32-inv", + "flashinfer-k123-copy-fp32-inv-fp16-k", + "flashinfer-k123-copy-fp32-inv-wbeta", + "flashinfer-k123-incta-inv", + ): + raise NotImplementedError(f"{args.cutedsl_variant} CuTeDSL variant currently supports equal mode only.") + seq_lens = _build_seq_lens(args) + total_t = sum(seq_lens) + q, k, g, beta, a_log, dt_bias = _make_inputs(args, 1, total_t) + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=q.device) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + scale = args.K**-0.5 + + if args.cutedsl_variant == "flashinfer-k123-copy": + + def run(): + out = chunk_kda_fwd_intra_sm100_varlen( + q=q, + k=k, + g=g, + beta=beta, + A_log=a_log, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + dt_bias=dt_bias, + safe_gate=True, + lower_bound=args.lower_bound, + seq_lens=seq_lens, + ) + if args.with_recompute_wu: + w, u = recompute_w_u_from_preprocessed( + out[0], + k, + beta, + out[5], + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return (*out, w, u) + return out + + elif args.cutedsl_variant == "flashinfer-k123-copy-fp32-inv": + + def run(): + out = chunk_kda_fwd_intra_sm100_varlen( + q=q, + k=k, + g=g, + beta=beta, + A_log=a_log, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + dt_bias=dt_bias, + safe_gate=True, + lower_bound=args.lower_bound, + seq_lens=seq_lens, + fp32_akk_inv=True, + ) + if args.with_recompute_wu: + w, u = recompute_w_u_from_preprocessed( + out[0], + k, + beta, + out[5], + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return (*out, w, u) + return out + + elif args.cutedsl_variant == "flashinfer-k123-incta-inv": + + def run(): + out = chunk_kda_fwd_intra_sm100_varlen( + q=q, + k=k, + g=g, + beta=beta, + A_log=a_log, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + dt_bias=dt_bias, + safe_gate=True, + lower_bound=args.lower_bound, + seq_lens=seq_lens, + fused_akk_inv=True, + ) + if args.with_recompute_wu: + w, u = recompute_w_u_from_preprocessed( + out[0], + k, + beta, + out[5], + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return (*out, w, u) + return out + + elif args.cutedsl_variant == "flashinfer-k123-copy-fp32-inv-wbeta": + + def run(): + out = chunk_kda_fwd_intra_sm100_varlen( + q=q, + k=k, + g=g, + beta=beta, + A_log=a_log, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + dt_bias=dt_bias, + safe_gate=True, + lower_bound=args.lower_bound, + seq_lens=seq_lens, + fp32_akk_inv=True, + preprocess_w_beta=True, + ) + if args.with_recompute_wu: + w, u = recompute_w_u_from_preprocessed( + out[0], + k, + beta, + out[5], + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + k_includes_beta=True, + ) + return (*out, w, u) + return out + + elif args.cutedsl_variant == "flashinfer-k123-copy-fp32-inv-fp16-k": + + def run(): + out = chunk_kda_fwd_intra_sm100_varlen( + q=q, + k=k, + g=g, + beta=beta, + A_log=a_log, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + dt_bias=dt_bias, + safe_gate=True, + lower_bound=args.lower_bound, + seq_lens=seq_lens, + fp32_akk_inv=True, + kscaled_fp16=True, + ) + if args.with_recompute_wu: + w, u = recompute_w_u_from_preprocessed( + out[0], + k, + beta, + out[5], + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return (*out, w, u) + return out + + def run_fla_with_gk(): + gk = kda_gate_chunk_cumsum( + g=g.float(), + A_log=a_log, + dt_bias=dt_bias, + scale=RCP_LN2, + chunk_size=BT, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=args.lower_bound, + ) + fla = fla_chunk_kda_fwd_intra( + q=q, + k=k, + v=k, + gk=gk, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=BT, + safe_gate=True, + disable_recompute=True, + ) + return (gk, *fla) + + def run_csrc_with_gk(): + gk = kda_gate_chunk_cumsum( + g=g, + A_log=a_log, + dt_bias=dt_bias, + scale=RCP_LN2, + chunk_size=BT, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=args.lower_bound, + ) + csrc = cula_chunk_kda_fwd_intra( + q=q, + k=k, + v=k, + gk=gk, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=BT, + safe_gate=True, + ) + return (gk, *csrc) + + valid_mask = _valid_lower_mask_varlen(seq_lens, args.H, q.device) + + def compare(cutedsl: tuple[torch.Tensor, ...], fla_with_gk: tuple[torch.Tensor, ...]): + gk, _, _, _, kg_fla, aqk_fla, akk_fla = fla_with_gk + exp_gk = torch.exp2(gk) + if args.cutedsl_variant == "flashinfer-k123-copy-fp32-inv-wbeta": + k_scaled_ref = ((k.float() * beta.unsqueeze(-1).float()) * exp_gk).to(torch.bfloat16) + elif args.cutedsl_variant == "flashinfer-k123-copy-fp32-inv-fp16-k": + k_scaled_ref = (k.float() * exp_gk).to(torch.float16) + else: + k_scaled_ref = (k.float() * exp_gk).to(torch.bfloat16) + q_scaled_ref = (q.float() * exp_gk).to(torch.bfloat16) + stats = [ + ("k_scaled", accuracy_stats(k_scaled_ref, cutedsl[0])), + ("kg", accuracy_stats(kg_fla, cutedsl[1])), + ("q_scaled", accuracy_stats(q_scaled_ref, cutedsl[2])), + ("Aqk_valid", accuracy_stats(aqk_fla, cutedsl[4], valid_mask)), + ] + if args.cutedsl_variant in ( + "flashinfer-k123-copy-fp32-inv", + "flashinfer-k123-copy-fp32-inv-fp16-k", + "flashinfer-k123-copy-fp32-inv-wbeta", + "flashinfer-k123-incta-inv", + ): + stats.append(("Akk_valid", accuracy_stats(akk_fla, cutedsl[5], valid_mask))) + if args.with_recompute_wu: + stats.extend( + [ + ("w", accuracy_stats(fla_with_gk[1], cutedsl[6])), + ("u", accuracy_stats(fla_with_gk[2], cutedsl[7])), + ] + ) + return stats + + nt, t_launch, pure = _varlen_launch_summary(seq_lens) + seq_preview = ",".join(str(x) for x in seq_lens[:8]) + if len(seq_lens) > 8: + seq_preview += ",..." + detail = ( + f"seqs={len(seq_lens)} total={total_t} H={args.H} K={args.K} NT={nt} " + f"T_launch={t_launch} pure={pure} dist={args.dist} bias={not args.no_bias} seq_lens=[{seq_preview}]" + ) + run_ref = run_csrc_with_gk if args.baseline == "csrc" else run_fla_with_gk + return BenchCase("varlen", detail, run, run_ref, compare) + + +def time_case(fn: Callable[[], tuple[torch.Tensor, ...]], warmup: int, iters: int) -> tuple[float, float]: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + wall_start = time.perf_counter() + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + wall_us = (time.perf_counter() - wall_start) * 1_000_000.0 / iters + event_us = start.elapsed_time(end) / iters * 1000.0 + return event_us, wall_us + + +def run_bench(cases: list[BenchCase], args: argparse.Namespace) -> None: + print("=" * 100, flush=True) + print(f" SM100 KDA K123 benchmark vs {args.baseline} chunk_intra", flush=True) + print(f" Device: {torch.cuda.get_device_name(0)}", flush=True) + print(f" BT={BT} H={args.H} K={args.K} warmup={args.warmup} iters={args.iters}", flush=True) + print(f" Baseline: kda_gate_chunk_cumsum + {args.baseline} chunk_kda_fwd_intra", flush=True) + print("=" * 100, flush=True) + for case in cases: + print(f"\n {'-' * 144}", flush=True) + print(f" {case.name}: {case.detail}", flush=True) + print(f" {'-' * 144}", flush=True) + print( + f" {'Ref(us)':>10} {'CuTeDSL(us)':>12} {'wall/ev':>8} │ {'Ref/CuTe':>9} │ {'Aqk rel_rmse':>13} {'Akk rel_rmse':>13}", + flush=True, + ) + print(f" {'-' * 144}", flush=True) + + cutedsl_out = case.run() + torch.cuda.synchronize() + cutedsl_event_us, cutedsl_wall_us = time_case(case.run, args.warmup, args.iters) + # wall/event ~1 means the GPU is the bottleneck and the timing is + # trustworthy; >>1 means host-side launch overhead is leaking into the + # measured interval (still unstable). + cutedsl_ratio = cutedsl_wall_us / cutedsl_event_us if cutedsl_event_us > 0 else float("nan") + if args.no_fla: + print( + f" {'N/A':>10} {cutedsl_event_us:12.1f} {cutedsl_ratio:8.2f} │ {'N/A':>9} │ {'N/A':>13} {'N/A':>13}", + flush=True, + ) + continue + + fla_out = case.run_fla() + torch.cuda.synchronize() + fla_event_us, _ = time_case(case.run_fla, args.warmup, args.iters) + speedup = fla_event_us / cutedsl_event_us if cutedsl_event_us > 0 else float("inf") + stats = case.compare(cutedsl_out, fla_out) + stats_by_name = _stats_map(stats) + aqk_rel = stats_by_name["Aqk_valid"][0] + akk_rel = stats_by_name["Akk_valid"][0] if "Akk_valid" in stats_by_name else float("nan") + akk_str = f"{akk_rel:13.3e}" if akk_rel == akk_rel else f"{'N/A':>13}" + print( + f" {fla_event_us:10.1f} {cutedsl_event_us:12.1f} {cutedsl_ratio:8.2f} │ {speedup:8.2f}x │ {aqk_rel:13.3e} {akk_str}", + flush=True, + ) + print(" Accuracy details vs FLA/reference:", flush=True) + for name, values in stats: + print(f" {name:<26} {_format_stat(values)}", flush=True) + + +def run_ncu(case: BenchCase, args: argparse.Namespace) -> None: + print( + f"[NCU profiler] case={case.name} {case.detail} warmup={args.profile_warmup} profile_iters={args.profile_iters}", + flush=True, + ) + case.run() + torch.cuda.synchronize() + for _ in range(args.profile_warmup): + case.run() + torch.cuda.synchronize() + torch.cuda.cudart().cudaProfilerStart() + for _ in range(args.profile_iters): + case.run() + torch.cuda.synchronize() + torch.cuda.cudart().cudaProfilerStop() + print("[NCU profiler] done", flush=True) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--mode", choices=("equal", "varlen", "both"), default="both") + parser.add_argument("--B", type=int, default=1) + parser.add_argument("--T", type=int, default=8192) + parser.add_argument("--H", type=int, default=64) + parser.add_argument("--K", type=int, default=K_DIM) + parser.add_argument("--device", default="cuda") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=20) + parser.add_argument("--profile-warmup", type=int, default=2) + parser.add_argument("--profile-iters", type=int, default=1) + parser.add_argument("--num-seqs", type=int, default=8) + parser.add_argument("--dist", choices=("uniform", "random", "skewed"), default="random") + parser.add_argument("--seq-lens", help="Comma/space separated varlen sequence lengths; overrides --T/--num-seqs/--dist.") + parser.add_argument("--lower-bound", type=float, default=-5.0) + parser.add_argument("--g-scale", type=float, default=1.0) + parser.add_argument("--alog-scale", type=float, default=1.0) + parser.add_argument("--dt-bias-scale", type=float, default=1.0) + parser.add_argument("--no-bias", action="store_true") + parser.add_argument( + "--no-fla", action="store_true", help="Only time CuTeDSL; skip baseline timing and accuracy comparison." + ) + parser.add_argument("--baseline", choices=("csrc", "fla"), default="csrc") + parser.add_argument( + "--with-recompute-wu", + action="store_true", + help="Include the specialized w/u recompute after fused intra in CuTeDSL timing.", + ) + parser.add_argument( + "--cutedsl-variant", + choices=( + "flashinfer-k123-copy", + "flashinfer-k123-copy-fp32-inv", + "flashinfer-k123-copy-fp32-inv-fp16-k", + "flashinfer-k123-copy-fp32-inv-wbeta", + "flashinfer-k123-incta-inv", + ), + default="flashinfer-k123-copy-fp32-inv-fp16-k", + help=("CuTeDSL variant to benchmark. Variants without Akk inverse skip Akk accuracy."), + ) + parser.add_argument("--ncu", action="store_true", help="Run one case under cudaProfilerStart/Stop for Nsight Compute.") + args = parser.parse_args() + + if args.K != K_DIM: + raise NotImplementedError(f"SM100 training intra currently supports K={K_DIM}, got {args.K}.") + if args.with_recompute_wu and args.cutedsl_variant not in ( + "flashinfer-k123-copy-fp32-inv", + "flashinfer-k123-copy-fp32-inv-fp16-k", + "flashinfer-k123-copy-fp32-inv-wbeta", + "flashinfer-k123-incta-inv", + ): + raise ValueError("--with-recompute-wu requires an Akk-inverse variant.") + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required.") + + cases: list[BenchCase] = [] + if args.mode in ("equal", "both"): + cases.append(make_equal_case(args)) + if args.mode in ("varlen", "both"): + cases.append(make_varlen_case(args)) + + if args.ncu: + if len(cases) != 1: + raise ValueError("--ncu requires --mode equal or --mode varlen so the report captures one target case.") + run_ncu(cases[0], args) + else: + run_bench(cases, args) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/bench_kda_sm100_recompute_wu.py b/benchmarks/bench_kda_sm100_recompute_wu.py new file mode 100644 index 00000000..1552ab11 --- /dev/null +++ b/benchmarks/bench_kda_sm100_recompute_wu.py @@ -0,0 +1,118 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Compare SM100 C++ and CuTe DSL recompute-WU kernels on the same inputs.""" + +import argparse +import pathlib +import sys + +import torch + +sys.path.insert(0, str(pathlib.Path(__file__).parent.parent)) + +import cula.cudac as cula_cuda +from benchmarks.bench_recompute_wu import prepare_recompute_wu_inputs +from benchmarks.utils import relative_rms_error_rel_max_mean_abs_rhs, triton_bench_fn +from cula.ops.kda.sm100 import recompute_wu as recompute_wu_module +from cula.ops.kda.sm100.recompute_wu import recompute_w_u_fwd + + +def _run_cpp(k, v, beta, A, gk, cu_seqlens, chunk_indices): + w = torch.empty_like(v) + u = torch.empty_like(v) + kg = torch.empty_like(v) + cula_cuda.recompute_w_u_cuda( + k, + v, + beta, + A, + gk, + cu_seqlens, + chunk_indices, + w, + u, + kg, + A.shape[-1], + None, + None, + ) + return w, u, None, kg + + +def _max_error(ref, out): + stats = [relative_rms_error_rel_max_mean_abs_rhs(a, b) for a, b in zip(ref, out) if a is not None] + return tuple(max(values) for values in zip(*stats)) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--heads", type=int, default=64) + parser.add_argument("--beta-bf16", action="store_true") + parser.add_argument( + "--force-varlen", + action="store_true", + help="Keep packed uniform inputs on the varlen kernel path (diagnostic only)", + ) + parser.add_argument("--lengths", type=int, nargs="+", default=[512, 1024, 4096, 8192, 16384, 32768]) + parser.add_argument( + "--profile", + choices=("cpp", "ws"), + help="Warm up, then launch exactly one selected kernel between CUDA profiler markers", + ) + args = parser.parse_args() + + if args.force_varlen: + recompute_wu_module._uniform_problem = lambda _cu_seqlens: None + + import benchmarks.bench_recompute_wu as common + + common.H = args.heads + common.HV = args.heads + device = torch.device("cuda") + if args.profile: + if len(args.lengths) != 1: + parser.error("--profile requires exactly one value in --lengths") + T = args.lengths[0] + cu_seqlens = torch.tensor([0, T, 2 * T], dtype=torch.int32, device=device) + _q, k, v, cu_gk, beta, A, cu_seqlens, chunk_indices = prepare_recompute_wu_inputs( + 2, + T, + device, + cu_seqlens=cu_seqlens, + ) + if args.beta_bf16: + beta = beta.bfloat16() + runners = { + "cpp": lambda: _run_cpp(k, v, beta, A, cu_gk, cu_seqlens, chunk_indices), + "ws": lambda: recompute_w_u_fwd(k, v, beta, A, cu_gk, cu_seqlens, chunk_indices), + } + runner = runners[args.profile] + runner() + torch.cuda.synchronize() + torch.cuda.cudart().cudaProfilerStart() + runner() + torch.cuda.synchronize() + torch.cuda.cudart().cudaProfilerStop() + print(f"profiled {args.profile} at T={T}, H={args.heads}") + return + + print(f"{'T':>8} {'C++ (ms)':>12} {'CuTeDSL (ms)':>14} {'C++/CuTeDSL':>14} {'rel_rmse':>12}") + for T in args.lengths: + cu_seqlens = torch.tensor([0, T, 2 * T], dtype=torch.int32, device=device) + _q, k, v, cu_gk, beta, A, cu_seqlens, chunk_indices = prepare_recompute_wu_inputs(2, T, device, cu_seqlens=cu_seqlens) + if args.beta_bf16: + beta = beta.bfloat16() + cpp = _run_cpp(k, v, beta, A, cu_gk, cu_seqlens, chunk_indices) + ws = recompute_w_u_fwd(k, v, beta, A, cu_gk, cu_seqlens, chunk_indices) + ws_err = _max_error(cpp, ws)[0] + if not torch.isfinite(torch.tensor(ws_err)): + raise AssertionError(f"non-finite error at T={T}: ws={ws_err}") + + cpp_ms = triton_bench_fn(lambda: _run_cpp(k, v, beta, A, cu_gk, cu_seqlens, chunk_indices)) + ws_ms = triton_bench_fn(lambda: recompute_w_u_fwd(k, v, beta, A, cu_gk, cu_seqlens, chunk_indices)) + print(f"{T:8d} {cpp_ms:12.4f} {ws_ms:14.4f} {cpp_ms / ws_ms:14.3f} {ws_err:12.6g}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/stress_kda_sm100_csrc_boundary_determinism.py b/benchmarks/stress_kda_sm100_csrc_boundary_determinism.py new file mode 100644 index 00000000..f3883a75 --- /dev/null +++ b/benchmarks/stress_kda_sm100_csrc_boundary_determinism.py @@ -0,0 +1,174 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Stress the bitwise-aligned SM100 CuTeDSL KDA forward path. + +The script first requires complete tensor equality with the csrc boundary for +Aqk, Akk, KG, W, and U. It then captures the CuTeDSL intra, Akk inverse, +recompute-WU, and an exact comparison against those csrc outputs in one CUDA +graph. Every replay therefore validates every output element. +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import sys +import time + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +import torch +import torch.nn.functional as F + +from cula.kda.chunk_intra import chunk_kda_fwd_intra as csrc_chunk_kda_fwd_intra +from cula.ops.kda.sm100.intra_fused import BT, K_DIM, chunk_kda_fwd_intra_sm100_from_gk +from cula.ops.kda.sm100.recompute_wu import recompute_w_u_fwd + +OUTPUT_NAMES = ("W", "U", "KG", "Aqk", "Akk") + + +def _bitwise_stats( + outputs: tuple[torch.Tensor, ...], references: tuple[torch.Tensor, ...] +) -> dict[str, dict[str, float | int | bool]]: + stats = {} + for name, output, reference in zip(OUTPUT_NAMES, outputs, references, strict=True): + mismatch_count = torch.count_nonzero(output != reference).item() + max_abs = (output.float() - reference.float()).abs().max().item() + stats[name] = { + "equal": mismatch_count == 0, + "mismatches": mismatch_count, + "max_abs": max_abs, + } + return stats + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--iterations", type=int, default=10_000_000) + parser.add_argument("--checkpoint", type=int, default=1_000_000) + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--seqlen", type=int, default=256) + parser.add_argument("--heads", type=int, default=4) + parser.add_argument("--beta-dtype", choices=("bfloat16", "float32"), default="bfloat16") + parser.add_argument("--seed", type=int, default=20260825) + parser.add_argument("--report-json") + args = parser.parse_args() + + if args.iterations <= 0 or args.checkpoint <= 0: + raise ValueError("--iterations and --checkpoint must be positive") + if args.seqlen <= 0 or args.seqlen % (4 * BT) != 0: + raise ValueError(f"--seqlen must be a positive multiple of {4 * BT}") + if args.batch <= 0 or args.heads <= 0: + raise ValueError("--batch and --heads must be positive") + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): + raise RuntimeError("an SM100 CUDA device is required") + + torch.manual_seed(args.seed) + device = torch.device("cuda") + shape = (args.batch, args.seqlen, args.heads, K_DIM) + q = F.normalize(torch.randn(*shape, device=device).float(), dim=-1).bfloat16() + k = F.normalize(torch.randn(*shape, device=device).float(), dim=-1).bfloat16() + gk = torch.randn(*shape, device=device, dtype=torch.float32) * 0.02 + beta_dtype = torch.bfloat16 if args.beta_dtype == "bfloat16" else torch.float32 + beta = torch.randn(*shape[:-1], device=device).sigmoid().to(beta_dtype) + scale = K_DIM**-0.5 + + def run_cutedsl() -> tuple[torch.Tensor, ...]: + aqk, akk = chunk_kda_fwd_intra_sm100_from_gk( + q=q, + k=k, + gk=gk, + beta=beta, + scale=scale, + fp32_akk_inv=True, + ) + w, u, _, kg = recompute_w_u_fwd(k, k, beta, akk, gk) + return w, u, kg, aqk, akk + + print( + f"device={torch.cuda.get_device_name(0)} shape={shape} beta_dtype={args.beta_dtype} iterations={args.iterations}", + flush=True, + ) + + w_ref, u_ref, _, kg_ref, aqk_ref, akk_ref = csrc_chunk_kda_fwd_intra( + q=q, + k=k, + v=k, + gk=gk, + beta=beta, + scale=scale, + chunk_size=BT, + safe_gate=True, + ) + references = (w_ref, u_ref, kg_ref, aqk_ref, akk_ref) + + outputs = run_cutedsl() + torch.cuda.synchronize() + bitwise = _bitwise_stats(outputs, references) + print("BITWISE_JSON=" + json.dumps(bitwise, sort_keys=True), flush=True) + if not all(values["equal"] for values in bitwise.values()): + raise AssertionError("CuTeDSL output is not bitwise equal to csrc") + if not all(torch.isfinite(output).all().item() for output in outputs): + raise AssertionError("CuTeDSL output contains NaN or Inf") + + # Warm compilation, allocator, and concatenation paths before capture. + for _ in range(3): + warm = run_cutedsl() + torch.cat([tensor.reshape(-1) for tensor in warm]) + torch.cuda.synchronize() + + reference_flat = torch.cat([tensor.reshape(-1) for tensor in references]) + mismatch_count = torch.zeros((), dtype=torch.int64, device=device) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run_cutedsl() + captured_flat = torch.cat([tensor.reshape(-1) for tensor in captured]) + mismatch_count.add_(torch.count_nonzero(captured_flat != reference_flat)) + + torch.cuda.synchronize() + if mismatch_count.item() != 0: + raise AssertionError(f"graph capture differed from csrc: mismatches={mismatch_count.item()}") + + started = time.perf_counter() + completed = 0 + while completed < args.iterations: + stop = min(completed + args.checkpoint, args.iterations) + for _ in range(completed, stop): + graph.replay() + torch.cuda.synchronize() + completed = stop + mismatches = mismatch_count.item() + elapsed = time.perf_counter() - started + print( + f"progress={completed}/{args.iterations} mismatches={mismatches} " + f"elapsed_s={elapsed:.3f} iterations_per_s={completed / elapsed:.1f}", + flush=True, + ) + if mismatches != 0: + raise AssertionError(f"non-deterministic output after {completed} iterations: mismatches={mismatches}") + + elapsed = time.perf_counter() - started + report = { + "status": "passed", + "device": torch.cuda.get_device_name(0), + "device_index_visible": torch.cuda.current_device(), + "shape": shape, + "beta_dtype": args.beta_dtype, + "iterations": args.iterations, + "mismatches": mismatch_count.item(), + "elapsed_seconds": elapsed, + "iterations_per_second": args.iterations / elapsed, + "bitwise_csrc": bitwise, + } + print("RESULT_JSON=" + json.dumps(report, sort_keys=True), flush=True) + if args.report_json: + with open(args.report_json, "w", encoding="utf-8") as handle: + json.dump(report, handle, indent=2, sort_keys=True) + handle.write("\n") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/stress_kda_sm100_varlen_determinism.py b/benchmarks/stress_kda_sm100_varlen_determinism.py new file mode 100644 index 00000000..9c9f0057 --- /dev/null +++ b/benchmarks/stress_kda_sm100_varlen_determinism.py @@ -0,0 +1,290 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Stress SM100 CuTeDSL KDA varlen forward determinism and csrc accuracy. + +The stress phase captures the complete CuTeDSL forward chain and an exact +comparison against a golden output in one CUDA graph. Every replay therefore +checks every output element; a device-side mismatch counter is inspected at +each checkpoint and after the requested number of iterations. +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import sys +import time + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +import torch +import torch.nn.functional as F +from fla.ops.kda.gate import kda_gate_chunk_cumsum +from fla.ops.utils import prepare_chunk_indices +from fla.ops.utils.constant import RCP_LN2 + +from benchmarks.utils import exclusive_cumsum, set_seed +from cula.kda.chunk_intra import chunk_kda_fwd_intra as csrc_chunk_kda_fwd_intra +from cula.ops.kda.sm100.intra_fused import BT, K_DIM, chunk_kda_fwd_intra_sm100_varlen +from cula.ops.kda.sm100.recompute_wu import recompute_w_u_from_preprocessed + +OUTPUT_NAMES = ("k_scaled", "kg", "q_scaled", "gk_last_exp", "Aqk", "Akk", "w", "u") +ACCURACY_LIMITS = { + "k_scaled": (1e-4, 2e-3), + "kg": (5e-4, 2e-3), + "q_scaled": (1e-4, 2e-3), + "gk_last_exp": (1e-3, 2e-2), + "Aqk": (2e-3, 2e-3), + "Akk": (1e-4, 2e-3), + "w": (1e-2, 4e-3), + "u": (1e-3, 2e-3), +} + + +def _parse_seq_lens(text: str) -> list[int]: + seq_lens = [int(value) for value in text.replace(",", " ").split()] + if len(seq_lens) < 2 or any(length <= 0 for length in seq_lens): + raise ValueError("--seq-lens must contain at least two positive lengths") + return seq_lens + + +def _make_lower_mask(seq_lens: list[int], heads: int, device: torch.device) -> torch.Tensor: + total_t = sum(seq_lens) + mask = torch.zeros((1, total_t, heads, BT), dtype=torch.bool, device=device) + cols = torch.arange(BT, device=device) + offset = 0 + for seq_len in seq_lens: + rows = torch.arange(seq_len, device=device) % BT + seq_mask = (cols[None, :] <= rows[:, None]).view(1, seq_len, 1, BT) + mask[:, offset : offset + seq_len] = seq_mask + offset += seq_len + return mask + + +def _gk_last_exp(gk: torch.Tensor, seq_lens: list[int]) -> torch.Tensor: + rows = [] + offset = 0 + for seq_len in seq_lens: + for chunk_start in range(0, seq_len, BT): + last_row = offset + min(chunk_start + BT, seq_len) - 1 + rows.append(gk[:, last_row]) + offset += seq_len + return torch.stack(rows, dim=1).contiguous().exp2() + + +def _accuracy_stats(ref: torch.Tensor, out: torch.Tensor) -> dict[str, float]: + ref_f = ref.float() + out_f = out.float() + diff = (out_f - ref_f).abs() + rmse = diff.square().mean().sqrt() + ref_rms = ref_f.square().mean().sqrt() + return { + "rel_rmse": (rmse / (ref_rms + 1e-8)).item(), + "max_abs": diff.max().item(), + "mean_abs": diff.mean().item(), + } + + +def _check_accuracy( + outputs: tuple[torch.Tensor, ...], + references: tuple[torch.Tensor, ...], + lower_mask: torch.Tensor, +) -> dict[str, dict[str, float]]: + stats = {} + for name, out, ref in zip(OUTPUT_NAMES, outputs, references): + if name in ("Aqk", "Akk"): + out = out[lower_mask] + ref = ref[lower_mask] + values = _accuracy_stats(ref, out) + rel_limit, abs_limit = ACCURACY_LIMITS[name] + if values["rel_rmse"] > rel_limit or values["max_abs"] > abs_limit: + raise AssertionError( + f"{name} accuracy failed: rel_rmse={values['rel_rmse']:.6e} " + f"(limit {rel_limit:.1e}), max_abs={values['max_abs']:.6e} " + f"(limit {abs_limit:.1e})" + ) + stats[name] = values + + for name, output in (("Aqk_upper", outputs[4]), ("Akk_upper", outputs[5])): + upper_max = output[~lower_mask].abs().max().item() + if upper_max != 0.0: + raise AssertionError(f"{name} is not exactly zero: max_abs={upper_max}") + stats[name] = {"rel_rmse": 0.0, "max_abs": upper_max, "mean_abs": 0.0} + return stats + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--iterations", type=int, default=10_000_000) + parser.add_argument("--checkpoint", type=int, default=1_000_000) + parser.add_argument("--seq-lens", default="65,127,193,255") + parser.add_argument("--heads", type=int, default=8) + parser.add_argument("--seed", type=int, default=20260824) + parser.add_argument("--report-json") + args = parser.parse_args() + + if args.iterations <= 0 or args.checkpoint <= 0: + raise ValueError("--iterations and --checkpoint must be positive") + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): + raise RuntimeError("an SM100 CUDA device is required") + + seq_lens = _parse_seq_lens(args.seq_lens) + total_t = sum(seq_lens) + device = torch.device("cuda") + set_seed(args.seed) + q = F.normalize(torch.randn(1, total_t, args.heads, K_DIM, device=device).float(), dim=-1).bfloat16() + k = F.normalize(torch.randn(1, total_t, args.heads, K_DIM, device=device).float(), dim=-1).bfloat16() + g = torch.randn(1, total_t, args.heads, K_DIM, device=device, dtype=torch.bfloat16) + beta = torch.randn(1, total_t, args.heads, device=device).sigmoid().bfloat16() + a_log = torch.randn(args.heads, device=device) + dt_bias = torch.randn(args.heads * K_DIM, device=device) + cu_seqlens = torch.tensor(exclusive_cumsum(seq_lens), dtype=torch.int32, device=device) + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + scale = K_DIM**-0.5 + lower_bound = -5.0 + + def run_cutedsl() -> tuple[torch.Tensor, ...]: + intra = chunk_kda_fwd_intra_sm100_varlen( + q=q, + k=k, + g=g, + beta=beta, + A_log=a_log, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + scale=scale, + dt_bias=dt_bias, + safe_gate=True, + lower_bound=lower_bound, + seq_lens=seq_lens, + fp32_akk_inv=True, + ) + w, u = recompute_w_u_from_preprocessed( + intra[0], + k, + beta, + intra[5], + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + ) + return (*intra, w, u) + + print( + f"device={torch.cuda.get_device_name(0)} seq_lens={seq_lens} total_t={total_t} " + f"heads={args.heads} iterations={args.iterations}", + flush=True, + ) + + gk = kda_gate_chunk_cumsum( + g=g, + A_log=a_log, + dt_bias=dt_bias, + scale=RCP_LN2, + chunk_size=BT, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + lower_bound=lower_bound, + ) + w_ref, u_ref, _, kg_ref, aqk_ref, akk_ref = csrc_chunk_kda_fwd_intra( + q=q, + k=k, + v=k, + gk=gk, + beta=beta, + scale=scale, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_size=BT, + safe_gate=True, + ) + references = ( + (k.float() * gk.exp2()).bfloat16(), + kg_ref, + (q.float() * gk.exp2()).bfloat16(), + _gk_last_exp(gk, seq_lens), + aqk_ref, + akk_ref, + w_ref, + u_ref, + ) + + outputs = run_cutedsl() + torch.cuda.synchronize() + lower_mask = _make_lower_mask(seq_lens, args.heads, device) + accuracy = _check_accuracy(outputs, references, lower_mask) + for name, values in accuracy.items(): + print( + f"accuracy {name:<12} rel_rmse={values['rel_rmse']:.6e} " + f"max_abs={values['max_abs']:.6e} mean_abs={values['mean_abs']:.6e}", + flush=True, + ) + + golden = tuple(output.clone() for output in outputs) + if not all(torch.isfinite(output).all().item() for output in golden): + raise AssertionError("golden CuTeDSL output contains NaN or Inf") + + # Warm all allocator and concatenation paths before graph capture. + for _ in range(3): + warm = run_cutedsl() + torch.cat([tensor.reshape(-1) for index, tensor in enumerate(warm) if index != 3]) + torch.cuda.synchronize() + + golden_bf16 = torch.cat([tensor.reshape(-1) for index, tensor in enumerate(golden) if index != 3]) + golden_fp32 = golden[3].reshape(-1) + mismatch_count = torch.zeros((), dtype=torch.int64, device=device) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run_cutedsl() + captured_bf16 = torch.cat([tensor.reshape(-1) for index, tensor in enumerate(captured) if index != 3]) + mismatch_count.add_(torch.count_nonzero(captured_bf16 != golden_bf16)) + mismatch_count.add_(torch.count_nonzero(captured[3].reshape(-1) != golden_fp32)) + + torch.cuda.synchronize() + if mismatch_count.item() != 0: + raise AssertionError(f"graph capture already differed from golden: mismatches={mismatch_count.item()}") + + started = time.perf_counter() + completed = 0 + while completed < args.iterations: + stop = min(completed + args.checkpoint, args.iterations) + for _ in range(completed, stop): + graph.replay() + torch.cuda.synchronize() + completed = stop + mismatches = mismatch_count.item() + elapsed = time.perf_counter() - started + rate = completed / elapsed + print( + f"progress={completed}/{args.iterations} mismatches={mismatches} " + f"elapsed_s={elapsed:.3f} iterations_per_s={rate:.1f}", + flush=True, + ) + if mismatches != 0: + raise AssertionError(f"non-deterministic output after {completed} iterations: mismatches={mismatches}") + + elapsed = time.perf_counter() - started + report = { + "status": "passed", + "device": torch.cuda.get_device_name(0), + "device_index_visible": torch.cuda.current_device(), + "seq_lens": seq_lens, + "total_t": total_t, + "heads": args.heads, + "iterations": args.iterations, + "mismatches": mismatch_count.item(), + "elapsed_seconds": elapsed, + "iterations_per_second": args.iterations / elapsed, + "accuracy": accuracy, + } + print("RESULT_JSON=" + json.dumps(report, sort_keys=True), flush=True) + if args.report_json: + with open(args.report_json, "w", encoding="utf-8") as handle: + json.dump(report, handle, indent=2, sort_keys=True) + handle.write("\n") + + +if __name__ == "__main__": + main() diff --git a/cula/ops/kda/sm100/akk_inv_fp32.py b/cula/ops/kda/sm100/akk_inv_fp32.py new file mode 100644 index 00000000..d6848022 --- /dev/null +++ b/cula/ops/kda/sm100/akk_inv_fp32.py @@ -0,0 +1,268 @@ +# 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 + +"""Standalone fp32-workspace Akk inverse for K123. + +Input: + A_phys [B, T, H, 64] fp32, K123 physical block-transposed pre-inverse layout. +Output: + A_out [B, T, H, 64] bf16, logical lower_tri((I + L)^-1), no beta[col] epilogue. + +This is a baseline kernel for K123 -> fp32 GMEM workspace -> standalone inverse. +It converts K123's physical layout to logical lower-triangular fp32 in SMEM, +then reuses the same TF32 Schur helper structure as kda_akk_inv_tf32.py. +""" + +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils + +from cula.ops.kda.sm100.akk_inv_tf32 import ( + _add_store_C16_smem, + _invert_diag_forward16, + _matmul16_smem_smem, + _matmul16_tmp_smem, + _matmul32_smem_smem, + _store_C16_smem, + _store_C16_tmp, +) + +BS = 64 +SB = 16 +THREADS = 128 +AKK_STRIDE = BS + 4 +TMP_STRIDE = SB + 4 +TMP_SLOTS = 4 +LOWER_TILE_ROWS = (0, 1, 1, 2, 2, 2, 3, 3, 3, 3) +LOWER_TILE_COLS = (0, 0, 1, 0, 1, 2, 0, 1, 2, 3) +UPPER_TILE_ROWS = (0, 0, 0, 1, 1, 2) +UPPER_TILE_COLS = (1, 2, 3, 2, 3, 3) + + +@cute.kernel +def akk_inv_fp32_physical_kernel( + mA_phys: cute.Tensor, + mA_out: cute.Tensor, + smat_layout: cute.Layout, + stmp_layout: cute.Layout, + NT: int, + H: int, + mCuSeqlens: cute.Tensor, + mChunkIndices: cute.Tensor, + IS_VARLEN: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = tidx % 32 + h_idx, nt_idx, b_idx = cute.arch.block_idx() + + smem = utils.SmemAllocator() + sAkk = smem.allocate_tensor(cutlass.Float32, smat_layout, 128) + sTmp = smem.allocate_tensor(cutlass.Float32, stmp_layout, 128) + + chunk_start = nt_idx * BS + eos = cutlass.Int32(chunk_start + BS) + if IS_VARLEN: + seq_id = cutlass.Int32(mChunkIndices[nt_idx, 0]) + local = cutlass.Int32(mChunkIndices[nt_idx, 1]) + bos = cutlass.Int32(mCuSeqlens[seq_id]) + eos = cutlass.Int32(mCuSeqlens[seq_id + 1]) + chunk_start = bos + local * BS + + # K123 stores the ten lower 16x16 tiles in block-transposed physical + # positions, while values inside each tile remain row-contiguous. Two + # groups of 64 threads load two tiles per iteration as float4 vectors. + rNorm = cute.make_rmem_tensor((4,), cutlass.Float32) + rNorm.fill(cutlass.Float32(0.0)) + tile_slot = tidx // 64 + tile_thread = tidx % 64 + tile_row = tile_thread // 4 + tile_vec = tile_thread % 4 + # The ten lower tiles are overwritten by the physical workspace below. + # Clear only the six upper tiles that the block inverse may read. + for tile_pair in cutlass.range_constexpr(3): + zero_row_blk = cutlass.select_( + tile_slot == 0, + cutlass.Int32(UPPER_TILE_ROWS[tile_pair * 2]), + cutlass.Int32(UPPER_TILE_ROWS[tile_pair * 2 + 1]), + ) + zero_col_blk = cutlass.select_( + tile_slot == 0, + cutlass.Int32(UPPER_TILE_COLS[tile_pair * 2]), + cutlass.Int32(UPPER_TILE_COLS[tile_pair * 2 + 1]), + ) + zero_row = zero_row_blk * SB + tile_row + zero_vec = zero_col_blk * (SB // 4) + tile_vec + sZeroRow = sAkk[zero_row, None] + sZeroVec = cute.local_tile(sZeroRow, (4,), (zero_vec,)) + cute.autovec_copy(rNorm, sZeroVec) + cute.arch.barrier() + + for tile_pair in cutlass.range_constexpr(5): + row_blk = cutlass.select_( + tile_slot == 0, + cutlass.Int32(LOWER_TILE_ROWS[tile_pair * 2]), + cutlass.Int32(LOWER_TILE_ROWS[tile_pair * 2 + 1]), + ) + col_blk = cutlass.select_( + tile_slot == 0, + cutlass.Int32(LOWER_TILE_COLS[tile_pair * 2]), + cutlass.Int32(LOWER_TILE_COLS[tile_pair * 2 + 1]), + ) + row = row_blk * SB + tile_row + col = col_blk * SB + tile_vec * 4 + src_row = col_blk * SB + tile_row + src_col = row_blk * SB + tile_vec * 4 + t_row = chunk_start + row + + if IS_VARLEN: + if t_row < eos: + gNormRow = mA_phys[b_idx, chunk_start + src_row, h_idx, None] + gNormVec = cute.local_tile(gNormRow, (4,), (src_col // 4,)) + cute.autovec_copy(gNormVec, rNorm) + else: + rNorm.fill(cutlass.Float32(0.0)) + else: + gNormRow = mA_phys[b_idx, chunk_start + src_row, h_idx, None] + gNormVec = cute.local_tile(gNormRow, (4,), (src_col // 4,)) + cute.autovec_copy(gNormVec, rNorm) + + if row_blk == col_blk: + for elem in cutlass.range_constexpr(4): + logical_col = col + elem + if row == logical_col: + rNorm[elem] = cutlass.Float32(1.0) + if row < logical_col: + rNorm[elem] = cutlass.Float32(0.0) + + sNormRow = sAkk[row, None] + sNormVec = cute.local_tile(sNormRow, (4,), (col // 4,)) + cute.autovec_copy(rNorm, sNormVec) + cute.arch.barrier() + + if tidx < 64: + _invert_diag_forward16(sAkk, tidx // 16, tidx) + cute.arch.barrier() + + if tidx < 64: + block32 = tidx // 32 + row_base = block32 * 32 + c0, c1, c2, c3, c4, c5, c6, c7 = _matmul16_smem_smem( + sAkk, row_base + 16, row_base + 16, sAkk, row_base + 16, row_base, lane_id + ) + _store_C16_tmp(sTmp, block32, -c0, -c1, -c2, -c3, -c4, -c5, -c6, -c7, lane_id) + cute.arch.barrier() + + if tidx < 64: + block32 = tidx // 32 + row_base = block32 * 32 + c0, c1, c2, c3, c4, c5, c6, c7 = _matmul16_tmp_smem(sTmp, block32, sAkk, row_base, row_base, lane_id) + _store_C16_smem(sAkk, row_base + 16, row_base, c0, c1, c2, c3, c4, c5, c6, c7, lane_id) + cute.arch.barrier() + + x = warp_idx // 2 + y = warp_idx % 2 + slot = warp_idx + row_o = 32 + y * 16 + col_c = x * 16 + + # Match CollectiveInverseTF32::blockwise_diagonal_inversed_32x32_to_64x64: + # accumulate the complete K=32 product in one FP32 accumulator. Splitting + # this into two K=16 accumulators followed by p + q changes the FP32 + # reduction tree and therefore cannot be bitwise equal to csrc. + p0, p1, p2, p3, p4, p5, p6, p7 = _matmul32_smem_smem(sAkk, row_o, 32, sAkk, 32, col_c, lane_id) + _store_C16_tmp( + sTmp, + slot, + -p0, + -p1, + -p2, + -p3, + -p4, + -p5, + -p6, + -p7, + lane_id, + ) + cute.arch.barrier() + + o0, o1, o2, o3, o4, o5, o6, o7 = _matmul16_tmp_smem(sTmp, slot, sAkk, x * 16, 0, lane_id) + r0, r1, r2, r3, r4, r5, r6, r7 = _matmul16_tmp_smem(sTmp, slot, sAkk, x * 16, 16, lane_id) + if x == 0: + _store_C16_smem(sAkk, row_o, 0, o0, o1, o2, o3, o4, o5, o6, o7, lane_id) + _store_C16_smem(sAkk, row_o, 16, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + cute.arch.barrier() + if x == 1: + _add_store_C16_smem(sAkk, row_o, 0, o0, o1, o2, o3, o4, o5, o6, o7, lane_id) + _add_store_C16_smem(sAkk, row_o, 16, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + cute.arch.barrier() + + # Vectorized inverse epilogue: 512 contiguous 8-element row segments are + # distributed evenly across the 128 threads. The input normalization + # cleared the upper triangle, so no per-element mask is needed here. + rOutFp32 = cute.make_rmem_tensor((8,), cutlass.Float32) + rOutBf16 = cute.make_rmem_tensor((8,), cutlass.BFloat16) + for store_iter in cutlass.range_constexpr((BS * BS) // (THREADS * 8)): + linear_vec = tidx + store_iter * THREADS + row = linear_vec // (BS // 8) + vec_idx = linear_vec % (BS // 8) + t_row = chunk_start + row + sOutRow = sAkk[row, None] + sOutVec = cute.local_tile(sOutRow, (8,), (vec_idx,)) + cute.autovec_copy(sOutVec, rOutFp32) + for elem in cutlass.range_constexpr(8): + if row < vec_idx * 8 + elem: + rOutFp32[elem] = cutlass.Float32(0.0) + rOutBf16.store(rOutFp32.load().to(cutlass.BFloat16)) + if IS_VARLEN: + if t_row < eos: + gOutRow = mA_out[b_idx, t_row, h_idx, None] + gOutVec = cute.local_tile(gOutRow, (8,), (vec_idx,)) + cute.autovec_copy(rOutBf16, gOutVec) + else: + gOutRow = mA_out[b_idx, t_row, h_idx, None] + gOutVec = cute.local_tile(gOutRow, (8,), (vec_idx,)) + cute.autovec_copy(rOutBf16, gOutVec) + + +@cute.jit +def akk_inv_fp32_physical_host( + A_phys: cute.Tensor, + A_out: cute.Tensor, + B: cutlass.Constexpr[int], + NT: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + mCuSeqlens: cute.Tensor, + mChunkIndices: cute.Tensor, + IS_VARLEN: cutlass.Constexpr[int], + T_VAL: cutlass.Constexpr[int], +): + in_layout = cute.make_layout((B, T_VAL, H, BS), stride=(T_VAL * H * BS, H * BS, BS, 1)) + out_layout = cute.make_layout((B, T_VAL, H, BS), stride=(T_VAL * H * BS, H * BS, BS, 1)) + gA_phys = cute.make_tensor(A_phys.iterator, in_layout) + gA_out = cute.make_tensor(A_out.iterator, out_layout) + + smat_layout = cute.make_layout((BS, BS), stride=(AKK_STRIDE, 1)) + stmp_layout = cute.make_layout((TMP_SLOTS, SB, SB), stride=(SB * TMP_STRIDE, TMP_STRIDE, 1)) + smem_bytes = BS * AKK_STRIDE * 4 + TMP_SLOTS * SB * TMP_STRIDE * 4 + 256 + + akk_inv_fp32_physical_kernel( + gA_phys, + gA_out, + smat_layout, + stmp_layout, + NT, + H, + mCuSeqlens, + mChunkIndices, + IS_VARLEN, + ).launch( + grid=(H, NT, B), + block=(THREADS, 1, 1), + smem=smem_bytes, + ) diff --git a/cula/ops/kda/sm100/akk_inv_tf32.py b/cula/ops/kda/sm100/akk_inv_tf32.py new file mode 100644 index 00000000..6c737abf --- /dev/null +++ b/cula/ops/kda/sm100/akk_inv_tf32.py @@ -0,0 +1,589 @@ +# 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 + +"""Akk 64x64 lower triangular inverse using the CUDA C TF32 Schur structure. + +This is an experimental CuTeDSL port of +`kerutils::CollectiveInverseTF32`. + +Input: + A_kk [B, T, H, 64] bf16, logical lower-triangular form, diag may be garbage. +Output: + A_kk [B, T, H, 64] bf16, lower_tri((I + L)^-1), no beta epilogue by default. + +Precision choices: + - SMEM is fp32 with the CUDA C padded stride 68. + - Diagonal 16x16 blocks use forward substitution in fp32. + - Off-diagonal Schur products use mma.sync f32.tf32.tf32.f32. + - Intermediate Schur products are stored/reloaded as fp32 scratch in this + first CuTeDSL version; TF32 truncation is still done only by MMA hardware. +""" + +import cutlass +import cutlass.cute as cute +import cutlass.utils as utils +import torch +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import T, dsl_user_op + +BS = 64 +SB = 16 +THREADS = 128 +AKK_STRIDE = BS + 4 +TMP_STRIDE = SB + 4 +TMP_SLOTS = 4 + + +@dsl_user_op +def mma_tf32_m16n8k8( + a0, + a1, + a2, + a3, + b0, + b1, + c0, + c1, + c2, + c3, + *, + loc=None, + ip=None, +): + """Register-level TF32 MMA with fp32 accumulator, shape m16n8k8.""" + a0b = llvm.bitcast(T.i32(), a0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a1b = llvm.bitcast(T.i32(), a1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a2b = llvm.bitcast(T.i32(), a2.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a3b = llvm.bitcast(T.i32(), a3.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b0b = llvm.bitcast(T.i32(), b0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b1b = llvm.bitcast(T.i32(), b1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [ + a0b, + a1b, + a2b, + a3b, + b0b, + b1b, + c0.ir_value(loc=loc, ip=ip), + c1.ir_value(loc=loc, ip=ip), + c2.ir_value(loc=loc, ip=ip), + c3.ir_value(loc=loc, ip=ip), + ], + """{ + mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 + {$0, $1, $2, $3}, + {$4, $5, $6, $7}, + {$8, $9}, + {$10, $11, $12, $13}; + }""", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +@dsl_user_op +def _invert_diag_forward16(sAkk: cute.Tensor, block_idx, tidx, *, loc=None, ip=None): + """CUDA C compute_diagonal_inverse_NxN<16>, in fp32.""" + tid_in_group = tidx % 16 + group_base = (block_idx % 2) * 16 + base = block_idx * 16 + row = cute.make_rmem_tensor(cute.make_layout((SB,), stride=(1,)), cutlass.Float32) + + for i in range(SB): + val = cutlass.Float32(sAkk[base + tid_in_group, base + i]) + is_lower = cutlass.Float32(i < tid_in_group) + is_diag = cutlass.Float32(i == tid_in_group) + row[i] = val * is_lower + is_diag + + for src_row in range(SB - 1): + row_scale = -row[src_row] + target_lane = group_base + src_row + active = cutlass.Float32(tid_in_group > src_row) + for i in range(src_row): + src_row_value = cute.arch.shuffle_sync_op( + value=row[i], + offset=target_lane, + mask=0xFFFFFFFF, + mask_and_clamp=31, + ) + row[i] = row[i] + active * row_scale * src_row_value + row[src_row] = active * row_scale + (cutlass.Float32(1.0) - active) * row[src_row] + + for i in range(SB): + sAkk[base + tid_in_group, base + i] = row[i] + + +@dsl_user_op +def _matmul16_smem_smem( + sA: cute.Tensor, + a_row, + a_col, + sB: cute.Tensor, + b_row, + b_col, + lane_id, + *, + loc=None, + ip=None, +): + """Return a 16x16 tile C = A@B using two m16n8k8 TF32 MMAs.""" + gid = lane_id // 4 + tid = lane_id % 4 + z = cutlass.Float32(0.0) + + a0 = cutlass.Float32(sA[a_row + gid, a_col + 2 * tid]) + a1 = cutlass.Float32(sA[a_row + gid + 8, a_col + 2 * tid]) + a2 = cutlass.Float32(sA[a_row + gid, a_col + 2 * tid + 1]) + a3 = cutlass.Float32(sA[a_row + gid + 8, a_col + 2 * tid + 1]) + b0n0 = cutlass.Float32(sB[b_row + 2 * tid, b_col + gid]) + b1n0 = cutlass.Float32(sB[b_row + 2 * tid + 1, b_col + gid]) + b0n1 = cutlass.Float32(sB[b_row + 2 * tid, b_col + 8 + gid]) + b1n1 = cutlass.Float32(sB[b_row + 2 * tid + 1, b_col + 8 + gid]) + c0, c1, c2, c3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n0, b1n0, z, z, z, z) + c4, c5, c6, c7 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n1, b1n1, z, z, z, z) + + a0 = cutlass.Float32(sA[a_row + gid, a_col + 8 + 2 * tid]) + a1 = cutlass.Float32(sA[a_row + gid + 8, a_col + 8 + 2 * tid]) + a2 = cutlass.Float32(sA[a_row + gid, a_col + 8 + 2 * tid + 1]) + a3 = cutlass.Float32(sA[a_row + gid + 8, a_col + 8 + 2 * tid + 1]) + b0n0 = cutlass.Float32(sB[b_row + 8 + 2 * tid, b_col + gid]) + b1n0 = cutlass.Float32(sB[b_row + 8 + 2 * tid + 1, b_col + gid]) + b0n1 = cutlass.Float32(sB[b_row + 8 + 2 * tid, b_col + 8 + gid]) + b1n1 = cutlass.Float32(sB[b_row + 8 + 2 * tid + 1, b_col + 8 + gid]) + c0, c1, c2, c3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n0, b1n0, c0, c1, c2, c3) + c4, c5, c6, c7 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n1, b1n1, c4, c5, c6, c7) + return c0, c1, c2, c3, c4, c5, c6, c7 + + +@dsl_user_op +def _matmul32_smem_smem( + sA: cute.Tensor, + a_row, + a_col, + sB: cute.Tensor, + b_row, + b_col, + lane_id, + *, + loc=None, + ip=None, +): + """Return a 16x16 tile over K=32 with csrc's four-MMA order.""" + gid = lane_id // 4 + tid = lane_id % 4 + c0 = cutlass.Float32(0.0) + c1 = cutlass.Float32(0.0) + c2 = cutlass.Float32(0.0) + c3 = cutlass.Float32(0.0) + c4 = cutlass.Float32(0.0) + c5 = cutlass.Float32(0.0) + c6 = cutlass.Float32(0.0) + c7 = cutlass.Float32(0.0) + + for k_base in (0, 8, 16, 24): + a0 = cutlass.Float32(sA[a_row + gid, a_col + k_base + 2 * tid]) + a1 = cutlass.Float32(sA[a_row + gid + 8, a_col + k_base + 2 * tid]) + a2 = cutlass.Float32(sA[a_row + gid, a_col + k_base + 2 * tid + 1]) + a3 = cutlass.Float32(sA[a_row + gid + 8, a_col + k_base + 2 * tid + 1]) + b0n0 = cutlass.Float32(sB[b_row + k_base + 2 * tid, b_col + gid]) + b1n0 = cutlass.Float32(sB[b_row + k_base + 2 * tid + 1, b_col + gid]) + b0n1 = cutlass.Float32(sB[b_row + k_base + 2 * tid, b_col + 8 + gid]) + b1n1 = cutlass.Float32(sB[b_row + k_base + 2 * tid + 1, b_col + 8 + gid]) + c0, c1, c2, c3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n0, b1n0, c0, c1, c2, c3) + c4, c5, c6, c7 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n1, b1n1, c4, c5, c6, c7) + return c0, c1, c2, c3, c4, c5, c6, c7 + + +@dsl_user_op +def _matmul16_tmp_smem( + sTmp: cute.Tensor, + slot, + sB: cute.Tensor, + b_row, + b_col, + lane_id, + *, + loc=None, + ip=None, +): + gid = lane_id // 4 + tid = lane_id % 4 + z = cutlass.Float32(0.0) + + a0 = cutlass.Float32(sTmp[slot, gid, 2 * tid]) + a1 = cutlass.Float32(sTmp[slot, gid + 8, 2 * tid]) + a2 = cutlass.Float32(sTmp[slot, gid, 2 * tid + 1]) + a3 = cutlass.Float32(sTmp[slot, gid + 8, 2 * tid + 1]) + b0n0 = cutlass.Float32(sB[b_row + 2 * tid, b_col + gid]) + b1n0 = cutlass.Float32(sB[b_row + 2 * tid + 1, b_col + gid]) + b0n1 = cutlass.Float32(sB[b_row + 2 * tid, b_col + 8 + gid]) + b1n1 = cutlass.Float32(sB[b_row + 2 * tid + 1, b_col + 8 + gid]) + c0, c1, c2, c3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n0, b1n0, z, z, z, z) + c4, c5, c6, c7 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n1, b1n1, z, z, z, z) + + a0 = cutlass.Float32(sTmp[slot, gid, 8 + 2 * tid]) + a1 = cutlass.Float32(sTmp[slot, gid + 8, 8 + 2 * tid]) + a2 = cutlass.Float32(sTmp[slot, gid, 8 + 2 * tid + 1]) + a3 = cutlass.Float32(sTmp[slot, gid + 8, 8 + 2 * tid + 1]) + b0n0 = cutlass.Float32(sB[b_row + 8 + 2 * tid, b_col + gid]) + b1n0 = cutlass.Float32(sB[b_row + 8 + 2 * tid + 1, b_col + gid]) + b0n1 = cutlass.Float32(sB[b_row + 8 + 2 * tid, b_col + 8 + gid]) + b1n1 = cutlass.Float32(sB[b_row + 8 + 2 * tid + 1, b_col + 8 + gid]) + c0, c1, c2, c3 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n0, b1n0, c0, c1, c2, c3) + c4, c5, c6, c7 = mma_tf32_m16n8k8(a0, a1, a2, a3, b0n1, b1n1, c4, c5, c6, c7) + return c0, c1, c2, c3, c4, c5, c6, c7 + + +@dsl_user_op +def _store_C16_tmp( + sTmp: cute.Tensor, + slot, + c0, + c1, + c2, + c3, + c4, + c5, + c6, + c7, + lane_id, + *, + loc=None, + ip=None, +): + gid = lane_id // 4 + tid = lane_id % 4 + sTmp[slot, gid, 2 * tid] = c0 + sTmp[slot, gid, 2 * tid + 1] = c1 + sTmp[slot, gid + 8, 2 * tid] = c2 + sTmp[slot, gid + 8, 2 * tid + 1] = c3 + sTmp[slot, gid, 8 + 2 * tid] = c4 + sTmp[slot, gid, 8 + 2 * tid + 1] = c5 + sTmp[slot, gid + 8, 8 + 2 * tid] = c6 + sTmp[slot, gid + 8, 8 + 2 * tid + 1] = c7 + + +@dsl_user_op +def _store_C16_smem( + sDst: cute.Tensor, + row, + col, + c0, + c1, + c2, + c3, + c4, + c5, + c6, + c7, + lane_id, + *, + loc=None, + ip=None, +): + gid = lane_id // 4 + tid = lane_id % 4 + sDst[row + gid, col + 2 * tid] = c0 + sDst[row + gid, col + 2 * tid + 1] = c1 + sDst[row + gid + 8, col + 2 * tid] = c2 + sDst[row + gid + 8, col + 2 * tid + 1] = c3 + sDst[row + gid, col + 8 + 2 * tid] = c4 + sDst[row + gid, col + 8 + 2 * tid + 1] = c5 + sDst[row + gid + 8, col + 8 + 2 * tid] = c6 + sDst[row + gid + 8, col + 8 + 2 * tid + 1] = c7 + + +@dsl_user_op +def _add_store_C16_smem( + sDst: cute.Tensor, + row, + col, + c0, + c1, + c2, + c3, + c4, + c5, + c6, + c7, + lane_id, + *, + loc=None, + ip=None, +): + gid = lane_id // 4 + tid = lane_id % 4 + sDst[row + gid, col + 2 * tid] = cutlass.Float32(sDst[row + gid, col + 2 * tid]) + c0 + sDst[row + gid, col + 2 * tid + 1] = cutlass.Float32(sDst[row + gid, col + 2 * tid + 1]) + c1 + sDst[row + gid + 8, col + 2 * tid] = cutlass.Float32(sDst[row + gid + 8, col + 2 * tid]) + c2 + sDst[row + gid + 8, col + 2 * tid + 1] = cutlass.Float32(sDst[row + gid + 8, col + 2 * tid + 1]) + c3 + sDst[row + gid, col + 8 + 2 * tid] = cutlass.Float32(sDst[row + gid, col + 8 + 2 * tid]) + c4 + sDst[row + gid, col + 8 + 2 * tid + 1] = cutlass.Float32(sDst[row + gid, col + 8 + 2 * tid + 1]) + c5 + sDst[row + gid + 8, col + 8 + 2 * tid] = cutlass.Float32(sDst[row + gid + 8, col + 8 + 2 * tid]) + c6 + sDst[row + gid + 8, col + 8 + 2 * tid + 1] = cutlass.Float32(sDst[row + gid + 8, col + 8 + 2 * tid + 1]) + c7 + + +@cute.kernel +def akk_inv_tf32_kernel( + mA_in: cute.Tensor, + mA_out: cute.Tensor, + mBeta: cute.Tensor, + smat_layout: cute.Layout, + stmp_layout: cute.Layout, + NT: int, + H: int, + mCuSeqlens: cute.Tensor, + mChunkIndices: cute.Tensor, + IS_VARLEN: cutlass.Constexpr[int], + APPLY_BETA_EPILOGUE: cutlass.Constexpr[int], +): + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + lane_id = tidx % 32 + h_idx, nt_idx, b_idx = cute.arch.block_idx() + + smem = utils.SmemAllocator() + sAkk = smem.allocate_tensor(cutlass.Float32, smat_layout, 128) + sTmp = smem.allocate_tensor(cutlass.Float32, stmp_layout, 128) + sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout((BS,), stride=(1,)), 128) + + chunk_start = nt_idx * BS + eos = cutlass.Int32(chunk_start + BS) + if IS_VARLEN: + seq_id = cutlass.Int32(mChunkIndices[nt_idx, 0]) + local = cutlass.Int32(mChunkIndices[nt_idx, 1]) + bos = cutlass.Int32(mCuSeqlens[seq_id]) + eos = cutlass.Int32(mCuSeqlens[seq_id + 1]) + chunk_start = bos + local * BS + + for i in range((BS * BS) // THREADS): + linear = tidx + i * THREADS + row = linear // BS + col = linear % BS + t_row = chunk_start + row + value = cutlass.Float32(0.0) + if IS_VARLEN: + if t_row < eos: + value = cutlass.Float32(mA_in[b_idx, t_row, h_idx, col]) + else: + value = cutlass.Float32(mA_in[b_idx, t_row, h_idx, col]) + if row == col: + value = cutlass.Float32(1.0) + if row < col: + value = cutlass.Float32(0.0) + sAkk[row, col] = value + + if tidx < BS: + beta_t = chunk_start + tidx + beta_val = cutlass.Float32(0.0) + if IS_VARLEN: + if beta_t < eos: + beta_val = cutlass.Float32(mBeta[b_idx, beta_t, h_idx]) + else: + beta_val = cutlass.Float32(mBeta[b_idx, beta_t, h_idx]) + sBeta[tidx] = beta_val + cute.arch.barrier() + + if tidx < 64: + _invert_diag_forward16(sAkk, tidx // 16, tidx) + cute.arch.barrier() + + # 16x16 -> 32x32 for the two diagonal 32x32 blocks. + if tidx < 64: + block32 = tidx // 32 + row_base = block32 * 32 + c0, c1, c2, c3, c4, c5, c6, c7 = _matmul16_smem_smem( + sAkk, row_base + 16, row_base + 16, sAkk, row_base + 16, row_base, lane_id + ) + _store_C16_tmp(sTmp, block32, -c0, -c1, -c2, -c3, -c4, -c5, -c6, -c7, lane_id) + cute.arch.barrier() + + if tidx < 64: + block32 = tidx // 32 + row_base = block32 * 32 + c0, c1, c2, c3, c4, c5, c6, c7 = _matmul16_tmp_smem(sTmp, block32, sAkk, row_base, row_base, lane_id) + _store_C16_smem(sAkk, row_base + 16, row_base, c0, c1, c2, c3, c4, c5, c6, c7, lane_id) + cute.arch.barrier() + + # 32x32 -> 64x64. Four warps compute partials over x and reduce by y. + x = warp_idx // 2 + y = warp_idx % 2 + slot = warp_idx + row_o = 32 + y * 16 + col_c = x * 16 + + p0, p1, p2, p3, p4, p5, p6, p7 = _matmul32_smem_smem(sAkk, row_o, 32, sAkk, 32, col_c, lane_id) + _store_C16_tmp( + sTmp, + slot, + -p0, + -p1, + -p2, + -p3, + -p4, + -p5, + -p6, + -p7, + lane_id, + ) + cute.arch.barrier() + + o0, o1, o2, o3, o4, o5, o6, o7 = _matmul16_tmp_smem(sTmp, slot, sAkk, x * 16, 0, lane_id) + r0, r1, r2, r3, r4, r5, r6, r7 = _matmul16_tmp_smem(sTmp, slot, sAkk, x * 16, 16, lane_id) + if x == 0: + _store_C16_smem(sAkk, row_o, 0, o0, o1, o2, o3, o4, o5, o6, o7, lane_id) + _store_C16_smem(sAkk, row_o, 16, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + cute.arch.barrier() + if x == 1: + _add_store_C16_smem(sAkk, row_o, 0, o0, o1, o2, o3, o4, o5, o6, o7, lane_id) + _add_store_C16_smem(sAkk, row_o, 16, r0, r1, r2, r3, r4, r5, r6, r7, lane_id) + cute.arch.barrier() + + row_start = warp_idx * SB + for ri in range(SB): + row = row_start + ri + col0 = lane_id * 2 + col1 = col0 + 1 + t_row = chunk_start + row + val0 = cutlass.Float32(sAkk[row, col0]) + val1 = cutlass.Float32(sAkk[row, col1]) + if row < col0: + val0 = cutlass.Float32(0.0) + if row < col1: + val1 = cutlass.Float32(0.0) + if APPLY_BETA_EPILOGUE: + val0 = val0 * cutlass.Float32(sBeta[col0]) + val1 = val1 * cutlass.Float32(sBeta[col1]) + if IS_VARLEN: + if t_row < eos: + mA_out[b_idx, t_row, h_idx, col0] = val0.to(cutlass.BFloat16) + mA_out[b_idx, t_row, h_idx, col1] = val1.to(cutlass.BFloat16) + else: + mA_out[b_idx, t_row, h_idx, col0] = val0.to(cutlass.BFloat16) + mA_out[b_idx, t_row, h_idx, col1] = val1.to(cutlass.BFloat16) + + +@cute.jit +def akk_inv_tf32_host( + A_in: cute.Tensor, + A_out: cute.Tensor, + Beta_in: cute.Tensor, + B: cutlass.Constexpr[int], + NT: cutlass.Constexpr[int], + H: cutlass.Constexpr[int], + mCuSeqlens: cute.Tensor, + mChunkIndices: cute.Tensor, + IS_VARLEN: cutlass.Constexpr[int], + T_VAL: cutlass.Constexpr[int], + APPLY_BETA_EPILOGUE: cutlass.Constexpr[int] = 0, +): + in_layout = cute.make_layout((B, T_VAL, H, BS), stride=(T_VAL * H * BS, H * BS, BS, 1)) + out_layout = cute.make_layout((B, T_VAL, H, BS), stride=(T_VAL * H * BS, H * BS, BS, 1)) + gA_in = cute.make_tensor(A_in.iterator, in_layout) + gA_out = cute.make_tensor(A_out.iterator, out_layout) + smat_layout = cute.make_layout((BS, BS), stride=(AKK_STRIDE, 1)) + stmp_layout = cute.make_layout((TMP_SLOTS, SB, SB), stride=(SB * TMP_STRIDE, TMP_STRIDE, 1)) + smem_bytes = BS * AKK_STRIDE * 4 + TMP_SLOTS * SB * TMP_STRIDE * 4 + BS * 4 + 256 + + akk_inv_tf32_kernel( + gA_in, + gA_out, + Beta_in, + smat_layout, + stmp_layout, + NT, + H, + mCuSeqlens, + mChunkIndices, + IS_VARLEN, + APPLY_BETA_EPILOGUE, + ).launch( + grid=(H, NT, B), + block=(THREADS, 1, 1), + smem=smem_bytes, + ) + + +_compile_cache = {} + + +def _wrap_tensor(t: torch.Tensor, element_type, *, dynamic: bool, assumed_align: int = 16): + ct = from_dlpack(t, assumed_align=assumed_align) + if dynamic: + ct = ct.mark_layout_dynamic() + ct.element_type = element_type + return ct + + +def _make_eqlen_metadata(device: torch.device): + cu = torch.empty(2, dtype=torch.int64, device=device) + ci = torch.empty(1, 2, dtype=torch.int64, device=device) + cu_ct = _wrap_tensor(cu, cutlass.Int64, dynamic=True, assumed_align=4) + ci_ct = _wrap_tensor(ci, cutlass.Int64, dynamic=True, assumed_align=4) + return cu, ci, cu_ct, ci_ct + + +def akk_inv_tf32( + a_logical_lower: torch.Tensor, + beta: torch.Tensor, + *, + apply_beta_epilogue: bool = False, +) -> torch.Tensor: + if a_logical_lower.dtype is not torch.bfloat16: + raise TypeError("a_logical_lower must be torch.bfloat16") + if beta.dtype is not torch.bfloat16: + raise TypeError("beta must be torch.bfloat16") + if a_logical_lower.ndim != 4 or a_logical_lower.shape[-1] != BS: + raise ValueError("a_logical_lower must have shape [B, T, H, 64]") + B, T_VAL, H, _ = a_logical_lower.shape + if T_VAL % BS != 0: + raise ValueError("equal-length path requires T to be a multiple of 64") + out = torch.empty_like(a_logical_lower) + cu, ci, cu_ct, ci_ct = _make_eqlen_metadata(a_logical_lower.device) + a_in = _wrap_tensor(a_logical_lower, cutlass.BFloat16, dynamic=False) + a_out = _wrap_tensor(out, cutlass.BFloat16, dynamic=False) + beta_ct = _wrap_tensor(beta, cutlass.BFloat16, dynamic=True) + key = ( + "tf32", + a_logical_lower.device.index or 0, + B, + T_VAL // BS, + H, + bool(apply_beta_epilogue), + ) + if key not in _compile_cache: + _compile_cache[key] = cute.compile( + akk_inv_tf32_host, + a_in, + a_out, + beta_ct, + B, + T_VAL // BS, + H, + cu_ct, + ci_ct, + 0, + T_VAL, + int(apply_beta_epilogue), + ) + _compile_cache[key](a_in, a_out, beta_ct, cu_ct, ci_ct) + del cu, ci + return out diff --git a/cula/ops/kda/sm100/intra_fused.py b/cula/ops/kda/sm100/intra_fused.py new file mode 100644 index 00000000..cdb7bc87 --- /dev/null +++ b/cula/ops/kda/sm100/intra_fused.py @@ -0,0 +1,3112 @@ +# ruff: noqa +""" +Copyright (c) 2025 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +""" +Persistent Fused K1+K2+K3 Kernel for KDA. + +Fuses gate activation + cumsum + scaling (K1), intra sub-chunk Aqk/Akk (K2), +and inter sub-chunk solve + merged inverse (K3) into a single persistent kernel. + +Grid: (NUM_SMS, 1, 1) — 148 persistent blocks, each loops over work units + Total work units = (NT/4) * H * B, distributed round-robin across SMs + Block i processes work units i, i+NUM_SMS, i+2*NUM_SMS, ... +Block: 992 threads (31 warps), warp-specialized by role: + Warps 0-15: K1 gate activation/cumsum/scaling (8×2, vec2) + Warps 16-25: K2/K3 MMA compute (one lower-triangular tile per warp) + Warp 26: dedicated TMA producer for Q/K/G + Warps 27-30: vectorized Aqk/Akk stores + +Pipeline (single for_generate, warp groups separated by if-blocks): + per work unit: + Warp 26: prefetch Q/K/G and reuse the two stages as MMA releases them + Warps 0-15: wait for TMA, run K1, arrive(k1_done) + Warps 16-25: wait(k1_done)+wait(store_done), MMA, arrive(mma_done+stage_reuse) + Warps 27-30: wait(mma_done), store sAqk/sAkk→GMEM, arrive(store_done) + All warp-group invariants are computed inside each group's if-block (not hoisted) + to eliminate cross-group register pressure — same budget as the _all version. + Mbarrier phases self-reset after 4 iterations (2 stages × 2 phases). + +Mbarriers: + tma_mbars[2]: count=1, warp 0 lane 0 → K1+MMA wait for TMA data + stage_reuse_mbars[2]: count=320, MMA(10 warps) → TMA waits before stage reuse + k1_done_mbars[2]: count=512, K1(16 warps) → MMA waits for g_cumsum ready + mma_done_mbars[2]: count=320, MMA(10 warps) → Store waits for sAqk/sAkk ready + store_done_mbars[2]: count=128, Store(4 warps) → MMA waits for sAqk/sAkk stage free + +SMEM: ~220KB (Q/K/G double buffers, fp32 gate cumsum, prefix scratch, + Aqk/Akk double buffers, staged beta, and pipeline barriers). + +Inputs: + g [B,T,H,K] bf16 raw gate + k [B,T,H,K] bf16 + q [B,T,H,K] bf16 + A_log [H] fp32 per-head log decay + beta [B,T,H] bf16 used for Akk unit lower triangular + scale fp32 1/sqrt(K) + +Outputs (g_cumsum stays in SMEM, not written to GMEM): + k_scaled [B,T,H,K] bf16 + q_scaled [B,T,H,K] bf16 q*exp2(g), scale deferred to K4/fwd_o + kg [B,T,H,K] bf16 + gk_last_exp[B,NT,H,K] fp32 + A_qk [B,T,H,BT] bf16 full merged (diagonal + off-diagonal) + A_kk [B,T,H,BT] bf16 pre-inverse output by default; bf16 inverse + output when fp32_akk_inv uses an fp32 workspace +""" + + +import weakref + +import cutlass +import cutlass.cute as cute +import torch +from cutlass import for_generate, yield_out +from cutlass._mlir import ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cutlass_dsl import T, dsl_user_op +from fla.ops.utils import prepare_chunk_indices + +from cula.ops.kda.sm100.akk_inv_fp32 import akk_inv_fp32_physical_host as _akk_inv_fp32_physical_host +from cula.ops.kda.sm100.akk_inv_tf32 import _add_store_C16_smem as _tf32_add_store_C16_smem +from cula.ops.kda.sm100.akk_inv_tf32 import _invert_diag_forward16 as _tf32_invert_diag_forward16 +from cula.ops.kda.sm100.akk_inv_tf32 import _matmul16_smem_smem as _tf32_matmul16_smem_smem +from cula.ops.kda.sm100.akk_inv_tf32 import _matmul32_smem_smem as _tf32_matmul32_smem_smem +from cula.ops.kda.sm100.akk_inv_tf32 import _matmul16_tmp_smem as _tf32_matmul16_tmp_smem +from cula.ops.kda.sm100.akk_inv_tf32 import _store_C16_smem as _tf32_store_C16_smem +from cula.ops.kda.sm100.akk_inv_tf32 import _store_C16_tmp as _tf32_store_C16_tmp + +# (Test-only references to upstream act_cumsum / intra_parallel / fla triton +# kernels were stripped during flashinfer migration — the fused fast path +# does not depend on them.) + +BT = 64 +BC = 16 +K_DIM = 128 +K_PAD = 8 +K_STRIDE = K_DIM + K_PAD # 136, padded row stride to avoid bank conflicts +CHUNKS_PER_BLOCK = 4 +NUM_SMS = 148 # Persistent kernel: one resident block per SM + +NUM_K1_TMA_WARPS = 16 # Warps 0-15: K1 compute (4 warpgroups, 8×2) -- TMA offloaded +NUM_MMA_WARPS = 11 # Warps 16-26: MMA (10 active + 1 TMA producer, dropped idle warp 27) +NUM_MMA_ACTIVE = 10 # mma_warp 0..9: actual MMA work +TMA_WARP_ID = NUM_K1_TMA_WARPS + NUM_MMA_ACTIVE # warp 26 = dedicated TMA producer +NUM_STORE_WARPS = 4 # Warps 27-30: vectorized stores / optional in-CTA inverse +STORE_WARP_BASE = NUM_K1_TMA_WARPS + NUM_MMA_WARPS +NUM_WARPS = STORE_WARP_BASE + NUM_STORE_WARPS # 31 +THREADS = NUM_WARPS * 32 # 992 + +NUM_SUB_CHUNKS = BT // BC # 4 +NUM_TILES = NUM_SUB_CHUNKS * (NUM_SUB_CHUNKS + 1) // 2 # 10 lower-tri tiles +MMA_K_TILE = 16 +NUM_MMA_K_TILES = K_DIM // MMA_K_TILE # 8 (bf16 m16n8k16) +AQK_TILE_PAD = 8 +AQK_TILE_STRIDE = BT + AQK_TILE_PAD # 72 — sAqk now 64x72 row-major (same shape as sAkk) + +# sAkk (fp32) and sG (bf16) are independently allocated — no physical alias. +# - sAkk FP32: row stride = AKK_STRIDE*4 bytes +# - sG BF16: row stride = K_DIM*2 bytes (= 256 B), independent of sAkk +# AKK_PAD=4 → stride 68 to avoid SMEM bank conflicts on MMA reads of sAkk +# (was 0/stride 64 to co-align bytes with sG under an alias scheme that is +# no longer in use). +2 KB total over 2 stages, fits the ~5 KB SMEM slack. +AKK_PAD = 4 +AKK_STRIDE = BT + AKK_PAD # 68 (padded, bank-conflict-free for MMA) +# sG independent layout constants (bf16 units) +G_ROW_STRIDE_BF16 = K_DIM # 128 bf16 per row +G_STAGE_STRIDE_BF16 = BT * K_DIM # 8192 bf16 per stage (one stage = BT*K_DIM) + +K1_ROW_GROUPS = 8 +K1_COL_GROUPS = 2 +ROWS_PER_K1_WARP = BT // K1_ROW_GROUPS # 8 +K1_COLS_PER_WARP = K_DIM // K1_COL_GROUPS # 64 +VEC = K1_COLS_PER_WARP // 32 # 2 +K_VEC = K_DIM // VEC # 64 +NUM_STAGES = 2 +PARTIAL_COLS = K_DIM + 4 # 132 +PARTIAL_COLS_PER_WARP = K_DIM // NUM_K1_TMA_WARPS # 8 +INV_TMP_STRIDE = BC + 4 + +_TILE_IQ = [0, 1, 1, 2, 2, 2, 3, 3, 3, 3] +_TILE_IK = [0, 0, 1, 0, 1, 2, 0, 1, 2, 3] + +LOG2E = 1.4426950408889634 +LN2 = 0.6931471805599453 +RCP_LN2 = LOG2E + + +@dsl_user_op +def k1_internal_barrier(*, loc=None, ip=None): + """Named barrier for K1+TMA warps (0-15, 512 threads). barrier_id=2.""" + llvm.inline_asm( + T.i32(), + [], + "membar.cta; bar.sync 2, 512; mov.u32 $0, 0;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def store_internal_barrier(*, loc=None, ip=None): + """Named barrier for the four store/inverse warps. barrier_id=3.""" + llvm.inline_asm( + T.i32(), + [], + "membar.cta; bar.sync 3, 128; mov.u32 $0, 0;", + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def pack_bf16x2_f32(hi_f32, lo_f32, *, loc=None, ip=None): + """Pack two fp32 values into a bf16x2 u32 register. + Returns u32 with bf16(hi) in bits [31:16] and bf16(lo) in bits [15:0]. + PTX: cvt.rn.bf16x2.f32 d, a, b -> d[31:16]=bf16(a), d[15:0]=bf16(b) + """ + result = llvm.inline_asm( + T.i32(), + [hi_f32.ir_value(loc=loc, ip=ip), lo_f32.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $1, $2;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Int32(result) + + +@dsl_user_op +def mma_bf16_m16n8k16( + a0, + a1, + a2, + a3, # 4 u32 (bf16x2 packed) A operands + b0, + b1, # 2 u32 (bf16x2 packed) B operands + c0, + c1, + c2, + c3, # 4 fp32 accumulators + *, + loc=None, + ip=None, +): + """bf16 MMA with fp32 accumulator, shape m16n8k16. + D_fp32 = A_bf16 * B_bf16 + C_fp32 + """ + # a/b already i32 (from pack_bf16x2_f32) -> no bitcast needed + result = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [ + a0.ir_value(loc=loc, ip=ip), + a1.ir_value(loc=loc, ip=ip), + a2.ir_value(loc=loc, ip=ip), + a3.ir_value(loc=loc, ip=ip), + b0.ir_value(loc=loc, ip=ip), + b1.ir_value(loc=loc, ip=ip), + c0.ir_value(loc=loc, ip=ip), + c1.ir_value(loc=loc, ip=ip), + c2.ir_value(loc=loc, ip=ip), + c3.ir_value(loc=loc, ip=ip), + ], + """{ + mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 + {$0, $1, $2, $3}, + {$4, $5, $6, $7}, + {$8, $9}, + {$10, $11, $12, $13}; + }""", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +@dsl_user_op +def mma_tf32_m16n8k8(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): + """Register-level TF32 MMA with fp32 accumulator, shape m16n8k8.""" + a0b = llvm.bitcast(T.i32(), a0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a1b = llvm.bitcast(T.i32(), a1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a2b = llvm.bitcast(T.i32(), a2.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a3b = llvm.bitcast(T.i32(), a3.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b0b = llvm.bitcast(T.i32(), b0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b1b = llvm.bitcast(T.i32(), b1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [ + a0b, + a1b, + a2b, + a3b, + b0b, + b1b, + c0.ir_value(loc=loc, ip=ip), + c1.ir_value(loc=loc, ip=ip), + c2.ir_value(loc=loc, ip=ip), + c3.ir_value(loc=loc, ip=ip), + ], + """{ + mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 + {$0, $1, $2, $3}, + {$4, $5, $6, $7}, + {$8, $9}, + {$10, $11, $12, $13}; + }""", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +SHFL_W8_CLAMP = 0x1800 + + +@dsl_user_op +def fast_rcp(x, *, loc=None, ip=None): + """Hardware fast reciprocal: rcp.approx.ftz.f32 (~2 cycles vs ~20 for div).""" + result = llvm.inline_asm( + T.f32(), + [x.ir_value(loc=loc, ip=ip)], + "rcp.approx.ftz.f32 $0, $1;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(result) + + +# ============================================================ +# Akk inverse helpers (fused from Akk_inverse_lower_triangle_bf16.py) +# ============================================================ +@dsl_user_op +def _ak_pack_bf16x2(lo_f32, hi_f32, *, loc=None, ip=None): + """cvt.rn.bf16x2.f32 -- pack two fp32 into bf16x2 (as fp32 bitcast view).""" + result = llvm.inline_asm( + T.i32(), + [lo_f32.ir_value(loc=loc, ip=ip), hi_f32.ir_value(loc=loc, ip=ip)], + "cvt.rn.bf16x2.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _ak_movmatrix_trans(src, *, loc=None, ip=None): + src_b = llvm.bitcast(T.i32(), src.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [src_b], + "movmatrix.sync.aligned.m8n8.trans.b16 $0, $1;", + "=r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _ak_mask_packed_ltri(packed, row, pair, *, loc=None, ip=None): + """Mask packed bf16x2 to lower-tri: zero elements where row < col.""" + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p, row.ir_value(loc=loc, ip=ip), pair.ir_value(loc=loc, ip=ip)], + """{ + .reg .b32 %c0, %c1, %mlo, %mhi, %mask; + .reg .pred %p0, %p1; + shl.b32 %c0, $3, 1; + add.u32 %c1, %c0, 1; + setp.ge.s32 %p0, $2, %c0; + setp.ge.s32 %p1, $2, %c1; + selp.b32 %mlo, 0xFFFF, 0, %p0; + selp.b32 %mhi, 0xFFFF0000, 0, %p1; + or.b32 %mask, %mlo, %mhi; + and.b32 $0, $1, %mask; + }""", + "=r,r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _ak_unpack_bf16x2_lo(packed, *, loc=None, ip=None): + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p], + "shl.b32 $0, $1, 16;", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _ak_unpack_bf16x2_hi(packed, *, loc=None, ip=None): + p = llvm.bitcast(T.i32(), packed.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + T.i32(), + [p], + "and.b32 $0, $1, 0xFFFF0000;", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Float32(llvm.bitcast(T.f32(), result, loc=loc, ip=ip)) + + +@dsl_user_op +def _ak_mma(a0, a1, a2, a3, b0, b1, c0, c1, c2, c3, *, loc=None, ip=None): + """BF16 MMA m16n8k16, args are Float32 (bitcast-viewed as i32 for bf16x2).""" + a0b = llvm.bitcast(T.i32(), a0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a1b = llvm.bitcast(T.i32(), a1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a2b = llvm.bitcast(T.i32(), a2.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + a3b = llvm.bitcast(T.i32(), a3.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b0b = llvm.bitcast(T.i32(), b0.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + b1b = llvm.bitcast(T.i32(), b1.ir_value(loc=loc, ip=ip), loc=loc, ip=ip) + result = llvm.inline_asm( + ir.Type.parse("!llvm.struct<(f32, f32, f32, f32)>"), + [ + a0b, + a1b, + a2b, + a3b, + b0b, + b1b, + c0.ir_value(loc=loc, ip=ip), + c1.ir_value(loc=loc, ip=ip), + c2.ir_value(loc=loc, ip=ip), + c3.ir_value(loc=loc, ip=ip), + ], + """{ + mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 + {$0, $1, $2, $3}, + {$4, $5, $6, $7}, + {$8, $9}, + {$10, $11, $12, $13}; + }""", + "=f,=f,=f,=f,r,r,r,r,r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + d0 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)) + d1 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)) + d2 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)) + d3 = cutlass.Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)) + return d0, d1, d2, d3 + + +@dsl_user_op +def _ak_invert_diag_neumann(sAkk, block_idx, lane_id, *, loc=None, ip=None): + """Invert 16x16 diag block via Neumann series (I+L)^-1 = (I-L)(I+L^2)(I+L^4)(I+L^8). + Reads/writes sAkk as fp32 packed bf16x2 (stride 36). Each block handled by 1 warp.""" + r_off = block_idx * 16 + c_off = block_idx * 8 # 16 cols = 8 pairs + gid = lane_id // 4 + tid = lane_id % 4 + + packed0 = cutlass.Float32(sAkk[r_off + gid, c_off + tid]) + packed1 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + tid]) + packed2 = cutlass.Float32(sAkk[r_off + gid, c_off + 4 + tid]) + packed3 = cutlass.Float32(sAkk[r_off + gid + 8, c_off + 4 + tid]) + + A_f0 = _ak_unpack_bf16x2_lo(packed0) + A_f1 = _ak_unpack_bf16x2_hi(packed0) + A_f2 = _ak_unpack_bf16x2_lo(packed1) + A_f3 = _ak_unpack_bf16x2_hi(packed1) + A_f4 = _ak_unpack_bf16x2_lo(packed2) + A_f5 = _ak_unpack_bf16x2_hi(packed2) + A_f6 = _ak_unpack_bf16x2_lo(packed3) + A_f7 = _ak_unpack_bf16x2_hi(packed3) + + _one = cutlass.Float32(1.0) + _zero = cutlass.Float32(0.0) + I_f0 = _one * cutlass.Float32(gid == 2 * tid) + _zero * cutlass.Float32(gid != 2 * tid) + I_f1 = _one * cutlass.Float32(gid == 2 * tid + 1) + _zero * cutlass.Float32(gid != 2 * tid + 1) + I_f2 = _one * cutlass.Float32(gid + 8 == 2 * tid) + _zero * cutlass.Float32(gid + 8 != 2 * tid) + I_f3 = _one * cutlass.Float32(gid + 8 == 2 * tid + 1) + _zero * cutlass.Float32(gid + 8 != 2 * tid + 1) + I_f4 = _one * cutlass.Float32(gid == 8 + 2 * tid) + _zero * cutlass.Float32(gid != 8 + 2 * tid) + I_f5 = _one * cutlass.Float32(gid == 8 + 2 * tid + 1) + _zero * cutlass.Float32(gid != 8 + 2 * tid + 1) + I_f6 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid) + _zero * cutlass.Float32(gid + 8 != 8 + 2 * tid) + I_f7 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid + 1) + _zero * cutlass.Float32(gid + 8 != 8 + 2 * tid + 1) + + L_f0 = A_f0 - I_f0 + L_f1 = A_f1 - I_f1 + L_f2 = A_f2 - I_f2 + L_f3 = A_f3 - I_f3 + L_f4 = A_f4 - I_f4 + L_f5 = A_f5 - I_f5 + L_f6 = A_f6 - I_f6 + L_f7 = A_f7 - I_f7 + INV_f0 = I_f0 - L_f0 + INV_f1 = I_f1 - L_f1 + INV_f2 = I_f2 - L_f2 + INV_f3 = I_f3 - L_f3 + INV_f4 = I_f4 - L_f4 + INV_f5 = I_f5 - L_f5 + INV_f6 = I_f6 - L_f6 + INV_f7 = I_f7 - L_f7 + + _zf = cutlass.Float32(0.0) + L_a0 = _ak_pack_bf16x2(L_f0, L_f1) + L_a1 = _ak_pack_bf16x2(L_f2, L_f3) + L_a2 = _ak_pack_bf16x2(L_f4, L_f5) + L_a3 = _ak_pack_bf16x2(L_f6, L_f7) + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 1: L^2, INV += INV*L^2 + L_b0 = _ak_movmatrix_trans(L_a0) + L_b1 = _ak_movmatrix_trans(L_a1) + L_b2 = _ak_movmatrix_trans(L_a2) + L_b3 = _ak_movmatrix_trans(L_a3) + Lp_c0, Lp_c1, Lp_c2, Lp_c3 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b0, L_b1, _zf, _zf, _zf, _zf) + Lp_c4, Lp_c5, Lp_c6, Lp_c7 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b2, L_b3, _zf, _zf, _zf, _zf) + Lp_a0 = _ak_pack_bf16x2(Lp_c0, Lp_c1) + Lp_a1 = _ak_pack_bf16x2(Lp_c2, Lp_c3) + Lp_a2 = _ak_pack_bf16x2(Lp_c4, Lp_c5) + Lp_a3 = _ak_pack_bf16x2(Lp_c6, Lp_c7) + Lp_b0 = _ak_movmatrix_trans(Lp_a0) + Lp_b1 = _ak_movmatrix_trans(Lp_a1) + Lp_b2 = _ak_movmatrix_trans(Lp_a2) + Lp_b3 = _ak_movmatrix_trans(Lp_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma(INV_a0, INV_a1, INV_a2, INV_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma(INV_a0, INV_a1, INV_a2, INV_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 2: L^4, INV += INV*L^4 + L4_c0, L4_c1, L4_c2, L4_c3 = _ak_mma(Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf) + L4_c4, L4_c5, L4_c6, L4_c7 = _ak_mma(Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf) + L4_a0 = _ak_pack_bf16x2(L4_c0, L4_c1) + L4_a1 = _ak_pack_bf16x2(L4_c2, L4_c3) + L4_a2 = _ak_pack_bf16x2(L4_c4, L4_c5) + L4_a3 = _ak_pack_bf16x2(L4_c6, L4_c7) + L4_b0 = _ak_movmatrix_trans(L4_a0) + L4_b1 = _ak_movmatrix_trans(L4_a1) + L4_b2 = _ak_movmatrix_trans(L4_a2) + L4_b3 = _ak_movmatrix_trans(L4_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma(INV_a0, INV_a1, INV_a2, INV_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma(INV_a0, INV_a1, INV_a2, INV_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 3: L^8, INV += INV*L^8 + L8_c0, L8_c1, L8_c2, L8_c3 = _ak_mma(L4_a0, L4_a1, L4_a2, L4_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf) + L8_c4, L8_c5, L8_c6, L8_c7 = _ak_mma(L4_a0, L4_a1, L4_a2, L4_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf) + L8_a0 = _ak_pack_bf16x2(L8_c0, L8_c1) + L8_a1 = _ak_pack_bf16x2(L8_c2, L8_c3) + L8_a2 = _ak_pack_bf16x2(L8_c4, L8_c5) + L8_a3 = _ak_pack_bf16x2(L8_c6, L8_c7) + L8_b0 = _ak_movmatrix_trans(L8_a0) + L8_b1 = _ak_movmatrix_trans(L8_a1) + L8_b2 = _ak_movmatrix_trans(L8_a2) + L8_b3 = _ak_movmatrix_trans(L8_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma(INV_a0, INV_a1, INV_a2, INV_a3, L8_b0, L8_b1, _zf, _zf, _zf, _zf) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma(INV_a0, INV_a1, INV_a2, INV_a3, L8_b2, L8_b3, _zf, _zf, _zf, _zf) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + + sAkk[r_off + gid, c_off + tid] = _ak_pack_bf16x2(INV_f0, INV_f1) + sAkk[r_off + gid + 8, c_off + tid] = _ak_pack_bf16x2(INV_f2, INV_f3) + sAkk[r_off + gid, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f4, INV_f5) + sAkk[r_off + gid + 8, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f6, INV_f7) + + +@dsl_user_op +def _ak_invert_diag_neumann_inreg( + sAkk, + block_idx, + lane_id, + raw_f0, + raw_f1, + raw_f2, + raw_f3, + raw_f4, + raw_f5, + raw_f6, + raw_f7, + *, + loc=None, + ip=None, +): + """In-register variant: invert 16x16 diag block from K2 acc fragment values + (no SMEM read). Applies I+L mask, runs Neumann iterations, writes INV to sAkk. + + raw_f0..raw_f7 are K2 fp32 acc*beta values in C-fragment layout for m16n8k16: + raw_f0,raw_f1 = (gid, 2*tid), (gid, 2*tid+1) + raw_f2,raw_f3 = (gid+8, 2*tid), (gid+8, 2*tid+1) + raw_f4,raw_f5 = (gid, 2*tid+8), (gid, 2*tid+9) + raw_f6,raw_f7 = (gid+8, 2*tid+8), (gid+8, 2*tid+9) + """ + r_off = block_idx * 16 + c_off = block_idx * 8 + gid = lane_id // 4 + tid = lane_id % 4 + + _one = cutlass.Float32(1.0) + _zero = cutlass.Float32(0.0) + + # Identity pattern (matches the SMEM-read variant's I_f computation) + I_f0 = _one * cutlass.Float32(gid == 2 * tid) + _zero * cutlass.Float32(gid != 2 * tid) + I_f1 = _one * cutlass.Float32(gid == 2 * tid + 1) + _zero * cutlass.Float32(gid != 2 * tid + 1) + I_f2 = _one * cutlass.Float32(gid + 8 == 2 * tid) + _zero * cutlass.Float32(gid + 8 != 2 * tid) # always 0 + I_f3 = _one * cutlass.Float32(gid + 8 == 2 * tid + 1) + _zero * cutlass.Float32(gid + 8 != 2 * tid + 1) # always 0 + I_f4 = _one * cutlass.Float32(gid == 8 + 2 * tid) + _zero * cutlass.Float32(gid != 8 + 2 * tid) # always 0 + I_f5 = _one * cutlass.Float32(gid == 8 + 2 * tid + 1) + _zero * cutlass.Float32(gid != 8 + 2 * tid + 1) # always 0 + I_f6 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid) + _zero * cutlass.Float32(gid + 8 != 8 + 2 * tid) + I_f7 = _one * cutlass.Float32(gid + 8 == 8 + 2 * tid + 1) + _zero * cutlass.Float32(gid + 8 != 8 + 2 * tid + 1) + + # Apply I+L mask to raw K2 values: + # diag (row==col): replace with 1 + # strict upper (rowcol within block): keep raw value (= L value) + m_gt_a = cutlass.Float32(gid > 2 * tid) # for cols 2*tid (with row=gid) + m_gt_b = cutlass.Float32(gid > 2 * tid + 1) # for cols 2*tid+1 + # (gid, 2*tid) and (gid, 2*tid+1): mask conditional + A_f0 = m_gt_a * raw_f0 + I_f0 + A_f1 = m_gt_b * raw_f1 + I_f1 + # (gid+8, 2*tid) and (gid+8, 2*tid+1): always strict lower (gid+8 > 2*tid+1 always) + A_f2 = raw_f2 + A_f3 = raw_f3 + # (gid, 2*tid+8) and (gid, 2*tid+9): always strict upper (gid <= 7 < 8 <= 2*tid+8) + A_f4 = _zero + A_f5 = _zero + # (gid+8, 2*tid+8): row-col offset = gid+8 - (2*tid+8) = gid - 2*tid -> same as A_f0 mask + # (gid+8, 2*tid+9): similar -> same as A_f1 mask + A_f6 = m_gt_a * raw_f6 + I_f6 + A_f7 = m_gt_b * raw_f7 + I_f7 + + # L = A - I, INV = I - L (first two Neumann terms) + L_f0 = A_f0 - I_f0 + L_f1 = A_f1 - I_f1 + L_f2 = A_f2 - I_f2 + L_f3 = A_f3 - I_f3 + L_f4 = A_f4 - I_f4 + L_f5 = A_f5 - I_f5 + L_f6 = A_f6 - I_f6 + L_f7 = A_f7 - I_f7 + INV_f0 = I_f0 - L_f0 + INV_f1 = I_f1 - L_f1 + INV_f2 = I_f2 - L_f2 + INV_f3 = I_f3 - L_f3 + INV_f4 = I_f4 - L_f4 + INV_f5 = I_f5 - L_f5 + INV_f6 = I_f6 - L_f6 + INV_f7 = I_f7 - L_f7 + + _zf = cutlass.Float32(0.0) + L_a0 = _ak_pack_bf16x2(L_f0, L_f1) + L_a1 = _ak_pack_bf16x2(L_f2, L_f3) + L_a2 = _ak_pack_bf16x2(L_f4, L_f5) + L_a3 = _ak_pack_bf16x2(L_f6, L_f7) + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 1: L^2, INV += INV*L^2 + L_b0 = _ak_movmatrix_trans(L_a0) + L_b1 = _ak_movmatrix_trans(L_a1) + L_b2 = _ak_movmatrix_trans(L_a2) + L_b3 = _ak_movmatrix_trans(L_a3) + Lp_c0, Lp_c1, Lp_c2, Lp_c3 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b0, L_b1, _zf, _zf, _zf, _zf) + Lp_c4, Lp_c5, Lp_c6, Lp_c7 = _ak_mma(L_a0, L_a1, L_a2, L_a3, L_b2, L_b3, _zf, _zf, _zf, _zf) + Lp_a0 = _ak_pack_bf16x2(Lp_c0, Lp_c1) + Lp_a1 = _ak_pack_bf16x2(Lp_c2, Lp_c3) + Lp_a2 = _ak_pack_bf16x2(Lp_c4, Lp_c5) + Lp_a3 = _ak_pack_bf16x2(Lp_c6, Lp_c7) + Lp_b0 = _ak_movmatrix_trans(Lp_a0) + Lp_b1 = _ak_movmatrix_trans(Lp_a1) + Lp_b2 = _ak_movmatrix_trans(Lp_a2) + Lp_b3 = _ak_movmatrix_trans(Lp_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma(INV_a0, INV_a1, INV_a2, INV_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma(INV_a0, INV_a1, INV_a2, INV_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 2: L^4, INV += INV*L^4 + L4_c0, L4_c1, L4_c2, L4_c3 = _ak_mma(Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b0, Lp_b1, _zf, _zf, _zf, _zf) + L4_c4, L4_c5, L4_c6, L4_c7 = _ak_mma(Lp_a0, Lp_a1, Lp_a2, Lp_a3, Lp_b2, Lp_b3, _zf, _zf, _zf, _zf) + L4_a0 = _ak_pack_bf16x2(L4_c0, L4_c1) + L4_a1 = _ak_pack_bf16x2(L4_c2, L4_c3) + L4_a2 = _ak_pack_bf16x2(L4_c4, L4_c5) + L4_a3 = _ak_pack_bf16x2(L4_c6, L4_c7) + L4_b0 = _ak_movmatrix_trans(L4_a0) + L4_b1 = _ak_movmatrix_trans(L4_a1) + L4_b2 = _ak_movmatrix_trans(L4_a2) + L4_b3 = _ak_movmatrix_trans(L4_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma(INV_a0, INV_a1, INV_a2, INV_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma(INV_a0, INV_a1, INV_a2, INV_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + INV_a0 = _ak_pack_bf16x2(INV_f0, INV_f1) + INV_a1 = _ak_pack_bf16x2(INV_f2, INV_f3) + INV_a2 = _ak_pack_bf16x2(INV_f4, INV_f5) + INV_a3 = _ak_pack_bf16x2(INV_f6, INV_f7) + + # Iter 3: L^8, INV += INV*L^8 + L8_c0, L8_c1, L8_c2, L8_c3 = _ak_mma(L4_a0, L4_a1, L4_a2, L4_a3, L4_b0, L4_b1, _zf, _zf, _zf, _zf) + L8_c4, L8_c5, L8_c6, L8_c7 = _ak_mma(L4_a0, L4_a1, L4_a2, L4_a3, L4_b2, L4_b3, _zf, _zf, _zf, _zf) + L8_a0 = _ak_pack_bf16x2(L8_c0, L8_c1) + L8_a1 = _ak_pack_bf16x2(L8_c2, L8_c3) + L8_a2 = _ak_pack_bf16x2(L8_c4, L8_c5) + L8_a3 = _ak_pack_bf16x2(L8_c6, L8_c7) + L8_b0 = _ak_movmatrix_trans(L8_a0) + L8_b1 = _ak_movmatrix_trans(L8_a1) + L8_b2 = _ak_movmatrix_trans(L8_a2) + L8_b3 = _ak_movmatrix_trans(L8_a3) + mm_c0, mm_c1, mm_c2, mm_c3 = _ak_mma(INV_a0, INV_a1, INV_a2, INV_a3, L8_b0, L8_b1, _zf, _zf, _zf, _zf) + mm_c4, mm_c5, mm_c6, mm_c7 = _ak_mma(INV_a0, INV_a1, INV_a2, INV_a3, L8_b2, L8_b3, _zf, _zf, _zf, _zf) + INV_f0 = INV_f0 + mm_c0 + INV_f1 = INV_f1 + mm_c1 + INV_f2 = INV_f2 + mm_c2 + INV_f3 = INV_f3 + mm_c3 + INV_f4 = INV_f4 + mm_c4 + INV_f5 = INV_f5 + mm_c5 + INV_f6 = INV_f6 + mm_c6 + INV_f7 = INV_f7 + mm_c7 + + sAkk[r_off + gid, c_off + tid] = _ak_pack_bf16x2(INV_f0, INV_f1) + sAkk[r_off + gid + 8, c_off + tid] = _ak_pack_bf16x2(INV_f2, INV_f3) + sAkk[r_off + gid, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f4, INV_f5) + sAkk[r_off + gid + 8, c_off + 4 + tid] = _ak_pack_bf16x2(INV_f6, INV_f7) + + +@dsl_user_op +def _ak_matmul_AB(sAkk, br_A, bc_A, br_B, bc_B, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 8 + rB = br_B * 16 + cB = bc_B * 8 + a0 = cutlass.Float32(sAkk[rA + gid, cA + tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid]) + bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid]) + bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid]) + bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid]) + bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid]) + b0 = _ak_movmatrix_trans(bA0) + b1 = _ak_movmatrix_trans(bA1) + b2 = _ak_movmatrix_trans(bA2) + b3 = _ak_movmatrix_trans(bA3) + cn0_0, cn0_1, cn0_2, cn0_3 = _ak_mma(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf) + cn1_0, cn1_1, cn1_2, cn1_3 = _ak_mma(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _ak_chain_mma_B(sAkk, br_B, bc_B, a0, a1, a2, a3, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rB = br_B * 16 + cB = bc_B * 8 + bA0 = cutlass.Float32(sAkk[rB + gid, cB + tid]) + bA1 = cutlass.Float32(sAkk[rB + gid + 8, cB + tid]) + bA2 = cutlass.Float32(sAkk[rB + gid, cB + 4 + tid]) + bA3 = cutlass.Float32(sAkk[rB + gid + 8, cB + 4 + tid]) + b0 = _ak_movmatrix_trans(bA0) + b1 = _ak_movmatrix_trans(bA1) + b2 = _ak_movmatrix_trans(bA2) + b3 = _ak_movmatrix_trans(bA3) + cn0_0, cn0_1, cn0_2, cn0_3 = _ak_mma(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf) + cn1_0, cn1_1, cn1_2, cn1_3 = _ak_mma(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _ak_chain_mma_A(sAkk, br_A, bc_A, b0, b1, b2, b3, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + _zf = cutlass.Float32(0.0) + rA = br_A * 16 + cA = bc_A * 8 + a0 = cutlass.Float32(sAkk[rA + gid, cA + tid]) + a1 = cutlass.Float32(sAkk[rA + gid + 8, cA + tid]) + a2 = cutlass.Float32(sAkk[rA + gid, cA + 4 + tid]) + a3 = cutlass.Float32(sAkk[rA + gid + 8, cA + 4 + tid]) + cn0_0, cn0_1, cn0_2, cn0_3 = _ak_mma(a0, a1, a2, a3, b0, b1, _zf, _zf, _zf, _zf) + cn1_0, cn1_1, cn1_2, cn1_3 = _ak_mma(a0, a1, a2, a3, b2, b3, _zf, _zf, _zf, _zf) + return cn0_0, cn0_1, cn0_2, cn0_3, cn1_0, cn1_1, cn1_2, cn1_3 + + +@dsl_user_op +def _ak_store_neg_C(sAkk, br, bc, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + r = br * 16 + c = bc * 8 + sAkk[r + gid, c + tid] = _ak_pack_bf16x2(-c0, -c1) + sAkk[r + gid + 8, c + tid] = _ak_pack_bf16x2(-c2, -c3) + sAkk[r + gid, c + 4 + tid] = _ak_pack_bf16x2(-c4, -c5) + sAkk[r + gid + 8, c + 4 + tid] = _ak_pack_bf16x2(-c6, -c7) + + +@dsl_user_op +def _ak_pack_C_to_A(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None): + a0 = _ak_pack_bf16x2(c0, c1) + a1 = _ak_pack_bf16x2(c2, c3) + a2 = _ak_pack_bf16x2(c4, c5) + a3 = _ak_pack_bf16x2(c6, c7) + return a0, a1, a2, a3 + + +@dsl_user_op +def _ak_pack_C_to_B(c0, c1, c2, c3, c4, c5, c6, c7, *, loc=None, ip=None): + a0 = _ak_pack_bf16x2(c0, c1) + a1 = _ak_pack_bf16x2(c2, c3) + a2 = _ak_pack_bf16x2(c4, c5) + a3 = _ak_pack_bf16x2(c6, c7) + b0 = _ak_movmatrix_trans(a0) + b1 = _ak_movmatrix_trans(a1) + b2 = _ak_movmatrix_trans(a2) + b3 = _ak_movmatrix_trans(a3) + return b0, b1, b2, b3 + + +@dsl_user_op +def _ak_store_C_temp(sT, buf, c0, c1, c2, c3, c4, c5, c6, c7, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + sT[gid, 2 * tid, buf] = c0 + sT[gid, 2 * tid + 1, buf] = c1 + sT[gid + 8, 2 * tid, buf] = c2 + sT[gid + 8, 2 * tid + 1, buf] = c3 + sT[gid, 8 + 2 * tid, buf] = c4 + sT[gid, 8 + 2 * tid + 1, buf] = c5 + sT[gid + 8, 8 + 2 * tid, buf] = c6 + sT[gid + 8, 8 + 2 * tid + 1, buf] = c7 + + +@dsl_user_op +def _ak_load_C_temp(sT, buf, lane_id, *, loc=None, ip=None): + gid = lane_id // 4 + tid = lane_id % 4 + c0 = cutlass.Float32(sT[gid, 2 * tid, buf]) + c1 = cutlass.Float32(sT[gid, 2 * tid + 1, buf]) + c2 = cutlass.Float32(sT[gid + 8, 2 * tid, buf]) + c3 = cutlass.Float32(sT[gid + 8, 2 * tid + 1, buf]) + c4 = cutlass.Float32(sT[gid, 8 + 2 * tid, buf]) + c5 = cutlass.Float32(sT[gid, 8 + 2 * tid + 1, buf]) + c6 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid, buf]) + c7 = cutlass.Float32(sT[gid + 8, 8 + 2 * tid + 1, buf]) + return c0, c1, c2, c3, c4, c5, c6, c7 + + +@dsl_user_op +def opaque_zero_from_work_id(*, loc=None, ip=None): + """ + Return 0 via opaque side-effectful asm (no inputs needed). + + Because has_side_effects=True, MLIR LICM treats this as having memory + effects and will NOT hoist it outside the for_generate loop. Any value + computed from the result (_oz) therefore appears loop-variant to LICM, + preventing get_slice() and scalar layout-invariant computations from being + hoisted to the kernel prologue. This keeps prologue register pressure < 64 + and eliminates the 440-byte stack frame that caused ~300-cycle L2 LDL + penalties per iteration. + """ + result = llvm.inline_asm( + T.i32(), + [], + "mov.b32 $0, 0;", # output = 0 (opaque to compiler constant-folding) + "=r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return cutlass.Int32(result) + + +@cute.kernel +def fused_kernel123( + 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, + mA_log: cute.Tensor, + mBeta: cute.Tensor, + scale: cutlass.Float32, + mKscaled: cute.Tensor, + mKg: cute.Tensor, + mQscaled: cute.Tensor, + mGkLast: cute.Tensor, + mAqk: cute.Tensor, # 4D (B, T, H, BT) — used by per-tile pure path + mAkk: cute.Tensor, # 4D (B, T, H, BT) — used by per-tile pure path + mAqk_v2: cute.Tensor, # 5D (B, T, H, BT/2, 2) — used by vec autovec non-pure path + mAkk_v2: cute.Tensor, # 5D (B, T, H, BT/2, 2) — used by vec autovec non-pure path + tiled_copy_qk_k1, + tiled_mma_k2, + tiled_copy_mma_A, + tiled_copy_mma_B, + tiled_copy_Gcum_norm, + tiled_copy_Gcum_gate, + qk_smem_layout, + g_smem_layout, + g_cumsum_layout, + num_chunks: int, + num_heads: int, + batch_size: int, + mCuSeqlens: cute.Tensor, + mChunkIndices: cute.Tensor, + IS_VARLEN: cutlass.Constexpr[int], + mDtBias: cute.Tensor, + lower_bound: cutlass.Float32, + HAS_BIAS: cutlass.Constexpr[int], + USE_SAFE_GATE: cutlass.Constexpr[int], + VARLEN_PURE: cutlass.Constexpr[int] = 0, + KSCALED_FP16: cutlass.Constexpr[int] = 0, + PREPROCESS_W_BETA: cutlass.Constexpr[int] = 0, + FUSE_AKK_INV: cutlass.Constexpr[int] = 0, + FP32_AKK_WORKSPACE: cutlass.Constexpr[int] = 0, + INPUT_GK_FP32: cutlass.Constexpr[int] = 0, + BETA_FP32: cutlass.Constexpr[int] = 0, +): + block_id, _, _ = cute.arch.block_idx() + tidx = cute.arch.thread_idx()[0] + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + lane_id = tidx % 32 + + total_cgs_per_head = cutlass.Int32(0) + cgs_per_head = cutlass.Int32(0) + total_cgs = cutlass.Int32(0) + if IS_VARLEN: + total_cgs_per_head = (num_chunks + CHUNKS_PER_BLOCK - 1) // CHUNKS_PER_BLOCK + total_cgs = total_cgs_per_head * num_heads + else: + cgs_per_head = num_chunks // CHUNKS_PER_BLOCK + total_cgs = cgs_per_head * num_heads * batch_size + + # ===================================================================== + # SMEM allocation (不使用 G/Akk alias) + # ===================================================================== + smem = cutlass.utils.SmemAllocator() + sQ = smem.allocate_tensor(cutlass.BFloat16, qk_smem_layout.outer, 128, swizzle=qk_smem_layout.inner) + sK = smem.allocate_tensor(cutlass.BFloat16, qk_smem_layout.outer, 128, swizzle=qk_smem_layout.inner) + sG = smem.allocate_tensor(cutlass.BFloat16, g_smem_layout, 128) + if cutlass.const_expr(INPUT_GK_FP32 != 0): + sGcum = smem.allocate_tensor(cutlass.Float32, g_cumsum_layout.outer, 128, swizzle=g_cumsum_layout.inner) + else: + sGcum = smem.allocate_tensor(cutlass.Float32, g_cumsum_layout, 128) + partial_last_layout = cute.make_layout((K1_ROW_GROUPS, PARTIAL_COLS), stride=(PARTIAL_COLS, 1)) + sPartialLast = smem.allocate_tensor(cutlass.Float32, partial_last_layout, 128) + + # sAqk: 64x72 row-major (BT rows, BT+pad cols), same shape as sAkk. + # Each sub-tile (i_q, i_k) sits at SMEM rows [i_q*BC, (i_q+1)*BC) cols + # [i_k*BC, (i_k+1)*BC) — directly mirrors the 64x64 attention matrix. + aqk_tile_layout = cute.make_layout((BT, AQK_TILE_STRIDE, NUM_STAGES), stride=(AQK_TILE_STRIDE, 1, BT * AQK_TILE_STRIDE)) + sAqk = smem.allocate_tensor(cutlass.BFloat16, aqk_tile_layout, 128) + + # Akk FP32 with stride=32 (4-way bank conflict vs 16-way with stride=64) + # shape [64, 32, 2] - stores two 64x32 sub-tiles, 64x64 Akk read via two MMAs + akk_tile_layout = cute.make_layout((BT, AKK_STRIDE, NUM_STAGES), stride=(AKK_STRIDE, 1, BT * AKK_STRIDE)) + sAkk = smem.allocate_tensor(cutlass.Float32, akk_tile_layout, 128) + + # K1 loads each beta value once; all ten MMA tile warps reuse it. + beta_smem_layout = cute.make_layout((BT, NUM_STAGES), stride=(1, BT)) + if cutlass.const_expr(BETA_FP32 != 0): + sBeta = smem.allocate_tensor(cutlass.Float32, beta_smem_layout, 128) + else: + sBeta = smem.allocate_tensor(cutlass.BFloat16, beta_smem_layout, 128) + + # ===================================================================== + # Mbarrier allocation & init + # ===================================================================== + tma_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + stage_reuse_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + k1_done_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + mma_done_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + store_done_mbars = smem.allocate_array(cutlass.Int64, NUM_STAGES) + + if cutlass.const_expr(INPUT_GK_FP32 != 0): + bytes_per_stage = BT * K_DIM * (2 + 2 + 4) + else: + bytes_per_stage = BT * K_DIM * 2 * 3 + + if tidx == 0: + for s in range(NUM_STAGES): + cute.arch.mbarrier_init(tma_mbars + s, 1) + # TMA warp is waiter (not arriver) on stage_reuse; and it skips mma_done arrive + cute.arch.mbarrier_init(stage_reuse_mbars + s, (NUM_MMA_WARPS - 1) * 32) + cute.arch.mbarrier_init(k1_done_mbars + s, NUM_K1_TMA_WARPS * 32) + cute.arch.mbarrier_init(mma_done_mbars + s, (NUM_MMA_WARPS - 1) * 32) + cute.arch.mbarrier_init(store_done_mbars + s, NUM_STORE_WARPS * 32) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + + # ===================================================================== + # SMEM init: zero out sAqk and sAkk valid 64x64 region (cols 64..71 are + # padding for SMEM bank-conflict avoidance, never read or written by MMA + # or store warps). Required for downstream row-major store optimizations + # — positions outside MMA-written sub-tiles stay at 0. + # + # Cooperative pattern (31 warps × 32 lanes = 992 threads): + # - Each warp owns 2 contiguous rows (warp_id*2, warp_id*2+1) + # - Warp 0 additionally owns the final two rows (62, 63) + # - For sAqk bf16: each lane owns 2 contiguous bf16 cols + # - For sAkk fp32: each lane owns 1 fp32 col (stride 68) + # ===================================================================== + _warp_id_in_cta = tidx >> 5 # tidx // 32, range 0..31 + _lane_id_warp = tidx & 31 # tidx % 32, range 0..31 + _row_base = _warp_id_in_cta * 2 # this warp owns rows [_row_base, _row_base+1] + for _s in cutlass.range_constexpr(NUM_STAGES): + for _ri in cutlass.range_constexpr(2): + _row = _row_base + _ri + # sAqk bf16: 2 cols per lane + _col_lo_bf16 = _lane_id_warp * 2 + _col_hi_bf16 = _col_lo_bf16 + 1 + sAqk[_row, _col_lo_bf16, _s] = cutlass.BFloat16(0.0) + sAqk[_row, _col_hi_bf16, _s] = cutlass.BFloat16(0.0) + # sAkk fp32: 1 col per lane (stride 68) + _col_fp32 = _lane_id_warp + sAkk[_row, _col_fp32, _s] = cutlass.Float32(0.0) + if _warp_id_in_cta == 0: + for _ri in cutlass.range_constexpr(2): + _row = 62 + _ri + _col_lo_bf16 = _lane_id_warp * 2 + _col_hi_bf16 = _col_lo_bf16 + 1 + sAqk[_row, _col_lo_bf16, _s] = cutlass.BFloat16(0.0) + sAqk[_row, _col_hi_bf16, _s] = cutlass.BFloat16(0.0) + sAkk[_row, _lane_id_warp, _s] = cutlass.Float32(0.0) + cute.arch.barrier() + + # ===================================================================== + # Pre-arrive (MMA warps only) + # stage_reuse_mbars: TMA waits before MMA arrives → pre-arrive all 10 MMA warps + # store_done_mbars: MMA waits before Store arrives → pre-arrive first 4 MMA warps + # ===================================================================== + if warp_idx >= NUM_K1_TMA_WARPS and warp_idx < NUM_K1_TMA_WARPS + NUM_MMA_WARPS and warp_idx != TMA_WARP_ID: + mma_warp_tmp = warp_idx - NUM_K1_TMA_WARPS + for s in range(NUM_STAGES): + cute.arch.mbarrier_arrive(stage_reuse_mbars + s) + if mma_warp_tmp < NUM_STORE_WARPS: + cute.arch.mbarrier_arrive(store_done_mbars + s) + + # ================================================================= + # Persistent outer loop. Single for_generate at top level (required). + # Opaque asm barrier on work_id prevents MLIR LICM from hoisting + # get_slice() and scalar layout invariants to the kernel prologue, + # keeping register pressure < 64 and eliminating prologue spill. + # ================================================================= + for work_id in for_generate(block_id, total_cgs, NUM_SMS): + i_cg = cutlass.Int32(0) + i_h = cutlass.Int32(0) + i_b = cutlass.Int32(0) + chunk_base = cutlass.Int32(0) + if IS_VARLEN: + i_cg = work_id % total_cgs_per_head + i_h = work_id // total_cgs_per_head + i_b = cutlass.Int32(0) + chunk_base = i_cg * CHUNKS_PER_BLOCK + else: + i_cg = work_id % cgs_per_head + i_h = (work_id // cgs_per_head) % num_heads + i_b = work_id // (cgs_per_head * num_heads) + chunk_base = i_cg * CHUNKS_PER_BLOCK + + # Anti-LICM barrier: _oz is always 0 but appears to depend on work_id. + # Because this asm has side_effects=True, it stays inside the loop. + # Any value computed from _oz/_lane/_warp is also loop-variant from + # LICM's perspective → get_slice() and scalar invariants stay in-loop. + _oz = opaque_zero_from_work_id() + _lane = lane_id + _oz + _warp = warp_idx + _oz + + # ============================================================= + # Warps 0-15: Fused TMA + K1 + # ============================================================= + if warp_idx < NUM_K1_TMA_WARPS: + # Warp-layout invariants (scope-local → no cross-group register spill) + k1_warp = _warp + warp_row_group = k1_warp % K1_ROW_GROUPS + warp_col_group = k1_warp // K1_ROW_GROUPS + k1_row_start = warp_row_group * ROWS_PER_K1_WARP + col_base = warp_col_group * K1_COLS_PER_WARP + _lane * VEC + col_vec_idx = warp_col_group * (K1_COLS_PER_WARP // VEC) + _lane + cumsum_scale = cutlass.Float32(RCP_LN2) + thr_copy_k1 = tiled_copy_qk_k1.get_slice(_lane) + + rAcc = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + rPrefix = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + rGkLast = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + if cutlass.const_expr(KSCALED_FP16 != 0): + rKsOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float16) + else: + rKsOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + rQsOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + rKgOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + rGkOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + + # exp_A depends only on the head. Compute it once per warp instead + # of issuing the same special-function operation in every lane. + exp_A = cutlass.Float32(0.0) + if _lane == 0: + exp_A = cute.exp(mA_log[i_h], fastmath=True) + exp_A = cute.arch.shuffle_sync(exp_A, 0) + + # Load dt_bias per (head, col) — broadcast across all rows + rBias = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + if HAS_BIAS: + for vi in cutlass.range_constexpr(VEC): + rBias[vi] = mDtBias[i_h, col_base + vi] + else: + for vi in cutlass.range_constexpr(VEC): + rBias[vi] = cutlass.Float32(0.0) + + # 3D TMA head slices (fixed for this work unit's head) + gQ_head = tma_tensor_Q[(None, None, i_h)] + gK_head = tma_tensor_K[(None, None, i_h)] + gG_head = tma_tensor_G[(None, None, i_h)] + + for chunk_iter in cutlass.range_constexpr(CHUNKS_PER_BLOCK): + cur_stage = chunk_iter % NUM_STAGES + cur_phase = chunk_iter // NUM_STAGES % 2 + chunk_idx = chunk_base + chunk_iter + chunk_start = cutlass.Int32(0) + ci_eos = cutlass.Int32(0) + if IS_VARLEN: + if chunk_idx < num_chunks: + _sid = cutlass.Int32(mChunkIndices[chunk_idx, 0]) + chunk_start = cutlass.Int32(mCuSeqlens[_sid]) + cutlass.Int32(mChunkIndices[chunk_idx, 1]) * BT + ci_eos = cutlass.Int32(mCuSeqlens[_sid + 1]) + else: + chunk_start = chunk_idx * BT + + cute.arch.mbarrier_wait(tma_mbars + cur_stage, cur_phase) + + csG = sG[(None, None, cur_stage)] + csGcum = sGcum[(None, None, cur_stage)] + csQ = sQ[(None, None, cur_stage)] + csK = sK[(None, None, cur_stage)] + csBeta = sBeta[(None, cur_stage)] + + if k1_warp == 0: + for beta_iter in cutlass.range_constexpr(BT // 32): + beta_row = beta_iter * 32 + _lane + csBeta[beta_row] = mBeta[i_b, chunk_start + beta_row, i_h] + + rGact = cute.make_rmem_tensor(cute.make_layout((ROWS_PER_K1_WARP, VEC)), cutlass.Float32) + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = cutlass.Float32(0.0) + + if cutlass.const_expr(INPUT_GK_FP32 != 0): + # One-to-one csrc port: consume the same FP32 cumulative + # gate tensor that csrc receives. Do not recompute the + # activation or cumsum inside this kernel. + for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP): + row = k1_row_start + ri + for vi in cutlass.range_constexpr(VEC): + rGact[ri, vi] = csGcum[row, col_base + vi] + for vi in cutlass.range_constexpr(VEC): + rGkLast[vi] = csGcum[BT - 1, col_base + vi] + cute.arch.mbarrier_arrive(k1_done_mbars + cur_stage) + else: + for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP): + row = k1_row_start + ri + for vi in cutlass.range_constexpr(VEC): + c = col_base + vi + g_val = csG[row, c].to(cutlass.Float32) + if HAS_BIAS: + g_val = g_val + rBias[vi] + g_activated = cutlass.Float32(0.0) + if USE_SAFE_GATE: + sigmoid_g = fast_rcp(cutlass.Float32(1.0) + cute.exp2(-exp_A * g_val * LOG2E, fastmath=True)) + g_activated = lower_bound * sigmoid_g + else: + softplus_g = ( + cute.log2( + cutlass.Float32(1.0) + cute.exp2(g_val * LOG2E, fastmath=True), + fastmath=True, + ) + * LN2 + ) + g_activated = -exp_A * softplus_g + # Varlen: zero gate for out-of-bounds rows so cumsum + # stays flat beyond the last valid position. + # VARLEN_PURE=1 elides this at compile time — caller + # guarantees all seq lengths are multiples of BT so no + # chunk has OOB rows. + if IS_VARLEN and not VARLEN_PURE: + if chunk_start + row >= ci_eos: + g_activated = cutlass.Float32(0.0) + rGact[ri, vi] = g_activated + rAcc[vi] = rAcc[vi] + g_activated + + for vi in cutlass.range_constexpr(VEC): + sPartialLast[warp_row_group, col_base + vi] = rAcc[vi] + + k1_internal_barrier() + + prefix_col_start = k1_warp * PARTIAL_COLS_PER_WARP + row_in_prefix = lane_id % K1_ROW_GROUPS + col_in_group = lane_id // K1_ROW_GROUPS + + for j in cutlass.range_constexpr(PARTIAL_COLS_PER_WARP // 4): + col = prefix_col_start + j * 4 + col_in_group + val = cutlass.Float32(sPartialLast[row_in_prefix, col]) + tmp = cute.arch.shuffle_sync_up(val, 1, mask=-1, mask_and_clamp=SHFL_W8_CLAMP) + if row_in_prefix >= 1: + val = val + tmp + tmp = cute.arch.shuffle_sync_up(val, 2, mask=-1, mask_and_clamp=SHFL_W8_CLAMP) + if row_in_prefix >= 2: + val = val + tmp + tmp = cute.arch.shuffle_sync_up(val, 4, mask=-1, mask_and_clamp=SHFL_W8_CLAMP) + if row_in_prefix >= 4: + val = val + tmp + sPartialLast[row_in_prefix, col] = val + + k1_internal_barrier() + + for vi in cutlass.range_constexpr(VEC): + rGkLast[vi] = sPartialLast[K1_ROW_GROUPS - 1, col_base + vi] + + for vi in cutlass.range_constexpr(VEC): + rPrefix[vi] = cutlass.Float32(0.0) + if warp_row_group > 0: + for vi in cutlass.range_constexpr(VEC): + rPrefix[vi] = sPartialLast[warp_row_group - 1, col_base + vi] + + # ---- Pass 2a: ONLY cumsum + write csGcum (critical path, minimal work) ---- + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = rPrefix[vi] + + for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP): + row = k1_row_start + ri + for vi in cutlass.range_constexpr(VEC): + rAcc[vi] = rAcc[vi] + rGact[ri, vi] + cs = rAcc[vi] * cumsum_scale + rGact[ri, vi] = cs + csGcum[row, col_base + vi] = cs + + # Signal MMA early: csGcum is ready + cute.arch.mbarrier_arrive(k1_done_mbars + cur_stage) + + for vi in cutlass.range_constexpr(VEC): + rGkLast[vi] = rGkLast[vi] * cumsum_scale + + if cutlass.const_expr(INPUT_GK_FP32 == 0): + # ---- Pass 2b: write GMEM (overlaps with MMA, off critical path) ---- + for ri in cutlass.range_constexpr(ROWS_PER_K1_WARP): + row = k1_row_start + ri + t = chunk_start + row + if cutlass.const_expr(PREPROCESS_W_BETA != 0): + beta_for_w = csBeta[row].to(cutlass.Float32) + + sK_tile = cute.local_tile(csK, tiler=(1, K1_COLS_PER_WARP), coord=(row, warp_col_group)) + tCsK = thr_copy_k1.partition_S(sK_tile) + tCrK = cute.make_fragment_like(tCsK) + cute.copy(tiled_copy_qk_k1, tCsK, thr_copy_k1.retile(tCrK)) + + sQ_tile = cute.local_tile(csQ, tiler=(1, K1_COLS_PER_WARP), coord=(row, warp_col_group)) + tCsQ = thr_copy_k1.partition_S(sQ_tile) + tCrQ = cute.make_fragment_like(tCsQ) + cute.copy(tiled_copy_qk_k1, tCsQ, thr_copy_k1.retile(tCrQ)) + + for vi in cutlass.range_constexpr(VEC): + cs = rGact[ri, vi] + + k_val = tCrK[vi].to(cutlass.Float32) + q_val = tCrQ[vi].to(cutlass.Float32) + + exp2_cs = cute.exp2(cs, fastmath=True) + exp2_kg = cute.exp2(rGkLast[vi] - cs, fastmath=True) + + if cutlass.const_expr(PREPROCESS_W_BETA != 0): + # Match the csrc W operand association exactly: + # bf16((k * beta) * exp2(gk)). WU consumes this + # value directly and must not multiply beta again. + rKsOut[vi] = ((k_val * beta_for_w) * exp2_cs).to(cutlass.BFloat16) + elif cutlass.const_expr(KSCALED_FP16 != 0): + rKsOut[vi] = (k_val * exp2_cs).to(cutlass.Float16) + else: + rKsOut[vi] = (k_val * exp2_cs).to(cutlass.BFloat16) + rQsOut[vi] = (q_val * exp2_cs).to(cutlass.BFloat16) + rKgOut[vi] = (k_val * exp2_kg).to(cutlass.BFloat16) + + if IS_VARLEN and not VARLEN_PURE: + if t < ci_eos: + cute.autovec_copy(rKsOut, mKscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rQsOut, mQscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rKgOut, mKg[i_b, t, i_h, col_vec_idx, None]) + else: + cute.autovec_copy(rKsOut, mKscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rQsOut, mQscaled[i_b, t, i_h, col_vec_idx, None]) + cute.autovec_copy(rKgOut, mKg[i_b, t, i_h, col_vec_idx, None]) + + if warp_row_group == 0: + for vi in cutlass.range_constexpr(VEC): + rGkOut[vi] = cute.exp2(rGkLast[vi], fastmath=True) + if IS_VARLEN: + if ci_eos > cutlass.Int32(0): + cute.autovec_copy(rGkOut, mGkLast[i_b, chunk_idx, i_h, col_vec_idx, None]) + else: + cute.autovec_copy(rGkOut, mGkLast[i_b, chunk_idx, i_h, col_vec_idx, None]) + + # ============================================================= + # Warp 26 (TMA_WARP_ID): dedicated TMA producer. + # Waits stage_reuse (gated by MMA arrives), issues TMA for Q/K/G, + # signals tma_mbar. Decouples MMA -> TMA dependency from K1 compute. + # ============================================================= + if warp_idx == TMA_WARP_ID: + gQ_head = tma_tensor_Q[(None, None, i_h)] + gK_head = tma_tensor_K[(None, None, i_h)] + gG_head = tma_tensor_G[(None, None, i_h)] + + # Prefetch chunk 0 -> stage 0 (stage_reuse[0] pre-arrived) + pf_cs = cutlass.Int32(0) + if IS_VARLEN: + pf_seq_id_0 = cutlass.Int32(mChunkIndices[chunk_base, 0]) + pf_local_0 = cutlass.Int32(mChunkIndices[chunk_base, 1]) + pf_bos_0 = cutlass.Int32(mCuSeqlens[pf_seq_id_0]) + pf_cs = pf_bos_0 + pf_local_0 * BT + else: + pf_cs = i_b * num_chunks * BT + chunk_base * BT + cute.arch.mbarrier_wait(stage_reuse_mbars, 0) + if lane_id == 0: + cute.arch.mbarrier_expect_tx(tma_mbars, bytes_per_stage) + sQ_pf = sQ[(None, None, 0)] + gQ_pf = cute.local_tile(cute.domain_offset((pf_cs, 0), gQ_head), (BT, K_DIM), (0, 0)) + ts_pf, tg_pf = cpasync.tma_partition( + tma_atom_Q, 0, cute.make_layout(1), cute.group_modes(sQ_pf, 0, 2), cute.group_modes(gQ_pf, 0, 2) + ) + cute.copy(tma_atom_Q, tg_pf, ts_pf, tma_bar_ptr=tma_mbars) + sK_pf = sK[(None, None, 0)] + gK_pf = cute.local_tile(cute.domain_offset((pf_cs, 0), gK_head), (BT, K_DIM), (0, 0)) + ts_pf, tg_pf = cpasync.tma_partition( + tma_atom_K, 0, cute.make_layout(1), cute.group_modes(sK_pf, 0, 2), cute.group_modes(gK_pf, 0, 2) + ) + cute.copy(tma_atom_K, tg_pf, ts_pf, tma_bar_ptr=tma_mbars) + if cutlass.const_expr(INPUT_GK_FP32 != 0): + sG_pf = sGcum[(None, None, 0)] + else: + sG_pf = sG[(None, None, 0)] + gG_pf = cute.local_tile(cute.domain_offset((pf_cs, 0), gG_head), (BT, K_DIM), (0, 0)) + ts_pf, tg_pf = cpasync.tma_partition( + tma_atom_G, 0, cute.make_layout(1), cute.group_modes(sG_pf, 0, 2), cute.group_modes(gG_pf, 0, 2) + ) + cute.copy(tma_atom_G, tg_pf, ts_pf, tma_bar_ptr=tma_mbars) + if lane_id == 0: + cute.arch.mbarrier_arrive(tma_mbars) + + # Issue TMAs for chunks 1..CHUNKS_PER_BLOCK-1 + for next_i in cutlass.range_constexpr(1, CHUNKS_PER_BLOCK): + next_stage = next_i % NUM_STAGES + next_phase = next_i // NUM_STAGES % 2 + next_cs = cutlass.Int32(0) + if IS_VARLEN: + next_chunk_idx = chunk_base + next_i + if next_chunk_idx < num_chunks: + _nsid = cutlass.Int32(mChunkIndices[next_chunk_idx, 0]) + next_cs = cutlass.Int32(mCuSeqlens[_nsid]) + cutlass.Int32(mChunkIndices[next_chunk_idx, 1]) * BT + else: + next_cs = i_b * num_chunks * BT + (chunk_base + next_i) * BT + cute.arch.mbarrier_wait(stage_reuse_mbars + next_stage, next_phase) + if lane_id == 0: + cute.arch.mbarrier_expect_tx(tma_mbars + next_stage, bytes_per_stage) + sQ_ns = sQ[(None, None, next_stage)] + gQ_ns = cute.local_tile(cute.domain_offset((next_cs, 0), gQ_head), (BT, K_DIM), (0, 0)) + ts_ns, tg_ns = cpasync.tma_partition( + tma_atom_Q, 0, cute.make_layout(1), cute.group_modes(sQ_ns, 0, 2), cute.group_modes(gQ_ns, 0, 2) + ) + cute.copy(tma_atom_Q, tg_ns, ts_ns, tma_bar_ptr=tma_mbars + next_stage) + sK_ns = sK[(None, None, next_stage)] + gK_ns = cute.local_tile(cute.domain_offset((next_cs, 0), gK_head), (BT, K_DIM), (0, 0)) + ts_ns, tg_ns = cpasync.tma_partition( + tma_atom_K, 0, cute.make_layout(1), cute.group_modes(sK_ns, 0, 2), cute.group_modes(gK_ns, 0, 2) + ) + cute.copy(tma_atom_K, tg_ns, ts_ns, tma_bar_ptr=tma_mbars + next_stage) + if cutlass.const_expr(INPUT_GK_FP32 != 0): + sG_ns = sGcum[(None, None, next_stage)] + else: + sG_ns = sG[(None, None, next_stage)] + gG_ns = cute.local_tile(cute.domain_offset((next_cs, 0), gG_head), (BT, K_DIM), (0, 0)) + ts_ns, tg_ns = cpasync.tma_partition( + tma_atom_G, 0, cute.make_layout(1), cute.group_modes(sG_ns, 0, 2), cute.group_modes(gG_ns, 0, 2) + ) + cute.copy(tma_atom_G, tg_ns, ts_ns, tma_bar_ptr=tma_mbars + next_stage) + if lane_id == 0: + cute.arch.mbarrier_arrive(tma_mbars + next_stage) + + # ============================================================= + # Warps 16-25: K2/K3 MMA compute + # ============================================================= + if warp_idx >= NUM_K1_TMA_WARPS and warp_idx < NUM_K1_TMA_WARPS + NUM_MMA_WARPS and warp_idx != TMA_WARP_ID: + # Warp-layout invariants (scope-local → no cross-group register spill) + _tid_in_group = _lane % 4 + _group_id = _lane // 4 + mma_warp = _warp - NUM_K1_TMA_WARPS + my_i_q = cutlass.Int32(0) + my_i_k = cutlass.Int32(0) + if mma_warp < 1: + my_i_q = cutlass.Int32(0) + my_i_k = mma_warp + elif mma_warp < 3: + my_i_q = cutlass.Int32(1) + my_i_k = mma_warp - 1 + elif mma_warp < 6: + my_i_q = cutlass.Int32(2) + my_i_k = mma_warp - 3 + elif mma_warp < NUM_MMA_ACTIVE: + my_i_q = cutlass.Int32(3) + my_i_k = mma_warp - 6 + q_row_base = my_i_q * BC + k_row_base = my_i_k * BC + akk_row_base = k_row_base + akk_col_base = q_row_base + norm_row = q_row_base + if my_i_q == my_i_k: + norm_row = q_row_base + cutlass.Int32(BC // 2) + row0 = _group_id + row1 = _group_id + 8 + col0 = _tid_in_group * 2 + col1 = _tid_in_group * 2 + 1 + col2 = 8 + _tid_in_group * 2 + col3 = 8 + _tid_in_group * 2 + 1 + thr_mma = tiled_mma_k2.get_slice(_lane) + thr_copy_A = tiled_copy_mma_A.get_slice(_lane) + thr_copy_B = tiled_copy_mma_B.get_slice(_lane) + thr_copy_Gn = tiled_copy_Gcum_norm.get_slice(_tid_in_group) + thr_copy_Ggate = tiled_copy_Gcum_gate.get_slice(_lane) + + for chunk_iter in cutlass.range_constexpr(CHUNKS_PER_BLOCK): + s = chunk_iter % NUM_STAGES + phase = chunk_iter // NUM_STAGES % 2 + chunk_idx = chunk_base + chunk_iter + chunk_start = cutlass.Int32(0) + if IS_VARLEN: + if chunk_idx < num_chunks: + _sid = cutlass.Int32(mChunkIndices[chunk_idx, 0]) + chunk_start = cutlass.Int32(mCuSeqlens[_sid]) + cutlass.Int32(mChunkIndices[chunk_idx, 1]) * BT + else: + chunk_start = chunk_idx * BT + + cute.arch.mbarrier_wait(k1_done_mbars + s, phase) + cute.arch.mbarrier_wait(store_done_mbars + s, phase) + + if mma_warp < NUM_MMA_ACTIVE: + csQ = sQ[(None, None, s)] + csK = sK[(None, None, s)] + csGcum = sGcum[(None, None, s)] + csAqk = sAqk[(None, None, s)] + csAkk = sAkk[(None, None, s)] + csBeta = sBeta[(None, s)] + + _z = cutlass.Float32(0.0) + + beta_row0 = csBeta[q_row_base + row0].to(cutlass.Float32) + beta_row1 = csBeta[q_row_base + row1].to(cutlass.Float32) + + acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = _z, _z, _z, _z + acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = _z, _z, _z, _z + acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = _z, _z, _z, _z + acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = _z, _z, _z, _z + + # bf16 m16n8k16 MMA: each k_block covers k=16. + for k_block in cutlass.range_constexpr(NUM_MMA_K_TILES): + # ---- Load Q/Kq bf16 fragments (16x16, 8 bf16/thread) ---- + sQ_tile = cute.local_tile(csQ, tiler=(16, 16), coord=(my_i_q, k_block)) + tCrQ = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sQ_tile)) + cute.copy(tiled_copy_mma_A, thr_copy_A.partition_S(sQ_tile), thr_copy_A.retile(tCrQ)) + + sKq_tile = cute.local_tile(csK, tiler=(16, 16), coord=(my_i_q, k_block)) + tCrKq = tiled_mma_k2.make_fragment_A(thr_mma.partition_A(sKq_tile)) + cute.copy(tiled_copy_mma_A, thr_copy_A.partition_S(sKq_tile), thr_copy_A.retile(tCrKq)) + + # ---- Issue K n0/n1 LDSMs early for better ILP ---- + sK_tile_n0 = cute.local_tile(csK, tiler=(8, 16), coord=(my_i_k * 2, k_block)) + tCrK_n0 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n0)) + cute.copy(tiled_copy_mma_B, thr_copy_B.partition_S(sK_tile_n0), thr_copy_B.retile(tCrK_n0)) + + sK_tile_n1 = cute.local_tile(csK, tiler=(8, 16), coord=(my_i_k * 2 + 1, k_block)) + tCrK_n1 = tiled_mma_k2.make_fragment_B(thr_mma.partition_B(sK_tile_n1)) + cute.copy(tiled_copy_mma_B, thr_copy_B.partition_S(sK_tile_n1), thr_copy_B.retile(tCrK_n1)) + + # ---- Gate norm (2x k=8 covers k=16) ---- + sGn_a = cute.local_tile(csGcum, tiler=(1, 8), coord=(norm_row, k_block * 2)) + tCsGn_a = thr_copy_Gn.partition_S(sGn_a) + tCrGn_a = cute.make_fragment_like(tCsGn_a, cutlass.Float32) + cute.copy(tiled_copy_Gcum_norm, tCsGn_a, thr_copy_Gn.retile(tCrGn_a)) + gn_a0 = tCrGn_a[0] + gn_a1 = tCrGn_a[1] + + sGn_b = cute.local_tile(csGcum, tiler=(1, 8), coord=(norm_row, k_block * 2 + 1)) + tCsGn_b = thr_copy_Gn.partition_S(sGn_b) + tCrGn_b = cute.make_fragment_like(tCsGn_b, cutlass.Float32) + cute.copy(tiled_copy_Gcum_norm, tCsGn_b, thr_copy_Gn.retile(tCrGn_b)) + gn_b0 = tCrGn_b[0] + gn_b1 = tCrGn_b[1] + + # ---- Gate Q (2x (16,8) partition_C covers m=16,k=16) ---- + sGq_a = cute.local_tile(csGcum, tiler=(16, 8), coord=(my_i_q, k_block * 2)) + tCrGq_a = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGq_a)) + cute.copy(tiled_copy_Gcum_gate, thr_copy_Ggate.partition_S(sGq_a), thr_copy_Ggate.retile(tCrGq_a)) + + sGq_b = cute.local_tile(csGcum, tiler=(16, 8), coord=(my_i_q, k_block * 2 + 1)) + tCrGq_b = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGq_b)) + cute.copy(tiled_copy_Gcum_gate, thr_copy_Ggate.partition_S(sGq_b), thr_copy_Ggate.retile(tCrGq_b)) + + # 8 Q gate values per thread (matching A bf16 m16n8k16 layout): + # first half k=0..7 (tCrGq_a): a0=(r0,c0) a1=(r0,c0+1) a2=(r0+8,c0) a3=(r0+8,c0+1) + # second half k=8..15 (tCrGq_b): a4=(r0,c0+8) a5=(r0,c0+9) a6=(r0+8,c0+8) a7=(r0+8,c0+9) + gate_q_0 = cute.exp2(tCrGq_a[0] - gn_a0, fastmath=True) + gate_q_1 = cute.exp2(tCrGq_a[1] - gn_a1, fastmath=True) + gate_q_2 = cute.exp2(tCrGq_a[2] - gn_a0, fastmath=True) + gate_q_3 = cute.exp2(tCrGq_a[3] - gn_a1, fastmath=True) + gate_q_4 = cute.exp2(tCrGq_b[0] - gn_b0, fastmath=True) + gate_q_5 = cute.exp2(tCrGq_b[1] - gn_b1, fastmath=True) + gate_q_6 = cute.exp2(tCrGq_b[2] - gn_b0, fastmath=True) + gate_q_7 = cute.exp2(tCrGq_b[3] - gn_b1, fastmath=True) + + # qa fp32 = Q*gate (8 per thread). tCrQ indexing assumption: + # [0..3] first k-chunk (k=0..7), [4..7] second k-chunk (k=8..15). + qa0 = tCrQ[0].to(cutlass.Float32) * gate_q_0 + qa1 = tCrQ[1].to(cutlass.Float32) * gate_q_1 + qa2 = tCrQ[2].to(cutlass.Float32) * gate_q_2 + qa3 = tCrQ[3].to(cutlass.Float32) * gate_q_3 + qa4 = tCrQ[4].to(cutlass.Float32) * gate_q_4 + qa5 = tCrQ[5].to(cutlass.Float32) * gate_q_5 + qa6 = tCrQ[6].to(cutlass.Float32) * gate_q_6 + qa7 = tCrQ[7].to(cutlass.Float32) * gate_q_7 + + ka0 = tCrKq[0].to(cutlass.Float32) * gate_q_0 + ka1 = tCrKq[1].to(cutlass.Float32) * gate_q_1 + ka2 = tCrKq[2].to(cutlass.Float32) * gate_q_2 + ka3 = tCrKq[3].to(cutlass.Float32) * gate_q_3 + ka4 = tCrKq[4].to(cutlass.Float32) * gate_q_4 + ka5 = tCrKq[5].to(cutlass.Float32) * gate_q_5 + ka6 = tCrKq[6].to(cutlass.Float32) * gate_q_6 + ka7 = tCrKq[7].to(cutlass.Float32) * gate_q_7 + + # ---- Gate K (2x (16,8) partition_C covers m=16,k=16) ---- + sGk_a = cute.local_tile(csGcum, tiler=(16, 8), coord=(my_i_k, k_block * 2)) + tCrGk_a = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGk_a)) + cute.copy(tiled_copy_Gcum_gate, thr_copy_Ggate.partition_S(sGk_a), thr_copy_Ggate.retile(tCrGk_a)) + + sGk_b = cute.local_tile(csGcum, tiler=(16, 8), coord=(my_i_k, k_block * 2 + 1)) + tCrGk_b = tiled_mma_k2.make_fragment_C(thr_mma.partition_C(sGk_b)) + cute.copy(tiled_copy_Gcum_gate, thr_copy_Ggate.partition_S(sGk_b), thr_copy_Ggate.retile(tCrGk_b)) + + # n0 uses rows 0..7 of (16,*) tile (tCrGk_*[0,1]) + # n1 uses rows 8..15 of (16,*) tile (tCrGk_*[2,3]) + gk_n0_0 = cute.exp2(gn_a0 - tCrGk_a[0], fastmath=True) + gk_n0_1 = cute.exp2(gn_a1 - tCrGk_a[1], fastmath=True) + gk_n0_2 = cute.exp2(gn_b0 - tCrGk_b[0], fastmath=True) + gk_n0_3 = cute.exp2(gn_b1 - tCrGk_b[1], fastmath=True) + + gk_n1_0 = cute.exp2(gn_a0 - tCrGk_a[2], fastmath=True) + gk_n1_1 = cute.exp2(gn_a1 - tCrGk_a[3], fastmath=True) + gk_n1_2 = cute.exp2(gn_b0 - tCrGk_b[2], fastmath=True) + gk_n1_3 = cute.exp2(gn_b1 - tCrGk_b[3], fastmath=True) + + # tCrK_n0 bf16 fragment: 4 elems/thread at (n_row, c0), (n_row, c0+1), + # (n_row, c0+8), (n_row, c0+9) — k-adjacent pairs + k_n0_b0 = tCrK_n0[0].to(cutlass.Float32) * gk_n0_0 + k_n0_b1 = tCrK_n0[1].to(cutlass.Float32) * gk_n0_1 + k_n0_b2 = tCrK_n0[2].to(cutlass.Float32) * gk_n0_2 + k_n0_b3 = tCrK_n0[3].to(cutlass.Float32) * gk_n0_3 + + k_n1_b0 = tCrK_n1[0].to(cutlass.Float32) * gk_n1_0 + k_n1_b1 = tCrK_n1[1].to(cutlass.Float32) * gk_n1_1 + k_n1_b2 = tCrK_n1[2].to(cutlass.Float32) * gk_n1_2 + k_n1_b3 = tCrK_n1[3].to(cutlass.Float32) * gk_n1_3 + + # TF32 m16n8k8 consumes one k8 half per call. The + # loaded BF16-fragment value order is (r0,k0), + # (r0,k1), (r8,k0), (r8,k1), so reorder A to the + # register MMA order used by kda_fuse_k4.py: + # (r0,k0), (r8,k0), (r0,k1), (r8,k1). + acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = mma_tf32_m16n8k8( + qa0, qa2, qa1, qa3, k_n0_b0, k_n0_b1, acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 + ) + acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 = mma_tf32_m16n8k8( + qa4, qa6, qa5, qa7, k_n0_b2, k_n0_b3, acc_aqk_n0_0, acc_aqk_n0_1, acc_aqk_n0_2, acc_aqk_n0_3 + ) + + acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = mma_tf32_m16n8k8( + qa0, qa2, qa1, qa3, k_n1_b0, k_n1_b1, acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 + ) + acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 = mma_tf32_m16n8k8( + qa4, qa6, qa5, qa7, k_n1_b2, k_n1_b3, acc_aqk_n1_0, acc_aqk_n1_1, acc_aqk_n1_2, acc_aqk_n1_3 + ) + + acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = mma_tf32_m16n8k8( + ka0, ka2, ka1, ka3, k_n0_b0, k_n0_b1, acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 + ) + acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 = mma_tf32_m16n8k8( + ka4, ka6, ka5, ka7, k_n0_b2, k_n0_b3, acc_akk_n0_0, acc_akk_n0_1, acc_akk_n0_2, acc_akk_n0_3 + ) + + acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = mma_tf32_m16n8k8( + ka0, ka2, ka1, ka3, k_n1_b0, k_n1_b1, acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 + ) + acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 = mma_tf32_m16n8k8( + ka4, ka6, ka5, ka7, k_n1_b2, k_n1_b3, acc_akk_n1_0, acc_akk_n1_1, acc_akk_n1_2, acc_akk_n1_3 + ) + + # sQ/sK/sG reads done, signal TMA before SMEM writes + cute.arch.mbarrier_arrive(stage_reuse_mbars + s) + + # Dual-path MMA write (constexpr-gated): + # - non-pure: apply causal + diag=1 inline so SMEM matches + # final GMEM layout (pairs with row-major vec autovec + # store warp that does no causal). + # - pure: write all 16x16 unconditionally (pairs with + # per-tile store warp that applies causal + diag in + # store; this is the original baseline behavior — no + # extra MMA-write cost). + _z16 = cutlass.BFloat16(0.0) + _one16 = cutlass.BFloat16(1.0) + _z32 = cutlass.Float32(0.0) + _one32 = cutlass.Float32(1.0) + if IS_VARLEN and not VARLEN_PURE: + if my_i_q == my_i_k: + # Diagonal sub-tile: causal-mask sAqk and write + # diag=1 / strict-lower=MMA*beta / strict-upper=0 + # to sAkk so SMEM is in final form. + _v_q00 = (acc_aqk_n0_0 * scale).to(cutlass.BFloat16) + if row0 < col0: + _v_q00 = _z16 + csAqk[q_row_base + row0, k_row_base + col0] = _v_q00 + _v_q01 = (acc_aqk_n0_1 * scale).to(cutlass.BFloat16) + if row0 < col1: + _v_q01 = _z16 + csAqk[q_row_base + row0, k_row_base + col1] = _v_q01 + _v_q02 = (acc_aqk_n0_2 * scale).to(cutlass.BFloat16) + if row1 < col0: + _v_q02 = _z16 + csAqk[q_row_base + row1, k_row_base + col0] = _v_q02 + _v_q03 = (acc_aqk_n0_3 * scale).to(cutlass.BFloat16) + if row1 < col1: + _v_q03 = _z16 + csAqk[q_row_base + row1, k_row_base + col1] = _v_q03 + _v_q04 = (acc_aqk_n1_0 * scale).to(cutlass.BFloat16) + if row0 < col2: + _v_q04 = _z16 + csAqk[q_row_base + row0, k_row_base + col2] = _v_q04 + _v_q05 = (acc_aqk_n1_1 * scale).to(cutlass.BFloat16) + if row0 < col3: + _v_q05 = _z16 + csAqk[q_row_base + row0, k_row_base + col3] = _v_q05 + _v_q06 = (acc_aqk_n1_2 * scale).to(cutlass.BFloat16) + if row1 < col2: + _v_q06 = _z16 + csAqk[q_row_base + row1, k_row_base + col2] = _v_q06 + _v_q07 = (acc_aqk_n1_3 * scale).to(cutlass.BFloat16) + if row1 < col3: + _v_q07 = _z16 + csAqk[q_row_base + row1, k_row_base + col3] = _v_q07 + + if cutlass.const_expr(FP32_AKK_WORKSPACE != 0 or FUSE_AKK_INV != 0): + _v_k00 = acc_akk_n0_0 * beta_row0 + if row0 == col0: + _v_k00 = _one32 + if row0 < col0: + _v_k00 = _z32 + csAkk[akk_row_base + row0, akk_col_base + col0] = _v_k00 + else: + _v_k00 = (acc_akk_n0_0 * beta_row0).to(cutlass.BFloat16) + if row0 == col0: + _v_k00 = _one16 + if row0 < col0: + _v_k00 = _z16 + csAkk[akk_row_base + row0, akk_col_base + col0] = _v_k00.to(cutlass.Float32) # BF16 -> FP32 + _v_k01 = acc_akk_n0_1 * beta_row0 # Keep FP32 + if row0 == col1: + _v_k01 = _one32 + if row0 < col1: + _v_k01 = _z32 + csAkk[akk_row_base + row0, akk_col_base + col1] = _v_k01 + _v_k02 = acc_akk_n0_2 * beta_row1 # Keep FP32 + if row1 == col0: + _v_k02 = _one32 + if row1 < col0: + _v_k02 = _z32 + csAkk[akk_row_base + row1, akk_col_base + col0] = _v_k02 + _v_k03 = acc_akk_n0_3 * beta_row1 # Keep FP32 + if row1 == col1: + _v_k03 = _one32 + if row1 < col1: + _v_k03 = _z32 + csAkk[akk_row_base + row1, akk_col_base + col1] = _v_k03 + _v_k04 = acc_akk_n1_0 * beta_row0 # Keep FP32 + if row0 == col2: + _v_k04 = _one32 + if row0 < col2: + _v_k04 = _z32 + csAkk[akk_row_base + row0, akk_col_base + col2] = _v_k04 + _v_k05 = acc_akk_n1_1 * beta_row0 # Keep FP32 + if row0 == col3: + _v_k05 = _one32 + if row0 < col3: + _v_k05 = _z32 + csAkk[akk_row_base + row0, akk_col_base + col3] = _v_k05 + _v_k06 = acc_akk_n1_2 * beta_row1 # Keep FP32 + if row1 == col2: + _v_k06 = _one32 + if row1 < col2: + _v_k06 = _z32 + csAkk[akk_row_base + row1, akk_col_base + col2] = _v_k06 + _v_k07 = acc_akk_n1_3 * beta_row1 # Keep FP32 + if row1 == col3: + _v_k07 = _one32 + if row1 < col3: + _v_k07 = _z32 + csAkk[akk_row_base + row1, akk_col_base + col3] = _v_k07 + else: + # Non-diag (i_q > i_k): write all 16x16 unchanged. + csAqk[q_row_base + row0, k_row_base + col0] = (acc_aqk_n0_0 * scale).to(cutlass.BFloat16) + csAqk[q_row_base + row0, k_row_base + col1] = (acc_aqk_n0_1 * scale).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col0] = (acc_aqk_n0_2 * scale).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col1] = (acc_aqk_n0_3 * scale).to(cutlass.BFloat16) + csAqk[q_row_base + row0, k_row_base + col2] = (acc_aqk_n1_0 * scale).to(cutlass.BFloat16) + csAqk[q_row_base + row0, k_row_base + col3] = (acc_aqk_n1_1 * scale).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col2] = (acc_aqk_n1_2 * scale).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col3] = (acc_aqk_n1_3 * scale).to(cutlass.BFloat16) + if cutlass.const_expr(FP32_AKK_WORKSPACE != 0 or FUSE_AKK_INV != 0): + csAkk[akk_row_base + row0, akk_col_base + col0] = acc_akk_n0_0 * beta_row0 + csAkk[akk_row_base + row0, akk_col_base + col1] = acc_akk_n0_1 * beta_row0 + csAkk[akk_row_base + row1, akk_col_base + col0] = acc_akk_n0_2 * beta_row1 + csAkk[akk_row_base + row1, akk_col_base + col1] = acc_akk_n0_3 * beta_row1 + csAkk[akk_row_base + row0, akk_col_base + col2] = acc_akk_n1_0 * beta_row0 + csAkk[akk_row_base + row0, akk_col_base + col3] = acc_akk_n1_1 * beta_row0 + csAkk[akk_row_base + row1, akk_col_base + col2] = acc_akk_n1_2 * beta_row1 + csAkk[akk_row_base + row1, akk_col_base + col3] = acc_akk_n1_3 * beta_row1 + else: + csAkk[akk_row_base + row0, akk_col_base + col0] = (acc_akk_n0_0 * beta_row0).to( + cutlass.BFloat16 + ) + csAkk[akk_row_base + row0, akk_col_base + col1] = (acc_akk_n0_1 * beta_row0).to( + cutlass.BFloat16 + ) + csAkk[akk_row_base + row1, akk_col_base + col0] = (acc_akk_n0_2 * beta_row1).to( + cutlass.BFloat16 + ) + csAkk[akk_row_base + row1, akk_col_base + col1] = (acc_akk_n0_3 * beta_row1).to( + cutlass.BFloat16 + ) + csAkk[akk_row_base + row0, akk_col_base + col2] = (acc_akk_n1_0 * beta_row0).to( + cutlass.BFloat16 + ) + csAkk[akk_row_base + row0, akk_col_base + col3] = (acc_akk_n1_1 * beta_row0).to( + cutlass.BFloat16 + ) + csAkk[akk_row_base + row1, akk_col_base + col2] = (acc_akk_n1_2 * beta_row1).to( + cutlass.BFloat16 + ) + csAkk[akk_row_base + row1, akk_col_base + col3] = (acc_akk_n1_3 * beta_row1).to( + cutlass.BFloat16 + ) + else: + # PURE: write all 16x16 unconditionally — store warp + # applies causal+diag in its per-tile loop. This is + # the original baseline behavior (no extra MMA cost). + csAqk[q_row_base + row0, k_row_base + col0] = (acc_aqk_n0_0 * scale).to(cutlass.BFloat16) + csAqk[q_row_base + row0, k_row_base + col1] = (acc_aqk_n0_1 * scale).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col0] = (acc_aqk_n0_2 * scale).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col1] = (acc_aqk_n0_3 * scale).to(cutlass.BFloat16) + csAqk[q_row_base + row0, k_row_base + col2] = (acc_aqk_n1_0 * scale).to(cutlass.BFloat16) + csAqk[q_row_base + row0, k_row_base + col3] = (acc_aqk_n1_1 * scale).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col2] = (acc_aqk_n1_2 * scale).to(cutlass.BFloat16) + csAqk[q_row_base + row1, k_row_base + col3] = (acc_aqk_n1_3 * scale).to(cutlass.BFloat16) + csAkk[akk_row_base + row0, akk_col_base + col0] = acc_akk_n0_0 * beta_row0 # Keep FP32 + csAkk[akk_row_base + row0, akk_col_base + col1] = acc_akk_n0_1 * beta_row0 # Keep FP32 + csAkk[akk_row_base + row1, akk_col_base + col0] = acc_akk_n0_2 * beta_row1 # Keep FP32 + csAkk[akk_row_base + row1, akk_col_base + col1] = acc_akk_n0_3 * beta_row1 # Keep FP32 + csAkk[akk_row_base + row0, akk_col_base + col2] = acc_akk_n1_0 * beta_row0 # Keep FP32 + csAkk[akk_row_base + row0, akk_col_base + col3] = acc_akk_n1_1 * beta_row0 # Keep FP32 + csAkk[akk_row_base + row1, akk_col_base + col2] = acc_akk_n1_2 * beta_row1 # Keep FP32 + csAkk[akk_row_base + row1, akk_col_base + col3] = acc_akk_n1_3 * beta_row1 # Keep FP32 + else: + cute.arch.mbarrier_arrive(stage_reuse_mbars + s) + + cute.arch.mbarrier_arrive(mma_done_mbars + s) + + # ============================================================= + # Warps 27-30: Store warps / optional in-CTA inverse workers + # ============================================================= + if warp_idx >= STORE_WARP_BASE: + store_warp = warp_idx - STORE_WARP_BASE + for chunk_iter in cutlass.range_constexpr(CHUNKS_PER_BLOCK): + s = chunk_iter % NUM_STAGES + phase = chunk_iter // NUM_STAGES % 2 + chunk_idx = chunk_base + chunk_iter + chunk_start = cutlass.Int32(0) + st_eos = cutlass.Int32(0) + if IS_VARLEN: + if chunk_idx < num_chunks: + _sid = cutlass.Int32(mChunkIndices[chunk_idx, 0]) + chunk_start = cutlass.Int32(mCuSeqlens[_sid]) + cutlass.Int32(mChunkIndices[chunk_idx, 1]) * BT + st_eos = cutlass.Int32(mCuSeqlens[_sid + 1]) + else: + chunk_start = chunk_idx * BT + st_eos = chunk_start + BT + + cute.arch.mbarrier_wait(mma_done_mbars + s, phase) + + csAqk = sAqk[(None, None, s)] + csAkk = sAkk[(None, None, s)] + + # Dual-path row-major store. SMEM is in final GMEM layout + # already (causal mask + diag=1 applied at MMA write; upper-tri + # SMEM is zero from CTA-startup init). Both paths use vec2 + # autovec_copy → STG.E.32 (4-byte coalesced). + # + # mAqk_v2 / mAkk_v2 shape (B, T, H, BT/2, 2): last dim is vec2. + # + # NON-PURE: full 64-col row-major + row mask (chunk may + # overflow seq end). 32 lanes/warp × 1 vec2 = full row. + # PURE: reduced cols. Warp s writes (s+1)*16 cols/row (no row + # mask, all rows in seq). Saves ~37% GMEM bandwidth vs full. + col_lo = lane_id * 2 + col_hi = col_lo + 1 + col_vec_idx = lane_id + rAqkOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + if cutlass.const_expr(FP32_AKK_WORKSPACE != 0): + rAkkFp32Out = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.Float32) + else: + rAkkOut = cute.make_rmem_tensor(cute.make_layout((VEC,)), cutlass.BFloat16) + + if cutlass.const_expr(INPUT_GK_FP32 != 0 or (IS_VARLEN and not VARLEN_PURE)): + # NON-PURE: row-major full-row vec autovec + row mask. SMEM + # upper-tri is zero (MMA-masked at write), so writing all 64 + # cols is correct. STG.E.32 (32 lanes × 4 bytes per row). + row_base_warp = store_warp * (BT // NUM_STORE_WARPS) + for ri in cutlass.range_constexpr(BT // NUM_STORE_WARPS): + full_row = row_base_warp + ri + full_abs_row = chunk_start + full_row + rAqkOut[0] = csAqk[full_row, col_lo] + rAqkOut[1] = csAqk[full_row, col_hi] + if cutlass.const_expr(INPUT_GK_FP32 != 0): + if full_row < col_lo: + rAqkOut[0] = cutlass.BFloat16(0.0) + if full_row < col_hi: + rAqkOut[1] = cutlass.BFloat16(0.0) + if full_abs_row < st_eos: + cute.autovec_copy(rAqkOut, mAqk_v2[i_b, full_abs_row, i_h, col_vec_idx, None]) + if cutlass.const_expr(FUSE_AKK_INV != 0): + pass + elif cutlass.const_expr(FP32_AKK_WORKSPACE != 0): + rAkkFp32Out[0] = cutlass.Float32(csAkk[full_row, col_lo]) + rAkkFp32Out[1] = cutlass.Float32(csAkk[full_row, col_hi]) + cute.autovec_copy(rAkkFp32Out, mAkk_v2[i_b, full_abs_row, i_h, col_vec_idx, None]) + else: + rAkkOut[0] = csAkk[full_row, col_lo].to(cutlass.BFloat16) # FP32 -> BF16 + rAkkOut[1] = csAkk[full_row, col_hi].to(cutlass.BFloat16) + cute.autovec_copy(rAkkOut, mAkk_v2[i_b, full_abs_row, i_h, col_vec_idx, None]) + else: + # PURE: keep the reduced lower-tile bandwidth, but store + # each 16-column row as two vec8 segments. Eight lanes per + # warp cover four rows, reducing the scalar epilogue's + # instruction count without writing the upper six tiles. + rAqkVec = cute.make_rmem_tensor((8,), cutlass.BFloat16) + if cutlass.const_expr(FP32_AKK_WORKSPACE != 0): + rAkkFp32Vec = cute.make_rmem_tensor((8,), cutlass.Float32) + else: + rAkkVec = cute.make_rmem_tensor((8,), cutlass.BFloat16) + for tile_idx in cutlass.range_constexpr(NUM_TILES): + i_q = _TILE_IQ[tile_idx] + i_k = _TILE_IK[tile_idx] + is_diag = _TILE_IQ[tile_idx] == _TILE_IK[tile_idx] + gmem_aqk_row_base = chunk_start + i_q * BC + gmem_aqk_col_base = i_k * BC + gmem_akk_row_base = chunk_start + i_k * BC + gmem_akk_col_base = i_q * BC + smem_aqk_row_base = i_q * BC + smem_aqk_col_base = i_k * BC + smem_akk_row_base = i_k * BC + smem_akk_col_base = i_q * BC + if lane_id < 2 * (BC // NUM_STORE_WARPS): + local_row = store_warp * (BC // NUM_STORE_WARPS) + lane_id // 2 + vec_idx = lane_id % 2 + for elem in cutlass.range_constexpr(8): + local_col = vec_idx * 8 + elem + aqk_val = csAqk[ + smem_aqk_row_base + local_row, + smem_aqk_col_base + local_col, + ] + if is_diag and local_row < local_col: + aqk_val = cutlass.BFloat16(0.0) + rAqkVec[elem] = aqk_val + + if cutlass.const_expr(FUSE_AKK_INV == 0): + akk_val = cutlass.Float32( + csAkk[ + smem_akk_row_base + local_row, + smem_akk_col_base + local_col, + ] + ) + if is_diag and local_row < local_col: + akk_val = cutlass.Float32(0.0) + if is_diag and local_row == local_col: + akk_val = cutlass.Float32(1.0) + if cutlass.const_expr(FP32_AKK_WORKSPACE != 0): + rAkkFp32Vec[elem] = akk_val + else: + rAkkVec[elem] = akk_val.to(cutlass.BFloat16) + + aqk_row = mAqk[i_b, gmem_aqk_row_base + local_row, i_h, None] + aqk_vec = cute.local_tile(aqk_row, (8,), (2 * i_k + vec_idx,)) + cute.autovec_copy(rAqkVec, aqk_vec) + + if cutlass.const_expr(FUSE_AKK_INV == 0): + akk_row = mAkk[i_b, gmem_akk_row_base + local_row, i_h, None] + akk_vec = cute.local_tile(akk_row, (8,), (2 * i_q + vec_idx,)) + if cutlass.const_expr(FP32_AKK_WORKSPACE != 0): + cute.autovec_copy(rAkkFp32Vec, akk_vec) + else: + cute.autovec_copy(rAkkVec, akk_vec) + + if cutlass.const_expr(FUSE_AKK_INV != 0): + # Aqk has reached GMEM, so its current-stage storage can + # safely become the 4x16x16 FP32 Schur workspace. The + # stage remains owned by these four warps until + # store_done, hence the next MMA stage cannot race this + # alias. 64x72 bf16 = 9216 B; the workspace needs 5120 B. + store_internal_barrier() + inv_tmp_layout = cute.make_layout( + (NUM_STORE_WARPS, BC, BC), + stride=(BC * INV_TMP_STRIDE, INV_TMP_STRIDE, 1), + ) + sInvTmp = cute.make_tensor( + cute.recast_ptr(csAqk.iterator, dtype=cutlass.Float32), + inv_tmp_layout, + ) + store_tid = store_warp * 32 + lane_id + + # Convert the block-transposed K123 accumulator into a + # logical unit-lower matrix before the TF32 Schur inverse. + for norm_i in cutlass.range_constexpr((BT * BT) // (NUM_STORE_WARPS * 32)): + linear = store_tid + norm_i * (NUM_STORE_WARPS * 32) + norm_row = linear // BT + norm_col = linear % BT + if norm_row >= norm_col: + norm_val = cutlass.Float32(0.0) + if norm_row == norm_col: + norm_val = cutlass.Float32(1.0) + else: + norm_row_blk = norm_row // BC + norm_col_blk = norm_col // BC + src_row = norm_row + src_col = norm_col + if norm_row_blk != norm_col_blk: + src_row = norm_col_blk * BC + norm_row % BC + src_col = norm_row_blk * BC + norm_col % BC + norm_val = cutlass.Float32(csAkk[src_row, src_col]) + if IS_VARLEN and not VARLEN_PURE: + if chunk_start + norm_row >= st_eos: + norm_val = cutlass.Float32(0.0) + csAkk[norm_row, norm_col] = norm_val + store_internal_barrier() + + if store_tid < 64: + _tf32_invert_diag_forward16(csAkk, store_tid // 16, store_tid) + store_internal_barrier() + + # Invert each diagonal 32x32 block. + if store_tid < 64: + block32 = store_tid // 32 + row_base = block32 * 32 + c0, c1, c2, c3, c4, c5, c6, c7 = _tf32_matmul16_smem_smem( + csAkk, + row_base + 16, + row_base + 16, + csAkk, + row_base + 16, + row_base, + lane_id, + ) + _tf32_store_C16_tmp( + sInvTmp, + block32, + -c0, + -c1, + -c2, + -c3, + -c4, + -c5, + -c6, + -c7, + lane_id, + ) + store_internal_barrier() + + if store_tid < 64: + block32 = store_tid // 32 + row_base = block32 * 32 + c0, c1, c2, c3, c4, c5, c6, c7 = _tf32_matmul16_tmp_smem( + sInvTmp, + block32, + csAkk, + row_base, + row_base, + lane_id, + ) + _tf32_store_C16_smem( + csAkk, + row_base + 16, + row_base, + c0, + c1, + c2, + c3, + c4, + c5, + c6, + c7, + lane_id, + ) + store_internal_barrier() + + # Schur solve for the lower-left 32x32 block. Each warp + # computes one 16x16 quadrant. + inv_x = store_warp // 2 + inv_y = store_warp % 2 + inv_slot = store_warp + inv_row_o = 32 + inv_y * 16 + inv_col_c = inv_x * 16 + + p0, p1, p2, p3, p4, p5, p6, p7 = _tf32_matmul32_smem_smem( + csAkk, inv_row_o, 32, csAkk, 32, inv_col_c, lane_id + ) + _tf32_store_C16_tmp( + sInvTmp, + inv_slot, + -p0, + -p1, + -p2, + -p3, + -p4, + -p5, + -p6, + -p7, + lane_id, + ) + store_internal_barrier() + + o0, o1, o2, o3, o4, o5, o6, o7 = _tf32_matmul16_tmp_smem( + sInvTmp, + inv_slot, + csAkk, + inv_x * 16, + 0, + lane_id, + ) + r0, r1, r2, r3, r4, r5, r6, r7 = _tf32_matmul16_tmp_smem( + sInvTmp, + inv_slot, + csAkk, + inv_x * 16, + 16, + lane_id, + ) + if inv_x == 0: + _tf32_store_C16_smem( + csAkk, + inv_row_o, + 0, + o0, + o1, + o2, + o3, + o4, + o5, + o6, + o7, + lane_id, + ) + _tf32_store_C16_smem( + csAkk, + inv_row_o, + 16, + r0, + r1, + r2, + r3, + r4, + r5, + r6, + r7, + lane_id, + ) + store_internal_barrier() + if inv_x == 1: + _tf32_add_store_C16_smem( + csAkk, + inv_row_o, + 0, + o0, + o1, + o2, + o3, + o4, + o5, + o6, + o7, + lane_id, + ) + _tf32_add_store_C16_smem( + csAkk, + inv_row_o, + 16, + r0, + r1, + r2, + r3, + r4, + r5, + r6, + r7, + lane_id, + ) + store_internal_barrier() + + row_base_warp = store_warp * (BT // NUM_STORE_WARPS) + for ri in cutlass.range_constexpr(BT // NUM_STORE_WARPS): + local_row = row_base_warp + ri + abs_row = chunk_start + local_row + val0 = cutlass.Float32(csAkk[local_row, col_lo]) + val1 = cutlass.Float32(csAkk[local_row, col_hi]) + if local_row < col_lo: + val0 = cutlass.Float32(0.0) + if local_row < col_hi: + val1 = cutlass.Float32(0.0) + rAkkOut[0] = val0.to(cutlass.BFloat16) + rAkkOut[1] = val1.to(cutlass.BFloat16) + if IS_VARLEN and not VARLEN_PURE: + if abs_row < st_eos: + cute.autovec_copy(rAkkOut, mAkk_v2[i_b, abs_row, i_h, col_vec_idx, None]) + else: + cute.autovec_copy(rAkkOut, mAkk_v2[i_b, abs_row, i_h, col_vec_idx, None]) + + cute.arch.mbarrier_arrive(store_done_mbars + s) + + yield_out() + + +# ========================================================================= +# Host function +# ========================================================================= +def make_host_function( + B, + NT, + H, + is_varlen=False, + T_padded=None, + has_bias=False, + use_safe_gate=False, + varlen_pure=False, + kscaled_fp16=False, + preprocess_w_beta=False, + fuse_akk_inv=False, + fp32_akk_workspace=False, + input_gk_fp32=False, + beta_fp32=False, +): + """ + `varlen_pure=True` asserts that all seq lengths in the batch are multiples + of BT (= 64). Under that assumption every chunk has 64 valid rows so the + four mask sites (K1 row mask, K1 store mask, MMA accumulator zero-fill, + Store row mask) are guaranteed to never fire and are dead-code eliminated + at compile time. Caller's data layout is unchanged — this is a hint only. + """ + _B, _NT, _H = B, NT, H + _IS_VARLEN = 1 if is_varlen else 0 + _HAS_BIAS = 1 if has_bias else 0 + _USE_SAFE_GATE = 1 if use_safe_gate else 0 + _VARLEN_PURE = 1 if (is_varlen and varlen_pure) else 0 + _KSCALED_FP16 = 1 if kscaled_fp16 else 0 + _PREPROCESS_W_BETA = 1 if preprocess_w_beta else 0 + if _KSCALED_FP16 and _PREPROCESS_W_BETA: + raise ValueError("kscaled_fp16 and preprocess_w_beta are mutually exclusive") + _FUSE_AKK_INV = 1 if fuse_akk_inv else 0 + _FP32_AKK_WORKSPACE = 1 if fp32_akk_workspace else 0 + _INPUT_GK_FP32 = 1 if input_gk_fp32 else 0 + _BETA_FP32 = 1 if beta_fp32 else 0 + if _INPUT_GK_FP32 and (_HAS_BIAS or _PREPROCESS_W_BETA): + raise ValueError("input_gk_fp32 does not accept gate bias or beta preprocessing") + if _FUSE_AKK_INV and _FP32_AKK_WORKSPACE: + raise ValueError("fuse_akk_inv and fp32_akk_workspace are mutually exclusive") + if is_varlen: + assert _B == 1, "Varlen requires B=1" + assert T_padded is not None, "T_padded required for varlen" + _T = T_padded + else: + _T = _NT * BT + + if is_varlen: + _total_cgs_val = ((_NT + CHUNKS_PER_BLOCK - 1) // CHUNKS_PER_BLOCK) * _H + else: + _total_cgs_val = (_NT // CHUNKS_PER_BLOCK) * _H * _B + + # 3D TMA view: (B*T, K_DIM, H) — domain_offset handles non-aligned addressing + _T_total = _B * _T + s_row = _H * K_DIM + s_col = 1 + s_h = K_DIM + + @cute.jit + def host_fn( + mQ, + mK, + mG, + mA_log, + mBeta, + scale, + mKscaled, + mKg, + mQscaled, + mGkLast, + mAqk, + mAkk, + mCuSeqlens, + mChunkIndices, + mDtBias, + lower_bound_val, + ): + # 3D TMA view: (B*T, K_DIM, H). domain_offset in kernel shifts to + # arbitrary chunk_start — no BT alignment required for varlen. + view_layout_3d = cute.make_layout( + (_T_total, K_DIM, _H), + stride=(s_row, s_col, s_h), + ) + mQ_view = cute.make_tensor(mQ.iterator, view_layout_3d) + mK_view = cute.make_tensor(mK.iterator, view_layout_3d) + mG_view = cute.make_tensor(mG.iterator, view_layout_3d) + + smem_atom_qk = tcgen05.make_smem_layout_atom(tcgen05.SmemLayoutAtomKind.K_SW128, cutlass.BFloat16) + qk_smem_2d = cute.tile_to_shape(smem_atom_qk, (BT, K_DIM), order=(0, 1)) + qk_smem_3d = cute.tile_to_shape(smem_atom_qk, (BT, K_DIM, NUM_STAGES), order=(0, 1, 2)) + + g_smem_2d = cute.make_layout((BT, K_DIM), stride=(G_ROW_STRIDE_BF16, 1)) + g_smem_3d = cute.make_layout((BT, K_DIM, NUM_STAGES), stride=(G_ROW_STRIDE_BF16, 1, G_STAGE_STRIDE_BF16)) + smem_atom_gk = tcgen05.make_smem_layout_atom(tcgen05.SmemLayoutAtomKind.K_SW128, cutlass.Float32) + gk_smem_2d = cute.tile_to_shape(smem_atom_gk, (BT, K_DIM), order=(0, 1)) + tma_op = cpasync.CopyBulkTensorTileG2SOp(cpasync.CtaGroup.ONE) + tma_atom_Q, tma_tensor_Q = cpasync.make_tiled_tma_atom( + tma_op, mQ_view, qk_smem_2d, cute.product_each(qk_smem_2d.shape), num_multicast=1 + ) + tma_atom_K, tma_tensor_K = cpasync.make_tiled_tma_atom( + tma_op, mK_view, qk_smem_2d, cute.product_each(qk_smem_2d.shape), num_multicast=1 + ) + if cutlass.const_expr(_INPUT_GK_FP32 != 0): + tma_atom_G, tma_tensor_G = cpasync.make_tiled_tma_atom( + tma_op, mG_view, gk_smem_2d, cute.product_each(gk_smem_2d.shape), num_multicast=1 + ) + else: + tma_atom_G, tma_tensor_G = cpasync.make_tiled_tma_atom( + tma_op, mG_view, g_smem_2d, cute.product_each(g_smem_2d.shape), num_multicast=1 + ) + + if cutlass.const_expr(_INPUT_GK_FP32 != 0): + g_cumsum_layout = cute.tile_to_shape(smem_atom_gk, (BT, K_DIM, NUM_STAGES), order=(0, 1, 2)) + else: + g_cumsum_layout = cute.make_layout((BT, K_DIM, NUM_STAGES), stride=(K_STRIDE, 1, BT * K_STRIDE)) + + copy_atom_qk_k1 = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.BFloat16, num_bits_per_copy=32) + tiled_copy_qk_k1 = cute.make_tiled_copy_tv( + copy_atom_qk_k1, thr_layout=cute.make_layout((1, 32)), val_layout=cute.make_layout((1, 2)) + ) + + out_v2_layout = cute.make_layout( + (_B, _T, _H, K_VEC, VEC), + stride=(_T * _H * K_DIM, _H * K_DIM, K_DIM, VEC, 1), + ) + mKscaled_v2 = cute.make_tensor(mKscaled.iterator, out_v2_layout) + mQscaled_v2 = cute.make_tensor(mQscaled.iterator, out_v2_layout) + mKg_v2 = cute.make_tensor(mKg.iterator, out_v2_layout) + + # vec2 views for mAqk / mAkk (BT dimension instead of K_DIM). + # Shape: (B, T, H, BT/2, 2). Each VEC=2 slot = 2 contiguous bf16 = 1 + # fp32-aligned 4-byte unit. Used by store warp's autovec_copy → + # STG.E.32 with 32 lanes coalesced to 1 cache line per row. + BT_VEC = BT // VEC # 32 + akk_v2_layout = cute.make_layout( + (_B, _T, _H, BT_VEC, VEC), + stride=(_T * _H * BT, _H * BT, BT, VEC, 1), + ) + mAqk_v2 = cute.make_tensor(mAqk.iterator, akk_v2_layout) + mAkk_v2 = cute.make_tensor(mAkk.iterator, akk_v2_layout) + + gklast_v2_layout = cute.make_layout( + (_B, _NT, _H, K_VEC, VEC), + stride=(_NT * _H * K_DIM, _H * K_DIM, K_DIM, VEC, 1), + ) + mGkLast_v2 = cute.make_tensor(mGkLast.iterator, gklast_v2_layout) + + mma_op = cute.nvgpu.warp.MmaF16BF16Op(cutlass.BFloat16, cutlass.Float32, (16, 8, 16)) + tiled_mma_k2 = cute.make_tiled_mma(mma_op, cute.make_layout((1, 1, 1)), permutation_mnk=(16, 8, 16)) + + tiled_copy_mma_A = cute.make_tiled_copy_A( + cute.make_copy_atom(cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 4), cutlass.BFloat16), tiled_mma_k2 + ) + tiled_copy_mma_B = cute.make_tiled_copy_B( + cute.make_copy_atom(cute.nvgpu.warp.LdMatrix8x8x16bOp(False, 2), cutlass.BFloat16), tiled_mma_k2 + ) + + copy_atom_Gcum = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=64) + tiled_copy_Gcum_norm = cute.make_tiled_copy_tv( + copy_atom_Gcum, thr_layout=cute.make_layout((1, 4)), val_layout=cute.make_layout((1, 2)) + ) + tiled_copy_Gcum_gate = cute.make_tiled_copy_C(copy_atom_Gcum, tiled_mma_k2) + + # SMEM layout: sG and sAkk are independent double-buffered allocations. + # Q/K bf16: 64 * 128 * 2 * 2 * NUM_STAGES = 65,536 B + # sG bf16: 64 * 128 * 2 * NUM_STAGES = 32,768 B. + # sAkk fp32: 64 * 68 * 4 * NUM_STAGES = 34,816 B + # G cumsum: 64 * 136 * 4 * NUM_STAGES = 69,632 B + # K1 partial prefix: 8 * 132 * 4 = 4,224 B + # Aqk SMEM: 64 * 72 * 2 * NUM_STAGES = 18,432 B + # Barriers: 5 * 2 * 8 = 80 B + # Subtotal: 225,744 B (~220 KB) + # SM100 max: 228 KB (margin ~2 KB) + smem_size = ( + BT * K_DIM * 2 * 2 * NUM_STAGES # Q/K bf16: 65,536 B + + BT * AKK_STRIDE * 4 * NUM_STAGES # sAkk fp32: 34,816 B + + BT * K_STRIDE * 4 * NUM_STAGES # G cumsum fp32: 69,632 B + + K1_ROW_GROUPS * PARTIAL_COLS * 4 # K1 partial prefix: 4,224 B + + BT * AQK_TILE_STRIDE * 2 * NUM_STAGES # sAqk bf16: 18,432 B + + BT * NUM_STAGES * (4 if _BETA_FP32 else 2) + + 5 * NUM_STAGES * 8 # Barriers: 80 B + + BT * K_DIM * 2 * NUM_STAGES # sG bf16 independent allocation + ) + smem_size = max(smem_size, 225 * 1024) + + _grid_x = min(NUM_SMS, _total_cgs_val) + + fused_kernel123( + tma_atom_Q, + tma_tensor_Q, + tma_atom_K, + tma_tensor_K, + tma_atom_G, + tma_tensor_G, + mA_log, + mBeta, + scale, + mKscaled_v2, + mKg_v2, + mQscaled_v2, + mGkLast_v2, + mAqk, + mAkk, + mAqk_v2, + mAkk_v2, + tiled_copy_qk_k1, + tiled_mma_k2, + tiled_copy_mma_A, + tiled_copy_mma_B, + tiled_copy_Gcum_norm, + tiled_copy_Gcum_gate, + qk_smem_3d, + g_smem_3d, + g_cumsum_layout, + _NT, + _H, + _B, + mCuSeqlens, + mChunkIndices, + _IS_VARLEN, + mDtBias, + lower_bound_val, + _HAS_BIAS, + _USE_SAFE_GATE, + _VARLEN_PURE, + _KSCALED_FP16, + _PREPROCESS_W_BETA, + _FUSE_AKK_INV, + _FP32_AKK_WORKSPACE, + _INPUT_GK_FP32, + _BETA_FP32, + ).launch( + grid=(_grid_x, 1, 1), + block=(THREADS, 1, 1), + smem=smem_size, + ) + + return host_fn + + +# ============================================================================ +# High-level Python wrappers for training usage +# ============================================================================ + +from fla.ops.utils import prepare_chunk_indices + +_K123_CACHE = {} +_VARLEN_PAD_CACHE = {} + + +def _tensor_version(t: torch.Tensor) -> int: + try: + return t._version + except RuntimeError: + return 0 + + +def _ct_cached(t: torch.Tensor, dtype) -> torch.Tensor: + """Cache CuTe tensor wrappers to avoid repeated wrapping overhead.""" + version = _tensor_version(t) + key = ("ct", id(t), dtype) + entry = _K123_CACHE.get(key) + if entry is None or entry[0]() is not t or entry[1] != version: + from cutlass.cute.runtime import from_dlpack + + # Use assumed_align and mark_layout_dynamic for optimal performance + wrapper = from_dlpack(t, assumed_align=16).mark_layout_dynamic() + wrapper.element_type = dtype + entry = (weakref.ref(t), version, wrapper) + _K123_CACHE[key] = entry + return entry[2] + + +def _get_bias_ct(dt_bias, H, K, device): + """Handle dt_bias reshape and wrapping for CuTe.""" + if dt_bias is not None: + # Reshape to [H, K] and wrap with assumed_align and mark_layout_dynamic + bias = dt_bias.contiguous().view(H, K) + return _ct_cached(bias, cutlass.Float32) + else: + # Return empty bias tensor with mark_layout_dynamic + key = ("empty_bias", device.index if device.index is not None else 0) + entry = _K123_CACHE.get(key) + if entry is None: + from cutlass.cute.runtime import from_dlpack + + bias = torch.empty(1, 1, dtype=torch.float32, device=device) + wrapper = from_dlpack(bias, assumed_align=16).mark_layout_dynamic() + wrapper.element_type = cutlass.Float32 + entry = (bias, wrapper) + _K123_CACHE[key] = entry + return entry[1] + + +def _get_out_buffers_equal(device, B, T_total, H, kscaled_dtype=torch.bfloat16): + """Cache and reuse the six K123 output buffers by shape.""" + key = (device.index if device.index is not None else 0, B, T_total, H, kscaled_dtype) + entry = _K123_CACHE.get(("buffers", key)) + if entry is None: + nt = T_total // BT + k_scaled = torch.empty(B, T_total, H, K_DIM, device=device, dtype=kscaled_dtype) + kg = torch.empty(B, T_total, H, K_DIM, device=device, dtype=torch.bfloat16) + q_scaled = torch.empty(B, T_total, H, K_DIM, device=device, dtype=torch.bfloat16) + gk_last_exp = torch.empty(B, nt, H, K_DIM, device=device, dtype=torch.float32) + A_qk = torch.empty(B, T_total, H, BT, device=device, dtype=torch.bfloat16) + A_kk = torch.empty(B, T_total, H, BT, device=device, dtype=torch.bfloat16) + entry = (k_scaled, kg, q_scaled, gk_last_exp, A_qk, A_kk) + _K123_CACHE[("buffers", key)] = entry + return entry + + +def _get_out_buffers_varlen(device, T_padded, NT, H, kscaled_dtype=torch.bfloat16): + """Varlen counterpart keyed on launch shape.""" + key = (device.index if device.index is not None else 0, T_padded, NT, H, kscaled_dtype) + entry = _K123_CACHE.get(("buffers_varlen", key)) + if entry is None: + k_scaled = torch.empty(1, T_padded, H, K_DIM, device=device, dtype=kscaled_dtype) + kg = torch.empty(1, T_padded, H, K_DIM, device=device, dtype=torch.bfloat16) + q_scaled = torch.empty(1, T_padded, H, K_DIM, device=device, dtype=torch.bfloat16) + gk_last_exp = torch.empty(1, NT, H, K_DIM, device=device, dtype=torch.float32) + A_qk = torch.empty(1, T_padded, H, BT, device=device, dtype=torch.bfloat16) + A_kk = torch.empty(1, T_padded, H, BT, device=device, dtype=torch.bfloat16) + entry = (k_scaled, kg, q_scaled, gk_last_exp, A_qk, A_kk) + _K123_CACHE[("buffers_varlen", key)] = entry + return entry + + +def _get_akk_fp32_workspace_equal(device, B, T_total, H): + key = (device.index if device.index is not None else 0, B, T_total, H) + entry = _K123_CACHE.get(("akk_fp32_workspace", key)) + if entry is None: + entry = torch.empty(B, T_total, H, BT, device=device, dtype=torch.float32) + _K123_CACHE[("akk_fp32_workspace", key)] = entry + return entry + + +def _get_akk_fp32_workspace_varlen(device, T_padded, NT, H): + key = (device.index if device.index is not None else 0, T_padded, NT, H) + entry = _K123_CACHE.get(("akk_fp32_workspace_varlen", key)) + if entry is None: + entry = torch.empty(1, T_padded, H, BT, device=device, dtype=torch.float32) + _K123_CACHE[("akk_fp32_workspace_varlen", key)] = entry + return entry + + +def _pad_varlen_inputs(q, k, g, beta, T_padded): + """Zero-pad packed varlen inputs to 4*BT boundary.""" + inputs = (q, k, g, beta) + if all(t.shape[1] == T_padded for t in inputs): + return inputs + + versions = tuple(_tensor_version(t) for t in inputs) + key = (*map(id, inputs), T_padded) + entry = _VARLEN_PAD_CACHE.get(key) + if entry is None or any(ref() is not tensor for ref, tensor in zip(entry[0], inputs, strict=True)) or entry[1] != versions: + + def _pad(t): + if t.shape[1] == T_padded: + return t + pad = t.new_zeros((t.shape[0], T_padded, *t.shape[2:])) + pad[:, : t.shape[1]].copy_(t) + return pad + + entry = (tuple(weakref.ref(t) for t in inputs), versions, tuple(_pad(t) for t in inputs)) + _VARLEN_PAD_CACHE[key] = entry + return entry[2] + + +def _get_eqlen_dummies(device): + """Dummy varlen tensors for equal-length launches.""" + key = ("eqlen_dummies", device.index if device.index is not None else 0) + entry = _K123_CACHE.get(key) + if entry is None: + cu_seqlens = torch.tensor([0, 1], dtype=torch.int32, device=device) + chunk_indices = torch.tensor([[0, 0]], dtype=torch.int32, device=device) + cu_ct = _ct_cached(cu_seqlens, cutlass.Int32) + ci_ct = _ct_cached(chunk_indices, cutlass.Int32) + entry = (cu_seqlens, chunk_indices, cu_ct, ci_ct) + _K123_CACHE[key] = entry + return entry[2], entry[3] + + +def _get_csrc_port_dummy_alog(device, heads): + """Return an unused A_log argument for the FP32-gk specialization.""" + key = ("csrc_port_dummy_alog", device.index if device.index is not None else 0, heads) + entry = _K123_CACHE.get(key) + if entry is None: + entry = torch.zeros(heads, dtype=torch.float32, device=device) + _K123_CACHE[key] = entry + return entry + + +def _get_csrc_port_dummy_bias(device, heads): + """Return an unused dt-bias argument for the FP32-gk specialization.""" + key = ("csrc_port_dummy_bias", device.index if device.index is not None else 0, heads) + entry = _K123_CACHE.get(key) + if entry is None: + entry = torch.zeros((heads, K_DIM), dtype=torch.float32, device=device) + _K123_CACHE[key] = entry + return entry + + +def _run_akk_inv_fp32_physical( + A_phys: torch.Tensor, + A_out: torch.Tensor, + B: int, + NT: int, + H: int, + cu_ct, + ci_ct, + *, + is_varlen: bool, + T_val: int, +): + dev = A_phys.device.index if A_phys.device.index is not None else 0 + phys_ct = _ct_cached(A_phys, cutlass.Float32) + out_ct = _ct_cached(A_out, cutlass.BFloat16) + cache_key = ("akk_inv_fp32_physical", dev, B, NT, H, bool(is_varlen), T_val) + kernel = _K123_CACHE.get(cache_key) + if kernel is None: + kernel = cute.compile( + _akk_inv_fp32_physical_host, + phys_ct, + out_ct, + B, + NT, + H, + cu_ct, + ci_ct, + 1 if is_varlen else 0, + T_val, + ) + _K123_CACHE[cache_key] = kernel + kernel(phys_ct, out_ct, cu_ct, ci_ct) + + +def chunk_kda_fwd_intra_sm100_equal( + q: torch.Tensor, + k: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + scale: float | None = None, + *, + dt_bias: torch.Tensor | None = None, + safe_gate: bool = True, + lower_bound: float | None = None, + chunk_size: int = BT, + fp32_akk_inv: bool = False, + fused_akk_inv: bool = False, + preprocess_w_beta: bool = False, + kscaled_fp16: bool = False, +): + """Launch FlashInfer fused K1+K2+K3 kernel for equal-length sequences.""" + if chunk_size != BT: + raise NotImplementedError(f"chunk_size must be {BT}, got {chunk_size}.") + if q.dim() != 4 or k.dim() != 4 or g.dim() != 4 or beta.dim() != 3: + raise ValueError("Expected q/k/g [B,T,H,K] and beta [B,T,H].") + if q.shape != k.shape or q.shape != g.shape: + raise ValueError(f"q/k/g shapes must match, got q={q.shape}, k={k.shape}, g={g.shape}.") + B, T_total, H, K = q.shape + if K != K_DIM: + raise NotImplementedError(f"K must be {K_DIM}, got {K}.") + if T_total % (4 * BT) != 0: + raise NotImplementedError(f"T must be a multiple of {4 * BT}, got {T_total}.") + if beta.shape != (B, T_total, H): + raise ValueError(f"beta shape must be {(B, T_total, H)}, got {tuple(beta.shape)}.") + if A_log.shape != (H,): + raise ValueError(f"A_log shape must be {(H,)}, got {tuple(A_log.shape)}.") + if scale is None: + scale = K_DIM**-0.5 + if safe_gate and lower_bound is None: + lower_bound = -5.0 + if fp32_akk_inv and fused_akk_inv: + raise ValueError("fp32_akk_inv and fused_akk_inv are mutually exclusive") + if preprocess_w_beta and kscaled_fp16: + raise ValueError("preprocess_w_beta and kscaled_fp16 are mutually exclusive") + + NT = T_total // BT + dev = q.device.index if q.device.index is not None else 0 + has_bias = dt_bias is not None + cache_key = ( + dev, + B, + NT, + H, + has_bias, + safe_gate, + fp32_akk_inv, + fused_akk_inv, + preprocess_w_beta, + kscaled_fp16, + ) + + kscaled_dtype = torch.float16 if kscaled_fp16 else torch.bfloat16 + k_scaled, kg, q_scaled, gk_last_exp, A_qk, A_kk = _get_out_buffers_equal(q.device, B, T_total, H, kscaled_dtype) + A_kk_workspace = _get_akk_fp32_workspace_equal(q.device, B, T_total, H) if fp32_akk_inv else A_kk + + q_ct = _ct_cached(q, cutlass.BFloat16) + k_ct = _ct_cached(k, cutlass.BFloat16) + g_ct = _ct_cached(g, cutlass.BFloat16) + alog_ct = _ct_cached(A_log, cutlass.Float32) + beta_ct = _ct_cached(beta, cutlass.BFloat16) + ks_ct = _ct_cached(k_scaled, cutlass.Float16 if kscaled_fp16 else cutlass.BFloat16) + kg_ct = _ct_cached(kg, cutlass.BFloat16) + qs_ct = _ct_cached(q_scaled, cutlass.BFloat16) + gk_ct = _ct_cached(gk_last_exp, cutlass.Float32) + aqk_ct = _ct_cached(A_qk, cutlass.BFloat16) + akk_ct = _ct_cached(A_kk_workspace, cutlass.Float32 if fp32_akk_inv else cutlass.BFloat16) + cu_ct, ci_ct = _get_eqlen_dummies(q.device) + + if has_bias: + bias_ct = _get_bias_ct(dt_bias, H, K_DIM, q.device) + else: + bias_ct = _get_bias_ct(None, H, K_DIM, q.device) + + lb_val = float(lower_bound) if lower_bound is not None else 0.0 + + ct_args = ( + q_ct, + k_ct, + g_ct, + alog_ct, + beta_ct, + float(scale), + ks_ct, + kg_ct, + qs_ct, + gk_ct, + aqk_ct, + akk_ct, + cu_ct, + ci_ct, + bias_ct, + lb_val, + ) + + kernel = _K123_CACHE.get(("kernel", cache_key)) + if kernel is None: + import cutlass.cute as cute + + host_fn = make_host_function( + B, + NT, + H, + is_varlen=False, + has_bias=has_bias, + use_safe_gate=safe_gate, + kscaled_fp16=kscaled_fp16, + preprocess_w_beta=preprocess_w_beta, + fuse_akk_inv=fused_akk_inv, + fp32_akk_workspace=fp32_akk_inv, + ) + kernel = cute.compile(host_fn, *ct_args) + _K123_CACHE[("kernel", cache_key)] = kernel + + kernel(*ct_args) + if fp32_akk_inv: + _run_akk_inv_fp32_physical( + A_kk_workspace, + A_kk, + B, + NT, + H, + cu_ct, + ci_ct, + is_varlen=False, + T_val=T_total, + ) + return k_scaled, kg, q_scaled, gk_last_exp, A_qk, A_kk + + +def chunk_kda_fwd_intra_sm100_from_gk( + q: torch.Tensor, + k: torch.Tensor, + gk: torch.Tensor, + beta: torch.Tensor, + scale: float | None = None, + *, + chunk_size: int = BT, + use_tf32_inverse: bool = True, + fp32_akk_inv: bool = True, +): + """CuTeDSL port of the equal-length csrc intra API. + + Unlike :func:`chunk_kda_fwd_intra_sm100_equal`, this entry point consumes + the already activated and chunk-cumsummed FP32 ``gk`` tensor, exactly like + ``chunk_kda_fwd_intra_cuda``. It deliberately preserves the csrc kernel + boundary; gate fusion is not part of this correctness baseline. + + ``fp32_akk_inv=True`` keeps the pre-inverse Akk workspace in FP32 and uses + the same single-accumulator K=32 TF32 Schur reduction order as csrc. The + experimental in-CTA inverse is intentionally excluded from this bitwise + correctness boundary. + """ + if chunk_size != BT: + raise NotImplementedError(f"chunk_size must be {BT}, got {chunk_size}.") + if q.dim() != 4 or k.dim() != 4 or gk.dim() != 4 or beta.dim() != 3: + raise ValueError("Expected q/k/gk [B,T,H,K] and beta [B,T,H].") + if q.shape != k.shape or q.shape != gk.shape: + raise ValueError(f"q/k/gk shapes must match, got q={q.shape}, k={k.shape}, gk={gk.shape}.") + if q.dtype != torch.bfloat16 or k.dtype != torch.bfloat16: + raise TypeError("q and k must be bfloat16, matching the csrc specialization.") + if gk.dtype != torch.float32: + raise TypeError(f"gk must be float32, got {gk.dtype}.") + if beta.dtype not in (torch.bfloat16, torch.float32): + raise TypeError(f"beta must be bfloat16 or float32, got {beta.dtype}.") + B, T_total, H, K = q.shape + if K != K_DIM: + raise NotImplementedError(f"K must be {K_DIM}, got {K}.") + if T_total % (CHUNKS_PER_BLOCK * BT) != 0: + raise NotImplementedError(f"T must be a multiple of {CHUNKS_PER_BLOCK * BT}, got {T_total}.") + if beta.shape != (B, T_total, H): + raise ValueError(f"beta shape must be {(B, T_total, H)}, got {tuple(beta.shape)}.") + if not use_tf32_inverse: + raise NotImplementedError("The one-to-one baseline currently ports the default csrc TF32 inverse only.") + if not fp32_akk_inv: + raise NotImplementedError( + "The csrc-boundary port requires fp32_akk_inv=True; the experimental in-CTA inverse is not bitwise aligned." + ) + if scale is None: + scale = K_DIM**-0.5 + + NT = T_total // BT + dev = q.device.index if q.device.index is not None else 0 + cache_key = ("csrc_port", dev, B, NT, H, beta.dtype, fp32_akk_inv) + k_scaled, kg, q_scaled, gk_last_exp, A_qk, A_kk = _get_out_buffers_equal(q.device, B, T_total, H, torch.bfloat16) + A_kk_workspace = _get_akk_fp32_workspace_equal(q.device, B, T_total, H) if fp32_akk_inv else A_kk + + q_ct = _ct_cached(q, cutlass.BFloat16) + k_ct = _ct_cached(k, cutlass.BFloat16) + gk_in_ct = _ct_cached(gk, cutlass.Float32) + alog_ct = _ct_cached(_get_csrc_port_dummy_alog(q.device, H), cutlass.Float32) + beta_type = cutlass.Float32 if beta.dtype == torch.float32 else cutlass.BFloat16 + beta_ct = _ct_cached(beta, beta_type) + ks_ct = _ct_cached(k_scaled, cutlass.BFloat16) + kg_ct = _ct_cached(kg, cutlass.BFloat16) + qs_ct = _ct_cached(q_scaled, cutlass.BFloat16) + gk_last_ct = _ct_cached(gk_last_exp, cutlass.Float32) + aqk_ct = _ct_cached(A_qk, cutlass.BFloat16) + akk_ct = _ct_cached(A_kk_workspace, cutlass.Float32 if fp32_akk_inv else cutlass.BFloat16) + cu_ct, ci_ct = _get_eqlen_dummies(q.device) + bias_ct = _ct_cached(_get_csrc_port_dummy_bias(q.device, H), cutlass.Float32) + + ct_args = ( + q_ct, + k_ct, + gk_in_ct, + alog_ct, + beta_ct, + float(scale), + ks_ct, + kg_ct, + qs_ct, + gk_last_ct, + aqk_ct, + akk_ct, + cu_ct, + ci_ct, + bias_ct, + 0.0, + ) + kernel = _K123_CACHE.get(("kernel", cache_key)) + if kernel is None: + host_fn = make_host_function( + B, + NT, + H, + is_varlen=False, + fuse_akk_inv=not fp32_akk_inv, + fp32_akk_workspace=fp32_akk_inv, + input_gk_fp32=True, + beta_fp32=beta.dtype == torch.float32, + ) + kernel = cute.compile(host_fn, *ct_args) + _K123_CACHE[("kernel", cache_key)] = kernel + + kernel(*ct_args) + if fp32_akk_inv: + _run_akk_inv_fp32_physical( + A_kk_workspace, + A_kk, + B, + NT, + H, + cu_ct, + ci_ct, + is_varlen=False, + T_val=T_total, + ) + return A_qk, A_kk + + +def chunk_kda_fwd_intra_sm100_varlen( + q: torch.Tensor, + k: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + cu_seqlens: torch.Tensor, + scale: float | None = None, + *, + chunk_indices: torch.Tensor | None = None, + dt_bias: torch.Tensor | None = None, + safe_gate: bool = True, + lower_bound: float | None = None, + seq_lens: list[int] | None = None, + fp32_akk_inv: bool = False, + fused_akk_inv: bool = False, + preprocess_w_beta: bool = False, + kscaled_fp16: bool = False, +): + """Launch FlashInfer fused K1+K2+K3 kernel for varlen sequences. + + Inputs are packed as [1, T_total, H, K] and indexed via cu_seqlens + chunk_indices. + """ + if q.dim() != 4 or k.dim() != 4 or g.dim() != 4 or beta.dim() != 3: + raise ValueError("Expected q/k/g [1,T_total,H,K] and beta [1,T_total,H].") + if q.shape != k.shape or q.shape != g.shape: + raise ValueError(f"q/k/g shapes must match, got q={q.shape}, k={k.shape}, g={g.shape}.") + B, T_total, H, K = q.shape + if B != 1: + raise ValueError("varlen launch expects packed tensors with B=1.") + if K != K_DIM: + raise NotImplementedError(f"K must be {K_DIM}, got {K}.") + if scale is None: + scale = K_DIM**-0.5 + if safe_gate and lower_bound is None: + lower_bound = -5.0 + if fp32_akk_inv and fused_akk_inv: + raise ValueError("fp32_akk_inv and fused_akk_inv are mutually exclusive") + if preprocess_w_beta and kscaled_fp16: + raise ValueError("preprocess_w_beta and kscaled_fp16 are mutually exclusive") + has_bias = dt_bias is not None + + # A packed uniform batch with full 4-chunk groups has exactly the same + # storage as the equal-length path. Route it there to remove per-work-unit + # chunk-index decoding and varlen predicates; flatten the outputs back to + # the packed ABI. The WU wrapper applies the same uniform fast path. + if seq_lens and len(set(seq_lens)) == 1 and seq_lens[0] % (4 * BT) == 0: + uniform_batch = len(seq_lens) + uniform_t = seq_lens[0] + if uniform_batch * uniform_t != T_total: + raise ValueError(f"seq_lens sum must equal packed T={T_total}, got {sum(seq_lens)}") + equal_out = chunk_kda_fwd_intra_sm100_equal( + q.view(uniform_batch, uniform_t, H, K), + k.view(uniform_batch, uniform_t, H, K), + g.view(uniform_batch, uniform_t, H, K), + beta.view(uniform_batch, uniform_t, H), + A_log, + scale, + dt_bias=dt_bias, + safe_gate=safe_gate, + lower_bound=lower_bound, + fp32_akk_inv=fp32_akk_inv, + fused_akk_inv=fused_akk_inv, + preprocess_w_beta=preprocess_w_beta, + kscaled_fp16=kscaled_fp16, + ) + return ( + equal_out[0].flatten(0, 1).unsqueeze(0), + equal_out[1].flatten(0, 1).unsqueeze(0), + equal_out[2].flatten(0, 1).unsqueeze(0), + equal_out[3].flatten(0, 1).unsqueeze(0), + equal_out[4].flatten(0, 1).unsqueeze(0), + equal_out[5].flatten(0, 1).unsqueeze(0), + ) + + cu_seqlens = cu_seqlens.contiguous().to(torch.int32) + if chunk_indices is None: + chunk_indices = prepare_chunk_indices(cu_seqlens, BT) + chunk_indices = chunk_indices.contiguous().to(torch.int32) + + NT = chunk_indices.shape[0] + + # Calculate varlen_pure: all seq_lens are multiples of BT + if seq_lens is not None: + varlen_pure = all(sl % BT == 0 for sl in seq_lens) + else: + cu_cpu = cu_seqlens.detach().cpu().tolist() + varlen_pure = all((cu_cpu[i + 1] - cu_cpu[i]) % BT == 0 for i in range(len(cu_cpu) - 1)) + + # Pad to 4*BT boundary + T_padded = ((T_total + 4 * BT - 1) // (4 * BT)) * (4 * BT) + q_pad, k_pad, g_pad, beta_pad = _pad_varlen_inputs(q, k, g, beta, T_padded) + + dev = q.device.index if q.device.index is not None else 0 + cache_key = ( + dev, + T_padded, + NT, + H, + has_bias, + safe_gate, + varlen_pure, + fp32_akk_inv, + fused_akk_inv, + preprocess_w_beta, + kscaled_fp16, + ) + + kscaled_dtype = torch.float16 if kscaled_fp16 else torch.bfloat16 + k_scaled, kg, q_scaled, gk_last_exp, A_qk, A_kk = _get_out_buffers_varlen(q.device, T_padded, NT, H, kscaled_dtype) + A_kk_workspace = _get_akk_fp32_workspace_varlen(q.device, T_padded, NT, H) if fp32_akk_inv else A_kk + + q_ct = _ct_cached(q_pad, cutlass.BFloat16) + k_ct = _ct_cached(k_pad, cutlass.BFloat16) + g_ct = _ct_cached(g_pad, cutlass.BFloat16) + alog_ct = _ct_cached(A_log, cutlass.Float32) + beta_ct = _ct_cached(beta_pad, cutlass.BFloat16) + ks_ct = _ct_cached(k_scaled, cutlass.Float16 if kscaled_fp16 else cutlass.BFloat16) + kg_ct = _ct_cached(kg, cutlass.BFloat16) + qs_ct = _ct_cached(q_scaled, cutlass.BFloat16) + gk_ct = _ct_cached(gk_last_exp, cutlass.Float32) + aqk_ct = _ct_cached(A_qk, cutlass.BFloat16) + akk_ct = _ct_cached(A_kk_workspace, cutlass.Float32 if fp32_akk_inv else cutlass.BFloat16) + cu_ct = _ct_cached(cu_seqlens, cutlass.Int32) + ci_ct = _ct_cached(chunk_indices, cutlass.Int32) + + if has_bias: + bias_ct = _get_bias_ct(dt_bias, H, K_DIM, q.device) + else: + bias_ct = _get_bias_ct(None, H, K_DIM, q.device) + + lb_val = float(lower_bound) if lower_bound is not None else 0.0 + + ct_args = ( + q_ct, + k_ct, + g_ct, + alog_ct, + beta_ct, + float(scale), + ks_ct, + kg_ct, + qs_ct, + gk_ct, + aqk_ct, + akk_ct, + cu_ct, + ci_ct, + bias_ct, + lb_val, + ) + + kernel = _K123_CACHE.get(("kernel_varlen", cache_key)) + if kernel is None: + import cutlass.cute as cute + + host_fn = make_host_function( + 1, + NT, + H, + is_varlen=True, + T_padded=T_padded, + has_bias=has_bias, + use_safe_gate=safe_gate, + varlen_pure=varlen_pure, + kscaled_fp16=kscaled_fp16, + preprocess_w_beta=preprocess_w_beta, + fuse_akk_inv=fused_akk_inv, + fp32_akk_workspace=fp32_akk_inv, + ) + kernel = cute.compile(host_fn, *ct_args) + _K123_CACHE[("kernel_varlen", cache_key)] = kernel + + kernel(*ct_args) + if fp32_akk_inv: + _run_akk_inv_fp32_physical( + A_kk_workspace, + A_kk, + 1, + NT, + H, + cu_ct, + ci_ct, + is_varlen=True, + T_val=T_padded, + ) + + # Slice back to real T_total + return ( + k_scaled[:, :T_total], + kg[:, :T_total], + q_scaled[:, :T_total], + gk_last_exp, + A_qk[:, :T_total], + A_kk[:, :T_total], + ) + + +chunk_kda_fwd_intra_sm100_training = chunk_kda_fwd_intra_sm100_equal + + +def _get_dt_bias_ct(dt_bias, H, K): + return _get_bias_ct(dt_bias, H, K, dt_bias.device) + + +def _get_empty_bias_ct(device): + return _get_bias_ct(None, 1, 1, device) diff --git a/cula/ops/kda/sm100/recompute_wu.py b/cula/ops/kda/sm100/recompute_wu.py new file mode 100644 index 00000000..317100bd --- /dev/null +++ b/cula/ops/kda/sm100/recompute_wu.py @@ -0,0 +1,1331 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +""" +CuTeDSL kernel for KDA recompute_w_u_fwd. + +Computes (always with gk, always recompute): + w = A @ (k * beta * exp2(gk)) — [BT, BK] per tile + u = A @ (v * beta) — [BT, BV] per tile + kg = k * exp2(gn - gk) — [BT, BK] per tile + +MMA layout (tcgen05 SS): + C[M, N] = A_mma[M, K] @ B_mma[N, K]^T + A_mma(SMEM) = A_mat[BT, BT], B_mma(SMEM) = B_proc^T[BN, BT] + Result = output[BT, BN], matching the time-major GMEM layout. + MMA tiler = (BT, BN, BT). + +The implementation mirrors the csrc SM100 persistent kernel: one CTA per SM, +three warp-groups, double-buffered input TMA pipelines, auxiliary beta loads, +and CUDA-core W/U stores directly from TMEM. + +Varlen mode: variable sequence lengths. cu_seqlens[N+1] gives token + offsets; chunk_indices[total_nt*2] gives (batch_idx, chunk_in_seq) + pairs for each global chunk index. TMA uses domain_offset for per-WU + alignment, matching the fwd_o.py pattern. + +Warp assignment: + 0-3: K prologue + kg output warpgroup + 4-7: V prologue + w/u epilogue warpgroup + 8: MMA warp + 9: Load warp + 10-11: beta-load auxiliary warps +""" + +import weakref + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import torch +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.runtime import make_fake_compact_tensor, make_fake_stream +from cutlass.cute.tensor import TensorSSA +from cutlass.cute.typing import Int32, Int64 +from fla.ops.utils import prepare_chunk_indices + +from cula.ops.ptx import reinterpret_cast, store_256b, subvec +from cula.ops.sm100.ptx import tcgen05_fence_after, tcgen05_fence_before, tcgen05_ld_32x32b +from cula.utils import USE_FAST_MATH, assert_blackwell + + +def _make_coop_group(size: int): + return pipeline.CooperativeGroup(pipeline.Agent.Thread, size) + + +class KDARecomputeWU: + def __init__( + self, + K: int = 128, + V: int = 128, + chunk_size: int = 64, + block_k: int = None, + block_v: int = None, + io_dtype: type[cutlass.Numeric] = cutlass.BFloat16, + acc_dtype: type[cutlass.Numeric] = cutlass.Float32, + beta_dtype: type[cutlass.Numeric] = cutlass.Float32, + is_varlen: bool = False, + preprocessed_k: bool = False, + persistent: bool = False, + use_fast_math: bool = True, + ): + assert K == 128 and V == 128, f"K and V must both be 128, got K={K}, V={V}" + assert_blackwell() + self.use_fast_math = use_fast_math + self.K = K + self.V = V + self.BT = chunk_size + self.BK = block_k if block_k is not None else K + self.BV = block_v if block_v is not None else V + self.NK = (K + self.BK - 1) // self.BK + self.NV = (V + self.BV - 1) // self.BV + self.io_dtype = io_dtype + self.acc_dtype = acc_dtype + self.beta_dtype = beta_dtype + self.is_varlen = is_varlen + self.preprocessed_k = preprocessed_k + + self.threads_per_warp = 32 + self.prologue_warp_ids = (0, 1, 2, 3) + self.epilogue_warp_ids = (4, 5, 6, 7) + self.mma_warp_id = 8 + self.load_warp_id = 9 + self.aux_warp_ids = (10, 11) + self.threads_per_cta = self.threads_per_warp * 12 # 384 + self.num_cuda_warps = 4 + self.num_cuda_threads = self.threads_per_warp * self.num_cuda_warps # 128 per compute WG + + self.BN = max(self.BK, self.BV) + self.mma_tiler = (self.BT, self.BN, self.BT) + + self.cta_group = tcgen05.CtaGroup.ONE + self.cluster_shape_mnk = (1, 1, 1) + self.buffer_align_bytes = 1024 + + self.bproc_stage = 1 + self.acc_pipe_stage = 1 + + # Match the csrc 384-thread, one-CTA-per-SM register split. + self.min_occupancy = 1 + self.num_regs_prologue = 224 + self.num_regs_epilogue = 200 + self.num_regs_others = 80 + self.a_stage = 2 + self.k_stage = 2 + self.g_stage = 2 + self.v_tma_stage = 2 + self.beta_stage = 2 + self.num_sm = utils.HardwareInfo().get_device_multiprocessor_count() + + @staticmethod + def _plan_tmem(tiled_mma, mma_tiler, acc_stages): + SM100_TMEM_CAPACITY_COLS = 512 + acc_shape = tiled_mma.partition_shape_C(mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, acc_stages)) + num_acc = tcgen05.find_tmem_tensor_col_offset(tCtAcc_fake) + total = 1 + while total < num_acc: + total *= 2 + assert total <= SM100_TMEM_CAPACITY_COLS + return total + + @cute.jit + def _tma_partition_A(self, tma_atom, tma_tensor, smem, tile_shape, tiled_mma, batch_idx, hidx): + coord = (None, 0, None) + gX = cute.local_tile(tma_tensor, cute.slice_(tile_shape, coord), (None, None, (hidx, batch_idx))) + thr_mma = tiled_mma.get_slice(0) + tCgX = thr_mma.partition_A(gX) + tXsX, tXgX = cpasync.tma_partition( + tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(smem, 0, 3), + cute.group_modes(tCgX, 0, 3), + ) + return tXsX, tXgX + + @cute.jit + def _tma_partition_B(self, tma_atom, tma_tensor, smem, tile_shape, tiled_mma, batch_idx, hidx): + coord = (0, None, None) + gX = cute.local_tile(tma_tensor, cute.slice_(tile_shape, coord), (None, None, (hidx, batch_idx))) + thr_mma = tiled_mma.get_slice(0) + tCgX = thr_mma.partition_B(gX) + tXsX, tXgX = cpasync.tma_partition( + tma_atom, + 0, + cute.make_layout(1), + cute.group_modes(smem, 0, 3), + cute.group_modes(tCgX, 0, 3), + ) + return tXsX, tXgX + + @cute.jit + def _data_tma_partition(self, atom, tma_tensor_3d, tile_shape, smem, head_idx, batch_idx): + """Partition for non-MMA TMA load (epilog-style).""" + gmem_2d = tma_tensor_3d[None, None, (head_idx, batch_idx)] + gC_tiled = cute.local_tile(gmem_2d, tile_shape, (None, None)) + sC_g = cute.group_modes(smem, 0, 2) + gC_g = cute.group_modes(gC_tiled, 0, 2) + bSG_sC, bSG_gC = cpasync.tma_partition( + atom, + 0, + cute.make_layout(1), + sC_g, + gC_g, + ) + return bSG_sC, bSG_gC + + @cute.jit + def _epilog_partition_varlen(self, atom, gC_2d, epi_tile, sC): + """Partition for varlen epilog TMA load (2D tensor with domain_offset). + Uses local_tile to correctly handle domain_offset coordinates. + """ + gC_tiled = cute.local_tile(gC_2d, epi_tile, (None, None)) + sC_g = cute.group_modes(sC, 0, 2) + gC_g = cute.group_modes(gC_tiled, 0, 2) + bSG_sC, bSG_gC = cpasync.tma_partition( + atom, + 0, + cute.make_layout(1), + sC_g, + gC_g, + ) + return bSG_sC, bSG_gC + + @cute.jit + def _decode_persistent_work(self, work_idx, total_nt, H, T, BT, cu_seqlens, chunk_indices): + chunk_global = work_idx // H + i_h = work_idx - chunk_global * H + if cutlass.const_expr(self.is_varlen): + i_b = chunk_indices[chunk_global * 2] + i_t = chunk_indices[chunk_global * 2 + 1] + tok_offset = cu_seqlens[i_b] + data_bidx = Int32(0) + seq_end = cu_seqlens[i_b + 1] + remaining = seq_end - (tok_offset + i_t * BT) + remaining = cutlass.select_(remaining > BT, Int32(BT), remaining) + else: + NT = (T + BT - 1) // BT + i_b = chunk_global // NT + i_t = chunk_global - i_b * NT + tok_offset = i_b * T + data_bidx = i_b + remaining = Int32(BT) + return i_b, i_t, i_h, tok_offset, data_bidx, remaining + + @cute.jit + def __call__( + self, + k_in: cute.Tensor, + v_in: cute.Tensor, + beta_in: cute.Tensor, + A_in: cute.Tensor, + gk_in: cute.Tensor, + w_in: cute.Tensor, + u_in: cute.Tensor, + kg_in: cute.Tensor, + cu_seqlens_in: cute.Tensor, + chunk_indices_in: cute.Tensor, + problem_size: tuple[Int32, Int32, Int32, Int32, Int32], + total_nt: Int32, + stream, + ): + k_ptr = k_in.iterator + v_ptr = v_in.iterator + beta_ptr = beta_in.iterator + A_ptr = A_in.iterator + gk_ptr = gk_in.iterator + w_ptr = w_in.iterator + u_ptr = u_in.iterator + kg_ptr = kg_in.iterator + cu_seqlens_ptr = cu_seqlens_in.iterator + chunk_indices_ptr = chunk_indices_in.iterator + + B, T, H, K, V = problem_size + BT = self.BT + + # For varlen: data_B=1, T=T_total + if cutlass.const_expr(self.is_varlen): + data_B = Int32(1) + else: + data_B = B + + # ---------- MMA setup ---------- + tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.io_dtype, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.MN, + self.acc_dtype, + self.cta_group, + self.mma_tiler[:2], + ) + + # csrc reserves the full 512 TMEM columns and aliases W/U by dp-lane: + # W starts at lane 0, U starts at lane 16. + self.tmem_total = 512 + + # Both MMA operands are resident in SMEM. A is the triangular Akk + # matrix; B is the time-transposed preprocessed K/V tile. + a_smem_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.io_dtype, + self.a_stage, + ) + b_smem_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.io_dtype, + self.bproc_stage, + ) + b_epi_staged = sm100_utils.make_smem_layout_epi( + self.io_dtype, + utils.LayoutEnum.ROW_MAJOR, + (self.BT, self.BN), + self.bproc_stage, + ) + assert cute.cosize(b_smem_staged) == cute.cosize(b_epi_staged) + + # ---------- TMA load op ---------- + tma_load_op = cpasync.CopyBulkTensorTileG2SOp(self.cta_group) + + cluster_layout = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (tiled_mma.thr_id.shape,), + ) + + # ---------- SMEM layouts: k (bf16), v (bf16), gk (fp32) ---------- + k_epi_staged = sm100_utils.make_smem_layout_epi( + self.io_dtype, + utils.LayoutEnum.ROW_MAJOR, + (self.BT, self.BK), + self.k_stage, + ) + v_epi_staged = sm100_utils.make_smem_layout_epi( + self.io_dtype, + utils.LayoutEnum.ROW_MAJOR, + (self.BT, self.BV), + self.v_tma_stage, + ) + gk_epi_staged = sm100_utils.make_smem_layout_epi( + self.acc_dtype, + utils.LayoutEnum.ROW_MAJOR, + (self.BT, self.BK), + self.g_stage, + ) + + # ---------- GMEM tensors (token-indexed) ---------- + # varlen: T=T_total, data_B=1 + # non-varlen: T=seq_len, data_B=B + A_layout = cute.make_layout( + (T, BT, (H, data_B)), + stride=(H * BT, 1, (BT, T * H * BT)), + ) + A_gmem = cute.make_tensor(A_ptr, A_layout) + + k_layout = cute.make_layout( + (T, K, (H, data_B)), + stride=(H * K, 1, (K, T * H * K)), + ) + k_gmem = cute.make_tensor(k_ptr, k_layout) + + v_layout = cute.make_layout( + (T, V, (H, data_B)), + stride=(H * V, 1, (V, T * H * V)), + ) + v_gmem = cute.make_tensor(v_ptr, v_layout) + + gk_layout = cute.make_layout( + (T, K, (H, data_B)), + stride=(H * K, 1, (K, T * H * K)), + ) + gk_gmem = cute.make_tensor(gk_ptr, gk_layout) + + # ---------- TMA descriptors ---------- + a_smem_one = cute.select(a_smem_staged, mode=[0, 1, 2]) + tma_atom_A, tma_tensor_A = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + A_gmem, + a_smem_one, + self.mma_tiler, + tiled_mma, + cluster_layout.shape, + ) + self.tma_A_bytes = cute.size_in_bytes(self.io_dtype, a_smem_one) + + k_epi_smem = cute.select(k_epi_staged, mode=[0, 1]) + tma_atom_k, tma_tensor_k = cpasync.make_tiled_tma_atom( + tma_load_op, + k_gmem, + k_epi_smem, + (self.BT, self.BK), + ) + + v_epi_smem = cute.select(v_epi_staged, mode=[0, 1]) + tma_atom_v, tma_tensor_v = cpasync.make_tiled_tma_atom( + tma_load_op, + v_gmem, + v_epi_smem, + (self.BT, self.BV), + ) + + gk_epi_smem = cute.select(gk_epi_staged, mode=[0, 1]) + if cutlass.const_expr(self.preprocessed_k): + # The preprocessed path does not consume gk. Reusing the k + # descriptor also lets the public wrapper pass k as the unused + # placeholder without constructing an fp32 tensor. + tma_atom_gk, tma_tensor_gk = tma_atom_k, tma_tensor_k + else: + tma_atom_gk, tma_tensor_gk = cpasync.make_tiled_tma_atom( + tma_load_op, + gk_gmem, + gk_epi_smem, + (self.BT, self.BK), + ) + + # ---------- TMA byte counts ---------- + self.tma_bytes_k = cute.size_in_bytes(self.io_dtype, k_epi_smem) + self.tma_bytes_v = cute.size_in_bytes(self.io_dtype, v_epi_smem) + self.tma_bytes_gk = cute.size_in_bytes(self.acc_dtype, gk_epi_smem) + # ---------- CUDA-core KG output staging ---------- + output_epi_staged = sm100_utils.make_smem_layout_epi( + self.io_dtype, + utils.LayoutEnum.ROW_MAJOR, + (self.BT, self.BN), + 1, + ) + + # ---------- SharedStorage ---------- + @cute.struct + class SharedStorage: + load_A_mbar: cute.struct.MemRange[Int64, self.a_stage * 2] + load_k_mbar: cute.struct.MemRange[Int64, self.k_stage * 2] + load_g_mbar: cute.struct.MemRange[Int64, self.g_stage * 2] + load_v_mbar: cute.struct.MemRange[Int64, self.v_tma_stage * 2] + beta_mbar: cute.struct.MemRange[Int64, self.beta_stage * 2] + prologue_ready_mbar: cute.struct.MemRange[Int64, self.bproc_stage * 2] + acc_mbar: cute.struct.MemRange[Int64, self.acc_pipe_stage * 2] + tmem_holding_buf: Int32 + sA: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(a_smem_staged)], + self.buffer_align_bytes, + ] + sB: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(b_smem_staged)], + self.buffer_align_bytes, + ] + sBV: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(b_smem_staged)], + self.buffer_align_bytes, + ] + sK: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(k_epi_staged)], + self.buffer_align_bytes, + ] + sV: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(v_epi_staged)], + self.buffer_align_bytes, + ] + sGK: cute.struct.Align[ + cute.struct.MemRange[self.acc_dtype, cute.cosize(gk_epi_staged)], + self.buffer_align_bytes, + ] + sOut: cute.struct.Align[ + cute.struct.MemRange[self.io_dtype, cute.cosize(output_epi_staged)], + self.buffer_align_bytes, + ] + sBeta: cute.struct.Align[ + cute.struct.MemRange[self.acc_dtype, self.beta_stage * self.BT], + 128, + ] + + self.shared_storage = SharedStorage + + # ---------- cu_seqlens / chunk_indices tensors ---------- + cu_seqlens = cute.make_tensor(cu_seqlens_ptr, cute.make_layout((B + 1,))) + chunk_indices = cute.make_tensor(chunk_indices_ptr, cute.make_layout((total_nt * 2,))) + + # ---------- Grid ---------- + # csrc launches exactly one persistent CTA per SM. + grid = (self.num_sm, 1, 1) + + self.kernel( + tiled_mma, + tma_atom_A, + tma_tensor_A, + tma_atom_k, + tma_tensor_k, + tma_atom_v, + tma_tensor_v, + tma_atom_gk, + tma_tensor_gk, + a_smem_staged, + b_smem_staged, + b_epi_staged, + k_epi_staged, + v_epi_staged, + gk_epi_staged, + output_epi_staged, + beta_ptr, + w_ptr, + u_ptr, + kg_ptr, + cu_seqlens, + chunk_indices, + problem_size, + total_nt, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + stream=stream, + min_blocks_per_mp=self.min_occupancy, + ) + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tma_atom_A: cute.CopyAtom, + tma_tensor_A: cute.Tensor, + tma_atom_k: cute.CopyAtom, + tma_tensor_k: cute.Tensor, + tma_atom_v: cute.CopyAtom, + tma_tensor_v: cute.Tensor, + tma_atom_gk: cute.CopyAtom, + tma_tensor_gk: cute.Tensor, + a_smem_staged: cute.ComposedLayout, + b_smem_staged: cute.ComposedLayout, + b_epi_staged: cute.ComposedLayout, + k_epi_staged: cute.ComposedLayout, + v_epi_staged: cute.ComposedLayout, + gk_epi_staged: cute.ComposedLayout, + output_epi_staged: cute.ComposedLayout, + beta_ptr: cute.Pointer, + w_ptr: cute.Pointer, + u_ptr: cute.Pointer, + kg_ptr: cute.Pointer, + cu_seqlens: cute.Tensor, + chunk_indices: cute.Tensor, + problem_size: tuple[Int32, Int32, Int32, Int32, Int32], + total_nt: Int32, + ): + B, T, H, K, V = problem_size + BT = self.BT + + # Match csrc StaticPersistentTileScheduler: one CTA per SM, + # tile_id = blockIdx.x + iteration * gridDim.x. + block_idx_x = cute.arch.block_idx()[0] + grid_dim_x = cute.arch.grid_dim()[0] + total_work_units = total_nt * H + num_iters = (total_work_units - block_idx_x + grid_dim_x - 1) // grid_dim_x + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + tidx, _, _ = cute.arch.thread_idx() + + if warp_idx == self.load_warp_id: + cpasync.prefetch_descriptor(tma_atom_A) + cpasync.prefetch_descriptor(tma_atom_k) + cpasync.prefetch_descriptor(tma_atom_v) + if cutlass.const_expr(not self.preprocessed_k): + cpasync.prefetch_descriptor(tma_atom_gk) + + # ---------- SMEM ---------- + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + sA = storage.sA.get_tensor(a_smem_staged.outer, swizzle=a_smem_staged.inner) + sBK = storage.sB.get_tensor(b_smem_staged.outer, swizzle=b_smem_staged.inner) + sBKTime = storage.sB.get_tensor(b_epi_staged.outer, swizzle=b_epi_staged.inner) + sBV = storage.sBV.get_tensor(b_smem_staged.outer, swizzle=b_smem_staged.inner) + sBVTime = storage.sBV.get_tensor(b_epi_staged.outer, swizzle=b_epi_staged.inner) + sK = storage.sK.get_tensor(k_epi_staged.outer, swizzle=k_epi_staged.inner) + sV = storage.sV.get_tensor(v_epi_staged.outer, swizzle=v_epi_staged.inner) + sGK = storage.sGK.get_tensor(gk_epi_staged.outer, swizzle=gk_epi_staged.inner) + sOut = storage.sOut.get_tensor(output_epi_staged.outer, swizzle=output_epi_staged.inner) + sBeta = cute.make_tensor( + cute.make_ptr(self.acc_dtype, storage.sBeta.data_ptr().toint(), cute.AddressSpace.smem), + cute.make_layout((self.BT, self.beta_stage), stride=(1, self.BT)), + ) + + # ---------- Pipelines ---------- + load_A_P, load_A_C = pipeline.PipelineTmaUmma.create( + num_stages=self.a_stage, + producer_group=_make_coop_group(1), + consumer_group=_make_coop_group(1), + tx_count=self.tma_A_bytes, + barrier_storage=storage.load_A_mbar.data_ptr(), + defer_sync=True, + ).make_participants() + + load_k_P, load_k_C = pipeline.PipelineTmaAsync.create( + num_stages=self.k_stage, + producer_group=_make_coop_group(1), + consumer_group=_make_coop_group(self.num_cuda_warps), + tx_count=self.tma_bytes_k, + barrier_storage=storage.load_k_mbar.data_ptr(), + defer_sync=True, + ).make_participants() + + load_g_P, load_g_C = pipeline.PipelineTmaAsync.create( + num_stages=self.g_stage, + producer_group=_make_coop_group(1), + consumer_group=_make_coop_group(self.num_cuda_warps), + tx_count=self.tma_bytes_gk, + barrier_storage=storage.load_g_mbar.data_ptr(), + defer_sync=True, + ).make_participants() + + load_v_P, load_v_C = pipeline.PipelineTmaAsync.create( + num_stages=self.v_tma_stage, + producer_group=_make_coop_group(1), + consumer_group=_make_coop_group(self.num_cuda_warps), + tx_count=self.tma_bytes_v, + barrier_storage=storage.load_v_mbar.data_ptr(), + defer_sync=True, + ).make_participants() + + beta_P, beta_C = pipeline.PipelineAsync.create( + num_stages=self.beta_stage, + producer_group=_make_coop_group(2 * self.threads_per_warp), + consumer_group=_make_coop_group(2 * self.num_cuda_threads), + barrier_storage=storage.beta_mbar.data_ptr(), + defer_sync=True, + ).make_participants() + + prologue_ready_P, prologue_ready_C = pipeline.PipelineAsyncUmma.create( + num_stages=self.bproc_stage, + producer_group=_make_coop_group(2 * self.num_cuda_threads), + consumer_group=_make_coop_group(1), + barrier_storage=storage.prologue_ready_mbar.data_ptr(), + defer_sync=True, + ).make_participants() + + acc_done_P, acc_done_C = pipeline.PipelineUmmaAsync.create( + num_stages=self.acc_pipe_stage, + producer_group=_make_coop_group(1), + consumer_group=_make_coop_group(self.num_cuda_threads), + barrier_storage=storage.acc_mbar.data_ptr(), + ).make_participants() + + # ---------- TMEM ---------- + tmem_alloc_bar = pipeline.NamedBarrier(barrier_id=1, num_threads=self.threads_per_cta) + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_bar, + allocator_warp_id=self.load_warp_id, + ) + tmem.allocate(self.tmem_total) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + tCrA = tiled_mma.make_fragment_A(sA) + tCrBK = tiled_mma.make_fragment_B(sBK) + tCrBV = tiled_mma.make_fragment_B(sBV) + + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(acc_shape) + tCtAccW = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + tCtAccU = cute.make_tensor(tmem_ptr + 16 * 65536, tCtAcc_fake.layout) + # Layout-only staged view used to derive the same 128-thread R2S + # mapping as the MMA output. W/U stores do not read through this view. + tCtAccMap_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, 2)) + tCtAccMap = cute.make_tensor(tmem_ptr, tCtAccMap_fake.layout) + + # ===================================================================== + # LOAD WARP + # ===================================================================== + if warp_idx == self.load_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_others) + + for wu_iter in cutlass.range(0, num_iters, unroll=0): + work_idx = block_idx_x + wu_iter * grid_dim_x + i_b, i_t, i_h, tok_offset, data_bidx, remaining = self._decode_persistent_work( + work_idx, total_nt, H, T, BT, cu_seqlens, chunk_indices + ) + # --- Domain offset (varlen) or alias (non-varlen) --- + if cutlass.const_expr(self.is_varlen): + tma_k_v = cute.domain_offset((tok_offset, 0, (0, 0)), tma_tensor_k) + tma_v_v = cute.domain_offset((tok_offset, 0, (0, 0)), tma_tensor_v) + tma_gk_v = cute.domain_offset((tok_offset, 0, (0, 0)), tma_tensor_gk) + tma_A_v = cute.domain_offset((tok_offset, 0, (0, 0)), tma_tensor_A) + else: + tma_k_v = tma_tensor_k + tma_v_v = tma_tensor_v + tma_gk_v = tma_tensor_gk + tma_A_v = tma_tensor_A + + # --- TMA partitions --- + if cutlass.const_expr(self.is_varlen): + bSG_sK, bSG_gK = self._epilog_partition_varlen( + tma_atom_k, + tma_k_v[None, None, (i_h, data_bidx)], + (self.BT, self.BK), + sK, + ) + bSG_sV, bSG_gV = self._epilog_partition_varlen( + tma_atom_v, + tma_v_v[None, None, (i_h, data_bidx)], + (self.BT, self.BV), + sV, + ) + if cutlass.const_expr(not self.preprocessed_k): + bSG_sGK, bSG_gGK = self._epilog_partition_varlen( + tma_atom_gk, + tma_gk_v[None, None, (i_h, data_bidx)], + (self.BT, self.BK), + sGK, + ) + tAsA, tAgA = self._tma_partition_A( + tma_atom_A, + tma_A_v, + sA, + self.mma_tiler, + tiled_mma, + data_bidx, + i_h, + ) + else: + bSG_sK, bSG_gK = self._data_tma_partition( + tma_atom_k, + tma_k_v, + (self.BT, self.BK), + sK, + i_h, + data_bidx, + ) + bSG_sV, bSG_gV = self._data_tma_partition( + tma_atom_v, + tma_v_v, + (self.BT, self.BV), + sV, + i_h, + data_bidx, + ) + if cutlass.const_expr(not self.preprocessed_k): + bSG_sGK, bSG_gGK = self._data_tma_partition( + tma_atom_gk, + tma_gk_v, + (self.BT, self.BK), + sGK, + i_h, + data_bidx, + ) + tAsA, tAgA = self._tma_partition_A( + tma_atom_A, + tma_A_v, + sA, + self.mma_tiler, + tiled_mma, + data_bidx, + i_h, + ) + + # --- Issue TMA loads --- + h_a = load_A_P.acquire_and_advance() + cute.copy( + tma_atom_A, + tAgA[(None, i_t, 0)], + tAsA[(None, h_a.index)], + tma_bar_ptr=h_a.barrier, + ) + + for i_kv in cutlass.range(0, self.NK): + k_h = load_k_P.acquire_and_advance() + cute.copy( + tma_atom_k, + bSG_gK[(None, i_t, i_kv)], + bSG_sK[None, k_h.index], + tma_bar_ptr=k_h.barrier, + ) + if cutlass.const_expr(not self.preprocessed_k): + g_h = load_g_P.acquire_and_advance() + cute.copy( + tma_atom_gk, + bSG_gGK[(None, i_t, i_kv)], + bSG_sGK[None, g_h.index], + tma_bar_ptr=g_h.barrier, + ) + v_h = load_v_P.acquire_and_advance() + cute.copy( + tma_atom_v, + bSG_gV[(None, i_t, i_kv)], + bSG_sV[None, v_h.index], + tma_bar_ptr=v_h.barrier, + ) + + # ===================================================================== + # AUX WARPS — 64-thread beta producer, matching csrc LoadAux + # ===================================================================== + elif warp_idx in self.aux_warp_ids: + cute.arch.warpgroup_reg_dealloc(self.num_regs_others) + aux_tidx = tidx % (2 * self.threads_per_warp) + for wu_iter in cutlass.range(0, num_iters, unroll=0): + work_idx = block_idx_x + wu_iter * grid_dim_x + i_b, i_t, i_h, tok_offset, data_bidx, remaining = self._decode_persistent_work( + work_idx, total_nt, H, T, BT, cu_seqlens, chunk_indices + ) + beta_h = beta_P.acquire_and_advance() + beta_base = (tok_offset + i_t * BT) * H + i_h + beta_gmem = cute.make_tensor( + cute.make_ptr( + self.beta_dtype, + (beta_ptr + beta_base).toint(), + cute.AddressSpace.gmem, + assumed_align=2, + ), + cute.make_layout((self.BT,), stride=(H,)), + ) + sBeta[(aux_tidx, beta_h.index)] = cutlass.select_( + aux_tidx < remaining, + beta_gmem[aux_tidx].to(self.acc_dtype), + self.acc_dtype(0.0), + ) + cute.arch.fence_proxy("async.shared", space="cta") + beta_h.commit() + + # ===================================================================== + # MMA WARP + # ===================================================================== + elif warp_idx == self.mma_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_others) + + num_kblks = cute.size(tCrBK, mode=[2]) + + for wu_iter in cutlass.range(0, num_iters, unroll=0): + work_idx = block_idx_x + wu_iter * grid_dim_x + i_b, i_t, i_h, tok_offset, data_bidx, remaining = self._decode_persistent_work( + work_idx, total_nt, H, T, BT, cu_seqlens, chunk_indices + ) + # Wait for A_mat — hold handle until all GEMMs finish reading sA + a_h = load_A_C.wait_and_advance() + + for i_kv in cutlass.range(0, self.NK): + # Match csrc: dispatch W and U into two fixed TMEM regions and + # publish a single completion generation for the pair. + bp_h = prologue_ready_C.wait_and_advance() + acc_h = acc_done_P.acquire_and_advance() + for kblk in cutlass.range(num_kblks, unroll_full=True): + tiled_mma.set(tcgen05.Field.ACCUMULATE, cutlass.Boolean(kblk != 0)) + cute.gemm( + tiled_mma, + tCtAccW, + tCrA[(None, None, kblk, a_h.index)], + tCrBK[(None, None, kblk, bp_h.index)], + tCtAccW, + ) + + for kblk in cutlass.range(num_kblks, unroll_full=True): + tiled_mma.set(tcgen05.Field.ACCUMULATE, cutlass.Boolean(kblk != 0)) + cute.gemm( + tiled_mma, + tCtAccU, + tCrA[(None, None, kblk, a_h.index)], + tCrBV[(None, None, kblk, bp_h.index)], + tCtAccU, + ) + acc_h.commit() + bp_h.release() + + # Release A after all GEMMs that read sA are dispatched + a_h.release() + + # ===================================================================== + # K PROLOGUE + KG OUTPUT WARPGROUP + # ===================================================================== + elif warp_idx in self.prologue_warp_ids: + cute.arch.warpgroup_reg_alloc(self.num_regs_prologue) + local_tidx = tidx % self.num_cuda_threads + t2r_atom = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition(8), tcgen05.Pack.NONE), self.acc_dtype) + tCtAcc_flat = tCtAccMap[((None, None), 0, 0, None)] + fake_out = cute.make_tensor( + cute.make_ptr(self.io_dtype, 0, cute.AddressSpace.smem), cute.dice(self.mma_tiler, (1, 1, None)) + ) + tiled_t2r = tcgen05.make_tmem_copy(t2r_atom, tCtAcc_flat[(None, None, 0)]) + thr_t2r = tiled_t2r.get_slice(local_tidx) + tTR_sOut = thr_t2r.partition_D(fake_out) + tTR_cM = thr_t2r.partition_D(cute.make_identity_tensor(cute.dice(self.mma_tiler, (1, 1, None)))) + r2s_atom = sm100_utils.get_smem_store_op(utils.LayoutEnum.ROW_MAJOR, self.io_dtype, self.acc_dtype, tiled_t2r) + tiled_r2s = cute.make_tiled_copy_D(r2s_atom, tiled_t2r) + thr_r2s = tiled_r2s.get_slice(local_tidx) + tRS_sBK = thr_r2s.partition_D(sBKTime) + tRS_sOut = thr_r2s.partition_D(sOut[(None, None, 0)]) + r_bproc = cute.make_rmem_tensor(tTR_sOut.shape, self.io_dtype) + r_kg = cute.make_rmem_tensor(tTR_sOut.shape, self.io_dtype) + + # csrc GmemTiledCopyO: 128 threads, 8 bf16 values per vector. + async_copy_elems = 128 // self.io_dtype.width + atom_ucopy = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.io_dtype, num_bits_per_copy=128) + n_threads_bn = self.BN // async_copy_elems + n_threads_bt = self.num_cuda_threads // n_threads_bn + r2g_tiled_copy = cute.make_tiled_copy_tv( + atom_ucopy, + cute.make_ordered_layout((n_threads_bt, n_threads_bn), order=(1, 0)), + cute.make_layout((1, async_copy_elems)), + ) + r2g_thr_copy = r2g_tiled_copy.get_slice(local_tidx) + tOsOut = r2g_thr_copy.partition_S(sOut[(None, None, 0)]) + cOut = cute.make_identity_tensor((self.BT, self.BN)) + tOcOut = r2g_thr_copy.partition_S(cOut) + r2g_frag = cute.make_fragment_like(tOsOut[None, 0, None], self.io_dtype) + kg_sync = pipeline.NamedBarrier(barrier_id=2, num_threads=self.num_cuda_threads) + + for wu_iter in cutlass.range(0, num_iters, unroll=0): + work_idx = block_idx_x + wu_iter * grid_dim_x + i_b, i_t, i_h, tok_offset, data_bidx, remaining = self._decode_persistent_work( + work_idx, total_nt, H, T, BT, cu_seqlens, chunk_indices + ) + beta_h = beta_C.wait_and_advance() + token_base = tok_offset + i_t * BT + kg_base = token_base * H * K + i_h * K + gKg = cute.make_tensor( + cute.make_ptr(self.io_dtype, (kg_ptr + kg_base).toint(), cute.AddressSpace.gmem, assumed_align=16), + cute.make_layout( + (self.BT, self.BK), + stride=(cute.assume(H * K, divby=async_copy_elems), 1), + ), + ) + tOgKg = r2g_thr_copy.partition_D(gKg) + + for i_kv in cutlass.range(0, self.NK): + k_h = load_k_C.wait_and_advance() + for ei in cutlass.range_constexpr(cute.size(tTR_cM)): + m_coord, n_coord = tTR_cM[ei] + k_val = sK[(m_coord, n_coord, k_h.index)].to(self.acc_dtype) + beta_val = sBeta[(m_coord, beta_h.index)] + if cutlass.const_expr(self.preprocessed_k): + r_bproc[ei] = (k_val * beta_val).to(self.io_dtype) + else: + # csrc loads K into registers before waiting for G, + # overlapping the independent TMA pipelines. + r_bproc[ei] = k_val.to(self.io_dtype) + + if cutlass.const_expr(not self.preprocessed_k): + g_h = load_g_C.wait_and_advance() + for ei in cutlass.range_constexpr(cute.size(tTR_cM)): + m_coord, n_coord = tTR_cM[ei] + k_val = r_bproc[ei].to(self.acc_dtype) + beta_val = sBeta[(m_coord, beta_h.index)] + g_val = sGK[(m_coord, n_coord, g_h.index)] + last_row = cutlass.select_(remaining < Int32(BT), remaining - 1, Int32(BT - 1)) + gn_val = sGK[(last_row, n_coord, g_h.index)] + bp = (k_val * beta_val * cute.exp2(g_val, fastmath=self.use_fast_math)).to(self.io_dtype) + kg = (k_val * cute.exp2(gn_val - g_val, fastmath=self.use_fast_math)).to(self.io_dtype) + if cutlass.const_expr(self.is_varlen): + r_bproc[ei] = cutlass.select_(m_coord < remaining, bp, self.io_dtype(0.0)) + r_kg[ei] = cutlass.select_(m_coord < remaining, kg, self.io_dtype(0.0)) + else: + r_bproc[ei] = bp + r_kg[ei] = kg + + bp_h = prologue_ready_P.acquire_and_advance() + r2s_b = tiled_r2s.retile(r_bproc) + cute.copy(tiled_r2s, r2s_b, tRS_sBK[(None, None, None, bp_h.index)]) + cute.arch.fence_proxy("async.shared", space="cta") + bp_h.commit() + k_h.release() + + if cutlass.const_expr(not self.preprocessed_k): + kg_sync.arrive_and_wait() + r2s_kg = tiled_r2s.retile(r_kg) + cute.copy(tiled_r2s, r2s_kg, tRS_sOut) + g_h.release() + kg_sync.arrive_and_wait() + for m1 in cutlass.range_constexpr(cute.size(tOsOut.shape[1])): + row = tOcOut[(0, 0), m1, 0][0] + if row < remaining: + cute.autovec_copy(tOsOut[None, m1, None], r2g_frag) + cute.autovec_copy(r2g_frag, tOgKg[None, m1, None]) + + beta_h.release() + + # ===================================================================== + # V PROLOGUE + W/U EPILOGUE WARPGROUP + # ===================================================================== + elif warp_idx in self.epilogue_warp_ids: + cute.arch.warpgroup_reg_alloc(self.num_regs_epilogue) + local_tidx = tidx % self.num_cuda_threads + t2r_atom = cute.make_copy_atom(tcgen05.Ld16x256bOp(tcgen05.Repetition(8), tcgen05.Pack.NONE), self.acc_dtype) + tCtAcc_flat = tCtAccMap[((None, None), 0, 0, None)] + fake_out = cute.make_tensor( + cute.make_ptr(self.io_dtype, 0, cute.AddressSpace.smem), cute.dice(self.mma_tiler, (1, 1, None)) + ) + tiled_t2r = tcgen05.make_tmem_copy(t2r_atom, tCtAcc_flat[(None, None, 0)]) + thr_t2r = tiled_t2r.get_slice(local_tidx) + tTR_sOut = thr_t2r.partition_D(fake_out) + tTR_cM = thr_t2r.partition_D(cute.make_identity_tensor(cute.dice(self.mma_tiler, (1, 1, None)))) + r2s_atom = sm100_utils.get_smem_store_op(utils.LayoutEnum.ROW_MAJOR, self.io_dtype, self.acc_dtype, tiled_t2r) + tiled_r2s = cute.make_tiled_copy_D(r2s_atom, tiled_t2r) + thr_r2s = tiled_r2s.get_slice(local_tidx) + tRS_sBV = thr_r2s.partition_D(sBVTime) + r_bproc = cute.make_rmem_tensor(tTR_sOut.shape, self.io_dtype) + + # Exact csrc dp-lane mapping: alternating 16-lane groups store W/U, + # and four warps cover the 64 output rows. + acc_idx = (local_tidx // 16) & 1 + output_row = (local_tidx // 32) * 16 + (local_tidx % 16) + quar_k = self.BN // 4 + + for wu_iter in cutlass.range(0, num_iters, unroll=0): + work_idx = block_idx_x + wu_iter * grid_dim_x + i_b, i_t, i_h, tok_offset, data_bidx, remaining = self._decode_persistent_work( + work_idx, total_nt, H, T, BT, cu_seqlens, chunk_indices + ) + beta_h = beta_C.wait_and_advance() + for i_kv in cutlass.range(0, self.NK): + v_h = load_v_C.wait_and_advance() + for ei in cutlass.range_constexpr(cute.size(tTR_cM)): + m_coord, n_coord = tTR_cM[ei] + value = (sV[(m_coord, n_coord, v_h.index)].to(self.acc_dtype) * sBeta[(m_coord, beta_h.index)]).to( + self.io_dtype + ) + if cutlass.const_expr(self.is_varlen): + r_bproc[ei] = cutlass.select_(m_coord < remaining, value, self.io_dtype(0.0)) + else: + r_bproc[ei] = value + bp_h = prologue_ready_P.acquire_and_advance() + r2s_b = tiled_r2s.retile(r_bproc) + cute.copy(tiled_r2s, r2s_b, tRS_sBV[(None, None, None, bp_h.index)]) + cute.arch.fence_proxy("async.shared", space="cta") + bp_h.commit() + v_h.release() + + wu_h = acc_done_C.wait_and_advance() + token_base = tok_offset + i_t * BT + output_addr = Int64(0) + if acc_idx == 0: + output_addr = (w_ptr + (token_base + output_row) * H * K + i_h * K + i_kv * self.BK).toint() + else: + output_addr = (u_ptr + (token_base + output_row) * H * V + i_h * V + i_kv * self.BV).toint() + + tcgen05_fence_after() + for quar in cutlass.range_constexpr(4): + acc_i32 = tcgen05_ld_32x32b( + quar_k, + acc_idx * (16 * 65536) + wu_h.index * 256 + quar * quar_k, + ) + cute.arch.fence_view_async_tmem_load() + if output_row < remaining: + acc_f32 = reinterpret_cast(acc_i32, Int32, quar_k, self.acc_dtype) + acc_bf16 = TensorSSA(acc_f32, (quar_k,), self.acc_dtype).to(self.io_dtype) + out_i32 = reinterpret_cast(acc_bf16, self.io_dtype, quar_k, Int32) + for store_idx in cutlass.range_constexpr(quar_k // 16): + store_256b( + output_addr + quar * quar_k * 2 + store_idx * 32, + subvec(out_i32, store_idx * 8, 8), + ) + tcgen05_fence_before() + wu_h.release() + + beta_h.release() + + # ---------- TMEM cleanup ---------- + tmem.relinquish_alloc_permit() + pipeline.sync(barrier_id=1) + tmem.free(tmem_ptr) + + +# ============================================================================ +# Compile cache +# ============================================================================ + +_recompute_wu_cache = {} +_dummy_cu_seqlens = None +_dummy_chunk_indices = None +_uniform_cu_cache = {} + + +def _uniform_problem(cu_seqlens: torch.Tensor) -> tuple[int, int] | None: + """Return ``(batch, sequence_length)`` for a uniform packed batch.""" + key = id(cu_seqlens) + entry = _uniform_cu_cache.get(key) + if entry is None or entry[0]() is not cu_seqlens or entry[1] != cu_seqlens._version: + offsets = cu_seqlens.detach().cpu().tolist() + lengths = [end - start for start, end in zip(offsets, offsets[1:])] + result = (len(lengths), lengths[0]) if lengths and len(set(lengths)) == 1 else None + _uniform_cu_cache[key] = (weakref.ref(cu_seqlens), cu_seqlens._version, result) + return result + return entry[2] + + +def _compile_recompute_wu( + H, + K, + V, + chunk_size=64, + block_k=None, + block_v=None, + persistent=True, + is_varlen=False, + beta_dtype=cutlass.Float32, + preprocessed_k=False, +): + key = ( + H, + K, + V, + chunk_size, + block_k, + block_v, + persistent, + is_varlen, + beta_dtype, + preprocessed_k, + USE_FAST_MATH, + ) + if key in _recompute_wu_cache: + return _recompute_wu_cache[key] + + kernel_obj = KDARecomputeWU( + K=K, + V=V, + chunk_size=chunk_size, + block_k=block_k, + block_v=block_v, + beta_dtype=beta_dtype, + is_varlen=is_varlen, + preprocessed_k=preprocessed_k, + use_fast_math=USE_FAST_MATH, + ) + + sym_a = cute.sym_int() + sym_b = cute.sym_int() + sym_cu = cute.sym_int() + sym_ci = cute.sym_int() + BT = chunk_size + + if is_varlen: + k_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_a, H, K), stride_order=(2, 1, 0), assumed_align=128) + v_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_a, H, V), stride_order=(2, 1, 0), assumed_align=128) + beta_fake = make_fake_compact_tensor(beta_dtype, (sym_a, H), stride_order=(1, 0), assumed_align=128) + A_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_a, H, BT), stride_order=(2, 1, 0), assumed_align=128) + gk_fake = make_fake_compact_tensor( + cutlass.BFloat16 if preprocessed_k else cutlass.Float32, + (sym_a, H, K), + stride_order=(2, 1, 0), + assumed_align=128, + ) + w_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_a, H, K), stride_order=(2, 1, 0), assumed_align=128) + u_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_a, H, V), stride_order=(2, 1, 0), assumed_align=128) + kg_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_a, H, K), stride_order=(2, 1, 0), assumed_align=128) + else: + k_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_a, sym_b, H, K), stride_order=(3, 2, 1, 0), assumed_align=128) + v_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_a, sym_b, H, V), stride_order=(3, 2, 1, 0), assumed_align=128) + beta_fake = make_fake_compact_tensor(beta_dtype, (sym_a, sym_b, H), stride_order=(2, 1, 0), assumed_align=128) + A_fake = make_fake_compact_tensor( + cutlass.BFloat16, (sym_a, sym_b, H, BT), stride_order=(3, 2, 1, 0), assumed_align=128 + ) + gk_fake = make_fake_compact_tensor( + cutlass.BFloat16 if preprocessed_k else cutlass.Float32, + (sym_a, sym_b, H, K), + stride_order=(3, 2, 1, 0), + assumed_align=128, + ) + w_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_a, sym_b, H, K), stride_order=(3, 2, 1, 0), assumed_align=128) + u_fake = make_fake_compact_tensor(cutlass.BFloat16, (sym_a, sym_b, H, V), stride_order=(3, 2, 1, 0), assumed_align=128) + kg_fake = make_fake_compact_tensor( + cutlass.BFloat16, (sym_a, sym_b, H, K), stride_order=(3, 2, 1, 0), assumed_align=128 + ) + + cu_fake = make_fake_compact_tensor(cutlass.Int32, (sym_cu,), assumed_align=128) + ci_fake = make_fake_compact_tensor(cutlass.Int32, (sym_ci,), assumed_align=128) + stream_fake = make_fake_stream(use_tvm_ffi_env_stream=True) + + compiled_fn = cute.compile( + kernel_obj, + k_fake, + v_fake, + beta_fake, + A_fake, + gk_fake, + w_fake, + u_fake, + kg_fake, + cu_fake, + ci_fake, + (Int32(1), Int32(1), Int32(H), Int32(K), Int32(V)), + Int32(1), + stream_fake, + options="--enable-tvm-ffi", + ) + _recompute_wu_cache[key] = compiled_fn + return compiled_fn + + +# ============================================================================ +# Public API +# ============================================================================ + + +def recompute_w_u_fwd(k, v, beta, A, gk, cu_seqlens=None, chunk_indices=None, block_k=None, block_v=None): + is_varlen = cu_seqlens is not None + packed_4d = is_varlen and k.dim() == 4 + restore_packed = False + if packed_4d: + if k.shape[0] != 1: + raise ValueError("varlen inputs must be packed with batch dimension 1") + k, v, beta, A, gk = (x.squeeze(0) for x in (k, v, beta, A, gk)) + + uniform = _uniform_problem(cu_seqlens) + if uniform is not None and uniform[1] % A.shape[-1] == 0: + batch, seq_len = uniform + k = k.view(batch, seq_len, *k.shape[1:]) + v = v.view(batch, seq_len, *v.shape[1:]) + beta = beta.view(batch, seq_len, *beta.shape[1:]) + A = A.view(batch, seq_len, *A.shape[1:]) + gk = gk.view(batch, seq_len, *gk.shape[1:]) + is_varlen = False + packed_4d = False + restore_packed = True + + if is_varlen: + BT = A.shape[-1] + T_total, H, K = k.shape + V = v.shape[2] + num_seqs = cu_seqlens.shape[0] - 1 + + # Single-seq varlen with aligned T → dispatch as non-varlen for TMA S2G speed + if num_seqs == 1 and T_total % BT == 0: + k_4d = k.unsqueeze(0) + v_4d = v.unsqueeze(0) + beta_4d = beta.unsqueeze(0) + A_4d = A.unsqueeze(0) + gk_4d = gk.unsqueeze(0) + w_4d, u_4d, _, kg_4d = recompute_w_u_fwd( + k_4d, + v_4d, + beta_4d, + A_4d, + gk_4d, + block_k=block_k, + block_v=block_v, + ) + return w_4d.squeeze(0), u_4d.squeeze(0), None, kg_4d.squeeze(0) + + if chunk_indices is not None: + ci_s = chunk_indices.reshape(-1) + else: + ci_s = prepare_chunk_indices(cu_seqlens, BT).reshape(-1) + + total_nt = ci_s.shape[0] // 2 + ps = (Int32(num_seqs), Int32(T_total), Int32(H), Int32(K), Int32(V)) + cu_s = cu_seqlens + else: + B, T, H, K = k.shape + V = v.shape[-1] + BT = A.shape[-1] + NT = (T + BT - 1) // BT + total_nt = B * NT + ps = (Int32(B), Int32(T), Int32(H), Int32(K), Int32(V)) + global _dummy_cu_seqlens, _dummy_chunk_indices + if _dummy_cu_seqlens is None or _dummy_cu_seqlens.device != k.device: + _dummy_cu_seqlens = torch.zeros(2, dtype=torch.int32, device=k.device) + if _dummy_chunk_indices is None or _dummy_chunk_indices.device != k.device: + _dummy_chunk_indices = torch.zeros(2, dtype=torch.int32, device=k.device) + cu_s = _dummy_cu_seqlens + ci_s = _dummy_chunk_indices + + w = torch.empty_like(k) + u = torch.empty_like(v) + kg = torch.empty_like(k) + + compiled_fn = _compile_recompute_wu( + H, + K, + V, + chunk_size=BT, + block_k=block_k, + block_v=block_v, + is_varlen=is_varlen, + beta_dtype=cutlass.Float32 if beta.dtype == torch.float32 else cutlass.BFloat16, + ) + + compiled_fn(k, v, beta, A, gk, w, u, kg, cu_s, ci_s, ps, Int32(total_nt)) + + if restore_packed: + w, u, kg = (x.flatten(0, 1).unsqueeze(0) for x in (w, u, kg)) + elif packed_4d: + w, u, kg = (x.unsqueeze(0) for x in (w, u, kg)) + return w, u, None, kg + + +def recompute_w_u_from_preprocessed( + k_scaled, + v, + beta, + A, + cu_seqlens=None, + chunk_indices=None, + block_k=None, + block_v=None, +): + """Compute only ``w`` and ``u`` when fused intra already produced scaled k. + + ``k_scaled`` is ``k * exp2(gk)`` and ``A`` is the inverted intra-chunk + matrix. The companion fused intra kernel already produced ``kg``, so this + path avoids loading the fp32 cumulative gate and writing ``kg`` again. + """ + is_varlen = cu_seqlens is not None + packed_4d = is_varlen and k_scaled.dim() == 4 + restore_packed = False + if packed_4d: + if k_scaled.shape[0] != 1: + raise ValueError("varlen inputs must be packed with batch dimension 1") + k_scaled, v, beta, A = (x.squeeze(0) for x in (k_scaled, v, beta, A)) + + uniform = _uniform_problem(cu_seqlens) + if uniform is not None and uniform[1] % A.shape[-1] == 0: + batch, seq_len = uniform + k_scaled = k_scaled.view(batch, seq_len, *k_scaled.shape[1:]) + v = v.view(batch, seq_len, *v.shape[1:]) + beta = beta.view(batch, seq_len, *beta.shape[1:]) + A = A.view(batch, seq_len, *A.shape[1:]) + is_varlen = False + packed_4d = False + restore_packed = True + + if is_varlen: + BT = A.shape[-1] + T_total, H, K = k_scaled.shape + V = v.shape[2] + num_seqs = cu_seqlens.shape[0] - 1 + ci_s = chunk_indices.reshape(-1) if chunk_indices is not None else prepare_chunk_indices(cu_seqlens, BT).reshape(-1) + total_nt = ci_s.shape[0] // 2 + ps = (Int32(num_seqs), Int32(T_total), Int32(H), Int32(K), Int32(V)) + cu_s = cu_seqlens + else: + B, T, H, K = k_scaled.shape + V = v.shape[-1] + BT = A.shape[-1] + total_nt = B * ((T + BT - 1) // BT) + ps = (Int32(B), Int32(T), Int32(H), Int32(K), Int32(V)) + global _dummy_cu_seqlens, _dummy_chunk_indices + if _dummy_cu_seqlens is None or _dummy_cu_seqlens.device != k_scaled.device: + _dummy_cu_seqlens = torch.zeros(2, dtype=torch.int32, device=k_scaled.device) + if _dummy_chunk_indices is None or _dummy_chunk_indices.device != k_scaled.device: + _dummy_chunk_indices = torch.zeros(2, dtype=torch.int32, device=k_scaled.device) + cu_s = _dummy_cu_seqlens + ci_s = _dummy_chunk_indices + + w = torch.empty_like(k_scaled) + u = torch.empty_like(v) + compiled_fn = _compile_recompute_wu( + H, + K, + V, + chunk_size=BT, + block_k=block_k, + block_v=block_v, + is_varlen=is_varlen, + beta_dtype=cutlass.Float32 if beta.dtype == torch.float32 else cutlass.BFloat16, + preprocessed_k=True, + ) + + # gk and kg are unused compile-signature placeholders for this + # specialization. Reusing existing bf16 buffers avoids extra allocations. + compiled_fn(k_scaled, v, beta, A, k_scaled, w, u, w, cu_s, ci_s, ps, Int32(total_nt)) + + if restore_packed: + w, u = (x.flatten(0, 1).unsqueeze(0) for x in (w, u)) + elif packed_4d: + w, u = (x.unsqueeze(0) for x in (w, u)) + return w, u diff --git a/docs/kda_sm100_cutedsl_fwd_status.md b/docs/kda_sm100_cutedsl_fwd_status.md new file mode 100644 index 00000000..64a6cb7e --- /dev/null +++ b/docs/kda_sm100_cutedsl_fwd_status.md @@ -0,0 +1,250 @@ +# SM100 CuTeDSL KDA forward status + +Branch: `icavan/cutedsl-sm100-fwd` + +Validation hardware: NVIDIA GB200 (SM100). + +## Implemented candidates + +- A csrc-aligned persistent CuTeDSL `recompute_w_u`: 384 threads, separate + A/K/G/V and beta pipelines, co-produced K/V MMA-ready barrier, dp-lane + aliased W/U TMEM accumulators, direct CUDA-core W/U stores, and vectorized + KG stores. Equal-length, packed-uniform, and varlen dispatch are supported. +- A preprocessed specialization that consumes `k_scaled` and skips the fp32 + cumulative-gate load and duplicate `kg` store. +- An FP16 `k * exp2(gk)` workspace specialization. It keeps the same two-byte + footprint as BF16, lets WU apply beta before the final BF16 MMA-operand + rounding, and avoids adding a multiply to the K123 critical path. +- Fused raw-gate K1/K2/K3 intra candidate plus standalone fp32 Akk inverse. +- A csrc-boundary CuTeDSL intra entry point that consumes the same precomputed + FP32 `gk` tensor as `chunk_kda_fwd_intra_cuda`; no gate activation or cumsum + is fused into this correctness baseline. Its FP32-workspace inverse uses the + csrc single-accumulator K=32 TF32 Schur order. +- A zero-copy packed-uniform route through the equal-length kernel. +- Same-input benchmarks against the repository SM100 csrc kernels. + +The public csrc dispatch remains unchanged. The standalone recompute-WU port is +bitwise equal to csrc for W, U, and KG in the FP32- and BF16-beta tests. At the +representative T=8192 shape it is about 4% faster than csrc. + +The csrc-boundary intra plus recompute-WU path is also bitwise equal for Aqk, +Akk, KG, W, and U. At `B=2,T=8192,H=64,K=V=128` it runs in 0.8645 ms versus +0.9157 ms for csrc, or 1.059x csrc throughput. + +## GB200 results + +The representative shape is BF16 `B=2, H=64, K=V=128, chunk_size=64`. + +### Modular recompute WU + +| T | csrc (ms) | CuTeDSL (ms) | csrc / CuTeDSL | +|---:|---:|---:|---:| +| 512 | 0.0328 | 0.0704 | 0.466x | +| 1024 | 0.0517 | 0.0782 | 0.662x | +| 4096 | 0.1702 | 0.1613 | 1.055x | +| 8192 | 0.3261 | 0.3129 | 1.042x | +| 16384 | 0.6376 | 0.6222 | 1.025x | +| 32768 | 1.2618 | 1.2534 | 1.007x | + +The standalone API computes the fp32 gate transform and writes `kg`, exactly as +the csrc baseline. The CuTeDSL runtime has a visible fixed launch/descriptor +cost at T=512 and T=1024; from T=4096 through T=32768 the aligned persistent +kernel meets or exceeds csrc throughput. The former staged-output/store-warp +CuTeDSL implementation has been removed. + +### Csrc-boundary intra plus recompute WU + +This comparison consumes the same precomputed FP32 `gk` tensor on both sides +and includes intra, the Akk inverse, and recompute WU. It uses three warmup +iterations and 20 CUDA-Event-timed iterations. + +| Shape | csrc (ms) | CuTeDSL (ms) | csrc / CuTeDSL | +|---|---:|---:|---:| +| B=2, T=8192, H=64 | 0.9157 | 0.8645 | 1.059x | + +For the measured input, `torch.equal` passes for every complete Aqk, Akk, KG, +W, and U tensor; each tensor has zero mismatched elements and zero maximum +absolute difference. + +### Raw-gate fused intra plus specialized WU + +The table below measures the complete csrc gate + intra + recompute chain +against CuTeDSL fused K1/K2/K3 + fp32 inverse + FP16-workspace preprocessed +W/U. Each point uses 10 warmup iterations and 30 measured iterations. + +| Shape | csrc (us) | CuTeDSL (us) | csrc / CuTeDSL | +|---|---:|---:|---:| +| equal, T=4096 | 554.7 | 554.3 | 1.001x | +| equal, T=8192 | 1072.9 | 1065.4 | 1.007x | +| equal, T=16384 | 2101.4 | 2139.0 | 0.982x | +| varlen, lengths=4096,4096 | 561.6 | 558.9 | 1.005x | + +Representative relative RMSE against csrc at T=8192 is: + +- `Aqk`: 4.112e-4 +- `Akk`: 7.065e-6 +- `w`: 1.020e-3 +- `u`: 4.838e-5 + +The W improvement comes from replacing only the internal `k_scaled` workspace +with FP16. Computing `bf16((k * beta) * exp2(gk))` directly in K123 improves W +parity to 4.233e-5 but costs about 30 us at T=8192, so it is retained only as +an experimental accuracy variant. The selected path does not add arithmetic +or bytes to K123. + +### FP64 precision criterion + +`tests/test_kda_sm100_intra_fused_cutedsl.py` builds Aqk, the unit-lower +inverse, W, and U from the same BF16 inputs in FP64. It compares both csrc and +CuTeDSL against that shared oracle and requires every CuTeDSL relative RMSE to +be no greater than the csrc error (apart from a 1e-6 comparison epsilon). + +The strict check uses the same precomputed FP32 `gk` for both implementations +and now requires bitwise equality to csrc for Aqk, Akk, KG, W, and U. It passes +with both FP32 and BF16 beta at `B=1,T=256,H=4` and +`B=2,T=512,H=8`, and the representative BF16-beta benchmark passes at +`B=2,T=8192,H=64`. The standalone recompute-WU test independently requires +bitwise equality for W, U, and KG. + +The source audit found why the earlier fused version missed the precision +target: it recomputed gate activation and the chunk scan with a different +reduction tree before intra. That changed the FP32 cumulative gate before any +TF32 MMA. Fusion and FP16 workspace rounding are therefore excluded from the +new baseline. + +The former mismatch was isolated to the final lower-left 32x32 Schur block. +The CuTeDSL baseline split the K=32 product into two independently initialized +K=16 accumulators and then added their FP32 results. csrc instead issues the +four K=8 TF32 MMA steps into one accumulator. Replacing the split reduction +with the csrc order removes the different FP32 rounding point and makes the +complete csrc-boundary chain bitwise equal. The experimental in-CTA inverse is +still excluded from this claim and the boundary API rejects selecting it. + +The csrc-boundary specialization also writes the complete 64x64 Aqk tile and +explicitly zeros the causal upper triangle. The earlier lower-tile-only store +left those global elements dependent on `torch.empty` allocator contents. +CuTe tensor-wrapper and varlen-padding caches now validate weak-reference +identity as well as tensor version, preventing Python object-ID reuse from +selecting a stale device pointer during long sequential workloads. + +### Profile decomposition + +The aligned recompute-WU kernel uses one 384-thread CTA per SM, 168 registers +per thread at launch, and about 199 KB dynamic shared memory. Nsight Compute +confirmed that the csrc launch is not a CUDA cluster; removing the accidental +single-CTA cluster launch from CuTeDSL removed about 4% at T=8192. Pipeline +barrier initialization is deferred so all seven pipelines share one CTA sync, +matching the single post-construction `__syncthreads()` in csrc. + +## Appendix A: csrc-boundary bitwise and determinism stress + +The bitwise-aligned csrc-boundary path was replayed 10,000,000 times on an +NVIDIA GB200. The case uses BF16 beta at `B=1,T=256,H=4,K=V=128`. Before stress +replay, the harness compares every complete Aqk, Akk, KG, W, and U tensor +against csrc with `torch.equal`. It then captures CuTeDSL intra, the Akk +inverse, recompute WU, and the complete output comparison against csrc in one +CUDA Graph. Every replay therefore checks every output element. + +```bash +python benchmarks/stress_kda_sm100_csrc_boundary_determinism.py \ + --iterations 10000000 --checkpoint 1000000 \ + --report-json /tmp/kda_sm100_csrc_boundary_10m.json +``` + +Bitwise alignment before graph replay: + +| Output | `torch.equal` | Mismatched elements | Max absolute difference | +|---|---:|---:|---:| +| Aqk | true | 0 | 0 | +| Akk | true | 0 | 0 | +| KG | true | 0 | 0 | +| W | true | 0 | 0 | +| U | true | 0 | 0 | + +Determinism result: + +- Iterations: 10,000,000 +- Exact element mismatches accumulated across all replays: 0 +- Elapsed time: 271.881 seconds +- Throughput: 36,780.9 iterations/second +- Status: passed + +The FP32-beta specialization is covered separately by the strict bitwise +pytest cases. The 10,000,000-replay stress above uses the representative BF16 +beta specialization. + +## Appendix B: experimental raw-gate varlen determinism stress + +The complete CuTeDSL varlen forward chain was replayed 10,000,000 times on an +NVIDIA GB200. The case uses four deliberately unaligned sequence lengths +`[65, 127, 193, 255]`, `H=8`, and therefore exercises partial chunks and the +non-pure varlen path. + +```bash +python benchmarks/stress_kda_sm100_varlen_determinism.py \ + --iterations 10000000 --checkpoint 1000000 \ + --report-json /tmp/kda_sm100_varlen_10m.json +``` + +The complete forward plus exact-output comparison is captured in one CUDA +Graph. Every replay compares every element of `k_scaled`, `kg`, `q_scaled`, +`gk_last_exp`, `Aqk`, `Akk`, `w`, and `u` against the first-run golden output; +the device-side mismatch counter is checked every 1,000,000 iterations. + +- Iterations: 10,000,000 +- Exact element mismatches: 0 +- Elapsed time: 471.316 seconds +- Throughput: 21,217.2 iterations/second +- Aqk/Akk upper-triangle max absolute value: 0 + +Accuracy against the csrc path for the same inputs: + +| Output | Relative RMSE | Max absolute error | +|---|---:|---:| +| k_scaled | 3.894e-8 | 2.384e-7 | +| kg | 4.572e-5 | 4.883e-4 | +| q_scaled | 6.323e-7 | 3.815e-6 | +| gk_last_exp | 4.398e-9 | 4.470e-8 | +| Aqk | 4.397e-4 | 1.221e-4 | +| Akk | 1.466e-5 | 4.883e-4 | +| w | 2.163e-3 | 9.766e-4 | +| u | 9.090e-5 | 9.766e-4 | + +## Rejected experiments + +- High-occupancy small-tile WU: slower than the warp-specialized baseline. +- Fully persistent WU scheduling: correct after alias fences, but slower due to + insufficient latency hiding. +- Shared chunk-end exponent with reciprocal: exact division was much slower; + approximate reciprocal plus synchronization also regressed. +- PDL chaining of K123, inverse, and WU: no measurable overlap benefit. +- In-CTA fused inverse with a single G buffer: serialized G staging and + regressed the full chain to about 2.04 ms. +- In-CTA fused inverse while aliasing the Aqk stage: numerically correct, but + serializes inverse work across four persistent chunks and regresses T=8192 + to 2068.8 us. +- Eight chunks per persistent workgroup: T=16384 regresses to 3014.9 us. +- Two chunks per workgroup: violates the current phase/pre-arrive protocol and + fails the dependent inverse launch; four remains required. +- Fused pairwise cumsum plus exact W operand: reduces the fused Aqk FP64 + regression from 0.164% to 0.104% and makes W equal to csrc, but still changes + the csrc gate boundary and runs at 1143.7 us. It is no longer the correctness + baseline. +- Precomputed-FP32-gk diagnostic with the old fused shell: passes the strict + FP64 gate but took 1675.2 us at T=8192. The completed csrc-boundary port + skips fused K1 output work and now measures 864.5 us at the same target + shape. + +## Remaining validation + +- Replace the experimental fused K2 warp-MMA/shared-memory path with the csrc + TMEM/UMMA residency before claiming that fused raw-gate candidate is also a + one-to-one implementation. +- Extend the performance matrix to additional batch/head combinations before + changing the default public dispatch. +- Keep the csrc fallback for shapes that do not satisfy the current K=V=128, + chunk-size=64 specialization constraints. + +The SM100 tests cover float32/bf16 beta, the preprocessed path, the FP64 oracle +criterion, and bitwise equality of the packed-uniform fast path with the equal +kernel. diff --git a/tests/test_kda_sm100_intra_fused_cutedsl.py b/tests/test_kda_sm100_intra_fused_cutedsl.py new file mode 100644 index 00000000..d42d2a3e --- /dev/null +++ b/tests/test_kda_sm100_intra_fused_cutedsl.py @@ -0,0 +1,230 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +import pathlib +import sys + +import pytest +import torch +import torch.nn.functional as F +from fla.ops.kda.gate import kda_gate_chunk_cumsum +from fla.ops.utils.constant import RCP_LN2 + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +from cula.kda.chunk_intra import chunk_kda_fwd_intra as csrc_chunk_kda_fwd_intra +from cula.ops.kda.sm100.intra_fused import ( + chunk_kda_fwd_intra_sm100_equal, + chunk_kda_fwd_intra_sm100_from_gk, + chunk_kda_fwd_intra_sm100_varlen, +) +from cula.ops.kda.sm100.recompute_wu import recompute_w_u_fwd + + +def _requires_sm100(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): + pytest.skip("SM100 CUDA device required") + + +def _chunk_fp64_oracle(q, k, v, gk, beta, scale, chunk_size): + """Small-shape oracle for ranking csrc and CuTeDSL rounding error.""" + batch, seqlen, heads, dim = q.shape + chunks = seqlen // chunk_size + + def chunked(x): + return x.double().view(batch, chunks, chunk_size, heads, -1).permute(0, 1, 3, 2, 4) + + q_c = chunked(q) + k_c = chunked(k) + v_c = chunked(v) + g_c = chunked(gk) + beta_c = beta.double().view(batch, chunks, chunk_size, heads).permute(0, 1, 3, 2) + exp_g = torch.exp2(g_c) + q_g = q_c * exp_g + k_g = k_c * exp_g + k_inv_g = k_c / exp_g + aqk = torch.tril(torch.matmul(q_g, k_inv_g.transpose(-1, -2)) * scale) + eye = torch.eye(chunk_size, dtype=torch.float64, device=q.device) + mat = eye + torch.tril(torch.matmul(k_g, k_inv_g.transpose(-1, -2)), diagonal=-1) * beta_c.unsqueeze(-1) + akk = torch.linalg.inv(mat) + w = torch.matmul(akk, (k_c * beta_c.unsqueeze(-1)) * exp_g) + u = torch.matmul(akk, v_c * beta_c.unsqueeze(-1)) + + def unchunk(x): + return x.permute(0, 1, 3, 2, 4).reshape(batch, seqlen, heads, x.shape[-1]) + + return unchunk(aqk), unchunk(akk), unchunk(w), unchunk(u) + + +def _rel_rmse_fp64(actual, oracle): + diff = actual.double() - oracle + return torch.sqrt(torch.mean(diff.square()) / torch.mean(oracle.square())).item() + + +@pytest.mark.parametrize("beta_dtype", [torch.bfloat16, torch.float32]) +def test_csrc_boundary_port_intra_wu_strict_fp64_precision(beta_dtype): + _requires_sm100() + torch.manual_seed(2) + device = torch.device("cuda") + batch, seqlen, heads, dim, chunk_size = 1, 256, 4, 128, 64 + scale = dim**-0.5 + lower_bound = -5.0 + + q = F.normalize(torch.randn(batch, seqlen, heads, dim, device=device).float(), dim=-1).bfloat16() + k = F.normalize(torch.randn(batch, seqlen, heads, dim, device=device).float(), dim=-1).bfloat16() + g = torch.randn(batch, seqlen, heads, dim, device=device, dtype=torch.bfloat16) + beta = torch.randn(batch, seqlen, heads, device=device).sigmoid().to(beta_dtype) + A_log = torch.randn(heads, device=device) + dt_bias = torch.randn(heads * dim, device=device) + + gk = kda_gate_chunk_cumsum( + g=g, + A_log=A_log, + dt_bias=dt_bias, + scale=RCP_LN2, + chunk_size=chunk_size, + lower_bound=lower_bound, + ) + w_ref, u_ref, _, kg_ref, Aqk_ref, Akk_ref = csrc_chunk_kda_fwd_intra( + q=q, + k=k, + v=k, + gk=gk, + beta=beta, + scale=scale, + chunk_size=chunk_size, + safe_gate=True, + ) + + Aqk, Akk = chunk_kda_fwd_intra_sm100_from_gk( + q=q, + k=k, + gk=gk, + beta=beta, + scale=scale, + fp32_akk_inv=True, + ) + w, u, _, kg = recompute_w_u_fwd(k, k, beta, Akk, gk) + assert torch.equal(kg, kg_ref), f"kg differs bitwise: max_abs={(kg.float() - kg_ref.float()).abs().max().item()}" + + row = torch.arange(seqlen, device=device) % chunk_size + col = torch.arange(chunk_size, device=device) + lower = (col[None, :] <= row[:, None]).view(1, seqlen, 1, chunk_size).expand_as(Akk) + assert torch.equal(Aqk, Aqk_ref), f"Aqk differs bitwise: max_abs={(Aqk.float() - Aqk_ref.float()).abs().max().item()}" + assert torch.equal(Akk, Akk_ref), f"Akk differs bitwise: max_abs={(Akk.float() - Akk_ref.float()).abs().max().item()}" + torch.testing.assert_close(Aqk[~lower], torch.zeros_like(Aqk[~lower]), rtol=0, atol=0) + torch.testing.assert_close(Akk[~lower], torch.zeros_like(Akk[~lower]), rtol=0, atol=0) + assert torch.equal(w, w_ref), f"w differs bitwise: max_abs={(w.float() - w_ref.float()).abs().max().item()}" + assert torch.equal(u, u_ref), f"u differs bitwise: max_abs={(u.float() - u_ref.float()).abs().max().item()}" + + aqk_oracle, akk_oracle, w_oracle, u_oracle = _chunk_fp64_oracle(q, k, k, gk, beta, scale, chunk_size) + precision_rows = [] + for name, candidate, baseline, oracle in ( + ("Aqk", Aqk, Aqk_ref, aqk_oracle), + ("Akk", Akk, Akk_ref, akk_oracle), + ("w", w, w_ref, w_oracle), + ("u", u, u_ref, u_oracle), + ): + candidate_error = _rel_rmse_fp64(candidate, oracle) + baseline_error = _rel_rmse_fp64(baseline, oracle) + precision_rows.append((name, candidate_error, baseline_error)) + regressions = [row for row in precision_rows if row[1] > row[2] * (1.0 + 1e-6)] + assert not regressions, "; ".join( + f"{name}: CuTeDSL={candidate_error:.6e}, csrc={baseline_error:.6e}" + for name, candidate_error, baseline_error in precision_rows + ) + + +@pytest.mark.parametrize("beta_dtype", [torch.bfloat16, torch.float32]) +def test_csrc_boundary_port_intra_wu_larger_shape_bitwise(beta_dtype): + _requires_sm100() + torch.manual_seed(11) + device = torch.device("cuda") + batch, seqlen, heads, dim, chunk_size = 2, 512, 8, 128, 64 + scale = dim**-0.5 + + q = F.normalize(torch.randn(batch, seqlen, heads, dim, device=device).float(), dim=-1).bfloat16() + k = F.normalize(torch.randn(batch, seqlen, heads, dim, device=device).float(), dim=-1).bfloat16() + gk = torch.randn(batch, seqlen, heads, dim, device=device, dtype=torch.float32) * 0.02 + beta = torch.randn(batch, seqlen, heads, device=device).sigmoid().to(beta_dtype) + + w_ref, u_ref, _, kg_ref, Aqk_ref, Akk_ref = csrc_chunk_kda_fwd_intra( + q=q, + k=k, + v=k, + gk=gk, + beta=beta, + scale=scale, + chunk_size=chunk_size, + safe_gate=True, + ) + Aqk, Akk = chunk_kda_fwd_intra_sm100_from_gk( + q=q, + k=k, + gk=gk, + beta=beta, + scale=scale, + fp32_akk_inv=True, + ) + w, u, _, kg = recompute_w_u_fwd(k, k, beta, Akk, gk) + + for name, actual, expected in ( + ("Aqk", Aqk, Aqk_ref), + ("Akk", Akk, Akk_ref), + ("kg", kg, kg_ref), + ("w", w, w_ref), + ("u", u, u_ref), + ): + assert torch.equal(actual, expected), ( + f"{name} differs bitwise: max_abs={(actual.float() - expected.float()).abs().max().item()}" + ) + + # The csrc-boundary kernel must overwrite the complete Aqk/Akk output, + # including the causal upper triangle; correctness cannot depend on the + # initial contents of the cached output buffers. + Aqk.fill_(float("nan")) + Akk.fill_(float("nan")) + Aqk_repeat, Akk_repeat = chunk_kda_fwd_intra_sm100_from_gk( + q=q, + k=k, + gk=gk, + beta=beta, + scale=scale, + fp32_akk_inv=True, + ) + assert torch.equal(Aqk_repeat, Aqk_ref) + assert torch.equal(Akk_repeat, Akk_ref) + + +def test_uniform_varlen_uses_equal_fp16_path_bitwise(): + _requires_sm100() + torch.manual_seed(7) + device = torch.device("cuda") + batch, seqlen, heads, dim = 2, 256, 4, 128 + q = F.normalize(torch.randn(batch, seqlen, heads, dim, device=device).float(), dim=-1).bfloat16() + k = F.normalize(torch.randn(batch, seqlen, heads, dim, device=device).float(), dim=-1).bfloat16() + g = torch.randn(batch, seqlen, heads, dim, device=device, dtype=torch.bfloat16) + beta = torch.randn(batch, seqlen, heads, device=device).sigmoid().bfloat16() + a_log = torch.full((heads,), -4.0, device=device) + dt_bias = torch.zeros(heads * dim, device=device) + kwargs = dict( + A_log=a_log, + dt_bias=dt_bias, + safe_gate=True, + lower_bound=-5.0, + fp32_akk_inv=True, + kscaled_fp16=True, + ) + equal = chunk_kda_fwd_intra_sm100_equal(q=q, k=k, g=g, beta=beta, **kwargs) + packed = chunk_kda_fwd_intra_sm100_varlen( + q=q.flatten(0, 1).unsqueeze(0), + k=k.flatten(0, 1).unsqueeze(0), + g=g.flatten(0, 1).unsqueeze(0), + beta=beta.flatten(0, 1).unsqueeze(0), + cu_seqlens=torch.tensor([0, seqlen, 2 * seqlen], dtype=torch.int32, device=device), + seq_lens=[seqlen, seqlen], + **kwargs, + ) + for equal_tensor, packed_tensor in zip(equal, packed, strict=True): + expected = equal_tensor.flatten(0, 1).unsqueeze(0) + torch.testing.assert_close(packed_tensor, expected, rtol=0, atol=0, equal_nan=True) diff --git a/tests/test_kda_sm100_recompute_wu_cutedsl.py b/tests/test_kda_sm100_recompute_wu_cutedsl.py new file mode 100644 index 00000000..81b60c54 --- /dev/null +++ b/tests/test_kda_sm100_recompute_wu_cutedsl.py @@ -0,0 +1,88 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +import pathlib +import sys + +import pytest +import torch +from fla.ops.utils import prepare_chunk_indices + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent.parent)) + +import cula.cudac as cula_cuda +from cula.ops.kda.sm100.recompute_wu import recompute_w_u_from_preprocessed, recompute_w_u_fwd + + +def _requires_sm100(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0): + pytest.skip("SM100 CUDA device required") + + +@pytest.mark.parametrize("beta_dtype", [torch.float32, torch.bfloat16]) +def test_recompute_wu_matches_csrc(beta_dtype): + _requires_sm100() + torch.manual_seed(0) + device = torch.device("cuda") + batch, seqlen, heads, dim, chunk_size = 2, 256, 16, 128, 64 + + k = torch.randn(1, batch * seqlen, heads, dim, device=device, dtype=torch.bfloat16) * 0.1 + v = torch.randn_like(k) * 0.1 + beta = torch.rand(1, batch * seqlen, heads, device=device, dtype=beta_dtype) + gk = torch.randn(1, batch * seqlen, heads, dim, device=device, dtype=torch.float32) * 0.02 + A = torch.randn(1, batch * seqlen, heads, chunk_size, device=device, dtype=torch.bfloat16) * 0.02 + cu_seqlens = torch.tensor([0, seqlen, 2 * seqlen], dtype=torch.int32, device=device) + chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) + + w_ref = torch.empty_like(k) + u_ref = torch.empty_like(v) + kg_ref = torch.empty_like(k) + cula_cuda.recompute_w_u_cuda( + k, + v, + beta, + A, + gk, + cu_seqlens, + chunk_indices, + w_ref, + u_ref, + kg_ref, + chunk_size, + None, + None, + ) + w, u, _, kg = recompute_w_u_fwd(k, v, beta, A, gk, cu_seqlens, chunk_indices) + + assert torch.equal(w, w_ref), f"w differs bitwise: max_abs={(w.float() - w_ref.float()).abs().max().item()}" + assert torch.equal(u, u_ref), f"u differs bitwise: max_abs={(u.float() - u_ref.float()).abs().max().item()}" + assert torch.equal(kg, kg_ref), f"kg differs bitwise: max_abs={(kg.float() - kg_ref.float()).abs().max().item()}" + + +def test_preprocessed_recompute_wu_matches_torch(): + _requires_sm100() + torch.manual_seed(1) + device = torch.device("cuda") + batch, seqlen, heads, dim, chunk_size = 1, 256, 4, 128, 64 + + k_scaled = torch.randn(batch, seqlen, heads, dim, device=device, dtype=torch.bfloat16) * 0.1 + v = torch.randn_like(k_scaled) * 0.1 + beta = torch.rand(batch, seqlen, heads, device=device, dtype=torch.bfloat16) + A = torch.randn(batch, seqlen, heads, chunk_size, device=device, dtype=torch.bfloat16) * 0.02 + row = torch.arange(seqlen, device=device) % chunk_size + col = torch.arange(chunk_size, device=device) + A.masked_fill_((col[None, :] > row[:, None]).view(1, seqlen, 1, chunk_size), 0) + + w, u = recompute_w_u_from_preprocessed(k_scaled, v, beta, A) + w_ref = torch.empty_like(w) + u_ref = torch.empty_like(u) + k_beta = (k_scaled.float() * beta.float().unsqueeze(-1)).bfloat16() + v_beta = (v.float() * beta.float().unsqueeze(-1)).bfloat16() + for start in range(0, seqlen, chunk_size): + end = start + chunk_size + A_tile = A[:, start:end].float() + w_ref[:, start:end] = torch.einsum("bmhk,bkhd->bmhd", A_tile, k_beta[:, start:end].float()) + u_ref[:, start:end] = torch.einsum("bmhk,bkhd->bmhd", A_tile, v_beta[:, start:end].float()) + + torch.testing.assert_close(w, w_ref, rtol=1e-2, atol=2e-3) + torch.testing.assert_close(u, u_ref, rtol=1e-2, atol=2e-3)