From a24b8fc02a4907c36d367d500871554fda8c096c Mon Sep 17 00:00:00 2001 From: Will Lin Date: Mon, 17 Aug 2026 17:39:20 +0000 Subject: [PATCH] [perf] Add fine-grained FA4 VSA kernels --- fastvideo-kernel/README.md | 72 +- .../benchmarks/bench_vsa_blackwell.py | 984 ++++ fastvideo-kernel/fa4/.flake8 | 4 + fastvideo-kernel/fa4/AUTHORS | 8 + fastvideo-kernel/fa4/LICENSE | 29 + fastvideo-kernel/fa4/MANIFEST.in | 6 + fastvideo-kernel/fa4/README.md | 33 + fastvideo-kernel/fa4/UPSTREAM.md | 37 + fastvideo-kernel/fa4/__init__.py | 18 + fastvideo-kernel/fa4/ampere_helpers.py | 103 + fastvideo-kernel/fa4/barrier.py | 71 + fastvideo-kernel/fa4/bench_utils.py | 243 + fastvideo-kernel/fa4/benchmark.py | 268 + .../fa4/benchmark_flash_attention_fp8.py | 434 ++ fastvideo-kernel/fa4/blackwell_helpers.py | 1115 ++++ fastvideo-kernel/fa4/block_info.py | 156 + fastvideo-kernel/fa4/block_sparse_utils.py | 2035 +++++++ fastvideo-kernel/fa4/block_sparsity.py | 676 +++ fastvideo-kernel/fa4/cache_utils.py | 281 + .../fa4/compute_block_sparsity.py | 551 ++ fastvideo-kernel/fa4/copy_utils.py | 372 ++ fastvideo-kernel/fa4/cute_dsl_ptxas.py | 151 + fastvideo-kernel/fa4/cute_dsl_utils.py | 158 + fastvideo-kernel/fa4/fa_logging.py | 97 + fastvideo-kernel/fa4/fast_math.py | 21 + fastvideo-kernel/fa4/flash_bwd.py | 1325 +++++ .../fa4/flash_bwd_mla_dk_sm100.py | 1183 ++++ .../fa4/flash_bwd_mla_dq_dqv_sm100.py | 1297 +++++ fastvideo-kernel/fa4/flash_bwd_mla_sm100.py | 2133 ++++++++ fastvideo-kernel/fa4/flash_bwd_postprocess.py | 587 ++ fastvideo-kernel/fa4/flash_bwd_preprocess.py | 465 ++ fastvideo-kernel/fa4/flash_bwd_sm100.py | 4041 ++++++++++++++ fastvideo-kernel/fa4/flash_bwd_sm120.py | 55 + fastvideo-kernel/fa4/flash_bwd_sm90.py | 1926 +++++++ fastvideo-kernel/fa4/flash_fwd.py | 1236 +++++ fastvideo-kernel/fa4/flash_fwd_combine.py | 698 +++ fastvideo-kernel/fa4/flash_fwd_mla_sm100.py | 3160 +++++++++++ fastvideo-kernel/fa4/flash_fwd_sm100.py | 4840 +++++++++++++++++ fastvideo-kernel/fa4/flash_fwd_sm120.py | 59 + fastvideo-kernel/fa4/flash_fwd_sm90.py | 1545 ++++++ fastvideo-kernel/fa4/interface.py | 3204 +++++++++++ fastvideo-kernel/fa4/mask.py | 1709 ++++++ fastvideo-kernel/fa4/mma_sm100_desc.py | 296 + fastvideo-kernel/fa4/named_barrier.py | 68 + fastvideo-kernel/fa4/pack_gqa.py | 263 + fastvideo-kernel/fa4/paged_kv.py | 247 + fastvideo-kernel/fa4/pipeline.py | 402 ++ fastvideo-kernel/fa4/pyproject.toml | 68 + fastvideo-kernel/fa4/seqlen_info.py | 302 + .../fa4/sm100_hd256_2cta_fmha_backward.py | 298 + ...100_hd256_2cta_fmha_backward_dkdvkernel.py | 3411 ++++++++++++ ...sm100_hd256_2cta_fmha_backward_dqkernel.py | 2317 ++++++++ .../fa4/sm100_hd256_2cta_fmha_forward.py | 1889 +++++++ fastvideo-kernel/fa4/sm90_config_search.py | 402 ++ fastvideo-kernel/fa4/softmax.py | 740 +++ fastvideo-kernel/fa4/testing.py | 487 ++ fastvideo-kernel/fa4/tile_scheduler.py | 1650 ++++++ fastvideo-kernel/fa4/topk_gather_kv.py | 278 + fastvideo-kernel/fa4/utils.py | 957 ++++ .../fastvideo_kernel/block_sparse_attn_256.py | 90 +- .../block_sparse_attn_cute_fwd.py | 46 +- .../tests/test_fa4_vsa_block_shapes.py | 526 ++ handoff.md | 170 + pyproject.toml | 11 +- 64 files changed, 52248 insertions(+), 56 deletions(-) create mode 100644 fastvideo-kernel/benchmarks/bench_vsa_blackwell.py create mode 100644 fastvideo-kernel/fa4/.flake8 create mode 100644 fastvideo-kernel/fa4/AUTHORS create mode 100644 fastvideo-kernel/fa4/LICENSE create mode 100644 fastvideo-kernel/fa4/MANIFEST.in create mode 100644 fastvideo-kernel/fa4/README.md create mode 100644 fastvideo-kernel/fa4/UPSTREAM.md create mode 100644 fastvideo-kernel/fa4/__init__.py create mode 100644 fastvideo-kernel/fa4/ampere_helpers.py create mode 100644 fastvideo-kernel/fa4/barrier.py create mode 100644 fastvideo-kernel/fa4/bench_utils.py create mode 100644 fastvideo-kernel/fa4/benchmark.py create mode 100644 fastvideo-kernel/fa4/benchmark_flash_attention_fp8.py create mode 100644 fastvideo-kernel/fa4/blackwell_helpers.py create mode 100644 fastvideo-kernel/fa4/block_info.py create mode 100644 fastvideo-kernel/fa4/block_sparse_utils.py create mode 100644 fastvideo-kernel/fa4/block_sparsity.py create mode 100644 fastvideo-kernel/fa4/cache_utils.py create mode 100644 fastvideo-kernel/fa4/compute_block_sparsity.py create mode 100644 fastvideo-kernel/fa4/copy_utils.py create mode 100644 fastvideo-kernel/fa4/cute_dsl_ptxas.py create mode 100644 fastvideo-kernel/fa4/cute_dsl_utils.py create mode 100644 fastvideo-kernel/fa4/fa_logging.py create mode 100644 fastvideo-kernel/fa4/fast_math.py create mode 100644 fastvideo-kernel/fa4/flash_bwd.py create mode 100644 fastvideo-kernel/fa4/flash_bwd_mla_dk_sm100.py create mode 100644 fastvideo-kernel/fa4/flash_bwd_mla_dq_dqv_sm100.py create mode 100644 fastvideo-kernel/fa4/flash_bwd_mla_sm100.py create mode 100644 fastvideo-kernel/fa4/flash_bwd_postprocess.py create mode 100644 fastvideo-kernel/fa4/flash_bwd_preprocess.py create mode 100644 fastvideo-kernel/fa4/flash_bwd_sm100.py create mode 100644 fastvideo-kernel/fa4/flash_bwd_sm120.py create mode 100644 fastvideo-kernel/fa4/flash_bwd_sm90.py create mode 100644 fastvideo-kernel/fa4/flash_fwd.py create mode 100644 fastvideo-kernel/fa4/flash_fwd_combine.py create mode 100644 fastvideo-kernel/fa4/flash_fwd_mla_sm100.py create mode 100644 fastvideo-kernel/fa4/flash_fwd_sm100.py create mode 100644 fastvideo-kernel/fa4/flash_fwd_sm120.py create mode 100644 fastvideo-kernel/fa4/flash_fwd_sm90.py create mode 100644 fastvideo-kernel/fa4/interface.py create mode 100644 fastvideo-kernel/fa4/mask.py create mode 100644 fastvideo-kernel/fa4/mma_sm100_desc.py create mode 100644 fastvideo-kernel/fa4/named_barrier.py create mode 100644 fastvideo-kernel/fa4/pack_gqa.py create mode 100644 fastvideo-kernel/fa4/paged_kv.py create mode 100644 fastvideo-kernel/fa4/pipeline.py create mode 100644 fastvideo-kernel/fa4/pyproject.toml create mode 100644 fastvideo-kernel/fa4/seqlen_info.py create mode 100644 fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_backward.py create mode 100644 fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_backward_dkdvkernel.py create mode 100644 fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_backward_dqkernel.py create mode 100644 fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_forward.py create mode 100644 fastvideo-kernel/fa4/sm90_config_search.py create mode 100644 fastvideo-kernel/fa4/softmax.py create mode 100644 fastvideo-kernel/fa4/testing.py create mode 100644 fastvideo-kernel/fa4/tile_scheduler.py create mode 100644 fastvideo-kernel/fa4/topk_gather_kv.py create mode 100644 fastvideo-kernel/fa4/utils.py create mode 100644 fastvideo-kernel/tests/test_fa4_vsa_block_shapes.py create mode 100644 handoff.md diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index 16df2e7ef..65c28650e 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -64,28 +64,42 @@ cd fastvideo-kernel ./build.sh --rocm ``` -### Optional: FA4 CuTe block-sparse backend (VSA-256 fastpath) +### Optional: FA4 CuTe block-sparse backend -The VSA-256 fastpath (tile volume 256, on NVIDIA Blackwell / sm_100) routes to the -FlashAttention-4 CuTe-DSL block-sparse kernel exposed as `flash_attn.cute`. This is -an **optional** dependency: it is imported lazily, and `video_sparse_attn` -transparently falls back to the Triton backend when it is absent (so the package is -fully usable without it). +The NVIDIA Blackwell / sm_100 VSA fastpath routes to FastVideo's in-tree +FlashAttention-4 CuTe-DSL source fork, exposed as `flash_attn.cute`. This is an +**optional** dependency: it is imported lazily, and `video_sparse_attn` +transparently falls back to the Triton backend when it is absent (so the package +is fully usable without it). -The symbols the fastpath needs (`flash_attn.cute.block_sparsity.BlockSparseTensorsTorch`, -`flash_attn.cute.interface._flash_attn_fwd`) are provided upstream by -[Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention). Pin to -commit `940cd9680f3315f2f06b43ab5bea2c2cf2d96806`, the revision FastVideo pins as -the `flash-attn-4` source in the repo-root `pyproject.toml`; other revisions may -have an incompatible `_flash_attn_fwd` signature. +The fork lives in `fastvideo-kernel/fa4`, retains the `flash-attn-4` distribution +name, and records its exact upstream revision and refresh procedure in +`fa4/UPSTREAM.md`. Install it editable while developing the kernels: ```bash -pip install "nvidia-cutlass-dsl>=4.5.0" torchvision -pip install "git+https://github.com/Dao-AILab/flash-attention.git@940cd9680f3315f2f06b43ab5bea2c2cf2d96806#subdirectory=flash_attn/cute" +uv pip install --no-deps --editable fastvideo-kernel/fa4 ``` +The production VSA-256 entrypoint keeps its existing logical Q256/KV256 mask. +When the CuTe backend is enabled, an optional shape override expands that mask +to the finer physical schedules without changing selected token pairs: + +```bash +FASTVIDEO_VSA_CUTEDSL=1 FASTVIDEO_VSA_FA4_BLOCK_SHAPE=128x64 ... +# At the original Q256 API boundary, logical Q64 children are safely paired +# onto the optimized physical Q128/KV64 schedule: +FASTVIDEO_VSA_CUTEDSL=1 FASTVIDEO_VSA_FA4_BLOCK_SHAPE=64x64 ... +``` + +The default `FASTVIDEO_VSA_FA4_BLOCK_SHAPE=256x256` preserves the historical +route (its logical KV256 edges are implemented with physical KV128 tiles). +Direct callers with arbitrary, distinct Q64 masks still use the native +Q64/KV64 kernel. The Q128/KV64 path enables score/probability double buffering +by default; set `FASTVIDEO_FA4_VSA_SP_DOUBLE_BUFFER=0` only to run its retained +legacy schedule for comparison or rollback. + The CuTe kernel JIT-compiles on first use. Verified on Blackwell (sm_100) against -`tests/test_vsa256_forward*.py`. +`tests/test_fa4_vsa_block_shapes.py` and `tests/test_vsa256_forward*.py`. ## Usage @@ -137,7 +151,33 @@ statistics and bitwise-compatible `dV`. ### VSA (block-sparse) TFLOPs -After building/installing `fastvideo-kernel`, run: +For the GB200 FA4 comparison, install the in-tree source fork and run the +issue-#4554-derived harness. Its default `exact256` mask mode preserves the +same selected token pairs across Q×KV block shapes, reports sparse-aware +TFLOP/s, MFU against the 2.5-PFLOP/s dense-BF16 GB200 peak, and efficiency +relative to the raw FA4 256×256 path: + +```bash +uv pip install --no-deps --editable fastvideo-kernel/fa4 +CUDA_VISIBLE_DEVICES=0 CUTE_DSL_ENABLE_TVM_FFI=1 \ + uv run --no-sync python fastvideo-kernel/benchmarks/bench_vsa_blackwell.py \ + --seq_lens 32768 --sparsities dense 90 \ + --block_shapes 256x256 128x64 64x64 +``` + +To exercise every GPU in a four-GB200 tray, launch one independent replica per +device. Each rank writes its own suffixed JSON (for example, +`/tmp/fa4_tray.rank0.json`): + +```bash +CUTE_DSL_ENABLE_TVM_FFI=1 uv run --no-sync torchrun --standalone --nproc-per-node=4 \ + fastvideo-kernel/benchmarks/bench_vsa_blackwell.py \ + --seq_lens 32768 --sparsities dense 90 \ + --block_shapes 256x256 128x64 64x64 --out /tmp/fa4_tray.json +``` + +The older generic benchmark below measures FastVideo's 64-token block-sparse +wrapper; it does not exercise the FA4 fine-grained source fork: ```bash cd fastvideo-kernel diff --git a/fastvideo-kernel/benchmarks/bench_vsa_blackwell.py b/fastvideo-kernel/benchmarks/bench_vsa_blackwell.py new file mode 100644 index 000000000..d9f9d9654 --- /dev/null +++ b/fastvideo-kernel/benchmarks/bench_vsa_blackwell.py @@ -0,0 +1,984 @@ +#!/usr/bin/env python3 +"""Block-sparse attention micro-benchmark: flashinfer ``vsa_blackwell`` and +FA4 CuTe-DSL BSA, on identical block masks where the backends can represent +them (GB200 / SM100). + +This harness originated in FlashInfer issue #4554 and its benchmark gist: + https://github.com/flashinfer-ai/flashinfer/issues/4554 + https://gist.github.com/SolitaryThinker/90a1d1447929fc38dc509c1852e76532 +The imported baseline is pinned to immutable gist revision +``e15ac9066f23ef3690e33e1cc1fdac45b4b9099f``. + +Measures achieved sparse-aware algorithmic TFLOP/s over a sequence-length x +sparsity x (Q block, KV block) grid. Every successful arm is checked against +the same fp32 masked-SDPA reference for that cell. FlashInfer's plan()/run() +split is timed separately (VSA masks are data-dependent and change every +layer x step, so mask-per-call deployments pay both). + +The default ``--mask_mode exact256`` samples one Q256/KV256 mask and expands +it into each finer block shape. Thus all shapes select exactly the same token +pairs and their output and relative-efficiency comparisons are apples-to- +apples. ``--mask_mode native`` samples independently at each requested shape. + +Arms: + flashinfer — BlockSparseAttentionWrapper(backend="vsa_blackwell"), + R=C=128. It runs only when the requested logical mask is + exactly representable by 128-token tiles. + cutedsl256 — original FastVideo 256x256 wrapper arm from the pinned gist; + used as a measured relative-efficiency reference. + fa4_wrapper — FastVideo's direct fine-grained BSHD wrapper, including + block-map conversion and variable-block-size mask plumbing. + fa4 — direct ``flash_attn.cute.interface._flash_attn_fwd`` call. + Sparse Q and KV block sizes are configured independently. + +For the stock 256x256 FA4 path, a logical KV256 edge expands into two physical +KV128 edges, matching ``block_sparse_attn_256``. The 128x64 and 64x64 paths +call FA4 with physical tiles (128, 64) and (64, 64), respectively. Mask-index +construction and output allocation are outside the raw FA4 timing region. +The direct ``fa4_wrapper`` 64x64 row is therefore the native Q64/KV64 kernel; +it is not the public VSA-256 ``FASTVIDEO_VSA_FA4_BLOCK_SHAPE=64x64`` adapter, +which safely pairs identical Q64 children onto physical Q128/KV64 and is +measured by selecting that environment variable on the ``cutedsl256`` arm. + +The vendored FA4 source is a flat ``flash_attn.cute`` package. It can be used +without installation by putting that directory on ``PYTHONPATH``; this script +detects the flat layout and mounts it under the expected import namespace: + + PYTHONPATH="$PWD/fastvideo-kernel/fa4${PYTHONPATH:+:$PYTHONPATH}" \\ + uv run --no-sync python \\ + fastvideo-kernel/benchmarks/bench_vsa_blackwell.py --quick + +MFU is reported separately from TFLOP/s. Its default denominator is the +official 2,500 TFLOP/s dense BF16 tensor-core peak for one GB200 GPU and is +explicitly overrideable with ``--peak_bf16_tflops``. Relative efficiency +against measured 256x256 baselines does not depend on that nominal peak. + +Examples: + python bench_vsa_blackwell.py + python bench_vsa_blackwell.py --quick + python bench_vsa_blackwell.py --seq_lens 32768 --sparsities dense 90 --block_shapes 256x256 128x64 64x64 + python bench_vsa_blackwell.py --mask_mode native --block_shapes 128x64 64x64 + torchrun --standalone --nproc-per-node=4 bench_vsa_blackwell.py --quick --out /tmp/fa4_tray.json +""" + +from __future__ import annotations + +import argparse +import importlib +import importlib.util +import inspect +import json +import math +import os +from pathlib import Path +import random +import subprocess +import sys +import time +import traceback +import types + +os.environ.setdefault("FASTVIDEO_VSA_CUTEDSL", "1") + +import numpy as np +import torch + +FI_BLOCK = 128 # flashinfer vsa_blackwell R=C +DEFAULT_BLOCK_SHAPES = ((256, 256), (128, 64), (64, 64)) +DEFAULT_GB200_BF16_DENSE_TFLOPS = 2500.0 +GB200_DATASHEET_URL = ( + "https://dam-cdn.nvd.orangelogic.com/AssetLink/" + "y441155802qub41q118b2852i557jem5.pdf" +) +ISSUE_URL = "https://github.com/flashinfer-ai/flashinfer/issues/4554" +GIST_URL = "https://gist.github.com/SolitaryThinker/90a1d1447929fc38dc509c1852e76532" +GIST_REVISION = "e15ac9066f23ef3690e33e1cc1fdac45b4b9099f" + +KEEP_FRAC = { + "dense": 1.0, + "50": 0.5, + "60": 0.4, + "70": 0.3, + "75": 0.25, + "80": 0.2, + "87.5": 0.125, + "90": 0.1, +} + + +class UnsupportedConfig(ValueError): + """The backend cannot exactly represent the requested logical mask.""" + + +def parse_block_shape(value: str) -> tuple[int, int]: + parts = value.lower().replace(",", "x").split("x") + if len(parts) != 2: + raise argparse.ArgumentTypeError(f"expected QxKV, got {value!r}") + try: + q_block, kv_block = (int(part) for part in parts) + except ValueError as exc: + raise argparse.ArgumentTypeError(f"expected integer QxKV, got {value!r}") from exc + if q_block <= 0 or kv_block <= 0: + raise argparse.ArgumentTypeError("Q and KV block sizes must be positive") + return q_block, kv_block + + +def flops_sparse_attention( + bs: int, + d: int, + selected_edges: int, + q_block: int, + kv_block: int, +) -> float: + """QK^T + PV FLOPs over selected token pairs only. + + ``selected_edges`` already includes the head dimension. Each selected + logical edge covers ``q_block * kv_block`` token pairs, and each of QK^T + and PV costs two FLOPs per head-dimension element. + """ + return 4.0 * bs * d * selected_edges * q_block * kv_block + + +def set_seed(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def make_logical_mask(h: int, nq: int, nkv: int, keep: int, seed: int) -> torch.Tensor: + """[H, NQ, NKV] bool with one diagonal anchor and random other edges. + + CPU generator so the pattern is device-independent and fully seeded. + """ + if not 1 <= keep <= nkv: + raise ValueError(f"keep must be in [1, {nkv}], got {keep}") + g = torch.Generator(device="cpu").manual_seed(seed) + scores = torch.rand(h, nq, nkv, generator=g) + q_rows = torch.arange(nq) + # For equal sequence lengths this is the KV block containing the first + # token of the Q block. It keeps every softmax row nonempty even when + # Q and KV block sizes differ. + diagonal_anchor = torch.div(q_rows * nkv, nq, rounding_mode="floor").clamp_max(nkv - 1) + scores[:, q_rows, diagonal_anchor] = 2.0 # rand() < 1, so this always wins topk + idx = torch.topk(scores, keep, dim=-1).indices + m = torch.zeros(h, nq, nkv, dtype=torch.bool) + m.scatter_(-1, idx, True) + assert int(m.sum(-1).min()) == keep and int(m.sum(-1).max()) == keep + return m.cuda() + + +def expand_exact256_mask(base_mask: torch.Tensor, q_block: int, kv_block: int) -> torch.Tensor: + """Expand a Q256/KV256 mask without changing its selected token pairs.""" + if 256 % q_block or 256 % kv_block: + raise ValueError(f"QxKV block shape {q_block}x{kv_block} must divide 256") + return base_mask.repeat_interleave(256 // q_block, dim=1).repeat_interleave( + 256 // kv_block, dim=2 + ) + + +def ref_masked_sdpa_fp32(q, k, v, logical_mask, q_block, kv_block, scale, qchunk=2048): + """fp32 SDPA with the expanded boolean token mask, chunked over q. + + q,k,v: [1,S,H,D] bf16. Returns [S,H,D] fp32. + TF32 is disabled in main(). + """ + _, seq_len, heads, head_dim = q.shape + if seq_len % q_block or k.shape[1] % kv_block: + raise ValueError("reference requires sequence lengths divisible by their logical block sizes") + qchunk = max(q_block, qchunk // q_block * q_block) + q32 = q[0].permute(1, 0, 2).float() # [H,S,D] + k32 = k[0].permute(1, 0, 2).float() + v32 = v[0].permute(1, 0, 2).float() + out = torch.empty(heads, seq_len, head_dim, dtype=torch.float32, device=q.device) + for i in range(0, seq_len, qchunk): + j = min(i + qchunk, seq_len) + scores = torch.matmul(q32[:, i:j], k32.transpose(-1, -2)).mul_(scale) + mask_rows = logical_mask[:, i // q_block:j // q_block] + mask_tokens = mask_rows.repeat_interleave(q_block, dim=1) + mask_tokens = mask_tokens.repeat_interleave(kv_block, dim=2) # [H,c,S] + scores.masked_fill_(~mask_tokens, float("-inf")) + out[:, i:j] = torch.matmul(torch.softmax(scores, dim=-1), v32) + del scores, mask_tokens + return out.permute(1, 0, 2).contiguous() # [S,H,D] + + +def arm_cutedsl(q, k, v, logical_mask, q_block, kv_block): + """FA4-lineage CuTe-DSL 256 path. Returns (bench_fn, out [S,H,D], extra). + + Timed as the full wrapper call (mask expansion + map->index included) — + that is its per-call hot path in a mask-per-call deployment. + """ + if (q_block, kv_block) != (256, 256): + raise UnsupportedConfig("cutedsl256 is defined only for logical Q256/KV256") + _mount_flat_fa4_from_pythonpath() + from fastvideo_kernel import block_sparse_attn_256 as wrapper_module + + block_sparse_attn_256_bshd = wrapper_module.block_sparse_attn_256_bshd + + nkv = logical_mask.shape[-1] + vbs = torch.full((nkv, ), kv_block, dtype=torch.int32, device=q.device) + mask = logical_mask.unsqueeze(0) # [1,H,NQ,NKV] + + def fn(): + return block_sparse_attn_256_bshd(q, k, v, mask, vbs) + + out, _lse = fn() + interface = importlib.import_module("flash_attn.cute.interface") + return fn, out[0], { + "fa4_interface": str(Path(interface.__file__).resolve()), + "fastvideo_wrapper": str(Path(wrapper_module.__file__).resolve()), + "fa4_import_mode": _FA4_IMPORT_MODE, + "route_semantics": "public_vsa256", + "requested_vsa256_block_shape": os.environ.get( + "FASTVIDEO_VSA_FA4_BLOCK_SHAPE", "256x256" + ), + } + + +def arm_fa4_wrapper(q, k, v, logical_mask, q_block, kv_block): + """FastVideo's direct fine-grained BSHD wrapper. + + This arm deliberately includes boolean-map conversion and the VBS + ``mask_mod``/aux-tensor plumbing in the timed region. Fixed-length cells + use fully valid KV blocks. In particular, its 64x64 row stays on native + Q64/KV64 rather than entering the public VSA-256 coalescing adapter. + """ + if (q_block, kv_block) not in ((128, 64), (64, 64)): + raise UnsupportedConfig("fa4_wrapper supports Q128/KV64 and Q64/KV64") + + _mount_flat_fa4_from_pythonpath() + from fastvideo_kernel import block_sparse_attn_cute_fwd as wrapper_module + + block_sparse_attn_cute_fwd_bshd = wrapper_module.block_sparse_attn_cute_fwd_bshd + + nkv = logical_mask.shape[-1] + vbs = torch.full((nkv,), kv_block, dtype=torch.int32, device=q.device) + mask = logical_mask.unsqueeze(0) + + def fn(): + return block_sparse_attn_cute_fwd_bshd(q, k, v, mask, vbs) + + out, _lse = fn() + interface = importlib.import_module("flash_attn.cute.interface") + return fn, out[0], { + "fa4_interface": str(Path(interface.__file__).resolve()), + "fastvideo_wrapper": str(Path(wrapper_module.__file__).resolve()), + "fa4_import_mode": _FA4_IMPORT_MODE, + "route_semantics": "direct_native_fine_grained", + "physical_tile": [q_block, kv_block], + } + + +def arm_flashinfer(q, k, v, logical_mask, q_block, kv_block): + """flashinfer vsa_blackwell (blk128). Returns (bench_fn, out, extra). + + run() timed alone (flashinfer's documented plan-once/run-many model); + steady-state plan() wall time recorded separately per cell. + """ + if q_block % FI_BLOCK or kv_block % FI_BLOCK: + raise UnsupportedConfig( + "flashinfer vsa_blackwell R=C=128 cannot exactly represent this logical mask" + ) + from flashinfer.sparse import BlockSparseAttentionWrapper + + if "block_mask" not in inspect.signature(BlockSparseAttentionWrapper.plan).parameters: + raise UnsupportedConfig( + "installed FlashInfer lacks the per-head block_mask planning API used by " + "issue #4554 (the issue used flashinfer 0.6.16.post2)" + ) + + _, seq_len, heads, head_dim = q.shape + mask128 = logical_mask.repeat_interleave(q_block // FI_BLOCK, dim=1).repeat_interleave( + kv_block // FI_BLOCK, dim=2 + ) + workspace = torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device=q.device) + wrapper = BlockSparseAttentionWrapper(workspace, backend="vsa_blackwell") + + def do_plan(): + wrapper.plan( + None, + None, + M=seq_len, + N=seq_len, + R=FI_BLOCK, + C=FI_BLOCK, + num_qo_heads=heads, + num_kv_heads=heads, + head_dim=head_dim, + block_mask=mask128, + q_data_type=torch.bfloat16, + o_data_type=torch.bfloat16, + ) + + torch.cuda.synchronize() + t0 = time.perf_counter() + do_plan() + torch.cuda.synchronize() + plan_first = (time.perf_counter() - t0) * 1e3 # includes JIT on first cell + plan_steady = [] + for _ in range(3): + torch.cuda.synchronize() + t0 = time.perf_counter() + do_plan() + torch.cuda.synchronize() + plan_steady.append((time.perf_counter() - t0) * 1e3) + qn, kn, vn = q[0], k[0], v[0] # NHD views of the same storage + + def fn(): + return wrapper.run(qn, kn, vn) + + out = fn() + return fn, out, { + "plan_ms_first": round(plan_first, 2), + "plan_ms": round(min(plan_steady), 2), + "keep_alive": wrapper, + } + + +_FA4_IMPORT_MODE = "normal" + + +def _flat_fa4_pythonpath() -> Path | None: + """Find FastVideo's flat ``flash_attn.cute`` source on PYTHONPATH.""" + for entry in os.environ.get("PYTHONPATH", "").split(os.pathsep): + if not entry: + continue + candidate = Path(entry).resolve() + if all((candidate / name).is_file() for name in ("UPSTREAM.md", "interface.py", "block_sparsity.py")): + return candidate + return None + + +def _mount_flat_fa4_from_pythonpath() -> None: + """Mount the flat vendored tree as ``flash_attn.cute`` without installing. + + A regular upstream checkout already has ``flash_attn/cute`` and needs no + special handling. FastVideo intentionally vendors only that subpackage, + so its source root itself is the package directory. + """ + global _FA4_IMPORT_MODE + source = _flat_fa4_pythonpath() + if source is None: + return + + try: + current = importlib.util.find_spec("flash_attn.cute.interface") + except (ImportError, AttributeError, ValueError): + current = None + if current is not None and current.origin is not None and Path(current.origin).resolve().parent == source: + _FA4_IMPORT_MODE = "local-pythonpath" + return + + # Load the package under its real namespace so the vendored source's + # absolute ``flash_attn.cute.*`` imports resolve back to the same tree. + try: + import flash_attn + except ModuleNotFoundError as exc: + if exc.name != "flash_attn": + raise + # The fork is a standalone ``flash_attn.cute`` distribution. Create + # its namespace parent when FA2's top-level package is not installed. + flash_attn = types.ModuleType("flash_attn") + flash_attn.__path__ = [] + sys.modules["flash_attn"] = flash_attn + + for module_name in tuple(sys.modules): + if module_name == "flash_attn.cute" or module_name.startswith("flash_attn.cute."): + del sys.modules[module_name] + spec = importlib.util.spec_from_file_location( + "flash_attn.cute", + source / "__init__.py", + submodule_search_locations=[str(source)], + ) + if spec is None or spec.loader is None: + raise ImportError(f"could not create an import spec for local FA4 source {source}") + package = importlib.util.module_from_spec(spec) + sys.modules["flash_attn.cute"] = package + setattr(flash_attn, "cute", package) + spec.loader.exec_module(package) + _FA4_IMPORT_MODE = "flat-pythonpath" + + +def _load_fa4_symbols(): + _mount_flat_fa4_from_pythonpath() + block_sparsity = importlib.import_module("flash_attn.cute.block_sparsity") + interface = importlib.import_module("flash_attn.cute.interface") + return block_sparsity.BlockSparseTensorsTorch, interface._flash_attn_fwd, interface + + +def _fa4_physical_tiles(q_block: int, kv_block: int) -> tuple[int, int]: + """Map logical sparse blocks to FA4's physical MMA tiles. + + KV256 is the historical logical VSA block and expands to KV128. Smaller + requested KV blocks remain physical. Q256 is two Q128 stages; Q128 and + Q64 select the single-stage paths being optimized. + """ + tile_m = min(q_block, 128) + tile_n = min(kv_block, 128) + if q_block % tile_m or kv_block % tile_n: + raise UnsupportedConfig( + f"logical QxKV={q_block}x{kv_block} is not divisible by physical " + f"FA4 tile {tile_m}x{tile_n}" + ) + return tile_m, tile_n + + +def arm_fa4(q, k, v, logical_mask, q_block, kv_block): + """Direct FA4 block-sparse forward; mask preparation is not timed.""" + BlockSparseTensorsTorch, _flash_attn_fwd, interface = _load_fa4_symbols() + tile_m, tile_n = _fa4_physical_tiles(q_block, kv_block) + kv_expansion = kv_block // tile_n + physical_mask = logical_mask.repeat_interleave(kv_expansion, dim=-1).contiguous() + heads, nq, nkv_physical = physical_mask.shape + selected_per_row = int(physical_mask[0, 0].sum().item()) + if selected_per_row <= 0 or not bool((physical_mask.sum(dim=-1) == selected_per_row).all()): + raise ValueError("direct FA4 arm requires the same positive selected-block count per row") + + physical_indices = torch.arange(nkv_physical, dtype=torch.int32, device=q.device) + physical_indices = physical_indices.view(1, 1, -1).expand_as(physical_mask) + full_idx = physical_indices.masked_select(physical_mask).view(1, heads, nq, selected_per_row) + full_cnt = torch.full((1, heads, nq), selected_per_row, dtype=torch.int32, device=q.device) + # Every selected logical block is fully valid in this fixed-length benchmark. + # Keep a compact empty partial-block representation to exercise FA4's full path. + mask_cnt = torch.zeros((1, heads, nq), dtype=torch.int32, device=q.device) + mask_idx = torch.zeros((1, heads, nq, 1), dtype=torch.int32, device=q.device) + sparse_tensors = BlockSparseTensorsTorch( + mask_block_cnt=mask_cnt, + mask_block_idx=mask_idx, + full_block_cnt=full_cnt, + full_block_idx=full_idx.contiguous(), + block_size=(q_block, tile_n), + ) + + batch, seq_len, _, _ = q.shape + out_buffer = torch.empty_like(q) + lse_buffer = torch.empty((batch, heads, seq_len), dtype=torch.float32, device=q.device) + scale = 1.0 / math.sqrt(q.shape[-1]) + + def fn(): + return _flash_attn_fwd( + q, + k, + v, + out=out_buffer, + lse=lse_buffer, + softmax_scale=scale, + tile_mn=(tile_m, tile_n), + num_splits=1, + pack_gqa=False, + block_sparse_tensors=sparse_tensors, + causal=False, + return_lse=True, + )[:2] + + out, _lse = fn() + return fn, out[0], { + "fa4_interface": str(Path(interface.__file__).resolve()), + "fa4_import_mode": _FA4_IMPORT_MODE, + "fa4_tile_m": tile_m, + "fa4_tile_n": tile_n, + "fa4_sparse_q_block": q_block, + "fa4_sparse_kv_block": tile_n, + "fa4_kv_expansion": kv_expansion, + "keep_alive": (sparse_tensors, out_buffer, lse_buffer), + } + + +ARMS = { + "flashinfer": arm_flashinfer, + "cutedsl256": arm_cutedsl, + "fa4_wrapper": arm_fa4_wrapper, + "fa4": arm_fa4, +} + + +def detect_arms(requested): + if requested != ["auto"]: + return requested + arms = [] + try: + from flashinfer.sparse import BlockSparseAttentionWrapper # noqa: F401 + if "block_mask" not in inspect.signature(BlockSparseAttentionWrapper.plan).parameters: + raise RuntimeError( + "per-head block_mask planning API unavailable; issue #4554 used " + "flashinfer 0.6.16.post2" + ) + arms.append("flashinfer") + except Exception as exc: # noqa: BLE001 + print(f"# arm flashinfer unavailable: {exc!r}", flush=True) + fa4_available = False + try: + _load_fa4_symbols() + fa4_available = True + except Exception as exc: # noqa: BLE001 + print(f"# arm fa4 unavailable: {exc!r}", flush=True) + if fa4_available: + try: + from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_256_bshd # noqa: F401 + + arms.append("cutedsl256") + except Exception as exc: # noqa: BLE001 + print(f"# arm cutedsl256 unavailable (optional): {exc!r}", flush=True) + try: + from fastvideo_kernel.block_sparse_attn_cute_fwd import ( # noqa: F401 + block_sparse_attn_cute_fwd_bshd, + ) + + arms.append("fa4_wrapper") + except Exception as exc: # noqa: BLE001 + print(f"# arm fa4_wrapper unavailable (optional): {exc!r}", flush=True) + arms.append("fa4") + if not arms: + raise RuntimeError( + "no benchmark arm importable — need flashinfer-python, fastvideo_kernel, and/or FA4" + ) + return arms + + +def print_table(rows, arms): + header = "| seq | sparsity | Q blk | KV blk | keep/NKV |" + separator = "|---|---|---|---|---|" + for arm in arms: + header += f" {arm} ms | {arm} TFLOP/s | {arm} MFU % |" + separator += "---|---|---|" + if arm == "flashinfer": + header += " plan ms |" + separator += "---|" + if arm == "fa4": + header += " vs raw 256 % | vs cutedsl256 % |" + separator += "---|---|" + if arm == "fa4_wrapper": + header += " vs cutedsl256 % |" + separator += "---|" + print("\n" + header + "\n" + separator) + for row in rows: + line = ( + f"| {row['seq_len']} | {row['sparsity']} | {row['q_block']} | {row['kv_block']} | " + f"{row['keep_kv_blocks']}/{row['nkv_blocks']} |" + ) + for arm in arms: + cell = row.get(arm, {}) + if cell.get("status") == "ok": + line += f" {cell['fwd_ms']:.3f} | {cell['tflops']:.1f} | {cell['mfu_pct']:.1f} |" + if arm == "flashinfer": + line += f" {cell.get('plan_ms', float('nan')):.2f} |" + if arm == "fa4": + raw_relative = cell.get("vs_fa4_256_pct", float("nan")) + wrapper_relative = cell.get("vs_cutedsl256_pct", float("nan")) + line += f" {raw_relative:.1f} | {wrapper_relative:.1f} |" + if arm == "fa4_wrapper": + wrapper_relative = cell.get("vs_cutedsl256_pct", float("nan")) + line += f" {wrapper_relative:.1f} |" + else: + status = cell.get("status", "FAILED") + line += f" {status} | — | — |" + if arm == "flashinfer": + line += " — |" + if arm == "fa4": + line += " — | — |" + if arm == "fa4_wrapper": + line += " — |" + print(line) + print() + + +def add_relative_efficiency(cell_rows): + """Attach denominator-free efficiency relative to the 256x256 baselines.""" + baseline = next( + (row for row in cell_rows if (row["q_block"], row["kv_block"]) == (256, 256)), + None, + ) + if baseline is None: + return + raw_tflops = baseline.get("fa4", {}).get("tflops") + wrapper_tflops = baseline.get("cutedsl256", {}).get("tflops") + for row in cell_rows: + fa4_result = row.get("fa4", {}) + if fa4_result.get("status") == "ok": + if raw_tflops: + fa4_result["vs_fa4_256_pct"] = round( + 100.0 * fa4_result["tflops"] / raw_tflops, + 2, + ) + if wrapper_tflops: + fa4_result["vs_cutedsl256_pct"] = round( + 100.0 * fa4_result["tflops"] / wrapper_tflops, + 2, + ) + fine_wrapper_result = row.get("fa4_wrapper", {}) + if fine_wrapper_result.get("status") == "ok" and wrapper_tflops: + fine_wrapper_result["vs_cutedsl256_pct"] = round( + 100.0 * fine_wrapper_result["tflops"] / wrapper_tflops, + 2, + ) + + +def main(): + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if world_size > 1: + if local_rank >= torch.cuda.device_count(): + raise RuntimeError( + f"LOCAL_RANK={local_rank} exceeds visible CUDA device count={torch.cuda.device_count()}" + ) + torch.cuda.set_device(local_rank) + + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--seq_lens", type=int, nargs="+", default=[4096, 8192, 16384, 32768, 49152, 65536]) + parser.add_argument("--sparsities", type=str, nargs="+", default=list(KEEP_FRAC), choices=list(KEEP_FRAC)) + parser.add_argument("--arms", type=str, nargs="+", default=["auto"], choices=["auto", *ARMS]) + parser.add_argument( + "--block_shapes", + type=parse_block_shape, + nargs="+", + default=list(DEFAULT_BLOCK_SHAPES), + metavar="QxKV", + help="independent logical QxKV sparse block pairs (default: 256x256 128x64 64x64)", + ) + parser.add_argument( + "--mask_mode", + choices=("exact256", "native"), + default="exact256", + help=( + "exact256 expands one shared Q256/KV256 mask into every shape; " + "native samples each shape independently (default: exact256)" + ), + ) + parser.add_argument( + "--peak_bf16_tflops", + type=float, + default=DEFAULT_GB200_BF16_DENSE_TFLOPS, + help=( + "per-GPU dense BF16 tensor-core peak used only for MFU %% " + f"(default: official GB200 {DEFAULT_GB200_BF16_DENSE_TFLOPS:.0f} TFLOP/s)" + ), + ) + parser.add_argument( + "--quick", + action="store_true", + help="sanity grid: {8192,32768} x {dense,90} x requested block shapes", + ) + parser.add_argument("--out", default="bsa_bench_results.json") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--rep", type=int, default=20) + parser.add_argument("--num_heads", type=int, default=12) + parser.add_argument("--head_dim", type=int, default=128) + args = parser.parse_args() + if args.quick: + args.seq_lens, args.sparsities = [8192, 32768], ["dense", "90"] + if "auto" in args.arms and args.arms != ["auto"]: + parser.error("--arms auto cannot be combined with explicit arms") + if args.peak_bf16_tflops <= 0: + parser.error("--peak_bf16_tflops must be positive") + args.block_shapes = list(dict.fromkeys(args.block_shapes)) + for seq_len in args.seq_lens: + if args.mask_mode == "exact256" and seq_len % 256: + parser.error(f"exact256 mask mode requires sequence length divisible by 256; got {seq_len}") + for q_block, kv_block in args.block_shapes: + if seq_len % q_block or seq_len % kv_block: + parser.error( + f"sequence length {seq_len} must be divisible by QxKV block shape " + f"{q_block}x{kv_block}" + ) + if args.mask_mode == "exact256" and (256 % q_block or 256 % kv_block): + parser.error( + "exact256 mask mode requires both block sizes to divide 256; " + f"got {q_block}x{kv_block}" + ) + + from triton.testing import do_bench + + # The fp32 reference must be true fp32, not TF32. + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + + arms = detect_arms(args.arms) + batch, heads, head_dim = 1, args.num_heads, args.head_dim + scale = 1.0 / math.sqrt(head_dim) + + meta = { + "device": torch.cuda.get_device_name(), + "cuda_device_index": torch.cuda.current_device(), + "rank": rank, + "local_rank": local_rank, + "world_size": world_size, + "capability": list(torch.cuda.get_device_capability()), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "seed": args.seed, + "warmup": args.warmup, + "rep": args.rep, + "batch": batch, + "heads": heads, + "head_dim": head_dim, + "dtype": "bfloat16", + "arms": arms, + "block_shapes": [list(shape) for shape in args.block_shapes], + "mask_mode": args.mask_mode, + "mask_base_block": 256 if args.mask_mode == "exact256" else None, + "peak_bf16_dense_tflops": args.peak_bf16_tflops, + "peak_note": ( + "official single-GB200 dense BF16 tensor-core peak; override for clock-derived diagnostics" + ), + "peak_source": GB200_DATASHEET_URL, + "flop_formula": "4*bs*d*selected_logical_edges*q_block*kv_block", + "mfu_formula": "100*sparse_aware_algorithmic_tflops/peak_bf16_dense_tflops", + "correctness_gate": "finite and max_abs_vs_fp32 <= max(1e-2, 8*bf16_rounding_floor)", + "timing_scope": { + "flashinfer": "run only; plan reported separately", + "cutedsl256": "full FastVideo wrapper including mask expansion and map-to-index", + "fa4_wrapper": ( + "FastVideo direct fine-grained BSHD wrapper including map-to-index and VBS mask " + "plumbing; 64x64 is native Q64/KV64, not the public VSA-256 coalescing adapter" + ), + "fa4": "raw _flash_attn_fwd dispatch; mask/index preparation and output allocation excluded", + }, + "provenance": { + "issue": ISSUE_URL, + "gist": GIST_URL, + "gist_revision": GIST_REVISION, + }, + "fa4_env": { + name: os.environ.get(name) + for name in ( + "FASTVIDEO_FA4_VSA_DUAL_STREAM", + "FASTVIDEO_FA4_VSA_SP_DOUBLE_BUFFER", + "FASTVIDEO_VSA_FA4_BLOCK_SHAPE", + ) + }, + } + try: + meta["driver"] = subprocess.check_output( + ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"], text=True + ).split()[0] + except Exception: # noqa: BLE001 + meta["driver"] = None + for module in ("flashinfer", "triton", "fastvideo_kernel"): + try: + meta[module] = __import__(module).__version__ + except Exception: # noqa: BLE001 + pass + if any(arm in arms for arm in ("cutedsl256", "fa4_wrapper", "fa4")): + try: + _, _, fa4_interface = _load_fa4_symbols() + meta["fa4_interface"] = str(Path(fa4_interface.__file__).resolve()) + meta["fa4_import_mode"] = _FA4_IMPORT_MODE + except Exception as exc: # noqa: BLE001 + meta["fa4_import_error"] = f"{type(exc).__name__}: {exc}" + print("META " + json.dumps(meta), flush=True) + + rows = [] + for seq_len in args.seq_lens: + # Same q,k,v for every sparsity cell of this seq len (isolates sparsity). + set_seed(args.seed) + q = torch.randn(batch, seq_len, heads, head_dim, dtype=torch.bfloat16, device="cuda") + k = torch.randn(batch, seq_len, heads, head_dim, dtype=torch.bfloat16, device="cuda") + v = torch.randn(batch, seq_len, heads, head_dim, dtype=torch.bfloat16, device="cuda") + + for label in args.sparsities: + fraction = KEEP_FRAC[label] + base_mask_seed = args.seed + seq_len + int(fraction * 1000) + if args.mask_mode == "exact256": + base_nblocks_256 = seq_len // 256 + base_keep_256 = ( + base_nblocks_256 + if label == "dense" + else max(1, round(fraction * base_nblocks_256)) + ) + base_mask_256 = make_logical_mask( + heads, + base_nblocks_256, + base_nblocks_256, + base_keep_256, + base_mask_seed, + ) + else: + base_nblocks_256 = None + base_keep_256 = None + base_mask_256 = None + exact_ref = ( + ref_masked_sdpa_fp32( + q, + k, + v, + base_mask_256, + 256, + 256, + scale, + ) + if base_mask_256 is not None + else None + ) + cell_rows = [] + fa4_shape_outputs = {} + for q_block, kv_block in args.block_shapes: + nq_blocks, nkv_blocks = seq_len // q_block, seq_len // kv_block + if args.mask_mode == "exact256": + if base_mask_256 is None: # pragma: no cover - guarded above + raise AssertionError("exact256 mode requires a base mask") + keep = base_keep_256 * (256 // kv_block) + mask_seed = base_mask_seed + logical_mask = expand_exact256_mask(base_mask_256, q_block, kv_block) + else: + keep = nkv_blocks if label == "dense" else max(1, round(fraction * nkv_blocks)) + mask_seed = ( + base_mask_seed + + q_block * 1_000_003 + + kv_block * 9_176 + ) + logical_mask = make_logical_mask( + heads, + nq_blocks, + nkv_blocks, + keep, + mask_seed, + ) + selected_edges = heads * nq_blocks * keep + flops = flops_sparse_attention( + batch, + head_dim, + selected_edges, + q_block, + kv_block, + ) + row = { + "seq_len": seq_len, + "sparsity": label, + "mask_mode": args.mask_mode, + "base_nblocks_256": base_nblocks_256, + "base_keep_256": base_keep_256, + "base_mask_seed": base_mask_seed, + "q_block": q_block, + "kv_block": kv_block, + "nq_blocks": nq_blocks, + "nkv_blocks": nkv_blocks, + "keep_kv_blocks": keep, + "selected_logical_edges": selected_edges, + "actual_sparsity": round(1.0 - keep / nkv_blocks, 4), + "mask_seed": mask_seed, + "flops": flops, + } + + ref = ( + exact_ref + if exact_ref is not None + else ref_masked_sdpa_fp32( + q, + k, + v, + logical_mask, + q_block, + kv_block, + scale, + ) + ) + # Irreducible bf16 output-rounding floor: even a bit-exact kernel + # emitting bf16 cannot beat this vs the fp32 reference. + row["bf16_floor_max_abs"] = float( + (ref - ref.to(torch.bfloat16).float()).abs().max() + ) + row["correctness_max_abs_limit"] = max( + 1e-2, + 8.0 * row["bf16_floor_max_abs"], + ) + + outputs = {} + for name in arms: + result = {} + try: + fn, out, extra = ARMS[name]( + q, + k, + v, + logical_mask, + q_block, + kv_block, + ) + torch.cuda.synchronize() + result["max_abs_vs_fp32"] = float((out.float() - ref).abs().max()) + if not bool(torch.isfinite(out).all()): + raise RuntimeError(f"{name} produced NaN or Inf") + if result["max_abs_vs_fp32"] > row["correctness_max_abs_limit"]: + raise RuntimeError( + f"{name} max_abs_vs_fp32={result['max_abs_vs_fp32']:.6g} " + f"exceeds limit={row['correctness_max_abs_limit']:.6g}" + ) + outputs[name] = out + elapsed_ms = float( + do_bench(fn, warmup=args.warmup, rep=args.rep, quantiles=None) + ) + achieved_tflops = flops / elapsed_ms * 1e-9 + result["fwd_ms"] = round(elapsed_ms, 4) + result["tflops"] = round(achieved_tflops, 2) + result["mfu_pct"] = round( + 100.0 * achieved_tflops / args.peak_bf16_tflops, + 2, + ) + for key, value in extra.items(): + if key != "keep_alive": + result[key] = value + result["status"] = "ok" + del fn, extra + except UnsupportedConfig as exc: + result["status"] = "SKIPPED" + result["error"] = str(exc) + except Exception as exc: # noqa: BLE001 + result["status"] = "FAILED" + result["error"] = f"{type(exc).__name__}: {exc}"[:500] + traceback.print_exc() + row[name] = result + torch.cuda.empty_cache() + + output_names = list(outputs) + if len(output_names) > 1: + row["cross_max_abs"] = {} + for lhs_idx, lhs in enumerate(output_names): + for rhs in output_names[lhs_idx + 1:]: + row["cross_max_abs"][f"{lhs}_vs_{rhs}"] = float( + (outputs[lhs].float() - outputs[rhs].float()).abs().max() + ) + if args.mask_mode == "exact256" and "fa4" in outputs: + fa4_shape_outputs[(q_block, kv_block)] = outputs["fa4"] + outputs.clear() + del ref, logical_mask + torch.cuda.empty_cache() + cell_rows.append(row) + + if args.mask_mode == "exact256": + baseline_output = fa4_shape_outputs.get((256, 256)) + if baseline_output is not None: + for row in cell_rows: + shape = (row["q_block"], row["kv_block"]) + target_output = fa4_shape_outputs.get(shape) + if target_output is not None: + row["fa4"]["max_abs_vs_fa4_256"] = float( + (target_output.float() - baseline_output.float()).abs().max() + ) + fa4_shape_outputs.clear() + add_relative_efficiency(cell_rows) + for row in cell_rows: + print("ROW " + json.dumps(row), flush=True) + rows.extend(cell_rows) + del base_mask_256, exact_ref + del q, k, v + torch.cuda.empty_cache() + + output_path = Path(args.out) + if world_size > 1: + output_path = output_path.with_name( + f"{output_path.stem}.rank{rank}{output_path.suffix}" + ) + with output_path.open("w") as output_file: + json.dump({"meta": meta, "rows": rows}, output_file, indent=2) + print_table(rows, arms) + print(f"wrote {output_path}") + return 0 + + +if __name__ == "__main__": + if not torch.cuda.is_available(): + raise RuntimeError("CUDA required") + sys.exit(main()) diff --git a/fastvideo-kernel/fa4/.flake8 b/fastvideo-kernel/fa4/.flake8 new file mode 100644 index 000000000..bae5b85c0 --- /dev/null +++ b/fastvideo-kernel/fa4/.flake8 @@ -0,0 +1,4 @@ +[flake8] +max-line-length = 100 +# W503: line break before binary operator +ignore = E731, E741, F841, W503 diff --git a/fastvideo-kernel/fa4/AUTHORS b/fastvideo-kernel/fa4/AUTHORS new file mode 100644 index 000000000..055e75b66 --- /dev/null +++ b/fastvideo-kernel/fa4/AUTHORS @@ -0,0 +1,8 @@ +Tri Dao +Jay Shah +Ted Zadouri +Markus Hoehnerbach +Vijay Thakkar +Timmy Liu +Driss Guessous +Reuben Stern \ No newline at end of file diff --git a/fastvideo-kernel/fa4/LICENSE b/fastvideo-kernel/fa4/LICENSE new file mode 100644 index 000000000..5860e4b33 --- /dev/null +++ b/fastvideo-kernel/fa4/LICENSE @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +* Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/fastvideo-kernel/fa4/MANIFEST.in b/fastvideo-kernel/fa4/MANIFEST.in new file mode 100644 index 000000000..29bd8a50d --- /dev/null +++ b/fastvideo-kernel/fa4/MANIFEST.in @@ -0,0 +1,6 @@ +include UPSTREAM.md +global-exclude *.egg-info/* +prune flash_attn_4.egg-info +prune flash_attn.egg-info +prune build +prune dist diff --git a/fastvideo-kernel/fa4/README.md b/fastvideo-kernel/fa4/README.md new file mode 100644 index 000000000..c7f1b32eb --- /dev/null +++ b/fastvideo-kernel/fa4/README.md @@ -0,0 +1,33 @@ +# FlashAttention-4 (CuTeDSL) + +FlashAttention-4 is a CuTeDSL-based implementation of FlashAttention for Hopper and Blackwell GPUs. + +## Installation + +```sh +pip install flash-attn-4 +``` + +If you're on CUDA 13, install with the `cu13` extra for best performance: + +```sh +pip install "flash-attn-4[cu13]" +``` + +## Usage + +```python +from flash_attn.cute import flash_attn_func, flash_attn_varlen_func + +out = flash_attn_func(q, k, v, causal=True) +``` + +## Development + +```sh +git clone https://github.com/Dao-AILab/flash-attention.git +cd flash-attention +pip install -e "flash_attn/cute[dev]" # CUDA 12.x +pip install -e "flash_attn/cute[dev,cu13]" # CUDA 13.x (e.g. B200) +pytest tests/cute/ +``` diff --git a/fastvideo-kernel/fa4/UPSTREAM.md b/fastvideo-kernel/fa4/UPSTREAM.md new file mode 100644 index 000000000..83bed5a05 --- /dev/null +++ b/fastvideo-kernel/fa4/UPSTREAM.md @@ -0,0 +1,37 @@ +# FastVideo FA4 source fork + +This directory vendors the `flash_attn.cute` Python package from +[Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention) at +commit [`82d6441eec5d4dfec120153db2c0145ae855a083`](https://github.com/Dao-AILab/flash-attention/commit/82d6441eec5d4dfec120153db2c0145ae855a083). +The upstream source path is `flash_attn/cute/`, and its BSD-3-Clause license is +preserved in `LICENSE`. + +FastVideo keeps this as a first-class source fork because its video sparse +attention kernels need FA4 scheduling and tile-shape changes that are developed +and benchmarked together with `fastvideo-kernel`. The distribution intentionally +retains the upstream name `flash-attn-4` and import namespace `flash_attn.cute`. + +The forward-kernel delta currently contains FastVideo's single-Q-block, +two-KV-stream Q128/KV64 schedule, four-slot score/probability double buffering, +and FP32 online-softmax merge. The Q64/KV64 primitives were adapted from +upstream's guarded smaller-tile commit +[`526c18d25bcbc7fc7d6740ab3c7c84ed2d42cb0b`](https://github.com/Dao-AILab/flash-attention/commit/526c18d25bcbc7fc7d6740ab3c7c84ed2d42cb0b), +then extended with the same two-stream sparse traversal and merge. These are +forward-only specializations; the existing upstream paths remain the fallback +outside their explicit SM100/SM110 dtype, shape, and metadata gates. + +Install the working tree without resolving dependencies: + +```bash +uv pip install --no-deps --editable fastvideo-kernel/fa4 +``` + +The initial performance comparison came from FlashInfer issue +[#4554](https://github.com/flashinfer-ai/flashinfer/issues/4554) and its +[benchmark harness](https://gist.github.com/SolitaryThinker/90a1d1447929fc38dc509c1852e76532) +at gist revision `e15ac9066f23ef3690e33e1cc1fdac45b4b9099f`. + +When refreshing from upstream, copy only `flash_attn/cute/` from an explicitly +checked-out commit, retain this file and FastVideo-specific changes, update the +commit and package version above, and rerun the FA4 correctness and GB200 +performance gates before accepting the refresh. diff --git a/fastvideo-kernel/fa4/__init__.py b/fastvideo-kernel/fa4/__init__.py new file mode 100644 index 000000000..a3481f263 --- /dev/null +++ b/fastvideo-kernel/fa4/__init__.py @@ -0,0 +1,18 @@ +"""Flash Attention CUTE (CUDA Template Engine) implementation.""" + +from importlib.metadata import PackageNotFoundError, version + +try: + __version__ = version("flash-attn-4") +except PackageNotFoundError: + __version__ = "0.0.0" + +from .interface import ( + flash_attn_func, + flash_attn_varlen_func, +) + +__all__ = [ + "flash_attn_func", + "flash_attn_varlen_func", +] diff --git a/fastvideo-kernel/fa4/ampere_helpers.py b/fastvideo-kernel/fa4/ampere_helpers.py new file mode 100644 index 000000000..e3072d8ce --- /dev/null +++ b/fastvideo-kernel/fa4/ampere_helpers.py @@ -0,0 +1,103 @@ +# Copyright (c) 2025, Tri Dao. +from typing import Type, Callable, Optional + +import cutlass +import cutlass.cute as cute + + +def get_smem_layout_atom(dtype: Type[cutlass.Numeric], k_dim: int) -> cute.ComposedLayout: + dtype_byte = cutlass.const_expr(dtype.width // 8) + bytes_per_row = cutlass.const_expr(k_dim * dtype_byte) + smem_k_block_size = ( + cutlass.const_expr( + 128 + if bytes_per_row % 128 == 0 + else (64 if bytes_per_row % 64 == 0 else (32 if bytes_per_row % 32 == 0 else 16)) + ) + // dtype_byte + ) + swizzle_bits = ( + 4 + if smem_k_block_size == 128 + else (3 if smem_k_block_size == 64 else (2 if smem_k_block_size == 32 else 1)) + ) + swizzle_base = 2 if dtype_byte == 4 else (3 if dtype_byte == 2 else 4) + return cute.make_composed_layout( + cute.make_swizzle(swizzle_bits, swizzle_base, swizzle_base), + 0, + cute.make_ordered_layout( + (8 if cutlass.const_expr(k_dim % 32 == 0) else 16, smem_k_block_size), order=(1, 0) + ), + ) + + +@cute.jit +def gemm( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + tCsA: cute.Tensor, + tCsB: cute.Tensor, + smem_thr_copy_A: cute.TiledCopy, + smem_thr_copy_B: cute.TiledCopy, + hook_fn: Optional[Callable] = None, + A_in_regs: cutlass.Constexpr[bool] = False, + B_in_regs: cutlass.Constexpr[bool] = False, + swap_AB: cutlass.Constexpr[bool] = False, +) -> None: + if cutlass.const_expr(swap_AB): + gemm( + tiled_mma, + acc, + tCrB, + tCrA, + tCsB, + tCsA, + smem_thr_copy_B, + smem_thr_copy_A, + hook_fn, + A_in_regs=B_in_regs, + B_in_regs=A_in_regs, + swap_AB=False, + ) + else: + tCrA_copy_view = smem_thr_copy_A.retile(tCrA) + tCrB_copy_view = smem_thr_copy_B.retile(tCrB) + if cutlass.const_expr(not A_in_regs): + cute.copy(smem_thr_copy_A, tCsA[None, None, 0], tCrA_copy_view[None, None, 0]) + if cutlass.const_expr(not B_in_regs): + cute.copy(smem_thr_copy_B, tCsB[None, None, 0], tCrB_copy_view[None, None, 0]) + for k in cutlass.range_constexpr(cute.size(tCsA.shape[2])): + if k < cute.size(tCsA.shape[2]) - 1: + if cutlass.const_expr(not A_in_regs): + cute.copy( + smem_thr_copy_A, tCsA[None, None, k + 1], tCrA_copy_view[None, None, k + 1] + ) + if cutlass.const_expr(not B_in_regs): + cute.copy( + smem_thr_copy_B, tCsB[None, None, k + 1], tCrB_copy_view[None, None, k + 1] + ) + cute.gemm(tiled_mma, acc, tCrA[None, None, k], tCrB[None, None, k], acc) + if cutlass.const_expr(k == 0 and hook_fn is not None): + hook_fn() + + +@cute.jit +def gemm_rs( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + tCsB: cute.Tensor, + smem_thr_copy_B: cute.TiledCopy, + hook_fn: Optional[Callable] = None, +) -> None: + tCrB_copy_view = smem_thr_copy_B.retile(tCrB) + cute.copy(smem_thr_copy_B, tCsB[None, None, 0], tCrB_copy_view[None, None, 0]) + for k in cutlass.range_constexpr(cute.size(tCrA.shape[2])): + if cutlass.const_expr(k < cute.size(tCrA.shape[2]) - 1): + cute.copy(smem_thr_copy_B, tCsB[None, None, k + 1], tCrB_copy_view[None, None, k + 1]) + cute.gemm(tiled_mma, acc, tCrA[None, None, k], tCrB[None, None, k], acc) + if cutlass.const_expr(k == 0 and hook_fn is not None): + hook_fn() diff --git a/fastvideo-kernel/fa4/barrier.py b/fastvideo-kernel/fa4/barrier.py new file mode 100644 index 000000000..c999b1801 --- /dev/null +++ b/fastvideo-kernel/fa4/barrier.py @@ -0,0 +1,71 @@ +import cutlass +import cutlass.cute as cute +from cutlass import Int32 +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass._mlir.dialects import llvm + + +@dsl_user_op +def ld_acquire(lock_ptr: cute.Pointer, *, loc=None, ip=None) -> cutlass.Int32: + lock_ptr_i64 = lock_ptr.toint(loc=loc, ip=ip).ir_value() + state = llvm.inline_asm( + T.i32(), + [lock_ptr_i64], + "ld.global.acquire.gpu.b32 $0, [$1];", + "=r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + return cutlass.Int32(state) + + +@dsl_user_op +def red_relaxed( + lock_ptr: cute.Pointer, val: cutlass.Constexpr[Int32], *, loc=None, ip=None +) -> None: + lock_ptr_i64 = lock_ptr.toint(loc=loc, ip=ip).ir_value() + llvm.inline_asm( + None, + [lock_ptr_i64, Int32(val).ir_value(loc=loc, ip=ip)], + "red.relaxed.gpu.global.add.s32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def red_release( + lock_ptr: cute.Pointer, val: cutlass.Constexpr[Int32], *, loc=None, ip=None +) -> None: + lock_ptr_i64 = lock_ptr.toint(loc=loc, ip=ip).ir_value() + llvm.inline_asm( + None, + [lock_ptr_i64, Int32(val).ir_value(loc=loc, ip=ip)], + "red.release.gpu.global.add.s32 [$0], $1;", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@cute.jit +def wait_eq(lock_ptr: cute.Pointer, thread_idx: int | Int32, flag_offset: int, val: Int32) -> None: + flag_ptr = lock_ptr + flag_offset + if thread_idx == 0: + read_val = Int32(0) + while read_val != val: + read_val = ld_acquire(flag_ptr) + + +@cute.jit +def arrive_inc( + lock_ptr: cute.Pointer, thread_idx: int | Int32, flag_offset: int, val: cutlass.Constexpr[Int32] +) -> None: + flag_ptr = lock_ptr + flag_offset + if thread_idx == 0: + red_release(flag_ptr, val) + # red_relaxed(flag_ptr, val) diff --git a/fastvideo-kernel/fa4/bench_utils.py b/fastvideo-kernel/fa4/bench_utils.py new file mode 100644 index 000000000..d83f1c4bb --- /dev/null +++ b/fastvideo-kernel/fa4/bench_utils.py @@ -0,0 +1,243 @@ +"""Shared benchmark utilities: attention_ref, cuDNN helpers, flops calculation.""" + +import math +import torch + +try: + import cudnn +except ImportError: + cudnn = None + + +# ── FLOPS calculation ──────────────────────────────────────────────────────── + + +def flops( + batch, + nheads, + seqlen_q, + seqlen_k, + headdim, + headdim_v, + causal=False, + window_size=(None, None), + has_qv=False, +): + if causal: + avg_seqlen = (max(0, seqlen_k - seqlen_q) + seqlen_k) / 2 + else: + if window_size == (None, None): + avg_seqlen = seqlen_k + else: + row_idx = torch.arange(seqlen_q, device="cuda") + col_left = ( + torch.maximum(row_idx + seqlen_k - seqlen_q - window_size[0], torch.tensor(0)) + if window_size[0] is not None + else torch.zeros_like(row_idx) + ) + col_right = ( + torch.minimum( + row_idx + seqlen_k - seqlen_q + window_size[1], torch.tensor(seqlen_k - 1) + ) + if window_size[1] is not None + else torch.full_like(row_idx, seqlen_k - 1) + ) + avg_seqlen = (col_right - col_left + 1).float().mean().item() + eff_headdim = headdim + headdim_v if has_qv else headdim + return batch * nheads * 2 * seqlen_q * avg_seqlen * (eff_headdim + headdim_v) + + +# ── Bandwidth calculation ──────────────────────────────────────────────────── + + +def bandwidth_fwd_bytes( + batch, + nheads, + nheads_kv, + seqlen_q, + seqlen_k, + headdim, + headdim_v, + dtype_bytes=2, + has_qv=False, + shared_kv=False, +): + """HBM traffic for one attention pass: read Q,K,V + write O.""" + q = batch * nheads * seqlen_q * headdim + qv = batch * nheads * seqlen_q * headdim_v if has_qv else 0 + k = batch * nheads_kv * seqlen_k * headdim + v = batch * nheads_kv * seqlen_k * headdim_v if not shared_kv else 0 + o = batch * nheads * seqlen_q * headdim_v + return (q + qv + k + v + o) * dtype_bytes + + +def bandwidth_bwd_bytes( + batch, nheads, nheads_kv, seqlen_q, seqlen_k, headdim, headdim_v, dtype_bytes=2 +): + """HBM traffic for one attention pass: read Q,K,V,dO + write dQ,dK,dV.""" + q = batch * nheads * seqlen_q * headdim + k = batch * nheads_kv * seqlen_k * headdim + v = batch * nheads_kv * seqlen_k * headdim_v + do = batch * nheads * seqlen_q * headdim_v + dq = q + dk = k + dv = v + return (q + k + v + do + dq + dk + dv) * dtype_bytes + + +# ── Reference attention ───────────────────────────────────────────────────── + +_attention_ref_mask_cache = {} + + +def attention_ref(q, k, v, causal=False): + """Standard attention reference implementation. + + Args: + q, k, v: (batch, seqlen, nheads, headdim) tensors. + causal: whether to apply causal mask. + """ + softmax_scale = 1.0 / math.sqrt(q.shape[-1]) + scores = torch.einsum("bthd,bshd->bhts", q * softmax_scale, k) + if causal: + if scores.shape[-2] not in _attention_ref_mask_cache: + mask = torch.tril( + torch.ones(scores.shape[-2:], device=scores.device, dtype=torch.bool), diagonal=0 + ) + _attention_ref_mask_cache[scores.shape[-2]] = mask + else: + mask = _attention_ref_mask_cache[scores.shape[-2]] + scores = scores.masked_fill(mask, float("-inf")) + attn = torch.softmax(scores, dim=-1) + return torch.einsum("bhts,bshd->bthd", attn, v) + + +# ── cuDNN graph helpers ───────────────────────────────────────────────────── + +_TORCH_TO_CUDNN_DTYPE = { + torch.float16: "HALF", + torch.bfloat16: "BFLOAT16", + torch.float32: "FLOAT", + torch.int32: "INT32", + torch.int64: "INT64", +} + + +def _build_cudnn_graph(io_dtype, tensors, build_fn): + """Build a cuDNN graph. Returns (graph, variant_pack, workspace).""" + assert cudnn is not None, "cuDNN is not available" + cudnn_dtype = getattr(cudnn.data_type, _TORCH_TO_CUDNN_DTYPE[io_dtype]) + graph = cudnn.pygraph( + io_data_type=cudnn_dtype, + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + graph_tensors = {name: graph.tensor_like(t.detach()) for name, t in tensors.items()} + variant_pack = build_fn(graph, graph_tensors) + graph.validate() + graph.build_operation_graph() + graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + graph.check_support() + graph.build_plans() + workspace = torch.empty(graph.get_workspace_size(), device="cuda", dtype=torch.uint8) + return graph, variant_pack, workspace + + +def cudnn_fwd_setup(q, k, v, causal=False, window_size_left=None): + """Build a cuDNN forward SDPA graph. + + Args: + q, k, v: (batch, nheads, seqlen, headdim) tensors (cuDNN layout). + causal: whether to apply causal mask. + window_size_left: sliding window size (None for no window). + + Returns: + (fwd_fn, o_gpu, stats_gpu) where fwd_fn is a zero-arg callable. + """ + b, nheads, seqlen_q, headdim = q.shape + headdim_v = v.shape[-1] + o_gpu = torch.empty(b, nheads, seqlen_q, headdim_v, dtype=q.dtype, device=q.device) + stats_gpu = torch.empty(b, nheads, seqlen_q, 1, dtype=torch.float32, device=q.device) + + def build(graph, gt): + o, stats = graph.sdpa( + name="sdpa", + q=gt["q"], + k=gt["k"], + v=gt["v"], + is_inference=False, + attn_scale=1.0 / math.sqrt(headdim), + use_causal_mask=causal or window_size_left is not None, + sliding_window_length=window_size_left + if window_size_left is not None and not causal + else None, + ) + o.set_output(True).set_dim(o_gpu.shape).set_stride(o_gpu.stride()) + stats.set_output(True).set_data_type(cudnn.data_type.FLOAT) + return {gt["q"]: q, gt["k"]: k, gt["v"]: v, o: o_gpu, stats: stats_gpu} + + graph, variant_pack, workspace = _build_cudnn_graph(q.dtype, {"q": q, "k": k, "v": v}, build) + + def fwd_fn(): + graph.execute(variant_pack, workspace) + return o_gpu + + return fwd_fn, o_gpu, stats_gpu + + +def cudnn_bwd_setup(q, k, v, o, g, lse, causal=False, window_size_left=None): + """Build a cuDNN backward SDPA graph. + + Args: + q, k, v, o, g, lse: (batch, nheads, seqlen, dim) tensors (cuDNN layout). + causal: whether to apply causal mask. + window_size_left: sliding window size (None for no window). + + Returns: + bwd_fn: zero-arg callable that returns (dq, dk, dv). + """ + headdim = q.shape[-1] + dq_gpu, dk_gpu, dv_gpu = torch.empty_like(q), torch.empty_like(k), torch.empty_like(v) + + def build(graph, gt): + dq, dk, dv = graph.sdpa_backward( + name="sdpa_backward", + q=gt["q"], + k=gt["k"], + v=gt["v"], + o=gt["o"], + dO=gt["g"], + stats=gt["lse"], + attn_scale=1.0 / math.sqrt(headdim), + use_causal_mask=causal or window_size_left is not None, + sliding_window_length=window_size_left + if window_size_left is not None and not causal + else None, + use_deterministic_algorithm=False, + ) + dq.set_output(True).set_dim(dq_gpu.shape).set_stride(dq_gpu.stride()) + dk.set_output(True).set_dim(dk_gpu.shape).set_stride(dk_gpu.stride()) + dv.set_output(True).set_dim(dv_gpu.shape).set_stride(dv_gpu.stride()) + return { + gt["q"]: q, + gt["k"]: k, + gt["v"]: v, + gt["o"]: o, + gt["g"]: g, + gt["lse"]: lse, + dq: dq_gpu, + dk: dk_gpu, + dv: dv_gpu, + } + + graph, variant_pack, workspace = _build_cudnn_graph( + q.dtype, + {"q": q, "k": k, "v": v, "o": o, "g": g, "lse": lse}, + build, + ) + + def bwd_fn(): + graph.execute(variant_pack, workspace) + return dq_gpu, dk_gpu, dv_gpu + + return bwd_fn diff --git a/fastvideo-kernel/fa4/benchmark.py b/fastvideo-kernel/fa4/benchmark.py new file mode 100644 index 000000000..9a7820e7b --- /dev/null +++ b/fastvideo-kernel/fa4/benchmark.py @@ -0,0 +1,268 @@ +# Copyright (c) 2023, Tri Dao. +"""Useful functions for writing test code.""" + +import torch +import torch.utils.benchmark as benchmark + + +def benchmark_forward( + fn, *inputs, repeats=10, desc="", verbose=True, amp=False, amp_dtype=torch.float16, **kwinputs +): + """Use Pytorch Benchmark on the forward pass of an arbitrary function.""" + if verbose: + print(desc, "- Forward pass") + + def amp_wrapper(*inputs, **kwinputs): + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + fn(*inputs, **kwinputs) + + t = benchmark.Timer( + stmt="fn_amp(*inputs, **kwinputs)", + globals={"fn_amp": amp_wrapper, "inputs": inputs, "kwinputs": kwinputs}, + num_threads=torch.get_num_threads(), + ) + m = t.timeit(repeats) + if verbose: + print(m) + return t, m + + +def benchmark_backward( + fn, + *inputs, + grad=None, + repeats=10, + desc="", + verbose=True, + amp=False, + amp_dtype=torch.float16, + **kwinputs, +): + """Use Pytorch Benchmark on the backward pass of an arbitrary function.""" + if verbose: + print(desc, "- Backward pass") + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + y = fn(*inputs, **kwinputs) + if type(y) is tuple: + y = y[0] + if grad is None: + grad = torch.randn_like(y) + else: + if grad.shape != y.shape: + raise RuntimeError("Grad shape does not match output shape") + + def f(*inputs, y, grad): + # Set .grad to None to avoid extra operation of gradient accumulation + for x in inputs: + if isinstance(x, torch.Tensor): + x.grad = None + y.backward(grad, retain_graph=True) + + t = benchmark.Timer( + stmt="f(*inputs, y=y, grad=grad)", + globals={"f": f, "inputs": inputs, "y": y, "grad": grad}, + num_threads=torch.get_num_threads(), + ) + m = t.timeit(repeats) + if verbose: + print(m) + return t, m + + +def benchmark_combined( + fn, + *inputs, + grad=None, + repeats=10, + desc="", + verbose=True, + amp=False, + amp_dtype=torch.float16, + **kwinputs, +): + """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function.""" + if verbose: + print(desc, "- Forward + Backward pass") + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + y = fn(*inputs, **kwinputs) + if type(y) is tuple: + y = y[0] + if grad is None: + grad = torch.randn_like(y) + else: + if grad.shape != y.shape: + raise RuntimeError("Grad shape does not match output shape") + + def f(grad, *inputs, **kwinputs): + for x in inputs: + if isinstance(x, torch.Tensor): + x.grad = None + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + y = fn(*inputs, **kwinputs) + if type(y) is tuple: + y = y[0] + y.backward(grad, retain_graph=True) + + t = benchmark.Timer( + stmt="f(grad, *inputs, **kwinputs)", + globals={"f": f, "fn": fn, "inputs": inputs, "grad": grad, "kwinputs": kwinputs}, + num_threads=torch.get_num_threads(), + ) + m = t.timeit(repeats) + if verbose: + print(m) + return t, m + + +def benchmark_fwd_bwd( + fn, + *inputs, + grad=None, + repeats=10, + desc="", + verbose=True, + amp=False, + amp_dtype=torch.float16, + **kwinputs, +): + """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function.""" + return ( + benchmark_forward( + fn, + *inputs, + repeats=repeats, + desc=desc, + verbose=verbose, + amp=amp, + amp_dtype=amp_dtype, + **kwinputs, + ), + benchmark_backward( + fn, + *inputs, + grad=grad, + repeats=repeats, + desc=desc, + verbose=verbose, + amp=amp, + amp_dtype=amp_dtype, + **kwinputs, + ), + ) + + +def benchmark_all( + fn, + *inputs, + grad=None, + repeats=10, + desc="", + verbose=True, + amp=False, + amp_dtype=torch.float16, + **kwinputs, +): + """Use Pytorch Benchmark on the forward+backward pass of an arbitrary function.""" + return ( + benchmark_forward( + fn, + *inputs, + repeats=repeats, + desc=desc, + verbose=verbose, + amp=amp, + amp_dtype=amp_dtype, + **kwinputs, + ), + benchmark_backward( + fn, + *inputs, + grad=grad, + repeats=repeats, + desc=desc, + verbose=verbose, + amp=amp, + amp_dtype=amp_dtype, + **kwinputs, + ), + benchmark_combined( + fn, + *inputs, + grad=grad, + repeats=repeats, + desc=desc, + verbose=verbose, + amp=amp, + amp_dtype=amp_dtype, + **kwinputs, + ), + ) + + +def pytorch_profiler( + fn, + *inputs, + trace_filename=None, + backward=False, + amp=False, + amp_dtype=torch.float16, + cpu=False, + verbose=True, + **kwinputs, +): + """Wrap benchmark functions in Pytorch profiler to see CUDA information.""" + if backward: + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + out = fn(*inputs, **kwinputs) + if type(out) is tuple: + out = out[0] + g = torch.randn_like(out) + for _ in range(30): # Warm up + if backward: + for x in inputs: + if isinstance(x, torch.Tensor): + x.grad = None + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + out = fn(*inputs, **kwinputs) + if type(out) is tuple: + out = out[0] + # Backward should be done outside autocast + if backward: + out.backward(g, retain_graph=True) + activities = ([torch.profiler.ProfilerActivity.CPU] if cpu else []) + [ + torch.profiler.ProfilerActivity.CUDA + ] + with torch.profiler.profile( + activities=activities, + record_shapes=True, + # profile_memory=True, + with_stack=True, + ) as prof: + if backward: + for x in inputs: + if isinstance(x, torch.Tensor): + x.grad = None + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + out = fn(*inputs, **kwinputs) + if type(out) is tuple: + out = out[0] + if backward: + out.backward(g, retain_graph=True) + if verbose: + # print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=50)) + print(prof.key_averages().table(row_limit=50)) + if trace_filename is not None: + prof.export_chrome_trace(trace_filename) + + +def benchmark_memory(fn, *inputs, desc="", verbose=True, **kwinputs): + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + torch.cuda.synchronize() + fn(*inputs, **kwinputs) + torch.cuda.synchronize() + mem = torch.cuda.max_memory_allocated() / ((2**20) * 1000) + if verbose: + print(f"{desc} max memory: {mem}GB") + torch.cuda.empty_cache() + return mem diff --git a/fastvideo-kernel/fa4/benchmark_flash_attention_fp8.py b/fastvideo-kernel/fa4/benchmark_flash_attention_fp8.py new file mode 100644 index 000000000..e50777608 --- /dev/null +++ b/fastvideo-kernel/fa4/benchmark_flash_attention_fp8.py @@ -0,0 +1,434 @@ +# Benchmark FP8 attention for FA4 (CuTe-DSL) on SM100. +# +# Run (recommended): +# python -m flash_attn.cute.benchmark_flash_attention_fp8 +# +# Notes: +# - This is intended to be used while bringing up FP8 support for SM100. +# - FP8 correctness depends on descales + max-offset scaling being implemented in the SM100 kernel. +# This script optionally checks output vs a BF16 PyTorch baseline on dequantized FP8 inputs. +# +# Adapted from: `hopper/benchmark_flash_attention_fp8.py` + +from __future__ import annotations + +import argparse +import inspect +import math +import time +from typing import Iterable + +import torch +from einops import rearrange + +from flash_attn.cute.benchmark import benchmark_forward +from flash_attn.cute.interface import _flash_attn_fwd as flash_attn_cute_fwd + +try: + import cudnn +except ImportError: + cudnn = None + + +def _torch_float8_dtype(name: str) -> torch.dtype: + if name in ("fp8", "fp8_e4m3", "fp8_e4m3fn"): + return torch.float8_e4m3fn + if name in ("fp8_e5m2", "fp8_e5m2fn"): + return torch.float8_e5m2 + raise ValueError(f"Unsupported fp8 dtype name: {name}") + + +def _parse_int_list(csv: str) -> list[int]: + out: list[int] = [] + for part in csv.split(","): + part = part.strip() + if not part: + continue + out.append(int(part)) + return out + + +def attention_pytorch(qkv: torch.Tensor, causal: bool) -> torch.Tensor: + """ + qkv: (batch, seqlen, 3, nheads, headdim) + out: (batch, seqlen, nheads, headdim) + """ + batch_size, seqlen, _, nheads, d = qkv.shape + q, k, v = qkv.unbind(dim=2) + q = rearrange(q, "b t h d -> (b h) t d") + k = rearrange(k, "b s h d -> (b h) d s") + softmax_scale = 1.0 / math.sqrt(d) + scores = torch.empty(batch_size * nheads, seqlen, seqlen, dtype=qkv.dtype, device=qkv.device) + scores = rearrange( + torch.baddbmm(scores, q, k, beta=0, alpha=softmax_scale), "(b h) t s -> b h t s", h=nheads + ) + if causal: + causal_mask = torch.triu(torch.full((seqlen, seqlen), -10000.0, device=scores.device), 1) + scores = scores + causal_mask.to(dtype=scores.dtype) + attention = torch.softmax(scores, dim=-1) + output = torch.einsum("bhts,bshd->bthd", attention, v) + return output.to(dtype=qkv.dtype) + + +def flops(batch: int, seqlen: int, headdim: int, nheads: int, causal: bool) -> int: + # Matches the hopper benchmark’s convention. + return 4 * batch * seqlen**2 * nheads * headdim // (2 if causal else 1) + + +def efficiency(flop: int, seconds: float) -> float: + return (flop / seconds / 1e12) if not math.isnan(seconds) else 0.0 + + +def time_fwd(fn, *args, repeats: int, **kwargs) -> float: + time.sleep(1) # reduce residual throttling effects between benchmarks + _, m = benchmark_forward(fn, *args, repeats=repeats, verbose=False, **kwargs) + return float(m.mean) + + +def convert_to_cudnn_type(torch_type): + if torch_type == torch.float16: + return cudnn.data_type.HALF + if torch_type == torch.bfloat16: + return cudnn.data_type.BFLOAT16 + if torch_type == torch.float32: + return cudnn.data_type.FLOAT + if torch_type == torch.int32: + return cudnn.data_type.INT32 + if torch_type == torch.int64: + return cudnn.data_type.INT64 + if torch_type == torch.float8_e4m3fn: + return cudnn.data_type.FP8_E4M3 + if torch_type == torch.float8_e5m2: + return cudnn.data_type.FP8_E5M2 + raise ValueError("Unsupported tensor data type.") + + +def cudnn_sdpa_fp8_setup(qkv: torch.Tensor, seqlen_q: int, seqlen_k: int, causal: bool): + """Minimal cudnn.fp8 sdpa runner (optional).""" + assert cudnn is not None, "cudnn python bindings not available" + b, _, _, nheads, headdim = qkv.shape + o_gpu = torch.zeros(b, seqlen_q, nheads, headdim, dtype=qkv.dtype, device=qkv.device) + o_gpu_transposed = torch.as_strided( + o_gpu, + [b, nheads, seqlen_q, headdim], + [nheads * seqlen_q * headdim, headdim, nheads * headdim, 1], + ) + amax_s_gpu = torch.empty(1, 1, 1, 1, dtype=torch.float32, device=qkv.device) + amax_o_gpu = torch.empty(1, 1, 1, 1, dtype=torch.float32, device=qkv.device) + + graph = cudnn.pygraph( + io_data_type=convert_to_cudnn_type(qkv.dtype), + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + + new_q = torch.as_strided( + qkv, + [b, nheads, seqlen_q, headdim], + [seqlen_q * nheads * headdim * 3, headdim, headdim * nheads * 3, 1], + storage_offset=0, + ) + q = graph.tensor( + name="Q", + dim=list(new_q.shape), + stride=list(new_q.stride()), + data_type=convert_to_cudnn_type(qkv.dtype), + ) + + new_k = torch.as_strided( + qkv, + [b, nheads, seqlen_k, headdim], + [seqlen_k * nheads * headdim * 3, headdim, headdim * nheads * 3, 1], + storage_offset=nheads * headdim, + ) + k = graph.tensor( + name="K", + dim=list(new_k.shape), + stride=list(new_k.stride()), + data_type=convert_to_cudnn_type(qkv.dtype), + ) + + new_v = torch.as_strided( + qkv, + [b, nheads, seqlen_k, headdim], + [seqlen_k * nheads * headdim * 3, headdim, headdim * nheads * 3, 1], + storage_offset=nheads * headdim * 2, + ) + v = graph.tensor( + name="V", + dim=list(new_v.shape), + stride=list(new_v.stride()), + data_type=convert_to_cudnn_type(qkv.dtype), + ) + + def _scale_tensor(): + return graph.tensor(dim=[1, 1, 1, 1], stride=[1, 1, 1, 1], data_type=cudnn.data_type.FLOAT) + + default_scale_gpu = torch.ones(1, 1, 1, 1, dtype=torch.float32, device="cuda") + descale_q = _scale_tensor() + descale_k = _scale_tensor() + descale_v = _scale_tensor() + descale_s = _scale_tensor() + scale_s = _scale_tensor() + scale_o = _scale_tensor() + + o, _, amax_s, amax_o = graph.sdpa_fp8( + q=q, + k=k, + v=v, + descale_q=descale_q, + descale_k=descale_k, + descale_v=descale_v, + descale_s=descale_s, + scale_s=scale_s, + scale_o=scale_o, + is_inference=True, + attn_scale=1.0 / math.sqrt(headdim), + use_causal_mask=causal, + name="sdpa", + ) + o.set_output(True).set_dim(o_gpu_transposed.shape).set_stride(o_gpu_transposed.stride()) + amax_s.set_output(False).set_dim(amax_s_gpu.shape).set_stride(amax_s_gpu.stride()) + amax_o.set_output(False).set_dim(amax_o_gpu.shape).set_stride(amax_o_gpu.stride()) + + graph.validate() + graph.build_operation_graph() + graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK]) + graph.check_support() + graph.build_plans() + + variant_pack = { + q: new_q, + k: new_k, + v: new_v, + descale_q: default_scale_gpu, + descale_k: default_scale_gpu, + descale_v: default_scale_gpu, + descale_s: default_scale_gpu, + scale_s: default_scale_gpu, + scale_o: default_scale_gpu, + o: o_gpu_transposed, + amax_s: amax_s_gpu, + amax_o: amax_o_gpu, + } + workspace = torch.empty(graph.get_workspace_size(), device="cuda", dtype=torch.uint8) + + def run(): + graph.execute(variant_pack, workspace) + return o_gpu + + return run + + +def _maybe_pass_descales(callable_, **kwargs): + sig = inspect.signature(callable_) + return {k: v for k, v in kwargs.items() if k in sig.parameters} + + +def main(argv: Iterable[str] | None = None) -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--repeats", type=int, default=30) + parser.add_argument("--dim", type=int, default=2048) + parser.add_argument("--headdims", default="64,128") + parser.add_argument("--dtype", default="fp8_e4m3fn") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--check", + action=argparse.BooleanOptionalAction, + default=True, + help="Enable correctness checks vs BF16 PyTorch baseline.", + ) + parser.add_argument( + "--check-quantization-only", + action="store_true", + help="Check FP8 kernel vs dequantized-FP8 baseline (quantization error only).", + ) + parser.add_argument("--atol-bf16", type=float, default=0.10) + parser.add_argument("--rtol-bf16", type=float, default=0.10) + parser.add_argument("--atol-fp8", type=float, default=0.50) + parser.add_argument("--rtol-fp8", type=float, default=0.50) + parser.add_argument("--run-cudnn", action="store_true") + args = parser.parse_args(list(argv) if argv is not None else None) + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + major, minor = torch.cuda.get_device_capability() + if major != 10: + raise RuntimeError( + f"This benchmark is for SM100 (compute capability 10.x). Got {major}.{minor}." + ) + + torch.manual_seed(args.seed) + device = "cuda" + fp8_dtype = _torch_float8_dtype(args.dtype) + headdim_vals = _parse_int_list(args.headdims) + bs_seqlen_vals = [(32, 512), (16, 1024), (8, 2048), (4, 4096), (2, 8192), (1, 16384)] + + methods = ["Pytorch", "FA4-CuTe-BF16", "FA4-CuTe-FP8"] + ( + ["cuDNN-FP8"] if args.run_cudnn and cudnn is not None else [] + ) + + fp8_failures = [] + + for headdim in headdim_vals: + for causal in (False, True): + for batch, seqlen in bs_seqlen_vals: + torch.cuda.empty_cache() + nheads = args.dim // headdim + if args.dim % headdim != 0: + raise ValueError(f"--dim must be divisible by headdim ({args.dim=} {headdim=})") + + q_bf16 = torch.randn( + batch, seqlen, nheads, headdim, device=device, dtype=torch.bfloat16 + ) + k_bf16 = torch.randn( + batch, seqlen, nheads, headdim, device=device, dtype=torch.bfloat16 + ) + v_bf16 = torch.randn( + batch, seqlen, nheads, headdim, device=device, dtype=torch.bfloat16 + ) + qkv_bf16 = torch.stack([q_bf16, k_bf16, v_bf16], dim=2) + + times = {} + speeds = {} + + out_ref_bf16 = None + try: + out_ref_bf16 = attention_pytorch(qkv_bf16, causal=causal) # warmup / reference + t = time_fwd(attention_pytorch, qkv_bf16, causal=causal, repeats=args.repeats) + times["Pytorch"] = t + except RuntimeError as e: + if "out of memory" in str(e).lower(): + times["Pytorch"] = float("nan") + out_ref_bf16 = None + else: + raise + + # FA4 / CuTe BF16 baseline + try: + softmax_scale = headdim**-0.5 + out_fa4_bf16, *_ = flash_attn_cute_fwd( + q_bf16, k_bf16, v_bf16, softmax_scale=softmax_scale, causal=causal + ) # warmup / compile + t = time_fwd( + flash_attn_cute_fwd, + q_bf16, + k_bf16, + v_bf16, + softmax_scale=softmax_scale, + causal=causal, + repeats=args.repeats, + ) + times["FA4-CuTe-BF16"] = t + if args.check and out_ref_bf16 is not None: + torch.testing.assert_close( + out_fa4_bf16, + out_ref_bf16, + atol=args.atol_bf16, + rtol=args.rtol_bf16, + ) + except Exception as e: + # Treat as fatal: BF16 kernel should be usable for basic sanity checking. + raise RuntimeError("FA4-CuTe BF16 baseline failed") from e + + # FA4 / CuTe FP8 + q_fp8 = q_bf16.to(fp8_dtype) + k_fp8 = k_bf16.to(fp8_dtype) + v_fp8 = v_bf16.to(fp8_dtype) + + # Placeholder descales (FA3-style: per-(batch, kv_head)). + q_descale = torch.ones(batch, nheads, device=device, dtype=torch.float32) + k_descale = torch.ones(batch, nheads, device=device, dtype=torch.float32) + v_descale = torch.ones(batch, nheads, device=device, dtype=torch.float32) + + # Optional: FP8 reference baseline (dequantized FP8 -> PyTorch) for quantization-error-only checks + out_ref_fp8 = None + if args.check and args.check_quantization_only: + try: + # Dequantize FP8 inputs back to BF16 (applying descales) + q_ref_fp8 = (q_fp8.to(torch.bfloat16) * q_descale[:, None, :, None]).to( + torch.bfloat16 + ) + k_ref_fp8 = (k_fp8.to(torch.bfloat16) * k_descale[:, None, :, None]).to( + torch.bfloat16 + ) + v_ref_fp8 = (v_fp8.to(torch.bfloat16) * v_descale[:, None, :, None]).to( + torch.bfloat16 + ) + qkv_ref_fp8 = torch.stack([q_ref_fp8, k_ref_fp8, v_ref_fp8], dim=2) + out_ref_fp8 = attention_pytorch(qkv_ref_fp8, causal=causal) + except RuntimeError as e: + if "out of memory" in str(e).lower(): + out_ref_fp8 = None + else: + raise + + fa4_kwargs = dict(softmax_scale=softmax_scale, causal=causal) + fa4_kwargs.update( + _maybe_pass_descales( + flash_attn_cute_fwd, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + ) + ) + + try: + # Warmup/compile (will raise until FP8 is implemented) + out_fa4_fp8, *_ = flash_attn_cute_fwd(q_fp8, k_fp8, v_fp8, **fa4_kwargs) + t = time_fwd( + flash_attn_cute_fwd, + q_fp8, + k_fp8, + v_fp8, + repeats=args.repeats, + **fa4_kwargs, + ) + times["FA4-CuTe-FP8"] = t + if args.check: + # Choose baseline: quantization-only (dequantized FP8) or full (BF16) + if args.check_quantization_only: + ref_baseline = out_ref_fp8 + else: + ref_baseline = out_ref_bf16 + + if ref_baseline is not None: + torch.testing.assert_close( + out_fa4_fp8, + ref_baseline, + atol=args.atol_fp8, + rtol=args.rtol_fp8, + ) + except Exception as e: + fp8_failures.append((causal, headdim, batch, seqlen, repr(e))) + times["FA4-CuTe-FP8"] = float("nan") + + if args.run_cudnn and cudnn is not None: + qkv_fp8 = qkv_bf16.to(fp8_dtype) + runner = cudnn_sdpa_fp8_setup(qkv_fp8, seqlen, seqlen, causal=causal) + _ = runner() # warmup + t = time_fwd(lambda: runner(), repeats=args.repeats) + times["cuDNN-FP8"] = t + + print(f"### causal={causal}, headdim={headdim}, batch={batch}, seqlen={seqlen} ###") + for method in methods: + t = times.get(method, float("nan")) + speeds[method] = efficiency(flops(batch, seqlen, headdim, nheads, causal), t) + if math.isnan(t): + print(f"{method} fwd: (skipped)") + else: + print(f"{method} fwd: {speeds[method]:.2f} TFLOPs/s, {t * 1e3:.3f} ms") + if math.isnan(times.get("FA4-CuTe-FP8", float("nan"))): + print("FA4-CuTe-FP8 status: FAILED") + + if fp8_failures: + print(f"\nFP8 failures: {len(fp8_failures)} (showing first 5)") + for causal, headdim, batch, seqlen, err in fp8_failures[:5]: + print(f"- causal={causal} headdim={headdim} batch={batch} seqlen={seqlen}: {err}") + return 1 + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/fastvideo-kernel/fa4/blackwell_helpers.py b/fastvideo-kernel/fa4/blackwell_helpers.py new file mode 100644 index 000000000..3a38f6352 --- /dev/null +++ b/fastvideo-kernel/fa4/blackwell_helpers.py @@ -0,0 +1,1115 @@ +# Copyright (c) 2025, Tri Dao. +from typing import Optional, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Boolean, const_expr +from cutlass.cute.nvgpu import tcgen05 +from cutlass._mlir.dialects import llvm + +import flash_attn.cute.mma_sm100_desc as sm100_desc + + +def _tcgen05_mma_kind(op: cute.nvgpu.tcgen05.mma.MmaOp) -> str: + if isinstance(op, tcgen05.mma.MmaF16BF16Op): + return "f16" + if isinstance(op, tcgen05.mma.MmaTF32Op): + return "tf32" + if isinstance(op, tcgen05.mma.MmaI8Op): + return "i8" + # cutlass-dsl >=4.5.2 builds plain FP8 MMAs as MmaF8F6F4Op (make_trivial_tiled_mma's + # _F8F6F4_TYPES branch); <4.4.x returned the now-legacy MmaFP8Op. Both map to kind::f8f6f4. + if isinstance(op, (tcgen05.mma.MmaFP8Op, tcgen05.mma.MmaF8F6F4Op)): + return "f8f6f4" + if isinstance(op, tcgen05.mma.MmaMXF8Op): + return "mxf8f6f4" + if isinstance(op, tcgen05.mma.MmaMXF4Op): + return "mxf4" + if isinstance(op, tcgen05.mma.MmaMXF4NVF4Op): + return "mxf4nvf4" + raise TypeError(f"Unsupported tcgen05 MMA op kind: {type(op).__name__}") + + +@cute.jit +def gemm_w_idx( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + A_idx: Optional[Int32] = None, + B_idx: Optional[Int32] = None, + zero_init: bool | Boolean = False, + swap_AB: bool = False, + num_unroll_groups: int = 1, +) -> None: + if const_expr(swap_AB): + return gemm_w_idx( + tiled_mma, acc, tCrB, tCrA, B_idx, A_idx, zero_init=zero_init, swap_AB=False + ) + else: + rA = tCrA if const_expr(A_idx is None) else tCrA[None, None, None, A_idx] + rB = tCrB if const_expr(B_idx is None) else tCrB[None, None, None, B_idx] + + mma_atom = cute.make_mma_atom(tiled_mma.op) + for k in cutlass.range( + cute.size(tCrA.shape[2]), unroll=cute.size(tCrA.shape[2]) // num_unroll_groups + ): + mma_atom.set(tcgen05.Field.ACCUMULATE, not zero_init or k != 0) + cute.gemm(mma_atom, acc, rA[None, None, k], rB[None, None, k], acc) + + +@cute.jit +def gemm_ptx_w_idx( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + sA: Optional[cute.Tensor], + sB: cute.Tensor, + A_idx: Optional[Int32] = None, + B_idx: Optional[Int32] = None, + zero_init: bool | Boolean = False, + cta_group: int = 1, + **kwargs, +) -> None: + rA = tCrA if const_expr(A_idx is None) else tCrA[None, None, None, A_idx] + rB = tCrB if const_expr(B_idx is None) else tCrB[None, None, None, B_idx] + sA_cur = None + if const_expr(sA is not None): + sA_cur = sA if const_expr(A_idx is None) else sA[None, None, None, A_idx] + sB_cur = sB if const_expr(B_idx is None) else sB[None, None, None, B_idx] + mma_atom = cute.make_mma_atom(tiled_mma.op) + acc_tmem_addr = acc.iterator.toint() + gemm_ptx_partial( + mma_atom.op, + acc_tmem_addr, + rA, + rB, + sA_cur, + sB_cur, + zero_init=zero_init, + cta_group=cta_group, + **kwargs, + ) + + +@cute.jit +def gemm( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + zero_init: bool | Boolean = False, +) -> None: + mma_atom = cute.make_mma_atom(tiled_mma.op) + for k in cutlass.range_constexpr(cute.size(tCrA.shape[2])): + mma_atom.set(tcgen05.Field.ACCUMULATE, not zero_init or k != 0) + cute.gemm(mma_atom, acc, tCrA[None, None, k], tCrB[None, None, k], acc) + + +def i64_to_i32x2(i: int) -> Tuple[int, int]: + """Convert a 64-bit integer to a tuple of two 32-bit integers.""" + return i & 0xFFFF_FFFF, (i >> 32) & 0xFFFF_FFFF + + +@cute.jit +def gemm_ptx( + op: cute.nvgpu.tcgen05.mma.MmaOp, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + sA: Optional[cute.Tensor], + sB: cute.Tensor, + zero_init: bool | Boolean = False, +) -> None: + is_ts = op.a_src == cute.nvgpu.tcgen05.OperandSource.TMEM + if const_expr(not is_ts): + assert sA is not None, "sA must be provided when a_src is not TMEM" + sA_layout = sA.layout if sA is not None else None + sB_layout = sB.layout + idesc: int = const_expr(sm100_desc.mma_op_to_idesc(op)) + kind = _tcgen05_mma_kind(op) + if const_expr(not is_ts): + sA_swizzle = sA.iterator.type.swizzle_type + smem_desc_base_a: int = const_expr( + sm100_desc.make_smem_desc_base( + cute.recast_layout(128, op.a_dtype.width, sA_layout[0]), + sA_swizzle, + sm100_desc.Major.K + if const_expr(op.a_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K) + else sm100_desc.Major.MN, + ) + ) + smem_desc_base_a_lo, smem_desc_a_hi = i64_to_i32x2(smem_desc_base_a) + smem_desc_base_a_lo = const_expr(smem_desc_base_a_lo) + smem_desc_a_hi = const_expr(smem_desc_a_hi) + else: + smem_desc_base_a = None + smem_desc_base_a_lo, smem_desc_a_hi = None, None + sB_swizzle = sB.iterator.type.swizzle_type + smem_desc_base_b: int = const_expr( + sm100_desc.make_smem_desc_base( + cute.recast_layout(128, op.b_dtype.width, sB_layout[0]), + sB_swizzle, + sm100_desc.Major.K + if const_expr(op.b_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K) + else sm100_desc.Major.MN, + ) + ) + smem_desc_base_b_lo, smem_desc_b_hi = i64_to_i32x2(smem_desc_base_b) + smem_desc_base_b_lo = const_expr(smem_desc_base_b_lo) + smem_desc_b_hi = const_expr(smem_desc_b_hi) + + if const_expr(not is_ts): + smem_desc_start_a_lo = Int32(smem_desc_base_a_lo) | sm100_desc.make_smem_desc_start_addr( + sA[None, None, 0].iterator + ) + else: + smem_desc_start_a_lo = None + smem_desc_start_b_lo = Int32(smem_desc_base_b_lo) | sm100_desc.make_smem_desc_start_addr( + sB[None, None, 0].iterator + ) + for k in cutlass.range_constexpr(cute.size(tCrA.shape[2])): + if const_expr(not is_ts): + smem_desc_a_lo = smem_desc_start_a_lo + ( + (cute.crd2idx((0, 0, k), sA_layout) * sA.element_type.width // 8) >> 4 + ) + smem_desc_b_lo = smem_desc_start_b_lo + ( + (cute.crd2idx((0, 0, k), sB_layout) * sB.element_type.width // 8) >> 4 + ) + # with cute.arch.elect_one(): + # cute.printf("smem_desc_a_lo = {}, smem_desc_b_lo = {}", smem_desc_a_lo, smem_desc_b_lo) + # cute.printf("smem_desc_a_lo_correct = {}, smem_desc_b_lo_correct = {}", smem_desc_a_lo_correct, smem_desc_b_lo_correct) + with cute.arch.elect_one(): + if const_expr(not is_ts): + llvm.inline_asm( + None, + [ + acc.iterator.toint().ir_value(), + smem_desc_a_lo.ir_value(), + smem_desc_b_lo.ir_value(), + Int32(not zero_init or k != 0).ir_value(), + ], + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b64 smem_desc_a, smem_desc_b;\n\t" + ".reg .b32 idesc;\n\t" + f"mov.b32 idesc, {hex(idesc)};\n\t" + f"mov.b64 smem_desc_a, {{$1, {hex(smem_desc_a_hi)}}};\n\t" + f"mov.b64 smem_desc_b, {{$2, {hex(smem_desc_b_hi)}}};\n\t" + "setp.ne.b32 p, $3, 0;\n\t" + f"tcgen05.mma.cta_group::1.kind::{kind} [$0], smem_desc_a, smem_desc_b, idesc, p;\n\t" + "}\n", + "r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + else: + llvm.inline_asm( + None, + [ + acc.iterator.toint().ir_value(), + tCrA[None, None, k].iterator.toint().ir_value(), + smem_desc_b_lo.ir_value(), + Int32(not zero_init or k != 0).ir_value(), + ], + "{\n\t" + ".reg .pred p;\n\t" + ".reg .b64 smem_desc_b;\n\t" + f"mov.b64 smem_desc_b, {{$2, {hex(smem_desc_b_hi)}}};\n\t" + "setp.ne.b32 p, $3, 0;\n\t" + f"tcgen05.mma.cta_group::1.kind::{kind} [$0], [$1], smem_desc_b, {hex(idesc)}, p;\n\t" + "}\n", + "r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@cute.jit +def gemm_ptx_loop( + op: cute.nvgpu.tcgen05.mma.MmaOp, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + sA: Optional[cute.Tensor], + sB: cute.Tensor, + zero_init: bool | Boolean = False, +) -> None: + is_ts = op.a_src == cute.nvgpu.tcgen05.OperandSource.TMEM + if const_expr(not is_ts): + assert sA is not None, "sA must be provided when a_src is not TMEM" + sA_layout = sA.layout if sA is not None else tCrA.layout + sB_layout = sB.layout + idesc: int = const_expr(sm100_desc.mma_op_to_idesc(op)) + kind = _tcgen05_mma_kind(op) + if const_expr(not is_ts): + sA_swizzle = sA.iterator.type.swizzle_type + smem_desc_base_a: int = const_expr( + sm100_desc.make_smem_desc_base( + cute.recast_layout(128, op.a_dtype.width, sA_layout[0]), + sA_swizzle, + sm100_desc.Major.K + if const_expr(op.a_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K) + else sm100_desc.Major.MN, + ) + ) + smem_desc_base_a_lo, smem_desc_a_hi = i64_to_i32x2(smem_desc_base_a) + smem_desc_base_a_lo = const_expr(smem_desc_base_a_lo) + smem_desc_a_hi = const_expr(smem_desc_a_hi) + else: + smem_desc_base_a = None + smem_desc_base_a_lo, smem_desc_a_hi = None, None + sB_swizzle = sB.iterator.type.swizzle_type + smem_desc_base_b: int = const_expr( + sm100_desc.make_smem_desc_base( + cute.recast_layout(128, op.b_dtype.width, sB_layout[0]), + sB_swizzle, + sm100_desc.Major.K + if const_expr(op.b_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K) + else sm100_desc.Major.MN, + ) + ) + smem_desc_base_b_lo, smem_desc_b_hi = i64_to_i32x2(smem_desc_base_b) + smem_desc_base_b_lo = const_expr(smem_desc_base_b_lo) + smem_desc_b_hi = const_expr(smem_desc_b_hi) + + if const_expr(not is_ts): + offset_a = [ + (cute.crd2idx((0, 0, k), sA_layout) * sA.element_type.width // 8) >> 4 + for k in cutlass.range_constexpr(cute.size(tCrA.shape[2])) + ] + else: + offset_a = [ + cute.crd2idx((0, 0, k), sA_layout) * op.a_dtype.width // 32 + for k in cutlass.range_constexpr(cute.size(tCrA.shape[2])) + ] + offset_a_diff = [ + offset_a[k] - offset_a[k - 1] for k in cutlass.range_constexpr(1, cute.size(tCrA.shape[2])) + ] + offset_b = [ + (cute.crd2idx((0, 0, k), sB_layout) * sB.element_type.width // 8) >> 4 + for k in cutlass.range_constexpr(cute.size(tCrB.shape[2])) + ] + offset_b_diff = [ + offset_b[k] - offset_b[k - 1] for k in cutlass.range_constexpr(1, cute.size(tCrB.shape[2])) + ] + + if const_expr(not is_ts): + smem_desc_start_a_lo = Int32( + smem_desc_base_a_lo | sm100_desc.make_smem_desc_start_addr(sA[None, None, 0].iterator) + ) + else: + smem_desc_start_a_lo = None + smem_desc_start_b_lo = Int32( + smem_desc_base_b_lo | sm100_desc.make_smem_desc_start_addr(sB[None, None, 0].iterator) + ) + pred_str = "p" if isinstance(zero_init, Boolean) else "0" if zero_init else "1" + if const_expr(not is_ts): + llvm.inline_asm( + None, + [ + acc.iterator.toint().ir_value(), + Int32(cute.arch.make_warp_uniform(smem_desc_start_a_lo)).ir_value(), + Int32(cute.arch.make_warp_uniform(smem_desc_start_b_lo)).ir_value(), + Int32(not zero_init).ir_value(), + ], + "{\n\t" + ".reg .pred leader_thread;\n\t" + ".reg .pred p;\n\t" + ".reg .b32 idesc;\n\t" + ".reg .b32 smem_desc_a_lo, smem_desc_b_lo;\n\t" + ".reg .b32 smem_desc_a_hi, smem_desc_b_hi;\n\t" + ".reg .b64 smem_desc_a, smem_desc_b;\n\t" + "elect.sync _|leader_thread, -1;\n\t" + f"mov.b32 idesc, {hex(idesc)};\n\t" + "mov.b32 smem_desc_a_lo, $1;\n\t" + "mov.b32 smem_desc_b_lo, $2;\n\t" + f"mov.b32 smem_desc_a_hi, {hex(smem_desc_a_hi)};\n\t" + f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t" + f"mov.b64 smem_desc_a, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + "setp.ne.b32 p, $3, 0;\n\t" + f"@leader_thread tcgen05.mma.cta_group::1.kind::{kind} [$0], smem_desc_a, smem_desc_b, idesc, {pred_str};\n\t" + + "".join( + ( + f"add.u32 smem_desc_a_lo, smem_desc_a_lo, {hex(offset_a_diff[k - 1])};\n\t" + f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t" + f"mov.b64 smem_desc_a, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + f"@leader_thread tcgen05.mma.cta_group::1.kind::{kind} [$0], smem_desc_a, smem_desc_b, idesc, 1;\n\t" + ) + for k in cutlass.range_constexpr(1, cute.size(tCrA.shape[2])) + ) + + "}\n", + "r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + else: + llvm.inline_asm( + None, + [ + acc.iterator.toint().ir_value(), + Int32(tCrA[None, None, 0].iterator.toint()).ir_value(), + Int32(smem_desc_start_b_lo).ir_value(), + Int32(not zero_init).ir_value(), + ], + "{\n\t" + ".reg .pred leader_thread;\n\t" + ".reg .pred p;\n\t" + ".reg .b32 idesc;\n\t" + ".reg .b32 tmem_a;\n\t" + ".reg .b32 smem_desc_b_lo;\n\t" + ".reg .b32 smem_desc_b_hi;\n\t" + ".reg .b64 smem_desc_b;\n\t" + "elect.sync _|leader_thread, -1;\n\t" + f"mov.b32 idesc, {hex(idesc)};\n\t" + "mov.b32 tmem_a, $1;\n\t" + "mov.b32 smem_desc_b_lo, $2;\n\t" + f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + "setp.ne.b32 p, $3, 0;\n\t" + f"@leader_thread tcgen05.mma.cta_group::1.kind::{kind} [$0], [tmem_a], smem_desc_b, idesc, {pred_str};\n\t" + + "".join( + ( + # f"add.u32 tmem_a, tmem_a, {hex(offset_a_diff[k - 1])};\n\t" + f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + # f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [$0], [tmem_a], smem_desc_b, idesc, 1;\n\t" + f"@leader_thread tcgen05.mma.cta_group::1.kind::{kind} [$0], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t" + ) + for k in cutlass.range_constexpr(1, cute.size(tCrA.shape[2])) + ) + + "}\n", + "r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@cute.jit +def gemm_ptx_partial( + op: cute.nvgpu.tcgen05.mma.MmaOp, + acc_tmem_addr: Int32, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + sA: Optional[cute.Tensor], + sB: cute.Tensor, + mbar_ptr: Optional[cutlass.Pointer] = None, + mbar_phase: Optional[Int32] = None, + split_arrive: Optional[int] = None, + zero_init: bool | Boolean = False, + # sA_offset: Int32 = 0, + # acc_offset: Int32 = 0, + tA_addr: Optional[Int32] = None, + cta_group: int = 1, +) -> None: + # acc_tmem_addr += acc_offset + is_ts = op.a_src == cute.nvgpu.tcgen05.OperandSource.TMEM + if const_expr(not is_ts): + assert sA is not None, "sA must be provided when a_src is not TMEM" + sA_layout = sA.layout if sA is not None else tCrA.layout + sB_layout = sB.layout + idesc: int = const_expr(sm100_desc.mma_op_to_idesc(op)) + kind = _tcgen05_mma_kind(op) + if const_expr(not is_ts): + sA_swizzle = sA.iterator.type.swizzle_type + smem_desc_base_a: int = const_expr( + sm100_desc.make_smem_desc_base( + cute.recast_layout(128, op.a_dtype.width, sA_layout[0]), + sA_swizzle, + sm100_desc.Major.K + if const_expr(op.a_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K) + else sm100_desc.Major.MN, + ) + ) + smem_desc_base_a_lo, smem_desc_a_hi = i64_to_i32x2(smem_desc_base_a) + smem_desc_base_a_lo = const_expr(smem_desc_base_a_lo) + smem_desc_a_hi = const_expr(smem_desc_a_hi) + else: + smem_desc_base_a = None + smem_desc_base_a_lo, smem_desc_a_hi = None, None + sB_swizzle = sB.iterator.type.swizzle_type + smem_desc_base_b: int = const_expr( + sm100_desc.make_smem_desc_base( + cute.recast_layout(128, op.b_dtype.width, sB_layout[0]), + sB_swizzle, + sm100_desc.Major.K + if const_expr(op.b_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K) + else sm100_desc.Major.MN, + ) + ) + smem_desc_base_b_lo, smem_desc_b_hi = i64_to_i32x2(smem_desc_base_b) + smem_desc_base_b_lo = const_expr(smem_desc_base_b_lo) + smem_desc_b_hi = const_expr(smem_desc_b_hi) + + tCrA_layout = ( + tCrA.layout + if const_expr(not is_ts) + else cute.recast_layout(32, tCrA.element_type.width, tCrA.layout) + ) + offset_a = [cute.crd2idx((0, 0, k), tCrA_layout) for k in range(cute.size(tCrA.shape[2]))] + offset_a_diff = [offset_a[k] - offset_a[k - 1] for k in range(1, cute.size(tCrA.shape[2]))] + offset_b = [cute.crd2idx((0, 0, k), tCrB.layout) for k in range(cute.size(tCrB.shape[2]))] + offset_b_diff = [offset_b[k] - offset_b[k - 1] for k in range(1, cute.size(tCrB.shape[2]))] + + if const_expr(not is_ts): + smem_desc_start_a_lo = Int32( + smem_desc_base_a_lo | sm100_desc.make_smem_desc_start_addr(sA[None, None, 0].iterator) + ) + # ) + sA_offset + else: + smem_desc_start_a_lo = None + smem_desc_start_b_lo = Int32( + smem_desc_base_b_lo | sm100_desc.make_smem_desc_start_addr(sB[None, None, 0].iterator) + ) + pred_str = "p" if isinstance(zero_init, Boolean) else "0" if zero_init else "1" + if const_expr(not is_ts): + assert mbar_ptr is None, "mbar_ptr must be None when a_src is not TMEM" + llvm.inline_asm( + None, + [ + # acc.iterator.toint().ir_value(), + Int32(cute.arch.make_warp_uniform(smem_desc_start_a_lo)).ir_value(), + Int32(cute.arch.make_warp_uniform(smem_desc_start_b_lo)).ir_value(), + Int32(not zero_init).ir_value(), + Int32(cute.arch.make_warp_uniform(acc_tmem_addr)).ir_value(), + ], + "{\n\t" + ".reg .pred leader_thread;\n\t" + ".reg .pred p;\n\t" + ".reg .b32 idesc;\n\t" + ".reg .b32 tmem_acc;\n\t" + ".reg .b32 smem_desc_a_lo_start, smem_desc_b_lo_start;\n\t" + ".reg .b32 smem_desc_a_lo, smem_desc_b_lo;\n\t" + ".reg .b32 smem_desc_a_hi, smem_desc_b_hi;\n\t" + ".reg .b64 smem_desc_a, smem_desc_b;\n\t" + "elect.sync _|leader_thread, -1;\n\t" + f"mov.b32 idesc, {hex(idesc)};\n\t" + # f"mov.b32 tmem_acc, {hex(acc_tmem_addr)};\n\t" + f"mov.b32 tmem_acc, $3;\n\t" + "mov.b32 smem_desc_a_lo_start, $0;\n\t" + "mov.b32 smem_desc_b_lo_start, $1;\n\t" + f"mov.b32 smem_desc_a_hi, {hex(smem_desc_a_hi)};\n\t" + f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t" + f"mov.b64 smem_desc_a, {{smem_desc_a_lo_start, smem_desc_a_hi}};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo_start, smem_desc_b_hi}};\n\t" + "setp.ne.b32 p, $2, 0;\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], smem_desc_a, smem_desc_b, idesc, {pred_str};\n\t" + + "".join( + ( + # f"add.u32 smem_desc_a_lo, smem_desc_a_lo, {hex(offset_a_diff[k - 1])};\n\t" + # f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t" + f"add.u32 smem_desc_a_lo, smem_desc_a_lo_start, {hex(offset_a[k])};\n\t" + f"add.u32 smem_desc_b_lo, smem_desc_b_lo_start, {hex(offset_b[k])};\n\t" + f"mov.b64 smem_desc_a, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], smem_desc_a, smem_desc_b, idesc, 1;\n\t" + ) + for k in range(1, cute.size(tCrA.shape[2])) + ) + + "}\n", + # "r,r,r", + "r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + else: + # For TS gemm, somehow tCrA.iterator.toint() returns 0 no matter what, so we need to + # explicitly pass in the tA_addr for correctness. + tA_addr = tCrA[None, None, 0].iterator.toint() if tA_addr is None else tA_addr + input_args = [ + # Int32(cute.arch.make_warp_uniform(tCrA[None, None, 0].iterator.toint())).ir_value(), + Int32(cute.arch.make_warp_uniform(tA_addr)).ir_value(), + Int32(cute.arch.make_warp_uniform(smem_desc_start_b_lo)).ir_value(), + Int32(not zero_init).ir_value(), + Int32(cute.arch.make_warp_uniform(acc_tmem_addr)).ir_value(), + ] + if const_expr(mbar_ptr is not None): + assert mbar_phase is not None, "mbar_phase must be provided when mbar_ptr is not None" + assert split_arrive is not None, ( + "split_arrive must be provided when mbar_ptr is not None" + ) + split_arrive_idx = split_arrive // op.shape_mnk[2] + input_args.append(mbar_ptr.toint().ir_value()) + input_args.append(Int32(mbar_phase).ir_value()) + mbar_wait_str = ( + ".reg .pred P1; \n\t" + "LAB_WAIT: \n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [$4], $5, 10000000; \n\t" + "@P1 bra DONE; \n\t" + "bra LAB_WAIT; \n\t" + "DONE: \n\t" + ) + else: + mbar_wait_str = "" + llvm.inline_asm( + None, + # [ + # # acc.iterator.toint().ir_value(), + # Int32(tCrA[None, None, 0].iterator.toint()).ir_value(), + # Int32(smem_desc_start_b_lo).ir_value(), + # Int32(not zero_init).ir_value(), + # ], + input_args, + "{\n\t" + ".reg .pred leader_thread;\n\t" + ".reg .pred p;\n\t" + ".reg .b32 idesc;\n\t" + ".reg .b32 tmem_acc;\n\t" + ".reg .b32 tmem_a;\n\t" + ".reg .b32 smem_desc_b_lo_start;\n\t" + ".reg .b32 smem_desc_b_lo;\n\t" + ".reg .b32 smem_desc_b_hi;\n\t" + ".reg .b64 smem_desc_b;\n\t" + "elect.sync _|leader_thread, -1;\n\t" + f"mov.b32 idesc, {hex(idesc)};\n\t" + # f"mov.b32 tmem_acc, {hex(acc_tmem_addr)};\n\t" + f"mov.b32 tmem_acc, $3;\n\t" + f"mov.b32 tmem_a, $0;\n\t" + f"mov.b32 smem_desc_b_lo_start, $1;\n\t" + f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo_start, smem_desc_b_hi}};\n\t" + "setp.ne.b32 p, $2, 0;\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], [tmem_a], smem_desc_b, idesc, {pred_str};\n\t" + + "".join( + ( + # f"add.u32 tmem_a, tmem_a, {hex(offset_a_diff[k - 1])};\n\t" + # f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t" + f"add.u32 smem_desc_b_lo, smem_desc_b_lo_start, {hex(offset_b[k])};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + # f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [tmem_acc], [tmem_a], smem_desc_b, idesc, 1;\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t" + ) + for k in range( + 1, + cute.size(tCrA.shape[2]) if const_expr(mbar_ptr is None) else split_arrive_idx, + ) + ) + + mbar_wait_str + + ( + "".join( + ( + f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t" + ) + for k in range(split_arrive_idx, cute.size(tCrA.shape[2])) + ) + if const_expr(mbar_ptr is not None) + else "" + ) + + "}\n", + "r,r,r,r" if const_expr(mbar_ptr is None) else "r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@cute.jit +def gemm_ptx_partial1( + op: cute.nvgpu.tcgen05.mma.MmaOp, + acc_tmem_addr: cutlass.Constexpr[int], + tCrA: cute.Tensor, + tCrB: cute.Tensor, + sA_base_addr_for_desc: Int32, + sA_addr_offset_for_desc: cutlass.Constexpr[int], + sA_stage: Int32, + sB_base_addr_for_desc: Int32, + sB_addr_offset_for_desc: cutlass.Constexpr[int], + sB_stage: Int32, + sA_layout: Optional[cute.Layout], + sB_layout: Optional[cute.Layout], + sA_swizzle: Optional[cute.Swizzle], + sB_swizzle: cute.Swizzle, + zero_init: bool | Boolean = False, +) -> None: + is_ts = op.a_src == cute.nvgpu.tcgen05.OperandSource.TMEM + if const_expr(not is_ts): + assert sA_layout is not None, "sA_layout must be provided when a_src is not TMEM" + assert sA_swizzle is not None, "sA_swizzle must be provided when a_src is not TMEM" + idesc: int = const_expr(sm100_desc.mma_op_to_idesc(op)) + kind = _tcgen05_mma_kind(op) + if const_expr(not is_ts): + smem_desc_base_a: int = const_expr( + sm100_desc.make_smem_desc_base( + cute.recast_layout(128, op.a_dtype.width, sA_layout[0]), + sA_swizzle, + sm100_desc.Major.K + if const_expr(op.a_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K) + else sm100_desc.Major.MN, + ) + ) + smem_desc_base_a_lo, smem_desc_a_hi = i64_to_i32x2(smem_desc_base_a) + smem_desc_base_a_lo = const_expr(smem_desc_base_a_lo) + smem_desc_a_hi = const_expr(smem_desc_a_hi) + else: + smem_desc_base_a = None + smem_desc_base_a_lo, smem_desc_a_hi = None, None + smem_desc_base_b: int = const_expr( + sm100_desc.make_smem_desc_base( + cute.recast_layout(128, op.b_dtype.width, sB_layout[0]), + sB_swizzle, + sm100_desc.Major.K + if const_expr(op.b_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K) + else sm100_desc.Major.MN, + ) + ) + smem_desc_base_b_lo, smem_desc_b_hi = i64_to_i32x2(smem_desc_base_b) + smem_desc_base_b_lo = const_expr(smem_desc_base_b_lo) + smem_desc_b_hi = const_expr(smem_desc_b_hi) + mask = [Int32(0)] * 4 + + if const_expr(not is_ts): + offset_a = [ + (cute.crd2idx((0, 0, k), sA_layout) * op.a_dtype.width // 8) >> 4 + for k in range(cute.size(tCrA.shape[2])) + ] + else: + offset_a = [ + cute.crd2idx((0, 0, k), sA_layout) * op.a_dtype.width // 32 + for k in range(cute.size(tCrA.shape[2])) + ] + offset_a_diff = [offset_a[k] - offset_a[k - 1] for k in range(1, cute.size(tCrA.shape[2]))] + offset_b = [ + (cute.crd2idx((0, 0, k), sB_layout) * op.b_dtype.width // 8) >> 4 + for k in range(cute.size(tCrB.shape[2])) + ] + offset_b_diff = [offset_b[k] - offset_b[k - 1] for k in range(1, cute.size(tCrB.shape[2]))] + + if const_expr(not is_ts): + # smem_desc_start_a_lo = Int32(smem_desc_base_a_lo | sm100_desc.make_smem_desc_start_addr(sA[None, None, 0].iterator)) + smem_desc_start_a_lo = const_expr(smem_desc_base_a_lo) + else: + smem_desc_start_a_lo = None + # smem_desc_start_b_lo = Int32(smem_desc_base_b_lo | sm100_desc.make_smem_desc_start_addr(sB[None, None, 0].iterator)) + smem_desc_start_b_lo = const_expr(smem_desc_base_b_lo) + pred_str = "p" if isinstance(zero_init, Boolean) else "0" if zero_init else "1" + if const_expr(not is_ts): + llvm.inline_asm( + None, + [ + # acc.iterator.toint().ir_value(), + # Int32(cute.arch.make_warp_uniform(smem_desc_start_a_lo)).ir_value(), + Int32(sA_base_addr_for_desc).ir_value(), + Int32(sA_stage).ir_value(), + # Int32(cute.arch.make_warp_uniform(smem_desc_start_b_lo)).ir_value(), + Int32(sB_base_addr_for_desc).ir_value(), + Int32(sB_stage).ir_value(), + Int32(not zero_init).ir_value(), + mask[0].ir_value(), + mask[1].ir_value(), + mask[2].ir_value(), + mask[3].ir_value(), + ], + "{\n\t" + ".reg .pred leader_thread;\n\t" + ".reg .pred p;\n\t" + ".reg .b32 idesc;\n\t" + ".reg .b32 tmem_acc;\n\t" + ".reg .b32 smem_desc_a_lo, smem_desc_b_lo;\n\t" + ".reg .b32 smem_desc_a_hi, smem_desc_b_hi;\n\t" + ".reg .b64 smem_desc_a, smem_desc_b;\n\t" + "elect.sync _|leader_thread, -1;\n\t" + f"mov.b32 idesc, {hex(idesc)};\n\t" + f"mov.b32 tmem_acc, {hex(acc_tmem_addr)};\n\t" + # "mov.b32 smem_desc_a_lo, $0;\n\t" + # f"add.u32 smem_desc_a_lo, $0, {hex(smem_desc_start_a_lo)};\n\t" + f"mad.lo.u32 smem_desc_a_lo, $1, {hex(sA_addr_offset_for_desc)}, $0;\n\t" + # "mov.b32 smem_desc_b_lo, $2;\n\t" + f"mad.lo.u32 smem_desc_b_lo, $3, {hex(sB_addr_offset_for_desc)}, $2;\n\t" + f"mov.b32 smem_desc_a_hi, {hex(smem_desc_a_hi)};\n\t" + f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t" + f"mov.b64 smem_desc_a, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + "setp.ne.b32 p, $4, 0;\n\t" + f"@leader_thread tcgen05.mma.cta_group::1.kind::{kind} [tmem_acc], smem_desc_a, smem_desc_b, idesc, {{$5, $6, $7, $8}}, {pred_str};\n\t" + + "".join( + ( + f"add.u32 smem_desc_a_lo, smem_desc_a_lo, {hex(offset_a_diff[k - 1])};\n\t" + f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t" + f"mov.b64 smem_desc_a, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + f"@leader_thread tcgen05.mma.cta_group::1.kind::{kind} [tmem_acc], smem_desc_a, smem_desc_b, idesc, {{$5, $6, $7, $8}}, 1;\n\t" + ) + for k in range(1, cute.size(tCrA.shape[2])) + ) + + "}\n", + "r,r,r,r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + else: + llvm.inline_asm( + None, + [ + # acc.iterator.toint().ir_value(), + Int32(tCrA[None, None, 0].iterator.toint()).ir_value(), + Int32(smem_desc_start_b_lo).ir_value(), + Int32(not zero_init).ir_value(), + mask[0].ir_value(), + mask[1].ir_value(), + mask[2].ir_value(), + mask[3].ir_value(), + ], + "{\n\t" + ".reg .pred leader_thread;\n\t" + ".reg .pred p;\n\t" + ".reg .b32 idesc;\n\t" + ".reg .b32 tmem_a;\n\t" + ".reg .b32 smem_desc_b_lo;\n\t" + ".reg .b32 smem_desc_b_hi;\n\t" + ".reg .b64 smem_desc_b;\n\t" + "elect.sync _|leader_thread, -1;\n\t" + f"mov.b32 idesc, {hex(idesc)};\n\t" + f"mov.b32 tmem_a, $1;\n\t" + f"mov.b32 smem_desc_b_lo, $2;\n\t" + f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + "setp.ne.b32 p, $3, 0;\n\t" + f"@leader_thread tcgen05.mma.cta_group::1.kind::{kind} [$0], [tmem_a], smem_desc_b, idesc, {{$4, $5, $6, $7}}, {pred_str};\n\t" + + "".join( + ( + f"add.u32 tmem_a, tmem_a, {hex(offset_a_diff[k - 1])};\n\t" + f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + f"@leader_thread tcgen05.mma.cta_group::1.kind::{kind} [$0], [tmem_a], smem_desc_b, idesc, {{$4, $5, $6, $7}}, 1;\n\t" + ) + for k in range(1, cute.size(tCrA.shape[2])) + ) + + "}\n", + "r,r,r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@cute.jit +def gemm_ptx_precomputed( + acc_tmem_addr: Int32, + smem_desc_start_a: Int32, # If TS, then this is the tmem start address for A + smem_desc_start_b: Int32, + idesc: int, + smem_desc_base_a: Optional[int], + smem_desc_base_b: int, + tCrA_layout: cute.Layout, + tCrB_layout: cute.Layout, + mbar_ptr: Optional[cutlass.Pointer] = None, + mbar_phase: Optional[Int32] = None, + zero_init: bool | Boolean = False, + cta_group: int = 1, + kind: str = "f16", +) -> None: + # acc_tmem_addr += acc_offset + is_ts = const_expr(smem_desc_base_a is None) + num_k_tile = cute.size(tCrA_layout.shape[2]) + if const_expr(not is_ts): + smem_desc_base_a_lo, smem_desc_a_hi = i64_to_i32x2(smem_desc_base_a) + else: + smem_desc_base_a_lo, smem_desc_a_hi = None, None + smem_desc_base_b_lo, smem_desc_b_hi = i64_to_i32x2(smem_desc_base_b) + + tCrA_layout = ( + tCrA_layout + if const_expr(not is_ts) + # else cute.recast_layout(32, tCrA.element_type.width, tCrA_layout) + # currently hard-coding the width to 16 + else cute.recast_layout(32, 16, tCrA_layout) + ) + offset_a = [cute.crd2idx((0, 0, k), tCrA_layout) for k in range(num_k_tile)] + offset_a_diff = [offset_a[k] - offset_a[k - 1] for k in range(1, num_k_tile)] + offset_b = [cute.crd2idx((0, 0, k), tCrB_layout) for k in range(num_k_tile)] + offset_b_diff = [offset_b[k] - offset_b[k - 1] for k in range(1, num_k_tile)] + + smem_desc_start_a_lo = None + if const_expr(not is_ts): + smem_desc_start_a_lo = Int32(smem_desc_base_a_lo | smem_desc_start_a) + # smem_desc_start_a_lo = smem_desc_start_a + smem_desc_start_b_lo = Int32(smem_desc_base_b_lo | smem_desc_start_b) + pred_str = "p" if isinstance(zero_init, Boolean) else "0" if zero_init else "1" + if const_expr(not is_ts): + assert mbar_ptr is None, "mbar_ptr must be None when a_src is not TMEM" + llvm.inline_asm( + None, + [ + # acc.iterator.toint().ir_value(), + Int32(cute.arch.make_warp_uniform(smem_desc_start_a_lo)).ir_value(), + Int32(cute.arch.make_warp_uniform(smem_desc_start_b_lo)).ir_value(), + Int32(not zero_init).ir_value(), + Int32(cute.arch.make_warp_uniform(acc_tmem_addr)).ir_value(), + ], + "{\n\t" + ".reg .pred leader_thread;\n\t" + ".reg .pred p;\n\t" + ".reg .b32 idesc;\n\t" + ".reg .b32 tmem_acc;\n\t" + ".reg .b32 smem_desc_a_lo_start, smem_desc_b_lo_start;\n\t" + ".reg .b32 smem_desc_a_lo, smem_desc_b_lo;\n\t" + ".reg .b32 smem_desc_a_hi, smem_desc_b_hi;\n\t" + ".reg .b64 smem_desc_a, smem_desc_b;\n\t" + "elect.sync _|leader_thread, -1;\n\t" + f"mov.b32 idesc, {hex(idesc)};\n\t" + # f"mov.b32 tmem_acc, {hex(acc_tmem_addr)};\n\t" + f"mov.b32 tmem_acc, $3;\n\t" + "mov.b32 smem_desc_a_lo_start, $0;\n\t" + "mov.b32 smem_desc_b_lo_start, $1;\n\t" + f"mov.b32 smem_desc_a_hi, {hex(smem_desc_a_hi)};\n\t" + f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t" + f"mov.b64 smem_desc_a, {{smem_desc_a_lo_start, smem_desc_a_hi}};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo_start, smem_desc_b_hi}};\n\t" + "setp.ne.b32 p, $2, 0;\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], smem_desc_a, smem_desc_b, idesc, {pred_str};\n\t" + + "".join( + ( + # f"add.u32 smem_desc_a_lo, smem_desc_a_lo, {hex(offset_a_diff[k - 1])};\n\t" + # f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t" + f"add.s32 smem_desc_a_lo, smem_desc_a_lo_start, {hex(offset_a[k])};\n\t" + f"add.s32 smem_desc_b_lo, smem_desc_b_lo_start, {hex(offset_b[k])};\n\t" + f"mov.b64 smem_desc_a, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], smem_desc_a, smem_desc_b, idesc, 1;\n\t" + ) + for k in range(1, num_k_tile) + ) + + "}\n", + # "r,r,r", + "r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + else: + input_args = [ + Int32(cute.arch.make_warp_uniform(smem_desc_start_a)).ir_value(), + Int32(cute.arch.make_warp_uniform(smem_desc_start_b_lo)).ir_value(), + Int32(not zero_init).ir_value(), + Int32(cute.arch.make_warp_uniform(acc_tmem_addr)).ir_value(), + ] + if const_expr(mbar_ptr is not None): + assert mbar_phase is not None, "mbar_phase must be provided when mbar_ptr is not None" + input_args.append(mbar_ptr.toint().ir_value()) + input_args.append(Int32(mbar_phase).ir_value()) + mbar_wait_str = ( + ".reg .pred P1; \n\t" + "LAB_WAIT: \n\t" + "mbarrier.try_wait.parity.shared::cta.b64 P1, [$4], $5, 10000000; \n\t" + "@P1 bra DONE; \n\t" + "bra LAB_WAIT; \n\t" + "DONE: \n\t" + ) + else: + mbar_wait_str = "" + llvm.inline_asm( + None, + # [ + # # acc.iterator.toint().ir_value(), + # Int32(tCrA_layout[None, None, 0].iterator.toint()).ir_value(), + # Int32(smem_desc_start_b_lo).ir_value(), + # Int32(not zero_init).ir_value(), + # ], + input_args, + "{\n\t" + ".reg .pred leader_thread;\n\t" + ".reg .pred p;\n\t" + ".reg .b32 idesc;\n\t" + ".reg .b32 tmem_acc;\n\t" + ".reg .b32 tmem_a;\n\t" + ".reg .b32 smem_desc_b_lo_start;\n\t" + ".reg .b32 smem_desc_b_lo;\n\t" + ".reg .b32 smem_desc_b_hi;\n\t" + ".reg .b64 smem_desc_b;\n\t" + "elect.sync _|leader_thread, -1;\n\t" + f"mov.b32 idesc, {hex(idesc)};\n\t" + # f"mov.b32 tmem_acc, {hex(acc_tmem_addr)};\n\t" + f"mov.b32 tmem_acc, $3;\n\t" + f"mov.b32 tmem_a, $0;\n\t" + f"mov.b32 smem_desc_b_lo_start, $1;\n\t" + f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo_start, smem_desc_b_hi}};\n\t" + "setp.ne.b32 p, $2, 0;\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], [tmem_a], smem_desc_b, idesc, {pred_str};\n\t" + + "".join( + ( + # f"add.u32 tmem_a, tmem_a, {hex(offset_a_diff[k - 1])};\n\t" + # f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t" + f"add.u32 smem_desc_b_lo, smem_desc_b_lo_start, {hex(offset_b[k])};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + # f"@leader_thread tcgen05.mma.cta_group::1.kind::f16 [tmem_acc], [tmem_a], smem_desc_b, idesc, 1;\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t" + ) + for k in range( + 1, + num_k_tile if const_expr(mbar_ptr is None) else num_k_tile // 4 * 3, + ) + ) + + mbar_wait_str + + ( + "".join( + ( + # f"add.u32 smem_desc_b_lo, smem_desc_b_lo, {hex(offset_b_diff[k - 1])};\n\t" + f"add.u32 smem_desc_b_lo, smem_desc_b_lo_start, {hex(offset_b[k])};\n\t" + f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], [tmem_a + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t" + ) + for k in range(num_k_tile // 4 * 3, num_k_tile) + ) + if const_expr(mbar_ptr is not None) + else "" + ) + + "}\n", + "r,r,r,r" if const_expr(mbar_ptr is None) else "r,r,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@cute.jit +def declare_ptx_smem_desc( + smem_desc_start_a: Int32, # If TS, then this is the tmem start address for A + smem_desc_base_a: Optional[int], + tCrA_layout: cute.Layout, + var_name_prefix: str = "smem_desc", +) -> None: + is_ts = const_expr(smem_desc_base_a is None) + num_k_tile = cute.size(tCrA_layout.shape[2]) + smem_desc_base_a_lo, smem_desc_a_hi = None, None + if const_expr(not is_ts): + smem_desc_base_a_lo, smem_desc_a_hi = i64_to_i32x2(smem_desc_base_a) + tCrA_layout = ( + tCrA_layout + if const_expr(not is_ts) + # else cute.recast_layout(32, tCrA.element_type.width, tCrA_layout) + # currently hard-coding the width to 16 + else cute.recast_layout(32, 16, tCrA_layout) + ) + offset_a = [cute.crd2idx((0, 0, k), tCrA_layout) for k in range(num_k_tile)] + smem_desc_start_a_lo = None + if const_expr(not is_ts): + smem_desc_start_a_lo = Int32(smem_desc_base_a_lo | smem_desc_start_a) + if const_expr(not is_ts): + llvm.inline_asm( + None, + [Int32(cute.arch.make_warp_uniform(smem_desc_start_a_lo)).ir_value()], + f".reg .b32 {var_name_prefix}_lo;\n\t" + f".reg .b64 {var_name_prefix}_<{num_k_tile}>;\n\t" + f"mov.b64 {var_name_prefix}_0, {{$0, {hex(smem_desc_a_hi)}}};\n\t" + + "".join( + ( + f"add.s32 {var_name_prefix}_lo, $0, {hex(offset_a[k])};\n\t" + f"mov.b64 {var_name_prefix}_{k}, {{{var_name_prefix}_lo, {hex(smem_desc_a_hi)}}};\n\t" + ) + for k in range(1, num_k_tile) + ), + "r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@cute.jit +def declare_ptx_idesc(op: cute.nvgpu.tcgen05.mma.MmaOp, var_name: str = "idesc") -> None: + idesc = const_expr(sm100_desc.mma_op_to_idesc(op)) + llvm.inline_asm( + None, + [], + f".reg .b32 {var_name};\n\t" # noqa + f"mov.b32 {var_name}, {hex(idesc)};\n\t", + constraints="", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@cute.jit +def gemm_ptx_precomputed_varname( + acc_tmem_addr: Int32, + smem_desc_start_b: Int32, + # idesc: int, + smem_desc_base_b: int, + tCrB_layout: cute.Layout, + smem_var_name_prefix: str, + idesc_var_name: str, + smem_offset: int, + zero_init: bool | Boolean = False, + cta_group: int = 1, + kind: str = "f16", +) -> None: + is_ts = False + num_k_tile = cute.size(tCrB_layout.shape[2]) + smem_desc_base_b_lo, smem_desc_b_hi = i64_to_i32x2(smem_desc_base_b) + offset_b = [cute.crd2idx((0, 0, k), tCrB_layout) for k in range(num_k_tile)] + + smem_desc_start_b_lo = Int32(smem_desc_base_b_lo | smem_desc_start_b) + pred_str = "p" if isinstance(zero_init, Boolean) else "0" if zero_init else "1" + if const_expr(not is_ts): + llvm.inline_asm( + None, + [ + Int32(cute.arch.make_warp_uniform(smem_desc_start_b_lo)).ir_value(), + Int32(not zero_init).ir_value(), + Int32(cute.arch.make_warp_uniform(acc_tmem_addr)).ir_value(), + ], + "{\n\t" + ".reg .pred leader_thread;\n\t" + ".reg .pred p;\n\t" + # ".reg .b32 idesc;\n\t" + ".reg .b32 tmem_acc;\n\t" + ".reg .b32 smem_desc_b_lo_start;\n\t" + ".reg .b32 smem_desc_a_lo, smem_desc_b_lo;\n\t" + ".reg .b32 smem_desc_a_hi, smem_desc_b_hi;\n\t" + # ".reg .b64 smem_desc_b;\n\t" + f".reg .b64 smem_desc_b_<{num_k_tile}>;\n\t" + "elect.sync _|leader_thread, -1;\n\t" + # f"mov.b32 idesc, {hex(idesc)};\n\t" + # f"mov.b32 tmem_acc, {hex(acc_tmem_addr)};\n\t" + f"mov.b32 tmem_acc, $2;\n\t" + "mov.b32 smem_desc_b_lo_start, $0;\n\t" + f"mov.b32 smem_desc_b_hi, {hex(smem_desc_b_hi)};\n\t" + f"mov.b64 {{smem_desc_a_lo, smem_desc_a_hi}}, {smem_var_name_prefix}_0;\n\t" + f"add.s32 smem_desc_a_lo, smem_desc_a_lo, {smem_offset};\n\t" + f"mov.b64 {smem_var_name_prefix}_0, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t" + f"mov.b64 smem_desc_b_0, {{smem_desc_b_lo_start, smem_desc_b_hi}};\n\t" + + "".join( + ( + f"mov.b64 {{smem_desc_a_lo, smem_desc_a_hi}}, {smem_var_name_prefix}_{k};\n\t" + f"add.s32 smem_desc_a_lo, smem_desc_a_lo, {smem_offset};\n\t" + f"add.s32 smem_desc_b_lo, smem_desc_b_lo_start, {hex(offset_b[k])};\n\t" + f"mov.b64 {smem_var_name_prefix}_{k}, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t" + f"mov.b64 smem_desc_b_{k}, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + ) + for k in range(1, num_k_tile) + ) + + "setp.ne.b32 p, $1, 0;\n\t" + # f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_0, smem_desc_b, idesc, {pred_str};\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], {smem_var_name_prefix}_0, smem_desc_b_0, {idesc_var_name}, {pred_str};\n\t" + + "".join( + ( + # f"mov.b64 {{smem_desc_a_lo, smem_desc_a_hi}}, {smem_var_name_prefix}_{k};\n\t" + # f"add.s32 smem_desc_a_lo, smem_desc_a_lo, {smem_offset};\n\t" + # f"add.s32 smem_desc_b_lo, smem_desc_b_lo_start, {hex(offset_b[k])};\n\t" + # f"mov.b64 {smem_var_name_prefix}_{k}, {{smem_desc_a_lo, smem_desc_a_hi}};\n\t" + # f"mov.b64 smem_desc_b, {{smem_desc_b_lo, smem_desc_b_hi}};\n\t" + # f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_{k}, smem_desc_b, idesc, 1;\n\t" + # f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::f16 [tmem_acc], {smem_var_name_prefix}_{k}, smem_desc_b, {idesc_var_name}, 1;\n\t" + f"@leader_thread tcgen05.mma.cta_group::{cta_group}.kind::{kind} [tmem_acc], {smem_var_name_prefix}_{k}, smem_desc_b_{k}, {idesc_var_name}, 1;\n\t" + ) + for k in range(1, num_k_tile) + ) + + "}\n", + "r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) diff --git a/fastvideo-kernel/fa4/block_info.py b/fastvideo-kernel/fa4/block_info.py new file mode 100644 index 000000000..35bb4365f --- /dev/null +++ b/fastvideo-kernel/fa4/block_info.py @@ -0,0 +1,156 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +from typing import Tuple, Optional +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, const_expr + +from flash_attn.cute.seqlen_info import SeqlenInfoQK, SeqlenInfoQKNewK + + +@dataclass(frozen=True) +class BlockInfo: + tile_m: cutlass.Constexpr[int] + tile_n: cutlass.Constexpr[int] + is_causal: cutlass.Constexpr[bool] + is_local: cutlass.Constexpr[bool] = False + is_split_kv: cutlass.Constexpr[bool] = False + window_size_left: Optional[Int32] = None + window_size_right: Optional[Int32] = None + qhead_per_kvhead_packgqa: cutlass.Constexpr[int] = 1 + + @cute.jit + def get_n_block_min_max( + self, + seqlen_info: SeqlenInfoQK, + m_block: Int32, + split_idx: Int32 = 0, + num_splits: Int32 = 1, + ) -> Tuple[Int32, Int32]: + n_block_max = cute.ceil_div(seqlen_info.seqlen_k, self.tile_n) + if const_expr(self.is_causal or (self.is_local and self.window_size_right is not None)): + m_idx_max = (m_block + 1) * self.tile_m + if const_expr(self.qhead_per_kvhead_packgqa > 1): + m_idx_max = cute.ceil_div(m_idx_max, self.qhead_per_kvhead_packgqa) + n_idx = m_idx_max + seqlen_info.seqlen_k - seqlen_info.seqlen_q + n_idx_right = n_idx if const_expr(self.is_causal) else n_idx + self.window_size_right + n_block_max = min(n_block_max, cute.ceil_div(n_idx_right, self.tile_n)) + n_block_min = 0 + if const_expr(self.is_local and self.window_size_left is not None): + m_idx_min = m_block * self.tile_m + if const_expr(self.qhead_per_kvhead_packgqa > 1): + m_idx_min = m_idx_min // self.qhead_per_kvhead_packgqa + n_idx = m_idx_min + seqlen_info.seqlen_k - seqlen_info.seqlen_q + n_idx_left = n_idx - self.window_size_left + n_block_min = cutlass.max(n_idx_left // self.tile_n, 0) + if cutlass.const_expr(self.is_split_kv): + num_n_blocks_per_split = ( + Int32(0) + if n_block_max <= n_block_min + else (n_block_max - n_block_min + num_splits - 1) // num_splits + ) + n_block_min = n_block_min + split_idx * num_n_blocks_per_split + n_block_max = cutlass.min(n_block_min + num_n_blocks_per_split, n_block_max) + return n_block_min, n_block_max + + @cute.jit + def get_m_block_min_max(self, seqlen_info: SeqlenInfoQK, n_block: Int32) -> Tuple[Int32, Int32]: + m_block_max = cute.ceil_div(seqlen_info.seqlen_q, self.tile_m) + m_block_min = 0 + if const_expr(self.is_causal or (self.is_local and self.window_size_right is not None)): + n_idx_min = n_block * self.tile_n + m_idx = n_idx_min + seqlen_info.seqlen_q - seqlen_info.seqlen_k + m_idx_right = m_idx if const_expr(self.is_causal) else m_idx - self.window_size_right + m_block_min = max(m_block_min, m_idx_right // self.tile_m) + if const_expr(self.is_local and self.window_size_left is not None): + n_idx_max = (n_block + 1) * self.tile_n + m_idx = n_idx_max + seqlen_info.seqlen_q - seqlen_info.seqlen_k + m_idx_left = m_idx + self.window_size_left + m_block_max = min(m_block_max, cute.ceil_div(m_idx_left, self.tile_m)) + return m_block_min, m_block_max + + @cute.jit + def get_n_block_k_new_min_max( + self, + seqlen_info: SeqlenInfoQKNewK, + m_block: Int32, + split_idx: Int32 = 0, + num_splits: Int32 = 1, + ) -> Tuple[Int32, Int32]: + """Get the block range for new K tokens (append KV). + + First computes the full n_block range via get_n_block_min_max, then maps + those blocks into the new-K index space by subtracting seqlen_k_og. + """ + n_block_min, n_block_max = self.get_n_block_min_max( + seqlen_info, + m_block, + split_idx, + num_splits, + ) + idx_k_new_min = cutlass.max(n_block_min * self.tile_n - seqlen_info.seqlen_k_og, 0) + idx_k_new_max = cutlass.min( + n_block_max * self.tile_n - seqlen_info.seqlen_k_og, seqlen_info.seqlen_k_new + ) + n_block_new_min = idx_k_new_min // self.tile_n + n_block_new_max = ( + cute.ceil_div(idx_k_new_max, self.tile_n) + if idx_k_new_max > idx_k_new_min + else n_block_new_min + ) + return n_block_new_min, n_block_new_max + + @cute.jit + def get_n_block_min_causal_local_mask( + self, + seqlen_info: SeqlenInfoQK, + m_block: Int32, + n_block_min: Int32, + ) -> Int32: + """If we have separate iterations with causal or local masking at the start, where do we stop""" + m_idx_min = m_block * self.tile_m + if const_expr(self.qhead_per_kvhead_packgqa > 1): + m_idx_min = m_idx_min // self.qhead_per_kvhead_packgqa + n_idx = m_idx_min + seqlen_info.seqlen_k - seqlen_info.seqlen_q + n_idx_right = ( + n_idx + if const_expr(not self.is_local or self.window_size_right is None) + else n_idx + self.window_size_right + ) + return cutlass.max(n_block_min, n_idx_right // self.tile_n) + + @cute.jit + def get_n_block_min_before_local_mask( + self, + seqlen_info: SeqlenInfoQK, + m_block: Int32, + n_block_min: Int32, + ) -> Int32: + """If we have separate iterations with local masking at the end, where do we stop the non-masked iterations""" + if const_expr(not self.is_local or self.window_size_left is None): + return n_block_min + else: + m_idx_max = (m_block + 1) * self.tile_m + if const_expr(self.qhead_per_kvhead_packgqa > 1): + m_idx_max = cute.ceil_div(m_idx_max, self.qhead_per_kvhead_packgqa) + n_idx = m_idx_max + seqlen_info.seqlen_k - seqlen_info.seqlen_q + n_idx_left = n_idx - self.window_size_left + return cutlass.max(n_block_min, cute.ceil_div(n_idx_left, self.tile_n)) + + @cute.jit + def get_n_block_max_for_m_block( + self, + seqlen_info: SeqlenInfoQK, + m_block: Int32, + ) -> Int32: + n_block_max = cute.ceil_div(seqlen_info.seqlen_k, self.tile_n) + if const_expr(self.is_causal or self.window_size_right is not None): + m_idx_max = (m_block + 1) * self.tile_m + if const_expr(self.qhead_per_kvhead_packgqa > 1): + m_idx_max = cute.ceil_div(m_idx_max, self.qhead_per_kvhead_packgqa) + n_idx_right = m_idx_max + seqlen_info.seqlen_k - seqlen_info.seqlen_q + if const_expr(self.window_size_right is not None): + n_idx_right += self.window_size_right + n_block_max = min(n_block_max, cute.ceil_div(n_idx_right, self.tile_n)) + return n_block_max diff --git a/fastvideo-kernel/fa4/block_sparse_utils.py b/fastvideo-kernel/fa4/block_sparse_utils.py new file mode 100644 index 000000000..dd254fc6b --- /dev/null +++ b/fastvideo-kernel/fa4/block_sparse_utils.py @@ -0,0 +1,2035 @@ +""" +Block-sparse runtime utilities for CUTE DSL kernels. + +This module contains runtime execution functions for block-sparse attention kernels. +These utilities are used by CUTE DSL kernels to produce and consume block-sparse loads. +""" + +from typing import Callable, Optional, Tuple +from functools import partial +import math +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, const_expr + +from quack import copy_utils + +# Import data structures from block_sparsity +from flash_attn.cute.block_sparsity import BlockSparseTensors +from flash_attn.cute.named_barrier import NamedBarrierBwd +from flash_attn.cute.seqlen_info import SeqlenInfoQK +from flash_attn.cute.utils import AuxData + + +@cute.jit +def _get_curr_blocksparse_tensors_varlen( + head_idx: cutlass.Int32, + m_block: cutlass.Int32, + blocksparse_tensors: BlockSparseTensors, + seqlen_info: SeqlenInfoQK, +) -> Tuple[cutlass.Int32, cute.Tensor, cutlass.Int32, Optional[cute.Tensor]]: + """Varlen path: tensors are 2D [nheads, total_m_blocks] / [nheads, total_n_blocks].""" + mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx, *_ = blocksparse_tensors + curr_m_block = seqlen_info.m_block_offset + m_block + curr_block_idx_offset = seqlen_info.block_idx_offset + m_block * seqlen_info.num_n_blocks + curr_mask_block_cnt = mask_block_cnt[head_idx, curr_m_block] + curr_mask_block_idx = cute.domain_offset(curr_block_idx_offset, mask_block_idx[head_idx, None]) + if const_expr(full_block_cnt is not None): + curr_full_block_cnt = full_block_cnt[head_idx, curr_m_block] + curr_full_block_idx = cute.domain_offset( + curr_block_idx_offset, full_block_idx[head_idx, None] + ) + else: + curr_full_block_cnt = Int32(0) + curr_full_block_idx = None + return (curr_mask_block_cnt, curr_mask_block_idx, curr_full_block_cnt, curr_full_block_idx) + + +@cute.jit +def _get_curr_blocksparse_tensors( + batch_idx: cutlass.Int32, + head_idx: cutlass.Int32, + m_block: cutlass.Int32, + blocksparse_tensors: BlockSparseTensors, +) -> Tuple[cutlass.Int32, cute.Tensor, cutlass.Int32, Optional[cute.Tensor]]: + """Fixed-length path: tensors are 4D [batch, nheads, m_block, n_block].""" + mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx, *_ = blocksparse_tensors + curr_mask_block_cnt = mask_block_cnt[batch_idx, head_idx, m_block] + curr_mask_block_idx = mask_block_idx[batch_idx, head_idx, m_block, None] + if const_expr(full_block_cnt is not None): + curr_full_block_cnt = full_block_cnt[batch_idx, head_idx, m_block] + curr_full_block_idx = full_block_idx[batch_idx, head_idx, m_block, None] + else: + curr_full_block_cnt = Int32(0) + curr_full_block_idx = None + return (curr_mask_block_cnt, curr_mask_block_idx, curr_full_block_cnt, curr_full_block_idx) + + +@cute.jit +def get_curr_blocksparse_tensors( + batch_idx: cutlass.Int32, + head_idx: cutlass.Int32, + m_block: cutlass.Int32, + blocksparse_tensors: BlockSparseTensors, + seqlen_info: SeqlenInfoQK, +) -> Tuple[cutlass.Int32, cute.Tensor, cutlass.Int32, Optional[cute.Tensor]]: + """Extract head, m_block, and batch-local blocksparsity data from blocksparse_tensors""" + if const_expr(len(blocksparse_tensors.mask_block_cnt.shape) == 2): + return _get_curr_blocksparse_tensors_varlen( + head_idx, m_block, blocksparse_tensors, seqlen_info + ) + return _get_curr_blocksparse_tensors(batch_idx, head_idx, m_block, blocksparse_tensors) + + +# NOTE [SM100 block-sparse empty tiles: mbarrier contract] +# +# For block-sparse SM100 forward, a given (m_block, stage) Q tile can have zero active +# KV blocks (total_block_cnt == 0). In that case there is no seqlen_kv iteration, so +# the softmax warp-group has no row stats to publish. +# +# The correction warp-group seeds fully-masked-row stats and runs the usual correction +# epilogue so output/LSE have well-defined values. Both warp-groups must still perform +# the softmax<->correction mbarrier handshake so phases advance correctly across +# empty->empty and empty->non-empty tile sequences. +# +# In the no-sink case, this corresponds to the usual fully-masked-row convention: +# output is zero and LSE is -inf. +# +# Barrier contract (each is `mbar_ptr + + stage`): +# +# Producer/consumer pairs: +# - `mbar_softmax_corr_full` : softmax arrive -> correction wait +# - `mbar_softmax_corr_empty` : correction arrive -> softmax wait +# - `mbar_P_full_O_rescaled` : softmax arrive (+ correction arrive) -> MMA wait +# - `mbar_P_full_2` : softmax arrive -> MMA wait +# - `mbar_corr_epi_full_/empty` : correction <-> epilogue (only when epilogue is separate) +# +# Empty tile (`total_block_cnt == 0`): +# - Softmax: skips the seqlen_kv softmax path entirely (no P stores, no `mbar_P_full_*`). +# It only arrives `mbar_softmax_corr_full` once per stage as a synthetic "no work" signal. +# At the `softmax_loop` level, softmax unconditionally waits `mbar_softmax_corr_empty` +# before each tile (when block-sparse) to drain a prior correction arrival and keep +# phases aligned across non-empty -> empty transitions. +# - Correction: waits `mbar_softmax_corr_full`, seeds stats + runs `correction_epilogue(scale=0)`, +# and arrives `mbar_softmax_corr_empty` (and `mbar_corr_epi_full_/empty` when applicable). +# - No `mbar_P_full_*` barriers are arrived (no P, no MMA O); only the softmax<->correction +# (and correction<->epilogue) handshakes advance phases. +# +# Non-empty tile: +# - Softmax: runs `softmax_step` (produces P) and uses `mbar_softmax_corr_full/empty` to +# publish row_max (during seqlen_kv) and final row stats (once per tile), and to advance phases; +# arrives `mbar_P_full_*` when P is stored. +# - Correction: waits `mbar_softmax_corr_full`, may rescale/release O, arrives `mbar_softmax_corr_empty` +# to ack/advance, and arrives `mbar_P_full_O_rescaled` when MMA can proceed. +# +# Backward (SM100): +# - Empty KV tile: for a given `n_block`, `total_m_block_cnt == 0` means no Q tiles contribute. +# - Both the load and compute loops guard all pipeline work on `process_tile`, so empty tiles +# skip producer/consumer operations entirely (no per-tile mbarrier phase handshake like forward). +# - In the `not dKV_postprocess` path, dK/dV for empty KV tiles are explicitly written as zeros +# even when `process_tile == False` (see `flash_bwd_sm100.py` `should_zero_dKV`). + + +@cute.jit +def load_block_list( + block_indices: cute.Tensor, + block_count, + first_block_preloaded: cutlass.Constexpr, + kv_producer_state, + load_K, + load_V, + pipeline_k, + pipeline_v, + intra_wg_overlap: cutlass.Constexpr, +): + """Iterate over the sparse blocks and load K, V into the pipeline. + For the intra_wg_overlap case, we overlap the loads of K and V. And this + means we need to pipeline the last V load from the partial block case, + with the loads for the full blocks. Set first_block_preloaded when the + caller has already issued the first K load for the list. + + Q is loaded separately on its own mbarrier before this function is called. + + Note: + we iterate along the block_n indices in reverse. + + Returns: + Updated kv_producer_state after processing the block list. + + """ + if block_count > 0: + if const_expr(not intra_wg_overlap): + for offset in cutlass.range(block_count): + n_block = block_indices[block_count - 1 - offset] + pipeline_k.producer_acquire(kv_producer_state) + load_K(src_idx=n_block, producer_state=kv_producer_state) + pipeline_v.producer_acquire(kv_producer_state) + load_V(src_idx=n_block, producer_state=kv_producer_state) + kv_producer_state.advance() + else: + n_block_first = block_indices[block_count - 1] + if const_expr(not first_block_preloaded): + pipeline_k.producer_acquire(kv_producer_state) + load_K(src_idx=n_block_first, producer_state=kv_producer_state) + + for idx in cutlass.range(block_count - 1, unroll=1): + n_block_prev = block_indices[block_count - 1 - idx] + n_block = block_indices[block_count - 2 - idx] + kv_producer_state_prev = kv_producer_state.clone() + kv_producer_state.advance() + pipeline_k.producer_acquire(kv_producer_state) + load_K(src_idx=n_block, producer_state=kv_producer_state) + pipeline_v.producer_acquire(kv_producer_state_prev) + load_V(src_idx=n_block_prev, producer_state=kv_producer_state_prev) + + return kv_producer_state + + +@cute.jit +def finish_overlap_v_load( + block_indices: cute.Tensor, + block_count, + load_V, + pipeline_v, + kv_producer_state, +): + """Load the final V block after overlapped K/V loads.""" + if block_count > 0: + n_block_last = block_indices[0] + pipeline_v.producer_acquire(kv_producer_state) + load_V(src_idx=n_block_last, producer_state=kv_producer_state) + kv_producer_state.advance() + + return kv_producer_state + + +@cute.jit +def sparse_tensor_m_block( + m_block, + qhead_per_kvhead: cutlass.Constexpr[int], + q_subtile_factor: cutlass.Constexpr[int], +): + """Map packed m_block indices to block-sparse tensor indices.""" + block = m_block + if const_expr(qhead_per_kvhead != 1): + block = block // qhead_per_kvhead + if const_expr(q_subtile_factor != 1): + block = block // q_subtile_factor + return block + + +@cute.jit +def produce_block_sparse_loads( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + m_block, + seqlen_info: SeqlenInfoQK, + kv_producer_state, + load_K, + load_V, + pipeline_k, + pipeline_v, + intra_wg_overlap: cutlass.Constexpr, + qhead_per_kvhead: cutlass.Constexpr[int] = 1, + q_subtile_factor: cutlass.Constexpr[int] = 1, +): + """Iterate over the mask and full block lists for a single tile. + + Q is loaded separately on its own mbarrier before this function is called. + + The masked (partial) list may leave the last V load pending when intra-warp-group + overlap is enabled. The first full block must consume that pending V while + issuing its own K load on the next pipeline stage. + + In the intra-wg-overlap path, the last masked block leaves its V copy in flight + while we advance the producer state to start the next full K. Either the full list + overlaps that pending V load, or, if no full blocks exist, we explicitly drain it. + + Args: + qhead_per_kvhead: Pack-GQA factor. When > 1, m_block is in packed space and + must be converted to unpacked for sparse tensor indexing. + """ + m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor) + + ( + curr_mask_block_cnt, + curr_mask_block_idx, + curr_full_block_cnt, + curr_full_block_idx, + ) = get_curr_blocksparse_tensors( + batch_idx, + head_idx, + m_block_sparse, + blocksparse_tensors, + seqlen_info, + ) + + mask_empty = curr_mask_block_cnt == 0 + full_empty = curr_full_block_cnt == 0 + + if mask_empty: + # No masked blocks: the full list owns the initial K load. + kv_producer_state = load_block_list( + curr_full_block_idx, + curr_full_block_cnt, + first_block_preloaded=False, + kv_producer_state=kv_producer_state, + load_K=load_K, + load_V=load_V, + pipeline_k=pipeline_k, + pipeline_v=pipeline_v, + intra_wg_overlap=intra_wg_overlap, + ) + + if const_expr(intra_wg_overlap) and curr_full_block_cnt > 0: + kv_producer_state = finish_overlap_v_load( + curr_full_block_idx, + curr_full_block_cnt, + load_V, + pipeline_v, + kv_producer_state, + ) + else: + # Masked blocks present. When overlap is disabled this fully drains the list. + kv_producer_state = load_block_list( + curr_mask_block_idx, + curr_mask_block_cnt, + first_block_preloaded=False, + kv_producer_state=kv_producer_state, + load_K=load_K, + load_V=load_V, + pipeline_k=pipeline_k, + pipeline_v=pipeline_v, + intra_wg_overlap=intra_wg_overlap, + ) + + if full_empty: + if const_expr(intra_wg_overlap): + kv_producer_state = finish_overlap_v_load( + curr_mask_block_idx, + curr_mask_block_cnt, + load_V, + pipeline_v, + kv_producer_state, + ) + else: + if const_expr(intra_wg_overlap): + # Bridge the masked list to the full list by overlapping the pending masked V + # with the first full K load. + n_block_mask_last = curr_mask_block_idx[0] + n_block_full_first = curr_full_block_idx[curr_full_block_cnt - 1] + kv_producer_state_prev = kv_producer_state.clone() + kv_producer_state.advance() + pipeline_k.producer_acquire(kv_producer_state) + load_K(src_idx=n_block_full_first, producer_state=kv_producer_state) + pipeline_v.producer_acquire(kv_producer_state_prev) + load_V(src_idx=n_block_mask_last, producer_state=kv_producer_state_prev) + + kv_producer_state = load_block_list( + curr_full_block_idx, + curr_full_block_cnt, + first_block_preloaded=True, + kv_producer_state=kv_producer_state, + load_K=load_K, + load_V=load_V, + pipeline_k=pipeline_k, + pipeline_v=pipeline_v, + intra_wg_overlap=intra_wg_overlap, + ) + + kv_producer_state = finish_overlap_v_load( + curr_full_block_idx, + curr_full_block_cnt, + load_V, + pipeline_v, + kv_producer_state, + ) + else: + # Non-overlap path with both lists: run the full list normally. + kv_producer_state = load_block_list( + curr_full_block_idx, + curr_full_block_cnt, + first_block_preloaded=False, + kv_producer_state=kv_producer_state, + load_K=load_K, + load_V=load_V, + pipeline_k=pipeline_k, + pipeline_v=pipeline_v, + intra_wg_overlap=intra_wg_overlap, + ) + + return kv_producer_state + + +@cute.jit +def consume_block_sparse_loads( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + m_block, + seqlen_info, + kv_consumer_state, + mma_pv_fn, + mma_one_n_block, + process_first_half_block, + process_last_half_block, + mask_fn, + score_mod_fn, + O_should_accumulate, + mask_mod, + fastdiv_mods, + intra_wg_overlap: cutlass.Constexpr, + warp_scheduler_barrier_sync: Callable, + warp_scheduler_barrier_arrive: Callable, + qhead_per_kvhead: cutlass.Constexpr[int] = 1, + q_subtile_factor: cutlass.Constexpr[int] = 1, +): + """Consume the mask and full block lists for a single tile on the consumer side. + + Mirrors `produce_block_sparse_loads` so that the consumer pipeline uses + the same sparse tensor indexing. + + Args: + qhead_per_kvhead: Pack-GQA factor. When > 1, m_block is in packed space and + must be converted to unpacked for sparse tensor indexing. + """ + m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor) + + ( + curr_mask_block_cnt, + curr_mask_block_idx, + curr_full_block_cnt, + curr_full_block_idx, + ) = get_curr_blocksparse_tensors( + batch_idx, + head_idx, + m_block_sparse, + blocksparse_tensors, + seqlen_info, + ) + + processed_any = curr_mask_block_cnt + curr_full_block_cnt > 0 + + if const_expr(not intra_wg_overlap): + if curr_mask_block_cnt > 0: + mask_n_block = curr_mask_block_idx[curr_mask_block_cnt - 1] + warp_scheduler_barrier_sync() + kv_consumer_state = mma_one_n_block( + kv_consumer_state, + n_block=mask_n_block, + mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate), + mask_fn=partial( + mask_fn, + mask_mod=mask_mod, + mask_seqlen=True, + fastdiv_mods=fastdiv_mods if cutlass.const_expr(mask_mod is not None) else None, + ), + is_first_n_block=True, + ) + O_should_accumulate = True + for i in cutlass.range(1, curr_mask_block_cnt): + mask_n_block = curr_mask_block_idx[curr_mask_block_cnt - 1 - i] + kv_consumer_state = mma_one_n_block( + kv_consumer_state, + n_block=mask_n_block, + mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate), + mask_fn=partial(mask_fn, mask_mod=mask_mod, mask_seqlen=False), + is_first_n_block=False, + ) + O_should_accumulate = True + if curr_full_block_cnt == 0: + warp_scheduler_barrier_arrive() + + if curr_full_block_cnt > 0: + full_n_block = curr_full_block_idx[curr_full_block_cnt - 1] + if curr_mask_block_cnt == 0: + warp_scheduler_barrier_sync() + kv_consumer_state = mma_one_n_block( + kv_consumer_state, + n_block=full_n_block, + mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate), + mask_fn=partial(mask_fn, mask_seqlen=True), + is_first_n_block=True, + ) + O_should_accumulate = True + for i in cutlass.range(1, curr_full_block_cnt): + full_n_block = curr_full_block_idx[curr_full_block_cnt - 1 - i] + kv_consumer_state = mma_one_n_block( + kv_consumer_state, + n_block=full_n_block, + mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate), + mask_fn=partial(mask_fn, mask_seqlen=False), + is_first_n_block=False, + ) + O_should_accumulate = True + else: + kv_consumer_state = mma_one_n_block( + kv_consumer_state, + n_block=full_n_block, + mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate), + mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=True), + is_first_n_block=False, + ) + O_should_accumulate = True + for i in cutlass.range(1, curr_full_block_cnt): + full_n_block = curr_full_block_idx[curr_full_block_cnt - 1 - i] + kv_consumer_state = mma_one_n_block( + kv_consumer_state, + n_block=full_n_block, + mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate), + mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=False), + is_first_n_block=False, + ) + O_should_accumulate = True + warp_scheduler_barrier_arrive() + else: + if curr_mask_block_cnt > 0: + mask_n_block = curr_mask_block_idx[curr_mask_block_cnt - 1] + kv_consumer_state = process_first_half_block( + n_block=mask_n_block, + seqlen=seqlen_info, + kv_consumer_state=kv_consumer_state, + mask_fn=partial( + mask_fn, + mask_mod=mask_mod, + mask_seqlen=True, + fastdiv_mods=fastdiv_mods if cutlass.const_expr(mask_mod is not None) else None, + ), + score_mod_fn=score_mod_fn, + is_first_block=True, + ) + for i in cutlass.range(1, curr_mask_block_cnt): + mask_n_block = curr_mask_block_idx[curr_mask_block_cnt - 1 - i] + kv_consumer_state = mma_one_n_block( + kv_consumer_state, + n_block=mask_n_block, + seqlen=seqlen_info, + mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate), + mask_fn=partial(mask_fn, mask_mod=mask_mod, mask_seqlen=False), + ) + O_should_accumulate = True + + if curr_full_block_cnt > 0: + full_n_block = curr_full_block_idx[curr_full_block_cnt - 1] + if curr_mask_block_cnt == 0: + kv_consumer_state = process_first_half_block( + n_block=full_n_block, + seqlen=seqlen_info, + kv_consumer_state=kv_consumer_state, + mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=True), + score_mod_fn=score_mod_fn, + is_first_block=True, + ) + else: + kv_consumer_state = mma_one_n_block( + kv_consumer_state, + n_block=full_n_block, + seqlen=seqlen_info, + mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate), + mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=True), + ) + O_should_accumulate = True + for i in cutlass.range(1, curr_full_block_cnt): + full_n_block = curr_full_block_idx[curr_full_block_cnt - 1 - i] + kv_consumer_state = mma_one_n_block( + kv_consumer_state, + n_block=full_n_block, + seqlen=seqlen_info, + mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate), + mask_fn=partial(mask_fn, mask_mod=None, mask_seqlen=False), + ) + O_should_accumulate = True + + if curr_mask_block_cnt + curr_full_block_cnt > 0: + kv_consumer_state = process_last_half_block( + kv_consumer_state=kv_consumer_state, + zero_init=not O_should_accumulate, + ) + O_should_accumulate = True + + return kv_consumer_state, O_should_accumulate, processed_any + + +@cute.jit +def split_block_range(block_count, split_idx: Int32, num_splits: Int32): + """Return the half-open block-list range assigned to one SplitKV partition.""" + blocks_per_split = cute.ceil_div(block_count, num_splits) + block_begin = cutlass.min(split_idx * blocks_per_split, block_count) + block_end = cutlass.min(block_begin + blocks_per_split, block_count) + return block_begin, block_end + + +@cute.jit +def dual_stream_block_count(total_block_count, stream_idx: Int32): + """Number of ordinals assigned to one stream by even/odd list position.""" + return (total_block_count + Int32(1) - stream_idx) // Int32(2) + + +@cute.jit +def _block_from_reverse_concat_ordinal( + ordinal: Int32, + mask_block_indices: cute.Tensor, + mask_begin: Int32, + mask_end: Int32, + full_block_indices: Optional[cute.Tensor], + full_begin: Int32, + full_end: Int32, +): + """Lookup an ordinal in ``reverse(masked) + reverse(full)``. + + This is the single canonical order used by the dual-stream loader, + softmax, and MMA block counts. Stream 0 owns even ordinals and stream 1 + owns odd ordinals. + """ + mask_count = mask_end - mask_begin + n_block = Int32(0) + if ordinal < mask_count: + n_block = mask_block_indices[mask_end - Int32(1) - ordinal] + else: + full_ordinal = ordinal - mask_count + if const_expr(full_block_indices is not None): + n_block = full_block_indices[full_end - Int32(1) - full_ordinal] + return n_block + + +@cute.jit +def get_dual_stream_block_counts( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + m_block, + split_idx: Int32, + num_splits: Int32, + qhead_per_kvhead: cutlass.Constexpr[int], + q_subtile_factor: cutlass.Constexpr[int], + seqlen_info: SeqlenInfoQK, +): + """Return stream-0/stream-1 counts for canonical ordinal parity.""" + total = get_total_block_count( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + split_idx, + num_splits, + qhead_per_kvhead, + q_subtile_factor, + seqlen_info, + ) + return dual_stream_block_count(total, Int32(0)), dual_stream_block_count(total, Int32(1)) + + +@cute.jit +def produce_block_sparse_loads_dual_stream_sm100( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + m_block, + seqlen_info: SeqlenInfoQK, + split_idx: Int32, + num_splits: Int32, + kv_producer_state, + load_Q, + load_K, + load_V, + q_producer_phase: Int32, + qhead_per_kvhead: cutlass.Constexpr[int], + q_subtile_factor: cutlass.Constexpr[int], +): + """Load two sparse streams through one interleaved KV FIFO. + + After the initial K0/K1, each stream pipelines ``V(i), K(i + 1)``. + This matches the consumer's PV/QK alternation and exposes the same load/MMA + overlap as the original single-stream FA4 schedule. + """ + m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor) + ( + curr_mask_block_cnt, + curr_mask_block_idx, + curr_full_block_cnt, + curr_full_block_idx, + ) = get_curr_blocksparse_tensors( + batch_idx, + head_idx, + m_block_sparse, + blocksparse_tensors, + seqlen_info, + ) + mask_begin, mask_end = split_block_range(curr_mask_block_cnt, split_idx, num_splits) + full_begin, full_end = split_block_range(curr_full_block_cnt, split_idx, num_splits) + total_block_count = mask_end - mask_begin + full_end - full_begin + + if total_block_count > Int32(0): + stream0_count = dual_stream_block_count(total_block_count, Int32(0)) + stream1_count = dual_stream_block_count(total_block_count, Int32(1)) + n_block0 = _block_from_reverse_concat_ordinal( + Int32(0), + curr_mask_block_idx, + mask_begin, + mask_end, + curr_full_block_idx, + full_begin, + full_end, + ) + load_K(block=n_block0, producer_state=kv_producer_state, page_idx=None) + load_Q(block=0, stage=0) + q_producer_phase ^= 1 + kv_producer_state.advance() + + if stream1_count > Int32(0): + n_block1 = _block_from_reverse_concat_ordinal( + Int32(1), + curr_mask_block_idx, + mask_begin, + mask_end, + curr_full_block_idx, + full_begin, + full_end, + ) + load_K(block=n_block1, producer_state=kv_producer_state, page_idx=None) + kv_producer_state.advance() + + for stream_idx in cutlass.range(stream0_count, unroll=1): + ordinal0 = stream_idx * Int32(2) + n_block0 = _block_from_reverse_concat_ordinal( + ordinal0, + curr_mask_block_idx, + mask_begin, + mask_end, + curr_full_block_idx, + full_begin, + full_end, + ) + load_V(block=n_block0, producer_state=kv_producer_state, page_idx=None) + kv_producer_state.advance() + + next_stream_idx = stream_idx + Int32(1) + if next_stream_idx < stream0_count: + n_block0_next = _block_from_reverse_concat_ordinal( + next_stream_idx * Int32(2), + curr_mask_block_idx, + mask_begin, + mask_end, + curr_full_block_idx, + full_begin, + full_end, + ) + load_K( + block=n_block0_next, + producer_state=kv_producer_state, + page_idx=None, + ) + kv_producer_state.advance() + + if stream_idx < stream1_count: + n_block1 = _block_from_reverse_concat_ordinal( + ordinal0 + Int32(1), + curr_mask_block_idx, + mask_begin, + mask_end, + curr_full_block_idx, + full_begin, + full_end, + ) + load_V(block=n_block1, producer_state=kv_producer_state, page_idx=None) + kv_producer_state.advance() + + if next_stream_idx < stream1_count: + n_block1_next = _block_from_reverse_concat_ordinal( + next_stream_idx * Int32(2) + Int32(1), + curr_mask_block_idx, + mask_begin, + mask_end, + curr_full_block_idx, + full_begin, + full_end, + ) + load_K( + block=n_block1_next, + producer_state=kv_producer_state, + page_idx=None, + ) + kv_producer_state.advance() + + return kv_producer_state, q_producer_phase + + + +@cute.jit +def produce_block_sparse_loads_dual_stream_spdb_sm100( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + m_block, + seqlen_info: SeqlenInfoQK, + split_idx: Int32, + num_splits: Int32, + kv_producer_state, + load_Q, + load_K, + load_V, + q_producer_phase: Int32, + qhead_per_kvhead: cutlass.Constexpr[int], + q_subtile_factor: cutlass.Constexpr[int], +): + """Load two sparse streams through one interleaved KV FIFO. + + After the initial K0/K1, each stream pipelines ``K(i + 1), V(i)``. + This matches the score-double-buffered consumer, which can issue the next + QK into an alternate score slot before waiting for the current P/PV. + """ + m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor) + ( + curr_mask_block_cnt, + curr_mask_block_idx, + curr_full_block_cnt, + curr_full_block_idx, + ) = get_curr_blocksparse_tensors( + batch_idx, + head_idx, + m_block_sparse, + blocksparse_tensors, + seqlen_info, + ) + mask_begin, mask_end = split_block_range(curr_mask_block_cnt, split_idx, num_splits) + full_begin, full_end = split_block_range(curr_full_block_cnt, split_idx, num_splits) + total_block_count = mask_end - mask_begin + full_end - full_begin + + if total_block_count > Int32(0): + stream0_count = dual_stream_block_count(total_block_count, Int32(0)) + stream1_count = dual_stream_block_count(total_block_count, Int32(1)) + n_block0 = _block_from_reverse_concat_ordinal( + Int32(0), + curr_mask_block_idx, + mask_begin, + mask_end, + curr_full_block_idx, + full_begin, + full_end, + ) + load_K(block=n_block0, producer_state=kv_producer_state, page_idx=None) + load_Q(block=0, stage=0) + q_producer_phase ^= 1 + kv_producer_state.advance() + + if stream1_count > Int32(0): + n_block1 = _block_from_reverse_concat_ordinal( + Int32(1), + curr_mask_block_idx, + mask_begin, + mask_end, + curr_full_block_idx, + full_begin, + full_end, + ) + load_K(block=n_block1, producer_state=kv_producer_state, page_idx=None) + kv_producer_state.advance() + + for stream_idx in cutlass.range(stream0_count, unroll=1): + ordinal0 = stream_idx * Int32(2) + next_stream_idx = stream_idx + Int32(1) + if next_stream_idx < stream0_count: + n_block0_next = _block_from_reverse_concat_ordinal( + next_stream_idx * Int32(2), + curr_mask_block_idx, + mask_begin, + mask_end, + curr_full_block_idx, + full_begin, + full_end, + ) + load_K( + block=n_block0_next, + producer_state=kv_producer_state, + page_idx=None, + ) + kv_producer_state.advance() + + n_block0 = _block_from_reverse_concat_ordinal( + ordinal0, + curr_mask_block_idx, + mask_begin, + mask_end, + curr_full_block_idx, + full_begin, + full_end, + ) + load_V(block=n_block0, producer_state=kv_producer_state, page_idx=None) + kv_producer_state.advance() + + if stream_idx < stream1_count: + if next_stream_idx < stream1_count: + n_block1_next = _block_from_reverse_concat_ordinal( + next_stream_idx * Int32(2) + Int32(1), + curr_mask_block_idx, + mask_begin, + mask_end, + curr_full_block_idx, + full_begin, + full_end, + ) + load_K( + block=n_block1_next, + producer_state=kv_producer_state, + page_idx=None, + ) + kv_producer_state.advance() + + n_block1 = _block_from_reverse_concat_ordinal( + ordinal0 + Int32(1), + curr_mask_block_idx, + mask_begin, + mask_end, + curr_full_block_idx, + full_begin, + full_end, + ) + load_V(block=n_block1, producer_state=kv_producer_state, page_idx=None) + kv_producer_state.advance() + + return kv_producer_state, q_producer_phase + + +@cute.jit +def load_block_list_sm100( + block_indices: cute.Tensor, + block_begin, + block_end, + load_q_with_first: cutlass.Constexpr, + q_stage: cutlass.Constexpr, + kv_producer_state, + load_Q, + load_K, + load_V, + pipeline_kv, +): + """SM100 version of load_block_list (no intra_wg_overlap, no extra_tx_count).""" + block_count = block_end - block_begin + if block_count > 0: + # First iteration: load Q alongside K if requested + n_block_first = block_indices[block_end - 1] + + if const_expr(load_q_with_first): + # SM100 loads Q0 and optionally Q1 + load_Q(block=0, stage=0) + if const_expr(q_stage == 2): + load_Q(block=1, stage=1) + + # SM100 doesn't use producer_acquire for pipeline_kv in load path + # The pipeline barriers are handled inside load_KV + load_K(block=n_block_first, producer_state=kv_producer_state, page_idx=None) + kv_producer_state.advance() + load_V(block=n_block_first, producer_state=kv_producer_state, page_idx=None) + kv_producer_state.advance() + + # Remaining blocks + for offset in cutlass.range(1, block_count): + n_block = block_indices[block_end - 1 - offset] + load_K(block=n_block, producer_state=kv_producer_state, page_idx=None) + kv_producer_state.advance() + load_V(block=n_block, producer_state=kv_producer_state, page_idx=None) + kv_producer_state.advance() + + return kv_producer_state + + +# SM100-specific tile processor using SM100 helpers +@cute.jit +def produce_block_sparse_loads_sm100( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + m_block, + seqlen_info: SeqlenInfoQK, + split_idx: Int32, + num_splits: Int32, + kv_producer_state, + load_Q, + load_K, + load_V, + pipeline_kv, + q_stage: cutlass.Constexpr, + q_producer_phase: Int32, + qhead_per_kvhead: cutlass.Constexpr, + q_subtile_factor: cutlass.Constexpr, +): + """SM100 entry point for sparse block iteration. + + SM100 uses PipelineTmaUmma which doesn't support extra_tx_count, so we use + simplified block processing that just calls producer_acquire without extras. + + Args: + m_block: which tile of m we are processing + qhead_per_kvhead: Constexpr pack factor + """ + m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor) + + ( + curr_mask_block_cnt, + curr_mask_block_idx, + curr_full_block_cnt, + curr_full_block_idx, + ) = get_curr_blocksparse_tensors( + batch_idx, + head_idx, + m_block_sparse, + blocksparse_tensors, + seqlen_info, + ) + + mask_begin, mask_end = split_block_range(curr_mask_block_cnt, split_idx, num_splits) + full_begin, full_end = split_block_range(curr_full_block_cnt, split_idx, num_splits) + mask_empty = mask_begin == mask_end + full_empty = full_begin == full_end + + q_phase_flipped = False + + if mask_empty: + # No masked blocks: process full list with Q loading + kv_producer_state = load_block_list_sm100( + curr_full_block_idx, + full_begin, + full_end, + load_q_with_first=True, + q_stage=q_stage, + kv_producer_state=kv_producer_state, + load_Q=load_Q, + load_K=load_K, + load_V=load_V, + pipeline_kv=pipeline_kv, + ) + q_phase_flipped = not full_empty + else: + # Process masked blocks with Q loading + kv_producer_state = load_block_list_sm100( + curr_mask_block_idx, + mask_begin, + mask_end, + load_q_with_first=True, + q_stage=q_stage, + kv_producer_state=kv_producer_state, + load_Q=load_Q, + load_K=load_K, + load_V=load_V, + pipeline_kv=pipeline_kv, + ) + q_phase_flipped = True + + if not full_empty: + # Process full blocks without Q loading + kv_producer_state = load_block_list_sm100( + curr_full_block_idx, + full_begin, + full_end, + load_q_with_first=False, + q_stage=q_stage, + kv_producer_state=kv_producer_state, + load_Q=load_Q, + load_K=load_K, + load_V=load_V, + pipeline_kv=pipeline_kv, + ) + + if q_phase_flipped: + q_producer_phase ^= 1 + + return kv_producer_state, q_producer_phase + + +@cute.jit +def get_total_block_count( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + m_block, + split_idx: Int32, + num_splits: Int32, + qhead_per_kvhead: cutlass.Constexpr, + q_subtile_factor: cutlass.Constexpr, + seqlen_info: SeqlenInfoQK, +): + m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor) + ( + curr_mask_block_cnt, + _, + curr_full_block_cnt, + _, + ) = get_curr_blocksparse_tensors( + batch_idx, + head_idx, + m_block_sparse, + blocksparse_tensors, + seqlen_info, + ) + + mask_begin, mask_end = split_block_range(curr_mask_block_cnt, split_idx, num_splits) + full_begin, full_end = split_block_range(curr_full_block_cnt, split_idx, num_splits) + return mask_end - mask_begin + full_end - full_begin + + +@cute.jit +def handle_block_sparse_empty_tile_correction_sm100( + tidx: Int32, + q_stage: cutlass.Constexpr, + m_block_size: cutlass.Constexpr, + qhead_per_kvhead, + pack_gqa: cutlass.Constexpr, + is_split_kv: cutlass.Constexpr, + learnable_sink, + mLSE, + seqlen_info, + m_block: Int32, + head_idx: Int32, + batch_idx: Int32, + split_idx: Int32, + stats: list, + correction_epilogue: Callable, + thr_mma_pv: cute.ThrMma, + tOtO: cute.Tensor, + sO: cute.Tensor, + pipeline_sm_stats: cutlass.pipeline.PipelineAsync, + sm_stats_barrier: cutlass.pipeline.NamedBarrier, + sm_stats_barrier_cta_wide: cutlass.Constexpr, + pipeline_o_epi: cutlass.pipeline.PipelineAsync, + sm_stats_consumer_phase: Int32, + o_corr_consumer_phase: Int32, + corr_epi_producer_phase: Int32, + softmax_scale_log2: Float32, + max_offset: Float32, + max_offset_scale: Float32, + mO_cur: Optional[cute.Tensor] = None, + gO: Optional[cute.Tensor] = None, + gmem_tiled_copy_O: Optional[cute.TiledCopy] = None, +): + """Handle SM100 forward block-sparse tiles with no active KV blocks. + + This path is taken when `total_block_cnt == 0`. The softmax warp-group still + arrives `mbar_softmax_corr_full` (synthetic "no work") so the correction + warp-group can: + + - seed fully-masked-row stats (row_sum=1; row_max=-inf when tracked) for LSE + - run `correction_epilogue` with `scale=0` so the output tile is written as zeros + (independent of any prior tmem contents) + - wait on `mbar_softmax_corr_full` and arrive `mbar_softmax_corr_empty` + (and `mbar_corr_epi_*` when applicable) so phases stay aligned across tiles + + This helper intentionally does not touch `mbar_P_full_*` since no P is produced. + See NOTE [SM100 block-sparse empty tiles: mbarrier contract]. + """ + LOG2_E = Float32(math.log2(math.e)) + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4 + + for stage in cutlass.range_constexpr(q_stage): + row_sum_value = Float32(1.0) + row_max_value = ( + -Float32.inf if const_expr(mLSE is not None or learnable_sink is not None) else None + ) + if const_expr(learnable_sink is not None): + sink_val = -Float32.inf + if const_expr(not pack_gqa): + sink_val = Float32(learnable_sink[head_idx]) + elif tidx < m_block_size: + q_head_idx = ( + (q_stage * m_block + stage) * m_block_size + tidx + ) % qhead_per_kvhead + head_idx * qhead_per_kvhead + sink_val = Float32(learnable_sink[q_head_idx]) + if sink_val != -Float32.inf and (const_expr(not is_split_kv) or split_idx == 0): + if row_max_value == -Float32.inf: + row_max_value = sink_val * (LOG2_E / softmax_scale_log2) + row_sum_value = max_offset_scale + else: + row_sum_value = row_sum_value + cute.math.exp2( + sink_val * LOG2_E - row_max_value * softmax_scale_log2 + max_offset, + fastmath=True, + ) + acc_flag = row_sum_value == Float32(0.0) or row_sum_value != row_sum_value + stats[stage] = (row_sum_value, row_max_value, acc_flag) + + # See NOTE [SM100 block-sparse empty tiles: mbarrier contract]. + # pipeline_sm_stats.consumer_wait_w_index_phase(stage, sm_stats_consumer_phase) + sm_stats_barrier.arrive_and_wait_w_index( + index=stage if const_expr(sm_stats_barrier_cta_wide) else stage * 4 + warp_idx + ) + pipeline_sm_stats.consumer_release_w_index(stage) + + if const_expr(gmem_tiled_copy_O is None): + pipeline_o_epi.producer_acquire_w_index_phase(stage, corr_epi_producer_phase) + + gO_stage = gO[None, None, stage] if const_expr(gO is not None) else None + correction_epilogue( + thr_mma_pv, + tOtO[None, None, None, stage], + tidx, + stage, + m_block, + seqlen_info.seqlen_q, + Float32(0.0), # zero scale ensures empty tile writes zeros into staged outputs + sO[None, None, stage], + mO_cur, + gO_stage, + gmem_tiled_copy_O, + ) + if const_expr(gmem_tiled_copy_O is None): + pipeline_o_epi.producer_commit_w_index(stage) + + sm_stats_consumer_phase ^= 1 + corr_epi_producer_phase ^= 1 + + return ( + sm_stats_consumer_phase, + o_corr_consumer_phase, + corr_epi_producer_phase, + ) + + +@cute.jit +def softmax_block_sparse_sm100( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + m_block, + seqlen_info: SeqlenInfoQK, + split_idx: Int32, + num_splits: Int32, + softmax_step: Callable, + mask_fn: Callable, + mask_fn_none: Callable, + mma_si_consumer_phase: Int32, + si_corr_producer_phase: Int32, + s0_s1_sequence_phase: Int32, + pipeline_sm_stats: cutlass.pipeline.PipelineAsync, + sm_stats_barrier: cutlass.pipeline.NamedBarrier, + sm_stats_barrier_cta_wide: cutlass.Constexpr, + q_stage: cutlass.Constexpr, + stage_idx: Int32, + check_m_boundary: bool, + qhead_per_kvhead: cutlass.Constexpr, + q_subtile_factor: cutlass.Constexpr[int] = 1, +): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4 + m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor) + + ( + curr_mask_block_cnt, + curr_mask_block_idx, + curr_full_block_cnt, + curr_full_block_idx, + ) = get_curr_blocksparse_tensors( + batch_idx, + head_idx, + m_block_sparse, + blocksparse_tensors, + seqlen_info, + ) + + mask_begin, mask_end = split_block_range(curr_mask_block_cnt, split_idx, num_splits) + full_begin, full_end = split_block_range(curr_full_block_cnt, split_idx, num_splits) + split_mask_block_cnt = mask_end - mask_begin + split_full_block_cnt = full_end - full_begin + total_block_cnt = split_mask_block_cnt + split_full_block_cnt + + if total_block_cnt == 0: + sm_stats_barrier.arrive_w_index( + index=( + stage_idx + if const_expr(sm_stats_barrier_cta_wide) + else stage_idx * 4 + warp_idx + ) + ) + else: + if split_mask_block_cnt > 0: + mask_n_block = curr_mask_block_idx[mask_end - 1] + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + ) = softmax_step( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + mask_n_block, + is_first=True, + mask_fn=partial(mask_fn, mask_seqlen=True, check_q_boundary=check_m_boundary), + ) + for i in cutlass.range(1, split_mask_block_cnt): + mask_n_block = curr_mask_block_idx[mask_end - 1 - i] + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + ) = softmax_step( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + mask_n_block, + mask_fn=partial(mask_fn, mask_seqlen=False, check_q_boundary=check_m_boundary), + ) + + if split_full_block_cnt > 0: + full_n_block = curr_full_block_idx[full_end - 1] + if split_mask_block_cnt == 0: + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + ) = softmax_step( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + full_n_block, + is_first=True, + mask_fn=partial( + mask_fn_none, mask_seqlen=True, check_q_boundary=check_m_boundary + ), + ) + else: + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + ) = softmax_step( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + full_n_block, + is_first=False, + mask_fn=partial( + mask_fn_none, mask_seqlen=True, check_q_boundary=check_m_boundary + ), + ) + for i in cutlass.range(1, split_full_block_cnt): + full_n_block = curr_full_block_idx[full_end - 1 - i] + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + ) = softmax_step( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + full_n_block, + mask_fn=partial( + mask_fn_none, mask_seqlen=False, check_q_boundary=check_m_boundary + ), + ) + + return ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + total_block_cnt == 0, + ) + + +@cute.jit +def softmax_block_sparse_dual_stream_sm100( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + m_block, + seqlen_info: SeqlenInfoQK, + split_idx: Int32, + num_splits: Int32, + softmax_step: Callable, + mask_fn: Callable, + mask_fn_none: Callable, + mma_si_consumer_phase: Int32, + si_corr_producer_phase: Int32, + s0_s1_sequence_phase: Int32, + sm_stats_barrier: cutlass.pipeline.NamedBarrier, + sm_stats_barrier_cta_wide: cutlass.Constexpr, + stage_idx: Int32, + check_m_boundary: bool, + qhead_per_kvhead: cutlass.Constexpr[int], + q_subtile_factor: cutlass.Constexpr[int] = 1, +): + """Softmax traversal for one canonical even/odd KV stream.""" + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4 + m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead, q_subtile_factor) + ( + curr_mask_block_cnt, + curr_mask_block_idx, + curr_full_block_cnt, + curr_full_block_idx, + ) = get_curr_blocksparse_tensors( + batch_idx, + head_idx, + m_block_sparse, + blocksparse_tensors, + seqlen_info, + ) + mask_begin, mask_end = split_block_range(curr_mask_block_cnt, split_idx, num_splits) + full_begin, full_end = split_block_range(curr_full_block_cnt, split_idx, num_splits) + split_mask_count = mask_end - mask_begin + split_full_count = full_end - full_begin + + mask_first_ordinal = stage_idx + stream_mask_count = dual_stream_block_count(split_mask_count, mask_first_ordinal) + # Full-list ordinals follow all masked-list ordinals, so carry parity over + # the masked/full boundary rather than restarting it for the full list. + full_first_ordinal = (stage_idx + split_mask_count) % Int32(2) + stream_full_count = dual_stream_block_count(split_full_count, full_first_ordinal) + stream_block_count = stream_mask_count + stream_full_count + + if stream_block_count == Int32(0): + sm_stats_barrier.arrive_w_index( + index=( + stage_idx + if const_expr(sm_stats_barrier_cta_wide) + else stage_idx * 4 + warp_idx + ) + ) + else: + if stream_mask_count > Int32(0): + local_ordinal = mask_first_ordinal + n_block = curr_mask_block_idx[mask_end - Int32(1) - local_ordinal] + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + ) = softmax_step( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + n_block, + is_first=True, + mask_fn=partial(mask_fn, mask_seqlen=True, check_q_boundary=check_m_boundary), + ) + for i in cutlass.range(1, stream_mask_count, unroll=1): + local_ordinal = mask_first_ordinal + i * Int32(2) + n_block = curr_mask_block_idx[mask_end - Int32(1) - local_ordinal] + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + ) = softmax_step( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + n_block, + mask_fn=partial( + mask_fn, mask_seqlen=False, check_q_boundary=check_m_boundary + ), + ) + + if const_expr(curr_full_block_idx is not None): + if stream_full_count > Int32(0): + local_ordinal = full_first_ordinal + n_block = curr_full_block_idx[full_end - Int32(1) - local_ordinal] + if stream_mask_count == Int32(0): + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + ) = softmax_step( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + n_block, + is_first=True, + mask_fn=partial( + mask_fn_none, + mask_seqlen=True, + check_q_boundary=check_m_boundary, + ), + ) + else: + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + ) = softmax_step( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + n_block, + is_first=False, + mask_fn=partial( + mask_fn_none, + mask_seqlen=True, + check_q_boundary=check_m_boundary, + ), + ) + for i in cutlass.range(1, stream_full_count, unroll=1): + local_ordinal = full_first_ordinal + i * Int32(2) + n_block = curr_full_block_idx[full_end - Int32(1) - local_ordinal] + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + ) = softmax_step( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + n_block, + mask_fn=partial( + mask_fn_none, + mask_seqlen=False, + check_q_boundary=check_m_boundary, + ), + ) + + return ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + stream_block_count == Int32(0), + stream_block_count, + ) + + +# ============================================================================= +# Backward-specific block-sparse helpers (SM100) +# ============================================================================= +# +# In backward, iteration is transposed compared to forward: +# - Forward: outer loop over m_blocks (Q tiles), inner loop over n_blocks (KV tiles) +# - Backward: outer loop over n_blocks (KV tiles), inner loop over m_blocks (Q tiles) +# +# The backward block-sparse tensors use "Q direction" indexing: +# - q_block_cnt[batch, head, n_block] → count of m_blocks to process for this KV tile +# - q_block_idx[batch, head, n_block, :] → indices of m_blocks to process +# + + +@cute.jit +def get_total_q_block_count_bwd( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + n_block, + q_subtile_factor: cutlass.Constexpr = 1, + m_block_max: int = 0, +): + """Count total tile iterations for given n_block (KV tile) in backward.""" + q_block_cnt, _, full_block_cnt, _, *_ = blocksparse_tensors + total = q_block_cnt[batch_idx, head_idx, n_block] + if const_expr(full_block_cnt is not None): + total = total + full_block_cnt[batch_idx, head_idx, n_block] + return total * q_subtile_factor + + +@cute.jit +def produce_block_sparse_q_loads_bwd_sm100( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + n_block, + # Pipeline states (will be returned after advancing) + producer_state_Q_LSE, + producer_state_dO_dPsum, + # Pipelines + pipeline_Q, + pipeline_LSE, + pipeline_dO, + pipeline_dPsum, + # Load functions + load_K, + load_V, + load_Q, + load_dO, + copy_stats, + # Global tensors for LSE/dPsum + gLSE, + sLSE, + gdPsum, + sdPsum, + # TMA copy bytes for extra_tx_count + tma_copy_bytes_K, + tma_copy_bytes_V, + # Flags for which loads to perform + should_load_Q: cutlass.Constexpr, + should_load_dO: cutlass.Constexpr, + # Subtiling factor and bounds + q_subtile_factor: cutlass.Constexpr = 1, + m_block_max: int = 0, +): + """SM100 backward block sparse loading with subtiling. + + Returns updated (producer_state_Q_LSE, producer_state_dO_dPsum). + First iteration loads K/V alongside Q/dO; subsequent iterations load only Q/dO. + """ + ( + curr_q_cnt, + curr_q_idx, + curr_full_cnt, + curr_full_idx, + loop_count, + ) = get_block_sparse_iteration_info_bwd( + blocksparse_tensors, batch_idx, head_idx, n_block, q_subtile_factor, m_block_max + ) + + for iter_idx in cutlass.range(loop_count, unroll=1): + m_block, _ = get_m_block_from_iter_bwd( + iter_idx, + curr_q_cnt, + curr_q_idx, + curr_full_cnt, + curr_full_idx, + q_subtile_factor, + m_block_max, + ) + m_block_safe = m_block + if m_block_max > 0: + m_block_safe = cutlass.min(m_block, m_block_max - 1) + + if iter_idx == 0: + # First block: load K/V alongside Q/dO + if const_expr(should_load_Q): + pipeline_Q.producer_acquire(producer_state_Q_LSE, extra_tx_count=tma_copy_bytes_K) + load_K(tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q_LSE)) + load_Q(m_block_safe, producer_state=producer_state_Q_LSE) + pipeline_Q.producer_commit(producer_state_Q_LSE) + pipeline_LSE.producer_acquire(producer_state_Q_LSE) + with cute.arch.elect_one(): + copy_stats( + gLSE[None, m_block_safe], + sLSE[None, producer_state_Q_LSE.index], + mbar_ptr=pipeline_LSE.producer_get_barrier(producer_state_Q_LSE), + ) + producer_state_Q_LSE.advance() + if const_expr(should_load_dO): + pipeline_dO.producer_acquire( + producer_state_dO_dPsum, extra_tx_count=tma_copy_bytes_V + ) + load_V(tma_bar_ptr=pipeline_dO.producer_get_barrier(producer_state_dO_dPsum)) + load_dO(m_block_safe, producer_state=producer_state_dO_dPsum) + pipeline_dO.producer_commit(producer_state_dO_dPsum) + pipeline_dPsum.producer_acquire(producer_state_dO_dPsum) + with cute.arch.elect_one(): + copy_stats( + gdPsum[None, m_block_safe], + sdPsum[None, producer_state_dO_dPsum.index], + mbar_ptr=pipeline_dPsum.producer_get_barrier(producer_state_dO_dPsum), + ) + producer_state_dO_dPsum.advance() + else: + # Subsequent blocks: just load Q/dO (K/V already loaded) + if const_expr(should_load_Q): + pipeline_Q.producer_acquire(producer_state_Q_LSE) + load_Q(m_block_safe, producer_state=producer_state_Q_LSE) + pipeline_Q.producer_commit(producer_state_Q_LSE) + pipeline_LSE.producer_acquire(producer_state_Q_LSE) + with cute.arch.elect_one(): + copy_stats( + gLSE[None, m_block_safe], + sLSE[None, producer_state_Q_LSE.index], + mbar_ptr=pipeline_LSE.producer_get_barrier(producer_state_Q_LSE), + ) + producer_state_Q_LSE.advance() + if const_expr(should_load_dO): + pipeline_dO.producer_acquire(producer_state_dO_dPsum) + load_dO(m_block_safe, producer_state=producer_state_dO_dPsum) + pipeline_dO.producer_commit(producer_state_dO_dPsum) + pipeline_dPsum.producer_acquire(producer_state_dO_dPsum) + with cute.arch.elect_one(): + copy_stats( + gdPsum[None, m_block_safe], + sdPsum[None, producer_state_dO_dPsum.index], + mbar_ptr=pipeline_dPsum.producer_get_barrier(producer_state_dO_dPsum), + ) + producer_state_dO_dPsum.advance() + + return producer_state_Q_LSE, producer_state_dO_dPsum + + +@cute.jit +def get_block_sparse_iteration_info_bwd( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + n_block, + q_subtile_factor: cutlass.Constexpr = 1, + m_block_max: int = 0, +): + """Extract block-sparse iteration info for backward pass. + + Returns (curr_q_cnt, curr_q_idx, curr_full_cnt, curr_full_idx, total_count). + """ + q_cnt, q_idx, full_cnt, full_idx, *_ = blocksparse_tensors + curr_q_cnt = q_cnt[batch_idx, head_idx, n_block] + curr_q_idx = q_idx[batch_idx, head_idx, n_block, None] + + if const_expr(full_cnt is not None): + curr_full_cnt = full_cnt[batch_idx, head_idx, n_block] + curr_full_idx = full_idx[batch_idx, head_idx, n_block, None] + else: + curr_full_cnt = Int32(0) + curr_full_idx = None + + sparse_block_count = curr_q_cnt + if const_expr(full_cnt is not None): + sparse_block_count = sparse_block_count + curr_full_cnt + total_count = sparse_block_count * q_subtile_factor + + return curr_q_cnt, curr_q_idx, curr_full_cnt, curr_full_idx, total_count + + +@cute.jit +def get_m_block_from_iter_bwd( + iter_idx, + curr_q_cnt, + curr_q_idx: cute.Tensor, + curr_full_cnt, + curr_full_idx: Optional[cute.Tensor], + q_subtile_factor: cutlass.Constexpr = 1, + m_block_max: int = 0, +): + """Derive m_block index and is_full_block flag from iteration index. + + Returns (m_block, is_full_block): + - m_block: The actual Q-tile block index + - is_full_block: True if this is a full block (no mask_mod needed) + """ + sparse_iter_idx = iter_idx // q_subtile_factor + subtile_offset = iter_idx % q_subtile_factor + + sparse_m_block = Int32(0) + is_full_block = False + if const_expr(curr_full_idx is not None): + if sparse_iter_idx < curr_q_cnt: + sparse_m_block = curr_q_idx[sparse_iter_idx] + else: + sparse_m_block = curr_full_idx[sparse_iter_idx - curr_q_cnt] + is_full_block = True + else: + sparse_m_block = curr_q_idx[sparse_iter_idx] + + return sparse_m_block * q_subtile_factor + subtile_offset, is_full_block + + +@cute.jit +def _load_q_do_block_sm90( + m_block, + producer_state_Q, + producer_state_dO, + pipeline_Q, + pipeline_dO, + load_K, + load_V, + load_Q, + load_dO, + load_LSE, + load_dPsum, + tma_copy_bytes_K, + tma_copy_bytes_V, + Q_stage_eq_dO_stage: cutlass.Constexpr, + load_kv: bool, +): + """Load one Q/dO block, optionally loading K/V on first iteration.""" + if load_kv: + pipeline_Q.producer_acquire(producer_state_Q, extra_tx_count=tma_copy_bytes_K) + load_K(tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q)) + else: + pipeline_Q.producer_acquire(producer_state_Q) + load_Q(m_block, producer_state=producer_state_Q) + load_LSE(m_block, producer_state=producer_state_Q) + + producer_state_dO_cur = ( + producer_state_dO if const_expr(not Q_stage_eq_dO_stage) else producer_state_Q + ) + if load_kv: + pipeline_dO.producer_acquire(producer_state_dO_cur, extra_tx_count=tma_copy_bytes_V) + load_V(tma_bar_ptr=pipeline_dO.producer_get_barrier(producer_state_dO_cur)) + else: + pipeline_dO.producer_acquire(producer_state_dO_cur) + load_dO(m_block, producer_state=producer_state_dO_cur) + load_dPsum(m_block, producer_state=producer_state_dO_cur) + + producer_state_Q.advance() + producer_state_dO.advance() + return producer_state_Q, producer_state_dO + + +@cute.jit +def produce_block_sparse_q_loads_bwd_sm90( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + n_block, + producer_state_Q, + producer_state_dO, + pipeline_Q, + pipeline_dO, + load_K, + load_V, + load_Q, + load_dO, + load_LSE, + load_dPsum, + tma_copy_bytes_K, + tma_copy_bytes_V, + Q_stage_eq_dO_stage: cutlass.Constexpr, + q_subtile_factor: cutlass.Constexpr, + m_block_max: int, +): + """SM90 backward block sparse loading with separate partial/full loops. + + K/V are loaded with the first valid block. Iterates partial blocks first, + then full blocks, matching consumer order. + + Returns updated (producer_state_Q, producer_state_dO). + """ + q_cnt, q_idx, full_cnt, full_idx, *_ = blocksparse_tensors + curr_q_cnt = q_cnt[batch_idx, head_idx, n_block] + curr_q_idx = q_idx[batch_idx, head_idx, n_block, None] + + if const_expr(full_cnt is not None): + curr_full_cnt = full_cnt[batch_idx, head_idx, n_block] + curr_full_idx = full_idx[batch_idx, head_idx, n_block, None] + else: + curr_full_cnt = Int32(0) + curr_full_idx = None + + kv_loaded = False + + for iter_idx in cutlass.range(curr_q_cnt * q_subtile_factor, unroll=1): + sparse_idx = iter_idx // q_subtile_factor + subtile_offset = iter_idx % q_subtile_factor + m_block = curr_q_idx[sparse_idx] * q_subtile_factor + subtile_offset + + if m_block < m_block_max: + producer_state_Q, producer_state_dO = _load_q_do_block_sm90( + m_block, + producer_state_Q, + producer_state_dO, + pipeline_Q, + pipeline_dO, + load_K, + load_V, + load_Q, + load_dO, + load_LSE, + load_dPsum, + tma_copy_bytes_K, + tma_copy_bytes_V, + Q_stage_eq_dO_stage, + load_kv=not kv_loaded, + ) + kv_loaded = True + + if const_expr(full_cnt is not None): + for iter_idx in cutlass.range(curr_full_cnt * q_subtile_factor, unroll=1): + sparse_idx = iter_idx // q_subtile_factor + subtile_offset = iter_idx % q_subtile_factor + m_block = curr_full_idx[sparse_idx] * q_subtile_factor + subtile_offset + + if m_block < m_block_max: + producer_state_Q, producer_state_dO = _load_q_do_block_sm90( + m_block, + producer_state_Q, + producer_state_dO, + pipeline_Q, + pipeline_dO, + load_K, + load_V, + load_Q, + load_dO, + load_LSE, + load_dPsum, + tma_copy_bytes_K, + tma_copy_bytes_V, + Q_stage_eq_dO_stage, + load_kv=not kv_loaded, + ) + kv_loaded = True + + return producer_state_Q, producer_state_dO + + +@cute.jit +def consume_block_sparse_mma_bwd_sm90( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + n_block, + consumer_state_Q, + consumer_state_dO, + mma_one_m_block_fn, + mask, + mask_mod, + is_causal: cutlass.Constexpr, + is_local: cutlass.Constexpr, + thr_mma_SdP, + score_mod_fn=None, + score_mod_bwd_fn=None, + q_subtile_factor: cutlass.Constexpr = 1, + m_block_max: int = 0, + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), +): + """SM90 backward block sparse MMA consumption with separate partial/full loops. + + Partial blocks are processed first (with mask_mod applied), then full blocks + (without mask_mod). This ensures mask_mod is only applied where needed. + + Returns updated (consumer_state_Q, consumer_state_dO). + """ + q_cnt, q_idx, full_cnt, full_idx, *_ = blocksparse_tensors + curr_q_cnt = q_cnt[batch_idx, head_idx, n_block] + curr_q_idx = q_idx[batch_idx, head_idx, n_block, None] + + if const_expr(full_cnt is not None): + curr_full_cnt = full_cnt[batch_idx, head_idx, n_block] + curr_full_idx = full_idx[batch_idx, head_idx, n_block, None] + else: + curr_full_cnt = Int32(0) + curr_full_idx = None + + dKV_accumulate = False + + mask_fn_partial = partial( + mask.apply_mask, + batch_idx=batch_idx, + head_idx=head_idx, + n_block=n_block, + thr_mma=thr_mma_SdP, + mask_seqlen=True, + mask_causal=is_causal, + mask_local=is_local, + mask_mod=mask_mod, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + + mask_fn_full = partial( + mask.apply_mask, + batch_idx=batch_idx, + head_idx=head_idx, + n_block=n_block, + thr_mma=thr_mma_SdP, + mask_seqlen=True, + mask_causal=is_causal, + mask_local=is_local, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + + for iter_idx in cutlass.range(curr_q_cnt * q_subtile_factor, unroll=1): + sparse_idx = iter_idx // q_subtile_factor + subtile_offset = iter_idx % q_subtile_factor + m_block = curr_q_idx[sparse_idx] * q_subtile_factor + subtile_offset + + if m_block < m_block_max: + consumer_state_Q, consumer_state_dO = mma_one_m_block_fn( + m_block, + consumer_state_Q, + consumer_state_dO, + mask_fn=mask_fn_partial, + score_mod_fn=score_mod_fn, + score_mod_bwd_fn=score_mod_bwd_fn, + dKV_accumulate=dKV_accumulate, + ) + dKV_accumulate = True + + if const_expr(full_cnt is not None): + for iter_idx in cutlass.range(curr_full_cnt * q_subtile_factor, unroll=1): + sparse_idx = iter_idx // q_subtile_factor + subtile_offset = iter_idx % q_subtile_factor + m_block = curr_full_idx[sparse_idx] * q_subtile_factor + subtile_offset + + if m_block < m_block_max: + consumer_state_Q, consumer_state_dO = mma_one_m_block_fn( + m_block, + consumer_state_Q, + consumer_state_dO, + mask_fn=mask_fn_full, + score_mod_fn=score_mod_fn, + score_mod_bwd_fn=score_mod_bwd_fn, + dKV_accumulate=dKV_accumulate, + ) + dKV_accumulate = True + + return consumer_state_Q, consumer_state_dO + + +@cute.jit +def _store_one_dQaccum_sm90( + m_block, + sdQaccum: cute.Tensor, + gdQaccum: cute.Tensor, + num_dQ_warp_groups: cutlass.Constexpr, + num_threads_per_warp_group: cutlass.Constexpr, + tma_copy_bytes_dQ, +): + """Store dQaccum for a single m_block.""" + for warp_group_idx in cutlass.range_constexpr(num_dQ_warp_groups): + cute.arch.cp_async_bulk_wait_group(num_dQ_warp_groups - 1 - warp_group_idx, read=True) + cute.arch.barrier_arrive( + barrier_id=int(NamedBarrierBwd.dQEmptyWG0) + warp_group_idx, + number_of_threads=num_threads_per_warp_group + cute.arch.WARP_SIZE, + ) + for warp_group_idx in cutlass.range_constexpr(num_dQ_warp_groups): + cute.arch.barrier( + barrier_id=int(NamedBarrierBwd.dQFullWG0) + warp_group_idx, + number_of_threads=num_threads_per_warp_group + cute.arch.WARP_SIZE, + ) + with cute.arch.elect_one(): + copy_utils.cpasync_reduce_bulk_add_f32( + sdQaccum[None, warp_group_idx].iterator, + gdQaccum[(None, warp_group_idx), m_block].iterator, + tma_copy_bytes_dQ, + ) + cute.arch.cp_async_bulk_commit_group() + + +@cute.jit +def dQaccum_store_block_sparse_bwd_sm90( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + n_block, + sdQaccum: cute.Tensor, + gdQaccum: cute.Tensor, + q_subtile_factor: cutlass.Constexpr, + m_block_max: int, + num_dQ_warp_groups: cutlass.Constexpr, + num_threads_per_warp_group: cutlass.Constexpr, + tma_copy_bytes_dQ, +): + """SM90 backward block sparse dQaccum store with separate partial/full loops. + + Iterates partial blocks first, then full blocks, matching producer/consumer order. + """ + q_cnt, q_idx, full_cnt, full_idx, *_ = blocksparse_tensors + curr_q_cnt = q_cnt[batch_idx, head_idx, n_block] + curr_q_idx = q_idx[batch_idx, head_idx, n_block, None] + + if const_expr(full_cnt is not None): + curr_full_cnt = full_cnt[batch_idx, head_idx, n_block] + curr_full_idx = full_idx[batch_idx, head_idx, n_block, None] + else: + curr_full_cnt = Int32(0) + curr_full_idx = None + + for iter_idx in cutlass.range(curr_q_cnt * q_subtile_factor, unroll=1): + sparse_idx = iter_idx // q_subtile_factor + subtile_offset = iter_idx % q_subtile_factor + m_block = curr_q_idx[sparse_idx] * q_subtile_factor + subtile_offset + + if m_block < m_block_max: + _store_one_dQaccum_sm90( + m_block, + sdQaccum, + gdQaccum, + num_dQ_warp_groups, + num_threads_per_warp_group, + tma_copy_bytes_dQ, + ) + + if const_expr(full_cnt is not None): + for iter_idx in cutlass.range(curr_full_cnt * q_subtile_factor, unroll=1): + sparse_idx = iter_idx // q_subtile_factor + subtile_offset = iter_idx % q_subtile_factor + m_block = curr_full_idx[sparse_idx] * q_subtile_factor + subtile_offset + + if m_block < m_block_max: + _store_one_dQaccum_sm90( + m_block, + sdQaccum, + gdQaccum, + num_dQ_warp_groups, + num_threads_per_warp_group, + tma_copy_bytes_dQ, + ) diff --git a/fastvideo-kernel/fa4/block_sparsity.py b/fastvideo-kernel/fa4/block_sparsity.py new file mode 100644 index 000000000..009886e83 --- /dev/null +++ b/fastvideo-kernel/fa4/block_sparsity.py @@ -0,0 +1,676 @@ +""" +Block-sparsity utilities for FlexAttention +""" + +from typing import Callable, NamedTuple, Tuple + +import cutlass.cute as cute +import torch + +from flash_attn.cute.cute_dsl_utils import get_broadcast_dims, to_cute_tensor + + +def ceildiv(a: int, b: int) -> int: + return (a + b - 1) // b + + +class BlockSparseTensors(NamedTuple): + mask_block_cnt: cute.Tensor + mask_block_idx: cute.Tensor + full_block_cnt: cute.Tensor | None = None + full_block_idx: cute.Tensor | None = None + cu_total_m_blocks: cute.Tensor | None = None + cu_block_idx_offsets: cute.Tensor | None = None + dq_write_order: cute.Tensor | None = None + dq_write_order_full: cute.Tensor | None = None + + def __new_from_mlir_values__(self, values): + new_fields = [] + idx = 0 + for original in self: + if original is None: + new_fields.append(None) + else: + new_fields.append(values[idx]) + idx += 1 + return BlockSparseTensors(*new_fields) + + +class BlockSparseTensorsTorch(NamedTuple): + mask_block_cnt: torch.Tensor + mask_block_idx: torch.Tensor + full_block_cnt: torch.Tensor | None = None + full_block_idx: torch.Tensor | None = None + cu_total_m_blocks: torch.Tensor | None = None + cu_block_idx_offsets: torch.Tensor | None = None + block_size: tuple[int, int] | None = None + dq_write_order: torch.Tensor | None = None + dq_write_order_full: torch.Tensor | None = None + spt: bool | None = None + + +def _ordered_to_dense_simple( + num_blocks: torch.Tensor, + indices: torch.Tensor, + num_cols: int, +) -> torch.Tensor: + """Convert ordered sparse representation to dense binary matrix. + + Args: + num_blocks: [B, H, num_rows] count of valid entries per row + indices: [B, H, num_rows, max_entries] column indices (valid entries packed left) + num_cols: total number of columns + + Returns: + dense: [B, H, num_rows, num_cols] binary int32 matrix + """ + B, H, num_rows, max_entries = indices.shape + device = indices.device + dense = torch.zeros(B, H, num_rows, num_cols + 1, dtype=torch.int32, device=device) + col_range = torch.arange(max_entries, device=device) + valid = col_range[None, None, None, :] < num_blocks[:, :, :, None] + safe_indices = torch.where(valid, indices.long(), num_cols) + row_idx = torch.arange(num_rows, device=device)[None, None, :, None].expand_as(indices) + b_idx = torch.arange(B, device=device)[:, None, None, None].expand_as(indices) + h_idx = torch.arange(H, device=device)[None, :, None, None].expand_as(indices) + dense[b_idx, h_idx, row_idx, safe_indices] = 1 + return dense[:, :, :, :num_cols] + + +def compute_dq_write_order( + fwd_mask_cnt: torch.Tensor, + fwd_mask_idx: torch.Tensor, + fwd_full_cnt: torch.Tensor | None, + fwd_full_idx: torch.Tensor | None, + bwd_mask_cnt: torch.Tensor, + bwd_mask_idx: torch.Tensor, + bwd_full_cnt: torch.Tensor | None, + bwd_full_idx: torch.Tensor | None, + spt: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Compute dQ write-order metadata for deterministic block-sparse backward. + + For each (n_block, i) in the backward iteration, computes the semaphore + lock value: the rank of n_block in the combined (partial + full) sorted + contributor list for the target m_block. + + Lock values are assigned in ascending n_block order (or descending if spt=True) + to guarantee deadlock-freedom with the CTA scheduling order. + + Args: + fwd_mask_cnt: [B, H, num_m_blocks] partial contributor counts per m_block + fwd_mask_idx: [B, H, num_m_blocks, max_kv] partial contributor n_block indices (ascending) + fwd_full_cnt: [B, H, num_m_blocks] full contributor counts per m_block (optional) + fwd_full_idx: [B, H, num_m_blocks, max_kv] full contributor n_block indices (optional) + bwd_mask_cnt: [B, H, num_n_blocks] partial iteration counts per n_block + bwd_mask_idx: [B, H, num_n_blocks, max_q] partial iteration m_block indices + bwd_full_cnt: [B, H, num_n_blocks] full iteration counts per n_block (optional) + bwd_full_idx: [B, H, num_n_blocks, max_q] full iteration m_block indices (optional) + spt: if True, reverse ordering (highest n_block gets lock_value=0) + + Returns: + (dq_write_order, dq_write_order_full): tensors parallel to bwd_mask_idx + and bwd_full_idx respectively, containing lock values. + """ + device = fwd_mask_idx.device + B, H, num_m, max_kv_partial = fwd_mask_idx.shape + _, _, num_n, max_q_partial = bwd_mask_idx.shape + + has_full = fwd_full_cnt is not None and fwd_full_idx is not None + + dense_partial = _ordered_to_dense_simple(fwd_mask_cnt, fwd_mask_idx, num_n) + if has_full: + dense_full = _ordered_to_dense_simple(fwd_full_cnt, fwd_full_idx, num_n) + dense = (dense_partial + dense_full).clamp(max=1) + else: + dense = dense_partial + + cumsum = dense.cumsum(dim=-1) + rank_table = (cumsum - dense).to(torch.int32) + + if spt: + total_per_m = cumsum[:, :, :, -1:] + rank_table = (total_per_m - 1 - rank_table).to(torch.int32) + + def _gather_write_order(bwd_idx, bwd_cnt): + b_i = torch.arange(B, device=device)[:, None, None, None].expand_as(bwd_idx) + h_i = torch.arange(H, device=device)[None, :, None, None].expand_as(bwd_idx) + n_i = torch.arange(bwd_idx.shape[2], device=device)[None, None, :, None].expand_as(bwd_idx) + m_vals = bwd_idx.long().clamp(0, num_m - 1) + return rank_table[b_i, h_i, m_vals, n_i].to(torch.int32) + + dq_write_order = _gather_write_order(bwd_mask_idx, bwd_mask_cnt) + + dq_write_order_full = None + if has_full and bwd_full_cnt is not None and bwd_full_idx is not None: + dq_write_order_full = _gather_write_order(bwd_full_idx, bwd_full_cnt) + + return dq_write_order, dq_write_order_full + + +def compute_dq_write_order_from_block_mask( + block_mask, + spt: bool = False, +) -> tuple[torch.Tensor, torch.Tensor | None]: + ( + _seq_q, + _seq_k, + kv_mask_cnt, + kv_mask_idx, + full_kv_cnt, + full_kv_idx, + q_mask_cnt, + q_mask_idx, + full_q_cnt, + full_q_idx, + *_, + ) = block_mask.as_tuple() + return compute_dq_write_order( + kv_mask_cnt, + kv_mask_idx, + full_kv_cnt, + full_kv_idx, + q_mask_cnt, + q_mask_idx, + full_q_cnt, + full_q_idx, + spt=spt, + ) + + +def get_sparse_q_block_size( + tensors: BlockSparseTensorsTorch | None, + seqlen_q: int, +) -> int | None: + """Return the Q sparse block size, or None when sparsity is unset or ambiguous.""" + if tensors is None: + return None + if tensors.block_size is not None: + return tensors.block_size[0] + num_m_blocks = tensors.mask_block_idx.shape[2] + min_block_size = ceildiv(seqlen_q, num_m_blocks) + max_block_size = seqlen_q if num_m_blocks == 1 else (seqlen_q - 1) // (num_m_blocks - 1) + if min_block_size != max_block_size: + return None + return min_block_size + + +def _expand_sparsity_tensor( + tensor: torch.Tensor, + expected_shape: Tuple[int, ...], + tensor_name: str, + context: str | None, + hint: str | Callable[[], str] | None, +) -> torch.Tensor: + """Check if we need to expand the tensor to expected shape, and do so if possible.""" + needs_expand = tensor.shape != expected_shape + if not needs_expand: + return tensor + can_expand = all(map(lambda cur, tgt: cur == tgt or cur == 1, tensor.shape, expected_shape)) + if not can_expand: + context_clause = f" ({context})" if context else "" + resolved_hint = hint() if callable(hint) else hint + hint_clause = f" Hint: {resolved_hint}" if resolved_hint else "" + raise ValueError( + f"{tensor_name}{context_clause} with shape {tensor.shape} cannot be expanded to expected shape {expected_shape}." + f"{hint_clause}" + ) + return tensor.expand(*expected_shape) + + +def _check_and_expand_block( + name: str, + cnt: torch.Tensor | None, + idx: torch.Tensor | None, + expected_count_shape: Tuple[int, ...], + expected_index_shape: Tuple[int, ...], + context: str | None, + hint: str | Callable[[], str] | None, +) -> Tuple[torch.Tensor | None, torch.Tensor | None]: + if (cnt is None) != (idx is None): + raise ValueError( + f"{name}_block_cnt and {name}_block_idx must both be provided or both be None" + ) + if cnt is None or idx is None: + return None, None + if cnt.dtype != torch.int32 or idx.dtype != torch.int32: + raise ValueError(f"{name}_block tensors must have dtype torch.int32") + if cnt.device != idx.device: + raise ValueError(f"{name}_block_cnt and {name}_block_idx must be on the same device") + if not cnt.is_cuda or not idx.is_cuda: + raise ValueError(f"{name}_block tensors must live on CUDA") + expanded_cnt = _expand_sparsity_tensor( + cnt, expected_count_shape, f"{name}_block_cnt", context, hint + ) + # [Note] Allow Compact block sparse indices + # Allow the last dimension (n_blocks) of idx to be <= expected, since + # FA4 only accesses indices 0..cnt-1 per query tile. This enables compact + # index tensors that avoid O(N^2) memory at long sequence lengths. + if idx.ndim == 4 and idx.shape[3] <= expected_index_shape[3]: + expected_index_shape = (*expected_index_shape[:3], idx.shape[3]) + expanded_idx = _expand_sparsity_tensor( + idx, expected_index_shape, f"{name}_block_idx", context, hint + ) + return expanded_cnt, expanded_idx + + +def _check_and_expand_metadata_tensor( + name: str, + tensor: torch.Tensor | None, + expected_shape: Tuple[int, ...], + context: str | None, + hint: str | Callable[[], str] | None, + device: torch.device, +) -> torch.Tensor | None: + if tensor is None: + return None + if tensor.dtype != torch.int32: + raise ValueError(f"{name} must have dtype torch.int32") + if tensor.device != device: + raise ValueError(f"{name} must be on the same device as block sparse tensors") + if not tensor.is_cuda: + raise ValueError(f"{name} must live on CUDA") + return _expand_sparsity_tensor(tensor, expected_shape, name, context, hint) + + +def get_block_sparse_expected_shapes( + batch_size: int, + num_head: int, + seqlen_q: int, + seqlen_k: int, + m_block_size: int, + n_block_size: int, + q_stage: int, +) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int]]: + """Return (expected_count_shape, expected_index_shape) for block sparse normalization.""" + m_block_size_effective = q_stage * m_block_size + expected_m_blocks = ceildiv(seqlen_q, m_block_size_effective) + expected_n_blocks = ceildiv(seqlen_k, n_block_size) + expected_count_shape = (batch_size, num_head, expected_m_blocks) + expected_index_shape = (batch_size, num_head, expected_m_blocks, expected_n_blocks) + return expected_count_shape, expected_index_shape + + +def infer_block_sparse_expected_shapes( + tensors: BlockSparseTensorsTorch, + *, + batch_size: int, + num_head: int, + seqlen_q: int, + seqlen_k: int, + m_block_size: int, + n_block_size: int, + q_stage: int, + context: str, + sparse_block_size_q: int | None = None, + sparse_block_size_kv: int | None = None, +) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int], int]: + """Infer shapes and scaling for block-sparse tensors. + + Expectations: + - mask_block_cnt is (B, H, M) and mask_block_idx is (B, H, M, N). + - Batch/head dims may be 1 for broadcast, or match the requested sizes. + - sparse_block_size_kv must match tile_n. + - sparse_block_size_q must be a multiple of q_stage * tile_m. + - If sparse_block_size_q is omitted and seqlen_q/num_m_blocks is ambiguous, + the caller must provide block_size to disambiguate. TODO will make this required in a future PR. + """ + base_m_block = q_stage * m_block_size + base_n_block = n_block_size + if sparse_block_size_kv is None: + sparse_block_size_kv = base_n_block + if sparse_block_size_kv != base_n_block: + raise ValueError(f"Block sparse tensors{context} require BLOCK_SIZE_KV={base_n_block}.") + if tensors.mask_block_idx is None: + raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.") + num_m_blocks = tensors.mask_block_idx.shape[2] + + if sparse_block_size_q is None: + sparse_block_size_q = get_sparse_q_block_size(tensors, seqlen_q) + if sparse_block_size_q is None and base_m_block != 1: + raise ValueError( + f"Block sparse tensors{context} require explicit sparse_block_size[0] " + f"to disambiguate block size for seqlen_q={seqlen_q} and num_m_blocks={num_m_blocks}." + ) + if sparse_block_size_q is None: + sparse_block_size_q = ceildiv(seqlen_q, num_m_blocks) + + if sparse_block_size_q % base_m_block != 0: + raise ValueError( + f"Block sparse tensors{context} have block size {sparse_block_size_q}, " + f"which must be a multiple of {base_m_block}." + ) + + expected_m_blocks = ceildiv(seqlen_q, sparse_block_size_q) + expected_n_blocks = ceildiv(seqlen_k, sparse_block_size_kv) + q_subtile_factor = sparse_block_size_q // base_m_block + expected_count_shape = (batch_size, num_head, expected_m_blocks) + expected_index_shape = (batch_size, num_head, expected_m_blocks, expected_n_blocks) + + mask_block_cnt = tensors.mask_block_cnt + mask_block_idx = tensors.mask_block_idx + if mask_block_cnt is None or mask_block_idx is None: + raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.") + if mask_block_cnt.ndim != 3 or mask_block_idx.ndim != 4: + raise ValueError( + f"Block sparse tensors{context} must have shapes (B, H, M) and (B, H, M, N)." + ) + for dim_name, cur, tgt in ( + ("batch", mask_block_cnt.shape[0], expected_count_shape[0]), + ("head", mask_block_cnt.shape[1], expected_count_shape[1]), + ): + if cur != tgt and cur != 1: + raise ValueError(f"Block sparse tensors{context} {dim_name} dim must be {tgt} or 1.") + for dim_name, cur, tgt in ( + ("batch", mask_block_idx.shape[0], expected_index_shape[0]), + ("head", mask_block_idx.shape[1], expected_index_shape[1]), + ): + if cur != tgt and cur != 1: + raise ValueError(f"Block sparse tensors{context} {dim_name} dim must be {tgt} or 1.") + if mask_block_cnt.shape[2] != mask_block_idx.shape[2]: + raise ValueError(f"Block sparse tensors{context} must share the same m-block dimension.") + # [Note] Allow Compact block sparse indices: FA4 only accesses indices 0..cnt-1 + # per query tile, so idx.shape[3] can be <= expected_n_blocks. + if mask_block_idx.shape[3] > expected_n_blocks: + raise ValueError( + f"Block sparse tensors{context} n-block dimension must be <= {expected_n_blocks}." + ) + if expected_m_blocks != num_m_blocks: + raise ValueError( + f"Block sparse tensors{context} m-block dimension {num_m_blocks} does not match " + f"sparse_block_size_q={sparse_block_size_q}. " + f"Set BlockSparseTensorsTorch.block_size to match the BlockMask BLOCK_SIZE." + ) + return expected_count_shape, expected_index_shape, q_subtile_factor + + +def get_block_sparse_expected_shapes_bwd( + batch_size: int, + num_head: int, + seqlen_q: int, + seqlen_k: int, + m_block_size: int, + n_block_size: int, + q_subtile_factor: int, +) -> Tuple[Tuple[int, int, int], Tuple[int, int, int, int]]: + """Return (expected_count_shape, expected_index_shape) for backward block sparse normalization. + + Backward uses Q-direction indexing (transposed from forward), where shapes are + indexed by N-blocks first, then M-blocks. The sparse_block_size_q is determined + by q_subtile_factor * m_block_size. + """ + sparse_block_size_q = q_subtile_factor * m_block_size + expected_m_blocks = ceildiv(seqlen_q, sparse_block_size_q) + expected_n_blocks = ceildiv(seqlen_k, n_block_size) + expected_count_shape = (batch_size, num_head, expected_n_blocks) + expected_index_shape = (batch_size, num_head, expected_n_blocks, expected_m_blocks) + return expected_count_shape, expected_index_shape + + +def normalize_block_sparse_tensors( + tensors: BlockSparseTensorsTorch, + *, + expected_count_shape: Tuple[int, ...], + expected_index_shape: Tuple[int, ...], + context: str | None = None, + hint: str | Callable[[], str] | None = None, +) -> BlockSparseTensorsTorch: + if tensors.mask_block_cnt is None or tensors.mask_block_idx is None: + raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.") + + mask_cnt, mask_idx = _check_and_expand_block( + "mask", + tensors.mask_block_cnt, + tensors.mask_block_idx, + expected_count_shape, + expected_index_shape, + context, + hint, + ) + if mask_cnt is None or mask_idx is None: + raise ValueError("mask_block_cnt and mask_block_idx must be provided for block sparsity.") + + full_cnt, full_idx = _check_and_expand_block( + "full", + tensors.full_block_cnt, + tensors.full_block_idx, + expected_count_shape, + expected_index_shape, + context, + hint, + ) + if full_cnt is not None and mask_cnt.device != full_cnt.device: + raise ValueError("All block sparse tensors must be on the same device") + + dq_write_order = _check_and_expand_metadata_tensor( + "dq_write_order", + tensors.dq_write_order, + tuple(mask_idx.shape), + context, + hint, + mask_cnt.device, + ) + dq_write_order_full = _check_and_expand_metadata_tensor( + "dq_write_order_full", + tensors.dq_write_order_full, + tuple(full_idx.shape) if full_idx is not None else expected_index_shape, + context, + hint, + mask_cnt.device, + ) + spt = tensors.spt + if spt is not None and not isinstance(spt, bool): + raise ValueError("spt must be a bool when provided") + if spt is not None and dq_write_order is None: + raise ValueError("spt requires dq_write_order to be provided") + + return BlockSparseTensorsTorch( + mask_block_cnt=mask_cnt, + mask_block_idx=mask_idx, + full_block_cnt=full_cnt, + full_block_idx=full_idx, + cu_total_m_blocks=tensors.cu_total_m_blocks, + cu_block_idx_offsets=tensors.cu_block_idx_offsets, + block_size=tensors.block_size, + dq_write_order=dq_write_order, + dq_write_order_full=dq_write_order_full, + spt=spt, + ) + + +def is_block_sparsity_enabled(tensors: BlockSparseTensorsTorch) -> bool: + return any(t is not None for t in (tensors.full_block_cnt, tensors.mask_block_cnt)) + + +def get_block_sparse_broadcast_pattern( + tensors: BlockSparseTensorsTorch, +) -> Tuple[Tuple[bool, ...], ...] | None: + """Return broadcast pattern for block sparse tensors by checking actual strides. + + Returns a tuple of broadcast patterns (one per tensor) where each pattern + is a tuple of bools indicating which dims have stride=0. + This is used in compile keys to ensure kernels are recompiled when + broadcast patterns change, since CuTe's mark_layout_dynamic() keeps + stride=0 as static. + + The tensors should already be expanded/normalized before calling this function. + + Returns None if block sparsity is not enabled. + """ + if not is_block_sparsity_enabled(tensors): + return None + + patterns = [] + for tensor in ( + tensors.mask_block_cnt, + tensors.mask_block_idx, + tensors.full_block_cnt, + tensors.full_block_idx, + tensors.dq_write_order, + tensors.dq_write_order_full, + ): + if tensor is not None: + patterns.append(get_broadcast_dims(tensor)) + else: + patterns.append(None) + return tuple(patterns) + + +def normalize_block_sparse_config( + tensors: BlockSparseTensorsTorch, + *, + batch_size: int, + num_head: int, + seqlen_q: int, + seqlen_k: int, + block_size: tuple[int, int], + q_stage: int, +) -> tuple[BlockSparseTensorsTorch, Tuple[Tuple[bool, ...], ...] | None, int]: + """Validate the block-sparse config, infer expected shapes, and normalize. + + Handles both fixed-length (3D `[B, H, M]` / 4D `[B, H, M, N]`) and varlen + (2D `[H, total_m_blocks]` / `[H, total_n_blocks]`) layouts. Varlen is + detected by `tensors.cu_total_m_blocks is not None` and forces + `q_subtile_factor == 1` (TODO: potentially remove this restriction). + """ + m_block_size, n_block_size = block_size + if tensors.block_size is None: + sparse_block_size_q, sparse_block_size_kv = None, n_block_size + else: + sparse_block_size_q, sparse_block_size_kv = tensors.block_size + if sparse_block_size_kv != n_block_size: + raise ValueError( + f"Block sparsity requires sparse_block_size[1]={n_block_size} to match tile_n." + ) + if tensors.cu_total_m_blocks is not None: + base_m_block = q_stage * m_block_size + if sparse_block_size_q is not None and sparse_block_size_q != base_m_block: + raise ValueError( + f"Varlen block sparsity requires sparse_block_size[0]={base_m_block} " + f"(= q_stage * tile_m); got {sparse_block_size_q}." + ) + total_m_blocks = tensors.mask_block_cnt.shape[-1] + total_n_blocks = tensors.mask_block_idx.shape[-1] + expected_count_shape = (num_head, total_m_blocks) + expected_index_shape = (num_head, total_n_blocks) + q_subtile_factor = 1 + else: + expected_count_shape, expected_index_shape, q_subtile_factor = ( + infer_block_sparse_expected_shapes( + tensors, + batch_size=batch_size, + num_head=num_head, + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + m_block_size=m_block_size, + n_block_size=n_block_size, + q_stage=q_stage, + context="forward", + sparse_block_size_q=sparse_block_size_q, + sparse_block_size_kv=sparse_block_size_kv, + ) + ) + normalized_tensors = normalize_block_sparse_tensors( + tensors, + expected_count_shape=expected_count_shape, + expected_index_shape=expected_index_shape, + ) + return ( + normalized_tensors, + get_block_sparse_broadcast_pattern(normalized_tensors), + q_subtile_factor, + ) + + +def normalize_block_sparse_config_bwd( + tensors: BlockSparseTensorsTorch, + *, + batch_size: int, + num_head: int, + seqlen_q: int, + seqlen_k: int, + block_size: tuple[int, int], + q_subtile_factor: int, +) -> tuple[BlockSparseTensorsTorch, Tuple[Tuple[bool, ...], ...] | None]: + m_block_size, n_block_size = block_size + if tensors.block_size is None: + sparse_block_size_q, sparse_block_size_kv = q_subtile_factor * m_block_size, n_block_size + else: + sparse_block_size_q, sparse_block_size_kv = tensors.block_size + if sparse_block_size_q != q_subtile_factor * m_block_size: + raise ValueError( + f"Block sparsity expects sparse_block_size_q={q_subtile_factor * m_block_size} " + f"for q_subtile_factor={q_subtile_factor}." + ) + if sparse_block_size_kv != n_block_size: + raise ValueError( + f"Block sparsity expects sparse_block_size[1]={n_block_size} to match tile_n." + ) + expected_count_shape, expected_index_shape = get_block_sparse_expected_shapes_bwd( + batch_size, + num_head, + seqlen_q, + seqlen_k, + m_block_size, + n_block_size, + q_subtile_factor, + ) + normalized_tensors = normalize_block_sparse_tensors( + tensors, + expected_count_shape=expected_count_shape, + expected_index_shape=expected_index_shape, + context="_flash_attn_bwd", + hint=lambda: ( + f"Backward expects Q-direction block-sparse tensors (q_mask_cnt/q_mask_idx, " + f"and optionally full_q_cnt/full_q_idx). Regenerate the backward BlockMask with " + f"BLOCK_SIZE=({q_subtile_factor * m_block_size}, {n_block_size})." + ), + ) + return normalized_tensors, get_block_sparse_broadcast_pattern(normalized_tensors) + + +def to_cute_block_sparse_tensors( + tensors: BlockSparseTensorsTorch, enable_tvm_ffi: bool = True +) -> BlockSparseTensors | None: + """Convert torch block sparsity tensors to CuTe tensors, optionally for tvm ffi""" + if not is_block_sparsity_enabled(tensors): + return None + mask_block_cnt_tensor, mask_block_idx_tensor = [ + to_cute_tensor(t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi) + for t in (tensors.mask_block_cnt, tensors.mask_block_idx) + ] + full_block_cnt_tensor, full_block_idx_tensor = [ + to_cute_tensor(t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi) + if t is not None + else None + for t in (tensors.full_block_cnt, tensors.full_block_idx) + ] + cu_total_m_blocks_tensor, cu_block_idx_offsets_tensor = [ + to_cute_tensor(t, assumed_align=4, leading_dim=0, enable_tvm_ffi=enable_tvm_ffi) + if t is not None + else None + for t in (tensors.cu_total_m_blocks, tensors.cu_block_idx_offsets) + ] + dq_write_order_tensor, dq_write_order_full_tensor = [ + to_cute_tensor(t, assumed_align=4, leading_dim=-1, enable_tvm_ffi=enable_tvm_ffi) + if t is not None + else None + for t in (tensors.dq_write_order, tensors.dq_write_order_full) + ] + + return BlockSparseTensors( + mask_block_cnt_tensor, + mask_block_idx_tensor, + full_block_cnt_tensor, + full_block_idx_tensor, + cu_total_m_blocks_tensor, + cu_block_idx_offsets_tensor, + dq_write_order_tensor, + dq_write_order_full_tensor, + ) + + +def fast_sampling(mask_mod): + """Convenience decorator to mark mask_mod as safe for 5-point fast sampling""" + mask_mod.use_fast_sampling = True + return mask_mod diff --git a/fastvideo-kernel/fa4/cache_utils.py b/fastvideo-kernel/fa4/cache_utils.py new file mode 100644 index 000000000..658a8d5b6 --- /dev/null +++ b/fastvideo-kernel/fa4/cache_utils.py @@ -0,0 +1,281 @@ +# Manage Ahead-of-Time (AOT) compiled kernels +import fcntl +import hashlib +import os +import pickle +import sys +import tempfile +import time +from functools import lru_cache +from getpass import getuser +from pathlib import Path +from typing import Hashable, TypeAlias + +import ctypes + +import cutlass +import cutlass.cute as cute +import tvm_ffi +from cutlass.cutlass_dsl import JitCompiledFunction +from flash_attn.cute.fa_logging import fa_log + +# Pre-load cute DSL runtime libraries with RTLD_GLOBAL so that their symbols +# (e.g. _cudaLibraryLoadData) are visible to .so modules loaded later via dlopen. +# Upstream cute.runtime.load_module loads these without RTLD_GLOBAL, which causes +# "undefined symbol" errors when loading cached kernels from disk. +for _lib_path in cute.runtime.find_runtime_libraries(enable_tvm_ffi=False): + if Path(_lib_path).exists(): + ctypes.CDLL(_lib_path, mode=ctypes.RTLD_GLOBAL) + +CompileKeyType: TypeAlias = tuple[Hashable, ...] +CallableFunction: TypeAlias = JitCompiledFunction | tvm_ffi.Function + +# Enable cache via `FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1` +CUTE_DSL_CACHE_ENABLED: bool = os.getenv("FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED", "0") == "1" + + +# Customize cache dir via `FLASH_ATTENTION_CUTE_DSL_CACHE_DIR`, default is +# `/tmp/${USER}/flash_attention_cute_dsl_cache`` +CUTE_DSL_CACHE_DIR: str | None = os.getenv("FLASH_ATTENTION_CUTE_DSL_CACHE_DIR", None) + + +def get_cache_path() -> Path: + if CUTE_DSL_CACHE_DIR is not None: + cache_dir = Path(CUTE_DSL_CACHE_DIR) + else: + cache_dir = Path(tempfile.gettempdir()) / getuser() / "flash_attention_cute_dsl_cache" + cache_dir.mkdir(parents=True, exist_ok=True) + return cache_dir + + +@lru_cache(maxsize=1) +def _compute_source_fingerprint() -> str: + """ + Hash all CuTe Python sources plus runtime ABI stamps into a short fingerprint. + + The fingerprint changes whenever: + - Any .py file under flash_attn/cute is added, removed, renamed, or modified. + - The Python minor version changes (e.g. 3.13 -> 3.14). + - The cutlass or tvm_ffi package version changes. + + Computed once per process and cached. + """ + cute_root = Path(__file__).resolve().parent + h = hashlib.sha256() + + h.update(f"py{sys.version_info.major}.{sys.version_info.minor}".encode()) + h.update(f"cutlass={cutlass.__version__}".encode()) + h.update(f"tvm_ffi={tvm_ffi.__version__}".encode()) + + for src in sorted(cute_root.rglob("*.py")): + if not src.is_file(): + continue + h.update(src.relative_to(cute_root).as_posix().encode()) + content = src.read_bytes() + h.update(len(content).to_bytes(8, "little")) + h.update(content) + + return h.hexdigest() + + +class FileLock: + """Context manager for advisory file locks using fcntl.flock. + + Supports exclusive (write) and shared (read) locks. + Always blocks with polling until the lock is acquired or timeout is reached. + + Usage: + with FileLock(lock_path, exclusive=True, timeout=15, label="abc"): + # do work under lock + """ + + def __init__( + self, + lock_path: Path, + exclusive: bool, + timeout: float = 15, + label: str = "", + ): + """ + Args: + lock_path: Path to the lock file on disk. + exclusive: True for exclusive (write) lock, False for shared (read) lock. + timeout: Max seconds to wait for lock acquisition before raising RuntimeError. + label: Optional human-readable label for error messages. + """ + self.lock_path: Path = lock_path + self.exclusive: bool = exclusive + self.timeout: float = timeout + self.label: str = label + self._fd: int = -1 + + @property + def _lock_label(self) -> str: + kind = "exclusive" if self.exclusive else "shared" + return f"{kind} {self.label}" if self.label else kind + + def __enter__(self) -> "FileLock": + open_flags = os.O_WRONLY | os.O_CREAT if self.exclusive else os.O_RDONLY | os.O_CREAT + lock_type = fcntl.LOCK_EX if self.exclusive else fcntl.LOCK_SH + + self._fd = os.open(str(self.lock_path), open_flags) + + deadline = time.monotonic() + self.timeout + acquired = False + while time.monotonic() < deadline: + try: + fcntl.flock(self._fd, lock_type | fcntl.LOCK_NB) + acquired = True + break + except OSError: + time.sleep(0.1) + if not acquired: + os.close(self._fd) + self._fd = None + raise RuntimeError( + f"Timed out after {self.timeout}s waiting for " + f"{self._lock_label} lock: {self.lock_path}" + ) + + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + if self._fd is not None: + fcntl.flock(self._fd, fcntl.LOCK_UN) + os.close(self._fd) + self._fd = None + + +class JITCache: + """ + In-memory cache for compiled functions. + """ + + def __init__(self): + self.cache: dict[CompileKeyType, CallableFunction] = {} + + def __setitem__(self, key: CompileKeyType, fn: JitCompiledFunction) -> None: + self.cache[key] = fn + + def __getitem__(self, key: CompileKeyType) -> CallableFunction: + return self.cache[key] + + def __contains__(self, key: CompileKeyType) -> bool: + return key in self.cache + + def clear(self) -> None: + """ + Clear in-memory cache of compiled functions + """ + self.cache.clear() + + +class JITPersistentCache(JITCache): + """ + In-memory cache for compiled functions, which is also backed by persistent storage. + Use cutedsl ahead-of-time (AOT) compilation, only supporting enable_tvm_ffi=True + """ + + EXPORT_FUNCTION_PREFIX = "func" + LOCK_TIMEOUT_SECONDS = 15 + + def __init__(self, cache_path: Path): + super().__init__() + cache_path.mkdir(parents=True, exist_ok=True) + self.cache_path: Path = cache_path + + def __setitem__(self, key: CompileKeyType, fn: JitCompiledFunction) -> None: + JITCache.__setitem__(self, key, fn) + self._try_export_to_storage(key, fn) + + def __getitem__(self, key: CompileKeyType) -> CallableFunction: + # Use __contains__ to try populating in-memory cache with persistent storage + self.__contains__(key) + return JITCache.__getitem__(self, key) + + def __contains__(self, key: CompileKeyType) -> bool: + # Checks in-memory cache first, then tries loading from storage. + # When returning True, guarantees the in-memory cache is populated. + if JITCache.__contains__(self, key): + return True + return self._try_load_from_storage(key) + + def _try_load_from_storage(self, key: CompileKeyType) -> bool: + """ + Try to load a function from persistent storage into in-memory cache. + Returns True if loaded successfully, False if not found on disk. + Holds a shared lock during loading to prevent concurrent writes. + """ + sha256_hex = self._key_to_hash(key) + obj_path = self.cache_path / f"{sha256_hex}.o" + with FileLock( + self._lock_path(sha256_hex), + exclusive=False, + timeout=self.LOCK_TIMEOUT_SECONDS, + label=sha256_hex, + ): + if obj_path.exists(): + fa_log(1, f"Loading compiled function from disk: {obj_path}") + m = cute.runtime.load_module(str(obj_path), enable_tvm_ffi=True) + fn = getattr(m, self.EXPORT_FUNCTION_PREFIX) + JITCache.__setitem__(self, key, fn) + return True + else: + fa_log(1, f"Cache miss on disk for key hash {sha256_hex}") + return False + + def _try_export_to_storage(self, key: CompileKeyType, fn: JitCompiledFunction) -> None: + """Export a compiled function to persistent storage under exclusive lock.""" + sha256_hex = self._key_to_hash(key) + with FileLock( + self._lock_path(sha256_hex), + exclusive=True, + timeout=self.LOCK_TIMEOUT_SECONDS, + label=sha256_hex, + ): + obj_path = self.cache_path / f"{sha256_hex}.o" + if obj_path.exists(): + # Another process already exported. + fa_log(1, f"Skipping export, already on disk: {obj_path}") + return + fa_log(1, f"Exporting compiled function to disk: {obj_path}") + fn.export_to_c( + object_file_path=str(obj_path), + function_name=self.EXPORT_FUNCTION_PREFIX, + ) + fa_log(1, f"Successfully exported compiled function to disk: {obj_path}") + + def _key_to_hash(self, key: CompileKeyType) -> str: + return hashlib.sha256(pickle.dumps(key)).hexdigest() + + def _lock_path(self, sha256_hex: str) -> Path: + return self.cache_path / f"{sha256_hex}.lock" + + def clear(self) -> None: + """ + Not only clear the in-memory cache. Also purge persistent compilation cache. + """ + fa_log(1, f"Clearing persistent cache at {self.cache_path}") + super().clear() + for child in self.cache_path.iterdir(): + child.unlink() + + +def get_jit_cache(name: str | None = None) -> JITCache: + """ + JIT cache factory. + `name` is an optional identifier to create subdirectories to manage cache. + + When persistent caching is enabled, artifacts are namespaced under a + source fingerprint directory so that code or dependency changes + automatically invalidate stale entries. + """ + if CUTE_DSL_CACHE_ENABLED: + path = get_cache_path() / _compute_source_fingerprint() + if name: + path = path / name + fa_log(1, f"Creating persistent JIT cache at {path}") + return JITPersistentCache(path) + else: + fa_log(1, "Persistent cache disabled, using in-memory JIT cache") + return JITCache() diff --git a/fastvideo-kernel/fa4/compute_block_sparsity.py b/fastvideo-kernel/fa4/compute_block_sparsity.py new file mode 100644 index 000000000..c32937b52 --- /dev/null +++ b/fastvideo-kernel/fa4/compute_block_sparsity.py @@ -0,0 +1,551 @@ +from functools import partial +from typing import Callable, Optional, Tuple + +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Boolean, Int8, Int32, const_expr + +from flash_attn.cute.block_sparsity import ( + BlockSparseTensors, + BlockSparseTensorsTorch, + to_cute_block_sparse_tensors, +) +from flash_attn.cute.block_sparse_utils import get_curr_blocksparse_tensors +from flash_attn.cute.testing import is_fake_mode +from flash_attn.cute.cute_dsl_utils import ( + get_aux_tensor_metadata, + to_cute_aux_tensor, + to_cute_tensor, +) +from flash_attn.cute.utils import ( + get_batch_from_cu_tensor, + hash_callable, + scalar_to_ssa, + ssa_to_scalar, +) +from flash_attn.cute.mask import call_mask_mod +from flash_attn.cute.seqlen_info import SeqlenInfoQK +from flash_attn.cute.utils import AuxData + + +class BlockSparsityKernel: + """Block sparsity kernel for FlexAttention. + + This kernel computes `mask_mod` for every token of each block + to determine if an n block is full, masked, or neither. + + Writes block counts and indices to a BlockSparseTensors object. + + When use_fast_sampling=True, uses 5-point sampling (4 corners + center) + which is much faster but only suitable for masks where this is sufficient. + + TODO: + - optimize mask_mod evaluation + - transposed tensors for bwd pass + """ + + def __init__( + self, + mask_mod: Callable, + tile_mn: Tuple[int, int], + compute_full_blocks: bool = True, + use_aux_tensors: bool = False, + use_fast_sampling: bool = False, + ): + self.mask_mod = mask_mod + self.tile_mn = tile_mn + self.compute_full_blocks = compute_full_blocks + self.use_aux_tensors = use_aux_tensors + self.use_fast_sampling = use_fast_sampling + + @cute.jit + def __call__( + self, + blocksparse_tensors: BlockSparseTensors, + seqlen_q: Int32, + seqlen_k: Int32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + aux_data: AuxData = AuxData(), + ): + mask_cnt, mask_idx, full_cnt, full_idx, mCuTotalMBlocks, mCuBlockIdxOffsets, *_ = ( + blocksparse_tensors + ) + + self.is_varlen_q = const_expr(mCuSeqlensQ is not None) + + if const_expr(self.compute_full_blocks): + assert full_cnt is not None and full_idx is not None, ( + "full block tensors must be provided when computing full blocks" + ) + if const_expr(not self.is_varlen_q): + batch_size, num_heads, num_m_blocks, _ = mask_idx.shape + total_m_blocks = batch_size * num_m_blocks + else: + assert const_expr(mCuTotalMBlocks is not None), ( + "mCuTotalMBlocks must be provided when varlen q" + ) + num_heads, total_m_blocks = mask_cnt.shape # num_m_blocks is total_m_blocks + batch_size = mCuSeqlensQ.shape[0] - 1 + + if const_expr(self.use_fast_sampling): + num_threads = 5 + self.num_warps = 1 + else: + num_threads = self.tile_mn[0] + self.num_warps = (num_threads + 32 - 1) // 32 + + if const_expr(not self.is_varlen_q): + grid = [num_m_blocks, num_heads, batch_size] + else: + grid = [total_m_blocks, num_heads, 1] + + self.kernel( + blocksparse_tensors, + seqlen_q, + seqlen_k, + batch_size, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + mCuTotalMBlocks, + mCuBlockIdxOffsets, + aux_data, + ).launch(grid=grid, block=[num_threads, 1, 1]) + + @cute.kernel + def kernel( + self, + blocksparse_tensors: BlockSparseTensors, + seqlen_q: Int32, + seqlen_k: Int32, + batch_size: Int32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + mCuTotalMBlocks: Optional[cute.Tensor] = None, + mCuBlockIdxOffsets: Optional[cute.Tensor] = None, + aux_data: AuxData = AuxData(), + ): + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.warp_idx() + lane_id = cute.arch.lane_idx() + + ssa = partial(scalar_to_ssa, dtype=Int32) + + @cute.struct + class SharedStorage: + reduction_buffer_smem: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int8, 2 * self.num_warps], 1024 + ] + + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage, 16) + + reduction_buffer = storage.reduction_buffer_smem.get_tensor( + cute.make_layout((self.num_warps, 2)) + ) + SeqlenInfoCls = partial( + SeqlenInfoQK.create, + seqlen_q_static=seqlen_q, + seqlen_k_static=seqlen_k, + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + mCuTotalMBlocks=mCuTotalMBlocks, + mCuBlockIdxOffsets=mCuBlockIdxOffsets, + tile_m=self.tile_mn[0], + tile_n=self.tile_mn[1], + ) + + if const_expr(not self.is_varlen_q): + m_block, head_idx, batch_idx = cute.arch.block_idx() + else: + global_m_block, head_idx, _ = cute.arch.block_idx() + batch_idx = get_batch_from_cu_tensor(global_m_block, mCuTotalMBlocks) + m_block = global_m_block - mCuTotalMBlocks[batch_idx] + + seqlen = SeqlenInfoCls(batch_idx) + seqlen_q = seqlen.seqlen_q + seqlen_k = seqlen.seqlen_k + global_m_block = seqlen.m_block_offset + m_block + + num_n_blocks = (seqlen_k + self.tile_mn[1] - 1) // self.tile_mn[1] + + _, curr_mask_idx, _, curr_full_idx = get_curr_blocksparse_tensors( + batch_idx, head_idx, m_block, blocksparse_tensors, seqlen + ) + + num_mask_blocks = Int32(0) + num_full_blocks = Int32(0) + + m_base = m_block * self.tile_mn[0] + if const_expr(self.use_fast_sampling): + # Loop-invariant per-thread q_idx for the 5 sample points + # (tidx 0, 1: top corners; 2, 3: bottom corners; 4: center). + q_idx_sample = m_base + if tidx == 2 or tidx == 3: + q_idx_sample = cutlass.min(m_base + self.tile_mn[0] - 1, seqlen_q - 1) + elif tidx == 4: + q_idx_sample = m_base + cutlass.min(seqlen_q - m_base, self.tile_mn[0]) // 2 + else: + q_idx_thread = m_base + tidx + thread_in_bounds = Boolean(tidx < self.tile_mn[0] and q_idx_thread < seqlen_q) + + for n_block in cutlass.range(num_n_blocks): + n_base = n_block * self.tile_mn[1] + + if const_expr(self.use_fast_sampling): + # 5-point sampling (4 corners + center). Interior n_blocks + # (n_base + tile_n <= seqlen_k) skip the OOB clamp on the right / + # center samples. + is_interior = (n_base + self.tile_mn[1]) <= seqlen_k + n_right = Int32(0) + n_mid = Int32(0) + if is_interior: + n_right = n_base + self.tile_mn[1] - 1 + n_mid = n_base + self.tile_mn[1] // 2 + else: + n_right = cutlass.min(n_base + self.tile_mn[1] - 1, seqlen_k - 1) + n_mid = n_base + cutlass.min(seqlen_k - n_base, self.tile_mn[1]) // 2 + + kv_idx = n_base + if tidx == 1 or tidx == 3: + kv_idx = n_right + elif tidx == 4: + kv_idx = n_mid + + thread_result = Boolean(False) + thread_is_valid = Boolean(False) + if tidx < 5: + thread_is_valid = Boolean(True) + thread_result = ssa_to_scalar( + call_mask_mod( + self.mask_mod, + ssa(batch_idx), + ssa(head_idx), + ssa(q_idx_sample), + ssa(kv_idx), + seqlen, + aux_data, + ) + ) + + has_unmasked = cute.arch.vote_any_sync(thread_result & thread_is_valid) + has_masked = cute.arch.vote_any_sync(Boolean(not thread_result) & thread_is_valid) + + else: + # Full path. Interior blocks (n_base + tile_n <= seqlen_k) drop the + # per-element bound check; the boundary block (at most one) keeps it. + thread_has_unmasked = Boolean(False) + thread_has_masked = Boolean(False) + kv_idx = Int32(0) + is_interior = (n_base + self.tile_mn[1]) <= seqlen_k + + if is_interior: + if thread_in_bounds: + for c in cutlass.range(self.tile_mn[1], unroll_full=True): + mask_val = ssa_to_scalar( + call_mask_mod( + self.mask_mod, + ssa(batch_idx), + ssa(head_idx), + ssa(q_idx_thread), + ssa(n_base + c), + seqlen, + aux_data, + ) + ) + thread_has_unmasked |= Boolean(mask_val) + thread_has_masked |= Boolean(not mask_val) + else: + if thread_in_bounds: + for c in cutlass.range(self.tile_mn[1], unroll_full=True): + kv_idx = n_base + c + if kv_idx < seqlen_k: + mask_val = ssa_to_scalar( + call_mask_mod( + self.mask_mod, + ssa(batch_idx), + ssa(head_idx), + ssa(q_idx_thread), + ssa(kv_idx), + seqlen, + aux_data, + ) + ) + thread_has_unmasked |= Boolean(mask_val) + thread_has_masked |= Boolean(not mask_val) + + warp_unmasked = cute.arch.vote_any_sync(thread_has_unmasked & thread_in_bounds) + warp_masked = cute.arch.vote_any_sync(thread_has_masked & thread_in_bounds) + if lane_id == 0: + reduction_buffer[warp_idx, 0] = Int8(1) if warp_unmasked else Int8(0) + reduction_buffer[warp_idx, 1] = Int8(1) if warp_masked else Int8(0) + cute.arch.sync_threads() + + # Cross-warp OR via warp 0; thread 0 (lane 0 of warp 0) holds the result. + has_unmasked = Boolean(False) + has_masked = Boolean(False) + if warp_idx == 0: + lane_unmasked = Boolean(False) + lane_masked = Boolean(False) + if lane_id < self.num_warps: + lane_unmasked = reduction_buffer[lane_id, 0] != Int8(0) + lane_masked = reduction_buffer[lane_id, 1] != Int8(0) + has_unmasked = cute.arch.vote_any_sync(lane_unmasked) + has_masked = cute.arch.vote_any_sync(lane_masked) + + # Only thread 0 updates the output arrays (common to both paths) + if tidx == 0: + # Block classification based on what we found: + # - If has_masked and has_unmasked: partial block (needs masking) + # - If only has_unmasked: full block (no masking needed) + # - If only has_masked: skip this block entirely + is_partial = Boolean(has_masked and has_unmasked) + is_full = Boolean(has_unmasked and (not has_masked)) + + if is_partial: + curr_mask_idx[num_mask_blocks] = n_block + num_mask_blocks += 1 + elif is_full and const_expr(self.compute_full_blocks): + curr_full_idx[num_full_blocks] = n_block + num_full_blocks += 1 + + # Only thread 0 writes back the counts + if tidx == 0: + mask_cnt, _, full_cnt, *_ = blocksparse_tensors + if const_expr(self.is_varlen_q): + mask_cnt[head_idx, global_m_block] = num_mask_blocks + if const_expr(self.compute_full_blocks): + full_cnt[head_idx, global_m_block] = num_full_blocks + else: + mask_cnt[batch_idx, head_idx, m_block] = num_mask_blocks + if const_expr(self.compute_full_blocks): + full_cnt[batch_idx, head_idx, m_block] = num_full_blocks + + +def compute_block_sparsity( + tile_m, + tile_n, + batch_size, + num_heads, + seqlen_q, + seqlen_k, + mask_mod: Callable, + aux_tensors: Optional[list], + device, + aux_scalars: Optional[tuple] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + cu_total_m_blocks: Optional[torch.Tensor] = None, + cu_block_idx_offsets: Optional[torch.Tensor] = None, + compute_full_blocks: bool = True, + use_fast_sampling: bool = False, +) -> BlockSparseTensorsTorch: + """ + Computes block sparsity for a given `mask_mod`. + + Args: + tile_m: The tile size for the m dimension. + tile_n: The tile size for the n dimension. + batch_size: The batch size. + num_heads: The number of heads. + seqlen_q: The sequence length for the query. + seqlen_k: The sequence length for the key. + mask_mod: The `mask_mod` callable to use. + aux_tensors: A list of auxiliary tensors. + device: The device to use. + cu_seqlens_q: Cumulative q sequence lengths for varlen + cu_seqlens_k: Cumulative k sequence lengths for varlen + seqused_q: Per-batch effective q sequence lengths + seqused_k: Per-batch effective k sequence lengths + cu_total_m_blocks: Cumulative total m blocks tensor for varlen q + cu_block_idx_offsets: Cumulative offsets into the packed mask_block_idx / + full_block_idx tensors per batch (== cumsum of M_b * N_b). + compute_full_blocks: Whether to compute full blocks. If False, only partially-masked blocks are computed. + use_fast_sampling: Whether to use 5-point sampling (4 corners + center). This is much faster, but only suitable for masks where this check is sufficient. + + Returns: + BlockSparseTensorsTorch + """ + aux_scalars = tuple(aux_scalars) if aux_scalars else None + + # Check if mask_mod is marked as suitable for 5-point sampling + use_fast_sampling = getattr(mask_mod, "use_fast_sampling", use_fast_sampling) + + num_m_blocks = (seqlen_q + tile_m - 1) // tile_m + num_n_blocks = (seqlen_k + tile_n - 1) // tile_n + + if cu_seqlens_q is not None: + assert cu_total_m_blocks is not None, "total m blocks must be provided when varlen q" + total_m_blocks = cu_total_m_blocks[-1].item() + if cu_block_idx_offsets is None and (cu_seqlens_k is not None or seqused_k is not None): + # Derive cu_block_idx_offsets from per-batch K seqlens. + cu_block_idx_offsets_list = [0] + for batch_idx in range(batch_size): + batch_seqlen_q = cu_seqlens_q[batch_idx + 1].item() - cu_seqlens_q[batch_idx].item() + if cu_seqlens_k is not None: + batch_seqlen_k = ( + cu_seqlens_k[batch_idx + 1].item() - cu_seqlens_k[batch_idx].item() + ) + else: + batch_seqlen_k = seqused_k[batch_idx].item() + num_m_blocks_batch = (batch_seqlen_q + tile_m - 1) // tile_m + num_n_blocks_batch = (batch_seqlen_k + tile_n - 1) // tile_n + cu_block_idx_offsets_list.append( + cu_block_idx_offsets_list[-1] + num_m_blocks_batch * num_n_blocks_batch + ) + cu_block_idx_offsets = torch.tensor( + cu_block_idx_offsets_list, dtype=torch.int32, device=device + ) + if cu_block_idx_offsets is not None: + total_n_blocks = cu_block_idx_offsets[-1].item() + else: + # Uniform-K varlen-Q: every batch has the same K seqlen. + total_n_blocks = total_m_blocks * num_n_blocks + + mask_block_cnt = torch.zeros((num_heads, total_m_blocks), device=device, dtype=torch.int32) + mask_block_idx = torch.zeros((num_heads, total_n_blocks), device=device, dtype=torch.int32) + full_block_cnt = ( + torch.zeros((num_heads, total_m_blocks), device=device, dtype=torch.int32) + if compute_full_blocks + else None + ) + full_block_idx = ( + torch.zeros((num_heads, total_n_blocks), device=device, dtype=torch.int32) + if compute_full_blocks + else None + ) + else: + total_m_blocks = batch_size * num_m_blocks + total_n_blocks = batch_size * num_m_blocks * num_n_blocks + + mask_block_cnt = torch.zeros( + (batch_size, num_heads, num_m_blocks), device=device, dtype=torch.int32 + ) + mask_block_idx = torch.zeros( + (batch_size, num_heads, num_m_blocks, num_n_blocks), device=device, dtype=torch.int32 + ) + full_block_cnt = ( + torch.zeros((batch_size, num_heads, num_m_blocks), device=device, dtype=torch.int32) + if compute_full_blocks + else None + ) + full_block_idx = ( + torch.zeros( + (batch_size, num_heads, num_m_blocks, num_n_blocks), + device=device, + dtype=torch.int32, + ) + if compute_full_blocks + else None + ) + + blocksparse_tensors_torch = BlockSparseTensorsTorch( + mask_block_cnt=mask_block_cnt, + mask_block_idx=mask_block_idx, + full_block_cnt=full_block_cnt, + full_block_idx=full_block_idx, + cu_total_m_blocks=cu_total_m_blocks, + cu_block_idx_offsets=cu_block_idx_offsets, + block_size=(tile_m, tile_n), + ) + + mask_mod_hash = hash_callable(mask_mod) + if aux_tensors is not None: + aux_tensor_metadata = get_aux_tensor_metadata(aux_tensors) + else: + aux_tensor_metadata = None + aux_scalar_metadata = tuple(type(s) for s in aux_scalars) if aux_scalars is not None else None + + compile_key = ( + tile_m, + tile_n, + mask_mod_hash, + aux_tensor_metadata, + aux_scalar_metadata, + compute_full_blocks, + cu_seqlens_q is None, + cu_seqlens_k is None, + seqused_q is None, + seqused_k is None, + aux_tensors is not None, + use_fast_sampling, + ) + if compile_key not in compute_block_sparsity.compile_cache: + ( + cu_seqlens_q_tensor, + cu_seqlens_k_tensor, + seqused_q_tensor, + seqused_k_tensor, + ) = [ + to_cute_tensor(t, assumed_align=4, leading_dim=0) if t is not None else None + for t in ( + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ) + ] + blocksparse_tensors = to_cute_block_sparse_tensors( + blocksparse_tensors_torch, enable_tvm_ffi=True + ) + if aux_tensors is not None: + cute_aux_tensors = [to_cute_aux_tensor(buf) for buf in aux_tensors] + else: + cute_aux_tensors = None + kernel = BlockSparsityKernel( + mask_mod, + tile_mn=(tile_m, tile_n), + compute_full_blocks=compute_full_blocks, + use_aux_tensors=aux_tensors is not None, + use_fast_sampling=use_fast_sampling, + ) + + compute_block_sparsity.compile_cache[compile_key] = cute.compile( + kernel, + blocksparse_tensors, + seqlen_q, + seqlen_k, + cu_seqlens_q_tensor, + cu_seqlens_k_tensor, + seqused_q_tensor, + seqused_k_tensor, + AuxData(cute_aux_tensors, aux_scalars), + options="--enable-tvm-ffi", + ) + + if not is_fake_mode(): + compute_block_sparsity.compile_cache[compile_key]( + ( + blocksparse_tensors_torch.mask_block_cnt, + blocksparse_tensors_torch.mask_block_idx, + blocksparse_tensors_torch.full_block_cnt, + blocksparse_tensors_torch.full_block_idx, + blocksparse_tensors_torch.cu_total_m_blocks, + blocksparse_tensors_torch.cu_block_idx_offsets, + blocksparse_tensors_torch.dq_write_order, + blocksparse_tensors_torch.dq_write_order_full, + ), + seqlen_q, + seqlen_k, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + AuxData(aux_tensors, aux_scalars), + ) + + return blocksparse_tensors_torch + + +compute_block_sparsity.compile_cache = {} diff --git a/fastvideo-kernel/fa4/copy_utils.py b/fastvideo-kernel/fa4/copy_utils.py new file mode 100644 index 000000000..d8c6083c8 --- /dev/null +++ b/fastvideo-kernel/fa4/copy_utils.py @@ -0,0 +1,372 @@ +# Copyright (c) 2025, Wentao Guo, Ted Zadouri, Tri Dao. + +import math +from typing import Optional, Type, Callable + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, const_expr +from cutlass.cute.nvgpu import cpasync +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass._mlir.dialects import llvm +import cutlass.pipeline + + +@dsl_user_op +def cvt_copy( + atom: cute.CopyAtom, + src: cute.Tensor, + dst: cute.Tensor, + *, + pred: Optional[cute.Tensor] = None, + loc=None, + ip=None, + **kwargs, +) -> None: + assert isinstance(src.iterator, cute.Pointer) and src.memspace == cute.AddressSpace.rmem + if const_expr(src.element_type != dst.element_type): + src_cvt = cute.make_fragment_like(src, dst.element_type, loc=loc, ip=ip) + src_cvt.store(src.load().to(dst.element_type)) + src = src_cvt + cute.copy(atom, src, dst, pred=pred, loc=loc, ip=ip, **kwargs) + + +@dsl_user_op +def load_s2r(src: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor: + dst = cute.make_fragment_like(src, src.element_type, loc=loc, ip=ip) + cute.autovec_copy(src, dst, loc=loc, ip=ip) + return dst + + +@dsl_user_op +def get_copy_atom( + dtype: Type[cutlass.Numeric], num_copy_elems: int, is_async: bool = False, *, loc=None, ip=None +) -> cute.CopyAtom: + num_copy_bits = const_expr(min(128, num_copy_elems * dtype.width)) + copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp() + return cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits) + + +@dsl_user_op +def make_tmem_copy( + tmem_copy_atom: cute.CopyAtom, num_wg: int = 1, *, loc=None, ip=None +) -> cute.CopyAtom: + num_dp, num_bits, num_rep, _ = sm100_utils.get_tmem_copy_properties(tmem_copy_atom) + assert num_dp == 32 + assert num_bits == 32 + tiler_mn = (cute.make_layout((128 * num_rep * num_wg // 32, 32), stride=(32, 1)),) + layout_tv = cute.make_layout( + ((32, 4, num_wg), (num_rep, 32)), stride=((0, 1, 4 * num_rep), (4, 4 * num_rep * num_wg)) + ) + return cute.make_tiled_copy(tmem_copy_atom, layout_tv, tiler_mn) + + +@dsl_user_op +def copy( + src: cute.Tensor, + dst: cute.Tensor, + *, + pred: Optional[cute.Tensor] = None, + num_copy_elems: int = 1, + is_async: bool = False, + loc=None, + ip=None, + **kwargs, +) -> None: + copy_atom = get_copy_atom(src.element_type, num_copy_elems, is_async) + cute.copy(copy_atom, src, dst, pred=pred, loc=loc, ip=ip, **kwargs) + + +def tiled_copy_1d( + dtype: Type[cutlass.Numeric], num_threads: int, num_copy_elems: int = 1, is_async: bool = False +) -> cute.TiledCopy: + num_copy_bits = num_copy_elems * dtype.width + copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp() + copy_atom = cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits) + thr_layout = cute.make_layout(num_threads) + val_layout = cute.make_layout(num_copy_elems) + return cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout) + + +def tiled_copy_2d( + dtype: Type[cutlass.Numeric], major_mode_size: int, num_threads: int, is_async: bool = False +) -> cute.TiledCopy: + num_copy_bits = math.gcd(major_mode_size, 128 // dtype.width) * dtype.width + copy_elems = num_copy_bits // dtype.width + copy_op = cpasync.CopyG2SOp() if is_async else cute.nvgpu.CopyUniversalOp() + copy_atom = cute.make_copy_atom(copy_op, dtype, num_bits_per_copy=num_copy_bits) + gmem_threads_per_row = major_mode_size // copy_elems + assert num_threads % gmem_threads_per_row == 0 + thr_layout = cute.make_ordered_layout( + (num_threads // gmem_threads_per_row, gmem_threads_per_row), + order=(1, 0), + ) + val_layout = cute.make_layout((1, copy_elems)) + return cute.make_tiled_copy_tv(copy_atom, thr_layout, val_layout) + + +@dsl_user_op +def atomic_add_fp32x4( + a: Float32, b: Float32, c: Float32, d: Float32, gmem_ptr: cute.Pointer, *, loc=None, ip=None +) -> None: + gmem_ptr_i64 = gmem_ptr.toint(loc=loc, ip=ip).ir_value() + # cache_hint = cutlass.Int64(0x12F0000000000000) + llvm.inline_asm( + None, + [ + gmem_ptr_i64, + Float32(a).ir_value(loc=loc, ip=ip), + Float32(b).ir_value(loc=loc, ip=ip), + Float32(c).ir_value(loc=loc, ip=ip), + Float32(d).ir_value(loc=loc, ip=ip), + ], + # [gmem_ptr_i64, Float32(a).ir_value(loc=loc, ip=ip), cache_hint.ir_value()], + "{\n\t" + # ".reg .b128 abcd;\n\t" + # "mov.b128 abcd, {$1, $2, $3, $4};\n\t" + ".reg .v4 .f32 abcd;\n\t" + # "mov.b128 abcd, {$1, $2, $3, $4};\n\t" + "mov.f32 abcd.x, $1;\n\t" + "mov.f32 abcd.y, $2;\n\t" + "mov.f32 abcd.z, $3;\n\t" + "mov.f32 abcd.w, $4;\n\t" + "red.global.add.v4.f32 [$0], abcd;\n\t" + # "red.global.add.L2::cache_hint.v4.f32 [$0], abcd, 0x14F0000000000000;\n\t" + "}\n", + # "red.global.add.L2::cache_hint.f32 [$0], $1, 0x12F0000000000000;", + # "red.global.add.L2::cache_hint.f32 [$0], $1, $2;", + "l,f,f,f,f", + # "l,f,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def set_block_rank( + smem_ptr: cute.Pointer, peer_cta_rank_in_cluster: Int32, *, loc=None, ip=None +) -> Int32: + """Map the given smem pointer to the address at another CTA rank in the cluster.""" + smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value() + return Int32( + llvm.inline_asm( + T.i32(), + [smem_ptr_i32, peer_cta_rank_in_cluster.ir_value()], + "mapa.shared::cluster.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def store_shared_remote_fp32x4( + a: Float32, + b: Float32, + c: Float32, + d: Float32, + smem_ptr: cute.Pointer, + mbar_ptr: cute.Pointer, + peer_cta_rank_in_cluster: Int32, + *, + loc=None, + ip=None, +) -> None: + remote_smem_ptr_i32 = set_block_rank( + smem_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip + ).ir_value() + remote_mbar_ptr_i32 = set_block_rank( + mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip + ).ir_value() + llvm.inline_asm( + None, + [ + remote_smem_ptr_i32, + remote_mbar_ptr_i32, + Float32(a).ir_value(loc=loc, ip=ip), + Float32(b).ir_value(loc=loc, ip=ip), + Float32(c).ir_value(loc=loc, ip=ip), + Float32(d).ir_value(loc=loc, ip=ip), + ], + "{\n\t" + ".reg .v4 .f32 abcd;\n\t" + "mov.f32 abcd.x, $2;\n\t" + "mov.f32 abcd.y, $3;\n\t" + "mov.f32 abcd.z, $4;\n\t" + "mov.f32 abcd.w, $5;\n\t" + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.f32 [$0], abcd, [$1];\n\t" + "}\n", + "r,r,f,f,f,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def cpasync_bulk_s2cluster( + smem_src_ptr: cute.Pointer, + smem_dst_ptr: cute.Pointer, + mbar_ptr: cute.Pointer, + size: int | Int32, + peer_cta_rank_in_cluster: Int32, + *, + loc=None, + ip=None, +): + smem_src_ptr_i32 = smem_src_ptr.toint(loc=loc, ip=ip).ir_value() + smem_dst_ptr_i32 = set_block_rank( + smem_dst_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip + ).ir_value() + mbar_ptr_i32 = set_block_rank(mbar_ptr, peer_cta_rank_in_cluster, loc=loc, ip=ip).ir_value() + llvm.inline_asm( + None, + [ + smem_dst_ptr_i32, + smem_src_ptr_i32, + mbar_ptr_i32, + Int32(size).ir_value(loc=loc, ip=ip), + ], + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes [$0], [$1], $3, [$2];", + "r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def cpasync_bulk_g2s( + gmem_ptr: cute.Pointer, + smem_ptr: cute.Pointer, + tma_bar_ptr: cute.Pointer, + size: int | Int32, + *, + loc=None, + ip=None, +): + gmem_ptr_i64 = gmem_ptr.toint(loc=loc, ip=ip).ir_value() + smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value() + mbar_ptr_i32 = tma_bar_ptr.toint(loc=loc, ip=ip).ir_value() + llvm.inline_asm( + None, + [gmem_ptr_i64, smem_ptr_i32, mbar_ptr_i32, Int32(size).ir_value()], + "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [$1], [$0], $3, [$2];", + "l,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def cpasync_reduce_bulk_add_f32( + smem_ptr: cute.Pointer, + gmem_ptr: cute.Pointer, + store_bytes: int | Int32, + *, + loc=None, + ip=None, +): + smem_ptr_i32 = smem_ptr.toint(loc=loc, ip=ip).ir_value() + # cache_hint = cutlass.Int64(0x14F0000000000000) # EVICT_LAST + llvm.inline_asm( + None, + [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value()], + "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 [$0], [$1], $2;", + "l,r,r", + # [gmem_ptr.llvm_ptr, smem_ptr_i32, Int32(store_bytes).ir_value(), cache_hint.ir_value()], + # "cp.reduce.async.bulk.global.shared::cta.bulk_group.L2::cache_hint.add.f32 [$0], [$1], $2, $3;", + # "l,r,r,l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +def cpasync_bulk_get_copy_fn( + src_tensor: cute.Tensor, + dst_tensor: cute.Tensor, + single_stage: bool = False, + **kwargs, +) -> Callable: + # src_is_smem = const_expr( + # isinstance(src_tensor.iterator, cute.Pointer) + # and src_tensor.memspace == cute.AddressSpace.smem + # ) + group_rank_src = const_expr(cute.rank(src_tensor) - (1 if not single_stage else 0)) + group_rank_dst = const_expr(cute.rank(dst_tensor) - (1 if not single_stage else 0)) + # ((atom_v, rest_v), STAGE), ((atom_v, rest_v), RestK) + src = cute.group_modes(src_tensor, 0, group_rank_src) + dst = cute.group_modes(dst_tensor, 0, group_rank_dst) + + def copy_bulk(src_idx, dst_idx, **new_kwargs): + size = const_expr(cute.size(src.shape[:-1]) * src.element_type.width // 8) + cpasync_bulk_g2s( + src[None, src_idx].iterator, + dst[None, dst_idx].iterator, + size=size, + **new_kwargs, + **kwargs, + ) + + def copy_bulk_single_stage(**new_kwargs): + size = const_expr(cute.size(src.shape) * src.element_type.width // 8) + cpasync_bulk_g2s(src.iterator, dst.iterator, size=size, **new_kwargs, **kwargs) + + return copy_bulk if const_expr(not single_stage) else copy_bulk_single_stage + + +def tma_get_copy_fn( + atom: cute.CopyAtom, + cta_coord: cute.Coord, + cta_layout: cute.Layout, + src_tensor: cute.Tensor, + dst_tensor: cute.Tensor, + filter_zeros: bool = False, + single_stage: bool = False, + **kwargs, +) -> Callable: + src_is_smem = const_expr( + isinstance(src_tensor.iterator, cute.Pointer) + and src_tensor.memspace == cute.AddressSpace.smem + ) + smem_tensor, gmem_tensor = (src_tensor, dst_tensor) if src_is_smem else (dst_tensor, src_tensor) + group_rank_smem = const_expr(cute.rank(smem_tensor) - (1 if not single_stage else 0)) + group_rank_gmem = const_expr(cute.rank(gmem_tensor) - (1 if not single_stage else 0)) + # ((atom_v, rest_v), STAGE), ((atom_v, rest_v), RestK) + s, g = cpasync.tma_partition( + atom, + cta_coord, + cta_layout, + cute.group_modes(smem_tensor, 0, group_rank_smem), + cute.group_modes(gmem_tensor, 0, group_rank_gmem), + ) + if const_expr(filter_zeros): + s = cute.filter_zeros(s) + g = cute.filter_zeros(g) + src, dst = (s, g) if src_is_smem else (g, s) + + def copy_tma(src_idx, dst_idx, **new_kwargs): + cute.copy(atom, src[None, src_idx], dst[None, dst_idx], **new_kwargs, **kwargs) + + def copy_tma_single_stage(**new_kwargs): + cute.copy(atom, src, dst, **new_kwargs, **kwargs) + + return (copy_tma if const_expr(not single_stage) else copy_tma_single_stage), s, g + + +def tma_producer_copy_fn(copy: Callable, pipeline: cutlass.pipeline.PipelineAsync): + def copy_fn(src_idx, producer_state: cutlass.pipeline.PipelineState, **new_kwargs): + copy( + src_idx=src_idx, + dst_idx=producer_state.index, + tma_bar_ptr=pipeline.producer_get_barrier(producer_state), + **new_kwargs, + ) + + return copy_fn diff --git a/fastvideo-kernel/fa4/cute_dsl_ptxas.py b/fastvideo-kernel/fa4/cute_dsl_ptxas.py new file mode 100644 index 000000000..4e00f3f00 --- /dev/null +++ b/fastvideo-kernel/fa4/cute_dsl_ptxas.py @@ -0,0 +1,151 @@ +""" +System ptxas replacement for CUTLASS DSL. +Environment variables: + CUTE_DSL_PTXAS_PATH - Path to ptxas (e.g., /usr/local/cuda/bin/ptxas) + CUTE_DSL_PTXAS_VERBOSE - Set to 1 for verbose output +""" + +import os +import sys +import re +import ctypes +import subprocess +from pathlib import Path + +import cutlass + + +CUTE_DSL_PTXAS_PATH = os.environ.get("CUTE_DSL_PTXAS_PATH", None) +VERBOSE = os.environ.get("CUTE_DSL_PTXAS_VERBOSE", "0") == "1" + +_original_load_cuda_library = None +_user_wanted_ptx = False # True if user originally set CUTE_DSL_KEEP_PTX=1 + + +def _log(msg): + if VERBOSE: + print(f"[ptxas] {msg}", file=sys.stderr) + + +def _get_ptx(compiled_func) -> tuple[str, Path] | None: + """Find and read PTX file, stripping null bytes.""" + func_name = getattr(compiled_func, "function_name", None) + if not func_name: + return None + + dump_dir = os.environ.get("CUTE_DSL_DUMP_DIR", Path.cwd()) + for ptx_path in Path(dump_dir).glob(f"*{func_name}*.ptx"): + content = ptx_path.read_text().rstrip("\x00") + if ".entry " in content and content.rstrip().endswith("}"): + _log(f"Found PTX: {ptx_path}") + return content, ptx_path + return None + + +def _compile_ptx(ptx_path: Path, ptx_content: str) -> bytes: + """Compile PTX to cubin using system ptxas.""" + # Extract arch from PTX + match = re.search(r"\.target\s+(sm_\d+[a-z]?)", ptx_content) + arch = match.group(1) if match else "sm_90a" + + # Write stripped content back if needed + if ptx_path.read_text() != ptx_content: + ptx_path.write_text(ptx_content) + + # Compile + cubin_tmp = ptx_path.with_suffix(".cubin.tmp") + try: + assert CUTE_DSL_PTXAS_PATH is not None + result = subprocess.run( + [CUTE_DSL_PTXAS_PATH, f"-arch={arch}", "-O3", "-o", str(cubin_tmp), str(ptx_path)], + capture_output=True, + text=True, + ) + if result.returncode != 0: + raise RuntimeError(f"ptxas failed: {result.stderr}") + + cubin_data = cubin_tmp.read_bytes() + _log(f"Compiled {ptx_path.name} -> {len(cubin_data)} bytes ({arch})") + + # Save cubin if CUTE_DSL_KEEP_CUBIN is set + if os.environ.get("CUTE_DSL_KEEP_CUBIN", "0") == "1": + cubin_out = ptx_path.with_suffix(".cubin") + cubin_out.write_bytes(cubin_data) + _log(f"Saved: {cubin_out}") + + return cubin_data + finally: + cubin_tmp.unlink(missing_ok=True) + + +def _patched_load_cuda_library(self): + """Replacement for _load_cuda_library that uses system ptxas.""" + + result = _get_ptx(self) + if not result: + _log("PTX not found, falling back to embedded ptxas") + return _original_load_cuda_library(self) + + ptx_content, ptx_path = result + + try: + cubin = _compile_ptx(ptx_path, ptx_content) + except Exception as e: + _log(f"Compilation failed ({e}), falling back to embedded ptxas") + return _original_load_cuda_library(self) + + # Load cubin + import cuda.bindings.runtime as cuda_runtime + + err, library = cuda_runtime.cudaLibraryLoadData(cubin, None, None, 0, None, None, 0) + if err != cuda_runtime.cudaError_t.cudaSuccess: + _log(f"cudaLibraryLoadData failed ({err}), falling back to embedded ptxas") + return _original_load_cuda_library(self) + + # Register kernels on all devices + _, cuda_load_to_device = self._get_cuda_init_and_load() + lib_ptr = ctypes.c_void_p(int(library)) + dev_id = ctypes.c_int32(0) + err_val = ctypes.c_int32(0) + args = (ctypes.c_void_p * 3)( + ctypes.cast(ctypes.pointer(lib_ptr), ctypes.c_void_p), + ctypes.cast(ctypes.pointer(dev_id), ctypes.c_void_p), + ctypes.cast(ctypes.pointer(err_val), ctypes.c_void_p), + ) + + for dev in range(self.num_devices): + dev_id.value = dev + cuda_load_to_device(args) + if err_val.value != 0: + _log("cuda_load_to_device failed, falling back to embedded ptxas") + return _original_load_cuda_library(self) + + _log(f"Loaded kernel from {ptx_path.name}") + + # Delete PTX if user didn't originally want it kept + if not _user_wanted_ptx: + ptx_path.unlink(missing_ok=True) + + return [cuda_runtime.cudaLibrary_t(lib_ptr.value)] + + +def patch(): + """Install system ptxas hook. Call before importing cutlass.""" + global _original_load_cuda_library, _user_wanted_ptx + + assert CUTE_DSL_PTXAS_PATH is not None + if not os.path.isfile(CUTE_DSL_PTXAS_PATH) or not os.access(CUTE_DSL_PTXAS_PATH, os.X_OK): + raise RuntimeError(f"ptxas not found: {CUTE_DSL_PTXAS_PATH}") + + # Track if user originally wanted PTX kept + _user_wanted_ptx = os.environ.get("CUTE_DSL_KEEP_PTX", "0") == "1" + # os.environ['CUTE_DSL_KEEP_PTX'] = '1' + assert os.environ.get("CUTE_DSL_KEEP_PTX", "0") == "1", ( + "Require CUTE_DSL_KEEP_PTX=1 to use system's ptxas" + ) + + cls = cutlass.cutlass_dsl.cuda_jit_executor.CudaDialectJitCompiledFunction + _original_load_cuda_library = cls._load_cuda_library + cls._load_cuda_library = _patched_load_cuda_library + _log("Patch applied") + return diff --git a/fastvideo-kernel/fa4/cute_dsl_utils.py b/fastvideo-kernel/fa4/cute_dsl_utils.py new file mode 100644 index 000000000..41976690c --- /dev/null +++ b/fastvideo-kernel/fa4/cute_dsl_utils.py @@ -0,0 +1,158 @@ +# Copyright (c) 2025, Tri Dao. + +from typing import Tuple +from functools import lru_cache + +import torch + +try: + from triton.tools.disasm import extract +except ImportError: + extract = None + +import cutlass +import cutlass.cute as cute +from cutlass.cutlass_dsl import NumericMeta +from cutlass.cute.runtime import from_dlpack + +StaticTypes = (cutlass.Constexpr, NumericMeta, int, bool, str, float, type(None)) + + +load_cubin_module_data_og = cutlass.base_dsl.runtime.cuda.load_cubin_module_data +cute_compile_og = cute.compile + + +torch2cute_dtype_map = { + torch.float16: cutlass.Float16, + torch.bfloat16: cutlass.BFloat16, + torch.float32: cutlass.Float32, + torch.float8_e4m3fn: cutlass.Float8E4M3FN, + torch.float8_e5m2: cutlass.Float8E5M2, +} + + +@lru_cache +def get_max_active_clusters(cluster_size): + return cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_size=cluster_size) + + +@lru_cache +def get_device_capacity(device: torch.device = None) -> Tuple[int, int]: + return torch.cuda.get_device_capability(device) + + +def assume_strides_aligned(t): + """Assume all strides except the last are divisible by 128 bits. + + Python int strides (e.g., stride=0 from GQA expand) are kept as-is + since they're static and don't need alignment assumptions. + """ + divby = 128 // t.element_type.width + strides = tuple(s if isinstance(s, int) else cute.assume(s, divby=divby) for s in t.stride[:-1]) + return (*strides, t.stride[-1]) + + +def assume_tensor_aligned(t): + """Rebuild a tensor with 128-bit aligned stride assumptions. Passes through None.""" + if t is None: + return None + return cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=assume_strides_aligned(t))) + + +def to_cute_tensor(t, assumed_align=16, leading_dim=-1, fully_dynamic=False, enable_tvm_ffi=True): + """Convert torch tensor to cute tensor for TVM FFI. leading_dim=-1 defaults to t.ndim-1.""" + if t is None: + return None + # NOTE: torch 2.9.1 doesn't support fp8 via DLPack but 2.11.0 nightly does + # currently export raw bytes as uint8 and tell cutlass correct type + # can directly export as fp8 when torch supports it + if t.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + tensor = from_dlpack( + t.view(torch.uint8).detach(), + assumed_align=assumed_align, + enable_tvm_ffi=enable_tvm_ffi, + ) + tensor.element_type = ( + cutlass.Float8E4M3FN if t.dtype == torch.float8_e4m3fn else cutlass.Float8E5M2 + ) + else: + tensor = from_dlpack(t.detach(), assumed_align=assumed_align, enable_tvm_ffi=enable_tvm_ffi) + if fully_dynamic: + return tensor.mark_layout_dynamic() + if leading_dim == -1: + leading_dim = t.ndim - 1 + return tensor.mark_layout_dynamic(leading_dim=leading_dim) + + +def to_cute_aux_tensor(t, enable_tvm_ffi=True): + """Convert torch tensor to cute tensor for TVM FFI, tailored to FlexAttention aux tensors. + This allows the user to specify alignment and leading dimension for aux tensors used in + custom score_mod callables. + """ + assumed_align: int = getattr(t, "__assumed_align__", None) + leading_dim: int = getattr(t, "__leading_dim__", None) + fully_dynamic: bool = leading_dim is None + + return to_cute_tensor( + t, + assumed_align=assumed_align, + leading_dim=leading_dim, + fully_dynamic=fully_dynamic, + enable_tvm_ffi=enable_tvm_ffi, + ) + + +def get_aux_tensor_metadata(aux_tensors): + return tuple( + ( + getattr(t, "__assumed_align__", 0), + getattr(t, "__leading_dim__", -1), + hasattr(t, "__leading_dim__"), + ) + for t in aux_tensors + ) + + +def get_broadcast_dims(tensor: torch.Tensor) -> Tuple[bool, ...]: + """Return tuple of bools indicating which dims have stride=0 (broadcast). + + This is useful for compile keys since CuTe's mark_layout_dynamic() keeps + stride=0 as static, meaning kernels compiled with different broadcast + patterns are not interchangeable. + """ + return tuple(s == 0 for s in tensor.stride()) + + +# credit: monellz (https://github.com/NVIDIA/cutlass/issues/2658#issuecomment-3630564264) +def dump_kernel_attributes(compiled_kernel): + from cuda.bindings import driver + from cutlass.utils import HardwareInfo + import torch + + device_id = torch.cuda.current_device() + hardware_info = HardwareInfo(device_id=device_id) + cubin_data = compiled_kernel.artifacts.CUBIN + assert cubin_data is not None, "cubin_data is None, need '--keep-cubin' option when compiling" + cuda_library = hardware_info._checkCudaErrors( + driver.cuLibraryLoadData(cubin_data, None, None, 0, None, None, 0) + ) + kernels = hardware_info._checkCudaErrors(driver.cuLibraryEnumerateKernels(1, cuda_library)) + kernel = hardware_info._checkCudaErrors(driver.cuKernelGetFunction(kernels[0])) + # more metrics: https://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__EXEC.html#group__CUDA__EXEC_1g5e92a1b0d8d1b82cb00dcfb2de15961b + local_size_bytes = hardware_info._checkCudaErrors( + driver.cuFuncGetAttribute( + driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, + kernel, + ) + ) + num_regs = hardware_info._checkCudaErrors( + driver.cuFuncGetAttribute( + driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NUM_REGS, + kernel, + ) + ) + + print("--- Kernel Info ---") + print(f"local_size_bytes: {local_size_bytes}") + print(f"num_regs: {num_regs}") + print("--- End Kernel Info ---") diff --git a/fastvideo-kernel/fa4/fa_logging.py b/fastvideo-kernel/fa4/fa_logging.py new file mode 100644 index 000000000..63189cd5d --- /dev/null +++ b/fastvideo-kernel/fa4/fa_logging.py @@ -0,0 +1,97 @@ +# Copyright (c) 2025, Tri Dao. + +"""Unified FlashAttention logging controlled by a single ``FA_LOG_LEVEL`` env var. + +Host-side messages go through Python ``logging`` (logger name ``flash_attn``). +A default ``StreamHandler`` is attached automatically when ``FA_LOG_LEVEL >= 1`` +so that standalone scripts get output without extra setup; applications that +configure their own logging can remove or replace it via the standard API. + +FA_LOG_LEVEL mapping:: + + 0 off nothing logged + 1 host host-side summaries only (no kernel printf) + 2 kernel host + curated kernel traces + 3 max host + all kernel traces (noisy, perf hit) + +Set via environment variable:: + + FA_LOG_LEVEL=1 python train.py + +Device-side ``cute.printf`` calls are compile-time eliminated via +``cutlass.const_expr`` when the log level is below the callsite threshold, +so there is zero performance cost when device logging is off. +Changing the log level after kernel compilation requires a recompile +(the level participates in the forward compile key). +""" + +import logging +import os +import sys + +import cutlass.cute as cute +from cutlass import const_expr + +_LOG_LEVEL_NAMES = {"off": 0, "host": 1, "kernel": 2, "max": 3} + + +def _parse_log_level(raw: str) -> int: + if raw in _LOG_LEVEL_NAMES: + return _LOG_LEVEL_NAMES[raw] + try: + level = int(raw) + except ValueError: + return 0 + return max(0, min(level, 3)) + + +_fa_log_level: int = _parse_log_level(os.environ.get("FA_LOG_LEVEL", "0")) + +_logger = logging.getLogger("flash_attn") +_logger.addHandler(logging.NullHandler()) +_default_handler: logging.Handler | None = None + + +def _configure_default_handler() -> None: + global _default_handler + if _fa_log_level >= 1: + if _default_handler is None: + _default_handler = logging.StreamHandler(sys.stdout) + _default_handler.setFormatter(logging.Formatter("[FA] %(message)s")) + _logger.addHandler(_default_handler) + _logger.setLevel(logging.DEBUG) + else: + if _default_handler is not None: + _logger.removeHandler(_default_handler) + _default_handler = None + _logger.setLevel(logging.WARNING) + + +_configure_default_handler() + + +def get_fa_log_level() -> int: + return _fa_log_level + + +def set_fa_log_level(level: int | str) -> None: + """Set the FA log level programmatically. + + Host logging takes effect immediately. Device logging changes only + affect kernels compiled after this call (new compile-key selection). + """ + global _fa_log_level + if isinstance(level, str): + level = _parse_log_level(level) + _fa_log_level = max(0, min(int(level), 3)) + _configure_default_handler() + + +def fa_log(level: int, msg: str): + if _fa_log_level >= level: + _logger.info(msg) + + +def fa_printf(level: int, fmt, *args): + if const_expr(_fa_log_level >= level): + cute.printf(fmt, *args) diff --git a/fastvideo-kernel/fa4/fast_math.py b/fastvideo-kernel/fa4/fast_math.py new file mode 100644 index 000000000..c56ea89e7 --- /dev/null +++ b/fastvideo-kernel/fa4/fast_math.py @@ -0,0 +1,21 @@ +# Copyright (c) 2025, Tri Dao. + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 + + +@cute.jit +def clz(x: Int32) -> Int32: + # for i in cutlass.range_constexpr(32): + # if (1 << (31 - i)) & x: + # return Int32(i) + # return Int32(32) + # Early exit is not supported yet + res = Int32(32) + done = False + for i in cutlass.range(32): + if ((1 << (31 - i)) & x) and not done: + res = Int32(i) + done = True + return res diff --git a/fastvideo-kernel/fa4/flash_bwd.py b/fastvideo-kernel/fa4/flash_bwd.py new file mode 100644 index 000000000..dfc56065e --- /dev/null +++ b/fastvideo-kernel/fa4/flash_bwd.py @@ -0,0 +1,1325 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +# A reimplementation of https://github.com/Dao-AILab/flash-attention/blob/main/hopper/mainloop_bwd_sm80.hpp +# from Cutlass C++ to Cute-DSL. +import math +from types import SimpleNamespace +from typing import Type, Callable, Optional +from functools import partial + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync, warp +from cutlass import Int32 +import cutlass.utils as utils_basic + +from quack import layout_utils +from flash_attn.cute import ampere_helpers as sm80_utils +from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned +from flash_attn.cute import utils +from flash_attn.cute.mask import AttentionMask +from flash_attn.cute.softmax import call_score_mod, call_score_mod_bwd +from flash_attn.cute.seqlen_info import SeqlenInfoQK +from quack.cute_dsl_utils import ParamsBase +from flash_attn.cute.tile_scheduler import SingleTileScheduler, SingleTileVarlenScheduler, TileSchedulerArguments +from flash_attn.cute.block_sparsity import BlockSparseTensors +from flash_attn.cute.utils import AuxData + + +class FlashAttentionBackwardSm80: + def __init__( + self, + dtype: Type[cutlass.Numeric], + head_dim: int, + head_dim_v: Optional[int] = None, + qhead_per_kvhead: int = 1, + m_block_size: int = 64, + n_block_size: int = 128, + num_stages_Q: int = 2, + num_stages_dO: int = 2, + num_threads: int = 256, + pack_gqa: bool = False, + is_causal: bool = False, + SdP_swapAB: bool = False, + dKV_swapAB: bool = False, + dQ_swapAB: bool = False, + AtomLayoutMSdP: int = 1, + AtomLayoutNdKV: int = 8, + AtomLayoutMdQ: int = 1, + V_in_regs: bool = False, + score_mod: cutlass.Constexpr | None = None, + score_mod_bwd: cutlass.Constexpr | None = None, + ): + """Initializes the configuration for a flash attention v2 kernel. + + All contiguous dimensions must be at least 16 bytes aligned which indicates the head dimension + should be a multiple of 8. + + :param head_dim: head dimension + :type head_dim: int + :param m_block_size: m block size + :type m_block_size: int + :param n_block_size: n block size + :type n_block_size: int + :param num_threads: number of threads + :type num_threads: int + :param is_causal: is causal + """ + self.dtype = dtype + # padding head_dim to a multiple of 32 (stricter than fwd's 16) due to + # backward kernel register layout requirements for dQ/dK/dV accumulation + hdim_multiple_of = 32 + self.head_dim_padded = int(math.ceil(head_dim / hdim_multiple_of) * hdim_multiple_of) + head_dim_v = head_dim_v if head_dim_v is not None else head_dim + self.same_hdim_kv = head_dim == head_dim_v + self.head_dim_v_padded = int(math.ceil(head_dim_v / hdim_multiple_of) * hdim_multiple_of) + # Can save registers (and hence be faster) if we don't have to check hdim predication + self.check_hdim_oob = head_dim != self.head_dim_padded + self.check_hdim_v_oob = head_dim_v != self.head_dim_v_padded + self.qhead_per_kvhead = qhead_per_kvhead + self.m_block_size = m_block_size + self.n_block_size = n_block_size + self.num_threads = num_threads + self.pack_gqa = pack_gqa + self.is_causal = is_causal + self.num_stages_Q = num_stages_Q + self.num_stages_dO = num_stages_dO + self.SdP_swapAB = SdP_swapAB + self.dKV_swapAB = dKV_swapAB + self.dQ_swapAB = dQ_swapAB + self.AtomLayoutMSdP = AtomLayoutMSdP + self.AtomLayoutNdKV = AtomLayoutNdKV + self.AtomLayoutMdQ = AtomLayoutMdQ + num_mma_warps = self.num_threads // cute.arch.WARP_SIZE + self.Mma_dKV_is_RS = AtomLayoutMSdP == 1 and AtomLayoutNdKV == num_mma_warps and SdP_swapAB and not dKV_swapAB + self.V_in_regs = V_in_regs + self.share_QV_smem = V_in_regs + self.score_mod = score_mod + self.score_mod_bwd = score_mod_bwd + + @staticmethod + def can_implement( + dtype, head_dim, head_dim_v, m_block_size, n_block_size, num_stages_Q, num_stages_dO, + num_threads, is_causal, + V_in_regs=False + ) -> bool: + """Check if the kernel can be implemented with the given parameters. + + :param dtype: data type + :type dtype: cutlass.Numeric + :param head_dim: head dimension + :type head_dim: int + :param m_block_size: m block size + :type m_block_size: int + :param n_block_size: n block size + :type n_block_size: int + :param num_threads: number of threads + :type num_threads: int + :param is_causal: is causal + :type is_causal: bool + + :return: True if the kernel can be implemented, False otherwise + :rtype: bool + """ + if dtype not in [cutlass.Float16, cutlass.BFloat16]: + return False + if head_dim % 8 != 0: + return False + if head_dim_v % 8 != 0: + return False + if n_block_size % 16 != 0: + return False + if num_threads % 32 != 0: + return False + # Check if block size setting is out of shared memory capacity + # Shared memory usage: Q tile + (K tile + V tile) where K and V use the same tile size + smem_usage_Q = m_block_size * head_dim * num_stages_Q * 2 + smem_usage_dO = m_block_size * head_dim_v * num_stages_dO * 2 + smem_usage_K = n_block_size * head_dim * 2 + smem_usage_V = n_block_size * head_dim_v * 2 + smem_usage_QV = (smem_usage_Q + smem_usage_V) if not V_in_regs else max(smem_usage_Q, smem_usage_V) + smem_usage = smem_usage_QV + smem_usage_dO + smem_usage_K + smem_capacity = utils_basic.get_smem_capacity_in_bytes("sm_80") + if smem_usage > smem_capacity: + return False + return True + + def _check_type( + self, + mQ_type: Type[cutlass.Numeric], + mK_type: Type[cutlass.Numeric], + mV_type: Type[cutlass.Numeric], + mdO_type: Type[cutlass.Numeric], + mLSE_type: Type[cutlass.Numeric], + mdPsum_type: Type[cutlass.Numeric], + mdQaccum_type: Type[cutlass.Numeric], + mdK_type: Type[cutlass.Numeric], + mdV_type: Type[cutlass.Numeric], + mCuSeqlensQ_type: Type[cutlass.Numeric] | None, + mCuSeqlensK_type: Type[cutlass.Numeric] | None, + mSeqUsedQ_type: Type[cutlass.Numeric] | None, + mSeqUsedK_type: Type[cutlass.Numeric] | None, + ): + if cutlass.const_expr(not (mQ_type == mK_type == mV_type == mdO_type)): + raise TypeError("All tensors must have the same data type") + if cutlass.const_expr(self.qhead_per_kvhead == 1): + if cutlass.const_expr(not (mdK_type == mdV_type == mQ_type)): + raise TypeError("mdK and mdV tensors must have the same data type as mQ") + else: + if cutlass.const_expr(not (mdK_type == mdV_type == cutlass.Float32)): + raise TypeError("mdKaccum and mdVaccum tensors must have the data type Float32") + if cutlass.const_expr(mQ_type not in [cutlass.Float16, cutlass.BFloat16]): + raise TypeError("Only Float16 or BFloat16 is supported") + if cutlass.const_expr(mLSE_type not in [cutlass.Float32]): + raise TypeError("LSE tensor must be Float32") + if cutlass.const_expr(mdPsum_type not in [cutlass.Float32]): + raise TypeError("dPsum tensor must be Float32") + if cutlass.const_expr(mdQaccum_type not in [cutlass.Float32]): + raise TypeError("dQaccum tensor must be Float32") + if cutlass.const_expr(mCuSeqlensQ_type not in [None, cutlass.Int32]): + raise TypeError("cuSeqlensQ tensor must be Int32") + if cutlass.const_expr(mCuSeqlensK_type not in [None, cutlass.Int32]): + raise TypeError("cuSeqlensK tensor must be Int32") + if cutlass.const_expr(mSeqUsedQ_type not in [None, cutlass.Int32]): + raise TypeError("SeqUsedQ tensor must be Int32") + if cutlass.const_expr(mSeqUsedK_type not in [None, cutlass.Int32]): + raise TypeError("SeqUsedK tensor must be Int32") + assert mQ_type == self.dtype + + def _setup_attributes(self): + # /////////////////////////////////////////////////////////////////////////////// + # Shared memory layout: Q/K/V + # /////////////////////////////////////////////////////////////////////////////// + sQ_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, self.head_dim_padded) + self.sQ_layout = cute.tile_to_shape( + sQ_layout_atom, (self.m_block_size, self.head_dim_padded, self.num_stages_Q), (0, 1, 2), + ) + sK_layout_atom = sQ_layout_atom + self.sK_layout = cute.tile_to_shape( + sK_layout_atom, (self.n_block_size, self.head_dim_padded), (0, 1), + ) + sV_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, self.head_dim_v_padded) + self.sV_layout = cute.tile_to_shape( + sV_layout_atom, (self.n_block_size, self.head_dim_v_padded), (0, 1), + ) + sdO_layout_atom = sV_layout_atom + self.sdO_layout = cute.tile_to_shape( + sdO_layout_atom, (self.m_block_size, self.head_dim_v_padded, self.num_stages_dO), (0, 1, 2), + ) + # TODO: do we set swizzle to be 3 here explicitly? + sPdS_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, self.n_block_size) + self.sPdS_layout = cute.tile_to_shape( + sPdS_layout_atom, (self.m_block_size, self.n_block_size), (0, 1), + ) + # We set stride to be multiple of 64 so that if ShuffleLSE, even if threads read from sLSE but out of bounds, + # it's still a valid smem address. + self.sLSE_layout = cute.make_layout( + (self.m_block_size, self.num_stages_Q), + stride=(1, cute.round_up(self.m_block_size, 64)), + ) + sLSEMma_layout = cute.make_layout( + (self.m_block_size, self.n_block_size, self.num_stages_Q), + stride=(1, 0, cute.round_up(self.m_block_size, 64)), + ) + sLSEMma_layout_transposed = cute.make_layout( + (self.n_block_size, self.m_block_size, self.num_stages_Q), + stride=(0, 1, cute.round_up(self.m_block_size, 64)), + ) + self.sLSEMma_layout = sLSEMma_layout if not self.SdP_swapAB else sLSEMma_layout_transposed + + # /////////////////////////////////////////////////////////////////////////////// + # GMEM Tiled copy: + # /////////////////////////////////////////////////////////////////////////////// + # Thread layouts for copies + universal_copy_bits = 128 + async_copy_elems = universal_copy_bits // self.dtype.width + # atom_async_copy: async copy atom for QKV load + atom_async_copy = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + self.dtype, + num_bits_per_copy=universal_copy_bits, + ) + # atom_universal_copy: universal copy atom for O store + atom_universal_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.dtype, num_bits_per_copy=universal_copy_bits, + ) + # tQK_layout: thread layout for QK load + tQK_shape_dim_1 = sQ_layout_atom.outer.shape[1] // async_copy_elems + assert self.num_threads % tQK_shape_dim_1 == 0, "num_threads must be divisible by tQK_shape_dim_1" + tQK_layout = cute.make_ordered_layout( + (self.num_threads // tQK_shape_dim_1, tQK_shape_dim_1), order=(1, 0), + ) + # Do we need to check if we overshot kBlockM when we load Q? + self.is_even_m_smem_q = self.m_block_size % tQK_layout.shape[0] == 0 + # Do we need to check if we overshot kBlockN when we load K? + self.is_even_n_smem_k = self.n_block_size % tQK_layout.shape[0] == 0 + tVdO_shape_dim_1 = sV_layout_atom.outer.shape[1] // async_copy_elems + assert self.num_threads % tVdO_shape_dim_1 == 0, "num_threads must be divisible by tVdO_shape_dim_1" + tVdO_layout = cute.make_ordered_layout( + (self.num_threads // tVdO_shape_dim_1, tVdO_shape_dim_1), order=(1, 0), + ) + # Do we need to check if we overshot kBlockN when we load V? + self.is_even_n_smem_v = self.n_block_size % tVdO_layout.shape[0] == 0 + self.is_even_m_smem_do = self.m_block_size % tVdO_layout.shape[0] == 0 + + # Value layouts for copies + vQKVdO_layout = cute.make_layout((1, async_copy_elems)) + + # gmem_tiled_copy_QK: tiled copy for QK load + self.gmem_tiled_copy_QK = cute.make_tiled_copy_tv(atom_async_copy, tQK_layout, vQKVdO_layout) + self.gmem_tiled_copy_VdO = cute.make_tiled_copy_tv(atom_async_copy, tVdO_layout, vQKVdO_layout) + self.gmem_tiled_copy_dK = cute.make_tiled_copy_tv(atom_universal_copy, tQK_layout, vQKVdO_layout) + self.gmem_tiled_copy_dV = cute.make_tiled_copy_tv(atom_universal_copy, tVdO_layout, vQKVdO_layout) + async_copy_elems_accum = universal_copy_bits // cutlass.Float32.width + + # I think we wouldn't require this with smarter padding + if cutlass.const_expr(not self.varlen_q): + async_copy_elems_accum = universal_copy_bits // cutlass.Float32.width + atom_async_copy_accum = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + cutlass.Float32, + num_bits_per_copy=universal_copy_bits, + ) + else: + async_copy_elems_accum = 1 + atom_async_copy_accum = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + cutlass.Float32, + num_bits_per_copy=cutlass.Float32.width, + ) + self.gmem_tiled_copy_LSE = cute.make_tiled_copy_tv( + atom_async_copy_accum, + cute.make_layout(self.num_threads), + cute.make_layout(async_copy_elems_accum), + ) + self.gmem_tiled_copy_dQaccum = cute.make_tiled_copy_tv( + cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), cutlass.Float32, num_bits_per_copy=cutlass.Float32.width + ), + cute.make_layout(self.num_threads), + cute.make_layout(1) + ) + if cutlass.const_expr(self.qhead_per_kvhead > 1): + self.gmem_tiled_copy_dK = self.gmem_tiled_copy_dQaccum + self.gmem_tiled_copy_dV = self.gmem_tiled_copy_dQaccum + + def _get_tiled_mma(self): + num_mma_warps = self.num_threads // 32 + AtomLayoutSdP = (self.AtomLayoutMSdP, num_mma_warps // self.AtomLayoutMSdP, 1) if cutlass.const_expr(not self.SdP_swapAB) else (num_mma_warps // self.AtomLayoutMSdP, self.AtomLayoutMSdP, 1) + tiled_mma_sdp = cute.make_tiled_mma( + warp.MmaF16BF16Op(self.dtype, cutlass.Float32, (16, 8, 16)), + AtomLayoutSdP, + permutation_mnk=(AtomLayoutSdP[0] * 16, AtomLayoutSdP[1] * 16, 16), + ) + AtomLayoutdKV = (self.AtomLayoutNdKV, num_mma_warps // self.AtomLayoutNdKV, 1) if cutlass.const_expr(not self.dKV_swapAB) else (num_mma_warps // self.AtomLayoutNdKV, self.AtomLayoutNdKV, 1) + tiled_mma_dkv = cute.make_tiled_mma( + warp.MmaF16BF16Op(self.dtype, cutlass.Float32, (16, 8, 16)), + AtomLayoutdKV, + permutation_mnk=(AtomLayoutdKV[0] * 16, AtomLayoutdKV[1] * 16, 16), + ) + AtomLayoutdQ = (self.AtomLayoutMdQ, num_mma_warps // self.AtomLayoutMdQ, 1) if cutlass.const_expr(not self.dQ_swapAB) else (num_mma_warps // self.AtomLayoutMdQ, self.AtomLayoutMdQ, 1) + tiled_mma_dq = cute.make_tiled_mma( + warp.MmaF16BF16Op(self.dtype, cutlass.Float32, (16, 8, 16)), + AtomLayoutdQ, + permutation_mnk=(AtomLayoutdQ[0] * 16, AtomLayoutdQ[1] * 16, 16), + ) + return tiled_mma_sdp, tiled_mma_dkv, tiled_mma_dq + + def _get_shared_storage_cls(self): + sQ_struct, sK_struct, sV_struct, sdO_struct = [ + cute.struct.Align[cute.struct.MemRange[self.dtype, cute.cosize(layout)], 1024] + for layout in (self.sQ_layout, self.sK_layout, self.sV_layout, self.sdO_layout) + ] + cosize_sQV = max(cute.cosize(self.sQ_layout), cute.cosize(self.sV_layout)) + sQV_struct = cute.struct.Align[cute.struct.MemRange[self.dtype, cosize_sQV], 1024] + sLSE_struct, sdPsum_struct = [ + cute.struct.Align[cute.struct.MemRange[cutlass.Float32, cute.cosize(layout)], 128] + for layout in (self.sLSE_layout, self.sLSE_layout) + ] + sP_struct, sdS_struct = [ + cute.struct.Align[cute.struct.MemRange[self.dtype, cute.cosize(layout)], 128] + for layout in (self.sPdS_layout, self.sPdS_layout) + ] + + @cute.struct + class SharedStorageSeparateQV: + sK: sK_struct + sV: sV_struct + sQ: sQ_struct + sdO: sdO_struct + sLSE: sLSE_struct + sdPsum: sdPsum_struct + sP: sP_struct + sdS: sdS_struct + # TODO: the case where there's no sP + + @cute.struct + class SharedStorageSharedQV: + sK: sK_struct + sV: sV_struct + sQ: sQV_struct + sdO: sdO_struct + sLSE: sLSE_struct + sdPsum: sdPsum_struct + sP: sP_struct + sdS: sdS_struct + + return SharedStorageSeparateQV if cutlass.const_expr(not self.share_QV_smem) else SharedStorageSharedQV + + @cute.jit + def __call__( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mdO: cute.Tensor, + mLSE: cute.Tensor, + mdPsum: cute.Tensor, + mdQaccum: cute.Tensor, + mdK: cute.Tensor, + mdV: cute.Tensor, + softmax_scale: cutlass.Float32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + window_size_left: Int32 | int | None = None, + window_size_right: Int32 | int | None = None, + mdQ_semaphore: Optional[cute.Tensor] = None, + mdK_semaphore: Optional[cute.Tensor] = None, + mdV_semaphore: Optional[cute.Tensor] = None, + aux_data: AuxData = AuxData(), + blocksparse_tensors: Optional[BlockSparseTensors] = None, + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + assert mdQ_semaphore is None and mdK_semaphore is None and mdV_semaphore is None, ( + "determinism not supported yet for Sm80" + ) + # Get the data type and check if it is fp16 or bf16 + self._check_type(*(t.element_type if t is not None else None + for t in (mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV, mCuSeqlensQ, mCuSeqlensK, mSeqUsedQ, mSeqUsedK))) + mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV = [ + assume_tensor_aligned(t) for t in (mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV) + ] + self.varlen_q = (mCuSeqlensQ is not None) + self._setup_attributes() + SharedStorage = self._get_shared_storage_cls() + tiled_mma_sdp, tiled_mma_dkv, tiled_mma_dq = self._get_tiled_mma() + + num_head = mQ.shape[1] if cutlass.const_expr(mCuSeqlensQ is not None) else mQ.shape[2] + + if cutlass.const_expr(mCuSeqlensK is not None): + TileScheduler = SingleTileVarlenScheduler + num_batch = mCuSeqlensK.shape[0] - 1 + else: + TileScheduler = SingleTileScheduler + num_batch = mK.shape[0] + + # Uses seqlen k, etc. since main bwd kernel's blocks are over n + tile_sched_args = TileSchedulerArguments( + num_block=cute.ceil_div(mK.shape[1], self.n_block_size), + num_head=num_head, + num_batch=num_batch, + num_splits=1, + seqlen_k=0, + headdim=mK.shape[2], + headdim_v=mV.shape[2], + total_q=mK.shape[0], + tile_shape_mn=(self.n_block_size, self.m_block_size), + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if cutlass.const_expr(self.pack_gqa) else 1, + mCuSeqlensQ=mCuSeqlensK, + mSeqUsedQ=mSeqUsedK, + ) + + tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args) + grid_dim = TileScheduler.get_grid_shape(tile_sched_params) + + softmax_scale_log2, softmax_scale = utils.compute_softmax_scale_log2(softmax_scale, self.score_mod) + self.kernel( + mQ, + mK, + mV, + mdO, + mLSE, + mdPsum, + mdQaccum, + mdK, + mdV, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + softmax_scale, + softmax_scale_log2, + self.sQ_layout, + self.sK_layout, + self.sV_layout, + self.sdO_layout, + self.sPdS_layout, + self.sLSE_layout, + self.sLSEMma_layout, + self.gmem_tiled_copy_QK, + self.gmem_tiled_copy_VdO, + self.gmem_tiled_copy_dK, + self.gmem_tiled_copy_dV, + self.gmem_tiled_copy_LSE, + self.gmem_tiled_copy_dQaccum, + tiled_mma_sdp, + tiled_mma_dkv, + tiled_mma_dq, + SharedStorage, + tile_sched_params, + TileScheduler, + aux_data, + ).launch( + grid=grid_dim, + block=[self.num_threads, 1, 1], + smem=SharedStorage.size_in_bytes(), + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mdO: cute.Tensor, + mLSE: cute.Tensor, + mdPsum: cute.Tensor, + mdQaccum: cute.Tensor, + mdK: cute.Tensor, + mdV: cute.Tensor, + mCuSeqlensQ: Optional[cute.Tensor], + mCuSeqlensK: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + softmax_scale: cutlass.Float32, + softmax_scale_log2: cutlass.Float32, + sQ_layout: cute.ComposedLayout, + sK_layout: cute.ComposedLayout, + sV_layout: cute.ComposedLayout, + sdO_layout: cute.ComposedLayout, + sPdS_layout: cute.ComposedLayout, + sLSE_layout: cute.Layout, + sLSEMma_layout: cute.Layout, + gmem_tiled_copy_QK: cute.TiledCopy, + gmem_tiled_copy_VdO: cute.TiledCopy, + gmem_tiled_copy_dK: cute.TiledCopy, + gmem_tiled_copy_dV: cute.TiledCopy, + gmem_tiled_copy_LSE: cute.TiledCopy, + gmem_tiled_copy_dQaccum: cute.TiledCopy, + tiled_mma_sdp: cute.TiledMma, + tiled_mma_dkv: cute.TiledMma, + tiled_mma_dq: cute.TiledMma, + SharedStorage: cutlass.Constexpr, + tile_sched_params: ParamsBase, + TileScheduler: cutlass.Constexpr[Callable], + aux_data: AuxData = AuxData(), + ): + # Thread index, block index + tidx, _, _ = cute.arch.thread_idx() + + tile_scheduler = TileScheduler.create(tile_sched_params) + work_tile = tile_scheduler.initial_work_tile_info() + + n_block, head_idx, batch_idx, _ = work_tile.tile_idx + + if work_tile.is_valid_tile: + seqlen = SeqlenInfoQK.create( + batch_idx, + mQ.shape[1], + mK.shape[1], + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + tile_m=self.m_block_size, + tile_n=self.n_block_size, + ) + + m_block_max = cute.ceil_div(seqlen.seqlen_q, self.m_block_size) + m_block_min = 0 + if cutlass.const_expr(self.is_causal): + m_block_min = max( + (n_block * self.n_block_size + seqlen.seqlen_q - seqlen.seqlen_k) // self.m_block_size, + m_block_min, + ) + # TODO: return early if m_block_max == 0 + + # /////////////////////////////////////////////////////////////////////////////// + # Get the appropriate tiles for this thread block. + # /////////////////////////////////////////////////////////////////////////////// + blkQ_shape = (self.m_block_size, self.head_dim_padded) + blkK_shape = (self.n_block_size, self.head_dim_padded) + blkV_shape = (self.n_block_size, self.head_dim_v_padded) + blkdO_shape = (self.m_block_size, self.head_dim_v_padded) + + if cutlass.const_expr(not seqlen.has_cu_seqlens_q): + mQ_cur = mQ[batch_idx, None, head_idx, None] + mLSE_cur = mLSE[batch_idx, head_idx, None] + mdO_cur = mdO[batch_idx, None, head_idx, None] + mdPsum_cur = mdPsum[batch_idx, head_idx, None] + mdQaccum_cur = mdQaccum[batch_idx, head_idx, None] + else: + padded_offset_q = seqlen.padded_offset_q + mQ_cur = cute.domain_offset((seqlen.offset_q, 0), mQ[None, head_idx, None]) + mLSE_cur = cute.domain_offset((padded_offset_q,), mLSE[head_idx, None]) + mdO_cur = cute.domain_offset((seqlen.offset_q, 0), mdO[None, head_idx, None]) + mdPsum_cur = cute.domain_offset((padded_offset_q,), mdPsum[head_idx, None]) + mdQaccum_cur = cute.domain_offset((padded_offset_q * self.head_dim_padded,), mdQaccum[head_idx, None]) + head_idx_kv = head_idx // self.qhead_per_kvhead if cutlass.const_expr(not self.pack_gqa) else head_idx + + if cutlass.const_expr(not seqlen.has_cu_seqlens_k): + mK_cur, mV_cur = [t[batch_idx, None, head_idx_kv, None] for t in (mK, mV)] + else: + mK_cur, mV_cur = [cute.domain_offset((seqlen.offset_k, 0), t[None, head_idx_kv, None]) for t in (mK, mV)] + + # (m_block_size, head_dim, m_block) + gQ = cute.local_tile(mQ_cur, blkQ_shape, (None, 0)) + # (n_block_size, head_dim) + gK = cute.local_tile(mK_cur, blkK_shape, (n_block, 0)) + # (n_block_size, head_dim_v) + gV = cute.local_tile(mV_cur, blkV_shape, (n_block, 0)) + # (m_block_size, head_dim_v, m_block) + gdO = cute.local_tile(mdO_cur, blkdO_shape, (None, 0)) + gLSE = cute.local_tile(mLSE_cur, (self.m_block_size,), (None,)) + gdPsum = cute.local_tile(mdPsum_cur, (self.m_block_size,), (None,)) + gdQaccum = cute.local_tile(mdQaccum_cur, (self.m_block_size * self.head_dim_padded,), (None,)) + + # /////////////////////////////////////////////////////////////////////////////// + # Get shared memory buffer + # /////////////////////////////////////////////////////////////////////////////// + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + sQ = storage.sQ.get_tensor(sQ_layout) + sK = storage.sK.get_tensor(sK_layout) + if cutlass.const_expr(not self.share_QV_smem): + sV = storage.sV.get_tensor(sV_layout) + else: + sV = cute.make_tensor(cute.recast_ptr(sQ.iterator, dtype=self.dtype), sV_layout) + sdO = storage.sdO.get_tensor(sdO_layout) + sP = storage.sP.get_tensor(sPdS_layout) + sdS = storage.sdS.get_tensor(sPdS_layout) + sLSE = storage.sLSE.get_tensor(sLSE_layout) + sdPsum = storage.sdPsum.get_tensor(sLSE_layout) + sLSEMma = storage.sLSE.get_tensor(sLSEMma_layout) + sdPsumMma = storage.sdPsum.get_tensor(sLSEMma_layout) + + # Transpose view of tensors for tiled mma + sQt, sdOt, sKt, sPt, sdSt = [layout_utils.transpose_view(t) for t in (sQ, sdO, sK, sP, sdS)] + + gmem_thr_copy_QK = gmem_tiled_copy_QK.get_slice(tidx) + gmem_thr_copy_VdO = gmem_tiled_copy_VdO.get_slice(tidx) + gmem_thr_copy_lse = gmem_tiled_copy_LSE.get_slice(tidx) + gmem_thr_copy_dQaccum = gmem_tiled_copy_dQaccum.get_slice(tidx) + # (CPY_Atom, CPY_M, CPY_K, m_block) + tQgQ = gmem_thr_copy_QK.partition_S(gQ) + tQsQ = gmem_thr_copy_QK.partition_D(sQ) + # (CPY_Atom, CPY_N, CPY_K) + tKgK = gmem_thr_copy_QK.partition_S(gK) + tKsK = gmem_thr_copy_QK.partition_D(sK) + # (CPY_Atom, CPY_N, CPY_K) + tVgV = gmem_thr_copy_VdO.partition_S(gV) + tVsV = gmem_thr_copy_VdO.partition_D(sV) + # (CPY_Atom, CPY_M, CPY_K, m_block) + tdOgdO = gmem_thr_copy_VdO.partition_S(gdO) + tdOsdO = gmem_thr_copy_VdO.partition_D(sdO) + tLSEgLSE = gmem_thr_copy_lse.partition_S(gLSE) + tLSEsLSE = gmem_thr_copy_lse.partition_D(sLSE) + tLSEgdPsum = gmem_thr_copy_lse.partition_S(gdPsum) + tLSEsdPsum = gmem_thr_copy_lse.partition_D(sdPsum) + tdQgdQaccum = gmem_thr_copy_dQaccum.partition_S(gdQaccum) + + # /////////////////////////////////////////////////////////////////////////////// + # Tile MMA compute thread partitions and allocate accumulators + # /////////////////////////////////////////////////////////////////////////////// + thr_mma_sdp = tiled_mma_sdp.get_slice(tidx) + thr_mma_dkv = tiled_mma_dkv.get_slice(tidx) + thr_mma_dq = tiled_mma_dq.get_slice(tidx) + acc_shape_dK = thr_mma_dkv.partition_shape_C((self.n_block_size, self.head_dim_padded)) + acc_shape_dV = thr_mma_dkv.partition_shape_C((self.n_block_size, self.head_dim_v_padded)) + acc_dK = cute.make_rmem_tensor(acc_shape_dK, cutlass.Float32) + acc_dV = cute.make_rmem_tensor(acc_shape_dV, cutlass.Float32) + acc_dK.fill(0.0) + acc_dV.fill(0.0) + + tSrQ = utils.mma_make_fragment_A(sQ[None, None, 0], thr_mma_sdp, swapAB=self.SdP_swapAB) + tSrK = utils.mma_make_fragment_B(sK, thr_mma_sdp, swapAB=self.SdP_swapAB) + tdPrdO = utils.mma_make_fragment_A(sdO[None, None, 0], thr_mma_sdp, swapAB=self.SdP_swapAB) + tdPrV = utils.mma_make_fragment_B(sV, thr_mma_sdp, swapAB=self.SdP_swapAB) + tdVrP = utils.mma_make_fragment_A(sPt, thr_mma_dkv, swapAB=self.dKV_swapAB) + tdVrdO = utils.mma_make_fragment_B(sdOt[None, None, 0], thr_mma_dkv, swapAB=self.dKV_swapAB) + tdKrdS = utils.mma_make_fragment_A(sdSt, thr_mma_dkv, swapAB=self.dKV_swapAB) + tdKrQ = utils.mma_make_fragment_B(sQt[None, None, 0], thr_mma_dkv, swapAB=self.dKV_swapAB) + tdQrdS = utils.mma_make_fragment_A(sdS, thr_mma_dq, swapAB=self.dQ_swapAB) + tdQrK = utils.mma_make_fragment_B(sKt, thr_mma_dq, swapAB=self.dQ_swapAB) + + LSEslice = (None, 0, None) if cutlass.const_expr(not self.SdP_swapAB) else (0, None, None) + tSsLSEMma = layout_utils.reshape_acc_to_mn(thr_mma_sdp.partition_C(sLSEMma))[LSEslice] + tSsdPsumMma = layout_utils.reshape_acc_to_mn(thr_mma_sdp.partition_C(sdPsumMma))[LSEslice] + + # /////////////////////////////////////////////////////////////////////////////// + # Smem copy atom tiling + # /////////////////////////////////////////////////////////////////////////////// + smem_copy_atom = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), self.dtype, + ) + smem_copy_atom_transposed = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4), self.dtype, + ) + smem_thr_copy_QdO = utils.make_tiled_copy_A( + smem_copy_atom, tiled_mma_sdp, swapAB=self.SdP_swapAB + ).get_slice(tidx) + smem_thr_copy_KV = utils.make_tiled_copy_B( + smem_copy_atom, tiled_mma_sdp, swapAB=self.SdP_swapAB + ).get_slice(tidx) + # TODO: should this be smem_copy_atom_transposed? + smem_thr_copy_PdSt = utils.make_tiled_copy_A( + smem_copy_atom_transposed, tiled_mma_dkv, swapAB=self.dKV_swapAB + ).get_slice(tidx) + smem_thr_copy_QdOt = utils.make_tiled_copy_B( + smem_copy_atom_transposed, tiled_mma_dkv, swapAB=self.dKV_swapAB + ).get_slice(tidx) + smem_thr_copy_dS = utils.make_tiled_copy_A( + smem_copy_atom, tiled_mma_dq, swapAB=self.dQ_swapAB + ).get_slice(tidx) + smem_thr_copy_Kt = utils.make_tiled_copy_B( + smem_copy_atom_transposed, tiled_mma_dq, swapAB=self.dQ_swapAB + ).get_slice(tidx) + # TODO: what's the number of bits? What if SdP_swapAB + r2s_thr_copy_PdS = cute.make_tiled_copy_C( + cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.dtype, num_bits_per_copy=2 * self.dtype.width + ), + tiled_mma_sdp, + ).get_slice(tidx) + + tSsQ = smem_thr_copy_QdO.partition_S(sQ) + tdPsdO = smem_thr_copy_QdO.partition_S(sdO) + tSsK = smem_thr_copy_KV.partition_S(sK) + tdPsV = smem_thr_copy_KV.partition_S(sV) + tdVsPt = smem_thr_copy_PdSt.partition_S(sPt) + tdKsdSt = smem_thr_copy_PdSt.partition_S(sdSt) + tdVsdOt = smem_thr_copy_QdOt.partition_S(sdOt) + tdKsQt = smem_thr_copy_QdOt.partition_S(sQt) + tdQsdS = smem_thr_copy_dS.partition_S(sdS) + tdQsKt = smem_thr_copy_Kt.partition_S(sKt) + tPsP = r2s_thr_copy_PdS.partition_D(sP) + tdSsdS = r2s_thr_copy_PdS.partition_D(sdS) + + # /////////////////////////////////////////////////////////////////////////////// + # Predicate: Mark indices that need to copy when problem_shape isn't a multiple + # of tile_shape + # /////////////////////////////////////////////////////////////////////////////// + # Construct identity layout for KV + cQ = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded)) + tQcQ = gmem_thr_copy_QK.partition_S(cQ) + t0QcQ = gmem_thr_copy_QK.get_slice(0).partition_S(cQ) + if cutlass.const_expr(self.head_dim_padded == self.head_dim_v_padded): + tdOcdO = tQcQ + t0dOcdO = t0QcQ + else: + cdO = cute.make_identity_tensor((self.m_block_size, self.head_dim_v_padded)) + tdOcdO = gmem_thr_copy_VdO.partition_S(cdO) + t0dOcdO = gmem_thr_copy_VdO.get_slice(0).partition_S(cdO) + cLSE = cute.make_identity_tensor((self.m_block_size,)) + tLSEcLSE = gmem_thr_copy_lse.partition_S(cLSE) + + # Allocate predicate tensors for m and n, here we only allocate the tile of k, and + # use "if" on the mn dimension. + # This is to reduce register pressure and gets 2-3% performance gain. + + d_head = mQ.shape[cute.rank(mQ) - 1] + d_head_v = mdO.shape[cute.rank(mdO) - 1] + + tQpQ = utils.predicate_k(tQcQ, limit=d_head) + if cutlass.const_expr(self.same_hdim_kv): + tdOpdO = tQpQ + else: + tdOpdO = utils.predicate_k(tdOcdO, limit=d_head_v) + + # group parameters for compute_one_m_block + mma_params = SimpleNamespace( + thr_mma_sdp=thr_mma_sdp, thr_mma_dkv=thr_mma_dkv, thr_mma_dq=thr_mma_dq, + tSrQ=tSrQ, tSrK=tSrK, tdPrdO=tdPrdO, tdPrV=tdPrV, + tdVrP=tdVrP, tdVrdO=tdVrdO, tdKrdS=tdKrdS, tdKrQ=tdKrQ, + tdQrdS=tdQrdS, tdQrK=tdQrK, + acc_dK=acc_dK, acc_dV=acc_dV, + ) + smem_copy_params = SimpleNamespace( + smem_thr_copy_QdO=smem_thr_copy_QdO, + smem_thr_copy_KV=smem_thr_copy_KV, + smem_thr_copy_PdSt=smem_thr_copy_PdSt, + smem_thr_copy_QdOt=smem_thr_copy_QdOt, + smem_thr_copy_dS=smem_thr_copy_dS, + smem_thr_copy_Kt=smem_thr_copy_Kt, + r2s_thr_copy_PdS=r2s_thr_copy_PdS, + tSsQ=tSsQ, tSsK=tSsK, tdPsdO=tdPsdO, tdPsV=tdPsV, + tSsLSEMma=tSsLSEMma, tSsdPsumMma=tSsdPsumMma, + tPsP=tPsP, tdSsdS=tdSsdS, + tdVsPt=tdVsPt, tdVsdOt=tdVsdOt, tdKsdSt=tdKsdSt, tdKsQt=tdKsQt, + tdQsdS=tdQsdS, tdQsKt=tdQsKt, + ) + gmem_copy_params = SimpleNamespace( + gmem_thr_copy_dQaccum=gmem_thr_copy_dQaccum, tdQgdQaccum=tdQgdQaccum + ) + load_Q_LSE = partial( + self.load_Q_LSE, gmem_tiled_copy_QK, gmem_tiled_copy_LSE, + tQgQ, tQsQ, tQcQ, t0QcQ, tQpQ, + tLSEgLSE, tLSEsLSE, tLSEcLSE, seqlen=seqlen.seqlen_q + ) + load_dO_dPsum = partial( + self.load_dO_dPsum, gmem_tiled_copy_VdO, gmem_tiled_copy_LSE, + tdOgdO, tdOsdO, tdOcdO, t0dOcdO, tdOpdO, + tLSEgdPsum, tLSEsdPsum, tLSEcLSE, seqlen=seqlen.seqlen_q + ) + compute_one_m_block = partial( + self.compute_one_m_block, mma_params=mma_params, + smem_copy_params=smem_copy_params, gmem_copy_params=gmem_copy_params, + load_Q_LSE=load_Q_LSE, load_dO_dPsum=load_dO_dPsum, + m_block_max=m_block_max, + softmax_scale=softmax_scale, + softmax_scale_log2=softmax_scale_log2, + aux_data=aux_data, + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Prologue + # /////////////////////////////////////////////////////////////////////////////// + # Start async loads of the last mn-tile, where we take care of the mn residue + self.load_V(gmem_thr_copy_VdO, tVgV, tVsV, n_block, seqlen=seqlen.seqlen_k, + headdim=d_head_v) + if cutlass.const_expr(self.V_in_regs): + cute.arch.cp_async_commit_group() + self.load_K(gmem_thr_copy_QK, tKgK, tKsK, n_block, seqlen=seqlen.seqlen_k, + headdim=d_head) + cute.arch.cp_async_commit_group() + + if cutlass.const_expr(self.V_in_regs): + cute.arch.cp_async_wait_group(1) + cute.arch.barrier() + tdPrV_copy_view = smem_thr_copy_KV.retile(tdPrV) + cute.copy(smem_thr_copy_KV, tdPsV, tdPrV_copy_view) + # Sync to avoid loading Q to smem_q, which overlaps with smem_v + cute.arch.barrier() + + m_block = m_block_min + assert self.num_stages_Q >= self.num_stages_dO + for stage in cutlass.range_constexpr(self.num_stages_Q): + if cutlass.const_expr(self.num_stages_Q == 1 or stage < self.num_stages_Q - 1): + if stage == 0 or m_block + stage < m_block_max: + load_Q_LSE(m_block + stage, smem_pipe_write_q=stage) + cute.arch.cp_async_commit_group() + if cutlass.const_expr(stage < self.num_stages_dO): + if stage == 0 or m_block + stage < m_block_max: + load_dO_dPsum(m_block + stage, smem_pipe_write_q=stage) + cute.arch.cp_async_commit_group() + + # /////////////////////////////////////////////////////////////////////////////// + # Mainloop + # /////////////////////////////////////////////////////////////////////////////// + # Start processing of the first n-block. + mask = AttentionMask(self.m_block_size, self.n_block_size, seqlen) + mask_fn = partial( + mask.apply_mask, n_block=n_block, thr_mma=thr_mma_sdp, + batch_idx=batch_idx, head_idx=head_idx, + mask_seqlen=True, mask_causal=self.is_causal + ) + smem_pipe_read_q = cutlass.Int32(0) + smem_pipe_read_do = cutlass.Int32(0) + smem_pipe_write_q = cutlass.Int32(self.num_stages_Q - 1) + smem_pipe_write_do = cutlass.Int32(0) + for m_tile in cutlass.range(m_block_min, m_block_max, unroll=1): + compute_one_m_block( + m_tile, smem_pipe_read_q, smem_pipe_read_do, smem_pipe_write_q, smem_pipe_write_do, + mask_fn=mask_fn, + ) + smem_pipe_read_q = self.advance_pipeline(smem_pipe_read_q, self.num_stages_Q) + smem_pipe_read_do = self.advance_pipeline(smem_pipe_read_do, self.num_stages_dO) + smem_pipe_write_q = self.advance_pipeline(smem_pipe_write_q, self.num_stages_Q) + smem_pipe_write_do = self.advance_pipeline(smem_pipe_write_do, self.num_stages_dO) + + # /////////////////////////////////////////////////////////////////////////////// + # Epilogue + # /////////////////////////////////////////////////////////////////////////////// + # If GQA, we scale dK in the postprocessing kernel instead + if cutlass.const_expr(self.qhead_per_kvhead == 1): + acc_dK.store(acc_dK.load() * softmax_scale) + # reuse sK and sV data iterator + sdK = cute.make_tensor(sK.iterator, sK_layout) + sdV = cute.make_tensor(sV.iterator, sV_layout) + self.epilogue( + acc_dK, acc_dV, mdK, mdV, sdK, sdV, + gmem_tiled_copy_dK, gmem_tiled_copy_dV, tiled_mma_dkv, + tidx, n_block, head_idx, batch_idx, seqlen, d_head, d_head_v + ) + + @cute.jit + def compute_one_m_block( + self, + m_block: cutlass.Int32, + smem_pipe_read_q: cutlass.Int32, + smem_pipe_read_do: cutlass.Int32, + smem_pipe_write_q: cutlass.Int32, + smem_pipe_write_do: cutlass.Int32, + mma_params: SimpleNamespace, + smem_copy_params: SimpleNamespace, + gmem_copy_params: SimpleNamespace, + load_Q_LSE: Callable, + load_dO_dPsum: Callable, + m_block_max: cutlass.Int32, + softmax_scale: cutlass.Float32, + softmax_scale_log2: cutlass.Float32, + aux_data: AuxData = AuxData(), + mask_fn: Optional[Callable] = None, + ): + def load_Q_next(): + m_block_next = m_block + (self.num_stages_Q - 1 if cutlass.const_expr(self.num_stages_Q > 1) else 1) + if m_block_next < m_block_max: + load_Q_LSE(m_block_next, smem_pipe_write_q) + cute.arch.cp_async_commit_group() + + def load_dO_next(): + if m_block + self.num_stages_dO < m_block_max: + load_dO_dPsum(m_block + self.num_stages_dO, smem_pipe_write_do) + cute.arch.cp_async_commit_group() + + # MMA S + acc_shape_SdP = mma_params.thr_mma_sdp.partition_shape_C( + (self.m_block_size, self.n_block_size) if cutlass.const_expr(not self.SdP_swapAB) else (self.n_block_size, self.m_block_size) + ) + acc_S = cute.make_rmem_tensor(acc_shape_SdP, cutlass.Float32) + acc_S.fill(0.0) + cute.arch.cp_async_wait_group(1 if cutlass.const_expr(self.num_stages_Q > 1) else 0) + cute.arch.barrier() + sm80_utils.gemm( + mma_params.thr_mma_sdp, acc_S, mma_params.tSrQ, mma_params.tSrK, + smem_copy_params.tSsQ[None, None, None, smem_pipe_read_q if cutlass.const_expr(self.num_stages_Q > 1) else 0], + smem_copy_params.tSsK, + smem_copy_params.smem_thr_copy_QdO, smem_copy_params.smem_thr_copy_KV, + swap_AB=self.SdP_swapAB, + ) + acc_S_pre = cute.make_fragment_like(acc_S) + acc_S_pre.store(acc_S.load()) + tLSErLSE = cute.make_fragment_like(smem_copy_params.tSsLSEMma[None, 0]) + cute.autovec_copy( + smem_copy_params.tSsLSEMma[None, smem_pipe_read_q if cutlass.const_expr(self.num_stages_Q > 1) else 0], tLSErLSE + ) + acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S) + acc_S_pre_mn = layout_utils.reshape_acc_to_mn(acc_S_pre) + if cutlass.const_expr(self.score_mod is not None): + for r in cutlass.range(cute.size(acc_S_mn, mode=[0]), unroll_full=True): + acc_S_mn[r, None].store( + call_score_mod( + self.score_mod, + acc_S_mn[r, None].load() * softmax_scale, + 0, + 0, + 0, + 0, + None, + aux_data, + ) + ) + if cutlass.const_expr(mask_fn is not None): + mask_fn(acc_S, m_block=m_block) + bidx = 0 + # if cute.arch.thread_idx()[0] == 0 and cute.arch.block_idx()[0] == bidx: cute.print_tensor(acc_S_mn) + # if cute.arch.thread_idx()[0] == 0 and cute.arch.block_idx()[0] == 1: cute.print_tensor(tLSErLSE) + assert cute.size(acc_S_mn, mode=[0]) == cute.size(tLSErLSE) + for r in cutlass.range(cute.size(acc_S_mn, mode=[0]), unroll_full=True): + acc_S_mn[r, None].store(cute.math.exp2(acc_S_mn[r, None].load() * softmax_scale_log2 - tLSErLSE[r], fastmath=True)) + # if cute.arch.thread_idx()[0] == 0 and cute.arch.block_idx()[0] == bidx: cute.print_tensor(acc_S_mn) + + # MMA dP + acc_dP = cute.make_rmem_tensor(acc_shape_SdP, cutlass.Float32) + acc_dP.fill(0.0) + cute.arch.cp_async_wait_group(1 if cutlass.const_expr(self.num_stages_dO > 1) else 0) + cute.arch.barrier() + sm80_utils.gemm( + mma_params.thr_mma_sdp, acc_dP, mma_params.tdPrdO, mma_params.tdPrV, + smem_copy_params.tdPsdO[None, None, None, smem_pipe_read_do if cutlass.const_expr(self.num_stages_dO > 1) else 0], + smem_copy_params.tdPsV, + smem_copy_params.smem_thr_copy_QdO, smem_copy_params.smem_thr_copy_KV, + hook_fn=load_Q_next if cutlass.const_expr(self.num_stages_Q > 1) else None, + swap_AB=self.SdP_swapAB, + ) + tLSErdPsum = cute.make_fragment_like(smem_copy_params.tSsdPsumMma[None, 0]) + cute.autovec_copy( + smem_copy_params.tSsdPsumMma[None, smem_pipe_read_do if cutlass.const_expr(self.num_stages_dO > 1) else 0], tLSErdPsum + ) + acc_dP_mn = layout_utils.reshape_acc_to_mn(acc_dP) + # if cute.arch.thread_idx()[0] == 0 and cute.arch.block_idx()[0] == bidx: cute.print_tensor(acc_dP_mn) + assert cute.size(acc_dP_mn, mode=[0]) == cute.size(tLSErdPsum) + for r in cutlass.range(cute.size(acc_dP_mn, mode=[0]), unroll_full=True): + grad_val = acc_S_mn[r, None].load() * (acc_dP_mn[r, None].load() - tLSErdPsum[r]) + if cutlass.const_expr(self.score_mod_bwd is not None): + grad_val = call_score_mod_bwd( + self.score_mod_bwd, + grad_val, + acc_S_pre_mn[r, None].load() * softmax_scale, + 0, + 0, + 0, + 0, + None, + aux_data, + ) + acc_dP_mn[r, None].store(grad_val) + # if cute.arch.thread_idx()[0] == 0 and cute.arch.block_idx()[0] == bidx: cute.print_tensor(acc_dP_mn) + rP = cute.make_fragment_like(acc_S, self.dtype) + rP.store(acc_S.load().to(self.dtype)) + if cutlass.const_expr(not self.Mma_dKV_is_RS): + tPrP = smem_copy_params.r2s_thr_copy_PdS.retile(rP) # ((Atom,AtomNum), MMA_N, MMA_N) + cute.copy(smem_copy_params.r2s_thr_copy_PdS, tPrP, smem_copy_params.tPsP) + rdS = cute.make_fragment_like(acc_dP, self.dtype) + rdS.store(acc_dP.load().to(self.dtype)) + if cutlass.const_expr(not self.Mma_dKV_is_RS): + cute.arch.barrier() # Make sure P is written + # For hdim 64, It's faster to write to smem_dS first before the dV gemm + if cutlass.const_expr(not self.Mma_dKV_is_RS): + tdSrdS = smem_copy_params.r2s_thr_copy_PdS.retile(rdS) + cute.copy(smem_copy_params.r2s_thr_copy_PdS, tdSrdS, smem_copy_params.tdSsdS) + if cutlass.const_expr(self.Mma_dKV_is_RS): + tdVrP = layout_utils.reshape_acc_to_frgA(rP) + else: + tdVrP = mma_params.tdVrP + + # MMA dK + sm80_utils.gemm( + mma_params.thr_mma_dkv, mma_params.acc_dV, tdVrP, mma_params.tdVrdO, + smem_copy_params.tdVsPt, + smem_copy_params.tdVsdOt[None, None, None, smem_pipe_read_do if cutlass.const_expr(self.num_stages_dO > 1) else 0], + smem_copy_params.smem_thr_copy_PdSt, smem_copy_params.smem_thr_copy_QdOt, + A_in_regs=self.Mma_dKV_is_RS, + swap_AB=self.dKV_swapAB, + ) + # if cute.arch.thread_idx()[0] == 0 and cute.arch.block_idx()[0] == bidx: cute.print_tensor(mma_params.acc_dV) + cute.arch.barrier() # Make sure dS is written + + # MMA dQ + def dQ_mma(hook_fn): + acc_shape_dQ = mma_params.thr_mma_dq.partition_shape_C( + (self.m_block_size, self.head_dim_padded) if cutlass.const_expr(not self.dQ_swapAB) else (self.head_dim_padded, self.m_block_size) + ) + acc_dQ = cute.make_rmem_tensor(acc_shape_dQ, cutlass.Float32) + acc_dQ.fill(0.0) + sm80_utils.gemm( + mma_params.thr_mma_dq, acc_dQ, mma_params.tdQrdS, mma_params.tdQrK, + smem_copy_params.tdQsdS, smem_copy_params.tdQsKt, + smem_copy_params.smem_thr_copy_dS, smem_copy_params.smem_thr_copy_Kt, + swap_AB=self.dQ_swapAB, + hook_fn=hook_fn + ) + # ((1, 1), num_elements) + acc_dQ_atomic = gmem_copy_params.gmem_thr_copy_dQaccum.retile(acc_dQ) + tdQgdQaccum_atomic = gmem_copy_params.tdQgdQaccum[None, None, m_block] + assert cute.size(acc_dQ_atomic) == cute.size(tdQgdQaccum_atomic) + for i in cutlass.range(cute.size(acc_dQ_atomic), unroll_full=True): + utils.atomic_add_fp32(acc_dQ_atomic[i], utils.elem_pointer(tdQgdQaccum_atomic, i)) + # utils.atomic_add_fp32(acc_dQ[i], tdQgdQaccum_atomic.iterator + i * tdQgdQaccum_atomic.stride[1]) + # if cute.arch.thread_idx()[0] == 64 and cute.arch.block_idx()[0] == bidx: cute.print_tensor(acc_dQ) + + # If num_stages_Q == 1, we want to do Mma_dK first so we can start loading Q for the next iteration + if cutlass.const_expr(self.num_stages_Q > 1): + dQ_mma(load_dO_next) + + # MMA dK + if cutlass.const_expr(self.Mma_dKV_is_RS): + tdKrdS = layout_utils.reshape_acc_to_frgA(rdS) + else: + tdKrdS = mma_params.tdKrdS + sm80_utils.gemm( + mma_params.thr_mma_dkv, mma_params.acc_dK, tdKrdS, mma_params.tdKrQ, + smem_copy_params.tdKsdSt, + smem_copy_params.tdKsQt[None, None, None, smem_pipe_read_q if cutlass.const_expr(self.num_stages_Q > 1) else 0], + smem_copy_params.smem_thr_copy_PdSt, smem_copy_params.smem_thr_copy_QdOt, + A_in_regs=self.Mma_dKV_is_RS, + swap_AB=self.dKV_swapAB, + hook_fn=load_dO_next if cutlass.const_expr(self.num_stages_Q == 1) else None, + ) + # if cute.arch.thread_idx()[0] == 0: cute.print_tensor(mma_params.acc_dK) + if cutlass.const_expr(self.num_stages_Q == 1): + cute.arch.barrier() + dQ_mma(load_Q_next) + + @cute.jit + def epilogue( + self, + acc_dK: cute.Tensor, + acc_dV: cute.Tensor, + mdK: cute.Tensor, + mdV: cute.Tensor, + sdK: cute.Tensor, + sdV: cute.Tensor, + gmem_tiled_copy_dK: cute.TiledCopy, + gmem_tiled_copy_dV: cute.TiledCopy, + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + n_block: cutlass.Int32, + num_head: cutlass.Int32, + batch_size: cutlass.Int32, + seqlen: SeqlenInfoQK, + d_head: cutlass.Int32, + d_head_v: cutlass.Int32 + ): + rdV = cute.make_fragment_like(acc_dV, self.dtype) + rdV.store(acc_dV.load().to(self.dtype)) + rdK = cute.make_fragment_like(acc_dK, self.dtype) + rdK.store(acc_dK.load().to(self.dtype)) + gmem_thr_copy_dK = gmem_tiled_copy_dK.get_slice(tidx) + gmem_thr_copy_dV = gmem_tiled_copy_dV.get_slice(tidx) + + batch_idx = batch_size + head_idx_kv = num_head // self.qhead_per_kvhead if cutlass.const_expr(not self.pack_gqa) else num_head + + if cutlass.const_expr(self.qhead_per_kvhead == 1): + # Make sure all threads have finished reading K and V, otherwise we get racy dQ + # because smem_q could be changed. + cute.arch.barrier() + # smem copy atom for dKV + smem_copy_atom_dKV = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.dtype, num_bits_per_copy=2 * self.dtype.width + ) + smem_thr_copy_dKV = cute.make_tiled_copy_C(smem_copy_atom_dKV, tiled_mma).get_slice(tidx) + taccdVrdV = smem_thr_copy_dKV.retile(rdV) + taccdKrdK = smem_thr_copy_dKV.retile(rdK) + taccdVsdV = smem_thr_copy_dKV.partition_D(sdV) + taccdKsdK = smem_thr_copy_dKV.partition_D(sdK) + # copy acc O from rmem to smem with the smem copy atom + cute.copy(smem_copy_atom_dKV, taccdVrdV, taccdVsdV) + cute.copy(smem_copy_atom_dKV, taccdKrdK, taccdKsdK) + + + if cutlass.const_expr(not seqlen.has_cu_seqlens_k): + mdK_cur, mdV_cur = [t[batch_idx, None, head_idx_kv, None] for t in (mdK, mdV)] + else: + mdK_cur, mdV_cur = [cute.domain_offset((seqlen.offset_k, 0), t[None, head_idx_kv, None]) for t in (mdK, mdV)] + + blkdK_shape = (self.n_block_size, self.head_dim_padded) + blkdV_shape = (self.n_block_size, self.head_dim_v_padded) + gdK = cute.local_tile(mdK_cur, blkdK_shape, (n_block, 0)) + gdV = cute.local_tile(mdV_cur, blkdV_shape, (n_block, 0)) + tdKsdK = gmem_thr_copy_dK.partition_S(sdK) + tdKgdK = gmem_thr_copy_dK.partition_D(gdK) + tdVsdV = gmem_thr_copy_dV.partition_S(sdV) + tdVgdV = gmem_thr_copy_dV.partition_D(gdV) + tdKrdK = cute.make_fragment_like(tdKgdK, self.dtype) + tdVrdV = cute.make_fragment_like(tdVgdV, self.dtype) + # sync before all smem stores are done. + cute.arch.barrier() + # load acc dK and dV from smem to rmem for wider vectorization + # Need to check OOB when reading from smem if kBlockN isn't evenly tiled + # TODO + cute.autovec_copy(tdKsdK, tdKrdK) + cute.autovec_copy(tdVsdV, tdVrdV) + + cdK = cute.make_identity_tensor((self.n_block_size, self.head_dim_padded)) + tdKcdK = gmem_thr_copy_dK.partition_S(cdK) + t0dKcdK = gmem_tiled_copy_dK.get_slice(0).partition_S(cdK) + if cutlass.const_expr(self.head_dim_padded == self.head_dim_v_padded): + tdVcdV = tdKcdK + t0dVcdV = t0dKcdK + else: + cdV = cute.make_identity_tensor((self.n_block_size, self.head_dim_v_padded)) + tdVcdV = gmem_thr_copy_dV.partition_S(cdV) + t0dVcdV = gmem_tiled_copy_dV.get_slice(0).partition_S(cdV) + tdKpdK = utils.predicate_k(tdKcdK, limit=d_head) + if cutlass.const_expr(self.same_hdim_kv): + tdVpdV = tdKpdK + else: + tdVpdV = utils.predicate_k(tdVcdV, limit=d_head_v) + # copy acc dK and acc_dV from rmem to gmem + for rest_m in cutlass.range_constexpr(cute.size(tdKrdK.shape[1])): + if t0dKcdK[0, rest_m, 0][0] < seqlen.seqlen_k - n_block * self.n_block_size - tdKcdK[0][0]: + cute.copy( + gmem_tiled_copy_dK, + tdKrdK[None, rest_m, None], + tdKgdK[None, rest_m, None], + pred=tdKpdK[None, rest_m, None] if cutlass.const_expr(self.check_hdim_oob) else None, + ) + for rest_m in cutlass.range_constexpr(cute.size(tdVrdV.shape[1])): + if t0dVcdV[0, rest_m, 0][0] < seqlen.seqlen_k - n_block * self.n_block_size - tdVcdV[0][0]: + cute.copy( + gmem_tiled_copy_dV, + tdVrdV[None, rest_m, None], + tdVgdV[None, rest_m, None], + pred=tdVpdV[None, rest_m, None] if cutlass.const_expr(self.check_hdim_v_oob) else None, + ) + + else: # qhead_per_kvhead > 1, do atomic add + # For Sm90, we need to sync to avoid racy writes to smem_q + # For Sm80, we don't need to sync since we're not touching smem + head_idx_kv = num_head // self.qhead_per_kvhead if cutlass.const_expr(not self.pack_gqa) else num_head + + if cutlass.const_expr(not seqlen.has_cu_seqlens_k): + mdK_cur, mdV_cur = [t[batch_idx, head_idx_kv, None] for t in (mdK, mdV)] + else: + padded_offset_k = seqlen.offset_k + batch_idx * self.n_block_size + mdK_cur = cute.domain_offset((padded_offset_k * self.head_dim_padded,), mdK[head_idx_kv, None]) + mdV_cur = cute.domain_offset((padded_offset_k * self.head_dim_v_padded,), mdV[head_idx_kv, None]) + + gdV = cute.local_tile(mdV_cur, (self.n_block_size * self.head_dim_v_padded,), (n_block,)) + gdK = cute.local_tile(mdK_cur, (self.n_block_size * self.head_dim_padded,), (n_block,)) + tdVgdVaccum = gmem_thr_copy_dV.partition_S(gdV) + tdKgdKaccum = gmem_thr_copy_dK.partition_S(gdK) + acc_dV_atomic = gmem_thr_copy_dV.retile(acc_dV) + acc_dK_atomic = gmem_thr_copy_dK.retile(acc_dK) + assert cute.size(acc_dV_atomic) == cute.size(tdVgdVaccum) + assert cute.size(acc_dK_atomic) == cute.size(tdKgdKaccum) + for i in cutlass.range(cute.size(acc_dV_atomic), unroll_full=True): + utils.atomic_add_fp32(acc_dV_atomic[i], utils.elem_pointer(tdVgdVaccum, i)) + for i in cutlass.range(cute.size(acc_dK_atomic), unroll_full=True): + utils.atomic_add_fp32(acc_dK_atomic[i], utils.elem_pointer(tdKgdKaccum, i)) + + @cute.jit + def advance_pipeline(self, pipeline_index, num_stages: cutlass.Constexpr): + return pipeline_index + 1 if pipeline_index < num_stages - 1 else 0 + + @cute.jit + def load_K( + self, + gmem_thr_copy: cute.TiledCopy, + tKgK: cute.Tensor, + tKsK: cute.Tensor, + block: cutlass.Int32, + seqlen: cutlass.Int32, + headdim: cutlass.Int32, + ): + cK = cute.make_identity_tensor((self.n_block_size, self.head_dim_padded)) + tKcK = gmem_thr_copy.partition_S(cK) + t0KcK = gmem_thr_copy.get_slice(0).partition_S(cK) + tKpK = utils.predicate_k(tKcK, limit=headdim) + for n in cutlass.range_constexpr(cute.size(tKsK.shape[1])): + # If kBlockN doesn't evenly divide the tiled copy, only the last `n` needs to be checked + if self.is_even_n_smem_k or n < cute.size(tKsK.shape[1]) - 1 or tKcK[0, n, 0][0] < self.n_block_size: + # Instead of using tKcK, we using t0KcK and subtract the offset from the limit + # (seqlen - block * kBlockN). This is because the entries of t0KcK are known at compile time. + predicate_n = t0KcK[0, n, 0][0] < seqlen - block * self.n_block_size - tKcK[0][0] + predicate = cute.make_fragment_like(tKpK[None, 0, None]) + for k in cutlass.range_constexpr(cute.size(predicate.shape[1])): + for i in cutlass.range_constexpr(cute.size(predicate.shape[0])): + predicate[i, k] = (tKpK[i, n, k] if cutlass.const_expr(self.check_hdim_oob) else True) and predicate_n + cute.copy( + gmem_thr_copy, tKgK[None, n, None], tKsK[None, n, None], pred=predicate, + ) + # We need to clear the sK smem tiles since we'll use sKt for mma_dq + + @cute.jit + def load_V( + self, + gmem_thr_copy: cute.TiledCopy, + tVgV: cute.Tensor, + tVsV: cute.Tensor, + block: cutlass.Int32, + seqlen: cutlass.Int32, + headdim: cutlass.Int32, + ): + cV = cute.make_identity_tensor((self.n_block_size, self.head_dim_v_padded)) + tVcV = gmem_thr_copy.partition_S(cV) + t0VcV = gmem_thr_copy.get_slice(0).partition_S(cV) + tVpV = utils.predicate_k(tVcV, limit=headdim) + for n in cutlass.range_constexpr(cute.size(tVsV.shape[1])): + # If kBlockN doesn't evenly divide the tiled copy, only the last `n` needs to be checked + if self.is_even_n_smem_v or n < cute.size(tVsV.shape[1]) - 1 or tVcV[0, n, 0][0] < self.n_block_size: + # Instead of using tVcV, we using t0VcV and subtract the offset from the limit + # (seqlen - block * kBlockN). This is because the entries of t0VcV are known at compile time. + predicate_n = t0VcV[0, n, 0][0] < seqlen - block * self.n_block_size - tVcV[0][0] + predicate = cute.make_fragment_like(tVpV[None, 0, None]) + for k in cutlass.range_constexpr(cute.size(predicate.shape[1])): + for i in cutlass.range_constexpr(cute.size(predicate.shape[0])): + predicate[i, k] = (tVpV[i, n, k] if cutlass.const_expr(self.check_hdim_oob) else True) and predicate_n + cute.copy( + gmem_thr_copy, tVgV[None, n, None], tVsV[None, n, None], pred=predicate, + ) + + @cute.jit + def load_Q_LSE( + self, + gmem_tiled_copy_Q: cute.TiledCopy, + gmem_tiled_copy_LSE: cute.TiledCopy, + tQgQ: cute.Tensor, + tQsQ: cute.Tensor, + tQcQ: cute.Tensor, + t0QcQ: cute.Tensor, + tQpQ: cute.Tensor, + tLSEgLSE: cute.Tensor, + tLSEsLSE: cute.Tensor, + tLSEcLSE: cute.Tensor, + block: cutlass.Int32, + smem_pipe_write_q: cutlass.Int32, + seqlen: cutlass.Int32, + ): + for m in cutlass.range_constexpr(cute.size(tQsQ.shape[1])): + # If kBlockM doesn't evenly divide the tiled copy, only the last `m` needs to be checked + if self.is_even_m_smem_q or m < cute.size(tQsQ.shape[1]) - 1 or tQcQ[0, m, 0][0] < self.m_block_size: + # Instead of using tQcQ, we using t0QcQ and subtract the offset from the limit + # (seqlen - block * kBlockM). This is because the entries of t0QcQ are known at compile time. + predicate_m = t0QcQ[0, m, 0][0] < seqlen - block * self.m_block_size - tQcQ[0][0] + predicate = cute.make_fragment_like(tQpQ[None, 0, None]) + for k in cutlass.range_constexpr(cute.size(predicate.shape[1])): + for i in cutlass.range_constexpr(cute.size(predicate.shape[0])): + predicate[i, k] = (tQpQ[i, m, k] if cutlass.const_expr(self.check_hdim_oob) else True) and predicate_m + cute.copy( + gmem_tiled_copy_Q, + tQgQ[None, m, None, block], + tQsQ[None, m, None, smem_pipe_write_q if cutlass.const_expr(self.num_stages_Q) > 1 else 0], + pred=predicate, + ) + # We need to clear the sQ smem tiles since we'll use sQt for mma_dK + # We made sure LSE length is padded so we read `kBlockM` elements so that all + # elements in sLSE are filled. Without this we might have uninitialized sLSE values. + for m in cutlass.range_constexpr(cute.size(tLSEsLSE.shape[1])): + if tLSEcLSE[0, m][0] < self.m_block_size: + cute.copy( + gmem_tiled_copy_LSE, + tLSEgLSE[None, m, block], + tLSEsLSE[None, m, smem_pipe_write_q if cutlass.const_expr(self.num_stages_Q > 1) else 0], + ) + + @cute.jit + def load_dO_dPsum( + self, + gmem_tiled_copy_dO: cute.TiledCopy, + gmem_tiled_copy_dPsum: cute.TiledCopy, + tdOgdO: cute.Tensor, + tdOsdO: cute.Tensor, + tdOcdO: cute.Tensor, + t0dOcdO: cute.Tensor, + tdOpdO: cute.Tensor, + tdPsumgdPsum: cute.Tensor, + tdPsumsdPsum: cute.Tensor, + tdPsumcdPsum: cute.Tensor, + block: cutlass.Int32, + smem_pipe_write_q: cutlass.Int32, + seqlen: cutlass.Int32, + ): + for m in cutlass.range_constexpr(cute.size(tdOsdO.shape[1])): + # If kBlockM doesn't evenly divide the tiled copy, only the last `m` needs to be checked + if self.is_even_m_smem_do or m < cute.size(tdOsdO.shape[1]) - 1 or tdOcdO[0, m, 0][0] < self.m_block_size: + # Instead of using tdOcdO, we using t0dOcdO and subtract the offset from the limit + # (seqlen - block * kBlockM). This is because the entries of t0dOcdO are known at compile time. + predicate_m = t0dOcdO[0, m, 0][0] < seqlen - block * self.m_block_size - tdOcdO[0][0] + predicate = cute.make_fragment_like(tdOpdO[None, 0, None]) + for k in cutlass.range_constexpr(cute.size(predicate.shape[1])): + for i in cutlass.range_constexpr(cute.size(predicate.shape[0])): + predicate[i, k] = (tdOpdO[i, m, k] if cutlass.const_expr(self.check_hdim_oob) else True) and predicate_m + cute.copy( + gmem_tiled_copy_dO, + tdOgdO[None, m, None, block], + tdOsdO[None, m, None, smem_pipe_write_q if cutlass.const_expr(self.num_stages_dO > 1) else 0], + pred=predicate, + ) + # We need to clear the sQ smem tiles since we'll use sQt for mma_dK + # We made sure LSE length is padded so we read `kBlockM` elements so that all + # elements in sLSE are filled. Without this we might have uninitialized sLSE values. + for m in cutlass.range_constexpr(cute.size(tdPsumgdPsum.shape[1])): + if tdPsumcdPsum[0, m][0] < self.m_block_size: + cute.copy( + gmem_tiled_copy_dPsum, + tdPsumgdPsum[None, m, block], + tdPsumsdPsum[None, m, smem_pipe_write_q if cutlass.const_expr(self.num_stages_dO > 1) else 0], + ) diff --git a/fastvideo-kernel/fa4/flash_bwd_mla_dk_sm100.py b/fastvideo-kernel/fa4/flash_bwd_mla_dk_sm100.py new file mode 100644 index 000000000..9083c126d --- /dev/null +++ b/fastvideo-kernel/fa4/flash_bwd_mla_dk_sm100.py @@ -0,0 +1,1183 @@ +# Modified from CUTLASS example file, original copyright: +# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from typing import Callable, Optional, Tuple, Union +import cuda.bindings.driver as cuda + +import cutlass +from cutlass import Int32, const_expr +from cutlass.cutlass_dsl import dsl_user_op +import cutlass.cute as cute +import cutlass.cute.testing as testing +import cutlass.utils as utils +from cutlass.utils.gemm.sm100 import ( + transform_partitioned_tensor_layout, + epilogue_smem_copy_and_partition, + epilogue_tmem_copy_and_partition, +) +import cutlass.pipeline as pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.cute.nvgpu import cpasync, tcgen05 + +from flash_attn.cute.utils import get_batch_from_cu_tensor + + +@dsl_user_op +def elem_pointer(x: cute.Tensor, coord, *, loc=None, ip=None) -> cute.Pointer: + """ + Get a pointer to an element at the specified coordinate in a tensor. + + Args: + x: The tensor (typically a shared memory tensor) + coord: The coordinate tuple, can be hierarchical like (row, (col, cluster_idx)) + + Returns: + Pointer to the element at the specified coordinate + """ + return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip) + + +class dKGemmKernel: + def __init__( + self, + topk: int, + heads: int, + dim: int, + varlen: bool, + ): + self.varlen = varlen + self.topk = topk + self.heads = heads + self.dim = dim + # A operand dS'^T: (total_q, heads, topk), topk-major + self.ab_dtype = cutlass.BFloat16 + self.a_major_mode = cute.nvgpu.OperandMajorMode.MN + # B operand Q: (total_q, heads, dim), dim-major + self.b_major_mode = cute.nvgpu.OperandMajorMode.MN + # Index operand I: (total_q, seqlen_k) + self.idx_dtype = cutlass.Int32 + + # Output dKaccum: (total_q, seqlen_k, dim), dim-major + self.c_dtype = cutlass.Float32 + self.c_layout = utils.LayoutEnum.ROW_MAJOR + + self.acc_dtype = cutlass.Float32 + + if self.topk % 256 == 1: + # The kernel schedule requires 2CTA instructions with this tile shape + self.cluster_shape_mn = (2, 1) + self.use_2cta_instrs = True + else: + self.cluster_shape_mn = (1, 1) + self.use_2cta_instrs = False + self.cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + self.atom_thr_size = 2 if self.use_2cta_instrs else 1 + self.cta_tile_shape_mnk = (128, 64, 1) + self.mma_tiler_dK = ( + self.cta_tile_shape_mnk[0] * self.atom_thr_size, + self.cta_tile_shape_mnk[1], + 1, + ) + self.arch = "sm_100" + + self.occupancy = 1 + # Set specialized warp ids + self.epilogue_warp_id = [0, 1, 2, 3] + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.sched_warp_id = 6 + self.threads_per_cta = 32 * len( + [ + self.mma_warp_id, + self.tma_warp_id, + self.sched_warp_id, + *self.epilogue_warp_id, + ] + ) + # Set barrier id for cta sync, epilogue sync and tmem ptr sync + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.I_load_bar_id = 3 + + @cute.jit + def __call__( + self, + dS: cute.Tensor, # (batch, seqlen_q, heads, topk) or (total_q, heads, topk) + I: cute.Tensor, # (batch, seqlen_q, heads, topk) or (total_q, heads, topk) + Q: cute.Tensor, # (batch, seqlen_q, heads, dim) or (total_q, heads, dim) + dKaccum: cute.Tensor, # (batch, seqlen_k, dim) or (total_k, dim) + cuSeqlensQ: Optional[cute.Tensor], # (batch + 1,) + cuSeqlensK: Optional[cute.Tensor], # (batch + 1,) + stream: cuda.CUstream, + ): + if const_expr(self.ab_dtype != dS.element_type): + raise TypeError(f"Type must match: {self.ab_dtype} != {dS.element_type}") + if const_expr(self.ab_dtype != Q.element_type): + raise TypeError(f"Type must match: {self.ab_dtype} != {Q.element_type}") + if const_expr(self.c_dtype != dKaccum.element_type): + raise TypeError(f"Type must match: {self.c_dtype} != {dKaccum.element_type}") + + if const_expr(self.varlen): + assert cuSeqlensQ is not None + assert cuSeqlensK is not None + + # For non-varlen, group batch and seqlen modes into token mode + if const_expr(not self.varlen): + batch_divmod = cute.FastDivmodDivisor(dS.shape[0]) + dS = cute.group_modes(dS, 0, 2) + I = cute.group_modes(I, 0, 2) + Q = cute.group_modes(Q, 0, 2) + dKaccum = cute.group_modes(dKaccum, 0, 2) + else: + batch_divmod = cute.FastDivmodDivisor(0) + + # Permute everything to (_, heads, tokens) for MMA + dS_mkl = cute.make_tensor(dS.iterator, cute.select(dS.layout, [2, 1, 0])) + I_ml = cute.make_tensor(I.iterator, cute.select(I.layout, [1, 0])) + Q_nkl = cute.make_tensor(Q.iterator, cute.select(Q.layout, [2, 1, 0])) + dKaccum_nl = cute.make_tensor(dKaccum.iterator, cute.select(dKaccum.layout, [1, 0])) + + # Configure tiled mma + self.tiled_mma_dK = utils.sm100.make_trivial_tiled_mma( + self.ab_dtype, + self.ab_dtype, + self.a_major_mode, + self.b_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler_dK[:2], + ) + + # Compute mma/cluster/tile shapes + mma_inst_shape_k = cute.size(self.tiled_mma_dK.shape_mnk, mode=[2]) + mma_inst_tile_k = 4 if self.heads == 64 else 8 + self.mma_tiler_dK = ( + self.mma_tiler_dK[0], + self.mma_tiler_dK[1], + mma_inst_shape_k * mma_inst_tile_k, + ) + self.cta_tile_shape_mnk_dK = ( + self.mma_tiler_dK[0] // self.atom_thr_size, + self.mma_tiler_dK[1], + self.mma_tiler_dK[2], + ) + + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (self.atom_thr_size,), + ) + + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + + # Compute epilogue subtile shape for TMA store + # self.epi_tile_dK = utils.sm100.compute_epilogue_tile_shape( + # self.cta_tile_shape_mnk_dK, + # self.use_2cta_instrs, + # self.c_layout, + # self.c_dtype, + # ) + self.epi_tile_dK = (cute.make_layout(128), cute.make_layout(32)) + self.epi_tile_dK_width = cute.size(self.epi_tile_dK[1].shape) + + self.dK_smem_layout = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile_dK, 1 + ) + + smem_capacity = utils.get_smem_capacity_in_bytes() + + # TMEM pipeline stages + # 64 TMEM columns per stage, 8 stages can fit in TMEM + self.num_acc_stage = 8 + self.num_tmem_alloc_cols = 512 + + # SMEM pipeline stages + # TMA load Q (persists for the whole token) + self.num_Q_stage = 1 + assert const_expr(self.mma_tiler_dK[1] * self.num_Q_stage == self.dim) + Q_smem_layout_stage_one = utils.sm100.make_smem_layout_b( + self.tiled_mma_dK, self.mma_tiler_dK, self.ab_dtype, 1 + ) + Q_bytes = cute.size_in_bytes(self.ab_dtype, Q_smem_layout_stage_one) * self.num_Q_stage + + # TMA store-reduce dK + self.num_c_stage = 2 + dK_bytes_per_stage = cute.size_in_bytes(self.c_dtype, self.dK_smem_layout) + dK_bytes = dK_bytes_per_stage * self.num_c_stage + + # cp.async load I + self.num_I_stage = 2 + self.I_smem_layout_staged = cute.make_layout( + (self.mma_tiler_dK[0] // self.atom_thr_size, self.num_I_stage) + ) + I_bytes = cute.size_in_bytes(self.idx_dtype, self.I_smem_layout_staged) + + # TMA load dS + dS_smem_layout_stage_one = utils.sm100.make_smem_layout_a( + self.tiled_mma_dK, self.mma_tiler_dK, self.ab_dtype, 1 + ) + dS_bytes_per_stage = cute.size_in_bytes(self.ab_dtype, dS_smem_layout_stage_one) + + mbar_helpers_bytes = 1024 + + # Increase dS stages to fill SMEM + self.num_dS_stage = ( + smem_capacity // self.occupancy - (mbar_helpers_bytes + dK_bytes + Q_bytes + I_bytes) + ) // dS_bytes_per_stage + dS_bytes = self.num_dS_stage * dS_bytes_per_stage + + # Increase dK stages to fill remainder + self.num_c_stage += ( + smem_capacity // self.occupancy + - (mbar_helpers_bytes + dK_bytes + dS_bytes + Q_bytes + I_bytes) + ) // dK_bytes_per_stage + + # Increase I stages to fill remainder + + # Compute shared memory layout + self.dS_smem_layout_staged = utils.sm100.make_smem_layout_a( + self.tiled_mma_dK, self.mma_tiler_dK, self.ab_dtype, self.num_dS_stage + ) + self.Q_smem_layout_staged = utils.sm100.make_smem_layout_b( + self.tiled_mma_dK, self.mma_tiler_dK, self.ab_dtype, self.num_Q_stage + ) + self.dK_smem_layout_staged = utils.sm100.make_smem_layout_epi( + self.c_dtype, self.c_layout, self.epi_tile_dK, self.num_c_stage + ) + + # TMA load for dS + dS_op = utils.sm100.cluster_shape_to_tma_atom_A( + self.cluster_shape_mn, self.tiled_mma_dK.thr_id + ) + self.dS_smem_layout = cute.slice_(self.dS_smem_layout_staged, (None, None, None, 0)) + tma_atom_dS, tma_tensor_dS = cute.nvgpu.make_tiled_tma_atom_A( + dS_op, + dS_mkl, + self.dS_smem_layout, + self.mma_tiler_dK, + self.tiled_mma_dK, + self.cluster_layout_vmnk.shape, + ) + + # TMA load for Q + Q_op = utils.sm100.cluster_shape_to_tma_atom_B( + self.cluster_shape_mn, self.tiled_mma_dK.thr_id + ) + self.Q_smem_layout = cute.slice_(self.Q_smem_layout_staged, (None, None, None, 0)) + tma_atom_Q, tma_tensor_Q = cute.nvgpu.make_tiled_tma_atom_B( + Q_op, + Q_nkl, + self.Q_smem_layout, + self.mma_tiler_dK, + self.tiled_mma_dK, + self.cluster_layout_vmnk.shape, + ) + self.Q_load_bytes = ( + cute.size_in_bytes(self.ab_dtype, self.Q_smem_layout) * self.atom_thr_size + ) + + self.dS_load_bytes = ( + cute.size_in_bytes(self.ab_dtype, self.dS_smem_layout) * self.atom_thr_size + ) + + # coalesced store for dKaccum: 1 warp copies 1 epi tile row + vector_width_dK = self.epi_tile_dK_width * self.c_dtype.width // cute.arch.WARP_SIZE + copy_atom_dK = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.c_dtype, + num_bits_per_copy=vector_width_dK, + ) + vector_elts_dK = vector_width_dK // self.c_dtype.width + thr_layout_dK = cute.make_layout((cute.arch.WARP_SIZE,)) + val_layout_dK = cute.make_layout((vector_elts_dK,)) + tiled_copy_dK = cute.make_tiled_copy_tv(copy_atom_dK, thr_layout_dK, val_layout_dK) + + # cp.async load I: 1 warp copies 32 values, 4 warps copy 128 values + copy_atom_I = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(), + self.idx_dtype, + num_bits_per_copy=32, + ) + vI = cute.make_layout((1,)) + tI = cute.make_layout((128,)) + tiled_copy_I = cute.make_tiled_copy_tv(copy_atom_I, tI, vI) + + # Setup clc stage by default + self.num_clc_stage = 1 + assert self.num_clc_stage == 1, "Only single-stage CLC pipeline is supported" + + # Response size is 4B * 4 elements + self.num_clc_response_bytes = 16 + + # Compute grid size and set up tile scheduler + total_q = cute.size(dS_mkl.shape[2]) + cluster_shape_mnl = (*self.cluster_shape_mn, 1) + num_ctas_mnl = (self.cluster_shape_mn[0] * total_q, self.cluster_shape_mn[1], 1) + self.tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl + ) + grid = utils.ClcDynamicPersistentTileScheduler.get_grid_shape(self.tile_sched_params) + + # Define shared storage for kernel + buffer_align_bytes = 1024 + + @cute.struct + class SharedStorage: + dS_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_dS_stage * 2] + Q_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_Q_stage * 2] + acc_full_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + tmem_dealloc_mbar: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2] + clc_response: cute.struct.MemRange[cutlass.Int32, 4] + sI: cute.struct.Align[ + cute.struct.MemRange[self.idx_dtype, cute.cosize(self.I_smem_layout_staged)], + buffer_align_bytes, + ] + sdS: cute.struct.Align[ + cute.struct.MemRange[self.ab_dtype, cute.cosize(self.dS_smem_layout_staged)], + buffer_align_bytes, + ] + sQ: cute.struct.Align[ + cute.struct.MemRange[self.ab_dtype, cute.cosize(self.Q_smem_layout_staged)], + buffer_align_bytes, + ] + sdK: cute.struct.Align[ + cute.struct.MemRange[self.c_dtype, cute.cosize(self.dK_smem_layout_staged)], + buffer_align_bytes, + ] + + # Launch the kernel synchronously + self.kernel( + self.tiled_mma_dK, + tma_atom_dS, + tma_tensor_dS, + tma_atom_Q, + tma_tensor_Q, + tiled_copy_dK, + dKaccum_nl, + tiled_copy_I, + I_ml, + batch_divmod, + cuSeqlensQ, + cuSeqlensK, + self.cluster_layout_vmnk, + self.dS_smem_layout_staged, + self.Q_smem_layout_staged, + self.dK_smem_layout_staged, + self.I_smem_layout_staged, + self.epi_tile_dK, + self.tile_sched_params, + SharedStorage, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + stream=stream, + ) + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma_dK: cute.TiledMma, + tma_atom_dS: cute.CopyAtom, + mdS_mkl: cute.Tensor, + tma_atom_Q: Optional[cute.CopyAtom], + mQ_nkl: Optional[cute.Tensor], + tiled_copy_dK: cute.TiledCopy, + mdKaccum_nl: Optional[cute.Tensor], + tiled_copy_I: cute.TiledCopy, + mI_ml: cute.Tensor, + batch_divmod: cute.FastDivmodDivisor, + cuSeqlensQ: Optional[cute.Tensor], + cuSeqlensK: Optional[cute.Tensor], + cluster_layout_vmnk: cute.Layout, + dS_smem_layout_staged: cute.ComposedLayout, + Q_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout], + dK_smem_layout_staged: Union[cute.ComposedLayout, cute.Layout], + I_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout], + epi_tile_dK: cute.Tile, + tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams, + SharedStorage: cutlass.Constexpr[Callable], + ): + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # + # Prefetch tma desc + # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_dS) + cpasync.prefetch_descriptor(tma_atom_Q) + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % self.atom_thr_size + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + is_first_cta_in_cluster = cta_rank_in_cluster == 0 + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + + smem = utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # Initialize mainloop ab_pipeline (barrier) and states + dS_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + dS_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + dS_producer, dS_consumer = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.dS_full_mbar_ptr.data_ptr(), + num_stages=self.num_dS_stage, + producer_group=dS_pipeline_producer_group, + consumer_group=dS_pipeline_consumer_group, + tx_count=self.dS_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + + Q_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + Q_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + Q_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.Q_full_mbar_ptr.data_ptr(), + num_stages=self.num_Q_stage, + producer_group=Q_pipeline_producer_group, + consumer_group=Q_pipeline_consumer_group, + tx_count=self.Q_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilogue_warp_id) * (2 if self.use_2cta_instrs else 1) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_full_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Initialize clc_pipeline (barrier) and states + clc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + cluster_size = cute.size(self.cluster_shape_mn) + num_clc_consumer_threads = 32 * (1 + cluster_size * (1 + len(self.epilogue_warp_id) + 1)) + clc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_clc_consumer_threads + ) + clc_pipeline = pipeline.PipelineClcFetchAsync.create( + barrier_storage=storage.clc_mbar_ptr.data_ptr(), + num_stages=self.num_clc_stage, + producer_group=clc_pipeline_producer_group, + consumer_group=clc_pipeline_consumer_group, + tx_count=self.num_clc_response_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_id)), + ) + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_id[0], + is_two_cta=self.use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar.ptr, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + + # Initial clc response pointer + clc_response_ptr = storage.clc_response.data_ptr() + + clc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_clc_stage + ) + + # + # Setup smem tensor A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + sdS = storage.sdS.get_tensor( + dS_smem_layout_staged.outer, + swizzle=dS_smem_layout_staged.inner, + ) + # (MMA_M, STAGE) + sI = storage.sI.get_tensor(I_smem_layout_staged) + + sQ = storage.sQ.get_tensor( + Q_smem_layout_staged.outer, + swizzle=Q_smem_layout_staged.inner, + ) + sdK = storage.sdK.get_tensor( + dK_smem_layout_staged.outer, + swizzle=dK_smem_layout_staged.inner, + ) + + # + # Compute multicast mask for A/B buffer full + # + a_full_mcast_mask = None + b_full_mcast_mask = None + if const_expr(self.is_a_mcast or self.is_b_mcast or self.use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + # + # Local_tile partition global tensors + # + # (bM, bK, RestM, RestK, RestL) + gdS_mkl = cute.local_tile( + mdS_mkl, cute.slice_(self.mma_tiler_dK, (None, 0, None)), (None, None, None) + ) + # (bM, RestM, RestL) + gI_ml = cute.local_tile(mI_ml, cute.slice_(self.mma_tiler_dK, (None, 0, 0)), (None, None)) + # (bN, bK, RestN, RestK, RestL) + gQ_nkl = cute.local_tile( + mQ_nkl, + cute.slice_(self.mma_tiler_dK, (0, None, None)), + (None, None, None), + ) + n_tile_cnt = cute.size(gQ_nkl, mode=[2]) + # (bN, RestN, RestL) + gdKaccum_nl = cute.local_tile( + mdKaccum_nl, + cute.slice_(self.mma_tiler_dK, (0, None, 0)), + (None, None), + ) + + m_tile_cnt = cute.size(gdS_mkl, mode=[2]) + k_tile_cnt = cute.size(gdS_mkl, mode=[3]) + assert const_expr(k_tile_cnt == 1) + + # + # Partition global tensor for TiledMMA_A/B/C + # + thr_mma_dK = tiled_mma_dK.get_slice(mma_tile_coord_v) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgQ = thr_mma_dK.partition_B(gQ_nkl) + + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tCgdKaccum_fake = thr_mma_dK.partition_C( + cute.make_identity_tensor( + cute.append(cute.slice_(self.mma_tiler_dK, (None, None, 0)), 1, up_to_rank=5) + ) + ) + + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tCgdS = thr_mma_dK.partition_A(gdS_mkl) + + # + # Partition global/shared tensor for TMA load A/B + # + # TMA load A partition_S/D + a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tAsdS, tAgdS = cpasync.tma_partition( + tma_atom_dS, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sdS, 0, 3), + cute.group_modes(tCgdS, 0, 3), + ) + # TMA load B partition_S/D + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + # ((atom_v, rest_v), RestM, RestK, RestL) + tBsQ, tBgQ = cpasync.tma_partition( + tma_atom_Q, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sQ, 0, 3), + cute.group_modes(tCgQ, 0, 3), + ) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrdS = tiled_mma_dK.make_fragment_A(sdS) + # (MMA, MMA_N, MMA_K, STAGE) + tCrQ = tiled_mma_dK.make_fragment_B(sQ) + # (MMA, MMA_M, MMA_N) + acc_shape_dK = tiled_mma_dK.partition_shape_C(self.mma_tiler_dK[:2]) + # (MMA, MMA_M, MMA_N, STAGE) + tCtdK_fake = tiled_mma_dK.make_fragment_C(cute.append(acc_shape_dK, self.num_acc_stage)) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + # + # Construct the scheduler + # + tile_sched = utils.ClcDynamicPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + clc_response_ptr, + ) + work_tile = tile_sched.initial_work_tile_info() + + # + # Specialized TMA load warp + # + + if warp_idx == self.tma_warp_id: + Q_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_Q_stage + ) + # + # Persistent tile scheduling loop + # + + dS_producer.reset() + peek_dS_empty_status = dS_producer.try_acquire() + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + total_q_coord = cur_tile_coord[0] // self.cluster_shape_mn[0] + + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), RestM, RestK) + tAgdS_slice = tAgdS[(None, None, None, total_q_coord)] + # ((atom_v, rest_v), RestN, RestK) + tBgQ_slice = tBgQ[(None, None, None, total_q_coord)] + + # + # Tma load loop -- fully unrolled as all sizes are static + # + for m_tile in cutlass.range_constexpr(m_tile_cnt): + # Conditionally wait for AB buffer empty + dS_handle = dS_producer.acquire_and_advance(peek_dS_empty_status) + + # TMA load dS + cute.copy( + tma_atom_dS, + tAgdS_slice[(None, m_tile, 0)], + tAsdS[(None, dS_handle.index)], + tma_bar_ptr=dS_handle.barrier, + mcast_mask=a_full_mcast_mask, + ) + + peek_dS_empty_status = dS_producer.try_acquire() + + if m_tile == 0: + # TMA load Q only on first m-tile + for n_tile in cutlass.range_constexpr(n_tile_cnt): + Q_pipeline.producer_acquire(Q_producer_state) + Q_load_barrier = Q_pipeline.producer_get_barrier(Q_producer_state) + cute.copy( + tma_atom_Q, + tBgQ_slice[(None, n_tile, 0)], + tBsQ[(None, n_tile)], + tma_bar_ptr=Q_load_barrier, + mcast_mask=b_full_mcast_mask, + ) + Q_pipeline.producer_commit(Q_producer_state) + Q_producer_state.advance() + + # + # Advance to next tile + # + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + # + # Wait A/B buffer empty + # + dS_producer.tail() + Q_pipeline.producer_tail(Q_producer_state) + + # + # Sched warp + # + elif warp_idx == self.sched_warp_id and is_first_cta_in_cluster: + # + # Persistent tile scheduling loop + # + clc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.ProducerConsumer, self.num_clc_stage + ) + + while work_tile.is_valid_tile: + # + # Advance to next tile + # + clc_pipeline.producer_acquire(clc_producer_state) + mbarrier_addr = clc_pipeline.producer_get_barrier(clc_producer_state) + tile_sched.advance_to_next_work(mbarrier_addr) + clc_producer_state.advance() + + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + clc_pipeline.producer_tail(clc_producer_state) + + # + # Specialized MMA warp + # + elif warp_idx == self.mma_warp_id: + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtdK_base = cute.make_tensor(tmem_ptr, tCtdK_fake.layout) + + # + # Persistent tile scheduling loop + # + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + Q_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_Q_stage + ) + + # Peek (try_wait) AB buffer full for k_tile = 0 + dS_consumer.reset() + peek_dS_full_status = cutlass.Boolean(1) + if is_leader_cta: + peek_dS_full_status = dS_consumer.try_wait() + + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + total_q_coord = cur_tile_coord[0] // self.cluster_shape_mn[0] + + # + # Mma mainloop + # + Q_consumer_state_wait = Q_consumer_state.clone() + for m_tile in cutlass.range(m_tile_cnt): + if is_leader_cta: + # Conditionally wait for AB buffer full + dS_handle = dS_consumer.wait_and_advance(peek_dS_full_status) + for n_tile in cutlass.range_constexpr(n_tile_cnt): + tCtdK = tCtdK_base[(None, None, None, acc_producer_state.index)] + tiled_mma_dK.set(tcgen05.Field.ACCUMULATE, False) + acc_pipeline.producer_acquire(acc_producer_state) + + if m_tile == 0: + Q_pipeline.consumer_wait(Q_consumer_state_wait) + Q_consumer_state_wait.advance() + + num_kblocks = cute.size(tCrdS, mode=[2]) + for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): + cute.gemm( + tiled_mma_dK, + tCtdK, + tCrdS[(None, None, kblk_idx, dS_handle.index)], + tCrQ[(None, None, kblk_idx, n_tile)], + tCtdK, + ) + # Enable accumulate on tCtdK after first kblock + tiled_mma_dK.set(tcgen05.Field.ACCUMULATE, True) + + if m_tile == m_tile_cnt - 1: + Q_pipeline.consumer_release(Q_consumer_state) + Q_consumer_state.advance() + + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + # Async arrive AB buffer empty + dS_handle.release() + + peek_dS_full_status = dS_consumer.try_wait() + else: + for n_tile in cutlass.range_constexpr(n_tile_cnt): + acc_producer_state.advance() + + # + # Advance to next tile + # + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + + # + # Specialized epilogue warps + # + elif warp_idx < self.mma_warp_id: + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtdK_base = cute.make_tensor(tmem_ptr, tCtdK_fake.layout) + + # Both gemms share accumulator and TMA store pipelines + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilogue_warp_id), + ) + + I_load_barrier = pipeline.NamedBarrier( + barrier_id=self.I_load_bar_id, + num_threads=32 * len(self.epilogue_warp_id), + ) + + # (EPI_TOPK, REST_TOPK, TOKENS) + gI_tile = cute.local_tile( + gI_ml, + (epi_tile_dK[0],), + (0 if is_leader_cta else 1, None, None), + ) + + thr_copy_I = tiled_copy_I.get_slice(tidx) + # (COPY_ATOM, EPI_TOPK, REST_TOPK, TOKENS) + tIgI = thr_copy_I.partition_S(gI_tile) + # (COPY_ATOM, EPI_TOPK, STAGE) + tIsI = thr_copy_I.partition_D(sI) + + num_subtiles_executed = 0 + while work_tile.is_valid_tile: + # Get tile coord from tile scheduler + cur_tile_coord = work_tile.tile_idx + token = cutlass.Int64(cur_tile_coord[0]) // self.cluster_shape_mn[0] + + # Get batch index from token_coord and cuSeqlensQ + # Get seqlen_k offset from cuSeqlensK + if const_expr(self.varlen): + batch = get_batch_from_cu_tensor(token, cuSeqlensQ) + seqlen_k_offset = cuSeqlensK[batch] + else: + _, batch = divmod(token, batch_divmod) + seqlen_k_offset = Int32(0) # unused + + sI_read_stage = 0 + sI_write_stage = 0 + + # Prefetch load I + I_load_barrier.arrive_and_wait() + for m_tile in cutlass.range_constexpr(min(m_tile_cnt, self.num_I_stage - 1)): + cute.copy( + tiled_copy_I, + tIgI[(None, None, m_tile, token)], + tIsI[(None, None, sI_write_stage)], + ) + cute.arch.cp_async_commit_group() + sI_write_stage = (sI_write_stage + 1) % self.num_I_stage + + for m_tile in cutlass.range(m_tile_cnt): + I_load_barrier.arrive_and_wait() + # cp.async load I + if m_tile < m_tile_cnt - 1: + cute.copy( + tiled_copy_I, + tIgI[(None, None, m_tile + 1, token)], + tIsI[(None, None, sI_write_stage)], + ) + cute.arch.cp_async_commit_group() + sI_write_stage = (sI_write_stage + 1) % self.num_I_stage + cute.arch.cp_async_wait_group(self.num_I_stage - 1) + I_load_barrier.arrive_and_wait() + + sI_tile = sI[(None, sI_read_stage)] + + for n_tile in cutlass.range_constexpr(n_tile_cnt): + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + mma_tile_coord_mnl = ( + m_tile, + n_tile, + token, + ) + acc_consumer_state = self.epilogue_scatter_reduce( + tidx, + 0 if is_leader_cta else 1, + tiled_copy_dK, + tCtdK_base, + sdK, + sI_tile, + gdKaccum_nl, + tCgdKaccum_fake, + batch, + seqlen_k_offset, + epi_tile_dK, + num_subtiles_executed, + mma_tile_coord_mnl, + acc_consumer_state, + acc_pipeline, + ) + num_subtiles_executed += self.mma_tiler_dK[1] // cute.size(epi_tile_dK[1]) + # Advance I consumer pipeline + sI_read_stage = (sI_read_stage + 1) % self.num_I_stage + + cute.arch.cp_async_wait_group(0) + # + # Advance to next tile + # + clc_pipeline.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + clc_pipeline.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + + # # Wait for C store complete + # dK_pipeline.producer_tail() + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + @cute.jit + def epilogue_scatter_reduce( + self, + epi_tidx: Int32, + cta_idx: Int32, + tiled_copy_c: cute.CopyAtom, + tCtAcc_base: cute.Tensor, + sC: cute.Tensor, + sI_tile: cute.Tensor, + gC_base: cute.Tensor, + tCgC_fake: cute.Tensor, + batch: Int32, + seqlen_k_offset: Int32, + epi_tile: cute.Tile, + num_subtiles_executed: Int32, + mma_tile_coord_mnl: Tuple[Int32, Int32, cutlass.Int64], + acc_consumer_state: pipeline.PipelineState, + acc_pipeline: pipeline.PipelineAsync, + ) -> pipeline.PipelineState: + warp_idx = cute.arch.make_warp_uniform(epi_tidx // 32) + + # Layout transformation for tCgC_base + # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, REST_M, REST_N, REST_L) + # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), REST_M, REST_N, REST_L) + tCgC_fake = transform_partitioned_tensor_layout(tCgC_fake) + + # Layout transformation for tCtAcc_base + # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, STAGE) + # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), STAGE) + tCtAcc = transform_partitioned_tensor_layout(tCtAcc_base) + + tiled_copy_t2r, tTR_tAcc_base, tTR_rAcc = epilogue_tmem_copy_and_partition( + self, + epi_tidx, + tCtAcc, + tCgC_fake, + epi_tile, + self.use_2cta_instrs, + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = epilogue_smem_copy_and_partition( + self, tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + + # (EPI_TILE_N, EPI_N, RestN, RestL) + gC_epi = cute.flat_divide(gC_base, (epi_tile[1],)) + # (EPI_TILE_N, EPI_N, SEQLEN_K) + gC_epi = gC_epi[None, None, mma_tile_coord_mnl[1], None] + # (EPI_TILE_N, MMA_M, STAGE) + sC_epi = cute.make_tensor( + sC.iterator, + cute.select(sC.layout, [1, 0, 2]), + # swizzle=sC.layout.inner, + ) + + thr_copy_c = tiled_copy_c.get_slice(epi_tidx % 32) + # (COPY_ATOM_N, COPY_N, MMA_M, STAGE) + tCsC = thr_copy_c.partition_S(sC_epi) + # (COPY_ATOM_N, COPY_N) + tCrC = cute.make_fragment_like(tCsC[None, None, 0, 0]) + # (COPY_ATOM_N, COPY_N, EPI_N, SEQLEN_K) + tCgC = thr_copy_c.partition_D(gC_epi) + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_consumer_state.index)] + + # + # Wait for accumulator buffer full + # + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + + epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.epilog_sync_bar_id, + num_threads=32 * len(self.epilogue_warp_id), + ) + + # + # Store accumulator to global memory in subtiles + # + subtile_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + for subtile_idx in range(subtile_cnt): + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn = tTR_tAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec = acc_vec.to(self.c_dtype) + tRS_rC.store(acc_vec) + + # + # Store C to shared memory + # + c_buffer = (num_subtiles_executed + subtile_idx) % self.num_c_stage + cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) + + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy("async.shared", space="cta") + epilog_sync_barrier.arrive_and_wait() + + # + # TMA store C to global memory (issued by lane 0 from all epi warps) + # + for topk_idx_in_warp in cutlass.range(32): + topk_idx = topk_idx_in_warp + warp_idx * 32 + seqlen_k_idx_in_batch = sI_tile[topk_idx] + if const_expr(self.varlen): + seqlen_k_idx = seqlen_k_idx_in_batch + seqlen_k_offset + else: + seqlen_k_idx = (batch, seqlen_k_idx_in_batch) + cute.copy(tiled_copy_c, tCsC[(None, None, topk_idx, c_buffer)], tCrC) + for j in cutlass.range_constexpr(cute.size(tCrC, mode=[1])): + for i in cutlass.range_constexpr(cute.size(tCrC, mode=[0])): + ptr = elem_pointer(tCgC, (i, j, subtile_idx, seqlen_k_idx)) + cute.arch.atomic_add( + ptr=ptr, + val=tCrC[i, j], + ) + epilog_sync_barrier.arrive_and_wait() + + epilog_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty + # + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + return acc_consumer_state + + def check_can_implement(self): + """Check if parameters are valid. + + :raises testing.CantImplementError: If the mma tiler, cluster shape, or alignments are invalid + """ + if self.dim != 64: + raise testing.CantImplementError(f"Only dim = 64 supported for dK gemm, got {self.dim}") + # Check valid MMA tile shape and cluster shape + if not ( + (not self.use_2cta_instrs and self.mma_tiler_dK[0] in [64, 128]) + or (self.use_2cta_instrs and self.mma_tiler_dK[0] in [128, 256]) + ): + raise testing.CantImplementError( + f"Invalid mma tiler & use_2cta_instrs: {self.mma_tiler_dK}, {self.use_2cta_instrs}" + ) + if self.mma_tiler_dK[1] not in range(32, 257, 32): + raise testing.CantImplementError(f"Invalid mma tiler N: {self.mma_tiler_dK[1]}") + # Skip illegal cluster shape + if self.cluster_shape_mn[0] % (2 if self.use_2cta_instrs else 1) != 0: + raise testing.CantImplementError(f"Invalid cluster shape M: {self.cluster_shape_mn[0]}") + # Skip invalid cluster shape + is_power_of_2 = lambda x: x > 0 and (x & (x - 1)) == 0 + if ( + self.cluster_shape_mn[0] * self.cluster_shape_mn[1] > 16 + or self.cluster_shape_mn[0] <= 0 + or self.cluster_shape_mn[1] <= 0 + or not is_power_of_2(self.cluster_shape_mn[0]) + or not is_power_of_2(self.cluster_shape_mn[1]) + ): + raise testing.CantImplementError(f"Invalid cluster shape: {self.cluster_shape_mn}") + + # Check that all tensors are 16B aligned for TMA + + def check_contiguous_16B_alignment(dtype, num_major_elements): + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contiguous_16B_alignment(self.ab_dtype, self.topk) + or not check_contiguous_16B_alignment(self.ab_dtype, self.dim) + or not check_contiguous_16B_alignment(self.c_dtype, self.dim) + ): + raise testing.CantImplementError( + f"Invalid tensor alignment: {self.ab_dtype=}, {self.c_dtype=}, {self.topk=}, {self.dim=}" + ) diff --git a/fastvideo-kernel/fa4/flash_bwd_mla_dq_dqv_sm100.py b/fastvideo-kernel/fa4/flash_bwd_mla_dq_dqv_sm100.py new file mode 100644 index 000000000..5014de6a1 --- /dev/null +++ b/fastvideo-kernel/fa4/flash_bwd_mla_dq_dqv_sm100.py @@ -0,0 +1,1297 @@ +# Copyright (c) 2026, Colfax International. + +""" +CuTe DSL implementation of dQ+dQv gemm for DSA backward. +Performs both dQ = dS @ K and dQv = dS @ V, where K and V are +gathered according to index tensor mIdxTopK. + +This uses MQA with 128 heads. + +Inputs: + - dS: [batch, seqlen_q, nheads, top_k] or [total_q, nheads, top_k] + - K: [batch, seqlen_k, hdim] or [total_k, hdim] + - V: [batch, seqlen_k, hdim_v] or [total_k, hdim_v] + - IdxTopK: [batch, seqlen_q, top_k] or [total_q, top_k] + +Outputs: + - dQ: [batch, seqlen_q, nheads, hdim] or [total_q, nheads, hdim] + - dQv: [batch, seqlen_q, nheads, hdim_v] or [total_q, nheads, hdim_v] + +All sizes are known at compile time except seqlen_q, which is the batch dimension. +Representative numbers are: + - nheads = 128 + - hdim = 64 + - hdim_v = 512 + - top_k = 2048 + +We launch a cluster of shape (1, 2) with mma tile 128x256, so that one cluster +covers the full dQv mma. dS is loaded via TMA and multicast across the CTAs. +Cluster 0 also performs the dQ mma with tile size 128x64. + +K and V are loaded via CpAsync, with logic according to CpasyncGatherKVManager. +""" + +from functools import partial +from typing import Optional, Tuple, Type + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.utils as utils +from cutlass import Int32, const_expr +from cutlass.cute import FastDivmodDivisor +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from flash_attn.cute.topk_gather_kv import CpasyncGatherKVManager +from flash_attn.cute.utils import get_batch_from_cu_tensor + + +class dQdQvGemmKernel: + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + nheads: int, + head_dim_k: Optional[int], + head_dim_v: int, + top_k: int, + ): + self.acc_dtype: Type[cutlass.Numeric] = acc_dtype + self.nheads = nheads + assert self.nheads == 128, ( + "only 128 heads supported; will expand to include 64 heads in a future PR." + ) + self.head_dim_k = head_dim_k or 0 # when head_dim_k not provided, dQ is not computed + self.head_dim_v = head_dim_v + self.top_k = top_k + self.tile_k = 128 + + self.cluster_shape_mn = (1, 2) + self.mma_tiler_dQ = (self.nheads, self.head_dim_k, self.tile_k) + self.mma_tiler_dQv = (self.nheads, self.head_dim_v // 2, self.tile_k) + self.num_mainloop_iters = self.top_k // self.tile_k + self.arch = "sm_100" + + self.cta_group = tcgen05.CtaGroup.ONE + + self.occupancy = 1 + self.threads_per_warp = cute.arch.WARP_SIZE + + # ---- Set specialized warp ids ---- + self.epilogue_warp_ids = (0, 1, 2, 3) + self.kv_load_warp_ids = (4, 5, 6, 7) + self.mma_warp_id = 8 + self.tma_warp_id = 9 + self.sched_warp_id = 10 + self.threads_per_cta = 32 * len( + ( + self.mma_warp_id, + self.tma_warp_id, + self.sched_warp_id, + *self.epilogue_warp_ids, + *self.kv_load_warp_ids, + ) + ) + # ---- Set barrier id for cta sync, epilogue sync and tmem ptr sync ---- + self.epilog_sync_bar_id = 1 + self.tmem_alloc_sync_bar_id = 2 + self.tmem_dealloc_sync_bar_id = 3 + self.kv_load_sync_bar_id = 4 + + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.epilog_sync_bar_id, + num_threads=self.threads_per_warp * len(self.epilogue_warp_ids), + ) + self.kv_load_sync_barrier = pipeline.NamedBarrier( + barrier_id=self.kv_load_sync_bar_id, + num_threads=self.threads_per_warp * len(self.kv_load_warp_ids), + ) + + self.is_persistent = False + + # ---- pipeline stages ---- TODO: tune these + self.num_stages_dS = 2 + self.num_stages_KV = 2 + self.num_stages_acc = 1 + self.num_stages_clc = 1 + + # ---- register allocation ---- + self.num_regs_KV = 224 + self.num_regs_epi = 128 + self.num_regs_other = 112 + + @cute.jit + def __call__( + self, + mdS: cute.Tensor, + mK: Optional[cute.Tensor], + mV: cute.Tensor, + mdQ: cute.Tensor, + mdQv: cute.Tensor, + mIdxTopK: cute.Tensor, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + self.compute_dQ = const_expr(mK is not None) + + # ---- dtype info ---- + self.ds_dtype: Type[cutlass.Numeric] = mdS.element_type + self.kv_dtype: Type[cutlass.Numeric] = mV.element_type + self.dq_dtype: Type[cutlass.Numeric] = mdQv.element_type + if const_expr(self.compute_dQ): + assert self.kv_dtype == mV.element_type + assert self.dq_dtype == mdQ.element_type + + varlen_q = const_expr(mCuSeqlensQ is not None) + varlen_k = const_expr(mCuSeqlensK is not None) + + # ------------------------------------------------------------------ # + # Reshape GMEM layouts for static strides # + # ------------------------------------------------------------------ # + seqlen_q = Int32(0) if const_expr(varlen_q) else mdS.shape[1] + seqlen_q_divmod = FastDivmodDivisor(seqlen_q) + seqlen_k = Int32(0) if const_expr(varlen_k) else mV.shape[1] + + # ---- group batch and seqlen modes in nonvarlen case ---- + def group_batch_seqlen(t: cute.Tensor, varlen: bool) -> cute.Tensor: + if const_expr(not varlen): + t = cute.make_tensor( + t.iterator, + cute.make_layout( + (t.shape[1], t.shape[0], *t.shape[2:]), + stride=(t.stride[1], t.stride[0], *t.stride[2:]), + ), + ) + t = cute.group_modes(t, 0, 2) + return t + + mdS = group_batch_seqlen(mdS, varlen_q) + mdQv = group_batch_seqlen(mdQv, varlen_q) + mV = group_batch_seqlen(mV, varlen_k) + mIdxTopK = group_batch_seqlen(mIdxTopK, varlen_q) + if const_expr(self.compute_dQ): + mdQ = group_batch_seqlen(mdQ, varlen_q) + mK = group_batch_seqlen(mK, varlen_k) + + # ---- transpose and make static modes static ---- + def static_reshape(t: cute.Tensor, *static_shapes) -> cute.Tensor: + static_modes = range(1, len(t.shape)) + return cute.make_tensor( + t.iterator, + cute.make_layout( + (*static_shapes, t.shape[0]), + stride=(*(t.stride[i] for i in static_modes), t.stride[0]), + ), + ) + + mdS = static_reshape(mdS, self.nheads, self.top_k) + mdQv = static_reshape(mdQv, self.nheads, self.head_dim_v) + mV = static_reshape(mV, self.head_dim_v) + mIdxTopK = static_reshape(mIdxTopK, self.top_k) + if const_expr(self.compute_dQ): + mdQ = static_reshape(mdQ, self.nheads, self.head_dim_k) + mK = static_reshape(mK, self.head_dim_k) + + # ---- layout info ---- + self.ds_major_mode = utils.LayoutEnum.from_tensor(mdS).mma_major_mode() + self.kv_major_mode = utils.LayoutEnum.from_tensor(mV).mma_major_mode() + self.dq_layout = utils.LayoutEnum.from_tensor(mdQv) + if const_expr(self.compute_dQ): + assert self.dq_layout == utils.LayoutEnum.from_tensor(mdQ) + + # ------------------------------------------------------------------ # + # Setup attributes that depend on kernel inputs # + # ------------------------------------------------------------------ # + tiled_mma_v = utils.sm100.make_trivial_tiled_mma( + self.ds_dtype, + self.ds_major_mode, + self.kv_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler_dQv[:2], + ) + if const_expr(self.compute_dQ): + tiled_mma_k = utils.sm100.make_trivial_tiled_mma( + self.ds_dtype, + self.ds_major_mode, + self.kv_major_mode, + self.acc_dtype, + self.cta_group, + self.mma_tiler_dQ[:2], + ) + + self.cta_tile_shape_dQv = ( + self.mma_tiler_dQv[0], + self.mma_tiler_dQv[1], + self.mma_tiler_dQv[2], + ) + + # ---- Compute cluster layout ---- + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma_v.thr_id.shape,), + ) + + # ---- Compute number of multicast CTAs for A/B ---- + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + + # ---- Compute epi tiles for dQ/dQv ---- + self.epi_tile_dQv = utils.sm100.compute_epilogue_tile_shape( + self.cta_tile_shape_dQv, + False, # use_2cta_instrs + self.dq_layout, + self.dq_dtype, + ) + self.epi_tile_dQ = None + if const_expr(self.compute_dQ): + self.epi_tile_dQ = utils.sm100.compute_epilogue_tile_shape( + self.mma_tiler_dQ, + False, + self.dq_layout, + self.dq_dtype, + ) + + # ---- Device-specific attributes ---- + self.smem_capacity = utils.get_smem_capacity_in_bytes() + + self.num_tmem_alloc_cols = 512 + + # ------------------------------------------------------------------ # + # Make SMEM layouts # + # ------------------------------------------------------------------ # + sdS_layout = utils.sm100.make_smem_layout_a( + tiled_mma_v, + self.mma_tiler_dQ, + self.ds_dtype, + self.num_stages_dS, + ) + sV_layout = utils.sm100.make_smem_layout_b( + tiled_mma_v, + self.mma_tiler_dQv, + self.kv_dtype, + self.num_stages_KV, + ) + sdQv_layout = utils.sm100.make_smem_layout_epi( + self.dq_dtype, + self.dq_layout, + self.epi_tile_dQv, + self.num_stages_acc, + ) + sK_layout, sdQ_layout = None, None + if const_expr(self.compute_dQ): + sK_layout = utils.sm100.make_smem_layout_b( + tiled_mma_k, + self.mma_tiler_dQ, + self.kv_dtype, + self.num_stages_KV, + ) + sdQ_layout = utils.sm100.make_smem_layout_epi( + self.dq_dtype, + self.dq_layout, + self.epi_tile_dQ, + self.num_stages_acc, + ) + + # ------------------------------------------------------------------ # + # Set up TMA load/stores # + # ------------------------------------------------------------------ # + atom_thr_size = cute.size(tiled_mma_v.thr_id.shape) + + # ---- Setup TMA load for dS ---- + dS_op = utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma_v.thr_id) + dS_smem_layout = cute.slice_(sdS_layout, (None, None, None, 0)) + tma_atom_dS, tma_tensor_dS = cute.nvgpu.make_tiled_tma_atom_A( + dS_op, + mdS, + dS_smem_layout, + self.mma_tiler_dQv, + tiled_mma_v, + self.cluster_layout_vmnk.shape, + ) + + dS_copy_size = cute.size_in_bytes(self.ds_dtype, dS_smem_layout) + self.num_tma_load_bytes = dS_copy_size * atom_thr_size + + # ---- Setup TMA store for dQ and dQV ---- + dQv_epi_smem_layout = cute.select(sdQv_layout, mode=[0, 1]) + tma_atom_dQv, tma_tensor_dQv = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), mdQv, dQv_epi_smem_layout, self.epi_tile_dQv + ) + tma_atom_dQ, tma_tensor_dQ = None, None + if const_expr(self.compute_dQ): + dQ_epi_smem_layout = cute.select(sdQ_layout, mode=[0, 1]) + tma_atom_dQ, tma_tensor_dQ = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), mdQ, dQ_epi_smem_layout, self.epi_tile_dQ + ) + + # ------------------------------------------------------------------ # + # Set up shared storage for SMEM # + # ------------------------------------------------------------------ # + + self.buffer_align_bytes = 1024 + + sdS_size = cute.cosize(sdS_layout) + sK_size = cute.cosize(sK_layout) if const_expr(self.compute_dQ) else 0 + sV_size = cute.cosize(sV_layout) + sdQ_size = cute.cosize(sdQ_layout) if const_expr(self.compute_dQ) else 0 + sdQv_size = cute.cosize(sdQv_layout) + assert sdQ_size <= sK_size, f"require {sdQ_size=} <= {sK_size=}" + assert sdQv_size <= sV_size, f"require {sdQv_size=} <= {sV_size=}" + + self.overlap_kv_epi = self.compute_dQ + if const_expr(self.overlap_kv_epi): + sdQ_size = 0 + sdQv_size = 0 + + @cute.struct + class SharedStorage: + mbar_ptr_dS: cute.struct.MemRange[cutlass.Int64, self.num_stages_dS * 2] + mbar_ptr_KV: cute.struct.MemRange[cutlass.Int64, self.num_stages_KV * 2] + mbar_ptr_dQ_dQv: cute.struct.MemRange[cutlass.Int64, self.num_stages_acc * 2] + mbar_ptr_KV_cpasync: cute.struct.MemRange[cutlass.Int64, self.num_stages_KV * 2] + mbar_ptr_load_kv_epi: cute.struct.MemRange[cutlass.Int64, 2] + # Tmem holding buffer + mbar_ptr_tmem_dealloc: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + # Clc pointers + clc_ptr: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int64, self.num_stages_clc * 2], 16 + ] + clc_response_ptr: cute.struct.Align[cute.struct.MemRange[cutlass.Int32, 4], 16] + # Smem tensors + sdS: cute.struct.Align[ + cute.struct.MemRange[self.ds_dtype, sdS_size], + self.buffer_align_bytes, + ] + sK: cute.struct.Align[ + cute.struct.MemRange[self.kv_dtype, sK_size], + self.buffer_align_bytes, + ] + sV: cute.struct.Align[ + cute.struct.MemRange[self.kv_dtype, sV_size], + self.buffer_align_bytes, + ] + sdQ: cute.struct.Align[ + cute.struct.MemRange[self.dq_dtype, sdQ_size], + self.buffer_align_bytes, + ] + sdQv: cute.struct.Align[ + cute.struct.MemRange[self.dq_dtype, sdQv_size], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + # ---- Compute grid size ---- + self.tile_sched_params, grid = self._compute_grid( + mdQv, + self.cta_tile_shape_dQv, + self.cluster_shape_mn, + ) + self.num_clc_response_bytes = 16 + # permute grid to conform to grid_dim_z <= 65536 constraint; + # this is undone in the kernel + grid = (grid[2], grid[1], grid[0]) + + # cute.printf("dQ/dQv grid: {}", grid) + # print("dQ/dQv SMEM: ", self.shared_storage.size_in_bytes()) + # ---- Launch the kernel synchronously ---- + self.kernel( + tiled_mma_k if const_expr(self.compute_dQ) else None, + tiled_mma_v, + tma_atom_dS, + tma_tensor_dS, + mK, + mV, + tma_atom_dQ, + tma_tensor_dQ, + tma_atom_dQv, + tma_tensor_dQv, + mIdxTopK, + mCuSeqlensQ, + mCuSeqlensK, + seqlen_q_divmod, + self.cluster_layout_vmnk, + sdS_layout, + sK_layout, + sV_layout, + sdQ_layout, + sdQv_layout, + self.epi_tile_dQ, + self.epi_tile_dQv, + self.tile_sched_params, + seqlen_k, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + ) + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma_k: Optional[cute.TiledMma], + tiled_mma_v: cute.TiledMma, + tma_atom_dS: cute.CopyAtom, + mdS: cute.Tensor, + mK: Optional[cute.Tensor], + mV: cute.Tensor, + tma_atom_dQ: Optional[cute.CopyAtom], + mdQ: Optional[cute.Tensor], + tma_atom_dQv: cute.CopyAtom, + mdQv: cute.Tensor, + mIdxTopK: cute.Tensor, + mCuSeqlensQ: cute.Tensor, + mCuSeqlensK: cute.Tensor, + seqlen_q_divmod: FastDivmodDivisor, + cluster_layout_vmnk: cute.Layout, + sdS_layout: cute.ComposedLayout, + sK_layout: Optional[cute.ComposedLayout], + sV_layout: cute.ComposedLayout, + sdQ_layout: Optional[cute.ComposedLayout], + sdQv_layout: cute.ComposedLayout, + epi_tile_dQ: Optional[cute.Tile], + epi_tile_dQv: cute.Tile, + tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams, + seqlen_k_static: Int32, + ): + """ + GPU device kernel performing the Persistent batched GEMM computation. + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # ------------------------------------------------------------------ # + # Prefetch TMA descriptors # + # ------------------------------------------------------------------ # + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_dS) + cpasync.prefetch_descriptor(tma_atom_dQv) + if const_expr(self.compute_dQ): + cpasync.prefetch_descriptor(tma_atom_dQ) + + # ------------------------------------------------------------------ # + # Cluster coordinates # + # ------------------------------------------------------------------ # + bidx, bidy, bidz = cute.arch.block_idx() + gridx, gridy, gridz = cute.arch.grid_dim() + mma_v_tile_coord_v = bidx % cute.size(tiled_mma_v.thr_id.shape) + if const_expr(self.compute_dQ): + mma_k_tile_coord_v = bidx % cute.size(tiled_mma_k.thr_id.shape) + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + is_first_cta = cta_rank_in_cluster == 0 + tidx, _, _ = cute.arch.thread_idx() + + # ------------------------------------------------------------------ # + # Shared storage allocation # + # ------------------------------------------------------------------ # + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # ------------------------------------------------------------------ # + # Initialize pipelines # + # ------------------------------------------------------------------ # + ThreadCooperativeGroup = partial(pipeline.CooperativeGroup, pipeline.Agent.Thread) + kv_commit_group = ThreadCooperativeGroup(1) + clc_producer_group = ThreadCooperativeGroup(1) + num_clc_consumer_threads = 32 * ( + 1 # sched warp on CTA0 only + + cute.size(self.cluster_shape_mn) + * (1 + len(self.epilogue_warp_ids) + len(self.kv_load_warp_ids) + 1) + # tma + epi + kv_load + mma, on BOTH CTAs + ) + clc_consumer_group = ThreadCooperativeGroup(num_clc_consumer_threads) + mma_warp = ThreadCooperativeGroup(1) + tma_warp = ThreadCooperativeGroup(cute.size(self.cluster_shape_mn)) + tma_warp_local = ThreadCooperativeGroup(1) + epilogue_warps = ThreadCooperativeGroup(len(self.epilogue_warp_ids)) + load_warps = ThreadCooperativeGroup(len(self.kv_load_warp_ids)) + + pipeline_dS = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.mbar_ptr_dS.data_ptr(), + num_stages=self.num_stages_dS, + producer_group=mma_warp, + consumer_group=tma_warp, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + pipeline_KV = pipeline.PipelineAsyncUmma.create( + barrier_storage=storage.mbar_ptr_KV.data_ptr(), + num_stages=self.num_stages_KV, + producer_group=kv_commit_group, + consumer_group=mma_warp, + defer_sync=True, + ) + + pipeline_dQ_dQv = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.mbar_ptr_dQ_dQv.data_ptr(), + num_stages=self.num_stages_acc, + producer_group=mma_warp, + consumer_group=epilogue_warps, + defer_sync=True, + ) + + pipeline_clc = pipeline.PipelineClcFetchAsync.create( + barrier_storage=storage.clc_ptr.data_ptr(), + num_stages=self.num_stages_clc, + producer_group=clc_producer_group, + consumer_group=clc_consumer_group, + tx_count=self.num_clc_response_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + pipeline_load_kv_epi = None + if const_expr(self.overlap_kv_epi): + pipeline_load_kv_epi = pipeline.PipelineAsync.create( + barrier_storage=storage.mbar_ptr_load_kv_epi.data_ptr(), + num_stages=1, + producer_group=epilogue_warps, + consumer_group=load_warps, + defer_sync=True, + ) + + # ------------------------------------------------------------------ # + # TMEM Allocation # + # ------------------------------------------------------------------ # + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=32 * len((self.mma_warp_id, *self.epilogue_warp_ids)), + ) + # ---- Tensor memory dealloc barrier init ---- + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_ids[0], + is_two_cta=False, + two_cta_tmem_dealloc_mbar_ptr=storage.mbar_ptr_tmem_dealloc, + ) + + # ---- Cluster arrive after barrier init ---- + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + + # ---- Initial clc response pointer ---- + clc_response_ptr = storage.clc_response_ptr.data_ptr() + + clc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_stages_clc + ) + + # ------------------------------------------------------------------ # + # SMEM tensors # + # ------------------------------------------------------------------ # + # (MMA, MMA_M, MMA_K, STAGE) + sdS = storage.sdS.get_tensor(sdS_layout.outer, swizzle=sdS_layout.inner) + if const_expr(self.compute_dQ): + sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner) + sV = storage.sV.get_tensor(sV_layout.outer, swizzle=sV_layout.inner) + # sdQ = storage.sdQ.get_tensor(sdQ_layout.outer, swizzle=sdQ_layout.inner) + if const_expr(self.compute_dQ): + sdQ = cute.make_tensor( + cute.recast_ptr(sK.iterator, sdQ_layout.inner, self.dq_dtype), sdQ_layout.outer + ) + if const_expr(not self.compute_dQ): + sdQv = storage.sdQv.get_tensor(sdQv_layout.outer, swizzle=sdQv_layout.inner) + else: + sdQv = cute.make_tensor( + cute.recast_ptr(sV.iterator, sdQv_layout.inner, self.dq_dtype), sdQv_layout.outer + ) + + dS_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + + # ------------------------------------------------------------------ # + # Global tile partitioning # + # ------------------------------------------------------------------ # + # (bM, bK, RestM, RestK, RestL) + gdS = cute.local_tile( + mdS, cute.slice_(self.mma_tiler_dQv, (None, 0, None)), (None, None, None) + ) + # (bM, bN, RestM, RestN, RestL) + if const_expr(self.compute_dQ): + gdQ = cute.local_tile( + mdQ, cute.slice_(self.mma_tiler_dQ, (None, None, 0)), (None, None, None) + ) + gdQv = cute.local_tile( + mdQv, cute.slice_(self.mma_tiler_dQv, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cute.size(gdS, mode=[3]) + + # ------------------------------------------------------------------ # + # TiledMMA partitioning # + # ------------------------------------------------------------------ # + thr_mma_v = tiled_mma_v.get_slice(mma_v_tile_coord_v) + # (MMA, MMA_M, MMA_K, RestM, RestK, RestL) + tdQvgdS = thr_mma_v.partition_A(gdS) + # (MMA, MMA_M, MMA_N, RestM, RestN, RestL) + tdQvgdQv = thr_mma_v.partition_C(gdQv) + if const_expr(self.compute_dQ): + thr_mma_k = tiled_mma_k.get_slice(mma_k_tile_coord_v) + tdQgdQ = thr_mma_k.partition_C(gdQ) + + # ------------------------------------------------------------------ # + # TMA partition for dS # + # ------------------------------------------------------------------ # + dS_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestM, RestK, RestL) + tdSsdS, tdSgdS = cpasync.tma_partition( + tma_atom_dS, + block_in_cluster_coord_vmnk[2], + dS_cta_layout, + cute.group_modes(sdS, 0, 3), + cute.group_modes(tdQvgdS, 0, 3), + ) + + # ------------------------------------------------------------------ # + # MMA fragments # + # ------------------------------------------------------------------ # + # (MMA, MMA_M, MMA_K, STAGE) + tdQvrdS = tiled_mma_v.make_fragment_A(sdS) + # (MMA, MMA_N, MMA_K, STAGE) + tdQvrV = tiled_mma_v.make_fragment_B(sV) + # (MMA, MMA_M, MMA_N) + acc_v_shape = tiled_mma_v.partition_shape_C(self.mma_tiler_dQv[:2]) + # (MMA, MMA_M, MMA_N, STAGE) + tdQvtAcc_fake = tiled_mma_v.make_fragment_C(cute.append(acc_v_shape, self.num_stages_acc)) + if const_expr(self.compute_dQ): + # (MMA, MMA_N, MMA_K, STAGE) + tdQrK = tiled_mma_k.make_fragment_B(sK) + # (MMA, MMA_M, MMA_N) + acc_k_shape = tiled_mma_k.partition_shape_C(self.mma_tiler_dQ[:2]) + # (MMA, MMA_M, MMA_N, STAGE) + tdQtAcc_fake = tiled_mma_k.make_fragment_C( + cute.append(acc_k_shape, self.num_stages_acc) + ) + + # ------------------------------------------------------------------ # + # Cluster wait before tensor memory alloc # + # ------------------------------------------------------------------ # + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + # ------------------------------------------------------------------ # + # Tile Scheduler # + # ------------------------------------------------------------------ # + tile_sched = utils.ClcDynamicPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + clc_response_ptr, + ) + work_tile = tile_sched.initial_work_tile_info() + + # ------------------------------------------------------------------ # + # TMA load warp # + # ------------------------------------------------------------------ # + if warp_idx == self.tma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + + producer_state_dS = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, stages=self.num_stages_dS + ) + + while work_tile.is_valid_tile: + # ---- Get tile coord from tile scheduler ---- + token, cta, _ = work_tile.tile_idx + + # ((atom_v, rest_v), RestK) + tdSgdS_slice = tdSgdS[(None, 0, None, token)] + + # ---- mainloop ---- + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + pipeline_dS.producer_acquire(producer_state_dS) + index_dS = producer_state_dS.index + + # ---- TMA load dS ---- + cute.copy( + tma_atom_dS, + tdSgdS_slice[(None, k_tile)], + tdSsdS[(None, index_dS)], + tma_bar_ptr=pipeline_dS.producer_get_barrier(producer_state_dS), + mcast_mask=dS_full_mcast_mask, + ) + + producer_state_dS.advance() + + # ---- Advance to next tile ---- + pipeline_clc.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + pipeline_clc.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + + # ---- Wait dS buffer empty ---- + pipeline_dS.producer_tail(producer_state_dS) + + # ------------------------------------------------------------------ # + # Clc Scheduler warp # + # ------------------------------------------------------------------ # + if warp_idx == self.sched_warp_id and is_first_cta: + cute.arch.setmaxregister_decrease(self.num_regs_other) + clc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.ProducerConsumer, self.num_stages_clc + ) + + while work_tile.is_valid_tile: + pipeline_clc.producer_acquire(clc_producer_state) + mbar_addr = pipeline_clc.producer_get_barrier(clc_producer_state) + tile_sched.advance_to_next_work(mbar_addr) + clc_producer_state.advance() + + pipeline_clc.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + pipeline_clc.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + + pipeline_clc.producer_tail(clc_producer_state) + + # ------------------------------------------------------------------ # + # CpAsync KV load warps # + # ------------------------------------------------------------------ # + if warp_idx >= self.kv_load_warp_ids[0] and warp_idx <= self.kv_load_warp_ids[-1]: + cute.arch.setmaxregister_increase(self.num_regs_KV) + find_batch = partial( + self.find_batch_from_q, seqlen_q_divmod=seqlen_q_divmod, mCuSeqlensQ=mCuSeqlensQ + ) + + load_epi_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, 1 + ) + producer_state_KV = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, stages=self.num_stages_KV + ) + + kv_tidx = tidx % (len(self.kv_load_warp_ids) * self.threads_per_warp) + kv_warp_idx = warp_idx % len(self.kv_load_warp_ids) + + mV_cta = cute.domain_offset((cta_rank_in_cluster * (self.head_dim_v // 2), 0), mV) + + while work_tile.is_valid_tile: + # ---- Get tile coord from tile scheduler ---- + token, cta, _ = work_tile.tile_idx + + batch_idx = find_batch(token) + k_batch_offset = ( + mCuSeqlensK[batch_idx] if const_expr(mCuSeqlensK is not None) else Int32(0) + ) + seqlen_k = ( + mCuSeqlensK[batch_idx + 1] - k_batch_offset + if const_expr(mCuSeqlensK is not None) + else seqlen_k_static + ) + if const_expr(mCuSeqlensK is not None): + if const_expr(self.compute_dQ): + mK_cur = cute.domain_offset((0, k_batch_offset), mK)[None, None] + mV_cur = cute.domain_offset((0, k_batch_offset), mV_cta)[None, None] + else: + if const_expr(self.compute_dQ): + mK_cur = cute.domain_offset((0, (0, batch_idx)), mK)[None, None] + mV_cur = cute.domain_offset((0, (0, batch_idx)), mV_cta)[None, None] + mIdxTopK_cur = mIdxTopK[None, token] + + cpasync_gather_kv_manager = CpasyncGatherKVManager.create( + mIdxTopK_cur, + 0, + kv_tidx, + kv_warp_idx, + self.top_k, + seqlen_k, + self.mma_tiler_dQv[2], + self.head_dim_k, + self.mma_tiler_dQv[1], + 1, + len(self.kv_load_warp_ids) * self.threads_per_warp, + mV.element_type, + 1, + ) + + # ---- K/V load mainloop ---- + for k_tile in cutlass.range_constexpr(k_tile_cnt): + # ---- Load top-k index tensor ---- + cpasync_gather_kv_manager.load_index_topk(k_tile, transpose=True) + + stage = producer_state_KV.index + pipeline_KV.producer_acquire(producer_state_KV) + + # ---- Load V (and optionally load K) ---- + cpasync_gather_kv_manager.load_X(mV_cur, sV[None, None, None, stage], True, "V") + if const_expr(self.compute_dQ): + if is_first_cta: + cpasync_gather_kv_manager.load_X( + mK_cur, sK[None, None, None, stage], True, "K" + ) + + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + self.kv_load_sync_barrier.arrive_and_wait() + + cute.arch.fence_proxy("async.shared", space="cta") + if kv_warp_idx == 0: + with cute.arch.elect_one(): + pipeline_KV.producer_commit(producer_state_KV) + producer_state_KV.advance() + + if const_expr(self.overlap_kv_epi): + pipeline_load_kv_epi.consumer_wait(load_epi_consumer_state) + with cute.arch.elect_one(): + pipeline_load_kv_epi.consumer_release(load_epi_consumer_state) + load_epi_consumer_state.advance() + + # ---- Advance to next tile ---- + pipeline_clc.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + pipeline_clc.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + + pipeline_KV.producer_tail(producer_state_KV) + + # ------------------------------------------------------------------ # + # MMA warp # + # ------------------------------------------------------------------ # + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + # --- Retrieve TMEM ptr and make accumulator tensors + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + if const_expr(self.compute_dQ): + tdQtAcc_base = cute.make_tensor(tmem_ptr, tdQtAcc_fake.layout) + tdQvtAcc_ptr = tmem_ptr + ( + tcgen05.find_tmem_tensor_col_offset(tdQtAcc_base) + if const_expr(self.compute_dQ) + else 0 + ) + tdQvtAcc_base = cute.make_tensor(tdQvtAcc_ptr, tdQvtAcc_fake.layout) + + consumer_state_dS = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, stages=self.num_stages_dS + ) + consumer_state_KV = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, stages=self.num_stages_KV + ) + producer_state_dQ_dQv = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_stages_acc + ) + + while work_tile.is_valid_tile: + # ---- Get tile coord from tile scheduler ---- + token, cta, _ = work_tile.tile_idx + + # ---- Set tensor memory buffer for current tile ---- + # (MMA, MMA_M, MMA_N) + tdQvtAcc = tdQvtAcc_base[(None, None, None, producer_state_dQ_dQv.index)] + if const_expr(self.compute_dQ): + tdQtAcc = tdQtAcc_base[(None, None, None, producer_state_dQ_dQv.index)] + + # ---- Wait for accumulator buffer empty ---- + pipeline_dQ_dQv.producer_acquire(producer_state_dQ_dQv) + + # ---- Reset the ACCUMULATE field for each tile ---- + tiled_mma_v.set(tcgen05.Field.ACCUMULATE, False) + if const_expr(self.compute_dQ): + tiled_mma_k.set(tcgen05.Field.ACCUMULATE, False) + + # ---- Mma mainloop ---- + for k_tile in cutlass.range_constexpr(k_tile_cnt): + pipeline_dS.consumer_wait(consumer_state_dS) + pipeline_KV.consumer_wait(consumer_state_KV) + dS_stage = consumer_state_dS.index + KV_stage = consumer_state_KV.index + + num_kblocks = cute.size(tdQvrdS, mode=[2]) + for kblk_idx in cutlass.range(num_kblocks, unroll_full=True): + # dQv += dS @ V + cute.gemm( + tiled_mma_v, + tdQvtAcc, + tdQvrdS[(None, None, kblk_idx, dS_stage)], + tdQvrV[(None, None, kblk_idx, KV_stage)], + tdQvtAcc, + ) + # Enable accumulate on tdQvtAcc after first kblock + tiled_mma_v.set(tcgen05.Field.ACCUMULATE, True) + + if const_expr(self.compute_dQ): + if is_first_cta: + # dQ += dS @ K + cute.gemm( + tiled_mma_k, + tdQtAcc, + tdQvrdS[(None, None, kblk_idx, dS_stage)], + tdQrK[(None, None, kblk_idx, KV_stage)], + tdQtAcc, + ) + # Enable accumulate on tdQtAcc after first kblock + tiled_mma_k.set(tcgen05.Field.ACCUMULATE, True) + + pipeline_dS.consumer_release(consumer_state_dS) + pipeline_KV.consumer_release(consumer_state_KV) + consumer_state_dS.advance() + consumer_state_KV.advance() + + pipeline_dQ_dQv.producer_commit(producer_state_dQ_dQv) + producer_state_dQ_dQv.advance() + + # ---- Advance to next tile ---- + pipeline_clc.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + pipeline_clc.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + + pipeline_dQ_dQv.producer_tail(producer_state_dQ_dQv) + + # ------------------------------------------------------------------ # + # Epilogue warps # + # ------------------------------------------------------------------ # + if warp_idx >= self.epilogue_warp_ids[0] and warp_idx <= self.epilogue_warp_ids[-1]: + cute.arch.setmaxregister_increase(self.num_regs_epi) + # ---- Alloc tensor memory buffer ---- + tmem.allocate(self.num_tmem_alloc_cols) + + # ---- Retrieving tensor memory ptr and make accumulator tensor ---- + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + if const_expr(self.compute_dQ): + tdQtAcc_base = cute.make_tensor(tmem_ptr, tdQtAcc_fake.layout) + tdQvtAcc_ptr = tmem_ptr + ( + tcgen05.find_tmem_tensor_col_offset(tdQtAcc_base) + if const_expr(self.compute_dQ) + else 0 + ) + tdQvtAcc_base = cute.make_tensor(tdQvtAcc_ptr, tdQvtAcc_fake.layout) + + epi_idx = tidx + # print(f"tdQvtAcc_base.layout = {tdQvtAcc_base.layout}") + # ---- TMEM -> RMEM -> SMEM -> GMEM copies + partitions ---- + tiled_copy_dQv_t2r, tTR_dQvtAcc_base, tTR_dQvrAcc = ( + self.epilogue_tmem_copy_and_partition( + epi_idx, + tdQvtAcc_base, + tdQvgdQv, + epi_tile_dQv, + self.mma_tiler_dQv, + self.dq_layout, + self.dq_dtype, + ) + ) + # print(f"tTR_tdQvtAcc_base.layout = {tTR_dQvtAcc_base.layout}") + tTR_rdQv = cute.make_rmem_tensor(tTR_dQvrAcc.shape, self.dq_dtype) + ( + tiled_copy_dQv_r2s, + tRS_rdQv, + tRS_sdQv, + ) = self.epilogue_smem_copy_and_partition( + self.dq_layout, + self.dq_dtype, + tiled_copy_dQv_t2r, + tTR_rdQv, + epi_idx, + sdQv, + ) + bSG_sdQv, bSG_gdQv_partitioned = self.epilogue_gmem_copy_and_partition( + tma_atom_dQv, + tdQvgdQv, + epi_tile_dQv, + sdQv, + ) + if const_expr(self.compute_dQ): + (tiled_copy_dQ_t2r, tTR_dQtAcc_base, tTR_dQrAcc) = ( + self.epilogue_tmem_copy_and_partition( + epi_idx, + tdQtAcc_base, + tdQgdQ, + epi_tile_dQ, + self.mma_tiler_dQ, + self.dq_layout, + self.dq_dtype, + ) + ) + + tTR_rdQ = cute.make_rmem_tensor(tTR_dQrAcc.shape, self.dq_dtype) + ( + tiled_copy_dQ_r2s, + tRS_rdQ, + tRS_sdQ, + ) = self.epilogue_smem_copy_and_partition( + self.dq_layout, + self.dq_dtype, + tiled_copy_dQ_t2r, + tTR_rdQ, + epi_idx, + sdQ, + ) + bSG_sdQ, bSG_gdQ_partitioned = self.epilogue_gmem_copy_and_partition( + tma_atom_dQ, + tdQgdQ, + epi_tile_dQ, + sdQ, + ) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_stages_acc + ) + load_epi_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, 1 + ) + + epi_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + cute.arch.WARP_SIZE, + ) + pipeline_epi = pipeline.PipelineTmaStore.create( + num_stages=self.num_stages_acc, + producer_group=epi_producer_group, + ) + + # ---- Persistent tile scheduling loop for epilogue ---- + while work_tile.is_valid_tile: + # ---- Get current work tile ---- + token, cta, bid_x = work_tile.tile_idx + + bSG_gdQv = bSG_gdQv_partitioned[ + (None, None, None, bid_x, cta_rank_in_cluster, token) + ] + tTR_dQvtAcc = tTR_dQvtAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + tTR_dQvtAcc = cute.group_modes(tTR_dQvtAcc, 3, cute.rank(tTR_dQvtAcc)) + bSG_gdQv = cute.group_modes(bSG_gdQv, 1, cute.rank(bSG_gdQv)) + + subtile_cnt_v = cute.size(tTR_dQvtAcc.shape, mode=[3]) + epi_subtile_counter_v = 0 + + if const_expr(self.compute_dQ): + bSG_gdQ = bSG_gdQ_partitioned[(None, None, None, bid_x, cta, token)] + tTR_dQtAcc = tTR_dQtAcc_base[ + (None, None, None, None, None, acc_consumer_state.index) + ] + tTR_dQtAcc = cute.group_modes(tTR_dQtAcc, 3, cute.rank(tTR_dQtAcc)) + bSG_gdQ = cute.group_modes(bSG_gdQ, 1, cute.rank(bSG_gdQ)) + + subtile_cnt_k = cute.size(tTR_dQtAcc.shape, mode=[3]) + epi_subtile_counter_k = 0 + + pipeline_dQ_dQv.consumer_wait(acc_consumer_state) + + for subtile_idx in cutlass.range(subtile_cnt_v, unroll_full=True): + store_dQ = ( + const_expr(self.compute_dQ) + and is_first_cta + and (subtile_idx < subtile_cnt_k) + ) + + if not store_dQ: + tTR_dQvtAcc_mn = tTR_dQvtAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_dQv_t2r, tTR_dQvtAcc_mn, tTR_dQvrAcc) + cute.arch.fence_view_async_tmem_load() + + # convert to output dtype + tRS_rdQv.store( + tiled_copy_dQv_r2s.retile(tTR_dQvrAcc).load().to(self.dq_dtype) + ) + + epi_buffer = epi_subtile_counter_v % self.num_stages_acc + cute.copy( + tiled_copy_dQv_r2s, tRS_rdQv, tRS_sdQv[(None, None, None, epi_buffer)] + ) + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier.arrive_and_wait() + + if warp_idx == self.epilogue_warp_ids[0]: + cute.copy( + tma_atom_dQv, + bSG_sdQv[(None, epi_buffer)], + bSG_gdQv[(None, subtile_idx)], + ) + pipeline_epi.producer_commit() + pipeline_epi.producer_acquire() + + self.epilog_sync_barrier.arrive_and_wait() + epi_subtile_counter_v += 1 + elif const_expr(self.compute_dQ): + tTR_dQtAcc_mn = tTR_dQtAcc[(None, None, None, subtile_idx)] + tTR_dQvtAcc_mn = tTR_dQvtAcc[(None, None, None, subtile_idx)] + cute.copy(tiled_copy_dQ_t2r, tTR_dQtAcc_mn, tTR_dQrAcc) + cute.copy(tiled_copy_dQv_t2r, tTR_dQvtAcc_mn, tTR_dQvrAcc) + + # convert to output dtype + tRS_rdQ.store(tiled_copy_dQ_r2s.retile(tTR_dQrAcc).load().to(self.dq_dtype)) + tRS_rdQv.store( + tiled_copy_dQv_r2s.retile(tTR_dQvrAcc).load().to(self.dq_dtype) + ) + + epi_buffer = epi_subtile_counter_v % self.num_stages_acc + epi_buffer_k = epi_subtile_counter_k % self.num_stages_acc + cute.copy( + tiled_copy_dQv_r2s, tRS_rdQv, tRS_sdQv[(None, None, None, epi_buffer)] + ) + cute.copy( + tiled_copy_dQ_r2s, tRS_rdQ, tRS_sdQ[(None, None, None, epi_buffer_k)] + ) + cute.arch.fence_proxy("async.shared", space="cta") + self.epilog_sync_barrier.arrive_and_wait() + + if warp_idx == self.epilogue_warp_ids[0]: + cute.copy( + tma_atom_dQv, + bSG_sdQv[(None, epi_buffer)], + bSG_gdQv[(None, subtile_idx)], + ) + cute.copy( + tma_atom_dQ, + bSG_sdQ[(None, epi_buffer_k)], + bSG_gdQ[(None, subtile_idx)], + ) + pipeline_epi.producer_commit() + pipeline_epi.producer_acquire() + + self.epilog_sync_barrier.arrive_and_wait() + epi_subtile_counter_v += 1 + epi_subtile_counter_k += 1 + + if const_expr(self.overlap_kv_epi): + pipeline_load_kv_epi.producer_acquire(load_epi_producer_state) + with cute.arch.elect_one(): + pipeline_load_kv_epi.producer_commit(load_epi_producer_state) + load_epi_producer_state.advance() + + with cute.arch.elect_one(): + pipeline_dQ_dQv.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # ---- Advance to next tile ---- + pipeline_clc.consumer_wait(clc_consumer_state) + work_tile = tile_sched.get_current_work() + pipeline_clc.consumer_release(clc_consumer_state) + clc_consumer_state.advance() + + # ---- Dealloc the tensor memory buffer ---- + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + pipeline_epi.producer_tail() + + def epilogue_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + mma_tiler_mnk, + c_layout, + c_dtype, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + copy_atom_t2r = utils.sm100.get_tmem_load_op( + mma_tiler_mnk, + c_layout, + c_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs=False, + ) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0, None)], epi_tile) + + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, RestM, RestN, RestL) + gC_mnl_epi = cute.flat_divide(gC_mnl[((None, None), 0, 0, None, None, None)], epi_tile) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, RestM, RestN, RestL) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) + + # (T2R, T2R_M, T2R_N) + rAcc_shape = tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape + tTR_rAcc = cute.make_rmem_tensor(rAcc_shape, self.acc_dtype) + + return (tiled_copy_t2r, tTR_tAcc, tTR_rAcc) + + def epilogue_smem_copy_and_partition( + self, + c_layout, + c_dtype: Type[cutlass.Numeric], + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: cutlass.Int32, + sC: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + copy_atom_r2s = utils.sm100.get_smem_store_op( + c_layout, c_dtype, self.acc_dtype, tiled_copy_t2r + ) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + # (R2S, R2S_M, R2S_N, PIPE_D) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sC = thr_copy_r2s.partition_D(sC) + # (R2S, R2S_M, R2S_N) + tRS_rC = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rC, tRS_sC + + def epilogue_gmem_copy_and_partition( + self, + tma_atom, + gC, + epi_tile, + sC, + ) -> Tuple[cute.Tensor, cute.Tensor]: + gC_epi = cute.flat_divide(gC[((None, None), 0, 0, None, None, None)], epi_tile) + + sC_for_tma_partition = cute.group_modes(sC, 0, 2) + + gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2) + + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_atom, + 0, + cute.make_layout(1), + sC_for_tma_partition, + gC_for_tma_partition, + ) + + return bSG_sC, bSG_gC + + @cute.jit + def find_batch_from_q( + self, + token: Int32, + seqlen_q_divmod: FastDivmodDivisor, + mCuSeqlensQ: Optional[cute.Tensor], + ) -> Int32: + """Find batch index from q token (binary search for varlen, divmod otherwise)""" + if const_expr(mCuSeqlensQ is not None): + return get_batch_from_cu_tensor(token, mCuSeqlensQ) + else: + batch, _ = divmod(token, seqlen_q_divmod) + return batch + + @staticmethod + def _compute_grid(c, cta_tile_shape_mnk, cluster_shape_mn): + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + c_logical_shape = gc[(0, (None, None, None))].shape + + num_ctas_mnl = ( + cute.size(c_logical_shape[0]), + cute.size(c_logical_shape[1]), + cute.size(c_logical_shape[2]), + ) + + tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams( + num_ctas_mnl, (*cluster_shape_mn, 1) + ) + grid = utils.ClcDynamicPersistentTileScheduler.get_grid_shape(tile_sched_params) + return tile_sched_params, grid diff --git a/fastvideo-kernel/fa4/flash_bwd_mla_sm100.py b/fastvideo-kernel/fa4/flash_bwd_mla_sm100.py new file mode 100644 index 000000000..0315578b5 --- /dev/null +++ b/fastvideo-kernel/fa4/flash_bwd_mla_sm100.py @@ -0,0 +1,2133 @@ +# Copyright (c) 2026, Colfax International. + +import math +from functools import partial +from typing import Callable, Optional + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int64, Int32, Boolean, const_expr +import cutlass.pipeline as pipeline +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.utils import ClcDynamicPersistentTileScheduler + +from quack import copy_utils, layout_utils + +from flash_attn.cute.pack_gqa import pack_gqa_layout +from flash_attn.cute.seqlen_info import SeqlenInfoQK +from flash_attn.cute.block_info import BlockInfo +import flash_attn.cute.blackwell_helpers as fa_sm100_utils +from flash_attn.cute.tile_scheduler import ( + ClcState, + SchedulingMode, + TileSchedulerArguments, + TileSchedulerProtocol, + SingleTileScheduler, + SingleTileLPTScheduler, + SingleTileVarlenScheduler, + ParamsBase, +) +from flash_attn.cute.fa_logging import fa_log, fa_printf +from flash_attn.cute.utils import smid, elem_pointer, get_batch_from_cu_tensor +from flash_attn.cute.copy_utils import tiled_copy_2d, atomic_add_fp32x4 + +from flash_attn.cute.topk_gather_kv import CpasyncGatherKVManager + + +from flash_attn.cute.named_barrier import NamedBarrierBwdSm100_MLA2CTA + + +class FlashAttentionSparseMLABackwardSm100: + def __init__( + self, + is_causal: bool = False, + topk_length: int = 2048, + qhead_per_kvhead: int = 1, + nheads_kv: int = 1, + hdim: int = 64, + hdimv: int = 512, + has_seqused_q: bool = False, + disable_bitmask: bool = False, + use_clc_scheduler: bool = True, + ): + use_cpasync_load_KV = True + self.is_causal = is_causal + self.is_local = False + self.pack_gqa = True + self.qhead_per_kvhead = qhead_per_kvhead + self.nheads_kv = nheads_kv + self.has_seqused_q = has_seqused_q + self.use_tma_O = True + self.use_cpasync_load_KV = True + self.use_tma_KV = False + self.topk_length = topk_length + self.is_topk_gather = True + assert qhead_per_kvhead == 128 or qhead_per_kvhead == 64 + + # user-provided option if topk indices guaranteed in bounds + self.disable_bitmask = disable_bitmask + + # ==== tile scheduler ==== + self.static_persistent = False + self.use_clc_scheduler = use_clc_scheduler + self.sched_stages = 1 + self.scheduling_mode = ( + SchedulingMode.CLC if self.use_clc_scheduler else SchedulingMode.STATIC + ) + + if const_expr(has_seqused_q): + self.TileScheduler = SingleTileVarlenScheduler + elif self.use_clc_scheduler: + self.TileScheduler = SingleTileLPTScheduler + else: + self.TileScheduler = SingleTileScheduler + + fa_log( + 1, + f"TileScheduler={self.TileScheduler.__name__}, scheduling_mode={self.scheduling_mode.name}", + ) + + # ==== thread info ==== + self.num_softmax_threads = 128 + self.num_epilogue_threads = 128 + self.num_load_threads = 32 + self.num_mma_threads = 32 + self.num_empty_threads = 0 + self.num_relay_threads = 32 + self.num_cpasync_load_threads = 128 + self.num_threads = 512 + self.num_warps = self.num_threads // 32 + self.softmax_warp_indices = (0, 1, 2, 3) + self.epilogue_warp_indices = (4, 5, 6, 7) + self.load_warp_id = 8 + self.mma_warp_id = 9 + self.clc_scheduler_warp_id = 10 + self.relay_warp_id = 11 + self.cpasync_load_warp_indices = (12, 13, 14, 15) + self.empty_warp_ids = () + + # ==== register usage ==== + assert self.num_warps == 16 + + self.num_regs_load = 128 + self.num_regs_mma = 128 + self.num_regs_softmax = 128 + self.num_regs_epilogue = 128 + self.num_regs_cpasync = 128 + self.num_regs_other = 128 + + # self.num_regs_load = 128 - 32 + # self.num_regs_mma = 128 - 32 + # self.num_regs_softmax = 128 + 32 + # self.num_regs_epilogue = 128 + 32 + # self.num_regs_cpasync = 128 - 32 + # self.num_regs_other = 48 + + self.num_regs_per_thread = 128 + self.num_regs_total = 512 + + assert ( + self.num_regs_mma + + self.num_regs_softmax + + self.num_regs_epilogue + + self.num_regs_cpasync + <= self.num_regs_total + ) + + # ==== 2cta info ==== + self.use_2cta_instrs = True + self.cta_group = tcgen05.CtaGroup.TWO + self.cta_group_size = 2 + self.cluster_shape_mn = (2, 1) + self.cluster_shape_mnk = (2, 1, 1) + + # ==== problem shape info ==== + self.hdim = hdim # ignored + self.hdimv = hdimv + self.tile_m = qhead_per_kvhead + self.tile_n = 64 + self.cta_tiler_mn = (self.tile_m // self.cta_group_size, self.tile_n) + self.cluster_tile_n = self.cta_group_size * self.tile_n + self.num_hdimv_splits = 2 # split hdimv in half for our Qv @ V^T and P @ V mmas. + + self.tile_P = (self.tile_m, self.tile_n) + self.tile_Pt = (self.tile_n, self.tile_m) + self.tile_dS = (self.tile_m, self.tile_n) + self.tile_dSt = (self.tile_n, self.tile_m) + self.tile_dV = (self.tile_n, 32) + + # ==== MMA info ==== + # dP.T = V @ dO.T , N x M x dv + # dV += P.T @ dO , N x dv x M + # dV += dS.T @ Qv , N x dv x M + self.mma_tiler_VdO = ( + self.cluster_tile_n, + self.tile_m, + self.hdimv // self.num_hdimv_splits, + ) + self.mma_tiler_PtdOt = ( + self.cluster_tile_n, + self.hdimv // self.num_hdimv_splits, + self.tile_m, + ) + self.mma_tiler_dStQvt = ( + self.cluster_tile_n, + self.hdimv // self.num_hdimv_splits, + self.tile_m, + ) + # note: store P.T, dS.T as tile_n major (i.e., as P and dS) + self.major_mode_V = tcgen05.OperandMajorMode.K + self.major_mode_dO = tcgen05.OperandMajorMode.K + self.major_mode_Pt = tcgen05.OperandMajorMode.MN + self.major_mode_dOt = tcgen05.OperandMajorMode.MN + self.major_mode_dSt = tcgen05.OperandMajorMode.MN + self.major_mode_Qvt = tcgen05.OperandMajorMode.MN + self.operand_source_V = tcgen05.OperandSource.SMEM + self.operand_source_Pt = tcgen05.OperandSource.SMEM + self.operand_source_dSt = tcgen05.OperandSource.SMEM + + # ==== pipeline info ==== + # stationary: dOi + # mainloop: + # *) P, scaleP => Pt + # *) dSt + # *) Vi, i = {0, 1} + # *) dOti => Qvi => dVi, i = {0, 1} + + # redundant names for ease-of-use + self.num_stages_V = 2 + self.num_stages_dO = 2 + self.num_stages_P = 1 + self.num_stages_Pt = 1 + self.num_stages_dS = 1 + self.num_stages_dSt = 1 + self.num_stages_dOt = 2 + self.num_stages_Qv = 2 + self.num_stages_Qvt = 2 + + self.num_stages_dP = 1 + self.num_stages_dPt = 1 + self.num_stages_dV = 2 # == hdimv splits, for Umma <-> Async + self.num_epi_stages_dV = 8 # == 2 splits x 4 slots/split + + self.num_stages_scaleP = 1 + self.num_stages_dPsum = 1 + + # ==== dtype info ==== + self.dtype_acc = Float32 + + # ==== TMEM info ==== + SM100_TMEM_CAPACITY_COLUMNS = 512 + self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + self.tmem_cols_dP = self.tile_m // self.cta_group_size + self.tmem_cols_dVi = (self.hdimv // self.num_hdimv_splits) // self.cta_group_size + self.tmem_offset_dV0 = 0 + self.tmem_offset_dV1 = self.tmem_offset_dV0 + self.tmem_cols_dVi + self.tmem_offsets_dV = [self.tmem_offset_dV0, self.tmem_offset_dV1] + self.tmem_offset_dP = self.tmem_offset_dV1 + self.tmem_cols_dVi + self.total_tmem = self.tmem_offset_dP + self.tmem_cols_dP + assert self.total_tmem <= self.tmem_alloc_cols, ( + f"Total TMEM columns allocated {self.total_tmem} exceeds capacity {self.tmem_alloc_cols}" + ) + + def _get_shared_storage_cls(self): + self.buffer_align_bytes = 1024 + + def smem_struct_align(dtype, staged_layout): + return cute.struct.Align[ + cute.struct.MemRange[dtype, cute.cosize(staged_layout)], + self.buffer_align_bytes, + ] + + def mbar_struct(num_stages): + return cute.struct.MemRange[Int64, 2 * num_stages] + + # sV, sdO, sP = sPt, sdSt = sdS, sdOt = sQvt = sdV + ( + sV_struct, + sdO_struct, + sP_struct, + sdS_struct, + sQvt_struct, + sScaleP_struct, + sdPsum_struct, + ) = ( + smem_struct_align(dtype, layout) + for dtype, layout in [ + (self.dtype, self.sV_layout_staged), + (self.dtype, self.sdO_layout_staged), + (self.dtype, self.sPt_layout_staged), + (self.dtype, self.sdSt_layout_staged), + (self.dtype, self.sQvt_layout_staged), + (self.dtype_scale, self.sScaleP_layout_staged), + (self.dtype_scale, self.sdPsum_layout_staged), + ] + ) + + ( + mbar_ptr_V_struct, # load V + mbar_ptr_dO_struct, # load dO + mbar_ptr_dOt_Qvt_struct, # load dOt => Qvt + mbar_ptr_dSt_struct, # store dS + mbar_ptr_P_struct, # load P + mbar_ptr_Pt_struct, # store Pt + mbar_ptr_dPt_struct, # dP mma + mbar_ptr_dV_struct, # dV mma + mbar_ptr_scaleP_struct, # load scaleP + mbar_ptr_dPsum_struct, # load dPsum + ) = ( + mbar_struct(n) + for n in [ + self.num_stages_V, + self.num_stages_dO, + self.num_stages_Qvt, + self.num_stages_dSt, + self.num_stages_P, + self.num_stages_Pt, + self.num_stages_dPt, + self.num_stages_dV, + self.num_stages_scaleP, + self.num_stages_dPsum, + ] + ) + mbar_ptr_tmem_dealloc_struct = Int64 + tmem_holding_buf_struct = Int32 + + self.sched_stages = 1 + clc_response_size = self.sched_stages * 4 if self.use_clc_scheduler else 0 + clc_mbar_size = self.sched_stages * 2 if self.use_clc_scheduler else 0 + + @cute.struct + class SharedStorage: + mbar_ptr_V: mbar_ptr_V_struct + mbar_ptr_V_cpasync: mbar_ptr_V_struct + mbar_ptr_dO: mbar_ptr_dO_struct + mbar_ptr_dOt_Qvt: mbar_ptr_dOt_Qvt_struct + mbar_ptr_P: mbar_ptr_P_struct + mbar_ptr_Pt: mbar_ptr_Pt_struct + mbar_ptr_dSt: mbar_ptr_dSt_struct + mbar_ptr_dPt: mbar_ptr_dPt_struct + mbar_ptr_dV: mbar_ptr_dV_struct + mbar_ptr_dV_epi: mbar_ptr_dV_struct + mbar_ptr_scaleP: mbar_ptr_scaleP_struct + mbar_ptr_dPsum: mbar_ptr_dPsum_struct + mbar_ptr_tmem_dealloc: mbar_ptr_tmem_dealloc_struct + tmem_holding_buf: tmem_holding_buf_struct + clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, clc_mbar_size] + clc_response: cute.struct.MemRange[Int32, clc_response_size] + + sScaleP: sScaleP_struct + sdPsum: sdPsum_struct + sV: sV_struct + sdO: sdO_struct + sP: sP_struct + sdS: sdS_struct + sQv: sQvt_struct + + # print("smem bytes = ", SharedStorage.size_in_bytes()) + + return SharedStorage + + # fmt: off + @cute.jit + def __call__( + self, + mdO: cute.Tensor, # (b, s_q, h, dv) or (total_q, h, dv) if there is cu_seqlens_q + mV: cute.Tensor, # (b_k, s_k, h_k, dv) or (total_k, h_k, dv) if there is cu_seqlens_k + mQv: cute.Tensor, # == mdO + mP: cute.Tensor, # (b, s_q, h, topk) or (total_q, h, topk) + mdV: cute.Tensor, # == mV + mdS: cute.Tensor, # == mP + mIndexTopk: cute.Tensor, # (b, s_q, topk) or (total_q, topk) if there is cu_seqlens_q + softmax_scale: Float32, + mScaleP: Optional[cute.Tensor] = None, # (b, s_q, topk//128, h) or (total_q, topk//128, h) + mdPsum: Optional[cute.Tensor] = None, # (b, s_q, h) or (total_q, h) if there is cu_seqlens_q + mCuSeqlensQ: Optional[cute.Tensor] = None, # (b + 1) + mCuSeqlensK: Optional[cute.Tensor] = None, # (b + 1) + mSeqUsedQ: Optional[cute.Tensor] = None, # (b) + mSeqUsedK: Optional[cute.Tensor] = None, # (b) + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + # fmt: on + # ==== dtype info ==== + self.dtype = mdO.element_type + self.dtype_dV = mdV.element_type + self.dtype_scale = Float32 + self.dtype_index = mIndexTopk.element_type + assert self.dtype.width == 16 + assert self.dtype_dV.width == 32 + assert self.dtype_index == Int32 + if const_expr(mScaleP is not None): + assert mScaleP.element_type == self.dtype_scale + if const_expr(mdPsum is not None): + assert mdPsum.element_type == self.dtype_scale + + # ==== Prepare Tensors ==== + new_stride = lambda mX: ( + *(cute.assume(s, divby=128 // mX.element_type.width) for s in mX.stride[:-1]), + mX.stride[-1], + ) + mQv, mV, mdV, mdO, mP, mdS, mScaleP = [ + cute.make_tensor(mX.iterator, cute.make_layout(mX.shape, stride=new_stride(mX))) + if mX is not None + else None + for mX in (mQv, mV, mdV, mdO, mP, mdS, mScaleP) + ] + # (b, s, h, d) -> (s, d, h, b) or + # (total, h, d) -> (total, d, h) + QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1] + KV_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensK is None) else [0, 2, 1] + mQv, mdO, mP, mdS = [ + cute.make_tensor(mX.iterator, cute.select(mX.layout, mode=QO_layout_transpose)) + if mX is not None + else None + for mX in (mQv, mdO, mP, mdS) + ] + mV, mdV = [ + cute.make_tensor(mX.iterator, cute.select(mX.layout, mode=KV_layout_transpose)) + if mX is not None + else None + for mX in (mV, mdV) + ] + + # (b, s, topk//128, h) -> (s, topk//128, h, b) or + # (total, topk//128, h) -> (total, topk//128, h) + ScaleP_layout_transpose = [1, 2, 3, 0] if const_expr(mCuSeqlensQ is None) else [0, 1, 2] + mScaleP = cute.make_tensor( + mScaleP.iterator, cute.select(mScaleP.layout, mode=ScaleP_layout_transpose) + ) + + # (b, s, h) -> (s, h, b) or + # (total, h) -> (total, h) + dPsum_layout_transpose = [1, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 1] + mdPsum = cute.make_tensor( + mdPsum.iterator, cute.select(mdPsum.layout, mode=dPsum_layout_transpose) + ) + + # (b, s_q, topk) -> (topk, s_q, b) or (total_q, topk) -> (topk, total_q) + topk_layout_transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0] + mIndexTopk = cute.make_tensor( + mIndexTopk.iterator, cute.select(mIndexTopk.layout, mode=topk_layout_transpose) + ) + topk_length_dynamic = mIndexTopk.shape[0] + + if const_expr(self.pack_gqa): + mQv, mdO, mP, mdS, mScaleP = [ + pack_gqa_layout(mX, self.qhead_per_kvhead, self.nheads_kv, head_idx=2) + if mX is not None + else None + for mX in (mQv, mdO, mP, mdS, mScaleP) + ] + if const_expr(mdPsum is not None): + mdPsum = pack_gqa_layout(mdPsum, self.qhead_per_kvhead, self.nheads_kv, head_idx=1) + + # ((h/h_k, s_q), dv, h_k, b) -> (dv, (h/h_k, s_q), h_k, b) + # or ((h/h_k, total_q), dv, h_k) -> (dv, (h/h_k, total_q), h_k) + mma_operand_layout_transpose = ( + [1, 0, 2, 3] if const_expr(mCuSeqlensQ is None) else [1, 0, 2] + ) + mQvt, mdOt = [ + cute.make_tensor(mX.iterator, cute.select(mX.layout, mode=mma_operand_layout_transpose)) + for mX in (mQv, mdO) + ] + + # fmt: off + # ==== Prepare MMAs ==== + # (local_var, dtype_a, major_a, major_b, mma_tiler, operand_source_a) + _mma_specs = [ + ("tiled_mma_VdO", self.dtype, self.major_mode_V, self.major_mode_dO, self.mma_tiler_VdO, self.operand_source_V), + ("tiled_mma_PtdOt", self.dtype, self.major_mode_Pt, self.major_mode_dOt, self.mma_tiler_PtdOt, self.operand_source_Pt), + ("tiled_mma_dStQvt", self.dtype, self.major_mode_dSt, self.major_mode_Qvt, self.mma_tiler_dStQvt, self.operand_source_dSt), + ] + tiled_mma_VdO, tiled_mma_PtdOt, tiled_mma_dStQvt = ( + sm100_utils.make_trivial_tiled_mma( + dtype_a, major_a, major_b, self.dtype_acc, self.cta_group, mma_tiler[:2], operand_source_a, + ) + for _, dtype_a, major_a, major_b, mma_tiler, operand_source_a in _mma_specs + ) + + # ==== Prepare SMEM layouts and TMAs ==== + # (attr, make_fn, tiled_mma, mma_tiler, dtype, num_stages) + _smem_layout_specs = [ + ("sV_layout", sm100_utils.make_smem_layout_a, tiled_mma_VdO, self.mma_tiler_VdO, self.dtype, self.num_stages_V), + ("sdO_layout", sm100_utils.make_smem_layout_b, tiled_mma_VdO, self.mma_tiler_VdO, self.dtype, self.num_stages_dO), + ("sPt_layout", sm100_utils.make_smem_layout_a, tiled_mma_PtdOt, self.mma_tiler_PtdOt, self.dtype, self.num_stages_Pt), + ("sdOt_layout", sm100_utils.make_smem_layout_b, tiled_mma_PtdOt, self.mma_tiler_PtdOt, self.dtype, self.num_stages_dOt), + ("sdSt_layout", sm100_utils.make_smem_layout_a, tiled_mma_dStQvt, self.mma_tiler_dStQvt, self.dtype, self.num_stages_dSt), + ("sQvt_layout", sm100_utils.make_smem_layout_b, tiled_mma_dStQvt, self.mma_tiler_dStQvt, self.dtype, self.num_stages_Qvt), + ] + for attr, make_fn, tiled_mma, mma_tiler, dtype, num_stages in _smem_layout_specs: + ab_kwarg = "a_dtype" if make_fn is sm100_utils.make_smem_layout_a else "b_dtype" + staged = make_fn( + tiled_mma=tiled_mma, + mma_tiler_mnk=mma_tiler, + num_stages=num_stages, + **{ab_kwarg: dtype}, + ) + setattr(self, f"{attr}_staged", staged) + setattr(self, attr, cute.select(staged, mode=[0, 1, 2])) + + # Prepare additional SMEM load layouts + self.P_layout_major = cutlass.utils.LayoutEnum.from_tensor(mP) + self.sP_layout_staged = sm100_utils.make_smem_layout_epi( + self.dtype, self.P_layout_major, self.tile_P, self.num_stages_P + ) + self.sP_layout = cute.select(self.sP_layout_staged, mode=[0, 1]) + self.sScaleP_layout_staged = cute.make_layout((self.tile_m, self.num_stages_scaleP)) + self.sScaleP_layout = cute.select(self.sScaleP_layout_staged, mode=[0]) + self.sdPsum_layout_staged = cute.make_layout((self.tile_m, self.num_stages_dPsum)) + self.sdPsum_layout = cute.select(self.sdPsum_layout_staged, mode=[0]) + + # ==== TMA load ==== + for attr, dtype, layout in [ + ("tma_copy_bytes_V", self.dtype, self.sV_layout), + ("tma_copy_bytes_dO", self.dtype, self.sdO_layout), + ("tma_copy_bytes_dOt", self.dtype, self.sdOt_layout), + ("tma_copy_bytes_Qvt", self.dtype, self.sQvt_layout), + ]: + setattr(self, attr, cute.size_in_bytes(dtype, layout) * self.cta_group_size) + + assert self.tma_copy_bytes_dOt == self.tma_copy_bytes_Qvt + self.tma_copy_bytes_P = cute.size_in_bytes(self.dtype, self.sP_layout) + self.tma_copy_bytes_scaleP = cute.size_in_bytes(self.dtype_scale, self.sScaleP_layout) + self.tma_copy_bytes_dPsum = cute.size_in_bytes(self.dtype_scale, self.sdPsum_layout) + + tma_load_op = cpasync.CopyBulkTensorTileG2SOp(self.cta_group) + cta_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), (tiled_mma_VdO.thr_id.shape,) + ) + cta_shape = cta_layout_vmnk.shape + + def make_tma(make_fn, mX, smem_layout, mma_tiler, tiled_mma): + return make_fn(tma_load_op, mX, smem_layout, mma_tiler, tiled_mma, cta_shape) + + A, B = cute.nvgpu.make_tiled_tma_atom_A, cute.nvgpu.make_tiled_tma_atom_B + + # (atom_name, tensor_name, make_fn, m, smem_layout, mma_tiler, tiled_mma) + _tma_specs = [ + ("tma_atom_dO", "tma_tensor_dO", B, mdO, self.sdO_layout, self.mma_tiler_VdO, tiled_mma_VdO), + ("tma_atom_dOt", "tma_tensor_dOt", B, mdOt, self.sdOt_layout, self.mma_tiler_PtdOt, tiled_mma_PtdOt), + ("tma_atom_Qvt", "tma_tensor_Qvt", B, mQvt, self.sQvt_layout, self.mma_tiler_dStQvt, tiled_mma_dStQvt), + ] + _tmas = {} + for atom_name, tensor_name, make_fn, m, smem_layout, mma_tiler, tiled_mma in _tma_specs: + _tmas[atom_name], _tmas[tensor_name] = ( + make_tma(make_fn, m, smem_layout, mma_tiler, tiled_mma) + ) + + (tma_atom_dO, tma_tensor_dO, + tma_atom_dOt, tma_tensor_dOt, + tma_atom_Qvt, tma_tensor_Qvt) = _tmas.values() + + # Make TMA load for P separately + tma_atom_P, tma_tensor_P = cute.nvgpu.cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + mP, + self.sP_layout, + self.tile_P, + ) + + # ==== TMA store ==== + tma_store_op = cpasync.CopyBulkTensorTileS2GOp() + + self.dS_layout_major = cutlass.utils.LayoutEnum.from_tensor(mdS) + self.dV_layout_major = cutlass.utils.LayoutEnum.from_tensor(mdV) + # (tile_m, tile_n, dS_stages) = (nheads, 64, dS_stage) + sdS_layout_staged = sm100_utils.make_smem_layout_epi( + self.dtype, self.dS_layout_major, self.tile_dS, self.num_stages_dSt + ) + # (tile_n, 32, dV_epi_stages) = (64, 32, 4 x 2) + sdV_layout_staged = sm100_utils.make_smem_layout_epi( + self.dtype_dV, self.dV_layout_major, self.tile_dV, self.num_epi_stages_dV + ) + tma_atom_dS, tma_tensor_dS = cpasync.make_tiled_tma_atom( + tma_store_op, mdS, cute.select(sdS_layout_staged, mode=[0, 1]), self.tile_dS + ) + # fmt: on + + # ==== Allocate shared memory ==== + SharedStorage = self._get_shared_storage_cls() + + # ==== Tile scheduler ==== + TileScheduler = self.TileScheduler + + fa_printf(1, "mdO = {}", mdO.layout) + batch_size_for_sched = cute.size(mdO.shape[3]) if const_expr(mCuSeqlensQ is None) else 1 + tile_sched_args = TileSchedulerArguments( + num_block=cute.ceil_div(cute.size(mdO.shape[0]), self.tile_m), + num_head=cute.size(mdO.shape[2]), + num_batch=batch_size_for_sched, + num_splits=1, + seqlen_k=cute.size(mV.shape[0]), + headdim=self.hdim, + headdim_v=self.hdimv, + total_q=cute.size(mdO.shape[0]) + if const_expr(mCuSeqlensQ is not None) + else cute.size(mdO.shape[0]) * cute.size(mdO.shape[3]), + tile_shape_mn=self.cta_tiler_mn, + mCuSeqlensQ=mCuSeqlensQ, + mSeqUsedQ=mSeqUsedQ, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + element_size=self.dtype.width // 8, + is_persistent=self.static_persistent, + lpt=False, + is_split_kv=False, + cluster_shape_mn=self.cluster_shape_mn, + use_cluster_idx=True, + ) + tile_sched_params = TileScheduler.to_underlying_arguments( + tile_sched_args, scheduling_mode=self.scheduling_mode + ) + self.tile_scheduler_cls = TileScheduler + grid_dim = TileScheduler.get_grid_shape(tile_sched_params) + fa_printf(1, "grid = {}", grid_dim) + + # ==== Named Barrier ==== + self.cpasync_barrier = cutlass.pipeline.NamedBarrier( + barrier_id=int(NamedBarrierBwdSm100_MLA2CTA.Cpasync), + num_threads=self.num_cpasync_load_threads, + ) + self.softmax_barrier = cutlass.pipeline.NamedBarrier( + barrier_id=int(NamedBarrierBwdSm100_MLA2CTA.Softmax), + num_threads=self.num_softmax_threads, + ) + self.epi_barrier = cutlass.pipeline.NamedBarrier( + barrier_id=int(NamedBarrierBwdSm100_MLA2CTA.Epilogue), + num_threads=self.num_epilogue_threads, + ) + + LOG2_E = math.log2(math.e) + softmax_scale_log2 = softmax_scale * LOG2_E + + # ==== Launch kernel ==== + block_dim = (self.num_threads, 1, 1) + self.kernel( + mV, + mdV, + tma_tensor_dO, + tma_tensor_dOt, + tma_tensor_Qvt, + tma_tensor_P, + tma_tensor_dS, + mScaleP, + mdPsum, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + mIndexTopk, + tma_atom_dO, + tma_atom_dOt, + tma_atom_Qvt, + tma_atom_P, + tma_atom_dS, + self.sV_layout_staged, + self.sdO_layout_staged, + self.sdOt_layout_staged, + self.sQvt_layout_staged, + self.sP_layout_staged, # load P + sdS_layout_staged, # store dS + sdV_layout_staged, + self.sPt_layout_staged, # mma Pt + self.sdSt_layout_staged, # mma dSt + self.sScaleP_layout_staged, + self.sdPsum_layout_staged, + tiled_mma_VdO, + tiled_mma_PtdOt, + tiled_mma_dStQvt, + softmax_scale, + softmax_scale_log2, + topk_length_dynamic, + tile_sched_params, + SharedStorage, + ).launch( + grid=grid_dim, + block=block_dim, + cluster=self.cluster_shape_mnk, + smem=SharedStorage.size_in_bytes(), + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mV: cute.Tensor, + mdV: cute.Tensor, + mdO: cute.Tensor, + mdOt: cute.Tensor, + mQvt: cute.Tensor, + mP: cute.Tensor, + mdS: cute.Tensor, + mScaleP: Optional[cute.Tensor], + mdPsum: Optional[cute.Tensor], + mCuSeqlensQ: Optional[cute.Tensor], + mCuSeqlensK: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + mIndexTopk: Optional[cute.Tensor], + tma_atom_dO: cute.CopyAtom, + tma_atom_dOt: cute.CopyAtom, + tma_atom_Qvt: cute.CopyAtom, + tma_atom_P: cute.CopyAtom, + tma_atom_dS: cute.CopyAtom, + sV_layout_staged: cute.ComposedLayout, + sdO_layout_staged: cute.ComposedLayout, + sdOt_layout_staged: cute.ComposedLayout, + sQvt_layout_staged: cute.ComposedLayout, + sP_layout_staged: cute.ComposedLayout, + sdS_layout_staged: cute.ComposedLayout, + sdV_layout_staged: cute.ComposedLayout, + sPt_layout_staged: cute.ComposedLayout, + sdSt_layout_staged: cute.ComposedLayout, + sScaleP_layout_staged: cute.Layout, + sdPsum_layout_staged: cute.Layout, + tiled_mma_VdO: cute.TiledMma, + tiled_mma_PtdOt: cute.TiledMma, + tiled_mma_dStQvt: cute.TiledMma, + softmax_scale: Float32, + softmax_scale_log2: Float32, + topk_length_dynamic: Optional[Int32], + tile_sched_params: ParamsBase, + SharedStorage: cutlass.Constexpr[Callable], + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + cta_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), (tiled_mma_VdO.thr_id.shape,) + ) + mma_tile_coord_v = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + is_leader_cta = mma_tile_coord_v == 0 + + # ==== Allocate SMEM ==== + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # ==== Prepare TMEM allocator ==== + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=int(NamedBarrierBwdSm100_MLA2CTA.TmemPtr), + num_threads=self.num_mma_threads + self.num_softmax_threads + self.num_epilogue_threads, + ) + tmem = cutlass.utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.mma_warp_id, + is_two_cta=self.use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.mbar_ptr_tmem_dealloc, + ) + + # ==== Prefetch TMA descriptors ==== + if warp_idx == self.load_warp_id: + cpasync.prefetch_descriptor(tma_atom_dO) + cpasync.prefetch_descriptor(tma_atom_dOt) + cpasync.prefetch_descriptor(tma_atom_Qvt) + cpasync.prefetch_descriptor(tma_atom_P) + cpasync.prefetch_descriptor(tma_atom_dS) + + # ==== Construct pipelines ==== + tma_warp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + mma_warp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + sm_warps = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_softmax_threads // 32) + store_warp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + sm_threads = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_softmax_threads) + epi_threads = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_epilogue_threads) + sm_threads_cluster = pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_softmax_threads * self.cta_group_size + ) + epi_threads_cluster = pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_epilogue_threads * self.cta_group_size + ) + cpasync_load_threads = pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_cpasync_load_threads + ) + relay_warps_cluster = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.cta_group_size) + relay_threads = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_relay_threads) + + TmaUmma = pipeline.PipelineTmaUmma + TmaAsync = pipeline.PipelineTmaAsync + AsyncUmma = pipeline.PipelineAsyncUmma + UmmaAsync = pipeline.PipelineUmmaAsync + Async = pipeline.PipelineAsync + + def make_pipeline(cls, mbar_ptr, num_stages, producer, consumer, tx_count=None): + return cls.create( + barrier_storage=mbar_ptr.data_ptr(), + num_stages=num_stages, + producer_group=producer, + consumer_group=consumer, + defer_sync=True, + **( + {"cta_layout_vmnk": cta_layout_vmnk} + if cls is not Async and cls is not TmaAsync + else {} + ), + **({"tx_count": tx_count} if tx_count is not None else {}), + ) + + # Unconditional pipelines + # fmt: off + # TmaUmma: dO, dOt & Qvt + pipeline_dO = make_pipeline(TmaUmma, storage.mbar_ptr_dO, self.num_stages_dO, tma_warp, mma_warp, self.tma_copy_bytes_dO) + pipeline_dOt_Qvt = make_pipeline(TmaUmma, storage.mbar_ptr_dOt_Qvt, self.num_stages_Qvt, tma_warp, mma_warp, self.tma_copy_bytes_dOt) + # TmaAsync: P, scaleP + pipeline_P = make_pipeline(TmaAsync, storage.mbar_ptr_P, self.num_stages_P, tma_warp, sm_warps, self.tma_copy_bytes_P) + pipeline_scaleP = make_pipeline(TmaAsync, storage.mbar_ptr_scaleP, self.num_stages_scaleP, tma_warp, sm_warps, self.tma_copy_bytes_scaleP) + pipeline_dPsum = make_pipeline(TmaAsync, storage.mbar_ptr_dPsum, self.num_stages_dPsum, tma_warp, sm_warps, self.tma_copy_bytes_dPsum) + # AsyncUmma: Pt => dV mma, dSt => dV mma + pipeline_Pt = make_pipeline(AsyncUmma, storage.mbar_ptr_Pt, self.num_stages_Pt, sm_threads_cluster, mma_warp) + pipeline_dSt = make_pipeline(AsyncUmma, storage.mbar_ptr_dSt, self.num_stages_dSt, sm_threads_cluster, mma_warp) + # UmmaAsync: dPt, dV + pipeline_dPt = make_pipeline(UmmaAsync, storage.mbar_ptr_dPt, self.num_stages_dPt, mma_warp, sm_threads_cluster) + pipeline_dV = make_pipeline(UmmaAsync, storage.mbar_ptr_dV, self.num_stages_dV, mma_warp, epi_threads_cluster) + # Async: dV_epi + pipeline_dV_epi = make_pipeline(Async, storage.mbar_ptr_dV_epi, self.num_stages_dV, tma_warp, store_warp) + + pipeline_V = make_pipeline(AsyncUmma, storage.mbar_ptr_V, self.num_stages_V, relay_warps_cluster, mma_warp) + pipeline_V_cpasync = make_pipeline(Async, storage.mbar_ptr_V_cpasync, self.num_stages_V, cpasync_load_threads, relay_threads) + # fmt: on + + pipeline.pipeline_init_arrive(cluster_shape_mn=cta_layout_vmnk, is_relaxed=True) + + # ==== Get SMEM tensors ==== + # fmt: off + sV, sdO, sP, sPt, sdOt, sdS, sdSt, sQvt = ( + store.get_tensor(layout.outer, swizzle=layout.inner) + for store, layout in [ + (storage.sV, sV_layout_staged), + (storage.sdO, sdO_layout_staged), + (storage.sP, sP_layout_staged), # P & Pt overlap + (storage.sP, sPt_layout_staged), # P & Pt overlap + (storage.sQv, sdOt_layout_staged), # {dOt, Qvt, dV} overlap + (storage.sdS, sdS_layout_staged), # dS & dSt overlap + (storage.sdS, sdSt_layout_staged), # dS & dSt overlap + (storage.sQv, sQvt_layout_staged), # {dOt, Qvt, dV} overlap + ] + ) + sdV = cute.make_tensor( + cute.recast_ptr(sdOt.iterator, sdV_layout_staged.inner, self.dtype_acc), sdV_layout_staged.outer + ) + assert cute.cosize(sdV) * self.dtype_acc.width // self.dtype.width == cute.cosize(sdOt) + + sScaleP = storage.sScaleP.get_tensor(sScaleP_layout_staged) + sdPsum = storage.sdPsum.get_tensor(sdPsum_layout_staged) + # fmt: on + + # ==== Get thread MMAs and accumulator fragments ==== + thr_mma_VdO = tiled_mma_VdO.get_slice(mma_tile_coord_v) + thr_mma_PtdOt = tiled_mma_PtdOt.get_slice(mma_tile_coord_v) + thr_mma_dStQvt = tiled_mma_dStQvt.get_slice(mma_tile_coord_v) + + acc_shape_dPt = thr_mma_VdO.partition_shape_C(self.mma_tiler_VdO[:2]) + acc_shape_dVi = thr_mma_PtdOt.partition_shape_C(self.mma_tiler_PtdOt[:2]) + tdPtdP_fake = thr_mma_VdO.make_fragment_C(acc_shape_dPt) + tdVtdV0_fake = thr_mma_PtdOt.make_fragment_C(acc_shape_dVi) + tdVtdV1_fake = thr_mma_PtdOt.make_fragment_C(acc_shape_dVi) + # tdPtdP = cute.make_tensor(tdPtdP.iterator + self.tmem_offset_dP, tdPtdP.layout) + # tdVtdV0 = cute.make_tensor(tdVtdV0.iterator + self.tmem_offset_dV0, tdVtdV0.layout) + # tdVtdV1 = cute.make_tensor(tdVtdV1.iterator + self.tmem_offset_dV1, tdVtdV1.layout) + + block_info = BlockInfo( + self.tile_m * self.cta_group_size, + self.tile_n, + is_causal=self.is_causal, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + SeqlenInfoCls = partial( + SeqlenInfoQK.create, + seqlen_q_static=mdO.shape[0] if const_expr(not self.pack_gqa) else mdO.shape[0][1], + seqlen_k_static=mV.shape[0], + tile_m=self.tile_m, + tile_n=self.tile_n, + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + ) + + if const_expr(self.use_clc_scheduler): + clc_response_ptr = storage.clc_response.data_ptr() + clc_mbar_ptr = storage.clc_mbar_ptr.data_ptr() + + clc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_clc_consumer_warps_per_cta = self.num_threads // cute.arch.WARP_SIZE + num_clc_consumer_warps = num_clc_consumer_warps_per_cta * self.cta_group_size + clc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, cute.arch.WARP_SIZE * num_clc_consumer_warps + ) + clc = ClcState.create( + hw_scheduler=ClcDynamicPersistentTileScheduler.create( + self.tile_scheduler_cls.clc_problem_shape(tile_sched_params), + cute.arch.block_idx(), + cute.arch.grid_dim(), + clc_response_ptr, + ), + pipeline=pipeline.PipelineClcFetchAsync.create( + barrier_storage=clc_mbar_ptr, + num_stages=self.sched_stages, + producer_group=clc_pipeline_producer_group, + consumer_group=clc_pipeline_consumer_group, + tx_count=16, + cta_layout_vmnk=cta_layout_vmnk, + ), + consumer_state=pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.sched_stages + ), + producer_state=pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.sched_stages + ), + ) + tile_scheduler = self.tile_scheduler_cls.create(tile_sched_params, clc=clc) + else: + tile_scheduler = self.tile_scheduler_cls.create(tile_sched_params) + assert isinstance(tile_scheduler, TileSchedulerProtocol), ( + f"tile_scheduler is not a TileSchedulerProtocol: {type(tile_scheduler)}" + ) + + pipeline.pipeline_init_wait(cluster_shape_mn=cta_layout_vmnk) + + if const_expr(self.use_clc_scheduler): + if warp_idx == self.clc_scheduler_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + if is_leader_cta: + self.clc_scheduler_warp(tile_scheduler) + else: + self.empty_warp(tile_scheduler) + for i in cutlass.range_constexpr(len(self.empty_warp_ids)): + if warp_idx == self.empty_warp_ids[i] and warp_idx != self.clc_scheduler_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + self.empty_warp(tile_scheduler) + else: + for i in cutlass.range_constexpr(len(self.empty_warp_ids)): + if warp_idx == self.empty_warp_ids[i]: + cute.arch.setmaxregister_decrease(self.num_regs_other) + + if const_expr(self.use_cpasync_load_KV): + if warp_idx == self.relay_warp_id: + if const_expr(self.num_regs_load < self.num_regs_per_thread): + cute.arch.setmaxregister_decrease(self.num_regs_load) + self.relay( + pipeline_V, + pipeline_V_cpasync, + topk_length_dynamic, + block_info, + SeqlenInfoCls, + tile_scheduler=tile_scheduler, + ) + + if warp_idx in self.cpasync_load_warp_indices: + if const_expr(self.num_regs_cpasync < self.num_regs_per_thread): + cute.arch.setmaxregister_decrease(self.num_regs_cpasync) + self.load_cpasync( + mIndexTopk, + mV, + sV, + pipeline_V, + pipeline_V_cpasync, + topk_length_dynamic, + block_info, + SeqlenInfoCls, + mCuSeqlensQ, + tile_scheduler=tile_scheduler, + ) + + if warp_idx == self.load_warp_id: + if const_expr(self.num_regs_load < self.num_regs_per_thread): + cute.arch.setmaxregister_decrease(self.num_regs_load) + self.load( + mdO, + mP, + mdOt, + mQvt, + mScaleP, + mdPsum, + sdO, + sP, + sdOt, + sQvt, + sScaleP, + sdPsum, + tma_atom_dO, + tma_atom_P, + tma_atom_dOt, + tma_atom_Qvt, + pipeline_dO, + pipeline_P, + pipeline_dOt_Qvt, + pipeline_Pt, + pipeline_dV_epi, + pipeline_scaleP, + pipeline_dPsum, + thr_mma_VdO, + thr_mma_PtdOt, + thr_mma_dStQvt, + topk_length_dynamic, + block_info, + SeqlenInfoCls, + mCuSeqlensQ, + tile_scheduler=tile_scheduler, + ) + + if warp_idx == self.mma_warp_id: + if const_expr(self.num_regs_mma < self.num_regs_per_thread): + cute.arch.setmaxregister_decrease(self.num_regs_mma) + # ==== Allocate TMEM ==== + tmem.allocate(self.tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.dtype_acc) + tdPtdP = cute.make_tensor(tmem_ptr + self.tmem_offset_dP, tdPtdP_fake.layout) + tdVtdV0 = cute.make_tensor(tmem_ptr + self.tmem_offset_dV0, tdVtdV0_fake.layout) + tdVtdV1 = cute.make_tensor(tmem_ptr + self.tmem_offset_dV1, tdVtdV1_fake.layout) + self.mma( + sV, + sdO, + sPt, + sdOt, + sdSt, + sQvt, + tdPtdP, + tdVtdV0, + tdVtdV1, + tiled_mma_VdO, + tiled_mma_PtdOt, + tiled_mma_dStQvt, + pipeline_V, + pipeline_dO, + pipeline_dPt, + pipeline_Pt, + pipeline_dOt_Qvt, + pipeline_dSt, + pipeline_dV, + is_leader_cta, + topk_length_dynamic, + block_info, + SeqlenInfoCls, + mCuSeqlensQ, + tile_scheduler=tile_scheduler, + ) + tmem.relinquish_alloc_permit() + tmem_alloc_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + + if warp_idx in self.softmax_warp_indices: + cute.arch.setmaxregister_increase(self.num_regs_softmax) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.dtype_acc) + tdPtdP = cute.make_tensor(tmem_ptr + self.tmem_offset_dP, tdPtdP_fake.layout) + self.compute_loop( + softmax_scale, + softmax_scale_log2, + thr_mma_VdO, + tdPtdP, + sP, + sdS, + sScaleP, + sdPsum, + mdS, + tma_atom_dS, + pipeline_P, + pipeline_Pt, + pipeline_dPt, + pipeline_dSt, + pipeline_scaleP, + pipeline_dPsum, + topk_length_dynamic, + block_info, + SeqlenInfoCls, + mCuSeqlensQ, + tile_scheduler=tile_scheduler, + ) + tmem_alloc_barrier.arrive() + + if warp_idx in self.epilogue_warp_indices: + if const_expr(self.num_regs_epilogue < self.num_regs_per_thread): + cute.arch.setmaxregister_decrease(self.num_regs_epilogue) + elif const_expr(self.num_regs_epilogue > self.num_regs_per_thread): + cute.arch.setmaxregister_increase(self.num_regs_epilogue) + + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.dtype_acc) + tdVtdV0 = cute.make_tensor(tmem_ptr + self.tmem_offset_dV0, tdVtdV0_fake.layout) + tdVtdV1 = cute.make_tensor(tmem_ptr + self.tmem_offset_dV1, tdVtdV1_fake.layout) + self.dVacc_store( + mIndexTopk, + mdV, + sdV, + tdVtdV0, + tdVtdV1, + thr_mma_PtdOt, + pipeline_dV, + pipeline_dV_epi, + topk_length_dynamic, + block_info, + SeqlenInfoCls, + mCuSeqlensQ, + tile_scheduler=tile_scheduler, + ) + tmem_alloc_barrier.arrive() + + @cute.jit + def clc_scheduler_warp( + self, + tile_scheduler: TileSchedulerProtocol, + ): + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + tile_scheduler.prefetch_next_work() + work_tile = tile_scheduler.advance_to_next_work() + cta_m_block, head_idx, batch_idx, _ = work_tile.tile_idx + if cute.arch.thread_idx()[0] == self.clc_scheduler_warp_id * cute.arch.WARP_SIZE: + fa_printf( + 3, + "[CLC] query sm={} cta={} (m_blk={},h={},b={},s={}) valid={}\n", + smid(), + cute.arch.block_idx()[0], + work_tile.tile_idx[0], + work_tile.tile_idx[1], + work_tile.tile_idx[2], + work_tile.tile_idx[3], + work_tile.is_valid_tile, + ) + tile_scheduler.producer_tail() + + @cute.jit + def empty_warp( + self, + tile_scheduler: TileSchedulerProtocol, + ): + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + work_tile = tile_scheduler.advance_to_next_work() + + @cute.jit + def relay( + self, + pipeline_V: pipeline.PipelineAsyncUmma, + pipeline_V_cpasync: pipeline.PipelineAsync, + topk_length_dynamic: Optional[Int32], + block_info: BlockInfo, + SeqlenInfoCls: Callable, + tile_scheduler: TileSchedulerProtocol, + ): + # ==== Make pipeline states ==== + # pipeline_V producer + # pipeline_V_cpasync consumer + producer_state_V = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, stages=self.num_stages_V + ) + consumer_state_V = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, stages=self.num_stages_V + ) + relay_V_fn = partial(self.relay_inner, pipeline_V_cpasync, pipeline_V) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + # m_block, head_idx, batch_idx, _ = work_tile.tile_idx + # seqlen = SeqlenInfoCls(batch_idx) + + num_n_block_groups = self.topk_length // self.cluster_tile_n + # num_n_block_groups = topk_length_dynamic // self.cluster_tile_n + + # ==== Mainloop ==== + for _ in cutlass.range(num_n_block_groups, unroll=1): + for _ in cutlass.range_constexpr(self.num_hdimv_splits): + consumer_state_V, producer_state_V = relay_V_fn( + consumer_state_V, producer_state_V + ) + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() + + pipeline_V.producer_tail(producer_state_V) + + @cute.jit + def relay_inner( + self, + pipeline_cpasync: pipeline.PipelineAsync, + pipeline_mma: pipeline.PipelineAsyncUmma, + consumer_state: pipeline.PipelineState, + producer_state: pipeline.PipelineState, + ): + pipeline_cpasync.consumer_wait(consumer_state) + with cute.arch.elect_one(): + pipeline_mma.producer_commit(producer_state) + consumer_state.advance() + producer_state.advance() + return consumer_state, producer_state + + @cute.jit + def load_cpasync( + self, + mIndexTopk: cute.Tensor, + mV: cute.Tensor, + sV: cute.Tensor, + pipeline_V: pipeline.PipelineAsyncUmma, + pipeline_V_cpasync: pipeline.PipelineAsync, + topk_length_dynamic: Optional[Int32], + block_info: BlockInfo, + SeqlenInfoCls: Callable, + mCuSeqlensQ: Optional[cute.Tensor], + tile_scheduler: TileSchedulerProtocol, + ): + # ==== cpasync load warpgroup ==== + # Description: loads tiles of V from gmem to smem using cpasync + # produces: V + # consumes: - + + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + tidx = cute.arch.thread_idx()[0] % self.num_cpasync_load_threads + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % ( + self.num_cpasync_load_threads // 32 + ) + + # ==== Make pipeline states ==== + # producer: acquire PipelineAsyncUmma <- mma + # producer: commit PipelineAsync -> relay + producer_state_V = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, stages=self.num_stages_V + ) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + m_block, head_idx, batch_idx, _ = work_tile.tile_idx + if const_expr(mCuSeqlensQ is not None): + batch_idx = get_batch_from_cu_tensor(m_block, mCuSeqlensQ) + seqlen = SeqlenInfoCls(batch_idx) + + num_n_block_groups = self.topk_length // self.cluster_tile_n + # num_n_block_groups = topk_length_dynamic // self.cluster_tile_n + + if const_expr(seqlen.has_cu_seqlens_q): + # m_block means absolute m_idx + mIndexTopk_cur = mIndexTopk[None, m_block] + else: + mIndexTopk_cur = mIndexTopk[None, m_block, batch_idx] + + if const_expr(self.is_causal): + m_local_idx = ( + m_block - seqlen.offset_q if const_expr(seqlen.has_cu_seqlens_q) else m_block + ) + seqlen_k_limit = m_local_idx + 1 + seqlen.seqlen_k - seqlen.seqlen_q + else: + seqlen_k_limit = seqlen.seqlen_k + cpasync_gather_kv_manager = CpasyncGatherKVManager.create( + mIndexTopk_cur, + cta_rank_in_cluster, + tidx, + warp_idx, + self.topk_length, + seqlen_k_limit, + self.cluster_tile_n, + self.hdim, + self.hdimv, + self.num_hdimv_splits, + self.num_cpasync_load_threads, + mV.element_type, + self.cta_group_size, + self.cpasync_barrier, + self.disable_bitmask, + ) + + # (seqlen_k, hdimv) + mV_cur = seqlen.offset_batch_K(mV, batch_idx, dim=3)[None, None, head_idx] + + load_V = partial( + self.cpasync_gather_load_KV, + cpasync_gather_kv_manager, + pipeline_V, + pipeline_V_cpasync, + sV, + False, + "V", + mV_cur, + ) + + # ==== Mainloop ==== + for n_block_group in cutlass.range(num_n_block_groups, unroll=1): + cpasync_gather_kv_manager.load_index_topk(n_block_group, transpose=False) + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_V = load_V(producer_state_V, d_offset=split * self.hdimv // 2) + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() + + pipeline_V.producer_tail(producer_state_V) + + @cute.jit + def cpasync_gather_load_KV( + self, + cpasync_gather_kv_manager: CpasyncGatherKVManager, + pipeline_mma: pipeline.PipelineAsyncUmma, + pipeline_cpasync: pipeline.PipelineAsync, + sX: cute.Tensor, + transpose: bool, + K_or_V: str, + mX: cute.Tensor, + producer_state: pipeline.PipelineState, + d_offset: int = 0, + ): + stage = producer_state.index + pipeline_mma.producer_acquire(producer_state) + cpasync_gather_kv_manager.load_X( + mX, sX[None, None, None, stage], transpose, K_or_V, d_offset + ) + cute.arch.cp_async_commit_group() + pipeline_cpasync.sync_object_full.arrive_cp_async_mbarrier(stage) + producer_state.advance() + return producer_state + + @cute.jit + def load( + self, + mdO: cute.Tensor, + mP: cute.Tensor, + mdOt: cute.Tensor, + mQvt: cute.Tensor, + mScaleP: Optional[cute.Tensor], + mdPsum: Optional[cute.Tensor], + sdO: cute.Tensor, + sP: cute.Tensor, + sdOt: cute.Tensor, + sQvt: cute.Tensor, + sScaleP: cute.Tensor, + sdPsum: cute.Tensor, + tma_atom_dO: cute.CopyAtom, + tma_atom_P: cute.CopyAtom, + tma_atom_dOt: cute.CopyAtom, + tma_atom_Qvt: cute.CopyAtom, + pipeline_dO: pipeline.PipelineAsync, # TmaUmma + pipeline_P: pipeline.PipelineAsync, # TmaAsync + pipeline_dOt_Qvt: pipeline.PipelineAsync, # TmaUmma + pipeline_Pt: pipeline.PipelineAsync, # AsyncUmma + pipeline_dV_epi: pipeline.PipelineAsync, # Async + pipeline_scaleP: pipeline.PipelineAsync, # TmaAsync + pipeline_dPsum: pipeline.PipelineAsync, # TmaAsync + thr_mma_VdO: cute.ThrMma, + thr_mma_PtdOt: cute.ThrMma, + thr_mma_dStQvt: cute.ThrMma, + topk_length_dynamic: Optional[Int32], + block_info: BlockInfo, + SeqlenInfoCls: Callable, + mCuSeqlensQ: Optional[cute.Tensor], + tile_scheduler: TileSchedulerProtocol, + ): + # ==== Load warp ==== + # Description: loads tiles of dO, P, dOt, Qvt from gmem to smem using TMA + # produces: dO, P, dOt, Qvt + # consumes: - + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + lane_idx = cute.arch.lane_idx() + + # ==== Make pipeline states ==== + Producer = pipeline.PipelineUserType.Producer + producer_state_dO = pipeline.make_pipeline_state(Producer, stages=self.num_stages_dO) + producer_state_P = pipeline.make_pipeline_state(Producer, stages=self.num_stages_P) + producer_state_dOt_Qvt = pipeline.make_pipeline_state(Producer, stages=self.num_stages_dOt) + producer_state_dV_epi = pipeline.make_pipeline_state(Producer, stages=self.num_stages_dV) + producer_state_scaleP = pipeline.make_pipeline_state( + Producer, stages=self.num_stages_scaleP + ) + producer_state_dPsum = pipeline.make_pipeline_state(Producer, stages=self.num_stages_dPsum) + + copy_atom_stats = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), Float32) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + m_block, head_idx, batch_idx, _ = work_tile.tile_idx + if const_expr(mCuSeqlensQ is not None): + batch_idx = get_batch_from_cu_tensor(m_block, mCuSeqlensQ) + seqlen = SeqlenInfoCls(batch_idx) + if const_expr(mCuSeqlensQ is not None): + m_block -= seqlen.offset_q + num_n_block_groups = self.topk_length // self.cluster_tile_n + # num_n_block_groups = topk_length_dynamic // self.cluster_tile_n + + # ==== Partition GMEM tensors ==== + # (seqlen_q, topk or hdimv) + mP_cur = seqlen.offset_batch_Q(mP, batch_idx, dim=3)[None, None, head_idx] + mdO_cur = seqlen.offset_batch_Q(mdO, batch_idx, dim=3)[None, None, head_idx] + + # (hdimv, seqlen_q) + offset = ( + (0, seqlen.offset_q) if const_expr(not self.pack_gqa) else (0, (0, seqlen.offset_q)) + ) + if const_expr(not seqlen.has_cu_seqlens_q): + mdOt_cur = mdOt[None, None, head_idx, batch_idx] + mQvt_cur = mQvt[None, None, head_idx, batch_idx] + else: + mdOt_cur = cute.domain_offset(offset, mdOt[None, None, head_idx]) + mQvt_cur = cute.domain_offset(offset, mQvt[None, None, head_idx]) + + gScaleP = None + if const_expr(mScaleP is not None): + mScaleP_cur = seqlen.offset_batch_Q(mScaleP, batch_idx, dim=3)[None, None, head_idx] + # (tile_m, topk//128) + gScaleP = cute.local_tile(mScaleP_cur, (self.tile_m,), (m_block, None)) + gdPsum = None + if const_expr(mdPsum is not None): + mdPsum_cur = seqlen.offset_batch_Q(mdPsum, batch_idx, dim=2)[None, head_idx] + # (tile_m) + gdPsum = cute.local_tile(mdPsum_cur, (self.tile_m,), (m_block,)) + + # (tile_m, tile_n, n_blocks) + gP = cute.local_tile( + mP_cur, + (self.tile_m, self.tile_n), + (m_block, None), + ) + # (tile_m, hdimv//2, 2) + gdO = cute.local_tile( + mdO_cur, + (self.mma_tiler_VdO[1], self.mma_tiler_VdO[2]), + (m_block, None), + ) + # (hdimv//2, tile_m, 2) + gdOt = cute.local_tile( + mdOt_cur, + (self.mma_tiler_PtdOt[1], self.mma_tiler_PtdOt[2]), + (None, m_block), + ) + gQvt = cute.local_tile( + mQvt_cur, + (self.mma_tiler_dStQvt[1], self.mma_tiler_dStQvt[2]), + (None, m_block), + ) + + tdPgdO = thr_mma_VdO.partition_B(gdO) + tdVgdOt = thr_mma_PtdOt.partition_B(gdOt) + tdVgQvt = thr_mma_dStQvt.partition_B(gQvt) + + # (V, REST) + tPsP, tPgP = cpasync.tma_partition( + atom=tma_atom_P, + cta_coord=0, + cta_layout=cute.make_layout(1), + smem_tensor=cute.group_modes(sP, 0, 2), + gmem_tensor=cute.group_modes(gP, 0, 2), + ) + tdOsdO, tdOgdO = cpasync.tma_partition( + atom=tma_atom_dO, + cta_coord=0, + cta_layout=cute.make_layout(1), + smem_tensor=cute.group_modes(sdO, 0, 3), + gmem_tensor=cute.group_modes(tdPgdO, 0, 3), + ) + tdOtsdOt, tdOtgdOt = cpasync.tma_partition( + atom=tma_atom_dOt, + cta_coord=0, + cta_layout=cute.make_layout(1), + smem_tensor=cute.group_modes(sdOt, 0, 3), + gmem_tensor=cute.group_modes(tdVgdOt, 0, 3), + ) + tQvtsQvt, tQvtgQvt = cpasync.tma_partition( + atom=tma_atom_Qvt, + cta_coord=0, + cta_layout=cute.make_layout(1), + smem_tensor=cute.group_modes(sQvt, 0, 3), + gmem_tensor=cute.group_modes(tdVgQvt, 0, 3), + ) + + load_P = partial(self.load_inner, tma_atom_P, tPgP, tPsP, pipeline_P) + load_dO = partial(self.load_inner, tma_atom_dO, tdOgdO, tdOsdO, pipeline_dO) + load_dOt = partial(self.load_inner, tma_atom_dOt, tdOtgdOt, tdOtsdOt, pipeline_dOt_Qvt) + load_Qvt = partial(self.load_inner, tma_atom_Qvt, tQvtgQvt, tQvtsQvt, pipeline_dOt_Qvt) + load_scaleP = partial( + self.load_inner, copy_atom_stats, gScaleP, sScaleP, pipeline_scaleP, bulk_copy=True + ) + load_dPsum = partial( + self.load_inner, copy_atom_stats, gdPsum, sdPsum, pipeline_dPsum, bulk_copy=True + ) + + # ==== Load stationary operands ==== + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_dO = load_dO(producer_state_dO, block=split) + + producer_state_dPsum = load_dPsum(producer_state_dPsum) + + # ==== Mainloop ==== + for n_block_group in cutlass.range(num_n_block_groups, unroll=1): + n_block = 2 * n_block_group + cta_rank_in_cluster + # load ScaleP + if const_expr(mScaleP is not None): + producer_state_scaleP = load_scaleP(producer_state_scaleP, block=n_block_group) + pipeline_Pt.producer_acquire(producer_state_P) + producer_state_P = load_P(producer_state_P, block=n_block) + for split in cutlass.range_constexpr(self.num_hdimv_splits): + pipeline_dV_epi.producer_acquire(producer_state_dV_epi) + producer_state_dV_epi.advance() + producer_state_dOt_Qvt = load_dOt(producer_state_dOt_Qvt, block=split) + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_dOt_Qvt = load_Qvt(producer_state_dOt_Qvt, block=split) + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() + + pipeline_P.producer_tail(producer_state_P) + pipeline_dO.producer_tail(producer_state_dO) + pipeline_dOt_Qvt.producer_tail(producer_state_dOt_Qvt) + + @cute.jit + def load_inner( + self, + copy_atom: cute.CopyAtom, + tXgX: cute.Tensor, + tXsX: cute.Tensor, + load_pipeline: pipeline.PipelineAsync, + producer_state: pipeline.PipelineState, + block: Optional[Int32] = None, + bulk_copy: bool = False, + ): + if const_expr(block is not None): + tXgX = tXgX[(None, block)] + if const_expr(cute.rank(tXsX) != 1): + assert cute.rank(tXsX) == 2, f"wrong rank for tXsX, got {cute.rank(tXsX)}" + stage = producer_state.index + tXsX = tXsX[(None, stage)] + + load_pipeline.producer_acquire(producer_state) + mbar_ptr = load_pipeline.producer_get_barrier(producer_state) + if const_expr(bulk_copy): + with cute.arch.elect_one(): + cute.copy(copy_atom, tXgX, tXsX, mbar_ptr=mbar_ptr) + else: + cute.copy(copy_atom, tXgX, tXsX, tma_bar_ptr=mbar_ptr) + producer_state.advance() + return producer_state + + @cute.jit + def mma( + self, + sV: cute.Tensor, + sdO: cute.Tensor, + sPt: cute.Tensor, + sdOt: cute.Tensor, + sdSt: cute.Tensor, + sQvt: cute.Tensor, + tdPtdP: cute.Tensor, + tdVtdV0: cute.Tensor, + tdVtdV1: cute.Tensor, + tiled_mma_VdO: cute.TiledMma, + tiled_mma_PtdOt: cute.TiledMma, + tiled_mma_dStQvt: cute.TiledMma, + pipeline_V: pipeline.PipelineAsync, # AsyncUmma + pipeline_dO: pipeline.PipelineAsync, # TmaUmma + pipeline_dPt: pipeline.PipelineAsync, # UmmaAsync + pipeline_Pt: pipeline.PipelineAsync, # AsyncUmma + pipeline_dOt_Qvt: pipeline.PipelineAsync, # TmaUmma + pipeline_dSt: pipeline.PipelineAsync, # AsyncUmma + pipeline_dV: pipeline.PipelineAsync, # UmmaAsync + is_leader_cta: Boolean, + topk_length_dynamic: Optional[Int32], + block_info: BlockInfo, + SeqlenInfoCls: Callable, + mCuSeqlensQ: Optional[cute.Tensor], + tile_scheduler: TileSchedulerProtocol, + ): + # ==== mma warp ==== + # Description: Computes dP = V @ dO^T, dV = P^T @ dO, and dV += dS^T @ Qv + # i.e. dP = gemm(V, dO), dV += gemm(P.T, dO.T), dV += gemm(dS.T, Qv.T) + # Produces: dP, dV + # Consumes: V, dO, P.T, dO.T, dS.T, Qv.T + lane_idx = cute.arch.lane_idx() + + tdVtdVs = [tdVtdV0, tdVtdV1] + + # Set accumulate = True for dS^T@Qv since we are accumulating on the P^T@dO result + tiled_mma_dStQvt.set(tcgen05.Field.ACCUMULATE, True) + + # Operands for dP=V@dO^T + tdPrV = tiled_mma_VdO.make_fragment_A(sV) + tdPrdO = tiled_mma_VdO.make_fragment_B(sdO) + + # Operands for dVi=P^T@dOi + tdVrPt = tiled_mma_PtdOt.make_fragment_A(sPt) + tdVrdOt = tiled_mma_PtdOt.make_fragment_B(sdOt) + + # Operands for dVi+=dS^T@Qvi + tdVrdSt = tiled_mma_dStQvt.make_fragment_A(sdSt) + tdVrQvt = tiled_mma_dStQvt.make_fragment_B(sQvt) + + use_ptx_gemm_VdO = False + use_ptx_gemm_PtdOt = False + use_ptx_gemm_dStQvt = False + + # GEMM functions + if const_expr(use_ptx_gemm_VdO): + gemm_VdO = partial( + fa_sm100_utils.gemm_ptx_partial, + tiled_mma_VdO.op, + self.tmem_offset_dP, + zero_init=True, + cta_group=self.cta_group_size, + ) + else: + gemm_VdO = partial( + fa_sm100_utils.gemm, + tiled_mma_VdO, + tdPtdP, + ) + if const_expr(use_ptx_gemm_PtdOt): + gemm_PtdOt = [ + partial( + fa_sm100_utils.gemm_ptx_partial, + tiled_mma_PtdOt.op, + self.tmem_offsets_dV[split], + zero_init=True, + cta_group=self.cta_group_size, + ) + for split in range(self.num_hdimv_splits) + ] + else: + gemm_PtdOt = [ + partial( + fa_sm100_utils.gemm, + tiled_mma_PtdOt, + tdVtdVs[split], + zero_init=True, + ) + for split in range(self.num_hdimv_splits) + ] + if const_expr(use_ptx_gemm_dStQvt): + gemm_dStQvt = [ + partial( + fa_sm100_utils.gemm_ptx_partial, + tiled_mma_dStQvt.op, + self.tmem_offsets_dV[split], + zero_init=False, + cta_group=self.cta_group_size, + ) + for split in range(self.num_hdimv_splits) + ] + else: + gemm_dStQvt = [ + partial( + fa_sm100_utils.gemm, + tiled_mma_dStQvt, + tdVtdVs[split], + zero_init=False, + ) + for split in range(self.num_hdimv_splits) + ] + + Consumer, Producer = pipeline.PipelineUserType.Consumer, pipeline.PipelineUserType.Producer + consumer_state_V = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_V) + consumer_state_dO = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_dO) + consumer_state_Pt = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_Pt) + consumer_state_dOt_Qvt = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_dOt) + consumer_state_dSt = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_dSt) + producer_state_dPt = pipeline.make_pipeline_state(Producer, stages=self.num_stages_dPt) + producer_state_dV = pipeline.make_pipeline_state(Producer, stages=self.num_stages_dV) + + mma_VdO = partial( + self.mma_inner, gemm_VdO, pipeline_V, tdPrV, sV, tdPrdO, sdO, swap_AB_stage=True, use_ptx=use_ptx_gemm_VdO + ) + mma_PtdOt = partial( + self.mma_inner, gemm_PtdOt, pipeline_dOt_Qvt, tdVrPt, sPt, tdVrdOt, sdOt, use_ptx=use_ptx_gemm_PtdOt + ) + mma_dStQvt = partial( + self.mma_inner, gemm_dStQvt, pipeline_dOt_Qvt, tdVrdSt, sdSt, tdVrQvt, sQvt, use_ptx=use_ptx_gemm_dStQvt + ) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + # m_block, head_idx, batch_idx, _ = work_tile.tile_idx + # if const_expr(mCuSeqlensQ is not None): + # batch_idx = get_batch_from_cu_tensor(m_block, mCuSeqlensQ) + # seqlen = SeqlenInfoCls(batch_idx) + # num_n_block_groups = self.topk_length // self.cluster_tile_n + num_n_block_groups = topk_length_dynamic // self.cluster_tile_n + + if is_leader_cta: + # ==== Prologue ==== + consumer_wait_state_dO = consumer_state_dO.clone() + for split in cutlass.range_constexpr(self.num_hdimv_splits): + pipeline_dO.consumer_wait(consumer_wait_state_dO) + consumer_wait_state_dO.advance() + + # ==== Mainloop ==== + for _ in cutlass.range(num_n_block_groups, unroll=1): + # 1. dP = V @ dO^T + # mma inner waits for V + pipeline_dPt.producer_acquire(producer_state_dPt) + for split in cutlass.range_constexpr(self.num_hdimv_splits): + consumer_state_V = mma_VdO( + consumer_state_V, a_stage=split, zero_init=split == 0 + ) + pipeline_dPt.producer_commit(producer_state_dPt) + producer_state_dPt.advance() + + # 2. dV = P^T @ dO + # mma inner waits for dOt + pipeline_Pt.consumer_wait(consumer_state_Pt) + producer_acquire_state_dV = producer_state_dV.clone() + for split in cutlass.range_constexpr(self.num_hdimv_splits): + pipeline_dV.producer_acquire(producer_acquire_state_dV) + producer_acquire_state_dV.advance() + consumer_state_dOt_Qvt = mma_PtdOt(consumer_state_dOt_Qvt, acc_stage=split) + pipeline_Pt.consumer_release(consumer_state_Pt) + consumer_state_Pt.advance() + + # 3. dV += dS^T @ Qv + # mma inner waits for Qvt + pipeline_dSt.consumer_wait(consumer_state_dSt) + for split in cutlass.range_constexpr(self.num_hdimv_splits): + consumer_state_dOt_Qvt = mma_dStQvt(consumer_state_dOt_Qvt, acc_stage=split) + pipeline_dV.producer_commit(producer_state_dV) + producer_state_dV.advance() + pipeline_dSt.consumer_release(consumer_state_dSt) + consumer_state_dSt.advance() + + # ==== Epilogue ==== + for _ in cutlass.range_constexpr(self.num_hdimv_splits): + pipeline_dO.consumer_release(consumer_state_dO) + consumer_state_dO.advance() + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() + + pipeline_dPt.producer_tail(producer_state_dPt) + pipeline_dV.producer_tail(producer_state_dV) + + @cute.jit + def mma_inner( + self, + gemm, + load_pipeline, + tCrA, + sA, + tCrB, + sB, + consumer_state: pipeline.PipelineState, + acc_stage: Optional[Int32] = None, + a_stage: Int32 = 0, + zero_init: Optional[bool] = None, + swap_AB_stage: bool = False, + use_ptx: bool = True, + ): + if const_expr(acc_stage is not None): + gemm = gemm[acc_stage] + + smem_stage = consumer_state.index + + if const_expr(not swap_AB_stage): + a_stage = a_stage + b_stage = smem_stage + else: + a_stage = smem_stage + b_stage = a_stage + + tCrA_cur = tCrA[None, None, None, a_stage] + sA_cur = sA[None, None, None, a_stage] + tCrB_cur = tCrB[None, None, None, b_stage] + sB_cur = sB[None, None, None, b_stage] + + kwargs = dict(tCrA=tCrA_cur, tCrB=tCrB_cur) + if const_expr(use_ptx): + kwargs |= dict(sA=sA_cur, sB=sB_cur) + if const_expr(zero_init is not None): + kwargs["zero_init"] = zero_init + + load_pipeline.consumer_wait(consumer_state) + gemm(**kwargs) + load_pipeline.consumer_release(consumer_state) + consumer_state.advance() + return consumer_state + + @cute.jit + def compute_loop( + self, + softmax_scale: Float32, + softmax_scale_log2: Float32, + thr_mma_VdO: cute.ThrMma, + tdPtdP: cute.Tensor, + sP: cute.Tensor, + sdS: cute.Tensor, + sScaleP: cute.Tensor, + sdPsum: cute.Tensor, + mdS: cute.Tensor, + tma_atom_dS: cute.CopyAtom, + pipeline_P: pipeline.PipelineAsync, # TmaAsync + pipeline_Pt: pipeline.PipelineAsync, # AsyncUmma + pipeline_dPt: pipeline.PipelineAsync, # UmmaAsync + pipeline_dSt: pipeline.PipelineAsync, # AsyncUmma + pipeline_scaleP: pipeline.PipelineAsync, # TmaAsync + pipeline_dPsum: pipeline.PipelineAsync, # TmaAsync + topk_length_dynamic: Optional[Int32], + block_info: BlockInfo, + SeqlenInfoCls: Callable, + mCuSeqlensQ: Optional[cute.Tensor], + tile_scheduler: TileSchedulerProtocol, + ): + tidx = cute.arch.thread_idx()[0] % self.num_softmax_threads + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % ( + self.num_softmax_threads // 32 + ) + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + leader_warp = warp_idx == 0 + + # 256b // 32 = 8 values, 128 mqa // 2 => tmem_rep = 8 + tmem_rep = self.tile_m // self.cta_group_size // 8 + copy_atom_t2r = cute.make_copy_atom( + tcgen05.copy.Ld16x256bOp(tcgen05.copy.Repetition(tmem_rep)), + self.dtype_acc, + ) + # ((64,(64,2)),1,1):((65536,(1,4194304 = 65536*64)),0,0) + tdPtdP = tdPtdP[(None, None), 0, 0] + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tdPtdP) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N) + # (((64,16),1),2,1):(((1,65536),0),1048576,0)>, 1048576/65536 = 16 + tdPtdP_t2r = thr_copy_t2r.partition_S(tdPtdP) + + cdP = cute.make_identity_tensor(self.mma_tiler_VdO[:2]) # (128, 128) + tdPcdP = thr_mma_VdO.partition_C(cdP)[(None, None), 0, 0] # (64,128):(1@0,1@1) + # (((2,2,8),1),2,1):(((1@1,8@0,8@1),0),16@0,0) + tdPcdP_t2r = thr_copy_t2r.partition_D(tdPcdP) + assert tdPcdP_t2r.shape[0][1] == 1, f"unexpected tdPcdP_t2r shape, got {tdPcdP_t2r.shape}" + + smem_load_op = cute.nvgpu.warp.LdMatrix8x8x16bOp(True, 4) # ldsm x num_matrices = ldsm x 4 + smem_store_op = cute.nvgpu.warp.StMatrix8x8x16bOp(True, 4) # stsm x num_matrices = stsm x 4 + smem_load_atom = cute.make_copy_atom(smem_load_op, self.dtype) + smem_store_atom = cute.make_copy_atom(smem_store_op, self.dtype) + tiled_copy_r2s = cute.make_tiled_copy_D(smem_store_atom, tiled_copy_t2r) + tiled_copy_s2r = cute.make_tiled_copy_D(smem_load_atom, tiled_copy_t2r) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + thr_copy_s2r = tiled_copy_s2r.get_slice(tidx) + + sPt_load_layout = cute.make_ordered_layout( + self.tile_Pt + (self.num_stages_P,), order=(1, 0, 2) + ) + # (tile_n, tile_m, stages_P) + sPt = cute.composition(sP, sPt_load_layout) + sdSt = cute.composition(sdS, sPt_load_layout) + + # (R2S, R2S_M, R2S_N, PIPE_D) + # ((8,4),2,1,1):((1,1024),16,0,0) + tSR_sPt = thr_copy_s2r.partition_S(sPt) + tRS_sdSt = thr_copy_r2s.partition_D(sdSt) + + # ((2,2),(2,8,1),stage):((0,0),(1,8,0),_) + tPsScaleP_nm = self.broadcast_tensor_nm_view(sScaleP, thr_mma_VdO, thr_copy_t2r) + tdPsdPsum_nm = self.broadcast_tensor_nm_view(sdPsum, thr_mma_VdO, thr_copy_t2r) + + Consumer, Producer = pipeline.PipelineUserType.Consumer, pipeline.PipelineUserType.Producer + + consumer_state_P = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_P) + consumer_state_dPt = pipeline.make_pipeline_state(Consumer, stages=1) + consumer_state_scaleP = pipeline.make_pipeline_state( + Consumer, stages=self.num_stages_scaleP + ) + consumer_state_dPsum = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_dPsum) + + producer_state_Pt = pipeline.make_pipeline_state(Producer, stages=self.num_stages_Pt) + producer_state_dSt = pipeline.make_pipeline_state(Producer, stages=self.num_stages_dSt) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + m_block, head_idx, batch_idx, _ = work_tile.tile_idx + if const_expr(mCuSeqlensQ is not None): + batch_idx = get_batch_from_cu_tensor(m_block, mCuSeqlensQ) + seqlen = SeqlenInfoCls(batch_idx) + if const_expr(mCuSeqlensQ is not None): + m_block -= seqlen.offset_q + num_n_block_groups = self.topk_length // self.cluster_tile_n + # num_n_block_groups = topk_length_dynamic // self.cluster_tile_n + + mdS_cur = seqlen.offset_batch_Q(mdS, batch_idx, dim=3)[None, None, head_idx] + + gdS = cute.local_tile(mdS_cur, (self.tile_m, self.tile_n), (m_block, None)) + store_dS, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_dS, + 0, + cute.make_layout(1), + sdS, + gdS, + ) + + pipeline_dPsum.consumer_wait(consumer_state_dPsum) + + tdPsdPsum_cur = tdPsdPsum_nm[0, None, consumer_state_dPsum.index] + tdPrdPsum_cur_f32 = cute.make_rmem_tensor(tdPsdPsum_cur.shape, dtype=self.dtype_scale) + cute.autovec_copy(tdPsdPsum_cur, tdPrdPsum_cur_f32) + + for n_block_group in cutlass.range(num_n_block_groups, unroll=1): + n_block = 2 * n_block_group + cta_rank_in_cluster + + pipeline_P.consumer_wait(consumer_state_P) + # todo: ablate wait -> try_wait + pipeline_scaleP.consumer_wait(consumer_state_scaleP) + + # (((2,2,8),1),2,1):(((1,2,4),0),32,0) + rPt = cute.make_rmem_tensor(tdPcdP_t2r.shape, self.dtype) + # (S2R, S2R_M, S2R_N) + rPt_copy_view = tiled_copy_s2r.retile(rPt) + tSR_sPt_cur = tSR_sPt[None, None, None, consumer_state_P.index] + cute.copy(tiled_copy_s2r, tSR_sPt_cur, rPt_copy_view) + + # ((2,2),(2,8,1)):((2,32),(1,4,0)) + rP_nm = layout_utils.reshape_acc_to_mn(rPt[(None, 0), None, None]) + + tPsScaleP_cur = tPsScaleP_nm[0, None, consumer_state_scaleP.index] + tPrScaleP_cur_f32 = cute.make_rmem_tensor( + tPsScaleP_cur.shape, dtype=self.dtype_scale + ) + tPrScaleP_cur = cute.make_rmem_tensor(tPsScaleP_cur.shape, dtype=self.dtype) + cute.autovec_copy(tPsScaleP_cur, tPrScaleP_cur_f32) + tPrScaleP_cur.store(tPrScaleP_cur_f32.load().to(self.dtype)) + + # scale P + for n in cutlass.range_constexpr(cute.size(rP_nm.shape[0])): + rP_cur = rP_nm[n, None] + rP_cur.store(rP_cur.load() * tPrScaleP_cur.load()) + cute.arch.sync_warp() + + cute.copy(tiled_copy_r2s, rPt_copy_view, tSR_sPt_cur) + cute.arch.fence_view_async_shared() + self.softmax_barrier.arrive_and_wait() + + pipeline_scaleP.consumer_release(consumer_state_scaleP) + consumer_state_scaleP.advance() + + pipeline_Pt.producer_commit(producer_state_Pt) + producer_state_Pt.advance() + + # note: mma also signals Pt free, signal acquired in tma warp + pipeline_P.consumer_release(consumer_state_P) + consumer_state_P.advance() + + pipeline_dPt.consumer_wait(consumer_state_dPt) + + # (((2,2,8),1),2,1):(((1,2,4),0),32,0) + tdPrdP_t2r = cute.make_rmem_tensor(tdPcdP_t2r.shape, self.dtype_acc) + cute.copy(tiled_copy_t2r, tdPtdP_t2r, tdPrdP_t2r) + cute.arch.fence_view_async_tmem_load() + self.softmax_barrier.arrive_and_wait() + + pipeline_dPt.consumer_release(consumer_state_dPt) + consumer_state_dPt.advance() + + # dS = P o (dP - dPsum) + rdP_nm = layout_utils.reshape_acc_to_mn(tdPrdP_t2r[(None, 0), None, None]) + for n in cutlass.range_constexpr(cute.size(rdP_nm.shape[0])): + rdP_cur = rdP_nm[n, None] + rdP_cur.store(rdP_cur.load() - tdPrdPsum_cur_f32.load()) + + rPt.store(rPt.load() * (tdPrdP_t2r.load() * softmax_scale).to(self.dtype)) + + # wait for tma store to free dSt buffer + if leader_warp: + cute.arch.cp_async_bulk_wait_group(1 - self.num_stages_dSt, read=True) + self.softmax_barrier.arrive_and_wait() + + # note: dS guaranteed free as mma operand + pipeline_dSt.producer_acquire(producer_state_dSt) + + tRS_sdSt_cur = tRS_sdSt[None, None, None, producer_state_dSt.index] + cute.copy(tiled_copy_r2s, rPt_copy_view, tRS_sdSt_cur) + + cute.arch.fence_view_async_shared() + self.softmax_barrier.arrive_and_wait() + pipeline_dSt.producer_commit(producer_state_dSt) + + # tma store + if leader_warp: + store_dS(src_idx=producer_state_dSt.index, dst_idx=n_block) + cute.arch.cp_async_bulk_commit_group() + + producer_state_dSt.advance() + + pipeline_dPsum.consumer_release(consumer_state_dPsum) + consumer_state_dPsum.advance() + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() + + # producer tails + + @cute.jit + def broadcast_tensor_nm_view( + self, + sX: cute.Tensor, # (tile_m, num_stages) + thr_mma: cute.ThrMma, + thr_copy_t2r: cute.ThrCopy, + ): + assert cute.size(sX.shape[0]) == self.tile_m + num_stages = sX.shape[1] if const_expr(cute.rank(sX) > 1) else 1 + sX_2D_cluster = cute.make_tensor( + sX.iterator, + cute.make_layout( + (self.tile_m, self.cluster_tile_n, num_stages), + stride=(1, 0, self.tile_m), + ), + ) + sXt_2D_cluster = layout_utils.transpose_view(sX_2D_cluster) + sXt_2D = thr_mma.partition_C(sXt_2D_cluster)[(None, None), 0, 0, None] + tXsXt_2D = thr_copy_t2r.partition_D(sXt_2D)[(None, 0), None, None, None] + tXsXt_nm = layout_utils.make_acc_tensor_mn_view(tXsXt_2D) + return tXsXt_nm + + @cute.jit + def dVacc_store( + self, + mIndexTopk: cute.Tensor, + mdV: cute.Tensor, + sdV: cute.Tensor, + tdVtdV0: cute.Tensor, + tdVtdV1: cute.Tensor, + thr_mma_PtdOt: cute.ThrMma, + pipeline_dV: pipeline.PipelineAsync, # UmmaAsync + pipeline_dV_epi: pipeline.PipelineAsync, # Async + topk_length_dynamic: Optional[Int32], + block_info: BlockInfo, + SeqlenInfoCls: Callable, + mCuSeqlensQ: Optional[cute.Tensor], + tile_scheduler: TileSchedulerProtocol, + ): + # ==== dVaccum store warpgroup ==== + # produces: - + # consumes: dV + + tdVtdV0 = tdVtdV0[(None, None), 0, 0] + tdVtdV1 = tdVtdV1[(None, None), 0, 0] + + num_epi_warps = self.num_epilogue_threads // 32 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + tidx = cute.arch.thread_idx()[0] % self.num_epilogue_threads + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % num_epi_warps + leader_warp = warp_idx == 0 + wg_half = warp_idx // 2 + + consumer_state_dV = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, stages=self.num_stages_dV + ) + + copy_atom_t2r = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), + self.dtype_acc, + ) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tdVtdV0) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tdVtdV0_t2r = thr_copy_t2r.partition_S(tdVtdV0) + tdVtdV1_t2r = thr_copy_t2r.partition_S(tdVtdV1) + tdVtdVs_t2r = [tdVtdV0_t2r, tdVtdV1_t2r] + + cdVmma = cute.make_identity_tensor(self.mma_tiler_PtdOt[:2]) + tdVcdVmma = thr_mma_PtdOt.partition_C(cdVmma)[(None, None), 0, 0] + tdVcdVmma_t2r = thr_copy_t2r.partition_D(tdVcdVmma) + + # 64 threads x 4 values to tile over tile_dV = (64, 32) + tiled_copy_r2s = tiled_copy_2d(self.dtype_acc, 4, 64) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx % 64) + + # ((4,1),1,8,(1,8)):((1,0),0,4,(0,2048)) + tRS_sdV = thr_copy_r2s.partition_D(sdV) + + tiled_copy_s2r = copy_utils.tiled_copy_2d(self.dtype_acc, 8, self.num_epilogue_threads, 4) + thr_copy_s2r = tiled_copy_s2r.get_slice(tidx) + # (V, M, N, STAGE) + tSR_sdV = thr_copy_s2r.partition_S(sdV) + + cdV = cute.make_identity_tensor(cute.product_each(sdV.shape[:2])) + # (V, M, N) + tdVcdV = thr_copy_s2r.partition_S(cdV) + + gmem_rows_per_thread = cute.size(tSR_sdV.shape[1]) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + m_block, head_idx, batch_idx, _ = work_tile.tile_idx + if const_expr(mCuSeqlensQ is not None): + batch_idx = get_batch_from_cu_tensor(m_block, mCuSeqlensQ) + seqlen = SeqlenInfoCls(batch_idx) + num_n_block_groups = self.topk_length // self.cluster_tile_n + # num_n_block_groups = topk_length_dynamic // self.cluster_tile_n + + # (seqlen_k, hdimv) + mdV_cur = seqlen.offset_batch_K(mdV, batch_idx, dim=3)[None, None, head_idx] + + # (topk, dv) + if const_expr(seqlen.has_cu_seqlens_q): + # m_block means absolute m_idx + mIndexTopk_cur = mIndexTopk[None, m_block] + else: + mIndexTopk_cur = mIndexTopk[None, m_block, batch_idx] + + # ==== Mainloop ==== + for n_block_group in cutlass.range(num_n_block_groups, unroll=1): + n_block = 2 * n_block_group + cta_rank_in_cluster + + rIdxTopK = cute.make_rmem_tensor((gmem_rows_per_thread,), dtype=self.dtype_index) + for j in cutlass.range_constexpr(gmem_rows_per_thread): + n_idx = n_block * self.tile_n + tdVcdV[0, j, 0][0] + rIdxTopK[j] = mIndexTopk_cur[n_idx] + + for split in cutlass.range_constexpr(self.num_hdimv_splits): + tdVtdV_t2r = tdVtdVs_t2r[split] + + pipeline_dV.consumer_wait(consumer_state_dV) + + # TODO: record meaning of hard-coded values + num_cols_per_store = self.tile_dV[1] * 2 + num_epi_subtiles = (self.hdimv // self.num_hdimv_splits) // num_cols_per_store + assert num_cols_per_store == 64 + assert num_epi_subtiles == 4 + assert cute.size(tdVtdV_t2r.shape[2]) == num_epi_subtiles + + tdVrdV_cur_shape = tdVcdVmma_t2r[None, None, 0].shape + tRS_rdV_cur_shape = tRS_sdV[None, None, None, 0].shape + assert cute.size(tdVrdV_cur_shape) == cute.size(tRS_rdV_cur_shape) + + tdVrdV_out_shape = tSR_sdV[None, None, None, 0].shape + (2,) + + for i in cutlass.range_constexpr(num_epi_subtiles): + tdVrdV_cur = cute.make_rmem_tensor(tdVrdV_cur_shape, self.dtype_acc) + cute.copy(tiled_copy_t2r, tdVtdV_t2r[None, None, i], tdVrdV_cur) + + tRS_rdV_cur = cute.make_tensor(tdVrdV_cur.iterator, tRS_rdV_cur_shape) + + stage = 4 * split + 2 * wg_half + (i % 2) + cute.copy(tiled_copy_r2s, tRS_rdV_cur, tRS_sdV[None, None, None, stage]) + cute.arch.fence_view_async_shared() + self.epi_barrier.arrive_and_wait() + + tSR_rdV = cute.make_rmem_tensor(tdVrdV_out_shape, dtype=self.dtype_acc) + + for w in cutlass.range_constexpr(2): + stage_out = 4 * split + 2 * w + (i % 2) + cute.copy( + tiled_copy_s2r, + tSR_sdV[None, None, None, stage_out], + tSR_rdV[None, None, None, w], + ) + + for j in cutlass.range_constexpr(gmem_rows_per_thread): + gmem_n_idx = rIdxTopK[j] + for w in cutlass.range_constexpr(2): + dv_offset = ( + self.hdimv // self.num_hdimv_splits * split # 256 * split + + (self.hdimv // self.num_hdimv_splits // 2) * w # 128 * w + + 32 * i + ) + dv_offset += tdVcdV[0, j, 0][1] + gmem_coord = (gmem_n_idx, dv_offset) + dV_gmem_ptr = elem_pointer(mdV_cur, gmem_coord) + + a = tSR_rdV[0, j, 0, w] + b = tSR_rdV[1, j, 0, w] + c = tSR_rdV[2, j, 0, w] + d = tSR_rdV[3, j, 0, w] + atomic_add_fp32x4(a, b, c, d, dV_gmem_ptr) + + cute.arch.fence_view_async_tmem_load() + self.epi_barrier.arrive_and_wait() + pipeline_dV.consumer_release(consumer_state_dV) + + if leader_warp: + with cute.arch.elect_one(): + pipeline_dV_epi.consumer_release(consumer_state_dV) + + consumer_state_dV.advance() + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() diff --git a/fastvideo-kernel/fa4/flash_bwd_postprocess.py b/fastvideo-kernel/fa4/flash_bwd_postprocess.py new file mode 100644 index 000000000..913b43d37 --- /dev/null +++ b/fastvideo-kernel/fa4/flash_bwd_postprocess.py @@ -0,0 +1,587 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +# A reimplementation of https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_bwd_postprocess_kernel.h +# from Cutlass C++ to Cute-DSL. +import math +from typing import Callable, Optional, Type + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.utils.hopper_helpers as sm90_utils_basic +import cutlass.utils.blackwell_helpers as sm100_utils_basic +from cutlass.cute.nvgpu import cpasync, warp, warpgroup +from cutlass import Float32, const_expr +from cutlass.utils import LayoutEnum + +from quack import copy_utils +from quack import layout_utils +from quack import sm90_utils + +from flash_attn.cute import utils +from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned +from flash_attn.cute import ampere_helpers as sm80_utils +from flash_attn.cute.seqlen_info import SeqlenInfoQK +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +from quack.cute_dsl_utils import ParamsBase +from flash_attn.cute.tile_scheduler import ( + SingleTileScheduler, + SingleTileVarlenScheduler, + TileSchedulerArguments, +) + + +class FlashAttentionBackwardPostprocess: + def __init__( + self, + dtype: Type[cutlass.Numeric], + head_dim: int, + arch: int, + tile_m: int = 128, + num_threads: int = 256, + AtomLayoutMdQ: int = 1, + dQ_swapAB: bool = False, + use_2cta_instrs: bool = False, + cluster_size: int = 1, # for varlen offsets + ): + """ + :param head_dim: head dimension + :type head_dim: int + :param tile_m: m block size + :type tile_m: int + """ + self.dtype = dtype + self.tile_m = tile_m + assert arch // 10 in [8, 9, 10, 11, 12], ( + "Only Ampere (8.x), Hopper (9.x), and Blackwell (10.x, 11.x, 12.x) are supported" + ) + self.arch = arch + # padding head_dim to a multiple of 32 as k_block_size + hdim_multiple_of = 32 + self.tile_hdim = int(math.ceil(head_dim / hdim_multiple_of) * hdim_multiple_of) + self.check_hdim_oob = head_dim != self.tile_hdim + self.num_threads = num_threads + self.AtomLayoutMdQ = AtomLayoutMdQ + self.dQ_swapAB = dQ_swapAB + self.use_2cta_instrs = use_2cta_instrs and arch // 10 in [10, 11] and head_dim != 64 + self.cluster_size = cluster_size + + @staticmethod + def can_implement(dtype, head_dim, tile_m, num_threads) -> bool: + """Check if the kernel can be implemented with the given parameters. + + :param dtype: data type + :type dtype: cutlass.Numeric + :param head_dim: head dimension + :type head_dim: int + :param tile_m: m block size + :type tile_m: int + + :return: True if the kernel can be implemented, False otherwise + :rtype: bool + """ + if dtype not in [cutlass.Float16, cutlass.BFloat16]: + return False + if head_dim % 8 != 0: + return False + if num_threads % 32 != 0: + return False + return True + + def _get_tiled_mma(self): + if const_expr(self.arch // 10 in [8, 12]): + num_mma_warps = self.num_threads // 32 + atom_layout_dQ = ( + (self.AtomLayoutMdQ, num_mma_warps // self.AtomLayoutMdQ, 1) + if const_expr(not self.dQ_swapAB) + else (num_mma_warps // self.AtomLayoutMdQ, self.AtomLayoutMdQ, 1) + ) + tiled_mma = cute.make_tiled_mma( + warp.MmaF16BF16Op(self.dtype, Float32, (16, 8, 16)), + atom_layout_dQ, + permutation_mnk=(atom_layout_dQ[0] * 16, atom_layout_dQ[1] * 16, 16), + ) + elif const_expr(self.arch // 10 == 9): + num_wg_mma = self.num_threads // 128 + atom_layout_dQ = (self.AtomLayoutMdQ, num_wg_mma // self.AtomLayoutMdQ) + tiler_mn_dQ = (self.tile_m // atom_layout_dQ[0], self.tile_hdim // atom_layout_dQ[1]) + tiled_mma = sm90_utils_basic.make_trivial_tiled_mma( + self.dtype, + self.dtype, + warpgroup.OperandMajorMode.K, # These don't matter, we only care about the accum + warpgroup.OperandMajorMode.K, + Float32, + atom_layout_mnk=(atom_layout_dQ if not self.dQ_swapAB else atom_layout_dQ[::-1]) + + (1,), + tiler_mn=tiler_mn_dQ if not self.dQ_swapAB else tiler_mn_dQ[::-1], + ) + else: + cta_group = tcgen05.CtaGroup.ONE + tiled_mma = sm100_utils_basic.make_trivial_tiled_mma( + self.dtype, + tcgen05.OperandMajorMode.MN, # dS_major_mode + tcgen05.OperandMajorMode.MN, # Kt_major_mode + Float32, + cta_group, + (self.tile_m, self.tile_hdim), + ) + if const_expr(self.arch // 10 in [8, 9, 12]): + assert self.num_threads == tiled_mma.size + return tiled_mma + + def _setup_attributes(self): + # /////////////////////////////////////////////////////////////////////////////// + # GMEM Tiled copy: + # /////////////////////////////////////////////////////////////////////////////// + # Thread layouts for copies + universal_copy_bits = 128 + async_copy_elems_accum = universal_copy_bits // Float32.width + atom_async_copy_accum = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + Float32, + num_bits_per_copy=universal_copy_bits, + ) + # We don't do bound checking for the gmem -> smem load so we just assert here. + assert (self.tile_m * self.tile_hdim // async_copy_elems_accum) % self.num_threads == 0 + self.g2s_tiled_copy_dQaccum = cute.make_tiled_copy_tv( + atom_async_copy_accum, + cute.make_layout(self.num_threads), + cute.make_layout(async_copy_elems_accum), + ) + num_s2r_copy_elems = 1 if const_expr(self.arch // 10 in [8, 12]) else 4 + if const_expr(self.arch // 10 in [8, 12]): + self.s2r_tiled_copy_dQaccum = copy_utils.tiled_copy_1d( + Float32, self.num_threads, num_s2r_copy_elems + ) + self.sdQaccum_layout = cute.make_layout(self.tile_m * self.tile_hdim) + elif const_expr(self.arch // 10 == 9): + num_threads_per_warp_group = 128 + num_wg_mma = self.num_threads // 128 + self.s2r_tiled_copy_dQaccum = cute.make_tiled_copy_tv( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128), + cute.make_layout((num_threads_per_warp_group, num_wg_mma)), # thr_layout + cute.make_layout(128 // Float32.width), # val_layout + ) + self.sdQaccum_layout = cute.make_layout( + (self.tile_m * self.tile_hdim // num_wg_mma, num_wg_mma) + ) + else: + self.dQ_reduce_ncol = 32 + dQaccum_reduce_stage = self.tile_hdim // self.dQ_reduce_ncol + assert self.num_threads == 128 # TODO: currently hard-coded + self.s2r_tiled_copy_dQaccum = copy_utils.tiled_copy_1d( + Float32, self.num_threads, num_s2r_copy_elems + ) + self.sdQaccum_layout = cute.make_layout( + (self.tile_m * self.tile_hdim // dQaccum_reduce_stage, dQaccum_reduce_stage) + ) + + num_copy_elems = 128 // self.dtype.width + threads_per_row = math.gcd(128, self.tile_hdim) // num_copy_elems + self.gmem_tiled_copy_dQ = copy_utils.tiled_copy_2d( + self.dtype, threads_per_row, self.num_threads, num_copy_elems + ) + # /////////////////////////////////////////////////////////////////////////////// + # Shared memory layout: dQ + # /////////////////////////////////////////////////////////////////////////////// + # We can't just use kHeadDim here. E.g. if MMA shape is 64 x 96 but split across 2 WGs, + # then setting kBlockKSmem to 32 will cause "Static shape_div failure". + # We want to treat it as 64 x 48, so kBlockKSmem should be 16. + mma_shape_n = self.tiled_mma.get_tile_size(1) + if const_expr(self.arch // 10 in [8, 12]): + sdQ_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, mma_shape_n) + self.sdQ_layout = cute.tile_to_shape( + sdQ_layout_atom, (self.tile_m, self.tile_hdim), (0, 1) + ) + elif const_expr(self.arch // 10 == 9): + wg_d_dQ = num_wg_mma // self.AtomLayoutMdQ + self.sdQ_layout = sm90_utils.make_smem_layout( + self.dtype, + LayoutEnum.ROW_MAJOR, + (self.tile_m, self.tile_hdim), + major_mode_size=self.tile_hdim // wg_d_dQ, + ) + else: + # TODO: this is hard-coded for hdim 128 + self.sdQ_layout = sm100_utils_basic.make_smem_layout_epi( + self.dtype, LayoutEnum.ROW_MAJOR, (self.tile_m, self.tile_hdim), 1 + ) + + @cute.jit + def __call__( + self, + mdQaccum: cute.Tensor, + mdQ: cute.Tensor, + scale: cutlass.Float32, + mCuSeqlensQ: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + # Get the data type and check if it is fp16 or bf16 + if const_expr(mdQ.element_type not in [cutlass.Float16, cutlass.BFloat16]): + raise TypeError("Only Float16 or BFloat16 is supported") + if const_expr(mdQaccum is not None): + if const_expr(mdQaccum.element_type not in [cutlass.Float32]): + raise TypeError("dQaccum tensor must be Float32") + + mdQaccum, mdQ = [assume_tensor_aligned(t) for t in (mdQaccum, mdQ)] + + self.tiled_mma = self._get_tiled_mma() + self._setup_attributes() + + smem_size = max( + cute.size_in_bytes(cutlass.Float32, self.sdQaccum_layout), + cute.size_in_bytes(self.dtype, self.sdQ_layout), + ) + + if const_expr(mCuSeqlensQ is not None): + TileScheduler = SingleTileVarlenScheduler + num_head = mdQ.shape[1] + num_batch = mCuSeqlensQ.shape[0] - 1 + num_block = cute.ceil_div(mdQ.shape[0], self.tile_m) + else: + TileScheduler = SingleTileScheduler + num_head = mdQ.shape[2] + num_batch = mdQ.shape[0] + num_block = cute.ceil_div(mdQ.shape[1], self.tile_m) + + tile_sched_args = TileSchedulerArguments( + num_block=num_block, + num_head=num_head, + num_batch=num_batch, + num_splits=1, + seqlen_k=0, + headdim=mdQ.shape[2], + headdim_v=0, + total_q=mdQ.shape[0], + tile_shape_mn=(self.tile_m, 1), + mCuSeqlensQ=mCuSeqlensQ, + mSeqUsedQ=mSeqUsedQ, + ) + + tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args) + grid_dim = TileScheduler.get_grid_shape(tile_sched_params) + + # grid_dim: (m_block, num_head, batch_size) + self.kernel( + mdQaccum, + mdQ, + mCuSeqlensQ, + mSeqUsedQ, + scale, + self.tiled_mma, + self.dQ_swapAB, + self.sdQaccum_layout, + self.sdQ_layout, + self.g2s_tiled_copy_dQaccum, + self.s2r_tiled_copy_dQaccum, + self.gmem_tiled_copy_dQ, + tile_sched_params, + TileScheduler, + ).launch( + grid=grid_dim, + block=[self.num_threads, 1, 1], + smem=smem_size, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mdQaccum: cute.Tensor, + mdQ: cute.Tensor, + mCuSeqlensQ: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + scale: cutlass.Float32, + tiled_mma: cute.TiledMma, + dQ_swapAB: cutlass.Constexpr, + sdQaccum_layout: cute.Layout, + sdQ_layout: cute.ComposedLayout, + g2s_tiled_copy_dQaccum: cute.TiledCopy, + s2r_tiled_copy_dQaccum: cute.TiledCopy, + gmem_tiled_copy_dQ: cute.TiledCopy, + tile_sched_params: ParamsBase, + TileScheduler: cutlass.Constexpr[Callable], + ): + # /////////////////////////////////////////////////////////////////////////////// + # Get shared memory buffer + # /////////////////////////////////////////////////////////////////////////////// + smem = cutlass.utils.SmemAllocator() + sdQaccum = smem.allocate_tensor(cutlass.Float32, sdQaccum_layout, byte_alignment=1024) + sdQaccum_flat = cute.make_tensor(sdQaccum.iterator, cute.make_layout(cute.size(sdQaccum))) + if const_expr(self.arch // 10 in [8, 9, 12]): + sdQ = cute.make_tensor(cute.recast_ptr(sdQaccum.iterator, dtype=self.dtype), sdQ_layout) + else: + # extra stage dimension + sdQ = cute.make_tensor( + cute.recast_ptr(sdQaccum.iterator, sdQ_layout.inner, dtype=self.dtype), + sdQ_layout.outer, + )[None, None, 0] + sdQt = layout_utils.transpose_view(sdQ) + + # Thread index, block index + tidx, _, _ = cute.arch.thread_idx() + + tile_scheduler = TileScheduler.create(tile_sched_params) + work_tile = tile_scheduler.initial_work_tile_info() + + m_block, head_idx, batch_idx, _ = work_tile.tile_idx + + if work_tile.is_valid_tile: + # /////////////////////////////////////////////////////////////////////////////// + # Get the appropriate tiles for this thread block. + # /////////////////////////////////////////////////////////////////////////////// + + seqlen = SeqlenInfoQK.create( + batch_idx, + mdQ.shape[1], + 0, + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=None, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=None, + tile_m=self.tile_m * self.cluster_size, + ) + if const_expr(not seqlen.has_cu_seqlens_q): + mdQ_cur = mdQ[batch_idx, None, head_idx, None] + mdQaccum_cur = mdQaccum[batch_idx, head_idx, None] + head_dim = mdQ.shape[3] + else: + padded_offset_q = seqlen.padded_offset_q + mdQ_cur = cute.domain_offset((seqlen.offset_q, 0), mdQ[None, head_idx, None]) + mdQaccum_cur = cute.domain_offset( + (padded_offset_q * self.tile_hdim,), mdQaccum[head_idx, None] + ) + head_dim = mdQ.shape[2] + + # HACK: Compiler doesn't seem to recognize that padding + # by padded_offset_q * self.tile_hdim keeps alignment + # since statically divisible by 4 + + mdQaccum_cur_ptr = cute.make_ptr( + dtype=mdQaccum_cur.element_type, + value=mdQaccum_cur.iterator.toint(), + mem_space=mdQaccum_cur.iterator.memspace, + assumed_align=mdQaccum.iterator.alignment, + ) + mdQaccum_cur = cute.make_tensor(mdQaccum_cur_ptr, mdQaccum_cur.layout) + + gdQaccum = cute.local_tile(mdQaccum_cur, (self.tile_m * self.tile_hdim,), (m_block,)) + gdQ = cute.local_tile(mdQ_cur, (self.tile_m, self.tile_hdim), (m_block, 0)) + + seqlen_q = seqlen.seqlen_q + seqlen_q_rounded = cute.round_up(seqlen_q, self.tile_m) + + if const_expr(self.arch // 10 in [10, 11] and self.use_2cta_instrs): + # 2-CTA: remap dQaccum layout into TMEM view before writing sdQ + num_reduce_threads = self.num_threads + thr_mma_dsk = tiled_mma.get_slice(tidx) + dQacc_shape = thr_mma_dsk.partition_shape_C((self.tile_m, self.tile_hdim)) + tdQtdQ = thr_mma_dsk.make_fragment_C(dQacc_shape) + tdQtdQ = cute.make_tensor(tdQtdQ.iterator, tdQtdQ.layout) + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.dQ_reduce_ncol)), Float32 + ) + tiled_tmem_ld = tcgen05.make_tmem_copy(tmem_load_atom, tdQtdQ) + thr_tmem_ld = tiled_tmem_ld.get_slice(tidx) + + cdQ = cute.make_identity_tensor((self.tile_m, self.tile_hdim)) + tdQcdQ = thr_mma_dsk.partition_C(cdQ) + tdQcdQ_tensor = cute.make_tensor(tdQcdQ.iterator, tdQcdQ.layout) + tdQrdQ = thr_tmem_ld.partition_D(tdQcdQ_tensor) + + tiled_copy_accum = s2r_tiled_copy_dQaccum + g2s_thr_copy = tiled_copy_accum.get_slice(tidx) + + # S -> R + tdQrdQ_fp32 = cute.make_rmem_tensor(tdQrdQ.shape, cutlass.Float32) + tdQrdQ_s2r = cute.make_tensor(tdQrdQ_fp32.iterator, tdQrdQ_fp32.shape) + + smem_copy_atom = sm100_utils_basic.get_smem_store_op( + LayoutEnum.ROW_MAJOR, self.dtype, cutlass.Float32, tiled_tmem_ld + ) + r2s_tiled_copy = cute.make_tiled_copy( + smem_copy_atom, + layout_tv=tiled_tmem_ld.layout_dst_tv_tiled, + tiler_mn=tiled_tmem_ld.tiler_mn, + ) + tdQsdQ_r2s = thr_tmem_ld.partition_D(thr_mma_dsk.partition_C(sdQ)) + tdQrdQ_r2s = cute.make_rmem_tensor(tdQsdQ_r2s.shape, self.dtype) + + num_stages = cute.size(tdQrdQ_fp32, mode=[1]) + stage_stride = self.dQ_reduce_ncol + row_groups = 2 + assert num_stages % row_groups == 0 + assert num_reduce_threads % row_groups == 0 + stage_groups = num_stages // row_groups + threads_per_row_group = num_reduce_threads // row_groups + stage_loads = tuple((row_group, row_group) for row_group in range(row_groups)) + stage_iters = tuple( + (row_group, row_group * threads_per_row_group) + for row_group in range(row_groups) + ) + s2r_lane = tidx % threads_per_row_group + s2r_buf = tidx // threads_per_row_group + + gdQaccum_layout_g2s = cute.make_layout( + shape=(self.tile_m * self.dQ_reduce_ncol, 1), stride=(1, 0) + ) + sdQaccum_g2s = g2s_thr_copy.partition_D(sdQaccum) + + # G -> S + for stage_group in cutlass.range_constexpr(stage_groups): + for stage_offset, smem_buf in stage_loads: + stage_idx = stage_group + stage_offset * stage_groups + gdQaccum_stage = cute.local_tile( + gdQaccum, + (self.tile_m * self.dQ_reduce_ncol,), + (stage_idx,), + ) + gdQaccum_stage_g2s = cute.make_tensor( + gdQaccum_stage.iterator, + gdQaccum_layout_g2s, + ) + tdQgdQ = g2s_thr_copy.partition_S(gdQaccum_stage_g2s) + cute.copy( + g2s_thr_copy, + tdQgdQ[None, None, 0], + sdQaccum_g2s[None, None, smem_buf], + ) + + cute.arch.fence_view_async_shared() + cute.arch.barrier(barrier_id=6, number_of_threads=num_reduce_threads) + + # S -> R + for stage_offset, lane_offset in stage_iters: + stage_idx = stage_group + stage_offset * stage_groups + s2r_src_tidx = s2r_lane + lane_offset + s2r_thr_copy = tiled_copy_accum.get_slice(s2r_src_tidx) + sdQaccum_src = s2r_thr_copy.partition_S(sdQaccum)[None, None, s2r_buf] + + tdQrdQ_s2r_cpy = tdQrdQ_s2r[None, stage_idx, None, None] + tdQrdQ_r2s_cpy = cute.make_tensor( + tdQrdQ_s2r_cpy.iterator, cute.make_layout(sdQaccum_src.shape) + ) + cute.copy(s2r_thr_copy, sdQaccum_src, tdQrdQ_r2s_cpy) + cute.arch.fence_view_async_shared() + cute.arch.barrier(barrier_id=7, number_of_threads=num_reduce_threads) + + # R -> S + stage_lo = stage_idx % stage_stride + stage_hi = stage_idx // stage_stride + tdQrdQ_r2s_cpy = cute.make_tensor( + cute.recast_ptr(tdQrdQ_r2s_cpy.iterator), + tdQrdQ_r2s[((None, 0), (stage_lo, stage_hi), 0, 0)].shape, + ) + dQ_vec = tdQrdQ_r2s_cpy.load() * scale + tdQrdQ_r2s[((None, 0), (stage_lo, stage_hi), 0, 0)].store( + dQ_vec.to(self.dtype) + ) + + # R -> S + cute.copy( + r2s_tiled_copy, + tdQrdQ_r2s[None, None, None, 0], + tdQsdQ_r2s[None, None, None, 0], + ) + cute.arch.fence_view_async_shared() + cute.arch.barrier(barrier_id=8, number_of_threads=num_reduce_threads) + else: + # Step 1: load dQaccum from gmem to smem + g2s_thr_copy_dQaccum = g2s_tiled_copy_dQaccum.get_slice(tidx) + tdQgdQaccum = g2s_thr_copy_dQaccum.partition_S(gdQaccum) + tdQsdQaccumg2s = g2s_thr_copy_dQaccum.partition_D(sdQaccum_flat) + cute.copy(g2s_tiled_copy_dQaccum, tdQgdQaccum, tdQsdQaccumg2s) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_wait_group(0) + cute.arch.barrier() + + # Step 2: load dQ from smem to rmem + s2r_thr_copy_dQaccum = s2r_tiled_copy_dQaccum.get_slice(tidx) + tdQsdQaccum = s2r_thr_copy_dQaccum.partition_S(sdQaccum) + tile_shape = (self.tile_m, self.tile_hdim) + acc = None + tiled_copy_t2r = None + if const_expr(self.arch // 10 in [8, 9, 12]): + acc_shape = tiled_mma.partition_shape_C( + tile_shape if const_expr(not dQ_swapAB) else tile_shape[::-1] + ) + acc = cute.make_rmem_tensor(acc_shape, cutlass.Float32) + assert cute.size(acc) == cute.size(tdQsdQaccum) + else: + thr_mma = tiled_mma.get_slice(0) # 1-CTA + dQacc_shape = tiled_mma.partition_shape_C((self.tile_m, self.tile_hdim)) + tdQtdQ = tiled_mma.make_fragment_C(dQacc_shape) + tdQcdQ = thr_mma.partition_C( + cute.make_identity_tensor((self.tile_m, self.tile_hdim)) + ) + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.dQ_reduce_ncol)), + Float32, + ) + tiled_copy_t2r = tcgen05.make_tmem_copy(tmem_load_atom, tdQtdQ) + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + tdQrdQ_t2r_shape = thr_copy_t2r.partition_D(tdQcdQ).shape + acc = cute.make_rmem_tensor(tdQrdQ_t2r_shape, Float32) + tdQrdQaccum = cute.make_tensor(acc.iterator, cute.make_layout(tdQsdQaccum.shape)) + cute.autovec_copy(tdQsdQaccum, tdQrdQaccum) + # Convert tdQrdQaccum from fp32 to fp16/bf16 + rdQ = cute.make_fragment_like(acc, self.dtype) + rdQ.store((acc.load() * scale).to(self.dtype)) + + # Step 3: Copy dQ from register to smem + cute.arch.barrier() # make sure all threads have finished loading dQaccum + if const_expr(self.arch // 10 in [8, 9, 12]): + copy_atom_r2s_dQ = utils.get_smem_store_atom( + self.arch, self.dtype, transpose=self.dQ_swapAB + ) + tiled_copy_r2s_dQ = cute.make_tiled_copy_C(copy_atom_r2s_dQ, tiled_mma) + else: + # copy_atom_r2s_dQ = sm100_utils_basic.get_smem_store_op( + # LayoutEnum.ROW_MAJOR, self.dtype, Float32, tiled_copy_t2r, + # ) + # tiled_copy_r2s_dQ = cute.make_tiled_copy_D(copy_atom_r2s_dQ, tiled_copy_t2r) + thr_layout_r2s_dQ = cute.make_layout((self.num_threads, 1)) # 128 threads + val_layout_r2s_dQ = cute.make_layout((1, 128 // self.dtype.width)) + copy_atom_r2s_dQ = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.dtype, + num_bits_per_copy=128, + ) + tiled_copy_r2s_dQ = cute.make_tiled_copy_tv( + copy_atom_r2s_dQ, thr_layout_r2s_dQ, val_layout_r2s_dQ + ) + thr_copy_r2s_dQ = tiled_copy_r2s_dQ.get_slice(tidx) + cdQ = cute.make_identity_tensor((self.tile_m, self.tile_hdim)) + if const_expr(self.arch // 10 in [8, 9, 12]): + taccdQrdQ = thr_copy_r2s_dQ.retile(rdQ) + else: + taccdQcdQ_shape = thr_copy_r2s_dQ.partition_S(cdQ).shape + taccdQrdQ = cute.make_tensor(rdQ.iterator, taccdQcdQ_shape) + taccdQsdQ = thr_copy_r2s_dQ.partition_D( + sdQ if const_expr(not self.dQ_swapAB) else sdQt + ) + cute.copy(thr_copy_r2s_dQ, taccdQrdQ, taccdQsdQ) + + # Step 4: Copy dQ from smem to register to prepare for coalesced write to gmem + cute.arch.barrier() # make sure all smem stores are done + gmem_thr_copy_dQ = gmem_tiled_copy_dQ.get_slice(tidx) + tdQgdQ = gmem_thr_copy_dQ.partition_S(gdQ) + tdQsdQ = gmem_thr_copy_dQ.partition_D(sdQ) + tdQrdQ = cute.make_fragment_like(tdQsdQ, self.dtype) + # TODO: check OOB when reading from smem if kBlockM isn't evenly tiled + cute.autovec_copy(tdQsdQ, tdQrdQ) + + # Step 5: Copy dQ from register to gmem + tdQcdQ = gmem_thr_copy_dQ.partition_S(cdQ) + tdQpdQ = utils.predicate_k(tdQcdQ, limit=head_dim) + for rest_m in cutlass.range(cute.size(tdQrdQ.shape[1]), unroll_full=True): + if tdQcdQ[0, rest_m, 0][0] < seqlen_q - m_block * self.tile_m: + cute.copy( + gmem_tiled_copy_dQ, + tdQrdQ[None, rest_m, None], + tdQgdQ[None, rest_m, None], + pred=tdQpdQ[None, rest_m, None], + ) diff --git a/fastvideo-kernel/fa4/flash_bwd_preprocess.py b/fastvideo-kernel/fa4/flash_bwd_preprocess.py new file mode 100644 index 000000000..8019a6038 --- /dev/null +++ b/fastvideo-kernel/fa4/flash_bwd_preprocess.py @@ -0,0 +1,465 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +# A reimplementation of https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_bwd_preprocess_kernel.h +# from Cutlass C++ to Cute-DSL. +# +# Computes D_i = (dO_i * O_i).sum(dim=-1), optionally adjusted for LSE gradient: +# D'_i = D_i - dLSE_i +# This works because in the backward pass: +# dS_ij = P_ij * (dP_ij - D_i) [standard] +# When LSE is differentiable, d(loss)/d(S_ij) gets an extra term dLSE_i * P_ij +# (since d(LSE_i)/d(S_ij) = P_ij), giving: +# dS_ij = P_ij * (dP_ij - D_i) + dLSE_i * P_ij +# = P_ij * (dP_ij - (D_i - dLSE_i)) +# So the main backward kernel is unchanged; we just replace D with D' = D - dLSE here. +import math +import operator +from functools import partial +from typing import Callable, Type, Optional + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, const_expr +from cutlass.cutlass_dsl import Arch, BaseDSL + +from quack import copy_utils, layout_utils + +from flash_attn.cute import utils +from flash_attn.cute.seqlen_info import SeqlenInfo +from quack.cute_dsl_utils import ParamsBase +from flash_attn.cute.tile_scheduler import ( + SingleTileScheduler, + SingleTileVarlenScheduler, + TileSchedulerArguments, +) +from flash_attn.cute.pack_gqa import pack_gqa_layout + + +class FlashAttentionBackwardPreprocess: + def __init__( + self, + dtype: Type[cutlass.Numeric], + head_dim: int, + head_dim_v: int, + tile_m: int = 128, + num_threads: int = 256, + use_padded_offsets: bool = True, + nheads_major: bool = False, + pack_gqa: bool = False, + qhead_per_kvhead: int = 1, + nheads_kv: int = 1, + ): + """ + All contiguous dimensions must be at least 16 bytes aligned which indicates the head dimension + should be a multiple of 8. + + :param head_dim: head dimension + :type head_dim: int + :param tile_m: m block size + :type tile_m: int + :param num_threads: number of threads + :type num_threads: int + """ + self.use_pdl = BaseDSL._get_dsl().get_arch_enum() >= Arch.sm_90a + self.dtype = dtype + self.tile_m = tile_m + # padding head_dim to a multiple of 32 as k_block_size + hdim_multiple_of = 32 + self.head_dim_padded = int(math.ceil(head_dim / hdim_multiple_of) * hdim_multiple_of) + self.head_dim_v_padded = int(math.ceil(head_dim_v / hdim_multiple_of) * hdim_multiple_of) + self.check_hdim_v_oob = head_dim_v != self.head_dim_v_padded + self.num_threads = num_threads + self.use_padded_offsets = use_padded_offsets + self.nheads_major = nheads_major + self.pack_gqa = pack_gqa + self.qhead_per_kvhead = qhead_per_kvhead + self.nheads_kv = nheads_kv + + @staticmethod + def can_implement(dtype, head_dim, tile_m, num_threads) -> bool: + """Check if the kernel can be implemented with the given parameters. + + :param dtype: data type + :type dtype: cutlass.Numeric + :param head_dim: head dimension + :type head_dim: int + :param tile_m: m block size + :type tile_m: int + :param num_threads: number of threads + :type num_threads: int + + :return: True if the kernel can be implemented, False otherwise + :rtype: bool + """ + if dtype not in [cutlass.Float16, cutlass.BFloat16]: + return False + if head_dim % 8 != 0: + return False + if num_threads % 32 != 0: + return False + if num_threads < tile_m: # For multiplying lse with log2 + return False + return True + + def _setup_attributes(self): + # /////////////////////////////////////////////////////////////////////////////// + # GMEM Tiled copy: + # /////////////////////////////////////////////////////////////////////////////// + # Thread layouts for copies + # We want kBlockKGmem to be a power of 2 so that when we do the summing, + # it's just between threads in the same warp + gmem_k_block_size = ( + 128 + if self.head_dim_v_padded % 128 == 0 + else ( + 64 + if self.head_dim_v_padded % 64 == 0 + else (32 if self.head_dim_v_padded % 32 == 0 else 16) + ) + ) + num_copy_elems = 128 // self.dtype.width + threads_per_row = gmem_k_block_size // num_copy_elems + self.gmem_tiled_copy_O = copy_utils.tiled_copy_2d( + self.dtype, threads_per_row, self.num_threads, num_copy_elems + ) + universal_copy_bits = 128 + num_copy_elems_dQaccum = universal_copy_bits // Float32.width + assert ( + self.tile_m * self.head_dim_padded // num_copy_elems_dQaccum + ) % self.num_threads == 0 + self.gmem_tiled_copy_dQaccum = copy_utils.tiled_copy_1d( + Float32, self.num_threads, num_copy_elems_dQaccum + ) + + @cute.jit + def __call__( + self, + mO: cute.Tensor, # (batch, seqlen, nheads, head_dim_v) or (total_q, nheads, head_dim_v) + mdO: cute.Tensor, # same shape as mO + mPdPsum: cute.Tensor, # (batch, nheads, seqlen_padded) or (nheads, total_q_padded) + mLSE: Optional[cute.Tensor], # (batch, nheads, seqlen) or (nheads, total_q) + mLSElog2: Optional[cute.Tensor], # same shape as mPdPsum + # (batch, nheads, seqlen_padded * head_dim_v) or (nheads, total_q_padded * head_dim_v) + mdQaccum: Optional[cute.Tensor], + mCuSeqlensQ: Optional[cute.Tensor], # (batch + 1,) + mSeqUsedQ: Optional[cute.Tensor], # (batch,) + mdLSE: Optional[cute.Tensor], # (batch, nheads, seqlen) or (nheads, total_q) + mRowMax: Optional[cute.Tensor], # (b, s, n, h) or (t, n, h) + mScaleP: Optional[cute.Tensor], # == mRowMax + softmax_scale: Float32, + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + # Get the data type and check if it is fp16 or bf16 + if const_expr(not (mO.element_type == mdO.element_type)): + raise TypeError("All tensors must have the same data type") + if const_expr(mO.element_type not in [cutlass.Float16, cutlass.BFloat16]): + raise TypeError("Only Float16 or BFloat16 is supported") + if const_expr(mPdPsum.element_type not in [Float32]): + raise TypeError("PdPsum tensor must be Float32") + if const_expr(mdQaccum is not None): + assert self.nheads_major is False + assert self.pack_gqa is False + assert self.use_padded_offsets is True + if const_expr(mdQaccum.element_type not in [Float32]): + raise TypeError("dQaccum tensor must be Float32") + if const_expr(mLSE is not None): + if const_expr(mLSE.element_type not in [Float32]): + raise TypeError("LSE tensor must be Float32") + if const_expr(mLSElog2 is not None): + if const_expr(mLSElog2.element_type not in [Float32]): + raise TypeError("LSElog2 tensor must be Float32") + if const_expr(mdLSE is not None): + if const_expr(mdLSE.element_type not in [Float32]): + raise TypeError("dLSE tensor must be Float32") + if const_expr(mScaleP is not None): + assert self.nheads_major is True + assert self.pack_gqa is True + assert mRowMax is not None + if const_expr(mScaleP.element_type not in [Float32]): + raise TypeError("ScaleP tensor must be Float32") + if const_expr(mRowMax.element_type not in [Float32]): + raise TypeError("RowMax tensor must be Float32") + + self._setup_attributes() + + # (b, s, h, d) -> (s, d, h, b) or + # (total, h, d) -> (total, d, h) + QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1] + mO, mdO = [ + cute.make_tensor(mX.iterator, cute.select(mX.layout, mode=QO_layout_transpose)) + for mX in (mO, mdO) + ] + + if const_expr(not self.nheads_major): + # (batch, nheads, seqlen) -> (seqlen, nheads, batch) or + # (nheads, total_q) -> (total_q, nheads) + transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0] + else: + # (batch, seqlen, nheads) -> (seqlen, nheads, batch) or + # (total_q, nheads) -> (total_q, nheads) + transpose = [1, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 1] + mPdPsum, mLSE, mLSElog2, mdLSE, mdQaccum = [ + layout_utils.select(mX, transpose) if mX is not None else None + for mX in (mPdPsum, mLSE, mLSElog2, mdLSE, mdQaccum) + ] + + # (b, s, n, h) => (s, n, h, b) or + # (total, n, h) == (total, n, h) + rowmax_layout_transpose = [1, 2, 3, 0] if const_expr(mCuSeqlensQ is None) else [0, 1, 2] + if const_expr(mRowMax is not None): + mRowMax = layout_utils.select(mRowMax, rowmax_layout_transpose) + if const_expr(mScaleP is not None): + mScaleP = layout_utils.select(mScaleP, rowmax_layout_transpose) + + # pack gqa + if const_expr(self.pack_gqa): + mO, mdO, mRowMax, mScaleP = [ + pack_gqa_layout(mX, self.qhead_per_kvhead, self.nheads_kv, head_idx=2) + if mX is not None + else None + for mX in (mO, mdO, mRowMax, mScaleP) + ] + mPdPsum, mLSE, mLSElog2, mdLSE = [ + pack_gqa_layout(mX, self.qhead_per_kvhead, self.nheads_kv, head_idx=1) + if mX is not None + else None + for mX in (mPdPsum, mLSE, mLSElog2, mdLSE) + ] + + # mO: (s, d, h, b) or (total, d, h) + if const_expr(mCuSeqlensQ is not None): + TileScheduler = SingleTileVarlenScheduler + num_head = mO.shape[2] + num_batch = mCuSeqlensQ.shape[0] - 1 + else: + TileScheduler = SingleTileScheduler + num_head = mO.shape[2] + num_batch = mO.shape[3] + + tile_sched_args = TileSchedulerArguments( + num_block=cute.ceil_div(mO.shape[0], self.tile_m), + num_head=num_head, + num_batch=num_batch, + num_splits=1, + seqlen_k=0, + headdim=0, + headdim_v=mO.shape[1], + total_q=cute.size(mO.shape[0]) + if const_expr(mCuSeqlensQ is not None) + else cute.size(mO.shape[0]) * cute.size(mO.shape[3]), + tile_shape_mn=(self.tile_m, 1), + mCuSeqlensQ=mCuSeqlensQ, + mSeqUsedQ=mSeqUsedQ, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + + tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args) + grid_dim = TileScheduler.get_grid_shape(tile_sched_params) + + LOG2_E = math.log2(math.e) + softmax_scale_log2 = softmax_scale * LOG2_E + + self.kernel( + mO, + mdO, + mPdPsum, + mLSE, + mLSElog2, + mdQaccum, + mCuSeqlensQ, + mSeqUsedQ, + mdLSE, + mRowMax, + mScaleP, + softmax_scale_log2, + self.gmem_tiled_copy_O, + self.gmem_tiled_copy_dQaccum, + tile_sched_params, + TileScheduler, + ).launch( + grid=grid_dim, + block=[self.num_threads, 1, 1], + stream=stream, + use_pdl=self.use_pdl, + ) + + @cute.kernel + def kernel( + self, + mO: cute.Tensor, + mdO: cute.Tensor, + mPdPsum: cute.Tensor, + mLSE: Optional[cute.Tensor], + mLSElog2: Optional[cute.Tensor], + mdQaccum: Optional[cute.Tensor], + mCuSeqlensQ: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mdLSE: Optional[cute.Tensor], + mRowMax: Optional[cute.Tensor], + mScaleP: Optional[cute.Tensor], + softmax_scale_log2: Float32, + gmem_tiled_copy_O: cute.TiledCopy, + gmem_tiled_copy_dQaccum: cute.TiledCopy, + tile_sched_params: ParamsBase, + TileScheduler: cutlass.Constexpr[Callable], + ): + # Thread index, block index + tidx, _, _ = cute.arch.thread_idx() + + tile_scheduler = TileScheduler.create(tile_sched_params) + work_tile = tile_scheduler.initial_work_tile_info() + m_block, head_idx, batch_idx, _ = work_tile.tile_idx + + # This kernel is launched with use_pdl=True, so the GPU may start executing it in + # "prologue" mode while the previous stream kernel is still running. We must wait + # before touching any upstream GMEM outputs (mO, mdO, mLSE); otherwise we risk + # reading a partially-written dout, which silently corrupts dpsum = sum(O * dO) and + # propagates to dQ/dK via dS = P * (dP - dpsum). + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_wait() + + if work_tile.is_valid_tile: + # /////////////////////////////////////////////////////////////////////////////// + # Get the appropriate tiles for this thread block. + # /////////////////////////////////////////////////////////////////////////////// + seqlen_static = mO.shape[0] if const_expr(not self.pack_gqa) else mO.shape[0][1] + seqlen = SeqlenInfo.create( + batch_idx, seqlen_static, mCuSeqlensQ, mSeqUsedQ, tile=self.tile_m + ) + # (seqlen, dv) + mO_cur, mdO_cur = [ + seqlen.offset_batch(mX, batch_idx, dim=3)[None, None, head_idx] for mX in (mO, mdO) + ] + mPdPsum_cur = seqlen.offset_batch( + mPdPsum, batch_idx, dim=2, padded=self.use_padded_offsets + )[None, head_idx] + headdim_v = mO_cur.shape[1] + seqlen_q = ( + seqlen.seqlen + if const_expr(not self.pack_gqa) + else seqlen.seqlen * self.qhead_per_kvhead + ) + seqlen_q_rounded = cute.round_up(seqlen_q, self.tile_m) + seqlen_limit = seqlen_q - m_block * self.tile_m + + lse = None + if const_expr(mLSE is not None): + mLSE_cur = seqlen.offset_batch(mLSE, batch_idx, dim=2)[None, head_idx] + gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (m_block,)) + lse = Float32.inf + if tidx < seqlen_limit: + lse = gLSE[tidx] + + blk_shape = (self.tile_m, self.head_dim_v_padded) + gO = cute.local_tile(mO_cur, blk_shape, (m_block, 0)) + gdO = cute.local_tile(mdO_cur, blk_shape, (m_block, 0)) + gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx) + # (CPY_Atom, CPY_M, CPY_K) + tOgO = gmem_thr_copy_O.partition_S(gO) + tOgdO = gmem_thr_copy_O.partition_S(gdO) + cO = cute.make_identity_tensor(blk_shape) + tOcO = gmem_thr_copy_O.partition_S(cO) + t0OcO = gmem_thr_copy_O.get_slice(0).partition_S(cO) + tOpO = None + if const_expr(self.check_hdim_v_oob): + tOpO = copy_utils.predicate_k(tOcO, limit=headdim_v) + # Each copy will use the same predicate + copy = partial(copy_utils.copy, pred=tOpO) + + tOrO = cute.make_rmem_tensor_like(tOgO) + tOrdO = cute.make_rmem_tensor_like(tOgdO) + if const_expr(self.check_hdim_v_oob): + tOrO.fill(0.0) + tOrdO.fill(0.0) + assert tOgO.shape == tOgdO.shape + for m in cutlass.range(cute.size(tOrO.shape[1]), unroll_full=True): + # Instead of using tOcO, we using t0OcO and subtract the offset from the limit. + # This is bc the entries of t0OcO are known at compile time. + if t0OcO[0, m, 0][0] < seqlen_limit - tOcO[0][0]: + copy(tOgO[None, m, None], tOrO[None, m, None]) + copy(tOgdO[None, m, None], tOrdO[None, m, None]) + # O and dO loads are done; signal that the next kernel can start. + # Correctness is ensured by griddepcontrol_wait() in bwd_sm90 before it reads our outputs. + if const_expr(self.use_pdl): + cute.arch.griddepcontrol_launch_dependents() + # Sum across the "k" dimension + pdpsum = (tOrO.load().to(Float32) * tOrdO.load().to(Float32)).reduce( + cute.ReductionOp.ADD, init_val=0.0, reduction_profile=(0, None, 1) + ) + threads_per_row = gmem_tiled_copy_O.layout_src_tv_tiled[0].shape[0] + assert cute.arch.WARP_SIZE % threads_per_row == 0 + pdpsum = utils.warp_reduce(pdpsum, operator.add, width=threads_per_row) + PdP_sum = cute.make_rmem_tensor(cute.size(tOrO, mode=[1]), Float32) + PdP_sum.store(pdpsum) + + # If dLSE is provided, compute D' = D - dLSE (see module docstring for derivation). + gdLSE = None + if const_expr(mdLSE is not None): + mdLSE_cur = seqlen.offset_batch(mdLSE, batch_idx, dim=2)[None, head_idx] + gdLSE = cute.local_tile(mdLSE_cur, (self.tile_m,), (m_block,)) + + # Write PdPsum from rmem -> gmem + gPdPsum = cute.local_tile(mPdPsum_cur, (self.tile_m,), (m_block,)) + # Only the thread corresponding to column 0 writes out the PdPsum to gmem + if tOcO[0, 0, 0][1] == 0: + for m in cutlass.range(cute.size(PdP_sum), unroll_full=True): + row = tOcO[0, m, 0][0] + PdPsum_val = 0.0 + if row < seqlen_limit: + PdPsum_val = PdP_sum[m] + if const_expr(mdLSE is not None): + PdPsum_val -= gdLSE[row] + gPdPsum[row] = PdPsum_val + + # Clear dQaccum + if const_expr(mdQaccum is not None): + mdQaccum_cur = seqlen.offset_batch( + mdQaccum, + batch_idx, + dim=2, + padded=self.use_padded_offsets, + multiple=self.head_dim_padded, + )[None, head_idx] + blkdQaccum_shape = (self.tile_m * self.head_dim_padded,) + gdQaccum = cute.local_tile(mdQaccum_cur, blkdQaccum_shape, (m_block,)) + gmem_thr_copy_dQaccum = gmem_tiled_copy_dQaccum.get_slice(tidx) + tdQgdQaccum = gmem_thr_copy_dQaccum.partition_S(gdQaccum) + zero = cute.make_rmem_tensor_like(tdQgdQaccum) + zero.fill(0.0) + cute.copy(gmem_tiled_copy_dQaccum, zero, tdQgdQaccum) + + LOG2_E = math.log2(math.e) + lse_log2 = lse * LOG2_E if lse != -Float32.inf else 0.0 + if const_expr(mLSElog2 is not None): + mLSElog2_cur = seqlen.offset_batch( + mLSElog2, batch_idx, dim=2, padded=self.use_padded_offsets + )[None, head_idx] + gLSElog2 = cute.local_tile(mLSElog2_cur, (self.tile_m,), (m_block,)) + LOG2_E = math.log2(math.e) + if tidx < seqlen_q_rounded - m_block * self.tile_m: + gLSElog2[tidx] = lse_log2 + + if const_expr(mRowMax is not None): + assert mLSE is not None + # (s, n) + mRowMax_cur, mScaleP_cur = [ + seqlen.offset_batch(mX, batch_idx, dim=3)[None, None, head_idx] + for mX in (mRowMax, mScaleP) + ] + # (tile_m, n) + gRowMax, gScaleP = [ + cute.local_tile(mX, (self.tile_m,), (m_block, None)) + for mX in (mRowMax_cur, mScaleP_cur) + ] + + assert self.tile_m <= self.num_threads + if const_expr(self.tile_m == self.num_threads) or tidx < self.tile_m: + for n in cutlass.range(gRowMax.shape[1], unroll=4): + row_max = gRowMax[tidx, n] + scale = 0.0 + if row_max != -Float32.inf and lse != -Float32.inf: + scale = softmax_scale_log2 * row_max - lse_log2 + scale = cute.math.exp2(scale, fastmath=True) + gScaleP[tidx, n] = scale diff --git a/fastvideo-kernel/fa4/flash_bwd_sm100.py b/fastvideo-kernel/fa4/flash_bwd_sm100.py new file mode 100644 index 000000000..f0d39f0c6 --- /dev/null +++ b/fastvideo-kernel/fa4/flash_bwd_sm100.py @@ -0,0 +1,4041 @@ +# Copyright (c) 2025, Ted Zadouri, Markus Hoehnerbach, Jay Shah, Tri Dao. +import math +from typing import Callable, Optional +from functools import partial + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute import FastDivmodDivisor +from cutlass import Float32, Int32, Int64, const_expr +from cutlass.utils import LayoutEnum +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.utils.blackwell_helpers as sm100_utils_basic +from cutlass.pipeline import PipelineAsync + +import quack.activation +from quack import layout_utils +from flash_attn.cute import utils +from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned +from flash_attn.cute import copy_utils +from flash_attn.cute import pipeline +from flash_attn.cute.blackwell_helpers import gemm_w_idx, gemm_ptx_w_idx # noqa +from flash_attn.cute.mask import AttentionMask +from flash_attn.cute.seqlen_info import SeqlenInfoQK +from flash_attn.cute.block_info import BlockInfo +from quack.cute_dsl_utils import ParamsBase +from flash_attn.cute.tile_scheduler import ( + TileSchedulerArguments, + SingleTileScheduler, + SingleTileLPTBwdScheduler, # noqa + SingleTileVarlenScheduler, +) + +from flash_attn.cute import barrier +from flash_attn.cute.named_barrier import NamedBarrierBwdSm100 +from flash_attn.cute.softmax import apply_score_mod_inner, apply_score_mod_bwd_inner +from flash_attn.cute.block_sparsity import BlockSparseTensors +from flash_attn.cute.utils import AuxData +from flash_attn.cute.block_sparse_utils import ( + get_total_q_block_count_bwd, + get_block_sparse_iteration_info_bwd, + get_m_block_from_iter_bwd, + produce_block_sparse_q_loads_bwd_sm100, +) + + +class FlashAttentionBackwardSm100: + arch = 100 + + def __init__( + self, + head_dim: int, + head_dim_v: Optional[int] = None, + is_causal: bool = False, + is_local: bool = False, + qhead_per_kvhead: cutlass.Constexpr[int] = 1, + tile_m: int = 128, + tile_n: int = 128, + is_persistent: bool = False, + deterministic: bool = False, + spt: Optional[bool] = None, + cluster_size: int = 1, + use_2cta_instrs: bool = False, + score_mod: cutlass.Constexpr | None = None, + score_mod_bwd: cutlass.Constexpr | None = None, + mask_mod: cutlass.Constexpr | None = None, + has_aux_tensors: cutlass.Constexpr = False, + q_subtile_factor: cutlass.Constexpr[int] = 1, + ): + # padding head_dim to a multiple of 16 as k_block_size + hdim_multiple_of = 16 + self.tile_hdim = int(math.ceil(head_dim / hdim_multiple_of) * hdim_multiple_of) + head_dim_v = head_dim_v if head_dim_v is not None else head_dim + self.same_hdim_kv = head_dim == head_dim_v + self.tile_hdimv = int(math.ceil(head_dim_v / hdim_multiple_of) * hdim_multiple_of) + self.check_hdim_oob = head_dim != self.tile_hdim + self.check_hdim_v_oob = head_dim_v != self.tile_hdimv + + self.tile_m = tile_m + self.tile_n = tile_n + + assert self.tile_hdim <= 128 or (self.tile_hdim == 192 and self.tile_hdimv == 128) + assert self.tile_hdimv <= 128 + + self.use_2cta_instrs = bool(use_2cta_instrs and cluster_size == 2) + self.cta_group_size = 2 if self.use_2cta_instrs else 1 + + assert self.tile_hdim != 192 or self.use_2cta_instrs, "Must use 2CTA for hdim 192" + + # CTA tiler + self.cta_tiler = (tile_n, tile_m, self.tile_hdim) + # S = K @ Q.T + self.mma_tiler_kq = (self.cta_group_size * tile_n, tile_m, self.tile_hdim) + # dP = V @ dO.T + self.mma_tiler_vdo = (self.cta_group_size * tile_n, tile_m, self.tile_hdimv) + # dV = P.T @ dO + self.mma_tiler_pdo = (self.cta_group_size * tile_n, self.tile_hdimv, tile_m) + # dK = dS.T @ Q + self.mma_tiler_dsq = (self.cta_group_size * tile_n, self.tile_hdim, tile_m) + # dQ = dS @ K + # 2-CTA: reduction dim is cluster-wide (tile_n * cta_group_size). + self.mma_tiler_dsk = (tile_m, self.tile_hdim, tile_n * self.cta_group_size) + + self.acc_dtype = Float32 + + assert cluster_size in (1, 2), "Only cluster_size=1 or 2 is supported" + self.cluster_shape_mn = (cluster_size, 1) + self.is_persistent = is_persistent + self.is_causal = is_causal + self.is_local = is_local + self.qhead_per_kvhead = qhead_per_kvhead + self.pack_gqa = False + self.deterministic = deterministic + self.spt_override = spt + + # Score mod and mask mod support + self.score_mod = score_mod + self.score_mod_bwd = score_mod_bwd + self.mask_mod = mask_mod + self.has_aux_tensors = has_aux_tensors + self.q_subtile_factor = q_subtile_factor + # For score_mod, use vec_size=1 (like forward) to handle per-element indices + if cutlass.const_expr(has_aux_tensors): + self.vec_size: cutlass.Constexpr = 1 + else: + self.vec_size: cutlass.Constexpr = 4 + self.qk_acc_dtype = Float32 + + # Speed optimizations, does not affect correctness + self.shuffle_LSE = False + self.shuffle_dPsum = False + # Generally slower to use store dS in smem for dK, and doesn't work for 2cta + self.use_smem_dS_for_mma_dK = False + + self.reduce_warp_ids = (0, 1, 2, 3) + self.compute_warp_ids = (4, 5, 6, 7, 8, 9, 10, 11) + self.mma_warp_id = 12 + self.load_warp_id = 13 + self.relay_warp_id = 14 + self.empty_warp_id = 15 + + # 16 warps -> 512 threads + self.threads_per_cta = cute.arch.WARP_SIZE * len( + ( + *self.reduce_warp_ids, + *self.compute_warp_ids, + self.mma_warp_id, + self.load_warp_id, + self.relay_warp_id, + self.empty_warp_id, + ) + ) + # NamedBarrier + self.compute_sync_barrier = cutlass.pipeline.NamedBarrier( + barrier_id=int(NamedBarrierBwdSm100.Compute), + num_threads=len(self.compute_warp_ids) * cute.arch.WARP_SIZE, + ) + # self.epilogue_sync_barrier = pipeline.NamedBarrier( + # barrier_id=2, + # num_threads=self.num_compute_warps * self.threads_per_warp, + # ) + self.reduce_sync_barrier = cutlass.pipeline.NamedBarrier( + barrier_id=int(NamedBarrierBwdSm100.dQaccReduce), + num_threads=len(self.reduce_warp_ids) * cute.arch.WARP_SIZE, + ) + # TMEM setup + self.tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") + # self.tmem_dK_offset = 0 + # self.tmem_dV_offset = self.tmem_dK_offset + self.tile_hdim + # self.tmem_dQ_offset = self.tmem_dV_offset + self.tile_hdimv + # self.tmem_dP_offset = self.tmem_dQ_offset # overlap with dQ + # self.tmem_S_offset = self.tmem_dQ_offset + max(self.tile_m, self.tile_hdim) + # self.tmem_P_offset = self.tmem_S_offset # overlap with S + # self.tmem_total = self.tmem_S_offset + self.tile_n + # assert self.tmem_total <= self.tmem_alloc_cols + + if self.use_2cta_instrs and self.tile_hdim == 192 and self.tile_hdimv == 128: + assert self.tile_m == 128 + assert self.tile_n == 128 + self.tmem_dV_offset = 0 + self.tmem_dK_offset = self.tmem_dV_offset + self.tile_hdimv + self.tmem_S_offset = self.tmem_dK_offset + self.tile_hdim + self.tmem_P_offset = self.tmem_S_offset # overlap with S + self.tmem_dP_offset = 512 - self.tile_m + self.tmem_dS_offset = self.tmem_dP_offset # overlaps with dP + self.tmem_dQ_offset = 512 - self.tile_hdim // 2 + else: + self.tmem_S_offset = 0 + self.tmem_P_offset = 0 # overlap with S + self.tmem_dV_offset = self.tmem_S_offset + self.tile_n + self.tmem_dP_offset = self.tmem_dV_offset + self.tile_hdimv + self.tmem_dQ_offset = ( + (self.tmem_S_offset + (self.tile_hdim // 2)) + if self.use_2cta_instrs + else self.tmem_dP_offset + ) + self.tmem_dK_offset = self.tmem_dP_offset + self.tile_m + self.tmem_dS_offset = self.tmem_dP_offset # overlap with dP + + if (not is_causal and not is_local) or deterministic: + self.num_regs_reduce = 136 if self.use_2cta_instrs else 152 + self.num_regs_compute = 136 + self.num_regs_load = 104 if self.use_2cta_instrs else 96 - 8 + self.num_regs_mma = 104 if self.use_2cta_instrs else self.num_regs_load + else: + self.num_regs_reduce = 136 if self.use_2cta_instrs else 136 + self.num_regs_compute = 136 if self.use_2cta_instrs else 144 + self.num_regs_load = 104 if self.use_2cta_instrs else 96 - 8 + self.num_regs_mma = 104 if self.use_2cta_instrs else self.num_regs_load + self.num_regs_empty = 24 + + if const_expr(self.tile_hdim == 192): + if not is_causal and not is_local: + self.num_regs_reduce = 128 + 8 + self.num_regs_compute = 128 + 8 + self.num_regs_load = 128 - 24 + self.num_regs_mma = self.num_regs_load + else: + self.num_regs_reduce = 128 + 8 + self.num_regs_compute = 128 + 8 + self.num_regs_load = 128 - 24 + self.num_regs_mma = self.num_regs_load + + assert ( + self.num_regs_reduce + + self.num_regs_compute * 2 + + max(self.num_regs_load, self.num_regs_mma) + <= 512 + ) + self.buffer_align_bytes = 1024 + + def _setup_attributes(self): + self.Q_stage = 1 if self.use_2cta_instrs else 2 + self.dO_stage = 1 + self.single_stage = 1 + # LSE_stage = Q_stage and dPsum_stage = dO_stage + self.sdKVaccum_stage = 2 + # number of tma reduce adds per dQacc mma + # todo: try 32/1 or 48/2 for 2cta d=192 dv=128 + if self.use_2cta_instrs and self.tile_hdim == 192: + self.dQ_reduce_ncol_t2r = 32 + self.dQ_reduce_ncol = 24 if not self.is_causal else 32 + self.sdQaccum_stage = 2 if not self.is_causal else 1 + else: + if self.use_2cta_instrs: + self.dQ_reduce_ncol = 16 if self.deterministic else 8 + self.sdQaccum_stage = 2 if self.deterministic else 4 + self.dQ_reduce_ncol_t2r = 32 + else: + self.dQ_reduce_ncol = 32 + self.sdQaccum_stage = 64 // self.dQ_reduce_ncol + self.dQ_reduce_ncol_t2r = self.dQ_reduce_ncol + assert (self.tile_hdim // self.cta_group_size) % self.dQ_reduce_ncol == 0 + self.dQaccum_reduce_stage = self.tile_hdim // self.dQ_reduce_ncol + self.dQaccum_reduce_stage_t2r = self.tile_hdim // self.dQ_reduce_ncol_t2r + self.cluster_reduce_dQ = False and cute.size(self.cluster_shape_mn) > 1 + # number of tma reduce adds for dKacc and dVacc epilogue (must divide hdim_per_wg) + self.dK_reduce_ncol = math.gcd(32, self.tile_hdim // 2) + # CTA group for MMA operations + self.cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + + def _get_tiled_mma(self): + # S.T = K @ Q.T + tiled_mma_S = sm100_utils_basic.make_trivial_tiled_mma( + self.q_dtype, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + self.acc_dtype, + self.cta_group, + self.mma_tiler_kq[:2], + ) + # dP.T = V @ dO.T + tiled_mma_dP = sm100_utils_basic.make_trivial_tiled_mma( + self.do_dtype, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + self.acc_dtype, + self.cta_group, + self.mma_tiler_vdo[:2], + ) + # dV += P.T @ dO --> (K, MN) major + tiled_mma_dV = sm100_utils_basic.make_trivial_tiled_mma( + self.do_dtype, + tcgen05.OperandMajorMode.K, # P_major_mode + tcgen05.OperandMajorMode.MN, # dO_major_mode + self.acc_dtype, + self.cta_group, + self.mma_tiler_pdo[:2], + a_source=tcgen05.OperandSource.TMEM, + ) + # dK += dS.T @ Q + if const_expr(self.use_smem_dS_for_mma_dK): + mma_dK_a_src = tcgen05.OperandSource.SMEM + else: + mma_dK_a_src = tcgen05.OperandSource.TMEM + tiled_mma_dK = sm100_utils_basic.make_trivial_tiled_mma( + self.do_dtype, + tcgen05.OperandMajorMode.K, # dS_major_mode + tcgen05.OperandMajorMode.MN, # Q_major_mode + self.acc_dtype, + self.cta_group, + self.mma_tiler_dsq[:2], + a_source=mma_dK_a_src, + ) + # dQ = dS @ K + tiled_mma_dQ = sm100_utils_basic.make_trivial_tiled_mma( + self.k_dtype, + tcgen05.OperandMajorMode.MN, # dS_major_mode + tcgen05.OperandMajorMode.MN, # Kt_major_mode + self.acc_dtype, + self.cta_group, + self.mma_tiler_dsk[:2], + ) + return tiled_mma_S, tiled_mma_dP, tiled_mma_dK, tiled_mma_dV, tiled_mma_dQ + + def _setup_smem_layout(self): + # S.T = K @ Q.T + sK_layout = sm100_utils_basic.make_smem_layout_a( + self.tiled_mma_S, + self.mma_tiler_kq, + self.k_dtype, + 1, + ) + self.sK_layout = cute.slice_(sK_layout, (None, None, None, 0)) + self.sQ_layout = sm100_utils_basic.make_smem_layout_b( + self.tiled_mma_S, + self.mma_tiler_kq, + self.q_dtype, + self.Q_stage, + ) + # dP.T = V @ dO.T + sV_layout = sm100_utils_basic.make_smem_layout_a( + self.tiled_mma_dP, + self.mma_tiler_vdo, + self.v_dtype, + 1, + ) + self.sV_layout = cute.slice_(sV_layout, (None, None, None, 0)) + self.sdOt_layout = sm100_utils_basic.make_smem_layout_b( + self.tiled_mma_dP, + self.mma_tiler_vdo, + self.do_dtype, + self.dO_stage, + ) + # dV += P.T @ dO + tP_layout = sm100_utils_basic.make_smem_layout_a( + self.tiled_mma_dV, + self.mma_tiler_pdo, + self.do_dtype, + 1, + ) + self.tP_layout = cute.slice_(tP_layout, (None, None, None, 0)) + self.sdO_layout = sm100_utils_basic.make_smem_layout_b( + self.tiled_mma_dV, + self.mma_tiler_pdo, + self.do_dtype, + self.dO_stage, + ) + # dK += dS.T @ Q + sdSt_layout = sm100_utils_basic.make_smem_layout_a( + self.tiled_mma_dK, + self.mma_tiler_dsq, + self.ds_dtype, + 1, + ) + self.sdSt_layout = cute.slice_(sdSt_layout, (None, None, None, 0)) + tdS_layout = sm100_utils_basic.make_smem_layout_a( + self.tiled_mma_dK, + self.mma_tiler_dsq, + self.ds_dtype, + 1, + ) + self.tdS_layout = cute.slice_(tdS_layout, (None, None, None, 0)) + self.sQt_layout = sm100_utils_basic.make_smem_layout_b( + self.tiled_mma_dK, + self.mma_tiler_dsq, + self.q_dtype, + self.Q_stage, + ) + # dQ = dS @ K + sdS_layout = sm100_utils_basic.make_smem_layout_a( + self.tiled_mma_dQ, + self.mma_tiler_dsk, + self.ds_dtype, + 1, + ) + self.sdS_layout = cute.slice_(sdS_layout, (None, None, None, 0)) + sKt_layout = sm100_utils_basic.make_smem_layout_b( + self.tiled_mma_dQ, + self.mma_tiler_dsk, + self.k_dtype, + 1, + ) + self.sKt_layout = cute.slice_(sKt_layout, (None, None, None, 0)) + self.sdS_xchg_layout = cute.make_layout(shape=(self.tile_n, self.tile_m // 2)) + + self.sdQaccum_layout = cute.make_layout( + (self.tile_m * self.dQ_reduce_ncol, self.sdQaccum_stage) + ) + self.sLSE_layout = cute.make_layout( + shape=(self.tile_m, self.Q_stage), stride=(1, cute.round_up(self.tile_m, 64)) + ) + self.sdPsum_layout = cute.make_layout( + shape=(self.tile_m, self.dO_stage), + stride=(1, cute.round_up(self.tile_m, 64)), + ) + self.sdK_epi_tile = ( + self.tile_n, + math.gcd(128 // (self.dk_dtype.width // 8), self.tile_hdim // 2), # 64 or 32 + ) # subtiles mma_tiler_dsq[:2] = mma_tiler_pdo[:2] + self.sdV_epi_tile = ( + self.tile_n, + math.gcd(128 // (self.dk_dtype.width // 8), self.tile_hdimv // 2), # 64 or 32 + ) # subtiles mma_tiler_dsq[:2] = mma_tiler_pdo[:2] + # headdim_64 gets 1 stage + self.num_epi_stages = max(1, (self.tile_hdim // 2) // self.sdK_epi_tile[1]) + self.num_epi_stages_v = max(1, (self.tile_hdimv // 2) // self.sdV_epi_tile[1]) + self.sdK_flat_epi_tile = self.tile_n * (self.tile_hdim // 2) // self.num_epi_stages + self.sdV_flat_epi_tile = self.tile_n * (self.tile_hdimv // 2) // self.num_epi_stages_v + if const_expr(not self.dKV_postprocess): + self.sdK_layout = sm100_utils_basic.make_smem_layout_epi( + self.dk_dtype, + LayoutEnum.ROW_MAJOR, + self.sdK_epi_tile, + 2, # num compute wgs + ) + self.sdV_layout = sm100_utils_basic.make_smem_layout_epi( + self.dv_dtype, + LayoutEnum.ROW_MAJOR, + self.sdV_epi_tile, + 2, # num compute wgs + ) + else: + self.sdK_layout = cute.make_layout((self.tile_n * self.dK_reduce_ncol, 2)) + # self.dK_reduce_ncol same for dV + self.sdV_layout = cute.make_layout((self.tile_n * self.dK_reduce_ncol, 2)) + + @cute.jit + def __call__( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mdO: cute.Tensor, + mLSE: cute.Tensor, + mdPsum: cute.Tensor, + mdQaccum: cute.Tensor, + mdK: cute.Tensor, + mdV: cute.Tensor, + softmax_scale: Float32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + window_size_left: Int32 | int | None = None, + window_size_right: Int32 | int | None = None, + mdQ_semaphore: Optional[cute.Tensor] = None, + mdK_semaphore: Optional[cute.Tensor] = None, + mdV_semaphore: Optional[cute.Tensor] = None, + aux_data: AuxData = AuxData(), + # Block-sparse tensors (Q direction - for iterating m_blocks per n_block): + blocksparse_tensors: Optional[BlockSparseTensors] = None, + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + self.q_dtype = mQ.element_type + self.k_dtype = mK.element_type + self.v_dtype = mV.element_type + self.do_dtype = mdO.element_type + self.lse_dtype = mLSE.element_type + self.dpsum_dtype = mdPsum.element_type + self.dqaccum_dtype = mdQaccum.element_type + self.dk_dtype = mdK.element_type + self.dv_dtype = mdV.element_type + self.ds_dtype = self.q_dtype + + self.is_varlen_k = mCuSeqlensK is not None or mSeqUsedK is not None + self.is_varlen_q = mCuSeqlensQ is not None or mSeqUsedQ is not None + self.use_tma_store = not (self.qhead_per_kvhead == 1 and mCuSeqlensK is not None) + # self.use_tma_store = not self.qhead_per_kvhead == 1 + self.dKV_postprocess = self.qhead_per_kvhead > 1 + + if const_expr(self.dKV_postprocess): + assert self.dk_dtype.width == 32, "Must accumulate dK in float precision for GQA" + assert self.dv_dtype.width == 32, "Must accumulate dV in float precision for GQA" + + mdQaccum, mdK, mdV = [assume_tensor_aligned(t) for t in (mdQaccum, mdK, mdV)] + + # (b, s, n, h) --> (s, h, n, b) or (t, n, h) -> (t, h, n) + QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1] + mQ, mdO = [layout_utils.select(t, mode=QO_layout_transpose) for t in (mQ, mdO)] + + KV_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensK is None) else [0, 2, 1] + mK, mV = [layout_utils.select(t, mode=KV_layout_transpose) for t in (mK, mV)] + + # (b, n, s) --> (s, n, b) or (n, t) --> (t, n) + LSE_dPsum_dQaccum_transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0] + mLSE, mdPsum, mdQaccum = [ + layout_utils.select(t, mode=LSE_dPsum_dQaccum_transpose) + for t in (mLSE, mdPsum, mdQaccum) + ] + + if const_expr(not self.dKV_postprocess): + layout_dKV_transpose = KV_layout_transpose + else: + layout_dKV_transpose = [2, 1, 0] if const_expr(mCuSeqlensK is None) else [1, 0] + mdK, mdV = [layout_utils.select(t, mode=layout_dKV_transpose) for t in (mdK, mdV)] + # (s, h, n, b) --> (h, s, n, b) or (t, h, n) -> (h, t, b) + dO_transpose = [1, 0, 2, 3] if const_expr(mCuSeqlensQ is None) else [1, 0, 2] + mdO = layout_utils.select(mdO, mode=dO_transpose) + + # Transposes for 2-CTA K/Q paths (Q follows Q seqlens, K follows K seqlens) + transpose_sh_q = dO_transpose + transpose_sh_k = [1, 0, 2, 3] if const_expr(mCuSeqlensK is None) else [1, 0, 2] + + # (b, n, block, stage) -> (block, stage, n, b) + semaphore_transpose = [2, 3, 1, 0] + if const_expr(self.deterministic): + assert mdQ_semaphore is not None + mdQ_semaphore = layout_utils.select(mdQ_semaphore, mode=semaphore_transpose) + + if const_expr(self.deterministic and self.qhead_per_kvhead > 1): + assert mdK_semaphore is not None + assert mdV_semaphore is not None + mdK_semaphore, mdV_semaphore = [ + layout_utils.select(t, mode=semaphore_transpose) + for t in (mdK_semaphore, mdV_semaphore) + ] + else: + mdK_semaphore = None + mdV_semaphore = None + + self._setup_attributes() + ( + self.tiled_mma_S, + self.tiled_mma_dP, + self.tiled_mma_dK, + self.tiled_mma_dV, + self.tiled_mma_dQ, + ) = self._get_tiled_mma() + self._setup_smem_layout() + + self.cluster_shape_mnk = (*self.cluster_shape_mn, 1) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (self.tiled_mma_S.thr_id.shape,), + ) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_q_do_mcast = self.num_mcast_ctas_b > 1 + + if const_expr(not self.dKV_postprocess): + self.mdK_layout_enum = LayoutEnum.from_tensor(mdK) + self.mdV_layout_enum = LayoutEnum.from_tensor(mdV) + dK_major_mode = self.mdK_layout_enum.mma_major_mode() + dV_major_mode = self.mdV_layout_enum.mma_major_mode() + if const_expr(dK_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of mdK is wrong") + if const_expr(dV_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of mdV is wrong") + + if const_expr(self.use_tma_store and not self.dKV_postprocess): + tma_copy_op_dKV = cpasync.CopyBulkTensorTileS2GOp() + tma_atom_dK, mdK_tma_tensor = cpasync.make_tiled_tma_atom( + tma_copy_op_dKV, + mdK, + cute.select(self.sdK_layout, mode=[0, 1]), + self.sdK_epi_tile, + 1, # no mcast + ) + tma_atom_dV, mdV_tma_tensor = cpasync.make_tiled_tma_atom( + tma_copy_op_dKV, + mdV, + cute.select(self.sdV_layout, mode=[0, 1]), + self.sdV_epi_tile, + 1, # no mcast + ) + else: + mdV_tma_tensor = mdV + mdK_tma_tensor = mdK + tma_atom_dV = None + tma_atom_dK = None + + if const_expr(not self.dKV_postprocess): + thr_layout_r2s_dKV = cute.make_ordered_layout((128, 1), order=(1, 0)) # 128 threads + val_layout_r2s_dKV = cute.make_ordered_layout( + (1, 128 // self.dk_dtype.width), order=(1, 0) + ) # 4 or 8 vals for 16 byte store + copy_atom_r2s_dKV = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.dk_dtype, + num_bits_per_copy=128, + ) + tiled_copy_r2s_dKV = cute.make_tiled_copy_tv( + copy_atom_r2s_dKV, thr_layout_r2s_dKV, val_layout_r2s_dKV + ) + else: + tiled_copy_r2s_dKV = copy_utils.tiled_copy_1d( + Float32, 128, num_copy_elems=128 // Float32.width + ) + + tma_load_op = cpasync.CopyBulkTensorTileG2SOp(self.cta_group) + # S.T = K @ Q.T + tma_atom_K, tma_tensor_K = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + mK, + cute.select(self.sK_layout, mode=[0, 1, 2]), + self.mma_tiler_kq, + self.tiled_mma_S, + self.cluster_layout_vmnk.shape, + ) + Q_tma_op = sm100_utils_basic.cluster_shape_to_tma_atom_B( + self.cluster_shape_mnk, self.tiled_mma_S.thr_id + ) + tma_atom_Q, tma_tensor_Q = cute.nvgpu.make_tiled_tma_atom_B( + Q_tma_op, + mQ, + cute.select(self.sQ_layout, mode=[0, 1, 2]), + self.mma_tiler_kq, + self.tiled_mma_S, + self.cluster_layout_vmnk.shape, + ) + # dP.T = V @ dO.T + tma_atom_V, tma_tensor_V = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + mV, + cute.select(self.sV_layout, mode=[0, 1, 2]), + self.mma_tiler_vdo, + self.tiled_mma_dP, + self.cluster_layout_vmnk.shape, + ) + # dV = P.T @ dO + dO_tma_op = sm100_utils_basic.cluster_shape_to_tma_atom_B( + self.cluster_shape_mnk, self.tiled_mma_dV.thr_id + ) + tma_atom_dO, tma_tensor_dO = cute.nvgpu.make_tiled_tma_atom_B( + dO_tma_op, + mdO, + cute.select(self.sdO_layout, mode=[0, 1, 2]), + self.mma_tiler_pdo, + self.tiled_mma_dV, + self.cluster_layout_vmnk.shape, + ) + # ------------------------------------------------------------ + # 2-CTA + # ------------------------------------------------------------ + tma_atom_dOt = tma_tensor_dOt = None + if const_expr(self.use_2cta_instrs): + tma_atom_dOt, tma_tensor_dOt = cute.nvgpu.make_tiled_tma_atom_B( + dO_tma_op, + layout_utils.select(mdO, mode=transpose_sh_q), + cute.select(self.sdOt_layout, mode=[0, 1, 2]), + self.mma_tiler_vdo, + self.tiled_mma_dP, + self.cluster_layout_vmnk.shape, + ) + tma_atom_Qt = tma_tensor_Qt = None + if const_expr(self.use_2cta_instrs): + tma_atom_Qt, tma_tensor_Qt = cute.nvgpu.make_tiled_tma_atom_B( + Q_tma_op, + layout_utils.select(mQ, mode=transpose_sh_q), + cute.select(self.sQt_layout, mode=[0, 1, 2]), + self.mma_tiler_dsq, + self.tiled_mma_dK, + self.cluster_layout_vmnk.shape, + ) + tma_atom_Kt = tma_tensor_Kt = None + if const_expr(self.use_2cta_instrs): + Kt_tma_op = sm100_utils_basic.cluster_shape_to_tma_atom_B( + self.cluster_shape_mnk, self.tiled_mma_dQ.thr_id + ) + tma_atom_Kt, tma_tensor_Kt = cute.nvgpu.make_tiled_tma_atom_B( + Kt_tma_op, + layout_utils.select(mK, mode=transpose_sh_k), + cute.select(self.sKt_layout, mode=[0, 1, 2]), + self.mma_tiler_dsk, + self.tiled_mma_dQ, + self.cluster_layout_vmnk.shape, + ) + + self.tma_copy_bytes = { + name: self.cta_group_size + * cute.size_in_bytes(mX.element_type, cute.select(layout, mode=[0, 1, 2])) + for name, mX, layout in [ + ("Q", mQ, self.sQ_layout), + ("K", mK, self.sK_layout), + ("V", mV, self.sV_layout), + ("dO", mdO, self.sdO_layout), + ] + } + self.tma_copy_bytes["LSE"] = self.tile_m * Float32.width // 8 + self.tma_copy_bytes["dPsum"] = self.tile_m * Float32.width // 8 + self.tma_copy_bytes["dQ"] = self.tile_m * self.dQ_reduce_ncol * Float32.width // 8 + self.tma_copy_bytes["dKacc"] = self.tile_n * self.dK_reduce_ncol * Float32.width // 8 + self.tma_copy_bytes["dS"] = cute.size_in_bytes(self.ds_dtype, self.sdS_layout) + self.tma_copy_bytes["sdS_xchg"] = self.tma_copy_bytes["dS"] // 2 # Half of dS for exchange + + # TileScheduler = SingleTileScheduler + if const_expr(self.is_varlen_k): + TileScheduler = SingleTileVarlenScheduler + elif const_expr(self.deterministic): + TileScheduler = SingleTileLPTBwdScheduler + else: + TileScheduler = SingleTileScheduler + if const_expr(self.spt_override is None): + self.spt = (self.is_causal or self.is_local) and self.deterministic + else: + assert self.spt_override is not None + self.spt = self.spt_override and self.deterministic + tile_sched_args = TileSchedulerArguments( + cute.ceil_div(cute.size(mK.shape[0]), self.cta_tiler[0]), # num_blocks + cute.size(mQ.shape[2]), # num_heads = num_query_heads + cute.size(mK.shape[3]) + if const_expr(mCuSeqlensK is None) + else cute.size(mCuSeqlensK.shape[0] - 1), # num_batches + 1, # num_splits + cute.size(mQ.shape[0]), # pass seqlen_q or total_q for seqlen_k + mQ.shape[1], # headdim + mV.shape[1], # headdim_v + total_q=cute.size(mK.shape[0]) # pass total_k for total_q + if const_expr(mCuSeqlensK is not None) + else cute.size(mK.shape[0]) * cute.size(mK.shape[3]), + tile_shape_mn=self.cta_tiler[:2], # (tile_n, tile_m) + cluster_shape_mn=self.cluster_shape_mnk[:2], + mCuSeqlensQ=mCuSeqlensK, + mSeqUsedQ=mSeqUsedK, + qhead_per_kvhead_packgqa=1, # pack_gqa disabled for bwd + element_size=self.k_dtype.width // 8, + is_persistent=self.is_persistent, # persistent mode not tested + lpt=self.spt, + head_swizzle=self.deterministic, + ) + + tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args) + self.tile_scheduler_cls = TileScheduler + grid_dim = TileScheduler.get_grid_shape(tile_sched_params) + + # Compute allocation sizes for shared buffers that are reused + # sQ is reused for sdK, sdO is reused for sdV + sQ_alloc_bytes = max( + cute.size_in_bytes(self.q_dtype, self.sQ_layout), + cute.size_in_bytes(self.dk_dtype, self.sdK_layout), + ) + sdO_alloc_bytes = max( + cute.size_in_bytes(self.dv_dtype, self.sdV_layout), + cute.size_in_bytes(self.do_dtype, self.sdO_layout), + ) + + sdK_bytes = cute.size_in_bytes(self.dk_dtype, self.sdK_layout) + sdV_bytes = cute.size_in_bytes(self.dv_dtype, self.sdV_layout) + assert sdV_bytes <= sdO_alloc_bytes, "sdV doesn't fit in sdO storage allocation" + assert sdK_bytes <= sQ_alloc_bytes, "sdK doesn't fit in sQ storage allocation" + # 2-CTA: sdV reuses sV, sdK reuses sK + sV_bytes = cute.size_in_bytes(self.v_dtype, self.sV_layout) + sK_bytes = cute.size_in_bytes(self.k_dtype, self.sK_layout) + if const_expr(self.use_2cta_instrs): + assert sdV_bytes <= sV_bytes, "sdV doesn't fit in sV storage allocation (2-CTA)" + assert sdK_bytes <= sK_bytes, "sdK doesn't fit in sK storage allocation (2-CTA)" + + if const_expr(self.use_2cta_instrs): + sQt_size = cute.cosize(self.sQt_layout) if const_expr(self.tile_hdim <= 128) else 0 + sdOt_size = cute.cosize(self.sdOt_layout) if const_expr(self.tile_hdim <= 128) else 0 + sdS_xchg_size = ( + cute.cosize(self.sdS_xchg_layout) if const_expr(self.tile_hdim <= 128) else 0 + ) + + @cute.struct + class SharedStorage: + Q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.Q_stage] + dO_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.dO_stage] + LSE_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.Q_stage] + dPsum_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.dO_stage] + S_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.single_stage] + dP_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.single_stage] + dS_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.single_stage] + dKV_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.sdKVaccum_stage] + dQ_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2] + dQ_cluster_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.dQaccum_reduce_stage // 2 + ] + dQ_cluster_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.dQaccum_reduce_stage // 2 + ] + tmem_holding_buf: Int32 + tmem_dealloc_mbar_ptr: cutlass.Int64 + + # 2-CTA + Qt_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.Q_stage] + Kt_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.single_stage] + dS_cluster_empty_mbar_ptr: cutlass.Int64 + dS_cluster_full_mbar_ptr: cutlass.Int64 + dS_cluster_leader_mbar_ptr: cutlass.Int64 + dQaccum_empty_mbar_ptr: cutlass.Int64 + + sQ: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, cute.cosize(self.sQ_layout)], + self.buffer_align_bytes, + ] + sK: cute.struct.Align[ + cute.struct.MemRange[self.k_dtype, cute.cosize(self.sK_layout)], + self.buffer_align_bytes, + ] + sV: cute.struct.Align[ + cute.struct.MemRange[self.v_dtype, cute.cosize(self.sV_layout)], + self.buffer_align_bytes, + ] + sdO: cute.struct.Align[ + cute.struct.MemRange[self.do_dtype, cute.cosize(self.sdO_layout)], + self.buffer_align_bytes, + ] + sQt: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, sQt_size], + self.buffer_align_bytes, + ] + sdOt: cute.struct.Align[ + cute.struct.MemRange[self.do_dtype, sdOt_size], + self.buffer_align_bytes, + ] + sdS_xchg: cute.struct.Align[ + cute.struct.MemRange[self.ds_dtype, sdS_xchg_size], + self.buffer_align_bytes, + ] + sKt: cute.struct.Align[ + cute.struct.MemRange[self.k_dtype, cute.cosize(self.sKt_layout)], + self.buffer_align_bytes, + ] + sdS: cute.struct.Align[ + cute.struct.MemRange[self.ds_dtype, cute.cosize(self.sdSt_layout)], + self.buffer_align_bytes, + ] + sLSE: cute.struct.Align[ + cute.struct.MemRange[self.lse_dtype, cute.cosize(self.sLSE_layout)], + 128, + ] + sdPsum: cute.struct.Align[ + cute.struct.MemRange[self.dpsum_dtype, cute.cosize(self.sdPsum_layout)], + 128, + ] + sdQaccum: cute.struct.Align[ + cute.struct.MemRange[self.dqaccum_dtype, cute.cosize(self.sdQaccum_layout)], + self.buffer_align_bytes if sdS_xchg_size == 0 else 128, + ] + + else: + + @cute.struct + class SharedStorage: + Q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.Q_stage] + dO_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.dO_stage] + LSE_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.Q_stage] + dPsum_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.dO_stage] + S_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.single_stage] + dP_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.single_stage] + dS_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.single_stage] + dKV_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * self.sdKVaccum_stage] + dQ_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2] + dQ_cluster_full_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.dQaccum_reduce_stage // 2 + ] + dQ_cluster_empty_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.dQaccum_reduce_stage // 2 + ] + tmem_holding_buf: Int32 + tmem_dealloc_mbar_ptr: Int64 + + sQ: cute.struct.Align[ + cute.struct.MemRange[cute.Uint8, sQ_alloc_bytes], + self.buffer_align_bytes, + ] + sK: cute.struct.Align[ + cute.struct.MemRange[self.k_dtype, cute.cosize(self.sK_layout)], + self.buffer_align_bytes, + ] + sV: cute.struct.Align[ + cute.struct.MemRange[self.v_dtype, cute.cosize(self.sV_layout)], + self.buffer_align_bytes, + ] + sdO: cute.struct.Align[ + cute.struct.MemRange[cute.Uint8, sdO_alloc_bytes], + self.buffer_align_bytes, + ] + sdS: cute.struct.Align[ + cute.struct.MemRange[self.ds_dtype, cute.cosize(self.sdSt_layout)], + 128, + ] + sLSE: cute.struct.Align[ + cute.struct.MemRange[self.lse_dtype, cute.cosize(self.sLSE_layout)], + 128, + ] + sdPsum: cute.struct.Align[ + cute.struct.MemRange[self.dpsum_dtype, cute.cosize(self.sdPsum_layout)], + 128, + ] + sdQaccum: cute.struct.Align[ + cute.struct.MemRange[self.dqaccum_dtype, cute.cosize(self.sdQaccum_layout)], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + LOG2_E = math.log2(math.e) + if const_expr(self.score_mod is None): + # Without score_mod: bake scale into log2 + softmax_scale_log2 = softmax_scale * LOG2_E + else: + # With score_mod: score_mod applied to S * softmax_scale, then use LOG2_E only + softmax_scale_log2 = LOG2_E + + if const_expr(window_size_left is not None): + window_size_left = Int32(window_size_left) + if const_expr(window_size_right is not None): + window_size_right = Int32(window_size_right) + + fastdiv_mods = None + if const_expr(aux_data.tensors is not None): + seqlen_q = cute.size(mQ.shape[0]) // ( + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1 + ) + seqlen_k = cute.size(mK.shape[0]) + seqlen_q_divmod = FastDivmodDivisor(seqlen_q) + seqlen_k_divmod = FastDivmodDivisor(seqlen_k) + fastdiv_mods = (seqlen_q_divmod, seqlen_k_divmod) + self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None) + + if const_expr(self.use_2cta_instrs): + assert blocksparse_tensors is None, ( + "2-CTA mode does not support block sparsity. " + "Please create kernel with use_2cta_instrs=False for block sparse attention." + ) + # 2-CTA: 231424 and 1-CTA: 232448 + # print("SMEM: ", self.shared_storage.size_in_bytes()) + if const_expr(self.use_block_sparsity or aux_data.tensors is not None): + assert all(x is None for x in (mCuSeqlensQ, mCuSeqlensK, mSeqUsedQ, mSeqUsedK)), ( + "Variable sequence length is not supported yet for blocksparse or aux tensors in bwd" + ) + + self.kernel( + tma_tensor_Q, + tma_tensor_Qt, + tma_tensor_K, + tma_tensor_Kt, + tma_tensor_V, + mLSE, + mdPsum, + tma_tensor_dO, + tma_tensor_dOt, + mdV, + mdK, + mdQaccum, + mdV_tma_tensor, + mdK_tma_tensor, + mdQ_semaphore, + mdK_semaphore, + mdV_semaphore, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + tma_atom_Q, + tma_atom_Qt, + tma_atom_K, + tma_atom_Kt, + tma_atom_V, + tma_atom_dO, + tma_atom_dOt, + tma_atom_dV, + tma_atom_dK, + self.sQ_layout, + self.sQt_layout, + self.sK_layout, + self.sKt_layout, + self.sV_layout, + self.sLSE_layout, + self.sdPsum_layout, + self.sdO_layout, + self.sdOt_layout, + self.sdSt_layout, + self.sdS_layout, + self.sdS_xchg_layout, + self.sdQaccum_layout, + self.sdK_layout, + self.sdV_layout, + self.tP_layout, + self.tdS_layout, + self.tiled_mma_S, + self.tiled_mma_dP, + self.tiled_mma_dV, + self.tiled_mma_dK, + self.tiled_mma_dQ, + tiled_copy_r2s_dKV, + softmax_scale, + softmax_scale_log2, + window_size_left, + window_size_right, + tile_sched_params, + aux_data, + fastdiv_mods, + blocksparse_tensors, + ).launch( + grid=grid_dim, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk if cute.size(self.cluster_shape_mnk) > 1 else None, + smem=self.shared_storage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + + def _generate_attention_mask_cls(self, window_size_left, window_size_right): + return partial( + AttentionMask, + self.tile_m, + self.tile_n * self.cta_group_size, + swap_AB=True, + window_size_left=window_size_left, + window_size_right=window_size_right, + ) + + @cute.kernel + def kernel( + self, + mQ: cute.Tensor, + mQt: Optional[cute.Tensor], + mK: cute.Tensor, + mKt: Optional[cute.Tensor], + mV: cute.Tensor, + mLSE: cute.Tensor, + mdPsum: cute.Tensor, + mdO: cute.Tensor, + mdOt: Optional[cute.Tensor], + mdV: cute.Tensor, + mdK: cute.Tensor, + mdQaccum: cute.Tensor, + mdV_tma_tensor: Optional[cute.Tensor], + mdK_tma_tensor: Optional[cute.Tensor], + mdQ_semaphore: Optional[cute.Tensor], + mdK_semaphore: Optional[cute.Tensor], + mdV_semaphore: Optional[cute.Tensor], + mCuSeqlensQ: Optional[cute.Tensor], + mCuSeqlensK: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + tma_atom_Q: cute.CopyAtom, + tma_atom_Qt: Optional[cute.CopyAtom], + tma_atom_K: cute.CopyAtom, + tma_atom_Kt: Optional[cute.CopyAtom], + tma_atom_V: cute.CopyAtom, + tma_atom_dO: cute.CopyAtom, + tma_atom_dOt: Optional[cute.CopyAtom], + tma_atom_dV: Optional[cute.CopyAtom], + tma_atom_dK: Optional[cute.CopyAtom], + sQ_layout: cute.ComposedLayout, + sQt_layout: cute.ComposedLayout, + sK_layout: cute.ComposedLayout, + sKt_layout: cute.ComposedLayout, + sV_layout: cute.ComposedLayout, + sLSE_layout: cute.Layout, + sdPsum_layout: cute.Layout, + sdO_layout: cute.ComposedLayout, + sdOt_layout: cute.ComposedLayout, + sdSt_layout: cute.ComposedLayout, + sdS_layout: cute.ComposedLayout, + sdS_xchg_layout: cute.Layout, + sdQaccum_layout: cute.Layout, + sdK_layout: cute.ComposedLayout | cute.Layout, + sdV_layout: cute.ComposedLayout | cute.Layout, + tP_layout: cute.ComposedLayout, + tdS_layout: cute.ComposedLayout, + tiled_mma_S: cute.TiledMma, + tiled_mma_dP: cute.TiledMma, + tiled_mma_dV: cute.TiledMma, + tiled_mma_dK: cute.TiledMma, + tiled_mma_dQ: cute.TiledMma, + tiled_copy_r2s_dKV: cute.TiledCopy, + softmax_scale: cutlass.Float32, + softmax_scale_log2: cutlass.Float32, + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + tile_sched_params: ParamsBase, + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + blocksparse_tensors: Optional[BlockSparseTensors] = None, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % self.cta_group_size + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + + # Prefetch tma descriptor + if warp_idx == self.load_warp_id: + with cute.arch.elect_one(): + cpasync.prefetch_descriptor(tma_atom_Q) + if const_expr(tma_atom_Qt is not None): + cpasync.prefetch_descriptor(tma_atom_Qt) + cpasync.prefetch_descriptor(tma_atom_K) + if const_expr(tma_atom_Kt is not None): + cpasync.prefetch_descriptor(tma_atom_Kt) + cpasync.prefetch_descriptor(tma_atom_V) + if const_expr(tma_atom_dOt is not None): + cpasync.prefetch_descriptor(tma_atom_dOt) + cpasync.prefetch_descriptor(tma_atom_dO) + if const_expr(tma_atom_dV is not None): + cpasync.prefetch_descriptor(tma_atom_dV) + if const_expr(tma_atom_dK is not None): + cpasync.prefetch_descriptor(tma_atom_dK) + + cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (tiled_mma_S.thr_id.shape,), + ) + + # Alloc + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + dQ_cluster_full_mbar_ptr = storage.dQ_cluster_full_mbar_ptr.data_ptr() + dQ_cluster_empty_mbar_ptr = storage.dQ_cluster_empty_mbar_ptr.data_ptr() + + if const_expr(self.use_2cta_instrs): + dS_cluster_full_mbar_ptr = storage.dS_cluster_full_mbar_ptr + dS_cluster_empty_mbar_ptr = storage.dS_cluster_empty_mbar_ptr + dS_cluster_leader_mbar_ptr = storage.dS_cluster_leader_mbar_ptr + dQaccum_empty_mbar_ptr = storage.dQaccum_empty_mbar_ptr + else: + dS_cluster_full_mbar_ptr = None + dS_cluster_empty_mbar_ptr = None + dS_cluster_leader_mbar_ptr = None + dQaccum_empty_mbar_ptr = None + + # Barrier initialization + if const_expr(self.use_2cta_instrs): + if const_expr(self.tile_hdim == 192): + if warp_idx == 2: + cute.arch.mbarrier_init( + dQaccum_empty_mbar_ptr, + len(self.reduce_warp_ids), + ) + if warp_idx == 4: + cute.arch.mbarrier_init(dS_cluster_full_mbar_ptr, 1) + cute.arch.mbarrier_init(dS_cluster_empty_mbar_ptr, 1) + cute.arch.mbarrier_init(dS_cluster_leader_mbar_ptr, 2) + + if const_expr(self.cluster_reduce_dQ): + if warp_idx == 4: + for i in range(self.dQaccum_reduce_stage // 2): + cute.arch.mbarrier_init(dQ_cluster_full_mbar_ptr + i, 1) + cute.arch.mbarrier_init(dQ_cluster_empty_mbar_ptr + i, 1) + + tmem_alloc_barrier = cutlass.pipeline.NamedBarrier( + barrier_id=int(NamedBarrierBwdSm100.TmemPtr), + num_threads=cute.arch.WARP_SIZE + * len((self.mma_warp_id, *self.compute_warp_ids, *self.reduce_warp_ids)), + ) + tmem = cutlass.utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.mma_warp_id, + is_two_cta=self.use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + # UMMA producers and AsyncThread consumers + pipeline_producer_group_MMA_AsyncThread = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + pipeline_consumer_group_MMA_AsyncThread = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, len(self.compute_warp_ids) * self.cta_group_size + ) + pipeline_S_P = cutlass.pipeline.PipelineUmmaAsync.create( + num_stages=1, + producer_group=pipeline_producer_group_MMA_AsyncThread, + consumer_group=pipeline_consumer_group_MMA_AsyncThread, + barrier_storage=storage.S_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + ) + pipeline_dP = cutlass.pipeline.PipelineUmmaAsync.create( + num_stages=1, + producer_group=pipeline_producer_group_MMA_AsyncThread, + consumer_group=pipeline_consumer_group_MMA_AsyncThread, + barrier_storage=storage.dP_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + ) + pipeline_dKV = cutlass.pipeline.PipelineUmmaAsync.create( + num_stages=2, + producer_group=pipeline_producer_group_MMA_AsyncThread, + consumer_group=pipeline_consumer_group_MMA_AsyncThread, + barrier_storage=storage.dKV_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + ) + pipeline_consumer_group_MMA_AsyncThread_dQ = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, + len(self.reduce_warp_ids) * self.cta_group_size, + ) # Compute + pipeline_dQ = cutlass.pipeline.PipelineUmmaAsync.create( + num_stages=1, + producer_group=pipeline_producer_group_MMA_AsyncThread, + consumer_group=pipeline_consumer_group_MMA_AsyncThread_dQ, + barrier_storage=storage.dQ_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + ) + + # AsyncThread producers and UMMA consumers + # Only 1 thread per warp will signal + pipeline_PdS_producer_group = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, + len(self.compute_warp_ids) * self.cta_group_size, + ) # Compute + pipeline_PdS_consumer_group = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, len([self.mma_warp_id]) + ) # MMA + pipeline_dS = cutlass.pipeline.PipelineAsyncUmma.create( + num_stages=1, + producer_group=pipeline_PdS_producer_group, + consumer_group=pipeline_PdS_consumer_group, + barrier_storage=storage.dS_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + ) + + # TMA producer and UMMA consumers + pipeline_producer_group = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, len([self.load_warp_id]) + ) + # The arrive count is the number of mcast size + pipeline_consumer_group = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, len([self.mma_warp_id]) * self.num_mcast_ctas_b + ) + pipeline_consumer_group_compute = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, + len(self.compute_warp_ids) * 1, + ) + pipeline_LSE = cutlass.pipeline.PipelineTmaAsync.create( + barrier_storage=storage.LSE_mbar_ptr.data_ptr(), + num_stages=self.Q_stage, + producer_group=pipeline_producer_group, + consumer_group=pipeline_consumer_group_compute, + tx_count=self.tma_copy_bytes["LSE"], + # cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + pipeline_dPsum = cutlass.pipeline.PipelineTmaAsync.create( + barrier_storage=storage.dPsum_mbar_ptr.data_ptr(), + num_stages=self.dO_stage, + producer_group=pipeline_producer_group, + consumer_group=pipeline_consumer_group_compute, + tx_count=self.tma_copy_bytes["dPsum"], + # cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + pipeline_Q = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.Q_mbar_ptr.data_ptr(), + num_stages=self.Q_stage, + producer_group=pipeline_producer_group, + consumer_group=pipeline_consumer_group, + tx_count=self.tma_copy_bytes["Q"], + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + if const_expr(self.use_2cta_instrs): + if const_expr(self.tile_hdim == 192): + pipeline_Qt = pipeline_Q + else: + pipeline_Qt = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.Qt_mbar_ptr.data_ptr(), + num_stages=self.Q_stage, + producer_group=pipeline_producer_group, + consumer_group=pipeline_consumer_group, + tx_count=self.tma_copy_bytes["Q"], + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + pipeline_Kt = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.Kt_mbar_ptr.data_ptr(), + num_stages=self.single_stage, + producer_group=pipeline_producer_group, + consumer_group=pipeline_consumer_group, + tx_count=self.tma_copy_bytes["K"], + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + else: + pipeline_Qt = pipeline_Kt = pipeline_Q + + pipeline_dO = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.dO_mbar_ptr.data_ptr(), + num_stages=self.dO_stage, + producer_group=pipeline_producer_group, + consumer_group=pipeline_consumer_group, + tx_count=self.tma_copy_bytes["dO"], + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=False, + ) + + sQ = storage.sQ.get_tensor(sQ_layout.outer, swizzle=sQ_layout.inner, dtype=self.q_dtype) + if const_expr(self.use_2cta_instrs and self.tile_hdim <= 128): + sQt = storage.sQt.get_tensor( + sQt_layout.outer, swizzle=sQt_layout.inner, dtype=self.q_dtype + ) + else: + sQt = cute.make_tensor( + cute.recast_ptr(sQ.iterator, sQt_layout.inner, dtype=self.q_dtype), sQt_layout.outer + ) + sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner) + if const_expr(self.use_2cta_instrs): + sKt = storage.sKt.get_tensor(sKt_layout.outer, swizzle=sKt_layout.inner) + else: + sKt = cute.make_tensor(cute.recast_ptr(sK.iterator, sKt_layout.inner), sKt_layout.outer) + sV = storage.sV.get_tensor(sV_layout.outer, swizzle=sV_layout.inner) + sdSt = storage.sdS.get_tensor(sdSt_layout.outer, swizzle=sdSt_layout.inner) + sdS = cute.make_tensor(cute.recast_ptr(sdSt.iterator, sdS_layout.inner), sdS_layout.outer) + if const_expr(self.use_2cta_instrs): + if const_expr(self.tile_hdim <= 128): + sdS_xchg = storage.sdS_xchg.get_tensor(sdS_xchg_layout) + else: + sdS_xchg = storage.sdQaccum.get_tensor(sdS_xchg_layout, dtype=self.ds_dtype) + else: + sdS_xchg = None + + sdO = storage.sdO.get_tensor( + sdO_layout.outer, swizzle=sdO_layout.inner, dtype=self.do_dtype + ) + if const_expr(self.use_2cta_instrs and self.tile_hdim <= 128): + sdOt = storage.sdOt.get_tensor( + sdOt_layout.outer, swizzle=sdOt_layout.inner, dtype=self.do_dtype + ) + else: + sdOt = cute.make_tensor( + cute.recast_ptr(sdO.iterator, sdOt_layout.inner, dtype=self.do_dtype), + sdOt_layout.outer, + ) + + sLSE = storage.sLSE.get_tensor(sLSE_layout) + sdPsum = storage.sdPsum.get_tensor(sdPsum_layout) + if const_expr(self.use_2cta_instrs): + if const_expr(not self.dKV_postprocess): + sdV = storage.sV.get_tensor( + sdV_layout.outer, swizzle=sdV_layout.inner, dtype=self.dv_dtype + ) + sdK = storage.sK.get_tensor( + sdK_layout.outer, swizzle=sdK_layout.inner, dtype=self.dk_dtype + ) + else: + sdV = storage.sV.get_tensor(sdV_layout, dtype=self.dv_dtype) + sdK = storage.sK.get_tensor(sdK_layout, dtype=self.dk_dtype) + elif const_expr(not self.dKV_postprocess): + sdV = storage.sdO.get_tensor( + sdV_layout.outer, swizzle=sdV_layout.inner, dtype=self.dv_dtype + ) + sdK = storage.sQ.get_tensor( + sdK_layout.outer, swizzle=sdK_layout.inner, dtype=self.dk_dtype + ) + else: + sdV = storage.sdO.get_tensor(sdV_layout, dtype=self.dv_dtype) + sdK = storage.sQ.get_tensor(sdK_layout, dtype=self.dk_dtype) + + # Buffer sizing is guaranteed by max(...) in SharedStorage declarations + # for both sQ (reused as sdK) and sdO (reused as sdV) + sdQaccum = storage.sdQaccum.get_tensor(sdQaccum_layout) + + # TMEM + # This is a fake tensor, by right need to retrieve tmem_ptr. But we know that we always + # request 512 columns of tmem, so we know that it starts at 0. + tmem_ptr = cute.make_ptr(Float32, 0, mem_space=cute.AddressSpace.tmem, assumed_align=16) + # S + thr_mma_S = tiled_mma_S.get_slice(mma_tile_coord_v) + Sacc_shape = thr_mma_S.partition_shape_C(self.mma_tiler_kq[:2]) # (M, N) + tStS = thr_mma_S.make_fragment_C(Sacc_shape) + # (MMA, MMA_M, MMA_N) + tStS = cute.make_tensor(tmem_ptr + self.tmem_S_offset, tStS.layout) + # dP + thr_mma_dP = tiled_mma_dP.get_slice(mma_tile_coord_v) + dPacc_shape = thr_mma_dP.partition_shape_C(self.mma_tiler_vdo[:2]) + tdPtdP = thr_mma_dP.make_fragment_C(dPacc_shape) + tdPtdP = cute.make_tensor(tmem_ptr + self.tmem_dP_offset, tdPtdP.layout) + # dV + thr_mma_dV = tiled_mma_dV.get_slice(mma_tile_coord_v) + dvacc_shape = thr_mma_dV.partition_shape_C(self.mma_tiler_pdo[:2]) + tdVtdV = thr_mma_dV.make_fragment_C(dvacc_shape) + tdVtdV = cute.make_tensor(tmem_ptr + self.tmem_dV_offset, tdVtdV.layout) + tP = cute.make_tensor( + cute.recast_ptr(tmem_ptr + self.tmem_P_offset, dtype=self.do_dtype), tP_layout.outer + ) + # dK + thr_mma_dK = tiled_mma_dK.get_slice(mma_tile_coord_v) + dkacc_shape = thr_mma_dK.partition_shape_C(self.mma_tiler_dsq[:2]) + tdKtdK = thr_mma_dK.make_fragment_C(dkacc_shape) + tdKtdK = cute.make_tensor(tmem_ptr + self.tmem_dK_offset, tdKtdK.layout) + tdS = cute.make_tensor( + cute.recast_ptr(tmem_ptr + self.tmem_dS_offset, dtype=self.ds_dtype), tdS_layout.outer + ) + # dQ + thr_mma_dQ = tiled_mma_dQ.get_slice(mma_tile_coord_v) + dQacc_shape = thr_mma_dQ.partition_shape_C(self.mma_tiler_dsk[:2]) + tdQtdQ = thr_mma_dQ.make_fragment_C(dQacc_shape) + tdQtdQ = cute.make_tensor(tmem_ptr + self.tmem_dQ_offset, tdQtdQ.layout) + + block_info = BlockInfo( + self.tile_m, + # self.tile_n, + self.tile_n * self.cluster_shape_mnk[0], # careful, this case is not very well-tested + self.is_causal, + self.is_local, + False, # is_split_kv + window_size_left, + window_size_right, + qhead_per_kvhead_packgqa=1, + ) + SeqlenInfoCls = partial( + SeqlenInfoQK.create, + seqlen_q_static=mQ.shape[0], + seqlen_k_static=mK.shape[0], + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + tile_m=self.tile_m, + tile_n=self.tile_n * self.cluster_shape_mnk[0], + ) + TileSchedulerCls = partial(self.tile_scheduler_cls.create, tile_sched_params) + + AttentionMaskCls = self._generate_attention_mask_cls(window_size_left, window_size_right) + # EMPTY + # (15) + if warp_idx == self.empty_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_empty) + + # RELAY + # (14) + if warp_idx == self.relay_warp_id: + cute.arch.setmaxregister_decrease( + self.num_regs_mma if self.use_2cta_instrs else self.num_regs_empty + ) + if const_expr(self.use_2cta_instrs): + self.relay( + dS_cluster_full_mbar_ptr, + dS_cluster_empty_mbar_ptr, + dS_cluster_leader_mbar_ptr, + cluster_layout_vmnk, + block_info, + SeqlenInfoCls, + TileSchedulerCls, + ) + + # LOAD + # (13) + if warp_idx == self.load_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_load) + self.load( + thr_mma_S, + thr_mma_dP, + thr_mma_dV, + thr_mma_dK, + thr_mma_dQ, + mQ, + mK, + mKt, + mV, + mdO, + mQt, + mdOt, + mLSE, + mdPsum, + sQ, + sK, + sKt, + sV, + sdO, + sQt, + sdOt, + sLSE, + sdPsum, + tma_atom_Q, + tma_atom_K, + tma_atom_Kt, + tma_atom_V, + tma_atom_dO, + tma_atom_Qt, + tma_atom_dOt, + pipeline_Q, + pipeline_Qt, + pipeline_Kt, + pipeline_dO, + pipeline_LSE, + pipeline_dPsum, + cluster_layout_vmnk, + block_info, + SeqlenInfoCls, + TileSchedulerCls, + blocksparse_tensors, + should_load_Q=True, + should_load_dO=True, + ) + + # MMA + # (12) + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_mma) + + # Alloc tmem buffer + tmem.allocate(self.tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(Float32) + + self.mma( + tiled_mma_S, + tiled_mma_dP, + tiled_mma_dV, + tiled_mma_dK, + tiled_mma_dQ, + sQ, + sQt, + sK, + sKt, + sV, + sdO, + sdOt, + tP, + sdSt, + sdS, + tdS, + tStS, + tdPtdP, + tdVtdV, + tdKtdK, + tdQtdQ, + dS_cluster_full_mbar_ptr, + dS_cluster_empty_mbar_ptr, + dS_cluster_leader_mbar_ptr, + pipeline_Q, + pipeline_Qt, + pipeline_Kt, + pipeline_dO, + pipeline_S_P, + pipeline_dS, + pipeline_dKV, + pipeline_dP, + pipeline_dQ, + block_info, + SeqlenInfoCls, + TileSchedulerCls, + is_leader_cta, + blocksparse_tensors, + ) + # Dealloc the tensor memory buffer + tmem.relinquish_alloc_permit() + tmem_alloc_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + + # Compute + # (4, 5, 6, 7, 8, 9, 10, 11) --> 8 warps + if warp_idx >= self.compute_warp_ids[0] and warp_idx <= self.compute_warp_ids[-1]: + cute.arch.setmaxregister_increase(self.num_regs_compute) # 8 warps + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(Float32) + self.compute_loop( + thr_mma_S, + thr_mma_dP, + thr_mma_dV, + thr_mma_dK, + tStS, + tdPtdP, + tdVtdV, + tdKtdK, + sLSE, + sdPsum, + mdV, + mdK, + sdS, + sdS_xchg, + pipeline_LSE, + pipeline_dPsum, + pipeline_S_P, + pipeline_dS, + pipeline_dKV, + pipeline_dP, + dS_cluster_empty_mbar_ptr, + dS_cluster_full_mbar_ptr, + dQaccum_empty_mbar_ptr, + softmax_scale, + softmax_scale_log2, + block_info, + SeqlenInfoCls, + AttentionMaskCls, + TileSchedulerCls, + sdV, + sdK, + mdV_tma_tensor, + mdK_tma_tensor, + tma_atom_dV, + tma_atom_dK, + tiled_copy_r2s_dKV, + mdK_semaphore, + mdV_semaphore, + aux_data, + fastdiv_mods, + blocksparse_tensors, + ) + tmem_alloc_barrier.arrive() + + # Reduce + # (0, 1, 2, 3) - dQ + if warp_idx >= self.reduce_warp_ids[0] and warp_idx <= self.reduce_warp_ids[-1]: + cute.arch.setmaxregister_increase(self.num_regs_reduce) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(Float32) + self.dQacc_reduce( + mdQaccum, + sdQaccum, + thr_mma_dQ, + tdQtdQ, + pipeline_dQ, + dQaccum_empty_mbar_ptr, + block_info, + SeqlenInfoCls, + TileSchedulerCls, + mdQ_semaphore, + blocksparse_tensors, + ) + tmem_alloc_barrier.arrive() + + return + + @cute.jit + def relay( + self, + dS_cluster_full_mbar_ptr: cute.Pointer, + dS_cluster_empty_mbar_ptr: cute.Pointer, + dS_cluster_leader_mbar_ptr: cute.Pointer, + cluster_layout_vmnk: cute.Layout, + block_info: BlockInfo, + SeqlenInfoCls: Callable, + TileSchedulerCls: Callable, + ): + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + dS_cluster_phase = Int32(0) + + tile_scheduler = TileSchedulerCls() + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + n_block, head_idx, batch_idx, _ = work_tile.tile_idx + seqlen = SeqlenInfoCls(batch_idx) + m_block_min, m_block_max = block_info.get_m_block_min_max( + seqlen, n_block // self.cluster_shape_mnk[0] + ) + head_idx_kv = head_idx // self.qhead_per_kvhead + + process_tile = ( + const_expr(not self.is_local and not self.is_varlen_q) or m_block_min < m_block_max + ) + + if process_tile: + num_iters = m_block_max - m_block_min + for _ in cutlass.range(num_iters, unroll=1): + # Wait for dS_xchg from peer CTA + cute.arch.mbarrier_wait(dS_cluster_full_mbar_ptr, phase=dS_cluster_phase) + + # Arrive on MMA leader warp + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(dS_cluster_leader_mbar_ptr, Int32(0)) + + dS_cluster_phase ^= 1 + + tile_scheduler.prefetch_next_work() + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + + @cute.jit + def load( + self, + thr_mma_S: cute.ThrMma, + thr_mma_dP: cute.ThrMma, + thr_mma_dV: cute.ThrMma, + thr_mma_dK: cute.ThrMma, + thr_mma_dQ: cute.ThrMma, + mQ: cute.Tensor, + mK: cute.Tensor, + mKt: Optional[cute.Tensor], + mV: cute.Tensor, + mdO: cute.Tensor, + mQt: Optional[cute.Tensor], + mdOt: Optional[cute.Tensor], + mLSE: cute.Tensor, + mdPsum: cute.Tensor, + sQ: cute.Tensor, + sK: cute.Tensor, + sKt: cute.Tensor, + sV: cute.Tensor, + sdO: cute.Tensor, + sQt: cute.Tensor, + sdOt: cute.Tensor, + sLSE: cute.Tensor, + sdPsum: cute.Tensor, + tma_atom_Q: cute.CopyAtom, + tma_atom_K: cute.CopyAtom, + tma_atom_Kt: Optional[cute.CopyAtom], + tma_atom_V: cute.CopyAtom, + tma_atom_dO: cute.CopyAtom, + tma_atom_Qt: Optional[cute.CopyAtom], + tma_atom_dOt: Optional[cute.CopyAtom], # 2-CTA only + pipeline_Q: PipelineAsync, + pipeline_Qt: PipelineAsync, + pipeline_Kt: PipelineAsync, + pipeline_dO: PipelineAsync, + pipeline_LSE: PipelineAsync, + pipeline_dPsum: PipelineAsync, + cluster_layout_vmnk: cute.Layout, + block_info: BlockInfo, + SeqlenInfoCls: Callable, + TileSchedulerCls: Callable, + blocksparse_tensors: Optional[BlockSparseTensors] = None, + should_load_Q: bool = True, + should_load_dO: bool = True, + ): + producer_state_Q_LSE = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Producer, self.Q_stage + ) + producer_state_Qt = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Producer, self.Q_stage + ) + producer_state_Kt = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Producer, self.single_stage + ) + producer_state_dO_dPsum = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Producer, self.dO_stage + ) + producer_state_Q_Qt = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Producer, self.Q_stage + ) + producer_state_O_Ot = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Producer, self.dO_stage + ) + producer_state_LSE = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Producer, self.Q_stage + ) + producer_state_dPsum = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Producer, self.dO_stage + ) + + # Compute multicast mask for Q & dO buffer full + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + q_do_mcast_mask = None + if const_expr(self.is_q_do_mcast): + q_do_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + + tile_scheduler = TileSchedulerCls() + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + n_block, head_idx, batch_idx, _ = work_tile.tile_idx + seqlen = SeqlenInfoCls(batch_idx) + m_block_min, m_block_max = block_info.get_m_block_min_max( + seqlen, n_block // self.cluster_shape_mnk[0] + ) + head_idx_kv = head_idx // self.qhead_per_kvhead + n_block_cta_group = n_block // self.cta_group_size + + # GMEM tensors (varlen-aware) + mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx] + mK_cur = seqlen.offset_batch_K(mK, batch_idx, dim=3)[None, None, head_idx_kv] + mV_cur = seqlen.offset_batch_K(mV, batch_idx, dim=3)[None, None, head_idx_kv] + if const_expr(not seqlen.has_cu_seqlens_q): + mdO_cur = mdO[None, None, head_idx, batch_idx] + else: + mdO_cur = cute.domain_offset((0, seqlen.offset_q), mdO[None, None, head_idx]) + mLSE_cur = seqlen.offset_batch_Q(mLSE, batch_idx, dim=2, padded=True)[None, head_idx] + mdPsum_cur = seqlen.offset_batch_Q(mdPsum, batch_idx, dim=2, padded=True)[ + None, head_idx + ] + + if const_expr(self.use_2cta_instrs): + if const_expr(not seqlen.has_cu_seqlens_q): + mQt_cur = mQt[None, None, head_idx, batch_idx] + mdOt_cur = mdOt[None, None, head_idx, batch_idx] + else: + mQt_cur = cute.domain_offset((0, seqlen.offset_q, 0), mQt)[None, None, head_idx] + mdOt_cur = cute.domain_offset((seqlen.offset_q, 0, 0), mdOt)[ + None, None, head_idx + ] + if const_expr(not seqlen.has_cu_seqlens_k): + mKt_cur = mKt[None, None, head_idx_kv, batch_idx] + else: + mKt_cur = cute.domain_offset((0, seqlen.offset_k, 0), mKt)[ + None, None, head_idx_kv + ] + + # (1) S.T = K @ Q.T + gK = cute.local_tile( + mK_cur, cute.select(self.mma_tiler_kq, mode=[0, 2]), (n_block_cta_group, 0) + ) + tSgK = thr_mma_S.partition_A(gK) + + gQ = cute.local_tile(mQ_cur, cute.select(self.mma_tiler_kq, mode=[1, 2]), (None, 0)) + tSgQ = thr_mma_S.partition_B(gQ) + gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (None,)) + gdPsum = cute.local_tile(mdPsum_cur, (self.tile_m,), (None,)) + gdO = cute.local_tile(mdO_cur, cute.select(self.mma_tiler_pdo, mode=[1, 2]), (0, None)) + tdPgdO = thr_mma_dV.partition_B(gdO) + + a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) + load_K, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_K, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + tSgK, + sK, + single_stage=True, + ) + + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + load_Q, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_Q, + cta_coord=block_in_cluster_coord_vmnk[1], + cta_layout=b_cta_layout, + src_tensor=tSgQ, + dst_tensor=sQ, + mcast_mask=q_do_mcast_mask, + ) + load_Q = copy_utils.tma_producer_copy_fn(load_Q, pipeline_Q) + + # (2) dP = V @ dO.T + gV = cute.local_tile( + mV_cur, cute.select(self.mma_tiler_vdo, mode=[0, 2]), (n_block_cta_group, 0) + ) + tdPgV = thr_mma_dP.partition_A(gV) + + load_V, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_V, + 0, + cute.make_layout(1), + tdPgV, + sV, + single_stage=True, + ) + + if const_expr(tma_atom_dOt is not None): + gdOt = cute.local_tile( + mdOt_cur, cute.select(self.mma_tiler_vdo, mode=[1, 2]), (None, 0) + ) + tdPgdO = thr_mma_dP.partition_B(gdOt) + load_dOt, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_dOt, + cta_coord=block_in_cluster_coord_vmnk[1], + cta_layout=b_cta_layout, + src_tensor=tdPgdO, + dst_tensor=sdOt, + mcast_mask=q_do_mcast_mask, + ) + load_dOt = copy_utils.tma_producer_copy_fn(load_dOt, pipeline_dO) + + # (3) dV += P.T @ dO + gdO = cute.local_tile(mdO_cur, cute.select(self.mma_tiler_pdo, mode=[1, 2]), (0, None)) + tdVgdO = thr_mma_dV.partition_B(gdO) + load_dO, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_dO, + cta_coord=block_in_cluster_coord_vmnk[1], + cta_layout=b_cta_layout, + src_tensor=tdVgdO, + dst_tensor=sdO, + mcast_mask=q_do_mcast_mask, + ) + load_dO = copy_utils.tma_producer_copy_fn(load_dO, pipeline_dO) + + # (4) dK += dS.T @ Q (2-CTA: needs separate Qt load) + if const_expr(tma_atom_Qt is not None): + gQt = cute.local_tile( + mQt_cur, cute.select(self.mma_tiler_dsq, mode=[1, 2]), (0, None) + ) + tdKgQt = thr_mma_dK.partition_B(gQt) + load_Qt, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_Qt, + cta_coord=block_in_cluster_coord_vmnk[1], + cta_layout=b_cta_layout, + src_tensor=tdKgQt, + dst_tensor=sQt, + mcast_mask=q_do_mcast_mask, + ) + load_Qt = copy_utils.tma_producer_copy_fn(load_Qt, pipeline_Qt) + + # (5) dQ = dS @ K + if const_expr(self.use_2cta_instrs): + gKt = cute.local_tile( + mKt_cur, cute.select(self.mma_tiler_dsk, mode=[1, 2]), (0, n_block_cta_group) + ) + tdQgK = thr_mma_dQ.partition_B(gKt) + + load_Kt, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_Kt, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + tdQgK, + sKt, + single_stage=True, + ) + + copy_atom_stats = cute.make_copy_atom(cpasync.CopyBulkG2SOp(), Float32) + copy_stats = partial(cute.copy, copy_atom_stats) + # copy_atom_stats = cute.make_copy_atom(cpasync.CopyBulkG2SMulticastOp(), Float32) + # sLSE = cute.logical_divide(sLSE, (64,))[(None, block_in_cluster_coord_vmnk[1]), None] + # gLSE = cute.logical_divide(gLSE, (64,))[(None, block_in_cluster_coord_vmnk[1]), None] + # sdPsum = cute.logical_divide(sdPsum, (64,))[(None, block_in_cluster_coord_vmnk[1]), None] + # gdPsum = cute.logical_divide(gdPsum, (64,))[(None, block_in_cluster_coord_vmnk[1]), None] + # copy_stats = partial(cute.copy, copy_atom_stats, mcast_mask=q_do_mcast_mask) + + # some tiles might be empty due to block sparsity + if const_expr(self.use_block_sparsity): + total_m_block_cnt = get_total_q_block_count_bwd( + blocksparse_tensors, + batch_idx, + head_idx, + n_block, + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + ) + process_tile = total_m_block_cnt > Int32(0) + else: + process_tile = ( + const_expr(not self.is_local and not self.is_varlen_q) + or m_block_min < m_block_max + ) + + if process_tile: + if const_expr(self.use_block_sparsity): + producer_state_Q_LSE, producer_state_dO_dPsum = ( + produce_block_sparse_q_loads_bwd_sm100( + blocksparse_tensors, + batch_idx, + head_idx, + n_block, + producer_state_Q_LSE, + producer_state_dO_dPsum, + pipeline_Q, + pipeline_LSE, + pipeline_dO, + pipeline_dPsum, + load_K, + load_V, + load_Q, + load_dO, + copy_stats, + gLSE, + sLSE, + gdPsum, + sdPsum, + self.tma_copy_bytes["K"], + self.tma_copy_bytes["V"], + should_load_Q=should_load_Q, + should_load_dO=should_load_dO, + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + ) + ) + else: + first_m_block = m_block_min + if const_expr(self.use_2cta_instrs and self.tile_hdim == 192): + #### Prologue #### + assert should_load_Q and should_load_dO + # K & Q (for S) + pipeline_Q.producer_acquire( + producer_state_Q_Qt, + extra_tx_count=self.tma_copy_bytes["K"], + ) + load_K(tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q_Qt)) + load_Q(first_m_block, producer_state=producer_state_Q_Qt) + pipeline_Q.producer_commit(producer_state_Q_Qt) + producer_state_Q_Qt.advance() + # LSE + pipeline_LSE.producer_acquire(producer_state_LSE) + with cute.arch.elect_one(): + copy_stats( + gLSE[None, first_m_block], + sLSE[None, producer_state_LSE.index], + mbar_ptr=pipeline_LSE.producer_get_barrier(producer_state_LSE), + ) + producer_state_LSE.advance() + + # dOt + V, for dP.T = V @ dO.T + pipeline_dO.producer_acquire( + producer_state_O_Ot, + extra_tx_count=self.tma_copy_bytes["V"], + ) + load_V(tma_bar_ptr=pipeline_dO.producer_get_barrier(producer_state_O_Ot)) + load_dOt(first_m_block, producer_state=producer_state_O_Ot) + pipeline_dO.producer_commit(producer_state_O_Ot) + producer_state_O_Ot.advance() + # dPsum + pipeline_dPsum.producer_acquire(producer_state_dPsum) + with cute.arch.elect_one(): + copy_stats( + gdPsum[None, first_m_block], + sdPsum[None, producer_state_dPsum.index], + mbar_ptr=pipeline_dPsum.producer_get_barrier(producer_state_dPsum), + ) + producer_state_dPsum.advance() + + # Qt, for dK = dS.T @ Q + pipeline_Qt.producer_acquire( + producer_state_Q_Qt, + extra_tx_count=self.tma_copy_bytes["K"], + ) + load_Qt(first_m_block, producer_state=producer_state_Q_Qt) + load_Kt(tma_bar_ptr=pipeline_Qt.producer_get_barrier(producer_state_Q_Qt)) + pipeline_Qt.producer_commit(producer_state_Q_Qt) + producer_state_Q_Qt.advance() + + # dO, for dV = P.T @ dO + pipeline_dO.producer_acquire(producer_state_O_Ot) + load_dO(first_m_block, producer_state=producer_state_O_Ot) + pipeline_dO.producer_commit(producer_state_O_Ot) + producer_state_O_Ot.advance() + + #### Mainloop #### + # 2CTA: [lse | Q | dOt | dPsum | Qt | dO] + for m_block in cutlass.range(m_block_min + 1, m_block_max, unroll=1): + # LSE + pipeline_LSE.producer_acquire(producer_state_LSE) + with cute.arch.elect_one(): + copy_stats( + gLSE[None, m_block], + sLSE[None, producer_state_LSE.index], + mbar_ptr=pipeline_LSE.producer_get_barrier(producer_state_LSE), + ) + producer_state_LSE.advance() + + # Q + pipeline_Q.producer_acquire(producer_state_Q_Qt) + load_Q(m_block, producer_state=producer_state_Q_Qt) + pipeline_Q.producer_commit(producer_state_Q_Qt) + producer_state_Q_Qt.advance() + + # dPsum + pipeline_dPsum.producer_acquire(producer_state_dPsum) + with cute.arch.elect_one(): + copy_stats( + gdPsum[None, m_block], + sdPsum[None, producer_state_dPsum.index], + mbar_ptr=pipeline_dPsum.producer_get_barrier( + producer_state_dPsum + ), + ) + producer_state_dPsum.advance() + + # dOt, for dP.T = V @ dO.T + pipeline_dO.producer_acquire(producer_state_O_Ot) + load_dOt(m_block, producer_state=producer_state_O_Ot) + pipeline_dO.producer_commit(producer_state_O_Ot) + producer_state_O_Ot.advance() + + # Qt, for dK = dS.T @ Q + pipeline_Qt.producer_acquire(producer_state_Q_Qt) + load_Qt(m_block, producer_state=producer_state_Q_Qt) + pipeline_Qt.producer_commit(producer_state_Q_Qt) + producer_state_Q_Qt.advance() + + # dO, for dV = P.T @ dO + pipeline_dO.producer_acquire(producer_state_O_Ot) + load_dO(m_block, producer_state=producer_state_O_Ot) + pipeline_dO.producer_commit(producer_state_O_Ot) + producer_state_O_Ot.advance() + + else: + #### Prologue #### + if const_expr(should_load_Q): + # K & Q (for S) + pipeline_Q.producer_acquire( + producer_state_Q_LSE, extra_tx_count=self.tma_copy_bytes["K"] + ) + load_K( + tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q_LSE) + ) + load_Q(first_m_block, producer_state=producer_state_Q_LSE) + pipeline_Q.producer_commit(producer_state_Q_LSE) + + # LSE + pipeline_LSE.producer_acquire(producer_state_Q_LSE) + with cute.arch.elect_one(): + copy_stats( + gLSE[None, first_m_block], + sLSE[None, producer_state_Q_LSE.index], + mbar_ptr=pipeline_LSE.producer_get_barrier( + producer_state_Q_LSE + ), + ) + producer_state_Q_LSE.advance() + + if const_expr(should_load_dO): + pipeline_dO.producer_acquire( + producer_state_dO_dPsum, + extra_tx_count=self.tma_copy_bytes["V"] + self.tma_copy_bytes["dO"] + if const_expr(tma_atom_dOt is not None) + else self.tma_copy_bytes["V"], + ) + load_V( + tma_bar_ptr=pipeline_dO.producer_get_barrier( + producer_state_dO_dPsum + ) + ) + load_dO(first_m_block, producer_state=producer_state_dO_dPsum) + if const_expr(tma_atom_dOt is not None): + load_dOt(first_m_block, producer_state=producer_state_dO_dPsum) + pipeline_dO.producer_commit(producer_state_dO_dPsum) + + # dPsum + pipeline_dPsum.producer_acquire(producer_state_dO_dPsum) + with cute.arch.elect_one(): + copy_stats( + gdPsum[None, first_m_block], + sdPsum[None, producer_state_dO_dPsum.index], + mbar_ptr=pipeline_dPsum.producer_get_barrier( + producer_state_dO_dPsum + ), + ) + producer_state_dO_dPsum.advance() + + if const_expr(self.use_2cta_instrs): + pipeline_Kt.producer_acquire(producer_state_Kt) + load_Kt(tma_bar_ptr=pipeline_Kt.producer_get_barrier(producer_state_Kt)) + pipeline_Kt.producer_commit(producer_state_Kt) + producer_state_Kt.advance() + #### Main Loop #### + for m_block in cutlass.range(m_block_min + 1, m_block_max, unroll=1): + if const_expr(should_load_Q): + if const_expr(tma_atom_Qt is not None): + pipeline_Qt.producer_acquire(producer_state_Qt) + load_Qt(m_block - 1, producer_state=producer_state_Qt) + pipeline_Qt.producer_commit(producer_state_Qt) + producer_state_Qt.advance() + + # Q (for S) + pipeline_Q.producer_acquire(producer_state_Q_LSE) + load_Q(m_block, producer_state=producer_state_Q_LSE) + pipeline_Q.producer_commit(producer_state_Q_LSE) + + # LSE + pipeline_LSE.producer_acquire(producer_state_Q_LSE) + with cute.arch.elect_one(): + copy_stats( + gLSE[None, m_block], + sLSE[None, producer_state_Q_LSE.index], + mbar_ptr=pipeline_LSE.producer_get_barrier( + producer_state_Q_LSE + ), + ) + producer_state_Q_LSE.advance() + + if const_expr(should_load_dO): + pipeline_dO.producer_acquire( + producer_state_dO_dPsum, + extra_tx_count=self.tma_copy_bytes["dO"] + if const_expr(tma_atom_dOt is not None) + else 0, + ) + load_dO(m_block, producer_state=producer_state_dO_dPsum) + if const_expr(tma_atom_dOt is not None): + load_dOt(m_block, producer_state=producer_state_dO_dPsum) + pipeline_dO.producer_commit(producer_state_dO_dPsum) + + # dPsum + pipeline_dPsum.producer_acquire(producer_state_dO_dPsum) + with cute.arch.elect_one(): + copy_stats( + gdPsum[None, m_block], + sdPsum[None, producer_state_dO_dPsum.index], + mbar_ptr=pipeline_dPsum.producer_get_barrier( + producer_state_dO_dPsum + ), + ) + producer_state_dO_dPsum.advance() + + #### Tail #### + if const_expr(should_load_Q): + if const_expr(tma_atom_Qt is not None): + pipeline_Qt.producer_acquire(producer_state_Qt) + load_Qt(m_block_max - 1, producer_state=producer_state_Qt) + pipeline_Qt.producer_commit(producer_state_Qt) + producer_state_Qt.advance() + + if const_expr(self.use_2cta_instrs and self.tile_hdim == 192): + pipeline_Q.producer_tail(producer_state_Q_Qt) + pipeline_LSE.producer_tail(producer_state_LSE) + pipeline_dO.producer_tail(producer_state_O_Ot) + pipeline_dPsum.producer_tail(producer_state_dPsum) + else: + if const_expr(should_load_Q): + pipeline_Q.producer_tail(producer_state_Q_LSE.clone()) + pipeline_LSE.producer_tail(producer_state_Q_LSE) + if const_expr(tma_atom_Qt is not None): + pipeline_Qt.producer_tail(producer_state_Qt) + if const_expr(should_load_dO): + pipeline_dO.producer_tail(producer_state_dO_dPsum.clone()) + pipeline_dPsum.producer_tail(producer_state_dO_dPsum) + + tile_scheduler.prefetch_next_work() + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + + @cute.jit + def mma( + self, + tiled_mma_S: cute.TiledMma, + tiled_mma_dP: cute.TiledMma, + tiled_mma_dV: cute.TiledMma, + tiled_mma_dK: cute.TiledMma, + tiled_mma_dQ: cute.TiledMma, + sQ: cute.Tensor, + sQt: cute.Tensor, + sK: cute.Tensor, + sKt: cute.Tensor, + sV: cute.Tensor, + sdO: cute.Tensor, + sdOt: cute.Tensor, + tP: cute.Tensor, + sdSt: cute.Tensor, + sdS: cute.Tensor, + tdS: cute.Tensor, + tStS: cute.Tensor, + tdPtdP: cute.Tensor, + tdVtdV: cute.Tensor, + tdKtdK: cute.Tensor, + tdQtdQ: cute.Tensor, + dS_cluster_full_mbar_ptr: cute.Pointer, + dS_cluster_empty_mbar_ptr: cute.Pointer, + dS_cluster_leader_mbar_ptr: cute.Pointer, + pipeline_Q: PipelineAsync, + pipeline_Qt: PipelineAsync, + pipeline_Kt: PipelineAsync, + pipeline_dO: PipelineAsync, + pipeline_S_P: PipelineAsync, + pipeline_dS: PipelineAsync, + pipeline_dKV: PipelineAsync, + pipeline_dP: PipelineAsync, + pipeline_dQ: PipelineAsync, + block_info: BlockInfo, + SeqlenInfoCls: Callable, + TileSchedulerCls: Callable, + is_leader_cta: cutlass.Boolean, + blocksparse_tensors: Optional[BlockSparseTensors] = None, + ): + # [2025-10-21] For reasons I don't understand, putting these partitioning in the main + # kernel (before warp specialization) is a lot slower tha putting them here. + # Partition smem / tmem tensors + # S = K @ Q.T + tSrK = tiled_mma_S.make_fragment_A(sK) + tSrQ = tiled_mma_S.make_fragment_B(sQ) + # dP = V @ dOt.T + tdPrV = tiled_mma_dP.make_fragment_A(sV) + tdPrdOt = tiled_mma_dP.make_fragment_B(sdOt) + # dK = dS.T @ Q + # For 2-CTA, dS (dK mma) MUST come from TMEM (cannot use SMEM) + if const_expr(self.use_smem_dS_for_mma_dK and not self.use_2cta_instrs): + tdKrdS = tiled_mma_dK.make_fragment_A(sdSt) # From SMEM + else: + tdKrdS = tiled_mma_dK.make_fragment_A(tdS) # From TMEM + + tdKrQ = tiled_mma_dK.make_fragment_B(sQt) + # dQ = dS @ K + tdQrdS = tiled_mma_dQ.make_fragment_A(sdS) + tdQrK = tiled_mma_dQ.make_fragment_B(sKt) + # dV = P @ dO.T + tdVrdO = tiled_mma_dV.make_fragment_B(sdO) + tdVrP = tiled_mma_dV.make_fragment_A(tP) + + # mma_qk_fn = partial(gemm_w_idx, tiled_mma_S, tStS, tSrK, tSrQ, zero_init=True) + mma_qk_fn = partial( + gemm_ptx_w_idx, + tiled_mma_S, + tStS, + tSrK, + tSrQ, + sA=sK, + sB=sQ, + zero_init=True, + cta_group=self.cta_group_size, + ) + # mma_dov_fn = partial(gemm_w_idx, tiled_mma_dP, tdPtdP, tdPrV, tdPrdOt, zero_init=True) + mma_dov_fn = partial( + gemm_ptx_w_idx, + tiled_mma_dP, + tdPtdP, + tdPrV, + tdPrdOt, + sA=sV, + sB=sdOt, + zero_init=True, + cta_group=self.cta_group_size, + ) + # mma_pdo_fn = partial(gemm_w_idx, tiled_mma_dV, tdVtdV, tdVrP, tdVrdO) + mma_pdo_fn = partial( + gemm_ptx_w_idx, + tiled_mma_dV, + tdVtdV, + tdVrP, + tdVrdO, + sA=None, + sB=sdO, + tA_addr=self.tmem_P_offset, + cta_group=self.cta_group_size, + ) + num_unroll_groups = 2 if const_expr(self.use_2cta_instrs) else 1 + mma_dsk_fn = partial( + gemm_w_idx, + tiled_mma_dQ, + tdQtdQ, + tdQrdS, + tdQrK, + zero_init=True, + num_unroll_groups=num_unroll_groups, + ) + # mma_dsk_fn = partial( + # gemm_ptx_w_idx, tiled_mma_dQ, tdQtdQ, tdQrdS, tdQrK, sA=sdS, sB=sKt, zero_init=True + # ) + if const_expr(self.use_smem_dS_for_mma_dK and not self.use_2cta_instrs): + mma_dsq_fn = partial(gemm_w_idx, tiled_mma_dK, tdKtdK, tdKrdS, tdKrQ) + else: + # Need to explicitly pass in tA_addr for correctness + mma_dsq_fn = partial( + gemm_ptx_w_idx, + tiled_mma_dK, + tdKtdK, + tdKrdS, + tdKrQ, + sA=None, + sB=sQt, + tA_addr=self.tmem_dS_offset, + cta_group=self.cta_group_size, + ) + + pipeline_Q_consumer = pipeline_Q.make_consumer() + + consumer_state_Qt = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Consumer, self.Q_stage + ) + consumer_state_Q = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Consumer, self.Q_stage + ) + consumer_state_Kt = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Consumer, self.single_stage + ) + consumer_state_dO = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Consumer, self.dO_stage + ) + producer_phase_acc = Int32(1) # For S & P, dP, dQ + producer_phase_dQ = Int32(1) # 2-CTA: separate phase for dQ pipeline + consumer_state_dS = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Consumer, 1 + ) + producer_phase_dKV = Int32(1) + cta_group = pipeline_S_P.cta_group + + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + dS_cluster_phase = Int32(0) + + tile_scheduler = TileSchedulerCls() + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + n_block, head_idx, batch_idx, _ = work_tile.tile_idx + seqlen = SeqlenInfoCls(batch_idx) # must be seqlen_k + m_block_min, m_block_max = block_info.get_m_block_min_max( + seqlen, n_block // self.cluster_shape_mnk[0] + ) + + if const_expr(self.use_block_sparsity): + block_iter_count = get_total_q_block_count_bwd( + blocksparse_tensors, + batch_idx, + head_idx, + n_block, + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + ) + process_tile = block_iter_count > Int32(0) + else: + block_iter_count = m_block_max - m_block_min + process_tile = ( + const_expr(not self.is_local and not self.is_varlen_q) + or m_block_min < m_block_max + ) + + if const_expr(self.use_2cta_instrs and self.tile_hdim == 192): + if is_leader_cta and process_tile: + accumulate_dK = False + accumulate_dV = False + + # ----------------------------------------------------------- + ###### MAIN LOOP + # ----------------------------------------------------------- + # 1. S.T = K @ Q.T + # 2. dP.T = V @ dO.T + # 3. dK = dS.T @ Q + # 4. dV = P.T @ dO + # 5. dQ = dS @ K + + main_loop_iters = m_block_max - m_block_min + + # empty waits + # pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc) + # pipeline_dP.sync_object_empty.wait(0, producer_phase_acc) + + for _ in cutlass.range(main_loop_iters, unroll=1): + # 1) S.T = K @ Q.T + pipeline_Q.consumer_wait(consumer_state_Q) + pipeline_dQ.sync_object_empty.wait( + 0, producer_phase_acc + ) # dQ tmem overlaps with S + mma_qk_fn(B_idx=consumer_state_Q.index) + pipeline_S_P.sync_object_full.arrive( + 0, pipeline_S_P.producer_mask, cta_group + ) + pipeline_Q.consumer_release(consumer_state_Q) + consumer_state_Q.advance() + + producer_phase_acc ^= 1 + + # 2) dP.T = V @ dO.T + pipeline_dO.consumer_wait(consumer_state_dO) + pipeline_S_P.sync_object_empty.wait( + 0, producer_phase_acc + ) # dP tmem overlaps with S + mma_dov_fn(B_idx=consumer_state_dO.index) + pipeline_dP.sync_object_full.arrive(0, pipeline_dP.producer_mask, cta_group) + pipeline_dO.consumer_release(consumer_state_dO) + consumer_state_dO.advance() + + # 3) dK = dS.T @ Q + pipeline_Q.consumer_wait(consumer_state_Q) + pipeline_dP.sync_object_empty.wait(0, producer_phase_acc) # dP -> dS + mma_dsq_fn(B_idx=consumer_state_Q.index, zero_init=not accumulate_dK) + pipeline_Q.consumer_release(consumer_state_Q) + consumer_state_Q.advance() + accumulate_dK = True + + # 4) dV = P.T @ dO + # Note: if dS is written to tmem, P must be written to tmem + pipeline_dO.consumer_wait(consumer_state_dO) + mma_pdo_fn(B_idx=consumer_state_dO.index, zero_init=not accumulate_dV) + pipeline_dO.consumer_release(consumer_state_dO) + consumer_state_dO.advance() + accumulate_dV = True + + # 5) dQ = dS @ K + pipeline_dS.consumer_wait(consumer_state_dS) + cute.arch.mbarrier_wait(dS_cluster_leader_mbar_ptr, phase=dS_cluster_phase) + mma_dsk_fn() + pipeline_dQ.sync_object_full.arrive(0, pipeline_dQ.producer_mask, cta_group) + pipeline_dS.consumer_release(consumer_state_dS) + consumer_state_dS.advance() + dS_cluster_phase ^= 1 + + # signal to the epilogue that dV is ready + pipeline_dKV.sync_object_empty.wait(0, producer_phase_dKV) + pipeline_dKV.sync_object_full.arrive(0, pipeline_dKV.producer_mask, cta_group) + # signal to the epilogue that dK is ready + pipeline_dKV.sync_object_empty.wait(1, producer_phase_dKV) + pipeline_dKV.sync_object_full.arrive(1, pipeline_dKV.producer_mask, cta_group) + producer_phase_dKV ^= 1 + elif const_expr(self.use_2cta_instrs): + if is_leader_cta and process_tile: + accumulate_dK = False + # ----------------------------------------------------------- + ###### Prologue + # ----------------------------------------------------------- + # 1. S = Q0 @ K.T + # 2. dP = V @ dOt.T + # 3. dV = P @ dO + + # 1) S = K @ Q + pipeline_Q.consumer_wait(consumer_state_Q) + pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc) + mma_qk_fn(B_idx=consumer_state_Q.index) + pipeline_S_P.sync_object_full.arrive(0, pipeline_S_P.producer_mask, cta_group) + pipeline_Q.consumer_release(consumer_state_Q) + consumer_state_Q.advance() + + # 2) dP = V @ dOt.T + pipeline_dO.consumer_wait(consumer_state_dO) + pipeline_dP.sync_object_empty.wait(0, producer_phase_acc) + mma_dov_fn(B_idx=consumer_state_dO.index) + pipeline_dP.sync_object_full.arrive(0, pipeline_dP.producer_mask, cta_group) + + # 3) dV = P.T @ dO + producer_phase_acc ^= 1 + pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc) + mma_pdo_fn(B_idx=consumer_state_dO.index, zero_init=True) + pipeline_dO.consumer_release(consumer_state_dO) + consumer_state_dO.advance() + + pipeline_Kt.consumer_wait(consumer_state_Kt) + # ----------------------------------------------------------- + ###### MAIN LOOP + # ----------------------------------------------------------- + # 1. S.T = K @ Q.T + # 2. dK = dS.T @ Q + # 3. dP.T = V @ dO.T + # 4. dQ = dS @ K + # 5. dV = P.T @ dO + + main_loop_iters = ( + block_iter_count - 1 + if const_expr(self.use_block_sparsity) + else m_block_max - m_block_min - 1 + ) + + for _ in cutlass.range(main_loop_iters, unroll=1): + # (1) S.T = K @ Q.T (next) + pipeline_Q.consumer_wait(consumer_state_Q) + pipeline_dQ.sync_object_empty.wait(0, producer_phase_dQ) + mma_qk_fn(B_idx=consumer_state_Q.index) + pipeline_S_P.sync_object_full.arrive( + 0, pipeline_S_P.producer_mask, cta_group + ) + pipeline_Q.consumer_release(consumer_state_Q) + consumer_state_Q.advance() + + # pipeline_dS.consumer_wait(consumer_state_dS) + # (2) dK += dS.T @ Q (cur) + pipeline_Qt.consumer_wait(consumer_state_Qt) + pipeline_dP.sync_object_empty.wait(0, producer_phase_acc) # dP -> dS + mma_dsq_fn(B_idx=consumer_state_Qt.index, zero_init=not accumulate_dK) + accumulate_dK = True + pipeline_Qt.consumer_release(consumer_state_Qt) + consumer_state_Qt.advance() + + # (3) dP.T = V @ dO.T (next) + pipeline_dO.consumer_wait(consumer_state_dO) + mma_dov_fn(B_idx=consumer_state_dO.index) + pipeline_dP.sync_object_full.arrive(0, pipeline_dP.producer_mask, cta_group) + + # (5) dQ = dS @ K (cur) + pipeline_dS.consumer_wait(consumer_state_dS) + cute.arch.mbarrier_wait(dS_cluster_leader_mbar_ptr, phase=dS_cluster_phase) + mma_dsk_fn() + pipeline_dQ.sync_object_full.arrive(0, pipeline_dQ.producer_mask, cta_group) + pipeline_dS.consumer_release(consumer_state_dS) + consumer_state_dS.advance() + dS_cluster_phase ^= 1 + producer_phase_dQ ^= 1 + + # (4) dV += P.T @ dO (next) + producer_phase_acc ^= 1 + pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc) # S -> P + mma_pdo_fn(B_idx=consumer_state_dO.index, zero_init=False) + pipeline_dO.consumer_release(consumer_state_dO) + consumer_state_dO.advance() + + pipeline_S_P.sync_object_full.arrive(0, pipeline_S_P.producer_mask, cta_group) + + # signal to the epilogue that dV is ready + pipeline_dKV.sync_object_empty.wait(0, producer_phase_dKV) + pipeline_dKV.sync_object_full.arrive(0, pipeline_dKV.producer_mask, cta_group) + pipeline_dKV.sync_object_empty.wait(1, producer_phase_dKV) + + # ----------------------------------------------------------- + # Tail: Remaining dK and dQ + # ----------------------------------------------------------- + # pipeline_dS.consumer_wait(consumer_state_dS) + # dK += dS.T @ Q + pipeline_Qt.consumer_wait(consumer_state_Qt) + pipeline_dP.sync_object_empty.wait(0, producer_phase_acc) # dP -> dS + mma_dsq_fn(B_idx=consumer_state_Qt.index, zero_init=not accumulate_dK) + pipeline_Qt.consumer_release(consumer_state_Qt) + consumer_state_Qt.advance() + # signal to the epilogue that dK is ready + pipeline_dKV.sync_object_full.arrive(1, pipeline_dKV.producer_mask, cta_group) + producer_phase_dKV ^= 1 + + # dQ = dS @ K + pipeline_dS.consumer_wait(consumer_state_dS) + cute.arch.mbarrier_wait(dS_cluster_leader_mbar_ptr, phase=dS_cluster_phase) + pipeline_dQ.sync_object_empty.wait(0, producer_phase_dQ) + mma_dsk_fn() + pipeline_dQ.sync_object_full.arrive(0, pipeline_dQ.producer_mask, cta_group) + pipeline_dS.consumer_release(consumer_state_dS) + pipeline_Kt.consumer_release(consumer_state_Kt) + consumer_state_dS.advance() + consumer_state_Kt.advance() + dS_cluster_phase ^= 1 + producer_phase_dQ ^= 1 + + producer_phase_acc ^= 1 + else: + if is_leader_cta and process_tile: + accumulate_dK = False + # ----------------------------------------------------------- + ###### Prologue + # ----------------------------------------------------------- + # 1. S = Q0 @ K.T + # 2. dP = V @ dOt.T + # 3. dV = P @ dO + + # 1) S = K @ Q + handle_Q = pipeline_Q_consumer.wait_and_advance() + pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc) + mma_qk_fn(B_idx=handle_Q.index) + pipeline_S_P.sync_object_full.arrive(0, pipeline_S_P.producer_mask, cta_group) + + # 2) dP = V @ dOt.T + pipeline_dO.consumer_wait(consumer_state_dO) + pipeline_dP.sync_object_empty.wait(0, producer_phase_acc) + pipeline_dQ.sync_object_empty.wait(0, producer_phase_acc) + mma_dov_fn(B_idx=consumer_state_dO.index) + pipeline_dP.sync_object_full.arrive(0, pipeline_dP.producer_mask, cta_group) + + producer_phase_acc ^= 1 + # 3) dV = P.T @ dO + pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc) + mma_pdo_fn(B_idx=consumer_state_dO.index, zero_init=True) + pipeline_dO.consumer_release(consumer_state_dO) + consumer_state_dO.advance() + + # ----------------------------------------------------------- + ###### MAIN LOOP + # ----------------------------------------------------------- + # 1. S = K @ Q.T + # 2. dQ = dS @ K + # 3. dK = dS.T @ Q + # 4. dP = V @ dOt.T + # 5. dV = P.T @ dO + + # For block sparsity, we use block_iter_count; for dense, use m_block range + # MMA doesn't need actual m_block indices, just the iteration count + main_loop_iters = ( + block_iter_count - 1 + if const_expr(self.use_block_sparsity) + else m_block_max - m_block_min - 1 + ) + + handle_Q_next = handle_Q + for _ in cutlass.range(main_loop_iters, unroll=1): + # (1) S.T = K @ Q.T + handle_Q_next = pipeline_Q_consumer.wait_and_advance() + mma_qk_fn(B_idx=handle_Q_next.index) + pipeline_S_P.sync_object_full.arrive( + 0, pipeline_S_P.producer_mask, cta_group + ) + + # (2) dK += dS.T @ Q + pipeline_dS.consumer_wait(consumer_state_dS) + mma_dsq_fn(B_idx=handle_Q.index, zero_init=not accumulate_dK) + accumulate_dK = True + handle_Q.release() + + # (3) dQ = dS @ K + mma_dsk_fn() + pipeline_dQ.sync_object_full.arrive(0, pipeline_dQ.producer_mask, cta_group) + pipeline_dS.consumer_release(consumer_state_dS) + consumer_state_dS.advance() + + # (4) dP = V @ dO.T + pipeline_dO.consumer_wait(consumer_state_dO) + pipeline_dQ.sync_object_empty.wait(0, producer_phase_acc) + mma_dov_fn(B_idx=consumer_state_dO.index) + pipeline_dP.sync_object_full.arrive(0, pipeline_dP.producer_mask, cta_group) + + # (5) dV += P.T @ dO + producer_phase_acc ^= 1 + pipeline_S_P.sync_object_empty.wait(0, producer_phase_acc) + mma_pdo_fn(B_idx=consumer_state_dO.index, zero_init=False) + pipeline_dO.consumer_release(consumer_state_dO) + consumer_state_dO.advance() + + handle_Q = handle_Q_next + + pipeline_S_P.sync_object_full.arrive(0, pipeline_S_P.producer_mask, cta_group) + + # signal to the epilogue that dV is ready + # pipeline_dKV.producer_acquire(producer_state_dKV) + pipeline_dKV.sync_object_empty.wait(0, producer_phase_dKV) + # pipeline_dKV.producer_commit(producer_state_dKV) + pipeline_dKV.sync_object_full.arrive(0, pipeline_dKV.producer_mask, cta_group) + # producer_state_dKV.advance() + # pipeline_dKV.producer_acquire(producer_state_dKV) + pipeline_dKV.sync_object_empty.wait(1, producer_phase_dKV) + + # ----------------------------------------------------------- + # Tail: Remaining dK and dQ + # ----------------------------------------------------------- + # 1) dK += dS.T @ Q + pipeline_dS.consumer_wait(consumer_state_dS) + mma_dsq_fn(B_idx=handle_Q.index, zero_init=not accumulate_dK) + # signal to the epilogue that dK is ready + pipeline_dKV.sync_object_full.arrive(1, pipeline_dKV.producer_mask, cta_group) + producer_phase_dKV ^= 1 + + # 2) dQ = dS @ K + mma_dsk_fn() + pipeline_dQ.sync_object_full.arrive(0, pipeline_dQ.producer_mask, cta_group) + handle_Q.release() + pipeline_dS.consumer_release(consumer_state_dS) + consumer_state_dS.advance() + + producer_phase_acc ^= 1 + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + # Currently it hangs if we have this S_P.producer_tail, will need to understand why + # pipeline_S_P.producer_tail(producer_state_S_P) + # pipeline_dP.producer_tail(producer_state_dP) + # pipeline_dKV.producer_tail(producer_state_dKV) + # pipeline_dQ.producer_tail(producer_state_dQ) + + @cute.jit + def split_wg( + self, + t: cute.Tensor, + wg_idx: cutlass.Int32, + num_wg: cutlass.Constexpr[int], + ): + reduced_shape = cute.product_each(t.shape) + rank = len(reduced_shape) + if const_expr(reduced_shape[1] > 1): + assert rank >= 2, "Need rank >= 2 for t in split_wg" + t = cute.logical_divide(t, (reduced_shape[0], reduced_shape[1] // num_wg)) + coord = (None, (None, wg_idx)) + (None,) * (rank - 2) + else: + assert rank >= 3, "Need rank >= 3 for t in split_wg" + if const_expr(rank == 3): + t = cute.logical_divide( + t, (reduced_shape[0], reduced_shape[1], reduced_shape[2] // num_wg) + ) + coord = ( + None, + None, + (None, wg_idx), + ) + (None,) * (rank - 3) + else: + t = cute.logical_divide( + t, + ( + reduced_shape[0], + reduced_shape[1], + reduced_shape[2], + reduced_shape[3] // num_wg, + ), + ) + coord = ( + None, + None, + None, + (None, wg_idx), + ) + (None,) * (rank - 4) + return t[coord] + + @cute.jit + def apply_score_mod( + self, + tSrS_t2r, + thr_copy_t2r, + thr_mma_S, + batch_idx, + head_idx, + m_block, + n_block, + softmax_scale, + seqlen_info, + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + ): + """Apply forward score modification for SM100 backward pass.""" + # In bwd, S is computed as K @ Q.T so dimensions are (tile_n, tile_m). + # With 2CTA, partition_C must see the full cluster tile so each CTA + # gets its own half of the tile. + cluster_tile_n = self.tile_n * self.cta_group_size + cluster_n_block = n_block // self.cta_group_size + cS = cute.make_identity_tensor((cluster_tile_n, self.tile_m)) + cS = cute.domain_offset((cluster_n_block * cluster_tile_n, m_block * self.tile_m), cS) + tScS = thr_mma_S.partition_C(cS) + tScS_idx = thr_copy_t2r.partition_D(tScS) + + apply_score_mod_inner( + tSrS_t2r, + tScS_idx, + self.score_mod, + batch_idx, + head_idx, + softmax_scale, + self.vec_size, + self.qk_acc_dtype, + aux_data, + fastdiv_mods, + seqlen_info, + constant_q_idx=None, + qhead_per_kvhead=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + transpose_indices=True, + ) + + @cute.jit + def apply_score_mod_bwd( + self, + grad_tensor, + score_tensor, + index_tensor, + batch_idx, + head_idx, + softmax_scale, + seqlen_info, + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + ): + """Apply backward score modification (joint graph) for SM100.""" + apply_score_mod_bwd_inner( + grad_tensor, + score_tensor, + index_tensor, + self.score_mod_bwd, + batch_idx, + head_idx, + softmax_scale, + self.vec_size, + self.qk_acc_dtype, + aux_data, + fastdiv_mods, + seqlen_info, + constant_q_idx=None, + qhead_per_kvhead=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + transpose_indices=True, + ) + + @cute.jit + def compute_loop( + self, + thr_mma_S: cute.ThrMma, + thr_mma_dP: cute.ThrMma, + thr_mma_dV: cute.ThrMma, + thr_mma_dK: cute.ThrMma, + tStS: cute.Tensor, + tdPtdP: cute.Tensor, + tdVtdV: cute.Tensor, + tdKtdK: cute.Tensor, + sLSE: cute.Tensor, + sdPsum: cute.Tensor, + mdV: cute.Tensor, + mdK: cute.Tensor, + sdS: cute.Tensor, + sdS_xchg: cute.Tensor, + pipeline_LSE: PipelineAsync, + pipeline_dPsum: PipelineAsync, + pipeline_S_P: PipelineAsync, + pipeline_dS: PipelineAsync, + pipeline_dKV: PipelineAsync, + pipeline_dP: PipelineAsync, + dS_cluster_empty_mbar_ptr: cute.Pointer, + dS_cluster_full_mbar_ptr: cute.Pointer, + dQaccum_empty_mbar_ptr: cute.Pointer, + softmax_scale: cutlass.Float32, + softmax_scale_log2: cutlass.Float32, + block_info: BlockInfo, + SeqlenInfoCls: Callable, + AttentionMaskCls: Callable, + TileSchedulerCls: Callable, + sdV: Optional[cute.Tensor], + sdK: Optional[cute.Tensor], + mdV_tma_tensor: Optional[cute.Tensor], + mdK_tma_tensor: Optional[cute.Tensor], + tma_atom_dV: Optional[cute.CopyAtom], + tma_atom_dK: Optional[cute.CopyAtom], + tiled_copy_r2s_dKV: Optional[cute.TiledCopy], + mdK_semaphore: Optional[cute.Tensor], + mdV_semaphore: Optional[cute.Tensor], + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + blocksparse_tensors: Optional[BlockSparseTensors] = None, + ): + sLSE_2D = cute.make_tensor( + sLSE.iterator, + cute.make_layout( + (self.tile_m, self.tile_n, self.Q_stage), + stride=(1, 0, cute.round_up(self.tile_m, 64)), + ), + ) + sdPsum_2D = cute.make_tensor( + sdPsum.iterator, + cute.make_layout( + (self.tile_m, self.tile_n, self.dO_stage), + stride=(1, 0, cute.round_up(self.tile_m, 64)), + ), + ) + # if const_expr(self.SdP_swapAB): + if const_expr(True): + sLSE_2D = layout_utils.transpose_view(sLSE_2D) + sdPsum_2D = layout_utils.transpose_view(sdPsum_2D) + + # tix: [128...384] 8 warps + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) # 4-11 + tidx = cute.arch.thread_idx()[0] % (cute.arch.WARP_SIZE * len(self.compute_warp_ids)) + # tidx = cute.arch.thread_idx()[0] - (cute.arch.WARP_SIZE * self.compute_warp_ids[0]) + dp_idx = tidx % 128 + num_wg = len(self.compute_warp_ids) // 4 # 2 + # wg_idx: + # 0: [256...384] + # 1: [128...256] + + tileP_f32_like = self.cta_tiler[1] // 32 * self.v_dtype.width + # tStS has shape ((128, 128), 1, 1), tStP has shape ((128, 64), 1, 1) + # tP overlap with tS + tStP = cute.composition(tStS, (cute.make_layout((self.tile_n, tileP_f32_like)), 1, 1)) + tStP = cute.make_tensor(tStS.iterator, tStP.layout) # Otherwise the tmem address is wrong + tScS = thr_mma_S.partition_C(cute.make_identity_tensor(self.mma_tiler_kq[:2])) + tScP = cute.composition(tScS, (cute.make_layout((self.tile_n, tileP_f32_like)), 1, 1)) + # tdS overlap with tdP + tdPtdS = cute.composition(tdPtdP, (cute.make_layout((self.tile_n, tileP_f32_like)), 1, 1)) + tdPcdP = thr_mma_dP.partition_C(cute.make_identity_tensor(self.mma_tiler_vdo[:2])) + tdPcdS = cute.composition(tdPcdP, (cute.make_layout((self.tile_n, tileP_f32_like)), 1, 1)) + + # 2-CTA assumes: repetiton should always be 32 & 16 + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), Float32 + ) + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(16)), Float32 + ) + + # tmem -> rmem + thr_copy_t2r = copy_utils.make_tmem_copy(tmem_load_atom, num_wg).get_slice(tidx) + tStS_t2r = thr_copy_t2r.partition_S(tStS) # (((32, 32), 1), 2, 1, 1) + tdPtdP_t2r = thr_copy_t2r.partition_S(tdPtdP) + tScS_t2r = thr_copy_t2r.partition_D(tScS) # ((32, 1), 2, 1, 1) + t0ScS_t2r = thr_copy_t2r.get_slice(0).partition_D(tScS) # ((32, 1), 2, 1, 1) + # ((32, 1), 2, 1, 1, STAGE) + tSsLSE = thr_copy_t2r.partition_D(thr_mma_S.partition_C(sLSE_2D)) + tSsdPsum = thr_copy_t2r.partition_D(thr_mma_dP.partition_C(sdPsum_2D)) + # rmem -> tmem + thr_copy_r2t = copy_utils.make_tmem_copy(tmem_store_atom, num_wg).get_slice(tidx) + tScP_r2t = thr_copy_r2t.partition_S(tScP) + tStP_r2t = thr_copy_r2t.partition_D(tStP) + tdPcdS_r2t = thr_copy_r2t.partition_S(tdPcdS) + tdPtdS_r2t = thr_copy_r2t.partition_D(tdPtdS) + # rmem -> smem + # This part is a bit iffy, we might be making a lot of assumptions here + copy_atom_r2s = sm100_utils_basic.get_smem_store_op( + LayoutEnum.ROW_MAJOR, self.ds_dtype, Float32, thr_copy_t2r + ) + thr_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, thr_copy_t2r).get_slice(tidx) + + # We assume the swizzle (i.e. layout.inner) stays the same + sdS_epi_layout = sm100_utils_basic.make_smem_layout_epi( + self.ds_dtype, LayoutEnum.ROW_MAJOR, (self.tile_n, self.tile_m), 1 + ) + sdS_layout = cute.slice_(sdS_epi_layout.outer, (None, None, 0)) # ((8,16), (64,2)) + # Need to group into 1 mode to be compatible w thr_copy_r2s + sdS_layout = cute.make_layout((sdS_layout.shape,), stride=(sdS_layout.stride,)) + sdS_epi = cute.make_tensor(sdS.iterator, sdS_layout) + tRS_sdS = thr_copy_r2s.partition_D(sdS_epi) + + if const_expr(self.use_2cta_instrs): + sdS_xchg_epi = cute.make_tensor( + cute.recast_ptr(sdS_xchg.iterator, sdS_epi_layout.inner), sdS_layout + ) + tRS_sdS_xchg = thr_copy_r2s.partition_D(sdS_xchg_epi) + + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + dS_cluster_empty_phase = Int32(1) + # 2-CTA: CTA 0 exchanges stage 1 (bottom half), CTA 1 exchanges stage 0 (top half) + exchange_stage = cta_rank_in_cluster ^ 1 if const_expr(self.use_2cta_instrs) else Int32(0) + + consumer_state_S_P_dP = pipeline.make_pipeline_state( # Our impl has shortcut for stage==1 + cutlass.pipeline.PipelineUserType.Consumer, 1 + ) + # consumer_phase_S_P_dP = Int32(0) + producer_state_dS = pipeline.make_pipeline_state( # Our impl has shortcut for stage==1 + cutlass.pipeline.PipelineUserType.Producer, 1 + ) + consumer_state_dKV = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Consumer, 2 + ) + consumer_state_LSE = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Consumer, self.Q_stage + ) + consumer_state_dPsum = pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Consumer, self.dO_stage + ) + + tile_scheduler = TileSchedulerCls() + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + n_block, head_idx, batch_idx, _ = work_tile.tile_idx + seqlen = SeqlenInfoCls(batch_idx) + m_block_min, m_block_max = block_info.get_m_block_min_max( + seqlen, n_block // self.cluster_shape_mnk[0] + ) + mask = AttentionMaskCls(seqlen) + cluster_n_block = n_block // self.cta_group_size + # TODO: condition mask_seqlen + mask_fn = partial( + mask.apply_mask_sm100_transposed, + tScS_t2r=tScS_t2r, + t0ScS_t2r=t0ScS_t2r, + n_block=cluster_n_block, + mask_seqlen=True, + mask_causal=self.is_causal, + mask_local=self.is_local, + mask_mod=self.mask_mod, + batch_idx=batch_idx, + head_idx=head_idx, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + + # prefetch_LSE = not self.is_causal + prefetch_LSE = False + + curr_q_cnt = Int32(0) + curr_q_idx = None + curr_full_cnt = Int32(0) + curr_full_idx = None + loop_count = m_block_max - m_block_min + process_tile = ( + const_expr(not self.is_local and not self.is_varlen_q) or m_block_min < m_block_max + ) + if const_expr(self.use_block_sparsity): + assert blocksparse_tensors is not None + ( + curr_q_cnt, + curr_q_idx, + curr_full_cnt, + curr_full_idx, + loop_count, + ) = get_block_sparse_iteration_info_bwd( + blocksparse_tensors, + batch_idx, + head_idx, + n_block, + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + ) + process_tile = loop_count > Int32(0) + + # Mainloop + # Block sparsity: iterate over sparse m_block count and derive actual m_block + # from Q_IDX/FULL_Q_IDX tensors. Dense: iterate m_block_min..m_block_max directly. + for iter_idx in cutlass.range(loop_count, unroll=1): + m_block = m_block_min + iter_idx + m_block_oob = False + is_full_block = False + if const_expr(self.use_block_sparsity): + m_block, is_full_block = get_m_block_from_iter_bwd( + iter_idx, + curr_q_cnt, + curr_q_idx, + curr_full_cnt, + curr_full_idx, + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + ) + m_block_oob = m_block >= m_block_max + # Prefetch 1 stage of LSE + pipeline_LSE.consumer_wait(consumer_state_LSE) + tSrLSE_s2r = cute.make_rmem_tensor(tScS_t2r[None, 0, 0, 0].shape, Float32) + if const_expr(prefetch_LSE and not self.shuffle_LSE): + cute.autovec_copy(tSsLSE[None, 0, 0, 0, consumer_state_LSE.index], tSrLSE_s2r) + + pipeline_S_P.consumer_wait(consumer_state_S_P_dP) + # pipeline_S_P.sync_object_full.wait(0, consumer_phase_S_P_dP) + #### TMEM->RMEM (Load S from TMEM) + tSrS_t2r = cute.make_rmem_tensor(tScS_t2r.shape, Float32) + cute.copy(thr_copy_t2r, tStS_t2r, tSrS_t2r) + + if const_expr(self.tile_hdim == 192): + # Signal S tmem load completion using pipeline_S_P when hdim 192 + # dP is overlapped with S + cute.arch.fence_view_async_tmem_load() + with cute.arch.elect_one(): + pipeline_S_P.consumer_release(consumer_state_S_P_dP) + elif const_expr(self.use_2cta_instrs and self.tile_hdim <= 128): + # Signal S tmem load completion using pipeline_dS when 2cta hdim 128 + # dQ is overlapped with S + if iter_idx > 0: + cute.arch.fence_view_async_tmem_load() + with cute.arch.elect_one(): + pipeline_dS.producer_commit(producer_state_dS) + producer_state_dS.advance() + + if const_expr(self.score_mod_bwd is not None): + tSrS_pre = cute.make_fragment_like(tSrS_t2r) + cute.autovec_copy(tSrS_t2r, tSrS_pre) + + if const_expr(self.score_mod is not None): + # Apply score_mod FIRST -> matches forward + self.apply_score_mod( + tSrS_t2r, + thr_copy_t2r, + thr_mma_S, + batch_idx, + head_idx, + m_block, + n_block, + softmax_scale, + seqlen, + aux_data, + fastdiv_mods, + ) + + #### APPLY MASK (after score_mod, matching forward pass order) + check_m_boundary = (m_block + 1) * self.tile_m > seqlen.seqlen_q + mask_fn( + tSrS_t2r, + m_block=m_block, + is_full_block=is_full_block, + check_m_boundary=check_m_boundary, + ) + num_stages = cute.size(tScS_t2r, mode=[1]) + # --------------------------------------------- + #### P = exp(S - LSE) + # --------------------------------------------- + lane_idx = cute.arch.lane_idx() + tSrP_r2t_f32 = cute.make_rmem_tensor(tScP_r2t.shape, Float32) # 64 + tSrP_r2t = cute.recast_tensor(tSrP_r2t_f32, self.q_dtype) + for stage in cutlass.range_constexpr(num_stages): + tSrS_cur = tSrS_t2r[None, stage, 0, 0] + tSsLSE_cur = tSsLSE[None, stage, 0, 0, consumer_state_LSE.index] + if const_expr(not self.shuffle_LSE): + if const_expr(stage > 0 or not prefetch_LSE): + cute.autovec_copy(tSsLSE_cur, tSrLSE_s2r) + tSrLSE = tSrLSE_s2r + else: + tSrLSE = tSsLSE_cur[lane_idx] + for v in cutlass.range_constexpr(cute.size(tSrS_t2r, mode=[0]) // 2): + if const_expr(not self.shuffle_LSE): + lse_pair = (tSrLSE[2 * v], tSrLSE[2 * v + 1]) + else: + lse_pair = ( + utils.shuffle_sync(tSrLSE, offset=2 * v), + utils.shuffle_sync(tSrLSE, offset=2 * v + 1), + ) + tSrS_cur[2 * v], tSrS_cur[2 * v + 1] = cute.arch.fma_packed_f32x2( + ((tSrS_cur[2 * v], tSrS_cur[2 * v + 1])), + (softmax_scale_log2, softmax_scale_log2), + (-lse_pair[0], -lse_pair[1]), + ) + tSrS_cur[2 * v] = cute.math.exp2(tSrS_cur[2 * v], fastmath=True) + tSrS_cur[2 * v + 1] = cute.math.exp2(tSrS_cur[2 * v + 1], fastmath=True) + utils.cvt_f16(tSrS_cur, tSrP_r2t[None, stage, 0, 0]) + if const_expr(stage == 0): + cute.arch.fence_view_async_tmem_load() + # Without this barrier, we could have 1 warp writing to P in tmem while + # another warp is still reading S from tmem. + self.compute_sync_barrier.arrive_and_wait() + cute.copy( + thr_copy_r2t, + tSrP_r2t_f32[None, stage, None, None], + tStP_r2t[None, stage, None, None], + ) + + cute.arch.fence_view_async_tmem_store() + cute.arch.fence_view_async_shared() + self.compute_sync_barrier.arrive_and_wait() + if const_expr(not self.tile_hdim == 192): + # Signal tmem store P completion with pipeline_S_P + with cute.arch.elect_one(): + pipeline_S_P.consumer_release(consumer_state_S_P_dP) + # pipeline_S_P.sync_object_empty.arrive(0, pipeline_S_P.consumer_mask) + # Normally we'd need syncwarp here since only 1 thread will signal in + # consumer_release, but we already have the self.compute_sync_barrier before this + pipeline_LSE.consumer_release(consumer_state_LSE) + consumer_state_LSE.advance() + # --------------------------------------------- + # dS.T = P.T * (dP.T - D) + # --------------------------------------------- + pipeline_dPsum.consumer_wait(consumer_state_dPsum) + pipeline_dP.consumer_wait(consumer_state_S_P_dP) + # pipeline_dP.sync_object_full.wait(0, consumer_phase_S_P_dP) + ### Now delayed to after loop + # consumer_state_S_P_dP.advance() + # consumer_phase_S_P_dP ^= 1 + + ##### dS.T = P.T * (dP.T - Psum) + for stage in cutlass.range_constexpr(num_stages): + tdPrdP_t2r = cute.make_rmem_tensor(tScS_t2r[None, 0, None, None].shape, Float32) + cute.copy(thr_copy_t2r, tdPtdP_t2r[None, stage, None, None], tdPrdP_t2r) + cute.arch.fence_view_async_tmem_load() + self.compute_sync_barrier.arrive_and_wait() + tdPrdP_cur = tdPrdP_t2r[None, 0, 0] + tSrS_cur = tSrS_t2r[None, stage, 0, 0] + tSsdPsum_cur = tSsdPsum[None, stage, 0, 0, consumer_state_dPsum.index] + if const_expr(not self.shuffle_dPsum): + tSrdPsum = cute.make_fragment_like(tSsdPsum_cur, Float32) + cute.autovec_copy(tSsdPsum_cur, tSrdPsum) + else: + tSrdPsum = tSsdPsum_cur[lane_idx] + for v in cutlass.range_constexpr(cute.size(tdPrdP_t2r, mode=[0]) // 2): + if const_expr(not self.shuffle_dPsum): + dPsum_pair = (tSrdPsum[2 * v], tSrdPsum[2 * v + 1]) + else: + dPsum_pair = ( + utils.shuffle_sync(tSrdPsum, offset=2 * v), + utils.shuffle_sync(tSrdPsum, offset=2 * v + 1), + ) + tdPrdP_cur[2 * v], tdPrdP_cur[2 * v + 1] = ( + quack.activation.sub_packed_f32x2( + (tdPrdP_cur[2 * v], tdPrdP_cur[2 * v + 1]), dPsum_pair + ) + ) + tdPrdP_cur[2 * v], tdPrdP_cur[2 * v + 1] = cute.arch.mul_packed_f32x2( + (tSrS_cur[2 * v], tSrS_cur[2 * v + 1]), + (tdPrdP_cur[2 * v], tdPrdP_cur[2 * v + 1]), + ) + + if const_expr(self.score_mod_bwd is not None): + tSrS_pre_cur = tSrS_pre[None, stage, 0, 0] + cluster_tile_n = self.tile_n * self.cta_group_size + cluster_n_block = n_block // self.cta_group_size + cS_bwd = cute.make_identity_tensor((cluster_tile_n, self.tile_m)) + cS_bwd = cute.domain_offset( + (cluster_n_block * cluster_tile_n, m_block * self.tile_m), + cS_bwd, + ) + tScS_bwd = thr_mma_S.partition_C(cS_bwd) + tScS_idx_bwd = thr_copy_t2r.partition_D(tScS_bwd) + tScS_idx_cur = tScS_idx_bwd[None, stage, 0, 0] + self.apply_score_mod_bwd( + tdPrdP_cur, + tSrS_pre_cur, + tScS_idx_cur, + batch_idx, + head_idx, + softmax_scale, + seqlen, + aux_data, + fastdiv_mods, + ) + # Zero out OOB positions (kv_idx >= seqlen_k) after score_mod_bwd + for i in cutlass.range(cute.size(tdPrdP_cur), unroll_full=True): + kv_idx = tScS_idx_cur[i][0] + tdPrdP_cur[i] = 0.0 if kv_idx >= seqlen.seqlen_k else tdPrdP_cur[i] + + tdPrdS_cvt = cute.make_fragment_like(tdPrdP_cur, self.ds_dtype) + utils.cvt_f16(tdPrdP_cur, tdPrdS_cvt) + if const_expr(stage == 0): + pipeline_dS.producer_acquire(producer_state_dS) + if const_expr(self.use_2cta_instrs): + tdPrdS_xchg = cute.make_fragment_like(tdPrdS_cvt, self.ds_dtype) + + # RMEM->TMEM: always write to TMEM for MMA + if const_expr(not self.use_smem_dS_for_mma_dK or self.use_2cta_instrs): + tdPrdS_r2t_f32 = cute.recast_tensor(tdPrdS_cvt, Float32) + cute.copy(thr_copy_r2t, tdPrdS_r2t_f32, tdPtdS_r2t[None, stage, 0, 0]) + + # RMEM->SMEM: For 2-CTA, keep exchange stage in registers, write non-exchange to sdS + if const_expr(self.use_2cta_instrs): + if exchange_stage == stage: + cute.autovec_copy(tdPrdS_cvt, tdPrdS_xchg) + else: + cute.autovec_copy(tdPrdS_cvt, tRS_sdS[None, stage]) + else: + cute.autovec_copy(tdPrdS_cvt, tRS_sdS[None, stage]) + + if const_expr(not self.use_smem_dS_for_mma_dK): + cute.arch.fence_view_async_tmem_store() + + if const_expr(self.use_2cta_instrs): + # use pipeline_dP to signal tmem store of dS + with cute.arch.elect_one(): + pipeline_dP.consumer_release(consumer_state_S_P_dP) + consumer_state_S_P_dP.advance() + + # After the loop: copy exchange registers to sdS_xchg buffer + if const_expr(self.use_2cta_instrs): + # when hdim 192, sdQaccum overlapped with sdS_xchg + if const_expr(self.tile_hdim == 192): + cute.arch.mbarrier_wait( + dQaccum_empty_mbar_ptr, phase=producer_state_dS.phase + ) + cute.autovec_copy(tdPrdS_xchg, tRS_sdS_xchg[None, 0]) + + cute.arch.fence_view_async_shared() + self.compute_sync_barrier.arrive_and_wait() + # Normally we'd need syncwarp here since only 1 thread will signal in + # consumer_release, but we already have the self.compute_sync_barrier before this + pipeline_dPsum.consumer_release(consumer_state_dPsum) + consumer_state_dPsum.advance() + # when 2cta hdim 128, pipeline_dS also signals S tmem load completion so is deferred + if const_expr(not (self.use_2cta_instrs and self.tile_hdim == 128)): + with cute.arch.elect_one(): + pipeline_dS.producer_commit(producer_state_dS) + producer_state_dS.advance() + + # 2-CTA: DSMEM copy from sdS_xchg to peer's sdS buffer + if const_expr(self.use_2cta_instrs): + stage_copy_bytes = const_expr(self.tma_copy_bytes["dS"] // 2) + stage_copy_elems = const_expr(stage_copy_bytes // (self.ds_dtype.width // 8)) + if tidx == 0: + peer_cta_rank_in_cluster = cta_rank_in_cluster ^ 1 + smem_src_ptr = sdS_xchg.iterator + # Destination is peer's sdS at our CTA's offset (exchange_stage position) + smem_dst_ptr = sdS.iterator + cta_rank_in_cluster * stage_copy_elems + cute.arch.mbarrier_arrive_and_expect_tx( + dS_cluster_full_mbar_ptr, + stage_copy_bytes, + peer_cta_rank_in_cluster=peer_cta_rank_in_cluster, + ) + copy_utils.cpasync_bulk_s2cluster( + smem_src_ptr, + smem_dst_ptr, + dS_cluster_full_mbar_ptr, + stage_copy_bytes, + peer_cta_rank_in_cluster=peer_cta_rank_in_cluster, + ) + + # Final signal for dS smem store completion + if const_expr(self.use_2cta_instrs and self.tile_hdim == 128): + if process_tile: + with cute.arch.elect_one(): + pipeline_dS.producer_commit(producer_state_dS) + producer_state_dS.advance() + + # Epilogue + # Run epilogue if we processed any m_blocks for this n_block + if process_tile: + if const_expr(not self.use_tma_store): + consumer_state_dKV = self.epilogue_dKV( + dp_idx, + warp_idx, + batch_idx, + head_idx, + n_block, + seqlen, + thr_mma_dV, + thr_mma_dK, + tdVtdV, + tdKtdK, + mdV, + mdK, + pipeline_dKV, + consumer_state_dKV, + softmax_scale, + ) + else: + thr_copy_r2s_dKV = tiled_copy_r2s_dKV.get_slice(dp_idx) + #### STORE dV + consumer_state_dKV = self.epilogue_dK_or_dV_tma( + dp_idx, + batch_idx, + head_idx, + n_block, + seqlen, + thr_mma_dV, + tdVtdV, + mdV_tma_tensor, + sdV, + tma_atom_dV, + thr_copy_r2s_dKV, + pipeline_dKV, + consumer_state_dKV, + None, # Don't scale + int(NamedBarrierBwdSm100.EpilogueWG1), # barrier_id + mdV_semaphore, + "V", + ) + #### STORE dK + consumer_state_dKV = self.epilogue_dK_or_dV_tma( + dp_idx, + batch_idx, + head_idx, + n_block, + seqlen, + thr_mma_dK, + tdKtdK, + mdK_tma_tensor, + sdK, + tma_atom_dK, + thr_copy_r2s_dKV, + pipeline_dKV, + consumer_state_dKV, + softmax_scale if const_expr(not self.dKV_postprocess) else None, + int(NamedBarrierBwdSm100.EpilogueWG1), # barrier_id + mdK_semaphore, + "K", + ) + # Zero dK/dV for empty tiles (local attention or block sparsity) + # When total_m_block_cnt == 0 for block sparsity, no Q tiles contribute to this KV tile + if const_expr(not self.dKV_postprocess): + should_zero_dKV = False + if const_expr(self.is_local or self.is_varlen_q): + should_zero_dKV = m_block_min >= m_block_max + if const_expr(self.use_block_sparsity): + # For block sparsity, zero when no m_blocks contribute to this n_block + if not process_tile: + should_zero_dKV = True + + if should_zero_dKV: + # For 2-CTA: use cluster-wide tile size (cta_group_size * tile_n) + cluster_tile_n = self.tile_n * self.cta_group_size + n_block_for_tile = n_block // self.cta_group_size + gmem_tiled_copy_zero_dK = copy_utils.tiled_copy_2d( + self.dk_dtype, + math.gcd(64, self.tile_hdim), + 128, # num_threads + ) + gmem_tiled_copy_zero_dV = copy_utils.tiled_copy_2d( + self.dv_dtype, + math.gcd(64, self.tile_hdimv), + 128, # num_threads + ) + gmem_thr_copy_zero_dK = gmem_tiled_copy_zero_dK.get_slice(dp_idx) + gmem_thr_copy_zero_dV = gmem_tiled_copy_zero_dV.get_slice(dp_idx) + mdV_cur = seqlen.offset_batch_K(mdV, batch_idx, dim=3)[None, None, head_idx] + mdK_cur = seqlen.offset_batch_K(mdK, batch_idx, dim=3)[None, None, head_idx] + gdK = cute.local_tile( + mdK_cur, (cluster_tile_n, self.tile_hdim), (n_block_for_tile, 0) + ) + gdV = cute.local_tile( + mdV_cur, (cluster_tile_n, self.tile_hdimv), (n_block_for_tile, 0) + ) + tdKgdK = gmem_thr_copy_zero_dK.partition_D(gdK) + tdVgdV = gmem_thr_copy_zero_dV.partition_D(gdV) + cdK = cute.make_identity_tensor((cluster_tile_n, self.tile_hdim)) + cdV = cute.make_identity_tensor((cluster_tile_n, self.tile_hdimv)) + tdKcdK = gmem_thr_copy_zero_dK.partition_D(cdK) + tdVcdV = gmem_thr_copy_zero_dV.partition_D(cdV) + assert cute.size(tdKgdK[None, 0, 0]) == cute.size(tdVgdV[None, 0, 0]) + zero = cute.make_fragment_like(tdKgdK[None, 0, 0]) + zero.fill(0.0) + if tidx < 128: + for i in cutlass.range_constexpr(tdKgdK.shape[1]): + row_idx = tdKcdK[0, i, 0][0] + if row_idx < seqlen.seqlen_k - cluster_tile_n * n_block_for_tile: + for j in cutlass.range_constexpr(tdKgdK.shape[2]): + cute.copy(gmem_tiled_copy_zero_dK, zero, tdKgdK[None, i, j]) + else: + for i in cutlass.range_constexpr(tdVgdV.shape[1]): + row_idx = tdVcdV[0, i, 0][0] + if row_idx < seqlen.seqlen_k - cluster_tile_n * n_block_for_tile: + for j in cutlass.range_constexpr(tdVgdV.shape[2]): + cute.copy(gmem_tiled_copy_zero_dV, zero, tdVgdV[None, i, j]) + + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + + @cute.jit + def _dq_semaphore_lock_value( + self, + iter_idx: Int32, + curr_q_cnt: Int32, + curr_dq_write_order: Optional[cute.Tensor], + curr_dq_write_order_full: Optional[cute.Tensor], + blocksparse_tensors: Optional[BlockSparseTensors], + block_info: BlockInfo, + seqlen, + m_block: Int32, + n_block: Int32, + ) -> Int32: + lock_value = n_block + if const_expr(self.spt): + n_block_max_for_m_block = block_info.get_n_block_max_for_m_block(seqlen, m_block) + lock_value = n_block_max_for_m_block - 1 - n_block + if const_expr(self.use_block_sparsity): + assert blocksparse_tensors is not None + if const_expr(blocksparse_tensors.dq_write_order is not None): + sparse_iter = iter_idx // self.q_subtile_factor + if sparse_iter < curr_q_cnt: + assert curr_dq_write_order is not None + lock_value = curr_dq_write_order[sparse_iter] + else: + assert curr_dq_write_order_full is not None + lock_value = curr_dq_write_order_full[sparse_iter - curr_q_cnt] + return lock_value + + @cute.jit + def dQacc_reduce( + self, + mdQaccum: cute.Tensor, + sdQaccum: cute.Tensor, + thr_mma_dQ: cute.ThrMma, + tdQtdQ: cute.Tensor, + pipeline_dQ: PipelineAsync, + dQaccum_empty_mbar_ptr: Optional[cute.Pointer], + block_info: BlockInfo, + SeqlenInfoCls: Callable, + TileSchedulerCls: Callable, + mdQ_semaphore: Optional[cute.Tensor], + blocksparse_tensors: Optional[BlockSparseTensors] = None, + ): + num_reduce_threads = cute.arch.WARP_SIZE * len(self.reduce_warp_ids) + tidx = cute.arch.thread_idx()[0] % num_reduce_threads + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx() % len(self.reduce_warp_ids)) + is_tma_warp = warp_idx == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + # TMEM -> RMEM + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.dQ_reduce_ncol_t2r)), Float32 + ) + thr_copy_t2r = tcgen05.make_tmem_copy(tmem_load_atom, tdQtdQ).get_slice(tidx) + tdQtdQ_t2r = thr_copy_t2r.partition_S(tdQtdQ) + tdQcdQ = thr_mma_dQ.partition_C(cute.make_identity_tensor(self.mma_tiler_dsk[:2])) + tdQrdQ_t2r_shape = thr_copy_t2r.partition_D(tdQcdQ).shape + # For 2-CTA: reduce_stage = dQaccum_reduce_stage_t2r / cta_group_size + expected_reduce_stages_t2r = self.dQaccum_reduce_stage_t2r // self.cta_group_size + assert cute.size(tdQrdQ_t2r_shape, mode=[1]) == expected_reduce_stages_t2r, ( + "dQaccum t2r reduce stage mismatch" + ) + expected_reduce_stages = self.dQaccum_reduce_stage // self.cta_group_size + # 2-CTA: CTA 0 -> (M/2, D) (stage 0, 1) & CTA 1 -> (M/2, D) (stage 2, 3) + stage_offset = ( + expected_reduce_stages * cta_rank_in_cluster if const_expr(self.use_2cta_instrs) else 0 + ) + + thr_copy_dQaccum_r2s = copy_utils.tiled_copy_1d( + self.dqaccum_dtype, num_reduce_threads, num_copy_elems=128 // self.dqaccum_dtype.width + ).get_slice(tidx) + tdQsdQ = thr_copy_dQaccum_r2s.partition_D(sdQaccum) + + read_flag = const_expr(not self.deterministic) + + tile_scheduler = TileSchedulerCls() + work_tile = tile_scheduler.initial_work_tile_info() + dQ_consumer_state = pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Consumer, 1 + ) + dQ_tma_store_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.sdQaccum_stage + ) + while work_tile.is_valid_tile: + n_block, head_idx, batch_idx, _ = work_tile.tile_idx + n_block_cta_group = n_block // self.cta_group_size # for 2cta + seqlen = SeqlenInfoCls(batch_idx) + m_block_min, m_block_max = block_info.get_m_block_min_max(seqlen, n_block_cta_group) + if const_expr(not seqlen.has_cu_seqlens_q): + mdQaccum_cur = mdQaccum[None, head_idx, batch_idx] + else: + mdQaccum_cur = cute.domain_offset( + (seqlen.padded_offset_q * self.tile_hdim,), mdQaccum[None, head_idx] + ) + gdQaccum_ = cute.local_tile(mdQaccum_cur, (self.tile_m * self.tile_hdim,), (None,)) + # (M * K / STAGE, STAGE, _) + gdQaccum = cute.flat_divide( + gdQaccum_, (self.tile_m * self.tile_hdim // self.dQaccum_reduce_stage,) + ) + + if const_expr(self.deterministic): + assert mdQ_semaphore is not None + mdQ_semaphore_cur = mdQ_semaphore[None, None, head_idx, batch_idx] + + delay_semaphore_release = not self.tile_hdim == 192 and not self.use_block_sparsity + + curr_q_cnt = Int32(0) + curr_q_idx = None + curr_full_cnt = Int32(0) + curr_full_idx = None + curr_dq_write_order = None + curr_dq_write_order_full = None + loop_count = m_block_max - m_block_min + process_tile = ( + const_expr(not self.is_local and not self.is_varlen_q) or m_block_min < m_block_max + ) + if const_expr(self.use_block_sparsity): + assert blocksparse_tensors is not None + ( + curr_q_cnt, + curr_q_idx, + curr_full_cnt, + curr_full_idx, + loop_count, + ) = get_block_sparse_iteration_info_bwd( + blocksparse_tensors, + batch_idx, + head_idx, + n_block, + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + ) + process_tile = loop_count > Int32(0) + if const_expr(self.deterministic and self.use_block_sparsity): + assert blocksparse_tensors is not None + if const_expr(blocksparse_tensors.dq_write_order is not None): + assert blocksparse_tensors.dq_write_order is not None + curr_dq_write_order = blocksparse_tensors.dq_write_order[ + batch_idx, head_idx, n_block, None + ] + if const_expr(blocksparse_tensors.dq_write_order_full is not None): + assert blocksparse_tensors.dq_write_order_full is not None + curr_dq_write_order_full = blocksparse_tensors.dq_write_order_full[ + batch_idx, head_idx, n_block, None + ] + + # dQacc_reduce mainloop + # Block sparsity: iterate over sparse m_block count and derive actual m_block + # from Q_IDX/FULL_Q_IDX tensors. Dense: iterate m_block_min..m_block_max directly. + for iter_idx in cutlass.range(loop_count, unroll=1): + m_block = m_block_min + iter_idx + m_block_oob_upper = False + if const_expr(self.use_block_sparsity): + m_block, _ = get_m_block_from_iter_bwd( + iter_idx, + curr_q_cnt, + curr_q_idx, + curr_full_cnt, + curr_full_idx, + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + ) + m_block_oob_upper = m_block >= m_block_max + pipeline_dQ.consumer_wait(dQ_consumer_state) + # TMEM -> RMEM + tdQrdQ_t2r = cute.make_rmem_tensor(tdQrdQ_t2r_shape, Float32) + cute.copy(thr_copy_t2r, tdQtdQ_t2r, tdQrdQ_t2r) + cute.arch.fence_view_async_tmem_load() + cute.arch.sync_warp() + with cute.arch.elect_one(): + pipeline_dQ.consumer_release(dQ_consumer_state) + dQ_consumer_state.advance() + + if m_block_max > 0: + m_block = cutlass.min(m_block, m_block_max - 1) + gdQaccum_cur = gdQaccum[None, None, m_block] + + tdQrdQ_shape = ( + self.dQ_reduce_ncol, + self.tile_hdim // self.cta_group_size // self.dQ_reduce_ncol, + ) + tdQrdQ = cute.make_tensor(tdQrdQ_t2r.iterator, tdQrdQ_shape) + + for stage in cutlass.range_constexpr(cute.size(tdQrdQ, mode=[1])): + smem_idx = dQ_tma_store_producer_state.index + tdQsdQ_r2s = tdQsdQ[None, None, smem_idx] + tdQrdQ_r2s = cute.make_tensor(tdQrdQ[None, stage].iterator, tdQsdQ_r2s.shape) + cute.copy(thr_copy_dQaccum_r2s, tdQrdQ_r2s, tdQsdQ_r2s) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_view_async_shared() + # semaphore acquire + if const_expr(self.deterministic and stage == 0): + if not m_block_oob_upper: + lock_value = self._dq_semaphore_lock_value( + iter_idx, + curr_q_cnt, + curr_dq_write_order, + curr_dq_write_order_full, + blocksparse_tensors, + block_info, + seqlen, + m_block, + n_block_cta_group, + ) + barrier.wait_eq( + mdQ_semaphore_cur[(m_block, None)].iterator, + tidx, + cta_rank_in_cluster, + lock_value, + ) + self.reduce_sync_barrier.arrive_and_wait() + # Copy from shared memory to global memory + if is_tma_warp and not m_block_oob_upper: + with cute.arch.elect_one(): + copy_utils.cpasync_reduce_bulk_add_f32( + sdQaccum[None, smem_idx].iterator, + gdQaccum_cur[None, stage + stage_offset].iterator, + self.tma_copy_bytes["dQ"] // 1, + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(self.sdQaccum_stage - 1, read=read_flag) + elif is_tma_warp: + # Drain pending TMA stores so SMEM buffers are safe to reuse + cute.arch.cp_async_bulk_wait_group(0, read=read_flag) + self.reduce_sync_barrier.arrive_and_wait() + dQ_tma_store_producer_state.advance() + + if const_expr(self.deterministic and stage == 0 and delay_semaphore_release): + if m_block > m_block_min: + barrier.arrive_inc( + mdQ_semaphore_cur[(m_block - 1, None)].iterator, + tidx, + cta_rank_in_cluster, + 1, + ) + + if const_expr(self.tile_hdim == 192): + if const_expr(self.sdQaccum_stage > 1): + if is_tma_warp: + cute.arch.cp_async_bulk_wait_group(0, read=read_flag) + self.reduce_sync_barrier.arrive_and_wait() + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(dQaccum_empty_mbar_ptr) + + # semaphore release + # NOTE: arrive_inc calls red_release which issues membar + if const_expr(self.deterministic and not delay_semaphore_release): + if const_expr(self.sdQaccum_stage > 1 and not self.tile_hdim == 192): + if is_tma_warp and not m_block_oob_upper: + cute.arch.cp_async_bulk_wait_group(0, read=read_flag) + self.reduce_sync_barrier.arrive_and_wait() + if not m_block_oob_upper: + barrier.arrive_inc( + mdQ_semaphore_cur[m_block, None].iterator, tidx, cta_rank_in_cluster, 1 + ) + + if process_tile: + if is_tma_warp: + cute.arch.cp_async_bulk_wait_group(0, read=read_flag) + self.reduce_sync_barrier.arrive_and_wait() + # final semaphore release + if const_expr(self.deterministic and delay_semaphore_release): + barrier.arrive_inc( + mdQ_semaphore_cur[(m_block_max - 1, None)].iterator, + tidx, + cta_rank_in_cluster, + 1, + ) + + if const_expr( + self.deterministic + and not self.spt + and not self.use_block_sparsity + and block_info.window_size_left is not None + ): + m_block_global_max = cute.ceil_div(seqlen.seqlen_q, self.tile_m) + for m_block in cutlass.range(m_block_max, m_block_global_max, unroll=1): + barrier.arrive_inc( + mdQ_semaphore_cur[(m_block, None)].iterator, tidx, cta_rank_in_cluster, 1 + ) + + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + + if const_expr(not self.deterministic): + cute.arch.cp_async_bulk_wait_group(0, read=True) + + @cute.jit + def epilogue_dKV( + self, + tidx: Int32, + warp_idx: Int32, + batch_idx: Int32, + head_idx: Int32, + n_block: Int32, + seqlen, + thr_mma_dV: cute.ThrMma, + thr_mma_dK: cute.ThrMma, + tdVtdV: cute.Tensor, + tdKtdK: cute.Tensor, + mdV: cute.Tensor, + mdK: cute.Tensor, + pipeline_dKV: PipelineAsync, + consumer_state_dKV: cutlass.pipeline.PipelineState, + softmax_scale: Float32, + ): + wg_idx = ( + cute.arch.thread_idx()[0] % (cute.arch.WARP_SIZE * len(self.compute_warp_ids)) + ) // 128 + num_wg = cute.arch.WARP_SIZE * len(self.compute_warp_ids) // 128 + + assert self.qhead_per_kvhead == 1, "This epilogue path is only for MHA" + mdV_cur = seqlen.offset_batch_K(mdV, batch_idx, dim=3)[None, None, head_idx] + mdK_cur = seqlen.offset_batch_K(mdK, batch_idx, dim=3)[None, None, head_idx] + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(16)), Float32 + ) + # dV + pipeline_dKV.consumer_wait(consumer_state_dKV) + + tiled_tmem_ld_dV = tcgen05.make_tmem_copy(tmem_load_atom, tdVtdV) + thr_tmem_ld_dV = tiled_tmem_ld_dV.get_slice(tidx) + + tdVtdV_t2r_p = thr_tmem_ld_dV.partition_S(tdVtdV) + tdVtdV_t2r = self.split_wg(tdVtdV_t2r_p, wg_idx, num_wg) + + cdV = cute.make_identity_tensor((self.mma_tiler_pdo[0], self.mma_tiler_pdo[1])) + tdVcdV = thr_mma_dV.partition_C(cdV) + tdVcdV_tensor = cute.make_tensor(tdVcdV.iterator, tdVcdV.layout) + + tdVcdV_t2r_p = thr_tmem_ld_dV.partition_D(tdVcdV_tensor) + tdVcdV_t2r = self.split_wg(tdVcdV_t2r_p, wg_idx, num_wg) + tdVrdV_t2r = cute.make_rmem_tensor(tdVcdV_t2r.shape, Float32) + + cute.copy(thr_tmem_ld_dV, tdVtdV_t2r, tdVrdV_t2r) + cute.arch.fence_view_async_tmem_load() + + universal_copy_bits = 128 + atom_universal_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.dv_dtype, + num_bits_per_copy=universal_copy_bits, + ) + tiled_gmem_store_dV = cute.make_tiled_copy( + atom_universal_copy, + layout_tv=tiled_tmem_ld_dV.layout_dst_tv_tiled, + tiler_mn=tiled_tmem_ld_dV.tiler_mn, + ) + + tdVrdV_r2s = cute.make_rmem_tensor(tdVrdV_t2r.shape, self.dv_dtype) + for i in cutlass.range_constexpr(cute.size(tdVrdV_t2r, mode=[1])): + dV_vec = tdVrdV_t2r[(None, i, 0, 0)].load() + tdVrdV_r2s[(None, i, 0, 0)].store(dV_vec.to(self.dv_dtype)) + + gdV = cute.local_tile(mdV_cur, (self.mma_tiler_pdo[0], self.tile_hdimv), (None, 0)) + gdV_tile = gdV[None, None, n_block // self.cta_group_size] + + tdVgdV = thr_mma_dV.partition_C(gdV_tile) + tdVgdV_r2g_p = thr_tmem_ld_dV.partition_D(tdVgdV) + tdVgdV_r2g = self.split_wg(tdVgdV_r2g_p, wg_idx, num_wg) + + if tidx < seqlen.seqlen_k - self.tile_n * n_block: + cute.copy(tiled_gmem_store_dV, tdVrdV_r2s, tdVgdV_r2g) + + cute.arch.sync_warp() + with cute.arch.elect_one(): + pipeline_dKV.consumer_release(consumer_state_dKV) + consumer_state_dKV.advance() + + # dK + pipeline_dKV.consumer_wait(consumer_state_dKV) + + tiled_tmem_ld_dK = tcgen05.make_tmem_copy(tmem_load_atom, tdKtdK) + thr_tmem_ld_dK = tiled_tmem_ld_dK.get_slice(tidx) + + tdKtdK_t2r_p = thr_tmem_ld_dK.partition_S(tdKtdK) + tdKtdK_t2r = self.split_wg(tdKtdK_t2r_p, wg_idx, num_wg) + + cdK = cute.make_identity_tensor((self.mma_tiler_dsq[0], self.mma_tiler_dsq[1])) + tdKcdK = thr_mma_dK.partition_C(cdK) + tdKcdK_tensor = cute.make_tensor(tdKcdK.iterator, tdKcdK.layout) + + tdKcdK_t2r_p = thr_tmem_ld_dK.partition_D(tdKcdK_tensor) + tdKcdK_t2r = self.split_wg(tdKcdK_t2r_p, wg_idx, num_wg) + tdKrdK_t2r = cute.make_rmem_tensor(tdKcdK_t2r.shape, Float32) + + cute.copy(tiled_tmem_ld_dK, tdKtdK_t2r, tdKrdK_t2r) + cute.arch.fence_view_async_tmem_load() + + universal_copy_bits = 128 + atom_universal_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.dk_dtype, + num_bits_per_copy=universal_copy_bits, + ) + + tiled_gmem_store_dK = cute.make_tiled_copy( + atom_universal_copy, + layout_tv=tiled_tmem_ld_dK.layout_dst_tv_tiled, + tiler_mn=tiled_tmem_ld_dK.tiler_mn, + ) + + tdKrdK_r2s = cute.make_rmem_tensor(tdKrdK_t2r.shape, self.dk_dtype) + + for i in cutlass.range_constexpr(cute.size(tdKrdK_t2r, mode=[1])): + dK_vec = tdKrdK_t2r[(None, i, 0, 0)].load() * softmax_scale + tdKrdK_r2s[(None, i, 0, 0)].store(dK_vec.to(self.dk_dtype)) + + gdK = cute.local_tile(mdK_cur, (self.mma_tiler_dsq[0], self.tile_hdim), (None, 0)) + gdK_tile = gdK[None, None, n_block // self.cta_group_size] + + tdKgdK = thr_mma_dK.partition_C(gdK_tile) + tdKgdK_r2g_p = thr_tmem_ld_dK.partition_D(tdKgdK) + tdKgdK_r2g = self.split_wg(tdKgdK_r2g_p, wg_idx, num_wg) + + if tidx < seqlen.seqlen_k - self.tile_n * n_block: + cute.copy(tiled_gmem_store_dK, tdKrdK_r2s, tdKgdK_r2g) + + cute.arch.sync_warp() + with cute.arch.elect_one(): + pipeline_dKV.consumer_release(consumer_state_dKV) + return consumer_state_dKV + + @cute.jit + def epilogue_dK_or_dV_tma( + self, + tidx: Int32, + batch_idx: Int32, + head_idx: Int32, + n_block: Int32, + seqlen, + thr_mma: cute.ThrMma, + tdKVtdKV: cute.Tensor, + mdKV: cute.Tensor, + sdKV: cute.Tensor, + tma_atom_dKV: cute.CopyAtom, + thr_copy_r2s_dKV: cute.TiledCopy, + pipeline_dKV: PipelineAsync, + consumer_state_dKV: cutlass.pipeline.PipelineState, + scale: Optional[Float32], + barrier_id: Int32, + mdKV_semaphore: Optional[cute.Tensor], + K_or_V: cutlass.Constexpr[str], + ) -> cutlass.pipeline.PipelineState: + assert K_or_V in ("K", "V") + tile_hdim = self.tile_hdim if const_expr(K_or_V == "K") else self.tile_hdimv + dtype = self.dk_dtype if const_expr(K_or_V == "K") else self.dv_dtype + epi_tile = self.sdK_epi_tile if const_expr(K_or_V == "K") else self.sdV_epi_tile + flat_epi_tile = ( + self.sdK_flat_epi_tile if const_expr(K_or_V == "K") else self.sdV_flat_epi_tile + ) + num_compute_threads = cute.arch.WARP_SIZE * len(self.compute_warp_ids) + wg_idx = (cute.arch.thread_idx()[0] % num_compute_threads) // 128 + num_wg = num_compute_threads // 128 + leader_warp = (cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4) == 0 + + cta_group_tile_n = const_expr(self.tile_n * self.cta_group_size) + + if const_expr(not self.dKV_postprocess): + sdKV = sdKV[None, None, wg_idx] # (tile_n, 64) for bf16 + else: + sdKV = sdKV[None, wg_idx] # (tile_n * 32) for fp32 + + # (8, tile_n / 128, 64 / 8) = (8, 1, 8) or (4, tile_n * 32 / (128 * 4)) = (4, 8) + tdKVsdKV_r2s = thr_copy_r2s_dKV.partition_D(sdKV) + + head_idx_kv = head_idx // self.qhead_per_kvhead + if const_expr(not self.dKV_postprocess): + assert not seqlen.has_cu_seqlens_k, "varlen uses non tma store path" + mdKV_cur = mdKV[None, None, head_idx_kv, batch_idx] # (seqlen, hdim) + gdKV_p = cute.local_tile( + mdKV_cur, (self.tile_n, tile_hdim), (n_block, 0) + ) # (tile_n, hdim) - per CTA + gdKV = self.split_wg(gdKV_p, wg_idx, num_wg) # (tile_n, hdim / 2) + gdKV_epi = cute.local_tile( + gdKV, epi_tile, (0, None) + ) # (tile_n, 64, epi_stage = (hdim / 2) / 64) + else: + # n_block_group = n_block // self.cta_group_size + if const_expr(not seqlen.has_cu_seqlens_k): + mdKV_cur = mdKV[None, head_idx_kv, batch_idx] # (seqlen * hdim) + else: + mdKV_cur = cute.domain_offset( + (seqlen.padded_offset_k * tile_hdim,), mdKV[None, head_idx_kv] + ) + gdKV_p = cute.local_tile( + mdKV_cur, (self.tile_n * tile_hdim,), (n_block,) + ) # (tile_n * hdim) + gdKV = cute.logical_divide(gdKV_p, (self.tile_n * tile_hdim // num_wg,))[ + ((None, wg_idx),) + ] # (tile_n * hdim / 2) + gdKV_epi = cute.flat_divide( + gdKV, (flat_epi_tile,) + ) # (tile_n * hdim / 2 / epi_stage, epi_stage) + + deterministic_KV = self.deterministic and self.qhead_per_kvhead > 1 + if const_expr(deterministic_KV): + assert mdKV_semaphore is not None + mdKV_semaphore_cur = mdKV_semaphore[n_block, None, head_idx_kv, batch_idx] + + if const_expr(not self.dKV_postprocess): + tdKVsdKV, tdKVgdKV = cpasync.tma_partition( + tma_atom_dKV, + 0, # no multicast + cute.make_layout(1), + cute.group_modes(sdKV, 0, 2), + cute.group_modes(gdKV_epi, 0, 2), + ) # (TMA) and (TMA, EPI_STAGE) + assert len(tdKVsdKV.shape) == 1, "Wrong rank for SMEM fragment tdKVsdKV" + assert len(tdKVgdKV.shape) == 2, "Wrong rank for GMEM fragment tdKVgdKV" + num_epi_stages = cute.size(tdKVgdKV.shape[1]) + if const_expr(K_or_V == "K"): + assert num_epi_stages == self.num_epi_stages, "Epi stage calculation is wrong (K)" + else: + assert num_epi_stages == self.num_epi_stages_v, "Epi stage calculation is wrong (V)" + else: + num_epi_stages = ( + self.num_epi_stages if const_expr(K_or_V == "K") else self.num_epi_stages_v + ) + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.dK_reduce_ncol)), Float32 + ) + + read_flag = const_expr(not deterministic_KV) + + pipeline_dKV.consumer_wait(consumer_state_dKV) + + # semaphore acquire + if const_expr(deterministic_KV): + barrier.wait_eq( + mdKV_semaphore_cur.iterator, tidx, wg_idx, head_idx % self.qhead_per_kvhead + ) + cute.arch.barrier(barrier_id=barrier_id + wg_idx, number_of_threads=128) + + for epi_stage in cutlass.range_constexpr(num_epi_stages): + # TMEM -> RMEM -- setup + thr_copy_t2r = tcgen05.make_tmem_copy(tmem_load_atom, tdKVtdKV).get_slice(tidx) + tdKVtdKV_t2r_p = thr_copy_t2r.partition_S(tdKVtdKV) + tdKVtdKV_t2r = self.split_wg(tdKVtdKV_t2r_p, wg_idx, num_wg)[None, None, 0, 0] + if const_expr(num_epi_stages > 1): + tdKVtdKV_t2r = tdKVtdKV_t2r[None, epi_stage] + + cdKV = cute.make_identity_tensor((cta_group_tile_n, tile_hdim)) + tdKVcdKV = thr_mma.partition_C(cdKV) + tdKVcdKV_t2r_p = thr_copy_t2r.partition_D(tdKVcdKV) + tdKVcdKV_t2r = self.split_wg(tdKVcdKV_t2r_p, wg_idx, num_wg)[None, None, 0, 0] + if const_expr(num_epi_stages > 1): + tdKVcdKV_t2r = tdKVcdKV_t2r[None, epi_stage] + + tdKVrdKV_t2r = cute.make_rmem_tensor(tdKVcdKV_t2r.shape, Float32) + + assert cute.size(tdKVrdKV_t2r) == cute.size(tdKVtdKV_t2r) // cute.arch.WARP_SIZE, ( + "RMEM<->TMEM fragment size mismatch" + ) + + # TMEM -> RMEM -- copy and fence + cute.copy(thr_copy_t2r, tdKVtdKV_t2r, tdKVrdKV_t2r) + cute.arch.fence_view_async_tmem_load() + + # RMEM -- scale and convert + if const_expr(scale is not None): + for i in cutlass.range(cute.size(tdKVrdKV_t2r.shape) // 2, unroll_full=True): + tdKVrdKV_t2r[2 * i], tdKVrdKV_t2r[2 * i + 1] = cute.arch.mul_packed_f32x2( + (tdKVrdKV_t2r[2 * i], tdKVrdKV_t2r[2 * i + 1]), (scale, scale) + ) + tdKVrdKV = cute.make_rmem_tensor(tdKVrdKV_t2r.shape, dtype) # (32 columns) + tdKVrdKV.store(tdKVrdKV_t2r.load().to(dtype)) + + # RMEM -> SMEM -- copy, fence and barrier + tdKVrdKV_r2s = cute.make_tensor(tdKVrdKV.iterator, tdKVsdKV_r2s.shape) + cute.copy(thr_copy_r2s_dKV, tdKVrdKV_r2s, tdKVsdKV_r2s) + cute.arch.fence_view_async_shared() + cute.arch.barrier(barrier_id=barrier_id + wg_idx, number_of_threads=128) + + # SMEM -> GMEM + if leader_warp: + if const_expr(not self.dKV_postprocess): + cute.copy(tma_atom_dKV, tdKVsdKV, tdKVgdKV[None, epi_stage]) + else: + with cute.arch.elect_one(): + copy_utils.cpasync_reduce_bulk_add_f32( + sdKV.iterator, + gdKV_epi[None, epi_stage].iterator, + self.tma_copy_bytes["dKacc"], + ) + if const_expr(epi_stage < num_epi_stages - 1): + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=read_flag) + cute.arch.barrier_arrive( + barrier_id=barrier_id + wg_idx, number_of_threads=128 + cute.arch.WARP_SIZE + ) + + # Barrier since all warps need to wait for SMEM to be freed + cute.arch.fence_view_async_shared() + cute.arch.barrier( + barrier_id=barrier_id + wg_idx, number_of_threads=128 + cute.arch.WARP_SIZE + ) + + # semaphore release + # NOTE: arrive_inc calls red_release which issues membar + if const_expr(deterministic_KV): + if leader_warp: + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=read_flag) + cute.arch.barrier(barrier_id=barrier_id + wg_idx, number_of_threads=128) + barrier.arrive_inc(mdKV_semaphore_cur.iterator, tidx, wg_idx, 1) + + cute.arch.sync_warp() + with cute.arch.elect_one(): + pipeline_dKV.consumer_release(consumer_state_dKV) + consumer_state_dKV.advance() + return consumer_state_dKV diff --git a/fastvideo-kernel/fa4/flash_bwd_sm120.py b/fastvideo-kernel/fa4/flash_bwd_sm120.py new file mode 100644 index 000000000..556c59e38 --- /dev/null +++ b/fastvideo-kernel/fa4/flash_bwd_sm120.py @@ -0,0 +1,55 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +# SM120 (Blackwell GeForce / DGX Spark) backward pass. +# +# SM120 uses the same SM80-era MMA instructions (mma.sync.aligned.m16n8k16) but has +# a smaller shared memory capacity (99 KB vs 163 KB on SM80). This module subclasses +# FlashAttentionBackwardSm80 and overrides the SMEM capacity check accordingly. + +import cutlass +import cutlass.utils as utils_basic + +from flash_attn.cute.flash_bwd import FlashAttentionBackwardSm80 + + +class FlashAttentionBackwardSm120(FlashAttentionBackwardSm80): + @staticmethod + def can_implement( + dtype, + head_dim, + head_dim_v, + m_block_size, + n_block_size, + num_stages_Q, + num_stages_dO, + num_threads, + is_causal, + V_in_regs=False, + ) -> bool: + """Check if the kernel can be implemented on SM120. + + Same logic as SM80 but uses SM120's shared memory capacity (99 KB). + """ + if dtype not in [cutlass.Float16, cutlass.BFloat16]: + return False + if head_dim % 8 != 0: + return False + if head_dim_v % 8 != 0: + return False + if n_block_size % 16 != 0: + return False + if num_threads % 32 != 0: + return False + # Shared memory usage: Q tile + dO tile + K tile + V tile + smem_usage_Q = m_block_size * head_dim * num_stages_Q * 2 + smem_usage_dO = m_block_size * head_dim_v * num_stages_dO * 2 + smem_usage_K = n_block_size * head_dim * 2 + smem_usage_V = n_block_size * head_dim_v * 2 + smem_usage_QV = ( + (smem_usage_Q + smem_usage_V) if not V_in_regs else max(smem_usage_Q, smem_usage_V) + ) + smem_usage = smem_usage_QV + smem_usage_dO + smem_usage_K + # SM120 has 99 KB shared memory (vs 163 KB on SM80) + smem_capacity = utils_basic.get_smem_capacity_in_bytes("sm_120") + if smem_usage > smem_capacity: + return False + return True diff --git a/fastvideo-kernel/fa4/flash_bwd_sm90.py b/fastvideo-kernel/fa4/flash_bwd_sm90.py new file mode 100644 index 000000000..7dbe85bee --- /dev/null +++ b/fastvideo-kernel/fa4/flash_bwd_sm90.py @@ -0,0 +1,1926 @@ +import math +from typing import Callable, Optional, Type +from functools import partial + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.utils.hopper_helpers as sm90_utils_basic +from cutlass.cute.nvgpu import cpasync, warpgroup +from cutlass.cute import FastDivmodDivisor +from cutlass import Float32, Int32, Boolean, const_expr +from cutlass.utils import LayoutEnum + +from quack import copy_utils +from quack import layout_utils +from quack import sm90_utils +from quack.sm90_utils import gemm_zero_init, gemm_w_idx + +from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned +from flash_attn.cute import utils +from flash_attn.cute.mask import AttentionMask +from flash_attn.cute.seqlen_info import SeqlenInfoQK +from flash_attn.cute.block_info import BlockInfo +from flash_attn.cute import pipeline +from quack.cute_dsl_utils import ParamsBase +from flash_attn.cute.tile_scheduler import ( + TileSchedulerArguments, + SingleTileScheduler, + SingleTileLPTBwdScheduler, + SingleTileVarlenScheduler, +) +from flash_attn.cute import barrier +from flash_attn.cute.named_barrier import NamedBarrierBwd +from flash_attn.cute.softmax import apply_score_mod_inner, apply_score_mod_bwd_inner +from flash_attn.cute.block_sparsity import BlockSparseTensors +from flash_attn.cute.utils import AuxData +from flash_attn.cute.block_sparse_utils import ( + get_total_q_block_count_bwd, + produce_block_sparse_q_loads_bwd_sm90, + consume_block_sparse_mma_bwd_sm90, + dQaccum_store_block_sparse_bwd_sm90, +) + + +class FlashAttentionBackwardSm90: + arch = 90 + + def __init__( + self, + dtype: Type[cutlass.Numeric], + head_dim: int, + head_dim_v: Optional[int] = None, + qhead_per_kvhead: int = 1, + is_causal: bool = False, + is_local: bool = False, + deterministic: bool = False, + tile_m: int = 64, + tile_n: int = 128, + Q_stage: int = 2, + dO_stage: int = 2, + PdS_stage: int = 2, + SdP_swapAB: bool = False, + dKV_swapAB: bool = False, + dQ_swapAB: bool = False, + AtomLayoutMSdP: int = 1, + AtomLayoutNdKV: int = 2, + AtomLayoutMdQ: int = 1, + num_threads: int = 384, + V_in_regs: bool = False, + score_mod: cutlass.Constexpr | None = None, + score_mod_bwd: cutlass.Constexpr | None = None, + mask_mod: cutlass.Constexpr | None = None, + has_aux_tensors: cutlass.Constexpr = False, + q_subtile_factor: cutlass.Constexpr[int] = 1, + dQ_single_wg: bool = False, + ): + self.dtype = dtype + # padding head_dim to a multiple of 16 as k_block_size + hdim_multiple_of = 16 + self.tile_hdim = int(math.ceil(head_dim / hdim_multiple_of) * hdim_multiple_of) + head_dim_v = head_dim_v if head_dim_v is not None else head_dim + self.same_hdim_kv = head_dim == head_dim_v + self.tile_hdimv = int(math.ceil(head_dim_v / hdim_multiple_of) * hdim_multiple_of) + # Can save registers (and hence be faster) if we don't have to check hdim predication + self.check_hdim_oob = head_dim != self.tile_hdim + self.check_hdim_v_oob = head_dim_v != self.tile_hdimv + self.qhead_per_kvhead = qhead_per_kvhead + self.is_causal = is_causal + self.is_local = is_local + self.deterministic = deterministic + self.tile_m = tile_m + self.tile_n = tile_n + self.num_threads = num_threads + self.Q_stage = Q_stage + self.dO_stage = dO_stage + self.PdS_stage = PdS_stage + assert self.dO_stage in [1, self.Q_stage] + assert self.PdS_stage in [1, self.Q_stage] + self.SdP_swapAB = SdP_swapAB + self.dKV_swapAB = dKV_swapAB + self.dQ_swapAB = dQ_swapAB + self.AtomLayoutMSdP = AtomLayoutMSdP + self.AtomLayoutNdKV = AtomLayoutNdKV + self.AtomLayoutMdQ = AtomLayoutMdQ + self.num_wg_mma = (self.num_threads // 128) - 1 + self.mma_dkv_is_rs = ( + AtomLayoutMSdP == 1 + and AtomLayoutNdKV == self.num_wg_mma + and SdP_swapAB + and not dKV_swapAB + ) + self.V_in_regs = V_in_regs + # May be overridden in __call__ for varlen inputs. + if qhead_per_kvhead > 1: + assert self.same_hdim_kv, "GQA backward requires head_dim == head_dim_v" + assert self.num_wg_mma == 2, "GQA backward assumes 2 warp groups" + # These are tuned for speed + # Do we keep the LSE and dPsum in each thread, or split them across 8 threads that share + # them and then shuffle to get the value whenever we need? This can reduce register + # pressure when SdP_swapAB, where each thread needs to keep statistics for (kBlockM / 4) + # rows. If !SdP_swapAB, each thread only needs to keep statistics for 2 rows. + self.shuffle_LSE = self.SdP_swapAB and self.tile_hdim <= 64 + self.shuffle_dPsum = self.SdP_swapAB and self.tile_hdim <= 64 + + self.buffer_align_bytes = 1024 + + self.score_mod = score_mod + self.score_mod_bwd = score_mod_bwd + self.mask_mod = mask_mod + self.has_aux_tensors = has_aux_tensors + self.q_subtile_factor = q_subtile_factor + if cutlass.const_expr(has_aux_tensors): + self.vec_size: cutlass.Constexpr = 1 + else: + self.vec_size: cutlass.Constexpr = 4 + self.qk_acc_dtype = Float32 + # dQ_single_wg: WG0 computes the full dQ GEMM, WG1 skips it. + # Only valid for 2 MMA warp groups. + # Credit: Ben Spector + if dQ_single_wg: + assert self.num_wg_mma == 2, "dQ_single_wg only supports 2 warp groups" + self.num_wg_dQ = 1 if dQ_single_wg else self.num_wg_mma + + @staticmethod + def can_implement( + dtype, + head_dim, + head_dim_v, + tile_m, + tile_n, + Q_stage, + num_threads, + V_in_regs=False, + ) -> bool: + if dtype not in [cutlass.Float16, cutlass.BFloat16]: + return False + if head_dim % 8 != 0: + return False + if head_dim_v % 8 != 0: + return False + if tile_n % 16 != 0: + return False + if num_threads % 32 != 0: + return False + if (tile_m * 2) % num_threads != 0: + return False + return True + + def _check_type( + self, + mQ_type: Type[cutlass.Numeric], + mK_type: Type[cutlass.Numeric], + mV_type: Type[cutlass.Numeric], + mdO_type: Type[cutlass.Numeric], + mLSE_type: Type[cutlass.Numeric], + mdPsum_type: Type[cutlass.Numeric], + mdQaccum_type: Type[cutlass.Numeric], + mdK_type: Type[cutlass.Numeric], + mdV_type: Type[cutlass.Numeric], + ): + # Get the data type and check if it is fp16 or bf16 + if const_expr(not (mQ_type == mK_type == mV_type == mdO_type)): + raise TypeError("All tensors must have the same data type") + if const_expr(mQ_type not in [cutlass.Float16, cutlass.BFloat16]): + raise TypeError("Only Float16 or BFloat16 is supported") + if const_expr(mLSE_type not in [Float32]): + raise TypeError("LSE tensor must be Float32") + if const_expr(mdPsum_type not in [Float32]): + raise TypeError("dPsum tensor must be Float32") + if const_expr(mdQaccum_type not in [Float32]): + raise TypeError("dQaccum tensor must be Float32") + if const_expr(self.qhead_per_kvhead == 1): + if const_expr(not (mdK_type == mdV_type == mQ_type)): + raise TypeError("mdK and mdV tensors must have the same data type as mQ") + else: + if const_expr(not (mdK_type == mdV_type == Float32)): + raise TypeError("mdKaccum and mdVaccum tensors must have the data type Float32") + assert mQ_type == self.dtype + + def _setup_attributes(self): + # We need to accommodate both Q and Q^T (and dO and dO^T) in shared memory. + # Q & dO are used in the SdP Mma and Q^T and dO^T are used in the dKV Mma. + # The M dimension (tile_m) doesn't matter for the layout, only the K dimension + wg_d_dKV = self.num_wg_mma // self.AtomLayoutNdKV + self.sQ_layout, self.sdO_layout = [ + # Need to set major_mode_size (mms) to accommodate Q and Q.T + sm90_utils.make_smem_layout(self.dtype, LayoutEnum.ROW_MAJOR, shape, stage, mms) + for shape, stage, mms in [ + ((self.tile_m, self.tile_hdim), self.Q_stage, self.tile_hdim // wg_d_dKV), + ((self.tile_m, self.tile_hdimv), self.dO_stage, self.tile_hdim // wg_d_dKV), + ] + ] + wg_d_dQ = self.num_wg_dQ // self.AtomLayoutMdQ + # Accomodate both K and K.T + self.sK_layout = sm90_utils.make_smem_layout( + self.dtype, + LayoutEnum.ROW_MAJOR, + (self.tile_n, self.tile_hdim), + stage=None, + major_mode_size=self.tile_hdim // wg_d_dQ, + ) + # There's only V, no V.T, so layout is normal + self.sV_layout = sm90_utils.make_smem_layout( + self.dtype, LayoutEnum.ROW_MAJOR, (self.tile_n, self.tile_hdimv), None + ) + # Accomodate both S and S.T + wg_n_SdP = self.num_wg_mma // self.AtomLayoutMSdP + wg_n_dKV = self.AtomLayoutNdKV + self.sPdS_layout = sm90_utils.make_smem_layout( + self.dtype, + LayoutEnum.ROW_MAJOR, + (self.tile_m, self.tile_n), + stage=self.PdS_stage, + major_mode_size=math.gcd(self.tile_n // wg_n_SdP, self.tile_n // wg_n_dKV), + ) + self.sdQaccum_layout = cute.make_layout( + (self.tile_m * self.tile_hdim // self.num_wg_dQ, self.num_wg_dQ) + ) + # dQaccum R->S + self.r2s_tiled_copy_dQaccum = cute.make_tiled_copy_tv( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128), + # thr_layout + cute.make_layout((self.num_threads_per_warp_group, self.num_wg_dQ)), + cute.make_layout(128 // Float32.width), # val_layout + ) + # dKVaccum for GQA epilogue - reuses sV+sK memory recast as f32 + # TODO: assert that sVaccum and sKaccum don't overflow smem + + def _get_tiled_mma(self): + maybe_swap_mn = lambda shape, swap: (shape[1], shape[0], *shape[2:]) if swap else shape + # S = Q @ K.T, dP = dO @ V.T + atom_layout_SdP = (self.AtomLayoutMSdP, self.num_wg_mma // self.AtomLayoutMSdP, 1) + tiler_mn_SdP = (self.tile_m // atom_layout_SdP[0], self.tile_n // atom_layout_SdP[1]) + tiled_mma_SdP = sm90_utils_basic.make_trivial_tiled_mma( + self.dtype, + self.dtype, + warpgroup.OperandMajorMode.K, + warpgroup.OperandMajorMode.K, + Float32, + atom_layout_mnk=maybe_swap_mn(atom_layout_SdP, self.SdP_swapAB), + tiler_mn=(64, tiler_mn_SdP[1] if not self.SdP_swapAB else tiler_mn_SdP[0]), + ) + # dV = P.T @ dO, dK = dS.T @ Q + atom_layout_dKV = (self.AtomLayoutNdKV, self.num_wg_mma // self.AtomLayoutNdKV, 1) + tiler_mn_dK = (self.tile_n // atom_layout_dKV[0], self.tile_hdim // atom_layout_dKV[1]) + tiler_mn_dV = (self.tile_n // atom_layout_dKV[0], self.tile_hdimv // atom_layout_dKV[1]) + tiled_mma_dK, tiled_mma_dV = [ + sm90_utils_basic.make_trivial_tiled_mma( + self.dtype, + self.dtype, + warpgroup.OperandMajorMode.MN + if not self.mma_dkv_is_rs + else warpgroup.OperandMajorMode.K, + warpgroup.OperandMajorMode.MN, + Float32, + atom_layout_mnk=maybe_swap_mn(atom_layout_dKV, self.dKV_swapAB), + tiler_mn=(64, tiler_mn_d[1] if not self.dKV_swapAB else tiler_mn_d[0]), + a_source=warpgroup.OperandSource.RMEM + if self.mma_dkv_is_rs + else warpgroup.OperandSource.SMEM, + ) + for tiler_mn_d in (tiler_mn_dK, tiler_mn_dV) + ] + # dQ = dS @ K + assert self.num_wg_dQ % self.AtomLayoutMdQ == 0 + atom_layout_dQ = (self.AtomLayoutMdQ, self.num_wg_dQ // self.AtomLayoutMdQ, 1) + tiler_mn_dQ = (self.tile_m // atom_layout_dQ[0], self.tile_hdim // atom_layout_dQ[1]) + tiled_mma_dQ = sm90_utils_basic.make_trivial_tiled_mma( + self.dtype, + self.dtype, + warpgroup.OperandMajorMode.K if not self.dQ_swapAB else warpgroup.OperandMajorMode.MN, + warpgroup.OperandMajorMode.MN if not self.dQ_swapAB else warpgroup.OperandMajorMode.K, + Float32, + atom_layout_mnk=maybe_swap_mn(atom_layout_dQ, self.dQ_swapAB), + tiler_mn=(64, tiler_mn_dQ[1] if not self.dQ_swapAB else tiler_mn_dQ[0]), + ) + return tiled_mma_SdP, tiled_mma_dK, tiled_mma_dV, tiled_mma_dQ + + def _get_shared_storage_cls(self): + sQ_struct, sK_struct, sV_struct, sdO_struct, sdQaccum_struct = [ + cute.struct.Align[cute.struct.MemRange[t, cute.cosize(layout)], self.buffer_align_bytes] + for (layout, t) in [ + (self.sQ_layout, self.dtype), + (self.sK_layout, self.dtype), + (self.sV_layout, self.dtype), + (self.sdO_layout, self.dtype), + (self.sdQaccum_layout, Float32), + ] + ] + + cosize_sdS = cute.cosize(self.sPdS_layout) + cosize_sP = cute.cosize(self.sPdS_layout) if const_expr(not self.mma_dkv_is_rs) else 0 + sLSE_struct = cute.struct.Align[ + cute.struct.MemRange[Float32, cute.round_up(self.tile_m, 64) * self.Q_stage], 128 + ] + sdPsum_struct = cute.struct.Align[ + cute.struct.MemRange[Float32, cute.round_up(self.tile_m, 64) * self.dO_stage], 128 + ] + + @cute.struct + class SharedStorageQKV: + mbar_ptr_Q: cute.struct.MemRange[cutlass.Int64, self.Q_stage * 2] + mbar_ptr_dO: cute.struct.MemRange[cutlass.Int64, self.dO_stage * 2] + sLSE: sLSE_struct + sdPsum: sdPsum_struct + sQ: sQ_struct + sV: sV_struct + sK: sK_struct + sdO: sdO_struct + sP: cute.struct.Align[cute.struct.MemRange[self.dtype, cosize_sP], 1024] + sdS: cute.struct.Align[cute.struct.MemRange[self.dtype, cosize_sdS], 1024] + sdQaccum: sdQaccum_struct + + return SharedStorageQKV + + @cute.jit + def __call__( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mdO: cute.Tensor, + mLSE: cute.Tensor, + mdPsum: cute.Tensor, + mdQaccum: cute.Tensor, + mdK: cute.Tensor, + mdV: cute.Tensor, + softmax_scale: Float32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + window_size_left: Int32 | int | None = None, + window_size_right: Int32 | int | None = None, + mdQ_semaphore: Optional[cute.Tensor] = None, + mdK_semaphore: Optional[cute.Tensor] = None, + mdV_semaphore: Optional[cute.Tensor] = None, + aux_data: AuxData = AuxData(), + blocksparse_tensors: Optional[BlockSparseTensors] = None, + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + # For GQA (qhead_per_kvhead > 1), multiple Q heads accumulate into the same dK/dV, + # so we need the float32 accum path + postprocess. + # For varlen_k with qhead_per_kvhead == 1, we use ragged TMA tensors. + self.varlen_k = mCuSeqlensK is not None or mSeqUsedK is not None + + self._check_type( + *( + t.element_type if t is not None else None + for t in (mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV) + ) + ) + + self.is_varlen_q = mCuSeqlensQ is not None or mSeqUsedQ is not None + + mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV = [ + assume_tensor_aligned(t) for t in (mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV) + ] + + # Non-varlen inputs are (b, s, n, h), varlen inputs are (s, n, h). + # We convert both to a seqlen-major view with head-dim second. + # Each tensor may have different rank when Q is padded (seqused_q) but K/V are unpadded (cu_seqlens_k). + def _qkv_transpose(t): + return layout_utils.select(t, [1, 3, 2, 0] if cute.rank(t.shape) == 4 else [0, 2, 1]) + + mQ, mK, mV, mdO = [_qkv_transpose(t) for t in (mQ, mK, mV, mdO)] + if const_expr(self.qhead_per_kvhead == 1): + mdK, mdV = [_qkv_transpose(t) for t in (mdK, mdV)] + else: + # Accum tensors are (b, n, s*h) for non-varlen and (n, s*h) for varlen. + accum_transpose = [2, 1, 0] if cute.rank(mdK.shape) == 3 else [1, 0] + mdK, mdV = [layout_utils.select(t, accum_transpose) for t in (mdK, mdV)] + # Non-varlen stats are (b, n, s), varlen stats are (n, s). + LSE_dPsum_dQaccum_transpose = [2, 1, 0] if cute.rank(mLSE.shape) == 3 else [1, 0] + mLSE, mdPsum, mdQaccum = [ + layout_utils.select(t, LSE_dPsum_dQaccum_transpose) for t in (mLSE, mdPsum, mdQaccum) + ] + + tiled_mma_SdP, tiled_mma_dK, tiled_mma_dV, tiled_mma_dQ = self._get_tiled_mma() + # (batch, num_head, num_m_blocks, cluster_size) -> (num_m_blocks, cluster_size, num_head, batch) + if const_expr(self.deterministic): + assert mdQ_semaphore is not None + mdQ_semaphore = layout_utils.select(mdQ_semaphore, mode=[2, 3, 1, 0]) + if const_expr(self.deterministic and self.qhead_per_kvhead > 1): + assert mdK_semaphore is not None + assert mdV_semaphore is not None + mdK_semaphore, mdV_semaphore = [ + layout_utils.select(t, mode=[2, 3, 1, 0]) for t in (mdK_semaphore, mdV_semaphore) + ] + else: + mdK_semaphore = None + mdV_semaphore = None + + self.num_mma_threads = tiled_mma_SdP.size + assert self.num_mma_threads + 128 == self.num_threads + + self.num_threads_per_warp_group = 128 + self.num_producer_threads = 32 + + REG_LIMIT = 504 if self.num_wg_mma == 2 else 512 + if const_expr(self.num_wg_mma == 2): + if const_expr(self.num_wg_dQ == 1): + self.num_mma_regs_wg0 = 256 + self.num_mma_regs_wg1 = 224 + else: + self.num_mma_regs_wg0 = 240 + self.num_mma_regs_wg1 = 240 + self.num_mma_regs = self.num_mma_regs_wg0 # for backward compat + self.num_producer_regs = 24 + assert ( + self.num_mma_regs_wg0 + self.num_mma_regs_wg1 + self.num_producer_regs <= REG_LIMIT + ) + else: # 3 warp groups + self.num_mma_regs_wg0 = 160 + self.num_mma_regs_wg1 = 160 + self.num_mma_regs = 160 + self.num_producer_regs = 32 + assert self.num_mma_regs_wg0 * self.num_wg_mma + self.num_producer_regs <= REG_LIMIT + + self._setup_attributes() + SharedStorage = self._get_shared_storage_cls() + + self.tma_copy_bytes = { + name: cute.size_in_bytes(mX.element_type, cute.select(layout, mode=[0, 1])) + for name, mX, layout in [ + ("Q", mQ, self.sQ_layout), + ("K", mK, self.sK_layout), + ("V", mV, self.sV_layout), + ("dO", mdO, self.sdO_layout), + ] + } + self.tma_copy_bytes["LSE"] = self.tile_m * Float32.width // 8 + self.tma_copy_bytes["dPsum"] = self.tile_m * Float32.width // 8 + self.tma_copy_bytes["dQ"] = ( + self.tile_m * self.tile_hdim * Float32.width // 8 // self.num_wg_dQ + ) + self.tma_copy_bytes["dKacc"] = self.tile_n * self.tile_hdim * Float32.width // 8 + self.tma_copy_bytes["dVacc"] = self.tile_n * self.tile_hdimv * Float32.width // 8 + + tma_atom_Q, tma_tensor_Q = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + mQ, + cute.select(self.sQ_layout, mode=[0, 1]), + (self.tile_m, self.tile_hdim), + ) + tma_atom_K, tma_tensor_K = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + mK, + cute.select(self.sK_layout, mode=[0, 1]), + (self.tile_n, self.tile_hdim), + ) + tma_atom_V, tma_tensor_V = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + mV, + cute.select(self.sV_layout, mode=[0, 1]), + (self.tile_n, self.tile_hdimv), + ) + tma_atom_dO, tma_tensor_dO = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + mdO, + cute.select(self.sdO_layout, mode=[0, 1]), + (self.tile_m, self.tile_hdimv), + ) + if const_expr(self.qhead_per_kvhead == 1): + mdK_tma = ( + copy_utils.create_ragged_tensor_for_tma(mdK, ragged_dim=0, ptr_shift=True) + if self.varlen_k + else mdK + ) + mdV_tma = ( + copy_utils.create_ragged_tensor_for_tma(mdV, ragged_dim=0, ptr_shift=True) + if self.varlen_k + else mdV + ) + tma_atom_dK, tma_tensor_dK = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + mdK_tma, + cute.select(self.sK_layout, mode=[0, 1]), + (self.tile_n, self.tile_hdim), + ) + tma_atom_dV, tma_tensor_dV = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + mdV_tma, + cute.select(self.sV_layout, mode=[0, 1]), + (self.tile_n, self.tile_hdimv), + ) + else: + tma_atom_dK = tma_atom_dV = tma_tensor_dK = tma_tensor_dV = None + + if const_expr(mCuSeqlensK is not None or mSeqUsedK is not None): + TileScheduler = SingleTileVarlenScheduler + elif const_expr(self.deterministic): + TileScheduler = SingleTileLPTBwdScheduler + else: + TileScheduler = SingleTileScheduler + self.spt = (self.is_causal or self.is_local) and self.deterministic + tile_sched_args = TileSchedulerArguments( + cute.ceil_div(cute.size(mK.shape[0]), self.tile_n), + cute.size(mQ.shape[2]), + cute.size(mK.shape[3]) + if const_expr(mCuSeqlensK is None) + else cute.size(mCuSeqlensK.shape[0] - 1), # num_batch + 1, # num_splits + cute.size(mQ.shape[0]), # pass seqlen_q or total_q for seqlen_k + mQ.shape[1], # headdim + mV.shape[1], # headdim_v + total_q=cute.size(mK.shape[0]) + if const_expr(mCuSeqlensK is not None) + else cute.size(mK.shape[0]) * cute.size(mK.shape[3]), + tile_shape_mn=(self.tile_n, self.tile_m), # Swapping the role of Q & K + mCuSeqlensQ=mCuSeqlensK, + mSeqUsedQ=mSeqUsedK, + qhead_per_kvhead_packgqa=1, + element_size=self.dtype.width // 8, + is_persistent=False, + lpt=self.spt, + head_swizzle=self.deterministic, + ) + + tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args) + grid_dim = TileScheduler.get_grid_shape(tile_sched_params) + + LOG2_E = math.log2(math.e) + if const_expr(self.score_mod is None): + softmax_scale_log2 = softmax_scale * LOG2_E + else: + softmax_scale_log2 = LOG2_E + + fastdiv_mods = None + if const_expr(aux_data.tensors is not None): + seqlen_q = cute.size(mQ.shape[0]) + seqlen_k = cute.size(mK.shape[0]) + seqlen_q_divmod = FastDivmodDivisor(seqlen_q) + seqlen_k_divmod = FastDivmodDivisor(seqlen_k) + fastdiv_mods = (seqlen_q_divmod, seqlen_k_divmod) + + qhead_per_kvhead_divmod = None + if const_expr(self.qhead_per_kvhead > 1): + qhead_per_kvhead_divmod = FastDivmodDivisor(self.qhead_per_kvhead) + + self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None) + + if const_expr(window_size_left is not None): + window_size_left = Int32(window_size_left) + if const_expr(window_size_right is not None): + window_size_right = Int32(window_size_right) + + self.kernel( + tma_tensor_Q, + tma_tensor_K, + tma_tensor_V, + tma_tensor_dO, + tma_tensor_dK if const_expr(self.qhead_per_kvhead == 1) else mdK, + tma_tensor_dV if const_expr(self.qhead_per_kvhead == 1) else mdV, + tma_atom_Q, + tma_atom_K, + tma_atom_V, + tma_atom_dO, + tma_atom_dK, + tma_atom_dV, + mLSE, + mdPsum, + mdQaccum, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + self.sQ_layout, + self.sK_layout, + self.sV_layout, + self.sPdS_layout, + self.sdO_layout, + self.sdQaccum_layout, + self.r2s_tiled_copy_dQaccum, + tiled_mma_SdP, + tiled_mma_dK, + tiled_mma_dV, + tiled_mma_dQ, + softmax_scale_log2, + softmax_scale, + tile_sched_params, + TileScheduler, + SharedStorage, + aux_data, + fastdiv_mods, + blocksparse_tensors, + qhead_per_kvhead_divmod, + mdQ_semaphore, + mdK_semaphore, + mdV_semaphore, + window_size_left, + window_size_right, + ).launch( + grid=grid_dim, + block=[self.num_threads, 1, 1], + stream=stream, + min_blocks_per_mp=1, + use_pdl=True, + ) + + @cute.kernel + def kernel( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mdO: cute.Tensor, + mdK: cute.Tensor, + mdV: cute.Tensor, + tma_atom_Q: cute.CopyAtom, + tma_atom_K: cute.CopyAtom, + tma_atom_V: cute.CopyAtom, + tma_atom_dO: cute.CopyAtom, + tma_atom_dK: cute.CopyAtom, + tma_atom_dV: cute.CopyAtom, + mLSE: cute.Tensor, + mdPsum: cute.Tensor, + mdQaccum: cute.Tensor, + mCuSeqlensQ: Optional[cute.Tensor], + mCuSeqlensK: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + sQ_layout: cute.ComposedLayout, + sK_layout: cute.ComposedLayout, + sV_layout: cute.ComposedLayout, + sPdS_layout: cute.ComposedLayout, + sdO_layout: cute.ComposedLayout, + sdQaccum_layout: cute.Layout, + r2s_tiled_copy_dQaccum: cute.TiledCopy, + tiled_mma_SdP: cute.TiledMma, + tiled_mma_dK: cute.TiledMma, + tiled_mma_dV: cute.TiledMma, + tiled_mma_dQ: cute.TiledMma, + softmax_scale_log2, + softmax_scale, + tile_sched_params: ParamsBase, + TileScheduler: cutlass.Constexpr[Callable], + SharedStorage: cutlass.Constexpr[Callable], + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + blocksparse_tensors: Optional[BlockSparseTensors] = None, + qhead_per_kvhead_divmod: Optional[FastDivmodDivisor] = None, + mdQ_semaphore: Optional[cute.Tensor] = None, + mdK_semaphore: Optional[cute.Tensor] = None, + mdV_semaphore: Optional[cute.Tensor] = None, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + # prefetch TMA descriptors + if warp_idx == 0: + for atom in [tma_atom_Q, tma_atom_K, tma_atom_V, tma_atom_dO, tma_atom_dK, tma_atom_dV]: + if const_expr(atom is not None): + cpasync.prefetch_descriptor(atom) + + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + pipeline_producer_group = cutlass.pipeline.CooperativeGroup(cutlass.pipeline.Agent.Thread) + pipeline_consumer_group = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, self.num_mma_threads // cute.arch.WARP_SIZE + ) + pipeline_Q = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.mbar_ptr_Q.data_ptr(), + num_stages=self.Q_stage, + producer_group=pipeline_producer_group, + consumer_group=pipeline_consumer_group, + tx_count=self.tma_copy_bytes["Q"] + self.tma_copy_bytes["LSE"], + defer_sync=True, + ) + pipeline_dO = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.mbar_ptr_dO.data_ptr(), + num_stages=self.dO_stage, + producer_group=pipeline_producer_group, + consumer_group=pipeline_consumer_group, + tx_count=self.tma_copy_bytes["dO"] + self.tma_copy_bytes["dPsum"], + defer_sync=False, + ) + + sQ = storage.sQ.get_tensor(sQ_layout.outer, swizzle=sQ_layout.inner) + sdO = storage.sdO.get_tensor(sdO_layout.outer, swizzle=sdO_layout.inner) + sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner) + sV = storage.sV.get_tensor(sV_layout.outer, swizzle=sV_layout.inner) + sP = None + if const_expr(not self.mma_dkv_is_rs): + sP = storage.sP.get_tensor(sPdS_layout.outer, swizzle=sPdS_layout.inner) + sdS = storage.sdS.get_tensor(sPdS_layout.outer, swizzle=sPdS_layout.inner) + sLSE = storage.sLSE.get_tensor( + cute.make_layout( + (self.tile_m, self.Q_stage), + stride=(1, cute.round_up(self.tile_m, 64)), + ) + ) + sdPsum = storage.sdPsum.get_tensor( + cute.make_layout( + (self.tile_m, self.dO_stage), + stride=(1, cute.round_up(self.tile_m, 64)), + ) + ) + sdQaccum = storage.sdQaccum.get_tensor(sdQaccum_layout) + + block_info = BlockInfo( + self.tile_m, + self.tile_n, + self.is_causal, + self.is_local, + False, # is_split_kv + window_size_left, + window_size_right, + qhead_per_kvhead_packgqa=1, + ) + SeqlenInfoCls = partial( + SeqlenInfoQK.create, + seqlen_q_static=mQ.shape[0], + seqlen_k_static=mK.shape[0], + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + tile_m=self.tile_m, + tile_n=self.tile_n, + ) + AttentionMaskCls = partial( + AttentionMask, + self.tile_m, + self.tile_n, + window_size_left=window_size_left, + window_size_right=window_size_right, + swap_AB=self.SdP_swapAB, + ) + TileSchedulerCls = partial(TileScheduler.create, tile_sched_params) + + if warp_idx < 4: + cute.arch.setmaxregister_decrease(self.num_producer_regs) + if warp_idx == 0: + self.load( + mQ, + mK, + mV, + mdO, + mLSE, + mdPsum, + sQ, + sK, + sV, + sdO, + sLSE, + sdPsum, + tma_atom_Q, + tma_atom_K, + tma_atom_V, + tma_atom_dO, + pipeline_Q, + pipeline_dO, + block_info, + SeqlenInfoCls, + TileSchedulerCls, + blocksparse_tensors, + qhead_per_kvhead_divmod, + ) + if warp_idx == 1: + self.dQaccum_store( + mdQaccum, + sdQaccum, + block_info, + TileSchedulerCls, + SeqlenInfoCls, + blocksparse_tensors, + mdQ_semaphore, + ) + else: + tidx, _, _ = cute.arch.thread_idx() + tidx = tidx - 128 + mma_args = ( + tiled_mma_SdP, + tiled_mma_dK, + tiled_mma_dV, + tiled_mma_dQ, + mdK, + mdV, + mdK_semaphore, + mdV_semaphore, + mdQaccum, + sQ, + sK, + sV, + sdO, + sP, + sdS, + sLSE, + sdPsum, + sdQaccum, + pipeline_Q, + pipeline_dO, + tidx, + tma_atom_dK, + tma_atom_dV, + r2s_tiled_copy_dQaccum, + softmax_scale_log2, + softmax_scale, + block_info, + SeqlenInfoCls, + AttentionMaskCls, + TileSchedulerCls, + aux_data, + fastdiv_mods, + blocksparse_tensors, + qhead_per_kvhead_divmod, + ) + if const_expr(self.num_wg_dQ == self.num_wg_mma): + # Both WGs compute dQ + cute.arch.setmaxregister_increase(self.num_mma_regs_wg0) + self.mma(*mma_args, is_dQ_wg=True) + else: + # WG0 computes dQ, WG1 skips it + warp_idx_in_mma = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - 4 + if warp_idx_in_mma < 4: + cute.arch.setmaxregister_increase(self.num_mma_regs_wg0) + self.mma(*mma_args, is_dQ_wg=True) + else: + cute.arch.setmaxregister_increase(self.num_mma_regs_wg1) + self.mma(*mma_args, is_dQ_wg=False) + + @cute.jit + def load( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mdO: cute.Tensor, + mLSE: cute.Tensor, + mdPsum: cute.Tensor, + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + sdO: cute.Tensor, + sLSE: cute.Tensor, + sdPsum: cute.Tensor, + tma_atom_Q: cute.CopyAtom, + tma_atom_K: cute.CopyAtom, + tma_atom_V: cute.CopyAtom, + tma_atom_dO: cute.CopyAtom, + pipeline_Q: cutlass.pipeline.PipelineAsync, + pipeline_dO: cutlass.pipeline.PipelineAsync, + block_info: BlockInfo, + SeqlenInfoCls: Callable, + TileSchedulerCls: Callable, + blocksparse_tensors: Optional[BlockSparseTensors] = None, + qhead_per_kvhead_divmod: Optional[FastDivmodDivisor] = None, + ): + warp_idx_in_wg = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4 + + if warp_idx_in_wg == 0: + producer_state_Q = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Producer, self.Q_stage + ) + producer_state_dO = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Producer, self.dO_stage + ) + tile_scheduler = TileSchedulerCls() + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + n_block, head_idx, batch_idx, _ = work_tile.tile_idx + seqlen = SeqlenInfoCls(batch_idx) + head_idx_kv = ( + head_idx + if const_expr(self.qhead_per_kvhead == 1) + else head_idx // qhead_per_kvhead_divmod + ) + mK_cur = seqlen.offset_batch_K(mK, batch_idx, dim=3)[None, None, head_idx_kv] + mV_cur = seqlen.offset_batch_K(mV, batch_idx, dim=3)[None, None, head_idx_kv] + gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (n_block, 0)) + gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (n_block, 0)) + + mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx] + mLSE_cur = seqlen.offset_batch_Q(mLSE, batch_idx, dim=2, padded=True)[ + None, head_idx + ] + mdO_cur = seqlen.offset_batch_Q(mdO, batch_idx, dim=3)[None, None, head_idx] + mdPsum_cur = seqlen.offset_batch_Q(mdPsum, batch_idx, dim=2, padded=True)[ + None, head_idx + ] + gQ = cute.local_tile(mQ_cur, (self.tile_m, self.tile_hdim), (None, 0)) + gdO = cute.local_tile(mdO_cur, (self.tile_m, self.tile_hdimv), (None, 0)) + gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (None,)) + gdPsum = cute.local_tile(mdPsum_cur, (self.tile_m,), (None,)) + + load_K, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_K, 0, cute.make_layout(1), gK, sK, single_stage=True + ) + load_V, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_V, 0, cute.make_layout(1), gV, sV, single_stage=True + ) + load_Q, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_Q, 0, cute.make_layout(1), gQ, sQ + ) + load_Q = copy_utils.tma_producer_copy_fn(load_Q, pipeline_Q) + load_dO, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_dO, 0, cute.make_layout(1), gdO, sdO + ) + load_dO = copy_utils.tma_producer_copy_fn(load_dO, pipeline_dO) + load_LSE = copy_utils.cpasync_bulk_get_copy_fn(gLSE, sLSE) + load_LSE = copy_utils.tma_producer_copy_fn(load_LSE, pipeline_Q) + load_dPsum = copy_utils.cpasync_bulk_get_copy_fn(gdPsum, sdPsum) + load_dPsum = copy_utils.tma_producer_copy_fn(load_dPsum, pipeline_dO) + + m_block_min, m_block_max = block_info.get_m_block_min_max(seqlen, n_block) + + if const_expr(not self.use_block_sparsity): + total_m_block_cnt = m_block_max - m_block_min + process_tile = ( + const_expr(not self.is_local and not self.is_varlen_q) + or m_block_min < m_block_max + ) + else: + total_m_block_cnt = get_total_q_block_count_bwd( + blocksparse_tensors, + batch_idx, + head_idx, + n_block, + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + ) + process_tile = total_m_block_cnt > Int32(0) + + if process_tile: + if const_expr(not self.use_block_sparsity): + first_m_block = m_block_min + pipeline_Q.producer_acquire( + producer_state_Q, extra_tx_count=self.tma_copy_bytes["K"] + ) + load_K(tma_bar_ptr=pipeline_Q.producer_get_barrier(producer_state_Q)) + load_Q(first_m_block, producer_state=producer_state_Q) + # Wait for bwd preprocess to finish writing LSE and dPsum + cute.arch.griddepcontrol_wait() + load_LSE(first_m_block, producer_state=producer_state_Q) + producer_state_dO_cur = ( + producer_state_dO + if const_expr(self.Q_stage != self.dO_stage) + else producer_state_Q + ) + pipeline_dO.producer_acquire( + producer_state_dO_cur, extra_tx_count=self.tma_copy_bytes["V"] + ) + load_V(tma_bar_ptr=pipeline_dO.producer_get_barrier(producer_state_dO_cur)) + load_dO(first_m_block, producer_state=producer_state_dO_cur) + load_dPsum(first_m_block, producer_state=producer_state_dO_cur) + producer_state_Q.advance() + producer_state_dO.advance() + + for m_block in cutlass.range(m_block_min + 1, m_block_max, unroll=1): + pipeline_Q.producer_acquire(producer_state_Q) + load_Q(m_block, producer_state=producer_state_Q) + load_LSE(m_block, producer_state=producer_state_Q) + producer_state_dO_cur = ( + producer_state_dO + if const_expr(self.Q_stage != self.dO_stage) + else producer_state_Q + ) + pipeline_dO.producer_acquire(producer_state_dO_cur) + load_dO(m_block, producer_state=producer_state_dO_cur) + load_dPsum(m_block, producer_state=producer_state_dO_cur) + producer_state_Q.advance() + producer_state_dO.advance() + else: + producer_state_Q, producer_state_dO = produce_block_sparse_q_loads_bwd_sm90( + blocksparse_tensors, + batch_idx, + head_idx, + n_block, + producer_state_Q, + producer_state_dO, + pipeline_Q, + pipeline_dO, + load_K, + load_V, + load_Q, + load_dO, + load_LSE, + load_dPsum, + self.tma_copy_bytes["K"], + self.tma_copy_bytes["V"], + Q_stage_eq_dO_stage=(self.Q_stage == self.dO_stage), + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + ) + + tile_scheduler.prefetch_next_work() + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + + @cute.jit + def apply_score_mod( + self, + acc_S: cute.Tensor, + thr_mma_SdP: cute.ThrMma, + batch_idx, + head_idx, + m_block, + n_block, + softmax_scale, + seqlen_info: SeqlenInfoQK, + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + ): + # [NOTE] SdP_swapAB: swapAB transposes the tile, so use (n, m) indexing + cS = cute.make_identity_tensor( + (self.tile_n, self.tile_m) if self.SdP_swapAB else (self.tile_m, self.tile_n) + ) + cS = cute.domain_offset( + (n_block * self.tile_n, m_block * self.tile_m) + if self.SdP_swapAB + else (m_block * self.tile_m, n_block * self.tile_n), + cS, + ) + tScS = thr_mma_SdP.partition_C(cS) + + apply_score_mod_inner( + acc_S, + tScS, + self.score_mod, + batch_idx, + head_idx, + softmax_scale, + self.vec_size, + self.qk_acc_dtype, + aux_data, + fastdiv_mods, + seqlen_info, + constant_q_idx=None, + qhead_per_kvhead=self.qhead_per_kvhead, + transpose_indices=self.SdP_swapAB, + ) + + @cute.jit + def apply_score_mod_bwd( + self, + grad_tensor: cute.Tensor, + score_tensor: cute.Tensor, + thr_mma_SdP: cute.ThrMma, + batch_idx, + head_idx, + m_block, + n_block, + softmax_scale, + seqlen_info: SeqlenInfoQK, + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + ): + cS = cute.make_identity_tensor( + (self.tile_n, self.tile_m) if self.SdP_swapAB else (self.tile_m, self.tile_n) + ) + cS = cute.domain_offset( + (n_block * self.tile_n, m_block * self.tile_m) + if self.SdP_swapAB + else (m_block * self.tile_m, n_block * self.tile_n), + cS, + ) + tScS = thr_mma_SdP.partition_C(cS) + + apply_score_mod_bwd_inner( + grad_tensor, + score_tensor, + tScS, + self.score_mod_bwd, + batch_idx, + head_idx, + softmax_scale, + self.vec_size, + self.qk_acc_dtype, + aux_data, + fastdiv_mods, + seqlen_info, + constant_q_idx=None, + qhead_per_kvhead=self.qhead_per_kvhead, + transpose_indices=self.SdP_swapAB, + ) + + @cute.jit + def mma( + self, + tiled_mma_SdP: cute.TiledMma, + tiled_mma_dK: cute.TiledMma, + tiled_mma_dV: cute.TiledMma, + tiled_mma_dQ: cute.TiledMma, + mdK: cute.Tensor, + mdV: cute.Tensor, + mdK_semaphore: Optional[cute.Tensor], + mdV_semaphore: Optional[cute.Tensor], + mdQaccum: cute.Tensor, + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + sdO: cute.Tensor, + sP: Optional[cute.Tensor], + sdS: cute.Tensor, + sLSE: cute.Tensor, + sdPsum: cute.Tensor, + sdQaccum: cute.Tensor, + pipeline_Q: cutlass.pipeline.PipelineAsync, + pipeline_dO: cutlass.pipeline.PipelineAsync, + tidx: Int32, + tma_atom_dK: cute.CopyAtom, + tma_atom_dV: cute.CopyAtom, + r2s_tiled_copy_dQaccum: cute.TiledCopy, + softmax_scale_log2: Float32, + softmax_scale: Float32, + block_info: BlockInfo, + SeqlenInfoCls: Callable, + AttentionMaskCls: Callable, + TileSchedulerCls: Callable, + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + blocksparse_tensors: Optional[BlockSparseTensors] = None, + qhead_per_kvhead_divmod: Optional[FastDivmodDivisor] = None, + is_dQ_wg: cutlass.Constexpr[bool] = True, + ): + warp_group_idx = cute.arch.make_warp_uniform(tidx // self.num_threads_per_warp_group) + warp_group_thread_layout = cute.make_layout( + self.num_wg_mma, stride=self.num_threads_per_warp_group + ) + thr_mma_SdP = tiled_mma_SdP.get_slice(tidx) + wg_mma_SdP = tiled_mma_SdP.get_slice(warp_group_thread_layout(warp_group_idx)) + wg_mma_dK = tiled_mma_dK.get_slice(warp_group_thread_layout(warp_group_idx)) + wg_mma_dV = tiled_mma_dV.get_slice(warp_group_thread_layout(warp_group_idx)) + wg_mma_dQ = None + if const_expr(is_dQ_wg): + wg_idx_dQ = warp_group_idx if const_expr(self.num_wg_dQ > 1) else 0 + wg_mma_dQ = tiled_mma_dQ.get_slice(warp_group_thread_layout(wg_idx_dQ)) + # S = Q @ K.T + shape_mnk_S = (self.tile_m, self.tile_n, self.tile_hdim) + _, tSrQ, tSrK = sm90_utils.partition_fragment_ABC( + wg_mma_SdP, shape_mnk_S, sQ, sK, swap_AB=self.SdP_swapAB + ) + mma_qk_fn = partial( + gemm_zero_init, tiled_mma_SdP, shape_mnk_S[:2], tSrQ, tSrK, swap_AB=self.SdP_swapAB + ) + # dP = dO @ V.T + shape_mnk_dP = (self.tile_m, self.tile_n, self.tile_hdimv) + _, tdPrdO, tdPrV = sm90_utils.partition_fragment_ABC( + wg_mma_SdP, shape_mnk_dP, sdO, sV, swap_AB=self.SdP_swapAB + ) + mma_dov_fn = partial( + gemm_zero_init, tiled_mma_SdP, shape_mnk_dP[:2], tdPrdO, tdPrV, swap_AB=self.SdP_swapAB + ) + # dV += P.T @ dO + sPt = layout_utils.transpose_view(sP) if sP is not None else None + sdOt = layout_utils.transpose_view(sdO) + shape_mnk_dV = (self.tile_n, self.tile_hdimv, self.tile_m) + acc_dV, tdVrPt, tdVrdOt = sm90_utils.partition_fragment_ABC( + wg_mma_dV, shape_mnk_dV, sPt, sdOt, swap_AB=self.dKV_swapAB + ) + if const_expr(not self.mma_dkv_is_rs): + mma_pdo_fn = partial( + gemm_w_idx, tiled_mma_dV, acc_dV, tdVrPt, tdVrdOt, swap_AB=self.dKV_swapAB + ) + else: + mma_pdo_fn = partial(gemm_w_idx, tiled_mma_dV, acc_dV, tCrB=tdVrdOt) + # dK += dS.T @ Q + sdSt = layout_utils.transpose_view(sdS) + sQt = layout_utils.transpose_view(sQ) + shape_mnk_dK = (self.tile_n, self.tile_hdim, self.tile_m) + acc_dK, tdKrdSt, tdKrQt = sm90_utils.partition_fragment_ABC( + wg_mma_dK, shape_mnk_dK, sdSt, sQt, swap_AB=self.dKV_swapAB + ) + if const_expr(not self.mma_dkv_is_rs): + mma_dsq_fn = partial( + gemm_w_idx, tiled_mma_dK, acc_dK, tdKrdSt, tdKrQt, swap_AB=self.dKV_swapAB + ) + else: + mma_dsq_fn = partial(gemm_w_idx, tiled_mma_dK, acc_dK, tCrB=tdKrQt) + # dQ = dS @ K + sKt = layout_utils.transpose_view(sK) + shape_mnk_dQ = (self.tile_m, self.tile_hdim, self.tile_n) + mma_dsk_fn = None + if const_expr(is_dQ_wg): + _, tdQrdS, tdQrKt = sm90_utils.partition_fragment_ABC( + wg_mma_dQ, shape_mnk_dQ, sdS, sKt, swap_AB=self.dQ_swapAB + ) + mma_dsk_fn = partial( + gemm_zero_init, + tiled_mma_dQ, + shape_mnk_dQ[:2], + tdQrdS, + tdQrKt, + swap_AB=self.dQ_swapAB, + ) + + # Smem copy atom tiling for P/dS R2S + copy_P_r2s = None + mms_PdS = self.tile_n // (self.num_wg_mma // self.AtomLayoutMSdP) + if const_expr(sP is not None): + sP_cpy = sP if const_expr(not self.SdP_swapAB) else sPt + copy_P_r2s, _, _ = copy_utils.get_smem_store_C( + tiled_mma_SdP, + sP_cpy, + tidx, + transpose=self.SdP_swapAB, + position_independent=True, + major_mode_size=mms_PdS, + ) + sdS_cpy = sdS if const_expr(not self.SdP_swapAB) else sdSt + copy_dS_r2s, _, _ = copy_utils.get_smem_store_C( + tiled_mma_SdP, + sdS_cpy, + tidx, + transpose=self.SdP_swapAB, + position_independent=True, + major_mode_size=mms_PdS, + ) + + tLSEsLSE = layout_utils.mma_partition_C_vec( + sLSE, thr_mma_SdP, expand_shape=self.tile_n, is_colvec=not self.SdP_swapAB + ) + tLSEsdPsum = layout_utils.mma_partition_C_vec( + sdPsum, thr_mma_SdP, expand_shape=self.tile_n, is_colvec=not self.SdP_swapAB + ) + # When shuffle=True, rows are distributed across 8 quads (4 threads each) within a warp. + # Each thread loads only ceil(num_rows/8) values; + shfl_copy = copy_utils.tiled_copy_1d(sLSE.element_type, num_threads=8, num_copy_elems=2) + if const_expr(self.shuffle_LSE): + tLSEsLSE = shfl_copy.get_slice(cute.arch.lane_idx() // 4).partition_S(tLSEsLSE) + # ((2, 1), 1, 2) -> (((2, 1), 1), 2) + tLSEsLSE = cute.group_modes(tLSEsLSE, 0, 2) + if const_expr(self.shuffle_dPsum): + tLSEsdPsum = shfl_copy.get_slice(cute.arch.lane_idx() // 4).partition_S(tLSEsdPsum) + tLSEsdPsum = cute.group_modes(tLSEsdPsum, 0, 2) + + tdQsdQaccum = None + if const_expr(is_dQ_wg): + smem_thr_copy_dQaccum = r2s_tiled_copy_dQaccum.get_slice(tidx) + tdQsdQaccum = smem_thr_copy_dQaccum.partition_D(sdQaccum) + + PdS_barrier = cutlass.pipeline.NamedBarrier( + barrier_id=int(NamedBarrierBwd.PdS), num_threads=self.num_mma_threads + ) + score_mod_fn = partial( + self.apply_score_mod, + thr_mma_SdP=thr_mma_SdP, + softmax_scale=softmax_scale, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + score_mod_bwd_fn = partial( + self.apply_score_mod_bwd, + thr_mma_SdP=thr_mma_SdP, + softmax_scale=softmax_scale, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + + mma_one_m_block_all = partial( + self.mma_one_m_block, + warp_group_idx=warp_group_idx, + mma_qk_fn=mma_qk_fn, + mma_dov_fn=mma_dov_fn, + mma_pdo_fn=mma_pdo_fn, + mma_dsq_fn=mma_dsq_fn, + mma_dsk_fn=mma_dsk_fn, + copy_P_r2s=copy_P_r2s, + copy_dS_r2s=copy_dS_r2s, + pipeline_Q=pipeline_Q, + pipeline_dO=pipeline_dO, + tLSEsLSE=tLSEsLSE, + tLSEsdPsum=tLSEsdPsum, + tdQsdQaccum=tdQsdQaccum, + softmax_scale_log2=softmax_scale_log2, + PdS_barrier=PdS_barrier, + # acc_dV=acc_dV, + # acc_dK=acc_dK, + is_dQ_wg=is_dQ_wg, + ) + + consumer_state_Q = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Consumer, self.Q_stage + ) + consumer_state_dO = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Consumer, self.dO_stage + ) + tile_scheduler = TileSchedulerCls() + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + n_block, head_idx, batch_idx, _ = work_tile.tile_idx + seqlen = SeqlenInfoCls(batch_idx) + mask = AttentionMaskCls(seqlen) + score_mod_fn_cur = partial( + score_mod_fn, + batch_idx=batch_idx, + head_idx=head_idx, + n_block=n_block, + seqlen_info=seqlen, + ) + score_mod_bwd_fn_cur = partial( + score_mod_bwd_fn, + batch_idx=batch_idx, + head_idx=head_idx, + n_block=n_block, + seqlen_info=seqlen, + ) + m_block_min, m_block_max = block_info.get_m_block_min_max(seqlen, n_block) + + if const_expr(not self.use_block_sparsity): + process_tile = ( + const_expr(not self.is_local and not self.is_varlen_q) + or m_block_min < m_block_max + ) + else: + total_m_block_cnt = get_total_q_block_count_bwd( + blocksparse_tensors, + batch_idx, + head_idx, + n_block, + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + ) + process_tile = total_m_block_cnt > Int32(0) + + if process_tile: + if const_expr(not self.use_block_sparsity): + mask_fn = partial( + mask.apply_mask, + batch_idx=batch_idx, + head_idx=head_idx, + n_block=n_block, + thr_mma=thr_mma_SdP, + mask_seqlen=True, + mask_causal=self.is_causal, + mask_local=self.is_local, + mask_mod=self.mask_mod, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + dKV_accumulate = False + for m_block in cutlass.range(m_block_min, m_block_max, unroll=1): + consumer_state_Q, consumer_state_dO = mma_one_m_block_all( + m_block, + consumer_state_Q, + consumer_state_dO, + mask_fn=mask_fn, + score_mod_fn=score_mod_fn_cur, + score_mod_bwd_fn=score_mod_bwd_fn_cur, + dKV_accumulate=dKV_accumulate, + ) + dKV_accumulate = True + else: + consumer_state_Q, consumer_state_dO = consume_block_sparse_mma_bwd_sm90( + blocksparse_tensors, + batch_idx, + head_idx, + n_block, + consumer_state_Q, + consumer_state_dO, + mma_one_m_block_all, + mask, + self.mask_mod, + is_causal=self.is_causal, + is_local=self.is_local, + thr_mma_SdP=thr_mma_SdP, + score_mod_fn=score_mod_fn_cur, + score_mod_bwd_fn=score_mod_bwd_fn_cur, + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + + if const_expr(self.qhead_per_kvhead == 1): + acc_dK.store(acc_dK.load() * softmax_scale) + self.epilogue_dKV( + acc_dV, + mdV, + sV, + acc_dK, + mdK, + sK, + seqlen, + tma_atom_dK, + tma_atom_dV, + tiled_mma_dK, + tiled_mma_dV, + tidx, + n_block, + head_idx, + batch_idx, + qhead_per_kvhead_divmod, + mdK_semaphore, + mdV_semaphore, + ) + else: + # KV tile with zero Q blocks produces no dK/dV; write zeros. + if const_expr(self.use_block_sparsity or self.is_local or self.is_varlen_q): + acc_dK.fill(0.0) + acc_dV.fill(0.0) + self.epilogue_dKV( + acc_dV, + mdV, + sV, + acc_dK, + mdK, + sK, + seqlen, + tma_atom_dK, + tma_atom_dV, + tiled_mma_dK, + tiled_mma_dV, + tidx, + n_block, + head_idx, + batch_idx, + qhead_per_kvhead_divmod, + mdK_semaphore, + mdV_semaphore, + ) + + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 4: + cute.arch.cp_async_bulk_wait_group(0, read=True) + + @staticmethod + @cute.jit + def _get_stat(tSrS: cute.Tensor, row: Int32, lane: Int32, shuffle: bool) -> Float32: + """Retrieve the statistic for a given accumulator row. + + When shuffle=False, direct register indexing. + When shuffle=True, warp shuffle from the thread group that holds the value. + """ + if const_expr(not shuffle): + return tSrS[row] + # tSrS: (((2, 1), 1), 1)), distributed across 8 threads in the warp + vecsize = cute.size(tSrS, mode=[0, 0]) # 2 + idx0, off, idx1 = cute.idx2crd(row, (vecsize, 8, cute.shape(tSrS, mode=[0, 1]))) + # register index: 0, 1, 0, 1, ..., 2, 3, 2, 3, ... + return utils.shuffle_sync(tSrS[idx0 + idx1 * vecsize], offset=off * 4 + (lane % 4)) + + @cute.jit + def mma_one_m_block( + self, + m_block: Int32, + consumer_state_Q: cutlass.pipeline.PipelineState | pipeline.PipelineStateSimple, + consumer_state_dO: cutlass.pipeline.PipelineState | pipeline.PipelineStateSimple, + warp_group_idx: Int32, + mma_qk_fn: Callable, + mma_dov_fn: Callable, + mma_pdo_fn: Callable, + mma_dsq_fn: Callable, + mma_dsk_fn: Callable, + copy_P_r2s: Optional[Callable], + copy_dS_r2s: Callable, + pipeline_Q: cutlass.pipeline.PipelineAsync, + pipeline_dO: cutlass.pipeline.PipelineAsync, + tLSEsLSE: cute.Tensor, + tLSEsdPsum: cute.Tensor, + tdQsdQaccum: Optional[cute.Tensor], + softmax_scale_log2: Float32, + PdS_barrier: cutlass.pipeline.NamedBarrier, + is_dQ_wg: cutlass.Constexpr[bool] = True, + mask_fn: Optional[Callable] = None, + score_mod_fn: Optional[Callable] = None, + score_mod_bwd_fn: Optional[Callable] = None, + dKV_accumulate: Boolean = True, + ): + consumer_state_dO_cur = ( + consumer_state_Q if const_expr(self.Q_stage == self.dO_stage) else consumer_state_dO + ) + smem_idx_Q = consumer_state_Q.index + smem_idx_dO = consumer_state_dO_cur.index if const_expr(self.dO_stage > 1) else 0 + smem_idx_PdS = smem_idx_Q if const_expr(self.PdS_stage > 1) else 0 + # (1) [GEMM 1] S = Q @ K^T + pipeline_Q.consumer_wait(consumer_state_Q, pipeline_Q.consumer_try_wait(consumer_state_Q)) + acc_S = mma_qk_fn(A_idx=smem_idx_Q, wg_wait=-1) + # If shuffle_LSE, OOB reads are OK since sLSE is already padded + tLSErLSE = copy_utils.load_s2r(tLSEsLSE[None, smem_idx_Q]) + # (2) [GEMM 2] dP = dO @ V.T + pipeline_dO.consumer_wait( + consumer_state_dO_cur, pipeline_dO.consumer_try_wait(consumer_state_dO_cur) + ) + acc_dP = mma_dov_fn(A_idx=smem_idx_Q, wg_wait=1) + + if const_expr(self.score_mod_bwd is not None): + acc_S_pre = cute.make_fragment_like(acc_S) + cute.autovec_copy(acc_S, acc_S_pre) + + if const_expr(self.score_mod is not None): + score_mod_fn(acc_S, m_block=m_block) + + # (3) [Pointwise 1] P = exp(S - LSE) + if cutlass.const_expr(mask_fn is not None): + mask_fn(acc_S, m_block=m_block) + acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S, transpose=self.SdP_swapAB) + lane_idx = cute.arch.lane_idx() + for r in cutlass.range_constexpr(cute.size(acc_S_mn, mode=[0])): + lse_val = self._get_stat(tLSErLSE, r, lane_idx, shuffle=self.shuffle_LSE) + for c in cutlass.range(cute.size(acc_S_mn, mode=[1]), unroll_full=True): + acc_S_mn[r, c] = cute.math.exp2( + acc_S_mn[r, c] * softmax_scale_log2 - lse_val, fastmath=True + ) + tLSErdPsum = copy_utils.load_s2r(tLSEsdPsum[None, smem_idx_dO]) + + # Convert P from f32 -> f16 + tdVrP = utils.cvt_f16(layout_utils.reshape_acc_to_frgA(acc_S), self.dtype) + # R2S for P + if const_expr(not self.mma_dkv_is_rs): + # sync to ensure P has already been used in the previous iteration before overwriting + if const_expr(self.PdS_stage == 1): + PdS_barrier.arrive_and_wait() + copy_P_r2s(tdVrP, dst_idx=smem_idx_PdS) + + # (4) [Pointwise 2] dS = P*(dP-dPsum) + warpgroup.wait_group(0) + acc_dP_mn = layout_utils.reshape_acc_to_mn(acc_dP, transpose=self.SdP_swapAB) + for r in cutlass.range_constexpr(cute.size(acc_dP_mn, mode=[0])): + dpsum_val = self._get_stat(tLSErdPsum, r, lane_idx, shuffle=self.shuffle_dPsum) + for c in cutlass.range(cute.size(acc_dP_mn, mode=[1]), unroll_full=True): + acc_dP_mn[r, c] = acc_S_mn[r, c] * (acc_dP_mn[r, c] - dpsum_val) + + if const_expr(self.score_mod_bwd is not None): + score_mod_bwd_fn(acc_dP, acc_S_pre, m_block=m_block) + + # Convert dS from f32 -> f16 + tdKrdS = utils.cvt_f16(layout_utils.reshape_acc_to_frgA(acc_dP), self.dtype) + + # If there's double buffering on dS, we don't need to sync here. + # Otherwise we might have WG1 writing to dS before WG2 is done reading from it during MmadQ. + # But because both WGs have to sync at the end of the loop and double buffering, + # this race condition is not possible. + # This sync is to ensure (1) P is written in case of !mma_dkv_is_rs and + # (2) dS is already read by the Mma in the previous iteration in case of mma_dkv_is_rs. + if const_expr(not self.mma_dkv_is_rs or (self.PdS_stage == 1 and self.mma_dkv_is_rs)): + cute.arch.fence_view_async_shared() + PdS_barrier.arrive_and_wait() + + # R2S for dS + copy_dS_r2s(tdKrdS, dst_idx=smem_idx_PdS) + + # (5) [GEMM 3] dV += P.T @ dO + if const_expr(not self.mma_dkv_is_rs): + mma_pdo_fn( + A_idx=smem_idx_PdS, B_idx=smem_idx_dO, zero_init=not dKV_accumulate, wg_wait=-1 + ) + else: + mma_pdo_fn(tCrA=tdVrP, B_idx=smem_idx_dO, zero_init=not dKV_accumulate, wg_wait=-1) + + # smem fence to make sure sdS is written before it's read by WGMMA + cute.arch.fence_view_async_shared() + PdS_barrier.arrive_and_wait() + + if const_expr(is_dQ_wg): + # (6) [GEMM 4] dQ = dS @ K + acc_dQ = mma_dsk_fn(A_idx=smem_idx_PdS, wg_wait=1) + pipeline_dO.consumer_release(consumer_state_dO_cur) # release dO as dV mma is done + + # (7) [GEMM 5] dK += dS.T @ Q + if const_expr(not self.mma_dkv_is_rs): + mma_dsq_fn( + A_idx=smem_idx_PdS, B_idx=smem_idx_Q, zero_init=not dKV_accumulate, wg_wait=1 + ) + else: + mma_dsq_fn(tCrA=tdKrdS, B_idx=smem_idx_Q, zero_init=not dKV_accumulate, wg_wait=1) + + # dQ R2S: wait for dQaccum_store to free the smem buffer, then write dQ to smem + # When dQ_single_wg, only WG0 enters here so warp_group_idx == 0 + cute.arch.barrier( + barrier_id=int(NamedBarrierBwd.dQEmptyWG0) + warp_group_idx, + number_of_threads=self.num_threads_per_warp_group + cute.arch.WARP_SIZE, + ) + tdQrdQaccum_flat = cute.make_tensor( + acc_dQ.iterator, cute.make_layout(tdQsdQaccum.shape) + ) + cute.autovec_copy(tdQrdQaccum_flat, tdQsdQaccum) + cute.arch.fence_view_async_shared() + cute.arch.barrier_arrive( + barrier_id=int(NamedBarrierBwd.dQFullWG0) + warp_group_idx, + number_of_threads=self.num_threads_per_warp_group + cute.arch.WARP_SIZE, + ) + + warpgroup.wait_group(0) + pipeline_Q.consumer_release(consumer_state_Q) + else: + # dQ_single_wg: WG1 skips dQ, only does dV wait + dK + # (7) [GEMM 5] dK += dS.T @ Q + if const_expr(not self.mma_dkv_is_rs): + mma_dsq_fn( + A_idx=smem_idx_PdS, B_idx=smem_idx_Q, zero_init=not dKV_accumulate, wg_wait=1 + ) + else: + mma_dsq_fn(tCrA=tdKrdS, B_idx=smem_idx_Q, zero_init=not dKV_accumulate, wg_wait=1) + pipeline_dO.consumer_release(consumer_state_dO_cur) + warpgroup.wait_group(0) + pipeline_Q.consumer_release(consumer_state_Q) + + consumer_state_Q.advance() + consumer_state_dO.advance() + return consumer_state_Q, consumer_state_dO + + @cute.jit + def epilogue_dKV( + self, + acc_dV: cute.Tensor, + mdV: cute.Tensor, + sV: cute.Tensor, + acc_dK: cute.Tensor, + mdK: cute.Tensor, + sK: cute.Tensor, + seqlen: SeqlenInfoQK, + tma_atom_dK: cute.CopyAtom, + tma_atom_dV: cute.CopyAtom, + tiled_mma_dK: cute.TiledMma, + tiled_mma_dV: cute.TiledMma, + tidx: Int32, + n_block: Int32, + head_idx: Int32, + batch_idx: Int32, + qhead_per_kvhead_divmod: Optional[FastDivmodDivisor] = None, + mdK_semaphore: Optional[cute.Tensor] = None, + mdV_semaphore: Optional[cute.Tensor] = None, + ): + epi_barrier = cutlass.pipeline.NamedBarrier( + barrier_id=int(NamedBarrierBwd.Epilogue), num_threads=self.num_mma_threads + ) + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + if const_expr(self.qhead_per_kvhead == 1): + mdK_cur = seqlen.offset_batch_K(mdK, batch_idx, dim=3, ragged=self.varlen_k)[ + None, None, head_idx + ] + mdV_cur = seqlen.offset_batch_K(mdV, batch_idx, dim=3, ragged=self.varlen_k)[ + None, None, head_idx + ] + gdK = cute.local_tile(mdK_cur, (self.tile_n, self.tile_hdim), (n_block, 0)) + gdV = cute.local_tile(mdV_cur, (self.tile_n, self.tile_hdimv), (n_block, 0)) + store_dK, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_dK, 0, cute.make_layout(1), sK, gdK, single_stage=True + ) + store_dV, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_dV, 0, cute.make_layout(1), sV, gdV, single_stage=True + ) + sdV = sV if const_expr(not self.dKV_swapAB) else layout_utils.transpose_view(sV) + sdK = sK if const_expr(not self.dKV_swapAB) else layout_utils.transpose_view(sK) + copy_dV_r2s, _, _ = copy_utils.get_smem_store_C( + tiled_mma_dV, + sdV, + tidx, + transpose=self.dKV_swapAB, + position_independent=True, + ) + copy_dK_r2s, _, _ = copy_utils.get_smem_store_C( + tiled_mma_dK, + sdK, + tidx, + transpose=self.dKV_swapAB, + position_independent=True, + ) + cute.arch.cp_async_bulk_wait_group(1, read=True) + epi_barrier.arrive_and_wait() + copy_dV_r2s(acc_dV, dst_idx=None) + cute.arch.fence_view_async_shared() + epi_barrier.arrive_and_wait() + if warp_idx == 4: + store_dV() + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(1, read=True) + epi_barrier.arrive_and_wait() + copy_dK_r2s(acc_dK, dst_idx=None) + cute.arch.fence_view_async_shared() + epi_barrier.arrive_and_wait() + if warp_idx == 4: + store_dK() + cute.arch.cp_async_bulk_commit_group() + else: + deterministic_KV = self.deterministic and self.qhead_per_kvhead > 1 + sdKaccum_shape0 = self.tile_n * self.tile_hdim // self.num_wg_mma + sdVaccum_shape0 = self.tile_n * self.tile_hdimv // self.num_wg_mma + sdKaccum_layout = cute.make_layout((sdKaccum_shape0, self.num_wg_mma)) + sdVaccum_layout = cute.make_layout((sdVaccum_shape0, self.num_wg_mma)) + head_idx_kv = head_idx // qhead_per_kvhead_divmod + if const_expr(deterministic_KV): + assert mdK_semaphore is not None + assert mdV_semaphore is not None + mdK_semaphore_cur = mdK_semaphore[n_block, None, head_idx_kv, batch_idx] + mdV_semaphore_cur = mdV_semaphore[n_block, None, head_idx_kv, batch_idx] + lock_value = head_idx % self.qhead_per_kvhead + mdKaccum_cur = seqlen.offset_batch_K( + mdK, batch_idx, dim=2, padded=True, multiple=self.tile_hdim + )[None, head_idx_kv] + mdVaccum_cur = seqlen.offset_batch_K( + mdV, batch_idx, dim=2, padded=True, multiple=self.tile_hdimv + )[None, head_idx_kv] + gdKaccum_ = cute.local_tile(mdKaccum_cur, (self.tile_n * self.tile_hdim,), (n_block,)) + gdKaccum = cute.flat_divide(gdKaccum_, (sdKaccum_shape0,)) + gdVaccum_ = cute.local_tile(mdVaccum_cur, (self.tile_n * self.tile_hdimv,), (n_block,)) + gdVaccum = cute.flat_divide(gdVaccum_, (sdVaccum_shape0,)) + # These two overlap each other + sVaccum_ptr = cute.recast_ptr(sV.iterator, dtype=Float32) + sdKaccum = cute.make_tensor(sVaccum_ptr, sdKaccum_layout) + sdVaccum = cute.make_tensor(sVaccum_ptr, sdVaccum_layout) + tiled_copy_dKVaccum_r2s = cute.make_tiled_copy_tv( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32, num_bits_per_copy=128), + cute.make_layout((self.num_threads_per_warp_group, self.num_wg_mma)), + cute.make_layout(128 // Float32.width), + ) + thr_copy_dKVaccum_r2s = tiled_copy_dKVaccum_r2s.get_slice(tidx) + tdKsdKaccum = thr_copy_dKVaccum_r2s.partition_D(sdKaccum) + tdVsdVaccum = thr_copy_dKVaccum_r2s.partition_D(sdVaccum) + + read_flag = const_expr(not deterministic_KV) + cute.arch.cp_async_bulk_wait_group(0, read=read_flag) + if const_expr(deterministic_KV): + barrier.wait_eq(mdK_semaphore_cur.iterator, tidx, 0, lock_value) + epi_barrier.arrive_and_wait() + tdKrdKaccum_flat = cute.make_tensor(acc_dK.iterator, tdKsdKaccum.shape) + cute.autovec_copy(tdKrdKaccum_flat, tdKsdKaccum) + cute.arch.fence_view_async_shared() + epi_barrier.arrive_and_wait() + if warp_idx == 4: + with cute.arch.elect_one(): + for wg_idx in cutlass.range_constexpr(self.num_wg_mma): + copy_utils.cpasync_reduce_bulk_add_f32( + sdKaccum[None, wg_idx].iterator, + gdKaccum[None, wg_idx].iterator, + self.tma_copy_bytes["dKacc"] // self.num_wg_mma, + ) + cute.arch.cp_async_bulk_commit_group() + + cute.arch.cp_async_bulk_wait_group(0, read=read_flag) + if const_expr(deterministic_KV): + barrier.arrive_inc(mdK_semaphore_cur.iterator, tidx, 0, 1) + barrier.wait_eq(mdV_semaphore_cur.iterator, tidx, 0, lock_value) + epi_barrier.arrive_and_wait() + tdVrdVaccum_flat = cute.make_tensor(acc_dV.iterator, tdVsdVaccum.shape) + cute.autovec_copy(tdVrdVaccum_flat, tdVsdVaccum) + cute.arch.fence_view_async_shared() + epi_barrier.arrive_and_wait() + if warp_idx == 4: + with cute.arch.elect_one(): + for wg_idx in cutlass.range_constexpr(self.num_wg_mma): + copy_utils.cpasync_reduce_bulk_add_f32( + sdVaccum[None, wg_idx].iterator, + gdVaccum[None, wg_idx].iterator, + self.tma_copy_bytes["dVacc"] // self.num_wg_mma, + ) + cute.arch.cp_async_bulk_commit_group() + if const_expr(deterministic_KV): + cute.arch.cp_async_bulk_wait_group(0, read=read_flag) + barrier.arrive_inc(mdV_semaphore_cur.iterator, tidx, 0, 1) + + @cute.jit + def dQaccum_store( + self, + mdQaccum: cute.Tensor, + sdQaccum: cute.Tensor, + block_info: BlockInfo, + TileSchedulerCls: cutlass.Constexpr[Callable], + SeqlenInfoCls: cutlass.Constexpr[Callable], + blocksparse_tensors: Optional[BlockSparseTensors] = None, + mdQ_semaphore: Optional[cute.Tensor] = None, + ): + tidx, _, _ = cute.arch.thread_idx() + # warp-local thread index (dQaccum_store runs on warp 1, global tidx 32-63) + warp_local_tidx = tidx % cute.arch.WARP_SIZE + read_flag = const_expr(not self.deterministic) + + tile_scheduler = TileSchedulerCls() + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + n_block, head_idx, batch_idx, _ = work_tile.tile_idx + seqlen = SeqlenInfoCls(batch_idx) + if const_expr(not seqlen.has_cu_seqlens_q): + mdQaccum_cur = mdQaccum[None, head_idx, batch_idx] + else: + mdQaccum_cur = cute.domain_offset( + (seqlen.padded_offset_q * self.tile_hdim,), mdQaccum[None, head_idx] + ) + # ((M * K / num_wg_dQ, num_wg_dQ), num_m_blocks) + gdQaccum = cute.local_tile( + mdQaccum_cur, + ( + cute.make_layout( + (self.tile_m * self.tile_hdim // self.num_wg_dQ, self.num_wg_dQ) + ), + ), + (None,), + ) + + if const_expr(mdQ_semaphore is not None): + # mdQ_semaphore is (num_m_blocks, cluster_size, num_head, batch) after transpose + mdQ_semaphore_cur = mdQ_semaphore[None, None, head_idx, batch_idx] + + m_block_min, m_block_max = block_info.get_m_block_min_max(seqlen, n_block) + if const_expr(not self.use_block_sparsity): + process_tile = ( + const_expr(not self.is_local and not self.is_varlen_q) + or m_block_min < m_block_max + ) + loop_count = m_block_max - m_block_min + else: + total_block_cnt = get_total_q_block_count_bwd( + blocksparse_tensors, + batch_idx, + head_idx, + n_block, + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + ) + process_tile = total_block_cnt > Int32(0) + + if process_tile: + if const_expr(not self.use_block_sparsity): + for iter_idx in cutlass.range(loop_count, unroll=1): + m_block = m_block_min + iter_idx + m_block_safe = m_block + + num_dQ_chunks = self.num_wg_dQ + for warp_group_idx in cutlass.range_constexpr(num_dQ_chunks): + if const_expr(not self.deterministic): + # If deterministic, we already waited at the end of the prev iter + cute.arch.cp_async_bulk_wait_group( + num_dQ_chunks - 1 - warp_group_idx, read=read_flag + ) + cute.arch.barrier_arrive( + barrier_id=int(NamedBarrierBwd.dQEmptyWG0) + warp_group_idx, + number_of_threads=self.num_threads_per_warp_group + + cute.arch.WARP_SIZE, + ) + + # Semaphore acquire: wait for prior n_blocks to finish writing this m_block + if const_expr(self.deterministic): + if const_expr(self.spt): + _, n_block_max_for_m_block = block_info.get_n_block_min_max( + seqlen, m_block_safe + ) + lock_value = n_block_max_for_m_block - 1 - n_block + else: + lock_value = n_block + barrier.wait_eq( + mdQ_semaphore_cur[(m_block_safe, None)].iterator, + warp_local_tidx, + 0, # flag_offset + lock_value, + ) + + for warp_group_idx in cutlass.range_constexpr(num_dQ_chunks): + cute.arch.barrier( + barrier_id=int(NamedBarrierBwd.dQFullWG0) + warp_group_idx, + number_of_threads=self.num_threads_per_warp_group + + cute.arch.WARP_SIZE, + ) + with cute.arch.elect_one(): + copy_utils.cpasync_reduce_bulk_add_f32( + sdQaccum[None, warp_group_idx].iterator, + gdQaccum[(None, warp_group_idx), m_block_safe].iterator, + self.tma_copy_bytes["dQ"], + ) + cute.arch.cp_async_bulk_commit_group() + + # Semaphore release: signal that this n_block is done with this m_block + if const_expr(self.deterministic): + cute.arch.cp_async_bulk_wait_group(0, read=read_flag) + barrier.arrive_inc( + mdQ_semaphore_cur[(m_block_safe, None)].iterator, + warp_local_tidx, + 0, # flag_offset + 1, + ) + else: + assert not self.deterministic, ( + "Deterministic not implemented for block-sparse backward" + ) + dQaccum_store_block_sparse_bwd_sm90( + blocksparse_tensors, + batch_idx, + head_idx, + n_block, + sdQaccum, + gdQaccum, + q_subtile_factor=self.q_subtile_factor, + m_block_max=m_block_max, + num_dQ_warp_groups=self.num_wg_dQ, + num_threads_per_warp_group=self.num_threads_per_warp_group, + tma_copy_bytes_dQ=self.tma_copy_bytes["dQ"], + ) + + # For local masking + deterministic (non-spt): signal remaining m_blocks + # that this n_block won't visit, so they don't deadlock waiting. + if const_expr( + self.deterministic and not self.spt and block_info.window_size_left is not None + ): + m_block_global_max = cute.ceil_div(seqlen.seqlen_q, self.tile_m) + for m_block in cutlass.range(m_block_max, m_block_global_max, unroll=1): + barrier.arrive_inc( + mdQ_semaphore_cur[(m_block, None)].iterator, + warp_local_tidx, + 0, # flag_offset + 1, + ) + + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + + if const_expr(not self.deterministic): + cute.arch.cp_async_bulk_wait_group(0, read=True) diff --git a/fastvideo-kernel/fa4/flash_fwd.py b/fastvideo-kernel/fa4/flash_fwd.py new file mode 100644 index 000000000..5d573b933 --- /dev/null +++ b/fastvideo-kernel/fa4/flash_fwd.py @@ -0,0 +1,1236 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +# A reimplementation of +# https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_fwd_kernel_sm80.h +# and https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_fwd_kernel_sm90.h +# from Cutlass C++ to Cute-DSL. +# Built on Cute-DSL example: https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/ampere/flash_attention_v2.py + +import math +from types import SimpleNamespace +from typing import Type, Callable, Optional +from functools import partial + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, const_expr +from cutlass.cute.nvgpu import cpasync, warp +import cutlass.utils as utils_basic +from cutlass.base_dsl.arch import Arch +from cutlass.cutlass_dsl import BaseDSL + +from quack import copy_utils +from quack import layout_utils + +from flash_attn.cute import ampere_helpers as sm80_utils +from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned +from flash_attn.cute import utils +from flash_attn.cute.mask import AttentionMask +from flash_attn.cute.softmax import Softmax, apply_score_mod_inner +from flash_attn.cute.seqlen_info import SeqlenInfoQK +from flash_attn.cute.block_info import BlockInfo +from flash_attn.cute.pack_gqa import PackGQA +from flash_attn.cute.named_barrier import NamedBarrierFwd +from flash_attn.cute.block_sparsity import BlockSparseTensors +from flash_attn.cute.tile_scheduler import SingleTileScheduler, SingleTileVarlenScheduler, TileSchedulerArguments +from flash_attn.cute.utils import AuxData + + +class FlashAttentionForwardBase: + + def __init__( + self, + dtype: Type[cutlass.Numeric], + head_dim: int, + head_dim_v: Optional[int] = None, + qhead_per_kvhead: int = 1, + is_causal: bool = False, + is_local: bool = False, + pack_gqa: bool = True, + tile_m: int = 128, + tile_n: int = 128, + num_stages: int = 1, + num_threads: int = 128, + Q_in_regs: bool = False, + score_mod: Optional[cutlass.Constexpr] = None, + mask_mod: Optional[cutlass.Constexpr] = None, + has_aux_tensors: bool = False, + q_subtile_factor: int = 1, + ): + """Initializes the configuration for a flash attention kernel. + + All contiguous dimensions must be at least 16 bytes aligned, which means that the head dimension + should be a multiple of 8. + + :param head_dim: head dimension + :type head_dim: int + :param tile_m: m block size + :type tile_m: int + :param tile_n: n block size + :type tile_n: int + :param num_threads: number of threads + :type num_threads: int + :param is_causal: is causal + :param score_mod: A callable that takes the attention scores and applies a modification. + Callable signature: ``score_mod(scores, batch_idx, head_idx, q_idx, kv_idx, aux_tensors) -> Any`` + :param mask_mod: A callable that takes the attention scores and returns a boolean representing whether that score should be masked. + Callable signature: ``mask_mod(batch_idx, head_idx, q_idx, kv_idx, aux_tensors) -> Boolean`` + """ + self.dtype = dtype + # padding head_dim to a multiple of 16 as k_block_size + hdim_multiple_of = 16 + self.tile_hdim = int(math.ceil(head_dim / hdim_multiple_of) * hdim_multiple_of) + head_dim_v = head_dim_v if head_dim_v is not None else head_dim + self.same_hdim_kv = head_dim == head_dim_v + self.tile_hdimv = int(math.ceil(head_dim_v / hdim_multiple_of) * hdim_multiple_of) + # Can save registers (and hence be faster) if we don't have to check hdim predication + self.check_hdim_oob = head_dim != self.tile_hdim + self.check_hdim_v_oob = head_dim_v != self.tile_hdimv + self.qhead_per_kvhead = qhead_per_kvhead + self.is_causal = is_causal + self.is_local = is_local + self.pack_gqa = pack_gqa + self.tile_m = tile_m + self.tile_n = tile_n + self.num_threads = num_threads + self.num_stages = num_stages + self.q_subtile_factor = q_subtile_factor + self.Q_in_regs = Q_in_regs + self.score_mod = score_mod + self.mask_mod = mask_mod + self.qk_acc_dtype = Float32 + self.score_vec_size: cutlass.Constexpr = getattr( + score_mod, "__vec_size__", 1 if cutlass.const_expr(has_aux_tensors) else 2 + ) + if self.score_vec_size > 2: + raise ValueError( + f"score_mod vec_size {self.score_vec_size} not supported on Sm80/90/120 " + "due to accumulator thread ownership pattern." + ) + self.mask_vec_size: cutlass.Constexpr = getattr(mask_mod, "__vec_size__", 1) + if self.mask_vec_size > 1: + raise ValueError( + f"mask_mod vec_size {self.mask_vec_size} not supported on Sm80/90/120 " + "due to accumulator thread ownership pattern." + ) + self.arch = BaseDSL._get_dsl().get_arch_enum() + + @staticmethod + def can_implement( + dtype, + head_dim, + head_dim_v, + tile_m, + tile_n, + num_stages, + num_threads, + is_causal, + Q_in_regs=False, + ) -> bool: + """Check if the kernel can be implemented with the given parameters. + + :param dtype: data type + :type dtype: cutlass.Numeric + :param head_dim: head dimension + :type head_dim: int + :param tile_m: m block size + :type tile_m: int + :param tile_n: n block size + :type tile_n: int + :param num_threads: number of threads + :type num_threads: int + :param is_causal: is causal + :type is_causal: bool + + :return: True if the kernel can be implemented, False otherwise + :rtype: bool + """ + if dtype not in [cutlass.Float16, cutlass.BFloat16]: + return False + if head_dim % 8 != 0: + return False + if head_dim_v % 8 != 0: + return False + if tile_n % 16 != 0: + return False + if num_threads % 32 != 0: + return False + # Check if block size setting is out of shared memory capacity + # Shared memory usage: Q tile + (K tile + V tile) where K and V use the same tile size + smem_usage_Q = tile_m * head_dim * 2 + smem_usage_K = tile_n * head_dim * num_stages * 2 + smem_usage_V = tile_n * head_dim_v * num_stages * 2 + smem_usage_QV = ( + (smem_usage_Q + smem_usage_V) if not Q_in_regs else max(smem_usage_Q, smem_usage_V) + ) + smem_usage = smem_usage_QV + smem_usage_K + # TODO: sm86 and sm89 + smem_capacity = utils_basic.get_smem_capacity_in_bytes("sm_80") + if smem_usage > smem_capacity: + return False + # Check if twice the block size is divisible by the number of threads + if (tile_m * 2) % num_threads != 0: + return False + return True + + def _check_type( + self, + mQ_type: Type[cutlass.Numeric], + mK_type: Type[cutlass.Numeric], + mV_type: Type[cutlass.Numeric], + mO_type: Type[cutlass.Numeric], + mLSE_type: Type[cutlass.Numeric] | None, + mCuSeqlensQ_type: Type[cutlass.Numeric] | None, + mCuSeqlensK_type: Type[cutlass.Numeric] | None, + mSeqUsedQ_type: Type[cutlass.Numeric] | None, + mSeqUsedK_type: Type[cutlass.Numeric] | None, + ): + # Get the data type and check if it is fp16 or bf16 + if const_expr(not (mQ_type == mK_type == mV_type == mO_type)): + raise TypeError("All tensors must have the same data type") + if const_expr(mQ_type not in [cutlass.Float16, cutlass.BFloat16]): + raise TypeError("Only Float16 or BFloat16 is supported") + if const_expr(mLSE_type not in [None, Float32]): + raise TypeError("LSE tensor must be Float32") + if const_expr(mCuSeqlensQ_type not in [None, Int32]): + raise TypeError("cu_seqlens_q tensor must be Int32") + if const_expr(mCuSeqlensK_type not in [None, Int32]): + raise TypeError("cu_seqlens_k tensor must be Int32") + if const_expr(mSeqUsedQ_type not in [None, Int32]): + raise TypeError("seqused_q tensor must be Int32") + if const_expr(mSeqUsedK_type not in [None, Int32]): + raise TypeError("seqused_k tensor must be Int32") + assert mQ_type == self.dtype + + def _setup_attributes(self): + # /////////////////////////////////////////////////////////////////////////////// + # Shared memory layout: Q/K/V + # /////////////////////////////////////////////////////////////////////////////// + sQ_layout_atom, sK_layout_atom, sV_layout_atom, sO_layout_atom, sP_layout_atom = ( + self._get_smem_layout_atom() + ) + self.sQ_layout = cute.tile_to_shape( + sQ_layout_atom, + (self.tile_m, self.tile_hdim), + (0, 1), + ) + self.sK_layout = cute.tile_to_shape( + sK_layout_atom, + (self.tile_n, self.tile_hdim, self.num_stages), + (0, 1, 2), + ) + self.sV_layout = cute.tile_to_shape( + sV_layout_atom, + (self.tile_n, self.tile_hdimv, self.num_stages), + (0, 1, 2), + ) + self.sO_layout = cute.tile_to_shape( + sO_layout_atom, + (self.tile_m, self.tile_hdimv), + (0, 1), + ) + if const_expr(sP_layout_atom is not None): + self.sP_layout = cute.tile_to_shape( + sP_layout_atom, + (self.tile_m, self.tile_n), + (0, 1), + ) + else: + self.sP_layout = None + + # /////////////////////////////////////////////////////////////////////////////// + # GMEM Tiled copy: + # /////////////////////////////////////////////////////////////////////////////// + # Thread layouts for copies + universal_copy_bits = 128 + async_copy_elems = universal_copy_bits // self.dtype.width + # atom_async_copy: async copy atom for QKV load + atom_async_copy = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + self.dtype, + num_bits_per_copy=universal_copy_bits, + ) + # atom_universal_copy: universal copy atom for O store + atom_universal_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.dtype, + num_bits_per_copy=universal_copy_bits, + ) + # tQ_layout and tK_layout: thread layout for QK load + tQK_shape_dim_1 = sQ_layout_atom.outer.shape[1] // async_copy_elems + assert self.num_Q_load_threads % tQK_shape_dim_1 == 0, ( + "num_threads must be divisible by tQK_shape_dim_1" + ) + assert self.num_producer_threads % tQK_shape_dim_1 == 0, ( + "num_threads must be divisible by tQK_shape_dim_1" + ) + tQ_layout = cute.make_ordered_layout( + (self.num_Q_load_threads // tQK_shape_dim_1, tQK_shape_dim_1), + order=(1, 0), + ) + tK_layout = cute.make_ordered_layout( + (self.num_producer_threads // tQK_shape_dim_1, tQK_shape_dim_1), + order=(1, 0), + ) + # So that we don't have to check if we overshoot kBlockM when we load Q + assert self.tile_m % tQ_layout.shape[0] == 0 + tV_shape_dim_1 = sV_layout_atom.outer.shape[1] // async_copy_elems + tV_layout = cute.make_ordered_layout( + (self.num_producer_threads // tV_shape_dim_1, tV_shape_dim_1), + order=(1, 0), + ) + # TODO: need a different layout for O if O dtype is not the same as V dtype + # tO_layout: thread layout for O store + tO_layout = cute.make_ordered_layout( + (self.num_epilogue_threads // tV_shape_dim_1, tV_shape_dim_1), + order=(1, 0), + ) + # So that we don't have to check if we overshoot kBlockM when we store O + assert self.tile_m % tO_layout.shape[0] == 0 + + # Value layouts for copies + vQKV_layout = cute.make_layout((1, async_copy_elems)) + vO_layout = vQKV_layout + + self.gmem_tiled_copy_Q = cute.make_tiled_copy_tv(atom_async_copy, tQ_layout, vQKV_layout) + self.gmem_tiled_copy_K = cute.make_tiled_copy_tv(atom_async_copy, tK_layout, vQKV_layout) + self.gmem_tiled_copy_V = cute.make_tiled_copy_tv(atom_async_copy, tV_layout, vQKV_layout) + # gmem_tiled_copy_O: tiled copy for O store + self.gmem_tiled_copy_O = cute.make_tiled_copy_tv(atom_universal_copy, tO_layout, vO_layout) + + def _get_smem_layout_atom(self): + raise NotImplementedError() + + def _get_tiled_mma(self): + raise NotImplementedError() + + def _get_shared_storage_cls(self): + raise NotImplementedError() + + @cute.jit + def __call__( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + softmax_scale: Float32, + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + """Configures and launches the flash attention kernel. + + mQ/mK/mV/mO has same data types(supports fp16 and bf16) and same layout: + (batch_size, seqlen_q, num_head, head_dim):(_, _, _, 1) + """ + raise NotImplementedError() + + @cute.jit + def epilogue( + self, + acc_O: cute.Tensor, + lse: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + sO: cute.Tensor, + seqlen: SeqlenInfoQK, + gmem_tiled_copy_O: cute.TiledCopy, + tma_atom_O: Optional[cute.CopyAtom], + tiled_mma: cute.TiledMma, + tidx: Int32, + m_block: Int32, + head_idx: Int32, + batch_idx: Int32, + ): + # store acc_O + rO = cute.make_fragment_like(acc_O, self.dtype) + rO.store(acc_O.load().to(self.dtype)) + # Make sure all threads have finished reading V + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_epilogue_threads + ) + smem_copy_atom_O = utils.get_smem_store_atom(self.arch.major * 10 + self.arch.minor, self.dtype) + smem_thr_copy_O = cute.make_tiled_copy_C(smem_copy_atom_O, tiled_mma).get_slice(tidx) + taccOrO = smem_thr_copy_O.retile(rO) + taccOsO = smem_thr_copy_O.partition_D(sO) + # taccOsO = copy_utils.partition_D_position_independent(smem_thr_copy_O, sO) + # copy acc O from rmem to smem with the smem copy atom + cute.copy(smem_copy_atom_O, taccOrO, taccOsO) + + cO = cute.make_identity_tensor((self.tile_m, self.tile_hdimv)) + pack_gqa = PackGQA( + self.tile_m, self.tile_hdimv, self.check_hdim_v_oob, self.qhead_per_kvhead + ) + + # Write LSE from rmem -> gmem + if const_expr(mLSE is not None): + mLSE_cur = seqlen.offset_batch_Q(mLSE, batch_idx, dim=2)[None, head_idx] + if const_expr(not self.pack_gqa): + gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (m_block,)) + gLSE_expanded_layout = cute.append( + gLSE.layout, cute.make_layout((self.tile_hdimv,), stride=(0,)) + ) + gLSE_expanded = cute.make_tensor(gLSE.iterator, gLSE_expanded_layout) + thr_mma = tiled_mma.get_slice(tidx) + taccOgLSE = layout_utils.reshape_acc_to_mn(thr_mma.partition_C(gLSE_expanded)) + assert cute.size(taccOgLSE, mode=[0]) == cute.size(lse) + taccOcO = layout_utils.reshape_acc_to_mn(thr_mma.partition_C(cO)) + t0accOcO = layout_utils.reshape_acc_to_mn(thr_mma.get_slice(0).partition_C(cO)) + # Only the thread corresponding to column 0 writes out the lse to gmem + if taccOcO[0][1] == 0: + for m in cutlass.range(cute.size(taccOgLSE.shape[1]), unroll_full=True): + if ( + t0accOcO[m, 0][0] + < seqlen.seqlen_q - m_block * self.tile_m - taccOcO[0][0] + ): + taccOgLSE[m, 0] = lse[m] + else: + pack_gqa.store_LSE(mLSE_cur, lse, tiled_mma, tidx, m_block, seqlen.seqlen_q) + + ragged = self.use_tma_O and (seqlen.has_cu_seqlens_q or seqlen.has_seqused_q) + mO_cur = seqlen.offset_batch_Q(mO, batch_idx, dim=3, ragged=ragged)[None, None, head_idx] + # thr_mma = tiled_mma.get_slice(tidx) + # taccOgO = thr_mma.partition_C(gO) + # cute.autovec_copy(rO, taccOgO) + # sync to make sure all smem stores are done + if const_expr(self.use_tma_O): + # ensure smem writes are visible to TMA + cute.arch.fence_view_async_shared() + cute.arch.barrier_arrive( + barrier_id=int(NamedBarrierFwd.Epilogue), + number_of_threads=self.num_epilogue_threads + cute.arch.WARP_SIZE, + ) + gO = cute.local_tile(mO_cur, (self.tile_m, self.tile_hdimv), (m_block, 0)) + store_O, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_O, 0, cute.make_layout(1), sO, gO, single_stage=True + ) + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + if warp_idx == 4: + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), + number_of_threads=self.num_epilogue_threads + cute.arch.WARP_SIZE, + ) + store_O() + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + else: + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), + number_of_threads=self.num_epilogue_threads, + ) + gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx) + tOsO = gmem_thr_copy_O.partition_S(sO) + tOrO = cute.make_fragment_like(tOsO, self.dtype) + # load acc O from smem to rmem for wider vectorization + cute.autovec_copy(tOsO, tOrO) + if const_expr(not self.pack_gqa): + gO = cute.local_tile(mO_cur, (self.tile_m, self.tile_hdimv), (m_block, 0)) + tOgO = gmem_thr_copy_O.partition_D(gO) + tOcO = gmem_thr_copy_O.partition_S(cO) + t0OcO = gmem_tiled_copy_O.get_slice(0).partition_S(cO) + tOpO = utils.predicate_k(tOcO, limit=mO.shape[1]) + # copy acc O from rmem to gmem + for rest_m in cutlass.range_constexpr(cute.size(tOrO.shape[1])): + if ( + t0OcO[0, rest_m, 0][0] + < seqlen.seqlen_q - m_block * self.tile_m - tOcO[0][0] + ): + cute.copy( + gmem_tiled_copy_O, + tOrO[None, rest_m, None], + tOgO[None, rest_m, None], + pred=tOpO[None, rest_m, None] + if const_expr(self.check_hdim_v_oob) + else None, + ) + else: + pack_gqa.store_O(mO_cur, tOrO, gmem_tiled_copy_O, tidx, m_block, seqlen.seqlen_q) + + @cute.jit + def advance_pipeline(self, pipeline_index): + return pipeline_index + 1 if pipeline_index < self.num_stages - 1 else 0 + + @cute.jit + def load_Q( + self, + gmem_thr_copy: cute.TiledCopy, + gQ: cute.Tensor, + sQ: cute.Tensor, + block: Int32, + seqlen: Int32, + headdim: Int32, + ): + tQsQ, tQgQ = gmem_thr_copy.partition_D(sQ), gmem_thr_copy.partition_S(gQ) + cQ = cute.make_identity_tensor((self.tile_m, self.tile_hdim)) + tQcQ = gmem_thr_copy.partition_S(cQ) + t0QcQ = gmem_thr_copy.get_slice(0).partition_S(cQ) + tQpQ = utils.predicate_k(tQcQ, limit=headdim) + for m in cutlass.range_constexpr(cute.size(tQsQ.shape[1])): + # Instead of using tQcQ, we using t0QcQ and subtract the offset from the limit + # (seqlen - block * kBlockM). This is because the entries of t0QcQ are known at compile time. + if t0QcQ[0, m, 0][0] < seqlen - block * self.tile_m - tQcQ[0][0]: + cute.copy( + gmem_thr_copy, + tQgQ[None, m, None], + tQsQ[None, m, None], + pred=tQpQ[None, m, None] if const_expr(self.check_hdim_oob) else None, + ) + # We don't need to clear the sQ smem tiles since we'll only write out the valid outputs + + @cute.jit + def load_K( + self, + gmem_tiled_copy: cute.TiledCopy, + tKgK: cute.Tensor, + tKsK: cute.Tensor, + tKcK: cute.Tensor, + t0KcK: cute.Tensor, + tKpK: cute.Tensor, + block: Int32, + smem_pipe_write: Int32, + seqlen: Int32, + need_predicates: cutlass.Constexpr, + ): + # Do we need to check if we overshoot kBlockN when we load K? + is_even_n_smem_k = self.tile_n % gmem_tiled_copy.tiler_mn[0].shape == 0 + if const_expr(need_predicates or not is_even_n_smem_k): + # Instead of using tKcK, we using t0KcK and subtract the offset from the limit + # (seqlen - block * kBlockN). This is because the entries of t0KcK are known at compile time. + if const_expr(is_even_n_smem_k): + seqlen_limit = seqlen - block * self.tile_n + else: + if const_expr(not need_predicates): + seqlen_limit = self.tile_n + else: + seqlen_limit = cutlass.min(seqlen - block * self.tile_n, self.tile_n) + seqlen_limit -= tKcK[0][0] + for n in cutlass.range_constexpr(cute.size(tKsK.shape[1])): + if t0KcK[0, n, 0][0] < seqlen_limit: + cute.copy( + gmem_tiled_copy, + tKgK[None, n, None, block], + tKsK[ + None, n, None, smem_pipe_write if const_expr(self.num_stages > 1) else 0 + ], + pred=tKpK[None, n, None] if const_expr(self.check_hdim_oob) else None, + ) + # We don't need to clear the sK smem tiles since we'll mask out the scores anyway. + else: + cute.copy( + gmem_tiled_copy, + tKgK[None, None, None, block], + tKsK[None, None, None, smem_pipe_write if const_expr(self.num_stages > 1) else 0], + pred=tKpK if const_expr(self.check_hdim_oob) else None, + ) + + @cute.jit + def load_V( + self, + gmem_tiled_copy: cute.TiledCopy, + tVgV: cute.Tensor, + tVsV: cute.Tensor, + tVcV: cute.Tensor, + t0VcV: cute.Tensor, + tVpV: cute.Tensor, + block: Int32, + smem_pipe_write: Int32, + seqlen: Int32, + need_predicates: cutlass.Constexpr, + ): + # Do we need to check if we overshoot kBlockN when we load V? + is_even_n_smem_v = self.tile_n % gmem_tiled_copy.tiler_mn[0].shape == 0 + if const_expr(need_predicates or not is_even_n_smem_v): + for n in cutlass.range_constexpr(cute.size(tVsV.shape[1])): + # If kBlockN doesn't evenly divide the tiled copy, only the last `n` needs to be checked + if ( + is_even_n_smem_v + or n < cute.size(tVsV.shape[1]) - 1 + or tVcV[0, n, 0][0] < self.tile_n + ): + predicate = tVpV[None, n, None] if const_expr(self.check_hdim_v_oob) else None + if const_expr(need_predicates): + seqlen_limit = seqlen - block * self.tile_n - tVcV[0][0] + predicate_n = t0VcV[0, n, 0][0] < seqlen_limit + predicate = cute.make_fragment_like(tVpV[None, 0, None]) + for k in cutlass.range_constexpr(cute.size(predicate.shape[1])): + for i in cutlass.range_constexpr(cute.size(predicate.shape[0])): + predicate[i, k] = ( + tVpV[i, n, k] if const_expr(self.check_hdim_v_oob) else True + ) and predicate_n + cute.copy( + gmem_tiled_copy, + tVgV[None, n, None, block], + tVsV[ + None, n, None, smem_pipe_write if const_expr(self.num_stages > 1) else 0 + ], + pred=predicate, + ) + else: + cute.copy( + gmem_tiled_copy, + tVgV[None, None, None, block], + tVsV[None, None, None, smem_pipe_write if const_expr(self.num_stages > 1) else 0], + pred=tVpV if const_expr(self.check_hdim_v_oob) else None, + ) + + +class FlashAttentionForwardSm80(FlashAttentionForwardBase): + def _get_smem_layout_atom(self): + sQ_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, self.tile_hdim) + sK_layout_atom = sQ_layout_atom + sV_layout_atom = sm80_utils.get_smem_layout_atom(self.dtype, self.tile_hdimv) + sO_layout_atom = sV_layout_atom + sP_layout_atom = None + return sQ_layout_atom, sK_layout_atom, sV_layout_atom, sO_layout_atom, sP_layout_atom + + def _get_tiled_mma(self): + tiled_mma_qk = cute.make_tiled_mma( + warp.MmaF16BF16Op(self.dtype, Float32, (16, 8, 16)), + (self.num_threads // 32, 1, 1), + permutation_mnk=(self.num_threads // 32 * 16, 16, 16), + ) + tiled_mma_pv = cute.make_tiled_mma( + warp.MmaF16BF16Op(self.dtype, Float32, (16, 8, 16)), + (self.num_threads // 32, 1, 1), + permutation_mnk=(self.num_threads // 32 * 16, 16, 16), + ) + return tiled_mma_qk, tiled_mma_pv + + def _get_shared_storage_cls(self): + sQ_struct, sK_struct, sV_struct = [ + cute.struct.Align[cute.struct.MemRange[self.dtype, cute.cosize(layout)], 1024] + for layout in (self.sQ_layout, self.sK_layout, self.sV_layout) + ] + cosize_sQV = max(cute.cosize(self.sQ_layout), cute.cosize(self.sV_layout)) + sQV_struct = cute.struct.Align[cute.struct.MemRange[self.dtype, cosize_sQV], 1024] + + @cute.struct + class SharedStorageQKV: + sV: sV_struct + sQ: sQ_struct + sK: sK_struct + + @cute.struct + class SharedStorageSharedQV: + sQ: sQV_struct + sK: sK_struct + + return SharedStorageQKV if const_expr(not self.Q_in_regs) else SharedStorageSharedQV + + @cute.jit + def __call__( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + softmax_scale: Float32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + mPageTable: Optional[cute.Tensor] = None, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + learnable_sink: Optional[cute.Tensor] = None, + blocksparse_tensors: Optional[BlockSparseTensors] = None, + aux_data: AuxData = AuxData(), + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + """Configures and launches the flash attention kernel. + + mQ/mK/mV/mO has same data types(supports fp16 and bf16) and same layout: + (batch_size, seqlen_q, num_head, head_dim):(_, _, _, 1) + """ + assert learnable_sink is None, "Learnable sink is not supported in this kernel" + self._check_type( + *(t.element_type if t is not None else None for t in (mQ, mK, mV, mO, mLSE, mCuSeqlensQ, mCuSeqlensK, mSeqUsedQ, mSeqUsedK)) + ) + tiled_mma_qk, tiled_mma_pv = self._get_tiled_mma() + self.num_mma_threads = tiled_mma_pv.size + self.num_producer_threads = self.num_threads + self.num_Q_load_threads = self.num_threads + self.num_epilogue_threads = self.num_threads + # self.use_tma_O = self.arch >= 90 and mCuSeqlensQ is None + self.use_tma_O = self.arch >= Arch.sm_90 + self._setup_attributes() + SharedStorage = self._get_shared_storage_cls() + mQ, mK, mV, mO = [assume_tensor_aligned(t) for t in (mQ, mK, mV, mO)] + # Layout permutation: 4D non-varlen vs 3D varlen + QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1] + KV_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensK is None) else [0, 2, 1] + mQ, mO = [ + cute.make_tensor(t.iterator, cute.select(t.layout, mode=QO_layout_transpose)) + for t in (mQ, mO) + ] + mK, mV = [ + cute.make_tensor(t.iterator, cute.select(t.layout, mode=KV_layout_transpose)) + for t in (mK, mV) + ] + if const_expr(mLSE is not None): + LSE_layout_transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0] + mLSE = cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose)) + # TileScheduler for varlen, simple grid for non-varlen + if const_expr(mCuSeqlensQ is not None or mSeqUsedQ is not None): + TileScheduler = SingleTileVarlenScheduler + else: + TileScheduler = SingleTileScheduler + num_batch = ( + mCuSeqlensQ.shape[0] - 1 + if const_expr(mCuSeqlensQ is not None) + else mQ.shape[3] + ) + tile_sched_args = TileSchedulerArguments( + num_block=cute.ceil_div(mQ.shape[0], self.tile_m), + num_head=cute.size(mQ.shape[2]), + num_batch=num_batch, + num_splits=1, + seqlen_k=0, + headdim=mQ.shape[1], + headdim_v=mV.shape[1], + total_q=cute.size(mQ.shape[0]) + if const_expr(mCuSeqlensQ is not None) + else cute.size(mQ.shape[0]) * cute.size(mQ.shape[3]), + tile_shape_mn=(self.tile_m, self.tile_n), + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + mCuSeqlensQ=mCuSeqlensQ, + mSeqUsedQ=mSeqUsedQ, + ) + tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args) + grid_dim = TileScheduler.get_grid_shape(tile_sched_params) + softmax_scale_log2, softmax_scale = utils.compute_softmax_scale_log2(softmax_scale, self.score_mod) + fastdiv_mods = utils.compute_fastdiv_mods(mQ, mK, self.qhead_per_kvhead, self.pack_gqa, aux_data.tensors) + + self.kernel( + mQ, + mK, + mV, + mO, + mLSE, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + softmax_scale_log2, + softmax_scale, + window_size_left, + window_size_right, + self.sQ_layout, + self.sK_layout, + self.sV_layout, + self.sO_layout, + self.sP_layout, + self.gmem_tiled_copy_Q, + self.gmem_tiled_copy_K, + self.gmem_tiled_copy_V, + self.gmem_tiled_copy_O, + tiled_mma_qk, + tiled_mma_pv, + SharedStorage, + tile_sched_params, + TileScheduler, + aux_data, + fastdiv_mods, + ).launch( + grid=grid_dim, + block=[self.num_threads, 1, 1], + smem=SharedStorage.size_in_bytes(), + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + mCuSeqlensQ: Optional[cute.Tensor], + mCuSeqlensK: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + softmax_scale_log2: Float32, + softmax_scale: Optional[Float32], + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + sQ_layout: cute.ComposedLayout, + sK_layout: cute.ComposedLayout, + sV_layout: cute.ComposedLayout, + sO_layout: cute.ComposedLayout, + sP_layout: cute.ComposedLayout | None, + gmem_tiled_copy_Q: cute.TiledCopy, + gmem_tiled_copy_K: cute.TiledCopy, + gmem_tiled_copy_V: cute.TiledCopy, + gmem_tiled_copy_O: cute.TiledCopy, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + SharedStorage: cutlass.Constexpr, + tile_sched_params, + TileScheduler: cutlass.Constexpr[Callable], + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + ): + # Thread index, block index + tidx, _, _ = cute.arch.thread_idx() + + tile_scheduler = TileScheduler.create(tile_sched_params) + work_tile = tile_scheduler.initial_work_tile_info() + m_block, num_head, batch_size, _ = work_tile.tile_idx + + block_info = BlockInfo( + self.tile_m, + self.tile_n, + self.is_causal, + self.is_local, + False, # is_split_kv + window_size_left, + window_size_right, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + seqlen = SeqlenInfoQK.create( + batch_idx=batch_size, + seqlen_q_static=mQ.shape[0], + seqlen_k_static=mK.shape[0], + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + ) + n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block) + # For varlen, wasted grid tiles (where batch_idx >= num_batch) will have + # seqlen_q=seqlen_k=0 and n_block_max=0. Clamp to 0 so we don't use a + # negative block index for K/V loads; the load/store predicates already + # guard all memory accesses when seqlen is 0. + n_block = cutlass.max(n_block_max - 1, 0) + + # /////////////////////////////////////////////////////////////////////////////// + # Get the appropriate tiles for this thread block. + # /////////////////////////////////////////////////////////////////////////////// + blkQ_shape = (self.tile_m, self.tile_hdim) + blkK_shape = (self.tile_n, self.tile_hdim) + blkV_shape = (self.tile_n, self.tile_hdimv) + num_head_kv = num_head // self.qhead_per_kvhead + if const_expr(not seqlen.has_cu_seqlens_q): + mQ_cur = mQ[None, None, num_head, batch_size] + else: + mQ_cur = cute.domain_offset((seqlen.offset_q, 0), mQ[None, None, num_head]) + if const_expr(not seqlen.has_cu_seqlens_k): + mK_cur = mK[None, None, num_head_kv, batch_size] + mV_cur = mV[None, None, num_head_kv, batch_size] + else: + mK_cur = cute.domain_offset((seqlen.offset_k, 0), mK[None, None, num_head_kv]) + mV_cur = cute.domain_offset((seqlen.offset_k, 0), mV[None, None, num_head_kv]) + gQ = cute.local_tile(mQ_cur, blkQ_shape, (m_block, 0)) + gK = cute.local_tile(mK_cur, blkK_shape, (None, 0)) + gV = cute.local_tile(mV_cur, blkV_shape, (None, 0)) + + # /////////////////////////////////////////////////////////////////////////////// + # Get shared memory buffer + # /////////////////////////////////////////////////////////////////////////////// + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + sQ = storage.sQ.get_tensor(sQ_layout) + sK = storage.sK.get_tensor(sK_layout) + if const_expr(not self.Q_in_regs): + sV = storage.sV.get_tensor(sV_layout) + else: + sV = cute.make_tensor(cute.recast_ptr(sQ.iterator, dtype=self.dtype), sV_layout) + # Transpose view of V to tensor with layout (head_dim_v, tile_n) for tiled mma + sVt = layout_utils.transpose_view(sV) + + gmem_thr_copy_K = gmem_tiled_copy_K.get_slice(tidx) + gmem_thr_copy_V = gmem_tiled_copy_V.get_slice(tidx) + # (CPY_Atom, CPY_N, CPY_K, n_block) + tKsK, tKgK = gmem_thr_copy_K.partition_D(sK), gmem_thr_copy_K.partition_S(gK) + # (CPY_Atom, CPY_N, CPY_K, n_block) + tVsV, tVgV = gmem_thr_copy_V.partition_D(sV), gmem_thr_copy_V.partition_S(gV) + + # /////////////////////////////////////////////////////////////////////////////// + # Tile MMA compute thread partitions and allocate accumulators + # /////////////////////////////////////////////////////////////////////////////// + thr_mma_qk = tiled_mma_qk.get_slice(tidx) + thr_mma_pv = tiled_mma_pv.get_slice(tidx) + tSrQ = thr_mma_qk.make_fragment_A(thr_mma_qk.partition_A(sQ)) + tSrK = thr_mma_qk.make_fragment_B(thr_mma_qk.partition_B(sK[None, None, 0])) + tOrVt = thr_mma_pv.make_fragment_B(thr_mma_pv.partition_B(sVt[None, None, 0])) + acc_shape_O = thr_mma_pv.partition_shape_C((self.tile_m, self.tile_hdimv)) + acc_O = cute.make_rmem_tensor(acc_shape_O, Float32) + acc_O.fill(0.0) + + # /////////////////////////////////////////////////////////////////////////////// + # Smem copy atom tiling + # /////////////////////////////////////////////////////////////////////////////// + smem_copy_atom_QK = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4), + self.dtype, + ) + smem_copy_atom_V = cute.make_copy_atom( + warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4), + self.dtype, + ) + smem_thr_copy_Q = utils.make_tiled_copy_A(smem_copy_atom_QK, tiled_mma_qk).get_slice(tidx) + smem_thr_copy_K = utils.make_tiled_copy_B(smem_copy_atom_QK, tiled_mma_qk).get_slice(tidx) + smem_thr_copy_V = utils.make_tiled_copy_B(smem_copy_atom_V, tiled_mma_pv).get_slice(tidx) + + tSsQ = smem_thr_copy_Q.partition_S(sQ) + tSsK = smem_thr_copy_K.partition_S(sK) + tOsVt = smem_thr_copy_V.partition_S(sVt) + + # /////////////////////////////////////////////////////////////////////////////// + # Predicate: Mark indices that need to copy when problem_shape isn't a multiple + # of tile_shape + # /////////////////////////////////////////////////////////////////////////////// + # Construct identity layout for KV + cK = cute.make_identity_tensor((self.tile_n, self.tile_hdim)) + tKcK = gmem_thr_copy_K.partition_S(cK) + t0KcK = gmem_thr_copy_K.get_slice(0).partition_S(cK) + if const_expr(self.tile_hdim == self.tile_hdimv): + tVcV = tKcK + t0VcV = t0KcK + else: + cV = cute.make_identity_tensor((self.tile_n, self.tile_hdimv)) + tVcV = gmem_thr_copy_V.partition_S(cV) + t0VcV = gmem_thr_copy_V.get_slice(0).partition_S(cV) + # Allocate predicate tensors for m and n, here we only allocate the tile of k, and + # use "if" on the mn dimension. + # This is to reduce register pressure and gets 2-3% performance gain. + tKpK = utils.predicate_k(tKcK, limit=mK.shape[1]) + if const_expr(self.same_hdim_kv): + tVpV = tKpK + else: + tVpV = utils.predicate_k(tVcV, limit=mV.shape[1]) + + # shape: (atom_v_m * rest_m) + softmax = Softmax.create( + softmax_scale_log2, + num_rows=acc_O.shape[0][0] * acc_O.shape[1], + softmax_scale=softmax_scale, + ) + softmax.reset() + + # group parameters for compute_one_n_block + mma_params = SimpleNamespace( + thr_mma_qk=thr_mma_qk, + thr_mma_pv=thr_mma_pv, + tSrQ=tSrQ, + tSrK=tSrK, + tOrVt=tOrVt, + acc_O=acc_O, + ) + smem_copy_params = SimpleNamespace( + smem_thr_copy_Q=smem_thr_copy_Q, + smem_thr_copy_K=smem_thr_copy_K, + smem_thr_copy_V=smem_thr_copy_V, + tSsQ=tSsQ, + tSsK=tSsK, + tOsVt=tOsVt, + ) + load_K = partial( + self.load_K, gmem_tiled_copy_K, tKgK, tKsK, tKcK, t0KcK, tKpK, seqlen=seqlen.seqlen_k + ) + load_V = partial( + self.load_V, gmem_tiled_copy_V, tVgV, tVsV, tVcV, t0VcV, tVpV, seqlen=seqlen.seqlen_k + ) + + compute_one_n_block = partial( + self.compute_one_n_block, + mma_params=mma_params, + smem_copy_params=smem_copy_params, + softmax=softmax, + load_K=load_K, + load_V=load_V, + score_mod=self.score_mod, + batch_idx=batch_size, + head_idx=num_head, + m_block=m_block, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Prologue + # /////////////////////////////////////////////////////////////////////////////// + # Start async loads of the last mn-tile, where we take care of the mn residue + gmem_thr_copy_Q = gmem_tiled_copy_Q.get_slice(tidx) + self.load_Q(gmem_thr_copy_Q, gQ, sQ, m_block, seqlen=seqlen.seqlen_q, headdim=mQ.shape[1]) + cute.arch.cp_async_commit_group() + + def preprocess_Q(): + cute.arch.cp_async_wait_group(self.num_stages * 2 - 1) + if const_expr(self.Q_in_regs): + cute.arch.barrier() + tSrQ_copy_view = smem_thr_copy_Q.retile(tSrQ) + cute.copy(smem_thr_copy_Q, tSsQ, tSrQ_copy_view) + + # If Q_in_regs, we load Q, then load 1 stage of K, then (optionally) rotate Q and + # read from smem_q to registers, then load V. + # If !Q_in_regs, we load Q, load all stages of K & V, then (optionally) rotate Q. + if const_expr(self.Q_in_regs): + load_K(n_block, smem_pipe_write=0, need_predicates=True) + cute.arch.cp_async_commit_group() + preprocess_Q() + cute.arch.barrier() # Make sure all threads have read smem_q before loading V + + for stage in cutlass.range_constexpr(self.num_stages): + if const_expr(not self.Q_in_regs or stage > 0): + if stage == 0 or n_block - stage >= 0: + load_K(n_block - stage, smem_pipe_write=stage, need_predicates=stage == 0) + cute.arch.cp_async_commit_group() + if const_expr(stage < self.num_stages - 1): + if stage == 0 or n_block - stage >= 0: + load_V(n_block - stage, smem_pipe_write=stage, need_predicates=stage == 0) + cute.arch.cp_async_commit_group() + if const_expr(not self.Q_in_regs): + preprocess_Q() + + # /////////////////////////////////////////////////////////////////////////////// + # Mainloop + # /////////////////////////////////////////////////////////////////////////////// + # Start processing of the first n-block. + # For performance reason, we separate out two kinds of iterations: + # those that need masking on S, and those that don't. + # We need masking on S for the very last block when K and V has length not multiple of tile_n. + # We also need masking on S if it's causal, for the last several blocks. + mask = AttentionMask( + self.tile_m, + self.tile_n, + seqlen, + window_size_left, + window_size_right, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + mask_fn = partial( + mask.apply_mask, + batch_idx=batch_size, + head_idx=num_head, + m_block=m_block, + thr_mma=thr_mma_qk, + mask_causal=self.is_causal, + mask_local=self.is_local, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods if const_expr(self.mask_mod is not None) else None, + ) + + # First iteration with seqlen masking + smem_pipe_read = Int32(0) + smem_pipe_write = Int32(self.num_stages - 1) + compute_one_n_block( + n_block, + smem_pipe_read, + smem_pipe_write, + is_first_n_block=True, + seqlen=seqlen, + mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=True), + ) + smem_pipe_read = self.advance_pipeline(smem_pipe_read) + smem_pipe_write = self.advance_pipeline(smem_pipe_write) + # Next couple of iterations with causal masking + if const_expr(self.is_causal or self.is_local): + n_block_min_causal_local_mask = block_info.get_n_block_min_causal_local_mask( + seqlen, m_block, n_block_min + ) + for n_tile in cutlass.range(n_block_max - 1 - n_block_min_causal_local_mask, unroll=1): + n_block = n_block_max - 2 - n_tile + compute_one_n_block( + n_block, + smem_pipe_read, + smem_pipe_write, + seqlen=seqlen, + mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=True), + ) + smem_pipe_read = self.advance_pipeline(smem_pipe_read) + smem_pipe_write = self.advance_pipeline(smem_pipe_write) + # The remaining iterations have no masking + for n_tile in cutlass.range(n_block, unroll=1): + compute_one_n_block( + n_block - n_tile - 1, smem_pipe_read, smem_pipe_write, + seqlen=seqlen, is_first_n_block=False, + mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=False) + ) + smem_pipe_read = self.advance_pipeline(smem_pipe_read) + smem_pipe_write = self.advance_pipeline(smem_pipe_write) + # TODO: local + + # normalize acc_O by row_sum and calculate the lse + row_scale = softmax.finalize() + softmax.rescale_O(acc_O, row_scale) + + # /////////////////////////////////////////////////////////////////////////////// + # Epilogue + # /////////////////////////////////////////////////////////////////////////////// + # reuse sQ's data iterator + sO = cute.make_tensor(sQ.iterator, sO_layout) + self.epilogue( + acc_O, + softmax.row_sum, + mO, + mLSE, + sO, + seqlen, + gmem_tiled_copy_O, + None, + tiled_mma_pv, + tidx, + m_block, + num_head, + batch_size, + ) + + @cute.jit + def compute_one_n_block( + self, + n_block: Int32, + smem_pipe_read: Int32, + smem_pipe_write: Int32, + mma_params: SimpleNamespace, + smem_copy_params: SimpleNamespace, + softmax: Softmax, + load_K: Callable, + load_V: Callable, + score_mod: Callable | None, + batch_idx: cutlass.Int32, + head_idx: cutlass.Int32, + m_block: cutlass.Int32, + seqlen: SeqlenInfoQK, + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + mask_fn: Optional[Callable] = None, + is_first_n_block: cutlass.Constexpr = False, + check_inf: cutlass.Constexpr = True, + ): + """Compute one n_block of S/O. + + This function provides different variants for processing the first n block versus + subsequent blocks. + """ + + def sync(): + cute.arch.cp_async_wait_group(self.num_stages * 2 - 2) + cute.arch.barrier() + + acc_shape_S = mma_params.thr_mma_qk.partition_shape_C((self.tile_m, self.tile_n)) + acc_S = cute.make_rmem_tensor(acc_shape_S, Float32) + acc_S.fill(0.0) + # wait for smem tile QK before mma calculation for S + sync() + + # need predicates for the first tile + def load_V_next(): + if self.num_stages == 1 or n_block - self.num_stages + 1 >= 0: + load_V( + n_block - self.num_stages + 1, + smem_pipe_write, + need_predicates=is_first_n_block and self.num_stages == 1, + ) + cute.arch.cp_async_commit_group() + + load_V_next() + sm80_utils.gemm( + mma_params.thr_mma_qk, + acc_S, + mma_params.tSrQ, + mma_params.tSrK, + smem_copy_params.tSsQ, + smem_copy_params.tSsK[ + None, None, None, smem_pipe_read if const_expr(self.num_stages > 1) else 0 + ], + smem_copy_params.smem_thr_copy_Q, + smem_copy_params.smem_thr_copy_K, + # hook_fn=load_V_next, + A_in_regs=self.Q_in_regs, + ) + if const_expr(score_mod is not None): + self.apply_score_mod( + mma_params.thr_mma_qk, + batch_idx, + head_idx, + m_block, + acc_S, + n_block, + softmax_scale=softmax.softmax_scale, + seqlen=seqlen, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + + smem_pipe_write = self.advance_pipeline(smem_pipe_write) + + def load_K_next(): + if n_block - self.num_stages >= 0: + load_K(n_block - self.num_stages, smem_pipe_write, need_predicates=False) + cute.arch.cp_async_commit_group() + + # wait for smem tile V for O + if const_expr(self.num_stages == 1): + sync() + load_K_next() + if const_expr(mask_fn is not None): + mask_fn(acc_S, n_block=n_block) + row_scale = softmax.online_softmax(acc_S, is_first=is_first_n_block, check_inf=check_inf) + softmax.rescale_O(mma_params.acc_O, row_scale) + rP = cute.make_fragment_like(acc_S, self.dtype) + rP.store(acc_S.load().to(self.dtype)) + tOrP = layout_utils.reshape_acc_to_frgA(rP) + if const_expr(self.num_stages > 1): + sync() + load_K_next() + sm80_utils.gemm_rs( + mma_params.thr_mma_pv, + mma_params.acc_O, + tOrP, + mma_params.tOrVt, + smem_copy_params.tOsVt[ + None, None, None, smem_pipe_read if const_expr(self.num_stages > 1) else 0 + ], + smem_copy_params.smem_thr_copy_V, + # hook_fn=load_K_next, + ) + # if const_expr(self.num_stages > 1): + # load_K_next() + @cute.jit + def apply_score_mod( + self, + thr_mma_qk, + batch_idx, + head_idx, + m_block, + acc_S, + n_block, + softmax_scale, + seqlen, + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + ): + # Prepare index tensor + cS = cute.make_identity_tensor((self.tile_m, self.tile_n)) + cS = cute.domain_offset((m_block * self.tile_m, n_block * self.tile_n), cS) + tScS = thr_mma_qk.partition_C(cS) + + apply_score_mod_inner( + acc_S, + tScS, + self.score_mod, + batch_idx, + head_idx, + softmax_scale, + self.score_vec_size, + self.qk_acc_dtype, + aux_data, + fastdiv_mods, + seqlen_info=seqlen, + constant_q_idx=None, + qhead_per_kvhead=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + + +# SM90 forward pass moved to flash_fwd_sm90.py; re-export for backward compatibility +def __getattr__(name): + if name == "FlashAttentionForwardSm90": + from flash_attn.cute.flash_fwd_sm90 import FlashAttentionForwardSm90 + return FlashAttentionForwardSm90 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/fastvideo-kernel/fa4/flash_fwd_combine.py b/fastvideo-kernel/fa4/flash_fwd_combine.py new file mode 100644 index 000000000..493620235 --- /dev/null +++ b/fastvideo-kernel/fa4/flash_fwd_combine.py @@ -0,0 +1,698 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +# A reimplementation of https://github.com/Dao-AILab/flash-attention/blob/main/hopper/flash_fwd_combine_kernel.h +# from Cutlass C++ to Cute-DSL. +import math +from typing import Type, Optional +from functools import partial + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync +from cutlass import Float32, Int32, Boolean, const_expr + +from flash_attn.cute import utils +from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned +from flash_attn.cute.seqlen_info import SeqlenInfo +from cutlass.cute import FastDivmodDivisor + + +class FlashAttentionForwardCombine: + def __init__( + self, + dtype: Type[cutlass.Numeric], + dtype_partial: Type[cutlass.Numeric], + head_dim: int, + tile_m: int = 8, + k_block_size: int = 64, + log_max_splits: int = 4, + num_threads: int = 256, + stages: int = 4, + ): + """ + Forward combine kernel for split attention computation. + + :param dtype: output data type + :param dtype_partial: partial accumulation data type + :param head_dim: head dimension + :param tile_m: m block size + :param k_block_size: k block size + :param log_max_splits: log2 of maximum splits + :param num_threads: number of threads + :param varlen: whether using variable length sequences + :param stages: number of pipeline stages + """ + self.dtype = dtype + self.dtype_partial = dtype_partial + self.head_dim = head_dim + self.tile_m = tile_m + self.k_block_size = k_block_size + self.max_splits = 1 << log_max_splits + self.num_threads = num_threads + self.is_even_k = head_dim % k_block_size == 0 + self.stages = stages + + @staticmethod + def can_implement( + dtype, + dtype_partial, + head_dim, + tile_m, + k_block_size, + log_max_splits, + num_threads, + ) -> bool: + """Check if the kernel can be implemented with the given parameters.""" + if dtype not in [cutlass.Float16, cutlass.BFloat16, cutlass.Float32]: + return False + if dtype_partial not in [cutlass.Float16, cutlass.BFloat16, Float32]: + return False + if head_dim % 8 != 0: + return False + if num_threads % 32 != 0: + return False + if tile_m % 8 != 0: + return False + max_splits = 1 << log_max_splits + if max_splits > 256: + return False + if (tile_m * max_splits) % num_threads != 0: + return False + return True + + def _setup_attributes(self): + # GMEM copy setup for O partial + universal_copy_bits = 128 + async_copy_elems = universal_copy_bits // self.dtype_partial.width + assert self.k_block_size % async_copy_elems == 0 + + k_block_gmem = ( + 128 if self.k_block_size % 128 == 0 else (64 if self.k_block_size % 64 == 0 else 32) + ) + gmem_threads_per_row = k_block_gmem // async_copy_elems + assert self.num_threads % gmem_threads_per_row == 0 + + # Async copy atom for O partial load + atom_async_copy_partial = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + self.dtype_partial, + num_bits_per_copy=universal_copy_bits, + ) + tOpartial_layout = cute.make_ordered_layout( + (self.num_threads // gmem_threads_per_row, gmem_threads_per_row), + order=(1, 0), + ) + vOpartial_layout = cute.make_layout((1, async_copy_elems)) # 4 vals per load + self.gmem_tiled_copy_O_partial = cute.make_tiled_copy_tv( + atom_async_copy_partial, tOpartial_layout, vOpartial_layout + ) + + # GMEM copy setup for final O (use universal copy for store) + atom_universal_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.dtype, + num_bits_per_copy=async_copy_elems * self.dtype.width, + ) + self.gmem_tiled_copy_O = cute.make_tiled_copy_tv( + atom_universal_copy, + tOpartial_layout, + vOpartial_layout, # 4 vals per store + ) + + # LSE copy setup with async copy (alignment = 1) + lse_copy_bits = Float32.width # 1 element per copy, width is in bits + m_block_smem = ( + 128 + if self.tile_m % 128 == 0 + else ( + 64 + if self.tile_m % 64 == 0 + else (32 if self.tile_m % 32 == 0 else (16 if self.tile_m % 16 == 0 else 8)) + ) + ) + gmem_threads_per_row_lse = m_block_smem + assert self.num_threads % gmem_threads_per_row_lse == 0 + + # Async copy atom for LSE load + atom_async_copy_lse = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.ALWAYS), + Float32, + num_bits_per_copy=lse_copy_bits, + ) + tLSE_layout = cute.make_ordered_layout( + (self.num_threads // gmem_threads_per_row_lse, gmem_threads_per_row_lse), + order=(1, 0), + ) + vLSE_layout = cute.make_layout(1) + self.gmem_tiled_copy_LSE = cute.make_tiled_copy_tv( + atom_async_copy_lse, tLSE_layout, vLSE_layout + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Shared memory + # /////////////////////////////////////////////////////////////////////////////// + + # Shared memory to register copy for LSE + self.smem_threads_per_col_lse = self.num_threads // m_block_smem + assert 32 % self.smem_threads_per_col_lse == 0 # Must divide warp size + + s2r_layout_atom_lse = cute.make_ordered_layout( + (self.smem_threads_per_col_lse, self.num_threads // self.smem_threads_per_col_lse), + order=(0, 1), + ) + self.s2r_tiled_copy_LSE = cute.make_tiled_copy_tv( + cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), Float32), + s2r_layout_atom_lse, + cute.make_layout(1), + ) + + # LSE shared memory layout with swizzling to avoid bank conflicts + # This works for kBlockMSmem = 8, 16, 32, 64, 128, no bank conflicts + if const_expr(m_block_smem == 8): + smem_lse_swizzle = cute.make_swizzle(5, 0, 5) + elif const_expr(m_block_smem == 16): + smem_lse_swizzle = cute.make_swizzle(4, 0, 4) + else: + smem_lse_swizzle = cute.make_swizzle(3, 2, 3) + smem_layout_atom_lse = cute.make_composed_layout( + smem_lse_swizzle, 0, cute.make_ordered_layout((8, m_block_smem), order=(1, 0)) + ) + self.smem_layout_lse = cute.tile_to_shape( + smem_layout_atom_lse, (self.max_splits, self.tile_m), (0, 1) + ) + + # O partial shared memory layout (simple layout for pipeline stages) + self.smem_layout_o = cute.make_ordered_layout( + (self.tile_m, self.k_block_size, self.stages), order=(1, 0, 2) + ) + + @cute.jit + def __call__( + self, + mO_partial: cute.Tensor, + mLSE_partial: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor] = None, + cu_seqlens: Optional[cute.Tensor] = None, + seqused: Optional[cute.Tensor] = None, + num_splits_dynamic_ptr: Optional[cute.Tensor] = None, + varlen_batch_idx: Optional[cute.Tensor] = None, + semaphore_to_reset: Optional[cute.Tensor] = None, + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + # Type checking + if const_expr(not (mO_partial.element_type == self.dtype_partial)): + raise TypeError("O partial tensor must match dtype_partial") + if const_expr(not (mO.element_type == self.dtype)): + raise TypeError("O tensor must match dtype") + if const_expr(mLSE_partial.element_type not in [Float32]): + raise TypeError("LSE partial tensor must be Float32") + if const_expr(mLSE is not None and mLSE.element_type not in [Float32]): + raise TypeError("LSE tensor must be Float32") + + # Shape validation - input tensors are in user format, need to be converted to kernel format + if const_expr(len(mO_partial.shape) not in [4, 5]): + raise ValueError( + "O partial tensor must have 4 or 5 dimensions: (num_splits, batch, seqlen, nheads, headdim) or (num_splits, total_q, nheads, headdim)" + ) + if const_expr(len(mLSE_partial.shape) not in [3, 4]): + raise ValueError( + "LSE partial tensor must have 3 or 4 dimensions: (num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads)" + ) + if const_expr(len(mO.shape) not in [3, 4]): + raise ValueError( + "O tensor must have 3 or 4 dimensions: (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim)" + ) + if const_expr(mLSE is not None and len(mLSE.shape) not in [2, 3]): + raise ValueError( + "LSE tensor must have 2 or 3 dimensions: (batch, seqlen, nheads) or (total_q, nheads)" + ) + + mO_partial, mO = [assume_tensor_aligned(t) for t in (mO_partial, mO)] + # (num_splits, b, seqlen, h, d) -> (seqlen, d, num_splits, h, b) + # or (num_splits, total_q, h, d) -> (total_q, d, num_splits, h) + O_partial_layout_transpose = ( + [2, 4, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 3, 0, 2] + ) + # (b, seqlen, h, d) -> (seqlen, d, h, b) or (total_q, h, d) -> (total_q, d, h) + mO_partial = cute.make_tensor( + mO_partial.iterator, cute.select(mO_partial.layout, mode=O_partial_layout_transpose) + ) + O_layout_transpose = [1, 3, 2, 0] if const_expr(cu_seqlens is None) else [0, 2, 1] + mO = cute.make_tensor(mO.iterator, cute.select(mO.layout, mode=O_layout_transpose)) + # (num_splits, b, seqlen, h) -> (seqlen, num_splits, h, b) + # or (num_splits, total_q, h) -> (total_q, num_splits, h) + LSE_partial_layout_transpose = [2, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 0, 2] + mLSE_partial = cute.make_tensor( + mLSE_partial.iterator, + cute.select(mLSE_partial.layout, mode=LSE_partial_layout_transpose), + ) + # (b, seqlen, h) -> (seqlen, h, b) or (total_q, h) -> (total_q, h) + LSE_layout_transpose = [1, 2, 0] if const_expr(cu_seqlens is None) else [0, 1] + mLSE = ( + cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose)) + if mLSE is not None + else None + ) + + # Determine if we have variable length sequences + varlen = const_expr(cu_seqlens is not None or seqused is not None) + + self._setup_attributes() + + @cute.struct + class SharedStorage: + sLSE: cute.struct.Align[ + cute.struct.MemRange[Float32, cute.cosize(self.smem_layout_lse)], 128 + ] + sMaxValidSplit: cute.struct.Align[cute.struct.MemRange[Int32, self.tile_m], 128] + sO: cute.struct.Align[ + cute.struct.MemRange[self.dtype_partial, cute.cosize(self.smem_layout_o)], 128 + ] + + smem_size = SharedStorage.size_in_bytes() + + # Grid dimensions: (ceil_div(seqlen, m_block), ceil_div(head_dim, k_block), num_head * batch) + seqlen = mO_partial.shape[0] + num_head = mO_partial.shape[3] + batch_size = ( + mO_partial.shape[4] + if const_expr(cu_seqlens is None) + else Int32(cu_seqlens.shape[0] - 1) + ) + + # Create FastDivmodDivisor objects for efficient division + seqlen_divmod = FastDivmodDivisor(seqlen) + head_divmod = FastDivmodDivisor(num_head) + + grid_dim = ( + cute.ceil_div(seqlen * num_head, self.tile_m), + cute.ceil_div(self.head_dim, self.k_block_size), + batch_size, + ) + + self.kernel( + mO_partial, + mLSE_partial, + mO, + mLSE, + cu_seqlens, + seqused, + num_splits_dynamic_ptr, + varlen_batch_idx, + semaphore_to_reset, + SharedStorage, + self.smem_layout_lse, + self.smem_layout_o, + self.gmem_tiled_copy_O_partial, + self.gmem_tiled_copy_O, + self.gmem_tiled_copy_LSE, + self.s2r_tiled_copy_LSE, + seqlen_divmod, + head_divmod, + varlen, + ).launch( + grid=grid_dim, + block=[self.num_threads, 1, 1], + smem=smem_size, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mO_partial: cute.Tensor, + mLSE_partial: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + cu_seqlens: Optional[cute.Tensor], + seqused: Optional[cute.Tensor], + num_splits_dynamic_ptr: Optional[cute.Tensor], + varlen_batch_idx: Optional[cute.Tensor], + semaphore_to_reset: Optional[cute.Tensor], + SharedStorage: cutlass.Constexpr, + smem_layout_lse: cute.Layout | cute.ComposedLayout, + smem_layout_o: cute.Layout, + gmem_tiled_copy_O_partial: cute.TiledCopy, + gmem_tiled_copy_O: cute.TiledCopy, + gmem_tiled_copy_LSE: cute.TiledCopy, + s2r_tiled_copy_LSE: cute.TiledCopy, + seqlen_divmod: FastDivmodDivisor, + head_divmod: FastDivmodDivisor, + varlen: cutlass.Constexpr[bool], + ): + # Thread and block indices + tidx, _, _ = cute.arch.thread_idx() + m_block, k_block, maybe_virtual_batch = cute.arch.block_idx() + + # Map virtual batch index to real batch index (for persistent tile schedulers) + batch_idx = ( + varlen_batch_idx[maybe_virtual_batch] + if const_expr(varlen_batch_idx is not None) + else maybe_virtual_batch + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Get shared memory buffer + # /////////////////////////////////////////////////////////////////////////////// + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + sLSE = storage.sLSE.get_tensor(smem_layout_lse) + sMaxValidSplit = storage.sMaxValidSplit.get_tensor((self.tile_m,)) + sO = storage.sO.get_tensor(smem_layout_o) + + # Handle semaphore reset — wait for dependent grids first + if const_expr(semaphore_to_reset is not None): + if ( + tidx == 0 + and m_block == cute.arch.grid_dim()[0] - 1 + and k_block == cute.arch.grid_dim()[1] - 1 + and maybe_virtual_batch == cute.arch.grid_dim()[2] - 1 + ): + cute.arch.griddepcontrol_wait() + semaphore_to_reset[0] = 0 + + # Get number of splits (use maybe_virtual_batch for per-batch-slot splits) + num_splits = ( + num_splits_dynamic_ptr[maybe_virtual_batch] + if const_expr(num_splits_dynamic_ptr is not None) + else mLSE_partial.shape[1] + ) + # Handle variable length sequences using SeqlenInfo + seqlen_info = SeqlenInfo.create( + batch_idx=batch_idx, + seqlen_static=mO_partial.shape[0], + cu_seqlens=cu_seqlens, + seqused=seqused, + # Don't need to pass in tile size since we won't use offset_padded + ) + seqlen, offset = seqlen_info.seqlen, seqlen_info.offset + + # Extract number of heads (head index will be determined dynamically) + num_head = mO_partial.shape[3] + max_idx = seqlen * num_head + + # Early exit for single split if dynamic + if (const_expr(num_splits_dynamic_ptr is None) or num_splits > 1) and ( + const_expr(not varlen) or m_block * self.tile_m < max_idx + ): + # Wait for dependent grids (e.g., the main attention kernel that produces O_partial/LSE_partial) + cute.arch.griddepcontrol_wait() + + # =============================== + # Step 1: Load LSE_partial from gmem to shared memory + # =============================== + + mLSE_partial_cur = seqlen_info.offset_batch(mLSE_partial, batch_idx, dim=3) + mLSE_partial_copy = cute.tiled_divide(mLSE_partial_cur, (1,)) + gmem_thr_copy_LSE = gmem_tiled_copy_LSE.get_slice(tidx) + tLSEsLSE = gmem_thr_copy_LSE.partition_D(sLSE) + # Create identity tensor for coordinate tracking + cLSE = cute.make_identity_tensor((self.max_splits, self.tile_m)) + tLSEcLSE = gmem_thr_copy_LSE.partition_S(cLSE) + + # Load LSE partial values + for m in cutlass.range(cute.size(tLSEcLSE, mode=[2]), unroll_full=True): + mi = tLSEcLSE[0, 0, m][1] # Get m coordinate + idx = m_block * self.tile_m + mi + if idx < max_idx: + # Calculate actual sequence position and head using FastDivmodDivisor + if const_expr(not varlen): + head_idx, m_idx = divmod(idx, seqlen_divmod) + else: + head_idx = idx // seqlen + m_idx = idx - head_idx * seqlen + mLSE_partial_cur_copy = mLSE_partial_copy[None, m_idx, None, head_idx] + for s in cutlass.range(cute.size(tLSEcLSE, mode=[1]), unroll_full=True): + si = tLSEcLSE[0, s, 0][0] # Get split coordinate + if si < num_splits: + cute.copy( + gmem_thr_copy_LSE, + mLSE_partial_cur_copy[None, si], + tLSEsLSE[None, s, m], + ) + else: + tLSEsLSE[None, s, m].fill(-Float32.inf) + # Don't need to zero out the rest of the LSEs, as we will not write the output to gmem + cute.arch.cp_async_commit_group() + + # =============================== + # Step 2: Load O_partial for pipeline stages + # =============================== + + gmem_thr_copy_O_partial = gmem_tiled_copy_O_partial.get_slice(tidx) + cO = cute.make_identity_tensor((self.tile_m, self.k_block_size)) + tOcO = gmem_thr_copy_O_partial.partition_D(cO) + tOsO_partial = gmem_thr_copy_O_partial.partition_D(sO) + mO_partial_cur = seqlen_info.offset_batch(mO_partial, batch_idx, dim=4) + + # Precompute these values to avoid recomputing them in the loop + num_rows = const_expr(cute.size(tOcO, mode=[1])) + tOmidx = cute.make_rmem_tensor(num_rows, cutlass.Int32) + tOhidx = cute.make_rmem_tensor(num_rows, cutlass.Int32) + tOrOptr = cute.make_rmem_tensor(num_rows, cutlass.Int64) + for m in cutlass.range(num_rows, unroll_full=True): + mi = tOcO[0, m, 0][0] # m coordinate + idx = m_block * self.tile_m + mi + if const_expr(not varlen): + tOhidx[m], tOmidx[m] = divmod(idx, seqlen_divmod) + else: + tOhidx[m] = idx // seqlen + tOmidx[m] = idx - tOhidx[m] * seqlen + tOrOptr[m] = utils.elem_pointer( + mO_partial_cur, (tOmidx[m], k_block * self.k_block_size, 0, tOhidx[m]) + ).toint() + if idx >= max_idx: + tOhidx[m] = -1 + + tOpO = None + if const_expr(not self.is_even_k): + tOpO = cute.make_rmem_tensor(cute.size(tOcO, mode=[2]), Boolean) + for k in cutlass.range(cute.size(tOpO), unroll_full=True): + tOpO[k] = tOcO[0, 0, k][1] < mO_partial.shape[1] - k_block * self.k_block_size + # if cute.arch.thread_idx()[0] == 0 and k_block == 1: cute.print_tensor(tOpO) + + load_O_partial = partial( + self.load_O_partial, + gmem_tiled_copy_O_partial, + tOrOptr, + tOsO_partial, + tOhidx, + tOpO, + tOcO, + mO_partial_cur.layout, + ) + + # Load first few stages of O_partial + for stage in cutlass.range(self.stages - 1, unroll_full=True): + if stage < num_splits: + load_O_partial(stage, stage) + cute.arch.cp_async_commit_group() + + # =============================== + # Step 3: Load and transpose LSE from smem to registers + # =============================== + + # Wait for LSE and initial O partial stages to complete + cute.arch.cp_async_wait_group(self.stages - 1) + cute.arch.sync_threads() + # if cute.arch.thread_idx()[0] == 0: + # # cute.print_tensor(sLSE) + # for i in range(64): + # cute.printf("sLSE[%d, 0] = %f", i, sLSE[i, 0]) + # cute.arch.sync_threads() + + s2r_thr_copy_LSE = s2r_tiled_copy_LSE.get_slice(tidx) + ts2rsLSE = s2r_thr_copy_LSE.partition_S(sLSE) + ts2rrLSE = cute.make_rmem_tensor_like(ts2rsLSE) + cute.copy(s2r_tiled_copy_LSE, ts2rsLSE, ts2rrLSE) + + # =============================== + # Step 4: Compute final LSE along split dimension + # =============================== + + lse_sum = cute.make_rmem_tensor(cute.size(ts2rrLSE, mode=[2]), Float32) + ts2rcLSE = s2r_thr_copy_LSE.partition_D(cLSE) + # We compute the max valid split for each row to short-circuit the computation later + max_valid_split = cute.make_rmem_tensor(cute.size(ts2rrLSE, mode=[2]), Int32) + assert cute.size(ts2rrLSE, mode=[0]) == 1 + # Compute max, scales, and final LSE for each row + for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True): + # Find max LSE value across splits + threads_per_col = const_expr(self.smem_threads_per_col_lse) + lse_max = cute.arch.warp_reduction_max( + ts2rrLSE[None, None, m] + .load() + .reduce(cute.ReductionOp.MAX, init_val=-Float32.inf, reduction_profile=0), + threads_in_group=threads_per_col, + ) + # if cute.arch.thread_idx()[0] == 0: cute.printf(lse_max) + # Find max valid split index + max_valid_idx = -1 + for s in cutlass.range(cute.size(ts2rrLSE, mode=[1]), unroll_full=True): + if ts2rrLSE[0, s, m] != -Float32.inf: + max_valid_idx = ts2rcLSE[0, s, 0][0] # Get split coordinate + # if cute.arch.thread_idx()[0] < 32: cute.printf(max_valid_idx) + max_valid_split[m] = cute.arch.warp_reduction_max( + max_valid_idx, threads_in_group=threads_per_col + ) + # Compute exp scales and sum + lse_max_cur = ( + 0.0 if lse_max == -Float32.inf else lse_max + ) # In case all local LSEs are -inf + LOG2_E = math.log2(math.e) + lse_sum_cur = 0.0 + for s in cutlass.range(cute.size(ts2rrLSE, mode=[1]), unroll_full=True): + scale = cute.math.exp2( + ts2rrLSE[0, s, m] * LOG2_E - (lse_max_cur * LOG2_E), fastmath=True + ) + lse_sum_cur += scale + ts2rrLSE[0, s, m] = scale # Store scale for later use + lse_sum_cur = cute.arch.warp_reduction_sum( + lse_sum_cur, threads_in_group=threads_per_col + ) + lse_sum[m] = cute.math.log(lse_sum_cur, fastmath=True) + lse_max + # Normalize scales + inv_sum = ( + 0.0 if (lse_sum_cur == 0.0 or lse_sum_cur != lse_sum_cur) else 1.0 / lse_sum_cur + ) + ts2rrLSE[None, None, m].store(ts2rrLSE[None, None, m].load() * inv_sum) + # Store the scales exp(lse - lse_logsum) back to smem + cute.copy(s2r_tiled_copy_LSE, ts2rrLSE, ts2rsLSE) + + # Store max valid split to smem + for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True): + if ts2rcLSE[0, 0, m][0] == 0: # Only thread responsible for s=0 writes + mi = ts2rcLSE[0, 0, m][1] + if mi < self.tile_m: + sMaxValidSplit[mi] = max_valid_split[m] + + # =============================== + # Step 5: Store final LSE to gmem + # =============================== + + if const_expr(mLSE is not None): + if const_expr(cu_seqlens is None): + mLSE_cur = mLSE[None, None, batch_idx] + else: + mLSE_cur = cute.domain_offset((offset, 0), mLSE) + if k_block == 0: # Only first k_block writes LSE when mLSE is provided + for m in cutlass.range(cute.size(ts2rrLSE, mode=[2]), unroll_full=True): + if ts2rcLSE[0, 0, m][0] == 0: # Only thread responsible for s=0 writes + mi = ts2rcLSE[0, 0, m][1] + idx = m_block * self.tile_m + mi + if idx < max_idx: + if const_expr(not varlen): + head_idx, m_idx = divmod(idx, seqlen_divmod) + else: + head_idx = idx // seqlen + m_idx = idx - head_idx * seqlen + mLSE_cur[m_idx, head_idx] = lse_sum[m] + + # =============================== + # Step 6: Read O_partial and accumulate final O + # =============================== + + cute.arch.sync_threads() + + # Get max valid split for this thread + thr_max_valid_split = sMaxValidSplit[tOcO[0, 0, 0][0]] + for m in cutlass.range(1, cute.size(tOcO, mode=[1]), unroll_full=True): + thr_max_valid_split = max(thr_max_valid_split, sMaxValidSplit[tOcO[0, m, 0][0]]) + + tOrO_partial = cute.make_rmem_tensor_like(tOsO_partial[None, None, None, 0]) + tOrO = cute.make_rmem_tensor_like(tOrO_partial, Float32) + tOrO.fill(0.0) + + stage_load = self.stages - 1 + stage_compute = 0 + + # Main accumulation loop + for s in cutlass.range(thr_max_valid_split + 1, unroll=4): + # Get scales for this split + scale = cute.make_rmem_tensor(num_rows, Float32) + for m in cutlass.range(num_rows, unroll_full=True): + scale[m] = sLSE[s, tOcO[0, m, 0][0]] # Get scale from smem + + # Load next stage if needed + split_to_load = s + self.stages - 1 + if split_to_load <= thr_max_valid_split: + load_O_partial(split_to_load, stage_load) + cute.arch.cp_async_commit_group() + stage_load = 0 if stage_load == self.stages - 1 else stage_load + 1 + + # Wait for the current stage to be ready + cute.arch.cp_async_wait_group(self.stages - 1) + # We don't need __syncthreads() because each thread is just reading its own data from smem + # Copy from smem to registers + cute.autovec_copy(tOsO_partial[None, None, None, stage_compute], tOrO_partial) + stage_compute = 0 if stage_compute == self.stages - 1 else stage_compute + 1 + + # Accumulate scaled partial results + for m in cutlass.range(num_rows, unroll_full=True): + if tOhidx[m] >= 0 and scale[m] > 0.0: + tOrO[None, m, None].store( + tOrO[None, m, None].load() + + scale[m] * tOrO_partial[None, m, None].load().to(Float32) + ) + + # =============================== + # Step 7: Write final O to gmem + # =============================== + + rO = cute.make_rmem_tensor_like(tOrO, self.dtype) + rO.store(tOrO.load().to(self.dtype)) + mO_cur = seqlen_info.offset_batch(mO, batch_idx, dim=3) + if const_expr(cu_seqlens is None): + mO_cur = mO[None, None, None, batch_idx] + else: + mO_cur = cute.domain_offset((offset, 0, 0), mO) + mO_cur = utils.domain_offset_aligned((0, k_block * self.k_block_size, 0), mO_cur) + elems_per_store = const_expr(cute.size(gmem_tiled_copy_O.layout_tv_tiled[1])) + # mO_cur_copy = cute.tiled_divide(mO_cur, (1, elems_per_store,)) + gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx) + # Write final results + for m in cutlass.range(num_rows, unroll_full=True): + if tOhidx[m] >= 0: + mO_cur_copy = cute.tiled_divide( + mO_cur[tOmidx[m], None, tOhidx[m]], (elems_per_store,) + ) + for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True): + k_idx = tOcO[0, 0, k][1] // elems_per_store + if const_expr(self.is_even_k) or tOpO[k]: + cute.copy(gmem_thr_copy_O, rO[None, m, k], mO_cur_copy[None, k_idx]) + + @cute.jit + def load_O_partial( + self, + gmem_tiled_copy_O_partial: cute.TiledCopy, + tOrOptr: cute.Tensor, + tOsO_partial: cute.Tensor, + tOhidx: cute.Tensor, + tOpO: Optional[cute.Tensor], + tOcO: cute.Tensor, + mO_cur_partial_layout: cute.Layout, + split: Int32, + stage: Int32, + ) -> None: + elems_per_load = const_expr(cute.size(gmem_tiled_copy_O_partial.layout_tv_tiled[1])) + tOsO_partial_cur = tOsO_partial[None, None, None, stage] + for m in cutlass.range(cute.size(tOcO, [1]), unroll_full=True): + if tOhidx[m] >= 0: + o_gmem_ptr = cute.make_ptr( + tOsO_partial.element_type, tOrOptr[m], cute.AddressSpace.gmem, assumed_align=16 + ) + mO_partial_cur = cute.make_tensor( + o_gmem_ptr, cute.slice_(mO_cur_partial_layout, (0, None, None, 0)) + ) + mO_partial_cur_copy = cute.tiled_divide(mO_partial_cur, (elems_per_load,)) + for k in cutlass.range(cute.size(tOcO, mode=[2]), unroll_full=True): + k_idx = tOcO[0, 0, k][1] // elems_per_load + if const_expr(tOpO is None) or tOpO[k]: + cute.copy( + gmem_tiled_copy_O_partial, + mO_partial_cur_copy[None, k_idx, split], + tOsO_partial_cur[None, m, k], + ) diff --git a/fastvideo-kernel/fa4/flash_fwd_mla_sm100.py b/fastvideo-kernel/fa4/flash_fwd_mla_sm100.py new file mode 100644 index 000000000..ccf9c049a --- /dev/null +++ b/fastvideo-kernel/fa4/flash_fwd_mla_sm100.py @@ -0,0 +1,3160 @@ +# Copyright (c) 2026, Colfax International. + +import math +from functools import partial +from typing import Callable, Optional + + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int64, Int32, Uint32, Boolean, const_expr +from cutlass.cute import FastDivmodDivisor +import cutlass.pipeline as pipeline +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.utils import ClcDynamicPersistentTileScheduler + +from quack import copy_utils + +from flash_attn.cute.pack_gqa import pack_gqa_layout, make_packgqa_tiled_tma_atom +from flash_attn.cute.paged_kv import PagedKVManager +from flash_attn.cute import utils as fa_utils +from flash_attn.cute.seqlen_info import SeqlenInfoQK +from flash_attn.cute.block_info import BlockInfo +from flash_attn.cute.mask import AttentionMask +import flash_attn.cute.blackwell_helpers as fa_sm100_utils +from flash_attn.cute.softmax import SoftmaxSm100 +from flash_attn.cute.tile_scheduler import ( + ClcState, + SchedulingMode, + TileSchedulerArguments, + TileSchedulerProtocol, + SingleTileScheduler, + SingleTileLPTScheduler, + SingleTileVarlenScheduler, + ParamsBase, +) +from flash_attn.cute.fa_logging import fa_log, fa_printf +from flash_attn.cute.utils import smid, get_batch_from_cu_tensor + +from flash_attn.cute.topk_gather_kv import CpasyncGatherKVManager + + +from flash_attn.cute.named_barrier import NamedBarrierFwdSm100_MLA2CTA + + +class FlashAttentionMLAForwardSm100: + def __init__( + self, + is_causal: bool = False, + use_cpasync_load_KV: bool = False, + topk_length: int = 2048, + is_topk_gather: bool = True, + pack_gqa: bool = False, + qhead_per_kvhead: int = 1, + nheads_kv: int = 1, + hdim: int = 64, + hdimv: int = 512, + has_seqused_q: bool = False, + has_cu_seqlens_q: bool = False, + disable_bitmask: bool = False, + use_clc_scheduler: bool = True, + has_qk: bool = True, + ): + self.is_causal = is_causal + self.is_local = False + self.pack_gqa = pack_gqa + self.qhead_per_kvhead = qhead_per_kvhead + assert qhead_per_kvhead <= 128 + self.nheads_kv = nheads_kv + self.use_tma_O = True + self.use_cpasync_load_KV = use_cpasync_load_KV + self.use_tma_KV = not use_cpasync_load_KV + self.topk_length = topk_length + self.is_topk_gather = is_topk_gather + if is_topk_gather: + assert pack_gqa + assert qhead_per_kvhead == 128, "require MQA 128 for DSA path" + assert use_cpasync_load_KV + # user-provided option if topk indices guaranteed in bounds + self.disable_bitmask = disable_bitmask + self.has_qk = has_qk + + # ==== tile scheduler ==== + self.is_persistent = False + self.use_clc_scheduler = use_clc_scheduler + self.sched_stages = 1 + self.scheduling_mode = ( + SchedulingMode.CLC if self.use_clc_scheduler else SchedulingMode.STATIC + ) + + self.is_varlen_q = has_seqused_q or has_cu_seqlens_q + self.use_packed_varlen_sched = has_cu_seqlens_q and qhead_per_kvhead == 128 and pack_gqa + self.use_varlen_scheduler = self.is_varlen_q and not self.use_packed_varlen_sched + + if const_expr(self.use_varlen_scheduler): + self.TileScheduler = SingleTileVarlenScheduler + elif self.use_clc_scheduler: + self.TileScheduler = SingleTileLPTScheduler + else: + self.TileScheduler = SingleTileScheduler + + fa_log( + 1, + f"TileScheduler={self.TileScheduler.__name__}, scheduling_mode={self.scheduling_mode.name}", + ) + + # ==== thread info ==== + self.num_softmax_threads = 128 + self.num_epilogue_threads = 128 + self.num_load_threads = 32 + self.num_mma_threads = 32 + self.num_empty_threads = 32 if use_cpasync_load_KV else 64 + self.num_relay_threads = 32 if use_cpasync_load_KV else 0 + self.num_cpasync_load_threads = 128 if use_cpasync_load_KV else 0 + self.num_threads = ( + self.num_softmax_threads + + self.num_epilogue_threads + + self.num_load_threads + + self.num_mma_threads + + self.num_empty_threads + + self.num_relay_threads + + self.num_cpasync_load_threads + ) + self.num_warps = self.num_threads // 32 + assert self.num_warps == 12 or self.num_warps == 16 + self.softmax_warp_indices = (0, 1, 2, 3) + self.epilogue_warp_indices = (4, 5, 6, 7) + self.load_warp_id = 8 + self.mma_warp_id = 9 + self.clc_scheduler_warp_id = 10 + self.relay_warp_id = 11 + self.empty_warp_ids = tuple( + w + for w, active in [ + (self.relay_warp_id, not use_cpasync_load_KV), + (self.clc_scheduler_warp_id, not self.use_clc_scheduler), + ] + if active + ) + self.cpasync_load_warp_indices = (12, 13, 14, 15) + + # ==== register usage ==== + if self.num_warps == 16: + self.num_regs_load = 112 + self.num_regs_mma = 112 + self.num_regs_softmax = 192 + self.num_regs_epilogue = 128 + self.num_regs_cpasync = 80 if self.use_cpasync_load_KV else 0 + self.num_regs_other = 48 + else: + self.num_regs_load = 168 - 40 + self.num_regs_mma = 168 - 40 + self.num_regs_softmax = 168 + 80 + self.num_regs_epilogue = 168 - 40 + self.num_regs_cpasync = 0 + self.num_regs_other = 48 + + self.num_regs_per_thread = 168 if self.num_warps == 12 else 128 + self.num_regs_total = 504 if self.num_warps == 12 else 512 + + assert ( + self.num_regs_mma + + self.num_regs_softmax + + self.num_regs_epilogue + + self.num_regs_cpasync + <= self.num_regs_total + ) + + # ==== 2cta info ==== + self.use_2cta_instrs = True + self.cta_group = tcgen05.CtaGroup.TWO + self.cta_group_size = 2 + self.cluster_shape_mn = (2, 1) + self.cluster_shape_mnk = (2, 1, 1) + + # ==== problem shape info ==== + self.hdim = hdim + self.hdimv = hdimv + self.cta_tile_m = 64 + self.cluster_tile_m = self.cta_group_size * self.cta_tile_m + self.tile_n = 128 + assert ( + pack_gqa is False + or self.cluster_tile_m % qhead_per_kvhead == 0 + or qhead_per_kvhead % self.cluster_tile_m == 0 + ) + self.num_hdimv_splits = 2 # split hdimv in half for our Qv @ V^T and P @ V mmas. + assert hdimv % 32 == 0 + assert self.topk_length % self.tile_n == 0 or not self.is_topk_gather + self.epi_tile = (self.cta_tile_m, self.hdimv // self.num_hdimv_splits) + self.tile_P = (self.cta_tile_m, self.tile_n) + + # ==== MMA info ==== + self.mma_tiler_QK = ( + self.cluster_tile_m, + self.tile_n, + self.hdim, + ) + self.mma_tiler_QvV = ( + self.cluster_tile_m, + self.tile_n, + self.hdimv // self.num_hdimv_splits, + ) + self.mma_tiler_PVt = ( + self.cluster_tile_m, + self.hdimv // self.num_hdimv_splits, + self.tile_n, + ) + self.major_mode_Q = tcgen05.OperandMajorMode.K + self.major_mode_Qvi = tcgen05.OperandMajorMode.K + self.major_mode_K = tcgen05.OperandMajorMode.K + self.major_mode_Vi = tcgen05.OperandMajorMode.K + self.major_mode_Vti = tcgen05.OperandMajorMode.MN + self.major_mode_P = tcgen05.OperandMajorMode.K + self.operand_source_Q = tcgen05.OperandSource.SMEM + self.operand_source_Qvi = tcgen05.OperandSource.SMEM + self.operand_source_P = tcgen05.OperandSource.SMEM + + # ==== pipeline info ==== + self.num_stages_Q = 1 + self.num_stages_K = 1 + self.num_stages_Qv = 2 + self.num_stages_V = 4 + self.num_stages_S = 2 + # self.num_stages_P = 1 if has_qk else 2 + self.num_stages_P = 1 + self.num_stages_Oi = 1 + self.num_stages_sm_stats = 2 + self.num_stages_bitmask = 2 + assert self.num_stages_S == 2, "mainloops expect 2 stages for S" + + # ==== dtype info ==== + self.dtype_acc = Float32 + + # ==== TMEM info ==== + SM100_TMEM_CAPACITY_COLUMNS = 512 + self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + self.tmem_cols_S = self.tile_n // self.cta_group_size + self.tmem_cols_Oi = (self.hdimv // self.num_hdimv_splits) // self.cta_group_size + self.tmem_offset_S = [ + self.tmem_cols_S * stage for stage in range(self.num_stages_S) + ] # allocate 64 TMEM columns for each stage of S + self.tmem_offset_O0 = self.tmem_cols_S * self.num_stages_S + self.tmem_offset_O1 = self.tmem_offset_O0 + self.tmem_cols_Oi + self.tmem_offsets_O = [self.tmem_offset_O0, self.tmem_offset_O1] + self.total_tmem = self.tmem_offset_O1 + self.tmem_cols_Oi + assert self.total_tmem <= self.tmem_alloc_cols, ( + f"Total TMEM columns allocated {self.total_tmem} exceeds capacity {self.tmem_alloc_cols}" + ) + + def _get_shared_storage_cls(self): + self.buffer_align_bytes = 1024 + + def smem_struct_align(dtype, staged_layout, disabled=False): + if disabled: + return cute.struct.MemRange[dtype, 0] + return cute.struct.Align[ + cute.struct.MemRange[dtype, cute.cosize(staged_layout)], + self.buffer_align_bytes, + ] + + def mbar_struct(num_stages): + return cute.struct.MemRange[Int64, 2 * num_stages] + + (sQ_struct, sK_struct, sQv_struct, sV_struct, sP_struct) = ( + smem_struct_align(dtype, layout, disabled) + for dtype, layout, disabled in [ + (self.dtype_Q, self.sQ_layout_staged, not self.has_qk), + (self.dtype_K, self.sK_layout_staged, not self.has_qk), + (self.dtype_Qv, self.sQv_layout_staged, False), + (self.dtype_V, self.sV_layout_staged, False), + (self.dtype_P, self.sP_layout_staged, False), + ] + ) + sStats_struct = cute.struct.MemRange[Float32, cute.cosize(self.sStats_layout)] + sScale_struct = cute.struct.MemRange[Float32, cute.cosize(self.sScale_layout)] + sBitmask_struct = cute.struct.MemRange[Uint32, cute.cosize(self.sBitmask_layout)] + + ( + mbar_ptr_Q_struct, + mbar_ptr_K_struct, + mbar_ptr_Qv_struct, + mbar_ptr_V_struct, + mbar_ptr_S_struct, + mbar_ptr_P_struct, + mbar_ptr_O0_struct, + mbar_ptr_O1_struct, + mbar_sm_stats_struct, + mbar_bitmask_struct, + ) = ( + mbar_struct(n) + for n in [ + self.num_stages_Q, + self.num_stages_K, + self.num_stages_Qv, + self.num_stages_V, + self.num_stages_S, + self.num_stages_P, + self.num_stages_Oi, + self.num_stages_Oi, + self.num_stages_sm_stats, + self.num_stages_bitmask, + ] + ) + mbar_ptr_tmem_dealloc_struct = Int64 + tmem_holding_buf_struct = Int32 + + self.sched_stages = 1 + clc_response_size = self.sched_stages * 4 if self.use_clc_scheduler else 0 + clc_mbar_size = self.sched_stages * 2 if self.use_clc_scheduler else 0 + + @cute.struct + class SharedStorage: + mbar_ptr_Q: mbar_ptr_Q_struct + mbar_ptr_K: mbar_ptr_K_struct + mbar_ptr_Qv: mbar_ptr_Qv_struct + mbar_ptr_V: mbar_ptr_V_struct + mbar_ptr_S: mbar_ptr_S_struct + mbar_ptr_P: mbar_ptr_P_struct + mbar_ptr_O0: mbar_ptr_O0_struct + mbar_ptr_O1: mbar_ptr_O1_struct + mbar_ptr_K_cpasync: mbar_ptr_K_struct + mbar_ptr_V_cpasync: mbar_ptr_V_struct + mbar_ptr_sm_stats: mbar_sm_stats_struct + mbar_ptr_bitmask: mbar_bitmask_struct + mbar_ptr_tmem_dealloc: mbar_ptr_tmem_dealloc_struct + tmem_holding_buf: tmem_holding_buf_struct + clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, clc_mbar_size] + clc_response: cute.struct.MemRange[Int32, clc_response_size] + sO_empty_mbar_ptr: cutlass.Int64 + + sRowMax: sStats_struct + sRowSum: sStats_struct + sScale: sScale_struct + sBitmask: sBitmask_struct + sQv: sQv_struct + sQ: sQ_struct + sK: sK_struct + sV: sV_struct + sP: sP_struct + + # print("smem bytes = ", SharedStorage.size_in_bytes()) + + return SharedStorage + + # fmt: off + @cute.jit + def __call__( + self, + mQ: Optional[cute.Tensor], # (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q + mQv: cute.Tensor, # (b, s_q, h, dv) or (total_q, h, d) if there is cu_seqlens_q + mK: Optional[cute.Tensor], # (b, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k or (num_pages, page_size, h_k, d) if there is page_table + mV: cute.Tensor, # (b, s_k, h_k, dv) or (total_k, h_k, dv) if there is cu_seqlens_k or (num_pages, page_size, h_k, dv) if there is page_table + mO: cute.Tensor, # (b, s_q, h, dv) or (total_q, h, dv) if there is cu_seqlens_q + mLSE: Optional[cute.Tensor], # (b, s_q, h) or (total_q, h) if there is cu_seqlens_q + softmax_scale: Float32, + mP: Optional[cute.Tensor] = None, # (b, s_q, h, topk) or (total_q, h, topk) if there is cu_seqlens_q + mRowMax: Optional[cute.Tensor] = None, # (b, s_q, topk // tile_n, h) or (total_q, topk // tile_n, h) if there is cu_seqlens_q + mCuSeqlensQ: Optional[cute.Tensor] = None, # (b + 1) + mCuSeqlensK: Optional[cute.Tensor] = None, # (b + 1) + mSeqUsedQ: Optional[cute.Tensor] = None, # (b) + mSeqUsedK: Optional[cute.Tensor] = None, # (b) + mIndexTopk: Optional[cute.Tensor] = None, # (b, s_q, topk) or (total_q, topk) if there is cu_seqlens_q + mPageTable: Optional[cute.Tensor] = None, + window_size_left: Int32 | int | None = None, + window_size_right: Int32 | int | None = None, + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + # fmt: on + self.store_P = mP is not None + self.store_row_max = mRowMax is not None + + if const_expr(self.has_qk): + assert mQ is not None and mK is not None, "has_qk requires mQ and mK" + else: + assert mQ is None and mK is None, "not has_qk disallows mQ and mK" + + # ==== dtype info ==== + self.dtype_Q = mQ.element_type if self.has_qk else cutlass.BFloat16 + self.dtype_K = mK.element_type if self.has_qk else cutlass.BFloat16 + self.dtype_Qv = mQv.element_type + self.dtype_V = mV.element_type + self.dtype_P = mV.element_type + self.dtype_O = mO.element_type + + if const_expr(self.store_P): + assert mP.element_type == self.dtype_P + + # ==== Prepare Tensors ==== + new_stride = lambda mX: ( + *(cute.assume(s, divby=128 // mX.element_type.width) for s in mX.stride[:-1]), + mX.stride[-1], + ) + mQ, mQv, mK, mV, mO, mP = [ + cute.make_tensor(mX.iterator, cute.make_layout(mX.shape, stride=new_stride(mX))) + if mX is not None + else None + for mX in (mQ, mQv, mK, mV, mO, mP) + ] + + # (b, s, h, d) -> (s, d, h, b) or + # (total, h, d) -> (total, d, h) or + # (num_pages, page_size, h_k, d) -> (page_size, d, h_k, num_pages) + QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1] + KV_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensK is None) else [0, 2, 1] + mQ, mQv, mO, mP = [ + cute.make_tensor(mX.iterator, cute.select(mX.layout, mode=QO_layout_transpose)) + if mX is not None + else None + for mX in (mQ, mQv, mO, mP) + ] + mK, mV = [ + cute.make_tensor(mX.iterator, cute.select(mX.layout, mode=KV_layout_transpose)) + if mX is not None + else None + for mX in (mK, mV) + ] + # (s_k, dv, h_k, b) -> (dv, s_k, h_k, b) or + # (total_k, dv, h_k) -> (dv, total_k, h_k) + V_layout_transpose = [1, 0, 2, 3] if const_expr(mCuSeqlensK is None) else [1, 0, 2] + mVt = cute.make_tensor(mV.iterator, cute.select(mV.layout, mode=V_layout_transpose)) + # (b, s_q, topk) -> (topk, s_q, b) or (total_q, topk) -> (topk, total_q) + topk_layout_transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0] + mIndexTopk = ( + cute.make_tensor( + mIndexTopk.iterator, cute.select(mIndexTopk.layout, mode=topk_layout_transpose) + ) + if mIndexTopk is not None + else None + ) + # (b, s_q, h) -> (s_q, h, b) or (total_q, h) -> (total_q, h) + LSE_layout_transpose = [1, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 1] + mLSE = ( + cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose)) + if mLSE is not None + else None + ) + # (b, s, topk//128, h) => (s, topk//128, h, b) or + # (total, topk//128, h) == (total, topk//128, h) + rowmax_layout_transpose = [1, 2, 3, 0] if const_expr(mCuSeqlensQ is None) else [0, 1, 2] + if const_expr(mRowMax is not None): + mRowMax = cute.make_tensor( + mRowMax.iterator, cute.select(mRowMax.layout, mode=rowmax_layout_transpose) + ) + + topk_length_dynamic = mIndexTopk.shape[0] if mIndexTopk is not None else None + + self.o_layout = cutlass.utils.LayoutEnum.from_tensor(mO) + self.p_layout = cutlass.utils.LayoutEnum.ROW_MAJOR + if const_expr(self.store_P): + assert cutlass.utils.LayoutEnum.from_tensor(mP) == self.p_layout + + mO_og = mO + mP_og = mP + if const_expr(self.pack_gqa): + mQ, mQv, mO, mP, mRowMax = [ + pack_gqa_layout(mX, self.qhead_per_kvhead, self.nheads_kv, head_idx=2) + if mX is not None + else None + for mX in (mQ, mQv, mO, mP, mRowMax) + ] + if const_expr(mLSE is not None): + mLSE = pack_gqa_layout(mLSE, self.qhead_per_kvhead, self.nheads_kv, head_idx=1) + + # ==== Prepare MMAs ==== + # (local_var, dtype_a, major_a, major_b, mma_tiler, operand_source_a) + # fmt: off + _mma_specs = [ + ("tiled_mma_QK", self.dtype_Q, self.major_mode_Q, self.major_mode_K, self.mma_tiler_QK, self.operand_source_Q), + ("tiled_mma_QvV", self.dtype_Qv, self.major_mode_Qvi, self.major_mode_Vi, self.mma_tiler_QvV, self.operand_source_Qvi), + ("tiled_mma_PVt", self.dtype_P, self.major_mode_P, self.major_mode_Vti, self.mma_tiler_PVt, self.operand_source_P), + ] + tiled_mma_QK, tiled_mma_QvV, tiled_mma_PVt = ( + sm100_utils.make_trivial_tiled_mma( + dtype_a, major_a, major_b, self.dtype_acc, self.cta_group, mma_tiler[:2], operand_source_a, + ) + for _, dtype_a, major_a, major_b, mma_tiler, operand_source_a in _mma_specs + ) + # fmt: on + + # ==== Prepare SMEM layouts and TMAs ==== + # (attr, make_fn, tiled_mma, mma_tiler, dtype, num_stages) + # fmt: off + _smem_layout_specs = [ + ("sQ_layout", sm100_utils.make_smem_layout_a, tiled_mma_QK, self.mma_tiler_QK, self.dtype_Q, self.num_stages_Q), + ("sK_layout", sm100_utils.make_smem_layout_b, tiled_mma_QK, self.mma_tiler_QK, self.dtype_K, self.num_stages_K), + ("sP_layout", sm100_utils.make_smem_layout_a, tiled_mma_PVt, self.mma_tiler_PVt, self.dtype_P, self.num_stages_P), + ("sQv_layout", sm100_utils.make_smem_layout_a, tiled_mma_QvV, self.mma_tiler_QvV, self.dtype_Qv, self.num_stages_Qv), + ("sV_layout", sm100_utils.make_smem_layout_b, tiled_mma_QvV, self.mma_tiler_QvV, self.dtype_V, self.num_stages_V), + ("sVt_layout", sm100_utils.make_smem_layout_b, tiled_mma_PVt, self.mma_tiler_PVt, self.dtype_V, self.num_stages_V), + ] + for attr, make_fn, tiled_mma, mma_tiler, dtype, num_stages in _smem_layout_specs: + ab_kwarg = "a_dtype" if make_fn is sm100_utils.make_smem_layout_a else "b_dtype" + staged = make_fn( + tiled_mma=tiled_mma, + mma_tiler_mnk=mma_tiler, + num_stages=num_stages, + **{ab_kwarg: dtype}, + ) + setattr(self, f"{attr}_staged", staged) + setattr(self, attr, cute.select(staged, mode=[0, 1, 2])) + # fmt: on + + self.sStats_layout = cute.make_layout((self.cta_tile_m, self.cta_group_size)) + self.sScale_layout = cute.make_layout((self.cta_tile_m, self.num_stages_sm_stats)) + self.sBitmask_layout = cute.make_layout((self.tile_n // 32, self.num_stages_bitmask)) + + # fmt: off + for attr, dtype, layout in [ + ("tma_copy_bytes_Q", self.dtype_Q, self.sQ_layout), + ("tma_copy_bytes_K", self.dtype_K, self.sK_layout), + ("tma_copy_bytes_Qvi", self.dtype_Qv, self.sQv_layout), + ("tma_copy_bytes_Vi", self.dtype_V, self.sV_layout), + ]: + setattr(self, attr, cute.size_in_bytes(dtype, layout) * self.cta_group_size) + # fmt: on + + tma_load_op = cpasync.CopyBulkTensorTileG2SOp(self.cta_group) + cta_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), (tiled_mma_QK.thr_id.shape,) + ) + cta_shape = cta_layout_vmnk.shape + + def make_tma(make_fn, mX, smem_layout, mma_tiler, tiled_mma): + return make_fn(tma_load_op, mX, smem_layout, mma_tiler, tiled_mma, cta_shape) + + A, B = cute.nvgpu.make_tiled_tma_atom_A, cute.nvgpu.make_tiled_tma_atom_B + + # (atom_name, tensor_name, make_fn, m, smem_layout, mma_tiler, tiled_mma, kv_only) + # fmt: off + _tma_specs = [ + ("tma_atom_Q", "tma_tensor_Q", A, mQ, self.sQ_layout, self.mma_tiler_QK, tiled_mma_QK, False), + ("tma_atom_Qv", "tma_tensor_Qv", A, mQv, self.sQv_layout, self.mma_tiler_QvV, tiled_mma_QvV, False), + ("tma_atom_K", "tma_tensor_K", B, mK, self.sK_layout, self.mma_tiler_QK, tiled_mma_QK, True), + ("tma_atom_V", "tma_tensor_V", B, mV, self.sV_layout, self.mma_tiler_QvV, tiled_mma_QvV, True), + ("tma_atom_Vt", "tma_tensor_Vt", B, mVt, self.sVt_layout, self.mma_tiler_PVt, tiled_mma_PVt, True), + ] + _tmas = {} + for atom_name, tensor_name, make_fn, m, smem_layout, mma_tiler, tiled_mma, kv_only in _tma_specs: + _tmas[atom_name], _tmas[tensor_name] = ( + make_tma(make_fn, m, smem_layout, mma_tiler, tiled_mma) + if const_expr((not kv_only or self.use_tma_KV) and m is not None) + else (None, None) + ) + + (tma_atom_Q, tma_tensor_Q, + tma_atom_Qv, tma_tensor_Qv, + tma_atom_K, tma_tensor_K, + tma_atom_V, tma_tensor_V, + tma_atom_Vt, tma_tensor_Vt) = _tmas.values() + # fmt: on + + tma_store_op = cpasync.CopyBulkTensorTileS2GOp() + self.ragged_tma_O = ( + self.use_tma_O + and self.is_varlen_q + and self.pack_gqa + and self.cta_tile_m % self.qhead_per_kvhead == 0 + ) + make_tiled_tma_atom_fn = ( + partial(make_packgqa_tiled_tma_atom, qhead_per_kvhead=self.qhead_per_kvhead, head_idx=2) + if const_expr(self.ragged_tma_O) + else cpasync.make_tiled_tma_atom + ) + + # ==== Set up P smem -> gmem tma store ==== + + # S<3,4,3> o 0 o ((8,8),(64,2),(1,1)):((64,512),(1,4096),(0,0)) + sP_layout_out = sm100_utils.make_smem_layout_epi( + self.dtype_P, self.p_layout, self.tile_P, self.num_stages_P + ) + + if const_expr(self.store_P): + # TODO: add asserts + mP_tma = mP_og if const_expr(self.ragged_tma_O) else mP + if const_expr(self.ragged_tma_O): + mP_tma = copy_utils.create_ragged_tensor_for_tma( + mP_tma, ragged_dim=0, ptr_shift=True + ) + tma_atom_P, tma_tensor_P = make_tiled_tma_atom_fn( + tma_store_op, mP_tma, cute.select(sP_layout_out, mode=[0, 1]), self.tile_P + ) + else: + tma_atom_P = None + tma_tensor_P = None + + # ==== Set up Oi smem -> gmem tma store ==== + + self.overlap_sO_sV = True + if const_expr(self.overlap_sO_sV): + num_stages_sO = self.num_stages_V + else: + num_stages_sO = self.num_hdimv_splits + sO_layout = sm100_utils.make_smem_layout_epi( + self.dtype_O, self.o_layout, self.epi_tile, num_stages_sO + ) + + if const_expr(self.use_tma_O): + mO_tma = mO_og if const_expr(self.ragged_tma_O) else mO + if const_expr(self.ragged_tma_O): + mO_tma = copy_utils.create_ragged_tensor_for_tma( + mO_tma, ragged_dim=0, ptr_shift=True + ) + + tma_atom_O, tma_tensor_O = make_tiled_tma_atom_fn( + tma_store_op, mO_tma, cute.select(sO_layout, mode=[0, 1]), self.epi_tile + ) + else: + tma_atom_O = None + tma_tensor_O = None + + # ==== Set up Oi rmem -> gmem copy ==== + universal_copy_bits = 128 + atom_universal_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.dtype_O, + num_bits_per_copy=universal_copy_bits, + ) + thread_layout_O_r2g = cute.make_layout((64, 2), stride=(1, 64)) + value_layout_O_r2g = cute.make_layout( + (1, self.hdimv // self.num_hdimv_splits // self.cta_group_size) + ) + tiled_copy_O_r2g = cute.make_tiled_copy_tv( + atom=atom_universal_copy, + thr_layout=thread_layout_O_r2g, + val_layout=value_layout_O_r2g, + ) + + # ==== Allocate shared memory ==== + SharedStorage = self._get_shared_storage_cls() + + # ==== Tile scheduler ==== + + TileScheduler = self.TileScheduler + + batch_size_for_sched = ( + cute.size(mQv.shape[3]) if const_expr(mCuSeqlensQ is None) + else cute.size(mCuSeqlensQ.shape[0] - 1) if self.use_varlen_scheduler + else 1 + ) + + tile_sched_args = TileSchedulerArguments( + num_block=cute.ceil_div(cute.size(mQv.shape[0]), self.cluster_tile_m), + num_head=cute.size(mQv.shape[2]), + num_batch=batch_size_for_sched, + num_splits=1, # todo: split_kv + seqlen_k=cute.size(mV.shape[0]) + if const_expr(mPageTable is None) + else cute.size(mV.shape[0]) * cute.size(mPageTable.shape[1]), + headdim=self.hdim, + headdim_v=self.hdimv, + total_q=cute.size(mQv.shape[0]) + if const_expr(mCuSeqlensQ is not None) + else cute.size(mQv.shape[0]) * cute.size(mQv.shape[3]), + tile_shape_mn=(self.cta_tile_m, self.tile_n), + mCuSeqlensQ=mCuSeqlensQ, + mSeqUsedQ=mSeqUsedQ, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + element_size=self.dtype_K.width // 8, + is_persistent=self.is_persistent, + # lpt=self.is_causal or self.is_local, + lpt=False, + is_split_kv=False, + cluster_shape_mn=self.cluster_shape_mn, + use_cluster_idx=True, + ) + tile_sched_params = TileScheduler.to_underlying_arguments( + tile_sched_args, scheduling_mode=self.scheduling_mode + ) + self.tile_scheduler_cls = TileScheduler + grid_dim = TileScheduler.get_grid_shape(tile_sched_params) + fa_printf(1, "grid = {}", grid_dim) + + # ==== Named Barrier ==== + self.cpasync_barrier = cutlass.pipeline.NamedBarrier( + barrier_id=int(NamedBarrierFwdSm100_MLA2CTA.Cpasync), + num_threads=self.num_cpasync_load_threads, + ) + self.softmax_barrier = cutlass.pipeline.NamedBarrier( + barrier_id=int(NamedBarrierFwdSm100_MLA2CTA.Softmax), + num_threads=self.num_softmax_threads, + ) + self.epi_barrier = cutlass.pipeline.NamedBarrier( + barrier_id=int(NamedBarrierFwdSm100_MLA2CTA.Epilogue), + num_threads=self.num_epilogue_threads, + ) + # softmax -> correction + self.sm_stats_barrier_full = cutlass.pipeline.NamedBarrier( + barrier_id=int(NamedBarrierFwdSm100_MLA2CTA.SoftmaxStatsFull), + num_threads=self.num_softmax_threads + self.num_epilogue_threads, + ) + self.sm_stats_barrier_empty = cutlass.pipeline.NamedBarrier( + barrier_id=int(NamedBarrierFwdSm100_MLA2CTA.SoftmaxStatsEmpty), + num_threads=self.num_softmax_threads + self.num_epilogue_threads, + ) + + LOG2_E = math.log2(math.e) + softmax_scale_log2 = softmax_scale * LOG2_E + + # ==== Launch kernel ==== + self.kernel( + tma_tensor_Q, + tma_tensor_Qv, + tma_tensor_K if self.use_tma_KV else mK, + tma_tensor_V if self.use_tma_KV else mV, + tma_tensor_Vt if self.use_tma_KV else mVt, + tma_tensor_O if self.use_tma_O else mO, + tma_tensor_P, + mLSE, + mRowMax, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + mIndexTopk, + mPageTable, + tma_atom_Q, + tma_atom_Qv, + tma_atom_K, + tma_atom_V, + tma_atom_Vt, + tma_atom_O, + tma_atom_P, + tiled_copy_O_r2g, + self.sQ_layout_staged, + self.sK_layout_staged, + self.sQv_layout_staged, + self.sV_layout_staged, + self.sVt_layout_staged, + self.sP_layout_staged, + self.sStats_layout, + self.sScale_layout, + self.sBitmask_layout, + sO_layout, + sP_layout_out, + tiled_mma_QK, + tiled_mma_QvV, + tiled_mma_PVt, + softmax_scale, + softmax_scale_log2, + topk_length_dynamic, + tile_sched_params, + SharedStorage, + ).launch( + grid=grid_dim, + block=( + self.num_threads, + 1, + 1, + ), + cluster=self.cluster_shape_mnk, + smem=SharedStorage.size_in_bytes(), + stream=stream, + ) + + @cute.kernel + def kernel( + self, + mQ: Optional[cute.Tensor], + mQv: cute.Tensor, + mK: Optional[cute.Tensor], + mV: cute.Tensor, + mVt: cute.Tensor, + mO: cute.Tensor, + mP: Optional[cute.Tensor], + mLSE: Optional[cute.Tensor], + mRowMax: Optional[cute.Tensor], + mCuSeqlensQ: Optional[cute.Tensor], + mCuSeqlensK: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + mIndexTopk: Optional[cute.Tensor], + mPageTable: Optional[cute.Tensor], + tma_atom_Q: cute.CopyAtom, + tma_atom_Qv: cute.CopyAtom, + tma_atom_K: Optional[cute.CopyAtom], + tma_atom_V: Optional[cute.CopyAtom], + tma_atom_Vt: Optional[cute.CopyAtom], + tma_atom_O: Optional[cute.CopyAtom], + tma_atom_P: Optional[cute.CopyAtom], + tiled_copy_O_r2g: cute.TiledCopy, + sQ_layout_staged: cute.ComposedLayout, + sK_layout_staged: cute.ComposedLayout, + sQv_layout_staged: cute.ComposedLayout, + sV_layout_staged: cute.ComposedLayout, + sVt_layout_staged: cute.ComposedLayout, + sP_layout_staged: cute.ComposedLayout, + sStats_layout: cute.Layout, + sScale_layout: cute.Layout, + sBitmask_layout: cute.Layout, + sO_layout: cute.ComposedLayout, + sP_layout_out: cute.ComposedLayout, + tiled_mma_QK: cute.TiledMma, + tiled_mma_QvV: cute.TiledMma, + tiled_mma_PVt: cute.TiledMma, + softmax_scale: Float32, + softmax_scale_log2: Float32, + topk_length_dynamic: Optional[Int32], + tile_sched_params: ParamsBase, + SharedStorage: cutlass.Constexpr[Callable], + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + cta_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), (tiled_mma_QvV.thr_id.shape,) + ) + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + mma_tile_coord_v = cta_rank_in_cluster % cute.size(tiled_mma_QvV.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + + # ==== Allocate SMEM ==== + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # ==== TMEM stuff ==== + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=int(NamedBarrierFwdSm100_MLA2CTA.TmemPtr), + num_threads=self.num_mma_threads + self.num_softmax_threads + self.num_epilogue_threads, + ) + tmem = cutlass.utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.mma_warp_id, + is_two_cta=self.use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.mbar_ptr_tmem_dealloc, + ) + + # ==== Prefetch TMA descriptors ==== + if warp_idx == self.load_warp_id: + if const_expr(self.has_qk): + cpasync.prefetch_descriptor(tma_atom_Q) + cpasync.prefetch_descriptor(tma_atom_Qv) + if const_expr(self.use_tma_KV): + if const_expr(self.has_qk): + cpasync.prefetch_descriptor(tma_atom_K) + cpasync.prefetch_descriptor(tma_atom_V) + cpasync.prefetch_descriptor(tma_atom_Vt) + if const_expr(self.use_tma_O): + cpasync.prefetch_descriptor(tma_atom_O) + + # ==== Construct pipelines ==== + tma_warp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + mma_warp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 1) + sm_threads = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_softmax_threads) + epi_threads = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_epilogue_threads) + sm_threads_cluster = pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_softmax_threads * self.cta_group_size + ) + epi_threads_cluster = pipeline.CooperativeGroup( + pipeline.Agent.Thread, self.num_epilogue_threads * self.cta_group_size + ) + + TmaUmma = pipeline.PipelineTmaUmma + AsyncUmma = pipeline.PipelineAsyncUmma + UmmaAsync = pipeline.PipelineUmmaAsync + Async = pipeline.PipelineAsync + + def make_pipeline(cls, mbar_ptr, num_stages, producer, consumer, tx_count=None): + return cls.create( + barrier_storage=mbar_ptr.data_ptr(), + num_stages=num_stages, + producer_group=producer, + consumer_group=consumer, + defer_sync=True, + **({"cta_layout_vmnk": cta_layout_vmnk} if cls is not Async else {}), + **({"tx_count": tx_count} if tx_count is not None else {}), + ) + + # Unconditional pipelines + # fmt: off + pipeline_Q = None + if const_expr(self.has_qk): + pipeline_Q = make_pipeline(TmaUmma, storage.mbar_ptr_Q, self.num_stages_Q, tma_warp, mma_warp, self.tma_copy_bytes_Q) + pipeline_Qv = make_pipeline(TmaUmma, storage.mbar_ptr_Qv, self.num_stages_Qv, tma_warp, mma_warp, self.tma_copy_bytes_Qvi) + pipeline_S = make_pipeline(UmmaAsync, storage.mbar_ptr_S, self.num_stages_S, mma_warp, sm_threads_cluster) + pipeline_P = make_pipeline(AsyncUmma, storage.mbar_ptr_P, self.num_stages_P, sm_threads_cluster, mma_warp) + pipeline_O0 = make_pipeline(UmmaAsync, storage.mbar_ptr_O0, self.num_stages_Oi, mma_warp, epi_threads_cluster) + pipeline_O1 = make_pipeline(UmmaAsync, storage.mbar_ptr_O1, self.num_stages_Oi, mma_warp, epi_threads_cluster) + pipeline_sm_stats = make_pipeline(Async, storage.mbar_ptr_sm_stats, self.num_stages_sm_stats, sm_threads, epi_threads) + + # K/V pipelines: type and producer depend on use_tma_KV + if const_expr(self.use_tma_KV): + pipeline_K = None + if const_expr(self.has_qk): + pipeline_K = make_pipeline(TmaUmma, storage.mbar_ptr_K, self.num_stages_K, tma_warp, mma_warp, self.tma_copy_bytes_K) + pipeline_V = make_pipeline(TmaUmma, storage.mbar_ptr_V, self.num_stages_V, tma_warp, mma_warp, self.tma_copy_bytes_Vi) + pipeline_K_cpasync = pipeline_V_cpasync = pipeline_bitmask = None + else: + cpasync_load_threads = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_cpasync_load_threads) + relay_warps_cluster = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.cta_group_size) + relay_threads = pipeline.CooperativeGroup(pipeline.Agent.Thread, self.num_relay_threads) + pipeline_K = pipeline_K_cpasync = None + if const_expr(self.has_qk): + pipeline_K = make_pipeline(AsyncUmma, storage.mbar_ptr_K, self.num_stages_K, relay_warps_cluster, mma_warp) + pipeline_K_cpasync = make_pipeline(Async, storage.mbar_ptr_K_cpasync, self.num_stages_K, cpasync_load_threads, relay_threads) + pipeline_V = make_pipeline(AsyncUmma, storage.mbar_ptr_V, self.num_stages_V, relay_warps_cluster, mma_warp) + pipeline_V_cpasync = make_pipeline(Async, storage.mbar_ptr_V_cpasync, self.num_stages_V, cpasync_load_threads, relay_threads) + pipeline_bitmask = ( + make_pipeline(Async, storage.mbar_ptr_bitmask, self.num_stages_bitmask, cpasync_load_threads, sm_threads) + if const_expr(self.is_topk_gather and not self.disable_bitmask) else None + ) + # fmt: on + + sO_empty_mbar_ptr = None + if const_expr(self.use_tma_O and self.overlap_sO_sV): + sO_empty_mbar_ptr = storage.sO_empty_mbar_ptr + if warp_idx == 0: + cute.arch.mbarrier_init(sO_empty_mbar_ptr, 1) + + pipeline.pipeline_init_arrive(cluster_shape_mn=cta_layout_vmnk, is_relaxed=True) + + # ==== Get SMEM tensors ==== + # fmt: off + sQ, sK, sQv, sV, sVt, sP, sP_out = ( + store.get_tensor(layout.outer, swizzle=layout.inner) + if const_expr(store._size > 0) else None + for store, layout in [ + (storage.sQ, sQ_layout_staged), + (storage.sK, sK_layout_staged), + (storage.sQv, sQv_layout_staged), + (storage.sV, sV_layout_staged), + (storage.sV, sVt_layout_staged), # sVt reuses sV storage + (storage.sP, sP_layout_staged), + (storage.sP, sP_layout_out), + ] + ) + # fmt: on + sRowMax = storage.sRowMax.get_tensor(sStats_layout) + sRowSum = storage.sRowSum.get_tensor(sStats_layout) + sScale = storage.sScale.get_tensor(sScale_layout) + sBitmask = None + if const_expr(self.is_topk_gather): + sBitmask = storage.sBitmask.get_tensor(sBitmask_layout) + + if const_expr(self.overlap_sO_sV): + sO_iterator = sV.iterator + assert cute.cosize(sO_layout) <= cute.cosize(sV_layout_staged) + else: + sO_iterator = sQv.iterator + assert cute.cosize(sO_layout) <= cute.cosize(sQv_layout_staged) + sO = cute.make_tensor( + cute.recast_ptr(sO_iterator, sO_layout.inner, self.dtype_O), sO_layout.outer + ) + + # ==== Get thread MMAs and accumulator fragments ==== + thr_mma_QK = tiled_mma_QK.get_slice(mma_tile_coord_v) + thr_mma_QvV = tiled_mma_QvV.get_slice(mma_tile_coord_v) + thr_mma_PVt = tiled_mma_PVt.get_slice(mma_tile_coord_v) + + acc_shape_S = thr_mma_QvV.partition_shape_C(self.mma_tiler_QvV[:2]) + tStS_fake = thr_mma_QvV.make_fragment_C(cute.append(acc_shape_S, self.num_stages_S)) + + acc_shape_Oi = thr_mma_PVt.partition_shape_C(self.mma_tiler_PVt[:2]) + tOtO0_fake = thr_mma_PVt.make_fragment_C(acc_shape_Oi) + tOtO1_fake = thr_mma_PVt.make_fragment_C(acc_shape_Oi) + + block_info = BlockInfo( + self.cta_tile_m * self.cta_group_size, + self.tile_n, + is_causal=self.is_causal, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + SeqlenInfoCls = partial( + SeqlenInfoQK.create, + seqlen_q_static=mQv.shape[0] if const_expr(not self.pack_gqa) else mQv.shape[0][1], + seqlen_k_static=mV.shape[0] + if const_expr(mPageTable is None) + else mV.shape[0] * mPageTable.shape[1], + tile_m=self.cta_tile_m, + tile_n=self.tile_n, + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + ) + AttentionMaskCls = partial( + AttentionMask, + self.cta_tile_m * self.cta_group_size, + self.tile_n, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + + if const_expr(self.use_clc_scheduler): + clc_response_ptr = storage.clc_response.data_ptr() + clc_mbar_ptr = storage.clc_mbar_ptr.data_ptr() + + clc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_clc_consumer_warps_per_cta = self.num_threads // cute.arch.WARP_SIZE + num_clc_consumer_warps = num_clc_consumer_warps_per_cta * self.cta_group_size + clc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, cute.arch.WARP_SIZE * num_clc_consumer_warps + ) + clc = ClcState.create( + hw_scheduler=ClcDynamicPersistentTileScheduler.create( + self.tile_scheduler_cls.clc_problem_shape(tile_sched_params), + cute.arch.block_idx(), + cute.arch.grid_dim(), + clc_response_ptr, + ), + pipeline=pipeline.PipelineClcFetchAsync.create( + barrier_storage=clc_mbar_ptr, + num_stages=self.sched_stages, + producer_group=clc_pipeline_producer_group, + consumer_group=clc_pipeline_consumer_group, + tx_count=16, + cta_layout_vmnk=cta_layout_vmnk, + ), + consumer_state=pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.sched_stages + ), + producer_state=pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.sched_stages + ), + ) + tile_scheduler = self.tile_scheduler_cls.create(tile_sched_params, clc=clc) + else: + tile_scheduler = self.tile_scheduler_cls.create(tile_sched_params) + assert isinstance(tile_scheduler, TileSchedulerProtocol), ( + f"tile_scheduler is not a TileSchedulerProtocol: {type(tile_scheduler)}" + ) + + pipeline.pipeline_init_wait(cluster_shape_mn=cta_layout_vmnk) + + if const_expr(self.use_clc_scheduler): + if warp_idx == self.clc_scheduler_warp_id: + if const_expr(self.num_regs_other < self.num_regs_per_thread): + cute.arch.setmaxregister_decrease(self.num_regs_other) + if is_leader_cta: + self.clc_scheduler_warp(tile_scheduler) + else: + self.empty_warp(tile_scheduler) + for i in cutlass.range_constexpr(len(self.empty_warp_ids)): + if warp_idx == self.empty_warp_ids[i] and warp_idx != self.clc_scheduler_warp_id: + if const_expr(self.num_regs_other < self.num_regs_per_thread): + cute.arch.setmaxregister_decrease(self.num_regs_other) + self.empty_warp(tile_scheduler) + else: + for i in cutlass.range_constexpr(len(self.empty_warp_ids)): + if warp_idx == self.empty_warp_ids[i]: + if const_expr(self.num_regs_other < self.num_regs_per_thread): + cute.arch.setmaxregister_decrease(self.num_regs_other) + + if const_expr(self.use_cpasync_load_KV): + if warp_idx == self.relay_warp_id: + if const_expr(self.num_regs_load < self.num_regs_per_thread): + cute.arch.setmaxregister_decrease(self.num_regs_load) + self.relay( + pipeline_K, + pipeline_V, + pipeline_K_cpasync, + pipeline_V_cpasync, + sO_empty_mbar_ptr, + topk_length_dynamic, + block_info, + SeqlenInfoCls, + tile_scheduler=tile_scheduler, + mCuSeqlensQ=mCuSeqlensQ, + ) + + if warp_idx in self.cpasync_load_warp_indices: + if const_expr(self.num_regs_cpasync < self.num_regs_per_thread): + cute.arch.setmaxregister_decrease(self.num_regs_cpasync) + self.load_cpasync( + mIndexTopk, + mK, + mV, + mVt, + sK, + sV, + sVt, + sBitmask, + pipeline_K, + pipeline_V, + pipeline_K_cpasync, + pipeline_V_cpasync, + pipeline_bitmask, + sO_empty_mbar_ptr, + topk_length_dynamic, + block_info, + SeqlenInfoCls, + tile_scheduler=tile_scheduler, + mPageTable=mPageTable, + mCuSeqlensQ=mCuSeqlensQ, + ) + + if warp_idx == self.load_warp_id: + if const_expr(self.num_regs_load < self.num_regs_per_thread): + cute.arch.setmaxregister_decrease(self.num_regs_load) + self.load( + mQ, + mK, + mQv, + mV, + mVt, + sQ, + sK, + sQv, + sV, + sVt, + tma_atom_Q, + tma_atom_K, + tma_atom_Qv, + tma_atom_V, + tma_atom_Vt, + pipeline_Q, + pipeline_K, + pipeline_Qv, + pipeline_V, + sO_empty_mbar_ptr, + thr_mma_QK, + thr_mma_QvV, + thr_mma_PVt, + topk_length_dynamic, + block_info, + SeqlenInfoCls, + tile_scheduler=tile_scheduler, + mPageTable=mPageTable, + mCuSeqlensQ=mCuSeqlensQ, + ) + + if warp_idx == self.mma_warp_id: + if const_expr(self.num_regs_mma < self.num_regs_per_thread): + cute.arch.setmaxregister_decrease(self.num_regs_mma) + # ==== Allocate TMEM ==== + tmem.allocate(self.tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.dtype_acc) + tStS = cute.make_tensor(tmem_ptr, tStS_fake.layout) + tOtO0 = cute.make_tensor(tmem_ptr + self.tmem_offset_O0, tOtO0_fake.layout) + tOtO1 = cute.make_tensor(tmem_ptr + self.tmem_offset_O1, tOtO1_fake.layout) + self.mma( + sQ, + sK, + sQv, + sV, + sVt, + sP, + tStS, + tOtO0, + tOtO1, + tiled_mma_QK, + tiled_mma_QvV, + tiled_mma_PVt, + pipeline_Q, + pipeline_K, + pipeline_Qv, + pipeline_V, + pipeline_S, + pipeline_P, + pipeline_O0, + pipeline_O1, + sO_empty_mbar_ptr, + is_leader_cta, + topk_length_dynamic, + block_info, + SeqlenInfoCls, + tile_scheduler=tile_scheduler, + mCuSeqlensQ=mCuSeqlensQ, + ) + tmem.relinquish_alloc_permit() + tmem_alloc_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + + if warp_idx in self.softmax_warp_indices: + if const_expr(self.num_regs_softmax > self.num_regs_per_thread): + cute.arch.setmaxregister_increase(self.num_regs_softmax) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.dtype_acc) + tStS = cute.make_tensor(tmem_ptr, tStS_fake.layout) + self.softmax_loop( + softmax_scale, + softmax_scale_log2, + mLSE, + mRowMax, + sRowMax, + sRowSum, + sScale, + sBitmask, + sP, + tStS, + thr_mma_QvV, + pipeline_S, + pipeline_P, + pipeline_sm_stats, + pipeline_bitmask, + sO_empty_mbar_ptr, + AttentionMaskCls, + topk_length_dynamic, + block_info, + SeqlenInfoCls, + tile_scheduler=tile_scheduler, + tma_atom_P=tma_atom_P, + mP=mP, + sP_out=sP_out, + mCuSeqlensQ=mCuSeqlensQ, + ) + tmem_alloc_barrier.arrive() + + if warp_idx in self.epilogue_warp_indices: + if const_expr(self.num_regs_epilogue < self.num_regs_per_thread): + cute.arch.setmaxregister_decrease(self.num_regs_epilogue) + elif const_expr(self.num_regs_epilogue > self.num_regs_per_thread): + cute.arch.setmaxregister_increase(self.num_regs_epilogue) + + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.dtype_acc) + tOtO0 = cute.make_tensor(tmem_ptr + self.tmem_offset_O0, tOtO0_fake.layout) + tOtO1 = cute.make_tensor(tmem_ptr + self.tmem_offset_O1, tOtO1_fake.layout) + self.correction_loop( + softmax_scale_log2, + mO, + mLSE, + tma_atom_O, + sRowMax, + sRowSum, + sScale, + sO, + tOtO0, + tOtO1, + pipeline_O0, + pipeline_O1, + pipeline_sm_stats, + sO_empty_mbar_ptr, + tiled_copy_O_r2g, + topk_length_dynamic, + block_info, + SeqlenInfoCls, + tile_scheduler=tile_scheduler, + mCuSeqlensQ=mCuSeqlensQ, + ) + tmem_alloc_barrier.arrive() + + @cute.jit + def clc_scheduler_warp( + self, + tile_scheduler: TileSchedulerProtocol, + ): + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + tile_scheduler.prefetch_next_work() + work_tile = tile_scheduler.advance_to_next_work() + # cluster_m_block, head_idx, batch_idx, _ = work_tile.tile_idx + if cute.arch.thread_idx()[0] == self.clc_scheduler_warp_id * cute.arch.WARP_SIZE: + fa_printf( + 3, + "[CLC] query sm={} cta={} (m_blk={},h={},b={},s={}) valid={}\n", + smid(), + cute.arch.block_idx()[0], + work_tile.tile_idx[0], + work_tile.tile_idx[1], + work_tile.tile_idx[2], + work_tile.tile_idx[3], + work_tile.is_valid_tile, + ) + tile_scheduler.producer_tail() + + @cute.jit + def empty_warp( + self, + tile_scheduler: TileSchedulerProtocol, + ): + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + work_tile = tile_scheduler.advance_to_next_work() + + @cute.jit + def relay( + self, + pipeline_K: Optional[pipeline.PipelineAsyncUmma], + pipeline_V: pipeline.PipelineAsyncUmma, + pipeline_K_cpasync: Optional[pipeline.PipelineAsync], + pipeline_V_cpasync: pipeline.PipelineAsync, + sO_empty_mbar_ptr: Optional[cute.Pointer], + topk_length_dynamic: Optional[Int32], + block_info: BlockInfo, + SeqlenInfoCls: Callable, + tile_scheduler: TileSchedulerProtocol, + mCuSeqlensQ: Optional[cute.Tensor] = None, + ): + # ==== Make pipeline states ==== + # pipeline_{K,V0,V1} producer + # pipeline_{K,V0,V1}_cpasync consumer + Producer, Consumer = pipeline.PipelineUserType.Producer, pipeline.PipelineUserType.Consumer + relay_K_fn = None + if const_expr(self.has_qk): + producer_state_K = pipeline.make_pipeline_state(Producer, stages=self.num_stages_K) + consumer_state_K = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_K) + relay_K_fn = partial(self.relay_inner, pipeline_K_cpasync, pipeline_K) + + producer_state_V = pipeline.make_pipeline_state(Producer, stages=self.num_stages_V) + consumer_state_V = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_V) + relay_V_fn = partial(self.relay_inner, pipeline_V_cpasync, pipeline_V) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + cluster_m_block, head_idx, batch_idx, _ = work_tile.tile_idx + if const_expr(self.use_packed_varlen_sched): + batch_idx = get_batch_from_cu_tensor(cluster_m_block, mCuSeqlensQ) + if const_expr(not self.is_topk_gather): + cluster_m_block -= mCuSeqlensQ[batch_idx] + + seqlen = SeqlenInfoCls(batch_idx) + if const_expr(self.is_topk_gather): + n_block_min = 0 + n_block_max = self.topk_length // self.tile_n + # n_block_max = topk_length_dynamic // self.tile_n + else: + n_block_min, n_block_max = block_info.get_n_block_min_max( + seqlen, + cluster_m_block, + ) + num_n_blocks = n_block_max - n_block_min + + # ==== Prologue ==== + # relay K, V0, V1 + if const_expr(self.has_qk): + consumer_state_K, producer_state_K = relay_K_fn(consumer_state_K, producer_state_K) + for _ in cutlass.range_constexpr(self.num_hdimv_splits): + consumer_state_V, producer_state_V = relay_V_fn(consumer_state_V, producer_state_V) + + # ==== Mainloop ==== + for _ in cutlass.range(num_n_blocks - 1, unroll=2): + # relay K, V0, V1, Vt0, Vt1 + if const_expr(self.has_qk): + consumer_state_K, producer_state_K = relay_K_fn( + consumer_state_K, producer_state_K + ) + for _ in cutlass.range_constexpr(2 * self.num_hdimv_splits): + consumer_state_V, producer_state_V = relay_V_fn( + consumer_state_V, producer_state_V + ) + + # ==== Epilogue === + # relay Vt0, Vt1 + for _ in cutlass.range_constexpr(self.num_hdimv_splits): + consumer_state_V, producer_state_V = relay_V_fn(consumer_state_V, producer_state_V) + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() + + if const_expr(self.has_qk): + pipeline_K.producer_tail(producer_state_K) + pipeline_V.producer_tail(producer_state_V) + + @cute.jit + def relay_inner( + self, + pipeline_cpasync: pipeline.PipelineAsync, + pipeline_mma: pipeline.PipelineAsyncUmma, + consumer_state: pipeline.PipelineState, + producer_state: pipeline.PipelineState, + ): + pipeline_cpasync.consumer_wait(consumer_state) + with cute.arch.elect_one(): + pipeline_mma.producer_commit(producer_state) + consumer_state.advance() + producer_state.advance() + return consumer_state, producer_state + + @cute.jit + def load_cpasync( + self, + mIndexTopk: cute.Tensor, + mK: Optional[cute.Tensor], + mV: cute.Tensor, + mVt: cute.Tensor, + sK: Optional[cute.Tensor], + sV: cute.Tensor, + sVt: cute.Tensor, + sBitmask: Optional[cute.Tensor], + pipeline_K: Optional[pipeline.PipelineAsyncUmma], + pipeline_V: pipeline.PipelineAsyncUmma, + pipeline_K_cpasync: Optional[pipeline.PipelineAsync], + pipeline_V_cpasync: pipeline.PipelineAsync, + pipeline_bitmask: pipeline.PipelineAsync, + sO_empty_mbar_ptr: Optional[cute.Pointer], + topk_length_dynamic: Optional[Int32], + block_info: BlockInfo, + SeqlenInfoCls: Callable, + tile_scheduler: TileSchedulerProtocol, + mPageTable: Optional[cute.Tensor] = None, + mCuSeqlensQ: Optional[cute.Tensor] = None, + ): + # ==== cpasync load warpgroup ==== + # Description: loads tiles of K, V, V0, V1 from gmem to smem using cpasync + # produces: K, V, V0, V1, bitmask + # consumes: - + + # cpasync load is used for both topk gather and paged KV with page_size != tile_n + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + tidx = cute.arch.thread_idx()[0] % self.num_cpasync_load_threads + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % ( + self.num_cpasync_load_threads // 32 + ) + + # ==== Make pipeline states ==== + # producer: acquire PipelineAsyncUmma <- mma + # producer: commit PipelineAsync -> relay + Producer = pipeline.PipelineUserType.Producer + if const_expr(self.has_qk): + producer_state_K = pipeline.make_pipeline_state(Producer, stages=self.num_stages_K) + producer_state_V = pipeline.make_pipeline_state(Producer, stages=self.num_stages_V) + if const_expr(self.is_topk_gather and not self.disable_bitmask): + producer_state_bitmask = pipeline.make_pipeline_state( + Producer, stages=self.num_stages_bitmask + ) + if const_expr(self.use_tma_O): + producer_phase_O = Int32(1) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + cluster_m_block, head_idx, batch_idx, _ = work_tile.tile_idx + if const_expr(self.use_packed_varlen_sched): + batch_idx = get_batch_from_cu_tensor(cluster_m_block, mCuSeqlensQ) + if const_expr(not self.is_topk_gather): + cluster_m_block -= mCuSeqlensQ[batch_idx] + head_idx_kv = ( + head_idx // self.qhead_per_kvhead if const_expr(not self.pack_gqa) else head_idx + ) + + seqlen = SeqlenInfoCls(batch_idx) + if const_expr(self.is_topk_gather): + n_block_min = 0 + n_block_max = self.topk_length // self.tile_n + # n_block_max = topk_length_dynamic // self.tile_n + else: + n_block_min, n_block_max = block_info.get_n_block_min_max( + seqlen, + cluster_m_block, + ) + num_n_blocks = n_block_max - n_block_min + + if const_expr(self.is_topk_gather): + # ==== Topk gather path ==== + # cluster_m_block == m_idx under MQA 128 assumption + m_idx = cluster_m_block + if const_expr(not seqlen.has_cu_seqlens_q): + mIndexTopk_cur = mIndexTopk[None, m_idx, batch_idx] + else: + offset_q = seqlen.offset_q if const_expr(not self.use_packed_varlen_sched) else 0 + mIndexTopk_cur = mIndexTopk[None, m_idx + offset_q] + + if const_expr(self.is_causal): + m_local_idx = ( + m_idx - seqlen.offset_q if const_expr(self.use_packed_varlen_sched) else m_idx + ) + seqlen_k_limit = m_local_idx + 1 + seqlen.seqlen_k - seqlen.seqlen_q + else: + seqlen_k_limit = seqlen.seqlen_k + cpasync_gather_kv_manager = CpasyncGatherKVManager.create( + mIndexTopk_cur, + cta_rank_in_cluster, + tidx, + warp_idx, + self.topk_length, + seqlen_k_limit, + self.tile_n, + self.hdim, + self.hdimv, + self.num_hdimv_splits, + self.num_cpasync_load_threads, + mV.element_type, + self.cta_group_size, + self.cpasync_barrier, + self.disable_bitmask, + sBitmask, + pipeline_bitmask, + ) + + # (seqlen_k, hdim) or (seqlen_k, hdimv) + if const_expr(self.has_qk): + mK_cur = seqlen.offset_batch_K(mK, batch_idx, dim=3)[None, None, head_idx_kv] + mV_cur = seqlen.offset_batch_K(mV, batch_idx, dim=3)[None, None, head_idx_kv] + # (hdimv, seqlen_k) + if const_expr(not seqlen.has_cu_seqlens_k): + mVt_cur = mVt[None, None, head_idx_kv, batch_idx] + else: + mVt_cur = cute.domain_offset((0, seqlen.offset_k), mVt[None, None, head_idx_kv]) + + hdimv_split_per_cta = self.hdimv // self.num_hdimv_splits // self.cta_group_size + mVt_cur = cute.tiled_divide(mVt_cur, (hdimv_split_per_cta,)) + mVt_cur = cute.logical_divide(mVt_cur, (1, self.cta_group_size, 1)) + mVt_cur = mVt_cur[(0, None), (cta_rank_in_cluster, None), (0, None)] + mVt_cur = cute.group_modes(mVt_cur, 0, 2) # ((hdimv//4, 2), seqlen_k) + + load_K = None + if const_expr(self.has_qk): + load_K = partial( + self.cpasync_gather_load_KV, + cpasync_gather_kv_manager, + pipeline_K, + pipeline_K_cpasync, + sK, + False, + "K", + mK_cur, + ) + load_V = partial( + self.cpasync_gather_load_KV, + cpasync_gather_kv_manager, + pipeline_V, + pipeline_V_cpasync, + sV, + False, + "V", + mV_cur, + ) + load_Vt = partial( + self.cpasync_gather_load_KV, + cpasync_gather_kv_manager, + pipeline_V, + pipeline_V_cpasync, + sVt, + True, + "V", + mVt_cur, + ) + + # process n_blocks in decreasing order + n_block = n_block_max - 1 + + # ==== Prologue ==== + # K, V0, V1 + cpasync_gather_kv_manager.load_index_topk(n_block, transpose=False) + if const_expr(self.has_qk): + producer_state_K = load_K(producer_state_K) + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_V = load_V(producer_state_V, d_offset=split * self.hdimv // 2) + if const_expr(not self.disable_bitmask): + producer_state_bitmask = cpasync_gather_kv_manager.compute_bitmask( + producer_state_bitmask + ) + + if const_expr(self.use_tma_O and self.overlap_sO_sV): + cute.arch.mbarrier_wait(sO_empty_mbar_ptr, phase=producer_phase_O) + producer_phase_O ^= 1 + + # ==== Mainloop ==== + for _ in cutlass.range(num_n_blocks - 1, unroll=2): + # K, V0, V1 + cpasync_gather_kv_manager.load_index_topk(n_block - 1, transpose=False) + if const_expr(self.has_qk): + producer_state_K = load_K(producer_state_K) + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_V = load_V( + producer_state_V, d_offset=split * self.hdimv // 2 + ) + if const_expr(not self.disable_bitmask): + producer_state_bitmask = cpasync_gather_kv_manager.compute_bitmask( + producer_state_bitmask + ) + # Vt0, Vt1 + cpasync_gather_kv_manager.load_index_topk(n_block, transpose=True) + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_V = load_Vt( + producer_state_V, d_offset=split * hdimv_split_per_cta + ) + # advance n_block + n_block -= 1 + + # ==== Epilogue ==== + + # Vt0, Vt1 for n_block = 0 + cpasync_gather_kv_manager.load_index_topk(0, transpose=True) + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_V = load_Vt( + producer_state_V, d_offset=split * hdimv_split_per_cta + ) + else: + # ==== Paged KV cp.async path (page_size != tile_n) ==== + page_size_divmod = FastDivmodDivisor(cute.size(mV.shape[0])) + hdimv_split = self.hdimv // self.num_hdimv_splits + hdimv_split_per_cta = hdimv_split // self.cta_group_size + + # CTA-split Vt: (dv, page_size, h_k, num_pages) -> ((dv/4, 2), page_size, h_k, num_pages) + mVt_cta = cute.tiled_divide(mVt, (hdimv_split_per_cta,)) + mVt_cta = cute.logical_divide(mVt_cta, (1, self.cta_group_size, 1, 1, 1)) + mVt_cta = mVt_cta[ + (0, None), (cta_rank_in_cluster, None), (0, None), (0, None), (0, None) + ] + mVt_cta = cute.group_modes(mVt_cta, 0, 2) + + # PagedKVManager for K (hdim=64): uses "K" mode only + if const_expr(self.has_qk): + paged_kv_K = PagedKVManager.create( + mPageTable, + mK, + mK, + page_size_divmod, + batch_idx, + head_idx_kv, + tidx, + seqlen.seqlen_k, + 0, + self.tile_n, + self.hdim, + self.hdim, + self.num_cpasync_load_threads, + mK.element_type, + arch=100, + ) + # PagedKVManager for V/Vt: "K" mode → V (non-transposed), "V" mode → Vt (transposed) + paged_kv_V = PagedKVManager.create( + mPageTable, + mV, + mVt_cta, + page_size_divmod, + batch_idx, + head_idx_kv, + tidx, + seqlen.seqlen_k, + 0, + self.tile_n, + hdimv_split, + hdimv_split_per_cta, + self.num_cpasync_load_threads, + mV.element_type, + arch=100, + ) + + if const_expr(self.has_qk): + load_K = partial( + self.cpasync_paged_load_KV, + paged_kv_K, + pipeline_K, + pipeline_K_cpasync, + sK, + False, + "K", + cta_rank_in_cluster, + ) + load_V = partial( + self.cpasync_paged_load_KV, + paged_kv_V, + pipeline_V, + pipeline_V_cpasync, + sV, + False, + "K", + cta_rank_in_cluster, + ) + load_Vt = partial( + self.cpasync_paged_load_KV, + paged_kv_V, + pipeline_V, + pipeline_V_cpasync, + sVt, + True, + "V", + cta_rank_in_cluster, + ) + + n_block_first = n_block_max - 1 + n_block = n_block_first + safe_n_block_first = n_block_first if num_n_blocks > 0 else 0 + + # ==== Prologue ==== + if const_expr(self.has_qk): + paged_kv_K.load_page_table(safe_n_block_first) + producer_state_K = load_K(n_block_first, producer_state_K) + paged_kv_V.load_page_table(safe_n_block_first) + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_V = load_V( + n_block_first, producer_state_V, d_offset=split * self.hdimv // 2 + ) + + if const_expr(self.use_tma_O and self.overlap_sO_sV): + cute.arch.mbarrier_wait(sO_empty_mbar_ptr, phase=producer_phase_O) + producer_phase_O ^= 1 + + # ==== Mainloop ==== + for n_block_idx in cutlass.range(num_n_blocks - 1, unroll=2): + n_block = n_block_first - n_block_idx + # K, V0, V1 for next block in descending order + if const_expr(self.has_qk): + paged_kv_K.load_page_table(n_block - 1) + producer_state_K = load_K(n_block - 1, producer_state_K) + paged_kv_V.load_page_table(n_block - 1) + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_V = load_V( + n_block - 1, producer_state_V, d_offset=split * self.hdimv // 2 + ) + # Vt0, Vt1 for current block + paged_kv_V.load_page_table(n_block) + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_V = load_Vt( + n_block, producer_state_V, d_offset=split * hdimv_split_per_cta + ) + + # ==== Epilogue ==== + paged_kv_V.load_page_table(n_block_min) + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_V = load_Vt( + n_block_min, producer_state_V, d_offset=split * hdimv_split_per_cta + ) + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() + + # note: don't use producer tail with pipeline_X_cpasync since we never use its producer_acquire. + if const_expr(self.is_topk_gather and not self.disable_bitmask): + # pipeline_bitmask invokes producer acquire in gather kv manager, + # so we should call its producer tail. + pipeline_bitmask.producer_tail(producer_state_bitmask) + + @cute.jit + def cpasync_gather_load_KV( + self, + cpasync_gather_kv_manager: CpasyncGatherKVManager, + pipeline_mma: pipeline.PipelineAsyncUmma, + pipeline_cpasync: pipeline.PipelineAsync, + sX: cute.Tensor, + transpose: bool, + K_or_V: str, + mX: cute.Tensor, + producer_state: pipeline.PipelineState, + d_offset: int = 0, + ): + stage = producer_state.index + pipeline_mma.producer_acquire(producer_state) + cpasync_gather_kv_manager.load_X( + mX, sX[None, None, None, stage], transpose, K_or_V, d_offset + ) + cute.arch.cp_async_commit_group() + pipeline_cpasync.sync_object_full.arrive_cp_async_mbarrier(stage) + producer_state.advance() + return producer_state + + @cute.jit + def cpasync_paged_load_KV( + self, + paged_kv_manager: PagedKVManager, + pipeline_mma: pipeline.PipelineAsyncUmma, + pipeline_cpasync: pipeline.PipelineAsync, + sX: cute.Tensor, + transpose: bool, + K_or_V: str, + cta_rank_in_cluster: Int32, + n_block: Int32, + producer_state: pipeline.PipelineState, + d_offset: int = 0, + ): + """Load one tile of K or V from paged gmem to smem using cp.async. + + Uses PagedKVManager for page table lookups and pointer computation, + with smem reshaping via cute.composition (same approach as CpasyncGatherKVManager). + + For non-transposed tensors (K, V0, V1): K_or_V="K", transpose=False + For transposed tensors (Vt0, Vt1): K_or_V="V", transpose=True + """ + stage = producer_state.index + pipeline_mma.producer_acquire(producer_state) + + # NOTE: load_page_table() must be called by the caller BEFORE this method. + # Calling it here (through a @cute.jit boundary) causes MLIR SSA verification + # errors because the rmem tensor writes inside load_page_table's dynamic + # cutlass.range loop cross region boundaries. This matches the SM90/SM100 + # pattern where load_page_table is called directly in the loop body. + + # Compute row pointers from cached page table entries + tPrXPtr = paged_kv_manager.compute_X_ptr(K_or_V, d_offset) + + # Reshape smem to flat 2D using composition (matches CpasyncGatherKVManager.load_X) + head_dim = ( + paged_kv_manager.head_dim_v_padded + if const_expr(K_or_V == "V") + else paged_kv_manager.head_dim_padded + ) + cta_tile_n = self.tile_n if const_expr(transpose) else self.tile_n // self.cta_group_size + order = (1, 0) if const_expr(transpose) else (0, 1) + + sX_stage = sX[None, None, None, stage] + sX_nd_layout = cute.make_ordered_layout((cta_tile_n, head_dim), order=order) + sX_nd = cute.composition(sX_stage, sX_nd_layout) + + cX = cute.make_identity_tensor((cta_tile_n, head_dim)) + tXsX = paged_kv_manager.gmem_thr_copy_KV.partition_D(sX_nd) + tXcX = paged_kv_manager.gmem_thr_copy_KV.partition_S(cX) + tXc0X = paged_kv_manager.gmem_thr_copy_KV.get_slice(0).partition_S(cX) + + base_offset = n_block * self.tile_n + if const_expr(not transpose): + base_offset += cta_tile_n * cta_rank_in_cluster + seqlenk_row_limit = ( + paged_kv_manager.seqlen_k - base_offset - tXcX[0][0] if n_block >= 0 else 0 + ) + + if const_expr(not transpose): + offset = cta_rank_in_cluster * ( + paged_kv_manager.gmem_threads_per_row // self.cta_group_size + ) + else: + offset = 0 + + for m in cutlass.range_constexpr(cute.size(tXsX, mode=[1])): + row_valid = tXc0X[0, m, 0][0] < seqlenk_row_limit + should_load = cute.make_fragment_like(tXsX[(0, None), m, 0], cute.Boolean) + should_load.fill(row_valid) + + x_ptr_i64 = fa_utils.shuffle_sync( + tPrXPtr[m // paged_kv_manager.gmem_threads_per_row], + (m + offset) % paged_kv_manager.gmem_threads_per_row, + width=paged_kv_manager.gmem_threads_per_row, + ) + x_gmem_ptr = cute.make_ptr( + paged_kv_manager.mK_paged.element_type, + x_ptr_i64, + cute.AddressSpace.gmem, + assumed_align=16, + ) + mX_cur = cute.make_tensor(x_gmem_ptr, cute.make_layout((head_dim,))) + mX_cur_copy = cute.tiled_divide(mX_cur, (paged_kv_manager.async_copy_elems,)) + + for k in cutlass.range_constexpr(cute.size(tXsX, mode=[2])): + ki = tXcX[0, 0, k][1] // paged_kv_manager.async_copy_elems + mX_cur_copy_ki = mX_cur_copy[None, ki] + tXsX_k = tXsX[None, m, k] + mX_cur_copy_ki = cute.make_tensor(mX_cur_copy_ki.iterator, tXsX_k.layout) + cute.copy( + paged_kv_manager.gmem_tiled_copy_KV, + mX_cur_copy_ki, + tXsX_k, + pred=should_load, + ) + + cute.arch.cp_async_commit_group() + pipeline_cpasync.sync_object_full.arrive_cp_async_mbarrier(stage) + producer_state.advance() + return producer_state + + @cute.jit + def load( + self, + mQ: Optional[cute.Tensor], + mK: Optional[cute.Tensor], + mQv: cute.Tensor, + mV: cute.Tensor, + mVt: cute.Tensor, + sQ: Optional[cute.Tensor], + sK: Optional[cute.Tensor], + sQv: cute.Tensor, + sV: cute.Tensor, + sVt: cute.Tensor, + tma_atom_Q: Optional[cute.CopyAtom], + tma_atom_K: Optional[cute.CopyAtom], + tma_atom_Qv: cute.CopyAtom, + tma_atom_V: cute.CopyAtom, + tma_atom_Vt: cute.CopyAtom, + pipeline_Q: Optional[pipeline.PipelineAsync], + pipeline_K: Optional[pipeline.PipelineAsync], + pipeline_Qv: pipeline.PipelineAsync, + pipeline_V: pipeline.PipelineAsync, + sO_empty_mbar_ptr: Optional[cute.Pointer], + thr_mma_QK: cute.ThrMma, + thr_mma_QvV: cute.ThrMma, + thr_mma_PVt: cute.ThrMma, + topk_length_dynamic: Optional[Int32], + block_info: BlockInfo, + SeqlenInfoCls: Callable, + tile_scheduler: TileSchedulerProtocol, + mPageTable: Optional[cute.Tensor] = None, + mCuSeqlensQ: Optional[cute.Tensor] = None, + ): + # ==== Load warp ==== + # Description: loads tiles of Q, Qv, K, V, V0, V1 from gmem to smem using TMA + # produces: Q, Qv, K, V, V0, V1 + # consumes: - + + # ==== Make pipeline states ==== + Producer = pipeline.PipelineUserType.Producer + if const_expr(self.has_qk): + producer_state_Q = pipeline.make_pipeline_state(Producer, stages=self.num_stages_Q) + producer_state_Qv = pipeline.make_pipeline_state(Producer, stages=self.num_stages_Qv) + if const_expr(self.use_tma_KV): + if const_expr(self.has_qk): + producer_state_K = pipeline.make_pipeline_state(Producer, stages=self.num_stages_K) + producer_state_V = pipeline.make_pipeline_state(Producer, stages=self.num_stages_V) + if const_expr(self.use_tma_O): + producer_phase_O = Int32(1) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + cluster_m_block, head_idx, batch_idx, _ = work_tile.tile_idx + if const_expr(self.use_packed_varlen_sched): + batch_idx = get_batch_from_cu_tensor(cluster_m_block, mCuSeqlensQ) + cluster_m_block -= mCuSeqlensQ[batch_idx] + head_idx_kv = ( + head_idx // self.qhead_per_kvhead if const_expr(not self.pack_gqa) else head_idx + ) + + seqlen = SeqlenInfoCls(batch_idx) + if const_expr(self.is_topk_gather): + n_block_min = 0 + n_block_max = self.topk_length // self.tile_n + # n_block_max = topk_length_dynamic // self.tile_n + else: + n_block_min, n_block_max = block_info.get_n_block_min_max( + seqlen, + cluster_m_block, + ) + num_n_blocks = n_block_max - n_block_min + even_n_blocks = num_n_blocks % 2 == 0 and num_n_blocks > 0 + num_n_block_groups = cute.ceil_div(num_n_blocks, self.num_stages_S) + + # ==== Partition GMEM tensors ==== + # (seqlen_q, hdim or hdimv//2) + if const_expr(self.has_qk): + mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx] + # (mma_tile_m, hdim or hdimv//2) + gQ = cute.local_tile( + mQ_cur, + (self.mma_tiler_QK[0], self.mma_tiler_QK[2]), + (cluster_m_block, 0), + ) + tSgQ = thr_mma_QK.partition_A(gQ) + tQsQ, tQgQ = cpasync.tma_partition( + atom=tma_atom_Q, + cta_coord=0, + cta_layout=cute.make_layout(1), + smem_tensor=cute.group_modes(sQ, 0, 3), + gmem_tensor=cute.group_modes(tSgQ, 0, 3), + ) + mQv_cur = seqlen.offset_batch_Q(mQv, batch_idx, dim=3)[None, None, head_idx] + gQv = cute.local_tile( + mQv_cur, + (self.mma_tiler_QvV[0], self.mma_tiler_QvV[2]), + (cluster_m_block, None), + ) + tSgQv = thr_mma_QvV.partition_A(gQv) + tQvsQv, tQvgQv = cpasync.tma_partition( + atom=tma_atom_Qv, + cta_coord=0, + cta_layout=cute.make_layout(1), + smem_tensor=cute.group_modes(sQv, 0, 3), + gmem_tensor=cute.group_modes(tSgQv, 0, 3), + ) + + if const_expr(self.use_tma_KV): + if const_expr(mPageTable is None): + mPageTable_cur = None + # Non-paged: select batch, tile over seqlen_k + if const_expr(self.has_qk): + # (seqlen_k, hdim) + mK_cur = seqlen.offset_batch_K(mK, batch_idx, dim=3)[ + None, None, head_idx_kv + ] + # (tile_n, hdim, num_n_blocks) + gK = cute.local_tile( + mK_cur, + (self.mma_tiler_QK[1], self.mma_tiler_QK[2]), + (None, 0), + ) + # (seqlen_k, hdimv) + mV_cur = seqlen.offset_batch_K(mV, batch_idx, dim=3)[None, None, head_idx_kv] + # (hdimv, seqlen_k) + if const_expr(not seqlen.has_cu_seqlens_k): + mVt_cur = mVt[None, None, head_idx_kv, batch_idx] + else: + mVt_cur = cute.domain_offset( + (0, seqlen.offset_k), mVt[None, None, head_idx_kv] + ) + # (tile_n, hdimv//4, num_n_blocks, num_d_blocks=4) + gV = cute.local_tile( + mV_cur, + (self.mma_tiler_QvV[1], self.mma_tiler_QvV[2]), + (None, None), + ) + # (tile_n, hdimv//4, num_d_blocks=4, num_n_blocks) + gV = cute.make_tensor(gV.iterator, cute.select(gV.layout, mode=[0, 1, 3, 2])) + # (hdimv//4, tile_n, num_d_blocks=4, num_n_blocks) + gVt = cute.local_tile( + mVt_cur, + (self.mma_tiler_PVt[1], self.mma_tiler_PVt[2]), + (None, None), + ) + else: + mPageTable_cur = mPageTable[batch_idx, None] + # Paged KV: keep pages dim, index by page_idx at load time + # TMA path assumes page_size == tile_n + if const_expr(self.has_qk): + # (page_size, hdim, num_pages) + mK_cur = mK[None, None, head_idx_kv, None] + # (tile_n, hdim, num_pages) + gK = cute.local_tile( + mK_cur, + (self.mma_tiler_QK[1], self.mma_tiler_QK[2]), + (0, 0, None), + ) + # (page_size, hdimv, num_pages) + mV_cur = mV[None, None, head_idx_kv, None] + # (hdimv, page_size, num_pages) + mVt_cur = mVt[None, None, head_idx_kv, None] + # (tile_n, hdimv//4, num_d_blocks=4, num_pages) + gV = cute.local_tile( + mV_cur, + (self.mma_tiler_QvV[1], self.mma_tiler_QvV[2]), + (0, None, None), + ) + # (hdimv//4, tile_n, num_d_blocks=4, num_pages) + gVt = cute.local_tile( + mVt_cur, + (self.mma_tiler_PVt[1], self.mma_tiler_PVt[2]), + (None, 0, None), + ) + + if const_expr(self.has_qk): + tSgK = thr_mma_QK.partition_B(gK) + tKsK, tKgK = cpasync.tma_partition( + atom=tma_atom_K, + cta_coord=0, + cta_layout=cute.make_layout(1), + smem_tensor=cute.group_modes(sK, 0, 3), + gmem_tensor=cute.group_modes(tSgK, 0, 3), + ) + + tSgV = thr_mma_QvV.partition_B(gV) + tOgVt = thr_mma_PVt.partition_B(gVt) + tVsV, tVgV = cpasync.tma_partition( + atom=tma_atom_V, + cta_coord=0, + cta_layout=cute.make_layout(1), + smem_tensor=cute.group_modes(sV, 0, 3), + gmem_tensor=cute.group_modes(tSgV, 0, 3), + ) + tVtsVt, tVtgVt = cpasync.tma_partition( + atom=tma_atom_Vt, + cta_coord=0, + cta_layout=cute.make_layout(1), + smem_tensor=cute.group_modes(sVt, 0, 3), + gmem_tensor=cute.group_modes(tOgVt, 0, 3), + ) + + if const_expr(self.has_qk): + load_Q = partial(self.load_inner, tma_atom_Q, tQgQ, tQsQ, pipeline_Q) + load_Qv = partial(self.load_inner, tma_atom_Qv, tQvgQv, tQvsQv, pipeline_Qv) + + if const_expr(self.use_tma_KV): + if const_expr(self.has_qk): + load_K = partial(self.load_inner, tma_atom_K, tKgK, tKsK, pipeline_K) + load_V = partial(self.load_inner, tma_atom_V, tVgV, tVsV, pipeline_V) + load_Vt = partial(self.load_inner, tma_atom_Vt, tVtgVt, tVtsVt, pipeline_V) + + # ==== Load stationary operands ==== + + # copy Q, Qvi gmem -> smem + if const_expr(self.has_qk): + producer_state_Q = load_Q(producer_state_Q) + for dv_split in cutlass.range_constexpr(2): + producer_state_Qv = load_Qv(producer_state_Qv, block=dv_split) + + if const_expr(self.use_tma_KV): + # ==== Prologue ==== + n_block_first = n_block_max - 1 if n_block_max > 0 else 0 + block = self._get_block_idx(n_block_first, mPageTable_cur) + # copy K gmem -> smem + if const_expr(self.has_qk): + producer_state_K = load_K(producer_state_K, block=block) + # copy Vi gmem -> smem + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_V = load_V(producer_state_V, block=block, split=split) + + if const_expr(self.use_tma_O and self.overlap_sO_sV): + cute.arch.mbarrier_wait(sO_empty_mbar_ptr, phase=producer_phase_O) + producer_phase_O ^= 1 + + # ==== Main loop ==== + for n_block_group in cutlass.range(num_n_block_groups - 1, unroll=1): + for stage in cutlass.range_constexpr(self.num_stages_S): + n_block = n_block_max - 1 - n_block_group * self.num_stages_S - stage + block_next = self._get_block_idx(n_block - 1, mPageTable_cur) + block = self._get_block_idx(n_block, mPageTable_cur) + if const_expr(self.has_qk): + # copy K gmem -> smem + producer_state_K = load_K(producer_state_K, block=block_next) + # copy Vi gmem -> smem + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_V = load_V( + producer_state_V, block=block_next, split=split + ) + # copy Vti gmem -> smem + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_V = load_Vt(producer_state_V, block=block, split=split) + + # ==== Epilogue ==== + num_final_n_blocks = self.num_stages_S if even_n_blocks else self.num_stages_S - 1 + for stage in cutlass.range(num_final_n_blocks, unroll_full=True): + n_block = num_final_n_blocks - 1 - stage + block = self._get_block_idx(n_block, mPageTable_cur) + if n_block > 0: + block_next = self._get_block_idx(n_block - 1, mPageTable_cur) + if const_expr(self.has_qk): + # copy K gmem -> smem + producer_state_K = load_K(producer_state_K, block=block_next) + # copy Vi gmem -> smem + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_V = load_V( + producer_state_V, block=block_next, split=split + ) + # copy Vti gmem -> smem + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_V = load_Vt(producer_state_V, block=block, split=split) + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() + + if const_expr(self.has_qk): + pipeline_Q.producer_tail(producer_state_Q) + pipeline_Qv.producer_tail(producer_state_Qv) + if const_expr(self.use_tma_KV): + if const_expr(self.has_qk): + pipeline_K.producer_tail(producer_state_K) + pipeline_V.producer_tail(producer_state_V) + + @cute.jit + def _get_block_idx( + self, + n_block, + mPageTable_cur: Optional[cute.Tensor], + ): + if const_expr(mPageTable_cur is not None): + return mPageTable_cur[n_block] + else: + return n_block + + @cute.jit + def load_inner( + self, + tma_atom: cute.CopyAtom, + tXgX: cute.Tensor, + tXsX: cute.Tensor, + load_pipeline: pipeline.PipelineAsync, + producer_state: pipeline.PipelineState, + block: Optional[Int32] = None, + split: Optional[Int32] = None, + ): + if const_expr(split is not None): + tXgX = tXgX[(None, split, None)] + if const_expr(block is not None): + tXgX = tXgX[(None, block)] + if const_expr(cute.rank(tXsX) != 1): + assert cute.rank(tXsX) == 2, f"wrong rank for tXsX, got {cute.rank(tXsX)}" + stage = producer_state.index + tXsX = tXsX[(None, stage)] + load_pipeline.producer_acquire(producer_state) + tma_bar_ptr = load_pipeline.producer_get_barrier(producer_state) + cute.copy(tma_atom, tXgX, tXsX, tma_bar_ptr=tma_bar_ptr) + producer_state.advance() + return producer_state + + @cute.jit + def mma( + self, + sQ: Optional[cute.Tensor], + sK: Optional[cute.Tensor], + sQv: cute.Tensor, + sV: cute.Tensor, + sVt: cute.Tensor, + sP: cute.Tensor, + tStS: cute.Tensor, + tOtO0: cute.Tensor, + tOtO1: cute.Tensor, + tiled_mma_QK: cute.TiledMma, + tiled_mma_QvV: cute.TiledMma, + tiled_mma_PVt: cute.TiledMma, + pipeline_Q: Optional[pipeline.PipelineAsync], + pipeline_K: Optional[pipeline.PipelineAsync], + pipeline_Qv: pipeline.PipelineAsync, + pipeline_V: pipeline.PipelineAsync, + pipeline_S: pipeline.PipelineAsync, + pipeline_P: pipeline.PipelineAsync, + pipeline_O0: pipeline.PipelineAsync, + pipeline_O1: pipeline.PipelineAsync, + sO_empty_mbar_ptr: Optional[cute.Pointer], + is_leader_cta: Boolean, + topk_length_dynamic: Optional[Int32], + block_info: BlockInfo, + SeqlenInfoCls: Callable, + tile_scheduler: TileSchedulerProtocol, + mCuSeqlensQ: Optional[cute.Tensor] = None, + ): + # ==== mma warp ==== + # Description: Computes Q @ K^T, Qv @ V^T, and P @ V + # Produces: S, O + # Consumes: Q, K, Qv, V, P + + pipelines_O = [pipeline_O0, pipeline_O1] + tOtOs = [tOtO0, tOtO1] + + use_ptx_gemm_QK = True + use_ptx_gemm_QvV = True + use_ptx_gemm_PVt = True + + # Operands for S = Q @ K^T + if const_expr(self.has_qk): + tSrQ = tiled_mma_QK.make_fragment_A(sQ) + tSrK = tiled_mma_QK.make_fragment_B(sK) + + # Operands for S += Qv @ V^T + tSrQv = tiled_mma_QvV.make_fragment_A(sQv) + tSrV = tiled_mma_QvV.make_fragment_B(sV) + + # Operands for Oi = P @ Vi + tOrP = tiled_mma_PVt.make_fragment_A(sP) + tOrVt = tiled_mma_PVt.make_fragment_B(sVt) + + # GEMM functions + if const_expr(self.has_qk): + if const_expr(use_ptx_gemm_QK): + gemm_QK = [ + partial( + fa_sm100_utils.gemm_ptx_partial, + tiled_mma_QK.op, + self.tmem_offset_S[stage], + zero_init=True, + cta_group=self.cta_group_size, + ) + for stage in range(self.num_stages_S) + ] + else: + gemm_QK = [ + partial( + fa_sm100_utils.gemm, + tiled_mma_QK, + tStS[None, None, None, stage], + zero_init=True, + ) + for stage in range(self.num_stages_S) + ] + if const_expr(use_ptx_gemm_QvV): + gemm_QvV = [ + partial( + fa_sm100_utils.gemm_ptx_partial, + tiled_mma_QvV.op, + self.tmem_offset_S[stage], + cta_group=self.cta_group_size, + ) + for stage in range(self.num_stages_S) + ] + else: + gemm_QvV = [ + partial( + fa_sm100_utils.gemm, + tiled_mma_QvV, + tStS[None, None, None, stage], + ) + for stage in range(self.num_stages_S) + ] + + if const_expr(use_ptx_gemm_PVt): + gemm_PVt = [ + partial( + fa_sm100_utils.gemm_ptx_partial, + tiled_mma_PVt.op, + self.tmem_offsets_O[split], + cta_group=self.cta_group_size, + ) + for split in range(self.num_hdimv_splits) + ] + else: + gemm_PVt = [ + partial( + fa_sm100_utils.gemm, + tiled_mma_PVt, + tOtOs[split], + ) + for split in range(self.num_hdimv_splits) + ] + + Consumer, Producer = pipeline.PipelineUserType.Consumer, pipeline.PipelineUserType.Producer + if const_expr(self.has_qk): + consumer_state_Q = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_Q) + consumer_state_K = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_K) + consumer_state_Qv = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_Qv) + consumer_state_V = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_V) + producer_state_S = pipeline.make_pipeline_state(Producer, stages=self.num_stages_S) + consumer_state_P = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_P) + producer_state_O0 = pipeline.make_pipeline_state(Producer, stages=self.num_stages_Oi) + producer_state_O1 = pipeline.make_pipeline_state(Producer, stages=self.num_stages_Oi) + + mma_fn = self.mma_inner + if const_expr(self.has_qk): + mma_QK = partial( + mma_fn, gemm_QK, pipeline_K, tSrQ, sQ, tSrK, sK, use_ptx=use_ptx_gemm_QK + ) + mma_QvV = partial( + mma_fn, gemm_QvV, pipeline_V, tSrQv, sQv, tSrV, sV, use_ptx=use_ptx_gemm_QvV + ) + mma_PVt = partial( + mma_fn, gemm_PVt, pipeline_V, tOrP, sP, tOrVt, sVt, use_ptx=use_ptx_gemm_PVt + ) + + work_tile = tile_scheduler.initial_work_tile_info() + O_should_accumulate = False + while work_tile.is_valid_tile: + cluster_m_block, head_idx, batch_idx, _ = work_tile.tile_idx + if const_expr(self.use_packed_varlen_sched): + batch_idx = get_batch_from_cu_tensor(cluster_m_block, mCuSeqlensQ) + cluster_m_block -= mCuSeqlensQ[batch_idx] + + seqlen = SeqlenInfoCls(batch_idx) + if const_expr(self.is_topk_gather): + n_block_min = 0 + # n_block_max = self.topk_length // self.tile_n + n_block_max = topk_length_dynamic // self.tile_n + else: + n_block_min, n_block_max = block_info.get_n_block_min_max( + seqlen, + cluster_m_block, + ) + num_n_blocks = n_block_max - n_block_min + even_n_blocks = num_n_blocks % 2 == 0 and num_n_blocks > 0 + num_n_block_groups = cute.ceil_div(num_n_blocks, self.num_stages_S) + + if is_leader_cta: + if const_expr(self.has_qk): + pipeline_Q.consumer_wait(consumer_state_Q) + + consumer_wait_state_Qv = consumer_state_Qv.clone() + for _ in cutlass.range_constexpr(self.num_hdimv_splits): + pipeline_Qv.consumer_wait(consumer_wait_state_Qv) + consumer_wait_state_Qv.advance() + + producer_states_O = [producer_state_O0, producer_state_O1] + + # ==== Prologue ==== + pipeline_S.producer_acquire(producer_state_S) + if const_expr(self.has_qk): + # S = Q @ K^T + consumer_state_K = mma_QK(consumer_state_K, acc_stage=0) + # S += Qvi @ Vi^T + for split in cutlass.range_constexpr(self.num_hdimv_splits): + consumer_state_V = mma_QvV( + consumer_state_V, + acc_stage=0, + a_stage=split, + zero_init=split == 0 and not self.has_qk, + ) + pipeline_S.producer_commit(producer_state_S) + producer_state_S.advance() + + # ==== Mainloop ==== + for _ in cutlass.range(num_n_block_groups - 1, unroll=1): + for stage in cutlass.range_constexpr(self.num_stages_S): + next_stage = const_expr((stage + 1) % self.num_stages_S) + pipeline_S.producer_acquire(producer_state_S) + if const_expr(self.has_qk): + # S = Q @ K^T + consumer_state_K = mma_QK(consumer_state_K, acc_stage=next_stage) + # S += Qvi @ Vi^T + for split in cutlass.range_constexpr(self.num_hdimv_splits): + consumer_state_V = mma_QvV( + consumer_state_V, + acc_stage=next_stage, + a_stage=split, + zero_init=split == 0 and not self.has_qk, + ) + pipeline_S.producer_commit(producer_state_S) + producer_state_S.advance() + # Oi += P @ Vi + pipeline_P.consumer_wait(consumer_state_P) + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_Oi = producer_states_O[split] + pipelines_O[split].producer_acquire(producer_state_Oi) + consumer_state_V = mma_PVt( + consumer_state_V, + acc_stage=split, + a_stage=consumer_state_P.index, + zero_init=not O_should_accumulate, + ) + pipelines_O[split].producer_commit(producer_state_Oi) + producer_state_Oi.advance() + producer_states_O[split] = producer_state_Oi + pipeline_P.consumer_release(consumer_state_P) + consumer_state_P.advance() + O_should_accumulate = True + + # ==== Epilogue ==== + num_final_n_blocks = self.num_stages_S if even_n_blocks else self.num_stages_S - 1 + for stage in cutlass.range_constexpr(self.num_stages_S): + n_block = num_final_n_blocks - 1 - stage + if const_expr(stage == 0): + if n_block > 0: + pipeline_S.producer_acquire(producer_state_S) + if const_expr(self.has_qk): + # S = Q @ K^T + consumer_state_K = mma_QK(consumer_state_K, acc_stage=stage + 1) + # S += Qvi @ Vi^T + for split in cutlass.range_constexpr(self.num_hdimv_splits): + consumer_state_V = mma_QvV( + consumer_state_V, + acc_stage=stage + 1, + a_stage=split, + zero_init=split == 0 and not self.has_qk, + ) + pipeline_S.producer_commit(producer_state_S) + producer_state_S.advance() + if n_block >= 0: + # Oi += P @ Vi + pipeline_P.consumer_wait(consumer_state_P) + for split in cutlass.range_constexpr(self.num_hdimv_splits): + producer_state_Oi = producer_states_O[split] + pipelines_O[split].producer_acquire(producer_state_Oi) + consumer_state_V = mma_PVt( + consumer_state_V, + acc_stage=split, + a_stage=consumer_state_P.index, + zero_init=not O_should_accumulate, + ) + pipelines_O[split].producer_commit(producer_state_Oi) + producer_state_Oi.advance() + producer_states_O[split] = producer_state_Oi + pipeline_P.consumer_release(consumer_state_P) + consumer_state_P.advance() + O_should_accumulate = True + + producer_state_O0, producer_state_O1 = producer_states_O + + if const_expr(self.has_qk): + pipeline_Q.consumer_release(consumer_state_Q) + consumer_state_Q.advance() + + # if we overlap sOi with sQvi for tma store, need to acquire signal + if const_expr(self.use_tma_O and not self.overlap_sO_sV): + pipeline_O0.producer_tail(producer_state_O0.clone()) + pipeline_O1.producer_tail(producer_state_O1.clone()) + + for _ in cutlass.range_constexpr(self.num_hdimv_splits): + pipeline_Qv.consumer_release(consumer_state_Qv) + consumer_state_Qv.advance() + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() + O_should_accumulate = False + + pipeline_S.producer_tail(producer_state_S) + pipeline_O0.producer_tail(producer_state_O0) + pipeline_O1.producer_tail(producer_state_O1) + + @cute.jit + def mma_inner( + self, + gemm, + load_pipeline, + tCrA, + sA, + tCrB, + sB, + consumer_state: pipeline.PipelineState, + acc_stage: Optional[Int32] = None, + a_stage: Int32 = 0, + zero_init: Optional[bool] = None, + use_ptx: bool = True, + ): + if const_expr(acc_stage is not None): + gemm = gemm[acc_stage] + + tCrA_cur = tCrA[None, None, None, a_stage] + sA_cur = sA[None, None, None, a_stage] + b_stage = consumer_state.index + tCrB_cur = tCrB[None, None, None, b_stage] + sB_cur = sB[None, None, None, b_stage] + + load_pipeline.consumer_wait(consumer_state) + + kwargs = dict(tCrA=tCrA_cur, tCrB=tCrB_cur) + if const_expr(use_ptx): + kwargs |= dict(sA=sA_cur, sB=sB_cur) + if const_expr(zero_init is not None): + kwargs["zero_init"] = zero_init + gemm(**kwargs) + + load_pipeline.consumer_release(consumer_state) + consumer_state.advance() + return consumer_state + + @cute.jit + def softmax_loop( + self, + softmax_scale: Float32, + softmax_scale_log2: Float32, + mLSE: Optional[cute.Tensor], + mRowMax: Optional[cute.Tensor], + sRowMax: cute.Tensor, + sRowSum: cute.Tensor, + sScale: cute.Tensor, + sBitmask: Optional[cute.Tensor], + sP: cute.Tensor, + tStS: cute.Tensor, + thr_mma_S: cute.ThrMma, + pipeline_S: pipeline.PipelineAsync, + pipeline_P: pipeline.PipelineAsync, + pipeline_sm_stats: pipeline.PipelineAsync, + pipeline_bitmask: Optional[pipeline.PipelineAsync], + sO_empty_mbar_ptr: Optional[cute.Pointer], + AttentionMaskCls: Callable, + topk_length_dynamic: Optional[Int32], + block_info: BlockInfo, + SeqlenInfoCls: Callable, + tile_scheduler: TileSchedulerProtocol, + tma_atom_P: Optional[cute.CopyAtom] = None, + mP: Optional[cute.Tensor] = None, + sP_out: Optional[cute.Tensor] = None, + mCuSeqlensQ: Optional[cute.Tensor] = None, + ): + # ==== softmax warpgroup ==== + # Description: computes softmax on S and writes the result to P + # Produces: P, softmax stats + # Consumes: S, bitmask (for topk sparsity) + + tidx = cute.arch.thread_idx()[0] % self.num_softmax_threads + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % ( + self.num_softmax_threads // 32 + ) + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + + tSAcc = tStS[(None, None), 0, 0, 0] + tSAcc_staged = [tStS[(None, None), 0, 0, stage] for stage in range(self.num_stages_S)] + + cS = cute.make_identity_tensor(self.mma_tiler_QK[:2]) # (128, 128) + tScS = thr_mma_S.partition_C(cS)[(None, None), 0, 0] # (64, 128) + + # S tmem -> rmem copy objects + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), + self.dtype_acc, + ) + tmem_load_tiled = tcgen05.make_tmem_copy(tmem_load_atom, tSAcc) + tmem_load_thr = tmem_load_tiled.get_slice(tidx) + # S tmem -> rmem copy operands + tStS_t2r = tmem_load_thr.partition_S(tSAcc) # (((32, 32), 1), 1, 2) + tStS_t2r_staged = [ + tmem_load_thr.partition_S(tSAcc_staged[stage]) for stage in range(self.num_stages_S) + ] + tScS_t2r = tmem_load_thr.partition_D(tScS) + tSrS_t2r = cute.make_rmem_tensor(tScS_t2r.shape, self.dtype_acc) + + # P rmem -> smem copy objects + universal_copy_bits = 128 + smem_store_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.dtype_P, + num_bits_per_copy=universal_copy_bits, + ) + smem_store_tiled = cute.make_tiled_copy_D(smem_store_atom, tmem_load_tiled) + smem_store_thr = smem_store_tiled.get_slice(tidx) + # P rmem -> smem copy operands + sP_mnp_layout = cute.make_ordered_layout( + self.tile_P + (self.num_stages_P,), order=(0, 1, 2) + ) + sP_mnp = cute.composition(sP, sP_mnp_layout) + sP_smem_view = smem_store_thr.partition_D(sP_mnp) + + Consumer, Producer = pipeline.PipelineUserType.Consumer, pipeline.PipelineUserType.Producer + consumer_state_S = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_S) + producer_state_P = pipeline.make_pipeline_state(Producer, stages=self.num_stages_P) + producer_state_sm_stats = pipeline.make_pipeline_state(Producer, stages=self.num_stages_sm_stats) + consumer_state_bitmask = None + if const_expr(self.is_topk_gather and not self.disable_bitmask): + consumer_state_bitmask = pipeline.make_pipeline_state( + Consumer, stages=self.num_stages_bitmask + ) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + cluster_m_block, head_idx, batch_idx, _ = work_tile.tile_idx + if const_expr(self.use_packed_varlen_sched): + batch_idx = get_batch_from_cu_tensor(cluster_m_block, mCuSeqlensQ) + cluster_m_block -= mCuSeqlensQ[batch_idx] + cta_m_block = cluster_m_block * self.cta_group_size + cta_rank_in_cluster + seqlen = SeqlenInfoCls(batch_idx) + if const_expr(self.is_topk_gather): + n_block_min = 0 + n_block_max = self.topk_length // self.tile_n + # n_block_max = topk_length_dynamic // self.tile_n + else: + n_block_min, n_block_max = block_info.get_n_block_min_max( + seqlen, + cluster_m_block, + ) + num_n_blocks = n_block_max - n_block_min + even_n_blocks = num_n_blocks % 2 == 0 and num_n_blocks > 0 + num_n_block_groups = cute.ceil_div(num_n_blocks, self.num_stages_S) + + gRowMax = None + if const_expr(mRowMax is not None): + # (seqlen_q, {seqlen_k_rounded, topk} / tile_n) + if const_expr(not seqlen.has_cu_seqlens_q): + mRowMax_cur = mRowMax[None, None, head_idx, batch_idx] + else: + q_offset = ( + seqlen.offset_q if const_expr(not self.pack_gqa) else (0, seqlen.offset_q) + ) + mRowMax_cur = cute.domain_offset((q_offset, 0), mRowMax[None, None, head_idx]) + # (cta_tile_m, {seqlen_k_rounded, topk} / tile_n) + gRowMax = cute.local_tile(mRowMax_cur, (self.cta_tile_m,), (cta_m_block, None)) + + store_P = None + if const_expr(self.store_P): + # (seqlen_q, seqlen_k) + mP_cur = seqlen.offset_batch_Q(mP, batch_idx, dim=3, ragged=self.ragged_tma_O)[ + None, None, head_idx + ] + # (cta_tile_m, tile_n, num_n_blocks) + gP = cute.local_tile(mP_cur, self.tile_P, (cta_m_block, None)) + store_P, tPsP, tPgP = copy_utils.tma_get_copy_fn( + tma_atom_P, + 0, + cute.make_layout(1), + sP_out, + gP, + ) + + mask = AttentionMaskCls(seqlen) + mask_fn = partial( + mask.apply_mask_sm100, + m_block=cluster_m_block, + thr_mma=thr_mma_S, + thr_tmem_load=tmem_load_thr, + mask_causal=self.is_causal, + mask_local=self.is_local, + batch_idx=batch_idx, + head_idx=head_idx, + r2p=False, # TODO: fix r2p for 2cta + ) + disable_mask = self.disable_bitmask and self.is_topk_gather + + softmax = SoftmaxSm100.create( + softmax_scale_log2, + rescale_threshold=8.0 if const_expr(self.dtype_Q.width == 16) else 0.0, + softmax_scale=softmax_scale, + ) + softmax.reset() + + softmax_step_fn = partial( + self.softmax_step, + softmax, + sRowMax, + sScale, + sBitmask, + tStS_t2r_staged, + tSrS_t2r, + sP_smem_view, + tmem_load_thr, + smem_store_thr, + pipeline_S, + pipeline_P, + pipeline_sm_stats, + pipeline_bitmask, + tidx, + warp_idx, + store_P=store_P, + gRowMax=gRowMax, + ) + + ### first iteration ### + n_block = n_block_max - 1 + ( + consumer_state_S, + producer_state_P, + producer_state_sm_stats, + consumer_state_bitmask, + ) = softmax_step_fn( + consumer_state_S, + producer_state_P, + producer_state_sm_stats, + consumer_state_bitmask, + 0, + n_block, + mask_fn=partial(mask_fn, mask_seqlen=True) + if not const_expr(disable_mask) + else None, + is_first=True, + ) + n_block -= 1 + + ### Separate iterations with causal masking + # note: For square mma tile, can mask at most 1 n_block_group + if const_expr((self.is_causal or self.is_local) and not self.is_topk_gather): + n_block_min_causal_local_mask = block_info.get_n_block_min_causal_local_mask( + seqlen, cluster_m_block, n_block_min + ) + num_masked_n_blocks = n_block_max - 1 - n_block_min_causal_local_mask + num_masked_n_block_groups = min( + num_n_block_groups - 1, cute.ceil_div(num_masked_n_blocks, self.num_stages_S) + ) + num_n_block_groups -= num_masked_n_block_groups + for _ in cutlass.range(num_masked_n_block_groups, unroll=1): + for stage in cutlass.range_constexpr(self.num_stages_S): + ( + consumer_state_S, + producer_state_P, + producer_state_sm_stats, + consumer_state_bitmask, + ) = softmax_step_fn( + consumer_state_S, + producer_state_P, + producer_state_sm_stats, + consumer_state_bitmask, + 1 - stage, + n_block, + mask_fn=partial(mask_fn, mask_seqlen=False), + ) + n_block -= 1 + + ### Mainloop ### + for n_block_group in cutlass.range(num_n_block_groups - 1, unroll=1): + for stage in cutlass.range_constexpr(self.num_stages_S): + ( + consumer_state_S, + producer_state_P, + producer_state_sm_stats, + consumer_state_bitmask, + ) = softmax_step_fn( + consumer_state_S, + producer_state_P, + producer_state_sm_stats, + consumer_state_bitmask, + 1 - stage, + n_block, + mask_fn=partial(mask_fn, mask_seqlen=False) + if const_expr(self.is_topk_gather and not self.disable_bitmask) + else None, + ) + n_block -= 1 + + ### last iteration if even ### + # always mask to simplify logic + if even_n_blocks: + ( + consumer_state_S, + producer_state_P, + producer_state_sm_stats, + consumer_state_bitmask, + ) = softmax_step_fn( + consumer_state_S, + producer_state_P, + producer_state_sm_stats, + consumer_state_bitmask, + 1, + n_block, + mask_fn=partial(mask_fn, mask_seqlen=False) + if not const_expr(disable_mask) + else None, + ) + n_block -= 1 + + # write row max and sum to smem + sRowSum[tidx % self.cta_tile_m, warp_idx // self.cta_group_size] = softmax.row_sum[0] + if const_expr(mLSE is not None): + if tidx < self.cta_tile_m: + sRowMax[tidx, 0] = softmax.row_max[0] + self.sm_stats_barrier_full.arrive() + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() + self.sm_stats_barrier_empty.arrive_and_wait() + + pipeline_P.producer_tail(producer_state_P) + pipeline_sm_stats.producer_tail(producer_state_sm_stats) + + @cute.jit + def softmax_step( + self, + softmax: SoftmaxSm100, + sRowMax: cute.Tensor, + sScale: cute.Tensor, + sBitmask: Optional[cute.Tensor], + tStS_t2r_staged: cute.Tensor, + tSrS_t2r: cute.Tensor, + sP_smem_view: cute.Tensor, + tmem_load_thr: cute.CopyAtom, + smem_store_thr: cute.CopyAtom, + pipeline_S: pipeline.PipelineAsync, + pipeline_P: pipeline.PipelineAsync, + pipeline_sm_stats: pipeline.PipelineAsync, + pipeline_bitmask: Optional[pipeline.PipelineAsync], + tidx: Int32, + warp_idx: Int32, + consumer_state_S: pipeline.PipelineState, + producer_state_P: pipeline.PipelineState, + producer_state_sm_stats: pipeline.PipelineState, + consumer_state_bitmask: Optional[pipeline.PipelineState], + stage: cutlass.Constexpr[Int32], + n_block: Int32, + mask_fn: Optional[Callable] = None, + is_first: Boolean = False, + store_P: Optional[Callable] = None, + gRowMax: Optional[cute.Tensor] = None, + ): + leader_warp = warp_idx == 0 + tSrP = cute.make_rmem_tensor(tSrS_t2r.shape, self.dtype_P) + rP_smem_view = smem_store_thr.retile(tSrP) + + pipeline_S.consumer_wait(consumer_state_S) + cute.copy(tmem_load_thr, tStS_t2r_staged[stage], tSrS_t2r) + cute.arch.fence_view_async_tmem_load() + pipeline_S.consumer_release(consumer_state_S) + + rBitmask = None + if const_expr(self.is_topk_gather and not self.disable_bitmask): + assert pipeline_bitmask is not None + assert consumer_state_bitmask is not None + pipeline_bitmask.consumer_wait(consumer_state_bitmask) + rBitmask = cute.make_rmem_tensor((self.tile_n // 64,), dtype=Uint32) + bitmask_col_offset = self.tile_n // 64 if warp_idx >= 2 else 0 + for i in cutlass.range_constexpr(cute.size(rBitmask)): + rBitmask[i] = sBitmask[bitmask_col_offset + i, consumer_state_bitmask.index] + + if const_expr(mask_fn is not None): + mask_fn(tSrS_t2r, n_block=n_block, rBitmask=rBitmask) + + # compute threadwise row_max + row_max = softmax.compute_row_max_local(tSrS_t2r.load(), is_first) + # 2-thread reduce row_max through smem + assert self.cta_tile_m * self.cta_group_size == 128 + sRowMax[tidx % self.cta_tile_m, warp_idx // self.cta_group_size] = row_max + self.softmax_barrier.arrive_and_wait() + # must release after barrier sync + if const_expr(self.is_topk_gather and not self.disable_bitmask): + pipeline_bitmask.consumer_release(consumer_state_bitmask) + row_max0 = sRowMax[tidx % self.cta_tile_m, 0] + row_max1 = sRowMax[tidx % self.cta_tile_m, 1] + row_max = max(row_max0, row_max1) + + row_max, acc_scale = softmax.update_row_max_from_local(row_max, is_first) + + if const_expr(gRowMax is not None): + if tidx < self.cta_tile_m: + gRowMax[tidx, n_block] = row_max + + # note: acc_scales agree for paired threads + pipeline_sm_stats.producer_acquire(producer_state_sm_stats) + if warp_idx < self.cta_group_size: + sScale[tidx % self.cta_tile_m, producer_state_sm_stats.index] = acc_scale + pipeline_sm_stats.producer_commit(producer_state_sm_stats) + + # x -> scale_log2*x-rowmax + softmax.scale_subtract_rowmax(tSrS_t2r, row_max) + + # x -> exp2(x) + softmax.apply_exp2_convert(tSrS_t2r, tSrP) + + if const_expr(self.store_P): + if leader_warp: + cute.arch.cp_async_bulk_wait_group(self.num_stages_P - 1, read=True) + self.softmax_barrier.arrive_and_wait() + + pipeline_P.producer_acquire(producer_state_P) + cute.copy( + smem_store_thr, rP_smem_view, sP_smem_view[None, None, None, producer_state_P.index] + ) + cute.arch.fence_view_async_shared() + pipeline_P.producer_commit(producer_state_P) + # unconditionally necessary for sRowMax read to complete before next iter's store + self.softmax_barrier.arrive_and_wait() + + if const_expr(self.store_P): + if leader_warp: + store_P(src_idx=producer_state_P.index, dst_idx=n_block) + cute.arch.cp_async_bulk_commit_group() + + consumer_state_S.advance() + producer_state_P.advance() + producer_state_sm_stats.advance() + if const_expr(self.is_topk_gather and not self.disable_bitmask): + consumer_state_bitmask.advance() + + softmax.update_row_sum(tSrS_t2r.load(), acc_scale, is_first) + + return consumer_state_S, producer_state_P, producer_state_sm_stats, consumer_state_bitmask + + @cute.jit + def correction_loop( + self, + softmax_scale_log2: Float32, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + tma_atom_O: Optional[cute.CopyAtom], + sRowMax: cute.Tensor, + sRowSum: cute.Tensor, + sScale: cute.Tensor, + sO: cute.Tensor, + tOtO0: cute.Tensor, + tOtO1: cute.Tensor, + pipeline_O0: pipeline.PipelineAsync, + pipeline_O1: pipeline.PipelineAsync, + pipeline_sm_stats: pipeline.PipelineAsync, + sO_empty_mbar_ptr: Optional[cute.Pointer], + tiled_copy_O_r2g: cute.TiledCopy, + topk_length_dynamic: Optional[Int32], + block_info: BlockInfo, + SeqlenInfoCls: Callable, + tile_scheduler: TileSchedulerProtocol, + mCuSeqlensQ: Optional[cute.Tensor] = None, + ): + ### ==== correction/epilogue warpgroup ==== + # Correction: copy scale smem -> rmem, copy O tmem -> rmem, rescale O, store O rmem -> tmem + # Epilogue: copy O tmem -> rmem, do final scaling of O, store O rmem -> gmem, + # optionally store LSE + # Produces: - + # Consumes: O, softmax stats + + tidx = cute.arch.thread_idx()[0] % self.num_epilogue_threads + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % ( + self.num_epilogue_threads // 32 + ) + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + leader_warp = warp_idx == 0 + + tOtO0 = tOtO0[(None, None), 0, 0] # (64, (128, 2)) + tOtO1 = tOtO1[(None, None), 0, 0] # (64, (128, 2)) + tOtOs = [tOtO0, tOtO1] + + # tuneable parameter + corr_tile_size = math.gcd(32, self.tmem_cols_Oi) + + tmem_load_atom_O = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), + self.dtype_acc, + ) + tmem_store_atom_O = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), + self.dtype_acc, + ) + thr_tmem_load_O = tcgen05.make_tmem_copy(tmem_load_atom_O, tOtO0).get_slice(tidx) + thr_tmem_store_O = tcgen05.make_tmem_copy(tmem_store_atom_O, tOtO0).get_slice(tidx) + + # ((32,1),1,4) + tOtOs_t2r = [ + thr_tmem_load_O.partition_S(tOtOs[split]) for split in range(self.num_hdimv_splits) + ] + tOtOs_r2t = [ + thr_tmem_store_O.partition_D(tOtOs[split]) for split in range(self.num_hdimv_splits) + ] + + cOi = cute.make_identity_tensor((self.cta_tile_m, self.hdimv // self.num_hdimv_splits)) + thr_tiled_copy_O_r2g = tiled_copy_O_r2g.get_slice(tidx) + tOicOi = thr_tiled_copy_O_r2g.partition_S(cOi) + + tOicOi_t2r = thr_tmem_load_O.partition_D(tOicOi[(None, None), 0, 0]) + + pipelines_O = [pipeline_O0, pipeline_O1] + + Consumer = pipeline.PipelineUserType.Consumer + consumer_state_O0 = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_Oi) + consumer_state_O1 = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_Oi) + consumer_state_sm_stats = pipeline.make_pipeline_state(Consumer, stages=self.num_stages_sm_stats) + + do_correction_rescale = partial( + self.correction_rescale, + thr_tmem_load_O, + thr_tmem_store_O, + tOicOi_t2r, + ) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + cluster_m_block, head_idx, batch_idx, _ = work_tile.tile_idx + if const_expr(self.use_packed_varlen_sched): + batch_idx = get_batch_from_cu_tensor(cluster_m_block, mCuSeqlensQ) + cluster_m_block -= mCuSeqlensQ[batch_idx] + cta_m_block = cluster_m_block * self.cta_group_size + cta_rank_in_cluster + + seqlen = SeqlenInfoCls(batch_idx) + if const_expr(self.is_topk_gather): + n_block_min = 0 + n_block_max = self.topk_length // self.tile_n + # n_block_max = topk_length_dynamic // self.tile_n + else: + n_block_min, n_block_max = block_info.get_n_block_min_max( + seqlen, + cluster_m_block, + ) + num_n_blocks = n_block_max - n_block_min + + consumer_states_O = [consumer_state_O0, consumer_state_O1] + + # acquire first signal and release immediately + pipeline_sm_stats.consumer_wait(consumer_state_sm_stats) + pipeline_sm_stats.consumer_release(consumer_state_sm_stats) + consumer_state_sm_stats.advance() + + for _ in cutlass.range(num_n_blocks - 1, unroll=1): + pipeline_sm_stats.consumer_wait(consumer_state_sm_stats) + scale = sScale[tidx % self.cta_tile_m, consumer_state_sm_stats.index] + should_rescale = cute.arch.vote_ballot_sync(scale < 1.0) != 0 + pipeline_sm_stats.consumer_release(consumer_state_sm_stats) + consumer_state_sm_stats.advance() + + for split in cutlass.range_constexpr(self.num_hdimv_splits): + consumer_state_Oi = consumer_states_O[split] + pipelines_O[split].consumer_wait(consumer_state_Oi) + if should_rescale: + do_correction_rescale( + tOtOs_t2r[split], + tOtOs_r2t[split], + scale, + ) + pipelines_O[split].consumer_release(consumer_state_Oi) + consumer_state_Oi.advance() + consumer_states_O[split] = consumer_state_Oi + + # (seqlen_q, hdimv) + mO_cur = seqlen.offset_batch_Q(mO, batch_idx, dim=3, ragged=self.ragged_tma_O)[ + None, None, head_idx + ] + # (cta_tile_m, hdimv//2, 2) + gO = cute.local_tile( + mO_cur, + (self.cta_tile_m, self.hdimv // self.num_hdimv_splits), + (cta_m_block, None), + ) + tOgO = thr_tiled_copy_O_r2g.partition_D(gO) + # ((32, 1), 1, 4) + tOrOs_t2r = [ + cute.make_rmem_tensor(tOicOi_t2r.shape, self.dtype_acc) + for split in range(self.num_hdimv_splits) + ] + tOrOs_r2g_f32 = [ + thr_tiled_copy_O_r2g.retile(tOrOs_t2r[split]) + for split in range(self.num_hdimv_splits) + ] + tOrOs_r2g = [ + cute.make_rmem_tensor_like(tOrOs_r2g_f32[split], self.dtype_O) + for split in range(self.num_hdimv_splits) + ] + if const_expr(self.use_tma_O): + tOsO = thr_tiled_copy_O_r2g.partition_D(sO) + store_O, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_O, + 0, + cute.make_layout(1), + sO, + gO, + ) + + self.sm_stats_barrier_full.arrive_and_wait() + + row_sum0 = sRowSum[tidx % self.cta_tile_m, 0] + row_sum1 = sRowSum[tidx % self.cta_tile_m, 1] + row_sum = row_sum0 + row_sum1 + acc_O_mn_row_is_zero_or_nan = row_sum == 0.0 or row_sum != row_sum + scale = cute.arch.rcp_approx(row_sum if not acc_O_mn_row_is_zero_or_nan else 1.0) + + row_max = 0.0 + if const_expr(mLSE is not None): + if tidx < self.cta_tile_m: + row_max = sRowMax[tidx, 0] + + self.sm_stats_barrier_empty.arrive() + + seqlen_q = ( + seqlen.seqlen_q + if const_expr(not self.pack_gqa) + else seqlen.seqlen_q * self.qhead_per_kvhead + ) + + # compute and store lse to gmem + if const_expr(mLSE is not None): + if const_expr(not seqlen.has_cu_seqlens_q): + mLSE_cur = mLSE[None, head_idx, batch_idx] + else: + lse_offset = ( + seqlen.offset_q if const_expr(not self.pack_gqa) else (0, seqlen.offset_q) + ) + mLSE_cur = cute.domain_offset((lse_offset,), mLSE[None, head_idx]) + gLSE = cute.local_tile(mLSE_cur, (self.cta_tile_m,), (cta_m_block,)) + if tidx < self.cta_tile_m: + LN2 = math.log(2.0) + lse = ( + (row_max * softmax_scale_log2 + cute.math.log2(row_sum, fastmath=True)) + * LN2 + if not acc_O_mn_row_is_zero_or_nan + else -Float32.inf + ) + if tidx < seqlen_q - cta_m_block * self.cta_tile_m: + gLSE[tidx] = lse + + row_idx = cta_m_block * self.cta_tile_m + tOicOi[0][0] + + for split in cutlass.range_constexpr(self.num_hdimv_splits): + consumer_state_Oi = consumer_states_O[split] + pipelines_O[split].consumer_wait(consumer_state_Oi) + # copy Oi tmem -> rmem + cute.copy( + thr_tmem_load_O, + tOtOs_t2r[split], + tOrOs_t2r[split], + ) + + # scale and downcast Oi + tOrOs_r2g[split].store((tOrOs_r2g_f32[split].load() * scale).to(self.dtype_O)) + + if const_expr(not self.use_tma_O): + # copy Oi rmem -> gmem + if row_idx < seqlen_q: + cute.copy( + thr_tiled_copy_O_r2g, + tOrOs_r2g[split], + tOgO[None, None, None, split], + ) + else: + # copy Oi rmem -> smem + if const_expr(self.overlap_sO_sV): + # last slot for Vti is always 2, 3 + sO_idx = 2 + split + else: + sO_idx = split + cute.copy( + thr_tiled_copy_O_r2g, + tOrOs_r2g[split], + tOsO[None, None, None, sO_idx], + ) + cute.arch.fence_view_async_shared() + self.epi_barrier.arrive_and_wait() + # tma store Oi smem -> gmem + if leader_warp: + store_O(src_idx=sO_idx, dst_idx=split) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(1 - split, read=True) + if const_expr(split == 1 and self.overlap_sO_sV): + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive(sO_empty_mbar_ptr) + + consumer_state_O0, consumer_state_O1 = consumer_states_O + + cute.arch.fence_view_async_tmem_load() + pipeline_O0.consumer_release(consumer_state_O0) + pipeline_O1.consumer_release(consumer_state_O1) + consumer_state_O0.advance() + consumer_state_O1.advance() + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() + + @cute.jit + def correction_rescale( + self, + thr_tmem_load: cute.CopyAtom, + thr_tmem_store: cute.CopyAtom, + tOcO_t2r: cute.Tensor, + tOtO_t2r: cute.Tensor, + tOtO_r2t: cute.Tensor, + scale: Float32, + ): + tOrO_t2r_frg = cute.make_rmem_tensor_like(tOcO_t2r[None, None, 0], self.dtype_acc) + + for i in cutlass.range_constexpr(cute.size(tOtO_t2r, mode=[2])): + tOtO_t2r_cur = tOtO_t2r[None, None, i] + tOtO_r2t_cur = tOtO_r2t[None, None, i] + + cute.copy(thr_tmem_load, tOtO_t2r_cur, tOrO_t2r_frg) + for j in cutlass.range(0, cute.size(tOrO_t2r_frg), 2, unroll_full=True): + tOrO_t2r_frg[j], tOrO_t2r_frg[j + 1] = cute.arch.mul_packed_f32x2( + (tOrO_t2r_frg[j], tOrO_t2r_frg[j + 1]), (scale, scale) + ) + cute.copy(thr_tmem_store, tOrO_t2r_frg, tOtO_r2t_cur) + cute.arch.fence_view_async_tmem_store() diff --git a/fastvideo-kernel/fa4/flash_fwd_sm100.py b/fastvideo-kernel/fa4/flash_fwd_sm100.py new file mode 100644 index 000000000..3467070c0 --- /dev/null +++ b/fastvideo-kernel/fa4/flash_fwd_sm100.py @@ -0,0 +1,4840 @@ +# Supported features: +# - BF16 & FP16 dtype +# - noncausal & causal attention +# - MHA, GQA, MQA +# - hdim 64, 96, 128, (192, 128). +# - varlen +# - sliding window +# - split-kv +# Unsupported features that will be added later: +# - page size != 128 +# - more hdim (192, 256) +# Based on the cutlass example and cute-dsl example: +# https://github.com/NVIDIA/cutlass/tree/main/examples/77_blackwell_fmha +# https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/blackwell/fmha.py + +import math +from typing import Tuple, Callable, Optional, Literal, NamedTuple +from functools import partial + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64, Boolean, const_expr +from cutlass.cute.nvgpu import cpasync +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +import cutlass.utils.blackwell_helpers as sm100_utils_basic +from cutlass import pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.utils import ClcDynamicPersistentTileScheduler +from cutlass.base_dsl.arch import Arch +from cutlass.cutlass_dsl import BaseDSL + +from quack import copy_utils, layout_utils + +from flash_attn.cute.paged_kv import PagedKVManager +from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned +from flash_attn.cute import utils +import flash_attn.cute.pipeline as pipeline_custom +import cutlass.pipeline as cutlass_pipeline +from flash_attn.cute.mask import AttentionMask +from flash_attn.cute.softmax import SoftmaxSm100, apply_score_mod_inner +from flash_attn.cute.seqlen_info import SeqlenInfoQK +from flash_attn.cute.block_info import BlockInfo +from flash_attn.cute.block_sparsity import BlockSparseTensors +from flash_attn.cute.block_sparse_utils import ( + dual_stream_block_count, + get_dual_stream_block_counts, + get_total_block_count, + produce_block_sparse_loads_dual_stream_sm100, + produce_block_sparse_loads_dual_stream_spdb_sm100, + produce_block_sparse_loads_sm100, + softmax_block_sparse_dual_stream_sm100, + softmax_block_sparse_sm100, + handle_block_sparse_empty_tile_correction_sm100, +) +from flash_attn.cute.pack_gqa import PackGQA, pack_gqa_layout +from flash_attn.cute import mma_sm100_desc as sm100_desc +from flash_attn.cute import blackwell_helpers as sm100_utils +from flash_attn.cute.named_barrier import NamedBarrierFwdSm100 +from cutlass.cute import FastDivmodDivisor +from quack.cute_dsl_utils import ParamsBase +from flash_attn.cute.tile_scheduler import ( + ClcState, + SchedulingMode, + TileSchedulerArguments, + TileSchedulerProtocol, + SingleTileScheduler, + StaticPersistentTileScheduler, + SingleTileLPTScheduler, + SingleTileVarlenScheduler, +) +from flash_attn.cute.fa_logging import fa_log, fa_printf +from flash_attn.cute.utils import smid +from flash_attn.cute.utils import AuxData + +# === TUNING KNOBS (agent-editable) === +# Keys: (use_2cta_instrs: bool, is_causal: bool, head_dim_padded: int, is_sm103: bool) +# Values: +# ex2_emu_freq: int — how often to use emulated exp2 (0=all hardware exp2, higher=more emulation). +# SM103 has fast native exp2, so set freq=0 there. +# ex2_emu_res: int — (hd256 only) number of fragment-pairs per freq period to emulate. +# ex2_emu_start_frg: int — fragment index to start emulation from +# num_regs_softmax: int — register count for softmax warps (multiple of 8) +# num_regs_correction: int — register count for correction warps (multiple of 8) +# num_regs_other is derived: 512 - num_regs_softmax * 2 - num_regs_correction +# (hd256 exception: num_regs_other is fixed at 32, not derived) +_TUNING_CONFIG = { + (True, False, 128, False): {"ex2_emu_freq": 10, "ex2_emu_start_frg": 1, "num_regs_softmax": 176, "num_regs_correction": 88}, + (False, True, 128, False): {"ex2_emu_freq": 16, "ex2_emu_start_frg": 1, "num_regs_softmax": 192, "num_regs_correction": 72}, + (True, False, 192, False): {"ex2_emu_freq": 16, "ex2_emu_start_frg": 0, "num_regs_softmax": 184, "num_regs_correction": 80}, + (False, True, 192, False): {"ex2_emu_freq": 32, "ex2_emu_start_frg": 1, "num_regs_softmax": 192, "num_regs_correction": 72}, + (True, False, 128, True): {"ex2_emu_freq": 0, "ex2_emu_start_frg": 0, "num_regs_softmax": 176, "num_regs_correction": 80}, + (False, True, 128, True): {"ex2_emu_freq": 0, "ex2_emu_start_frg": 0, "num_regs_softmax": 176, "num_regs_correction": 64}, + (True, False, 192, True): {"ex2_emu_freq": 0, "ex2_emu_start_frg": 0, "num_regs_softmax": 176, "num_regs_correction": 64}, + (False, True, 192, True): {"ex2_emu_freq": 0, "ex2_emu_start_frg": 0, "num_regs_softmax": 176, "num_regs_correction": 72}, + (True, False, 256, False): {"ex2_emu_freq": 14, "ex2_emu_res": 6, "ex2_emu_start_frg": 0, "num_regs_softmax": 256, "num_regs_correction": 160}, + (True, True, 256, False): {"ex2_emu_freq": 14, "ex2_emu_res": 6, "ex2_emu_start_frg": 0, "num_regs_softmax": 256, "num_regs_correction": 160}, +} +_FP8_TUNING_CONFIG = { + (True, False, 128, False): {'ex2_emu_freq': 10, 'ex2_emu_start_frg': 1, 'num_regs_softmax': 160, 'num_regs_correction': 72}, + # Causal hd128 FP8 previously inherited bf16's freq=16. FP8 fwd is MUFU/ex2-bound, so a more + # aggressive emulation freq=8 offloads more exp from MUFU: +3.4%(4k)..+5.5%(16k) on B200, + # MHA & GQA, accuracy-neutral (3-run validated, locked clocks). freq=8 would regress non-causal + # (0.94x), hence keyed on is_causal=True only. + (False, True, 128, False): {'ex2_emu_freq': 8, 'ex2_emu_start_frg': 1}, +} +_FP8_SMALL_HDIM_REGS = { + False: {"num_regs_softmax": 168, "num_regs_correction": 96, "num_regs_other": 80}, + True: {"num_regs_softmax": 152, "num_regs_correction": 96, "num_regs_other": 112}, +} +# === END TUNING KNOBS === + + +class DescaleTensors(NamedTuple): + q_descale: Optional[cute.Tensor] = None + k_descale: Optional[cute.Tensor] = None + v_descale: Optional[cute.Tensor] = None + + def __new_from_mlir_values__(self, values): + return DescaleTensors(*((*values, None, None, None)[:3])) + + +def supports_sm100_tile_m64_head_dims(head_dim: int, head_dim_v: int) -> bool: + return ( + 8 <= head_dim <= 128 + and 8 <= head_dim_v <= 128 + and head_dim % 8 == 0 + and head_dim_v % 8 == 0 + ) or (head_dim, head_dim_v) == (192, 128) + + +class FlashAttentionForwardSm100: + + def __init__( + self, + # dtype: Type[cutlass.Numeric], + head_dim: int, + head_dim_v: Optional[int] = None, + qhead_per_kvhead: cutlass.Constexpr[int] = 1, + is_causal: bool = False, + is_local: bool = False, + is_split_kv: bool = False, + pack_gqa: bool = False, + q_subtile_factor: int = 1, + m_block_size: int = 128, + n_block_size: int = 128, + q_stage: cutlass.Constexpr[int] = 2, + num_kv_streams: cutlass.Constexpr[int] = 1, + use_sp_double_buffer: cutlass.Constexpr[bool] = False, + is_persistent: bool = True, + score_mod: cutlass.Constexpr | None = None, + mask_mod: cutlass.Constexpr | None = None, + has_aux_tensors: cutlass.Constexpr = False, + paged_kv_non_tma: bool = False, + is_varlen_q: bool = False, + use_2cta_instrs: bool = False, + use_clc_scheduler: bool = False, + ): + self.use_tma_KV = not paged_kv_non_tma + # self.dtype = dtype + # padding head_dim to a multiple of 16 as k_block_size + hdim_multiple_of = 16 + self.head_dim_padded = int(math.ceil(head_dim / hdim_multiple_of) * hdim_multiple_of) + head_dim_v = head_dim_v if head_dim_v is not None else head_dim + self.same_hdim_kv = head_dim == head_dim_v + self.head_dim_v_padded = int(math.ceil(head_dim_v / hdim_multiple_of) * hdim_multiple_of) + self.same_hdim_kv_padded = self.head_dim_padded == self.head_dim_v_padded + self.check_hdim_oob = head_dim != self.head_dim_padded + self.check_hdim_v_oob = head_dim_v != self.head_dim_v_padded + self.m_block_size = m_block_size + self.n_block_size = n_block_size + self.p_in_smem = m_block_size == 64 + self.q_stage = 1 if self.p_in_smem else q_stage + assert self.q_stage in [1, 2] + if m_block_size not in (64, 128) or ( + m_block_size == 64 and n_block_size not in (64, 128) + ): + raise NotImplementedError( + "FlashAttentionForwardSm100 supports m_block_size=128, or m_block_size=64 " + f"with n_block_size in (64, 128); got ({m_block_size}, {n_block_size})." + ) + if self.p_in_smem and ( + not supports_sm100_tile_m64_head_dims(head_dim, head_dim_v) + or use_2cta_instrs + or use_clc_scheduler + or paged_kv_non_tma + ): + raise NotImplementedError( + "SM100 tile_m=64 forward requires supported head dimensions, 1CTA, " + "the static scheduler, and TMA KV loads" + ) + self.num_kv_streams = num_kv_streams + assert self.num_kv_streams in [1, 2] + self.num_compute_slots = self.q_stage * self.num_kv_streams + assert self.num_compute_slots in [1, 2] + self.use_sp_double_buffer = use_sp_double_buffer + if self.use_sp_double_buffer: + assert self.num_kv_streams == 2 + assert self.q_stage == 1 + assert self.m_block_size == 128 and self.n_block_size == 64 + assert self.head_dim_padded == 128 and self.head_dim_v_padded == 128 + assert not is_causal and not is_local + assert not is_varlen_q and not is_split_kv and not pack_gqa + assert self.use_tma_KV and not use_2cta_instrs and not use_clc_scheduler + self.num_sp_tmem_slots = 4 if self.use_sp_double_buffer else 2 + self.num_sp_pipeline_stages = ( + 4 if self.use_sp_double_buffer else self.num_compute_slots + ) + self.use_2cta_instrs = use_2cta_instrs + # If split_P_arrive, the softmax warps write some columns of P first, signal to the MMA warp + # to being the P @ V MMA, then write the rest of P and signal again. This allows some overlap + # between compute the last couple columns of P and the P @ V MMA. + self.split_P_arrive = n_block_size // 4 * 3 + self.split_P_arrive = int(self.split_P_arrive / 32) * 32 # multiple of 32 + assert self.split_P_arrive % 32 == 0 + assert self.split_P_arrive < self.n_block_size + if self.p_in_smem: + self.split_P_arrive = 0 + self.arch = BaseDSL._get_dsl().get_arch_enum() + assert self.arch.is_family_of(Arch.sm_100f) or self.arch.is_family_of(Arch.sm_110f), \ + "Only SM 10.x and 11.x are supported" + if self.p_in_smem and self.arch not in (Arch.sm_100a, Arch.sm_100f): + raise NotImplementedError("tile_m=64 forward is currently supported only on SM100") + + self.cta_group_size = 2 if self.use_2cta_instrs else 1 + # cta_tiler M includes only 1 CTA, the scheduler will take into account the cluster shape + self.cta_tiler = (self.q_stage * m_block_size, n_block_size, self.head_dim_padded) + # With 2CTA, the MMA tiler M covers both CTAs, so it's cta_group_size * m_block_size. + # Each CTA owns m_block_size rows; the 2CTA MMA instruction spans both. + self.mma_tiler_qk = (self.cta_group_size * m_block_size, n_block_size, self.head_dim_padded) + self.mma_tiler_pv = (self.cta_group_size * m_block_size, self.head_dim_v_padded, n_block_size) + self.qk_acc_dtype = Float32 + self.pv_acc_dtype = Float32 + self.cluster_shape_mn = (2, 1) if self.use_2cta_instrs else (1, 1) + self.is_persistent = is_persistent + self.is_causal = is_causal + self.is_local = is_local + self.is_varlen_q = is_varlen_q + self.qhead_per_kvhead = qhead_per_kvhead + self.is_split_kv = is_split_kv + self.pack_gqa = pack_gqa + self.use_tma_O = ( + not (self.pack_gqa and self.m_block_size % self.qhead_per_kvhead != 0) + and not (self.pack_gqa and self.is_split_kv) + and not is_varlen_q + ) + self.use_correction_warps_for_epi = not self.use_tma_O + self.q_subtile_factor = q_subtile_factor + assert not (self.is_split_kv and self.head_dim_v_padded >= 192), ( + "SplitKV is not supported for hdim >= 192" + ) + self.score_mod = score_mod + self.mask_mod = mask_mod + self.score_vec_size: cutlass.Constexpr = getattr( + score_mod, "__vec_size__", 1 if cutlass.const_expr(has_aux_tensors) else 2 + ) + self.mask_vec_size: cutlass.Constexpr = getattr(mask_mod, "__vec_size__", 1) + # Does S1 need to wait for S0 to finish + # self.s0_s1_barrier = self.head_dim_padded in [64, 96] and (not self.is_causal and not self.is_local) + # NOTE: is_family_of also matches any future sm_10x with x > 3 — intentional. + # The flag gates ex2 emulation; sm_103 (B300) has fast hardware ex2 and later + # Blackwell variants are assumed to inherit this, so forward-inclusion is correct + # despite the literal `is_sm103` name. + is_sm103 = self.arch.is_family_of(Arch.sm_103f) + self.is_sm103 = is_sm103 + # enable_ex2_emu is derived: True if tuning config has freq > 0, else fallback to default logic + _default_enable_ex2_emu = (self.head_dim_padded <= 128 or (self.head_dim_padded == 192 and self.use_2cta_instrs and not self.is_causal and not self.is_local)) and not is_sm103 + self.enable_ex2_emu = _default_enable_ex2_emu + self.s0_s1_barrier = False + self.overlap_sO_sQ = ( + (self.head_dim_padded == 192 and self.head_dim_v_padded >= 64) or + (self.head_dim_v_padded >= 128 and self.is_split_kv) + ) + if self.overlap_sO_sQ: + self.is_persistent = False + + assert self.use_tma_KV or not (self.check_hdim_oob or self.check_hdim_v_oob), ( + "Paged KV does not support irregular head dim" + ) + + # ClC does not compose with these other features, so disable even if requested + self.use_clc_scheduler = ( + use_clc_scheduler + and self.use_tma_KV + and not self.overlap_sO_sQ + ) + self.sched_stages = 1 + if self.use_clc_scheduler: + assert self.cluster_shape_mn[1] == 1, f"CLC requires cluster N == 1: {self.cluster_shape_mn}" + assert self.cluster_shape_mn[0] in (1, 2), f"bad CLC cluster M: {self.cluster_shape_mn}" + assert self.cluster_shape_mn[0] == self.cta_group_size, ( + f"CLC cluster M != cta_group_size: {self.cluster_shape_mn}, {self.cta_group_size}" + ) + + self.scheduling_mode = SchedulingMode.CLC if self.use_clc_scheduler else SchedulingMode.STATIC + + if is_varlen_q: + self.TileScheduler = SingleTileVarlenScheduler + elif self.is_causal or self.is_local or self.use_clc_scheduler: + self.TileScheduler = SingleTileLPTScheduler + elif self.is_persistent: + self.TileScheduler = StaticPersistentTileScheduler + else: + self.TileScheduler = SingleTileScheduler + + fa_log(1, f"TileScheduler={self.TileScheduler.__name__}, scheduling_mode={self.scheduling_mode.name}, USE_2CTA={self.use_2cta_instrs}") + + self.softmax0_warp_ids = (0, 1, 2, 3) + self.softmax1_warp_ids = (4, 5, 6, 7) + self.correction_warp_ids = (8, 9, 10, 11) + self.mma_warp_id = 12 + self.epilogue_warp_ids = (13,) + self.load_warp_ids = (14,) + self.empty_warp_ids = (15,) + self.tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols("sm_100") + + self.threads_per_cta = cute.arch.WARP_SIZE * len( + ( + *self.softmax0_warp_ids, + *self.softmax1_warp_ids, + *self.correction_warp_ids, + self.mma_warp_id, + *self.load_warp_ids, + *self.epilogue_warp_ids, + *self.empty_warp_ids, + ) + ) + + self.use_tma_Q = not (self.pack_gqa and self.m_block_size % self.qhead_per_kvhead != 0) + + if self.num_compute_slots == 1: + if not self.use_tma_KV or not self.use_tma_Q: + self.empty_warp_ids = self.empty_warp_ids + self.load_warp_ids + self.load_warp_ids = self.softmax1_warp_ids + else: + self.empty_warp_ids = self.empty_warp_ids + self.softmax1_warp_ids + self.softmax1_warp_ids = () + elif not self.use_tma_KV: + self.load_warp_ids = (14, 15) + self.empty_warp_ids = () + + if self.use_correction_warps_for_epi: + self.empty_warp_ids = self.empty_warp_ids + self.epilogue_warp_ids + self.epilogue_warp_ids = self.correction_warp_ids + + self.clc_scheduler_warp_id = self.empty_warp_ids[0] if self.use_clc_scheduler else None + + self.tmem_s_offset = ( + [i * self.n_block_size for i in range(self.num_sp_tmem_slots)] + if self.use_sp_double_buffer + else [0, self.n_block_size] + ) + self.tmem_o_offset = [ + self.tmem_s_offset[-1] + self.n_block_size + i * self.head_dim_v_padded + for i in range(self.num_compute_slots) + ] + self.tmem_total = self.tmem_o_offset[-1] + self.head_dim_v_padded + assert self.tmem_total <= self.tmem_alloc_cols + self.tmem_s_to_p_offset = self.n_block_size // 2 + self.tmem_p_offset = [ + self.tmem_s_offset[i] + self.tmem_s_to_p_offset + for i in range(self.num_sp_tmem_slots) + ] + + # vec buffer for row_max & row_sum + self.tmem_vec_offset = self.tmem_s_offset + + # Look up tuning config for register counts and ex2_emu params + _tune_key = (self.use_2cta_instrs, self.is_causal, self.head_dim_padded, self.is_sm103) + self._tune = _TUNING_CONFIG.get(_tune_key, {}) + if "ex2_emu_freq" in self._tune: + self.enable_ex2_emu = self._tune["ex2_emu_freq"] > 0 + if self.head_dim_padded < 96: + self.num_regs_softmax = 200 if not paged_kv_non_tma else 184 + self.num_regs_correction = 64 + self.num_regs_other = 48 if not paged_kv_non_tma else 80 + else: + if not paged_kv_non_tma and "num_regs_softmax" in self._tune: + self.num_regs_softmax = self._tune["num_regs_softmax"] + self.num_regs_correction = self._tune["num_regs_correction"] + elif not paged_kv_non_tma: + self.num_regs_softmax = 192 + self.num_regs_correction = 80 + else: + self.num_regs_softmax = 184 + self.num_regs_correction = 64 + self.num_regs_other = 512 - self.num_regs_softmax * 2 - self.num_regs_correction + if self.p_in_smem and self.num_kv_streams == 2: + self.num_regs_softmax = 176 + self.num_regs_correction = 112 + self.num_regs_other = 48 + + self.buffer_align_bytes = 1024 + + def _setup_attributes(self): + """Set up configurations and parameters for the FMHA kernel operation. + + This method initializes and configures various attributes required for the + execution of the fused multi-head attention kernel, mainly about the pipeline stages: + + - Sets up staging parameters for Q, K, V inputs and accumulator data + - Configures pipeline stages for softmax, correction, and epilogue operations + """ + + smem_size_q = self.q_stage * self.m_block_size * self.head_dim_padded * self.q_dtype.width // 8 + smem_size_o = self.q_stage * self.m_block_size * self.head_dim_v_padded * self.o_dtype.width // 8 + smem_size_q_o = smem_size_q + smem_size_o if not self.overlap_sO_sQ else max(smem_size_q, smem_size_o) + smem_size_k_per_stage = self.n_block_size * self.head_dim_padded * self.k_dtype.width // 8 + smem_size_v_per_stage = self.n_block_size * self.head_dim_v_padded * self.v_dtype.width // 8 + smem_size_kv_per_stage = max(smem_size_k_per_stage, smem_size_v_per_stage) // self.cta_group_size + self.s_stage = self.num_sp_tmem_slots + smem_size_p = ( + self.s_stage * self.m_block_size * self.n_block_size * self.q_dtype.width // 8 + if self.p_in_smem + else 0 + ) + # Cap small head_dim from over-staging: the 224*1024 budget undercounts + # per-stage state, so at hd_padded=16 the unbounded formula picks 52 stages + # and overflows the 227 KB SMEM cap. No-op for hd_padded >= 32 (max 26). + kv_stage = min( + (224 * 1024 - smem_size_q_o - smem_size_p) // smem_size_kv_per_stage, + 32, + ) + if self.head_dim_padded == 192 and self.head_dim_v_padded == 128 and kv_stage == 2: + # For hdim 192,128, we can fit 3 stages if we use uneven_kv_smem + kv_stage = 3 + self.kv_stage = kv_stage + # print("kv_stage", self.kv_stage) + assert self.s_stage >= self.num_compute_slots + # For hdim 192,128 1CTA, we don't have enough smem to store all 3 stages of KV: + # 128 x 192 x 2 bytes x 3 stages = 144KB, and we need 96KB for Q. + # Instead we store smem as [smem_large, smem_small, smem_large], where smem_large is + # 128 x 192 and smem_small is 128 x 128. We set the stride between the stages to be + # 128 * 160, so that indexing the 0th and 2nd stages will get the right address, + # but for the 1st stage we need to add or subtract (depending on phase) 128 x 64. + self.uneven_kv_smem = ( + self.head_dim_padded == 192 and self.head_dim_v_padded == 128 and self.kv_stage == 3 + ) + self.uneven_kv_smem_offset = ( + self.n_block_size * (self.head_dim_padded - self.head_dim_v_padded) // 2 + if self.uneven_kv_smem + else 0 + ) + assert self.uneven_kv_smem_offset % 1024 == 0 + + @cute.jit + def __call__( + self, + mQ: cute.Tensor, # (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q + mK: cute.Tensor, # (b_k, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k or (num_pages, page_size, h_k, d) if there is page_table + mV: cute.Tensor, # (b_k, s_k, h_k, dv) or (total_k, h_k, dv) if there is cu_seqlens_k or (num_pages, page_size, h_k, dv) if there is page_table + mO: cute.Tensor, # (b, s_q, h, dv) or (total_q, h, dv) if there is cu_seqlens_q + mLSE: Optional[cute.Tensor], + softmax_scale: Float32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + mPageTable: Optional[cute.Tensor] = None, # (b_k, max_num_pages_per_seq) + window_size_left: Int32 | int | None = None, + window_size_right: Int32 | int | None = None, + learnable_sink: Optional[cute.Tensor] = None, + descale_tensors: Optional[DescaleTensors] = None, + blocksparse_tensors: Optional[BlockSparseTensors] = None, + aux_data: AuxData = AuxData(), + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + """Execute the Fused Multi-Head Attention operation on the provided tensors. + + This method prepares the input tensors for processing, validates their shapes and types, + configures the computation parameters, and launches the CUDA kernel. + + The method handles: + 1. Tensor layout transformations for specific memory access patterns + 2. Validation of tensor shapes and data types + 3. Initialization of hardware-specific parameters and memory layouts + 4. Configuration of TMA (Tensor Memory Access) operations + 5. Grid and work scheduling computation + 6. Kernel launch with appropriate parameters + """ + # setup static attributes before smem/grid/tma computation + self.q_dtype = mQ.element_type + self.k_dtype = mK.element_type + self.v_dtype = mV.element_type + self.o_dtype = mO.element_type + mQ, mK, mV, mO = [assume_tensor_aligned(t) for t in (mQ, mK, mV, mO)] + Q_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1] + mQ = cute.make_tensor(mQ.iterator, cute.select(mQ.layout, mode=Q_layout_transpose)) + # (s_k, d, h_k, b_k) or (total_k, d, h_k) if there's cu_seqlens_k or (page_size, d, h_k, num_pages) if there's page_table + KV_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensK is None) else [0, 2, 1] + mK, mV = [ + cute.make_tensor(t.iterator, cute.select(t.layout, mode=KV_layout_transpose)) + for t in (mK, mV) + ] + if const_expr(self.is_split_kv): + O_layout_transpose = [2, 4, 3, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 3, 2, 0] + LSE_layout_transpose = [3, 2, 1, 0] if const_expr(mCuSeqlensQ is None) else [2, 1, 0] + num_splits = mO.shape[0] + else: + O_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1] + LSE_layout_transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0] + num_splits = Int32(1) + mO = cute.make_tensor(mO.iterator, cute.select(mO.layout, mode=O_layout_transpose)) + mLSE = ( + cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=LSE_layout_transpose)) + if const_expr(mLSE is not None) + else None + ) + # (s, d, h, b) -> (d, s, h, b) + V_layout_transpose = [1, 0, 2, 3] if const_expr(mCuSeqlensK is None) else [1, 0, 2] + mV = cute.make_tensor(mV.iterator, cute.select(mV.layout, mode=V_layout_transpose)) + + # check type consistency + if const_expr(self.q_dtype != self.k_dtype): + raise TypeError(f"Type mismatch: {self.q_dtype} != {self.k_dtype}") + if const_expr(self.q_dtype != self.v_dtype): + raise TypeError(f"Type mismatch: {self.q_dtype} != {self.v_dtype}") + if const_expr( + self.p_in_smem and self.q_dtype not in (cutlass.Float16, cutlass.BFloat16) + ): + raise TypeError("SM100 tile_m=64 forward requires FP16 or BF16 tensors") + if const_expr(self.q_dtype.width == 8): + paged_kv_non_tma = not self.use_tma_KV + if const_expr(self.head_dim_padded < 96): + fp8_regs = _FP8_SMALL_HDIM_REGS[paged_kv_non_tma] + self.num_regs_softmax = fp8_regs["num_regs_softmax"] + self.num_regs_correction = fp8_regs["num_regs_correction"] + self.num_regs_other = fp8_regs["num_regs_other"] + else: + fp8_tune = _FP8_TUNING_CONFIG.get( + (self.use_2cta_instrs, self.is_causal, self.head_dim_padded, self.is_sm103), {} + ) + if const_expr("ex2_emu_freq" in fp8_tune): + self._tune = {**self._tune, **fp8_tune} + self.enable_ex2_emu = self._tune["ex2_emu_freq"] > 0 + if const_expr(not paged_kv_non_tma and "num_regs_softmax" in fp8_tune): + self.num_regs_softmax = fp8_tune["num_regs_softmax"] + self.num_regs_correction = fp8_tune["num_regs_correction"] + self.num_regs_other = 512 - self.num_regs_softmax * 2 - self.num_regs_correction + self._setup_attributes() + self.ex2_emu_freq = 0 + self.ex2_emu_start_frg = self._tune.get("ex2_emu_start_frg", 1) + if const_expr( + self.num_kv_streams == 2 + and self.m_block_size == 128 + and self.n_block_size == 64 + and self.head_dim_padded == 128 + and not self.is_causal + and not self.is_local + and not self.is_sm103 + ): + self.ex2_emu_start_frg = 0 + if const_expr(self.enable_ex2_emu): + self.ex2_emu_freq = self._tune.get("ex2_emu_freq", 16) + if const_expr( + self.pack_gqa and self.head_dim_padded > 64 and not self.is_causal and not self.is_local + ): + self.ex2_emu_freq = 32 if mCuSeqlensQ is not None or mSeqUsedQ is not None else self._tune.get("ex2_emu_freq", 10) + if const_expr( + self.num_kv_streams == 2 + and self.m_block_size == 128 + and self.n_block_size == 64 + and self.head_dim_padded == 128 + and not self.is_causal + and not self.is_local + and not self.is_sm103 + ): + self.ex2_emu_freq = 24 if self.use_sp_double_buffer else 8 + + cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + q_major_mode = tcgen05.OperandMajorMode.K + k_major_mode = tcgen05.OperandMajorMode.K + v_major_mode = tcgen05.OperandMajorMode.MN + self.o_layout = cutlass.utils.LayoutEnum.from_tensor(mO) + p_major_mode = tcgen05.OperandMajorMode.K + tiled_mma_qk = sm100_utils_basic.make_trivial_tiled_mma( + self.q_dtype, + q_major_mode, + k_major_mode, + self.qk_acc_dtype, + cta_group, + self.mma_tiler_qk[:2], + ) + tiled_mma_pv = sm100_utils_basic.make_trivial_tiled_mma( + self.v_dtype, + p_major_mode, + v_major_mode, + self.pv_acc_dtype, + cta_group, + self.mma_tiler_pv[:2], + tcgen05.OperandSource.SMEM if self.p_in_smem else tcgen05.OperandSource.TMEM, + ) + + self.cluster_shape_mnk = (*self.cluster_shape_mn, 1) + cta_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), (tiled_mma_qk.thr_id.shape,) + ) + + # epi_tile is per-CTA (not full 2CTA) since each CTA writes its own O portion + self.epi_tile = (self.m_block_size, self.head_dim_v_padded) + + sQ_layout = sm100_utils_basic.make_smem_layout_a( + tiled_mma_qk, self.mma_tiler_qk, self.q_dtype, self.q_stage + ) + sK_layout = sm100_utils_basic.make_smem_layout_b( + tiled_mma_qk, self.mma_tiler_qk, self.k_dtype, self.kv_stage + ) + tP_layout = sm100_utils_basic.make_smem_layout_a( + tiled_mma_pv, self.mma_tiler_pv, self.q_dtype, self.s_stage + ) + sV_layout = sm100_utils_basic.make_smem_layout_b( + tiled_mma_pv, self.mma_tiler_pv, self.v_dtype, self.kv_stage + ) + sO_layout = sm100_utils_basic.make_smem_layout_epi( + self.o_dtype, self.o_layout, self.epi_tile, self.q_stage + ) + if const_expr(not self.same_hdim_kv_padded): + # sK and sV are using the same physical smem so we need to adjust the stride so that they line up + stride_sK = const_expr( + max(sK_layout.outer.stride[-1], 0) + ) # take max to turn tuple to Int32 + stride_sV = const_expr(max(sV_layout.outer.stride[-1], 0)) + stage_stride = const_expr( + max(stride_sK, stride_sV) + if not self.uneven_kv_smem + else (stride_sK + stride_sV) // 2 + ) + sK_layout = cute.make_composed_layout( + sK_layout.inner, + 0, + cute.make_layout( + (*sK_layout.outer.shape[:-1], self.kv_stage), + stride=(*sK_layout.outer.stride[:-1], stage_stride), + ), + ) + sV_layout = cute.make_composed_layout( + sV_layout.inner, + 0, + cute.make_layout( + (*sV_layout.outer.shape[:-1], self.kv_stage), + stride=(*sV_layout.outer.stride[:-1], stage_stride), + ), + ) + + if const_expr(self.pack_gqa): + nheads_kv = mK.shape[2] + mQ = pack_gqa_layout(mQ, self.qhead_per_kvhead, nheads_kv, head_idx=2) + mO = pack_gqa_layout(mO, self.qhead_per_kvhead, nheads_kv, head_idx=2) + if const_expr(mLSE is not None): + mLSE = pack_gqa_layout(mLSE, self.qhead_per_kvhead, nheads_kv, head_idx=1) + + self.tma_copy_bytes = { + name: cute.size_in_bytes(mX.element_type, cute.select(layout, mode=[0, 1, 2])) + for name, mX, layout in [ + ("Q", mQ, sQ_layout), + ("K", mK, sK_layout), + ("V", mV, sV_layout), + ] + } + for name in ("Q", "K", "V"): + self.tma_copy_bytes[name] *= self.cta_group_size + + # TMA load for Q + tma_load_op = cpasync.CopyBulkTensorTileG2SOp(cta_group) + tma_store_op = cpasync.CopyBulkTensorTileS2GOp() + + if const_expr(self.use_tma_Q): + tma_atom_Q, mQ = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + mQ, + cute.select(sQ_layout, mode=[0, 1, 2]), + self.mma_tiler_qk, + tiled_mma_qk, + cta_layout_vmnk.shape, + ) + gmem_tiled_copy_Q = None + else: + tma_atom_Q = None + async_copy_elems = 128 // self.q_dtype.width + num_load_threads = cute.arch.WARP_SIZE * len(self.load_warp_ids) + threads_per_row = math.gcd(self.head_dim_padded // async_copy_elems, num_load_threads) + gmem_tiled_copy_Q = copy_utils.tiled_copy_2d( + self.q_dtype, threads_per_row, num_load_threads, async_copy_elems, is_async=True + ) + + tma_atom_K = None + tma_atom_V = None + if const_expr(self.use_tma_KV): + # TMA load for K + tma_atom_K, mK = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + mK, + cute.select(sK_layout, mode=[0, 1, 2]), + self.mma_tiler_qk, + tiled_mma_qk, + cta_layout_vmnk.shape, + ) + # TMA load for V + tma_atom_V, mV = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + mV, + cute.select(sV_layout, mode=[0, 1, 2]), + self.mma_tiler_pv, + tiled_mma_pv, + cta_layout_vmnk.shape, + ) + + self.num_epilogue_threads = cute.arch.WARP_SIZE * len(self.epilogue_warp_ids) + if const_expr(self.use_tma_O): + tma_atom_O, mO = cpasync.make_tiled_tma_atom( + tma_store_op, mO, cute.select(sO_layout, mode=[0, 1]), self.epi_tile + ) + gmem_tiled_copy_O = None + else: + tma_atom_O = None + universal_copy_bits = 128 + async_copy_elems = universal_copy_bits // self.o_dtype.width + atom_universal_copy = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.o_dtype, + num_bits_per_copy=universal_copy_bits, + ) + tO_shape_dim_1 = sO_layout.outer.shape[1][0] // async_copy_elems + tO_layout = cute.make_ordered_layout( + (self.num_epilogue_threads // tO_shape_dim_1, tO_shape_dim_1), + order=(1, 0), + ) + # So that we don't have to check if we overshoot kBlockM when we store O + assert self.m_block_size % tO_layout.shape[0] == 0 + vO_layout = cute.make_layout((1, async_copy_elems)) + gmem_tiled_copy_O = cute.make_tiled_copy_tv(atom_universal_copy, tO_layout, vO_layout) + + TileScheduler = self.TileScheduler + _num_block_divisor = self.cta_tiler[0] * (self.cta_group_size if not self.is_persistent and self.cta_group_size > 1 else 1) + tile_sched_args = TileSchedulerArguments( + cute.ceil_div(cute.size(mQ.shape[0]), _num_block_divisor), + cute.size(mQ.shape[2]), + cute.size(mQ.shape[3]) + if const_expr(mCuSeqlensQ is None) + else cute.size(mCuSeqlensQ.shape[0] - 1), + num_splits, + cute.size(mK.shape[0]) + if const_expr(mPageTable is None) + else mK.shape[0] * mPageTable.shape[1], + mQ.shape[1], + mV.shape[0], # Note that this is different from Sm90 since we transpose mV in Sm100 + total_q=cute.size(mQ.shape[0]) + if const_expr(mCuSeqlensQ is not None) + else cute.size(mQ.shape[0]) * cute.size(mQ.shape[3]), + tile_shape_mn=self.cta_tiler[:2], + mCuSeqlensQ=mCuSeqlensQ, + mSeqUsedQ=mSeqUsedQ, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + element_size=self.k_dtype.width // 8, + is_persistent=self.is_persistent, + lpt=self.is_causal or self.is_local, + is_split_kv=self.is_split_kv, + cluster_shape_mn=self.cluster_shape_mn, + use_cluster_idx=not self.is_persistent and self.cta_group_size > 1, + ) + tile_sched_params = TileScheduler.to_underlying_arguments( + tile_sched_args, scheduling_mode=self.scheduling_mode + ) + self.tile_scheduler_cls = TileScheduler + grid_dim = TileScheduler.get_grid_shape(tile_sched_params) + + sO_size = cute.cosize(sO_layout) if const_expr(not self.overlap_sO_sQ) else 0 + sQ_size = ( + cute.cosize(sQ_layout) if const_expr(not self.overlap_sO_sQ) else + cutlass.max(cute.cosize(sQ_layout), cute.cosize(sO_layout) * self.o_dtype.width // self.q_dtype.width) + ) + + clc_response_size = self.sched_stages * 4 if self.use_clc_scheduler else 0 + clc_mbar_size = self.sched_stages * 2 if self.use_clc_scheduler else 0 + smem_p_size = ( + self.s_stage * self.m_block_size * self.n_block_size if self.p_in_smem else 0 + ) + + @cute.struct + class SharedStorage: + # m_barriers for pipelines + mbar_load_Q: cute.struct.MemRange[Int64, self.q_stage * 2] + mbar_load_KV: cute.struct.MemRange[Int64, self.kv_stage * 2] + mbar_S_full_P_full_O_rescaled: cute.struct.MemRange[ + Int64, self.num_sp_pipeline_stages * 2 + ] + mbar_P_full_lastsplit: cute.struct.MemRange[ + Int64, self.num_sp_pipeline_stages * 2 + ] + mbar_O_full: cute.struct.MemRange[Int64, self.num_compute_slots * 2] + mbar_softmax_stats: cute.struct.MemRange[Int64, self.num_compute_slots * 2] + # mbar_softmax_stats: cute.struct.MemRange[Int64, self.num_compute_slots * 4 * 2] + mbar_O_epi: cute.struct.MemRange[Int64, self.q_stage * 2] + mbar_s0_s1_sequence: cute.struct.MemRange[Int64, 2 * 2] + # Tmem dealloc cluster barrier + tmem_dealloc_mbar_ptr: Int64 + # Tmem holding buffer + tmem_holding_buf: Int32 + # Smem tensors + # store row max and row sum + sScale: cute.struct.MemRange[ + Float32, self.num_compute_slots * self.m_block_size * 2 + ] + # CLC buffers placed here to utilize padding before sO's 1024-byte alignment. + # This avoids adding bytes at the end when we're at the smem limit. + # PipelineClcFetchAsync expects 2 * sched_stages mbarriers (full + empty). + clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, clc_mbar_size] + # CLC response storage (16 bytes per stage, stored as 4 Int32s). + clc_response: cute.struct.MemRange[Int32, clc_response_size] + # Large TMA buffers with 1024-byte alignment + sO: cute.struct.Align[ + cute.struct.MemRange[self.o_dtype, sO_size], self.buffer_align_bytes + ] + sQ: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, sQ_size], self.buffer_align_bytes + ] + sK: cute.struct.Align[ + # cute.cosize(sK_layout) is correct even in the case of self.uneven_kv_smem + cute.struct.MemRange[self.k_dtype, cute.cosize(sK_layout)], + self.buffer_align_bytes, + ] + sP: cute.struct.Align[ + cute.struct.MemRange[self.q_dtype, smem_p_size], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + softmax_scale_log2, softmax_scale = utils.compute_softmax_scale_log2(softmax_scale, self.score_mod) + window_size_left = Int32(window_size_left) if window_size_left is not None else None + window_size_right = Int32(window_size_right) if window_size_right is not None else None + fastdiv_mods = utils.compute_fastdiv_mods(mQ, mK, self.qhead_per_kvhead, self.pack_gqa, aux_data.tensors, mPageTable) + + head_divmod = None + if cutlass.const_expr(self.pack_gqa): + head_divmod = FastDivmodDivisor(self.qhead_per_kvhead) + + self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None) + if cutlass.const_expr(self.use_block_sparsity and mPageTable is not None): + raise NotImplementedError("Block sparsity + paged KV not supported on SM100") + if cutlass.const_expr(self.use_block_sparsity and self.is_varlen_q): + assert const_expr(blocksparse_tensors.cu_total_m_blocks is not None), ( + "blocksparse_tensors.cu_total_m_blocks must be provided for varlen blocksparsity" + ) + + # Launch the kernel synchronously + self.kernel( + mQ, + mK, + mV, + mO, + mLSE, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + mPageTable, + tma_atom_Q, + tma_atom_K, + tma_atom_V, + tma_atom_O, + softmax_scale_log2, + softmax_scale, + window_size_left, + window_size_right, + learnable_sink, + descale_tensors, + blocksparse_tensors, + sQ_layout, + sK_layout, + tP_layout, + sV_layout, + sO_layout, + gmem_tiled_copy_Q, + gmem_tiled_copy_O, + tiled_mma_qk, + tiled_mma_pv, + tile_sched_params, + num_splits, + aux_data, + fastdiv_mods, + head_divmod, + ).launch( + grid=grid_dim, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk if cute.size(self.cluster_shape_mnk) > 1 else None, + stream=stream, + min_blocks_per_mp=1, + ) + + def _generate_attention_mask_cls(self, window_size_left, window_size_right): + return partial( + AttentionMask, + self.m_block_size, + self.n_block_size, + window_size_left=window_size_left, + window_size_right=window_size_right, + qhead_per_kvhead_packgqa=( + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1 + ), + ) + + # GPU device kernel + @cute.kernel + def kernel( + self, + mQ: cute.Tensor, # (s_q, d, h, b) or (total_q, d, h) if there is cu_seqlens_q + mK: cute.Tensor, # (s_k, d, h_k, b_k) or (total_k, d, h_k) if there is cu_seqlens_k or (page_size, d, h_k, num_pages) if there is page_table + mV: cute.Tensor, # (d, s_k, h_k, b_k) or (d, total_k, h_k) if there is cu_seqlens_k or (d, page_size, h_k, num_pages) if there is page_table + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + mCuSeqlensQ: Optional[cute.Tensor], + mCuSeqlensK: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + mPageTable: Optional[cute.Tensor], + tma_atom_Q: Optional[cute.CopyAtom], + tma_atom_K: Optional[cute.CopyAtom], + tma_atom_V: Optional[cute.CopyAtom], + tma_atom_O: Optional[cute.CopyAtom], + softmax_scale_log2: Float32, + softmax_scale: Float32 | None, + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + learnable_sink: Optional[cute.Tensor], + descale_tensors: Optional[DescaleTensors], + blocksparse_tensors: Optional[BlockSparseTensors], + sQ_layout: cute.ComposedLayout, + sK_layout: cute.ComposedLayout, + tP_layout: cute.ComposedLayout, + sV_layout: cute.ComposedLayout, + sO_layout: cute.ComposedLayout, + gmem_tiled_copy_Q: Optional[cute.TiledCopy], + gmem_tiled_copy_O: Optional[cute.TiledCopy], + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + tile_sched_params: ParamsBase, + num_splits: Int32, + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + head_divmod=None, + ): + """The device kernel implementation of the Fused Multi-Head Attention. + + This kernel coordinates multiple specialized warps to perform different phases of the FMHA computation: + 1. Load warp: Loads Q, K, V data from global memory to shared memory using TMA + 2. MMA warp: Performs matrix multiplications (Q*K^T and P*V) + 3. Softmax warps: Compute softmax normalization on attention scores + 4. Correction warps: Apply adjustments to intermediate results + 5. Epilogue warp: Handles final output transformation and storage + + The kernel implements a complex pipeline with overlapping computation and memory operations, + using tensor memory access (TMA) for efficient data loading, warp specialization for different + computation phases, and optional attention masking. + """ + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + # Prefetch tma descriptor + if warp_idx == 0: + for tma_atom in (tma_atom_Q, tma_atom_K, tma_atom_V, tma_atom_O): + if const_expr(tma_atom is not None): + cpasync.prefetch_descriptor(tma_atom) + + cta_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), (tiled_mma_qk.thr_id.shape,) + ) + # Setup cta/thread coordinates + bidx, _, _ = cute.arch.block_idx() + if const_expr(cute.size(tiled_mma_qk.thr_id.shape) == 1): + mma_tile_coord_v = 0 + else: + mma_tile_coord_v = bidx % cute.size(tiled_mma_qk.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + + # Alloc + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=int(NamedBarrierFwdSm100.TmemPtr), + num_threads=cute.arch.WARP_SIZE * len( + (self.mma_warp_id, + *self.softmax0_warp_ids, + *self.softmax1_warp_ids, + *self.correction_warp_ids) + ), + ) + # Tensor memory dealloc barrier init + tmem = cutlass.utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.mma_warp_id, + is_two_cta=self.use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + ThreadCooperativeGroup = partial(pipeline.CooperativeGroup, pipeline.Agent.Thread) + mma_warp = ThreadCooperativeGroup(len([self.mma_warp_id])) + tma_warp = ThreadCooperativeGroup(1) + load_threads = ThreadCooperativeGroup(len(self.load_warp_ids) * cute.arch.WARP_SIZE) + softmax_warps = ThreadCooperativeGroup(len(self.softmax0_warp_ids)) + softmax_threads = ThreadCooperativeGroup(cute.arch.WARP_SIZE * len(self.softmax0_warp_ids)) + # softmax_threads = ThreadCooperativeGroup(cute.arch.WARP_SIZE) + correction_threads = ThreadCooperativeGroup( + cute.arch.WARP_SIZE * len(self.correction_warp_ids) + ) + # correction_threads = ThreadCooperativeGroup(cute.arch.WARP_SIZE) + softmax_correction_threads = ThreadCooperativeGroup( + cute.arch.WARP_SIZE * len(self.softmax0_warp_ids + self.correction_warp_ids) + ) + epilogue_threads = ThreadCooperativeGroup(cute.arch.WARP_SIZE * len(self.epilogue_warp_ids)) + # For UMMA-bridging pipelines: the non-MMA side spans both CTAs in the cluster, + # so the thread count must include warps from both CTAs. + softmax_warps_cluster = ThreadCooperativeGroup( + len(self.softmax0_warp_ids) * self.cta_group_size + ) + correction_threads_cluster = ThreadCooperativeGroup( + cute.arch.WARP_SIZE * len(self.correction_warp_ids) * self.cta_group_size + ) + softmax_correction_threads_cluster = ThreadCooperativeGroup( + cute.arch.WARP_SIZE * len(self.softmax0_warp_ids + self.correction_warp_ids) * self.cta_group_size + ) + if const_expr(self.use_tma_Q): + pipeline_q = pipeline_custom.PipelineTmaUmma.create( + barrier_storage=storage.mbar_load_Q.data_ptr(), + num_stages=self.q_stage, + producer_group=tma_warp, + consumer_group=mma_warp, + tx_count=self.tma_copy_bytes["Q"], + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + else: + pipeline_q = pipeline_custom.PipelineAsyncUmma.create( + barrier_storage=storage.mbar_load_Q.data_ptr(), + num_stages=self.q_stage, + producer_group=load_threads, + consumer_group=mma_warp, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + if const_expr(self.use_tma_KV): + pipeline_kv = pipeline_custom.PipelineTmaUmma.create( + barrier_storage=storage.mbar_load_KV.data_ptr(), + num_stages=self.kv_stage, + producer_group=tma_warp, + consumer_group=mma_warp, + tx_count=self.tma_copy_bytes["K"], + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + else: + pipeline_kv = pipeline.PipelineAsyncUmma.create( + barrier_storage=storage.mbar_load_KV.data_ptr(), + num_stages=self.kv_stage, + producer_group=load_threads, + consumer_group=mma_warp, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + # This pipeline is not the typical producer-consumer pipeline. The "producer" mma warp + # uses it to signal that S is ready, and the softmax threads wait for S to be ready. + # When softmax threads write P to tmem and the correction threads have rescaled O, they + # signal as "consumer". The mma warp then waits for that signal to do the P @ V gemm. + pipeline_s_p_o = pipeline_custom.PipelineUmmaAsync.create( + barrier_storage=storage.mbar_S_full_P_full_O_rescaled.data_ptr(), + num_stages=self.num_sp_pipeline_stages, + producer_group=mma_warp, + consumer_group=softmax_correction_threads_cluster, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + pipeline_p_lastsplit = pipeline_custom.PipelineAsyncUmma.create( + barrier_storage=storage.mbar_P_full_lastsplit.data_ptr(), + num_stages=self.num_sp_pipeline_stages, + producer_group=softmax_warps_cluster, + consumer_group=mma_warp, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + # MMA warp uses this to signal to the correction warps that O is ready. + pipeline_o_acc = pipeline_custom.PipelineUmmaAsync.create( + barrier_storage=storage.mbar_O_full.data_ptr(), + num_stages=self.num_compute_slots, + producer_group=mma_warp, + consumer_group=correction_threads_cluster, + cta_layout_vmnk=cta_layout_vmnk, + defer_sync=True, + ) + pipeline_s0_s1_sequence = None + if const_expr(self.s0_s1_barrier and self.num_compute_slots > 1): + # This is not a typical producer-consumer pipeline. We will directly use + # pipeline_s0_s1_sequence.sync_object_full and will not use + # pipeline_s0_s1_sequence.sync_object_empty. + pipeline_s0_s1_sequence = pipeline_custom.PipelineAsync.create( + barrier_storage=storage.mbar_s0_s1_sequence.data_ptr(), + num_stages=2, + producer_group=softmax_threads, + consumer_group=softmax_threads, + defer_sync=True, + ) + pipeline_sm_stats = pipeline_custom.PipelineAsync.create( + barrier_storage=storage.mbar_softmax_stats.data_ptr(), + num_stages=self.num_compute_slots, + producer_group=softmax_threads, + consumer_group=correction_threads, + defer_sync=True, + ) + # Should put the NamedBarrier inside the pipeline class so we'll just have pipeline_sm_stats + sm_stats_barrier = pipeline_custom.NamedBarrier( + barrier_id=int(NamedBarrierFwdSm100.SoftmaxStatsW0), + num_threads=cute.arch.WARP_SIZE * (8 if self.p_in_smem else 2), + ) + smem_p_barrier = pipeline_custom.NamedBarrier( + barrier_id=int(NamedBarrierFwdSm100.SoftmaxSmemP0), + num_threads=cute.arch.WARP_SIZE * len(self.softmax0_warp_ids), + ) + pipeline_o_epi = None + if const_expr(not self.use_correction_warps_for_epi): + pipeline_o_epi = pipeline_custom.PipelineAsync.create( + barrier_storage=storage.mbar_O_epi.data_ptr(), + num_stages=self.q_stage, + producer_group=correction_threads, + consumer_group=epilogue_threads, + defer_sync=True, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=cta_layout_vmnk, is_relaxed=True) + + # Generate smem tensor Q/K/V/O + # (MMA, MMA_Q, MMA_D, PIPE) + sQ = storage.sQ.get_tensor(sQ_layout.outer, swizzle=sQ_layout.inner) + # (MMA, MMA_K, MMA_D, PIPE) + sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner) + # (MMA, MMA_K, MMA_D, PIPE) + # Strip swizzle info to reuse smem + sV = cute.make_tensor(cute.recast_ptr(sK.iterator, sV_layout.inner), sV_layout.outer) + if const_expr(not self.overlap_sO_sQ): + sO = storage.sO.get_tensor(sO_layout.outer, swizzle=sO_layout.inner) + else: + sO = cute.make_tensor(cute.recast_ptr(sQ.iterator, sO_layout.inner, self.o_dtype), sO_layout.outer) + + sScale = storage.sScale.get_tensor( + cute.make_layout(self.num_compute_slots * self.m_block_size * 2) + ) + + thr_mma_qk = tiled_mma_qk.get_slice(mma_tile_coord_v) + thr_mma_pv = tiled_mma_pv.get_slice(mma_tile_coord_v) + + qk_acc_shape = thr_mma_qk.partition_shape_C(self.mma_tiler_qk[:2]) + # This is a fake tensor, by right we need to retrieve tmem_ptr. But we know that we always + # request 512 columns of tmem, so we know that it starts at 0. + if const_expr(self.p_in_smem): + tStS = thr_mma_qk.make_fragment_C(qk_acc_shape) + tStS = cute.make_tensor( + tStS.iterator, + cute.append( + tStS.layout, + cute.make_layout( + (self.s_stage,), stride=(self.n_block_size,) + ), + ), + ) + else: + tStS = thr_mma_qk.make_fragment_C( + cute.append(qk_acc_shape, self.s_stage) + ) + pv_acc_shape = thr_mma_pv.partition_shape_C(self.mma_tiler_pv[:2]) + if const_expr(self.p_in_smem): + tOtO = thr_mma_pv.make_fragment_C(pv_acc_shape) + # The compact codomain of the M64 accumulator layout includes TMEM + # datapath bits. Letting cute.append infer a stage stride would move + # stage 1 across datapaths instead of to the next column range. + tOtO = cute.make_tensor( + tOtO.iterator, + cute.append( + tOtO.layout, + cute.make_layout( + (self.num_compute_slots,), + stride=(self.head_dim_v_padded,), + ), + ), + ) + tOtO_addr = (tOtO.iterator + self.tmem_o_offset[0]).toint() + tOtO_ptr = cute.make_ptr( + self.pv_acc_dtype, + tOtO_addr, + cute.AddressSpace.tmem, + assumed_align=8, + ) + tOtO = cute.make_tensor(tOtO_ptr, tOtO.layout) + else: + tOtO = thr_mma_pv.make_fragment_C( + cute.append(pv_acc_shape, self.num_compute_slots) + ) + tOtO = cute.make_tensor(tOtO.iterator + self.tmem_o_offset[0], tOtO.layout) + sP = None + if const_expr(self.p_in_smem): + sP = storage.sP.get_tensor(tP_layout.outer, swizzle=tP_layout.inner) + tOrP = thr_mma_pv.make_fragment_A(sP) + else: + tP = cute.make_tensor(tStS.iterator, tP_layout.outer) + tOrP = thr_mma_pv.make_fragment_A(tP)[None, None, None, 0] + tP_width_ratio = Float32.width // self.v_dtype.width + tP_stage_stride = ( + self.tmem_p_offset[1] - self.tmem_p_offset[0] + ) * tP_width_ratio + tOrP = cute.make_tensor( + tOrP.iterator + self.tmem_p_offset[0] * tP_width_ratio, + cute.append( + tOrP.layout, + cute.make_layout((self.s_stage,), stride=(tP_stage_stride,)), + ), + ) + + block_info = BlockInfo( + # This is cta_tiler, not mma_tiler_qk, since we move by block by (2 * mma_tiler[0], mma_tiler[1]) + self.cta_tiler[0], + self.cta_tiler[1], + self.is_causal, + self.is_local, + self.is_split_kv, + window_size_left, + window_size_right, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + SeqlenInfoCls = partial( + SeqlenInfoQK.create, + tile_m=self.m_block_size, + tile_n=self.n_block_size, + seqlen_q_static=mQ.shape[0] if const_expr(not self.pack_gqa) else mQ.shape[0][1], + seqlen_k_static=mK.shape[0] + if const_expr(mPageTable is None) + else mK.shape[0] * mPageTable.shape[1], + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + mCuTotalMBlocks=( + blocksparse_tensors.cu_total_m_blocks if blocksparse_tensors is not None else None + ), + mCuBlockIdxOffsets=( + blocksparse_tensors.cu_block_idx_offsets if blocksparse_tensors is not None else None + ), + ) + AttentionMaskCls = self._generate_attention_mask_cls( + window_size_left, window_size_right + ) + # Cluster wait before tensor memory alloc + pipeline_init_wait(cluster_shape_mn=cta_layout_vmnk) + + if const_expr(self.use_clc_scheduler): + clc_response_ptr = storage.clc_response.data_ptr() + clc_mbar_ptr = storage.clc_mbar_ptr.data_ptr() + + clc_pipeline_producer_group = cutlass_pipeline.CooperativeGroup( + cutlass_pipeline.Agent.Thread + ) + num_clc_consumer_warps_per_cta = self.threads_per_cta // cute.arch.WARP_SIZE + # NB on CTA0 warp15 == scheduler on CTA1 == empty but still both consume + num_clc_consumer_warps = num_clc_consumer_warps_per_cta * self.cta_group_size + clc_pipeline_consumer_group = cutlass_pipeline.CooperativeGroup( + cutlass_pipeline.Agent.Thread, cute.arch.WARP_SIZE * num_clc_consumer_warps + ) + + block_idx = cute.arch.block_idx() + clc = ClcState.create( + hw_scheduler=ClcDynamicPersistentTileScheduler.create( + self.tile_scheduler_cls.clc_problem_shape(tile_sched_params), + block_idx, + cute.arch.grid_dim(), + clc_response_ptr, + ), + pipeline=cutlass_pipeline.PipelineClcFetchAsync.create( + barrier_storage=clc_mbar_ptr, + num_stages=self.sched_stages, + producer_group=clc_pipeline_producer_group, + consumer_group=clc_pipeline_consumer_group, + tx_count=16, + cta_layout_vmnk=cta_layout_vmnk, + ), + consumer_state=cutlass_pipeline.make_pipeline_state( + cutlass_pipeline.PipelineUserType.Consumer, self.sched_stages + ), + producer_state=cutlass_pipeline.make_pipeline_state( + cutlass_pipeline.PipelineUserType.Producer, self.sched_stages + ), + ) + tile_scheduler = self.tile_scheduler_cls.create(tile_sched_params, clc=clc) + else: + tile_scheduler = self.tile_scheduler_cls.create(tile_sched_params) + assert isinstance(tile_scheduler, TileSchedulerProtocol), f"tile_scheduler is not a TileSchedulerProtocol: {type(tile_scheduler)}" + + # /////////////////////////////////////////////////////////////////////////////// + # EMPTY / CLC SCHEDULER WARP + # /////////////////////////////////////////////////////////////////////////////// + if const_expr(self.use_clc_scheduler): + if warp_idx == self.clc_scheduler_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + if is_leader_cta: + self.clc_scheduler_warp(tile_scheduler) + else: + self.empty_warp(tile_scheduler) + for i in cutlass.range_constexpr(len(self.empty_warp_ids)): + if warp_idx == self.empty_warp_ids[i] and warp_idx != self.clc_scheduler_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + self.empty_warp(tile_scheduler) + else: + for i in cutlass.range_constexpr(len(self.empty_warp_ids)): + if warp_idx == self.empty_warp_ids[i]: + cute.arch.setmaxregister_decrease(self.num_regs_other) + + # /////////////////////////////////////////////////////////////////////////////// + # LOAD + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx >= self.load_warp_ids[0] and warp_idx <= self.load_warp_ids[-1]: + cute.arch.setmaxregister_decrease(self.num_regs_other) + self.load( + thr_mma_qk, + thr_mma_pv, + mQ, + mK, + mV, + sQ, + sK, + sV, + mPageTable, + tma_atom_Q, + tma_atom_K, + tma_atom_V, + gmem_tiled_copy_Q, + pipeline_q, + pipeline_kv, + block_info, + num_splits, + SeqlenInfoCls, + blocksparse_tensors, + tile_scheduler=tile_scheduler, + ) + + # /////////////////////////////////////////////////////////////////////////////// + # MMA + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_other) + # Alloc tensor memory buffer + tmem.allocate(cute.arch.get_max_tmem_alloc_cols("sm_100")) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) + mma_fn = ( + self.mma_dual_stream_spdb + if const_expr(self.use_sp_double_buffer) + else self.mma_dual_stream + if const_expr(self.num_kv_streams == 2) + else self.mma + ) + mma_fn( + tiled_mma_qk, + tiled_mma_pv, + sQ, + sK, + sV, + sP, + tStS, + tOtO, + tOrP, + pipeline_q, + pipeline_kv, + pipeline_s_p_o, + pipeline_p_lastsplit, + pipeline_o_acc, + is_leader_cta, + block_info, + num_splits, + SeqlenInfoCls, + blocksparse_tensors, + tile_scheduler=tile_scheduler, + ) + # Dealloc the tensor memory buffer + tmem.relinquish_alloc_permit() + tmem_alloc_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + + # /////////////////////////////////////////////////////////////////////////////// + # Epilogue + # /////////////////////////////////////////////////////////////////////////////// + if const_expr(not self.use_correction_warps_for_epi): + if warp_idx >= self.epilogue_warp_ids[0] and warp_idx <= self.epilogue_warp_ids[-1]: + cute.arch.setmaxregister_decrease(self.num_regs_other) + self.epilogue_s2g( + mO, + sO, + gmem_tiled_copy_O, + tma_atom_O, + pipeline_o_epi, + block_info, + num_splits, + SeqlenInfoCls, + mma_tile_coord_v, + blocksparse_tensors=blocksparse_tensors, + tile_scheduler=tile_scheduler, + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Softmax + # /////////////////////////////////////////////////////////////////////////////// + if ( + ( + const_expr(self.num_compute_slots == 2) + and warp_idx <= self.softmax1_warp_ids[-1] + ) + or ( + const_expr(self.num_compute_slots == 1) + and warp_idx <= self.softmax0_warp_ids[-1] + ) + ): + # increase register after decreasing + cute.arch.setmaxregister_increase(self.num_regs_softmax) + # sync with mma warp before retrieving tmem ptr + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) + softmax_loop = partial( + self.softmax_loop, + softmax_scale_log2=softmax_scale_log2, + softmax_scale=softmax_scale, + descale_tensors=descale_tensors, + thr_mma_qk=thr_mma_qk, + sP=sP, + sScale=sScale, + mLSE=mLSE, + pipeline_s_p_o=pipeline_s_p_o, + pipeline_p_lastsplit=pipeline_p_lastsplit, + pipeline_sm_stats=pipeline_sm_stats, + sm_stats_barrier=sm_stats_barrier, + smem_p_barrier=smem_p_barrier, + pipeline_s0_s1_sequence=pipeline_s0_s1_sequence, + learnable_sink=learnable_sink, + block_info=block_info, + num_splits=num_splits, + SeqlenInfoCls=SeqlenInfoCls, + AttentionMaskCls=AttentionMaskCls, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + head_divmod=head_divmod, + blocksparse_tensors=blocksparse_tensors, + tile_scheduler=tile_scheduler, + ) + + if const_expr(not self.s0_s1_barrier): + stage = Int32( + 0 + if const_expr(self.num_compute_slots == 1) + or warp_idx < self.softmax1_warp_ids[0] + else 1 + ) + softmax_loop(stage=stage, tStS=tStS) + else: + # If there's s0_s1_barrier, it's faster to have 2 WGs having different code + if warp_idx < self.softmax1_warp_ids[0]: + softmax_loop(stage=0, tStS=tStS) + if warp_idx < self.correction_warp_ids[0] and warp_idx >= self.softmax1_warp_ids[0]: + softmax_loop(stage=1, tStS=tStS) + + tmem_alloc_barrier.arrive() + + # /////////////////////////////////////////////////////////////////////////////// + # Correction + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx >= self.correction_warp_ids[0] and warp_idx < self.mma_warp_id: + cute.arch.setmaxregister_decrease(self.num_regs_correction) + # sync with mma warp before retrieving tmem ptr + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) + correction_fn = ( + self.correction_loop_dual_stream_spdb + if const_expr(self.use_sp_double_buffer) + else self.correction_loop_dual_stream + if const_expr(self.num_kv_streams == 2) + else self.correction_loop + ) + correction_fn( + thr_mma_qk, + thr_mma_pv, + tStS, + tOtO, + sScale, + mO, + mLSE, + sO, + pipeline_s_p_o, + pipeline_o_acc, + pipeline_sm_stats, + sm_stats_barrier, + pipeline_o_epi, + learnable_sink, + descale_tensors, + gmem_tiled_copy_O, + tma_atom_O, + softmax_scale_log2, + block_info, + num_splits, + SeqlenInfoCls, + blocksparse_tensors, + tile_scheduler=tile_scheduler, + ) + tmem_alloc_barrier.arrive() + + return + + @cute.jit + def load( + self, + thr_mma_qk: cute.ThrMma, + thr_mma_pv: cute.ThrMma, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + mPageTable: Optional[cute.Tensor], + tma_atom_Q: Optional[cute.CopyAtom], + tma_atom_K: Optional[cute.CopyAtom], + tma_atom_V: Optional[cute.CopyAtom], + gmem_tiled_copy_Q: Optional[cute.TiledCopy], + pipeline_q: pipeline.PipelineAsync, + pipeline_kv: pipeline.PipelineAsync, + block_info: BlockInfo, + num_splits: Int32, + SeqlenInfoCls: Callable, + blocksparse_tensors: Optional[BlockSparseTensors], + tile_scheduler: TileSchedulerProtocol, + ): + num_load_threads = len(self.load_warp_ids) * cute.arch.WARP_SIZE + tidx = cute.arch.thread_idx()[0] % num_load_threads + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + issue_kv_for_this_warp = ( + const_expr(not self.use_tma_KV or len(self.load_warp_ids) == 1) or + warp_idx == self.load_warp_ids[0] + ) + issue_q_for_this_warp = ( + const_expr(not self.use_tma_Q or len(self.load_warp_ids) == 1) or + warp_idx == self.load_warp_ids[0] + ) + q_producer_phase = Int32(1) + kv_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.kv_stage + ) + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + seqlen = SeqlenInfoCls(batch_idx) + mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx] + + head_idx_kv = ( + head_idx // self.qhead_per_kvhead if const_expr(not self.pack_gqa) else head_idx + ) + if const_expr(mPageTable is None): + if const_expr(not seqlen.has_cu_seqlens_k): + mK_cur, mV_cur = [t[None, None, head_idx_kv, batch_idx] for t in (mK, mV)] + else: + mK_cur = cute.domain_offset((seqlen.offset_k, 0), mK[None, None, head_idx_kv]) + mV_cur = cute.domain_offset((0, seqlen.offset_k), mV[None, None, head_idx_kv]) + gK = cute.local_tile(mK_cur, cute.select(self.mma_tiler_qk, mode=[1, 2]), (None, 0)) + gV = cute.local_tile(mV_cur, cute.select(self.mma_tiler_pv, mode=[1, 2]), (0, None)) + else: + # Need to keep batch coord None since we'll index into it with page idx + mK_cur, mV_cur = [t[None, None, head_idx_kv, None] for t in (mK, mV)] + gK = cute.local_tile( + mK_cur, cute.select(self.mma_tiler_qk, mode=[1, 2]), (None, 0, None) + ) + gV = cute.local_tile( + mV_cur, cute.select(self.mma_tiler_pv, mode=[1, 2]), (0, None, None) + ) + tSgK = thr_mma_qk.partition_B(gK) + tOgV = thr_mma_pv.partition_B(gV) + if const_expr(self.use_tma_Q): + tiler_gQ = ((self.mma_tiler_qk[0] * self.q_stage), self.head_dim_padded) + gQ = cute.local_tile(mQ_cur, tiler_gQ, (m_block, 0)) # (128 * 2, 128) + gQ = layout_utils.select( + cute.flat_divide(gQ, (self.mma_tiler_qk[0],)), mode=[0, 2, 1] + ) # (128, 128, 2) + tSgQ = thr_mma_qk.partition_A(gQ) + load_Q_fn, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_Q, 0, cute.make_layout(1), tSgQ, sQ + ) + load_Q = partial(self.load_Q, load_Q_fn, pipeline_q=pipeline_q, phase=q_producer_phase) + else: + assert gmem_tiled_copy_Q is not None + load_Q = partial( + self.load_Q_non_tma, + mQ_cur, + sQ, + gmem_tiled_copy_Q, + pipeline_q, + tidx, + seqlen.seqlen_q, + m_block, + phase=q_producer_phase, + ) + + if const_expr(self.use_tma_KV): + tKsK, tKgK = cpasync.tma_partition( + tma_atom_K, + 0, # no multicast + cute.make_layout(1), + cute.group_modes(sK, 0, 3), + cute.group_modes(tSgK, 0, 3), + ) + tVsV, tVgV = cpasync.tma_partition( + tma_atom_V, + 0, # no multicast + cute.make_layout(1), + cute.group_modes(sV, 0, 3), + cute.group_modes(tOgV, 0, 3), + ) + paged_kv_manager = None + else: + page_size = mK.shape[0] + paged_kv_manager = PagedKVManager.create( + mPageTable, + mK, + mV, + FastDivmodDivisor(page_size), + batch_idx, + head_idx_kv, + tidx, + seqlen.seqlen_k, + 0, # leftpad_k + self.n_block_size, + self.head_dim_padded, + self.head_dim_v_padded, + num_load_threads, + mK.element_type, + ) + tKsK, tKgK = None, None + tVsV, tVgV = None, None + + load_K = partial( + self.load_KV, + tma_atom_K, + tKgK, + tKsK, + paged_kv_manager, + sK, + pipeline_kv=pipeline_kv, + K_or_V="K", + ) + load_V = partial( + self.load_KV, + tma_atom_V, + tVgV, + tVsV, + paged_kv_manager, + sV, + pipeline_kv=pipeline_kv, + K_or_V="V", + ) + if const_expr(not self.use_block_sparsity): + n_block_min, n_block_max = block_info.get_n_block_min_max( + seqlen, m_block, split_idx, num_splits + ) + if const_expr(not self.is_split_kv) or n_block_min < n_block_max: + n_block_first = n_block_max - 1 if n_block_max > 0 else 0 + page_idx = ( + mPageTable[batch_idx, n_block_first] + if const_expr(mPageTable is not None and self.use_tma_KV) + else None + ) + if const_expr(not self.use_tma_KV): + paged_kv_manager.load_page_table(n_block_first) + if issue_kv_for_this_warp: + load_K(block=n_block_max - 1, producer_state=kv_producer_state, page_idx=page_idx) # K0 + # load_K(block=n_block_max - 1, producer_state=kv_producer_state, page_idx=page_idx, extra_tx_count=self.tma_copy_bytes["Q"]) # K0 + if issue_q_for_this_warp: + load_Q(block=0, stage=0) + if issue_kv_for_this_warp: + kv_producer_state.advance() + if const_expr(self.q_stage == 2) and issue_q_for_this_warp: + load_Q(block=1, stage=1) + q_producer_phase ^= 1 + if issue_kv_for_this_warp: + load_V(block=n_block_max - 1, producer_state=kv_producer_state, page_idx=page_idx) # V0 + kv_producer_state.advance() + for i in cutlass.range(n_block_max - 1 - n_block_min, unroll=1): + n_block = n_block_max - 2 - i + page_idx = ( + mPageTable[batch_idx, n_block] + if const_expr(mPageTable is not None and self.use_tma_KV) + else None + ) + if const_expr(not self.use_tma_KV): + paged_kv_manager.load_page_table(n_block) + # if cute.arch.thread_idx()[0] % 32 == 0: cute.printf("n_block = {}, page_idx = {}", n_block, page_idx) + if issue_kv_for_this_warp: + load_K(block=n_block, producer_state=kv_producer_state, page_idx=page_idx) # Ki + kv_producer_state.advance() + load_V(block=n_block, producer_state=kv_producer_state, page_idx=page_idx) # Vi + kv_producer_state.advance() + + else: + if const_expr(self.num_kv_streams == 2): + dual_stream_load = ( + produce_block_sparse_loads_dual_stream_spdb_sm100 + if const_expr(self.use_sp_double_buffer) + else produce_block_sparse_loads_dual_stream_sm100 + ) + kv_producer_state, q_producer_phase = dual_stream_load( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + seqlen, + split_idx, + num_splits, + kv_producer_state, + load_Q, + load_K, + load_V, + q_producer_phase, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + self.q_subtile_factor, + ) + else: + kv_producer_state, q_producer_phase = produce_block_sparse_loads_sm100( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + seqlen, + split_idx, + num_splits, + kv_producer_state, + load_Q, + load_K, + load_V, + pipeline_kv, + self.q_stage, + q_producer_phase, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + self.q_subtile_factor, + ) + + + work_tile = tile_scheduler.advance_to_next_work() + # End of persistent scheduler loop + + if issue_kv_for_this_warp: + pipeline_kv.producer_tail(kv_producer_state) + # This is equivalent to pipeline_q.producer_tail for the TMA-Q producer warp. + if issue_q_for_this_warp: + pipeline_q.producer_acquire_w_index_phase(self.q_stage - 1, q_producer_phase) + + @cute.jit + def mma( + self, + tiled_mma_qk: cute.ThrMma, + tiled_mma_pv: cute.ThrMma, + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + sP: cute.Tensor | None, + tStS: cute.Tensor, + tOtO: cute.Tensor, + tOrP: cute.Tensor, + pipeline_q: pipeline.PipelineAsync, + pipeline_kv: pipeline.PipelineAsync, + pipeline_s_p_o: pipeline.PipelineAsync, + pipeline_p_lastsplit: pipeline.PipelineAsync, + pipeline_o_acc: pipeline.PipelineAsync, + is_leader_cta: Boolean, + block_info: BlockInfo, + num_splits: Int32, + SeqlenInfoCls: Callable, + blocksparse_tensors: Optional[BlockSparseTensors], + tile_scheduler=None, + ): + tSrQ = tiled_mma_qk.make_fragment_A(sQ) + tSrK = tiled_mma_qk.make_fragment_B(sK) + tOrV = tiled_mma_pv.make_fragment_B(sV) + if const_expr(self.q_stage == 2): + tSrQs = (tSrQ[None, None, None, 0], tSrQ[None, None, None, 1]) + else: + tSrQs = (tSrQ[None, None, None, 0],) + + qk_mma_op, pv_mma_op = tiled_mma_qk.op, tiled_mma_pv.op + qk_mma_idesc, pv_mma_idesc = sm100_desc.mma_op_to_idesc(qk_mma_op), sm100_desc.mma_op_to_idesc(pv_mma_op) + qk_mma_kind = sm100_utils._tcgen05_mma_kind(qk_mma_op) + q_smem_base = sm100_desc.smem_desc_base_from_tensor(sQ, sm100_desc.Major.K) + k_smem_base = sm100_desc.smem_desc_base_from_tensor(sK, sm100_desc.Major.K) + v_smem_base = sm100_desc.smem_desc_base_from_tensor(sV, sm100_desc.Major.MN) + q_smem_start = [sm100_desc.make_smem_desc_start_addr(sQ[None, None, None, stage].iterator) for stage in range(self.q_stage)] + + sm100_utils.declare_ptx_smem_desc(q_smem_start[self.q_stage - 1], q_smem_base, tSrQ[None, None, None, 0].layout, var_name_prefix="fa_fwd_q_smem_desc") + sm100_utils.declare_ptx_idesc(qk_mma_op, var_name="fa_fwd_qk_mma_idesc") + sm100_utils.declare_ptx_idesc(pv_mma_op, var_name="fa_fwd_pv_mma_idesc") + + sQ_stage_stride = (sQ.layout.stride[-1] * sQ.element_type.width // 8) >> 4 + if const_expr(self.q_stage == 1): + sQ_stage_stride = 0 + gemm_Si = [ + partial( + # sm100_utils.gemm_ptx_precomputed, + # self.tmem_s_offset[stage], + # smem_desc_start_a=q_smem_start[stage], + # idesc=qk_mma_idesc, + # smem_desc_base_a=q_smem_base, + # smem_desc_base_b=k_smem_base, + # tCrA_layout=tSrQ[None, None, None, 0].layout, + sm100_utils.gemm_ptx_precomputed_varname, + self.tmem_s_offset[slot], + # idesc=qk_mma_idesc, + smem_desc_base_b=k_smem_base, + tCrB_layout=tSrK[None, None, None, 0].layout, + smem_var_name_prefix="fa_fwd_q_smem_desc", + idesc_var_name="fa_fwd_qk_mma_idesc", + kind=qk_mma_kind, + smem_offset=( + -sQ_stage_stride + if slot // self.num_kv_streams == 0 + else sQ_stage_stride + ), + zero_init=True, + cta_group=self.cta_group_size, + ) + for slot in range(self.num_compute_slots) + ] + # gemm_Si = [ + # partial( + # sm100_utils.gemm, + # tiled_mma_qk, + # tStS[None, None, None, stage], + # tCrA=tSrQ[None, None, None, stage], + # zero_init=True, + # ) + # for stage in range(self.q_stage) + # ] + gemm_Pi = [ + partial( + # sm100_utils.gemm_ptx_precomputed, + sm100_utils.gemm_ptx_partial, + pv_mma_op, + self.tmem_o_offset[slot], + tOrP[None, None, None, slot], + sA=sP[None, None, None, slot] if self.p_in_smem else None, + split_arrive=self.split_P_arrive if self.split_P_arrive > 0 else None, + # smem_desc_start_a=tOrP[None, None, None, stage].iterator.toint(), + # smem_desc_start_a=self.tmem_p_offset[stage], + # idesc=pv_mma_idesc, + # smem_desc_base_a=None, + # smem_desc_base_b=v_smem_base, + # tCrA_layout=tOrP[None, None, None, 0].layout, + # tCrB_layout=tOrV[None, None, None, 0].layout + cta_group=self.cta_group_size, + ) + for slot in range(self.num_compute_slots) + ] + # gemm_Pi = [ + # partial( + # sm100_utils.gemm, tOtO[None, None, None, stage], tCrA=tOrP[None, None, None, stage] + # ) + # for stage in range(self.q_stage) + # ] + + mma_q_consumer_phase = Int32(0) + mma_kv_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.kv_stage + ) + P_full_O_rescaled_phase = Int32(0) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + seqlen = SeqlenInfoCls(batch_idx) + + block_iter_count = Int32(0) + process_tile = False + + if const_expr(self.use_block_sparsity): + block_iter_count = get_total_block_count( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + split_idx, + num_splits, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + self.q_subtile_factor, + seqlen_info=seqlen, + ) + process_tile = block_iter_count > Int32(0) + else: + n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block, split_idx, num_splits) + block_iter_count = n_block_max - n_block_min + if const_expr(not self.is_split_kv): + process_tile = True + else: + process_tile = n_block_min < n_block_max + + if process_tile and is_leader_cta: + for stage in cutlass.range_constexpr(self.num_compute_slots): + # GEMM_QK00 (Q0 * K0 -> S0) or GEMM_QK01 (Q1 * K0 -> S1) + # 1. wait for Q0 / Q1 + if const_expr(stage % self.num_kv_streams == 0): + pipeline_q.consumer_wait_w_index_phase( + stage // self.num_kv_streams, mma_q_consumer_phase + ) + # 2. wait for K0 + if const_expr(stage == 0): + pipeline_kv.consumer_wait(mma_kv_consumer_state) + Ki_index, Ki_phase = mma_kv_consumer_state.index, mma_kv_consumer_state.phase + tSrKi = tSrK[None, None, None, Ki_index] + # We don't need to acquire empty S0 / S1. + # For the first iteration, we don't need to wait as we're guaranteed S0 / S1 + # are empty. For subsequent iterations, the wait happened at the end + # of the while loop. + # 3. gemm + # sm100_utils.gemm(tiled_mma_qk, tStS[None, None, None, stage], tSrQ[None, None, None, stage], tSrKi, zero_init=True) + sK_cur = sK[None, None, None, Ki_index] + if const_expr(self.uneven_kv_smem): + sK_cur = self.offset_kv_smem(sK_cur, Ki_index, Ki_phase) + # gemm_Si[stage](tCrB=tSrKi, sB=sK_cur) + gemm_Si[stage]( + smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sK_cur.iterator) + ) + # gemm_Si[stage](tCrB=tSrKi) + # 4. release S0 / S1 + pipeline_s_p_o.producer_commit_w_index(stage) + mma_q_consumer_phase ^= 1 + # 5. release K0 + pipeline_kv.consumer_release(mma_kv_consumer_state) + mma_kv_consumer_state.advance() + # End of GEMM (Q1 * K0 -> S1) + # Note: Q0 & Q1 are still needed in the seqlen_kv loop + # so we need to release them after the seqlen_kv loop + + # O hasn't been accumulated yet, its first MMA calculation doesn't need to accumulate + block_loop_count = block_iter_count - 1 + O_should_accumulate = False + for i in cutlass.range(block_loop_count, unroll=1): + # GEMM_PV00 (P0 * V0 -> O0_partial), O0 needs to be accumulated in the seqlen_kv loop + # 1. wait for V0 + pipeline_kv.consumer_wait(mma_kv_consumer_state) + mma_kv_release_state = mma_kv_consumer_state.clone() + Vi_index, Vi_phase = mma_kv_consumer_state.index, mma_kv_consumer_state.phase + tOrVi = tOrV[None, None, None, Vi_index] + for stage in cutlass.range_constexpr(self.num_compute_slots): + # 2. acquire corrected O0/O1_partial and P0 / P1 + # For the first iteration in this work tile, waiting for O0/O1_partial + # means that the correction warps has finished reading tO during + # the last iteration of the previous work tile. + pipeline_s_p_o.producer_acquire_w_index_phase(stage, P_full_O_rescaled_phase) + # 3. gemm + # sm100_utils.gemm(tiled_mma_pv, tOtO0, tOrP0, tOrVi, zero_init=True) + # gemm_Pi[stage](tCrB=tOrVi, sB=sV[None, None, None, Vi_index], zero_init=not O_should_accumulate) + sV_cur = sV[None, None, None, Vi_index] + if const_expr(self.uneven_kv_smem): + sV_cur = self.offset_kv_smem(sV_cur, Vi_index, Vi_phase) + gemm_Pi[stage]( + tCrB=tOrVi, + sB=sV_cur, + # smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sV_cur.iterator), + zero_init=not O_should_accumulate, + mbar_ptr=pipeline_p_lastsplit.sync_object_full.get_barrier(stage) if self.split_P_arrive > 0 else None, + mbar_phase=P_full_O_rescaled_phase, + ) + # Don't need to signal O_full to the correction warps since the + # correction warps wait for the softmax warps anyway. By the time the softmax + # warps finished, S_i for the next iteration must have been done, so O_i-1 + # must have been done as well. + # pipeline_o_acc.producer_commit_w_index(stage) + # 4. release V(i-1) + if const_expr(stage == self.num_compute_slots - 1): + pipeline_kv.consumer_release(mma_kv_release_state) + mma_kv_release_state.advance() + # End of GEMM_PV00 (P0 * V0 -> O0_partial) + + # GEMM_QK0i (Q0 * Ki -> S0) + # 1. wait for Ki + if const_expr(stage == 0): + mma_kv_consumer_state.advance() + pipeline_kv.consumer_wait(mma_kv_consumer_state) + Ki_index, Ki_phase = mma_kv_consumer_state.index, mma_kv_consumer_state.phase + # 2. gemm + # Don't need to wait for the softmax warp to have finished reading the previous + # Si, since this gemm is scheduled after the PV gemm, which guaranteed that Si + # has been read and Pi has been written. + # sm100_utils.gemm(tiled_mma_qk, tStS[None, None, None, stage], tSrQ[None, None, None, stage], tSrK[None, None, None, Ki_index], zero_init=True) + sK_cur = sK[None, None, None, Ki_index] + if const_expr(self.uneven_kv_smem): + sK_cur = self.offset_kv_smem(sK_cur, Ki_index, Ki_phase) + # gemm_Si[stage](tCrB=tSrK[None, None, None, Ki_index], sB=sK_cur) + gemm_Si[stage]( + smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sK_cur.iterator) + ) + # gemm_Si[stage](tCrB=tSrK[None, None, None, Ki_index]) + # 3. release S0 / S1 + pipeline_s_p_o.producer_commit_w_index(stage) + # End of GEMM_QK0i (Q0 * Ki -> S0) + # 4. release Ki + pipeline_kv.consumer_release(mma_kv_consumer_state) + mma_kv_consumer_state.advance() + P_full_O_rescaled_phase ^= 1 + O_should_accumulate = True + # End of seqlen_kv loop + + # release Q0 & Q1 + for stage in cutlass.range(self.q_stage): + pipeline_q.consumer_release_w_index(stage) + + # GEMM_PV00 (P0 * V0 -> O0_partial), O0 needs to be accumulated in the seqlen_kv loop + # 1. wait for V0 + pipeline_kv.consumer_wait(mma_kv_consumer_state) + Vi_index, Vi_phase = mma_kv_consumer_state.index, mma_kv_consumer_state.phase + tOrVi = tOrV[None, None, None, Vi_index] + for stage in cutlass.range_constexpr(self.num_compute_slots): + # 2. acquire corrected Oi_partial and Pi + pipeline_s_p_o.producer_acquire_w_index_phase(stage, P_full_O_rescaled_phase) + # 3. gemm + # sm100_utils.gemm(tiled_mma_pv, tOtO0, tOrP0, tOrVi, zero_init=True) + # gemm_Pi[stage](tCrB=tOrVi, sB=sV[None, None, None, Vi_index], zero_init=not O_should_accumulate) + sV_cur = sV[None, None, None, Vi_index] + if const_expr(self.uneven_kv_smem): + sV_cur = self.offset_kv_smem(sV_cur, Vi_index, Vi_phase) + gemm_Pi[stage]( + tCrB=tOrVi, + sB=sV_cur, + # smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sV_cur.iterator), + zero_init=not O_should_accumulate, + mbar_ptr=pipeline_p_lastsplit.sync_object_full.get_barrier(stage) if self.split_P_arrive > 0 else None, + mbar_phase=P_full_O_rescaled_phase, + ) + # 4. release accumulated O0_partial + # We do need O_full here since for the last tile, by the time the softmax warp + # has signaled to the correction warps, the softmax warp has just finished + # computing the row sum of the current tile. It does not guarantee that the 1st + # tile of the next work tile has been computed yet. + pipeline_o_acc.producer_commit_w_index(stage) + # End of GEMM_PV00 (P0 * V0 -> O0_partial) + P_full_O_rescaled_phase ^= 1 + # 5. release Vi_end + pipeline_kv.consumer_release(mma_kv_consumer_state) + mma_kv_consumer_state.advance() + # End of GEMM_PV1(i_end) (P1 * Vi_end -> O1) + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() + # End of persistent scheduler loop + + # We don't need pipeline_s_p_o.producer_tail() since there's no dangling mbarrier at the end + # pipeline_s_p_o.producer_acquire_w_index_phase(self.q_stage - 1, P_full_O_rescaled_phase) + # We don't need pipeline_o_acc.producer_tail() since we don't call + # pipeline_o_acc.producer_acquire() inside the loop. + + @cute.jit + def mma_dual_stream( + self, + tiled_mma_qk: cute.ThrMma, + tiled_mma_pv: cute.ThrMma, + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + sP: cute.Tensor | None, + tStS: cute.Tensor, + tOtO: cute.Tensor, + tOrP: cute.Tensor, + pipeline_q: pipeline.PipelineAsync, + pipeline_kv: pipeline.PipelineAsync, + pipeline_s_p_o: pipeline.PipelineAsync, + pipeline_p_lastsplit: pipeline.PipelineAsync, + pipeline_o_acc: pipeline.PipelineAsync, + is_leader_cta: Boolean, + block_info: BlockInfo, + num_splits: Int32, + SeqlenInfoCls: Callable, + blocksparse_tensors: Optional[BlockSparseTensors], + tile_scheduler=None, + ): + """MMA consumer for the q-stage-1, two-KV-stream sparse schedule. + + The load warp writes one FIFO pair as K(even), K(odd), V(even), + V(odd). The two score/output slots share Q but keep independent P/O + barrier phases so an odd tail does not advance the inactive stream. + """ + assert const_expr(self.q_stage == 1) + assert const_expr(self.num_kv_streams == 2) + assert const_expr(self.use_block_sparsity) + + tSrQ = tiled_mma_qk.make_fragment_A(sQ) + tSrK = tiled_mma_qk.make_fragment_B(sK) + tOrV = tiled_mma_pv.make_fragment_B(sV) + + qk_mma_op, pv_mma_op = tiled_mma_qk.op, tiled_mma_pv.op + qk_mma_kind = sm100_utils._tcgen05_mma_kind(qk_mma_op) + q_smem_base = sm100_desc.smem_desc_base_from_tensor(sQ, sm100_desc.Major.K) + k_smem_base = sm100_desc.smem_desc_base_from_tensor(sK, sm100_desc.Major.K) + q_smem_start = sm100_desc.make_smem_desc_start_addr( + sQ[None, None, None, 0].iterator + ) + + sm100_utils.declare_ptx_smem_desc( + q_smem_start, + q_smem_base, + tSrQ[None, None, None, 0].layout, + var_name_prefix="fa_fwd_q_smem_desc", + ) + sm100_utils.declare_ptx_idesc(qk_mma_op, var_name="fa_fwd_qk_mma_idesc") + sm100_utils.declare_ptx_idesc(pv_mma_op, var_name="fa_fwd_pv_mma_idesc") + + gemm_Si = [ + partial( + sm100_utils.gemm_ptx_precomputed_varname, + self.tmem_s_offset[slot], + smem_desc_base_b=k_smem_base, + tCrB_layout=tSrK[None, None, None, 0].layout, + smem_var_name_prefix="fa_fwd_q_smem_desc", + idesc_var_name="fa_fwd_qk_mma_idesc", + kind=qk_mma_kind, + smem_offset=0, + zero_init=True, + cta_group=self.cta_group_size, + ) + for slot in range(2) + ] + gemm_Pi = [ + partial( + sm100_utils.gemm_ptx_partial, + pv_mma_op, + self.tmem_o_offset[slot], + tOrP[None, None, None, slot], + sA=sP[None, None, None, slot] if self.p_in_smem else None, + split_arrive=self.split_P_arrive if self.split_P_arrive > 0 else None, + cta_group=self.cta_group_size, + ) + for slot in range(2) + ] + + mma_q_consumer_phase = Int32(0) + mma_kv_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.kv_stage + ) + p_phase0 = Int32(0) + p_phase1 = Int32(0) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + seqlen = SeqlenInfoCls(batch_idx) + stream0_count, stream1_count = get_dual_stream_block_counts( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + split_idx, + num_splits, + 1, + self.q_subtile_factor, + seqlen, + ) + + if stream0_count > Int32(0) and is_leader_cta: + pipeline_q.consumer_wait_w_index_phase(0, mma_q_consumer_phase) + o0_should_accumulate = False + o1_should_accumulate = False + + # Initial score MMAs consume the two prefetched K blocks. + pipeline_kv.consumer_wait(mma_kv_consumer_state) + k0_idx, k0_phase = ( + mma_kv_consumer_state.index, + mma_kv_consumer_state.phase, + ) + sK0 = sK[None, None, None, k0_idx] + if const_expr(self.uneven_kv_smem): + sK0 = self.offset_kv_smem(sK0, k0_idx, k0_phase) + gemm_Si[0]( + smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sK0.iterator) + ) + pipeline_s_p_o.producer_commit_w_index(0) + pipeline_kv.consumer_release(mma_kv_consumer_state) + mma_kv_consumer_state.advance() + + if stream1_count > Int32(0): + pipeline_kv.consumer_wait(mma_kv_consumer_state) + k1_idx, k1_phase = ( + mma_kv_consumer_state.index, + mma_kv_consumer_state.phase, + ) + sK1 = sK[None, None, None, k1_idx] + if const_expr(self.uneven_kv_smem): + sK1 = self.offset_kv_smem(sK1, k1_idx, k1_phase) + gemm_Si[1]( + smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sK1.iterator) + ) + pipeline_s_p_o.producer_commit_w_index(1) + pipeline_kv.consumer_release(mma_kv_consumer_state) + mma_kv_consumer_state.advance() + + if stream0_count == Int32(1): + pipeline_q.consumer_release_w_index(0) + mma_q_consumer_phase ^= 1 + + for stream_idx in cutlass.range(stream0_count, unroll=1): + # PV(s0, i), followed immediately by QK(s0, i + 1). + pipeline_kv.consumer_wait(mma_kv_consumer_state) + v0_idx, v0_phase = ( + mma_kv_consumer_state.index, + mma_kv_consumer_state.phase, + ) + sV0 = sV[None, None, None, v0_idx] + if const_expr(self.uneven_kv_smem): + sV0 = self.offset_kv_smem(sV0, v0_idx, v0_phase) + pipeline_s_p_o.producer_acquire_w_index_phase(0, p_phase0) + gemm_Pi[0]( + tCrB=tOrV[None, None, None, v0_idx], + sB=sV0, + zero_init=not o0_should_accumulate, + mbar_ptr=( + pipeline_p_lastsplit.sync_object_full.get_barrier(0) + if self.split_P_arrive > 0 + else None + ), + mbar_phase=p_phase0, + ) + if stream_idx == stream0_count - Int32(1): + pipeline_o_acc.producer_commit_w_index(0) + pipeline_kv.consumer_release(mma_kv_consumer_state) + mma_kv_consumer_state.advance() + p_phase0 ^= 1 + o0_should_accumulate = True + + next_stream_idx = stream_idx + Int32(1) + if next_stream_idx < stream0_count: + pipeline_kv.consumer_wait(mma_kv_consumer_state) + k0_next_idx, k0_next_phase = ( + mma_kv_consumer_state.index, + mma_kv_consumer_state.phase, + ) + sK0_next = sK[None, None, None, k0_next_idx] + if const_expr(self.uneven_kv_smem): + sK0_next = self.offset_kv_smem( + sK0_next, k0_next_idx, k0_next_phase + ) + gemm_Si[0]( + smem_desc_start_b=sm100_desc.make_smem_desc_start_addr( + sK0_next.iterator + ) + ) + pipeline_s_p_o.producer_commit_w_index(0) + pipeline_kv.consumer_release(mma_kv_consumer_state) + mma_kv_consumer_state.advance() + if ( + next_stream_idx == stream0_count - Int32(1) + and stream0_count > stream1_count + ): + pipeline_q.consumer_release_w_index(0) + mma_q_consumer_phase ^= 1 + + # PV(s1, i), followed immediately by QK(s1, i + 1). + if stream_idx < stream1_count: + pipeline_kv.consumer_wait(mma_kv_consumer_state) + v1_idx, v1_phase = ( + mma_kv_consumer_state.index, + mma_kv_consumer_state.phase, + ) + sV1 = sV[None, None, None, v1_idx] + if const_expr(self.uneven_kv_smem): + sV1 = self.offset_kv_smem(sV1, v1_idx, v1_phase) + pipeline_s_p_o.producer_acquire_w_index_phase(1, p_phase1) + gemm_Pi[1]( + tCrB=tOrV[None, None, None, v1_idx], + sB=sV1, + zero_init=not o1_should_accumulate, + mbar_ptr=( + pipeline_p_lastsplit.sync_object_full.get_barrier(1) + if self.split_P_arrive > 0 + else None + ), + mbar_phase=p_phase1, + ) + if stream_idx == stream1_count - Int32(1): + pipeline_o_acc.producer_commit_w_index(1) + pipeline_kv.consumer_release(mma_kv_consumer_state) + mma_kv_consumer_state.advance() + p_phase1 ^= 1 + o1_should_accumulate = True + + if next_stream_idx < stream1_count: + pipeline_kv.consumer_wait(mma_kv_consumer_state) + k1_next_idx, k1_next_phase = ( + mma_kv_consumer_state.index, + mma_kv_consumer_state.phase, + ) + sK1_next = sK[None, None, None, k1_next_idx] + if const_expr(self.uneven_kv_smem): + sK1_next = self.offset_kv_smem( + sK1_next, k1_next_idx, k1_next_phase + ) + gemm_Si[1]( + smem_desc_start_b=sm100_desc.make_smem_desc_start_addr( + sK1_next.iterator + ) + ) + pipeline_s_p_o.producer_commit_w_index(1) + pipeline_kv.consumer_release(mma_kv_consumer_state) + mma_kv_consumer_state.advance() + if ( + next_stream_idx == stream1_count - Int32(1) + and stream0_count == stream1_count + ): + pipeline_q.consumer_release_w_index(0) + mma_q_consumer_phase ^= 1 + + work_tile = tile_scheduler.advance_to_next_work() + + + @cute.jit + def mma_dual_stream_spdb( + self, + tiled_mma_qk: cute.ThrMma, + tiled_mma_pv: cute.ThrMma, + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + sP: cute.Tensor | None, + tStS: cute.Tensor, + tOtO: cute.Tensor, + tOrP: cute.Tensor, + pipeline_q: pipeline.PipelineAsync, + pipeline_kv: pipeline.PipelineAsync, + pipeline_s_p_o: pipeline.PipelineAsync, + pipeline_p_lastsplit: pipeline.PipelineAsync, + pipeline_o_acc: pipeline.PipelineAsync, + is_leader_cta: Boolean, + block_info: BlockInfo, + num_splits: Int32, + SeqlenInfoCls: Callable, + blocksparse_tensors: Optional[BlockSparseTensors], + tile_scheduler=None, + ): + """MMA consumer for the q-stage-1, two-KV-stream sparse schedule. + + With score double buffering, each stream alternates between two S/P + slots while retaining one O accumulator. The load and MMA FIFOs both + order K(i + 1) before V(i), allowing QK(i + 1) to run ahead of the + P(i)/PV(i) dependency. + """ + assert const_expr(self.q_stage == 1) + assert const_expr(self.num_kv_streams == 2) + assert const_expr(self.use_sp_double_buffer) + assert const_expr(self.use_block_sparsity) + + tSrQ = tiled_mma_qk.make_fragment_A(sQ) + tSrK = tiled_mma_qk.make_fragment_B(sK) + tOrV = tiled_mma_pv.make_fragment_B(sV) + + qk_mma_op, pv_mma_op = tiled_mma_qk.op, tiled_mma_pv.op + qk_mma_kind = sm100_utils._tcgen05_mma_kind(qk_mma_op) + q_smem_base = sm100_desc.smem_desc_base_from_tensor(sQ, sm100_desc.Major.K) + k_smem_base = sm100_desc.smem_desc_base_from_tensor(sK, sm100_desc.Major.K) + q_smem_start = sm100_desc.make_smem_desc_start_addr( + sQ[None, None, None, 0].iterator + ) + + sm100_utils.declare_ptx_smem_desc( + q_smem_start, + q_smem_base, + tSrQ[None, None, None, 0].layout, + var_name_prefix="fa_fwd_q_smem_desc", + ) + sm100_utils.declare_ptx_idesc(qk_mma_op, var_name="fa_fwd_qk_mma_idesc") + sm100_utils.declare_ptx_idesc(pv_mma_op, var_name="fa_fwd_pv_mma_idesc") + + gemm_S = partial( + sm100_utils.gemm_ptx_precomputed_varname, + smem_desc_base_b=k_smem_base, + tCrB_layout=tSrK[None, None, None, 0].layout, + smem_var_name_prefix="fa_fwd_q_smem_desc", + idesc_var_name="fa_fwd_qk_mma_idesc", + kind=qk_mma_kind, + smem_offset=0, + zero_init=True, + cta_group=self.cta_group_size, + ) + gemm_P = partial( + sm100_utils.gemm_ptx_partial, + pv_mma_op, + split_arrive=self.split_P_arrive if self.split_P_arrive > 0 else None, + cta_group=self.cta_group_size, + ) + + mma_q_consumer_phase = Int32(0) + mma_kv_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.kv_stage + ) + qk_count0 = Int32(0) + qk_count1 = Int32(0) + pv_count0 = Int32(0) + pv_count1 = Int32(0) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + seqlen = SeqlenInfoCls(batch_idx) + stream0_count, stream1_count = get_dual_stream_block_counts( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + split_idx, + num_splits, + 1, + self.q_subtile_factor, + seqlen, + ) + + if stream0_count > Int32(0) and is_leader_cta: + pipeline_q.consumer_wait_w_index_phase(0, mma_q_consumer_phase) + o0_should_accumulate = False + o1_should_accumulate = False + + # Initial score MMAs consume the two prefetched K blocks. + pipeline_kv.consumer_wait(mma_kv_consumer_state) + k0_idx, k0_phase = ( + mma_kv_consumer_state.index, + mma_kv_consumer_state.phase, + ) + sK0 = sK[None, None, None, k0_idx] + if const_expr(self.uneven_kv_smem): + sK0 = self.offset_kv_smem(sK0, k0_idx, k0_phase) + sp_slot0 = Int32(0) + Int32(2) * (qk_count0 % Int32(2)) + gemm_S( + sp_slot0 * self.n_block_size, + smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sK0.iterator) + ) + pipeline_s_p_o.producer_commit_w_index(sp_slot0) + pipeline_kv.consumer_release(mma_kv_consumer_state) + mma_kv_consumer_state.advance() + qk_count0 += Int32(1) + + if stream1_count > Int32(0): + pipeline_kv.consumer_wait(mma_kv_consumer_state) + k1_idx, k1_phase = ( + mma_kv_consumer_state.index, + mma_kv_consumer_state.phase, + ) + sK1 = sK[None, None, None, k1_idx] + if const_expr(self.uneven_kv_smem): + sK1 = self.offset_kv_smem(sK1, k1_idx, k1_phase) + sp_slot1 = Int32(1) + Int32(2) * (qk_count1 % Int32(2)) + gemm_S( + sp_slot1 * self.n_block_size, + smem_desc_start_b=sm100_desc.make_smem_desc_start_addr(sK1.iterator) + ) + pipeline_s_p_o.producer_commit_w_index(sp_slot1) + pipeline_kv.consumer_release(mma_kv_consumer_state) + mma_kv_consumer_state.advance() + qk_count1 += Int32(1) + + if stream0_count == Int32(1): + pipeline_q.consumer_release_w_index(0) + mma_q_consumer_phase ^= 1 + + for stream_idx in cutlass.range(stream0_count, unroll=1): + next_stream_idx = stream_idx + Int32(1) + if next_stream_idx < stream0_count: + pipeline_kv.consumer_wait(mma_kv_consumer_state) + k0_next_idx, k0_next_phase = ( + mma_kv_consumer_state.index, + mma_kv_consumer_state.phase, + ) + sK0_next = sK[None, None, None, k0_next_idx] + if const_expr(self.uneven_kv_smem): + sK0_next = self.offset_kv_smem( + sK0_next, k0_next_idx, k0_next_phase + ) + sp_slot0 = Int32(0) + Int32(2) * (qk_count0 % Int32(2)) + gemm_S( + sp_slot0 * self.n_block_size, + smem_desc_start_b=sm100_desc.make_smem_desc_start_addr( + sK0_next.iterator + ) + ) + pipeline_s_p_o.producer_commit_w_index(sp_slot0) + pipeline_kv.consumer_release(mma_kv_consumer_state) + mma_kv_consumer_state.advance() + qk_count0 += Int32(1) + if ( + next_stream_idx == stream0_count - Int32(1) + and stream0_count > stream1_count + ): + pipeline_q.consumer_release_w_index(0) + mma_q_consumer_phase ^= 1 + + pipeline_kv.consumer_wait(mma_kv_consumer_state) + v0_idx, v0_phase = ( + mma_kv_consumer_state.index, + mma_kv_consumer_state.phase, + ) + sV0 = sV[None, None, None, v0_idx] + if const_expr(self.uneven_kv_smem): + sV0 = self.offset_kv_smem(sV0, v0_idx, v0_phase) + sp_slot0 = Int32(0) + Int32(2) * (pv_count0 % Int32(2)) + sp_phase0 = pv_count0 // Int32(2) + pipeline_s_p_o.producer_acquire_w_index_phase(sp_slot0, sp_phase0) + gemm_P( + self.tmem_o_offset[0], + tOrP[None, None, None, sp_slot0], + tCrB=tOrV[None, None, None, v0_idx], + sA=None, + sB=sV0, + zero_init=not o0_should_accumulate, + mbar_ptr=( + pipeline_p_lastsplit.sync_object_full.get_barrier(sp_slot0) + if self.split_P_arrive > 0 + else None + ), + mbar_phase=sp_phase0, + ) + pipeline_o_acc.producer_commit_w_index(0) + pipeline_kv.consumer_release(mma_kv_consumer_state) + mma_kv_consumer_state.advance() + pv_count0 += Int32(1) + o0_should_accumulate = True + + if stream_idx < stream1_count: + if next_stream_idx < stream1_count: + pipeline_kv.consumer_wait(mma_kv_consumer_state) + k1_next_idx, k1_next_phase = ( + mma_kv_consumer_state.index, + mma_kv_consumer_state.phase, + ) + sK1_next = sK[None, None, None, k1_next_idx] + if const_expr(self.uneven_kv_smem): + sK1_next = self.offset_kv_smem( + sK1_next, k1_next_idx, k1_next_phase + ) + sp_slot1 = Int32(1) + Int32(2) * (qk_count1 % Int32(2)) + gemm_S( + sp_slot1 * self.n_block_size, + smem_desc_start_b=sm100_desc.make_smem_desc_start_addr( + sK1_next.iterator + ) + ) + pipeline_s_p_o.producer_commit_w_index(sp_slot1) + pipeline_kv.consumer_release(mma_kv_consumer_state) + mma_kv_consumer_state.advance() + qk_count1 += Int32(1) + if ( + next_stream_idx == stream1_count - Int32(1) + and stream0_count == stream1_count + ): + pipeline_q.consumer_release_w_index(0) + mma_q_consumer_phase ^= 1 + + pipeline_kv.consumer_wait(mma_kv_consumer_state) + v1_idx, v1_phase = ( + mma_kv_consumer_state.index, + mma_kv_consumer_state.phase, + ) + sV1 = sV[None, None, None, v1_idx] + if const_expr(self.uneven_kv_smem): + sV1 = self.offset_kv_smem(sV1, v1_idx, v1_phase) + sp_slot1 = Int32(1) + Int32(2) * (pv_count1 % Int32(2)) + sp_phase1 = pv_count1 // Int32(2) + pipeline_s_p_o.producer_acquire_w_index_phase(sp_slot1, sp_phase1) + gemm_P( + self.tmem_o_offset[1], + tOrP[None, None, None, sp_slot1], + tCrB=tOrV[None, None, None, v1_idx], + sA=None, + sB=sV1, + zero_init=not o1_should_accumulate, + mbar_ptr=( + pipeline_p_lastsplit.sync_object_full.get_barrier(sp_slot1) + if self.split_P_arrive > 0 + else None + ), + mbar_phase=sp_phase1, + ) + pipeline_o_acc.producer_commit_w_index(1) + pipeline_kv.consumer_release(mma_kv_consumer_state) + mma_kv_consumer_state.advance() + pv_count1 += Int32(1) + o1_should_accumulate = True + + work_tile = tile_scheduler.advance_to_next_work() + + # for both softmax0 and softmax1 warp group + @cute.jit + def _kv_head_idx(self, head_idx: Int32) -> Int32: + """Map query-head tile index -> KV-head index (FA3 descale semantics).""" + if cutlass.const_expr(self.pack_gqa): + return head_idx + return head_idx // self.qhead_per_kvhead + + @cute.jit + def _load_effective_descales( + self, + descale_tensors: Optional[DescaleTensors], + batch_idx: Int32, + kv_head_idx: Int32, + ) -> Tuple[Float32, Float32]: + """Load effective QK and V descales, defaulting unspecified tensors to identity.""" + qk_descale = Float32(1.0) + v_descale = Float32(1.0) + if cutlass.const_expr(descale_tensors is not None): + if cutlass.const_expr(descale_tensors.q_descale is not None): + qk_descale = qk_descale * Float32(descale_tensors.q_descale[batch_idx, kv_head_idx]) + if cutlass.const_expr(descale_tensors.k_descale is not None): + qk_descale = qk_descale * Float32(descale_tensors.k_descale[batch_idx, kv_head_idx]) + if cutlass.const_expr(descale_tensors.v_descale is not None): + v_descale = Float32(descale_tensors.v_descale[batch_idx, kv_head_idx]) + return qk_descale, v_descale + + @cute.jit + def softmax_loop( + self, + stage: int | Int32, + softmax_scale_log2: Float32, + softmax_scale: Float32 | None, + descale_tensors: Optional[DescaleTensors], + thr_mma_qk: cute.ThrMma, + tStS: cute.Tensor, # ((TILE_M, TILE_N), 1, 1, num_sp_tmem_slots) + sP: cute.Tensor | None, + sScale: cute.Tensor, + mLSE: Optional[cute.Tensor], + pipeline_s_p_o: pipeline.PipelineAsync, + pipeline_p_lastsplit: pipeline.PipelineAsync, + pipeline_sm_stats: pipeline.PipelineAsync, + sm_stats_barrier: pipeline.NamedBarrier, + smem_p_barrier: pipeline.NamedBarrier, + pipeline_s0_s1_sequence: Optional[pipeline.PipelineAsync], + learnable_sink: Optional[cute.Tensor], + block_info: BlockInfo, + num_splits: Int32, + SeqlenInfoCls: Callable, + AttentionMaskCls: Callable, + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + head_divmod=None, + blocksparse_tensors: Optional[BlockSparseTensors] = None, + tile_scheduler=None, + ): + """Compute softmax on attention scores from QK matrix multiplication. + + This method handles the softmax computation for either the first or second half of the + attention matrix, depending on the 'stage' parameter. It calculates row-wise maximum + and sum values needed for stable softmax computation, applies optional masking, and + transforms raw attention scores into probability distributions. + + The implementation uses specialized memory access patterns and efficient math operations + for computing exp(x) using exp2 functions. It also coordinates pipeline + synchronization between MMA, correction, and sequence processing stages. + """ + tidx = cute.arch.thread_idx()[0] % ( + cute.arch.WARP_SIZE + # * (len(self.softmax0_warp_ids) if stage == 0 else len(self.softmax1_warp_ids) + * (len(self.softmax0_warp_ids)) + ) + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4 + q_stage_idx = stage // self.num_kv_streams + aux_tensors = aux_data.tensors + + cta_qk_tiler = (self.mma_tiler_qk[0] // thr_mma_qk.thr_id.shape, self.mma_tiler_qk[1]) + tSAcc = tStS[(None, None), 0, 0, stage] # (128, 128) + tScS = thr_mma_qk.partition_C(cute.make_identity_tensor(self.mma_tiler_qk[:2])) + tScS = tScS[(None, None), 0, 0] # (128, 128) + smem_store_thr = None + sP_smem_view = None + if const_expr(self.p_in_smem): + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld16x32bx2Op(tcgen05.copy.Repetition(16)), + self.qk_acc_dtype, + ) + tmem_load_tiled = tcgen05.make_tmem_copy(tmem_load_atom, tSAcc) + thr_tmem_load = tmem_load_tiled.get_slice(tidx) + tStS_t2r = thr_tmem_load.partition_S(tSAcc) + smem_store_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.q_dtype, + num_bits_per_copy=128, + ) + smem_store_tiled = cute.make_tiled_copy_D(smem_store_atom, tmem_load_tiled) + smem_store_thr = smem_store_tiled.get_slice(tidx) + sP_smem_view = smem_store_thr.partition_D( + cute.composition( + sP, + cute.make_ordered_layout( + (self.m_block_size, self.n_block_size, self.s_stage), + order=(0, 1, 2), + ), + ) + ) + thr_tmem_store = thr_tmem_load + tStP_r2t = tStS_t2r + else: + tilePlikeFP32 = self.mma_tiler_qk[1] // Float32.width * self.v_dtype.width + tStP_layout = cute.composition( + tSAcc.layout, cute.make_layout((self.m_block_size, tilePlikeFP32)) + ) + tStP_base = cute.make_tensor( + tSAcc.iterator + self.tmem_s_to_p_offset, tStP_layout + ) + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), + self.qk_acc_dtype, + ) + thr_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tSAcc).get_slice(tidx) + if const_expr(self.use_sp_double_buffer): + tSAcc_pair = cute.make_tensor( + tSAcc.iterator, + cute.append( + tSAcc.layout, + cute.make_layout( + (2,), stride=(2 * self.n_block_size,) + ), + ), + ) + tStS_t2r = thr_tmem_load.partition_S(tSAcc_pair) + else: + tStS_t2r = thr_tmem_load.partition_S(tSAcc) + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp( + tcgen05.copy.Repetition(8 if const_expr(self.q_dtype.width == 8) else 16) + ), + Float32, + ) + thr_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP_base).get_slice(tidx) + if const_expr(self.use_sp_double_buffer): + tStP = cute.make_tensor( + tStP_base.iterator, + cute.append( + tStP_layout, + cute.make_layout( + (2,), stride=(2 * self.n_block_size,) + ), + ), + ) + tStP_r2t = thr_tmem_store.partition_D(tStP) + else: + tStP_r2t = thr_tmem_store.partition_D(tStP_base) + + mma_si_consumer_phase = Int32(0) + sm_stats_producer_phase = Int32(1) + s0_s1_sequence_phase = Int32(1 if stage == 0 else 0) + + # self.warp_scheduler_barrier_init() + + warp_idx_in_wg = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4 + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + kv_head_idx = self._kv_head_idx(head_idx) + seqlen = SeqlenInfoCls(batch_idx) + n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block, split_idx, num_splits) + + mask = AttentionMaskCls(seqlen) + shared_mask_kwargs = dict( + m_block=(self.q_stage * m_block + q_stage_idx) * self.cta_group_size, + thr_mma=thr_mma_qk, + thr_tmem_load=thr_tmem_load, + mask_causal=self.is_causal, + mask_local=self.is_local, + batch_idx=batch_idx, + head_idx=head_idx, + aux_data=aux_data, + vec_size=self.mask_vec_size, + ) + + # Recompute fastdiv_mods if necessary + recompute_fastdiv_mods_q = cutlass.const_expr( + aux_tensors is not None and (seqlen.has_cu_seqlens_q or seqlen.has_seqused_q) + ) + recompute_fastdiv_mods_k = cutlass.const_expr( + aux_tensors is not None and (seqlen.has_cu_seqlens_k or seqlen.has_seqused_k) + ) + + if cutlass.const_expr(fastdiv_mods is not None): + seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods + fastdiv_mods = ( + seqlen_q_divmod + if not recompute_fastdiv_mods_q + else FastDivmodDivisor(seqlen.seqlen_q), + seqlen_k_divmod + if not recompute_fastdiv_mods_k + else FastDivmodDivisor(seqlen.seqlen_k), + ) + + mask_mod = self.mask_mod if const_expr(self.mask_mod is not None) else None + mask_fn = partial( + mask.apply_mask_sm100, + mask_mod=mask_mod, + fastdiv_mods=fastdiv_mods, + head_divmod=head_divmod, + r2p=not self.p_in_smem, + **shared_mask_kwargs, + ) + if const_expr(self.use_block_sparsity): + # Full blocks dont need mask_mod + mask_fn_none = partial( + mask.apply_mask_sm100, + mask_mod=None, + fastdiv_mods=fastdiv_mods, + head_divmod=head_divmod, + r2p=not self.p_in_smem, + **shared_mask_kwargs, + ) + else: + mask_fn_none = None + + qk_descale, _ = self._load_effective_descales(descale_tensors, batch_idx, kv_head_idx) + + max_offset = 8 if cutlass.const_expr(self.q_dtype.width == 8) else 0 + if const_expr(self.score_mod is None): + softmax_scale_log2_eff = softmax_scale_log2 * qk_descale + softmax_scale_eff = None + else: + softmax_scale_log2_eff = softmax_scale_log2 + softmax_scale_eff = softmax_scale * qk_descale + + rescale_threshold = ( + 0.0 + if const_expr(self.p_in_smem) + else 8.0 + if const_expr(self.q_dtype.width == 16) + else 4.0 + if const_expr(self.q_dtype.width == 8) + else 0.0 + ) + softmax = SoftmaxSm100.create( + softmax_scale_log2_eff, + rescale_threshold=rescale_threshold, + softmax_scale=softmax_scale_eff, + max_offset=max_offset, + row_pair_stride=16 if const_expr(self.p_in_smem) else 0, + ) + softmax.reset() + + if const_expr(self.use_block_sparsity): + if const_expr(self.num_kv_streams == 2): + stream0_count, stream1_count = get_dual_stream_block_counts( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + split_idx, + num_splits, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + self.q_subtile_factor, + seqlen, + ) + tile_block_count = ( + stream0_count if stage == Int32(0) else stream1_count + ) + else: + tile_block_count = get_total_block_count( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + split_idx, + num_splits, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + self.q_subtile_factor, + seqlen_info=seqlen, + ) + has_work = tile_block_count > Int32(0) + else: + tile_block_count = n_block_max - n_block_min + has_work = const_expr(not self.is_split_kv) or tile_block_count > Int32(0) + + softmax_step = partial( + self.softmax_step, + softmax=softmax, + thr_mma_qk=thr_mma_qk, + pipeline_s_p_o=pipeline_s_p_o, + pipeline_p_lastsplit=pipeline_p_lastsplit, + pipeline_sm_stats=pipeline_sm_stats, + sm_stats_barrier=sm_stats_barrier, + smem_p_barrier=smem_p_barrier, + pipeline_s0_s1_sequence=pipeline_s0_s1_sequence, + thr_tmem_load=thr_tmem_load, + thr_tmem_store=thr_tmem_store, + tStS_t2r=tStS_t2r, + tStP_r2t=tStP_r2t, + smem_store_thr=smem_store_thr, + sP_smem_view=sP_smem_view, + sScale=sScale, + stage=stage, + batch_idx=batch_idx, + head_idx=head_idx, + m_block=(self.q_stage * m_block + q_stage_idx) * self.cta_group_size, + seqlen=seqlen, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + head_divmod=head_divmod, + ) + + if const_expr(self.use_block_sparsity) or has_work: + pipeline_sm_stats.producer_acquire_w_index_phase(stage, sm_stats_producer_phase) + sm_stats_producer_phase ^= 1 + + # Block sparse or dense iteration + if const_expr(self.use_block_sparsity): + # When aux_tensors exist, Q indices beyond seqlen_q must be wrapped to avoid + # OOB aux_tensor access. Only edge tiles (where m_tile_end > seqlen_q) need this. + if const_expr(aux_tensors is not None): + m_tile_end = ( + (self.q_stage * m_block + q_stage_idx + 1) * self.cta_group_size + ) * self.m_block_size + check_m_boundary = m_tile_end > seqlen.seqlen_q + else: + check_m_boundary = False + if const_expr(self.num_kv_streams == 2): + ( + mma_si_consumer_phase, + sm_stats_producer_phase, + s0_s1_sequence_phase, + empty_tile, + tile_block_count, + ) = softmax_block_sparse_dual_stream_sm100( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + seqlen, + split_idx, + num_splits, + softmax_step, + mask_fn, + mask_fn_none, + mma_si_consumer_phase, + sm_stats_producer_phase, + s0_s1_sequence_phase, + sm_stats_barrier, + self.p_in_smem, + Int32(stage), + check_m_boundary, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + self.q_subtile_factor, + ) + else: + ( + mma_si_consumer_phase, + sm_stats_producer_phase, + s0_s1_sequence_phase, + empty_tile, + ) = softmax_block_sparse_sm100( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + seqlen, + split_idx, + num_splits, + softmax_step, + mask_fn, + mask_fn_none, + mma_si_consumer_phase, + sm_stats_producer_phase, + s0_s1_sequence_phase, + pipeline_sm_stats, + sm_stats_barrier, + self.p_in_smem, + self.num_compute_slots, + Int32(stage), + check_m_boundary, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + self.q_subtile_factor, + ) + if not empty_tile: + my_row = self._softmax_tidx_to_row(tidx) + sScale[my_row + stage * self.m_block_size] = softmax.row_sum[0] + if const_expr( + self.num_kv_streams == 2 + or mLSE is not None + or learnable_sink is not None + ): + sScale[ + my_row + + stage * self.m_block_size + + self.num_compute_slots * self.m_block_size + ] = softmax.row_max[0] + # if tidx == 0: + # cute.printf("softmax row sum stage %d: %f, row_max = %f\n", stage, softmax.row_sum[0], softmax.row_max[0]) + # See block_sparse_utils.py NOTE [SM100 block-sparse empty tiles: mbarrier contract]. + # pipeline_sm_stats.producer_commit_w_index(stage) + sm_stats_barrier.arrive_w_index( + index=stage if self.p_in_smem else stage * 4 + warp_idx + ) + # if tidx == 0: cute.printf("softmax row sum stage %d: %f\n", stage, softmax.row_sum[0]) + else: + if const_expr(not self.is_split_kv) or tile_block_count > Int32(0): + mma_si_consumer_phase, sm_stats_producer_phase, s0_s1_sequence_phase = softmax_step( + mma_si_consumer_phase, + sm_stats_producer_phase, + s0_s1_sequence_phase, + n_block_max - 1, + is_first=True, + mask_fn=partial(mask_fn, mask_seqlen=True), + ) + n_block_max -= 1 + # Next couple of iterations with causal masking + if const_expr(self.is_causal or self.is_local): + n_block_min_causal_local_mask = block_info.get_n_block_min_causal_local_mask( + seqlen, m_block, n_block_min + ) + for n_tile in cutlass.range(n_block_max - n_block_min_causal_local_mask, unroll=1): + n_block = n_block_max - 1 - n_tile + mma_si_consumer_phase, sm_stats_producer_phase, s0_s1_sequence_phase = ( + softmax_step( + mma_si_consumer_phase, + sm_stats_producer_phase, + s0_s1_sequence_phase, + n_block, + mask_fn=partial(mask_fn, mask_seqlen=False), + ) + ) + n_block_max = cutlass.min(n_block_max, n_block_min_causal_local_mask) + # The remaining iterations have no masking (but may still need mask_mod) + n_block_min_before_local_mask = block_info.get_n_block_min_before_local_mask( + seqlen, m_block, n_block_min + ) + for n_tile in cutlass.range(n_block_max - n_block_min_before_local_mask, unroll=1): + n_block = n_block_max - n_tile - 1 + if const_expr(self.mask_mod is not None): + mma_si_consumer_phase, sm_stats_producer_phase, s0_s1_sequence_phase = softmax_step( + mma_si_consumer_phase, sm_stats_producer_phase, s0_s1_sequence_phase, n_block, + mask_fn=partial(mask_fn, mask_seqlen=False), + ) + else: + mma_si_consumer_phase, sm_stats_producer_phase, s0_s1_sequence_phase = softmax_step( + mma_si_consumer_phase, sm_stats_producer_phase, s0_s1_sequence_phase, n_block, + ) + # Separate iterations with local masking on the left + if const_expr(self.is_local and block_info.window_size_left is not None): + n_block_max = cutlass.min(n_block_max, n_block_min_before_local_mask) + for n_tile in cutlass.range(0, n_block_max - n_block_min, unroll=1): + n_block = n_block_max - 1 - n_tile + mma_si_consumer_phase, sm_stats_producer_phase, s0_s1_sequence_phase = ( + softmax_step( + mma_si_consumer_phase, + sm_stats_producer_phase, + s0_s1_sequence_phase, + n_block, + mask_fn=partial(mask_fn, mask_seqlen=False), + ) + ) + # Now that we no longer already have the 1st iteration, need mask_seqlen=True here + + # Dense path always writes scale / signals + my_row = self._softmax_tidx_to_row(tidx) + sScale[my_row + stage * self.m_block_size] = softmax.row_sum[0] + if const_expr(mLSE is not None or learnable_sink is not None): + sScale[ + my_row + + stage * self.m_block_size + + self.num_compute_slots * self.m_block_size + ] = softmax.row_max[0] + # pipeline_sm_stats.producer_commit_w_index(stage) + sm_stats_barrier.arrive_w_index( + index=stage if self.p_in_smem else stage * 4 + warp_idx + ) + + # # Write LSE to gmem + # if const_expr(mLSE is not None): + # acc_O_mn_row_is_zero_or_nan = softmax.row_sum[0] == 0.0 or softmax.row_sum[0] != softmax.row_sum[0] + # scale = ( + # cute.arch.rcp_approx(softmax.row_sum[0] if not acc_O_mn_row_is_zero_or_nan else 1.0) + # ) + # LN2 = math.log(2.0) + # lse = ( + # (softmax.row_max[0] * softmax.scale_log2 + cute.math.log2(softmax.row_sum[0], fastmath=True)) * LN2 + # if not acc_O_mn_row_is_zero_or_nan else -Float32.inf + # ) + # if const_expr(not seqlen.has_cu_seqlens_q): + # mLSE_cur = mLSE[None, head_idx, batch_idx] + # else: + # mLSE_cur = cute.domain_offset((seqlen.offset_q,), mLSE[None, head_idx]) + # gLSE = cute.local_tile(mLSE_cur, (self.m_block_size,), (m_block * 2 + stage,)) + # if tidx < seqlen.seqlen_q - (m_block * 2 + stage) * self.m_block_size: + # gLSE[tidx] = lse + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() + # End of persistent scheduler loop + + # This is equivalent to pipeline_sm_stats.producer_tail + pipeline_sm_stats.producer_acquire_w_index_phase(stage, sm_stats_producer_phase) + # This is equivalent to pipeline_s0_s1.producer_tail + if const_expr(self.s0_s1_barrier): + if stage == 0: + pipeline_s0_s1_sequence.sync_object_full.wait(stage, s0_s1_sequence_phase) + + @cute.jit + def softmax_step( + self, + mma_si_consumer_phase: Int32, + sm_stats_producer_phase: Int32, + s0_s1_sequence_phase: Int32, + n_block: Int32, + softmax: SoftmaxSm100, + thr_mma_qk: cute.ThrMma, + pipeline_s_p_o: pipeline.PipelineAsync, + pipeline_p_lastsplit: pipeline.PipelineAsync, + pipeline_sm_stats: pipeline.PipelineAsync, + sm_stats_barrier: pipeline.NamedBarrier, + smem_p_barrier: pipeline.NamedBarrier, + pipeline_s0_s1_sequence: Optional[pipeline.PipelineAsync], + thr_tmem_load: cute.CopyAtom, + thr_tmem_store: cute.CopyAtom, + tStS_t2r: cute.Tensor, + tStP_r2t: cute.Tensor, + smem_store_thr: cute.ThrCopy | None, + sP_smem_view: cute.Tensor | None, + sScale: cute.Tensor, + stage: int | Int32, + batch_idx: Int32, + head_idx: Int32, + m_block: Int32, + seqlen, + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + head_divmod=None, + mask_fn: Optional[Callable] = None, + is_first: bool = False, + ) -> Tuple[cute.Int32, cute.Int32, cute.Int32]: + """Perform a single step of the softmax computation on a block of attention scores. + + This method processes one block of the attention matrix, computing numerically stable + softmax by first finding the row maximum, subtracting it from all elements, applying + exponential function, and then normalizing by the sum of exponentials. It also handles + optional masking of attention scores. + + The method involves several key operations: + 1. Loading attention scores from tensor memory + 2. Applying optional masking based on position + 3. Computing row-wise maximum values for numerical stability + 4. Transforming scores using exp2(x*scale - max*scale) + 5. Computing row sums for normalization + 6. Coordinating pipeline synchronization between different processing stages + """ + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4 + tilePlikeFP32 = self.mma_tiler_qk[1] // Float32.width * self.v_dtype.width + tScS = thr_mma_qk.partition_C(cute.make_identity_tensor(self.mma_tiler_qk[:2])) + tScS = tScS[(None, None), 0, 0] # (128, 128) + # tScScale = cute.composition(tScS, cute.make_layout((self.m_block_size, 1))) + cta_qk_tiler = (self.mma_tiler_qk[0] // thr_mma_qk.thr_id.shape, self.mma_tiler_qk[1]) + tScS_shape = cta_qk_tiler # (128, 128) + tScP_shape = (tScS_shape[0], tilePlikeFP32) # (128, 64) + + if const_expr(self.use_sp_double_buffer): + sp_parity = mma_si_consumer_phase % Int32(2) + sp_slot = stage + Int32(2) * sp_parity + sp_phase = mma_si_consumer_phase // Int32(2) + tStS_t2r_cur = tStS_t2r[None, None, None, sp_parity] + tStP_r2t_cur = tStP_r2t[None, None, None, sp_parity] + else: + sp_slot = stage + sp_phase = mma_si_consumer_phase + tStS_t2r_cur = tStS_t2r + tStP_r2t_cur = tStP_r2t + + # Wait for Si + pipeline_s_p_o.consumer_wait_w_index_phase(sp_slot, sp_phase) + tSrS_t2r = cute.make_rmem_tensor(thr_tmem_load.partition_D(tScS).shape, self.qk_acc_dtype) + cute.copy(thr_tmem_load, tStS_t2r_cur, tSrS_t2r) + # tSrS_t2r = copy_utils.load_t2r(thr_tmem_load, tScS_shape, tStS_t2r) + if cutlass.const_expr(self.score_mod is not None): + self.apply_score_mod( + tSrS_t2r, + thr_tmem_load, + thr_mma_qk, + batch_idx, + head_idx, + m_block, + n_block, + softmax, + seqlen, + aux_data, + fastdiv_mods, + head_divmod, + ) + + if const_expr(mask_fn is not None): + mask_fn(tSrS_t2r, n_block=n_block) + row_max, acc_scale = softmax.update_row_max(tSrS_t2r.load(), is_first) + + if const_expr(not is_first): + thread_idx = thr_tmem_load.thr_idx + sScale[ + self._softmax_tidx_to_row(thread_idx) + stage * self.m_block_size + ] = acc_scale + # if thread_idx == 0: cute.printf("softmax acc_scale stage %d: %f, row_max = %f\n", stage, acc_scale, row_max) + # Notify correction wg that row_max is ready + # pipeline_sm_stats.producer_commit_w_index(stage) + sm_stats_barrier.arrive_w_index( + index=stage if self.p_in_smem else stage * 4 + warp_idx + ) + + # if thread_idx == 0 and stage == 0: cute.print_tensor(tSrS_t2r) + softmax.scale_subtract_rowmax(tSrS_t2r, row_max) + # Sequence barrier wait + if const_expr(self.s0_s1_barrier): + pipeline_s0_s1_sequence.sync_object_full.wait(stage, s0_s1_sequence_phase) + if const_expr(self.p_in_smem): + tSrP_r2t = cute.make_rmem_tensor(tSrS_t2r.shape, self.q_dtype) + else: + tSrP_r2t_f32 = cute.make_rmem_tensor( + thr_tmem_store.partition_S(cute.make_identity_tensor(tScP_shape)).shape, + Float32, + ) + tSrP_r2t = cute.make_tensor( + cute.recast_ptr(tSrP_r2t_f32.iterator, dtype=self.q_dtype), + tSrS_t2r.layout, + ) + # softmax.scale_apply_exp2_convert(tSrS_t2r, row_max, tSrP_r2t) + softmax.apply_exp2_convert( + tSrS_t2r, + tSrP_r2t, + ex2_emu_freq=self.ex2_emu_freq, + ex2_emu_start_frg=self.ex2_emu_start_frg, + ) + # Sequence barrier arrive + if const_expr(self.s0_s1_barrier): + pipeline_s0_s1_sequence.sync_object_full.arrive(1 - stage, dst=None) + if const_expr(self.p_in_smem): + rP_smem_view = smem_store_thr.retile(tSrP_r2t) + cute.copy( + smem_store_thr, + rP_smem_view, + sP_smem_view[None, None, None, stage], + ) + cute.arch.fence_view_async_shared() + # The S/P/O pipeline already counts every softmax and correction + # consumer; its release is the visibility handshake for this P tile. + pipeline_s_p_o.consumer_release_w_index(sp_slot) + else: + for i in cutlass.range_constexpr(cute.size(tStP_r2t_cur.shape[2])): + cute.copy( + thr_tmem_store, + tSrP_r2t_f32[None, None, i], + tStP_r2t_cur[None, None, i], + ) + if const_expr(self.split_P_arrive > 0): + split_P_arrive_idx = ( + cute.size(tStP_r2t_cur.shape[2]) + * self.split_P_arrive + // self.n_block_size + ) + if const_expr(i + 1 == split_P_arrive_idx): + cute.arch.fence_view_async_tmem_store() + pipeline_s_p_o.consumer_release_w_index(sp_slot) + cute.arch.fence_view_async_tmem_store() + if const_expr(self.split_P_arrive > 0): + cute.arch.sync_warp() + with cute.arch.elect_one(): + pipeline_p_lastsplit.producer_commit_w_index(sp_slot) + else: + pipeline_s_p_o.consumer_release_w_index(sp_slot) + pipeline_sm_stats.producer_acquire_w_index_phase(stage, sm_stats_producer_phase) + softmax.update_row_sum(tSrS_t2r.load(), acc_scale, is_first) + # acc_scale = cute.math.exp2(acc_scale_, fastmath=True) + mma_si_consumer_phase = ( + mma_si_consumer_phase + Int32(1) + if const_expr(self.use_sp_double_buffer) + else mma_si_consumer_phase ^ 1 + ) + return mma_si_consumer_phase, sm_stats_producer_phase ^ 1, s0_s1_sequence_phase ^ 1 + + @cute.jit + def correction_loop( + self, + thr_mma_qk: cute.ThrMma, + thr_mma_pv: cute.ThrMma, + tStS: cute.Tensor, + tOtO: cute.Tensor, + sScale: cute.Tensor, + mO: cute.Tensor, + mLSE: cute.Tensor, + sO: cute.Tensor, + pipeline_s_p_o: pipeline.PipelineAsync, + pipeline_o_acc: pipeline.PipelineAsync, + pipeline_sm_stats: pipeline.PipelineAsync, + sm_stats_barrier: pipeline.NamedBarrier, + pipeline_o_epi: pipeline.PipelineAsync, + learnable_sink: Optional[cute.Tensor], + descale_tensors: Optional[DescaleTensors], + gmem_tiled_copy_O: cute.TiledCopy, + tma_atom_O: cute.CopyAtom, + softmax_scale_log2: Float32, + block_info: BlockInfo, + num_splits: Int32, + SeqlenInfoCls: Callable, + blocksparse_tensors: Optional[BlockSparseTensors] = None, + tile_scheduler=None, + ): + tidx = cute.arch.thread_idx()[0] % (cute.arch.WARP_SIZE * len(self.correction_warp_ids)) + row_tidx = tidx % self.m_block_size + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4 + mma_tile_coord_v = thr_mma_qk.thr_idx + + # First iter: no correction is required + # Notify mma warp that O has been rescaled + for stage in cutlass.range(self.num_compute_slots): + pipeline_s_p_o.consumer_release_w_index(stage) + + sm_stats_consumer_phase = Int32(0) + o_corr_consumer_phase = Int32(0) + corr_epi_producer_phase = Int32(1) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + kv_head_idx = self._kv_head_idx(head_idx) + qk_descale, v_descale = self._load_effective_descales(descale_tensors, batch_idx, kv_head_idx) + if const_expr(self.score_mod is None): + softmax_scale_log2_eff = softmax_scale_log2 * qk_descale + else: + softmax_scale_log2_eff = softmax_scale_log2 + + max_offset = Float32(8.0) if cutlass.const_expr(self.q_dtype.width == 8) else Float32(0.0) + max_offset_scale = ( + Float32(256.0) if cutlass.const_expr(self.q_dtype.width == 8) else Float32(1.0) + ) + seqlen = SeqlenInfoCls(batch_idx) + n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block, split_idx, num_splits) + + if const_expr(self.is_split_kv): + mO_cur = seqlen.offset_batch_Q(mO, batch_idx, dim=3)[None, None, head_idx, split_idx] + else: + mO_cur = seqlen.offset_batch_Q(mO, batch_idx, dim=3)[None, None, head_idx] + gO = None + if const_expr(self.use_tma_O or not self.pack_gqa): + tiler_gO = ((self.mma_tiler_pv[0] * self.q_stage), self.head_dim_v_padded) + gO = cute.local_tile(mO_cur, tiler_gO, (m_block, 0)) # (128 * 2, 128) + gO = layout_utils.select( + cute.flat_divide(gO, (self.mma_tiler_pv[0],)), mode=[0, 2, 1] + ) # (128, 128, 2) + gO = cute.flat_divide(gO, (self.mma_tiler_pv[0] // self.cta_group_size,))[None, mma_tile_coord_v, None, None] + + # Default LSE to -inf for invalid split_idx tiles + stats = [(0.0, -Float32.inf if const_expr(mLSE is not None or learnable_sink is not None) else None, True)] * self.q_stage + + if const_expr(self.use_block_sparsity): + total_block_count = get_total_block_count( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + split_idx, + num_splits, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + self.q_subtile_factor, + seqlen_info=seqlen, + ) + has_work = total_block_count > Int32(0) + else: + total_block_count = n_block_max - n_block_min + has_work = const_expr(not self.is_split_kv) or total_block_count > Int32(0) + + if has_work: + # Ignore first signal from softmax as no correction is required + # pipeline_sm_stats.consumer_wait_w_index_phase(0, sm_stats_consumer_phase) + sm_stats_barrier.arrive_and_wait_w_index( + index=0 if self.p_in_smem else 0 * 4 + warp_idx + ) + pipeline_sm_stats.consumer_release_w_index(0) + if const_expr(self.num_compute_slots == 2): + # pipeline_sm_stats.consumer_wait_w_index_phase(1, sm_stats_consumer_phase) + sm_stats_barrier.arrive_and_wait_w_index( + index=1 if self.p_in_smem else 1 * 4 + warp_idx + ) + sm_stats_consumer_phase ^= 1 + + for i in cutlass.range(total_block_count - 1, unroll=1): + for stage in cutlass.range_constexpr(self.num_compute_slots): + # wait for S0 / S1 + # pipeline_sm_stats.consumer_wait_w_index_phase(stage, sm_stats_consumer_phase) + sm_stats_barrier.arrive_and_wait_w_index( + index=stage if self.p_in_smem else stage * 4 + warp_idx + ) + # cute.copy(tiled_tmem_load_vec, tStScales_t2r[stage], tSrScale_t2r) + # cute.arch.fence_view_async_tmem_load() + # scale = tSrScale_t2r[0] + # should_rescale = True + # if tidx == 0: cute.printf("Correction scale i = %d, for stage %d: %f, should_rescale = %d\n", i, stage, scale, should_rescale) + # Don't need O_full anymore, since by the time softmax has signaled the correction + # warps, S_i must have been done, so O_i-1 must have been done as well. + # pipeline_o_acc.consumer_wait_w_index_phase(stage, o_corr_consumer_phase) + if const_expr(self.p_in_smem): + self.correction_rescale_smem_p( + thr_mma_pv, + tOtO[None, None, None, stage], + tidx, + sScale, + stage, + ) + else: + scale = sScale[row_tidx + stage * self.m_block_size] + if cute.arch.vote_ballot_sync(scale < 1.0) != 0: + self.correction_rescale( + thr_mma_pv, + tOtO[None, None, None, stage], + tidx, + scale, + ) + # Notify mma warp that O has been rescaled + pipeline_s_p_o.consumer_release_w_index(stage) + pipeline_sm_stats.consumer_release_w_index( + self.num_compute_slots - 1 - stage + ) + sm_stats_consumer_phase ^= 1 + # o_corr_consumer_phase ^= 1 + if const_expr(self.num_compute_slots == 2): + pipeline_sm_stats.consumer_release_w_index(1) + # End of seqlen_corr_loop_steps + + # Even in the case of self.overlap_sO_sQ, we can write to stage 0 of sO without + # additional sync because the MMA in the top half must have been done. + # Similarly we can write to stage 1 of sO without additional sync. + learnable_sink_val = [None] * self.q_stage + if const_expr(learnable_sink is not None): + if const_expr(not self.pack_gqa): + sink_val = Float32(learnable_sink[head_idx]) + learnable_sink_val = [sink_val] * self.q_stage + else: # Each thread might have a different sink value due to different q_head + for stage in cutlass.range_constexpr(self.q_stage): + q_head_idx = ( + ((m_block * self.q_stage + stage) * self.cta_group_size + mma_tile_coord_v) * self.m_block_size + row_tidx + ) % self.qhead_per_kvhead + head_idx * self.qhead_per_kvhead + learnable_sink_val[stage] = Float32(learnable_sink[q_head_idx]) + for stage in cutlass.range_constexpr(self.q_stage): + # pipeline_sm_stats.consumer_wait_w_index_phase(stage, sm_stats_consumer_phase) + sm_stats_barrier.arrive_and_wait_w_index( + index=stage if self.p_in_smem else stage * 4 + warp_idx + ) + # cute.copy(tiled_tmem_load_vec, tStScales_t2r[stage], tSrScale_t2r) + # cute.arch.fence_view_async_tmem_load() + # scale = tSrScale_t2r[0] + row_sum = sScale[row_tidx + stage * self.m_block_size] + if const_expr(mLSE is not None or learnable_sink is not None): + row_max = sScale[ + row_tidx + + (stage + self.num_compute_slots) * self.m_block_size + ] + else: + row_max = None + pipeline_sm_stats.consumer_release_w_index(stage) + if const_expr(learnable_sink is not None): + LOG2_E = math.log2(math.e) + sink_val = learnable_sink_val[stage] + if const_expr(not self.is_split_kv) or split_idx == 0: + if row_max == -Float32.inf: + # It's possible to have an empty row with splitKV. + row_max = sink_val * (LOG2_E / softmax_scale_log2_eff) + row_sum = max_offset_scale + else: + row_sum += cute.math.exp2( + sink_val * LOG2_E - row_max * softmax_scale_log2_eff + max_offset, fastmath=True + ) + acc_O_mn_row_is_zero_or_nan = row_sum == 0.0 or row_sum != row_sum + stats[stage] = (row_sum, row_max, acc_O_mn_row_is_zero_or_nan) + scale = cute.arch.rcp_approx(row_sum if not acc_O_mn_row_is_zero_or_nan else 1.0) + scale = scale * v_descale + if const_expr(self.p_in_smem): + correction_threads = ( + len(self.correction_warp_ids) * cute.arch.WARP_SIZE + ) + cute.arch.barrier( + barrier_id=int(NamedBarrierFwdSm100.CorrectionScale), + number_of_threads=correction_threads, + ) + if tidx < self.m_block_size: + sScale[row_tidx + stage * self.m_block_size] = scale + cute.arch.barrier( + barrier_id=int(NamedBarrierFwdSm100.CorrectionScale), + number_of_threads=correction_threads, + ) + # Wait for the last O to be ready from the MMA warp + pipeline_o_acc.consumer_wait_w_index_phase(stage, o_corr_consumer_phase) + if const_expr(not self.use_correction_warps_for_epi): + pipeline_o_epi.producer_acquire_w_index_phase(stage, corr_epi_producer_phase) + gO_stage = gO[None, None, stage] if const_expr(gO is not None) else None + self.correction_epilogue( + thr_mma_pv, + tOtO[None, None, None, stage], + tidx, + stage, + m_block, + seqlen.seqlen_q, + scale, + sO[None, None, stage], + mO_cur, + gO_stage, + gmem_tiled_copy_O, + sScale=sScale if const_expr(self.p_in_smem) else None, + ) + # Signal for the next work tile that O buffers in tmem are already read, so + # mma warp can write to them + pipeline_s_p_o.consumer_release_w_index(stage) + if const_expr(not self.use_correction_warps_for_epi): + pipeline_o_epi.producer_commit_w_index(stage) + # if tidx == 0: cute.printf("Correction final scale for stage %d: %f\n", stage, scale) + + o_corr_consumer_phase ^= 1 + sm_stats_consumer_phase ^= 1 + corr_epi_producer_phase ^= 1 + else: + gmem_tiled_copy_O_for_empty_tile = None + if const_expr(self.use_correction_warps_for_epi): + gmem_tiled_copy_O_for_empty_tile = gmem_tiled_copy_O + if const_expr(self.use_block_sparsity): + ( + sm_stats_consumer_phase, + o_corr_consumer_phase, + corr_epi_producer_phase, + ) = handle_block_sparse_empty_tile_correction_sm100( + tidx, + self.q_stage, + self.m_block_size, + self.qhead_per_kvhead, + self.pack_gqa, + self.is_split_kv, + learnable_sink, + mLSE, + seqlen, + m_block, + head_idx, + batch_idx, + split_idx, + stats, + partial(self.correction_epilogue, zero_output=True), + thr_mma_pv, + tOtO, + sO, + pipeline_sm_stats, + sm_stats_barrier, + self.p_in_smem, + pipeline_o_epi, + sm_stats_consumer_phase, + o_corr_consumer_phase, + corr_epi_producer_phase, + softmax_scale_log2_eff, + max_offset, + max_offset_scale, + mO_cur, + gO, + gmem_tiled_copy_O_for_empty_tile, + ) + + if const_expr(mLSE is not None): + if const_expr(not seqlen.has_cu_seqlens_q): + if const_expr(self.is_split_kv): + mLSE_cur = mLSE[None, head_idx, batch_idx, split_idx] + else: + mLSE_cur = mLSE[None, head_idx, batch_idx] + else: + offset = ( + seqlen.offset_q if const_expr(not self.pack_gqa) else (0, seqlen.offset_q) + ) + if const_expr(self.is_split_kv): + mLSE_cur = cute.domain_offset((offset,), mLSE[None, head_idx, split_idx]) + else: + mLSE_cur = cute.domain_offset((offset,), mLSE[None, head_idx]) + for stage in cutlass.range_constexpr(self.q_stage): + m_tile_idx = (m_block * self.q_stage + stage) * self.cta_group_size + mma_tile_coord_v + row_sum, row_max, acc_O_mn_row_is_zero_or_nan = stats[stage] + # if tidx == 0 and stage <= 1: + # cute.printf("row_sum = {}, row_max = {}, acc_O_mn_row_is_zero_or_nan = {}\n", row_sum, row_max, acc_O_mn_row_is_zero_or_nan) + LN2 = math.log(2.0) + lse = ( + (row_max * softmax_scale_log2_eff + (cute.math.log2(row_sum, fastmath=True) - max_offset)) * LN2 + if not acc_O_mn_row_is_zero_or_nan + else -Float32.inf + ) + seqlen_q = ( + seqlen.seqlen_q + if const_expr(not self.pack_gqa) + else seqlen.seqlen_q * self.qhead_per_kvhead + ) + if const_expr(not self.pack_gqa or self.m_block_size % self.qhead_per_kvhead == 0): + gLSE = cute.local_tile(mLSE_cur, (self.m_block_size,), (m_tile_idx,)) + if ( + tidx < self.m_block_size + and row_tidx < seqlen_q - m_tile_idx * self.m_block_size + ): + # This actually just works with PackGQA too + gLSE[row_tidx] = lse + else: + idx = m_tile_idx * self.m_block_size + row_tidx + if tidx < self.m_block_size and idx < seqlen_q: + m_idx = idx // self.qhead_per_kvhead + h_idx = idx - m_idx * self.qhead_per_kvhead + lse_ptr_i64 = utils.elem_pointer(mLSE_cur, ((h_idx, m_idx),)).toint() + lse_gmem_ptr = cute.make_ptr( + mLSE_cur.element_type, lse_ptr_i64, cute.AddressSpace.gmem, assumed_align=4 + ) + cute.make_tensor(lse_gmem_ptr, (1,))[0] = lse + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() + # End of persistent scheduler loop + + # This is equivalent to pipeline_o_epi.consumer_tail() for the correction warps + if const_expr(not self.use_correction_warps_for_epi): + pipeline_o_epi.producer_acquire_w_index_phase(self.q_stage - 1, corr_epi_producer_phase) + + @cute.jit + def correction_loop_dual_stream( + self, + thr_mma_qk: cute.ThrMma, + thr_mma_pv: cute.ThrMma, + tStS: cute.Tensor, + tOtO: cute.Tensor, + sScale: cute.Tensor, + mO: cute.Tensor, + mLSE: cute.Tensor, + sO: cute.Tensor, + pipeline_s_p_o: pipeline.PipelineAsync, + pipeline_o_acc: pipeline.PipelineAsync, + pipeline_sm_stats: pipeline.PipelineAsync, + sm_stats_barrier: pipeline.NamedBarrier, + pipeline_o_epi: pipeline.PipelineAsync, + learnable_sink: Optional[cute.Tensor], + descale_tensors: Optional[DescaleTensors], + gmem_tiled_copy_O: cute.TiledCopy, + tma_atom_O: cute.CopyAtom, + softmax_scale_log2: Float32, + block_info: BlockInfo, + num_splits: Int32, + SeqlenInfoCls: Callable, + blocksparse_tensors: Optional[BlockSparseTensors] = None, + tile_scheduler=None, + ): + """Correct and merge two independently-normalized sparse KV streams.""" + assert const_expr(self.num_kv_streams == 2) + assert const_expr(self.q_stage == 1) + assert const_expr(self.use_block_sparsity) + assert const_expr(self.use_tma_O) + assert const_expr(not self.use_correction_warps_for_epi) + assert const_expr(not self.pack_gqa) + assert const_expr(not self.is_split_kv) + assert const_expr(not self.is_varlen_q) + assert const_expr(learnable_sink is None) + + tidx = cute.arch.thread_idx()[0] % ( + cute.arch.WARP_SIZE * len(self.correction_warp_ids) + ) + row_tidx = tidx % self.m_block_size + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4 + + # Seed the first P/O acquire for each stream. Thereafter a stream is + # released only when it has a real PV event; empty tiles leave this + # seed untouched for the next active tile. + pipeline_s_p_o.consumer_release_w_index(0) + pipeline_s_p_o.consumer_release_w_index(1) + + o_corr_phase0 = Int32(0) + o_corr_phase1 = Int32(0) + corr_epi_producer_phase = Int32(1) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + kv_head_idx = self._kv_head_idx(head_idx) + qk_descale, v_descale = self._load_effective_descales( + descale_tensors, batch_idx, kv_head_idx + ) + softmax_scale_log2_eff = ( + softmax_scale_log2 * qk_descale + if const_expr(self.score_mod is None) + else softmax_scale_log2 + ) + max_offset = ( + Float32(8.0) + if const_expr(self.q_dtype.width == 8) + else Float32(0.0) + ) + seqlen = SeqlenInfoCls(batch_idx) + stream0_count, stream1_count = get_dual_stream_block_counts( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + split_idx, + num_splits, + 1, + self.q_subtile_factor, + seqlen, + ) + has_stream0 = stream0_count > Int32(0) + has_stream1 = stream1_count > Int32(0) + + # Every stream emits one first signal: the first score block when + # active, or a synthetic empty-stream signal otherwise. + sm_stats_barrier.arrive_and_wait_w_index( + index=0 if self.p_in_smem else 0 * 4 + warp_idx + ) + pipeline_sm_stats.consumer_release_w_index(0) + sm_stats_barrier.arrive_and_wait_w_index( + index=1 if self.p_in_smem else 1 * 4 + warp_idx + ) + pipeline_sm_stats.consumer_release_w_index(1) + + # Signals 1..count-1 carry the online-softmax correction scale. + if stream0_count > Int32(1): + for corr_idx in cutlass.range(stream0_count - Int32(1), unroll=1): + sm_stats_barrier.arrive_and_wait_w_index( + index=0 if self.p_in_smem else 0 * 4 + warp_idx + ) + if const_expr(self.p_in_smem): + self.correction_rescale_smem_p( + thr_mma_pv, + tOtO[None, None, None, 0], + tidx, + sScale, + Int32(0), + ) + else: + scale0 = sScale[row_tidx] + if cute.arch.vote_ballot_sync(scale0 < Float32(1.0)) != 0: + self.correction_rescale( + thr_mma_pv, + tOtO[None, None, None, 0], + tidx, + scale0, + ) + pipeline_s_p_o.consumer_release_w_index(0) + pipeline_sm_stats.consumer_release_w_index(0) + + if corr_idx + Int32(1) < stream1_count: + sm_stats_barrier.arrive_and_wait_w_index( + index=1 if self.p_in_smem else 1 * 4 + warp_idx + ) + if const_expr(self.p_in_smem): + self.correction_rescale_smem_p( + thr_mma_pv, + tOtO[None, None, None, 1], + tidx, + sScale, + Int32(1), + ) + else: + scale1 = sScale[row_tidx + self.m_block_size] + if cute.arch.vote_ballot_sync(scale1 < Float32(1.0)) != 0: + self.correction_rescale( + thr_mma_pv, + tOtO[None, None, None, 1], + tidx, + scale1, + ) + pipeline_s_p_o.consumer_release_w_index(1) + pipeline_sm_stats.consumer_release_w_index(1) + + # Active streams emit one additional final-stat signal. Empty + # streams have already completed their sole synthetic signal above. + row_sum0 = Float32(0.0) + row_sum1 = Float32(0.0) + row_max0 = -Float32.inf + row_max1 = -Float32.inf + if has_stream0: + sm_stats_barrier.arrive_and_wait_w_index( + index=0 if self.p_in_smem else 0 * 4 + warp_idx + ) + row_sum0 = sScale[row_tidx] + row_max0 = sScale[ + row_tidx + self.num_compute_slots * self.m_block_size + ] + pipeline_sm_stats.consumer_release_w_index(0) + pipeline_o_acc.consumer_wait_w_index_phase(0, o_corr_phase0) + if has_stream1: + sm_stats_barrier.arrive_and_wait_w_index( + index=1 if self.p_in_smem else 1 * 4 + warp_idx + ) + row_sum1 = sScale[row_tidx + self.m_block_size] + row_max1 = sScale[ + row_tidx + + self.m_block_size + + self.num_compute_slots * self.m_block_size + ] + pipeline_sm_stats.consumer_release_w_index(1) + pipeline_o_acc.consumer_wait_w_index_phase(1, o_corr_phase1) + + valid0 = has_stream0 and row_sum0 != Float32(0.0) and row_sum0 == row_sum0 + valid1 = has_stream1 and row_sum1 != Float32(0.0) and row_sum1 == row_sum1 + merged_row_max = -Float32.inf + if valid0: + merged_row_max = row_max0 + if valid1: + merged_row_max = ( + max(merged_row_max, row_max1) if valid0 else row_max1 + ) + + alpha0 = Float32(0.0) + alpha1 = Float32(0.0) + if valid0: + alpha0 = cute.math.exp2( + (row_max0 - merged_row_max) * softmax_scale_log2_eff, + fastmath=True, + ) + if valid1: + alpha1 = cute.math.exp2( + (row_max1 - merged_row_max) * softmax_scale_log2_eff, + fastmath=True, + ) + merged_row_sum = alpha0 * row_sum0 + alpha1 * row_sum1 + merged_invalid = ( + merged_row_sum == Float32(0.0) or merged_row_sum != merged_row_sum + ) + inv_sum = cute.arch.rcp_approx( + merged_row_sum if not merged_invalid else Float32(1.0) + ) + output_scale0 = ( + alpha0 * inv_sum * v_descale if not merged_invalid else Float32(0.0) + ) + output_scale1 = ( + alpha1 * inv_sum * v_descale if not merged_invalid else Float32(0.0) + ) + if const_expr(self.p_in_smem): + correction_threads = ( + len(self.correction_warp_ids) * cute.arch.WARP_SIZE + ) + cute.arch.barrier( + barrier_id=int(NamedBarrierFwdSm100.CorrectionScale), + number_of_threads=correction_threads, + ) + if tidx < self.m_block_size: + sScale[row_tidx] = output_scale0 + sScale[row_tidx + self.m_block_size] = output_scale1 + cute.arch.barrier( + barrier_id=int(NamedBarrierFwdSm100.CorrectionScale), + number_of_threads=correction_threads, + ) + + pipeline_o_epi.producer_acquire_w_index_phase( + 0, corr_epi_producer_phase + ) + self.correction_merge_epilogue( + thr_mma_pv, + tOtO[None, None, None, 0], + tOtO[None, None, None, 1], + tidx, + output_scale0, + output_scale1, + has_stream0, + has_stream1, + sO[None, None, 0], + sScale=sScale if const_expr(self.p_in_smem) else None, + ) + if has_stream0: + pipeline_s_p_o.consumer_release_w_index(0) + o_corr_phase0 ^= 1 + if has_stream1: + pipeline_s_p_o.consumer_release_w_index(1) + o_corr_phase1 ^= 1 + pipeline_o_epi.producer_commit_w_index(0) + corr_epi_producer_phase ^= 1 + + if const_expr(mLSE is not None): + mLSE_cur = mLSE[None, head_idx, batch_idx] + gLSE = cute.local_tile(mLSE_cur, (self.m_block_size,), (m_block,)) + lse = -Float32.inf + if not merged_invalid: + lse = ( + merged_row_max * softmax_scale_log2_eff + + cute.math.log2(merged_row_sum, fastmath=True) + - max_offset + ) * Float32(math.log(2.0)) + if ( + tidx < self.m_block_size + and row_tidx < seqlen.seqlen_q - m_block * self.m_block_size + ): + gLSE[row_tidx] = lse + + work_tile = tile_scheduler.advance_to_next_work() + + pipeline_o_epi.producer_acquire_w_index_phase( + 0, corr_epi_producer_phase + ) + + + @cute.jit + def correction_loop_dual_stream_spdb( + self, + thr_mma_qk: cute.ThrMma, + thr_mma_pv: cute.ThrMma, + tStS: cute.Tensor, + tOtO: cute.Tensor, + sScale: cute.Tensor, + mO: cute.Tensor, + mLSE: cute.Tensor, + sO: cute.Tensor, + pipeline_s_p_o: pipeline.PipelineAsync, + pipeline_o_acc: pipeline.PipelineAsync, + pipeline_sm_stats: pipeline.PipelineAsync, + sm_stats_barrier: pipeline.NamedBarrier, + pipeline_o_epi: pipeline.PipelineAsync, + learnable_sink: Optional[cute.Tensor], + descale_tensors: Optional[DescaleTensors], + gmem_tiled_copy_O: cute.TiledCopy, + tma_atom_O: cute.CopyAtom, + softmax_scale_log2: Float32, + block_info: BlockInfo, + num_splits: Int32, + SeqlenInfoCls: Callable, + blocksparse_tensors: Optional[BlockSparseTensors] = None, + tile_scheduler=None, + ): + """Correct and merge two independently-normalized sparse KV streams.""" + assert const_expr(self.num_kv_streams == 2) + assert const_expr(self.q_stage == 1) + assert const_expr(self.use_block_sparsity) + assert const_expr(self.use_tma_O) + assert const_expr(not self.use_correction_warps_for_epi) + assert const_expr(not self.pack_gqa) + assert const_expr(not self.is_split_kv) + assert const_expr(not self.is_varlen_q) + assert const_expr(learnable_sink is None) + + tidx = cute.arch.thread_idx()[0] % ( + cute.arch.WARP_SIZE * len(self.correction_warp_ids) + ) + row_tidx = tidx % self.m_block_size + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4 + + # Seed the first P/O acquire for each stream. Thereafter a stream is + # released only when it has a real PV event; empty tiles leave this + # seed untouched for the next active tile. + pipeline_s_p_o.consumer_release_w_index(0) + pipeline_s_p_o.consumer_release_w_index(1) + + o_corr_phase0 = Int32(0) + o_corr_phase1 = Int32(0) + sp_corr_count0 = Int32(0) + sp_corr_count1 = Int32(0) + corr_epi_producer_phase = Int32(1) + + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + kv_head_idx = self._kv_head_idx(head_idx) + qk_descale, v_descale = self._load_effective_descales( + descale_tensors, batch_idx, kv_head_idx + ) + softmax_scale_log2_eff = ( + softmax_scale_log2 * qk_descale + if const_expr(self.score_mod is None) + else softmax_scale_log2 + ) + max_offset = ( + Float32(8.0) + if const_expr(self.q_dtype.width == 8) + else Float32(0.0) + ) + seqlen = SeqlenInfoCls(batch_idx) + stream0_count, stream1_count = get_dual_stream_block_counts( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + split_idx, + num_splits, + 1, + self.q_subtile_factor, + seqlen, + ) + has_stream0 = stream0_count > Int32(0) + has_stream1 = stream1_count > Int32(0) + + # Every stream emits one first signal: the first score block when + # active, or a synthetic empty-stream signal otherwise. + sm_stats_barrier.arrive_and_wait_w_index( + index=0 if self.p_in_smem else 0 * 4 + warp_idx + ) + pipeline_sm_stats.consumer_release_w_index(0) + if has_stream0: + sp_corr_count0 += Int32(1) + sm_stats_barrier.arrive_and_wait_w_index( + index=1 if self.p_in_smem else 1 * 4 + warp_idx + ) + pipeline_sm_stats.consumer_release_w_index(1) + if has_stream1: + sp_corr_count1 += Int32(1) + + # Signals 1..count-1 carry the online-softmax correction scale. + if stream0_count > Int32(1): + for corr_idx in cutlass.range(stream0_count - Int32(1), unroll=1): + sm_stats_barrier.arrive_and_wait_w_index( + index=0 if self.p_in_smem else 0 * 4 + warp_idx + ) + pipeline_o_acc.consumer_wait_w_index_phase(0, o_corr_phase0) + o_corr_phase0 ^= 1 + if const_expr(self.p_in_smem): + self.correction_rescale_smem_p( + thr_mma_pv, + tOtO[None, None, None, 0], + tidx, + sScale, + Int32(0), + ) + else: + scale0 = sScale[row_tidx] + if cute.arch.vote_ballot_sync(scale0 < Float32(1.0)) != 0: + self.correction_rescale( + thr_mma_pv, + tOtO[None, None, None, 0], + tidx, + scale0, + ) + sp_slot0 = Int32(0) + Int32(2) * (sp_corr_count0 % Int32(2)) + pipeline_s_p_o.consumer_release_w_index(sp_slot0) + sp_corr_count0 += Int32(1) + pipeline_sm_stats.consumer_release_w_index(0) + + if corr_idx + Int32(1) < stream1_count: + sm_stats_barrier.arrive_and_wait_w_index( + index=1 if self.p_in_smem else 1 * 4 + warp_idx + ) + pipeline_o_acc.consumer_wait_w_index_phase(1, o_corr_phase1) + o_corr_phase1 ^= 1 + if const_expr(self.p_in_smem): + self.correction_rescale_smem_p( + thr_mma_pv, + tOtO[None, None, None, 1], + tidx, + sScale, + Int32(1), + ) + else: + scale1 = sScale[row_tidx + self.m_block_size] + if cute.arch.vote_ballot_sync(scale1 < Float32(1.0)) != 0: + self.correction_rescale( + thr_mma_pv, + tOtO[None, None, None, 1], + tidx, + scale1, + ) + sp_slot1 = Int32(1) + Int32(2) * (sp_corr_count1 % Int32(2)) + pipeline_s_p_o.consumer_release_w_index(sp_slot1) + sp_corr_count1 += Int32(1) + pipeline_sm_stats.consumer_release_w_index(1) + + # Active streams emit one additional final-stat signal. Empty + # streams have already completed their sole synthetic signal above. + row_sum0 = Float32(0.0) + row_sum1 = Float32(0.0) + row_max0 = -Float32.inf + row_max1 = -Float32.inf + if has_stream0: + sm_stats_barrier.arrive_and_wait_w_index( + index=0 if self.p_in_smem else 0 * 4 + warp_idx + ) + row_sum0 = sScale[row_tidx] + row_max0 = sScale[ + row_tidx + self.num_compute_slots * self.m_block_size + ] + pipeline_sm_stats.consumer_release_w_index(0) + pipeline_o_acc.consumer_wait_w_index_phase(0, o_corr_phase0) + o_corr_phase0 ^= 1 + if has_stream1: + sm_stats_barrier.arrive_and_wait_w_index( + index=1 if self.p_in_smem else 1 * 4 + warp_idx + ) + row_sum1 = sScale[row_tidx + self.m_block_size] + row_max1 = sScale[ + row_tidx + + self.m_block_size + + self.num_compute_slots * self.m_block_size + ] + pipeline_sm_stats.consumer_release_w_index(1) + pipeline_o_acc.consumer_wait_w_index_phase(1, o_corr_phase1) + o_corr_phase1 ^= 1 + + valid0 = has_stream0 and row_sum0 != Float32(0.0) and row_sum0 == row_sum0 + valid1 = has_stream1 and row_sum1 != Float32(0.0) and row_sum1 == row_sum1 + merged_row_max = -Float32.inf + if valid0: + merged_row_max = row_max0 + if valid1: + merged_row_max = ( + max(merged_row_max, row_max1) if valid0 else row_max1 + ) + + alpha0 = Float32(0.0) + alpha1 = Float32(0.0) + if valid0: + alpha0 = cute.math.exp2( + (row_max0 - merged_row_max) * softmax_scale_log2_eff, + fastmath=True, + ) + if valid1: + alpha1 = cute.math.exp2( + (row_max1 - merged_row_max) * softmax_scale_log2_eff, + fastmath=True, + ) + merged_row_sum = alpha0 * row_sum0 + alpha1 * row_sum1 + merged_invalid = ( + merged_row_sum == Float32(0.0) or merged_row_sum != merged_row_sum + ) + inv_sum = cute.arch.rcp_approx( + merged_row_sum if not merged_invalid else Float32(1.0) + ) + output_scale0 = ( + alpha0 * inv_sum * v_descale if not merged_invalid else Float32(0.0) + ) + output_scale1 = ( + alpha1 * inv_sum * v_descale if not merged_invalid else Float32(0.0) + ) + if const_expr(self.p_in_smem): + correction_threads = ( + len(self.correction_warp_ids) * cute.arch.WARP_SIZE + ) + cute.arch.barrier( + barrier_id=int(NamedBarrierFwdSm100.CorrectionScale), + number_of_threads=correction_threads, + ) + if tidx < self.m_block_size: + sScale[row_tidx] = output_scale0 + sScale[row_tidx + self.m_block_size] = output_scale1 + cute.arch.barrier( + barrier_id=int(NamedBarrierFwdSm100.CorrectionScale), + number_of_threads=correction_threads, + ) + + pipeline_o_epi.producer_acquire_w_index_phase( + 0, corr_epi_producer_phase + ) + self.correction_merge_epilogue( + thr_mma_pv, + tOtO[None, None, None, 0], + tOtO[None, None, None, 1], + tidx, + output_scale0, + output_scale1, + has_stream0, + has_stream1, + sO[None, None, 0], + sScale=sScale if const_expr(self.p_in_smem) else None, + ) + if has_stream0: + next_sp_slot0 = Int32(0) + Int32(2) * (sp_corr_count0 % Int32(2)) + pipeline_s_p_o.consumer_release_w_index(next_sp_slot0) + if has_stream1: + next_sp_slot1 = Int32(1) + Int32(2) * (sp_corr_count1 % Int32(2)) + pipeline_s_p_o.consumer_release_w_index(next_sp_slot1) + pipeline_o_epi.producer_commit_w_index(0) + corr_epi_producer_phase ^= 1 + + if const_expr(mLSE is not None): + mLSE_cur = mLSE[None, head_idx, batch_idx] + gLSE = cute.local_tile(mLSE_cur, (self.m_block_size,), (m_block,)) + lse = -Float32.inf + if not merged_invalid: + lse = ( + merged_row_max * softmax_scale_log2_eff + + cute.math.log2(merged_row_sum, fastmath=True) + - max_offset + ) * Float32(math.log(2.0)) + if ( + tidx < self.m_block_size + and row_tidx < seqlen.seqlen_q - m_block * self.m_block_size + ): + gLSE[row_tidx] = lse + + work_tile = tile_scheduler.advance_to_next_work() + + pipeline_o_epi.producer_acquire_w_index_phase( + 0, corr_epi_producer_phase + ) + + @cute.jit + def correction_merge_epilogue( + self, + thr_mma: cute.ThrMma, + tOtO0: cute.Tensor, + tOtO1: cute.Tensor, + tidx: Int32, + scale0: Float32, + scale1: Float32, + has_stream0: Boolean, + has_stream1: Boolean, + sO: cute.Tensor, + sScale: cute.Tensor | None = None, + ): + """Merge two FP32 TMEM numerators and convert/store only once.""" + corr_tile_size = 8 * 32 // self.o_dtype.width + tOsO = thr_mma.get_slice(0).partition_C(sO) + tOcO = thr_mma.partition_C(cute.make_identity_tensor(self.mma_tiler_pv[:2])) + + tile_layout = cute.make_layout((self.m_block_size, corr_tile_size)) + tOtO0_i = cute.logical_divide(tOtO0, tile_layout) + tOtO1_i = cute.logical_divide(tOtO1, tile_layout) + tOcO_i = cute.logical_divide(tOcO, tile_layout) + tOsO_i = cute.logical_divide(tOsO, tile_layout) + + epi_subtile = (self.epi_tile[0], corr_tile_size) + tmem_copy_atom = sm100_utils_basic.get_tmem_load_op( + self.mma_tiler_pv, + self.o_layout, + self.o_dtype, + self.pv_acc_dtype, + epi_subtile, + use_2cta_instrs=self.use_2cta_instrs, + ) + tiled_tmem_load = tcgen05.make_tmem_copy( + tmem_copy_atom, tOtO0_i[(None, None), 0] + ) + thr_tmem_load = tiled_tmem_load.get_slice(tidx) + if const_expr(self.p_in_smem): + smem_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.o_dtype + ) + else: + smem_copy_atom = sm100_utils_basic.get_smem_store_op( + self.o_layout, self.o_dtype, self.pv_acc_dtype, tiled_tmem_load + ) + tiled_smem_store = cute.make_tiled_copy_D(smem_copy_atom, tiled_tmem_load) + + tOtO0_t2r = thr_tmem_load.partition_S(tOtO0_i[(None, None), None]) + tOtO1_t2r = thr_tmem_load.partition_S(tOtO1_i[(None, None), None]) + tOsO_s2r = copy_utils.partition_D_position_independent( + thr_tmem_load, tOsO_i[(None, None), None] + ) + tOcO_t2r = thr_tmem_load.partition_D(tOcO_i[(None, None), None]) + use_per_elem_scale = const_expr(self.p_in_smem) and sScale is not None + + for i in cutlass.range( + self.head_dim_v_padded // corr_tile_size, unroll_full=True + ): + tOsO_r2s_i = tOsO_s2r[None, 0, 0, i] + tOrO = cute.make_rmem_tensor( + tOcO_t2r[None, 0, 0, i].shape, self.pv_acc_dtype + ) + tOcO_t2r_i = tOcO_t2r[None, 0, 0, i] + tOrO.fill(Float32(0.0)) + if has_stream0: + cute.copy( + tiled_tmem_load, tOtO0_t2r[None, 0, 0, i], tOrO + ) + if const_expr(use_per_elem_scale): + for k in cutlass.range(cute.size(tOrO), unroll_full=True): + row_idx = tOcO_t2r_i[k][0] + tOrO[k] = tOrO[k] * sScale[row_idx] + else: + for j in cutlass.range(0, cute.size(tOrO), 2, unroll_full=True): + tOrO[j], tOrO[j + 1] = cute.arch.mul_packed_f32x2( + (tOrO[j], tOrO[j + 1]), (scale0, scale0) + ) + if has_stream1: + tOrO1 = cute.make_rmem_tensor_like(tOrO, self.pv_acc_dtype) + cute.copy( + tiled_tmem_load, + tOtO1_t2r[None, 0, 0, i], + tOrO1, + ) + if const_expr(use_per_elem_scale): + for k in cutlass.range(cute.size(tOrO), unroll_full=True): + row_idx = tOcO_t2r_i[k][0] + tOrO[k] = tOrO[k] + tOrO1[k] * sScale[ + row_idx + self.m_block_size + ] + else: + for j in cutlass.range(0, cute.size(tOrO), 2, unroll_full=True): + tOrO[j], tOrO[j + 1] = cute.arch.fma_packed_f32x2( + (tOrO1[j], tOrO1[j + 1]), + (scale1, scale1), + (tOrO[j], tOrO[j + 1]), + ) + elif has_stream1: + cute.copy( + tiled_tmem_load, tOtO1_t2r[None, 0, 0, i], tOrO + ) + if const_expr(use_per_elem_scale): + for k in cutlass.range(cute.size(tOrO), unroll_full=True): + row_idx = tOcO_t2r_i[k][0] + tOrO[k] = tOrO[k] * sScale[ + row_idx + self.m_block_size + ] + else: + for j in cutlass.range(0, cute.size(tOrO), 2, unroll_full=True): + tOrO[j], tOrO[j + 1] = cute.arch.mul_packed_f32x2( + (tOrO[j], tOrO[j + 1]), (scale1, scale1) + ) + copy_utils.cvt_copy(tiled_smem_store, tOrO, tOsO_r2s_i) + + cute.arch.fence_view_async_tmem_load() + cute.arch.fence_view_async_shared() + + @cute.jit + def _softmax_tidx_to_row(self, tidx): + if cutlass.const_expr(self.p_in_smem): + return (tidx // 32) * 16 + tidx % 16 + return tidx + + @cute.jit + def correction_rescale_smem_p( + self, + thr_mma: cute.ThrMma, + tOtO: cute.Tensor, + tidx: Int32, + sScale: cute.Tensor, + stage: Int32, + ): + tOcO = thr_mma.partition_C(cute.make_identity_tensor(self.mma_tiler_pv[:2])) + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld16x256bOp(tcgen05.copy.Repetition(2)), + self.pv_acc_dtype, + ) + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St16x256bOp(tcgen05.copy.Repetition(2)), + self.pv_acc_dtype, + ) + thr_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tOtO).get_slice(tidx) + thr_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tOtO).get_slice(tidx) + tOtO_t2r = thr_tmem_load.partition_S(tOtO) + tOcO_t2r = thr_tmem_load.partition_D(tOcO) + tOtO_r2t = thr_tmem_store.partition_D(tOtO) + # An M64 Ld16x256 fragment spans exactly two unique output rows. + # Online-softmax correction scales are <= 1, so skip the TMEM + # round-trip only when every lane's two row scales remain identity. + row_scale0 = sScale[tOcO_t2r[0][0] + stage * self.m_block_size] + row_scale1 = sScale[tOcO_t2r[2][0] + stage * self.m_block_size] + should_rescale = cute.arch.vote_ballot_sync( + min(row_scale0, row_scale1) < Float32(1.0) + ) != 0 + if should_rescale: + tOrO = cute.make_rmem_tensor(tOcO_t2r.shape, self.pv_acc_dtype) + cute.copy(thr_tmem_load, tOtO_t2r, tOrO) + cute.arch.fence_view_async_tmem_load() + for i in cutlass.range(cute.size(tOrO), unroll_full=True): + row_idx = tOcO_t2r[i][0] + tOrO[i] = tOrO[i] * sScale[ + row_idx + stage * self.m_block_size + ] + cute.copy(thr_tmem_store, tOrO, tOtO_r2t) + cute.arch.fence_view_async_tmem_store() + + @cute.jit + def correction_rescale( + self, + thr_mma: cute.ThrMma, + tOtO: cute.Tensor, + tidx: Int32, + scale: Float32, + ): + """Rescale intermediate attention results based on softmax normalization factor. + + This method performs a crucial correction step in the attention computation pipeline. + When processing attention in blocks, the softmax normalization factors may change + as new blocks are processed. This method rescales previously computed partial + output values to account for updated normalization factors. + + The implementation uses efficient tensor memory operations to: + 1. Load existing partial attention output from tensor memory + 2. Apply the scaling factor to all elements + 3. Store the rescaled results back to tensor memory + """ + tOcO = thr_mma.partition_C(cute.make_identity_tensor(self.mma_tiler_pv[:2])) + corr_tile_size = 16 # tuneable parameter + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.pv_acc_dtype + ) + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), + self.pv_acc_dtype, + ) + tOtO_i = cute.composition(tOtO, cute.make_layout((self.m_block_size, corr_tile_size))) + tOcO_i = cute.composition(tOcO, cute.make_layout((self.m_block_size, corr_tile_size))) + thr_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_i).get_slice(tidx) + thr_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tOtO_i).get_slice(tidx) + tOtO_t2r = thr_tmem_load.partition_S(tOtO_i) + tOrO_t2r_shape = thr_tmem_load.partition_D(tOcO_i).shape + tOtO_r2t = thr_tmem_store.partition_D(tOtO_i) + + frg_count = self.head_dim_v_padded // corr_tile_size + tOrO_frg = cute.make_rmem_tensor((tOrO_t2r_shape, frg_count), self.pv_acc_dtype) + for i in cutlass.range_constexpr(frg_count): + tOrO_frg = cute.make_rmem_tensor(tOrO_t2r_shape, self.pv_acc_dtype) + tOtO_t2r_i = cute.make_tensor(tOtO_t2r.iterator + i * corr_tile_size, tOtO_t2r.layout) + cute.copy(thr_tmem_load, tOtO_t2r_i, tOrO_frg) + for j in cutlass.range(0, cute.size(tOrO_frg), 2, unroll_full=True): + tOrO_frg[j], tOrO_frg[j + 1] = cute.arch.mul_packed_f32x2( + (tOrO_frg[j], tOrO_frg[j + 1]), (scale, scale) + ) + tOtO_r2t_i = cute.make_tensor(tOtO_r2t.iterator + i * corr_tile_size, tOtO_r2t.layout) + cute.copy(thr_tmem_store, tOrO_frg, tOtO_r2t_i) + cute.arch.fence_view_async_tmem_store() + + @cute.jit + def correction_epilogue( + self, + thr_mma: cute.ThrMma, + tOtO: cute.Tensor, + tidx: Int32, + stage: Int32, + m_block: Int32, + seqlen_q: Int32, + scale: Float32, + sO: cute.Tensor, + mO_cur: Optional[cute.Tensor] = None, + gO: Optional[cute.Tensor] = None, + gmem_tiled_copy_O: Optional[cute.TiledCopy] = None, + sScale: cute.Tensor | None = None, + zero_output: bool = False, + ): + """Apply final scaling and transformation to attention output before writing to global memory. + + This correction_epilogue function handles the final processing step for attention output values. + It applies a scaling factor to the accumulated attention results and prepares the + data for efficient transfer back to global memory. + + The method performs: + 1. Loading of accumulated attention results from tensor memory + 2. Application of the final output scaling factor + 3. Type conversion if necessary (typically from higher precision accumulator to output precision) + 4. Reorganization of data for optimal memory access patterns + 5. Preparation for efficient TMA store operations + + :param thr_mma: Thread MMA operation for the computation + :type thr_mma: cute.ThrMma + :param tOtO: Tensor containing accumulated attention output + :type tOtO: cute.Tensor + :param scale: Final scaling factor to apply to the output + :type scale: Float32 + :param sO: Shared memory tensor for the final output + :type sO: cute.Tensor + """ + + corr_tile_size = 8 * 32 // self.o_dtype.width + # Use CTA 0 mapping for smem partitioning since sO is per-CTA sized + tOsO = thr_mma.get_slice(0).partition_C(sO) + tOcO = thr_mma.partition_C(cute.make_identity_tensor(self.mma_tiler_pv[:2])) + + tOtO_i = cute.logical_divide(tOtO, cute.make_layout((self.m_block_size, corr_tile_size))) + tOcO_i = cute.logical_divide(tOcO, cute.make_layout((self.m_block_size, corr_tile_size))) + tOsO_i = cute.logical_divide(tOsO, cute.make_layout((self.m_block_size, corr_tile_size))) + + epi_subtile = (self.epi_tile[0], corr_tile_size) + tmem_copy_atom = sm100_utils_basic.get_tmem_load_op( + self.mma_tiler_pv, + self.o_layout, + self.o_dtype, + self.pv_acc_dtype, + epi_subtile, + use_2cta_instrs=self.use_2cta_instrs, + ) + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_i[(None, None), 0]) + thr_tmem_load = tiled_tmem_load.get_slice(tidx) + if const_expr(self.p_in_smem): + smem_copy_atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), self.o_dtype + ) + else: + smem_copy_atom = sm100_utils_basic.get_smem_store_op( + self.o_layout, self.o_dtype, self.pv_acc_dtype, tiled_tmem_load + ) + tiled_smem_store = cute.make_tiled_copy_D(smem_copy_atom, tiled_tmem_load) + + tOtO_t2r = thr_tmem_load.partition_S(tOtO_i[(None, None), None]) + tOsO_s2r = copy_utils.partition_D_position_independent(thr_tmem_load, tOsO_i[(None, None), None]) + tOcO_t2r = thr_tmem_load.partition_D(tOcO_i[(None, None), None]) + use_per_elem_scale = const_expr(self.p_in_smem) and sScale is not None + for i in cutlass.range(self.head_dim_v_padded // corr_tile_size, unroll_full=True): + tOtO_t2r_i = tOtO_t2r[None, 0, 0, i] + tOsO_r2s_i = tOsO_s2r[None, 0, 0, i] + tOcO_t2r_i = tOcO_t2r[None, 0, 0, i] + tOrO_frg = cute.make_rmem_tensor(tOcO_t2r_i.shape, self.pv_acc_dtype) + if const_expr(zero_output): + tOrO_frg.fill(0.0) + else: + cute.copy(tiled_tmem_load, tOtO_t2r_i, tOrO_frg) + if const_expr(use_per_elem_scale): + for k in cutlass.range(cute.size(tOrO_frg), unroll_full=True): + row_idx = tOcO_t2r_i[k][0] + tOrO_frg[k] = ( + tOrO_frg[k] + * sScale[row_idx + stage * self.m_block_size] + ) + else: + for j in cutlass.range(0, cute.size(tOrO_frg), 2, unroll_full=True): + tOrO_frg[j], tOrO_frg[j + 1] = cute.arch.mul_packed_f32x2( + (tOrO_frg[j], tOrO_frg[j + 1]), (scale, scale) + ) + copy_utils.cvt_copy(tiled_smem_store, tOrO_frg, tOsO_r2s_i) + cute.arch.fence_view_async_shared() + + if const_expr(self.use_correction_warps_for_epi): + assert(not self.use_tma_O) + assert(gmem_tiled_copy_O is not None) + cute.arch.barrier(barrier_id=int(NamedBarrierFwdSm100.Epilogue), + number_of_threads=len(self.epilogue_warp_ids) * cute.arch.WARP_SIZE) + mma_tile_coord_v = thr_mma.thr_idx + m_tile_idx = (m_block * self.q_stage + stage) * self.cta_group_size + mma_tile_coord_v + self._store_O_to_gmem( + sO, gO, mO_cur, gmem_tiled_copy_O, tidx, seqlen_q, m_tile_idx + ) + + @cute.jit + def _store_O_to_gmem( + self, + sO_stage: cute.Tensor, + gO: Optional[cute.Tensor], + mO_cur: cute.Tensor, + gmem_tiled_copy_O: cute.TiledCopy, + tidx: Int32, + seqlen_q: Int32, + m_tile_idx: Int32, + ): + """Copy a single stage of O from smem to gmem via registers.""" + gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx) + tOsO = gmem_thr_copy_O.partition_S(sO_stage) + cO = cute.make_identity_tensor((self.m_block_size, self.head_dim_v_padded)) + tOcO = gmem_thr_copy_O.partition_S(cO) + t0OcO = gmem_tiled_copy_O.get_slice(0).partition_S(cO) + tOpO = copy_utils.predicate_k(tOcO, limit=mO_cur.shape[1]) + pack_gqa = PackGQA( + self.m_block_size, + self.head_dim_v_padded, + self.check_hdim_v_oob, + self.qhead_per_kvhead, + ) + + # load acc O from smem to rmem for wider vectorization + tOrO = cute.make_fragment_like(tOsO, self.o_dtype) + cute.autovec_copy(tOsO, tOrO) + # copy acc O from rmem to gmem + if const_expr(not self.pack_gqa): + assert gO is not None + tOgO = gmem_thr_copy_O.partition_D(gO) + for rest_m in cutlass.range_constexpr(cute.size(tOrO.shape[1])): + if ( + t0OcO[0, rest_m, 0][0] < seqlen_q - m_tile_idx * self.m_block_size - tOcO[0][0] + ): + cute.copy( + gmem_tiled_copy_O, + tOrO[None, rest_m, None], + tOgO[None, rest_m, None], + pred=tOpO[None, rest_m, None] + if const_expr(self.check_hdim_v_oob) + else None, + ) + else: + pack_gqa.store_O( + mO_cur, tOrO, gmem_tiled_copy_O, tidx, m_tile_idx, seqlen_q + ) + + @cute.jit + def epilogue_s2g( + self, + mO: cute.Tensor, + sO: cute.Tensor, + gmem_tiled_copy_O: cute.TiledCopy, + tma_atom_O: Optional[cute.CopyAtom], + pipeline_o_epi: pipeline.PipelineAsync, + block_info: BlockInfo, + num_splits: int, + SeqlenInfoCls: Callable, + mma_tile_coord_v: Int32 = 0, + blocksparse_tensors: Optional[BlockSparseTensors] = None, + tile_scheduler=None, + ): + epi_consumer_phase = Int32(0) + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + m_block, head_idx, batch_idx, split_idx = work_tile.tile_idx + seqlen = SeqlenInfoCls(batch_idx) + n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block, split_idx, num_splits) + has_work = const_expr(self.use_block_sparsity or not self.is_split_kv) or n_block_min < n_block_max + + if has_work: + if const_expr(self.is_split_kv): + mO_cur = seqlen.offset_batch_Q(mO, batch_idx, dim=3)[None, None, head_idx, split_idx] + else: + mO_cur = seqlen.offset_batch_Q(mO, batch_idx, dim=3)[None, None, head_idx] + gO = None + if const_expr(self.use_tma_O or not self.pack_gqa): + tiler_gO = ((self.mma_tiler_pv[0] * self.q_stage), self.head_dim_v_padded) + gO = cute.local_tile(mO_cur, tiler_gO, (m_block, 0)) # (128 * 2, 128) + gO = layout_utils.select( + cute.flat_divide(gO, (self.mma_tiler_pv[0],)), mode=[0, 2, 1] + ) # (128, 128, 2) + gO = cute.flat_divide(gO, (self.mma_tiler_pv[0] // self.cta_group_size,))[None, mma_tile_coord_v, None, None] + + if const_expr(self.use_tma_O): + store_O, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_O, 0, cute.make_layout(1), sO, gO + ) + for stage in cutlass.range(self.q_stage, unroll_full=True): + # wait from corr, issue tma store on smem + # 1. wait for O0 / O1 final + pipeline_o_epi.consumer_wait_w_index_phase(stage, epi_consumer_phase) + # 2. copy O0 / O1 to gmem + store_O(src_idx=stage, dst_idx=stage) + cute.arch.cp_async_bulk_commit_group() + for stage in cutlass.range_constexpr(self.q_stage): + # Ensure O0 / O1 buffer is ready to be released + cute.arch.cp_async_bulk_wait_group(self.q_stage - 1 - stage, read=True) + pipeline_o_epi.consumer_release_w_index(stage) + else: + tidx = cute.arch.thread_idx()[0] % ( + cute.arch.WARP_SIZE * len(self.epilogue_warp_ids) + ) + for stage in cutlass.range_constexpr(self.q_stage): + # wait from corr, issue tma store on smem + # 1. wait for O0 / O1 final + pipeline_o_epi.consumer_wait_w_index_phase(stage, epi_consumer_phase) + # 2. copy O0 / O1 to gmem + m_tile_idx = (m_block * self.q_stage + stage) * self.cta_group_size + mma_tile_coord_v + gO_stage = gO[None, None, stage] if const_expr(gO is not None) else None + self._store_O_to_gmem( + sO[None, None, stage], gO_stage, mO_cur, gmem_tiled_copy_O, + tidx, seqlen.seqlen_q, m_tile_idx, + ) + pipeline_o_epi.consumer_release_w_index(stage) + + epi_consumer_phase ^= 1 + + # Advance to next tile + work_tile = tile_scheduler.advance_to_next_work() + + @cute.jit + def clc_scheduler_warp( + self, + tile_scheduler: TileSchedulerProtocol, + ): + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + tile_scheduler.prefetch_next_work() + work_tile = tile_scheduler.advance_to_next_work() + if cute.arch.thread_idx()[0] == self.clc_scheduler_warp_id * cute.arch.WARP_SIZE: + fa_printf( + 3, + "[CLC] query sm={} cta={} (m_blk={},h={},b={},s={}) valid={}\n", + smid(), + cute.arch.block_idx()[0], + work_tile.tile_idx[0], + work_tile.tile_idx[1], + work_tile.tile_idx[2], + work_tile.tile_idx[3], + work_tile.is_valid_tile, + ) + tile_scheduler.producer_tail() + + @cute.jit + def empty_warp( + self, + tile_scheduler: TileSchedulerProtocol, + ): + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + work_tile = tile_scheduler.advance_to_next_work() + + def load_Q( + self, + load_Q_fn: Callable, + pipeline_q: pipeline.PipelineAsync, + block: Int32, + stage: int, + phase: Int32, + ): + pipeline_q.producer_acquire_w_index_phase(stage, phase) + load_Q_fn(src_idx=block, dst_idx=stage, tma_bar_ptr=pipeline_q.sync_object_full.get_barrier(stage)) + + def load_Q_non_tma( + self, + mQ: cute.Tensor, + sQ: cute.Tensor, + gmem_tiled_copy_Q: cute.TiledCopy, + pipeline_q: pipeline.PipelineAsync, + tidx: Int32, + seqlen_q: Int32, + m_block: Int32, + block: Int32, + stage: int, + phase: Int32, + ): + assert self.cta_group_size == 1, "cta_group_size must be 1 for non-tma Q load" + pipeline_q.producer_acquire_w_index_phase(stage, phase) + pack_gqa = PackGQA( + self.m_block_size, + self.head_dim_padded, + self.check_hdim_oob, + self.qhead_per_kvhead, + ) + sQ_stage = sQ[None, None, None, stage] + sQ_pi = cute.make_tensor( + sQ_stage.iterator, + cute.make_layout( + (sQ_stage.shape[0][0], (sQ_stage.shape[0][1], sQ_stage.shape[2])), + stride=(sQ_stage.stride[0][0], (sQ_stage.stride[0][1], sQ_stage.stride[2])), + ), + ) + pack_gqa.load_Q(mQ, sQ_pi, gmem_tiled_copy_Q, tidx, m_block * self.q_stage + block, seqlen_q) + cute.arch.cp_async_commit_group() + pipeline_q.sync_object_full.arrive_cp_async_mbarrier(stage) + + @cute.jit + def load_KV( + self, + tma_atom: Optional[cute.CopyAtom], + tXgX: Optional[cute.Tensor], + tXsX: Optional[cute.Tensor], + paged_kv_manager: Optional[PagedKVManager], + sX: cute.Tensor, + block: Int32, + pipeline_kv: pipeline.PipelineAsync, + producer_state: pipeline.PipelineState, + K_or_V: Literal["K", "V"], + page_idx: Optional[Int32] = None, + extra_tx_count: Optional[Int32] = None, + ): + assert K_or_V in ("K", "V") + stage, phase = producer_state.index, producer_state.phase + extra_tx_count_kv = self.tma_copy_bytes[K_or_V] - self.tma_copy_bytes["K"] + extra_tx_count = ( + extra_tx_count_kv + (extra_tx_count if extra_tx_count is not None else 0) if const_expr(self.use_tma_KV) + else None + ) + extra_kwargs = {"extra_tx_count": extra_tx_count} if const_expr(self.use_tma_KV) else {} + pipeline_kv.producer_acquire(producer_state, **extra_kwargs) + if const_expr(K_or_V == "K" and self.uneven_kv_smem): + # Before this round, the smem location was occupied by V, which is smaller than + # K. So we need to wait for the stage after that (stage 1) to be empty as well. + if stage == 0: + pipeline_kv.sync_object_empty.wait(1, phase) + + if const_expr(self.use_tma_KV): + assert tXgX is not None and tXsX is not None and tma_atom is not None + tXsX_cur = tXsX[None, stage] + if const_expr(self.uneven_kv_smem): + # Since this is the producer_state, the phase starts at 1, so we have to invert it + tXsX_cur = self.offset_kv_smem(tXsX_cur, stage, phase ^ 1) + # Currently we assume that page_size == n_block_size so we index into tXgX with block = 0 + tXgX_cur = tXgX[None, block] if const_expr(page_idx is None) else tXgX[None, 0, page_idx] + cute.copy(tma_atom, tXgX_cur, tXsX_cur, tma_bar_ptr=pipeline_kv.producer_get_barrier(producer_state)) + else: + assert paged_kv_manager is not None + assert extra_tx_count is None + sX_cur = sX[None, None, None, stage] + if const_expr(self.uneven_kv_smem): + sX_cur = self.offset_kv_smem(sX_cur, stage, phase ^ 1) + paged_kv_manager.load_KV(block, sX_cur, K_or_V) + cute.arch.cp_async_commit_group() + pipeline_kv.sync_object_full.arrive_cp_async_mbarrier(stage) + + @cute.jit + def offset_kv_smem(self, sX: cute.Tensor, stage: Int32, phase: Int32): + if const_expr(self.uneven_kv_smem): + # smem layout is [smem_large, smem_small, smem_large], and the current stride is + # (smem_large + smem_small) // 2. So for stage == 1, move right by offset if + # phase == 0, or left by offset if phase == 1. + offset = 0 if stage != 1 else self.uneven_kv_smem_offset * (1 - 2 * phase) + # Hint that the offset is 128-bit aligned so that + # ptr + offset preserves the alignment needed by cp.async. + offset = cute.assume(offset, divby=128 // self.k_dtype.width) + return cute.make_tensor(sX.iterator + offset, sX.layout) + else: + return sX + + # @cute.jit + # def warp_scheduler_barrier_init(self): + # warp_group_idx = utils.canonical_warp_group_idx(sync=False) + # if warp_group_idx == 0: + # cute.arch.barrier_arrive( + # barrier_id=int(NamedBarrierFwdSm100.WarpSchedulerWG1), number_of_threads=2 * 128, + # ) + + # def warp_scheduler_barrier_sync(self): + # cute.arch.barrier( + # barrier_id=int(NamedBarrierFwdSm100.WarpSchedulerWG1) + utils.canonical_warp_group_idx(sync=False), + # number_of_threads=2 * 128 + # ) + + # def warp_scheduler_barrier_arrive(self): + # cur_wg = utils.canonical_warp_group_idx(sync=False) + # next_wg = 1 - cur_wg + # cute.arch.barrier_arrive( + # barrier_id=int(NamedBarrierFwdSm100.WarpSchedulerWG1) + next_wg, number_of_threads=2 * 128, + # ) + + @cute.jit + def apply_score_mod( + self, + tSrS_t2r, + thr_tmem_load, + thr_mma_qk, + batch_idx, + head_idx, + m_block, + n_block, + softmax, + seqlen: SeqlenInfoQK, + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + head_divmod=None, + ): + """Apply score modification for SM100 (constant q_idx).""" + # Prepare index tensor with extra partition + cS = cute.make_identity_tensor((self.m_block_size, self.n_block_size)) + cS = cute.domain_offset((m_block * self.m_block_size, n_block * self.n_block_size), cS) + tScS = thr_mma_qk.partition_C(cS) + tScS = tScS[(None, None), 0, 0] + tScS_t2r = thr_tmem_load.partition_D(tScS) + + # Shared q_idx for all scores + q_idx_logical = tScS_t2r[0][0] + + # For Pack-GQA, compute the logical head index for this tile + if cutlass.const_expr(self.pack_gqa): + assert head_divmod is not None + # Building up the logical q_head idx: final_q_head = kv_head * qhead_per_kvhead + (q_physical % qhead_per_kvhead) + q_physical = q_idx_logical + q_idx_logical, head_offset = divmod(q_physical, head_divmod) + head_idx = head_idx * self.qhead_per_kvhead + head_offset + + if cutlass.const_expr(aux_data.tensors is not None): + seqlen_q_divmod, _ = fastdiv_mods + _, q_idx_logical = divmod(q_idx_logical, seqlen_q_divmod) + + apply_score_mod_inner( + tSrS_t2r, + tScS_t2r, + self.score_mod, + batch_idx, + head_idx, + softmax.softmax_scale, + self.score_vec_size, + self.qk_acc_dtype, + aux_data, + fastdiv_mods, + seqlen_info=seqlen, + constant_q_idx=q_idx_logical, + qhead_per_kvhead=self.qhead_per_kvhead if cutlass.const_expr(self.pack_gqa) else 1, + ) diff --git a/fastvideo-kernel/fa4/flash_fwd_sm120.py b/fastvideo-kernel/fa4/flash_fwd_sm120.py new file mode 100644 index 000000000..08d219acf --- /dev/null +++ b/fastvideo-kernel/fa4/flash_fwd_sm120.py @@ -0,0 +1,59 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +# SM120 (Blackwell GeForce / DGX Spark) forward pass. +# +# SM120 uses the same SM80-era MMA instructions (mma.sync.aligned.m16n8k16) but has +# a smaller shared memory capacity (99 KB vs 163 KB on SM80). This module subclasses +# FlashAttentionForwardSm80 and overrides the SMEM capacity check accordingly. + +import cutlass +import cutlass.utils as utils_basic + +from flash_attn.cute.flash_fwd import FlashAttentionForwardSm80 + + +class FlashAttentionForwardSm120(FlashAttentionForwardSm80): + # Keep arch = 80 to use CpAsync code paths (no TMA for output). + # The compilation target is determined by the GPU at compile time, not this field. + arch = 80 + + @staticmethod + def can_implement( + dtype, + head_dim, + head_dim_v, + tile_m, + tile_n, + num_stages, + num_threads, + is_causal, + Q_in_regs=False, + ) -> bool: + """Check if the kernel can be implemented on SM120. + + Same logic as SM80 but uses SM120's shared memory capacity (99 KB). + """ + if dtype not in [cutlass.Float16, cutlass.BFloat16]: + return False + if head_dim % 8 != 0: + return False + if head_dim_v % 8 != 0: + return False + if tile_n % 16 != 0: + return False + if num_threads % 32 != 0: + return False + # Shared memory usage: Q tile + (K tile + V tile) + smem_usage_Q = tile_m * head_dim * 2 + smem_usage_K = tile_n * head_dim * num_stages * 2 + smem_usage_V = tile_n * head_dim_v * num_stages * 2 + smem_usage_QV = ( + (smem_usage_Q + smem_usage_V) if not Q_in_regs else max(smem_usage_Q, smem_usage_V) + ) + smem_usage = smem_usage_QV + smem_usage_K + # SM120 has 99 KB shared memory (vs 163 KB on SM80) + smem_capacity = utils_basic.get_smem_capacity_in_bytes("sm_120") + if smem_usage > smem_capacity: + return False + if (tile_m * 2) % num_threads != 0: + return False + return True diff --git a/fastvideo-kernel/fa4/flash_fwd_sm90.py b/fastvideo-kernel/fa4/flash_fwd_sm90.py new file mode 100644 index 000000000..91acd286d --- /dev/null +++ b/fastvideo-kernel/fa4/flash_fwd_sm90.py @@ -0,0 +1,1545 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +# SM90 (Hopper) forward pass for flash attention, extracted from flash_fwd.py. + +from types import SimpleNamespace +from typing import Callable, Literal, Optional +from functools import partial + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, const_expr +from cutlass.cute.nvgpu import cpasync, warpgroup +from cutlass.utils import LayoutEnum +import cutlass.utils.hopper_helpers as sm90_utils_basic +from cutlass import pipeline +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.base_dsl.arch import Arch + +from quack import copy_utils +from quack import layout_utils +from quack import sm90_utils + +from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned +from flash_attn.cute import utils +from flash_attn.cute.mask import AttentionMask +from flash_attn.cute.softmax import Softmax, apply_score_mod_inner +from flash_attn.cute.seqlen_info import SeqlenInfoQK +from flash_attn.cute.block_info import BlockInfo +from flash_attn.cute.block_sparsity import BlockSparseTensors +from flash_attn.cute.block_sparse_utils import ( + produce_block_sparse_loads, + consume_block_sparse_loads, +) +from flash_attn.cute import pipeline as pipeline_custom +from flash_attn.cute.pack_gqa import PackGQA, pack_gqa_layout, make_packgqa_tiled_tma_atom +from flash_attn.cute.paged_kv import PagedKVManager +from flash_attn.cute.named_barrier import NamedBarrierFwd +from quack.cute_dsl_utils import ParamsBase +from flash_attn.cute.tile_scheduler import ( + TileSchedulerArguments, + SingleTileScheduler, + SingleTileLPTScheduler, + SingleTileVarlenScheduler, +) +from cutlass.cute import FastDivmodDivisor + +from flash_attn.cute.flash_fwd import FlashAttentionForwardBase +from flash_attn.cute.utils import AuxData + + +class FlashAttentionForwardSm90(FlashAttentionForwardBase): + def __init__( + self, + *args, + intra_wg_overlap: bool = True, + mma_pv_is_rs: bool = True, + paged_kv_non_tma: bool = False, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.intra_wg_overlap = intra_wg_overlap + self.mma_pv_is_rs = mma_pv_is_rs + self.buffer_align_bytes = 1024 + self.use_tma_KV = not paged_kv_non_tma + assert self.use_tma_KV or not (self.check_hdim_oob or self.check_hdim_v_oob), ( + "Paged KV does not support irregular head dim" + ) + self.cluster_shape_mn = (1, 1) + assert self.arch.is_family_of(Arch.sm_90a), "Only SM 9.x is supported" + + def _get_smem_layout_atom(self): + sQ_layout_atom = warpgroup.make_smem_layout_atom( + sm90_utils_basic.get_smem_layout_atom(LayoutEnum.ROW_MAJOR, self.dtype, self.tile_hdim), + self.dtype, + ) + sK_layout_atom = sQ_layout_atom + sV_layout_atom = warpgroup.make_smem_layout_atom( + sm90_utils_basic.get_smem_layout_atom( + LayoutEnum.ROW_MAJOR, self.dtype, self.tile_hdimv + ), + self.dtype, + ) + sO_layout_atom = sV_layout_atom + if not self.mma_pv_is_rs: + sP_layout_atom = warpgroup.make_smem_layout_atom( + sm90_utils_basic.get_smem_layout_atom( + LayoutEnum.ROW_MAJOR, self.dtype, self.tile_n + ), + self.dtype, + ) + else: + sP_layout_atom = None + return sQ_layout_atom, sK_layout_atom, sV_layout_atom, sO_layout_atom, sP_layout_atom + + def _get_tiled_mma(self): + tiled_mma_qk = sm90_utils_basic.make_trivial_tiled_mma( + self.dtype, + self.dtype, + warpgroup.OperandMajorMode.K, + warpgroup.OperandMajorMode.K, + Float32, + atom_layout_mnk=(self.tile_m // 64, 1, 1), + tiler_mn=(64, self.tile_n), + ) + tiled_mma_pv = sm90_utils_basic.make_trivial_tiled_mma( + self.dtype, + self.dtype, + warpgroup.OperandMajorMode.K, + warpgroup.OperandMajorMode.MN, + Float32, + atom_layout_mnk=(self.tile_m // 64, 1, 1), # Might need (1, 2, 1) for hdim 512 + tiler_mn=(64, self.tile_hdimv), + a_source=warpgroup.OperandSource.RMEM + if self.mma_pv_is_rs + else warpgroup.OperandSource.SMEM, + ) + return tiled_mma_qk, tiled_mma_pv + + def _get_shared_storage_cls(self): + sQ_struct, sK_struct, sV_struct = [ + cute.struct.Align[ + cute.struct.MemRange[self.dtype, cute.cosize(layout)], self.buffer_align_bytes + ] + for layout in (self.sQ_layout, self.sK_layout, self.sV_layout) + ] + cosize_sQV = max(cute.cosize(self.sQ_layout), cute.cosize(self.sV_layout)) + sQV_struct = cute.struct.Align[cute.struct.MemRange[self.dtype, cosize_sQV], 1024] + cosize_sP = cute.cosize(self.sP_layout) if const_expr(self.sP_layout is not None) else 0 + sP_struct = cute.struct.Align[cute.struct.MemRange[self.dtype, cosize_sP], 1024] + # 1 stage * 2 for Q pipeline (full + empty), self.num_stages*2 for K, self.num_stages*2 for V, + mbar_ptr_Q_struct = cute.struct.MemRange[cutlass.Int64, 1 * 2] + mbar_ptr_K_struct = cute.struct.MemRange[cutlass.Int64, self.num_stages * 2] + mbar_ptr_V_struct = cute.struct.MemRange[cutlass.Int64, self.num_stages * 2] + + @cute.struct + class SharedStorageQKV: + mbar_ptr_Q: mbar_ptr_Q_struct + mbar_ptr_K: mbar_ptr_K_struct + mbar_ptr_V: mbar_ptr_V_struct + sV: sV_struct + sQ: sQ_struct + sK: sK_struct + sP: sP_struct + + @cute.struct + class SharedStorageSharedQV: + mbar_ptr_Q: mbar_ptr_Q_struct + mbar_ptr_K: mbar_ptr_K_struct + mbar_ptr_V: mbar_ptr_V_struct + sQ: sQV_struct + sK: sK_struct + sP: sP_struct + + return SharedStorageQKV if const_expr(not self.Q_in_regs) else SharedStorageSharedQV + + @cute.jit + def __call__( + self, + mQ: cute.Tensor, # (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q + mK: cute.Tensor, # (b_k, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k or (num_pages, page_size, h_k, d) if there is page_table + mV: cute.Tensor, # (b_k, s_k, h_k, dv) or (total_k, h_k, dv) if there is cu_seqlens_k or (num_pages, page_size, h_k, dv) if there is page_table + mO: cute.Tensor, # (b, s_q, h, dv) or (total_q, h, dv) if there is cu_seqlens_q + mLSE: Optional[cute.Tensor], + softmax_scale: Float32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + mPageTable: Optional[cute.Tensor] = None, # (b_k, max_num_pages_per_seq) + window_size_left: Int32 | int | None = None, + window_size_right: Int32 | int | None = None, + learnable_sink: Optional[cute.Tensor] = None, + blocksparse_tensors: Optional[BlockSparseTensors] = None, + aux_data: AuxData = AuxData(), + # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI). + stream: cuda.CUstream = None, + ): + """Configures and launches the flash attention kernel. + + mQ/mK/mV/mO has same data types(supports fp16 and bf16) and same layout: + (batch_size, seqlen_q, num_head, head_dim):(_, _, _, 1) + """ + + self._check_type( + *( + t.element_type if t is not None else None + for t in (mQ, mK, mV, mO, mLSE, mCuSeqlensQ, mCuSeqlensK, mSeqUsedQ, mSeqUsedK) + ) + ) + + self.varlen_q = mCuSeqlensQ is not None or mSeqUsedQ is not None + + mQ, mK, mV, mO = [assume_tensor_aligned(t) for t in (mQ, mK, mV, mO)] + QO_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensQ is None) else [0, 2, 1] + mQ, mO = [layout_utils.select(t, QO_layout_transpose) for t in (mQ, mO)] + KV_layout_transpose = [1, 3, 2, 0] if const_expr(mCuSeqlensK is None) else [0, 2, 1] + mK, mV = [layout_utils.select(t, KV_layout_transpose) for t in (mK, mV)] + LSE_layout_transpose = [2, 1, 0] if const_expr(mCuSeqlensQ is None) else [1, 0] + mLSE = ( + layout_utils.select(mLSE, LSE_layout_transpose) + if const_expr(mLSE is not None) + else None + ) + + tiled_mma_qk, tiled_mma_pv = self._get_tiled_mma() + self.num_mma_threads = tiled_mma_qk.size + self.num_threads_per_warp_group = 128 + self.num_wg_mma = self.num_mma_threads // self.num_threads_per_warp_group + assert self.num_wg_mma in [1, 2, 3] + self.num_threads = self.num_threads_per_warp_group * (self.num_wg_mma + 1) + self.num_producer_threads = 32 + self.num_Q_load_threads = self.num_threads_per_warp_group # If not TMA_Q + self.num_epilogue_threads = self.num_mma_threads + self.num_mma_regs, self.num_producer_regs = {1: (256, 56), 2: (240, 24), 3: (160, 32)}[ + self.num_wg_mma + ] + self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None) + + self.use_scheduler_barrier = ( + (self.num_wg_mma >= 2 and self.tile_hdim <= 128) + if const_expr(self.intra_wg_overlap) + else (self.num_wg_mma == 2) + ) + self.use_tma_Q = self.arch >= Arch.sm_90 and not ( + self.pack_gqa and self.tile_m % self.qhead_per_kvhead != 0 + ) + self.use_tma_O = self.use_tma_Q + # Producer needs more registers when doing cp.async Q or KV loads + if const_expr(self.num_wg_mma == 2 and (not self.use_tma_Q or not self.use_tma_KV)): + self.num_mma_regs, self.num_producer_regs = 224, 40 + self.rescale_O_before_gemm = self.tile_hdimv > 128 and self.intra_wg_overlap + self._setup_attributes() + # TODO: we prob don't need most of what's in _setup_attributes + self.sQ_layout, self.sK_layout, self.sV_layout, self.sO_layout = [ + sm90_utils.make_smem_layout(mX.element_type, LayoutEnum.ROW_MAJOR, shape, stage) + for mX, shape, stage in [ + (mQ, (self.tile_m, self.tile_hdim), None), + (mK, (self.tile_n, self.tile_hdim), self.num_stages), + (mV, (self.tile_n, self.tile_hdimv), self.num_stages), + (mO, (self.tile_m, self.tile_hdimv), None), + ] + ] + self.sP_layout = None + if const_expr(not self.mma_pv_is_rs): + self.sP_layout = sm90_utils.make_smem_layout( + mV.element_type, LayoutEnum.ROW_MAJOR, (self.tile_m, self.tile_n) + ) + + SharedStorage = self._get_shared_storage_cls() + + mQ_og, mO_og = mQ, mO + if const_expr(self.pack_gqa): + nheads_kv = mK.shape[2] + mQ = pack_gqa_layout(mQ, self.qhead_per_kvhead, nheads_kv, head_idx=2) + mO = pack_gqa_layout(mO, self.qhead_per_kvhead, nheads_kv, head_idx=2) + if const_expr(mLSE is not None): + mLSE = pack_gqa_layout(mLSE, self.qhead_per_kvhead, nheads_kv, head_idx=1) + + # TMA + gmem_tiled_copy_Q = cpasync.CopyBulkTensorTileG2SOp() + gmem_tiled_copy_KV = cpasync.CopyBulkTensorTileG2SOp() # Might multicast + gmem_tiled_copy_O = cpasync.CopyBulkTensorTileS2GOp() + self.tma_copy_bytes = { + name: cute.size_in_bytes(mX.element_type, cute.select(layout, mode=[0, 1])) + for name, mX, layout in [ + ("Q", mQ, self.sQ_layout), + ("K", mK, self.sK_layout), + ("V", mV, self.sV_layout), + ] + } + make_tiled_tma_atom_fn = ( + partial(make_packgqa_tiled_tma_atom, qhead_per_kvhead=self.qhead_per_kvhead, head_idx=2) + if const_expr(self.pack_gqa) + else cpasync.make_tiled_tma_atom + ) + tma_atom_Q, tma_tensor_Q = None, None + if const_expr(self.use_tma_Q): + tma_atom_Q, tma_tensor_Q = make_tiled_tma_atom_fn( + gmem_tiled_copy_Q, + mQ_og if const_expr(self.pack_gqa) else mQ, + self.sQ_layout, + (self.tile_m, self.tile_hdim), # No mcast + ) + tma_atom_K, tma_tensor_K = None, None + tma_atom_V, tma_tensor_V = None, None + if const_expr(self.use_tma_KV): + tma_atom_K, tma_tensor_K = cpasync.make_tiled_tma_atom( + gmem_tiled_copy_KV, + mK, + cute.select(self.sK_layout, mode=[0, 1]), + (self.tile_n, self.tile_hdim), + 1, # No mcast for now + ) + tma_atom_V, tma_tensor_V = cpasync.make_tiled_tma_atom( + gmem_tiled_copy_KV, + mV, + cute.select(self.sV_layout, mode=[0, 1]), + (self.tile_n, self.tile_hdimv), + 1, # No mcast for now + ) + tma_atom_O, tma_tensor_O = None, None + if const_expr(self.use_tma_O): + mO_tma = mO_og if const_expr(self.pack_gqa) else mO + if const_expr(self.varlen_q): + mO_tma = copy_utils.create_ragged_tensor_for_tma( + mO_tma, ragged_dim=0, ptr_shift=True + ) + tma_atom_O, tma_tensor_O = make_tiled_tma_atom_fn( + gmem_tiled_copy_O, + mO_tma, + self.sO_layout, + (self.tile_m, self.tile_hdimv), # No mcast + ) + if const_expr(mCuSeqlensQ is not None or mSeqUsedQ is not None): + TileScheduler = SingleTileVarlenScheduler + else: + TileScheduler = ( + SingleTileScheduler + if const_expr(not self.is_causal or self.is_local) + else SingleTileLPTScheduler + ) + tile_sched_args = TileSchedulerArguments( + cute.ceil_div(cute.size(mQ.shape[0]), self.tile_m), + cute.size(mQ.shape[2]), + cute.size(mQ.shape[3]) + if const_expr(mCuSeqlensQ is None) + else cute.size(mCuSeqlensQ.shape[0] - 1), + 1, # num_splits + cute.size(mK.shape[0]) + if const_expr(mPageTable is None) + else mK.shape[0] * mPageTable.shape[1], + mQ.shape[1], + mV.shape[1], + total_q=cute.size(mQ.shape[0]) + if const_expr(mCuSeqlensQ is not None) + else cute.size(mQ.shape[0]) * cute.size(mQ.shape[3]), + tile_shape_mn=(self.tile_m, self.tile_n), + mCuSeqlensQ=mCuSeqlensQ, + mSeqUsedQ=mSeqUsedQ, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + element_size=self.dtype.width // 8, + is_persistent=False, + lpt=self.is_causal or self.is_local, + ) + tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args) + grid_dim = TileScheduler.get_grid_shape(tile_sched_params) + softmax_scale_log2, softmax_scale = utils.compute_softmax_scale_log2( + softmax_scale, self.score_mod + ) + window_size_left = Int32(window_size_left) if window_size_left is not None else None + window_size_right = Int32(window_size_right) if window_size_right is not None else None + fastdiv_mods = utils.compute_fastdiv_mods( + mQ, mK, self.qhead_per_kvhead, self.pack_gqa, aux_data.tensors, mPageTable + ) + + self.kernel( + tma_tensor_Q if const_expr(self.use_tma_Q) else mQ, + tma_tensor_K if const_expr(self.use_tma_KV) else mK, + tma_tensor_V if const_expr(self.use_tma_KV) else mV, + tma_tensor_O if const_expr(self.use_tma_O) else mO, + mLSE, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + mPageTable, + tma_atom_Q, + tma_atom_K, + tma_atom_V, + tma_atom_O, + softmax_scale_log2, + softmax_scale, + window_size_left, + window_size_right, + learnable_sink, + blocksparse_tensors, + self.sQ_layout, + self.sK_layout, + self.sV_layout, + self.sO_layout, + self.sP_layout, + self.gmem_tiled_copy_Q, + self.gmem_tiled_copy_K, + self.gmem_tiled_copy_V, + self.gmem_tiled_copy_O, + tiled_mma_qk, + tiled_mma_pv, + tile_sched_params, + TileScheduler, + SharedStorage, + aux_data, + fastdiv_mods, + ).launch( + grid=grid_dim, + block=[self.num_threads, 1, 1], + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.kernel + def kernel( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + mCuSeqlensQ: Optional[cute.Tensor], + mCuSeqlensK: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + mSeqUsedK: Optional[cute.Tensor], + mPageTable: Optional[cute.Tensor], + tma_atom_Q: Optional[cute.CopyAtom], + tma_atom_K: Optional[cute.CopyAtom], + tma_atom_V: Optional[cute.CopyAtom], + tma_atom_O: Optional[cute.CopyAtom], + softmax_scale_log2: Float32, + softmax_scale: Optional[Float32], + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + learnable_sink: Optional[cute.Tensor], + blocksparse_tensors: Optional[BlockSparseTensors], + sQ_layout: cute.ComposedLayout, + sK_layout: cute.ComposedLayout, + sV_layout: cute.ComposedLayout, + sO_layout: cute.ComposedLayout, + sP_layout: cute.ComposedLayout | None, + gmem_tiled_copy_Q: cute.TiledCopy, + gmem_tiled_copy_K: cute.TiledCopy, + gmem_tiled_copy_V: cute.TiledCopy, + gmem_tiled_copy_O: cute.TiledCopy, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + tile_sched_params: ParamsBase, + TileScheduler: cutlass.Constexpr[Callable], + SharedStorage: cutlass.Constexpr[Callable], + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + # Prefetch tma descriptor + if warp_idx == 0: + for tma_atom in (tma_atom_Q, tma_atom_K, tma_atom_V, tma_atom_O): + if const_expr(tma_atom is not None): + cpasync.prefetch_descriptor(tma_atom) + + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(SharedStorage) + + # Mbarrier / pipeline init + mbar_ptr_Q = storage.mbar_ptr_Q.data_ptr() + + ThreadCooperativeGroup = partial(pipeline.CooperativeGroup, pipeline.Agent.Thread) + tma_warp = ThreadCooperativeGroup(1) + load_threads = ThreadCooperativeGroup(self.num_threads_per_warp_group) + mma_warps = ThreadCooperativeGroup(self.num_mma_threads // cute.arch.WARP_SIZE) + if const_expr(self.use_tma_Q): + pipeline_q = pipeline_custom.PipelineTmaAsync.create( + barrier_storage=mbar_ptr_Q, + num_stages=1, + producer_group=tma_warp, + consumer_group=mma_warps, + tx_count=self.tma_copy_bytes["Q"], + defer_sync=True, + ) + else: + pipeline_q = pipeline_custom.PipelineCpAsync.create( + barrier_storage=mbar_ptr_Q, + num_stages=1, + producer_group=load_threads, + consumer_group=mma_warps, + defer_sync=True, + elect_one_release=True, + syncwarp_before_release=False, + ) + + if const_expr(self.use_tma_KV): + pipeline_k = pipeline_custom.PipelineTmaAsync.create( + barrier_storage=storage.mbar_ptr_K.data_ptr(), + num_stages=self.num_stages, + producer_group=tma_warp, + consumer_group=mma_warps, + tx_count=self.tma_copy_bytes["K"], + defer_sync=True, + ) + pipeline_v = pipeline_custom.PipelineTmaAsync.create( + barrier_storage=storage.mbar_ptr_V.data_ptr(), + num_stages=self.num_stages, + producer_group=tma_warp, + consumer_group=mma_warps, + tx_count=self.tma_copy_bytes["V"], + defer_sync=True, + ) + else: + pipeline_k = pipeline_custom.PipelineCpAsync.create( + barrier_storage=storage.mbar_ptr_K.data_ptr(), + num_stages=self.num_stages, + producer_group=load_threads, + consumer_group=mma_warps, + defer_sync=True, + elect_one_release=True, + syncwarp_before_release=False, + ) + pipeline_v = pipeline_custom.PipelineCpAsync.create( + barrier_storage=storage.mbar_ptr_V.data_ptr(), + num_stages=self.num_stages, + producer_group=load_threads, + consumer_group=mma_warps, + defer_sync=True, + elect_one_release=True, + syncwarp_before_release=False, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # /////////////////////////////////////////////////////////////////////////////// + # Get shared memory buffer + # /////////////////////////////////////////////////////////////////////////////// + sQ = storage.sQ.get_tensor(sQ_layout.outer, swizzle=sQ_layout.inner) + sK = storage.sK.get_tensor(sK_layout.outer, swizzle=sK_layout.inner) + if const_expr(not self.Q_in_regs): + sV = storage.sV.get_tensor(sV_layout.outer, swizzle=sV_layout.inner) + else: + sV = storage.sQ.get_tensor( + sV_layout.outer, swizzle=sV_layout.inner, dtype=mV.element_type + ) + # Transpose view of V to tensor with layout (head_dim_v, tile_n) for tiled mma + sVt = layout_utils.transpose_view(sV) + sP = None + if const_expr(sP_layout is not None): + sP = storage.sP.get_tensor(sP_layout.outer, swizzle=sP_layout.inner) + # reuse sQ's data iterator + sO = storage.sQ.get_tensor(sO_layout.outer, swizzle=sO_layout.inner, dtype=self.dtype) + + block_info = BlockInfo( + self.tile_m, + self.tile_n, + self.is_causal, + self.is_local, + False, # is_split_kv + window_size_left, + window_size_right, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + SeqlenInfoCls = partial( + SeqlenInfoQK.create, + seqlen_q_static=mQ.shape[0] if const_expr(not self.pack_gqa) else mQ.shape[0][1], + seqlen_k_static=mK.shape[0] + if const_expr(mPageTable is None) + else mK.shape[0] * mPageTable.shape[1], + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + mCuTotalMBlocks=( + blocksparse_tensors.cu_total_m_blocks if blocksparse_tensors is not None else None + ), + mCuBlockIdxOffsets=( + blocksparse_tensors.cu_block_idx_offsets + if blocksparse_tensors is not None + else None + ), + # Don't need to pass in tile_mn because we won't access offset_padded + ) + AttentionMaskCls = partial( + AttentionMask, + self.tile_m, + self.tile_n, + window_size_left=window_size_left, + window_size_right=window_size_right, + qhead_per_kvhead_packgqa=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + TileSchedulerCls = partial(TileScheduler.create, tile_sched_params) + + # Cluster wait before starting + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + if warp_idx < 4: # Producer + cute.arch.setmaxregister_decrease(self.num_producer_regs) + self.load( + mQ, + mK, + mV, + sQ, + sK, + sV, + tma_atom_Q, + tma_atom_K, + tma_atom_V, + pipeline_k, + pipeline_v, + pipeline_q, + gmem_tiled_copy_Q, + mPageTable, + blocksparse_tensors, + block_info, + SeqlenInfoCls, + TileSchedulerCls, + ) + + else: # Consumer + cute.arch.setmaxregister_increase(self.num_mma_regs) + # /////////////////////////////////////////////////////////////////////////////// + # Tile MMA compute thread partitions and allocate accumulators + # /////////////////////////////////////////////////////////////////////////////// + tidx, _, _ = cute.arch.thread_idx() + tidx = tidx - 128 + self.mma( + tiled_mma_qk, + tiled_mma_pv, + mO, + mLSE, + sQ, + sK, + sVt, + sP, + sO, + learnable_sink, + pipeline_k, + pipeline_v, + pipeline_q, + gmem_tiled_copy_O, + tma_atom_O, + tidx, + softmax_scale_log2, + softmax_scale, + block_info, + SeqlenInfoCls, + AttentionMaskCls, + TileSchedulerCls, + blocksparse_tensors, + aux_data, + fastdiv_mods, + ) + + @cute.jit + def load( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + tma_atom_Q: Optional[cute.CopyAtom], + tma_atom_K: Optional[cute.CopyAtom], + tma_atom_V: Optional[cute.CopyAtom], + pipeline_k: pipeline.PipelineAsync, + pipeline_v: pipeline.PipelineAsync, + pipeline_q: pipeline.PipelineAsync, + gmem_tiled_copy_Q: cute.TiledCopy, + mPageTable: Optional[cute.Tensor], + blocksparse_tensors: Optional[BlockSparseTensors], + block_info: BlockInfo, + SeqlenInfoCls: Callable, + TileSchedulerCls: Callable, + ): + warp_idx_in_wg = cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4 + tidx, _, _ = cute.arch.thread_idx() + + # TMA: only warp 0 loads. cp_async: all warps load. + # When not use_tma_Q, all 128 producer threads participate in Q loading. + is_load_warp = warp_idx_in_wg == 0 or const_expr(not self.use_tma_KV or not self.use_tma_Q) + # KV loading restricted to warp 0 for TMA, all warps for non-TMA KV + is_kv_load_warp = warp_idx_in_wg == 0 or const_expr(not self.use_tma_KV) + + if is_load_warp: + q_producer_phase = Int32(1) + kv_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_stages + ) + tile_scheduler = TileSchedulerCls() + work_tile = tile_scheduler.initial_work_tile_info() + while work_tile.is_valid_tile: + # if work_tile.is_valid_tile: + m_block, head_idx, batch_idx, _ = work_tile.tile_idx + seqlen = SeqlenInfoCls(batch_idx) + mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx] + head_idx_kv = ( + head_idx // self.qhead_per_kvhead if const_expr(not self.pack_gqa) else head_idx + ) + + load_Q = None + if const_expr(self.use_tma_Q): + gQ = cute.local_tile(mQ_cur, (self.tile_m, self.tile_hdim), (m_block, 0)) + load_Q, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_Q, 0, cute.make_layout(1), gQ, sQ, single_stage=True + ) + + paged_kv_manager = None + tma_load_K_fn = None + tma_load_V_fn = None + if const_expr(self.use_tma_KV): + # === TMA path (non-paged and paged with page_size == n_block_size) === + if const_expr(mPageTable is not None): + # Paged TMA: keep page dimension indexable + mK_cur = mK[None, None, head_idx_kv, None] + mV_cur = mV[None, None, head_idx_kv, None] + gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (0, 0, None)) + gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (0, 0, None)) + else: + # Non-paged TMA + mK_cur = seqlen.offset_batch_K(mK, batch_idx, dim=3)[ + None, None, head_idx_kv + ] + mV_cur = seqlen.offset_batch_K(mV, batch_idx, dim=3)[ + None, None, head_idx_kv + ] + gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (None, 0)) + gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (None, 0)) + # TODO: mcast + tma_load_K_fn, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_K, 0, cute.make_layout(1), gK, sK + ) + tma_load_K_fn = copy_utils.tma_producer_copy_fn(tma_load_K_fn, pipeline_k) + tma_load_V_fn, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_V, 0, cute.make_layout(1), gV, sV + ) + tma_load_V_fn = copy_utils.tma_producer_copy_fn(tma_load_V_fn, pipeline_v) + else: + # === cp_async path (paged KV with page_size != n_block_size) === + paged_kv_manager = PagedKVManager.create( + mPageTable, + mK, + mV, + FastDivmodDivisor(mK.shape[0]), + batch_idx, + head_idx_kv, + tidx, + seqlen.seqlen_k, + 0, # leftpad_k + self.tile_n, + self.tile_hdim, + self.tile_hdimv, + self.num_threads_per_warp_group, + mK.element_type, + arch=self.arch.major * 10 + self.arch.minor, + ) + + load_K = partial( + self.load_KV, + tma_load_K_fn, + paged_kv_manager, + sK, + pipeline_kv=pipeline_k, + K_or_V="K", + ) + load_V = partial( + self.load_KV, + tma_load_V_fn, + paged_kv_manager, + sV, + pipeline_kv=pipeline_v, + K_or_V="V", + ) + + pack_gqa = None + if const_expr(not self.use_tma_Q): + pack_gqa = PackGQA( + self.tile_m, self.tile_hdim, self.check_hdim_oob, self.qhead_per_kvhead + ) + + if const_expr(not self.use_block_sparsity): + n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block) + # if cute.arch.thread_idx()[0] == 0: + # cute.printf("m_block = %d, n_block_min: %d, n_block_max: %d", m_block, n_block_min, n_block_max) + # Clamp n_block to 0 when n_block_max == 0 (can happen with causal + # + pack_gqa when seqlen_k < tile_n). TMA handles n_block=-1 + # gracefully (fills zeros), but cp.async would crash on + # out-of-bounds page table access. + n_block = ( + n_block_max - 1 + if const_expr(self.use_tma_KV) + else cutlass.max(n_block_max - 1, 0) + ) + page_idx = ( + mPageTable[batch_idx, n_block] + if const_expr(mPageTable is not None and self.use_tma_KV) + else None + ) + + # First iteration: load K on pipeline_k, Q on pipeline_q + if is_kv_load_warp: + pipeline_k.producer_acquire(kv_producer_state) + if const_expr(not self.use_tma_KV): + paged_kv_manager.load_page_table(n_block) + load_K(block=n_block, producer_state=kv_producer_state, page_idx=page_idx) + if const_expr(self.use_tma_Q): + if warp_idx_in_wg == 0: + pipeline_q.producer_acquire_w_index_phase(0, q_producer_phase) + load_Q(tma_bar_ptr=pipeline_q.sync_object_full.get_barrier(0)) + q_producer_phase ^= 1 + else: + pipeline_q.producer_acquire_w_index_phase(0, q_producer_phase) + pack_gqa.load_Q( + mQ_cur, sQ, gmem_tiled_copy_Q, tidx, m_block, seqlen.seqlen_q + ) + cute.arch.cp_async_commit_group() + pipeline_q.producer_commit_w_index(0) + q_producer_phase ^= 1 + + if is_kv_load_warp: + if const_expr(not self.intra_wg_overlap or not self.use_tma_KV): + pipeline_v.producer_acquire(kv_producer_state) + load_V( + block=n_block, producer_state=kv_producer_state, page_idx=page_idx + ) + kv_producer_state.advance() + for i in cutlass.range(n_block_max - 1 - n_block_min, unroll=1): + n_block = n_block_max - 1 - i - 1 + page_idx = ( + mPageTable[batch_idx, n_block] + if const_expr(mPageTable is not None and self.use_tma_KV) + else None + ) + if const_expr(not self.use_tma_KV): + paged_kv_manager.load_page_table(n_block) + pipeline_k.producer_acquire(kv_producer_state) + load_K( + block=n_block, + producer_state=kv_producer_state, + page_idx=page_idx, + ) + pipeline_v.producer_acquire(kv_producer_state) + load_V( + block=n_block, + producer_state=kv_producer_state, + page_idx=page_idx, + ) + kv_producer_state.advance() + else: + for i in cutlass.range(n_block_max - 1 - n_block_min, unroll=1): + n_block_prev = n_block_max - i - 1 + n_block = n_block_prev - 1 + page_idx = ( + mPageTable[batch_idx, n_block] + if const_expr(mPageTable is not None) + else None + ) + page_idx_prev = ( + mPageTable[batch_idx, n_block_prev] + if const_expr(mPageTable is not None) + else None + ) + kv_producer_state_prev = kv_producer_state.clone() + kv_producer_state.advance() + pipeline_k.producer_acquire(kv_producer_state) + load_K( + block=n_block, + producer_state=kv_producer_state, + page_idx=page_idx, + ) + pipeline_v.producer_acquire(kv_producer_state_prev) + load_V( + block=n_block_prev, + producer_state=kv_producer_state_prev, + page_idx=page_idx_prev, + ) + n_block = n_block_min + page_idx = ( + mPageTable[batch_idx, n_block] + if const_expr(mPageTable is not None) + else None + ) + pipeline_v.producer_acquire(kv_producer_state) + load_V( + block=n_block, producer_state=kv_producer_state, page_idx=page_idx + ) + kv_producer_state.advance() + else: + # Block sparsity: use TMA closures directly (not paged) + # Load Q on pipeline_q, separate from K/V pipeline + if const_expr(self.use_tma_Q): + if warp_idx_in_wg == 0: + pipeline_q.producer_acquire_w_index_phase(0, q_producer_phase) + load_Q(tma_bar_ptr=pipeline_q.sync_object_full.get_barrier(0)) + q_producer_phase ^= 1 + else: + pipeline_q.producer_acquire_w_index_phase(0, q_producer_phase) + pack_gqa.load_Q( + mQ_cur, sQ, gmem_tiled_copy_Q, tidx, m_block, seqlen.seqlen_q + ) + cute.arch.cp_async_commit_group() + pipeline_q.producer_commit_w_index(0) + q_producer_phase ^= 1 + if is_kv_load_warp: + kv_producer_state = produce_block_sparse_loads( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + seqlen, + kv_producer_state, + tma_load_K_fn, + tma_load_V_fn, + pipeline_k, + pipeline_v, + self.intra_wg_overlap, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + self.q_subtile_factor, + ) + + tile_scheduler.prefetch_next_work() + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + # End of persistent scheduler loop + + # Producer tail is only useful for cluster to avoid early exit of blocks. + # We only need producer_tail on V since that's the last that's loaded, we don't + # need it for Q (no cluster) and K. + if is_kv_load_warp: + pipeline_v.producer_tail(kv_producer_state) + + @cute.jit + def load_KV( + self, + tma_load_fn: Optional[Callable], + paged_kv_manager: Optional[PagedKVManager], + sX: cute.Tensor, + block: Int32, + pipeline_kv: pipeline.PipelineAsync, + producer_state: pipeline.PipelineState, + K_or_V: Literal["K", "V"], + page_idx: Optional[Int32] = None, + ): + if const_expr(self.use_tma_KV): + src_idx = block if const_expr(page_idx is None) else page_idx + tma_load_fn(src_idx=src_idx, producer_state=producer_state) + else: + paged_kv_manager.load_KV(block, sX[None, None, producer_state.index], K_or_V) + cute.arch.cp_async_commit_group() + pipeline_kv.producer_commit(producer_state) + + @cute.jit + def mma( + self, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + sQ: cute.Tensor, + sK: cute.Tensor, + sVt: cute.Tensor, + sP: Optional[cute.Tensor], + sO: cute.Tensor, + learnable_sink: Optional[cute.Tensor], + pipeline_k: pipeline.PipelineAsync, + pipeline_v: pipeline.PipelineAsync, + pipeline_q: pipeline.PipelineAsync, + gmem_tiled_copy_O: cute.TiledCopy, + tma_atom_O: Optional[cute.CopyAtom], + tidx: Int32, + softmax_scale_log2: Float32, + softmax_scale: Optional[Float32], + block_info: BlockInfo, + SeqlenInfoCls: Callable, + AttentionMaskCls: Callable, + TileSchedulerCls: Callable, + blocksparse_tensors: Optional[BlockSparseTensors], + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + ): + aux_tensors = aux_data.tensors + warp_group_idx = cute.arch.make_warp_uniform(tidx // self.num_threads_per_warp_group) + warp_group_thread_layout = cute.make_layout( + self.num_wg_mma, stride=self.num_threads_per_warp_group + ) + thr_mma_qk = tiled_mma_qk.get_slice(tidx) + wg_mma_qk = tiled_mma_qk.get_slice(warp_group_thread_layout(warp_group_idx)) + wg_mma_pv = tiled_mma_pv.get_slice(warp_group_thread_layout(warp_group_idx)) + _, tSrQ, tSrK = sm90_utils.partition_fragment_ABC( + wg_mma_qk, (self.tile_m, self.tile_n, self.tile_hdim), sQ, sK + ) + mma_qk_fn = partial( + sm90_utils.gemm_zero_init, tiled_mma_qk, (self.tile_m, self.tile_n), tSrQ, tSrK + ) + acc_O, tOrP, tOrVt = sm90_utils.partition_fragment_ABC( + wg_mma_pv, (self.tile_m, self.tile_hdimv, self.tile_n), sP, sVt + ) + mma_pv_fn = partial(sm90_utils.gemm_w_idx, tiled_mma_pv, acc_O, tOrP, tOrVt) + + # /////////////////////////////////////////////////////////////////////////////// + # Smem copy atom tiling + # /////////////////////////////////////////////////////////////////////////////// + smem_copy_atom_P = utils.get_smem_store_atom( + self.arch.major * 10 + self.arch.minor, self.dtype + ) + smem_thr_copy_P = cute.make_tiled_copy_C(smem_copy_atom_P, tiled_mma_qk).get_slice(tidx) + tPsP = smem_thr_copy_P.partition_D(sP) if const_expr(sP is not None) else None + smem_copy_params = SimpleNamespace(smem_thr_copy_P=smem_thr_copy_P, tPsP=tPsP) + + self.mma_init() + + q_consumer_phase = Int32(0) + kv_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_stages + ) + + tile_scheduler = TileSchedulerCls() + work_tile = tile_scheduler.initial_work_tile_info() + softmax = Softmax.create( + softmax_scale_log2, + num_rows=acc_O.shape[0][0] * acc_O.shape[1], + softmax_scale=softmax_scale, + ) + + # For RescaleOBeforeGemm: persistent scores_scale across iterations + scores_scale = None + if const_expr(self.rescale_O_before_gemm): + scores_scale = cute.make_rmem_tensor_like(softmax.row_max, Float32) + + mma_one_n_block_all = partial( + self.mma_one_n_block_intrawg_overlap + if const_expr(self.intra_wg_overlap) + else self.mma_one_n_block, + mma_qk_fn=mma_qk_fn, + pipeline_k=pipeline_k, + pipeline_v=pipeline_v, + acc_O=acc_O, + tOrP=tOrP, + smem_copy_params=smem_copy_params, + check_inf=True, + scores_scale=scores_scale, + ) + + process_first_half_block = partial( + self.first_half_block_overlap, + mma_qk_fn=mma_qk_fn, + pipeline_k=pipeline_k, + tOrP=tOrP, + smem_copy_params=smem_copy_params, + scores_scale=scores_scale, + softmax=softmax, + acc_O=acc_O, + ) + process_last_half_block = partial( + self.last_half_block_overlap, + pipeline_v=pipeline_v, + mma_pv_fn=mma_pv_fn, + scores_scale=scores_scale, + softmax=softmax, + acc_O=acc_O, + ) + while work_tile.is_valid_tile: + # if work_tile.is_valid_tile: + + # shape: (atom_v_m * rest_m) + m_block, head_idx, batch_idx, _ = work_tile.tile_idx + seqlen = SeqlenInfoCls(batch_idx) + + # Recompute fastdiv_mods if necessary for varlen with aux_tensors + recompute_fastdiv_mods_q = cutlass.const_expr( + aux_tensors is not None and (seqlen.has_cu_seqlens_q or seqlen.has_seqused_q) + ) + recompute_fastdiv_mods_k = cutlass.const_expr( + aux_tensors is not None and (seqlen.has_cu_seqlens_k or seqlen.has_seqused_k) + ) + if cutlass.const_expr(fastdiv_mods is not None): + seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods + fastdiv_mods = ( + seqlen_q_divmod + if not recompute_fastdiv_mods_q + else FastDivmodDivisor(seqlen.seqlen_q), + seqlen_k_divmod + if not recompute_fastdiv_mods_k + else FastDivmodDivisor(seqlen.seqlen_k), + ) + + mask = AttentionMaskCls(seqlen) + mask_fn = partial( + mask.apply_mask, + batch_idx=batch_idx, + head_idx=head_idx, + m_block=m_block, + thr_mma=thr_mma_qk, + mask_causal=self.is_causal, + mask_local=self.is_local, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + score_mod_fn = None + if const_expr(self.score_mod is not None): + score_mod_fn = partial( + self.apply_score_mod, + thr_mma_qk, + batch_idx, + head_idx, + m_block, + softmax_scale=softmax_scale, + aux_data=aux_data, + fastdiv_mods=fastdiv_mods, + ) + mma_one_n_block = partial( + mma_one_n_block_all, seqlen=seqlen, softmax=softmax, score_mod_fn=score_mod_fn + ) + n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block) + pipeline_q.consumer_wait_w_index_phase(0, q_consumer_phase) + # For performance reason, we separate out two kinds of iterations: + # those that need masking on S, and those that don't. + # We need masking on S for the very last block when K and V has length not multiple of tile_n. + # We also need masking on S if it's causal, for the last several blocks. + # softmax.reset() # Don't need reset as we explicitly call softmax w is_first=True + O_should_accumulate = False + + # ========================================== + # MAINLOOP + # ========================================== + if const_expr(not self.use_block_sparsity): + # ========================================== + # No block-sparsity (original path) + # ========================================== + # First iteration with seqlen masking + if const_expr(self.intra_wg_overlap): + kv_consumer_state = process_first_half_block( + n_block=n_block_max - 1, + seqlen=seqlen, + kv_consumer_state=kv_consumer_state, + mask_fn=partial(mask_fn, mask_mod=self.mask_mod), + score_mod_fn=score_mod_fn, + is_first_block=True, + ) + else: + self.warp_scheduler_barrier_sync() + kv_consumer_state = mma_one_n_block( + kv_consumer_state, + n_block=n_block_max - 1, + seqlen=seqlen, + mma_pv_fn=partial(mma_pv_fn, zero_init=True), + is_first_n_block=True, + mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=True), + ) + O_should_accumulate = True + # if cute.arch.thread_idx()[0] == 128: cute.printf("m_block = {}, n_block_max = {}, n_block_min = {}", m_block, n_block_max, n_block_min) + n_block_max -= 1 + # Next couple of iterations with causal masking + if const_expr(self.is_causal or self.is_local): + n_block_min_causal_local_mask = block_info.get_n_block_min_causal_local_mask( + seqlen, m_block, n_block_min + ) + # if cute.arch.thread_idx()[0] == 128: cute.printf("n_block_min_causal_local_mask = {}", n_block_min_causal_local_mask) + for n_tile in cutlass.range( + n_block_max - n_block_min_causal_local_mask, unroll=1 + ): + kv_consumer_state = mma_one_n_block( + kv_consumer_state, + n_block=n_block_max - 1 - n_tile, + seqlen=seqlen, + mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate), + mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=False), + ) + O_should_accumulate = True + n_block_max = cutlass.min(n_block_max, n_block_min_causal_local_mask) + # The remaining iterations have no masking + n_block_min_before_local_mask = block_info.get_n_block_min_before_local_mask( + seqlen, m_block, n_block_min + ) + # if cute.arch.thread_idx()[0] == 128: cute.printf("n_block_min_before_local_mask = {}, n_block_min = {}", n_block_min_before_local_mask, n_block_min) + for n_tile in cutlass.range(n_block_max - n_block_min_before_local_mask, unroll=1): + kv_consumer_state = mma_one_n_block( + kv_consumer_state, + n_block=n_block_max - 1 - n_tile, + seqlen=seqlen, + mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate), + mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=False), + ) + O_should_accumulate = True + # Separate iterations with local masking on the left + if const_expr(self.is_local and block_info.window_size_left is not None): + n_block_max = cutlass.min(n_block_max, n_block_min_before_local_mask) + for n_tile in cutlass.range(n_block_max - n_block_min, unroll=1): + kv_consumer_state = mma_one_n_block( + kv_consumer_state, + n_block=n_block_max - 1 - n_tile, + seqlen=seqlen, + mma_pv_fn=partial(mma_pv_fn, zero_init=not O_should_accumulate), + mask_fn=partial(mask_fn, mask_mod=self.mask_mod, mask_seqlen=False), + ) + O_should_accumulate = True + # Release Q pipeline so the producer can load the next tile's Q + pipeline_q.consumer_release_w_index(0) + # Last "half" iteration + if const_expr(self.intra_wg_overlap): + kv_consumer_state = process_last_half_block( + kv_consumer_state=kv_consumer_state, + zero_init=not O_should_accumulate, + ) + O_should_accumulate = True + else: + self.warp_scheduler_barrier_arrive() + + else: + # ========================================== + # Block sparsity + # ========================================== + kv_consumer_state, O_should_accumulate, processed_any = consume_block_sparse_loads( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + seqlen, + kv_consumer_state, + mma_pv_fn, + mma_one_n_block, + process_first_half_block, + process_last_half_block, + mask_fn, + score_mod_fn, + O_should_accumulate, + self.mask_mod, + fastdiv_mods, + self.intra_wg_overlap, + self.warp_scheduler_barrier_sync, + self.warp_scheduler_barrier_arrive, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + self.q_subtile_factor, + ) + + # Release Q pipeline so the producer can load the next tile's Q + pipeline_q.consumer_release_w_index(0) + + # Handle empty case (when no blocks to process) + if not processed_any: + softmax.reset() + acc_O.fill(0.0) + + q_consumer_phase ^= 1 + + sink_val = None + if const_expr(learnable_sink is not None): + if const_expr(not self.pack_gqa): + sink_val = Float32(learnable_sink[head_idx]) + else: # Each thread might have a different sink value due to different q_head + sink_val = cute.make_rmem_tensor_like(softmax.row_max, Float32) + cS = cute.make_identity_tensor((self.tile_m, self.tile_n)) + tScS_mn = layout_utils.reshape_acc_to_mn(thr_mma_qk.partition_C(cS)) + for r in cutlass.range(cute.size(sink_val), unroll_full=True): + row = m_block * self.tile_m + tScS_mn[r][0] + q_head_idx = row % self.qhead_per_kvhead + head_idx * self.qhead_per_kvhead + sink_val[r] = Float32(learnable_sink[q_head_idx]) + + # normalize acc_O by row_sum and calculate the lse + row_scale = softmax.finalize(sink_val=sink_val) + softmax.rescale_O(acc_O, row_scale) + + # /////////////////////////////////////////////////////////////////////////////// + # Epilogue + # /////////////////////////////////////////////////////////////////////////////// + self.epilogue( + acc_O, + softmax.row_sum, + mO, + mLSE, + sO, + seqlen, + gmem_tiled_copy_O, + tma_atom_O, + tiled_mma_pv, + tidx, + m_block, + head_idx, + batch_idx, + ) + + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + + @cute.jit + def first_half_block_overlap( + self, + n_block: Int32, + mma_qk_fn: Callable, + kv_consumer_state, + pipeline_k, + tOrP: cute.Tensor, + smem_copy_params: SimpleNamespace, + softmax: Softmax, + seqlen: SeqlenInfoQK, + scores_scale: Optional[cute.Tensor] = None, + acc_O: Optional[cute.Tensor] = None, + mask_fn: Callable = None, + score_mod_fn: Optional[Callable] = None, + is_first_block: bool = False, + ): + """Processes the first half block when using intra-warpgroup-overlap""" + + pipeline_k.consumer_wait(kv_consumer_state, pipeline_k.consumer_try_wait(kv_consumer_state)) + acc_S = mma_qk_fn(B_idx=kv_consumer_state.index, wg_wait=0) + pipeline_k.consumer_release(kv_consumer_state) + + # Apply score modification if present + if const_expr(score_mod_fn is not None): + score_mod_fn(acc_S, n_block=n_block, seqlen=seqlen) + + # Apply mask; mask_seqlen always True for first block + # Caveat: if full block further right than mask block, seqlen masking is redundant; + # however, masking is being applied anyway, so essentially no perf hit + mask_fn(acc_S, n_block=n_block, mask_seqlen=True) + + row_scale = softmax.online_softmax(acc_S, is_first=is_first_block) + + tOrP_acc = layout_utils.reshape_acc_to_frgA(acc_S) + tOrP_cur = ( + tOrP + if const_expr(self.mma_pv_is_rs) + else cute.make_rmem_tensor_like(tOrP_acc, self.dtype) + ) + tOrP_cur.store(tOrP_acc.load().to(self.dtype)) + + if const_expr(not self.mma_pv_is_rs): + tPrP = smem_copy_params.smem_thr_copy_P.retile(tOrP_cur) + cute.copy(smem_copy_params.smem_thr_copy_P, tPrP, smem_copy_params.tPsP) + # Fence and barrier to make smem store visible to WGMMA + cute.arch.fence_view_async_shared() + cute.arch.sync_warp() + + # For RescaleOBeforeGemm: initialize acc_O + if const_expr(self.rescale_O_before_gemm): + acc_O.fill(0.0) + scores_scale.store(row_scale.load()) + + return kv_consumer_state + + @cute.jit + def last_half_block_overlap( + self, + kv_consumer_state, + pipeline_v, + mma_pv_fn: Callable, + zero_init: bool, + scores_scale: Optional[cute.Tensor] = None, + softmax: Optional[Softmax] = None, + acc_O: Optional[cute.Tensor] = None, + ): + """Processes the final PV GEMM when using intra-warpgroup-overlap""" + + # For RescaleOBeforeGemm: rescale O before the final PV GEMM + if const_expr(self.rescale_O_before_gemm): + softmax.rescale_O(acc_O, scores_scale) + + pipeline_v.consumer_wait(kv_consumer_state, pipeline_v.consumer_try_wait(kv_consumer_state)) + mma_pv_fn(B_idx=kv_consumer_state.index, zero_init=zero_init, wg_wait=0) + pipeline_v.consumer_release(kv_consumer_state) + kv_consumer_state.advance() + return kv_consumer_state + + @cute.jit + def mma_one_n_block( + self, + smem_pipe_read: pipeline.PipelineState | pipeline_custom.PipelineStateSimple, + n_block: Int32, + mma_qk_fn: Callable, + mma_pv_fn: Callable, + pipeline_k: pipeline.PipelineAsync, + pipeline_v: pipeline.PipelineAsync, + acc_O: cute.Tensor, + tOrP: cute.Tensor, + smem_copy_params: SimpleNamespace, + softmax: Softmax, + seqlen: SeqlenInfoQK, + scores_scale: Optional[cute.Tensor] = None, # not used + score_mod_fn: Optional[Callable] = None, + mask_fn: Optional[Callable] = None, + is_first_n_block: cutlass.Constexpr = False, + check_inf: cutlass.Constexpr = True, + ): + pipeline_k.consumer_wait(smem_pipe_read, pipeline_k.consumer_try_wait(smem_pipe_read)) + # S = Q @ K.T + acc_S = mma_qk_fn(B_idx=smem_pipe_read.index, wg_wait=-1) + self.warp_scheduler_barrier_arrive() + warpgroup.wait_group(0) + pipeline_k.consumer_release(smem_pipe_read) + + # handle score mods and masking + if const_expr(score_mod_fn is not None): + score_mod_fn(acc_S, n_block=n_block, seqlen=seqlen) + if const_expr(mask_fn is not None): + mask_fn(acc_S=acc_S, n_block=n_block) + + row_scale = softmax.online_softmax(acc_S, is_first=is_first_n_block, check_inf=check_inf) + # if cute.arch.thread_idx()[0] == 0: cute.print_tensor(layout_utils.reshape_acc_to_mn(acc_S)) + tOrP_acc = layout_utils.reshape_acc_to_frgA(acc_S) + tOrP_cur = ( + tOrP + if const_expr(self.mma_pv_is_rs) + else cute.make_rmem_tensor_like(tOrP_acc, self.dtype) + ) + # tOrP.store(tOrP_acc.load().to(self.dtype)) + # the "to(self.dtype)" conversion fails to vectorize for block sizes other + # than 128 x 128, i.e. it calls convert on 1 fp32 element at a time instead of + # 2 elements. So we just call ptx directly. + utils.cvt_f16(tOrP_acc, tOrP_cur) + if const_expr(not self.mma_pv_is_rs): + tPrP = smem_copy_params.smem_thr_copy_P.retile(tOrP_cur) + cute.copy(smem_copy_params.smem_thr_copy_P, tPrP, smem_copy_params.tPsP) + softmax.rescale_O(acc_O, row_scale) + if const_expr(not self.mma_pv_is_rs): + # Fence and barrier to make sure smem store is visible to WGMMA + cute.arch.fence_view_async_shared() + cute.arch.sync_warp() # Only need syncwarp since each warp is using its own P values for MmaPV + pipeline_v.consumer_wait(smem_pipe_read, pipeline_v.consumer_try_wait(smem_pipe_read)) + self.warp_scheduler_barrier_sync() + # O += P @ V + mma_pv_fn(B_idx=smem_pipe_read.index, wg_wait=0) + pipeline_v.consumer_release(smem_pipe_read) + smem_pipe_read.advance() + return smem_pipe_read + + @cute.jit + def mma_one_n_block_intrawg_overlap( + self, + smem_pipe_read: pipeline.PipelineState | pipeline_custom.PipelineStateSimple, + n_block: Int32, + mma_qk_fn: Callable, + mma_pv_fn: Callable, + pipeline_k: pipeline.PipelineAsync, + pipeline_v: pipeline.PipelineAsync, + acc_O: cute.Tensor, + tOrP: cute.Tensor, + smem_copy_params: SimpleNamespace, + softmax: Softmax, + seqlen: SeqlenInfoQK, + scores_scale: Optional[cute.Tensor] = None, + score_mod_fn: Optional[Callable] = None, + mask_fn: Optional[Callable] = None, + check_inf: cutlass.Constexpr = True, + ): + smem_pipe_read_v = smem_pipe_read.clone() + smem_pipe_read.advance() + pipeline_k.consumer_wait(smem_pipe_read, pipeline_k.consumer_try_wait(smem_pipe_read)) + self.warp_scheduler_barrier_sync() + # S = Q @ K.T + acc_S = mma_qk_fn(B_idx=smem_pipe_read.index, wg_wait=-1) + # RescaleOBeforeGemm: rescale O while QK GEMM is in flight, before PV GEMM + if const_expr(self.rescale_O_before_gemm): + softmax.rescale_O(acc_O, scores_scale) + pipeline_v.consumer_wait(smem_pipe_read_v, pipeline_v.consumer_try_wait(smem_pipe_read_v)) + # O += P @ V + mma_pv_fn(B_idx=smem_pipe_read_v.index, wg_wait=-1) + self.warp_scheduler_barrier_arrive() + warpgroup.wait_group(1) + pipeline_k.consumer_release(smem_pipe_read) + + # handle score mods and masking + if const_expr(score_mod_fn is not None): + score_mod_fn(acc_S, n_block=n_block, seqlen=seqlen) + if const_expr(mask_fn is not None): + mask_fn(acc_S=acc_S, n_block=n_block) + # if cute.arch.thread_idx()[0] == 128: cute.print_tensor(layout_utils.reshape_acc_to_mn(acc_S)) + + row_scale = softmax.online_softmax(acc_S, check_inf=check_inf) + warpgroup.wait_group(0) + pipeline_v.consumer_release(smem_pipe_read_v) + tOrP_acc = layout_utils.reshape_acc_to_frgA(acc_S) + tOrP_cur = ( + tOrP + if const_expr(self.mma_pv_is_rs) + else cute.make_rmem_tensor_like(tOrP_acc, self.dtype) + ) + # tOrP_cur.store(tOrP_acc.load().to(self.dtype)) + # the "to(self.dtype)" conversion fails to vectorize for block sizes other + # than 128 x 128, i.e. it calls convert on 1 fp32 element at a time instead of + # 2 elements. So we just call ptx directly. + utils.cvt_f16(tOrP_acc, tOrP_cur) + if const_expr(not self.mma_pv_is_rs): + tPrP = smem_copy_params.smem_thr_copy_P.retile(tOrP_cur) + cute.copy(smem_copy_params.smem_thr_copy_P, tPrP, smem_copy_params.tPsP) + if const_expr(not self.rescale_O_before_gemm): + softmax.rescale_O(acc_O, row_scale) + if const_expr(self.rescale_O_before_gemm): + scores_scale.store(row_scale.load()) + if const_expr(not self.mma_pv_is_rs): + # Fence and barrier to make sure smem store is visible to WGMMA + cute.arch.fence_view_async_shared() + cute.arch.sync_warp() # Only need syncwarp since each warp is using its own P values for MmaPV + return smem_pipe_read + + @cute.jit + def mma_init(self): + warp_group_idx = utils.canonical_warp_group_idx(sync=False) + if const_expr(self.use_scheduler_barrier): + if warp_group_idx == 1: + cute.arch.barrier_arrive( + barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1), + number_of_threads=2 * self.num_threads_per_warp_group, + ) + + @cute.jit + def apply_score_mod( + self, + thr_mma_qk, + batch_idx, + head_idx, + m_block, + acc_S, + n_block, + softmax_scale, + seqlen, + aux_data: AuxData = AuxData(), + fastdiv_mods=None, + ): + # Prepare index tensor + cS = cute.make_identity_tensor((self.tile_m, self.tile_n)) + cS = cute.domain_offset((m_block * self.tile_m, n_block * self.tile_n), cS) + tScS = thr_mma_qk.partition_C(cS) + + apply_score_mod_inner( + acc_S, + tScS, + self.score_mod, + batch_idx, + head_idx, + softmax_scale, + self.score_vec_size, + self.qk_acc_dtype, + aux_data, + fastdiv_mods, + seqlen_info=seqlen, + constant_q_idx=None, + qhead_per_kvhead=self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + + def warp_scheduler_barrier_sync(self): + if const_expr(self.use_scheduler_barrier): + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1) + - 1 + + utils.canonical_warp_group_idx(sync=False), + number_of_threads=2 * self.num_threads_per_warp_group, + ) + + def warp_scheduler_barrier_arrive(self): + if const_expr(self.use_scheduler_barrier): + assert self.num_wg_mma in [2, 3] + cur_wg = utils.canonical_warp_group_idx(sync=False) - 1 + if const_expr(self.num_wg_mma == 2): + next_wg = 1 - cur_wg + else: + t = cur_wg + 1 + next_wg = t % self.num_wg_mma + cute.arch.barrier_arrive( + barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1) + next_wg, + number_of_threads=2 * self.num_threads_per_warp_group, + ) diff --git a/fastvideo-kernel/fa4/interface.py b/fastvideo-kernel/fa4/interface.py new file mode 100644 index 000000000..e8dd4a76b --- /dev/null +++ b/fastvideo-kernel/fa4/interface.py @@ -0,0 +1,3204 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. +# [2025-07-04] Version in Cute-DSL, for Hopper and Blackwell. You'll need install nvidia-cutlass-dsl==4.2.0. + +import os +import math +from dataclasses import dataclass +from functools import lru_cache +from typing import Optional, Tuple, Callable + +import torch + + + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Float32 +from quack.compile_utils import make_fake_tensor as fake_tensor +from flash_attn.cute.cache_utils import get_jit_cache +from flash_attn.cute.testing import is_fake_mode + + +if os.environ.get("CUTE_DSL_PTXAS_PATH", None) is not None: + from flash_attn.cute import cute_dsl_ptxas # noqa: F401 + + # Patch to dump ptx and then use system ptxas to compile to cubin + cute_dsl_ptxas.patch() + + +from flash_attn.cute import utils +from flash_attn.cute import fa_logging +from flash_attn.cute.cute_dsl_utils import ( + get_aux_tensor_metadata, + get_broadcast_dims, + to_cute_aux_tensor, + to_cute_tensor, +) +from flash_attn.cute.flash_fwd import FlashAttentionForwardSm80 +from flash_attn.cute.flash_fwd_sm90 import FlashAttentionForwardSm90 +from flash_attn.cute.flash_fwd_sm100 import ( + DescaleTensors, + FlashAttentionForwardSm100, + supports_sm100_tile_m64_head_dims, +) +from flash_attn.cute.flash_fwd_sm120 import FlashAttentionForwardSm120 +from flash_attn.cute.flash_bwd_preprocess import FlashAttentionBackwardPreprocess +from flash_attn.cute.flash_bwd import FlashAttentionBackwardSm80 +from flash_attn.cute.flash_bwd_sm90 import FlashAttentionBackwardSm90 +from flash_attn.cute.flash_bwd_sm100 import FlashAttentionBackwardSm100 +from flash_attn.cute.flash_bwd_sm120 import FlashAttentionBackwardSm120 +from flash_attn.cute.flash_bwd_postprocess import FlashAttentionBackwardPostprocess +from flash_attn.cute.flash_fwd_combine import FlashAttentionForwardCombine +from flash_attn.cute.flash_fwd_mla_sm100 import FlashAttentionMLAForwardSm100 +from flash_attn.cute.flash_bwd_mla_sm100 import FlashAttentionSparseMLABackwardSm100 +from flash_attn.cute.flash_bwd_mla_dq_dqv_sm100 import dQdQvGemmKernel +from flash_attn.cute.flash_bwd_mla_dk_sm100 import dKGemmKernel + +# SM100 head_dim=256 2CTA kernel imports +from flash_attn.cute.sm100_hd256_2cta_fmha_forward import BlackwellFusedMultiHeadAttentionForward +from flash_attn.cute.sm100_hd256_2cta_fmha_backward import BlackwellFusedMultiHeadAttentionBackward + +from flash_attn.cute.utils import AuxData +from flash_attn.cute.block_sparsity import ( + BlockSparseTensorsTorch, + get_sparse_q_block_size, + to_cute_block_sparse_tensors, + normalize_block_sparse_config, + normalize_block_sparse_config_bwd, +) + +def _parse_arch_str(arch_str): + """Parse arch string (e.g. 'sm_80', 'sm_90a', '80', '100') to int (e.g. 80, 90, 100).""" + import re + match = re.match(r"^(?:sm_?|SM_?)?(\d+)(\d)([af]?)$", arch_str) + if not match: + raise ValueError(f"Invalid arch format: {arch_str}") + major, minor, _ = match.groups() + return int(major) * 10 + int(minor) + + +@lru_cache(maxsize=None) +def _get_device_arch(): + """Cached device arch check. + + Override with FLASH_ATTENTION_ARCH (e.g. 'sm_80' or '80') to select which + kernel path to use (SM80/SM90/SM100/SM120) independently of the compilation + target (CUTE_DSL_ARCH). + + For CPU-only compilation (no GPU), set both: + FLASH_ATTENTION_ARCH=sm_80 (kernel selection) + CUTE_DSL_ARCH=sm_80 (compilation target) + """ + arch_override = os.environ.get("FLASH_ATTENTION_ARCH", None) + if arch_override is not None: + return _parse_arch_str(arch_override) + major, minor = torch.cuda.get_device_capability() + return major * 10 + int(minor) + + +def _requests_sm100_sparse_64x64_tiles( + arch: int, block_sparse_tensors: BlockSparseTensorsTorch | None +) -> bool: + return ( + arch == 100 + and block_sparse_tensors is not None + and block_sparse_tensors.block_size == (64, 64) + ) + + +def _validate_head_dims(head_dim: int, head_dim_v: int, compute_capability: int, alignment: int) -> None: + """Validate head dimension constraints based on compute capability.""" + is_deepseek_shape = head_dim == 192 and head_dim_v == 128 + is_deepseek_mla_absorbed_shape = (head_dim == 64 or head_dim == head_dim_v) and head_dim_v == 512 + is_dedicate_kernel_shape = head_dim == 256 and head_dim_v == 256 + is_standard_range = 8 <= head_dim <= 128 and 8 <= head_dim_v <= 128 + + is_sm90_range = 8 <= head_dim <= 256 and 8 <= head_dim_v <= 256 + if compute_capability == 9: + assert is_sm90_range and head_dim % alignment == 0 and head_dim_v % alignment == 0, ( + f"(head_dim, head_dim_v)=({head_dim}, {head_dim_v}) is not supported on SM90. " + f"head_dim and head_dim_v must be between 8 and 256 and divisible by {alignment}." + ) + elif compute_capability in [10, 11]: + assert (is_standard_range or is_deepseek_shape or is_deepseek_mla_absorbed_shape or is_dedicate_kernel_shape) and head_dim % alignment == 0 and head_dim_v % alignment == 0, ( + f"(head_dim, head_dim_v)=({head_dim}, {head_dim_v}) is not supported on SM100/SM110. " + f"head_dim and head_dim_v must be between 8 and 128 and divisible by {alignment}, or (192, 128) for DeepSeek, or (256, 256) for hd256." + ) + + +@dataclass(frozen=True) +class FwdConfig: + m_block_size: int + n_block_size: int + mma_pv_is_rs: bool + intra_wg_overlap: bool + + +def _tile_size_fwd_sm90(head_dim, head_dim_v, is_causal, is_local, sparse_block_size_q=None): + """Return FwdConfig for SM90 forward. + + Tile sizes and flags based on tile_size_fwd_sm90 in hopper/tile_size.h, adjusted + for the Python kernel's different register/smem tradeoffs (benchmarked on H100 SXM). + + When sparse_block_size_q is set, tile_m must divide it. For head_dim <= 96 the + optimal tile_m=192 is used when compatible, otherwise we fall back to 128. + """ + if head_dim <= 64: + # C++: 192×192 non-causal, 192×128 causal/local. + # Python: 192×128 RS+OL is consistently best across seqlens. + if sparse_block_size_q is not None and sparse_block_size_q % 192 != 0: + return FwdConfig(128, 128, True, True) + return FwdConfig(192, 128, True, True) + elif head_dim <= 96: + # C++: 192×144 noRS+OL for all cases. + # Python: RS is catastrophic with 192× tiles (~300 vs ~600 TFLOPS). + # noRS+OL is always required. Causal: 192×128 slightly better short seqlen. + if sparse_block_size_q is not None and sparse_block_size_q % 192 != 0: + return FwdConfig(128, 128, False, True) + if is_causal or is_local: + return FwdConfig(192, 128, False, True) + else: + return FwdConfig(192, 144, False, True) + elif head_dim <= 128: + return FwdConfig(128, 128, True, True) + elif head_dim <= 192: + tile_n = 96 if is_local else (128 if head_dim_v <= 128 else 112) + return FwdConfig(128, tile_n, True, True) + else: # hdim 256 + tile_n = 64 if is_local else 80 + return FwdConfig(128, tile_n, True, True) + +@dataclass(frozen=True) +class BwdConfig: + m_block_size: int + n_block_size: int + num_stages_Q: int + num_stages_dO: int + num_stages_PdS: int + SdP_swapAB: bool + dKV_swapAB: bool + dQ_swapAB: bool + AtomLayoutMSdP: int + AtomLayoutNdKV: int + AtomLayoutMdQ: int + num_wg: int = 2 # MMA warp groups (total threads = (num_wg + 1) * 128) + dQ_single_wg: bool = False + + +def _tile_size_bwd_sm90(head_dim, head_dim_v, causal, local, sparse_block_size_q=None): + """Return BwdConfig for SM90. + + Configs based on C++ FA3 hopper/flash_bwd_launch_template.h, + benchmarked on H100 SXM. + """ + if head_dim <= 64: + # C++ FA3: 128, 128, 64, ..., 2, 2, true, false, false, 2, 1, 2, 2 + return BwdConfig( + m_block_size=128, n_block_size=128, + num_stages_Q=2, num_stages_dO=2, num_stages_PdS=2, + SdP_swapAB=True, dKV_swapAB=False, dQ_swapAB=False, + AtomLayoutMSdP=1, AtomLayoutNdKV=2, AtomLayoutMdQ=2, + ) + elif head_dim <= 96: + # C++ FA3: 64, 128, 96, dQ_swapAB=False + return BwdConfig( + m_block_size=64, n_block_size=128, + num_stages_Q=2, num_stages_dO=2, num_stages_PdS=2, + SdP_swapAB=True, dKV_swapAB=False, dQ_swapAB=False, + AtomLayoutMSdP=1, AtomLayoutNdKV=2, AtomLayoutMdQ=1, + dQ_single_wg=True, + ) + elif head_dim <= 128: + # C++ FA3: causal/local: 64, 128; non-causal: 80, 128 with dQ_swapAB + is_causal_or_local = causal or local + m_block_size = 64 if is_causal_or_local else 80 + if sparse_block_size_q is not None and sparse_block_size_q % m_block_size != 0: + m_block_size = 64 + return BwdConfig( + m_block_size=m_block_size, + n_block_size=128, + num_stages_Q=2, num_stages_dO=2, num_stages_PdS=2, + SdP_swapAB=True, dKV_swapAB=False, + dQ_swapAB=m_block_size % 64 != 0, + AtomLayoutMSdP=1, AtomLayoutNdKV=2, AtomLayoutMdQ=1, + ) + elif head_dim <= 192: + hdimv128 = head_dim_v <= 128 + if hdimv128: + return BwdConfig( + m_block_size=64, n_block_size=96, + num_stages_Q=2, num_stages_dO=2, num_stages_PdS=1, + SdP_swapAB=False, dKV_swapAB=True, dQ_swapAB=False, + AtomLayoutMSdP=1, AtomLayoutNdKV=2, AtomLayoutMdQ=1, + num_wg=2, + ) + else: + return BwdConfig( + m_block_size=64, n_block_size=96, + num_stages_Q=2, num_stages_dO=1, num_stages_PdS=1, + SdP_swapAB=False, dKV_swapAB=True, dQ_swapAB=False, + AtomLayoutMSdP=1, AtomLayoutNdKV=2, AtomLayoutMdQ=1, + num_wg=2, + ) + else: + # hdim 256 + return BwdConfig( + m_block_size=64, n_block_size=64, + num_stages_Q=1, num_stages_dO=1, num_stages_PdS=1, + SdP_swapAB=False, dKV_swapAB=False, dQ_swapAB=False, + AtomLayoutMSdP=1, AtomLayoutNdKV=1, AtomLayoutMdQ=1, + ) + + + +def maybe_contiguous(x): + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def _validate_tensor(t, name, expected_shape, expected_dtype, expected_device): + assert t.shape == expected_shape, f"{name} shape {t.shape} != expected {expected_shape}" + assert t.dtype == expected_dtype, f"{name} dtype {t.dtype} != expected {expected_dtype}" + assert t.device == expected_device, f"{name} device {t.device} != expected {expected_device}" + if not is_fake_mode(): + assert t.is_cuda, f"{name} must be on CUDA" + +torch2cute_dtype_map = { + torch.float16: cutlass.Float16, + torch.bfloat16: cutlass.BFloat16, + torch.float32: cutlass.Float32, + torch.float8_e4m3fn: cutlass.Float8E4M3FN, + torch.float8_e5m2: cutlass.Float8E5M2, +} + + +def num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, max_splits): + # If num_n_blocks is too small, use 1 split. For example, we never split for hdim = 128 and seqlen_k = 512. + if num_n_blocks <= 4: + return 1 + # Avoid ZeroDivisionError when batch_size or seqlen_q is 0. The empty-Q + # early-exit in _flash_attn_fwd handles correctness for those shapes; this + # guard just keeps the heuristic safe if called in other contexts. + if total_mblocks == 0: + return 1 + + # NOTE: We should revisit this heuristic after persistence is supported for split KV. + # Sometimes, it's ideal to over-schedule splits for better efficiency. + return min(num_SMs // total_mblocks, max_splits, num_n_blocks) + + +def _resolve_causal_local_window(causal, window_size_left, window_size_right, mask_mod=None): + """Resolve causal/local/window settings into canonical form. + + Returns (causal, local, window_size_left, window_size_right). + """ + if mask_mod is not None: + return False, False, window_size_left, window_size_right + if causal: + window_size_right = 0 + if window_size_left is not None and window_size_right is not None and window_size_left + window_size_right < 0: + window_size_left = None + window_size_right = None + if window_size_left is not None or window_size_right is not None: + if window_size_left is None and window_size_right == 0: + causal, local = True, False + window_size_right = None + else: + causal, local = False, True + else: + local = False + return causal, local, window_size_left, window_size_right + +def _flash_attn_fwd( + q: Optional[torch.Tensor], + k: Optional[torch.Tensor], + v: torch.Tensor, + qv: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + min_seqlen_k: Optional[int] = None, + page_table: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + softcap: Optional[float] = None, + window_size_left: Optional[int] = None, + window_size_right: Optional[int] = None, + learnable_sink: Optional[torch.Tensor] = None, + tile_mn: Optional[Tuple[int, int]] = None, + mma_pv_is_rs: Optional[bool] = None, + intra_wg_overlap: Optional[bool] = None, + num_threads: int = 384, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + _arch: Optional[int] = None, + score_mod: Optional[Callable] = None, + mask_mod: Optional[Callable] = None, + block_sparse_tensors: Optional[BlockSparseTensorsTorch] = None, + return_lse: bool = False, + out: Optional[torch.Tensor] = None, + lse: Optional[torch.Tensor] = None, + aux_tensors: Optional[list[torch.Tensor]] = None, + aux_scalars: Optional[tuple] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + gather_kv_indices: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """Forward pass for FlashAttention. + + Args: + ... + score_mod: A callable that takes the attention scores and applies a modification. + mask_mod: A callable that takes token position information and selectively masks + block_sparse_tensors: A tuple of tensors used for block sparsity. + return_lse: Whether to return the log softmax of the attention scores. If set to True will always calculate + The returned LSE supports taking gradient. + out: Optional pre-allocated output tensor. If None, will be allocated internally. + lse: Optional pre-allocated log-sum-exp tensor. If None, will be allocated when needed. + aux_tensors: Some score_mods will want to read from global aux_tensors. This is how we thread them through to the inner kernel. + aux_scalars: Runtime scalar captures used by score_mod or mask_mod. + """ + aux_scalars = tuple(aux_scalars) if aux_scalars else None + q, k, v, qv = [maybe_contiguous(t) for t in (q, k, v, qv)] + assert q is not None or qv is not None + assert v is not None + q_descale, k_descale, v_descale = [maybe_contiguous(t) for t in (q_descale, k_descale, v_descale)] + q_shape = q.shape if q is not None else qv.shape + num_head, head_dim = q_shape[-2:] + if cu_seqlens_q is None: + batch_size, seqlen_q = q_shape[:2] + total_q = batch_size * seqlen_q + else: + batch_size = cu_seqlens_q.shape[0] - 1 + seqlen_q = None + total_q = q_shape[0] + if page_table is not None: + assert cu_seqlens_k is None, "page_table is not supported with cu_seqlens_k" + assert page_table.dtype == torch.int32, "page_table must be int32" + assert page_table.stride(-1) == 1, "page_table must be contiguous in the last dimension" + max_num_pages_per_seq = page_table.shape[1] + assert page_table.shape == (batch_size, max_num_pages_per_seq) + num_pages, page_size = v.shape[:2] + seqlen_k = num_pages * page_size + else: + num_pages, page_size = None, None + seqlen_k = v.shape[-3] + num_head_kv = v.shape[-2] + head_dim_v = v.shape[-1] + if cu_seqlens_k is None: + if page_table is None: + assert k is None or k.shape == (batch_size, seqlen_k, num_head_kv, head_dim) + assert v.shape == (batch_size, seqlen_k, num_head_kv, head_dim_v) + else: + assert k is None or k.shape == (num_pages, page_size, num_head_kv, head_dim) + assert v.shape == (num_pages, page_size, num_head_kv, head_dim_v) + else: + assert k is None or k.shape == (seqlen_k, num_head_kv, head_dim) + assert v.shape == (seqlen_k, num_head_kv, head_dim_v) + assert cu_seqlens_k.shape == (batch_size + 1,), ( + "cu_seqlens_k must have shape (batch_size + 1,)" + ) + + if cu_seqlens_q is not None: + assert cu_seqlens_q.shape == (batch_size + 1,), ( + "cu_seqlens_q must have shape (batch_size + 1,)" + ) + assert seqused_q is None or seqused_q.shape == (batch_size,), ( + "seqused_q must have shape (batch_size,)" + ) + assert seqused_k is None or seqused_k.shape == (batch_size,), ( + "seqused_k must have shape (batch_size,)" + ) + assert v.dtype in [torch.float16, torch.bfloat16, torch.float8_e4m3fn, torch.float8_e5m2], ( + "inputs must be float16, bfloat16, fp8 e4m3fn, or fp8 e5m2" + ) + + input_tensors = {"q": q, "k": k, "v": v, "qv": qv} + present = {name: t for name, t in input_tensors.items() if t is not None} + names = list(present.keys()) + for i in range(len(names)): + for j in range(i + 1, len(names)): + a, b = names[i], names[j] + assert present[a].dtype == present[b].dtype, f"{a}.dtype {present[a].dtype} != {b}.dtype {present[b].dtype}" + + q_dtype = q.dtype if q is not None else qv.dtype + + for t in [cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k]: + if t is not None: + assert t.dtype == torch.int32, ( + "cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k must be int32" + ) + assert t.stride(0) == 1, ( + "cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k must be contiguous" + ) + if learnable_sink is not None: + assert learnable_sink.shape == (num_head,) + assert learnable_sink.dtype == torch.bfloat16, "learnable_sink must be bfloat16" + + if not is_fake_mode(): + assert all( + t is None or t.is_cuda + for t in ( + q, + k, + v, + qv, + q_descale, + k_descale, + v_descale, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + page_table, + learnable_sink, + ) + ), "inputs must be on CUDA device" + arch = _get_device_arch() if _arch is None else _arch + assert arch // 10 in [8, 9, 10, 11, 12], "Unsupported compute capability. Supported: 8.x, 9.x, 10.x, 11.x, 12.x" + assert num_head % num_head_kv == 0, "num_head must be divisible by num_head_kv" + alignment = 16 // v.element_size() + if arch // 10 not in [8, 12]: + _validate_head_dims(head_dim, head_dim_v, arch // 10, alignment) + if softmax_scale is None: + softmax_scale = ( + 1.0 / math.sqrt(head_dim) if qv is None or q is None + else 1.0 / math.sqrt(head_dim + head_dim_v) + ) + if softcap == 0.0: + softcap = None + qhead_per_kvhead = num_head // num_head_kv + if pack_gqa is None: + pack_gqa = qhead_per_kvhead > 1 + + is_fp8 = v.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) + requires_grad = any(t is not None and t.requires_grad for t in [q, k, v, qv]) + if is_fp8 and requires_grad: + raise NotImplementedError("FA4 CuTe FP8 backward is not supported yet (forward-only).") + out_torch_dtype = torch.bfloat16 if is_fp8 else q_dtype + device = v.device + q_batch_seqlen_shape = (batch_size, seqlen_q) if cu_seqlens_q is None else (total_q,) + + if qv is None: + lse_shape = (batch_size, num_head, seqlen_q) if cu_seqlens_q is None else (num_head, total_q) + else: + # num_head contiguous better for MQA in MLA absorbed + lse_shape = (batch_size, seqlen_q, num_head) if cu_seqlens_q is None else (total_q, num_head) + + if out is None: + out = torch.empty( + *q_batch_seqlen_shape, num_head, head_dim_v, dtype=out_torch_dtype, device=device + ) + else: + _validate_tensor(out, "out", (*q_batch_seqlen_shape, num_head, head_dim_v), out_torch_dtype, device) + + if lse is None: + lse = ( + torch.empty(lse_shape, dtype=torch.float32, device=device) + if requires_grad or return_lse + else None + ) + elif lse is not None: + _validate_tensor(lse, "lse", lse_shape, torch.float32, device) + + if seqlen_k == 0 or total_q == 0: + out.zero_() + if lse is not None: + lse.fill_(float("-inf")) + return out, lse, None, None + + if is_fp8: + for t, name in ((q_descale, "q_descale"), (k_descale, "k_descale"), (v_descale, "v_descale")): + if t is not None: + _validate_tensor(t, name, (batch_size, num_head_kv), torch.float32, device) + else: + assert q_descale is None and k_descale is None and v_descale is None, ( + "q_descale/k_descale/v_descale are only supported for FP8 inputs" + ) + + dtype = torch2cute_dtype_map[q_dtype] + if is_fp8: + assert arch // 10 == 10, "FP8 is only supported on SM100 (compute capability 10.x) for FA4 CuTe." + use_block_sparsity = block_sparse_tensors is not None + + causal, local, window_size_left, window_size_right = _resolve_causal_local_window( + causal, window_size_left, window_size_right, mask_mod + ) + + requested_use_clc_scheduler = utils._get_use_clc_scheduler_default() + requested_disable_2cta = utils._get_disable_2cta_default(is_fwd=True) + + current_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + + # SM80/SM120: uses SM80 MMA, 128 threads (4 warps) + if arch // 10 in [8, 12]: + num_threads = 128 + + fwd_cfg = FwdConfig(128, 128, True, True) # default + if tile_mn is None: + if arch // 10 == 12: + # SM120 tile sizes tuned for 99 KB SMEM capacity: + # D<=64: 128x128 → 48 KB (good occupancy) + # D>64: 128x64 → 64 KB (128x128 would use 96 KB, hurting occupancy) + if head_dim <= 64: + fwd_cfg = FwdConfig(128, 128, True, True) + else: + fwd_cfg = FwdConfig(128, 64, True, True) + elif arch // 10 == 8: + fwd_cfg = FwdConfig(128, 64, True, True) # SM80, should tune + elif arch // 10 == 9: + sparse_q = get_sparse_q_block_size(block_sparse_tensors, seqlen_q) + fwd_cfg = _tile_size_fwd_sm90(head_dim, head_dim_v, causal, local, sparse_block_size_q=sparse_q) + elif _requests_sm100_sparse_64x64_tiles(arch, block_sparse_tensors): + fwd_cfg = FwdConfig(64, 64, True, True) + else: + fwd_cfg = FwdConfig(tile_mn[0], tile_mn[1], fwd_cfg.mma_pv_is_rs, fwd_cfg.intra_wg_overlap) + tile_m, tile_n = fwd_cfg.m_block_size, fwd_cfg.n_block_size + p_in_smem = arch // 10 in [10, 11] and tile_m == 64 + if p_in_smem: + if arch != 100: + raise NotImplementedError("tile_m=64 forward is currently supported only on SM100") + if tile_n not in (64, 128): + raise NotImplementedError( + f"SM100 tile_m=64 forward requires tile_n in (64, 128), got tile_n={tile_n}" + ) + if is_fp8: + raise NotImplementedError("SM100 tile_m=64 forward currently supports only FP16 and BF16") + if qv is not None or not supports_sm100_tile_m64_head_dims(head_dim, head_dim_v): + raise NotImplementedError( + "SM100 tile_m=64 forward supports head dimensions in [8, 128] or (192, 128)" + ) + if page_table is not None and head_dim != head_dim_v: + raise NotImplementedError( + "SM100 tile_m=64 paged KV requires matching K and V head dimensions" + ) + if page_table is not None and page_size != tile_n: + raise NotImplementedError( + "SM100 tile_m=64 forward requires paged KV page_size to equal tile_n" + ) + if mma_pv_is_rs is None: + mma_pv_is_rs = fwd_cfg.mma_pv_is_rs + if intra_wg_overlap is None: + intra_wg_overlap = fwd_cfg.intra_wg_overlap + + if max_seqlen_q is None: + max_seqlen_q = seqlen_q if cu_seqlens_q is None else total_q + if max_seqlen_k is None: + max_seqlen_k = seqlen_k + if cu_seqlens_k is None and seqused_k is None: + min_seqlen_k = seqlen_k + seqlen_q_packgqa = max_seqlen_q * qhead_per_kvhead + if arch // 10 in [10, 11]: + sparse_q = get_sparse_q_block_size(block_sparse_tensors, seqlen_q) + # Dense FA4 amortizes one CTA across two 128-row Q stages. A sparse + # mask whose declared Q block is exactly one tile must retain that + # granularity: combining two rows would union distinct VSA selections. + # The SM100 kernel already has a q_stage=1 specialization; selecting it + # here also creates the base schedule for FastVideo's dual-KV-stream + # single-Q-block path. + single_q_sparse_tile = use_block_sparsity and sparse_q == tile_m + q_stage = 1 if p_in_smem or single_q_sparse_tile or seqlen_q_packgqa <= tile_m else 2 + else: + q_stage = 1 + m_block_size_effective = q_stage * tile_m + seqlen_k_loaded = max_seqlen_k if not local else max(0, min(max_seqlen_k, (window_size_right or max_seqlen_k) + (window_size_left or max_seqlen_k) + 1 + tile_m)) + num_m_blocks = (seqlen_q_packgqa + m_block_size_effective - 1) // m_block_size_effective + total_mblocks = batch_size * num_head_kv * num_m_blocks + num_n_blocks = (seqlen_k_loaded + tile_n - 1) // tile_n + num_SMs = 132 if is_fake_mode() else torch.cuda.get_device_properties(device).multi_processor_count + if num_splits < 1: + num_splits = num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, 128) + + # SplitKV uses float32 partial output, which doubles the O buffer size + # in shared memory, causing OOM for diff-headdim (192, 128) + if arch // 10 in [10, 11] and head_dim != head_dim_v and num_splits > 1: + if num_n_blocks >= 64 and head_dim_v != 512: + tile_n = 64 + num_n_blocks = (seqlen_k_loaded + tile_n - 1) // tile_n + num_splits = num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, 128) + else: + num_splits = 1 + + is_split_kv = num_splits > 1 + if is_split_kv: + out_partial = torch.empty(num_splits, *q_batch_seqlen_shape, num_head, head_dim_v, dtype=torch.float32, device=device) + lse_partial = torch.empty(num_splits, *lse_shape, dtype=torch.float32, device=device) + + use_2cta_instrs = ( + arch // 10 in [10, 11] + and not requested_disable_2cta + and not causal + and not local + and not is_split_kv + and cu_seqlens_q is None + and seqused_q is None + and not use_block_sparsity + and not p_in_smem + and page_size in [None, 128] + and int(math.ceil(head_dim / 16) * 16) in [128, 192] + and int(math.ceil(head_dim_v / 16) * 16) == 128 + and seqlen_q_packgqa > 2 * tile_m + and (tile_m % qhead_per_kvhead == 0 or not pack_gqa) + ) + + # hd=256 2CTA forward uses dedicated kernel (Blackwell family) + use_dedicated_hd256_kernel = arch // 10 in [10, 11] and head_dim == 256 and head_dim_v == 256 + use_2cta_instrs = use_2cta_instrs or use_dedicated_hd256_kernel + + if softcap is not None: + assert score_mod is None, "softcap and score_mod cannot be used together" + score_mod = utils.create_softcap_scoremod(softcap) + elif score_mod is not None: + if arch // 10 == 8: + raise NotImplementedError("Custom user-provided score_mod is not supported on SM8x architectures.") + + # hash score and mask mods for compile cache + score_mod_hash = utils.hash_callable(score_mod) if score_mod is not None else False + mask_mod_hash = utils.hash_callable(mask_mod) if mask_mod is not None else False + + is_varlen = ( + cu_seqlens_q is not None + or cu_seqlens_k is not None + or seqused_q is not None + or seqused_k is not None + ) + + # CLC regressed for varlen MHA and dense noncausal. Imbalanced varlen shapes + # keep more K/V blocks in flight and hurt L2; dense noncausal mostly just + # pays work-stealing overhead. + is_varlen_mha = is_varlen and qhead_per_kvhead == 1 + is_dense_noncausal = not is_varlen and not causal and not local + use_clc_scheduler = ( + requested_use_clc_scheduler + and not is_varlen_mha + and not is_dense_noncausal + and not p_in_smem + ) + + if use_block_sparsity: + # NB: pack_gqa requires block sparse head dim == 1 (broadcasted) + head_dim_idx = 0 if block_sparse_tensors.mask_block_cnt.ndim == 2 else 1 + if pack_gqa and block_sparse_tensors.mask_block_cnt.shape[head_dim_idx] != 1: + pack_gqa = False + if cu_seqlens_q is not None: + assert block_sparse_tensors.cu_total_m_blocks is not None, ( + "Varlen block sparsity requires block_sparse_tensors.cu_total_m_blocks." + ) + + dual_stream_enabled = os.environ.get("FASTVIDEO_FA4_VSA_DUAL_STREAM", "1") != "0" + sp_double_buffer_enabled = ( + os.environ.get("FASTVIDEO_FA4_VSA_SP_DOUBLE_BUFFER", "1") != "0" + ) + num_kv_streams = ( + 2 + if dual_stream_enabled + and arch // 10 in [10, 11] + and use_block_sparsity + and q_stage == 1 + and tile_m in (64, 128) + and tile_n == 64 + and q_dtype in [torch.float16, torch.bfloat16] + and page_table is None + and not pack_gqa + and not is_varlen + and not is_split_kv + and learnable_sink is None + else 1 + ) + + use_sp_double_buffer = ( + sp_double_buffer_enabled + and num_kv_streams == 2 + and arch == 100 + and tile_m == 128 + and tile_n == 64 + and int(math.ceil(head_dim / 16) * 16) == 128 + and int(math.ceil(head_dim_v / 16) * 16) == 128 + and not causal + and not local + and page_table is None + and not pack_gqa + and not is_varlen + and not is_split_kv + and learnable_sink is None + and score_mod is None + ) + + # See get_broadcast_dims for why this is needed in compile key + block_sparse_broadcast_pattern = None + normalized_block_sparse_tensors = None + q_subtile_factor = 1 + if block_sparse_tensors is not None: + ( + normalized_block_sparse_tensors, + block_sparse_broadcast_pattern, + q_subtile_factor, + ) = normalize_block_sparse_config( + block_sparse_tensors, + batch_size=batch_size, + num_head=num_head, + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + block_size=(tile_m, tile_n), + q_stage=q_stage, + ) + if aux_tensors is not None: + aux_tensor_metadata = get_aux_tensor_metadata(aux_tensors) + else: + aux_tensor_metadata = None + aux_scalar_metadata = tuple(type(s) for s in aux_scalars) if aux_scalars is not None else None + + if qv is not None: + assert arch // 10 in [10, 11], "only support Blackwell arch with qv" + assert q is None or qv.shape[:-1] == q.shape[:-1] + assert qv.shape[-1] == head_dim_v + assert head_dim_v == 512 + assert q is None or head_dim == 64 + assert not local, "local not yet supported with qv" + assert q_descale is None and k_descale is None and v_descale is None, ( + "q_descale/k_descale/v_descale are not yet supported with qv" + ) + assert tile_n == 128 + + assert not is_split_kv, "split kv not supported with qv" + assert learnable_sink is None + assert softcap is None + assert score_mod is None + assert mask_mod is None + + if page_table is not None: + assert gather_kv_indices is None, "paged KV + topk sparsity not yet supported together" + + qv = maybe_contiguous(qv) + + gather_kv_length = 2048 # dummy value + sparse_kv = gather_kv_indices is not None + # always use kv bitmask by default (handles -1 sentinel) + disable_sparse_kv_bitmask = False + if sparse_kv: + assert gather_kv_indices.shape[:-1] == qv.shape[:-2] + gather_kv_length = gather_kv_indices.shape[-1] + assert gather_kv_length % 128 == 0 + # if min_seqlen_k is None or causal: + # disable_sparse_kv_bitmask = False + # else: + # # seqlen_k_boundary = min_seqlen_k - max_seqlen_q + 1 if causal else min_seqlen_k + # seqlen_k_boundary = min_seqlen_k + # disable_sparse_kv_bitmask = seqlen_k_boundary >= gather_kv_length + + if requires_grad and sparse_kv: + if cu_seqlens_q is None: + p = torch.empty(batch_size, seqlen_q, num_head, gather_kv_length, dtype=q_dtype, device=device) + row_max = torch.empty(batch_size, seqlen_q, gather_kv_length//128, num_head, dtype=torch.float32, device=device) + else: + p = torch.empty(total_q, num_head, gather_kv_length, dtype=q_dtype, device=device) + row_max = torch.empty(total_q, gather_kv_length//128, num_head, dtype=torch.float32, device=device) + else: + p = row_max = None + else: + assert gather_kv_indices is None, "gather_kv_indices is only supported with qv" + gather_kv_length = None + sparse_kv = None + disable_sparse_kv_bitmask = None + p = row_max = None + + compile_key = ( + dtype, + head_dim, + head_dim_v, + qhead_per_kvhead, + causal, + score_mod_hash, + mask_mod_hash, + use_block_sparsity, + block_sparse_broadcast_pattern, + aux_tensor_metadata, + aux_scalar_metadata, + lse is None, + cu_seqlens_q is None, + cu_seqlens_k is None, + seqused_q is None, + seqused_k is None, + page_table is not None, + window_size_left is not None, + window_size_right is not None, + learnable_sink is not None, + q_descale is not None, + k_descale is not None, + v_descale is not None, + block_sparse_tensors is None or block_sparse_tensors.cu_total_m_blocks is None, + block_sparse_tensors is None or block_sparse_tensors.cu_block_idx_offsets is None, + tile_m, + tile_n, + q_stage, + num_kv_streams, + use_sp_double_buffer, + num_threads, + is_split_kv, + pack_gqa, + arch, + page_size not in [None, tile_n], # paged KV non-TMA + use_2cta_instrs, + q_subtile_factor, + mma_pv_is_rs, + intra_wg_overlap, + use_clc_scheduler, + q is not None, + qv is not None, + p is not None, + row_max is not None, + gather_kv_length, + sparse_kv, + disable_sparse_kv_bitmask, + fa_logging.get_fa_log_level(), + ) + + if compile_key not in _flash_attn_fwd.compile_cache: + ( + cu_seqlens_q_tensor, + cu_seqlens_k_tensor, + seqused_q_tensor, + seqused_k_tensor, + learnable_sink_tensor, + ) = [ + to_cute_tensor(t, assumed_align=4, leading_dim=0) + if t is not None + else None + for t in (cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, learnable_sink) + ] + page_table_tensor = ( + to_cute_tensor(page_table, assumed_align=4, leading_dim=1) + if page_table is not None + else None + ) + q_tensor, k_tensor, v_tensor, o_tensor = [ + to_cute_tensor(t) for t in (q, k, v, out if not is_split_kv else out_partial) + ] + if is_split_kv: + lse_tensor = to_cute_tensor(lse_partial, assumed_align=4) + else: + lse_tensor = to_cute_tensor(lse, assumed_align=4) + + q_descale_tensor, k_descale_tensor, v_descale_tensor = ( + to_cute_tensor(t, assumed_align=4, leading_dim=1) + for t in (q_descale, k_descale, v_descale) + ) + descale_tensors_tensor = ( + DescaleTensors( + q_descale=q_descale_tensor, + k_descale=k_descale_tensor, + v_descale=v_descale_tensor, + ) + if q_descale_tensor is not None + or k_descale_tensor is not None + or v_descale_tensor is not None + else None + ) + + sparse_tensors = None + if normalized_block_sparse_tensors is not None: + sparse_tensors = to_cute_block_sparse_tensors(normalized_block_sparse_tensors) + + cute_aux_tensors = None + aux_tensor_metadata = None + if aux_tensors is not None: + cute_aux_tensors = [to_cute_aux_tensor(buf) for buf in aux_tensors] + + qv_tensor = to_cute_tensor(qv) + gather_kv_indices_tensor = to_cute_tensor(gather_kv_indices) + p_tensor = to_cute_tensor(p) + row_max_tensor = to_cute_tensor(row_max) + + if arch // 10 == 8: + assert page_table is None, "paged KV not supported on SM 8.0" + assert not is_split_kv, "SplitKV not supported on SM 8.0" + fa_fwd = FlashAttentionForwardSm80( + dtype, + head_dim, + head_dim_v, + qhead_per_kvhead, + is_causal=causal, + is_local=local, + pack_gqa=pack_gqa, + tile_m=tile_m, + tile_n=tile_n, + num_stages=1, + num_threads=num_threads, + Q_in_regs=False, + score_mod=score_mod, + mask_mod=mask_mod, + has_aux_tensors=aux_tensors is not None, + ) + elif arch // 10 == 9: + assert not is_split_kv, "SplitKV not supported on SM 9.0" + fa_fwd = FlashAttentionForwardSm90( + dtype, + head_dim, + head_dim_v, + qhead_per_kvhead, + is_causal=causal, + is_local=local, + pack_gqa=pack_gqa, + tile_m=tile_m, + tile_n=tile_n, + # num_stages=1, + num_stages=2, + num_threads=num_threads, + Q_in_regs=False, + intra_wg_overlap=intra_wg_overlap, + mma_pv_is_rs=mma_pv_is_rs, + mask_mod=mask_mod, + score_mod=score_mod, + has_aux_tensors=aux_tensors is not None, + q_subtile_factor=q_subtile_factor, + paged_kv_non_tma=page_size not in [None, tile_n], + ) + elif arch // 10 in [10, 11]: + if qv is not None: + paged_kv_cpasync = page_table is not None and page_size != tile_n + has_qk = q is not None + fa_fwd = FlashAttentionMLAForwardSm100( + is_causal=causal, + use_cpasync_load_KV=sparse_kv or paged_kv_cpasync, + topk_length=gather_kv_length, + is_topk_gather=sparse_kv, + pack_gqa=pack_gqa, + qhead_per_kvhead=qhead_per_kvhead, + nheads_kv=num_head_kv, + has_seqused_q=seqused_q is not None, + has_cu_seqlens_q=cu_seqlens_q is not None, + disable_bitmask=disable_sparse_kv_bitmask, + has_qk=has_qk, + ) + else: + if use_dedicated_hd256_kernel: + # hd=256 2CTA forward: check for currently unsupported features + assert softcap is None, "SM100 forward with head_dim=256 does not support softcap" + assert not use_block_sparsity, \ + "SM100 forward with head_dim=256 does not support block sparsity" + assert learnable_sink is None, \ + "SM100 forward with head_dim=256 does not support learnable_sink" + assert seqused_q is None and seqused_k is None, \ + "SM100 forward with head_dim=256 does not support seqused_q/seqused_k" + if page_table is not None: + assert max_seqlen_k % page_size == 0, ( + f"SM100 hd256 2CTA paged KV requires max_seqlen_k divisible by " + f"page_size ({page_size}), got max_seqlen_k={max_seqlen_k}" + ) + assert page_table.shape[1] == max_seqlen_k // page_size, ( + f"SM100 hd256 2CTA paged KV requires page_table.shape[1] == " + f"max_seqlen_k // page_size ({max_seqlen_k} // {page_size} = " + f"{max_seqlen_k // page_size}), got {page_table.shape[1]}; " + f"pass page_table[:, :{max_seqlen_k // page_size}] to slice to " + f"the actual sequence length" + ) + assert page_table.stride(0) == page_table.shape[1], ( + f"SM100 hd256 2CTA paged KV requires a fully contiguous page_table " + f"(stride(0)={page_table.stride(0)} must equal " + f"shape[1]={page_table.shape[1]})" + ) + # pack_gqa is an auto-selected optimization; disable it for hd256 kernel + pack_gqa = False + # The hd256 dedicated kernel builds tensor layouts with hardcoded + # contiguous strides computed from shape dimensions, so non-contiguous + # inputs (e.g. from .transpose()) produce wrong memory accesses. + # maybe_contiguous() above only guarantees stride(-1)==1; make fully + # contiguous here before the compile key is derived from shapes. + q = q.contiguous() if not q.is_contiguous() else q + k = k.contiguous() if not k.is_contiguous() else k + v = v.contiguous() if not v.is_contiguous() else v + + flash_fwd_obj_cls = ( + BlackwellFusedMultiHeadAttentionForward + if use_dedicated_hd256_kernel + else FlashAttentionForwardSm100 + ) + kv_stream_kwargs = ( + {} + if use_dedicated_hd256_kernel + else { + "num_kv_streams": num_kv_streams, + "use_sp_double_buffer": use_sp_double_buffer, + } + ) + + fa_fwd = flash_fwd_obj_cls( + head_dim, + head_dim_v, + qhead_per_kvhead=qhead_per_kvhead, + is_causal=causal, + is_local=local, + is_split_kv=is_split_kv, + pack_gqa=pack_gqa, + m_block_size=tile_m, + n_block_size=tile_n, + q_stage=q_stage, + is_persistent=not causal + and not local + and cu_seqlens_q is None + and seqused_q is None + and not is_split_kv, + score_mod=score_mod, + mask_mod=mask_mod, + has_aux_tensors=aux_tensors is not None, + paged_kv_non_tma=page_size not in [None, tile_n], + is_varlen_q=cu_seqlens_q is not None or seqused_q is not None, + q_subtile_factor=q_subtile_factor, + use_2cta_instrs=use_2cta_instrs, + use_clc_scheduler=use_clc_scheduler, + **kv_stream_kwargs, + ) + elif arch // 10 == 12: + # SM120 (Blackwell GeForce / DGX Spark): uses SM80 MMA with SM120 SMEM capacity + assert not use_block_sparsity, "Block sparsity not supported on SM 12.0" + assert page_table is None, "Paged KV not supported on SM 12.0 in this PR" + assert not is_split_kv, "SplitKV not supported on SM 12.0 in this PR" + fa_fwd = FlashAttentionForwardSm120( + dtype, + head_dim, + head_dim_v, + qhead_per_kvhead, + is_causal=causal, + is_local=local, + pack_gqa=pack_gqa, + tile_m=tile_m, + tile_n=tile_n, + num_stages=1, + num_threads=num_threads, + Q_in_regs=False, + score_mod=score_mod, + mask_mod=mask_mod, + has_aux_tensors=aux_tensors is not None, + ) + else: + raise ValueError( + f"Unsupported compute capability: {arch}. Supported: 8.x, 9.x, 10.x, 11.x, 12.x" + ) + # TODO: check @can_implement + if qv is not None: + _flash_attn_fwd.compile_cache[compile_key] = cute.compile( + fa_fwd, + q_tensor, + qv_tensor, + k_tensor, + v_tensor, + o_tensor, + lse_tensor, + softmax_scale, + p_tensor, + row_max_tensor, + cu_seqlens_q_tensor, + cu_seqlens_k_tensor, + seqused_q_tensor, + seqused_k_tensor, + gather_kv_indices_tensor, + page_table_tensor, + window_size_left, + window_size_right, + current_stream, + options="--enable-tvm-ffi", + ) + else: + compile_args = [ + fa_fwd, + q_tensor, + k_tensor, + v_tensor, + o_tensor, + lse_tensor, + softmax_scale, + cu_seqlens_q_tensor, + cu_seqlens_k_tensor, + seqused_q_tensor, + seqused_k_tensor, + page_table_tensor, + window_size_left, + window_size_right, + learnable_sink_tensor, + ] + if arch // 10 in [10, 11]: + compile_args.append(descale_tensors_tensor) + compile_args.extend([ + sparse_tensors, + AuxData(cute_aux_tensors, aux_scalars), + ]) + compile_args.append(current_stream) + _flash_attn_fwd.compile_cache[compile_key] = cute.compile( + *compile_args, options="--enable-tvm-ffi" + ) + + if not is_fake_mode(): + q_call, k_call, v_call, qv_call = [ + t.detach() if t is not None else None + for t in (q, k, v, qv) + ] + if is_fp8: + # need uint8 workaround until we pin torch >= 2.11.0 where fp8 export is supported + q_call, k_call, v_call, qv_call = [ + t.view(torch.uint8) if t is not None else None + for t in (q_call, k_call, v_call, qv_call) + ] + descale_tensors = ( + DescaleTensors(q_descale=q_descale, k_descale=k_descale, v_descale=v_descale) + if q_descale is not None or k_descale is not None or v_descale is not None + else None + ) + if qv is not None: + _flash_attn_fwd.compile_cache[compile_key]( + q_call, + qv_call, + k_call, + v_call, + out.detach(), + lse, + softmax_scale, + p, + row_max, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + gather_kv_indices, + page_table, + window_size_left, + window_size_right, + ) + else: + call_args = [ + q_call, + k_call, + v_call, + out.detach() if not is_split_kv else out_partial, + lse_partial if is_split_kv else lse, + softmax_scale, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + page_table, + window_size_left, + window_size_right, + learnable_sink, + ] + if arch // 10 in [10, 11]: + call_args.append(descale_tensors) + call_args.extend([ + ( + normalized_block_sparse_tensors.mask_block_cnt, + normalized_block_sparse_tensors.mask_block_idx, + normalized_block_sparse_tensors.full_block_cnt, + normalized_block_sparse_tensors.full_block_idx, + normalized_block_sparse_tensors.cu_total_m_blocks, + normalized_block_sparse_tensors.cu_block_idx_offsets, + normalized_block_sparse_tensors.dq_write_order, + normalized_block_sparse_tensors.dq_write_order_full, + ) + if normalized_block_sparse_tensors is not None + else None, + AuxData(aux_tensors, aux_scalars), + ]) + _flash_attn_fwd.compile_cache[compile_key](*call_args) + if is_split_kv: + _flash_attn_fwd_combine( + out_partial, + lse_partial.transpose(-1, -2), + out, + lse.transpose(-1, -2) if lse is not None else None, + cu_seqlens_q, + seqused_q, + ) + return out, lse, p, row_max + + +_flash_attn_fwd.compile_cache = get_jit_cache("fwd") + + +def make_fake_bwd_tensors(dtype, has_gqa, varlen_q, varlen_k, nheads_major=False): + sym = cute.sym_int + # divisibility in elements: assumed_align_bytes = divisibility * dtype.width // 8 + # For 16-byte align: fp16/bf16 → divisibility=8, float32 → divisibility=4 + div = 128 // dtype.width # 8 for fp16/bf16 + # Shared sym_ints for dimensions that must match across tensors + b, seqlen_q, seqlen_k, h_q, d, d_v = sym(), sym(), sym(), sym(), sym(), sym() + topk = sym() + h_kv = h_q if not has_gqa else sym() + seqlen_q_rounded, seqlen_k_rounded = sym(), sym() + seqlen_q_d_rounded, seqlen_k_d_rounded, seqlen_k_dv_rounded = sym(), sym(), sym() + total_q, total_k, total_q_rounded, total_k_rounded = sym(), sym(), sym(), sym() + total_q_d_rounded, total_k_d_rounded, total_k_dv_rounded = sym(), sym(), sym() + b_seqlenq = (b, seqlen_q) if not varlen_q else (total_q,) + b_seqlenk = (b, seqlen_k) if not varlen_k else (total_k,) + mQ = fake_tensor(dtype, (*b_seqlenq, h_q, d), divisibility=div) + mO = fake_tensor(dtype, (*b_seqlenq, h_q, d_v), divisibility=div) + mdO = fake_tensor(dtype, (*b_seqlenq, h_q, d_v), divisibility=div) + mK = fake_tensor(dtype, (*b_seqlenk, h_kv, d), divisibility=div) + mV = fake_tensor(dtype, (*b_seqlenk, h_kv, d_v), divisibility=div) + mdQ = fake_tensor(dtype, (*b_seqlenq, h_q, d), divisibility=div) + mdK = fake_tensor(dtype, (*b_seqlenk, h_kv, d), divisibility=div) + mdV = fake_tensor(dtype, (*b_seqlenk, h_kv, d_v), divisibility=div) + + sq = seqlen_q if not varlen_q else total_q + sq_r = seqlen_q_rounded if not varlen_q else total_q_rounded + sq_dr = seqlen_q_d_rounded if not varlen_q else total_q_d_rounded + + def shape(*dims): + batch = (b,) if not varlen_q else () + return (*batch, h_q, *dims) if not nheads_major else (*batch, *dims, h_q) + + mLSE = fake_tensor(Float32, shape(sq), divisibility=1) + mLSElog2 = fake_tensor(Float32, shape(sq_r), divisibility=4) + mPdPsum = fake_tensor(Float32, shape(sq_r), divisibility=4) + dQaccum = fake_tensor(Float32, shape(sq_dr), divisibility=4) + mScaleP = fake_tensor(Float32, shape(sq, topk), divisibility=4) + + if not has_gqa: + mdKaccum, mdVaccum = None, None + else: + if not varlen_k: + mdKaccum = fake_tensor(Float32, (b, h_kv, seqlen_k_rounded), divisibility=4) + mdVaccum = fake_tensor(Float32, (b, h_kv, seqlen_k_dv_rounded), divisibility=4) + else: + mdKaccum = fake_tensor(Float32, (h_kv, total_k_rounded), divisibility=4) + mdVaccum = fake_tensor(Float32, (h_kv, total_k_dv_rounded), divisibility=4) + return mQ, mK, mV, mO, mdO, mdQ, mdK, mdV, mLSE, mLSElog2, mPdPsum, dQaccum, mdKaccum, mdVaccum, mScaleP + + +def _compile_bwd_preprocess( + dtype, + head_dim, + head_dim_v, + m_block_size, + has_cuseqlens_q, + has_seqused_q, + has_dlse, + has_dq_accum, + has_scaleP, + use_padded_offsets, + nheads_major, + pack_gqa, + qhead_per_kvhead, + nheads_kv, +): + """Compile bwd preprocess kernel using cute fake tensors (no real GPU tensors needed).""" + mQ, mK, mV, mO, mdO, mdQ, mdK, mdV, mLSE, mLSElog2, mPdPsum, mdQaccum, mdKaccum, mdVaccum, mScaleP = make_fake_bwd_tensors( + dtype, has_gqa=True, varlen_q=has_cuseqlens_q, varlen_k=False, nheads_major=nheads_major, + ) + batch = mQ.shape[0] if not has_cuseqlens_q else cute.sym_int() + batchp1 = cute.sym_int() + mCuSeqlensQ = fake_tensor(Int32, (batchp1,), divisibility=1) if has_cuseqlens_q else None + mSequsedQ = fake_tensor(Int32, (batch,), divisibility=1) if has_seqused_q else None + mdLSE = fake_tensor(Float32, mLSE.shape, divisibility=1) if has_dlse else None + mLSElog2 = None if has_scaleP else mLSElog2 + mdQaccum = mdQaccum if has_dq_accum else None + mRowMax = fake_tensor(Float32, mScaleP.shape, divisibility=1) if has_scaleP else None + mScaleP = fake_tensor(Float32, mScaleP.shape, divisibility=1) if has_scaleP else None + softmax_scale = Float32(1.0) + fa_bwd_pre = FlashAttentionBackwardPreprocess( + dtype, head_dim, head_dim_v, m_block_size, + use_padded_offsets=use_padded_offsets, + nheads_major=nheads_major, + pack_gqa=pack_gqa, + qhead_per_kvhead=qhead_per_kvhead, + nheads_kv=nheads_kv, + ) + return cute.compile( + fa_bwd_pre, mO, mdO, mPdPsum, mLSE, mLSElog2, mdQaccum, mCuSeqlensQ, mSequsedQ, mdLSE, + mRowMax, mScaleP, softmax_scale, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +def _bwd_preprocess( + out, dout, dpsum, lse, lse_log2, dq_accum, + cu_seqlens_q, seqused_q, dlse, + dtype, head_dim, head_dim_v, m_block_size, + row_max=None, + scale_p=None, + use_padded_offsets=True, + nheads_major=False, + pack_gqa=False, + qhead_per_kvhead=1, # only used with pack_gqa + nheads_kv=1, # only used with pack_gqa + softmax_scale=1.0, # only used with scale_p +): + """Backward preprocess: compute (o * dout).sum(dim=-1) - dLSE, lse * log2_e, and zero out dq_accum.""" + if row_max is not None: + assert scale_p is not None + compile_key = ( + dtype, head_dim, head_dim_v, m_block_size, + cu_seqlens_q is not None, + seqused_q is not None, + dlse is not None, + dq_accum is not None, + row_max is not None, + use_padded_offsets, + nheads_major, + pack_gqa, + qhead_per_kvhead, + nheads_kv, + ) + if compile_key not in _bwd_preprocess.compile_cache: + _bwd_preprocess.compile_cache[compile_key] = _compile_bwd_preprocess(*compile_key) + if not is_fake_mode(): + _bwd_preprocess.compile_cache[compile_key]( + out, dout, dpsum, lse, lse_log2, dq_accum, cu_seqlens_q, seqused_q, dlse, + row_max, scale_p, + softmax_scale, + ) + + +_bwd_preprocess.compile_cache = get_jit_cache("bwd_pre") + + +def _compile_bwd_postprocess( + dtype, hdim, block_size, num_threads, atom_layout, swap_ab, + has_cuseqlens_q, has_seqused_q, + use_2cta_instrs, cluster_size, arch, +): + """Compile bwd postprocess kernel using cute fake tensors.""" + mQ, mK, mV, mO, mdO, mdQ, mdK, mdV, mLSE, mLSElog2, mPdPsum, mdQaccum, mdKaccum, mdVaccum, mScaleP = make_fake_bwd_tensors( + dtype, has_gqa=True, varlen_q=has_cuseqlens_q, varlen_k=False + ) + batch = mQ.shape[0] if not has_cuseqlens_q else cute.sym_int() + batchp1 = cute.sym_int() + mCuSeqlensQ = fake_tensor(Int32, (batchp1,), divisibility=1) if has_cuseqlens_q else None + mSeqUsedQ = fake_tensor(Int32, (batch,), divisibility=1) if has_seqused_q else None + fa_bwd_post = FlashAttentionBackwardPostprocess( + dtype, hdim, arch, block_size, num_threads, atom_layout, swap_ab, + use_2cta_instrs=use_2cta_instrs, + cluster_size=cluster_size, + ) + return cute.compile( + fa_bwd_post, mdQaccum, mdQ, Float32(0.0), mCuSeqlensQ, mSeqUsedQ, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +def _bwd_postprocess_convert( + accum, output, scale, + cu_seqlens, seqused, + arch, dtype, hdim, block_size, num_threads, + atom_layout, swap_ab, + use_2cta_instrs=False, cluster_size=1, +): + """Backward postprocess: convert float32 accumulator to bf16/fp16 output.""" + compile_key = ( + dtype, hdim, block_size, num_threads, atom_layout, swap_ab, + cu_seqlens is not None, seqused is not None, + use_2cta_instrs, cluster_size, arch, + ) + if compile_key not in _bwd_postprocess_convert.compile_cache: + _bwd_postprocess_convert.compile_cache[compile_key] = _compile_bwd_postprocess(*compile_key) + if not is_fake_mode(): + _bwd_postprocess_convert.compile_cache[compile_key]( + accum, output, scale, cu_seqlens, seqused, + ) + + +_bwd_postprocess_convert.compile_cache = get_jit_cache("bwd_post") + + +def _flash_attn_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + dout: torch.Tensor, + lse: torch.Tensor, + softmax_scale: Optional[float] = None, + causal: bool = False, + softcap: float = 0.0, + window_size_left: Optional[int] = None, + window_size_right: Optional[int] = None, + m_block_size: int = 64, + n_block_size: int = 128, + num_threads: int = 256, + pack_gqa: bool = False, + num_stages_Q: int = 2, + num_stages_dO: int = 2, + SdP_swapAB: bool = False, + dKV_swapAB: bool = False, + dQ_swapAB: bool = False, + AtomLayoutMSdP: int = 2, + AtomLayoutNdKV: int = 2, + AtomLayoutMdQ: int = 2, + V_in_regs: bool = False, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + deterministic: bool = False, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + score_mod: Optional[Callable] = None, + score_mod_bwd: Optional[Callable] = None, + mask_mod: Optional[Callable] = None, + aux_tensors: Optional[list[torch.Tensor]] = None, + aux_scalars: Optional[tuple] = None, + block_sparse_tensors: Optional[BlockSparseTensorsTorch] = None, + dlse: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + aux_scalars = tuple(aux_scalars) if aux_scalars else None + arch = _get_device_arch() + assert arch // 10 in [9, 10, 11, 12], "Unsupported compute capability. Supported: 9.x, 10.x, 11.x, 12.x" + sparse_q = None + if block_sparse_tensors is not None and arch // 10 == 9: + sparse_q = block_sparse_tensors.block_size[0] if block_sparse_tensors.block_size is not None else 128 + + num_head, head_dim = q.shape[-2:] + head_dim_v = v.shape[-1] + + window_size = [window_size_left, window_size_right] + causal, local, window_size_left, window_size_right = _resolve_causal_local_window( + causal, window_size_left, window_size_right + ) + + if arch // 10 == 12: + # SM120: uses SM80 MMA with 99 KB SMEM, 128 threads (4 warps). + m_block_size = 64 + n_block_size = 64 + if head_dim <= 64: + num_stages_Q = 2 + num_stages_dO = 2 + else: + num_stages_Q = 1 + num_stages_dO = 1 + SdP_swapAB = False + dKV_swapAB = False + dQ_swapAB = False + AtomLayoutMSdP = 4 + AtomLayoutNdKV = 4 + AtomLayoutMdQ = 4 + V_in_regs = False + cluster_size = 1 + use_2cta_instrs = False + num_threads = 128 + assert not (block_sparse_tensors is not None), "Block sparsity backward not supported on SM 12.0" + assert score_mod is None and score_mod_bwd is None, "score_mod backward not supported on SM 12.0" + assert mask_mod is None, "mask_mod backward not supported on SM 12.0" + assert deterministic is False, "deterministic backward not supported on SM 12.0" + elif arch // 10 == 9: + cfg = _tile_size_bwd_sm90( + head_dim, + head_dim_v, + causal, + local, + sparse_block_size_q=sparse_q, + ) + m_block_size = cfg.m_block_size + n_block_size = cfg.n_block_size + num_stages_Q = cfg.num_stages_Q + num_stages_dO = cfg.num_stages_dO + num_stages_PdS = cfg.num_stages_PdS + SdP_swapAB = cfg.SdP_swapAB + dKV_swapAB = cfg.dKV_swapAB + dQ_swapAB = cfg.dQ_swapAB + AtomLayoutMSdP = cfg.AtomLayoutMSdP + AtomLayoutNdKV = cfg.AtomLayoutNdKV + AtomLayoutMdQ = cfg.AtomLayoutMdQ + num_threads = (cfg.num_wg + 1) * 128 + dQ_single_wg = cfg.dQ_single_wg + cluster_size = 1 + use_2cta_instrs = False + is_varlen = ( + cu_seqlens_q is not None + or cu_seqlens_k is not None + or seqused_q is not None + or seqused_k is not None + ) + else: + m_block_size = 128 + n_block_size = 128 + dQ_swapAB = False + dKV_swapAB = False + AtomLayoutMdQ = 1 + AtomLayoutNdKV = 1 + requested_disable_2cta = utils._get_disable_2cta_default() + disable_2cta = ( + requested_disable_2cta + or block_sparse_tensors is not None + ) + cluster_size = 2 if head_dim >= 128 and not disable_2cta else 1 + use_2cta_instrs = cluster_size==2 + + use_dedicated_hd256_kernel = arch // 10 in [10, 11] and head_dim == 256 and head_dim_v == 256 + use_2cta_instrs = use_2cta_instrs or use_dedicated_hd256_kernel + + q, k, v, out, dout, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = [ + maybe_contiguous(t) + for t in (q, k, v, out, dout, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ] + if cu_seqlens_q is None: + batch_size, seqlen_q = q.shape[:2] + total_q = batch_size * seqlen_q + else: + batch_size = cu_seqlens_q.shape[0] - 1 + total_q = q.shape[0] + seqlen_q = max_seqlen_q if max_seqlen_q is not None else total_q + + if cu_seqlens_k is None: + batch_size, seqlen_k = k.shape[:2] + total_k = batch_size * seqlen_k + else: + batch_size = cu_seqlens_k.shape[0] - 1 + total_k = k.shape[0] + seqlen_k = max_seqlen_k if max_seqlen_k is not None else total_k + + num_head_kv = k.shape[-2] + + use_block_sparsity = block_sparse_tensors is not None + q_subtile_factor = sparse_q // m_block_size if sparse_q is not None else 2 + seqlen_q_rounded = (seqlen_q + m_block_size - 1) // m_block_size * m_block_size + seqlen_k_rounded = (seqlen_k + n_block_size - 1) // n_block_size * n_block_size + num_n_blocks = seqlen_k_rounded // n_block_size + if cluster_size == 2 and num_n_blocks % cluster_size != 0: + seqlen_k_rounded = seqlen_k_rounded + n_block_size + + if cu_seqlens_k is None: + assert k.shape == (batch_size, seqlen_k, num_head_kv, head_dim) + assert v.shape == (batch_size, seqlen_k, num_head_kv, head_dim_v) + else: + assert k.shape == (total_k, num_head_kv, head_dim) + assert v.shape == (total_k, num_head_kv, head_dim_v) + assert cu_seqlens_k.shape == (batch_size + 1,), ( + "cu_seqlens_k must have shape (batch_size + 1,)" + ) + + if cu_seqlens_q is not None: + assert cu_seqlens_q.shape == (batch_size + 1,), ( + "cu_seqlens_q must have shape (batch_size + 1,)" + ) + + assert out.shape == (total_q, num_head, head_dim_v) + assert dout.shape == (total_q, num_head, head_dim_v) + assert lse.shape == (num_head, total_q), "lse must have shape (num_head, total_q)" + else: + assert out.shape == (batch_size, seqlen_q, num_head, head_dim_v) + assert dout.shape == (batch_size, seqlen_q, num_head, head_dim_v) + assert lse.shape == (batch_size, num_head, seqlen_q), ( + "lse must have shape (batch_size, num_head, seqlen_q)" + ) + + assert q.dtype in [torch.float16, torch.bfloat16], "inputs must be float16 or bfloat16" + assert q.dtype == k.dtype == v.dtype == out.dtype == dout.dtype, ( + "inputs must have the same dtype" + ) + for t in [cu_seqlens_q, cu_seqlens_k]: + if t is not None: + assert t.dtype == torch.int32, "cu_seqlens_q, cu_seqlens_k must be int32" + assert lse.dtype == torch.float32, "lse must be float32" + if dlse is not None: + dlse = maybe_contiguous(dlse) + if not is_fake_mode(): + assert all( + t is None or t.is_cuda for t in (q, k, v, out, dout, lse, cu_seqlens_q, cu_seqlens_k) + ), "inputs must be on CUDA device" + assert num_head % num_head_kv == 0, "num_head must be divisible by num_head_kv" + alignment = 16 // q.element_size() + if arch // 10 != 12: + _validate_head_dims(head_dim, head_dim_v, arch // 10, alignment) + if softmax_scale is None: + softmax_scale = 1.0 / math.sqrt(head_dim) + qhead_per_kvhead = num_head // num_head_kv + if pack_gqa is None: + pack_gqa = qhead_per_kvhead > 1 + # pack_gqa backward not yet supported in bwd + pack_gqa = False + + if softcap != 0.0: + assert score_mod is None and score_mod_bwd is None, ( + "softcap and score_mod/score_mod_bwd cannot be used together" + ) + score_mod = utils.create_softcap_scoremod(softcap) + score_mod_bwd = utils.create_softcap_scoremod_bwd(softcap) + if score_mod is not None: + assert score_mod_bwd is not None, "score_mod_bwd is required when score_mod is provided" + assert cu_seqlens_q is None and cu_seqlens_k is None, ( + "varlen + score_mod not supported in bwd yet" + ) + if arch // 10 == 8: + raise NotImplementedError("Custom user-provided score_mod is not supported on SM8x architectures.") + + device = q.device + out_torch_dtype = q.dtype + + if dq is None: + dq = torch.empty_like(q) + else: + _validate_tensor(dq, "dq", q.shape, out_torch_dtype, device) + + if dk is None: + dk = torch.empty_like(k) + else: + _validate_tensor(dk, "dk", k.shape, out_torch_dtype, device) + + if dv is None: + dv = torch.empty_like(v) + else: + _validate_tensor(dv, "dv", v.shape, out_torch_dtype, device) + + head_dim_rounded = (head_dim + 32 - 1) // 32 * 32 + + if cu_seqlens_q is None: + dq_accum = ( + None + if use_dedicated_hd256_kernel + else torch.empty( + batch_size, + num_head, + seqlen_q_rounded * head_dim_rounded, + dtype=torch.float32, + device=device, + ) + ) + dpsum = torch.empty( + batch_size, num_head, seqlen_q_rounded, dtype=torch.float32, device=device + ) + lse_log2 = torch.empty( + batch_size, num_head, seqlen_q_rounded, dtype=torch.float32, device=device + ) + else: + total_q_rounded_padded = ( + (total_q + cu_seqlens_q.shape[0] * m_block_size - 1) // m_block_size * m_block_size + ) + dq_accum = ( + None + if use_dedicated_hd256_kernel + else torch.empty( + num_head, total_q_rounded_padded * head_dim_rounded, dtype=torch.float32, device=device + ) + ) + dpsum = torch.empty(num_head, total_q_rounded_padded, dtype=torch.float32, device=device) + lse_log2 = torch.empty(num_head, total_q_rounded_padded, dtype=torch.float32, device=device) + + # GQA (qhead_per_kvhead > 1) needs dK/dV accum+postprocess since multiple Q heads + # accumulate into the same dK/dV. SM90 varlen_k with qhead_per_kvhead==1 now uses + # ragged TMA tensors for direct store, so no longer needs accum+postprocess. + # hd=256 2CTA backward has its own internal postprocess for dK/dV. + dKV_postprocess = qhead_per_kvhead > 1 and not use_dedicated_hd256_kernel + if dKV_postprocess: + head_dim_v_rounded = (head_dim_v + 32 - 1) // 32 * 32 + if cu_seqlens_k is None: + dk_accum = torch.zeros( + batch_size, + num_head_kv, + seqlen_k_rounded * head_dim_rounded, + dtype=torch.float32, + device=device, + ) + dv_accum = torch.zeros( + batch_size, + num_head_kv, + seqlen_k_rounded * head_dim_v_rounded, + dtype=torch.float32, + device=device, + ) + else: + cluster_tile_n = cluster_size * n_block_size + total_k_rounded_padded = ( + (total_k + cu_seqlens_k.shape[0] * cluster_tile_n - 1) // cluster_tile_n * cluster_tile_n + ) + dk_accum = torch.zeros( + num_head_kv, + total_k_rounded_padded * head_dim_rounded, + dtype=torch.float32, + device=device, + ) + dv_accum = torch.zeros( + num_head_kv, + total_k_rounded_padded * head_dim_v_rounded, + dtype=torch.float32, + device=device, + ) + + dtype = torch2cute_dtype_map[q.dtype] + current_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + + if deterministic: + dQ_semaphore = torch.zeros(batch_size, num_head, seqlen_q_rounded // m_block_size, cluster_size, dtype=torch.int32, device=device) + else: + dQ_semaphore = None + + if deterministic and qhead_per_kvhead > 1: + dK_semaphore = torch.zeros(batch_size, num_head_kv, seqlen_k_rounded // n_block_size, 2, dtype=torch.int32, device=device) + dV_semaphore = torch.zeros(batch_size, num_head_kv, seqlen_k_rounded // n_block_size, 2, dtype=torch.int32, device=device) + else: + dK_semaphore = None + dV_semaphore = None + + # Preprocess kernel: compute (o * dout).sum(dim=-1) - dLSE, lse * log2_e, and zero out dq_accum. + # For hd=256 dedicated path, dq_accum is None so preprocess only fills dpsum/lse_log2. + _bwd_preprocess( + out, dout, dpsum, lse, lse_log2, dq_accum, + cu_seqlens_q, seqused_q, dlse, + dtype, head_dim, head_dim_v, m_block_size, + ) + # num_threads: SM90 derives from BwdConfig.num_wg, SM120 is set to 128 above, + # SM100/SM110 uses default from function signature (384). + if arch // 10 not in [9, 12]: + num_threads = 384 + + # Backward kernel: compute dk, dv, dq_accum. + score_mod_hash = utils.hash_callable(score_mod) if score_mod else False + score_mod_bwd_hash = utils.hash_callable(score_mod_bwd) if score_mod_bwd else False + mask_mod_hash = utils.hash_callable(mask_mod) if mask_mod else False + num_aux_tensors = len(aux_tensors) if aux_tensors else 0 + aux_scalar_metadata = tuple(type(s) for s in aux_scalars) if aux_scalars is not None else None + cute_aux_tensors = None + if aux_tensors is not None: + cute_aux_tensors = [to_cute_tensor(buf, assumed_align=None, fully_dynamic=True) for buf in aux_tensors] + + block_sparse_broadcast_pattern = None + normalized_block_sparse_tensors = None + if block_sparse_tensors is not None: + ( + normalized_block_sparse_tensors, + block_sparse_broadcast_pattern, + ) = normalize_block_sparse_config_bwd( + block_sparse_tensors, + batch_size=batch_size, + num_head=num_head, + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + block_size=(m_block_size, n_block_size), + q_subtile_factor=q_subtile_factor, + ) + if deterministic: + if normalized_block_sparse_tensors.dq_write_order is None: + raise ValueError( + "deterministic block-sparse backward requires dq_write_order in block_sparse_tensors" + ) + if ( + normalized_block_sparse_tensors.full_block_cnt is not None + and normalized_block_sparse_tensors.dq_write_order_full is None + ): + raise ValueError( + "deterministic block-sparse backward requires dq_write_order_full when full blocks are present" + ) + if normalized_block_sparse_tensors.spt is None: + raise ValueError( + "deterministic block-sparse backward requires block_sparse_tensors.spt " + "to match dq_write_order direction" + ) + if ( + normalized_block_sparse_tensors is not None + and normalized_block_sparse_tensors.spt is not None + ): + spt = normalized_block_sparse_tensors.spt and deterministic + else: + spt = (causal or local) and deterministic + + if arch // 10 in [8, 9, 12]: + compile_key = ( + arch, + dtype, + head_dim, + head_dim_v, + qhead_per_kvhead, + causal, + window_size_left is not None, + window_size_right is not None, + m_block_size, + n_block_size, + num_threads, + pack_gqa, + num_stages_Q, + num_stages_dO, + SdP_swapAB, + dKV_swapAB, + dQ_swapAB, + AtomLayoutMSdP, + AtomLayoutNdKV, + AtomLayoutMdQ, + V_in_regs, + dQ_single_wg, + deterministic, + cu_seqlens_q is None, + cu_seqlens_k is None, + seqused_q is None, + seqused_k is None, + score_mod_hash, + score_mod_bwd_hash, + mask_mod_hash, + num_aux_tensors, + aux_scalar_metadata, + use_block_sparsity, + block_sparse_broadcast_pattern, + get_broadcast_dims(q), + get_broadcast_dims(k), + get_broadcast_dims(v), + get_broadcast_dims(dout), + # Prevent TVM stride poisoning when only one block is present. + (seqlen_q_rounded // m_block_size == 1), + (seqlen_k_rounded // n_block_size == 1), + ) + else: + compile_key = ( + arch, + dtype, + head_dim, + head_dim_v, + qhead_per_kvhead, + causal, + window_size_left is not None, + window_size_right is not None, + m_block_size, + n_block_size, + num_threads, + pack_gqa, + cluster_size, + use_2cta_instrs, + deterministic, + spt, + score_mod_hash, + score_mod_bwd_hash, + mask_mod_hash, + num_aux_tensors, + aux_scalar_metadata, + use_block_sparsity, + block_sparse_broadcast_pattern, + cu_seqlens_q is None, + cu_seqlens_k is None, + seqused_q is None, + seqused_k is None, + get_broadcast_dims(q), + get_broadcast_dims(k), + get_broadcast_dims(v), + get_broadcast_dims(dout), + # Prevent TVM stride poisoning when only one block is present. + (seqlen_q_rounded // m_block_size == 1), + (seqlen_k_rounded // n_block_size == 1), + ) + + if compile_key not in _flash_attn_bwd.compile_cache: + q_tensor, k_tensor, v_tensor, do_tensor, dq_tensor, dk_tensor, dv_tensor = [ + to_cute_tensor(t) for t in (q, k, v, dout, dq, dk, dv) + ] + lse_log2_tensor, dpsum_tensor = [to_cute_tensor(t) for t in (lse_log2, dpsum)] + dq_accum_tensor = to_cute_tensor(dq_accum) if dq_accum is not None else None + if dKV_postprocess: + dk_accum_tensor, dv_accum_tensor = [ + to_cute_tensor(t) for t in (dk_accum, dv_accum) + ] + cu_seqlens_q_tensor, cu_seqlens_k_tensor, seqused_q_tensor, seqused_k_tensor = [ + to_cute_tensor(t, assumed_align=4) if t is not None else None + for t in (cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + ] + dQ_semaphore_tensor, dK_semaphore_tensor, dV_semaphore_tensor = [ + utils.convert_from_dlpack_leading_static(t.detach(), leading_dim=3, alignment=4, stride_order=t.dim_order()) + if t is not None else None + for t in (dQ_semaphore, dK_semaphore, dV_semaphore) + ] + if arch // 10 in [8, 12]: + flash_bwd_obj_cls = FlashAttentionBackwardSm120 if arch // 10 == 12 else FlashAttentionBackwardSm80 + fa_bwd_obj = flash_bwd_obj_cls( + dtype, + head_dim, + head_dim_v, + qhead_per_kvhead, + m_block_size, + n_block_size, + num_stages_Q, + num_stages_dO, + num_threads, + pack_gqa, + causal, + SdP_swapAB, + dKV_swapAB, + dQ_swapAB, + AtomLayoutMSdP, + AtomLayoutNdKV, + AtomLayoutMdQ, + V_in_regs=V_in_regs, + score_mod=score_mod, + score_mod_bwd=score_mod_bwd, + ) + elif arch // 10 == 9: + fa_bwd_obj = FlashAttentionBackwardSm90( + dtype, + head_dim, + head_dim_v, + qhead_per_kvhead, + causal, + is_local=local, + deterministic=deterministic, + tile_m=m_block_size, + tile_n=n_block_size, + Q_stage=num_stages_Q, + dO_stage=num_stages_dO, + PdS_stage=num_stages_PdS, + SdP_swapAB=SdP_swapAB, + dKV_swapAB=dKV_swapAB, + dQ_swapAB=dQ_swapAB, + AtomLayoutMSdP=AtomLayoutMSdP, + AtomLayoutNdKV=AtomLayoutNdKV, + AtomLayoutMdQ=AtomLayoutMdQ, + num_threads=num_threads, + V_in_regs=V_in_regs, + score_mod=score_mod, + score_mod_bwd=score_mod_bwd, + mask_mod=mask_mod, + has_aux_tensors=aux_tensors is not None, + q_subtile_factor=q_subtile_factor, + dQ_single_wg=dQ_single_wg, + ) + else: + if use_dedicated_hd256_kernel: + assert softcap == 0.0, "SM100 backward with head_dim=256 does not support softcap" + assert block_sparse_tensors is None, \ + "SM100 backward with head_dim=256 does not support block sparsity" + assert dlse is None, \ + "SM100 backward with head_dim=256 does not support dlse" + assert seqused_q is None and seqused_k is None, \ + "SM100 backward with head_dim=256 does not support seqused_q/seqused_k" + # Same as forward: hd256 kernel uses hardcoded contiguous strides. + q = q.contiguous() if not q.is_contiguous() else q + k = k.contiguous() if not k.is_contiguous() else k + v = v.contiguous() if not v.is_contiguous() else v + out = out.contiguous() if not out.is_contiguous() else out + dout = dout.contiguous() if not dout.is_contiguous() else dout + + dq_tile_mn = (128, 128) + dkdv_tile_mn = (128, 64) + fa_bwd_obj = BlackwellFusedMultiHeadAttentionBackward( + head_dim, + head_dim_v, + is_causal=causal, + is_local=local, + qhead_per_kvhead=qhead_per_kvhead, + is_persistent=False, + deterministic=deterministic, + cluster_size=cluster_size, + use_2cta_instrs=use_2cta_instrs, + score_mod=score_mod, + score_mod_bwd=score_mod_bwd, + mask_mod=mask_mod, + has_aux_tensors=aux_tensors is not None, + q_subtile_factor=q_subtile_factor, + tile_m_dq=dq_tile_mn[0], + tile_n_dq=dq_tile_mn[1], + tile_m_dkdv=dkdv_tile_mn[0], + tile_n_dkdv=dkdv_tile_mn[1], + ) + else: + fa_bwd_obj = FlashAttentionBackwardSm100( + head_dim, + head_dim_v, + is_causal=causal, + is_local=local, + qhead_per_kvhead=qhead_per_kvhead, + tile_m=m_block_size, + tile_n=n_block_size, + cluster_size=cluster_size, + use_2cta_instrs=use_2cta_instrs, + deterministic=deterministic, + spt=spt, + score_mod=score_mod, + score_mod_bwd=score_mod_bwd, + mask_mod=mask_mod, + has_aux_tensors=aux_tensors is not None, + q_subtile_factor=q_subtile_factor, + ) + + # Block sparse tensors for backward use Q-direction indexing (transposed from forward). + sparse_tensors_compile = None + if normalized_block_sparse_tensors is not None: + sparse_tensors_compile = to_cute_block_sparse_tensors(normalized_block_sparse_tensors) + dq_accum_tensor = dq_tensor if use_dedicated_hd256_kernel else dq_accum_tensor + + # TODO: check @can_implement + _flash_attn_bwd.compile_cache[compile_key] = cute.compile( + fa_bwd_obj, + q_tensor, + k_tensor, + v_tensor, + do_tensor, + lse_log2_tensor, + dpsum_tensor, + dq_accum_tensor, + dk_tensor if not dKV_postprocess else dk_accum_tensor, + dv_tensor if not dKV_postprocess else dv_accum_tensor, + softmax_scale, + cu_seqlens_q_tensor, + cu_seqlens_k_tensor, + seqused_q_tensor, + seqused_k_tensor, + window_size_left, + window_size_right, + dQ_semaphore_tensor, + dK_semaphore_tensor, + dV_semaphore_tensor, + AuxData(cute_aux_tensors, aux_scalars), + sparse_tensors_compile, + current_stream, + options="--enable-tvm-ffi", + ) + if not is_fake_mode(): + dq_accum = dq if use_dedicated_hd256_kernel else dq_accum + _flash_attn_bwd.compile_cache[compile_key]( + q.detach(), + k.detach(), + v.detach(), + dout, + lse_log2, + dpsum, + dq_accum, + dk if not dKV_postprocess else dk_accum, + dv if not dKV_postprocess else dv_accum, + softmax_scale, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + window_size_left, + window_size_right, + dQ_semaphore, + dK_semaphore, + dV_semaphore, + AuxData(aux_tensors, aux_scalars), + ( + normalized_block_sparse_tensors.mask_block_cnt, + normalized_block_sparse_tensors.mask_block_idx, + normalized_block_sparse_tensors.full_block_cnt, + normalized_block_sparse_tensors.full_block_idx, + normalized_block_sparse_tensors.cu_total_m_blocks, + normalized_block_sparse_tensors.cu_block_idx_offsets, + normalized_block_sparse_tensors.dq_write_order, + normalized_block_sparse_tensors.dq_write_order_full, + ) + if normalized_block_sparse_tensors is not None + else None, + ) + # Postprocess: convert dq_accum from float32 to dq in bf16/fp16 + # hd=256 2CTA backward has its own internal postprocess, skip here. + if not use_dedicated_hd256_kernel: + if arch // 10 == 9: + # dQ postprocess: match main kernel's MMA WG count, unless dQ_single_wg + num_threads_post_dQ = 128 if dQ_single_wg else cfg.num_wg * 128 + num_threads_post_dKV = cfg.num_wg * 128 + else: + num_threads_post_dQ = 128 + num_threads_post_dKV = 128 + + _bwd_postprocess_convert( + dq_accum, dq, softmax_scale, + cu_seqlens_q, seqused_q, + arch, dtype, head_dim, m_block_size, num_threads_post_dQ, + AtomLayoutMdQ, dQ_swapAB, + use_2cta_instrs=use_2cta_instrs, cluster_size=1, + ) + + if dKV_postprocess: + # Postprocess: convert dk_accum from float32 to dk in bf16/fp16 + _bwd_postprocess_convert( + dk_accum, dk, softmax_scale, + cu_seqlens_k, seqused_k, + arch, dtype, head_dim, n_block_size, num_threads_post_dKV, + AtomLayoutNdKV, dKV_swapAB, + cluster_size=cluster_size, + ) + # Postprocess: convert dv_accum from float32 to dv in bf16/fp16 + _bwd_postprocess_convert( + dv_accum, dv, 1.0, + cu_seqlens_k, seqused_k, + arch, dtype, head_dim_v, n_block_size, num_threads_post_dKV, + AtomLayoutNdKV, dKV_swapAB, + cluster_size=cluster_size, + ) + + return dq, dk, dv + + +_flash_attn_bwd.compile_cache = get_jit_cache("bwd") + + +def _flash_attn_bwd_sparse_mla( + q: Optional[torch.Tensor], + k: Optional[torch.Tensor], + v: torch.Tensor, + qv: torch.Tensor, + out: torch.Tensor, + dout: torch.Tensor, + lse: torch.Tensor, + p: torch.Tensor, + row_max: torch.Tensor, + gather_kv_indices: torch.Tensor, + learnable_sink: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + m_block_size: int = 128, + n_block_size: int = 64, + num_threads: int = 256, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + min_seqlen_k: Optional[int] = None, + deterministic: bool = False, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + dqv: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + arch = _get_device_arch() + assert arch // 10 in [10, 11], "Unsupported compute capability. Supported: 10.x, 11.x" + assert gather_kv_indices is not None, "require gather kv indices for backward" + + q_shape = q.shape if q is not None else qv.shape + nheads, head_dim = q_shape[-2:] + nheads_kv, head_dim_v = v.shape[-2:] + qhead_per_kvhead = nheads // nheads_kv + gather_kv_length = gather_kv_indices.shape[-1] + assert nheads_kv == 1 and qhead_per_kvhead == 128, f"sparse MLA bwd: only MQA 128 supported for now" + assert gather_kv_length % 128 == 0, f"sparse MLA bwd: {gather_kv_length=} must be divisible by 128" + assert deterministic is False, "sparse MLA bwd: deterministic mode not yet supported" + assert learnable_sink is None, "sparse MLA bwd: learnable sink not yet supported" + assert seqused_q is None and seqused_k is None, "sparse MLA bwd: seqused_q,k not yet supported" + + if softmax_scale is None: + softmax_scale = ( + 1.0 / math.sqrt(head_dim) if qv is None or q is None + else 1.0 / math.sqrt(head_dim + head_dim_v) + ) + + q, k, v, qv, out, dout, lse, p, row_max = [ + maybe_contiguous(t) + for t in (q, k, v, qv, out, dout, lse, p, row_max) + ] + gather_kv_indices, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, learnable_sink = [ + maybe_contiguous(t) + for t in (gather_kv_indices, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, learnable_sink) + ] + device = v.device + + varlen_q = cu_seqlens_q is not None or seqused_q is not None + if cu_seqlens_q is None: + batch_size, seqlen_q = q_shape[:2] + total_q = batch_size * seqlen_q + p_shape = (batch_size, seqlen_q, nheads, gather_kv_length) + else: + batch_size = cu_seqlens_q.shape[0] - 1 + total_q = q_shape[0] + seqlen_q = max_seqlen_q if max_seqlen_q is not None else total_q + p_shape = (total_q, nheads, gather_kv_length) + + varlen_k = cu_seqlens_k is not None or seqused_k is not None + if cu_seqlens_k is None: + batch_size, seqlen_k = v.shape[:2] + total_k = batch_size * seqlen_k + else: + batch_size = cu_seqlens_k.shape[0] - 1 + total_k = v.shape[0] + seqlen_k = max_seqlen_k if max_seqlen_k is not None else total_k + if not varlen_k: + min_seqlen_k = seqlen_k + + assert varlen_q == varlen_k, "sparse MLA bwd: either q and k are both varlen or not" + + # always use kv bitmask by default (handles -1 sentinel) + disable_sparse_kv_bitmask = False + # if min_seqlen_k is None or causal: + # disable_sparse_kv_bitmask = False + # else: + # disable_sparse_kv_bitmask = min_seqlen_k >= gather_kv_length + + prealloc_dq = dq is not None + prealloc_dk = dk is not None + prealloc_dqv = dqv is not None + prealloc_dv = dv is not None + dq = dk = None + if not prealloc_dq and q is not None: + dq = torch.empty_like(q) + if not prealloc_dk and k is not None: + dk = torch.zeros_like(k, dtype=torch.float32) + if not prealloc_dv: + dv = torch.zeros_like(v, dtype=torch.float32) + if not prealloc_dqv: + dqv = torch.empty_like(qv) + ds = torch.empty_like(p) + + device = v.device + dtype = v.dtype + if q is not None: + _validate_tensor(dq, "dq", q.shape, dtype, device) + if k is not None: + _validate_tensor(dk, "dk", k.shape, torch.float32, device) + _validate_tensor(dv, "dv", v.shape, torch.float32, device) + _validate_tensor(dqv, "dqv", qv.shape, dtype, device) + _validate_tensor(p, "p", p_shape, dtype, device) + + if cu_seqlens_q is None: + dpsum = torch.empty(batch_size, seqlen_q, nheads, dtype=torch.float32, device=device) + else: + dpsum = torch.empty(total_q, nheads, dtype=torch.float32, device=device) + scale_p = torch.empty_like(row_max) + + dtype = torch2cute_dtype_map[dout.dtype] + current_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + + # Preprocess kernel: compute (o * dout).sum(dim=-1), scale_p. + _bwd_preprocess( + out, dout, dpsum, lse, None, None, + cu_seqlens_q, seqused_q, None, + dtype, head_dim, head_dim_v, m_block_size, + row_max=row_max, + scale_p=scale_p, + use_padded_offsets=False, + nheads_major=True, + pack_gqa=True, + qhead_per_kvhead=qhead_per_kvhead, + nheads_kv=nheads_kv, + softmax_scale=softmax_scale, + ) + + compile_key = ( + dtype, + head_dim, + head_dim_v, + qhead_per_kvhead, + causal, + cu_seqlens_q is None, + cu_seqlens_k is None, + seqused_q is None, + seqused_k is None, + q is not None, + gather_kv_length, + learnable_sink is not None, + disable_sparse_kv_bitmask, + ) + + if compile_key not in _flash_attn_bwd_sparse_mla.compile_cache: + ( + cu_seqlens_q_tensor, + cu_seqlens_k_tensor, + seqused_q_tensor, + seqused_k_tensor, + learnable_sink_tensor, + ) = [ + to_cute_tensor(t, assumed_align=4, leading_dim=0) + for t in (cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, learnable_sink) + ] + ( + v_tensor, + qv_tensor, + do_tensor, + p_tensor, + scale_p_tensor, + dpsum_tensor, + ds_tensor, + dv_tensor, + gather_kv_indices_tensor, + ) = [ + to_cute_tensor(t) for t in (v, qv, dout, p, scale_p, dpsum, ds, dv, gather_kv_indices) + ] + + fa_bwd_obj = FlashAttentionSparseMLABackwardSm100( + is_causal=causal, + topk_length=gather_kv_length, + qhead_per_kvhead=qhead_per_kvhead, + nheads_kv=nheads_kv, + has_seqused_q=seqused_q is not None, + disable_bitmask=disable_sparse_kv_bitmask, + ) + fa_bwd_kernel = cute.compile( + fa_bwd_obj, + do_tensor, + v_tensor, + qv_tensor, + p_tensor, + dv_tensor, + ds_tensor, + gather_kv_indices_tensor, + softmax_scale, + scale_p_tensor, + dpsum_tensor, + cu_seqlens_q_tensor, + cu_seqlens_k_tensor, + seqused_q_tensor, + seqused_k_tensor, + current_stream, + options="--enable-tvm-ffi", + ) + _flash_attn_bwd_sparse_mla.compile_cache[compile_key] = fa_bwd_kernel + + if not is_fake_mode(): + _flash_attn_bwd_sparse_mla.compile_cache[compile_key]( + dout, + v, + qv, + p, + dv, + ds, + gather_kv_indices, + softmax_scale, + scale_p, + dpsum, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + ) + + v = v.squeeze(-2) + if k is not None: + k = k.squeeze(-2) + + _sparse_mla_dq_dqv( + ds, k, v, dq, dqv, gather_kv_indices, cu_seqlens_q, cu_seqlens_k, + ) + + if k is not None: + dk = dk.squeeze(-2) + _sparse_mla_dk(ds, gather_kv_indices, q, dk, cu_seqlens_q, cu_seqlens_k) + dk = dk.unsqueeze(-2) + + # return dk, dv in float32: all-reduce across sequence-parallel ranks must happen + # before downcasting to avoid rounding error during inter-rank grad accumulation + return dq, dk, dv, dqv + +_flash_attn_bwd_sparse_mla.compile_cache = get_jit_cache("bwd_dsa") + + +def _compile_sparse_mla_dq_dqv( + dtype, nheads, head_dim, head_dim_v, top_k, varlen_q, varlen_k, compute_dq, +): + sym = cute.sym_int + b, b_plus_1, seqlen_q, seqlen_k = sym(), sym(), sym(), sym() + total_q, total_k = sym(), sym() + b_seqlenq = (b, seqlen_q) if not varlen_q else (total_q,) + b_seqlenk = (b, seqlen_k) if not varlen_k else (total_k,) + + div = 128 // dtype.width # 8 for fp16/bf16 + + mdS = fake_tensor(dtype, (*b_seqlenq, nheads, top_k), divisibility=div) + mK = fake_tensor(dtype, (*b_seqlenk, head_dim), divisibility=div) + mV = fake_tensor(dtype, (*b_seqlenk, head_dim_v), divisibility=div) + mdQ = fake_tensor(dtype, (*b_seqlenq, nheads, head_dim), divisibility=div) + mdQv = fake_tensor(dtype, (*b_seqlenq, nheads, head_dim_v), divisibility=div) + mIdxTopK = fake_tensor(Int32, (*b_seqlenq, top_k), divisibility=div) + + mCuSeqlensQ = fake_tensor(Int32, (b_plus_1,), divisibility=1) if varlen_q else None + mCuSeqlensK = fake_tensor(Int32, (b_plus_1,), divisibility=1) if varlen_k else None + + dq_dqv_gemm = dQdQvGemmKernel( + acc_dtype=Float32, + nheads=nheads, + head_dim_k=head_dim, + head_dim_v=head_dim_v, + top_k=top_k, + ) + + return cute.compile( + dq_dqv_gemm, + mdS, + mK if compute_dq else None, + mV, + mdQ if compute_dq else None, + mdQv, + mIdxTopK, + mCuSeqlensQ, + mCuSeqlensK, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +def _sparse_mla_dq_dqv( + ds, k, v, dq, dqv, gather_kv_indices, cu_seqlens_q, cu_seqlens_k, +): + """Compute dQ = dS @ K and dQv = dS @ V""" + *_, nheads, gather_kv_length = ds.shape + + head_dim_v = v.shape[-1] + head_dim = k.shape[-1] if k is not None else 0 + + dtype = ds.dtype + dtype_cute = torch2cute_dtype_map[dtype] + + varlen_q = cu_seqlens_q is not None + varlen_k = cu_seqlens_k is not None + + compile_key = ( + dtype_cute, nheads, head_dim, head_dim_v, gather_kv_length, varlen_q, varlen_k, k is not None, + ) + if compile_key not in _sparse_mla_dq_dqv.compile_cache: + _sparse_mla_dq_dqv.compile_cache[compile_key] = _compile_sparse_mla_dq_dqv( + *compile_key + ) + if not is_fake_mode(): + _sparse_mla_dq_dqv.compile_cache[compile_key]( + ds, k, v, dq, dqv, gather_kv_indices, cu_seqlens_q, cu_seqlens_k + ) + +_sparse_mla_dq_dqv.compile_cache = get_jit_cache("dq_dqv_gemm") + + +def _compile_sparse_mla_dk( + dtype, + dtype_acc, + nheads: int, + head_dim: int, + topk: int, + varlen: bool, +): + kernel = dKGemmKernel( + topk, + nheads, + head_dim, + varlen, + ) + # Check if configuration can be implemented + kernel.check_can_implement() + + div = 128 // dtype.width + + sym = cute.sym_int + batch_fake = sym() + batchp1_fake = sym() + seqlen_q_fake = sym() + seqlen_k_fake = sym() + total_q_fake = (batch_fake, seqlen_q_fake) if not varlen else (sym(),) + total_k_fake = (batch_fake, seqlen_k_fake) if not varlen else (sym(),) + + mdS = fake_tensor(dtype, (*total_q_fake, nheads, topk), divisibility=div) + mI = fake_tensor(Int32, (*total_q_fake, topk), divisibility=div) + mQ = fake_tensor(dtype, (*total_q_fake, nheads, head_dim), divisibility=div) + mdK = fake_tensor(dtype_acc, (*total_k_fake, head_dim), divisibility=div) + mCuSeqlensQ = fake_tensor(Int32, (batchp1_fake,), divisibility=1) if varlen else None + mCuSeqlensK = fake_tensor(Int32, (batchp1_fake,), divisibility=1) if varlen else None + + return cute.compile( + kernel, + mdS, + mI, + mQ, + mdK, + mCuSeqlensQ, + mCuSeqlensK, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +def _sparse_mla_dk( + dS: torch.Tensor, + index_topk: torch.Tensor, + q: torch.Tensor, + dk: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor], + cu_seqlens_k: Optional[torch.Tensor], +): + """Compute dKaccum = scatter(dS'^T @ Q, I). + + Args: + dS: (*total_q, heads, topk), bf16 + index_topk: (*total_q, topk), int32 + Q: (*total_q, heads, dim), bf16 + dK: (*total_q, dim), fp32 + cuSeqlensQ: (batch + 1,), int32, omit for non-varlen + cuSeqlensK: (batch + 1,), int32, omit for non-varlen + + Accumulates in place on top of dK. + + For varlen, total_q and total_k are 1-dimensional, and the seqlen indices per batch are + determined using the cuSeqlensQ and cuSeqlensK tensors. + For non-varlen, total_q and total_k are (batch, seqlen_q) and (batch, seqlen_k). + """ + dtype = dS.dtype + dtype_cute = torch2cute_dtype_map[dtype] + dtype_acc = dk.dtype + dtype_acc_cute = torch2cute_dtype_map[dtype_acc] + + varlen = cu_seqlens_q is not None + nheads, topk = dS.shape[-2], dS.shape[-1] + head_dim = q.shape[-1] if q is not None else 0 + + compile_key = ( + dtype_cute, dtype_acc_cute, nheads, head_dim, topk, varlen, + ) + + if compile_key not in _sparse_mla_dk.compile_cache: + _sparse_mla_dk.compile_cache[compile_key] = _compile_sparse_mla_dk(*compile_key) + + if not is_fake_mode(): + _sparse_mla_dk.compile_cache[compile_key](dS, index_topk, q, dk, cu_seqlens_q, cu_seqlens_k) + +_sparse_mla_dk.compile_cache = get_jit_cache("dk_gemm") + + +class FlashAttnFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + qv: Optional[torch.Tensor] = None, + gather_kv_indices: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size: Tuple[Optional[int], Optional[int]] = (None, None), + learnable_sink: Optional[torch.Tensor] = None, + softcap: float = 0.0, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + deterministic: bool = False, + score_mod: Optional[Callable] = None, + score_mod_bwd: Optional[Callable] = None, + mask_mod: Optional[Callable] = None, + aux_tensors: Optional[list] = None, + aux_scalars: Optional[tuple] = None, + block_sparse_tensors: Optional[BlockSparseTensorsTorch] = None, + block_sparse_tensors_bwd: Optional[BlockSparseTensorsTorch] = None, + return_lse: bool = False, + ): + aux_scalars = tuple(aux_scalars) if aux_scalars else None + shared_kv = k is v + if shared_kv and v.shape[-1] == 512: + # specialize MLA attention formula + # O = softmax(Q @ K.T + Qv @ V.T) @ V + # by setting q, k to None + qv = q if qv is None else qv + q = k = None + out, lse, p, row_max = _flash_attn_fwd( + q, + k, + v, + qv=qv, + softmax_scale=softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + learnable_sink=learnable_sink, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + score_mod=score_mod, + mask_mod=mask_mod, + aux_tensors=aux_tensors, + aux_scalars=aux_scalars, + block_sparse_tensors=block_sparse_tensors, + return_lse=return_lse, + gather_kv_indices=gather_kv_indices, + ) + ctx.save_for_backward(q, k, v, qv, out, lse, p, row_max, gather_kv_indices, *(aux_tensors or ())) + ctx.shared_kv = shared_kv + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.return_lse = return_lse + ctx.score_mod = score_mod + ctx.score_mod_bwd = score_mod_bwd + ctx.mask_mod = mask_mod + ctx.aux_scalars = aux_scalars + ctx.block_sparse_tensors_bwd = block_sparse_tensors_bwd + ctx.set_materialize_grads(False) + return out, lse + + @staticmethod + def backward(ctx, dout, dlse): + q, k, v, qv, out, lse, p, row_max, gather_kv_indices, *aux = ctx.saved_tensors + aux_tensors = aux if aux else None + if not ctx.return_lse: + dlse = None + if dout is None: + dout = torch.zeros_like(out) + if qv is not None: + dq, dk, dv, dqv = _flash_attn_bwd_sparse_mla( + q, + k, + v, + qv, + out, + dout, + lse, + p, + row_max, + gather_kv_indices, + softmax_scale=ctx.softmax_scale, + causal=ctx.causal, + ) + if ctx.shared_kv: + return dqv, dv, None, None, *((None,) * 30) + else: + return dq, dk, dv, dqv, *((None,) * 30) + else: + dq, dk, dv = _flash_attn_bwd( + q, + k, + v, + out, + dout, + lse, + ctx.softmax_scale, + ctx.causal, + ctx.softcap, + window_size_left=ctx.window_size[0], + window_size_right=ctx.window_size[1], + deterministic=ctx.deterministic, + score_mod=ctx.score_mod, + score_mod_bwd=ctx.score_mod_bwd, + mask_mod=ctx.mask_mod, + aux_tensors=aux_tensors, + aux_scalars=ctx.aux_scalars, + block_sparse_tensors=ctx.block_sparse_tensors_bwd, + dlse=dlse, + ) + return dq, dk, dv, *((None,) * 30) # Extra Nones is fine + + +class FlashAttnVarlenFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q: Optional[torch.Tensor], + k: Optional[torch.Tensor], + v: torch.Tensor, + qv: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + min_seqlen_k: Optional[int] = None, + gather_kv_indices: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size: Tuple[Optional[int], Optional[int]] = (None, None), + learnable_sink: Optional[torch.Tensor] = None, + softcap: float = 0.0, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + deterministic: bool = False, + score_mod: Optional[Callable] = None, + score_mod_bwd: Optional[Callable] = None, + mask_mod: Optional[Callable] = None, + block_sparse_tensors: Optional[list] = None, + aux_tensors: Optional[list] = None, + aux_scalars: Optional[tuple] = None, + return_lse: bool = False, + ): + aux_scalars = tuple(aux_scalars) if aux_scalars else None + shared_kv = k is v + if shared_kv and v.shape[-1] == 512: + # specialize MLA attention formula + # O = softmax(Q @ K.T + Qv @ V.T) @ V + # by setting q, k to None + qv = q if qv is None else qv + q = k = None + out, lse, p, row_max = _flash_attn_fwd( + q, + k, + v, + qv=qv, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + seqused_q=seqused_q, + seqused_k=seqused_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + min_seqlen_k=min_seqlen_k, + page_table=page_table, + softmax_scale=softmax_scale, + causal=causal, + window_size_left=window_size[0], + window_size_right=window_size[1], + learnable_sink=learnable_sink, + softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, + score_mod=score_mod, + mask_mod=mask_mod, + block_sparse_tensors=block_sparse_tensors, + aux_tensors=aux_tensors, + aux_scalars=aux_scalars, + return_lse=return_lse, + gather_kv_indices=gather_kv_indices, + ) + ctx.save_for_backward( + q, + k, + v, + qv, + out, + lse, + p, + row_max, + gather_kv_indices, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + *(aux_tensors or ()), + ) + ctx.shared_kv = shared_kv + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.softcap = softcap + ctx.deterministic = deterministic + ctx.max_seqlen_q = max_seqlen_q + ctx.max_seqlen_k = max_seqlen_k + ctx.min_seqlen_k = min_seqlen_k + ctx.return_lse = return_lse + ctx.score_mod = score_mod + ctx.score_mod_bwd = score_mod_bwd + ctx.mask_mod = mask_mod + ctx.aux_scalars = aux_scalars + ctx.set_materialize_grads(False) + return out, lse + + @staticmethod + def backward(ctx, dout, dlse): + q, k, v, qv, out, lse, p, row_max, gather_kv_indices, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k, *aux = ctx.saved_tensors + aux_tensors = aux if aux else None + if not ctx.return_lse: + dlse = None + if dout is None: + dout = torch.zeros_like(out) + if qv is not None: + dq, dk, dv, dqv = _flash_attn_bwd_sparse_mla( + q, + k, + v, + qv, + out, + dout, + lse, + p, + row_max, + gather_kv_indices, + softmax_scale=ctx.softmax_scale, + causal=ctx.causal, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + seqused_q=seqused_q, + seqused_k=seqused_k, + max_seqlen_q=ctx.max_seqlen_q, + max_seqlen_k=ctx.max_seqlen_k, + min_seqlen_k=ctx.min_seqlen_k, + ) + if ctx.shared_kv: + return dqv, dv, None, None, *((None,) * 31) + else: + return dq, dk, dv, dqv, *((None,) * 31) + else: + dq, dk, dv = _flash_attn_bwd( + q, + k, + v, + out, + dout, + lse, + ctx.softmax_scale, + ctx.causal, + ctx.softcap, + window_size_left=ctx.window_size[0], + window_size_right=ctx.window_size[1], + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + seqused_q=seqused_q, + seqused_k=seqused_k, + max_seqlen_q=ctx.max_seqlen_q, + max_seqlen_k=ctx.max_seqlen_k, + deterministic=ctx.deterministic, + score_mod=ctx.score_mod, + score_mod_bwd=ctx.score_mod_bwd, + aux_tensors=aux_tensors, + aux_scalars=ctx.aux_scalars, + mask_mod=ctx.mask_mod, + dlse=dlse, + ) + return dq, dk, dv, *((None,) * 31) + + +def flash_attn_func( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + qv: Optional[torch.Tensor] = None, + gather_kv_indices: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size: Tuple[Optional[int], Optional[int]] = (None, None), + learnable_sink: Optional[torch.Tensor] = None, + softcap: float = 0.0, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + deterministic: bool = False, + score_mod: Optional[Callable] = None, + score_mod_bwd: Optional[Callable] = None, + mask_mod: Optional[Callable] = None, + aux_tensors: Optional[list] = None, + aux_scalars: Optional[tuple] = None, + block_sparse_tensors: Optional[BlockSparseTensorsTorch] = None, + block_sparse_tensors_bwd: Optional[BlockSparseTensorsTorch] = None, + return_lse: bool = False, +): + return FlashAttnFunc.apply( + q, + k, + v, + qv, + gather_kv_indices, + softmax_scale, + causal, + window_size, + learnable_sink, + softcap, + num_splits, + pack_gqa, + deterministic, + score_mod, + score_mod_bwd, + mask_mod, + aux_tensors, + aux_scalars, + block_sparse_tensors, + block_sparse_tensors_bwd, + return_lse, + ) + + +def flash_attn_varlen_func( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + qv: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = None, + min_seqlen_k: Optional[int] = None, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + gather_kv_indices: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size: Tuple[Optional[int], Optional[int]] = (None, None), + learnable_sink: Optional[torch.Tensor] = None, + softcap: float = 0.0, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + deterministic: bool = False, + score_mod: Optional[Callable] = None, + score_mod_bwd: Optional[Callable] = None, + mask_mod: Optional[Callable] = None, + block_sparse_tensors: Optional[BlockSparseTensorsTorch] = None, + aux_tensors: Optional[list] = None, + aux_scalars: Optional[tuple] = None, + return_lse: bool = False, +): + """ + Tensor arguments: + q: (total_q, nheads, hdim) or (batch, seqlen_q, nheads, hdim) + k: (total_k, nheads_k, hdim) or (batch, seqlen_k, nheads_k, hdim) + v: (total_k, nheads_k, hdim_v) or (batch, seqlen_k, nheads_k, hdim_v) + qv: (total_q, nheads, hdim_v) or (batch, seqlen_q, nheads, hdim_v) + cu_seqlens_q: (batch + 1) or seqused_q: (batch) + cu_seqlens_k: (batch + 1) or seqused_k: (batch) + gather_kv_indices: (total_q, gather_kv_length) or + (batch, seqlen_q, gather_kv_length) + page_table: (batch, max_num_pages_per_seq) + + Return: + out: (total_q, nheads, hdim) or (batch, seqlen_q, nheads, hdim) + lse: (nheads, total_q) or (batch, nheads, seqlen_q) if not has_qv (standard) + (total_q, nheads) or (batch, seqlen_q, nheads) if has_qv + + Explanation of some optional arguments & decisions: + + qv: we write the MLA weight absorbed formula as + O = softmax(scale * (Q @ K.T + Qv @ V.T)) @ V + where Q = q_pe, Qv = q_nope, K = pe_cache, V = kv_cache. + + lse return shape: with Qv, MQA with nheads at least divisible by 4 is typical, + so we arrange for nheads as the contiguous mode for better vectorization. + + gather_kv_indices: used for topk sparsity with MLA absorption kernel. + """ + return FlashAttnVarlenFunc.apply( + q, + k, + v, + qv, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + min_seqlen_k, + gather_kv_indices, + page_table, + softmax_scale, + causal, + window_size, + learnable_sink, + softcap, + num_splits, + pack_gqa, + deterministic, + score_mod, + score_mod_bwd, + mask_mod, + block_sparse_tensors, + aux_tensors, + aux_scalars, + return_lse, + ) + + +def _compile_fwd_combine( + dtype, dtype_partial, head_dim, tile_m, k_block_size, log_max_splits, + has_cu_seqlens, has_seqused, has_lse, has_varlen_batch_idx, +): + """Compile fwd combine kernel using cute fake tensors (no real GPU tensors needed).""" + sym = cute.sym_int + div = 128 // dtype_partial.width # 16-byte alignment in elements + + fa_combine = FlashAttentionForwardCombine( + dtype=dtype, + dtype_partial=dtype_partial, + head_dim=head_dim, + tile_m=tile_m, + k_block_size=k_block_size, + log_max_splits=log_max_splits, + ) + if not fa_combine.can_implement( + dtype, dtype_partial, head_dim, tile_m, k_block_size, log_max_splits, + num_threads=256, + ): + raise RuntimeError( + "FlashAttention combine kernel cannot be implemented with given parameters" + ) + + if has_cu_seqlens: + # Varlen: (num_splits, total_q, nheads, headdim) + num_splits, total_q, nheads = sym(), sym(), sym() + mO_partial = fake_tensor(dtype_partial, (num_splits, total_q, nheads, head_dim), divisibility=div) + mLSE_partial = fake_tensor(Float32, (num_splits, total_q, nheads), divisibility=1, leading_dim=1) + mO = fake_tensor(dtype, (total_q, nheads, head_dim), divisibility=div) + mLSE = fake_tensor(Float32, (total_q, nheads), divisibility=1, leading_dim=0) if has_lse else None + else: + # Batched: (num_splits, batch, seqlen, nheads, headdim) + num_splits, batch, seqlen, nheads = sym(), sym(), sym(), sym() + mO_partial = fake_tensor(dtype_partial, (num_splits, batch, seqlen, nheads, head_dim), divisibility=div) + mLSE_partial = fake_tensor(Float32, (num_splits, batch, seqlen, nheads), divisibility=1, leading_dim=2) + mO = fake_tensor(dtype, (batch, seqlen, nheads, head_dim), divisibility=div) + mLSE = fake_tensor(Float32, (batch, seqlen, nheads), divisibility=1, leading_dim=1) if has_lse else None + batch = mO_partial.shape[1] + + batch_for_1d = batch if not has_cu_seqlens else sym() + batchp1 = sym() + mCuSeqlens = fake_tensor(Int32, (batchp1,), divisibility=1) if has_cu_seqlens else None + mSeqused = fake_tensor(Int32, (batch_for_1d,), divisibility=1) if has_seqused else None + mNumSplitsDynamic = None # Not parametrized in compile_key + mVarlenBatchIdx = fake_tensor(Int32, (batch_for_1d,), divisibility=1) if has_varlen_batch_idx else None + mSemaphore = None # Not parametrized in compile_key + + return cute.compile( + fa_combine, + mO_partial, mLSE_partial, mO, mLSE, + mCuSeqlens, mSeqused, mNumSplitsDynamic, mVarlenBatchIdx, mSemaphore, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options="--enable-tvm-ffi", + ) + + +def _flash_attn_fwd_combine( + out_partial: torch.Tensor, + lse_partial: torch.Tensor, + out: torch.Tensor, + lse: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, + seqused: Optional[torch.Tensor] = None, + num_splits_dynamic_ptr: Optional[torch.Tensor] = None, + varlen_batch_idx: Optional[torch.Tensor] = None, + semaphore_to_reset: Optional[torch.Tensor] = None, +) -> None: + """Forward combine kernel for split attention computation. + + Combines partial outputs and log-sum-exp values from multiple splits + of attention computation into final outputs. + + Args: + out_partial: Partial outputs tensor (num_splits, batch, seqlen, nheads, headdim) or + (num_splits, total_q, nheads, headdim) if there's cu_seqlens + lse_partial: Partial LSE tensor (num_splits, batch, seqlen, nheads) or + (num_splits, total_q, nheads) if there's cu_seqlens + out: Output tensor (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim) if there's cu_seqlens + lse: Output LSE tensor (batch, seqlen, nheads) or (total_q, nheads) if there's cu_seqlens. + cu_seqlens: Cumulative sequence lengths for variable length sequences + seqused: Used sequence lengths for each batch + num_splits_dynamic_ptr: Dynamic number of splits per batch + semaphore_to_reset: Semaphore for synchronization + k_block_size: Block size for head dimension + + Returns: + None + """ + assert out_partial.dtype in [torch.float16, torch.bfloat16, torch.float32], ( + "out_partial must be fp16, bf16, or fp32" + ) + if not is_fake_mode(): + assert out_partial.is_cuda and lse_partial.is_cuda, "tensors must be on CUDA device" + # Determine if this is variable length based on dimensions + is_varlen = out_partial.dim() == 4 + # Validate optional tensors + for t, name in [ + (cu_seqlens, "cu_seqlens"), + (seqused, "seqused"), + (num_splits_dynamic_ptr, "num_splits_dynamic_ptr"), + ]: + if t is not None: + if not is_fake_mode(): + assert t.is_cuda, f"{name} must be on CUDA device" + assert t.is_contiguous(), f"{name} must be contiguous" + head_dim = out_partial.shape[-1] + num_splits = out_partial.shape[0] + assert num_splits <= 256 + # If hdim is 96 or 192, it's faster to round them to 128 or 256 respectively + # so that kBlockM is smaller and we have more parallelism. + k_block_size = 64 if head_dim <= 64 else 128 + # We want kBlockM to be as small as possible to maximize parallelism. + # E.g., if hdim is 64, we want kBlockM to be 16 so that we can use 256 threads, each reading 4 elements (floats). + tile_m = 8 if k_block_size % 128 == 0 else (16 if k_block_size % 64 == 0 else 32) + log_max_splits = max(math.ceil(math.log2(num_splits)), 4) + if tile_m == 8: + # If kBlockM == 8 then the minimum number of splits is 32. + # TODO: we can deal w this by using 128 threads instead + log_max_splits = max(log_max_splits, 5) + + # Create combine kernel configuration + dtype = torch2cute_dtype_map[out.dtype] + dtype_partial = torch2cute_dtype_map[out_partial.dtype] + compile_key = ( + dtype, + dtype_partial, + head_dim, + tile_m, + k_block_size, + log_max_splits, + cu_seqlens is not None, + seqused is not None, + lse is not None, + varlen_batch_idx is not None, + ) + if compile_key not in _flash_attn_fwd_combine.compile_cache: + _flash_attn_fwd_combine.compile_cache[compile_key] = _compile_fwd_combine( + *compile_key + ) + if not is_fake_mode(): + _flash_attn_fwd_combine.compile_cache[compile_key]( + out_partial, lse_partial, out, lse, + cu_seqlens, seqused, num_splits_dynamic_ptr, varlen_batch_idx, + semaphore_to_reset, + ) + + +_flash_attn_fwd_combine.compile_cache = get_jit_cache("fwd_combine") + + +def flash_attn_combine( + out_partial: torch.Tensor, + lse_partial: torch.Tensor, + out: Optional[torch.Tensor] = None, + out_dtype: Optional[torch.dtype] = None, + cu_seqlens: Optional[torch.Tensor] = None, + seqused: Optional[torch.Tensor] = None, + varlen_batch_idx: Optional[torch.Tensor] = None, + return_lse: bool = True, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Flash Attention combine function for split attention computation. + + Combines partial outputs and log-sum-exp values from multiple splits + of attention computation into final outputs. This is the main user-facing + interface for the combine kernel. + + Args: + out_partial: Partial outputs tensor with shape: + - (num_splits, batch_size, seqlen, num_heads, head_size) for regular batched input + - (num_splits, total_q, num_heads, head_size) for variable length input + lse_partial: Partial LSE tensor with shape: + - (num_splits, batch_size, seqlen, num_heads) for regular batched input + - (num_splits, total_q, num_heads) for variable length input + out: Optional output tensor. If None, will be created automatically. + out_dtype: Optional output dtype. If None, will use fp16/bf16 based on input. + cu_seqlens: Cumulative sequence lengths for variable length sequences + seqused: Used sequence lengths for each batch + varlen_batch_idx: Optional mapping from virtual batch index to real batch index + (int32 tensor of shape (batch_size,)). Used by persistent tile schedulers + that reorder batch processing for load balancing. + return_lse: Whether to return the combined LSE tensor. Default is True. + + Returns: + Tuple of (out, lse) where: + - out: Combined output tensor with shape (batch_size, seqlen, num_heads, head_size) + or (total_q, num_heads, head_size) for varlen + - lse: Combined log-sum-exp tensor with shape (batch_size, seqlen, num_heads) + or (total_q, num_heads) for varlen. None if return_lse=False + + Note: + This function expects the input tensors to be in the format produced by + split attention computation, where the first dimension is num_splits. + The permuting from user format to kernel format is now done inside the kernel. + """ + # Input validation + assert out_partial.dim() in [4, 5], "out_partial must have 4 or 5 dimensions" + # Determine if this is variable length based on dimensions + is_varlen = out_partial.dim() == 4 + if is_varlen: + # Variable length: (num_splits, total_q, num_heads, head_size) + num_splits, total_q, num_heads, head_size = out_partial.shape + batch_size = 1 # Treat as single batch for varlen + seqlen = total_q + else: + # Regular batched: (num_splits, batch_size, seqlen, num_heads, head_size) + num_splits, batch_size, seqlen, num_heads, head_size = out_partial.shape + # Determine output dtype + if out_dtype is None: + out_dtype = out_partial.dtype + # Create output if not provided + device = out_partial.device + if out is None: + if is_varlen: + out = torch.empty(total_q, num_heads, head_size, dtype=out_dtype, device=device) + else: + out = torch.empty( + batch_size, seqlen, num_heads, head_size, dtype=out_dtype, device=device + ) + # Create lse output only if requested + if return_lse: + if is_varlen: + lse = torch.empty(num_heads, total_q, dtype=torch.float32, device=device) + else: + lse = torch.empty(batch_size, num_heads, seqlen, dtype=torch.float32, device=device) + lse = lse.transpose(-1, -2) + else: + lse = None + _flash_attn_fwd_combine( + out_partial, + lse_partial, + out, + lse, + cu_seqlens, + seqused, + varlen_batch_idx=varlen_batch_idx, + ) + return out, lse diff --git a/fastvideo-kernel/fa4/mask.py b/fastvideo-kernel/fa4/mask.py new file mode 100644 index 000000000..312cb0615 --- /dev/null +++ b/fastvideo-kernel/fa4/mask.py @@ -0,0 +1,1709 @@ +# Copyright (c) 2025, Tri Dao. + +from typing import Optional, Callable, TypeAlias, Tuple +from dataclasses import dataclass +import enum + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Uint32, const_expr +from cutlass.cutlass_dsl import min as dsl_min + +from quack import layout_utils +import flash_attn.cute.utils as utils +from flash_attn.cute.block_info import BlockInfo +from flash_attn.cute.seqlen_info import SeqlenInfoQK +from flash_attn.cute.utils import AuxData + +MaskGenFn: TypeAlias = Callable[[int], Uint32] +MASK_R2P_CHUNK_SIZE: int = 32 + + +@cute.jit +def call_mask_mod( + mask_mod: cutlass.Constexpr, + batch_idx, + head_idx, + q_idx, + kv_idx, + seqlen_info, + aux_data: AuxData, +): + # Compatibility shim for pre-aux_scalars mask_mod callables. + if const_expr(aux_data.scalars is not None): + return mask_mod( + batch_idx, + head_idx, + q_idx, + kv_idx, + seqlen_info, + aux_data.tensors, + aux_data.scalars, + ) + return mask_mod( + batch_idx, + head_idx, + q_idx, + kv_idx, + seqlen_info, + aux_data.tensors, + ) + + +@cute.jit +def r2p_bitmask_below(limit: Int32, s: int) -> Uint32: + """32-bit R2P bitmask keeping positions < limit (exclusive upper bound). + + Positions 0..limit-1 in chunk `s` get bit=1 (keep), the rest bit=0 (mask). + Uses inline PTX to avoid shift-by-type-width UB. + """ + m = max((s + 1) * MASK_R2P_CHUNK_SIZE - limit, 0) + return utils.shr_u32(Uint32(0xFFFFFFFF), Uint32(m)) + + +@cute.jit +def r2p_bitmask_above(limit: Int32, s: int) -> Uint32: + """32-bit R2P bitmask keeping positions >= limit (inclusive lower bound). + + Positions limit..31 in chunk `s` get bit=1 (keep), the rest bit=0 (mask). + Uses inline PTX to avoid shift-by-type-width UB. + """ + n = max(limit - s * MASK_R2P_CHUNK_SIZE, 0) + return utils.shl_u32(Uint32(0xFFFFFFFF), Uint32(n)) + + +@cute.jit +def mask_r2p_lambda( + X: cute.Tensor, + mask_gen_fn: cutlass.Constexpr[MaskGenFn], + rank1: bool = False, +) -> None: + """Apply R2P masking with a custom bitmask generator. + + mask_gen_fn(chunk_idx: constexpr int) -> Uint32: + Returns a 32-bit bitmask for the chunk. Bit i set means column + chunk_idx * chunk_size + i is KEPT; bit i clear means masked to -inf. + """ + ncol = const_expr(cute.size(X.shape[cute.rank(X) - 1]) if not rank1 else cute.size(X.shape)) + # 32-column chunks. The mask_gen_fn returns a Uint32 bitmask (1=keep). + CHUNK_SIZE = MASK_R2P_CHUNK_SIZE + for s in cutlass.range_constexpr(cute.ceil_div(ncol, CHUNK_SIZE)): + mask = mask_gen_fn(s) + # This needs to be range_constexpr, o/w the compiler can't generate the R2P instruction + for i in cutlass.range_constexpr(min(CHUNK_SIZE, ncol - s * CHUNK_SIZE)): + in_bound = cutlass.Boolean(mask & (Uint32(1) << i)) + c = s * CHUNK_SIZE + i + if const_expr(rank1): + X[c] = X[c] if in_bound else -Float32.inf + else: + for r in cutlass.range_constexpr(cute.size(X.shape[0])): + X[r, c] = X[r, c] if in_bound else -Float32.inf + + +@cute.jit +def sm90_col_to_r2p_idx(col_limit: Int32) -> Int32: + """Transform SM90 MMA column coordinate to R2P element index. + + SM90 MMA accumulator column indices are non-contiguous: 0, 1, 8, 9, 16, 17, ... + Element indices are contiguous: 0, 1, 2, 3, 4, 5, ... + This converts a column-space threshold to element-space for r2p_bitmask_below/above. + """ + return col_limit // 8 * 2 + min(col_limit % 8, 2) + + +@cute.jit +def row_to_r2p_idx(x: Int32, num_rep: int, num_wg: int) -> Int32: + """Convert a row coordinate to an R2P element index in the warp-group interleaved layout. + + In the SM100 backward pass, 2 warp groups share TMEM. The TMEM load atom + distributes rows in an interleaved pattern: elements 0..num_rep-1 map to + rows 0..num_rep-1 (warp group 0), elements num_rep..2*num_rep-1 map to + rows num_rep*num_wg..num_rep*num_wg+num_rep-1 (warp group 1), and so on. + Row-coordinate thresholds (causal limits, window bounds, uih_len) must be + converted to element indices before use with r2p_bitmask_above/below. + + Rows not owned by this thread (in the gap between warp groups) are clamped + to the boundary element index, which is safe because R2P thresholds are + monotonic. + + Example with num_rep=16, num_wg=2: + row 0 -> elem 0, row 15 -> elem 15, + row 16 -> elem 16 (clamped), row 31 -> elem 16 (clamped), + row 32 -> elem 16, row 33 -> elem 17, row 47 -> elem 31. + """ + return x // (num_rep * num_wg) * num_rep + min(x % (num_rep * num_wg), num_rep) + + +@cute.jit +def apply_packed_mask_chunk( + X: cute.Tensor, + chunk_idx: cutlass.Constexpr[int], + mask: Uint32, +) -> None: + """Apply one 32-bit keep mask to one 32-column chunk. + + The one-iteration chunk loop keeps the same lowering pattern as mask_r2p_lambda. + """ + ncol = const_expr(cute.size(X.shape)) + col_base = chunk_idx * MASK_R2P_CHUNK_SIZE + for s in cutlass.range_constexpr(1): + for i in cutlass.range_constexpr( + min(MASK_R2P_CHUNK_SIZE, ncol - col_base - s * MASK_R2P_CHUNK_SIZE) + ): + in_bound = cutlass.Boolean(mask & (Uint32(1) << i)) + c = col_base + s * MASK_R2P_CHUNK_SIZE + i + X[c] = X[c] if in_bound else -Float32.inf + + +@dataclass(frozen=True) +class AttentionMask: + tile_m: cutlass.Constexpr[int] + tile_n: cutlass.Constexpr[int] + seqlen_info: SeqlenInfoQK + window_size_left: Optional[Int32] = None + window_size_right: Optional[Int32] = None + qhead_per_kvhead_packgqa: cutlass.Constexpr[int] = 1 # only pass in if we're doing PackGQA + swap_AB: cutlass.Constexpr[bool] = False + + @property + def seqlen_q(self) -> Int32: + return self.seqlen_info.seqlen_q + + @property + def seqlen_k(self) -> Int32: + return self.seqlen_info.seqlen_k + + @cute.jit + def apply_mask( + self, + acc_S: cute.Tensor, + batch_idx: cutlass.Int32, + head_idx: cutlass.Int32, + m_block: cutlass.Int32, + n_block: cutlass.Int32, + thr_mma: cute.TiledMma, + mask_seqlen: cutlass.Constexpr[bool], + mask_causal: cutlass.Constexpr[bool], + mask_local: cutlass.Constexpr[bool] = False, + mask_mod: cutlass.Constexpr[Optional[Callable]] = None, + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + ) -> None: + assert not (mask_causal and mask_local), "mask_causal and mask_local cannot be both True" + acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S, transpose=self.swap_AB) + acc_shape = (self.tile_m, self.tile_n) + cS = cute.make_identity_tensor(acc_shape if not self.swap_AB else acc_shape[::-1]) + tScS_mn = layout_utils.reshape_acc_to_mn(thr_mma.partition_C(cS), transpose=self.swap_AB) + # We use t0ScS as these indices are known at compile time. We then must subtract the + # column limit by the thread column offset. + t0ScS_mn = layout_utils.reshape_acc_to_mn( + thr_mma.get_slice(0).partition_C(cS), transpose=self.swap_AB + ) + ROW = 0 if const_expr(not self.swap_AB) else 1 + COL = 1 if const_expr(not self.swap_AB) else 0 + thr_col_offset = tScS_mn[0][COL] + # To handle edge cases of completely masked out rows where n_block_max = 0, + # we treat negative n_blocks as 0th n_block + # TODO: find more transparent solution + if n_block < 0: + n_block = 0 + seqlenk_col_limit = self.seqlen_k - n_block * self.tile_n - thr_col_offset + if const_expr(not mask_causal and not mask_local and mask_mod is None): + if const_expr(mask_seqlen): + r2p = const_expr(not self.swap_AB) + if const_expr(not r2p): + # traverse column index. + for c in cutlass.range(cute.size(tScS_mn.shape[1]), unroll_full=True): + oob = t0ScS_mn[0, c][COL] >= seqlenk_col_limit + for r in cutlass.range(cute.size(tScS_mn.shape[0]), unroll_full=True): + acc_S_mn[r, c] = -Float32.inf if oob else acc_S_mn[r, c] + else: + seqlenk_col_limit_r2p = sm90_col_to_r2p_idx(seqlenk_col_limit) + mask_r2p_lambda(acc_S_mn, lambda s: r2p_bitmask_below(seqlenk_col_limit_r2p, s)) + + elif const_expr( + not mask_causal and not mask_local and mask_mod is not None + ): # FlexAttention mask mod + nrow = const_expr(cute.size(tScS_mn.shape[0])) + ncol = const_expr(cute.size(tScS_mn.shape[1])) + has_fastdiv = const_expr( + fastdiv_mods is not None + and fastdiv_mods[0] is not None + and fastdiv_mods[1] is not None + ) + wrap_aux_indices = const_expr( + has_fastdiv and mask_seqlen and const_expr(aux_data.tensors is not None) + ) + + for r in cutlass.range_constexpr(nrow): + # Respect swap_AB: ROW/COL determine which coordinate component corresponds to Q/KV. + local_row = tScS_mn[r, 0][ROW] + global_row_idx = local_row + m_block * self.tile_m + row_for_mod = global_row_idx + head_idx_for_mod = head_idx + if const_expr(self.qhead_per_kvhead_packgqa != 1): + head_offset = global_row_idx % self.qhead_per_kvhead_packgqa + head_idx_for_mod = head_idx * self.qhead_per_kvhead_packgqa + head_offset + row_for_mod = global_row_idx // self.qhead_per_kvhead_packgqa + row_for_seqlen = row_for_mod + if const_expr(wrap_aux_indices): + _, row_for_mod = divmod(row_for_mod, fastdiv_mods[0]) + + for col in cutlass.range_constexpr(ncol): + col_idx_local = t0ScS_mn[0, col][COL] + # Convert to absolute column index + global_col_idx = thr_col_offset + col_idx_local + n_block * self.tile_n + col_for_mod = global_col_idx + if const_expr(wrap_aux_indices): + _, col_for_mod = divmod(global_col_idx, fastdiv_mods[1]) + + batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32) + head_idx_ssa = utils.scalar_to_ssa(head_idx_for_mod, cutlass.Int32) + q_idx_ssa = utils.scalar_to_ssa(row_for_mod, cutlass.Int32) + kv_idx_ssa = utils.scalar_to_ssa(col_for_mod, cutlass.Int32) + mask_value = call_mask_mod( + mask_mod, + batch_idx_ssa, + head_idx_ssa, + q_idx_ssa, + kv_idx_ssa, + self.seqlen_info, + aux_data, + ) + cond = cutlass.Boolean(utils.ssa_to_scalar(mask_value)) + if const_expr(mask_seqlen): + out_of_bounds = (row_for_seqlen >= self.seqlen_q) or ( + global_col_idx >= self.seqlen_k + ) + if out_of_bounds: + acc_S_mn[r, col] = -cutlass.Float32.inf + else: + acc_S_mn[r, col] = acc_S_mn[r, col] if cond else -cutlass.Float32.inf + else: + acc_S_mn[r, col] = acc_S_mn[r, col] if cond else -cutlass.Float32.inf + + else: # Causal or local + if const_expr(not self.swap_AB): + # If PackGQA, we split the work of compute divmod among threads in the same row + threads_per_row = thr_mma.tv_layout_C.shape[0][0] + mma_m_idx = None + if const_expr(self.qhead_per_kvhead_packgqa != 1): + assert not self.swap_AB, "swap_AB with PackGQA not supported yet" + assert cute.arch.WARP_SIZE % threads_per_row == 0, ( + "threads_per_row must divide WARP_SIZE" + ) + assert cute.size(acc_S_mn.shape[0]) <= threads_per_row + tidx = thr_mma.thr_idx + mma_m_idx = ( + m_block * self.tile_m + tScS_mn[tidx % threads_per_row, 0][0] + ) // self.qhead_per_kvhead_packgqa + causal_row_offset = ( + 1 + self.seqlen_k - n_block * self.tile_n - self.seqlen_q - thr_col_offset + ) + if const_expr(mask_causal): + r2p = const_expr(not self.swap_AB) # R2P trick, see apply_mask_sm100 + for r in cutlass.range(cute.size(tScS_mn.shape[0]), unroll_full=True): + # get the column index limit based on current row. Only consider the row index, so the column index sets to 0. + if const_expr(self.qhead_per_kvhead_packgqa == 1): + row_idx = tScS_mn[r, 0][0] + m_block * self.tile_m + else: + row_idx = utils.shuffle_sync( + mma_m_idx, r % threads_per_row, width=threads_per_row + ) + col_limit_right = row_idx + causal_row_offset + if const_expr(mask_seqlen): + col_limit_right = cutlass.min(col_limit_right, seqlenk_col_limit) + if const_expr(not r2p): + # traverse column index. + for c in cutlass.range(cute.size(tScS_mn.shape[1]), unroll_full=True): + acc_S_mn[r, c] = ( + -Float32.inf + if t0ScS_mn[0, c][1] >= col_limit_right + else acc_S_mn[r, c] + ) + else: + col_limit_r2p = sm90_col_to_r2p_idx(col_limit_right) + mask_r2p_lambda( + acc_S_mn[r, None], + lambda s: r2p_bitmask_below(col_limit_r2p, s), + rank1=True, + ) + else: # Local + local_row_offset_right = ( + causal_row_offset + self.window_size_right + if const_expr(self.window_size_right is not None) + else None + ) + local_row_offset_left = ( + causal_row_offset - 1 - self.window_size_left + if const_expr(self.window_size_left is not None) + else None + ) + r2p_local = const_expr(not self.swap_AB) + for r in cutlass.range(cute.size(tScS_mn.shape[0]), unroll_full=True): + if const_expr(self.qhead_per_kvhead_packgqa == 1): + row_idx = tScS_mn[r, 0][0] + m_block * self.tile_m + else: + row_idx = utils.shuffle_sync( + mma_m_idx, r % threads_per_row, width=threads_per_row + ) + if const_expr(self.window_size_right is not None): + col_limit_right = row_idx + local_row_offset_right + else: + col_limit_right = self.tile_n + if const_expr(mask_seqlen): + col_limit_right = cutlass.min(col_limit_right, seqlenk_col_limit) + col_limit_left = ( + row_idx + local_row_offset_left + if const_expr(self.window_size_left is not None) + else 0 + ) + if const_expr(not r2p_local): + # traverse column index. + for c in cutlass.range(cute.size(tScS_mn.shape[1]), unroll_full=True): + col_idx = t0ScS_mn[0, c][1] + if col_idx >= col_limit_right or col_idx < col_limit_left: + acc_S_mn[r, c] = -Float32.inf + else: + col_limit_right_r2p = sm90_col_to_r2p_idx(col_limit_right) + col_limit_left_r2p = sm90_col_to_r2p_idx(col_limit_left) + + def mask_gen_fn(s: int) -> Uint32: + return r2p_bitmask_below( + col_limit_right_r2p, s + ) & r2p_bitmask_above(col_limit_left_r2p, s) + + mask_r2p_lambda(acc_S_mn[r, None], mask_gen_fn, rank1=True) + else: # swap_AB + assert self.qhead_per_kvhead_packgqa == 1 + thr_row_offset = tScS_mn[0][ROW] + causal_row_offset = ( + seqlenk_col_limit - self.seqlen_q + m_block * self.tile_m + thr_row_offset + ) + if const_expr(mask_causal): + for c in cutlass.range(cute.size(tScS_mn.shape[1]), unroll_full=True): + col0 = t0ScS_mn[0, c][COL] + # If col0 is beyond the column limit, we want to mask out the entire + # column, by setting row limit to be self.tile_m. + row_limit_top = ( + self.tile_m + if col0 >= seqlenk_col_limit and mask_seqlen + else col0 - causal_row_offset + ) + for r in cutlass.range(cute.size(tScS_mn.shape[0]), unroll_full=True): + acc_S_mn[r, c] = ( + -Float32.inf + if t0ScS_mn[r, 0][ROW] < row_limit_top + else acc_S_mn[r, c] + ) + else: + for c in cutlass.range(cute.size(tScS_mn.shape[1]), unroll_full=True): + col0 = t0ScS_mn[0, c][COL] + # If col0 is beyond the column limit, we want to mask out the entire + # column, by setting row limit to be self.tile_m. + row_limit_top = ( + self.tile_m + if col0 >= seqlenk_col_limit and mask_seqlen + else ( + col0 - causal_row_offset - self.window_size_right + if const_expr(self.window_size_right is not None) + else 0 + ) + ) + row_limit_bot = ( + col0 - causal_row_offset + self.window_size_left + if const_expr(self.window_size_left is not None) + else self.tile_m + ) + for r in cutlass.range(cute.size(tScS_mn.shape[0]), unroll_full=True): + row_idx = t0ScS_mn[r, 0][ROW] + acc_S_mn[r, c] = ( + -Float32.inf + if row_idx < row_limit_top or row_idx > row_limit_bot + else acc_S_mn[r, c] + ) + + @cute.jit + def apply_mask_mod_sm100_scalar( + self, + acc_S: cute.Tensor, + tScS_t2r: cute.Tensor, + m_block: Int32, + n_block: Int32, + mask_seqlen: cutlass.Constexpr[bool], + mask_mod: cutlass.Constexpr[Callable], + batch_idx: Int32, + head_idx: Int32, + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + head_divmod=None, + check_q_boundary: bool = False, + ) -> None: + """Apply a scalar FlexAttention mask_mod to an SM100 accumulator fragment. + + Each accumulator lane calls mask_mod once with logical (batch, head, q, kv) + indices. Pack-GQA rows are converted back to logical q/head indices before + the call. When aux tensors are present, indices are wrapped with fastdiv so + mask_mod never reads outside the per-example auxiliary storage. + """ + has_fastdiv = const_expr( + fastdiv_mods is not None and fastdiv_mods[0] is not None and fastdiv_mods[1] is not None + ) + batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32) + ncol = const_expr(cute.size(tScS_t2r.shape)) + + for i in cutlass.range_constexpr(ncol): + row_coord = tScS_t2r[i][0] if not self.swap_AB else tScS_t2r[i][1] + col_coord = tScS_t2r[i][1] if not self.swap_AB else tScS_t2r[i][0] + global_row = row_coord + m_block * self.tile_m + global_col = col_coord + n_block * self.tile_n + + if const_expr(self.qhead_per_kvhead_packgqa != 1): + assert head_divmod is not None + mask_row, head_offset = divmod(global_row, head_divmod) + head_idx_for_mod = head_idx * self.qhead_per_kvhead_packgqa + head_offset + else: + head_idx_for_mod = head_idx + mask_row = global_row + + mask_row_for_mod = mask_row + if const_expr(has_fastdiv and aux_data.tensors is not None): + if check_q_boundary: + _, mask_row_for_mod = divmod(mask_row, fastdiv_mods[0]) + global_col_for_mod = global_col + if const_expr(has_fastdiv and mask_seqlen and aux_data.tensors is not None): + _, global_col_for_mod = divmod(global_col, fastdiv_mods[1]) + + head_idx_ssa = utils.scalar_to_ssa(head_idx_for_mod, cutlass.Int32) + mask_row_ssa = utils.scalar_to_ssa(mask_row_for_mod, cutlass.Int32) + kv_idx_ssa = utils.scalar_to_ssa(global_col_for_mod, cutlass.Int32) + mask_value = call_mask_mod( + mask_mod, + batch_idx_ssa, + head_idx_ssa, + mask_row_ssa, + kv_idx_ssa, + self.seqlen_info, + aux_data, + ) + cond = cutlass.Boolean(utils.ssa_to_scalar(mask_value)) + acc_S[i] = acc_S[i] if cond else -Float32.inf + if const_expr(mask_seqlen): + acc_S[i] = -Float32.inf if global_col >= self.seqlen_k else acc_S[i] + if check_q_boundary: + acc_S[i] = -Float32.inf if mask_row >= self.seqlen_q else acc_S[i] + + @cute.jit + def apply_mask_mod_sm100_vector( + self, + acc_S: cute.Tensor, + tScS_t2r: cute.Tensor, + m_block: Int32, + n_block: Int32, + mask_seqlen: cutlass.Constexpr[bool], + mask_mod: cutlass.Constexpr[Callable], + batch_idx: Int32, + head_idx: Int32, + vec_size: cutlass.Constexpr[int], + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + head_divmod=None, + check_q_boundary: bool = False, + ) -> None: + """Apply a vectorized FlexAttention mask_mod to an SM100 fragment. + + mask_mod receives vec_size adjacent KV indices for one logical q row and + returns bit-packed Uint32 keep masks. Low bits correspond to lower KV + indices. The packed masks are combined with sequence-boundary checks, then + applied in 32-column chunks so the final masking lowers to R2P. + """ + has_fastdiv = const_expr( + fastdiv_mods is not None and fastdiv_mods[0] is not None and fastdiv_mods[1] is not None + ) + batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32) + ncol = const_expr(cute.size(tScS_t2r.shape)) + mask_vals_per_apply = const_expr(max(1, vec_size // 32)) + calls_per_apply = const_expr(max(1, 32 // vec_size)) + n_calls = const_expr(cute.ceil_div(ncol, vec_size)) + mask_vals = cute.make_rmem_tensor(mask_vals_per_apply, dtype=cutlass.Uint32) + + # Accumulate enough vector mask_mod calls to produce 32-bit chunks that + # apply_packed_mask_chunk can lower to R2P. + for s in cutlass.range_constexpr(n_calls): + if const_expr(s % calls_per_apply == 0): + for c in cutlass.range_constexpr(mask_vals_per_apply): + mask_vals[c] = cutlass.Uint32(0) + i = s * vec_size + row_coord = tScS_t2r[i][0] if not self.swap_AB else tScS_t2r[i][1] + col_coord = tScS_t2r[i][1] if not self.swap_AB else tScS_t2r[i][0] + global_row = row_coord + m_block * self.tile_m + global_col = col_coord + n_block * self.tile_n + if const_expr(self.qhead_per_kvhead_packgqa != 1): + assert head_divmod is not None + mask_row, head_offset = divmod(global_row, head_divmod) + head_idx_for_mod = head_idx * self.qhead_per_kvhead_packgqa + head_offset + else: + head_idx_for_mod = head_idx + mask_row = global_row + mask_row_for_mod = mask_row + if const_expr(has_fastdiv and aux_data.tensors is not None): + if check_q_boundary: + _, mask_row_for_mod = divmod(mask_row, fastdiv_mods[0]) + + head_idx_ssa = utils.scalar_to_ssa(head_idx_for_mod, cutlass.Int32).broadcast_to( + (vec_size,) + ) + mask_row_ssa = utils.scalar_to_ssa(mask_row_for_mod, cutlass.Int32).broadcast_to( + (vec_size,) + ) + batch_idx_ssa_call = batch_idx_ssa.broadcast_to((vec_size,)) + kv_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) + + # Build the per-lane KV indices for this vectorized mask_mod call. + for j in cutlass.range_constexpr(min(vec_size, ncol - i)): + col_j_coord = tScS_t2r[i + j][1] if not self.swap_AB else tScS_t2r[i + j][0] + col_j_global = col_j_coord + n_block * self.tile_n + col_j_for_mod = col_j_global + if const_expr(has_fastdiv and mask_seqlen and aux_data.tensors is not None): + _, col_j_for_mod = divmod(col_j_global, fastdiv_mods[1]) + kv_idx_vec[j] = col_j_for_mod + kv_idx_ssa = kv_idx_vec.load() + + # mask_value is already bit-packed by the vectorized mask_mod. + mask_value = call_mask_mod( + mask_mod, + batch_idx_ssa_call, + head_idx_ssa, + mask_row_ssa, + kv_idx_ssa, + self.seqlen_info, + aux_data, + ) + + # For vec_size < 32, multiple mask_mod calls fill one R2P chunk. + bit_offset = const_expr((s % calls_per_apply) * vec_size) + seqlen_thresh_call = ( + self.seqlen_k - global_col if const_expr(mask_seqlen) else cutlass.Int32(0) + ) + q_in_bounds = mask_row < self.seqlen_q if check_q_boundary else cutlass.Boolean(True) + for c in cutlass.range_constexpr(mask_vals_per_apply): + mask_val = mask_value[c] + if const_expr(vec_size < 32): + lane_keep = utils.shr_u32( + cutlass.Uint32(0xFFFFFFFF), + cutlass.Uint32(32 - vec_size), + ) + mask_val = mask_val & lane_keep + if const_expr(mask_seqlen): + mask_val = mask_val & r2p_bitmask_below(seqlen_thresh_call, c) + if check_q_boundary: + mask_val = mask_val if q_in_bounds else cutlass.Uint32(0) + mask_vals[c] = mask_vals[c] | (mask_val << bit_offset) + + # Apply only when the 32-bit chunk is complete, or at the tile tail. + is_last_in_apply = const_expr(s % calls_per_apply == calls_per_apply - 1) + is_last_overall = const_expr(s == n_calls - 1) + if const_expr(is_last_in_apply or is_last_overall): + apply_idx = s // calls_per_apply + for c in cutlass.range_constexpr(mask_vals_per_apply): + chunk_idx = apply_idx * mask_vals_per_apply + c + # Skip packed chunks that start past the accumulator fragment. + if const_expr(chunk_idx * 32 < ncol): + apply_packed_mask_chunk(acc_S, chunk_idx, mask_vals[c]) + + @cute.jit + def apply_mask_sm100( + self, + acc_S: cute.Tensor, + m_block: Int32, + n_block: Int32, + thr_mma: cute.TiledMma, + thr_tmem_load: cute.TiledCopy, + mask_seqlen: cutlass.Constexpr[bool], + mask_causal: cutlass.Constexpr[bool], + mask_local: cutlass.Constexpr[bool] = False, + mask_mod: cutlass.Constexpr[Optional[Callable]] = None, + batch_idx: Int32 = None, + head_idx: Int32 = None, + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + head_divmod=None, + vec_size: cutlass.Constexpr[int] = 1, + check_q_boundary: bool = False, + r2p: bool = True, + rBitmask: Optional[cute.Tensor] = None, + ) -> None: + assert not (mask_causal and mask_local), "mask_causal and mask_local cannot be both True" + acc_shape = (self.tile_m, self.tile_n) + cS = cute.make_identity_tensor(acc_shape if not self.swap_AB else acc_shape[::-1]) + tScS = thr_mma.partition_C(cS) + tScS = tScS[(None, None), 0, 0] + tScS_t2r = thr_tmem_load.partition_D(tScS) + # To handle edge cases of completely masked out rows where n_block_max = 0, + # we treat negative n_blocks as 0th n_block + # TODO: find more transparent solution + if n_block < 0: + n_block = 0 + seqlenk_col_limit = self.seqlen_k - n_block * self.tile_n + + if const_expr(rBitmask is not None): + ncol_packed = const_expr(cute.size(rBitmask.shape[0])) + for i in cutlass.range_constexpr(ncol_packed): + col_start = 32 * i # mask is bit-packed into uint32 + curr_mask_val = rBitmask[i] + for j in cutlass.range_constexpr(32): + curr_col = col_start + j + mask = (curr_mask_val >> j) & 1 + acc_S[curr_col] = acc_S[curr_col] if cutlass.Boolean(mask) else -Float32.inf + + elif const_expr(not mask_causal and not mask_local and mask_mod is None): + if const_expr(mask_seqlen): + if const_expr(not r2p): + for i in cutlass.range(cute.size(tScS_t2r.shape), unroll_full=True): + # if tScS_t2r[i][1] >= seqlenk_col_limit: + # acc_S[i] = -Float32.inf + # For some reason the 2 lines above generate really bad SASS + acc_S[i] = -Float32.inf if tScS_t2r[i][1] >= seqlenk_col_limit else acc_S[i] + else: + mask_r2p_lambda( + acc_S, + lambda s: r2p_bitmask_below(seqlenk_col_limit, s), + rank1=True, + ) + + elif const_expr(not mask_causal and not mask_local and mask_mod is not None): + # FlexAttention mask_mod vectorization is gated on `mask_mod.__vec_size__`. + # vec_size == 1 returns a scalar Boolean. vec_size > 1 returns packed + # Uint32 mask fragments: one word per 32 evaluated columns. + assert vec_size % 32 == 0 or 32 % vec_size == 0, ( + "vec_size must divide 32 or be a multiple of 32" + ) + if const_expr(vec_size == 1): + self.apply_mask_mod_sm100_scalar( + acc_S, + tScS_t2r, + m_block, + n_block, + mask_seqlen, + mask_mod, + batch_idx, + head_idx, + aux_data, + fastdiv_mods, + head_divmod, + check_q_boundary, + ) + else: + self.apply_mask_mod_sm100_vector( + acc_S, + tScS_t2r, + m_block, + n_block, + mask_seqlen, + mask_mod, + batch_idx, + head_idx, + vec_size, + aux_data, + fastdiv_mods, + head_divmod, + check_q_boundary, + ) + + else: # Causal or local + causal_row_offset = self.seqlen_k - n_block * self.tile_n - self.seqlen_q + row_idx = tScS_t2r[0][0] + m_block * self.tile_m + if const_expr(self.qhead_per_kvhead_packgqa != 1): + row_idx = row_idx // self.qhead_per_kvhead_packgqa + if const_expr(mask_causal): + col_limit_right = row_idx + causal_row_offset + 1 + if const_expr(mask_seqlen): + col_limit_right = cutlass.min(col_limit_right, seqlenk_col_limit) + # if cute.arch.thread_idx()[0] % 32 == 0: + # cute.printf("tidx = %d, tidx tmem = %d, row_idx = %d, col_limit_right = %d, causal_row_offset = %d\n", cute.arch.thread_idx()[0], thr_tmem_load.thr_idx, row_idx, col_limit_right, causal_row_offset) + ncol = const_expr(cute.size(tScS_t2r.shape)) + if const_expr(not r2p): + for i in cutlass.range(ncol, unroll_full=True): + acc_S[i] = -Float32.inf if tScS_t2r[i][1] >= col_limit_right else acc_S[i] + else: + mask_r2p_lambda( + acc_S, + lambda s: r2p_bitmask_below(col_limit_right, s), + rank1=True, + ) + else: + local_row_offset_right = ( + causal_row_offset + 1 + self.window_size_right + if const_expr(self.window_size_right is not None) + else None + ) + local_row_offset_left = ( + causal_row_offset - self.window_size_left + if const_expr(self.window_size_left is not None) + else None + ) + if const_expr(self.window_size_right is not None): + col_limit_right = row_idx + local_row_offset_right + else: + col_limit_right = self.tile_n + if const_expr(mask_seqlen): + col_limit_right = cutlass.min(col_limit_right, seqlenk_col_limit) + col_limit_left = ( + row_idx + local_row_offset_left + if const_expr(self.window_size_left is not None) + else 0 + ) + if const_expr(not r2p): + # if cute.arch.thread_idx()[0] == 0 or cute.arch.thread_idx()[0] == 128: cute.printf("m_block = {}, n_block = {}, row_idx = {}, causal_row_offset = {}, col_limit_right = {}, col_limit_left = {}", m_block, n_block, row_idx, causal_row_offset, col_limit_right, col_limit_left) + for i in cutlass.range(cute.size(tScS_t2r.shape), unroll_full=True): + col_idx = tScS_t2r[i][1] + acc_S[i] = ( + -Float32.inf + if col_idx >= col_limit_right or col_idx < col_limit_left + else acc_S[i] + ) + else: + # Dual-bound R2P masking for SM100. + # Masks elements where: NOT (col_limit_left <= col < col_limit_right) + + def mask_gen_fn(s: int) -> Uint32: + return r2p_bitmask_below(col_limit_right, s) & r2p_bitmask_above( + col_limit_left, s + ) + + mask_r2p_lambda(acc_S, mask_gen_fn, rank1=True) + + @cute.jit + def apply_mask_sm100_transposed( + self, + acc_S: cute.Tensor, + tScS_t2r: cute.Tensor, + t0ScS_t2r: cute.Tensor, + m_block: cutlass.Int32, + n_block: cutlass.Int32, + mask_seqlen: cutlass.Constexpr, + mask_causal: cutlass.Constexpr, + mask_local: cutlass.Constexpr, + mask_mod: cutlass.Constexpr[Optional[Callable]] = None, + batch_idx: Int32 = None, + head_idx: Int32 = None, + aux_data: AuxData = AuxData(), + fastdiv_mods=(None, None), + is_full_block: bool = False, + check_m_boundary: bool = True, + ) -> None: + """ + Backward pass: mask S = K @ Q.T where n_block tiles seqlen_k and m_block tiles seqlen_q. + + Coordinate convention: + - ROW corresponds to Q (m_block) + - COL corresponds to KV (n_block) + + is_full_block: If True, skip mask_mod (all elements valid). Only apply seqlen masking. + check_m_boundary: If False, skip seqlen_q boundary check (optimization for non-boundary m_blocks). + When iterating m_blocks in forward order, only the last m_block may be partial. + """ + assert not (mask_causal and mask_local), "mask_causal and mask_local cannot be both True" + ROW = 0 if const_expr(not self.swap_AB) else 1 + COL = 1 if const_expr(not self.swap_AB) else 0 + # assert t0ScS_t2r[0][COL] == 0, "col0 == 0" # tmp comment for 2-cta bwd + thr_col_offset = tScS_t2r[0][COL] + seqlenk_col_limit = self.seqlen_k - n_block * self.tile_n - thr_col_offset + + if const_expr(not mask_causal and not mask_local and mask_mod is not None): + # Block sparse case with mask_mod (backward) + # + # Coordinate convention: ROW → Q (m_block), COL → KV (n_block). + # These already account for swap_AB. + # + # FULL blocks: mask_mod returns True for all elements, so skip it. + # Still need seqlen bounds check (elements may be OOB on last m_block). + # PARTIAL blocks: apply mask_mod element-wise, then seqlen bounds. + if is_full_block: + if const_expr(mask_seqlen): + if seqlenk_col_limit <= 0: + # Entire tile is OOB for K + for i in cutlass.range(cute.size(acc_S.shape), unroll_full=True): + acc_S[i] = -cutlass.Float32.inf + elif check_m_boundary: + # Last m_block: check Q and K boundaries + ncol = const_expr(cute.size(tScS_t2r.shape)) + for i in cutlass.range_constexpr(ncol): + row_coord = tScS_t2r[i][ROW] + col_coord = tScS_t2r[i][COL] + global_q = row_coord + m_block * self.tile_m + global_kv = col_coord + n_block * self.tile_n + q_out_of_bounds = global_q >= self.seqlen_q + kv_out_of_bounds = global_kv >= self.seqlen_k + out_of_bounds = q_out_of_bounds or kv_out_of_bounds + acc_S[i] = -cutlass.Float32.inf if out_of_bounds else acc_S[i] + else: + # Partial block + has_fastdiv = const_expr( + fastdiv_mods is not None + and fastdiv_mods[0] is not None + and fastdiv_mods[1] is not None + ) + wrap_aux_indices = const_expr( + has_fastdiv and mask_seqlen and const_expr(aux_data.tensors is not None) + ) + batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32) + head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32) + + ncol = const_expr(cute.size(tScS_t2r.shape)) + for i in cutlass.range_constexpr(ncol): + row_coord = tScS_t2r[i][ROW] + col_coord = tScS_t2r[i][COL] + global_q = row_coord + m_block * self.tile_m + global_kv = col_coord + n_block * self.tile_n + + q_idx_for_mod = global_q + kv_idx_for_mod = global_kv + if const_expr(wrap_aux_indices): + _, q_idx_for_mod = divmod(global_q, fastdiv_mods[0]) + _, kv_idx_for_mod = divmod(global_kv, fastdiv_mods[1]) + + q_idx_ssa = utils.scalar_to_ssa(q_idx_for_mod, cutlass.Int32) + kv_idx_ssa = utils.scalar_to_ssa(kv_idx_for_mod, cutlass.Int32) + + mask_value = call_mask_mod( + mask_mod, + batch_idx_ssa, + head_idx_ssa, + q_idx_ssa, + kv_idx_ssa, + self.seqlen_info, + aux_data, + ) + cond = cutlass.Boolean(utils.ssa_to_scalar(mask_value)) + acc_S[i] = acc_S[i] if cond else -cutlass.Float32.inf + + if const_expr(mask_seqlen): + # check_m_boundary=False skips q check for non-boundary m_blocks + q_out_of_bounds = check_m_boundary and (global_q >= self.seqlen_q) + kv_out_of_bounds = global_kv >= self.seqlen_k + out_of_bounds = q_out_of_bounds or kv_out_of_bounds + acc_S[i] = -cutlass.Float32.inf if out_of_bounds else acc_S[i] + + elif const_expr(not mask_causal and not mask_local): + if const_expr(mask_seqlen): + if seqlenk_col_limit <= 0: + for i in cutlass.range(cute.size(acc_S.shape), unroll_full=True): + acc_S[i] = -cutlass.Float32.inf + else: # Causal or local + thr_row_offset = tScS_t2r[0][ROW] + seqlenq_row_limit = self.seqlen_q - m_block * self.tile_m - thr_row_offset + causal_offset = seqlenq_row_limit - seqlenk_col_limit + if const_expr(mask_causal): + # tidx = cute.arch.thread_idx()[0] % 256 + # if tidx < 32: + # cute.printf("tidx = {}, {} {}, {} {}", tidx, tScS_t2r[0][0], tScS_t2r[0][1], tScS_t2r[1][0], tScS_t2r[1][1]) + row_limit_top = causal_offset + if const_expr(mask_seqlen): + # If col is beyond the column limit, we want to mask out the entire + # column, by setting row limit to be self.tile_m. + if seqlenk_col_limit <= 0: + row_limit_top = self.tile_m + r2p = True + if const_expr(not r2p): + for i in cutlass.range(cute.size(acc_S.shape), unroll_full=True): + acc_S[i] = ( + -cutlass.Float32.inf if t0ScS_t2r[i][ROW] < row_limit_top else acc_S[i] + ) + else: + num_rep = cute.size(tScS_t2r, mode=[0]) # 16 or 32 + num_wg = 2 + row_limit = row_to_r2p_idx(row_limit_top, num_rep, num_wg) + mask_r2p_lambda( + acc_S, + lambda s: r2p_bitmask_above(row_limit, s), + rank1=True, + ) + else: + if const_expr(self.window_size_right is not None): + row_limit_top = causal_offset - self.window_size_right + else: + row_limit_top = 0 + if const_expr(self.window_size_left is not None): + row_limit_bot = causal_offset + self.window_size_left + if const_expr(mask_seqlen): + if seqlenk_col_limit <= 0: + row_limit_top = self.tile_m + r2p = True + if const_expr(not r2p): + for i in cutlass.range(cute.size(acc_S.shape), unroll_full=True): + row_idx = t0ScS_t2r[i][ROW] + local_mask = row_idx < row_limit_top + if const_expr(self.window_size_left is not None): + local_mask |= row_idx > row_limit_bot + acc_S[i] = -cutlass.Float32.inf if local_mask else acc_S[i] + else: + + def mask_gen_fn(s: int) -> Uint32: + num_rep = cute.size(tScS_t2r, mode=[0]) + num_wg = 2 + + row_limit = row_to_r2p_idx(row_limit_top, num_rep, num_wg) + mask = r2p_bitmask_above(row_limit, s) + + if const_expr(self.window_size_left is not None): + row_limit_bottom = row_to_r2p_idx(row_limit_bot + 1, num_rep, num_wg) + mask = mask & r2p_bitmask_below(row_limit_bottom, s) + + return mask + + mask_r2p_lambda( + acc_S, + mask_gen_fn, + rank1=True, + ) + + +# ----------------------------------------------------------------------------- +# SM100 FMHA fused-mask policy layer (separate from generic mask primitives). +# ----------------------------------------------------------------------------- + + +class Sm100MaskEnum(enum.Enum): + """Enumeration of mask types for FMHA operations. + + - RESIDUAL_MASK: Residual mask for handling variable sequence lengths + - WINDOW_MASK: Window mask for attention which also includes causal and no mask + - WINDOW_MASK_INFERENCE: Same as the window mask, but has the limitation that the end of q is aligned with the end of k + - WINDOW_MASK_BWD: Window mask for backward pass + - WINDOW_MASK_BWD_INFERENCE: Same as the window mask for backward pass, but has the limitation that the end of q is aligned with the end of k + """ + + NO_MASK = enum.auto() + RESIDUAL_MASK = enum.auto() + CAUSAL_MASK = enum.auto() + WINDOW_MASK = enum.auto() + WINDOW_MASK_INFERENCE = enum.auto() + # Deprecated the following types + WINDOW_MASK_BWD = enum.auto() + WINDOW_MASK_BWD_INFERENCE = enum.auto() + RESIDUAL_MASK_BWD = enum.auto() + + +class Sm100FusedMask: + """A fused mask implementation for FMHA operations. + + This class handles different types of attention masks including no mask, + residual mask for variable sequence lengths, and causal mask for + autoregressive attention patterns. + + The class provides methods to: + - Calculate trip counts for different mask types + - Apply masks to attention scores + - Handle masked and unmasked trip calculations + """ + + def get_trip_count( + mask_type: Sm100MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Calculate the number of trips needed for the current block. + + The trip count depends on the mask type and the block coordinates. + For causal masks, it considers the autoregressive constraint. + + :param mask_type: Type of mask to use + :type mask_type: utils.Sm100MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Number of trips needed. + :rtype: Int32 + """ + result = 0 + offset = 0 + if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + if cutlass.const_expr(mask_type == Sm100MaskEnum.RESIDUAL_MASK): + result = cute.ceil_div(seqlen_k, tile_shape[1]) + if cutlass.const_expr(mask_type is Sm100MaskEnum.RESIDUAL_MASK_BWD): + result = cute.ceil_div(seqlen_q, tile_shape[0]) + if cutlass.const_expr( + mask_type == Sm100MaskEnum.WINDOW_MASK + or mask_type == Sm100MaskEnum.WINDOW_MASK_INFERENCE + ): + if cutlass.const_expr(window_size_right is None): + result = cute.ceil_div(seqlen_k, tile_shape[1]) + else: + max_idx_q = (blk_coord[0] + 1) * tile_shape[0] + idx_k = max_idx_q + offset + window_size_right + tmp_blocks_k = cute.ceil_div(idx_k, tile_shape[1]) + max_blocks_k = cute.ceil_div(seqlen_k, tile_shape[1]) + result = dsl_min(max_blocks_k, tmp_blocks_k) + if cutlass.const_expr( + mask_type == Sm100MaskEnum.WINDOW_MASK_BWD + or mask_type == Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE + ): + if cutlass.const_expr(window_size_left is None): + result = cute.ceil_div(seqlen_q, tile_shape[0]) + else: + max_idx_k = (blk_coord[1] + 1) * tile_shape[1] + idx_k = max_idx_k + offset + window_size_left + tmp_blocks_q = cute.ceil_div(idx_k, tile_shape[0]) + max_blocks_q = cute.ceil_div(seqlen_q, tile_shape[0]) + result = dsl_min(max_blocks_q, tmp_blocks_q) + start_block = Sm100FusedMask.get_trip_start( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + result = result - start_block + return result + + @cute.jit + def get_trip_start_count_via_block_info( + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + is_causal: cutlass.Constexpr[bool] = False, + is_local: cutlass.Constexpr[bool] = False, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Tuple[Int32, Int32]: + block_info = BlockInfo( + tile_m=tile_shape[0], + tile_n=tile_shape[1], + is_causal=is_causal, + is_local=is_local and not is_causal, + window_size_left=window_size_left, + window_size_right=window_size_right, + ) + + seqlen_info = SeqlenInfoQK( + offset_q=Int32(0), + offset_k=Int32(0), + padded_offset_q=Int32(0), + padded_offset_k=Int32(0), + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + m_block_offset=Int32(0), + block_idx_offset=Int32(0), + num_n_blocks=cute.ceil_div(seqlen_k, tile_shape[1]), + has_cu_seqlens_q=False, + has_cu_seqlens_k=False, + has_seqused_q=False, + has_seqused_k=False, + ) + n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen_info, blk_coord[0]) + return n_block_min, n_block_max - n_block_min + + @cute.jit + def get_trip_mask_bounds_via_block_info( + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + is_causal: cutlass.Constexpr[bool] = False, + is_local: cutlass.Constexpr[bool] = False, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Tuple[Int32, Int32]: + """Return SM100-style mask boundaries for dense iteration. + + Returns: + - n_block_min_causal_local_mask: right-side masked region start + - n_block_min_before_local_mask: start of fully unmasked middle region + """ + block_info = BlockInfo( + tile_m=tile_shape[0], + tile_n=tile_shape[1], + is_causal=is_causal, + is_local=is_local and not is_causal, + window_size_left=window_size_left, + window_size_right=window_size_right, + ) + seqlen_info = SeqlenInfoQK( + offset_q=Int32(0), + offset_k=Int32(0), + padded_offset_q=Int32(0), + padded_offset_k=Int32(0), + seqlen_q=seqlen_q, + seqlen_k=seqlen_k, + m_block_offset=Int32(0), + block_idx_offset=Int32(0), + num_n_blocks=cute.ceil_div(seqlen_k, tile_shape[1]), + has_cu_seqlens_q=False, + has_cu_seqlens_k=False, + has_seqused_q=False, + has_seqused_k=False, + ) + n_block_min, _ = block_info.get_n_block_min_max(seqlen_info, blk_coord[0]) + n_block_min_causal_local_mask = block_info.get_n_block_min_causal_local_mask( + seqlen_info, blk_coord[0], n_block_min + ) + n_block_min_before_local_mask = block_info.get_n_block_min_before_local_mask( + seqlen_info, blk_coord[0], n_block_min + ) + return n_block_min_causal_local_mask, n_block_min_before_local_mask + + @cute.jit + def get_trip_start( + mask_type: Sm100MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Get the start of the trip for the current block. + + :param mask_type: Type of mask to use + :type mask_type: utils.Sm100MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + """ + result = 0 + offset = 0 + if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + if cutlass.const_expr( + mask_type is Sm100MaskEnum.WINDOW_MASK + or mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE + ): + if cutlass.const_expr(window_size_left is not None): + min_idx_q = blk_coord[0] * tile_shape[0] + idx_k = min_idx_q + offset - window_size_left + tmp_blocks_k = idx_k // tile_shape[1] + result = max(tmp_blocks_k, result) + if cutlass.const_expr( + mask_type is Sm100MaskEnum.WINDOW_MASK_BWD + or mask_type is Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE + ): + if cutlass.const_expr(window_size_right is not None): + min_idx_k = blk_coord[1] * tile_shape[1] + idx_q = min_idx_k + offset - window_size_right + tmp_blocks_q = idx_q // tile_shape[0] + result = max(tmp_blocks_q, result) + return result + + @cute.jit + def get_leading_mask_id( + mask_type: Sm100MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Tuple[Int32, Int32]: + """ + Get the begin and end tile idx for the leading mask. + + :param mask_type: Type of mask to use + :type mask_type: utils.Sm100MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Tuple of (begin, end) tile idx for the leading mask. + :rtype: Tuple[Int32, Int32] + """ + offset = 0 + if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + leading_mask_begin = Sm100FusedMask.get_trip_start( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + trip_count = Sm100FusedMask.get_trip_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + + leading_mask_end = leading_mask_begin + if cutlass.const_expr( + mask_type is Sm100MaskEnum.WINDOW_MASK + or mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE + ): + if cutlass.const_expr(window_size_left is not None): + min_idx_q = (blk_coord[0] + 1) * tile_shape[0] + offset - window_size_left + leading_mask_end = dsl_min( + cute.ceil_div(min_idx_q, tile_shape[1]) - 1, + trip_count + leading_mask_begin - 1, + ) + else: + leading_mask_end = leading_mask_begin - 1 + elif cutlass.const_expr( + mask_type is Sm100MaskEnum.WINDOW_MASK_BWD + or mask_type is Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE + ): + if cutlass.const_expr(window_size_right is not None): + min_idx_k = (blk_coord[1] + 1) * tile_shape[1] + offset - window_size_right + leading_mask_end = cute.ceil_div(min_idx_k, tile_shape[0]) - 1 + else: + leading_mask_end = leading_mask_begin - 1 + return leading_mask_begin, leading_mask_end + + @cute.jit + def get_trailing_mask_id( + mask_type: Sm100MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Tuple[Optional[Int32], Optional[Int32]]: + """ + Get the begin and end tile idx for the trailing mask. + + :param mask_type: Type of mask to use + :type mask_type: utils.Sm100MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Tuple of (begin, end) tile idx for the trailing mask. + :rtype: Tuple[Int32, Int32] + """ + offset = 0 + if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE): + offset = seqlen_k - seqlen_q + if cutlass.const_expr(mask_type is Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE): + offset = seqlen_q - seqlen_k + trip_start = Sm100FusedMask.get_trip_start( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + trip_count = Sm100FusedMask.get_trip_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + + trailing_mask_begin, trailing_mask_end = None, None + if cutlass.const_expr( + mask_type is Sm100MaskEnum.WINDOW_MASK + or mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE + ): + if cutlass.const_expr(window_size_right is not None): + min_idx_q = blk_coord[0] * tile_shape[0] + offset + window_size_right + trailing_mask_begin = dsl_min( + min_idx_q // tile_shape[1], trip_count + trip_start - 1 + ) + trailing_mask_end = trip_count + trip_start - 1 + else: + # last tile, we always apply mask on it regardless whether it's a residual tile + trailing_mask_begin = trip_count + trip_start - 1 + trailing_mask_end = trip_count + trip_start - 1 + else: + if cutlass.const_expr(window_size_left is not None): + min_idx_k = blk_coord[1] * tile_shape[1] + offset + window_size_left + 1 + max_idx_k = (blk_coord[1] + 1) * tile_shape[1] + offset + window_size_left + trailing_mask_begin = dsl_min( + cute.ceil_div(min_idx_k, tile_shape[0]) - 1, + trip_count + trip_start - 1, + ) + trailing_mask_end = dsl_min( + cute.ceil_div(max_idx_k, tile_shape[0]) - 1, + trip_count + trip_start - 1, + ) + else: + # last tile, we always apply mask on it regardless whether it's a residual tile + trailing_mask_begin = trip_count + trip_start - 1 + trailing_mask_end = trip_count + trip_start - 1 + + return trailing_mask_begin, trailing_mask_end + + @cute.jit + def get_masked_leading_count( + mask_type: Sm100MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Calculate the number of masked trips for the leading mask. + + This is used for blocks that need special handling due to masking. + + :param mask_type: Type of mask to use + :type mask_type: utils.Sm100MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Number of masked trips. + :rtype: Int32 + """ + result = 0 + if cutlass.const_expr( + mask_type is not Sm100MaskEnum.RESIDUAL_MASK + and mask_type is not Sm100MaskEnum.RESIDUAL_MASK_BWD + ): + if cutlass.const_expr(window_size_left is not None or window_size_right is not None): + leading_mask_begin, leading_mask_end = Sm100FusedMask.get_leading_mask_id( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + result = max(leading_mask_end - leading_mask_begin + 1, 0) + + return result + + @cute.jit + def get_masked_trailing_count( + mask_type: Sm100MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + rem_count: Optional[Int32] = 0, + ) -> Int32: + """ + Calculate the number of masked trips for the trailing mask. + + This is used for blocks that need special handling due to masking. + + :param mask_type: Type of mask to use + :type mask_type: utils.Sm100MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + :param rem_count: Remaining count from previous calculations. + :type rem_count: Int32 + + :return: Number of masked trips. + :rtype: Int32 + """ + result = 0 + + if cutlass.const_expr( + mask_type is not Sm100MaskEnum.RESIDUAL_MASK + and mask_type is not Sm100MaskEnum.RESIDUAL_MASK_BWD + ): + if cutlass.const_expr(window_size_left is not None or window_size_right is not None): + trailing_mask_begin, trailing_mask_end = Sm100FusedMask.get_trailing_mask_id( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + leading_mask_begin, leading_mask_end = Sm100FusedMask.get_leading_mask_id( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + if cutlass.const_expr( + trailing_mask_begin is not None and trailing_mask_end is not None + ): + if trailing_mask_begin <= leading_mask_end: + result = max(trailing_mask_end - leading_mask_end, 0) + else: + result = max(trailing_mask_end - trailing_mask_begin + 1, 0) + else: + if seqlen_k % tile_shape[1] != 0: + result = 1 + else: + result = 0 + + return result + rem_count + + @cute.jit + def get_unmasked_trip_count( + mask_type: Sm100MaskEnum, + blk_coord: cute.Coord, + tile_shape: cute.Shape, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + ) -> Int32: + """ + Calculate the number of unmasked trips for the current block. + + This represents the number of trips that don't require special + masking treatment. + + :param mask_type: Type of mask to use + :type mask_type: utils.Sm100MaskEnum + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param tile_shape: Shape of the tile. + :type tile_shape: cute.Shape + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Int32 + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[Int32] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[Int32] + + :return: Number of unmasked trips. + :rtype: Int32 + """ + result = ( + Sm100FusedMask.get_trip_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + - Sm100FusedMask.get_masked_leading_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + ) + - Sm100FusedMask.get_masked_trailing_count( + mask_type, + blk_coord, + tile_shape, + seqlen_q, + seqlen_k, + window_size_left, + window_size_right, + 0, + ) + ) + return result + + @cute.jit + def apply_mask( + mask_type: Sm100MaskEnum, + acc_qk: cute.Tensor, + index_qk: cute.Tensor, + seqlen_q: Int32, + seqlen_k: Int32, + window_size_left: Optional[int] = None, + window_size_right: Optional[int] = None, + index_transform: cutlass.Constexpr = lambda index_q, index_k: ( + index_q, + index_k, + ), + ): + """ + Apply the appropriate mask to the attention scores. + + This method modifies the attention scores (acc_qk) based on the mask type + and the positions in the index tensor. + + :param mask_type: Type of mask to use + :type mask_type: utils.Sm100MaskEnum + :param acc_qk: Accumulated QK attention scores tensor. + :type acc_qk: cute.Tensor + :param index_qk: Index tensor containing position information. + :type index_qk: cute.Tensor + :param seqlen_k: Key sequence length for attention computation. + :type seqlen_k: Int32 + :param seqlen_q: Query sequence length for attention computation. + :type seqlen_q: Optional[int] + :param window_size_left: Left-side sliding window size for attention masking. + :type window_size_left: Optional[int] + :param window_size_right: Right-side sliding window size for attention masking. + :type window_size_right: Optional[int] + """ + offset = 0 + # NOTE: causal masking in this repo aligns the *end* of Q with the *end* of K + # when seqlen_k != seqlen_q (same as the test/reference implementation): + # k_index <= q_index + (seqlen_k - seqlen_q) + window_right + # In our kernels, causal is represented by (window_left is None, window_right is not None). + if cutlass.const_expr(window_size_left is None and window_size_right is not None): + offset = seqlen_k - seqlen_q + elif cutlass.const_expr( + mask_type is Sm100MaskEnum.WINDOW_MASK_INFERENCE + or mask_type is Sm100MaskEnum.WINDOW_MASK_BWD_INFERENCE + ): + offset = seqlen_k - seqlen_q + for i in cutlass.range_constexpr(cute.size(acc_qk), unroll_full=True): + index_q, index_k = index_transform(*index_qk[i]) + if cutlass.const_expr(window_size_left is not None or window_size_right is not None): + if cutlass.const_expr(window_size_left is None): + if index_q + offset + window_size_right < index_k: + acc_qk[i] = -Float32.inf + if index_k >= seqlen_k or index_q >= seqlen_q: # residual mask + acc_qk[i] = -Float32.inf + elif cutlass.const_expr(window_size_right is None): + if index_q + offset - window_size_left > index_k: + acc_qk[i] = -Float32.inf + if index_k >= seqlen_k or index_q >= seqlen_q: # residual mask + acc_qk[i] = -Float32.inf + else: + max_K_index = dsl_min(index_q + offset + window_size_right, seqlen_k) + min_K_index = max(0, index_q + offset - window_size_left) + if index_k > max_K_index or index_k < min_K_index: + acc_qk[i] = -Float32.inf + if index_k >= seqlen_k or index_q >= seqlen_q: # residual mask + acc_qk[i] = -Float32.inf + + if cutlass.const_expr( + mask_type == Sm100MaskEnum.RESIDUAL_MASK + or mask_type == Sm100MaskEnum.RESIDUAL_MASK_BWD + ): + if index_k >= seqlen_k or index_q >= seqlen_q: + acc_qk[i] = -Float32.inf + + @cute.jit + def apply_mask_via_causal_local( + acc_qk: cute.Tensor, + index_qk: cute.Tensor, + seqlen_q: Int32, + seqlen_k: Int32, + apply_semantic_window: cutlass.Constexpr[bool] = True, + is_causal: cutlass.Constexpr[bool] = False, + is_local: cutlass.Constexpr[bool] = False, + window_size_left: Optional[int] = None, + window_size_right: Optional[int] = None, + index_transform: cutlass.Constexpr = lambda index_q, index_k: ( + index_q, + index_k, + ), + ): + """Apply forward mask without mask_type. + + - If apply_semantic_window=True, apply causal/local window constraints. + - Always apply residual OOB masking (index_k>=seqlen_k or index_q>=seqlen_q). + """ + offset = 0 + if cutlass.const_expr(apply_semantic_window): + # Match WINDOW_MASK_INFERENCE semantics: end-align Q/K when lengths differ. + offset = seqlen_k - seqlen_q + for i in cutlass.range_constexpr(cute.size(acc_qk), unroll_full=True): + index_q, index_k = index_transform(*index_qk[i]) + if cutlass.const_expr(apply_semantic_window): + if cutlass.const_expr(is_causal and not is_local): + # Pure causal; tolerate both external forms: + # - (None, None) from interface + # - (None, 0) from fused-mask-style callers + right = 0 if const_expr(window_size_right is None) else window_size_right + if index_q + offset + right < index_k: + acc_qk[i] = -Float32.inf + elif cutlass.const_expr( + is_local or window_size_left is not None or window_size_right is not None + ): + if cutlass.const_expr(window_size_left is None): + if index_q + offset + window_size_right < index_k: + acc_qk[i] = -Float32.inf + elif cutlass.const_expr(window_size_right is None): + if index_q + offset - window_size_left > index_k: + acc_qk[i] = -Float32.inf + else: + max_K_index = dsl_min(index_q + offset + window_size_right, seqlen_k) + min_K_index = max(0, index_q + offset - window_size_left) + if index_k > max_K_index or index_k < min_K_index: + acc_qk[i] = -Float32.inf + # Residual mask is always needed for boundary protection. + if index_k >= seqlen_k or index_q >= seqlen_q: + acc_qk[i] = -Float32.inf diff --git a/fastvideo-kernel/fa4/mma_sm100_desc.py b/fastvideo-kernel/fa4/mma_sm100_desc.py new file mode 100644 index 000000000..9a2adfca8 --- /dev/null +++ b/fastvideo-kernel/fa4/mma_sm100_desc.py @@ -0,0 +1,296 @@ +# Copyright (c) 2025, Tri Dao. +# Ported Cutlass code from C++ to Python: +# https://github.com/NVIDIA/cutlass/blob/main/include/cute/arch/mma_sm100_desc.hpp +# https://github.com/NVIDIA/cutlass/blob/main/include/cute/atom/mma_traits_sm100.hpp + +from enum import IntEnum + +import cutlass +import cutlass.cute as cute + +# --------------------------------------------------------------------------- +# Enumerations that match the HW encodings (values MUST stay identical) +# --------------------------------------------------------------------------- + + +class Major(IntEnum): # matrix “layout” in the ISA docs + K = 0 + MN = 1 + + +class ScaleIn(IntEnum): # negate flags + One = 0 + Neg = 1 + + +class Saturate(IntEnum): + False_ = 0 + True_ = 1 + + +class CFormat(IntEnum): # 2-bit field (bits 4-5) + F16 = 0 + F32 = 1 + S32 = 2 + + +class F16F32Format(IntEnum): # 3-bit field (A/B element type) + F16 = 0 + BF16 = 1 + TF32 = 2 + + +class S8Format(IntEnum): + UINT8 = 0 + INT8 = 1 + + +class MXF8F6F4Format(IntEnum): + E4M3 = 0 + E5M2 = 1 + E2M3 = 3 + E3M2 = 4 + E2M1 = 5 + + +class MaxShift(IntEnum): + NoShift = 0 + MaxShift8 = 1 + MaxShift16 = 2 + MaxShift32 = 3 + + +# --------------------------------------------------------------------------- +# CUTLASS-type → encoding helpers +# --------------------------------------------------------------------------- + + +def to_UMMA_format(cutlass_type) -> int: + """ + Map a CUTLASS scalar class to the 3-bit encoding for Matrix A/B. + """ + if cutlass_type is cutlass.Int8: + return S8Format.INT8 + # Unsigned 8-bit (if available in your CUTLASS build) + if cutlass_type is cutlass.Uint8: + return S8Format.UINT8 + # FP-16 / BF-16 + if cutlass_type is cutlass.Float16: + return F16F32Format.F16 + if cutlass_type is cutlass.BFloat16: + return F16F32Format.BF16 + # TensorFloat-32 (8-bit exponent, 10-bit mantissa packed in 19 bits) + if cutlass_type is cutlass.TFloat32: + return F16F32Format.TF32 + # Float-8 / Float-6 / Float-4 – add whenever CUTLASS exposes them + if cutlass_type is cutlass.Float8E4M3FN: + return MXF8F6F4Format.E4M3 + if cutlass_type is cutlass.Float8E5M2: + return MXF8F6F4Format.E5M2 + raise TypeError(f"Unsupported CUTLASS scalar type for A/B: {cutlass_type!r}") + + +def to_C_format(cutlass_type) -> int: + """ + Map a CUTLASS scalar class to the 2-bit accumulator encoding. + """ + if cutlass_type is cutlass.Float16: + return CFormat.F16 + if cutlass_type is cutlass.Float32: + return CFormat.F32 + if cutlass_type is cutlass.Int32: + return CFormat.S32 + raise TypeError(f"Unsupported CUTLASS scalar type for accumulator: {cutlass_type!r}") + + +# --------------------------------------------------------------------------- +# The constructor – accepts only CUTLASS scalar classes +# --------------------------------------------------------------------------- + + +def make_instr_desc( + a_type, # CUTLASS scalar class, e.g. cutlass.Int8 + b_type, + c_type, + M: int, # 64, 128 or 256 + N: int, # 8 … 256 (multiple of 8) + a_major: Major, + b_major: Major, + a_neg: ScaleIn = ScaleIn.One, + b_neg: ScaleIn = ScaleIn.One, + c_sat: Saturate = Saturate.False_, + is_sparse: bool = False, + max_shift: MaxShift = MaxShift.NoShift, +) -> int: + """ + Build the 32-bit instruction descriptor for Blackwell MMA. + All matrix/accumulator **types must be CUTLASS scalar classes** – + passing integers is forbidden. + """ + # --- encode element formats ------------------------------------------------- + a_fmt = int(to_UMMA_format(a_type)) + b_fmt = int(to_UMMA_format(b_type)) + c_fmt = int(to_C_format(c_type)) + + # --- range checks on M/N ----------------------------------------------------- + if M not in (64, 128, 256): + raise ValueError("M must be 64, 128 or 256") + if N < 8 or N > 256 or (N & 7): + raise ValueError("N must be a multiple of 8 in the range 8…256") + + m_dim = M >> 4 # 5-bit field + n_dim = N >> 3 # 6-bit field + + # fmt: off + # --- pack the bit-fields ----------------------------------------------------- + desc = 0 + desc |= (0 & 0x3) << 0 # sparse_id2 (always 0 here) + desc |= (int(is_sparse) & 0x1) << 2 # sparse_flag + desc |= (int(c_sat) & 0x1) << 3 # saturate + desc |= (c_fmt & 0x3) << 4 # c_format + desc |= (a_fmt & 0x7) << 7 # a_format + desc |= (b_fmt & 0x7) << 10 # b_format + desc |= (int(a_neg) & 0x1) << 13 # a_negate + desc |= (int(b_neg) & 0x1) << 14 # b_negate + desc |= (int(a_major) & 0x1) << 15 # a_major + desc |= (int(b_major) & 0x1) << 16 # b_major + desc |= (n_dim & 0x3F) << 17 # n_dim (6 bits) + desc |= (m_dim & 0x1F) << 24 # m_dim (5 bits) + desc |= (int(max_shift) & 0x3) << 30 # max_shift (2 bits) + # fmt: on + + return desc & 0xFFFF_FFFF # ensure 32-bit result + + +def mma_op_to_idesc(op: cute.nvgpu.tcgen05.mma.MmaOp): + return make_instr_desc( + op.a_dtype, + op.b_dtype, + op.acc_dtype, + op.shape_mnk[0], + op.shape_mnk[1], + Major.K if op.a_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K else Major.MN, + Major.K if op.b_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K else Major.MN, + ) + + +class LayoutType(IntEnum): # occupies the top-3 bits [61:64) + SWIZZLE_NONE = 0 # (a.k.a. “INTERLEAVE” in older docs) + SWIZZLE_128B_BASE32B = 1 + SWIZZLE_128B = 2 + SWIZZLE_64B = 4 + SWIZZLE_32B = 6 + # values 3,5,7 are reserved / illegal for UMMA + + +# --------------------------------------------------------------------------- +# Helpers – figure out the SWIZZLE_* family from the tensor layout +# --------------------------------------------------------------------------- + + +def _layout_type(swizzle: cute.Swizzle) -> LayoutType: + B, M, S = swizzle.num_bits, swizzle.num_base, swizzle.num_shift + + if M == 4: # Swizzle<*,4,3> + if S != 3: + raise ValueError("Unexpected swizzle shift – want S==3 for M==4") + return { + 0: LayoutType.SWIZZLE_NONE, + 1: LayoutType.SWIZZLE_32B, + 2: LayoutType.SWIZZLE_64B, + 3: LayoutType.SWIZZLE_128B, + }[B] # KeyError ⇒ invalid B→ raise + if M == 5: # Swizzle<2,5,2> (the only legal triple for M==5) + if (B, S) != (2, 2): + raise ValueError("Only Swizzle<2,5,2> supported for 128B_BASE32B") + return LayoutType.SWIZZLE_128B_BASE32B + + # Any other (M,B,S) triple is not a UMMA-legal shared-memory layout + raise ValueError("Unsupported swizzle triple for UMMA smem descriptor") + + +def make_smem_desc_base(layout: cute.Layout, swizzle: cute.Swizzle, major: Major) -> int: + """ + Convert a 2-D *shared-memory* Cute layout into the Blackwell 64-bit + smem-descriptor, without the smem start address. + layout must correspond to layout of an uint128 tensor. + """ + # ------------------------------------------------------------------ meta + layout_type = _layout_type(swizzle) # resolve SWIZZLE_* family + + VERSION = 1 # bits 46–47 + LBO_MODE = 0 # bit 52 + BASE_OFFSET = 0 # bits 49–51 (CUTLASS always 0) + + # ---------------------------------------------------------- strides (units: uint128_t = 16 B) + swizzle_atom_mn_size = { + LayoutType.SWIZZLE_NONE: 1, + LayoutType.SWIZZLE_32B: 2, + LayoutType.SWIZZLE_64B: 4, + LayoutType.SWIZZLE_128B: 8, + LayoutType.SWIZZLE_128B_BASE32B: 8, + }[layout_type] + + if major is Major.MN: + swizzle_atom_k_size = 4 if layout_type is LayoutType.SWIZZLE_128B_BASE32B else 8 + canonical_layout = cute.logical_divide(layout, (swizzle_atom_mn_size, swizzle_atom_k_size)) + if not cute.is_congruent(canonical_layout, ((1, 1), (1, 1))): + raise ValueError("Not a canonical UMMA_MN Layout: Expected profile failure.") + stride_00 = canonical_layout.stride[0][0] + if layout_type is not LayoutType.SWIZZLE_NONE and stride_00 != 1: + raise ValueError("Not a canonical UMMA_MN Layout: Expected stride failure.") + stride_10 = canonical_layout.stride[1][0] + if stride_10 != swizzle_atom_mn_size: + raise ValueError("Not a canonical UMMA_MN Layout: Expected stride failure.") + stride_01, stride_11 = canonical_layout.stride[0][1], canonical_layout.stride[1][1] + if layout_type is LayoutType.SWIZZLE_NONE: + stride_byte_offset, leading_byte_offset = stride_01, stride_11 + else: + stride_byte_offset, leading_byte_offset = stride_11, stride_01 + else: + if layout_type == LayoutType.SWIZZLE_128B_BASE32B: + raise ValueError("SWIZZLE_128B_BASE32B is invalid for Major-K") + if not cute.size(layout.shape[0]) % 8 == 0: + raise ValueError("Not a canonical UMMA_K Layout: Expected MN-size multiple of 8.") + canonical_layout = cute.logical_divide(layout, (8, 2)) + if not cute.is_congruent(canonical_layout, ((1, 1), (1, 1))): + raise ValueError("Not a canonical UMMA_K Layout: Expected profile failure.") + stride_00 = canonical_layout.stride[0][0] + if stride_00 != swizzle_atom_mn_size: + raise ValueError("Not a canonical UMMA_K Layout: Expected stride failure.") + stride_10 = canonical_layout.stride[1][0] + if layout_type is not LayoutType.SWIZZLE_NONE and stride_10 != 1: + raise ValueError("Not a canonical UMMA_K Layout: Expected stride failure.") + stride_01 = canonical_layout.stride[0][1] + stride_byte_offset, leading_byte_offset = stride_01, stride_10 + + # ------------------------------------------------------------------ pack + desc = 0 + # leading_byte_offset_ [16:30) + desc |= (leading_byte_offset & 0x3FFF) << 16 + # stride_byte_offset_ [32:46) + desc |= (stride_byte_offset & 0x3FFF) << 32 + # version_ [46:48) + desc |= (VERSION & 0x3) << 46 + # base_offset_ [49:52) + desc |= (BASE_OFFSET & 0x7) << 49 + # lbo_mode_ [52:53) + desc |= (LBO_MODE & 0x1) << 52 + # layout_type_ [61:64) + desc |= (int(layout_type) & 0x7) << 61 + + return desc & 0xFFFF_FFFF_FFFF_FFFF # force 64-bit width + + +def make_smem_desc_start_addr(start_addr: cute.Pointer) -> cutlass.Int32: + # 14 bits, remove 4 LSB (bits 0-13 in desc) + return (start_addr.toint() & 0x3FFFF) >> 4 + + +def smem_desc_base_from_tensor(sA: cute.Tensor, major: Major) -> int: + sA_swizzle = sA.iterator.type.swizzle_type + return make_smem_desc_base( + cute.recast_layout(128, sA.element_type.width, sA.layout[0]), + sA_swizzle, + major, + ) diff --git a/fastvideo-kernel/fa4/named_barrier.py b/fastvideo-kernel/fa4/named_barrier.py new file mode 100644 index 000000000..cf252f3db --- /dev/null +++ b/fastvideo-kernel/fa4/named_barrier.py @@ -0,0 +1,68 @@ +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. + +import enum + + +class NamedBarrierFwd(enum.IntEnum): + Epilogue = enum.auto() # starts from 1 as barrier 0 is reserved for sync_threads() + WarpSchedulerWG1 = enum.auto() + WarpSchedulerWG2 = enum.auto() + WarpSchedulerWG3 = enum.auto() + PFull = enum.auto() + PEmpty = enum.auto() + + +class NamedBarrierFwdSm100(enum.IntEnum): + Epilogue = enum.auto() # starts from 1 as barrier 0 is reserved for sync_threads() + TmemPtr = enum.auto() + SoftmaxStatsW0 = enum.auto() + SoftmaxStatsW1 = enum.auto() + SoftmaxStatsW2 = enum.auto() + SoftmaxStatsW3 = enum.auto() + SoftmaxStatsW4 = enum.auto() + SoftmaxStatsW5 = enum.auto() + SoftmaxStatsW6 = enum.auto() + SoftmaxStatsW7 = enum.auto() + # Reserve one independent SMEM-P visibility barrier per compute slot. A + # single shared barrier can mix arrivals from the two softmax warpgroups. + SoftmaxSmemP0 = enum.auto() + SoftmaxSmemP1 = enum.auto() + CorrectionScale = enum.auto() + + +class NamedBarrierBwd(enum.IntEnum): + Epilogue = enum.auto() + WarpSchedulerWG1 = enum.auto() + WarpSchedulerWG2 = enum.auto() + WarpSchedulerWG3 = enum.auto() + PdS = enum.auto() + dQFullWG0 = enum.auto() + dQFullWG1 = enum.auto() + dQFullWG2 = enum.auto() + dQEmptyWG0 = enum.auto() + dQEmptyWG1 = enum.auto() + dQEmptyWG2 = enum.auto() + + +class NamedBarrierBwdSm100(enum.IntEnum): + EpilogueWG1 = enum.auto() + EpilogueWG2 = enum.auto() + Compute = enum.auto() + dQaccReduce = enum.auto() + TmemPtr = enum.auto() + + +class NamedBarrierFwdSm100_MLA2CTA(enum.IntEnum): + Epilogue = enum.auto() + TmemPtr = enum.auto() + Cpasync = enum.auto() + Softmax = enum.auto() + SoftmaxStatsFull = enum.auto() + SoftmaxStatsEmpty = enum.auto() + + +class NamedBarrierBwdSm100_MLA2CTA(enum.IntEnum): + Epilogue = enum.auto() + TmemPtr = enum.auto() + Cpasync = enum.auto() + Softmax = enum.auto() diff --git a/fastvideo-kernel/fa4/pack_gqa.py b/fastvideo-kernel/fa4/pack_gqa.py new file mode 100644 index 000000000..5b481b5e6 --- /dev/null +++ b/fastvideo-kernel/fa4/pack_gqa.py @@ -0,0 +1,263 @@ +# Copyright (c) 2025, Tri Dao. + +from dataclasses import dataclass +from typing import Union, Tuple + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync + + +from quack import layout_utils +import flash_attn.cute.utils as utils + + +def pack_gqa_layout(T, qhead_per_kvhead, nheads_kv, head_idx): + """Reshape a tensor to fold qhead_per_kvhead into the seqlen dimension (mode 0). + + The head dimension is at mode ``head_idx``. Modes before it (1..head_idx-1) + are kept as-is (e.g. headdim for Q/O tensors), and modes after it are kept + as-is (e.g. batch). + + For Q/O tensors (head_idx=2): + (seqlen_q, headdim, nheads, batch, ...) -> ((qhead_per_kvhead, seqlen_q), headdim, nheads_kv, batch, ...) + For LSE tensors (head_idx=1): + (seqlen_q, nheads, batch, ...) -> ((qhead_per_kvhead, seqlen_q), nheads_kv, batch, ...) + """ + head_stride = T.stride[head_idx] + shape_packed = ( + (qhead_per_kvhead, T.shape[0]), + *[T.shape[i] for i in range(1, head_idx)], + nheads_kv, + *[T.shape[i] for i in range(head_idx + 1, len(T.shape))], + ) + stride_packed = ( + (head_stride, T.stride[0]), + *[T.stride[i] for i in range(1, head_idx)], + head_stride * qhead_per_kvhead, + *[T.stride[i] for i in range(head_idx + 1, len(T.shape))], + ) + return cute.make_tensor(T.iterator, cute.make_layout(shape_packed, stride=stride_packed)) + + +def make_packgqa_tiled_tma_atom( + op: cute.atom.CopyOp, + gmem_tensor: cute.Tensor, + smem_layout: Union[cute.Layout, cute.ComposedLayout], + cta_tiler: Tuple[int, int], + qhead_per_kvhead: int, + head_idx: int, +): + # This packing and unpacking of the layout is so that we keep the same TMA dimension as usual. + # e.g. for (seqlen, d, nheads, b) layout, we still have 4D TMA after packing to + # ((nheads, seqlen), d, b). + # If we instead pack directly to ((qhead_per_kvhead, seqlen), d, nheads_kv, b) we'd have 5D TMA. + # Pack headdim and seqlen dim into 1: (seqlen, d, nheads, b) -> ((nheads, seqlen), d, b) + gmem_tensor = layout_utils.select( + gmem_tensor, [head_idx, *range(head_idx), *range(head_idx + 1, cute.rank(gmem_tensor))] + ) + gmem_tensor = cute.group_modes(gmem_tensor, 0, 2) + assert cta_tiler[0] % qhead_per_kvhead == 0, ( + "CTA tile size in the seqlen dimension must be divisible by qhead_per_kvhead" + ) + tma_atom, tma_tensor = cpasync.make_tiled_tma_atom( + op, + gmem_tensor, + smem_layout, + ((qhead_per_kvhead, cta_tiler[0] // qhead_per_kvhead), cta_tiler[1]), # No mcast + ) + # Unpack from ((nheads, seqlen), d, b) -> ((qhead_per_kvhead, seqlen), d, nheads_kv, b) + T = tma_tensor + shape_packed = ( + (qhead_per_kvhead, T.shape[0][1]), + *[T.shape[i] for i in range(1, head_idx)], + T.shape[0][0] // qhead_per_kvhead, + *[T.shape[i] for i in range(head_idx, len(T.shape))], + ) + stride_packed = ( + *[T.stride[i] for i in range(head_idx)], + T.stride[0][0] * qhead_per_kvhead, + *[T.stride[i] for i in range(head_idx, len(T.shape))], + ) + tma_tensor = cute.make_tensor(T.iterator, cute.make_layout(shape_packed, stride=stride_packed)) + return tma_atom, tma_tensor + + +def unpack_gqa_layout(T, qhead_per_kvhead, head_idx): + """Reverse of pack_gqa_layout: unfold qhead_per_kvhead from the seqlen dimension (mode 0). + + The head dimension is at mode ``head_idx``. Modes before it (1..head_idx-1) + are kept as-is (e.g. headdim for Q/O tensors), and modes after it are kept + as-is (e.g. batch). + + For Q/O tensors (head_idx=2): + ((qhead_per_kvhead, seqlen_q), headdim, nheads_kv, batch, ...) -> (seqlen_q, headdim, nheads, batch, ...) + For LSE tensors (head_idx=1): + ((qhead_per_kvhead, seqlen_q), nheads_kv, batch, ...) -> (seqlen_q, nheads, batch, ...) + """ + seqlen_stride = T.stride[0][1] + head_stride = T.stride[0][0] + shape_unpacked = ( + T.shape[0][1], + *[T.shape[i] for i in range(1, head_idx)], + T.shape[head_idx] * qhead_per_kvhead, + *[T.shape[i] for i in range(head_idx + 1, len(T.shape))], + ) + stride_unpacked = ( + seqlen_stride, + *[T.stride[i] for i in range(1, head_idx)], + head_stride, + *[T.stride[i] for i in range(head_idx + 1, len(T.shape))], + ) + return cute.make_tensor(T.iterator, cute.make_layout(shape_unpacked, stride=stride_unpacked)) + + +@dataclass +class PackGQA: + m_block_size: cutlass.Constexpr[int] + head_dim_padded: cutlass.Constexpr[int] + check_hdim_oob: cutlass.Constexpr[bool] + qhead_per_kvhead: cutlass.Constexpr[bool] + + @cute.jit + def compute_ptr( + self, + tensor: cute.Tensor, + cRows: cute.Tensor, + tidx: cutlass.Int32, + block: cutlass.Int32, + threads_per_row: cutlass.Constexpr[int], + num_threads: cutlass.Constexpr[int], + ): + num_ptr_per_thread = cute.ceil_div(cute.size(cRows), threads_per_row) + tPrPtr = cute.make_rmem_tensor(num_ptr_per_thread, cutlass.Int64) + for i in cutlass.range_constexpr(num_ptr_per_thread): + row = i * num_threads + cRows[tidx % threads_per_row][0] + idx = block * self.m_block_size + row + m_idx = idx // self.qhead_per_kvhead + h_idx = idx - m_idx * self.qhead_per_kvhead + tPrPtr[i] = utils.elem_pointer(tensor, ((h_idx, m_idx),)).toint() + return tPrPtr + + @cute.jit + def load_Q( + self, + mQ: cute.Tensor, # ((qhead_per_kvhead, seqlen_q), headdim) + sQ: cute.Tensor, # (m_block_size, head_dim_padded) + gmem_tiled_copy: cute.TiledCopy, + tidx: cutlass.Int32, + block: cutlass.Int32, + seqlen: cutlass.Int32, + ): + gmem_thr_copy = gmem_tiled_copy.get_slice(tidx) + cQ = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded)) + tQsQ = gmem_thr_copy.partition_D(sQ) + tQcQ = gmem_thr_copy.partition_S(cQ) + t0QcQ = gmem_thr_copy.get_slice(0).partition_S(cQ) + tQpQ = utils.predicate_k(tQcQ, limit=mQ.shape[1]) + tQcQ_row = tQcQ[0, None, 0] + threads_per_row = gmem_tiled_copy.layout_tv_tiled.shape[0][0] + assert cute.arch.WARP_SIZE % threads_per_row == 0, "threads_per_row must divide WARP_SIZE" + num_threads = gmem_tiled_copy.size + tPrQPtr = self.compute_ptr(mQ[None, 0], tQcQ_row, tidx, block, threads_per_row, num_threads) + for m in cutlass.range_constexpr(cute.size(tQsQ.shape[1])): + q_ptr_i64 = utils.shuffle_sync( + tPrQPtr[m // threads_per_row], m % threads_per_row, width=threads_per_row + ) + q_gmem_ptr = cute.make_ptr( + mQ.element_type, q_ptr_i64, cute.AddressSpace.gmem, assumed_align=16 + ) + if ( + t0QcQ[0, m, 0][0] + < seqlen * self.qhead_per_kvhead - block * self.m_block_size - tQcQ_row[0][0] + ): + mQ_cur = cute.make_tensor(q_gmem_ptr, (self.head_dim_padded,)) + elems_per_load = cute.size(tQsQ.shape[0][0]) + mQ_cur_copy = cute.tiled_divide(mQ_cur, (elems_per_load,)) + for k in cutlass.range_constexpr(cute.size(tQsQ.shape[2])): + ki = tQcQ[0, 0, k][1] // elems_per_load + cute.copy( + gmem_thr_copy, + mQ_cur_copy[None, ki], + tQsQ[None, m, k], + pred=tQpQ[None, m, k] if cutlass.const_expr(self.check_hdim_oob) else None, + ) + # We don't need to clear the sQ smem tiles since we'll only write out the valid outputs + + @cute.jit + def store_LSE( + self, + mLSE: cute.Tensor, # (qhead_per_kvhead, seqlen_q) + tLSErLSE: cute.Tensor, # (m_block_size, head_dim_padded) + tiled_mma: cute.TiledMma, + tidx: cutlass.Int32, + block: cutlass.Int32, + seqlen: cutlass.Int32, + ): + thr_mma = tiled_mma.get_slice(tidx) + caccO = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded)) + taccOcO = thr_mma.partition_C(caccO) + taccOcO_row = layout_utils.reshape_acc_to_mn(taccOcO)[None, 0] + assert cute.size(tLSErLSE) == cute.size(taccOcO_row) + threads_per_row = tiled_mma.tv_layout_C.shape[0][0] + assert cute.arch.WARP_SIZE % threads_per_row == 0, "threads_per_row must divide WARP_SIZE" + assert cute.size(tLSErLSE) <= threads_per_row + num_threads = tiled_mma.size + tPrLSEPtr = self.compute_ptr(mLSE, taccOcO_row, tidx, block, threads_per_row, num_threads) + for m in cutlass.range_constexpr(cute.size(tLSErLSE)): + lse_ptr_i64 = utils.shuffle_sync( + tPrLSEPtr[m // threads_per_row], + m % threads_per_row, + width=threads_per_row, + ) + lse_gmem_ptr = cute.make_ptr( + mLSE.element_type, lse_ptr_i64, cute.AddressSpace.gmem, assumed_align=4 + ) + row = block * self.m_block_size + taccOcO_row[m][0] + # Only the thread corresponding to column 0 writes out the lse to gmem + if taccOcO[0][1] == 0 and row < seqlen * self.qhead_per_kvhead: + mLSE_copy = cute.make_tensor(lse_gmem_ptr, (1,)) + mLSE_copy[0] = tLSErLSE[m] + + @cute.jit + def store_O( + self, + mO: cute.Tensor, # ((qhead_per_kvhead, seqlen_q), headdim) + tOrO: cute.Tensor, # (m_block_size, head_dim_padded) split across threads according to gmem_tiled_copy + gmem_tiled_copy: cute.TiledCopy, + tidx: cutlass.Int32, + block: cutlass.Int32, + seqlen: cutlass.Int32, + ): + gmem_thr_copy = gmem_tiled_copy.get_slice(tidx) + cO = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded)) + tOcO = gmem_thr_copy.partition_S(cO) + t0OcO = gmem_thr_copy.get_slice(0).partition_S(cO) + tOpO = utils.predicate_k(tOcO, limit=mO.shape[1]) + tOcO_row = tOcO[0, None, 0] + threads_per_row = gmem_tiled_copy.layout_tv_tiled.shape[0][0] + assert cute.arch.WARP_SIZE % threads_per_row == 0, "threads_per_row must divide WARP_SIZE" + num_threads = gmem_tiled_copy.size + tPrOPtr = self.compute_ptr(mO[None, 0], tOcO_row, tidx, block, threads_per_row, num_threads) + for m in cutlass.range_constexpr(cute.size(tOrO.shape[1])): + o_ptr_i64 = utils.shuffle_sync( + tPrOPtr[m // threads_per_row], m % threads_per_row, width=threads_per_row + ) + o_gmem_ptr = cute.make_ptr( + mO.element_type, o_ptr_i64, cute.AddressSpace.gmem, assumed_align=16 + ) + if ( + t0OcO[0, m, 0][0] + < seqlen * self.qhead_per_kvhead - block * self.m_block_size - tOcO_row[0][0] + ): + mO_cur = cute.make_tensor(o_gmem_ptr, (self.head_dim_padded,)) + elems_per_load = cute.size(tOrO.shape[0][0]) + mO_cur_copy = cute.tiled_divide(mO_cur, (elems_per_load,)) + for k in cutlass.range_constexpr(cute.size(tOrO.shape[2])): + ki = tOcO[0, 0, k][1] // elems_per_load + cute.copy( + gmem_thr_copy, + tOrO[None, m, k], + mO_cur_copy[None, ki], + pred=tOpO[None, m, k] if cutlass.const_expr(self.check_hdim_oob) else None, + ) diff --git a/fastvideo-kernel/fa4/paged_kv.py b/fastvideo-kernel/fa4/paged_kv.py new file mode 100644 index 000000000..407a8b8c6 --- /dev/null +++ b/fastvideo-kernel/fa4/paged_kv.py @@ -0,0 +1,247 @@ +from typing import Type +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync +from cutlass import Int32, const_expr + +from flash_attn.cute import utils +from quack.cute_dsl_utils import ParamsBase +from cutlass.cute import FastDivmodDivisor + +import math + + +@dataclass +class PagedKVManager(ParamsBase): + mPageTable: cute.Tensor + mK_paged: cute.Tensor + mV_paged: cute.Tensor + thread_idx: Int32 + + page_size_divmod: FastDivmodDivisor + seqlen_k: Int32 + leftpad_k: Int32 + n_block_size: Int32 + num_threads: cutlass.Constexpr[Int32] + head_dim_padded: cutlass.Constexpr[Int32] + head_dim_v_padded: cutlass.Constexpr[Int32] + + arch: cutlass.Constexpr[Int32] + v_gmem_transposed: cutlass.Constexpr[bool] + + gmem_threads_per_row: cutlass.Constexpr[Int32] + page_entry_per_thread: Int32 + async_copy_elems: Int32 + + gmem_tiled_copy_KV: cute.TiledCopy + gmem_thr_copy_KV: cute.TiledCopy + tPrPage: cute.Tensor + tPrPageOffset: cute.Tensor + tKpK: cute.Tensor + tVpV: cute.Tensor + + @staticmethod + def create( + mPageTable: cute.Tensor, + mK_paged: cute.Tensor, + mV_paged: cute.Tensor, + page_size_divmod: FastDivmodDivisor, + bidb: Int32, + bidh: Int32, + thread_idx: Int32, + seqlen_k: Int32, + leftpad_k: Int32, + n_block_size: cutlass.Constexpr[Int32], + head_dim_padded: cutlass.Constexpr[Int32], + head_dim_v_padded: cutlass.Constexpr[Int32], + num_threads: cutlass.Constexpr[Int32], + dtype: Type[cutlass.Numeric], + arch: cutlass.Constexpr[int] = 100, + ): + # SM100 transposes V in gmem to (dv, page_size, num_pages); + # SM90 keeps V as (page_size, dv, num_pages), same layout as K. + v_gmem_transposed = arch != 90 + universal_copy_bits = 128 + async_copy_elems = universal_copy_bits // dtype.width + dtype_bytes = dtype.width // 8 + gmem_k_block_size = math.gcd( + head_dim_padded, + head_dim_v_padded, + 128 // dtype_bytes, + ) + assert gmem_k_block_size % async_copy_elems == 0 + gmem_threads_per_row = gmem_k_block_size // async_copy_elems + assert cute.arch.WARP_SIZE % gmem_threads_per_row == 0 + atom_async_copy = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + dtype, + num_bits_per_copy=universal_copy_bits, + ) + thr_layout = cute.make_ordered_layout( + (num_threads // gmem_threads_per_row, gmem_threads_per_row), + order=(1, 0), + ) + val_layout = cute.make_layout((1, async_copy_elems)) + gmem_tiled_copy_KV = cute.make_tiled_copy_tv(atom_async_copy, thr_layout, val_layout) + gmem_thr_copy_KV = gmem_tiled_copy_KV.get_slice(thread_idx) + page_entry_per_thread = n_block_size // num_threads + + tPrPage = cute.make_rmem_tensor((page_entry_per_thread,), Int32) + tPrPageOffset = cute.make_rmem_tensor((page_entry_per_thread,), Int32) + + mPageTable = mPageTable[bidb, None] + mK_paged = mK_paged[None, None, bidh, None] + mV_paged = mV_paged[None, None, bidh, None] + + cK = cute.make_identity_tensor((n_block_size, head_dim_padded)) + tKcK = gmem_thr_copy_KV.partition_S(cK) + tKpK = utils.predicate_k(tKcK, limit=mK_paged.shape[1]) + + if const_expr(head_dim_padded == head_dim_v_padded): + tVpV = tKpK + else: + cV = cute.make_identity_tensor((n_block_size, head_dim_v_padded)) + tVcV = gmem_thr_copy_KV.partition_S(cV) + # When V is transposed in gmem, dv is shape[0]; otherwise dv is shape[1] (same as K) + V_limit = cute.size(mV_paged.shape[0 if v_gmem_transposed else 1]) + tVpV = utils.predicate_k(tVcV, limit=V_limit) + + return PagedKVManager( + mPageTable, + mK_paged, + mV_paged, + thread_idx, + page_size_divmod, + seqlen_k, + leftpad_k, + n_block_size, + num_threads, + head_dim_padded, + head_dim_v_padded, + arch, + v_gmem_transposed, + gmem_threads_per_row, + page_entry_per_thread, + async_copy_elems, + gmem_tiled_copy_KV, + gmem_thr_copy_KV, + tPrPage, + tPrPageOffset, + tKpK, + tVpV, + ) + + @cute.jit + def load_page_table(self, n_block: Int32): + for i in cutlass.range(self.page_entry_per_thread, unroll=1): + row = ( + i * self.num_threads + + (self.thread_idx % self.gmem_threads_per_row) + * (self.num_threads // self.gmem_threads_per_row) + + (self.thread_idx // self.gmem_threads_per_row) + ) + row_idx = n_block * self.n_block_size + row + + page_idx, page_offset = divmod(row_idx + self.leftpad_k, self.page_size_divmod) + + is_valid = ( + (i + 1) * self.num_threads <= self.n_block_size or row < self.n_block_size + ) and row_idx < self.seqlen_k + page = self.mPageTable[page_idx] if is_valid else 0 + + self.tPrPage[i] = page + self.tPrPageOffset[i] = page_offset + + @cute.jit + def compute_X_ptr(self, K_or_V: str, d_offset: int = 0): + tPrXPtr = cute.make_rmem_tensor((self.page_entry_per_thread,), cutlass.Int64) + mX = self.mK_paged if const_expr(K_or_V == "K") else self.mV_paged + # K is always (page_size, d, num_pages). V matches K when not transposed, + # but is (dv, page_size, num_pages) when transposed (SM100). + transposed = const_expr(K_or_V == "V" and self.v_gmem_transposed) + for i in cutlass.range(self.page_entry_per_thread, unroll=1): + page = self.tPrPage[i] + page_offset = self.tPrPageOffset[i] + if const_expr(transposed): + tPrXPtr[i] = utils.elem_pointer(mX, (d_offset, page_offset, page)).toint() + else: + tPrXPtr[i] = utils.elem_pointer(mX, (page_offset, d_offset, page)).toint() + return tPrXPtr + + @cute.jit + def _flatten_smem_sm100(self, sX: cute.Tensor, K_or_V: str): + """Flatten SM100 smem ((a,b), cta_split, k) to (a,(b,k)); transpose V to (d,page_size).""" + sX_pi = cute.make_tensor( + sX.iterator, + cute.make_layout( + (sX.shape[0][0], (sX.shape[0][1], sX.shape[2])), + stride=(sX.stride[0][0], (sX.stride[0][1], sX.stride[2])), + ), + ) + if const_expr(K_or_V == "V"): + sX_pi = cute.make_tensor(sX_pi.iterator, cute.select(sX_pi.layout, mode=[1, 0])) + return sX_pi + + @cute.jit + def _copy_row_async( + self, + tXsX: cute.Tensor, + tXcX: cute.Tensor, + mX_paged_cur_copy: cute.Tensor, + m: Int32, + should_load: cute.Tensor, + ): + """Issue cp.async copies for one row across all k-tiles.""" + for k in cutlass.range_constexpr(cute.size(tXsX, mode=[2])): + ki = tXcX[0, 0, k][1] // self.async_copy_elems + mX_paged_cur_copy_ki = mX_paged_cur_copy[None, ki] + tXsX_k = tXsX[None, m, k] + mX_paged_cur_copy_ki = cute.make_tensor(mX_paged_cur_copy_ki.iterator, tXsX_k.layout) + cute.copy( + self.gmem_tiled_copy_KV, + mX_paged_cur_copy_ki, + tXsX_k, + pred=should_load, + ) + + @cute.jit + def load_KV(self, n_block: Int32, sX: cute.Tensor, K_or_V: str): + assert K_or_V in ("K", "V") + + tPrXPtr = self.compute_X_ptr(K_or_V) + + if const_expr(self.arch == 90): + # SM90: sX is already stage-sliced by caller (sK[None, None, stage]). + # Flatten hierarchical modes to get (n_block_size, head_dim). + sX_pi = cute.group_modes(sX, 0, 1) + # SM90 does NOT transpose V here (it's transposed via utils.transpose_view before MMA) + else: + sX_pi = self._flatten_smem_sm100(sX, K_or_V) + + head_dim = self.head_dim_v_padded if const_expr(K_or_V == "V") else self.head_dim_padded + cX = cute.make_identity_tensor((self.n_block_size, head_dim)) + tXsX = self.gmem_thr_copy_KV.partition_D(sX_pi) + tXcX = self.gmem_thr_copy_KV.partition_S(cX) + tXc0X = self.gmem_thr_copy_KV.get_slice(0).partition_S(cX) + + seqlenk_row_limit = ( + self.seqlen_k - n_block * self.n_block_size - tXcX[0][0] if n_block >= 0 else 0 + ) + for m in cutlass.range_constexpr(cute.size(tXsX, mode=[1])): + row_valid = tXc0X[0, m, 0][0] < seqlenk_row_limit + should_load = cute.make_fragment_like(tXsX[(0, None), m, 0], cute.Boolean) + should_load.fill(row_valid) + + x_ptr_i64 = utils.shuffle_sync( + tPrXPtr[m // self.gmem_threads_per_row], + m % self.gmem_threads_per_row, + width=self.gmem_threads_per_row, + ) + x_gmem_ptr = cute.make_ptr( + self.mK_paged.element_type, x_ptr_i64, cute.AddressSpace.gmem, assumed_align=16 + ) + mX_paged_cur = cute.make_tensor(x_gmem_ptr, cute.make_layout((head_dim,))) + mX_paged_cur_copy = cute.tiled_divide(mX_paged_cur, (self.async_copy_elems,)) + self._copy_row_async(tXsX, tXcX, mX_paged_cur_copy, m, should_load) diff --git a/fastvideo-kernel/fa4/pipeline.py b/fastvideo-kernel/fa4/pipeline.py new file mode 100644 index 000000000..f8fdc1e80 --- /dev/null +++ b/fastvideo-kernel/fa4/pipeline.py @@ -0,0 +1,402 @@ +# Copyright (c) 2025, Tri Dao. + +from typing import Optional +from dataclasses import dataclass + +import cutlass.cute as cute +from cutlass import Boolean, Int32, const_expr +from cutlass.cutlass_dsl import if_generate, dsl_user_op +from cutlass.pipeline import PipelineState +from cutlass.pipeline import PipelineUserType +from cutlass.pipeline import NamedBarrier as NamedBarrierOg +from cutlass.pipeline import PipelineAsync as PipelineAsyncOg +from cutlass.pipeline import PipelineCpAsync as PipelineCpAsyncOg +from cutlass.pipeline import PipelineTmaAsync as PipelineTmaAsyncOg +from cutlass.pipeline import PipelineTmaUmma as PipelineTmaUmmaOg +from cutlass.pipeline import PipelineUmmaAsync as PipelineUmmaAsyncOg +from cutlass.pipeline import PipelineAsyncUmma as PipelineAsyncUmmaOg + + +def _override_create(parent_cls, child_cls): + """Create a static factory that constructs parent_cls then re-classes to child_cls.""" + + @staticmethod + def create(*args, **kwargs): + obj = parent_cls.create(*args, **kwargs) + # Can't assign to __class__ directly since the dataclass is frozen + object.__setattr__(obj, "__class__", child_cls) + return obj + + return create + + +def _make_state(index: Int32, phase: Int32) -> PipelineState: + """Construct a PipelineState from index and phase (count/stages unused by callers).""" + return PipelineState(stages=0, count=Int32(0), index=index, phase=phase) + + +class PipelineStateSimple: + """ + Pipeline state contains an index and phase bit corresponding to the current position in the circular buffer. + Use a single Int32 to store both the index and phase bit, then we use divmod to get the + index and phase. If stages is a power of 2, divmod turns into bit twiddling. + """ + + def __init__(self, stages: int, phase_index: Int32): + self._stages = stages + self._phase_index = phase_index + + def clone(self) -> "PipelineStateSimple": + return PipelineStateSimple(self.stages, self._phase_index) + + @property + def stages(self) -> int: + return self._stages + + @property + def index(self) -> Int32: + if const_expr(self._stages == 1): + return Int32(0) + else: + return self._phase_index % self._stages + + @property + def phase(self) -> Int32: + # PTX docs say that the phase parity needs to be 0 or 1, so by right we need to + # take modulo 2. But in practice just passing the phase in without modulo works fine. + if const_expr(self._stages == 1): + return self._phase_index + else: + return self._phase_index // self._stages + + def advance(self): + if const_expr(self._stages == 1): + self._phase_index ^= 1 + else: + self._phase_index += 1 + + def __extract_mlir_values__(self): + phase_index = self._phase_index + return [phase_index.ir_value()] + + def __new_from_mlir_values__(self, values): + return PipelineStateSimple(self.stages, Int32(values[0])) + + +def make_pipeline_state(type: PipelineUserType, stages: int): + """ + Creates a pipeline state. Producers are assumed to start with an empty buffer and have a flipped phase bit of 1. + """ + if type is PipelineUserType.Producer: + return PipelineStateSimple(stages, Int32(stages)) + elif type is PipelineUserType.Consumer: + return PipelineStateSimple(stages, Int32(0)) + else: + assert False, "Error: invalid PipelineUserType specified for make_pipeline_state." + + +# ── Shared helpers ─────────────────────────────────────────────────────────── + + +def _call_with_elect_one(parent_method, self, state, elect_one, syncwarp, loc, ip): + """Optionally wrap a parent pipeline method call in sync_warp + elect_one.""" + if const_expr(elect_one): + if const_expr(syncwarp): + cute.arch.sync_warp() + with cute.arch.elect_one(): + parent_method(self, state, loc=loc, ip=ip) + else: + parent_method(self, state, loc=loc, ip=ip) + + +# ── Mixin: _w_index / _w_index_phase variants that delegate to parent ─────── +# Each parent class has PipelineState-based methods (producer_acquire, producer_commit, +# consumer_wait, consumer_release). The _w_index_phase variants just construct a +# PipelineState from (index, phase) and delegate. + + +class _PipelineIndexPhaseMixin: + """Mixin providing _w_index_phase / _w_index methods that delegate to PipelineState-based parents.""" + + @dsl_user_op + def producer_acquire_w_index_phase( + self, + index: Int32, + phase: Int32, + try_acquire_token: Optional[Boolean] = None, + *, + loc=None, + ip=None, + ): + state = _make_state(index, phase) + # Call the parent's producer_acquire (which takes PipelineState) + self.producer_acquire(state, try_acquire_token, loc=loc, ip=ip) + + @dsl_user_op + def producer_commit_w_index(self, index: Int32, *, loc=None, ip=None): + state = _make_state(index, Int32(0)) + self.producer_commit(state, loc=loc, ip=ip) + + @dsl_user_op + def consumer_wait_w_index_phase( + self, + index: Int32, + phase: Int32, + try_wait_token: Optional[Boolean] = None, + *, + loc=None, + ip=None, + ): + state = _make_state(index, phase) + self.consumer_wait(state, try_wait_token, loc=loc, ip=ip) + + @dsl_user_op + def consumer_release_w_index(self, index: Int32, *, loc=None, ip=None): + state = _make_state(index, Int32(0)) + self.consumer_release(state, loc=loc, ip=ip) + + +# ── NamedBarrier ───────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class NamedBarrier(NamedBarrierOg): + create = _override_create(NamedBarrierOg, None) # patched below + + @dsl_user_op + def arrive_w_index(self, index: Int32, *, loc=None, ip=None) -> None: + """ + The aligned flavor of arrive is used when all threads in the CTA will execute the + same instruction. See PTX documentation. + """ + cute.arch.barrier_arrive( + barrier_id=self.barrier_id + index, + number_of_threads=self.num_threads, + loc=loc, + ip=ip, + ) + + @dsl_user_op + def arrive_and_wait_w_index(self, index: Int32, *, loc=None, ip=None) -> None: + cute.arch.barrier( + barrier_id=self.barrier_id + index, + number_of_threads=self.num_threads, + loc=loc, + ip=ip, + ) + + +NamedBarrier.create = _override_create(NamedBarrierOg, NamedBarrier) + + +# ── PipelineAsync ──────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class PipelineAsync(_PipelineIndexPhaseMixin, PipelineAsyncOg): + """ + PipelineAsync with optional elect_one for producer_commit and consumer_release. + + When elect_one_*=True (set at create time), only one elected thread per warp + signals the barrier arrive. This is useful when the mask count is set to 1 per warp. + + Args (to create): + elect_one_commit: If True, only elected thread signals producer_commit. + syncwarp_before_commit: If True (default), issue syncwarp before elect_one. + elect_one_release: If True, only elected thread signals consumer_release. + syncwarp_before_release: If True (default), issue syncwarp before elect_one. + Set syncwarp to False when threads are already converged (e.g. after wgmma wait_group). + """ + + _elect_one_commit: bool = False + _syncwarp_before_commit: bool = True + _elect_one_release: bool = False + _syncwarp_before_release: bool = True + + @staticmethod + def create( + *args, + elect_one_commit: bool = False, + syncwarp_before_commit: bool = True, + elect_one_release: bool = False, + syncwarp_before_release: bool = True, + **kwargs, + ): + obj = PipelineAsyncOg.create(*args, **kwargs) + object.__setattr__(obj, "__class__", PipelineAsync) + object.__setattr__(obj, "_elect_one_commit", elect_one_commit) + object.__setattr__(obj, "_syncwarp_before_commit", syncwarp_before_commit) + object.__setattr__(obj, "_elect_one_release", elect_one_release) + object.__setattr__(obj, "_syncwarp_before_release", syncwarp_before_release) + return obj + + @dsl_user_op + def producer_commit(self, state: PipelineState, *, loc=None, ip=None): + _call_with_elect_one( + PipelineAsyncOg.producer_commit, + self, + state, + self._elect_one_commit, + self._syncwarp_before_commit, + loc, + ip, + ) + + @dsl_user_op + def consumer_release(self, state: PipelineState, *, loc=None, ip=None): + _call_with_elect_one( + PipelineAsyncOg.consumer_release, + self, + state, + self._elect_one_release, + self._syncwarp_before_release, + loc, + ip, + ) + + # _w_index variants inherited from _PipelineIndexPhaseMixin, which delegate + # to producer_commit / consumer_release above. + + +# ── PipelineCpAsync ────────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class PipelineCpAsync(_PipelineIndexPhaseMixin, PipelineCpAsyncOg): + _elect_one_release: bool = False + _syncwarp_before_release: bool = True + + @staticmethod + def create( + *args, + elect_one_release: bool = False, + syncwarp_before_release: bool = True, + **kwargs, + ): + obj = PipelineCpAsyncOg.create(*args, **kwargs) + object.__setattr__(obj, "__class__", PipelineCpAsync) + object.__setattr__(obj, "_elect_one_release", elect_one_release) + object.__setattr__(obj, "_syncwarp_before_release", syncwarp_before_release) + return obj + + @dsl_user_op + def consumer_release(self, state: PipelineState, *, loc=None, ip=None): + _call_with_elect_one( + PipelineCpAsyncOg.consumer_release, + self, + state, + self._elect_one_release, + self._syncwarp_before_release, + loc, + ip, + ) + + # _w_index variants inherited from _PipelineIndexPhaseMixin. + + +# ── PipelineTmaAsync ──────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class PipelineTmaAsync(_PipelineIndexPhaseMixin, PipelineTmaAsyncOg): + """Override producer_acquire to take in extra_tx_count parameter.""" + + @dsl_user_op + def producer_acquire( + self, + state: PipelineState, + try_acquire_token: Optional[Boolean] = None, + extra_tx_count: int = 0, + *, + loc=None, + ip=None, + ): + """ + TMA producer commit conditionally waits on buffer empty and sets the transaction barrier for leader threadblocks. + """ + if_generate( + try_acquire_token is None or try_acquire_token == 0, + lambda: self.sync_object_empty.wait(state.index, state.phase, loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + if const_expr(extra_tx_count == 0): + self.sync_object_full.arrive(state.index, self.producer_mask, loc=loc, ip=ip) + else: + tx_count = self.sync_object_full.tx_count + extra_tx_count + self.sync_object_full.arrive_and_expect_tx(state.index, tx_count, loc=loc, ip=ip) + + +PipelineTmaAsync.create = _override_create(PipelineTmaAsyncOg, PipelineTmaAsync) + + +# ── PipelineTmaUmma ───────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class PipelineTmaUmma(_PipelineIndexPhaseMixin, PipelineTmaUmmaOg): + """Override producer_acquire to take in extra_tx_count parameter.""" + + @dsl_user_op + def producer_acquire( + self, + state: PipelineState, + try_acquire_token: Optional[Boolean] = None, + extra_tx_count: int = 0, + *, + loc=None, + ip=None, + ): + """ + TMA producer commit conditionally waits on buffer empty and sets the transaction barrier for leader threadblocks. + """ + if_generate( + try_acquire_token is None or try_acquire_token == 0, + lambda: self.sync_object_empty.wait(state.index, state.phase, loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + if const_expr(extra_tx_count == 0): + if_generate( + self.is_leader_cta, + lambda: self.sync_object_full.arrive( + state.index, self.producer_mask, loc=loc, ip=ip + ), + loc=loc, + ip=ip, + ) + else: + tx_count = self.sync_object_full.tx_count + extra_tx_count + if_generate( + self.is_leader_cta, + lambda: self.sync_object_full.arrive_and_expect_tx( + state.index, tx_count, loc=loc, ip=ip + ), + loc=loc, + ip=ip, + ) + + +PipelineTmaUmma.create = _override_create(PipelineTmaUmmaOg, PipelineTmaUmma) + + +# ── PipelineUmmaAsync ─────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class PipelineUmmaAsync(_PipelineIndexPhaseMixin, PipelineUmmaAsyncOg): + pass + + +PipelineUmmaAsync.create = _override_create(PipelineUmmaAsyncOg, PipelineUmmaAsync) + + +# ── PipelineAsyncUmma ─────────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class PipelineAsyncUmma(_PipelineIndexPhaseMixin, PipelineAsyncUmmaOg): + pass + + +PipelineAsyncUmma.create = _override_create(PipelineAsyncUmmaOg, PipelineAsyncUmma) diff --git a/fastvideo-kernel/fa4/pyproject.toml b/fastvideo-kernel/fa4/pyproject.toml new file mode 100644 index 000000000..3a34e42b9 --- /dev/null +++ b/fastvideo-kernel/fa4/pyproject.toml @@ -0,0 +1,68 @@ +[build-system] +requires = ["setuptools>=75"] +build-backend = "setuptools.build_meta" + +[project] +name = "flash-attn-4" +version = "4.0.0b20.dev2+fastvideo.82d6441" +description = "Flash Attention CUTE (CUDA Template Engine) implementation" +readme = "README.md" +requires-python = ">=3.10" +license = "BSD-3-Clause" +authors = [ + {name = "Tri Dao"}, +] +classifiers = [ + "Development Status :: 3 - Alpha", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", +] + +dependencies = [ + "nvidia-cutlass-dsl==4.6.0.dev0", + "torch", + "einops", + "typing_extensions", + "apache-tvm-ffi>=0.1.5,<0.2", + "torch-c-dlpack-ext", + "quack-kernels>=0.5.0", +] + +[project.optional-dependencies] +cu13 = ["nvidia-cutlass-dsl[cu13]==4.6.0.dev0"] +dev = [ + "pytest", + "pytest-xdist", + "ruff", +] + +[project.urls] +Homepage = "https://github.com/Dao-AILab/flash-attention" +Repository = "https://github.com/Dao-AILab/flash-attention" + +[tool.setuptools] +packages = ["flash_attn.cute"] +package-dir = {"flash_attn.cute" = "."} + +[[tool.uv.index]] +name = "pytorch-cu130" +url = "https://download.pytorch.org/whl/cu130" +explicit = true + +[tool.uv.sources] +torch = [ + { index = "pytorch-cu130", marker = "extra == 'cu13'" }, +] + +[tool.ruff] +line-length = 100 + +[tool.ruff.lint] +ignore = [ + "E731", # do not assign a lambda expression, use a def + "E741", # Do not use variables named 'I', 'O', or 'l' + "F841", # local variable is assigned to but never used + "D102", # Missing docstring in public methods +] diff --git a/fastvideo-kernel/fa4/seqlen_info.py b/fastvideo-kernel/fa4/seqlen_info.py new file mode 100644 index 000000000..c8ba56726 --- /dev/null +++ b/fastvideo-kernel/fa4/seqlen_info.py @@ -0,0 +1,302 @@ +from typing import Optional +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, const_expr + +from quack import copy_utils + +""" +This consolidates all the info related to sequence length. This is so that we can do all +the gmem reads once at the beginning of each tile, rather than having to repeat these reads +to compute various things like n_block_min, n_block_max, etc. +""" + + +@dataclass(frozen=True) +class SeqlenInfo: + offset: Int32 + offset_padded: Int32 + seqlen: Int32 + has_cu_seqlens: cutlass.Constexpr[bool] = False + + @staticmethod + def create( + batch_idx: Int32, + seqlen_static: Int32, + cu_seqlens: Optional[cute.Tensor] = None, + seqused: Optional[cute.Tensor] = None, + tile: cutlass.Constexpr[int] = 128, + ): + offset = 0 if const_expr(cu_seqlens is None) else cu_seqlens[batch_idx] + offset_padded = ( + 0 + if const_expr(cu_seqlens is None) + # Add divby so that the compiler knows the alignment when moving by offset_padded + else cute.assume((offset + batch_idx * tile) // tile * tile, divby=tile) + ) + if const_expr(seqused is not None): + seqlen = seqused[batch_idx] + elif const_expr(cu_seqlens is not None): + seqlen = cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx] + else: + seqlen = seqlen_static + return SeqlenInfo(offset, offset_padded, seqlen, has_cu_seqlens=cu_seqlens is not None) + + def offset_batch( + self, + mT: cute.Tensor, + batch_idx: Int32, + dim: int, + padded: cutlass.Constexpr[bool] = False, + multiple: int = 1, + ) -> cute.Tensor: + """Offset a tensor by batch index. batch dim is at position `dim`, seqlen is at dim=0.""" + if const_expr(not self.has_cu_seqlens): + idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mT) - 1 - dim) + return mT[idx] + else: + off = multiple * (self.offset if const_expr(not padded) else self.offset_padded) + offset = off if const_expr(cute.rank(mT.shape[0]) == 1) else (0, off) + idx = (offset,) + (None,) * (cute.rank(mT) - 1) + return cute.domain_offset(idx, mT) + + +@dataclass(frozen=True) +class SeqlenInfoQK: + offset_q: Int32 + offset_k: Int32 + padded_offset_q: Int32 + padded_offset_k: Int32 + seqlen_q: Int32 + seqlen_k: Int32 + m_block_offset: Int32 + block_idx_offset: Int32 + num_n_blocks: Int32 + has_cu_seqlens_q: cutlass.Constexpr[bool] + has_cu_seqlens_k: cutlass.Constexpr[bool] + has_seqused_q: cutlass.Constexpr[bool] + has_seqused_k: cutlass.Constexpr[bool] + + @staticmethod + def create( + batch_idx: Int32, + seqlen_q_static: Int32, + seqlen_k_static: Int32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + mCuTotalMBlocks: Optional[cute.Tensor] = None, + mCuBlockIdxOffsets: Optional[cute.Tensor] = None, + tile_m: cutlass.Constexpr[Int32] = 128, + tile_n: cutlass.Constexpr[Int32] = 128, + ): + offset_q = 0 if const_expr(mCuSeqlensQ is None) else mCuSeqlensQ[batch_idx] + offset_k = 0 if const_expr(mCuSeqlensK is None) else mCuSeqlensK[batch_idx] + padded_offset_q = ( + 0 + if const_expr(mCuSeqlensQ is None) + else cute.assume((offset_q + batch_idx * tile_m) // tile_m * tile_m, divby=tile_m) + ) + padded_offset_k = ( + 0 + if const_expr(mCuSeqlensK is None) + else cute.assume((offset_k + batch_idx * tile_n) // tile_n * tile_n, divby=tile_n) + ) + if const_expr(mSeqUsedQ is not None): + seqlen_q = mSeqUsedQ[batch_idx] + else: + seqlen_q = ( + seqlen_q_static + if const_expr(mCuSeqlensQ is None) + else mCuSeqlensQ[batch_idx + 1] - offset_q + ) + if const_expr(mSeqUsedK is not None): + seqlen_k = mSeqUsedK[batch_idx] + else: + seqlen_k = ( + seqlen_k_static + if const_expr(mCuSeqlensK is None) + else mCuSeqlensK[batch_idx + 1] - offset_k + ) + m_block_offset = 0 if const_expr(mCuTotalMBlocks is None) else mCuTotalMBlocks[batch_idx] + num_n_blocks = (seqlen_k + tile_n - 1) // tile_n + block_idx_offset = ( + mCuBlockIdxOffsets[batch_idx] + if const_expr(mCuBlockIdxOffsets is not None) + else m_block_offset * num_n_blocks + ) + return SeqlenInfoQK( + offset_q, + offset_k, + padded_offset_q, + padded_offset_k, + seqlen_q, + seqlen_k, + m_block_offset, + block_idx_offset, + num_n_blocks, + has_cu_seqlens_q=mCuSeqlensQ is not None, + has_cu_seqlens_k=mCuSeqlensK is not None, + has_seqused_q=mSeqUsedQ is not None, + has_seqused_k=mSeqUsedK is not None, + ) + + def offset_batch_Q( + self, + mQ: cute.Tensor, + batch_idx: Int32, + dim: int, + padded: cutlass.Constexpr[bool] = False, + ragged: cutlass.Constexpr[bool] = False, + ) -> cute.Tensor: + """Seqlen must be the first dimension of mQ""" + if const_expr(not ragged): + if const_expr(not self.has_cu_seqlens_q): + idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mQ) - 1 - dim) + return mQ[idx] + else: + offset_q = self.offset_q if const_expr(not padded) else self.padded_offset_q + offset_q = offset_q if const_expr(cute.rank(mQ.shape[0]) == 1) else (None, offset_q) + idx = (offset_q,) + (None,) * (cute.rank(mQ) - 1) + return cute.domain_offset(idx, mQ) + else: + if const_expr(not self.has_cu_seqlens_q): + offset_q = 0 + idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mQ) - 1 - dim) + mQ = mQ[idx] + else: + offset_q = self.offset_q if const_expr(not padded) else self.padded_offset_q + if const_expr(cute.rank(mQ.shape[0]) == 1): + return copy_utils.offset_ragged_tensor( + mQ, offset_q, self.seqlen_q, ragged_dim=0, ptr_shift=True + ) + else: # PackGQA + assert cute.rank(mQ.shape[0]) == 2 + # Unpack before calling offset_ragged_tensor, then pack + idx = ((None, None),) + (None,) * (cute.rank(mQ) - 1) + mQ = mQ[idx] + mQ = copy_utils.offset_ragged_tensor( + mQ, offset_q, self.seqlen_q, ragged_dim=1, ptr_shift=True + ) + return cute.group_modes(mQ, 0, 2) + + def offset_batch_K( + self, + mK: cute.Tensor, + batch_idx: Int32, + dim: int, + padded: cutlass.Constexpr[bool] = False, + ragged: cutlass.Constexpr[bool] = False, + multiple: int = 1, + ) -> cute.Tensor: + """Seqlen must be the first dimension of mK""" + if const_expr(not ragged): + if const_expr(not self.has_cu_seqlens_k): + idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mK) - 1 - dim) + return mK[idx] + else: + offset_k = self.offset_k if const_expr(not padded) else self.padded_offset_k + offset_k *= multiple + idx = (offset_k,) + (None,) * (cute.rank(mK) - 1) + return cute.domain_offset(idx, mK) + else: + if const_expr(not self.has_cu_seqlens_k): + offset_k = 0 + idx = (None,) * dim + (batch_idx,) + (None,) * (cute.rank(mK) - 1 - dim) + mK = mK[idx] + else: + offset_k = self.offset_k if const_expr(not padded) else self.padded_offset_k + offset_k *= multiple + return copy_utils.offset_ragged_tensor( + mK, offset_k, self.seqlen_k, ragged_dim=0, ptr_shift=True + ) + + +@dataclass(frozen=True) +class SeqlenInfoQKNewK: + """Sequence length info for append-KV with left-padding and new K support. + + Extends SeqlenInfoQK with: + - leftpad_k: left padding for K (tokens to skip at the start of the KV cache) + - offset_k_new: offset into the new K tensor + - seqlen_k_og: original K length (before appending new K), excluding leftpad + - seqlen_k_new: length of new K to append + - seqlen_k: total K length (seqlen_k_og + seqlen_k_new) + - seqlen_rotary: position for rotary embedding computation + """ + + leftpad_k: Int32 + offset_q: Int32 + offset_k: Int32 + offset_k_new: Int32 + seqlen_q: Int32 + seqlen_k_og: Int32 + seqlen_k_new: Int32 + seqlen_k: Int32 + seqlen_rotary: Int32 + + @staticmethod + def create( + batch_idx: Int32, + seqlen_q_static: Int32, + seqlen_k_static: Int32, + shape_K_new_0: Int32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mCuSeqlensKNew: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + mLeftpadK: Optional[cute.Tensor] = None, + mSeqlensRotary: Optional[cute.Tensor] = None, + ): + leftpad_k = 0 if const_expr(mLeftpadK is None) else mLeftpadK[batch_idx] + offset_q = 0 if const_expr(mCuSeqlensQ is None) else mCuSeqlensQ[batch_idx] + if const_expr(mCuSeqlensK is not None): + offset_k = mCuSeqlensK[batch_idx] + leftpad_k + else: + offset_k = leftpad_k if const_expr(mCuSeqlensQ is not None) else 0 + offset_k_new = 0 if const_expr(mCuSeqlensKNew is None) else mCuSeqlensKNew[batch_idx] + # seqlen_q + if const_expr(mSeqUsedQ is not None): + seqlen_q = mSeqUsedQ[batch_idx] + elif const_expr(mCuSeqlensQ is not None): + seqlen_q = mCuSeqlensQ[batch_idx + 1] - mCuSeqlensQ[batch_idx] + else: + seqlen_q = seqlen_q_static + # seqlen_k_og: original K length (excluding leftpad) + if const_expr(mSeqUsedK is not None): + seqlen_k_og = mSeqUsedK[batch_idx] - leftpad_k + elif const_expr(mCuSeqlensK is not None): + seqlen_k_og = mCuSeqlensK[batch_idx + 1] - mCuSeqlensK[batch_idx] - leftpad_k + else: + seqlen_k_og = ( + seqlen_k_static - leftpad_k + if const_expr(mCuSeqlensQ is not None) + else seqlen_k_static + ) + # seqlen_k_new + if const_expr(mCuSeqlensKNew is None): + seqlen_k_new = 0 if const_expr(mCuSeqlensQ is None) else shape_K_new_0 + else: + seqlen_k_new = mCuSeqlensKNew[batch_idx + 1] - mCuSeqlensKNew[batch_idx] + seqlen_k = seqlen_k_og if const_expr(mCuSeqlensQ is None) else seqlen_k_og + seqlen_k_new + + # seqlen_rotary: defaults to seqlen_k_og + leftpad_k unless explicitly provided + if const_expr(mSeqlensRotary is not None): + seqlen_rotary = mSeqlensRotary[batch_idx] + else: + seqlen_rotary = seqlen_k_og + leftpad_k + return SeqlenInfoQKNewK( + leftpad_k, + offset_q, + offset_k, + offset_k_new, + seqlen_q, + seqlen_k_og, + seqlen_k_new, + seqlen_k, + seqlen_rotary, + ) diff --git a/fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_backward.py b/fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_backward.py new file mode 100644 index 000000000..376a48fbd --- /dev/null +++ b/fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_backward.py @@ -0,0 +1,298 @@ +# Copyright (c) 2025, Siyu Wang, Shengbin Di, Yuxi Chi, Johnsonms, Linfeng Zheng, Haoyan Huang, Lanbo Li, Yun Zhong, Man Yuan, Minmin Sun, Yong Li, Wei Lin. + + +"""Fused multi-head attention (FMHA) backward for the SM100 architecture using CUTE DSL. + +Constraints: +* Supported head dimensions: 256 only +* mma_tiler_mn must be 64,64 +* Batch size must be the same for Q, K, and V tensors +""" + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.typing import Int32 + +from flash_attn.cute.sm100_hd256_2cta_fmha_backward_dqkernel import ( + BlackwellFusedMultiHeadAttentionBackwardDQKernel, +) +from flash_attn.cute.sm100_hd256_2cta_fmha_backward_dkdvkernel import ( + BlackwellFusedMultiHeadAttentionBackwardDKDVKernel, +) +from flash_attn.cute.cute_dsl_utils import assume_tensor_aligned +from flash_attn.cute.utils import AuxData + + +def _as_bshkrd_tensor( + tensor: cute.Tensor, + h_k: Int32, + h_r: Int32, + varlen: bool, +) -> cute.Tensor: + """Normalize (B,S,H,D)/(S,H,D) tensors to (B,S,H_k,H_r,D) view.""" + if cutlass.const_expr(cute.rank(tensor.layout) == 5): + return tensor + if cutlass.const_expr(cute.rank(tensor.layout) == 4): + return cute.make_tensor( + tensor.iterator, + cute.make_layout( + (tensor.shape[0], tensor.shape[1], h_k, h_r, tensor.shape[3]), + stride=( + tensor.stride[0], + tensor.stride[1], + tensor.stride[2] * h_r, + tensor.stride[2], + tensor.stride[3], + ), + ), + ) + assert cutlass.const_expr(cute.rank(tensor.layout) == 3), "Expected rank-3 varlen tensor" + assert cutlass.const_expr(varlen), "Rank-3 input is only valid for varlen backward" + return cute.make_tensor( + tensor.iterator, + cute.make_layout( + (1, tensor.shape[0], h_k, h_r, tensor.shape[2]), + stride=( + 0, + tensor.stride[0], + tensor.stride[1] * h_r, + tensor.stride[1], + tensor.stride[2], + ), + ), + ) + + +def _as_shhb_tensor( + tensor: cute.Tensor, + h_k: Int32, + h_r: Int32, + b: Int32, + varlen: bool, +) -> cute.Tensor: + """Normalize (B,H,S)/(H,S) tensors to (S, ((H_r, H_k), B)) view.""" + if cutlass.const_expr(cute.rank(tensor.layout) == 3): + return cute.make_tensor( + tensor.iterator, + cute.make_layout( + (tensor.shape[2], ((h_r, h_k), tensor.shape[0])), + stride=( + tensor.stride[2], + ((tensor.stride[1], tensor.stride[1] * h_r), tensor.stride[0]), + ), + ), + ) + assert cutlass.const_expr(cute.rank(tensor.layout) == 2), "Expected rank-2 varlen tensor" + assert cutlass.const_expr(varlen), "Rank-2 input is only valid for varlen backward" + return cute.make_tensor( + tensor.iterator, + cute.make_layout( + (tensor.shape[1], ((h_r, h_k), b)), + stride=( + tensor.stride[1], + ((tensor.stride[0], tensor.stride[0] * h_r), 0), + ), + ), + ) + + +class BlackwellFusedMultiHeadAttentionBackward: + """FMHA backward class for executing CuTeDSL kernel.""" + + def __init__( + self, + head_dim: int, + head_dim_v: int | None = None, + is_causal: bool = False, + is_local: bool = False, + qhead_per_kvhead: cutlass.Constexpr[int] = 1, + is_persistent: bool = False, + deterministic: bool = False, + cluster_size: int = 1, + use_2cta_instrs: bool = False, + score_mod: cutlass.Constexpr | None = None, + score_mod_bwd: cutlass.Constexpr | None = None, + mask_mod: cutlass.Constexpr | None = None, + has_aux_tensors: cutlass.Constexpr = False, + q_subtile_factor: cutlass.Constexpr[int] = 1, + tile_m_dq: int = 128, + tile_n_dq: int = 128, + tile_m_dkdv: int = 128, + tile_n_dkdv: int = 64, + window_size_left: int | None = None, + window_size_right: int | None = None, + use_clc_scheduler: bool = False, + ): + """Initialization.""" + head_dim_v = head_dim if head_dim_v is None else head_dim_v + assert head_dim == 256 and head_dim_v == 256, ( + "SM100 dedicated backward kernel only supports (head_dim, head_dim_v) = (256, 256)" + ) + assert not is_local, "SM100 backward with head_dim=256 does not support local attention" + assert tile_m_dq == 128 and tile_n_dq == 128, ( + "SM100 dedicated backward kernel only supports tile_m_dq=128 and tile_n_dq=128" + ) + assert tile_m_dkdv == 128 and tile_n_dkdv == 64, ( + "SM100 dedicated backward kernel only supports tile_m_dkdv=128 and tile_n_dkdv=64" + ) + assert score_mod is None and score_mod_bwd is None and mask_mod is None, ( + "SM100 backward with head_dim=256 does not support score_mod/mask_mod" + ) + assert not deterministic, ( + "SM100 backward with head_dim=256 does not support deterministic mode" + ) + assert not has_aux_tensors, "SM100 backward with head_dim=256 does not support aux_tensors" + assert cluster_size in (1, 2), ( + "SM100 backward with head_dim=256 only supports cluster_size in {1, 2}" + ) + assert use_2cta_instrs, "SM100 backward with head_dim=256 requires use_2cta_instrs=True" + # q_subtile_factor is accepted for interface parity with FlashAttentionBackwardSm100, + # but this dedicated kernel uses fixed internal behavior. + + self.acc_dtype = cutlass.Float32 + self.is_causal = is_causal + self.window_size_left = ( + None if (window_size_left is None or window_size_left < 0) else window_size_left + ) + self.window_size_right = ( + None if (window_size_right is None or window_size_right < 0) else window_size_right + ) + self.tile_m_dq = tile_m_dq + self.tile_n_dq = tile_n_dq + self.tile_m_dkdv = tile_m_dkdv + self.tile_n_dkdv = tile_n_dkdv + self.use_clc_scheduler = use_clc_scheduler + + self.dq_kernel = BlackwellFusedMultiHeadAttentionBackwardDQKernel( + self.acc_dtype, + (self.tile_m_dq, self.tile_n_dq, 256), + self.is_causal, + self.window_size_left, + self.window_size_right, + False, # is_persistent + False, # split_head + use_clc_scheduler=self.use_clc_scheduler, + ) + self.dkdv_kernel = BlackwellFusedMultiHeadAttentionBackwardDKDVKernel( + self.acc_dtype, + (self.tile_m_dkdv, self.tile_n_dkdv, 256), + self.is_causal, + self.window_size_left, + self.window_size_right, + use_clc_scheduler=self.use_clc_scheduler, + ) + + @cute.jit + def __call__( + self, + Q: cute.Tensor, + K: cute.Tensor, + V: cute.Tensor, + dO: cute.Tensor, + lse_log2: cute.Tensor, + dpsum: cute.Tensor, + dQ_accum: cute.Tensor | None, + dK: cute.Tensor, + dV: cute.Tensor, + scale_softmax: cutlass.Float32, + cumulative_s_q: cute.Tensor | None, + cumulative_s_k: cute.Tensor | None, + seqused_q: cute.Tensor | None = None, + seqused_k: cute.Tensor | None = None, + window_size_left: Int32 | None = None, + window_size_right: Int32 | None = None, + dQ_semaphore: cute.Tensor | None = None, + dK_semaphore: cute.Tensor | None = None, + dV_semaphore: cute.Tensor | None = None, + aux_data: AuxData = AuxData(), + block_sparse_tensors: cute.Tensor | None = None, + stream: cuda.CUstream = None, + ): + """Host function to launch CuTeDSL kernel.""" + assert seqused_q is None and seqused_k is None, ( + "SM100 backward with head_dim=256 does not support seqused_q/seqused_k" + ) + assert window_size_left is None and window_size_right is None, ( + "SM100 backward with head_dim=256 uses constructor-provided window sizes" + ) + assert dQ_semaphore is None and dK_semaphore is None and dV_semaphore is None, ( + "SM100 backward with head_dim=256 does not use semaphores" + ) + assert block_sparse_tensors is None, ( + "SM100 backward with head_dim=256 does not support block sparse tensors" + ) + assert aux_data.tensors is None or len(aux_data.tensors) == 0, ( + "SM100 backward with head_dim=256 does not support aux_tensors" + ) + assert aux_data.scalars is None or len(aux_data.scalars) == 0, ( + "SM100 backward with head_dim=256 does not support aux_scalars" + ) + assert dQ_accum is not None, ( + "SM100 backward with head_dim=256 expects dQ tensor at dQ_accum slot" + ) + dQ = dQ_accum + varlen = cumulative_s_q is not None or cumulative_s_k is not None + q_rank = cute.rank(Q.layout) + k_rank = cute.rank(K.layout) + if cutlass.const_expr(q_rank == 5): + h_q = Q.shape[2] * Q.shape[3] + elif cutlass.const_expr(q_rank == 4): + h_q = Q.shape[2] + else: + h_q = Q.shape[1] + if cutlass.const_expr(k_rank == 5): + h_k = K.shape[2] + elif cutlass.const_expr(k_rank == 4): + h_k = K.shape[2] + else: + h_k = K.shape[1] + h_r = h_q // h_k + if cutlass.const_expr(cumulative_s_q is not None): + b = cumulative_s_q.shape[0] - 1 + elif cutlass.const_expr(cumulative_s_k is not None): + b = cumulative_s_k.shape[0] - 1 + else: + b = Q.shape[0] + + Q, K, V, dQ, dK, dV, dO = [assume_tensor_aligned(t) for t in (Q, K, V, dQ, dK, dV, dO)] + + Q = _as_bshkrd_tensor(Q, h_k, h_r, varlen) + K = _as_bshkrd_tensor(K, h_k, 1, varlen) + V = _as_bshkrd_tensor(V, h_k, 1, varlen) + dQ = _as_bshkrd_tensor(dQ, h_k, h_r, varlen) + dK = _as_bshkrd_tensor(dK, h_k, 1, varlen) + dV = _as_bshkrd_tensor(dV, h_k, 1, varlen) + dO = _as_bshkrd_tensor(dO, h_k, h_r, varlen) + scaled_LSE = _as_shhb_tensor(lse_log2, h_k, h_r, b, varlen) + sum_OdO = _as_shhb_tensor(dpsum, h_k, h_r, b, varlen) + + # Keep original order: dQ first, then dKdV. + self.dq_kernel( + Q, + K, + V, + dQ, + dO, + scaled_LSE, + sum_OdO, + cumulative_s_q, + cumulative_s_k, + scale_softmax, + stream, + ) + self.dkdv_kernel( + Q, + K, + V, + dK, + dV, + dO, + scaled_LSE, + sum_OdO, + cumulative_s_q, + cumulative_s_k, + scale_softmax, + stream, + ) diff --git a/fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_backward_dkdvkernel.py b/fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_backward_dkdvkernel.py new file mode 100644 index 000000000..885ae336f --- /dev/null +++ b/fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_backward_dkdvkernel.py @@ -0,0 +1,3411 @@ +# Copyright (c) 2025, Siyu Wang, Shengbin Di, Yuxi Chi, Johnsonms, Linfeng Zheng, Haoyan Huang, Lanbo Li, Yun Zhong, Man Yuan, Minmin Sun, Yong Li, Wei Lin. + +"""Fused multi-head attention (FMHA) backward for the SM100 architecture using CUTE DSL. + +Constraints: +* Supported head dimensions: 256 only +* cta_tiler_mn must be 64,128 +* Batch size must be the same for Q, K, and V tensors +""" + +import enum +import math + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass.cute.nvgpu import cpasync, tcgen05 +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.typing import Int32 +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait + +from cutlass.utils import ClcDynamicPersistentTileScheduler +from flash_attn.cute.tile_scheduler import ( + ClcState, + SM100_TMEM_CAPACITY_COLUMNS, + make_sm100_thread_cooperative_group as make_thread_cooperative_group, + Sm100FmhaClcDynamicTileSchedulerParams as FmhaClcDynamicTileSchedulerParams, + Sm100FmhaClcDynamicTileScheduler as FmhaClcDynamicTileScheduler, + Sm100FmhaStaticTileSchedulerParams as FmhaStaticTileSchedulerParams, +) + +import flash_attn.cute.copy_utils as fa_copy_utils + +LAYOUT_RANK_CONSTANT = 3 + + +@cute.jit +def split_wg( + t: cute.Tensor, + num_warp_groups: Int32, + wg_idx: Int32, +) -> cute.Tensor: + """Split warp group.""" + # dishengbin, TODO：need to double check if more efficient to split in other dimensions + ret = None + if cutlass.const_expr(cute.rank(t.layout) == LAYOUT_RANK_CONSTANT): + p = cute.composition( + t, + cute.make_layout( + ( + t.shape[0], + t.shape[1], + (num_warp_groups, cute.size(t, mode=[2]) // num_warp_groups), + ) + ), + ) + ret = p[None, None, (wg_idx, None)] + else: + p = cute.composition( + t, + cute.make_layout( + ( + t.shape[0], + t.shape[1], + t.shape[2], + (num_warp_groups, cute.size(t, mode=[3]) // num_warp_groups), + ) + ), + ) + ret = p[None, None, None, (wg_idx, None)] + return ret + + +class MaskType(enum.Enum): + """Mask type used in FMHA backward.""" + + NO_MASK = enum.auto() + RESIDUAL_MASK_FOR_BACKWARD = enum.auto() + CAUSAL_MASK_FOR_BACKWARD = enum.auto() + + +def Tmemory_offset(lane, col): + """Tensor memory offset.""" + return (lane << 16) + col + + +permute_order = (0, 1, 2, 3, 4) + + +class BlackwellFusedMultiHeadAttentionBackwardDKDVKernel: + """FMHA backward class for executing CuTeDSL kernel.""" + + def __init__( + self, + acc_dtype: type[cutlass.Numeric], + cta_tiler: tuple[int, int, int], + is_causal: bool, + window_size_left: int | None, + window_size_right: int | None, + use_clc_scheduler: bool = False, + ): + """Initialization.""" + self.acc_dtype = acc_dtype + self.cta_tiler = cta_tiler + self.use_clc_scheduler = use_clc_scheduler + self.sched_warp_id = 10 if use_clc_scheduler else None + # TODO: need check, not sure whether need to *2 if 2cta + self.tile_shape_Q = cta_tiler[0] + self.tile_shape_K = cta_tiler[1] + self.tile_shape_dQ_K = cta_tiler[2] + self.tile_shape_dV_dO = cta_tiler[2] + # For S + self.KQ_mma_tiler = ( + cta_tiler[1] * 2, + cta_tiler[0], + cta_tiler[2], + ) + # For dP + self.VdO_mma_tiler = ( + cta_tiler[1] * 2, + cta_tiler[0], + cta_tiler[2], + ) + # For dV + self.PdO_mma_tiler = ( + cta_tiler[1] * 2, + cta_tiler[2], + cta_tiler[0], + ) + # For dK + self.dSQ_mma_tiler = ( + cta_tiler[1] * 2, + cta_tiler[2], + cta_tiler[0], + ) + # For dQ, dishengbin, need to remove + self.dSK_mma_tiler = ( + cta_tiler[0] * 2, + cta_tiler[2], + cta_tiler[1], + ) + self.cluster_shape_mn = (2, 1) + self.is_causal = is_causal + self.window_size_left: int = -1 if window_size_left is None else window_size_left + self.window_size_right: int = -1 if window_size_right is None else window_size_right + self.has_sliding_window = False + if self.window_size_left > 0 or self.window_size_right > 0: + self.has_sliding_window = True + if self.is_causal: + self.window_size_right = 0 + + self.compute_warp_id = (0, 1, 2, 3, 4, 5, 6, 7) + self.mma_warp_id = 8 + self.load_warp_id = 9 + self.empty_warp_id = 10 + + self.num_compute_warps = 8 + + self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + + self.threads_per_warp = 32 + self.threads_per_cta = self.threads_per_warp * (self.num_compute_warps + 4) + + self.cta_sync_bar_id = 0 + self.tmem_alloc_sync_bar_id = 1 + self.compute_sync_bar_id = 2 + self.epilogue_sync_bar_id = 3 + self.reduce_sync_bar_id = 4 + + self.tmem_dK_offset = 0 + self.tmem_dV_offset = Tmemory_offset(0, cta_tiler[2] // 2) + self.tmem_dP_offset = Tmemory_offset(0, cta_tiler[2] + cta_tiler[0] // 2) + self.tmem_S_offset = Tmemory_offset(0, cta_tiler[2]) + + self.num_regs_reduce = 152 + self.num_regs_compute = 128 + self.num_regs_mma = 128 + self.num_regs_empty = 96 + self.num_regs_load = 96 + + self.buffer_align_bytes = 128 + + def _setup_attributes(self): + """Settings for pipeline stage.""" + self.load_mma_Q_stage = 1 + self.load_mma_K_stage = 1 + self.load_mma_V_stage = 1 + self.load_mma_QT_stage = 1 + self.load_mma_dO_stage = 1 + self.load_compute_LSE_stage = 1 + self.load_compute_sum_OdO_stage = 1 + self.mma_compute_S_stage = 1 + self.mma_compute_dP_stage = 1 + self.compute_mma_P_stage = 1 + self.compute_mma_dS_stage = 1 + self.mma_compute_dKdV_stage = 2 + + if cutlass.const_expr(self.use_clc_scheduler): + self.num_clc_stage = 1 + self.num_clc_response_bytes = 16 + + @cute.jit + def __call__( + self, + Q: cute.Tensor, + K: cute.Tensor, + V: cute.Tensor, + dK: cute.Tensor, + dV: cute.Tensor, + dO: cute.Tensor, + scaled_LSE: cute.Tensor, + sum_OdO: cute.Tensor, + cumulative_s_q: cute.Tensor | None, + cumulative_s_k: cute.Tensor | None, + scale_softmax: cutlass.Float32, + stream: cuda.CUstream, + ): + """Host function to launch CuTeDSL kernel.""" + varlen = cumulative_s_q is not None or cumulative_s_k is not None + # Infer shape metadata from normalized 5D tensors (B, S, H_k, H_r, D). + h_r = Q.shape[3] + h_k = Q.shape[2] + if cutlass.const_expr(cumulative_s_q is not None): + b = cumulative_s_q.shape[0] - 1 + elif cutlass.const_expr(cumulative_s_k is not None): + b = cumulative_s_k.shape[0] - 1 + else: + b = Q.shape[0] + problem_shape = ( + Q.shape[1], + K.shape[1], + Q.shape[4], + ((h_r, h_k), b), + ) + hb = ((h_r, h_k), b) + # (b, s, h_k, h_r, d) -> (s, d, ((h_r, h_k), b)) + Q = cute.make_tensor( + Q.iterator, + cute.make_layout( + (Q.shape[1], Q.shape[4], hb), + stride=( + cute.assume(Q.stride[1], divby=64), + Q.stride[4], + ( + (Q.stride[3], Q.stride[2]), + 0 if cumulative_s_q is not None else cute.assume(Q.stride[0], divby=64), + ), + ), + ), + ) + # (b, s, h_k, 1, d) -> (s, d, ((1, h_k), b)) + K = cute.make_tensor( + K.iterator, + cute.make_layout( + (K.shape[1], K.shape[4], hb), + stride=( + cute.assume(K.stride[1], divby=64), + K.stride[4], + ( + (0, K.stride[2]), + 0 if cumulative_s_k is not None else cute.assume(K.stride[0], divby=64), + ), + ), + ), + ) + # (b, s, h_k, 1, d) -> (s, d, ((1, h_k), b)) + V = cute.make_tensor( + V.iterator, + cute.make_layout( + (V.shape[1], V.shape[4], hb), + stride=( + cute.assume(V.stride[1], divby=64), + V.stride[4], + ( + (0, V.stride[2]), + 0 if cumulative_s_k is not None else cute.assume(V.stride[0], divby=64), + ), + ), + ), + ) + # (s, d, ((h_r, h_k), b)) -> (d, s, ((h_r, h_k), b)) + QT = cute.make_tensor( + Q.iterator, + cute.make_layout( + (Q.shape[1], Q.shape[0], Q.shape[2]), + stride=( + Q.stride[1], + Q.stride[0], + Q.stride[2], + ), + ), + ) + dK = cute.make_tensor( + dK.iterator, + cute.make_layout( + (dK.shape[1], dK.shape[4], hb), + stride=( + cute.assume(dK.stride[1], divby=64), + dK.stride[4], + ( + (0, dK.stride[2]), + 0 if cumulative_s_k is not None else cute.assume(dK.stride[0], divby=64), + ), + ), + ), + ) + dV = cute.make_tensor( + dV.iterator, + cute.make_layout( + (dV.shape[1], dV.shape[4], hb), + stride=( + cute.assume(dV.stride[1], divby=64), + dV.stride[4], + ( + (0, dV.stride[2]), + 0 if cumulative_s_k is not None else cute.assume(dV.stride[0], divby=64), + ), + ), + ), + ) + # (s, d, ((h_r, h_k), b)) + dO = cute.make_tensor( + dO.iterator, + cute.make_layout( + (dO.shape[1], dO.shape[4], hb), + stride=( + cute.assume(dO.stride[1], divby=64), + dO.stride[4], + ( + (dO.stride[3], dO.stride[2]), + 0 if cumulative_s_q is not None else cute.assume(dO.stride[0], divby=64), + ), + ), + ), + ) + + # (s, d, ((h_r, h_k), b)) -> (d, s, ((h_r, h_k), b)) + dOT = cute.make_tensor( + dO.iterator, + cute.make_layout( + (dO.shape[1], dO.shape[0], dO.shape[2]), + stride=( + dO.stride[1], + dO.stride[0], + dO.stride[2], + ), + ), + ) + + self.Q_major_mode = utils.LayoutEnum.from_tensor(Q).mma_major_mode() + self.K_major_mode = utils.LayoutEnum.from_tensor(K).mma_major_mode() + self.dK_major_mode = utils.LayoutEnum.from_tensor(dK).mma_major_mode() + self.V_major_mode = utils.LayoutEnum.from_tensor(V).mma_major_mode() + self.dV_major_mode = utils.LayoutEnum.from_tensor(dV).mma_major_mode() + self.dO_major_mode = utils.LayoutEnum.from_tensor(dO).mma_major_mode() + + if cutlass.const_expr(self.Q_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError(f"The layout of q is not supported: {self.Q_major_mode}") + if cutlass.const_expr(self.K_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of k is not supported") + if cutlass.const_expr(self.dK_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of dk is not supported") + if cutlass.const_expr(self.V_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of v is not supported") + if cutlass.const_expr(self.dV_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of dv is not supported") + + self._setup_attributes() + + cta_group = tcgen05.CtaGroup.TWO + PT_source = tcgen05.OperandSource.SMEM + + # compute S + KQ_tiled_mma = sm100_utils.make_trivial_tiled_mma( + K.element_type, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + self.acc_dtype, + cta_group, + self.KQ_mma_tiler[:2], + ) + # compute dP + VdO_tiled_mma = sm100_utils.make_trivial_tiled_mma( + V.element_type, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.K, + self.acc_dtype, + cta_group, + self.VdO_mma_tiler[:2], + ) + # compute dV + PdO_tiled_mma = sm100_utils.make_trivial_tiled_mma( + dO.element_type, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.MN, + self.acc_dtype, + cta_group, + self.PdO_mma_tiler[:2], + PT_source, + ) + # compute dK + dSQ_tiled_mma = sm100_utils.make_trivial_tiled_mma( + Q.element_type, + tcgen05.OperandMajorMode.K, + tcgen05.OperandMajorMode.MN, + self.acc_dtype, + cta_group, + self.dSQ_mma_tiler[:2], + ) + # compute dQ + # dishengbin, need to remove, but used in dS_mem_layout_staged + dSK_tiled_mma = sm100_utils.make_trivial_tiled_mma( + K.element_type, + tcgen05.OperandMajorMode.MN, + tcgen05.OperandMajorMode.MN, + self.acc_dtype, + cta_group, + self.dSK_mma_tiler[:2], + ) + + atom_thr_size = cute.size(KQ_tiled_mma.thr_id.shape) + self.cluster_shape_mnk = (*self.cluster_shape_mn, 1) # type: ignore[assignment] + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (atom_thr_size,), + ) + + K_smem_layout_staged = sm100_utils.make_smem_layout_a( + KQ_tiled_mma, + self.KQ_mma_tiler, + K.element_type, + 1, + ) + Q_smem_layout_staged = sm100_utils.make_smem_layout_b( + KQ_tiled_mma, + self.KQ_mma_tiler, + Q.element_type, + self.load_mma_Q_stage, + ) + V_smem_layout_staged = sm100_utils.make_smem_layout_a( + VdO_tiled_mma, + self.VdO_mma_tiler, + V.element_type, + 1, + ) + dO_smem_layout_staged = sm100_utils.make_smem_layout_b( + VdO_tiled_mma, + self.VdO_mma_tiler, + dO.element_type, + self.load_mma_dO_stage, + ) + # dishengbin, need to remove, but used for sPT, need to double check + dS_smem_layout_staged = sm100_utils.make_smem_layout_a( + dSK_tiled_mma, + self.dSK_mma_tiler, + Q.element_type, + self.compute_mma_dS_stage, + ) + dST_smem_layout_staged = sm100_utils.make_smem_layout_a( + dSQ_tiled_mma, + self.dSQ_mma_tiler, + Q.element_type, + self.compute_mma_dS_stage, + ) + tiled_mma = dSQ_tiled_mma + is_k_major = tiled_mma.op.a_major_mode == tcgen05.OperandMajorMode.K + a_major_mode = tcgen05.OperandMajorMode.K if is_k_major else tcgen05.OperandMajorMode.MN + tmp = cute.dice(self.dSQ_mma_tiler, (1, None, 1)) + a_smem_shape = tiled_mma.partition_shape_A( + cute.dice(self.dSQ_mma_tiler, (1, None, 1)), + ) + a_smem_shape_mn_k = ( + cute.size(a_smem_shape[0][0]) * a_smem_shape[1], + cute.size(a_smem_shape[0][1]) * a_smem_shape[2], + ) + smem_layout_atom_kind = sm100_utils.get_smem_layout_atom_ab( + a_major_mode, + K.element_type, + a_smem_shape_mn_k, + ) + a_smem_layout_atom = sm100_utils.make_smem_layout_atom( + smem_layout_atom_kind, + K.element_type, + ) + + a_smem_shape = cute.append( + a_smem_shape, + self.compute_mma_dS_stage, + ) + order = (2, 1, 3) if not is_k_major else (1, 2, 3) + dST_smem_layout_staged_tmp = sm100_utils.tile_to_mma_shape( + a_smem_layout_atom, + a_smem_shape, + order=order, + ) + QT_smem_layout_staged = sm100_utils.make_smem_layout_b( + dSQ_tiled_mma, + self.dSQ_mma_tiler, + Q.element_type, + self.load_mma_QT_stage, + ) + P_smem_layout_staged = sm100_utils.make_smem_layout_a( + PdO_tiled_mma, + self.PdO_mma_tiler, + Q.element_type, + self.compute_mma_P_stage, + ) + dOT_smem_layout_staged = sm100_utils.make_smem_layout_b( + PdO_tiled_mma, + self.PdO_mma_tiler, + dO.element_type, + self.load_mma_dO_stage, + ) + LSE_smem_layout = cute.make_layout((self.cta_tiler[0], self.load_compute_LSE_stage)) + sum_OdO_smem_layout = cute.make_layout((self.cta_tiler[0], self.load_compute_sum_OdO_stage)) + + tma_load_op = cpasync.CopyBulkTensorTileG2SOp(cta_group) + tma_reduce_op = cpasync.CopyReduceBulkTensorTileS2GOp() + + K_smem_layout = cute.select(K_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_K, tma_tensor_K = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + K, + K_smem_layout, + self.KQ_mma_tiler, + KQ_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + V_smem_layout = cute.select(V_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_V, tma_tensor_V = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + V, + V_smem_layout, + self.VdO_mma_tiler, + VdO_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + Q_smem_layout = cute.select(Q_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_Q, tma_tensor_Q = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + Q, + Q_smem_layout, + self.KQ_mma_tiler, + KQ_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + QT_smem_layout = cute.select(QT_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_QT, tma_tensor_QT = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + QT, + QT_smem_layout, + self.dSQ_mma_tiler, + dSQ_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + dO_smem_layout = cute.select(dO_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_dO, tma_tensor_dO = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + dO, + dO_smem_layout, + self.VdO_mma_tiler, + VdO_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + dOT_smem_layout = cute.select(dOT_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_dOT, tma_tensor_dOT = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + dOT, + dOT_smem_layout, + self.PdO_mma_tiler, + PdO_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # for 2cta, tma_copy_QT_bytes is same as the tma_copy_Q_bytes + self.tma_copy_Q_bytes = cute.size_in_bytes(Q.element_type, Q_smem_layout) * atom_thr_size + self.tma_copy_K_bytes = cute.size_in_bytes(K.element_type, K_smem_layout) * atom_thr_size + self.tma_copy_V_bytes = cute.size_in_bytes(V.element_type, V_smem_layout) * atom_thr_size + self.tma_copy_dO_bytes = cute.size_in_bytes(dO.element_type, dO_smem_layout) * atom_thr_size + + # Variant 3a epilogue: TMA store atoms (S2G) for dK / dV. + # Each compute warp group owns half the hd_v output via split_wg, so + # the per-WG epi tile is (cta_tiler[1], cta_tiler[2] / num_compute_wgs). + # SMEM staging will alias onto sP+sdST in subsequent commits; for now + # the atoms and layouts are built and threaded through but unused. + tma_store_op = cpasync.CopyBulkTensorTileS2GOp() + num_compute_wgs = self.num_compute_warps // 4 + # Variant 3a Path 2: CTA-shared epilogue SMEM. epi_tile is (M, gcd(128B, ...)) + # = (M, 64) for bf16 hd=256. Total stages = num_compute_wgs * num_epi_stages + # = 4 stages of (64, 64), virtually a per-CTA (64, 256) buffer aliased onto + # sP+sdST. Both warp-groups cooperatively populate this buffer; TMA fires + # one (64, 64) box per stage to the corresponding (64, 64) GMEM slice. + epi_cols_dKV = math.gcd( + 128 // (dK.element_type.width // 8), self.cta_tiler[2] // num_compute_wgs + ) + num_epi_stages_dKV = (self.cta_tiler[2] // num_compute_wgs) // epi_cols_dKV + epi_tile_dKV = (self.cta_tiler[1], epi_cols_dKV) + total_epi_stages = num_compute_wgs * num_epi_stages_dKV + dK_layout_enum = utils.LayoutEnum.from_tensor(dK) + dV_layout_enum = utils.LayoutEnum.from_tensor(dV) + sdK_epi_layout = sm100_utils.make_smem_layout_epi( + dK.element_type, + dK_layout_enum, + epi_tile_dKV, + total_epi_stages, + ) + sdV_epi_layout = sm100_utils.make_smem_layout_epi( + dV.element_type, + dV_layout_enum, + epi_tile_dKV, + total_epi_stages, + ) + tma_atom_dK, tma_tensor_dK = cpasync.make_tiled_tma_atom( + tma_store_op, + dK, + cute.select(sdK_epi_layout, mode=[0, 1]), + epi_tile_dKV, + ) + tma_atom_dV, tma_tensor_dV = cpasync.make_tiled_tma_atom( + tma_store_op, + dV, + cute.select(sdV_epi_layout, mode=[0, 1]), + epi_tile_dKV, + ) + + @cute.struct + class SharedStorage: + # Pipeline barriers + load_mma_Q_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_mma_Q_stage * 2] + load_mma_K_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_mma_K_stage * 2] + load_mma_V_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_mma_V_stage * 2] + load_mma_QT_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_mma_QT_stage * 2] + load_mma_dO_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_mma_dO_stage * 2] + load_mma_dOT_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_mma_dO_stage * 2] + load_compute_lse_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.load_compute_LSE_stage * 2 + ] + load_compute_sum_OdO_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.load_compute_sum_OdO_stage * 2 + ] + mma_compute_S_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.mma_compute_S_stage * 2 + ] + mma_compute_dP_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.mma_compute_dP_stage * 2 + ] + compute_mma_P_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.compute_mma_P_stage * 2 + ] + compute_mma_dS_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.compute_mma_dS_stage * 2 + ] + mma_compute_dKdV_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.mma_compute_dKdV_stage * 2 + ] + tmem_holding_buf: cutlass.Int32 + tmem_dealloc_mbar_ptr: cutlass.Int64 + clc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2] + clc_response: cute.struct.MemRange[Int32, 4] + # Smem tensors + sK: cute.struct.Align[ + cute.struct.MemRange[K.element_type, cute.cosize(K_smem_layout_staged)], + self.buffer_align_bytes, + ] + # only used in 2cta + sV: cute.struct.Align[ + cute.struct.MemRange[V.element_type, cute.cosize(V_smem_layout_staged)], + self.buffer_align_bytes, + ] + sQ: cute.struct.Align[ + cute.struct.MemRange[Q.element_type, cute.cosize(Q_smem_layout_staged)], + self.buffer_align_bytes, + ] + sQT: cute.struct.Align[ + cute.struct.MemRange[Q.element_type, cute.cosize(QT_smem_layout_staged)], + self.buffer_align_bytes, + ] + sdO: cute.struct.Align[ + cute.struct.MemRange[dO.element_type, cute.cosize(dO_smem_layout_staged)], + self.buffer_align_bytes, + ] + sdOT: cute.struct.Align[ + cute.struct.MemRange[dO.element_type, cute.cosize(dOT_smem_layout_staged)], + self.buffer_align_bytes, + ] + # only used in 2cta + # dishengbin checked whether we need sP + sP: cute.struct.Align[ + cute.struct.MemRange[Q.element_type, cute.cosize(P_smem_layout_staged)], + self.buffer_align_bytes, + ] + sdST: cute.struct.Align[ + cute.struct.MemRange[Q.element_type, cute.cosize(dST_smem_layout_staged)], + self.buffer_align_bytes, + ] + + sLSE: cute.struct.Align[ + cute.struct.MemRange[self.acc_dtype, cute.cosize(LSE_smem_layout)], + self.buffer_align_bytes, + ] + sSum_OdO: cute.struct.Align[ + cute.struct.MemRange[self.acc_dtype, cute.cosize(sum_OdO_smem_layout)], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + # =============================== bwd =============================== + K_val = problem_shape[1] + _, H_K = problem_shape[3][0] + B = problem_shape[3][1] + problem_shape_mbh = ( + cute.ceil_div(K_val, self.cta_tiler[1]), + cute.size(B), + cute.size(H_K), + ) + if cutlass.const_expr(self.use_clc_scheduler): + self.tile_sched_params = FmhaClcDynamicTileSchedulerParams( + problem_shape_mbh, + (*self.cluster_shape_mn, 1), + ) + bwd_grid = FmhaClcDynamicTileScheduler.get_grid_shape(self.tile_sched_params) + else: + self.tile_sched_params = FmhaStaticTileSchedulerParams( + is_persistent=False, + problem_shape_mbh=problem_shape_mbh, + ) + bwd_grid = self._compute_bwd_grid(problem_shape, self.cta_tiler[1]) + bwd_grid = cute.round_up(bwd_grid, self.cluster_shape_mnk) + + self.dkdv_bwd( + KQ_tiled_mma, + VdO_tiled_mma, + PdO_tiled_mma, + dSQ_tiled_mma, + tma_atom_K, + tma_tensor_K, + K, + tma_atom_V, + tma_tensor_V, + tma_atom_Q, + tma_tensor_Q, + Q, + tma_atom_QT, + tma_tensor_QT, + tma_atom_dO, + tma_tensor_dO, + tma_atom_dOT, + tma_tensor_dOT, + dK, + dV, + tma_atom_dK, + tma_tensor_dK, + tma_atom_dV, + tma_tensor_dV, + scaled_LSE, + scale_softmax, + sum_OdO, + problem_shape, + cumulative_s_q, + cumulative_s_k, + self.cluster_layout_vmnk, + K_smem_layout_staged, + Q_smem_layout_staged, + V_smem_layout_staged, + dO_smem_layout_staged, + dS_smem_layout_staged, + dST_smem_layout_staged, + QT_smem_layout_staged, + dOT_smem_layout_staged, + P_smem_layout_staged, + LSE_smem_layout, + sum_OdO_smem_layout, + sdK_epi_layout, + sdV_epi_layout, + self.tile_sched_params, + ).launch( + grid=bwd_grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + smem=self.shared_storage.size_in_bytes(), # type: ignore [attr-defined] + stream=stream, + min_blocks_per_mp=1, + ) + + @cute.kernel + def dkdv_bwd( + self, + KQ_tiled_mma: cute.TiledMma, + VdO_tiled_mma: cute.TiledMma, + PdO_tiled_mma: cute.TiledMma, + dSQ_tiled_mma: cute.TiledMma, + tma_atom_K: cute.CopyAtom, + K_in: cute.Tensor, + K_ref: cute.Tensor, + tma_atom_V: cute.CopyAtom, + V_in: cute.Tensor, + tma_atom_Q: cute.CopyAtom, + Q_in: cute.Tensor, + Q_ref: cute.Tensor, + tma_atom_QT: cute.CopyAtom, + QT_in: cute.Tensor, + tma_atom_dO: cute.CopyAtom, + dO_in: cute.Tensor, + tma_atom_dOT: cute.CopyAtom, + dOT_in: cute.Tensor, + dK: cute.Tensor, + dV: cute.Tensor, + tma_atom_dK: cute.CopyAtom, + dK_tma: cute.Tensor, + tma_atom_dV: cute.CopyAtom, + dV_tma: cute.Tensor, + LSE: cute.Tensor, + scale_softmax: cutlass.Float32, + sum_OdO: cute.Tensor, + problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]], + cumulative_s_q: cute.Tensor | None, + cumulative_s_k: cute.Tensor | None, + cluster_layout_vmnk: cute.Layout, + K_smem_layout_staged: cute.ComposedLayout, + Q_smem_layout_staged: cute.ComposedLayout, + V_smem_layout_staged: cute.ComposedLayout, + dO_smem_layout_staged: cute.ComposedLayout, + dS_smem_layout_staged: cute.ComposedLayout, + dST_smem_layout_staged: cute.ComposedLayout, + QT_smem_layout_staged: cute.ComposedLayout, + dOT_smem_layout_staged: cute.ComposedLayout, + P_smem_layout_staged: cute.ComposedLayout, + LSE_smem_layout: cute.Layout, + sum_OdO_smem_layout: cute.Layout, + sdK_epi_layout: cute.ComposedLayout, + sdV_epi_layout: cute.ComposedLayout, + tile_sched_params: FmhaStaticTileSchedulerParams | FmhaClcDynamicTileSchedulerParams, + ): + """Core CuTeDSL backward kernel.""" + bidx, bidy, bidz = cute.arch.block_idx() + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + varlen = cumulative_s_q is not None or cumulative_s_k is not None + + mma_tile_coord_v = bidx % cute.size(KQ_tiled_mma.thr_id.shape) + + if warp_idx == self.load_warp_id: + cpasync.prefetch_descriptor(tma_atom_K) + cpasync.prefetch_descriptor(tma_atom_Q) + cpasync.prefetch_descriptor(tma_atom_QT) + cpasync.prefetch_descriptor(tma_atom_V) + cpasync.prefetch_descriptor(tma_atom_dO) + cpasync.prefetch_descriptor(tma_atom_dOT) + + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + load_mma_Q_producer, load_mma_Q_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.load_mma_Q_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_Q_bytes, + barrier_storage=storage.load_mma_Q_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + load_mma_K_producer, load_mma_K_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.load_mma_K_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_K_bytes, + barrier_storage=storage.load_mma_K_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + load_mma_V_producer, load_mma_V_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.load_mma_V_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_V_bytes, + barrier_storage=storage.load_mma_V_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + load_mma_QT_producer, load_mma_QT_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.load_mma_QT_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_Q_bytes, + barrier_storage=storage.load_mma_QT_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + load_mma_dO_producer, load_mma_dO_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.load_mma_dO_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_dO_bytes, + barrier_storage=storage.load_mma_dO_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + load_mma_dOT_producer, load_mma_dOT_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.load_mma_dO_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_dO_bytes, + barrier_storage=storage.load_mma_dOT_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + load_compute_LSE_producer, load_compute_LSE_consumer = pipeline.PipelineCpAsync.create( + num_stages=self.load_compute_LSE_stage, + producer_group=make_thread_cooperative_group(self.threads_per_warp), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * self.num_compute_warps + ), + barrier_storage=storage.load_compute_lse_mbar_ptr.data_ptr(), + ).make_participants() + load_compute_sum_OdO_producer, load_compute_sum_OdO_consumer = ( + pipeline.PipelineCpAsync.create( + num_stages=self.load_compute_sum_OdO_stage, + producer_group=make_thread_cooperative_group(self.threads_per_warp), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * self.num_compute_warps + ), + barrier_storage=storage.load_compute_sum_OdO_mbar_ptr.data_ptr(), + ).make_participants() + ) + mma_compute_S_producer, mma_compute_S_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.mma_compute_S_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group( + self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0] + ), + barrier_storage=storage.mma_compute_S_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + mma_compute_dP_producer, mma_compute_dP_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.mma_compute_dP_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group( + self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0] + ), + barrier_storage=storage.mma_compute_dP_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + compute_mma_P_producer, compute_mma_P_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.compute_mma_P_stage, + producer_group=make_thread_cooperative_group( + self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0] + ), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + barrier_storage=storage.compute_mma_P_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + compute_mma_dS_producer, compute_mma_dS_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.compute_mma_dS_stage, + producer_group=make_thread_cooperative_group( + self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0] + ), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + barrier_storage=storage.compute_mma_dS_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + mma_compute_dKdV_producer, mma_compute_dKdV_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.mma_compute_dKdV_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group( + self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0] + ), + barrier_storage=storage.mma_compute_dKdV_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + + cute.arch.barrier(barrier_id=self.cta_sync_bar_id, number_of_threads=self.threads_per_cta) + + # setup mma + sQ = storage.sQ.get_tensor(Q_smem_layout_staged.outer, swizzle=Q_smem_layout_staged.inner) + sK = storage.sK.get_tensor(K_smem_layout_staged.outer, swizzle=K_smem_layout_staged.inner) + sV = storage.sV.get_tensor(V_smem_layout_staged.outer, swizzle=V_smem_layout_staged.inner) + sdO = storage.sdO.get_tensor( + dO_smem_layout_staged.outer, swizzle=dO_smem_layout_staged.inner + ) + sLSE = storage.sLSE.get_tensor(LSE_smem_layout) + sSum_OdO = storage.sSum_OdO.get_tensor(sum_OdO_smem_layout) + tmem_holding_buf = storage.tmem_holding_buf + # for 2cta, QT use different mem from Q + + sQT = storage.sQT.get_tensor( + QT_smem_layout_staged.outer, swizzle=QT_smem_layout_staged.inner + ) + sdST = storage.sdST.get_tensor( + dST_smem_layout_staged.outer, swizzle=dST_smem_layout_staged.inner + ) + tP_fake_ptr = cute.make_ptr(sQ.element_type, 0, cute.AddressSpace.tmem) + tP = cute.make_tensor(tP_fake_ptr, P_smem_layout_staged.outer) + + sP = storage.sP.get_tensor(P_smem_layout_staged.outer, swizzle=P_smem_layout_staged.inner) + + sdOT = storage.sdOT.get_tensor( + dOT_smem_layout_staged.outer, swizzle=dOT_smem_layout_staged.inner + ) + + # tSTrK shape : (MMA, MMA_M, MMA_K, STAGE) + tSTrK = KQ_tiled_mma.make_fragment_A(sK) + # tSTrQ shape : (MMA, MMA_N, MMA_K, STAGE) + tSTrQ = KQ_tiled_mma.make_fragment_B(sQ) + + # tdPTrV shape : (MMA, MMA_M, MMA_K, STAGE) + tdPTrV = VdO_tiled_mma.make_fragment_A(sV) + # tdPTrdO shape : (MMA, MMA_N, MMA_K, STAGE) + tdPTrdO = VdO_tiled_mma.make_fragment_B(sdO) + + # tdKrdST shape: (MMA, MMA_M, MMA_K, STAGE) + tdKrdST = dSQ_tiled_mma.make_fragment_A(sdST) + # tdKrQT shape : (MMA, MMA_N, MMA_K, STAGE) + tdKrQT = dSQ_tiled_mma.make_fragment_B(sQT) + + tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=self.tmem_alloc_sync_bar_id, + num_threads=self.threads_per_cta, + ) + + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=tmem_alloc_barrier, + allocator_warp_id=self.load_warp_id, + is_two_cta=True, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + + tmem.allocate(self.tmem_alloc_cols) + + # wait for tmem allocation and retrieve the pointer + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + # Cluster arrive after barrier init + # is_relaxed=False has memory consistency guarantee + pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=False) + + if cutlass.const_expr(self.use_clc_scheduler): + clc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + cluster_size = cute.size(self.cluster_shape_mnk) + num_clc_consumer_threads = self.threads_per_warp * ( + 1 # sched_warp (CTA 0 only) + + cluster_size + * ( + len(self.compute_warp_id) + + 1 # mma_warp + + 1 # load_warp + ) + ) + clc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_clc_consumer_threads + ) + clc_response_ptr = storage.clc_response.data_ptr() + clc = ClcState.create( + hw_scheduler=ClcDynamicPersistentTileScheduler.create( + self.tile_sched_params.clc_hw_params(), + cute.arch.block_idx(), + cute.arch.grid_dim(), + clc_response_ptr, + ), + pipeline=pipeline.PipelineClcFetchAsync.create( + barrier_storage=storage.clc_mbar_ptr.data_ptr(), + num_stages=self.num_clc_stage, + producer_group=clc_pipeline_producer_group, + consumer_group=clc_pipeline_consumer_group, + tx_count=self.num_clc_response_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ), + consumer_state=pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_clc_stage + ), + producer_state=pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_clc_stage + ), + ) + tile_sched = FmhaClcDynamicTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + clc_response_ptr, + clc, + ) + work_tile = tile_sched.initial_work_tile_info() + else: + clc = None + clc_response_ptr = None + + tSTtST_shape = KQ_tiled_mma.partition_shape_C(cute.select(self.KQ_mma_tiler, mode=[0, 1])) + tSTtST = KQ_tiled_mma.make_fragment_C(tSTtST_shape) + # tSTtST shape : (MMA, MMA_M, MMA_N) + tSTtST = cute.make_tensor(tmem_ptr + self.tmem_S_offset, tSTtST.layout) + + # tdVrP shape : (MMA, MMA_M, MMA_K, STAGE) + tdVrP = PdO_tiled_mma.make_fragment_A(sP) + # tdVrdOT shape : (MMA, MMA_N, MMA_K, STAGE) + tdVrdOT = PdO_tiled_mma.make_fragment_B(sdOT) + + tdPTtdPT_shape = VdO_tiled_mma.partition_shape_C( + cute.select(self.VdO_mma_tiler, mode=[0, 1]) + ) + tdPTtdPT = VdO_tiled_mma.make_fragment_C(tdPTtdPT_shape) + # tdPTtdPT shape : (MMA, MMA_M, MMA_N) + tdPTtdPT = cute.make_tensor(tmem_ptr + self.tmem_dP_offset, tdPTtdPT.layout) + + tdKtdK_shape = dSQ_tiled_mma.partition_shape_C(cute.select(self.dSQ_mma_tiler, mode=[0, 1])) + tdKtdK = dSQ_tiled_mma.make_fragment_C(tdKtdK_shape) + # tdKtdK shape : (MMA, MMA_M, MMA_N) + tdKtdK = cute.make_tensor(tmem_ptr + self.tmem_dK_offset, tdKtdK.layout) + + tdVtdV_shape = PdO_tiled_mma.partition_shape_C(cute.select(self.PdO_mma_tiler, mode=[0, 1])) + tdVtdV = PdO_tiled_mma.make_fragment_C(tdVtdV_shape) + # tdVtdV shape : (MMA, MMA_M, MMA_N) + tdVtdV = cute.make_tensor(tmem_ptr + self.tmem_dV_offset, tdVtdV.layout) + + # get the current batch problem shape + + if cutlass.const_expr(self.use_clc_scheduler): + # ===== CLC PERSISTENT PATH: per-warp while loops ===== + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + is_first_cta_in_cluster = cta_rank_in_cluster == 0 + is_sched_warp = warp_idx == self.sched_warp_id and is_first_cta_in_cluster + + # Register allocation ONCE (before any loop) + if warp_idx == self.load_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_load) + elif warp_idx == self.mma_warp_id: + cute.arch.warpgroup_reg_alloc(self.num_regs_mma) + elif warp_idx >= self.compute_warp_id[0] and warp_idx <= self.compute_warp_id[-1]: + cute.arch.warpgroup_reg_alloc(self.num_regs_compute) + else: + cute.arch.warpgroup_reg_dealloc(self.num_regs_empty) + + # Cluster wait + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + # ===== SCHEDULER WARP ===== + if is_sched_warp: + while work_tile.is_valid_tile: + tile_sched.prefetch_next_work() + work_tile = tile_sched.advance_to_next_work() + tile_sched.producer_tail() + + # ===== LOAD WARP ===== + elif warp_idx == self.load_warp_id: + while work_tile.is_valid_tile: + # Decode coordinates from CLC work tile + blk_coord_k = work_tile.tile_idx[0] + blk_coord_b = work_tile.tile_idx[2][0] + blk_coord_h_k = work_tile.tile_idx[2][1] + blk_coord = ( + Int32(0), + blk_coord_k, + Int32(0), + ((Int32(0), blk_coord_h_k), blk_coord_b), + ) + seqlen_q_cur_batch = Q_ref.shape[0] + seqlen_k_cur_batch = K_ref.shape[0] + blk_offset = (Int32(0), Int32(0), Int32(0), ((Int32(0), Int32(0)), Int32(0))) + if cutlass.const_expr(varlen): + assert isinstance(cumulative_s_q, cute.Tensor) + assert isinstance(cumulative_s_k, cute.Tensor) + seqlen_q_cur_batch = ( + cumulative_s_q[blk_coord_b + 1] - cumulative_s_q[blk_coord_b] + ) + seqlen_k_cur_batch = ( + cumulative_s_k[blk_coord_b + 1] - cumulative_s_k[blk_coord_b] + ) + blk_offset = ( + cumulative_s_q[blk_coord_b], + cumulative_s_k[blk_coord_b], + Int32(0), + ((Int32(0), Int32(0)), Int32(0)), + ) + iter_start, iter_end = self.get_Q_block_min_max( + seqlen_q_cur_batch, + seqlen_k_cur_batch, + blk_coord_k, + is_2cta=True, + ) + iter_count = (iter_end - iter_start) * problem_shape[3][0][0] + if iter_count <= 0: + if blk_coord_k * self.tile_shape_K < seqlen_k_cur_batch: + problem_shape_cur_batch = ( + seqlen_q_cur_batch, + seqlen_k_cur_batch, + problem_shape[2], + problem_shape[3], + ) + self.epilogue_clear( + blk_coord, blk_offset, problem_shape_cur_batch, dK, dV + ) + else: + problem_shape_cur_batch = ( + seqlen_q_cur_batch, + seqlen_k_cur_batch, + problem_shape[2], + problem_shape[3], + ) + ( + load_mma_Q_producer, + load_mma_K_producer, + load_mma_V_producer, + load_compute_LSE_producer, + load_mma_dO_producer, + load_mma_dOT_producer, + load_compute_sum_OdO_producer, + load_mma_QT_producer, + ) = self.load( + K_in, + V_in, + Q_in, + QT_in, + dO_in, + dOT_in, + LSE, + sum_OdO, + sK, + sQ, + sQT, + sV, + sdO, + sdOT, + sLSE, + sSum_OdO, + KQ_tiled_mma, + VdO_tiled_mma, + PdO_tiled_mma, + dSQ_tiled_mma, + tma_atom_K, + tma_atom_Q, + tma_atom_QT, + tma_atom_V, + tma_atom_dO, + tma_atom_dOT, + blk_offset, + problem_shape_cur_batch, + varlen, + iter_count, + iter_start, + iter_end, + load_mma_Q_producer, + load_mma_Q_consumer, + load_mma_K_producer, + load_mma_K_consumer, + load_mma_V_producer, + load_mma_V_consumer, + load_compute_LSE_producer, + load_compute_LSE_consumer, + load_mma_dO_producer, + load_mma_dO_consumer, + load_mma_dOT_producer, + load_mma_dOT_consumer, + load_compute_sum_OdO_producer, + load_compute_sum_OdO_consumer, + load_mma_QT_producer, + load_mma_QT_consumer, + blk_coord_k, + blk_coord_h_k, + blk_coord_b, + ) + # CLC advance + work_tile = tile_sched.advance_to_next_work() + # producer_tail after loop + load_mma_K_producer.tail() + load_mma_V_producer.tail() + load_mma_Q_producer.tail() + load_compute_LSE_producer.tail() + load_mma_dO_producer.tail() + load_mma_dOT_producer.tail() + load_compute_sum_OdO_producer.tail() + load_mma_QT_producer.tail() + + # ===== MMA WARP ===== + elif warp_idx == self.mma_warp_id: + while work_tile.is_valid_tile: + blk_coord_k = work_tile.tile_idx[0] + blk_coord_b = work_tile.tile_idx[2][0] + blk_coord_h_k = work_tile.tile_idx[2][1] + blk_coord = ( + Int32(0), + blk_coord_k, + Int32(0), + ((Int32(0), blk_coord_h_k), blk_coord_b), + ) + seqlen_q_cur_batch = Q_ref.shape[0] + seqlen_k_cur_batch = K_ref.shape[0] + blk_offset = (Int32(0), Int32(0), Int32(0), ((Int32(0), Int32(0)), Int32(0))) + if cutlass.const_expr(varlen): + assert isinstance(cumulative_s_q, cute.Tensor) + assert isinstance(cumulative_s_k, cute.Tensor) + seqlen_q_cur_batch = ( + cumulative_s_q[blk_coord_b + 1] - cumulative_s_q[blk_coord_b] + ) + seqlen_k_cur_batch = ( + cumulative_s_k[blk_coord_b + 1] - cumulative_s_k[blk_coord_b] + ) + blk_offset = ( + cumulative_s_q[blk_coord_b], + cumulative_s_k[blk_coord_b], + Int32(0), + ((Int32(0), Int32(0)), Int32(0)), + ) + iter_start, iter_end = self.get_Q_block_min_max( + seqlen_q_cur_batch, + seqlen_k_cur_batch, + blk_coord_k, + is_2cta=True, + ) + iter_count = (iter_end - iter_start) * problem_shape[3][0][0] + if iter_count <= 0: + if blk_coord_k * self.tile_shape_K < seqlen_k_cur_batch: + problem_shape_cur_batch = ( + seqlen_q_cur_batch, + seqlen_k_cur_batch, + problem_shape[2], + problem_shape[3], + ) + self.epilogue_clear( + blk_coord, blk_offset, problem_shape_cur_batch, dK, dV + ) + else: + ( + mma_compute_S_producer, + mma_compute_dP_producer, + mma_compute_dKdV_producer, + load_mma_Q_consumer, + load_mma_K_consumer, + load_mma_V_consumer, + load_mma_dO_consumer, + load_mma_dOT_consumer, + compute_mma_P_consumer, + compute_mma_dS_consumer, + load_mma_QT_consumer, + ) = self.mma_2cta( + KQ_tiled_mma, + VdO_tiled_mma, + PdO_tiled_mma, + dSQ_tiled_mma, + tSTtST, + tSTrQ, + tSTrK, + tdPTtdPT, + tdPTrV, + tdPTrdO, + tdVtdV, + tdVrP, + tdVrdOT, + tdKrdST, + tdKtdK, + tdKrQT, + iter_count, + load_mma_Q_consumer, + load_mma_K_consumer, + load_mma_V_consumer, + mma_compute_S_producer, + load_mma_dO_consumer, + mma_compute_dP_producer, + load_mma_dOT_consumer, + compute_mma_P_consumer, + compute_mma_dS_consumer, + load_mma_QT_consumer, + mma_compute_dKdV_producer, + ) + # CLC advance + work_tile = tile_sched.advance_to_next_work() + # producer_tail after loop + mma_compute_S_producer.tail() + mma_compute_dP_producer.tail() + mma_compute_dKdV_producer.tail() + + # ===== COMPUTE WARPS ===== + elif warp_idx >= self.compute_warp_id[0] and warp_idx <= self.compute_warp_id[-1]: + while work_tile.is_valid_tile: + blk_coord_k = work_tile.tile_idx[0] + blk_coord_b = work_tile.tile_idx[2][0] + blk_coord_h_k = work_tile.tile_idx[2][1] + blk_coord = ( + Int32(0), + blk_coord_k, + Int32(0), + ((Int32(0), blk_coord_h_k), blk_coord_b), + ) + seqlen_q_cur_batch = Q_ref.shape[0] + seqlen_k_cur_batch = K_ref.shape[0] + blk_offset = (Int32(0), Int32(0), Int32(0), ((Int32(0), Int32(0)), Int32(0))) + if cutlass.const_expr(varlen): + assert isinstance(cumulative_s_q, cute.Tensor) + assert isinstance(cumulative_s_k, cute.Tensor) + seqlen_q_cur_batch = ( + cumulative_s_q[blk_coord_b + 1] - cumulative_s_q[blk_coord_b] + ) + seqlen_k_cur_batch = ( + cumulative_s_k[blk_coord_b + 1] - cumulative_s_k[blk_coord_b] + ) + blk_offset = ( + cumulative_s_q[blk_coord_b], + cumulative_s_k[blk_coord_b], + Int32(0), + ((Int32(0), Int32(0)), Int32(0)), + ) + iter_start, iter_end = self.get_Q_block_min_max( + seqlen_q_cur_batch, + seqlen_k_cur_batch, + blk_coord_k, + is_2cta=True, + ) + iter_count = (iter_end - iter_start) * problem_shape[3][0][0] + if iter_count <= 0: + if blk_coord_k * self.tile_shape_K < seqlen_k_cur_batch: + problem_shape_cur_batch = ( + seqlen_q_cur_batch, + seqlen_k_cur_batch, + problem_shape[2], + problem_shape[3], + ) + self.epilogue_clear( + blk_coord, blk_offset, problem_shape_cur_batch, dK, dV + ) + else: + problem_shape_cur_batch = ( + seqlen_q_cur_batch, + seqlen_k_cur_batch, + problem_shape[2], + problem_shape[3], + ) + ( + compute_mma_P_producer, + compute_mma_dS_producer, + mma_compute_S_consumer, + compute_mma_P_consumer, + load_compute_LSE_consumer, + load_compute_sum_OdO_consumer, + mma_compute_dP_consumer, + compute_mma_dS_consumer, + mma_compute_dKdV_consumer, + ) = self.compute( + tSTtST, + tdPTtdPT, + tdVrP, + sP, + sLSE, + sdST, + sdOT, + sSum_OdO, + dK, + dV, + tdKtdK, + tdVtdV, + PdO_tiled_mma, + dSQ_tiled_mma, + blk_coord, + blk_offset, + problem_shape_cur_batch, + iter_count, + iter_start, + iter_end, + scale_softmax, + mma_compute_S_producer, + mma_compute_S_consumer, + compute_mma_P_producer, + compute_mma_P_consumer, + load_compute_LSE_producer, + load_compute_LSE_consumer, + load_compute_sum_OdO_producer, + load_compute_sum_OdO_consumer, + mma_compute_dP_producer, + mma_compute_dP_consumer, + compute_mma_dS_producer, + compute_mma_dS_consumer, + mma_compute_dKdV_producer, + mma_compute_dKdV_consumer, + varlen, + sK, + seqlen_k_cur_batch, + tma_atom_dK, + dK_tma, + tma_atom_dV, + dV_tma, + sdK_epi_layout, + sdV_epi_layout, + ) + cute.arch.barrier( + barrier_id=self.epilogue_sync_bar_id, + number_of_threads=self.num_compute_warps * self.threads_per_warp, + ) + # CLC advance + work_tile = tile_sched.advance_to_next_work() + # producer_tail after loop + compute_mma_P_producer.tail() + compute_mma_dS_producer.tail() + + else: + # ===== STATIC PATH: original non-persistent code ===== + blk_coord = (Int32(0), bidx, Int32(0), ((Int32(0), bidy), bidz)) + seqlen_q_cur_batch = Q_ref.shape[0] + seqlen_k_cur_batch = K_ref.shape[0] + blk_offset = (Int32(0), Int32(0), Int32(0), ((Int32(0), Int32(0)), Int32(0))) + if cutlass.const_expr(varlen): + assert isinstance(cumulative_s_q, cute.Tensor) + assert isinstance(cumulative_s_k, cute.Tensor) + seqlen_q_cur_batch = cumulative_s_q[bidz + 1] - cumulative_s_q[bidz] + seqlen_k_cur_batch = cumulative_s_k[bidz + 1] - cumulative_s_k[bidz] + blk_offset = ( + cumulative_s_q[bidz], + cumulative_s_k[bidz], + Int32(0), + ((Int32(0), Int32(0)), Int32(0)), + ) + + iter_start, iter_end = self.get_Q_block_min_max( + seqlen_q_cur_batch, + seqlen_k_cur_batch, + blk_coord[1], + is_2cta=True, + ) + + # Cluster wait + pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + iter_count = (iter_end - iter_start) * problem_shape[3][0][0] + problem_shape_cur_batch = ( + seqlen_q_cur_batch, + seqlen_k_cur_batch, + problem_shape[2], + problem_shape[3], + ) + if iter_count <= 0: + if bidx * self.tile_shape_K < seqlen_k_cur_batch: + self.epilogue_clear( + blk_coord, + blk_offset, + problem_shape_cur_batch, + dK, + dV, + ) + # /////////////////////////////////////////////////////////////////////////////// + # LOAD + # /////////////////////////////////////////////////////////////////////////////// + elif warp_idx == self.load_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_load) + + self.load( + K_in, + V_in, + Q_in, + QT_in, + dO_in, + dOT_in, + LSE, + sum_OdO, + sK, + sQ, + sQT, + sV, + sdO, + sdOT, + sLSE, + sSum_OdO, + KQ_tiled_mma, + VdO_tiled_mma, + PdO_tiled_mma, + dSQ_tiled_mma, + tma_atom_K, + tma_atom_Q, + tma_atom_QT, + tma_atom_V, + tma_atom_dO, + tma_atom_dOT, + blk_offset, + problem_shape_cur_batch, + varlen, + iter_count, + iter_start, + iter_end, + load_mma_Q_producer, + load_mma_Q_consumer, + load_mma_K_producer, + load_mma_K_consumer, + load_mma_V_producer, + load_mma_V_consumer, + load_compute_LSE_producer, + load_compute_LSE_consumer, + load_mma_dO_producer, + load_mma_dO_consumer, + load_mma_dOT_producer, + load_mma_dOT_consumer, + load_compute_sum_OdO_producer, + load_compute_sum_OdO_consumer, + load_mma_QT_producer, + load_mma_QT_consumer, + ) + + # /////////////////////////////////////////////////////////////////////////////// + # MMA + # /////////////////////////////////////////////////////////////////////////////// + elif warp_idx == self.mma_warp_id: + cute.arch.warpgroup_reg_alloc(self.num_regs_mma) + + self.mma_2cta( + KQ_tiled_mma, + VdO_tiled_mma, + PdO_tiled_mma, + dSQ_tiled_mma, + tSTtST, + tSTrQ, + tSTrK, + tdPTtdPT, + tdPTrV, + tdPTrdO, + tdVtdV, + tdVrP, + tdVrdOT, + tdKrdST, + tdKtdK, + tdKrQT, + iter_count, + load_mma_Q_consumer, + load_mma_K_consumer, + load_mma_V_consumer, + mma_compute_S_producer, + load_mma_dO_consumer, + mma_compute_dP_producer, + load_mma_dOT_consumer, + compute_mma_P_consumer, + compute_mma_dS_consumer, + load_mma_QT_consumer, + mma_compute_dKdV_producer, + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Compute + # /////////////////////////////////////////////////////////////////////////////// + elif warp_idx >= self.compute_warp_id[0] and warp_idx <= self.compute_warp_id[-1]: + cute.arch.warpgroup_reg_alloc(self.num_regs_compute) + + self.compute( + tSTtST, + tdPTtdPT, + tdVrP, + sP, + sLSE, + sdST, + sdOT, + sSum_OdO, + dK, + dV, + tdKtdK, + tdVtdV, + PdO_tiled_mma, + dSQ_tiled_mma, + blk_coord, + blk_offset, + problem_shape_cur_batch, + iter_count, + iter_start, + iter_end, + scale_softmax, + mma_compute_S_producer, + mma_compute_S_consumer, + compute_mma_P_producer, + compute_mma_P_consumer, + load_compute_LSE_producer, + load_compute_LSE_consumer, + load_compute_sum_OdO_producer, + load_compute_sum_OdO_consumer, + mma_compute_dP_producer, + mma_compute_dP_consumer, + compute_mma_dS_producer, + compute_mma_dS_consumer, + mma_compute_dKdV_producer, + mma_compute_dKdV_consumer, + varlen, + sK, + seqlen_k_cur_batch, + tma_atom_dK, + dK_tma, + tma_atom_dV, + dV_tma, + sdK_epi_layout, + sdV_epi_layout, + ) + + cute.arch.barrier( + barrier_id=self.epilogue_sync_bar_id, + number_of_threads=self.num_compute_warps * self.threads_per_warp, + ) + + else: + cute.arch.warpgroup_reg_dealloc(self.num_regs_empty) + + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + # dishengbin Deallocate tmem for early exit + # Dealloc the tensor memory + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + @cute.jit + def get_Q_block_min_max( + self, + seq_Q: Int32, + seq_K: Int32, + blk_coord_k: Int32, + is_2cta: bool, + ): + """Get Q tiles range.""" + Q_block_max = cute.ceil_div(seq_Q, self.tile_shape_Q) + Q_block_min = cutlass.Int32(0) + if cutlass.const_expr(self.has_sliding_window): + # For 2cta, use the last K block in the cluster so both CTAs get the same Q_block_max + blk_coord_k_for_max = (blk_coord_k // 2) * 2 + 1 + Q_block_max_tmp = cute.ceil_div( + (blk_coord_k_for_max + 1) * self.tile_shape_K + + seq_Q + - seq_K + + self.window_size_left, + self.tile_shape_Q, + ) + Q_block_max = min(Q_block_max, Q_block_max_tmp) + if cutlass.const_expr(self.is_causal or self.has_sliding_window): + # For 2cta, use the first K block in the cluster so both CTAs get the same Q_block_min. + # This ensures both CTAs in a cluster run the same number of pipeline iterations, + # avoiding hang from mismatched producer_commit / consumer_wait counts. + blk_coord_k_for_min = (blk_coord_k // 2) * 2 + Q_block_min_tmp = ( + blk_coord_k_for_min * self.tile_shape_K + seq_Q - seq_K - self.window_size_right + ) // self.tile_shape_Q + # Consider the case of 2cta, we need to ensure the K block is aligned to 2 + Q_block_min_tmp = Q_block_min_tmp - Q_block_min_tmp % 2 + Q_block_min = max(Q_block_min_tmp, Q_block_min) + return Q_block_min, Q_block_max + + @cute.jit + def load( + self, + K_in: cute.Tensor, + V_in: cute.Tensor, + Q_in: cute.Tensor, + QT_in: cute.Tensor, + dO_in: cute.Tensor, + dOT_in: cute.Tensor, + LSE_in: cute.Tensor, + sum_OdO_in: cute.Tensor, + sK: cute.Tensor, + sQ: cute.Tensor, + sQT: cute.Tensor, + sV: cute.Tensor, + sdO: cute.Tensor, + sdOT: cute.Tensor, + sLSE: cute.Tensor, + sSum_OdO: cute.Tensor, + KQ_tiled_mma: cute.TiledMma, + VdO_tiled_mma: cute.TiledMma, + PdO_tiled_mma: cute.TiledMma, + dSQ_tiled_mma: cute.TiledMma, + tma_atom_K: cute.CopyAtom, + tma_atom_Q: cute.CopyAtom, + tma_atom_QT: cute.CopyAtom, + tma_atom_V: cute.CopyAtom, + tma_atom_dO: cute.CopyAtom, + tma_atom_dOT: cute.CopyAtom, + blk_offset: cute.Shape, + problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]], + varlen: bool, + iter_count: Int32, + iter_start: Int32, + iter_end: Int32, + load_mma_Q_producer, + load_mma_Q_consumer, + load_mma_K_producer, + load_mma_K_consumer, + load_mma_V_producer, + load_mma_V_consumer, + load_compute_LSE_producer, + load_compute_LSE_consumer, + load_mma_dO_producer, + load_mma_dO_consumer, + load_mma_dOT_producer, + load_mma_dOT_consumer, + load_compute_sum_OdO_producer, + load_compute_sum_OdO_consumer, + load_mma_QT_producer, + load_mma_QT_consumer, + blk_coord_k_override: Int32 = Int32(-1), + blk_coord_h_k_override: Int32 = Int32(-1), + blk_coord_b_override: Int32 = Int32(-1), + ): + """TMA load.""" + tidx, _, _ = cute.arch.thread_idx() + if cutlass.const_expr(self.use_clc_scheduler): + blk_coord_k = blk_coord_k_override + blk_coord_h_k = blk_coord_h_k_override + blk_coord_b = blk_coord_b_override + else: + blk_coord_k, blk_coord_h_k, blk_coord_b = cute.arch.block_idx() + blk_coord_h_r = Int32(0) + blk_coord_h = (blk_coord_h_r, blk_coord_h_k) + iter_index = iter_start + mma_tile_coord_v = blk_coord_k % cute.size(KQ_tiled_mma.thr_id.shape) + mma_tile_coord_m = blk_coord_k // cute.size(KQ_tiled_mma.thr_id.shape) + + K = cute.domain_offset(cute.select(blk_offset, mode=[1, 2, 3]), K_in) + V = cute.domain_offset(cute.select(blk_offset, mode=[1, 2, 3]), V_in) + Q = cute.domain_offset(cute.select(blk_offset, mode=[0, 2, 3]), Q_in) + QT = cute.domain_offset(cute.select(blk_offset, mode=[2, 0, 3]), QT_in) + dO = cute.domain_offset(cute.select(blk_offset, mode=[0, 2, 3]), dO_in) + dOT = cute.domain_offset(cute.select(blk_offset, mode=[2, 0, 3]), dOT_in) + blk_offset_stats = blk_offset + if cutlass.const_expr(varlen): + cuseqlen_q_stats = cute.assume( + (blk_offset[0] + blk_coord_b * self.tile_shape_Q) + // self.tile_shape_Q + * self.tile_shape_Q, + divby=self.tile_shape_Q, + ) + blk_offset_stats = ( + cuseqlen_q_stats, + blk_offset[1], + blk_offset[2], + blk_offset[3], + ) + LSE = cute.domain_offset(cute.select(blk_offset_stats, mode=[0, 3]), LSE_in) + sum_OdO = cute.domain_offset(cute.select(blk_offset_stats, mode=[0, 3]), sum_OdO_in) + + gK = cute.local_tile(K, cute.select(self.KQ_mma_tiler, mode=[0, 2]), (None, None, None)) + gQ = cute.local_tile(Q, cute.select(self.KQ_mma_tiler, mode=[1, 2]), (None, None, None)) + gQT = cute.local_tile(QT, cute.select(self.dSQ_mma_tiler, mode=[1, 2]), (None, None, None)) + gV = cute.local_tile(V, cute.select(self.VdO_mma_tiler, mode=[0, 2]), (None, None, None)) + gdO = cute.local_tile(dO, cute.select(self.VdO_mma_tiler, mode=[1, 2]), (None, None, None)) + gdOT = cute.local_tile( + dOT, cute.select(self.PdO_mma_tiler, mode=[1, 2]), (None, None, None) + ) + + KQ_thr_mma = KQ_tiled_mma.get_slice(mma_tile_coord_v) + VdO_thr_mma = VdO_tiled_mma.get_slice(mma_tile_coord_v) + PdO_thr_mma = PdO_tiled_mma.get_slice(mma_tile_coord_v) + dSQ_thr_mma = dSQ_tiled_mma.get_slice(mma_tile_coord_v) + + tSTgK = KQ_thr_mma.partition_A(gK) + tSTgQ = KQ_thr_mma.partition_B(gQ) + tdKgQT = dSQ_thr_mma.partition_B(gQT) + tdPTgV = VdO_thr_mma.partition_A(gV) + tdPTgdO = VdO_thr_mma.partition_B(gdO) + tdVgdOT = PdO_thr_mma.partition_B(gdOT) + + cta_layout_mnk = cute.make_layout(self.cluster_shape_mnk) + cta_layout_vmnk = cute.tiled_divide(cta_layout_mnk, (KQ_tiled_mma.thr_id,)) + cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cute.arch.block_idx_in_cluster()) + + tKsK, tKgK_mkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_K, + cta_in_cluster_coord_vmnk[2], + cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])), + cute.group_modes(sK, 0, 3), + cute.group_modes(tSTgK, 0, 3), + ) + tQsQ, tQgQ_mkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_Q, + cta_in_cluster_coord_vmnk[1], + cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])), + cute.group_modes(sQ, 0, 3), + cute.group_modes(tSTgQ, 0, 3), + ) + tQTsQT, tQTgQT_mkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_QT, + cta_in_cluster_coord_vmnk[1], + cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])), + cute.group_modes(sQT, 0, 3), + cute.group_modes(tdKgQT, 0, 3), + ) + tVsV, tVgV_mkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_V, + cta_in_cluster_coord_vmnk[2], + cute.make_layout(cute.size(cta_layout_vmnk, mode=[2])), + cute.group_modes(sV, 0, 3), + cute.group_modes(tdPTgV, 0, 3), + ) + tdOsdO, tdOgdO_mkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_dO, + cta_in_cluster_coord_vmnk[1], + cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])), + cute.group_modes(sdO, 0, 3), + cute.group_modes(tdPTgdO, 0, 3), + ) + tdOTsdOT, tdOTgdOT_mkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_dOT, + cta_in_cluster_coord_vmnk[1], + cute.make_layout(cute.size(cta_layout_vmnk, mode=[1])), + cute.group_modes(sdOT, 0, 3), + cute.group_modes(tdVgdOT, 0, 3), + ) + + k_handle = load_mma_K_producer.acquire_and_advance() + cute.copy( + tma_atom_K, + tKgK_mkl[(None, mma_tile_coord_m, 0, (blk_coord_h, blk_coord_b))], + tKsK[None, 0], + tma_bar_ptr=k_handle.barrier, + ) + + q_handle = load_mma_Q_producer.acquire_and_advance() + cute.copy( + tma_atom_Q, + tQgQ_mkl[(None, iter_index, 0, (blk_coord_h, blk_coord_b))], + tQsQ[None, q_handle.index], + tma_bar_ptr=q_handle.barrier, + ) + + lse_handle = load_compute_LSE_producer.acquire_and_advance() + thread_idx = tidx % self.threads_per_warp + async_copy_num_elts = sLSE.shape[0] // self.threads_per_warp + atom_async_copy = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.ALWAYS), + self.acc_dtype, + num_bits_per_copy=self.acc_dtype.width, + ) + sLSE_for_copy = cute.flat_divide(sLSE, (1,)) + LSE_for_copy = cute.flat_divide(LSE, (1,)) + # Warp-coalesced: at each i, lane T accesses index `T + i*W` (stride-1 + # across the warp) instead of `T*N + i` (stride-N across the warp). + for i in cutlass.range_constexpr(async_copy_num_elts): + LSE_idx = self.tile_shape_Q * iter_index + thread_idx + i * self.threads_per_warp + sLSE_idx = thread_idx + i * self.threads_per_warp + if cute.elem_less(LSE_idx, problem_shape[0]): + cute.copy( + atom_async_copy, + LSE_for_copy[None, LSE_idx, (blk_coord_h, blk_coord_b)], + sLSE_for_copy[None, sLSE_idx, lse_handle.index], + ) + else: + sLSE_for_copy[None, sLSE_idx, lse_handle.index].fill(0.0) + lse_handle.commit() + + v_handle = load_mma_V_producer.acquire_and_advance() + cute.copy( + tma_atom_V, + tVgV_mkl[(None, mma_tile_coord_m, 0, (blk_coord_h, blk_coord_b))], + tVsV[(None, 0)], + tma_bar_ptr=v_handle.barrier, + ) + + do_handle = load_mma_dO_producer.acquire_and_advance() + cute.copy( + tma_atom_dO, + tdOgdO_mkl[(None, iter_index, 0, (blk_coord_h, blk_coord_b))], + tdOsdO[(None, do_handle.index)], + tma_bar_ptr=do_handle.barrier, + ) + + sum_odo_handle = load_compute_sum_OdO_producer.acquire_and_advance() + sSum_OdO_for_copy = cute.flat_divide(sSum_OdO, (1,)) + sum_OdO_for_copy = cute.flat_divide(sum_OdO, (1,)) + for i in cutlass.range_constexpr(async_copy_num_elts): + sum_OdO_idx = self.tile_shape_Q * iter_index + thread_idx + i * self.threads_per_warp + sSum_OdO_idx = thread_idx + i * self.threads_per_warp + if cute.elem_less(sum_OdO_idx, problem_shape[0]): + cute.copy( + atom_async_copy, + sum_OdO_for_copy[None, sum_OdO_idx, (blk_coord_h, blk_coord_b)], + sSum_OdO_for_copy[None, sSum_OdO_idx, sum_odo_handle.index], + ) + else: + sSum_OdO_for_copy[None, sSum_OdO_idx, sum_odo_handle.index].fill(0.0) + sum_odo_handle.commit() + + dot_handle = load_mma_dOT_producer.acquire_and_advance() + cute.copy( + tma_atom_dOT, + tdOTgdOT_mkl[(None, 0, iter_index, (blk_coord_h, blk_coord_b))], + tdOTsdOT[None, dot_handle.index], + tma_bar_ptr=dot_handle.barrier, + ) + + qt_handle = load_mma_QT_producer.acquire_and_advance() + cute.copy( + tma_atom_QT, + tQTgQT_mkl[(None, 0, iter_index, (blk_coord_h, blk_coord_b))], + tQTsQT[None, qt_handle.index], + tma_bar_ptr=qt_handle.barrier, + ) + + iter_count -= 1 + iter_index += 1 + + while iter_count > 0: + if iter_index == iter_end: + iter_index = iter_start + blk_coord_h_r += 1 + blk_coord_h = (blk_coord_h_r, blk_coord_h_k) + + q_handle = load_mma_Q_producer.acquire_and_advance() + cute.copy( + tma_atom_Q, + tQgQ_mkl[(None, iter_index, 0, (blk_coord_h, blk_coord_b))], + tQsQ[None, q_handle.index], + tma_bar_ptr=q_handle.barrier, + ) + + lse_handle = load_compute_LSE_producer.acquire_and_advance() + sLSE_for_copy = cute.flat_divide(sLSE, (1,)) + LSE_for_copy = cute.flat_divide(LSE, (1,)) + for i in cutlass.range_constexpr(async_copy_num_elts): + LSE_idx = self.tile_shape_Q * iter_index + thread_idx + i * self.threads_per_warp + sLSE_idx = thread_idx + i * self.threads_per_warp + if cute.elem_less(LSE_idx, problem_shape[0]): + cute.copy( + atom_async_copy, + LSE_for_copy[None, LSE_idx, (blk_coord_h, blk_coord_b)], + sLSE_for_copy[None, sLSE_idx, lse_handle.index], + ) + else: + sLSE_for_copy[None, sLSE_idx, lse_handle.index].fill(0.0) + lse_handle.commit() + + do_handle = load_mma_dO_producer.acquire_and_advance() + cute.copy( + tma_atom_dO, + tdOgdO_mkl[(None, iter_index, 0, (blk_coord_h, blk_coord_b))], + tdOsdO[None, do_handle.index], + tma_bar_ptr=do_handle.barrier, + ) + + sum_odo_handle = load_compute_sum_OdO_producer.acquire_and_advance() + sSum_OdO_for_copy = cute.flat_divide(sSum_OdO, (1,)) + sum_OdO_for_copy = cute.flat_divide(sum_OdO, (1,)) + for i in cutlass.range_constexpr(async_copy_num_elts): + sum_OdO_idx = ( + self.tile_shape_Q * iter_index + thread_idx + i * self.threads_per_warp + ) + sSum_OdO_idx = thread_idx + i * self.threads_per_warp + if cute.elem_less(sum_OdO_idx, problem_shape[0]): + cute.copy( + atom_async_copy, + sum_OdO_for_copy[None, sum_OdO_idx, (blk_coord_h, blk_coord_b)], + sSum_OdO_for_copy[None, sSum_OdO_idx, sum_odo_handle.index], + ) + else: + sSum_OdO_for_copy[None, sSum_OdO_idx, sum_odo_handle.index].fill(0.0) + sum_odo_handle.commit() + + dot_handle = load_mma_dOT_producer.acquire_and_advance() + cute.copy( + tma_atom_dOT, + tdOTgdOT_mkl[(None, 0, iter_index, (blk_coord_h, blk_coord_b))], + tdOTsdOT[None, dot_handle.index], + tma_bar_ptr=dot_handle.barrier, + ) + + qt_handle = load_mma_QT_producer.acquire_and_advance() + cute.copy( + tma_atom_QT, + tQTgQT_mkl[(None, 0, iter_index, (blk_coord_h, blk_coord_b))], + tQTsQT[None, qt_handle.index], + tma_bar_ptr=qt_handle.barrier, + ) + + iter_count -= 1 + iter_index += 1 + + if not cutlass.const_expr(self.use_clc_scheduler): + load_mma_K_producer.tail() + load_mma_V_producer.tail() + load_mma_Q_producer.tail() + load_compute_LSE_producer.tail() + load_mma_dO_producer.tail() + load_mma_dOT_producer.tail() + load_compute_sum_OdO_producer.tail() + load_mma_QT_producer.tail() + + return ( + load_mma_Q_producer, + load_mma_K_producer, + load_mma_V_producer, + load_compute_LSE_producer, + load_mma_dO_producer, + load_mma_dOT_producer, + load_compute_sum_OdO_producer, + load_mma_QT_producer, + ) + + @cute.jit + def mma_2cta( + self, + KQ_tiled_mma: cute.TiledMma, + VdO_tiled_mma: cute.TiledMma, + PdO_tiled_mma: cute.TiledMma, + dSQ_tiled_mma: cute.TiledMma, + tSTtST: cute.Tensor, + tSTrQ: cute.Tensor, + tSTrK: cute.Tensor, + tdPTtdPT: cute.Tensor, + tdPTrV: cute.Tensor, + tdPTrdO: cute.Tensor, + tdVtdV: cute.Tensor, + tdVrP: cute.Tensor, + tdVrdOT: cute.Tensor, + tdKrdST: cute.Tensor, + tdKtdK: cute.Tensor, + tdKrQT: cute.Tensor, + iter_count: Int32, + load_mma_Q_consumer, + load_mma_K_consumer, + load_mma_V_consumer, + mma_compute_S_producer, + load_mma_dO_consumer, + mma_compute_dP_producer, + load_mma_dOT_consumer, + compute_mma_P_consumer, + compute_mma_dS_consumer, + load_mma_QT_consumer, + mma_compute_dKdV_producer, + ): + """CuTeDSL kernel for mma pipeline.""" + load_mma_Q_releaser = load_mma_Q_consumer.clone() + load_mma_K_releaser = load_mma_K_consumer.clone() + load_mma_V_releaser = load_mma_V_consumer.clone() + + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + is_leader_cta = cta_rank_in_cluster % 2 == 0 + + if is_leader_cta: + s_handle = mma_compute_S_producer.acquire_and_advance() + k_handle = load_mma_K_consumer.wait_and_advance() + q_handle = load_mma_Q_consumer.wait_and_advance() + + # Compute S = K * Q + for k_block in cutlass.range(0, cute.size(tSTrQ, mode=[2]), unroll_full=True): + KQ_tiled_mma.set(tcgen05.Field.ACCUMULATE, k_block != 0) + cute.gemm( + KQ_tiled_mma, + tSTtST, + tSTrK[None, None, k_block, 0], + tSTrQ[None, None, k_block, q_handle.index], + tSTtST, + ) + q_handle.release() + + cute.arch.fence_view_async_tmem_store() + s_handle.commit() + + do_handle = load_mma_dO_consumer.wait_and_advance() + v_handle = load_mma_V_consumer.wait_and_advance() + + dp_handle = mma_compute_dP_producer.acquire_and_advance() + + # Compute dP = V * dO + VdO_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for k_block in cutlass.range(0, cute.size(tdPTrV, mode=[2]), unroll_full=True): + cute.gemm( + VdO_tiled_mma, + tdPTtdPT, + tdPTrV[None, None, k_block, 0], + tdPTrdO[None, None, k_block, do_handle.index], + tdPTtdPT, + ) + VdO_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + dp_handle.commit() + do_handle.release() + # V only produced once by load(); hold v_handle until end, release there via releaser + + p_handle = compute_mma_P_consumer.wait_and_advance() + dot_handle = load_mma_dOT_consumer.wait_and_advance() + + # Compute dV = P * dO (First iteration) + PdO_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for k_block in cutlass.range(0, cute.size(tdVrP, mode=[2]), unroll_full=True): + cute.gemm( + PdO_tiled_mma, + tdVtdV, + tdVrP[None, None, k_block, 0], + tdVrdOT[None, None, k_block, dot_handle.index], + tdVtdV, + ) + PdO_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + dot_handle.release() + p_handle.release() + + iter_count -= 1 + + dSQ_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + while iter_count > 0: + if is_leader_cta: + q_handle = load_mma_Q_consumer.wait_and_advance() + s_handle = mma_compute_S_producer.acquire_and_advance() + + # Compute S = K * Q + for k_block in cutlass.range(0, cute.size(tSTrQ, mode=[2]), unroll_full=True): + KQ_tiled_mma.set(tcgen05.Field.ACCUMULATE, k_block != 0) + cute.gemm( + KQ_tiled_mma, + tSTtST, + tSTrK[None, None, k_block, 0], + tSTrQ[None, None, k_block, q_handle.index], + tSTtST, + ) + q_handle.release() + s_handle.commit() + + if is_leader_cta: + qt_handle = load_mma_QT_consumer.wait_and_advance() + ds_handle = compute_mma_dS_consumer.wait_and_advance() + + # Compute dK = dS * QT + for k_block in cutlass.range(0, cute.size(tdKrdST, mode=[2]), unroll_full=True): + cute.gemm( + dSQ_tiled_mma, + tdKtdK, + tdKrdST[ + None, + None, + k_block, + ds_handle.index, + ], + tdKrQT[None, None, k_block, qt_handle.index], + tdKtdK, + ) + dSQ_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + qt_handle.release() + ds_handle.release() + + if is_leader_cta: + dp_handle = mma_compute_dP_producer.acquire_and_advance() + do_handle = load_mma_dO_consumer.wait_and_advance() + # V only produced once by load(); reuse same V (index 0) for all loop iterations + VdO_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for k_block in cutlass.range(0, cute.size(tdPTrV, mode=[2]), unroll_full=True): + cute.gemm( + VdO_tiled_mma, + tdPTtdPT, + tdPTrV[None, None, k_block, 0], + tdPTrdO[None, None, k_block, do_handle.index], + tdPTtdPT, + ) + VdO_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + dp_handle.commit() + do_handle.release() + + if is_leader_cta: + p_handle = compute_mma_P_consumer.wait_and_advance() + dot_handle = load_mma_dOT_consumer.wait_and_advance() + + # Compute dV = P * dO (Loop iterations) + for k_block in cutlass.range(0, cute.size(tdVrP, mode=[2]), unroll_full=True): + cute.gemm( + PdO_tiled_mma, + tdVtdV, + tdVrP[None, None, k_block, 0], + tdVrdOT[None, None, k_block, dot_handle.index], + tdVtdV, + ) + PdO_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + p_handle.release() + dot_handle.release() + + iter_count -= 1 + + if is_leader_cta: + dkdv_handle = mma_compute_dKdV_producer.acquire_and_advance() + dkdv_handle.commit() + + load_mma_K_releaser.release() + load_mma_K_releaser.advance() + load_mma_V_releaser.release() + load_mma_V_releaser.advance() + + if is_leader_cta: + dkdv_handle = mma_compute_dKdV_producer.acquire_and_advance() + + ds_handle = compute_mma_dS_consumer.wait_and_advance() + qt_handle = load_mma_QT_consumer.wait_and_advance() + + # Compute dK = dS * Q + for k_block in cutlass.range(0, cute.size(tdKrdST, mode=[2]), unroll_full=True): + cute.gemm( + dSQ_tiled_mma, + tdKtdK, + tdKrdST[None, None, k_block, ds_handle.index], + tdKrQT[None, None, k_block, qt_handle.index], + tdKtdK, + ) + dSQ_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + + dkdv_handle.commit() + qt_handle.release() + ds_handle.release() + + if not cutlass.const_expr(self.use_clc_scheduler): + mma_compute_S_producer.tail() + mma_compute_dP_producer.tail() + mma_compute_dKdV_producer.tail() + + return ( + mma_compute_S_producer, + mma_compute_dP_producer, + mma_compute_dKdV_producer, + load_mma_Q_consumer, + load_mma_K_consumer, + load_mma_V_consumer, + load_mma_dO_consumer, + load_mma_dOT_consumer, + compute_mma_P_consumer, + compute_mma_dS_consumer, + load_mma_QT_consumer, + ) + + @cute.jit + def reg_to_smem_mma64x64( + self, + regs: cute.Tensor, + smem: cute.Tensor, + index: Int32, + tiler_mn: tuple[Int32, Int32], + dp_idx: Int32, + wg_idx: Int32, + ): + smem_slice = smem[None, None, None, index] + # TODO: double check the layout of the data in reg. + # TODO: this may introduce additional smem transpose. + thread_layout = cute.make_ordered_layout( + tiler_mn, + (0, 1), # TODO: (0,1) or (1,0) ??? + ) + smem_slice_tmp = cute.composition(smem_slice, thread_layout) + + # TODO: temporary code for tile 64 x 64. + tmp_shape = ((8, 2, 4), (2, 4, 2, 2, 2)) + tmp_stride = ((64, 512, 1024), (1, 2, 8, 16, 32)) + smem_copy = cute.composition(smem_slice_tmp, cute.make_layout(tmp_shape, stride=tmp_stride)) + + # TODO: the following code is only for tile 64 x 64. + # TODO: need to modify the code for other tile sizes. + lane_idx = dp_idx % 32 + reg_shape = regs.shape + atom_loops = reg_shape[0][0][2] + block_loops = reg_shape[2] + # | 00 ~ 07 | 08 ~ 15 | 16 ~ 23 | 24 ~ 31 | 32 ~ 39 | 40 ~ 47 | 48 ~ 55 | 56 ~ 63 | + # |---- atom size ----|---- atom size ----|---- atom size ----|---- atom size ----| + # |---- wg0 ----|---- wg1 ----|---- wg0 ----|---- wg1 ----| + for ia in cutlass.range(atom_loops): + for ib in cutlass.range(block_loops): + # the lower 8 lines + regs_copy = regs[((None, 0, ia), 0), 0, ib] # two elements + smem_copy_slice = smem_copy[ + (lane_idx // 4, 0, dp_idx // 32), + (None, lane_idx % 4, ia, wg_idx, ib), + ] + cute.autovec_copy(regs_copy, smem_copy_slice) + # the upper 8 lines + regs_copy = regs[((None, 1, ia), 0), 0, ib] + smem_copy_slice = smem_copy[ + (lane_idx // 4, 1, dp_idx // 32), + (None, lane_idx % 4, ia, wg_idx, ib), + ] + cute.autovec_copy(regs_copy, smem_copy_slice) + + @cute.jit + def reg_to_smem_mma128x128_2cta( + self, + regs: cute.Tensor, + smem: cute.Tensor, + index: Int32, + tiler_mn: tuple[Int32, Int32], + dp_idx: Int32, + wg_idx: Int32, + smem_RowMajor: bool = True, + ): + smem_slice = smem[None, None, None, index] + # K>> smem_slice: tensor, S<3,4,3>> o ((64,16),1,(4,2)):((64,1),0,(16,4096))> + thread_layout = cute.make_ordered_layout( + # (tileN, tileM) + tiler_mn, + (0, 1), + ) + # K>> thread_layout: (64,128):(128,1) + smem_slice_tmp = cute.composition(smem_slice, thread_layout) + + # NOTE: hardcode for tcgen05.ld.32x32b.x8 & mma128x64+2cta + # tmp_shape = ((32, 2), (8, 2, 2, 2)) # for 64x64 tile + # tmp_stride = ((64, 32*64), (1, 8, 16, 32)) + # NOTE: hardcode for tcgen05.ld.32x32b.x16 & mma128x64+2cta + tmp_shape = ((32, 2), (16, 2, 2, 2)) # for 128x64 tile + tmp_stride = ((64, 32 * 64), (1, 16, 32, 64 * 64)) + # smem_copy = cute.composition(smem_slice_tmp, cute.make_layout(tmp_shape, stride=tmp_stride)) + smem_copy = cute.make_tensor( + smem_slice_tmp.iterator, cute.make_layout(tmp_shape, stride=tmp_stride) + ) + + warp_idx = dp_idx // 32 + warp_row_idx = warp_idx % 2 + warp_col_idx = warp_idx // 2 # corresponding to the second 64 cols in smem + lane_idx = dp_idx % 32 + reg_shape = ( + regs.shape + ) # ((8,1),1,2):((1,0),0,8) for 64x64, ((16,1),1,2):((1,0),0,16) for 128x64 + block_loops = reg_shape[2] + + # TODO: maybe can use cp.async for optimization + for ib in cutlass.range(block_loops): + regs_copy = regs[(None, 0), 0, ib] + smem_copy_slice = smem_copy[(lane_idx, warp_row_idx), (None, wg_idx, ib, warp_col_idx)] + cute.autovec_copy(regs_copy, smem_copy_slice) + + @cute.jit + def reg_to_smem_mma128x64_2cta( + self, + regs: cute.Tensor, + smem: cute.Tensor, + index: Int32, + tiler_mn: tuple[Int32, Int32], + dp_idx: Int32, + wg_idx: Int32, + smem_RowMajor: bool = True, + ): + smem_slice = smem[None, None, None, index] + thread_layout = cute.make_ordered_layout( + # (tileN, tileM) + tiler_mn, + (1, 0) if smem_RowMajor else (0, 1), + ) + smem_slice_tmp = cute.composition(smem_slice, thread_layout) + # NOTE: hardcode for tcgen05.ld.32x32b.x8 & mma128x64+2cta + tmp_shape = ((32, 2), (8, 2, 2, 2)) + tmp_stride = ((64, 32 * 64), (1, 8, 16, 32)) + smem_copy = cute.composition(smem_slice_tmp, cute.make_layout(tmp_shape, stride=tmp_stride)) + + warp_idx = dp_idx // 32 + warp_row_idx = warp_idx % 2 + warp_col_idx = warp_idx // 2 + lane_idx = dp_idx % 32 + reg_shape = regs.shape # ((8,1),1,2):((1,0),0,8) + block_loops = reg_shape[2] + + # TODO: maybe can use cp.async for optimization + for ib in cutlass.range(block_loops): + regs_copy = regs[(None, 0), 0, ib] + smem_copy_slice = smem_copy[(lane_idx, warp_row_idx), (None, wg_idx, ib, warp_col_idx)] + cute.autovec_copy(regs_copy, smem_copy_slice) + + @cute.jit + def compute( + self, + tSTtST: cute.Tensor, + tdPTtdPT: cute.Tensor, + tdVrP: cute.Tensor, + sP: cute.Tensor, + sLSE: cute.Tensor, + # sdS: cute.Tensor, + sdST: cute.Tensor, + sdOT: cute.Tensor, + sSum_OdO: cute.Tensor, + dK: cute.Tensor, + dV: cute.Tensor, + tdKtdK: cute.Tensor, + tdVtdV: cute.Tensor, + PdO_tiled_mma: cute.TiledMma, + dSQ_tiled_mma: cute.TiledMma, + blk_coord: cute.Coord, + blk_offset: cute.Shape, + problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]], + iter_count: Int32, + iter_start: Int32, + iter_end: Int32, + scale_softmax: cutlass.Float32, + mma_compute_S_producer, + mma_compute_S_consumer, + compute_mma_P_producer, + compute_mma_P_consumer, + load_compute_LSE_producer, + load_compute_LSE_consumer, + load_compute_sum_OdO_producer, + load_compute_sum_OdO_consumer, + mma_compute_dP_producer, + mma_compute_dP_consumer, + compute_mma_dS_producer, + compute_mma_dS_consumer, + mma_compute_dKdV_producer, + mma_compute_dKdV_consumer, + varlen: bool, + sK: cute.Tensor, + problem_shape_k_cur_batch: Int32, + tma_atom_dK: cute.CopyAtom, + dK_tma: cute.Tensor, + tma_atom_dV: cute.CopyAtom, + dV_tma: cute.Tensor, + sdK_epi_layout: cute.ComposedLayout, + sdV_epi_layout: cute.ComposedLayout, + ): + """CuTeDSL kernel for recomputing softmax and producing dk and dv.""" + tidx, _, _ = cute.arch.thread_idx() + Q, K, _, _ = problem_shape + _, blk_coord_k, _, _ = blk_coord + + iter_index = iter_start + + # adi: TMEM_ST, TMEM_DPT + tmem_load_op = tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(16)) + tmem_load_atom = cute.make_copy_atom( + tmem_load_op, + self.acc_dtype, + ) + + tSTtST = tSTtST[(None, None), 0, 0] + tdPTtdPT = tdPTtdPT[(None, None), 0, 0] + + cST = cute.make_identity_tensor(cute.select(self.cta_tiler, mode=[1, 0])) + cdPT = cute.make_identity_tensor(cute.select(self.cta_tiler, mode=[1, 0])) + + num_warp_groups = self.num_compute_warps // 4 + dp_idx = tidx % 128 + wg_idx = (tidx % (self.num_compute_warps * self.threads_per_warp)) // 128 + tiled_t2r = tcgen05.make_tmem_copy(tmem_load_atom, tSTtST) + thr_t2r = tiled_t2r.get_slice(dp_idx) + + tTR_cST = thr_t2r.partition_D(cST) + tTR_cST = split_wg(tTR_cST, num_warp_groups, wg_idx) + tTR_rST = cute.make_rmem_tensor(tTR_cST.shape, self.acc_dtype) + + tTR_tST = thr_t2r.partition_S(tSTtST) + tTR_tST = split_wg(tTR_tST, num_warp_groups, wg_idx) + + tTR_cdPT_p = thr_t2r.partition_D(cdPT) + tTR_cdPT = split_wg(tTR_cdPT_p, num_warp_groups, wg_idx) + tTR_rdPT = cute.make_rmem_tensor(tTR_cdPT.shape, self.acc_dtype) + + tTR_tdPT = thr_t2r.partition_S(tdPTtdPT) + tTR_tdPT = split_wg(tTR_tdPT, num_warp_groups, wg_idx) + + tdVcST = PdO_tiled_mma.get_slice(0).partition_A(cST) + + is_residual_k = blk_coord_k * self.tile_shape_K + self.tile_shape_K > K + last_iter = iter_end - 1 + + while iter_count > 0: + s_handle = mma_compute_S_consumer.wait_and_advance() + p_handle = compute_mma_P_producer.acquire_and_advance() + lse_handle = load_compute_LSE_consumer.wait_and_advance() + + leading_causal_masking = cutlass.Boolean(False) + if cutlass.const_expr(self.is_causal): + # TODO: could be optimized by specify an exact iter_index + # + # NOTE (causal + 2CTA correctness): + # `iter_start` can be rounded down to an even index for 2CTA. When the true + # causal boundary Q tile is odd, it becomes (iter_start + 1). In that case, + # we must treat both (iter_start, iter_start + 1) as "masked tiles" so the + # per-element causal mask is applied on the boundary tile too. + leading_causal_masking = iter_index == iter_start + q_block_min_unaligned = ( + blk_coord_k * self.tile_shape_K + Q - K - self.window_size_right + ) // self.tile_shape_Q + boundary_in_second = (q_block_min_unaligned % 2) == 1 + leading_causal_masking = leading_causal_masking or ( + boundary_in_second and (iter_index == iter_start + 1) + ) + offset_partial_tile = (K - Q) % self.tile_shape_K + need_additional_mask = offset_partial_tile and iter_index == iter_start + 1 + leading_causal_masking = leading_causal_masking or need_additional_mask + leading_causal_masking = cute.arch.shuffle_sync(leading_causal_masking, 0) + + trailing_residual_masking = cutlass.Boolean(False) + trailing_residual_masking = iter_index == last_iter or is_residual_k + trailing_residual_masking = cute.arch.shuffle_sync(trailing_residual_masking, 0) + + # For causal, every tile may contain (q,k) with k > q; we must apply per-element mask for all Q tiles. + is_masked_tile = ( + leading_causal_masking + or trailing_residual_masking + or self.has_sliding_window + or cutlass.const_expr(self.is_causal) + ) + + # Compute P = softmax(S, LSE) + cute.copy(tiled_t2r, tTR_tST, tTR_rST) + + if is_masked_tile: + for i in cutlass.range(cute.size(tTR_rST), unroll_full=True): + c_transpose = tTR_cST[i] + pos = ( + cute.get(c_transpose, mode=[1]) + iter_index * self.tile_shape_Q, + cute.get(c_transpose, mode=[0]) + blk_coord_k * self.tile_shape_K, + ) + if cutlass.const_expr(self.has_sliding_window): + if cutlass.const_expr(self.window_size_left < 0): + tTR_rST[i] = ( + -cutlass.Float32.inf + if pos[1] > pos[0] + K - Q + self.window_size_right + else tTR_rST[i] + ) + else: + max_K_index = min(pos[0] + K - Q + self.window_size_right, K) + min_K_index = max(0, pos[0] + K - Q - self.window_size_left) + tTR_rST[i] = ( + -cutlass.Float32.inf + if pos[1] > max_K_index or pos[1] < min_K_index + else tTR_rST[i] + ) + if cutlass.const_expr(self.is_causal) and ( + pos[0] + K - Q < pos[1] or not cute.elem_less(pos, (Q, K)) + ): + tTR_rST[i] = -cutlass.Float32.inf + if not cute.elem_less(pos, (Q, K)): + tTR_rST[i] = -cutlass.Float32.inf + + log2_e = cutlass.Float32(math.log2(math.e)) + softmax_scale_log2_e = scale_softmax * log2_e + + for i in cutlass.range(0, cute.size(tTR_rST), 2, unroll_full=True): + lse = ( + -sLSE[ + cute.get(tTR_cST[i], mode=[1]), + lse_handle.index, + ], + -sLSE[ + cute.get(tTR_cST[i + 1], mode=[1]), + lse_handle.index, + ], + ) + tTR_rST[i], tTR_rST[i + 1] = cute.arch.fma_packed_f32x2( + (tTR_rST[i], tTR_rST[i + 1]), + (softmax_scale_log2_e, softmax_scale_log2_e), + lse, + ) + tTR_rST[i] = cute.math.exp2(tTR_rST[i], fastmath=True) + tTR_rST[i + 1] = cute.math.exp2(tTR_rST[i + 1], fastmath=True) + + # convert fp32 P to fp16 P which will be used in the PdO + tTR_rPT = self.quantize(tTR_rST, dV.element_type) # tTR_rST is ST in fp32 in RF. + self.reg_to_smem_mma128x128_2cta( + tTR_rPT, + sP, + p_handle.index, + (self.tile_shape_K, self.tile_shape_Q), + dp_idx, + wg_idx, + ) + cute.arch.fence_view_async_shared() + cute.arch.barrier( + barrier_id=self.compute_sync_bar_id, + number_of_threads=self.num_compute_warps * self.threads_per_warp, + ) + + p_handle.commit() + + s_handle.release() + lse_handle.release() + + sum_odo_handle = load_compute_sum_OdO_consumer.wait_and_advance() + dp_handle = mma_compute_dP_consumer.wait_and_advance() + ds_handle = compute_mma_dS_producer.acquire_and_advance() + + # Compute dS = dsoftmax(P, dP, sum_OdO) + cute.copy(tiled_t2r, tTR_tdPT, tTR_rdPT) + + for i in cutlass.range(0, cute.size(tTR_rdPT), 2, unroll_full=True): + dpsum_0 = -sSum_OdO[ + cute.get(tTR_cdPT[i], mode=[1]), + sum_odo_handle.index, + ] + dpsum_1 = -sSum_OdO[ + cute.get(tTR_cdPT[i + 1], mode=[1]), + sum_odo_handle.index, + ] + if cutlass.const_expr(varlen): + if not cute.elem_less(cute.get(tTR_cdPT[i], mode=[1]), Q): + dpsum_0 = 0.0 + if not cute.elem_less(cute.get(tTR_cdPT[i + 1], mode=[1]), Q): + dpsum_1 = 0.0 + tTR_rdPT[i], tTR_rdPT[i + 1] = cute.arch.add_packed_f32x2( + (tTR_rdPT[i], tTR_rdPT[i + 1]), + (dpsum_0, dpsum_1), + ) + tTR_rdPT[i], tTR_rdPT[i + 1] = cute.arch.mul_packed_f32x2( + (tTR_rdPT[i], tTR_rdPT[i + 1]), (tTR_rST[i], tTR_rST[i + 1]) + ) + # For causal, force dS to zero at masked (q,k) so dK/dV accumulation is correct + if cutlass.const_expr(self.is_causal): + for i in cutlass.range(cute.size(tTR_rdPT), unroll_full=True): + c_transpose = tTR_cdPT[i] + pos = ( + cute.get(c_transpose, mode=[1]) + iter_index * self.tile_shape_Q, + cute.get(c_transpose, mode=[0]) + blk_coord_k * self.tile_shape_K, + ) + if pos[0] + K - Q < pos[1] or not cute.elem_less(pos, (Q, K)): + tTR_rdPT[i] = cutlass.Float32(0.0) + # convert fp32 dS to fp16 dS which will be used in the computation of dK and DQ + tTR_rdST = self.quantize(tTR_rdPT, dV.element_type) + + cute.arch.fence_view_async_tmem_load() + dp_handle.release() + + self.reg_to_smem_mma128x128_2cta( + tTR_rdST, + sdST, + ds_handle.index, + (self.tile_shape_K, self.tile_shape_Q), + dp_idx, + wg_idx, + ) + cute.arch.fence_view_async_shared() + cute.arch.barrier( + barrier_id=self.compute_sync_bar_id, + number_of_threads=self.num_compute_warps * self.threads_per_warp, + ) + + ds_handle.commit() + sum_odo_handle.release() + + iter_count -= 1 + iter_index += 1 + if iter_index == iter_end: + iter_index = iter_start + + # Epilogue + mma_compute_dKdV_consumer = self.epilogue( + blk_coord, + blk_offset, + problem_shape, + dK, + dV, + tdKtdK, + tdVtdV, + scale_softmax, + mma_compute_dKdV_producer, + mma_compute_dKdV_consumer, + problem_shape_k_cur_batch, + tma_atom_dK, + dK_tma, + tma_atom_dV, + dV_tma, + sdK_epi_layout, + sdV_epi_layout, + varlen, + sdOT, + sP, + ) + + if not cutlass.const_expr(self.use_clc_scheduler): + compute_mma_P_producer.tail() + compute_mma_dS_producer.tail() + + return ( + compute_mma_P_producer, + compute_mma_dS_producer, + mma_compute_S_consumer, + compute_mma_P_consumer, + load_compute_LSE_consumer, + load_compute_sum_OdO_consumer, + mma_compute_dP_consumer, + compute_mma_dS_consumer, + mma_compute_dKdV_consumer, + ) + + @cute.jit + def quantize( + self, + input_t: cute.Tensor, + element_dtype: type[cutlass.Numeric], + ) -> cute.Tensor: + """Convert Float32 to element dtype.""" + output = cute.make_rmem_tensor(input_t.shape, element_dtype) + output.store(input_t.load().to(element_dtype)) + return output + + @cute.jit + def store( + self, + gmem: cute.Tensor, + regs: cute.Tensor, + coord: cute.Tensor, + tensor_shape: cute.Shape, + ): + for i in cutlass.range(cute.size(coord, mode=[2]), unroll_full=True): + coord_i = coord[None, 0, i] + gmem_i = gmem[None, 0, i] + regs_i = regs[None, 0, i] + if cute.elem_less(coord_i[0], tensor_shape): + gmem_i.store(regs_i.load().to(gmem.element_type)) + + @cute.jit + def epilogue_clear( + self, + blk_coord: cute.Coord, + blk_offset: cute.Shape, + problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]], + dK: cute.Tensor, + dV: cute.Tensor, + ): + """Early stopping needs to clear dK and dV.""" + tidx, _, _ = cute.arch.thread_idx() + block_dim_x, _, _ = cute.arch.block_dim() + _, K, _, HB = problem_shape + _, blk_coord_k, _, blk_coord_batch = blk_coord + + mdK_offset = cute.assume(blk_offset[1] * dK.stride[0], divby=64) + mdK = cute.make_tensor( + dK.iterator + mdK_offset, + cute.make_layout((K, self.tile_shape_dQ_K, HB), stride=dK.stride), + ) + gdK = cute.local_tile(mdK, (self.cta_tiler[1], self.cta_tiler[2]), (None, None, None)) + gdK = gdK[None, None, blk_coord_k, 0, blk_coord_batch] + cdK = cute.domain_offset( + (blk_coord_k * self.tile_shape_K, 0), + cute.make_identity_tensor((self.cta_tiler[1], self.cta_tiler[2])), + ) + + mdV_offset = cute.assume(blk_offset[1] * dV.stride[0], divby=64) + mdV = cute.make_tensor( + dV.iterator + mdV_offset, + cute.make_layout((K, self.tile_shape_dV_dO, HB), stride=dV.stride), + ) + gdV = cute.local_tile(mdV, (self.cta_tiler[1], self.cta_tiler[2]), (None, None, None)) + gdV = gdV[None, None, blk_coord_k, 0, blk_coord_batch] + cdV = cute.domain_offset( + (blk_coord_k * self.tile_shape_K, 0), + cute.make_identity_tensor((self.cta_tiler[1], self.cta_tiler[2])), + ) + + num_zero_epi_threads = 256 + + tiled_copy_r2g = fa_copy_utils.tiled_copy_2d( + dK.element_type, self.cta_tiler[2], num_zero_epi_threads + ) + + thr_copy_r2g = tiled_copy_r2g.get_slice(tidx) + + tRG_gdK = thr_copy_r2g.partition_D(gdK) + tRG_cdK = thr_copy_r2g.partition_D(cdK) + tRG_gdV = thr_copy_r2g.partition_D(gdV) + tRG_cdV = thr_copy_r2g.partition_D(cdV) + + zero_frg = cute.make_rmem_tensor_like(tRG_gdK[None, 0, None]) + zero_frg.fill(dK.element_type(0.0)) + + # check we don't need zero fragment duplication + V_frg_size = cute.size(tRG_gdV[None, 0, None]) + assert cute.size(zero_frg) == V_frg_size + + if tidx < num_zero_epi_threads: + for n in cutlass.range(cute.size(tRG_gdK.shape[1]), unroll_full=True): + if cute.elem_less(tRG_cdK[0, n, 0][0], problem_shape[1]): + cute.copy(tiled_copy_r2g, zero_frg, tRG_gdK[None, n, None]) + + for n in cutlass.range(cute.size(tRG_gdV.shape[1]), unroll_full=True): + if cute.elem_less(tRG_cdV[0, n, 0][0], problem_shape[1]): + cute.copy(tiled_copy_r2g, zero_frg, tRG_gdV[None, n, None]) + + @cute.jit + def epilogue( + self, + blk_coord: cute.Coord, + blk_offset: cute.Shape, + problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]], + dK: cute.Tensor, + dV: cute.Tensor, + tdKtdK: cute.Tensor, + tdVtdV: cute.Tensor, + scale_softmax: cutlass.Float32, + mma_compute_dKdV_producer, + mma_compute_dKdV_consumer, + problem_shape_k_cur_batch: Int32, + tma_atom_dK: cute.CopyAtom, + dK_tma: cute.Tensor, + tma_atom_dV: cute.CopyAtom, + dV_tma: cute.Tensor, + sdK_epi_layout: cute.ComposedLayout, + sdV_epi_layout: cute.ComposedLayout, + varlen: bool, + sdOT: cute.Tensor, + sP: cute.Tensor, + ): + """Variant 3a (5/5) Path 2: CTA-shared SMEM with cooperative WG writes + TMA bulk store. + + Both warp-groups cooperatively populate a per-CTA (64, 256) virtual SMEM + buffer (4 stages of (64, 64) aliased onto sP+sdST). Per-thread t2r N + coverage is interleaved across the full hd=256, so per-WG TMA is not + viable — instead we treat SMEM as one shared per-CTA buffer and let + each thread's `self.store`-equivalent write into it via a (64, 256) + virtual tensor whose N axis maps (n%64, n//64) → (N_within, stage). + After an inter-WG barrier (256 threads), the leader warp fires 4 TMA + bulk stores, one per stage, to the corresponding (64, 64) GMEM slice. + Varlen falls back to per-thread self.store as in flash_bwd_sm100.py. + """ + tidx, _, _ = cute.arch.thread_idx() + _, K, D, HB = problem_shape + _, blk_coord_k, _, blk_coord_batch = blk_coord + + # adi: TMEM_DK, TMEM_DV + tmem_copy_op = tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)) + load_op = cute.make_copy_atom( + tmem_copy_op, + self.acc_dtype, + ) + + tdKtdK = tdKtdK[(None, None), 0, 0] + mdK_offset = cute.assume(blk_offset[1] * dK.stride[0], divby=64) + mdK = cute.make_tensor( + dK.iterator + mdK_offset, + cute.make_layout((K, self.tile_shape_dQ_K, HB), stride=dK.stride), + ) + gdK = cute.local_tile(mdK, (self.cta_tiler[1], self.cta_tiler[2]), (None, None, None)) + gdK = gdK[None, None, blk_coord_k, 0, blk_coord_batch] + cdK = cute.domain_offset( + (blk_coord_k * self.tile_shape_K, 0), + cute.make_identity_tensor((self.cta_tiler[1], self.cta_tiler[2])), + ) + + num_warp_groups = self.num_compute_warps // 4 + dp_idx = tidx % 128 + wg_idx = (tidx % (self.num_compute_warps * self.threads_per_warp)) // 128 + leader_warp = (cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4) == 0 + + # Path 2 SMEM staging. dV stages through sdOT (already-consumed by the + # dV MMA before the dV epilogue begins). dK stages through sP+sdST + # (dead after dK MMA completes, before dK epilogue runs). + s_epi_dK = cute.make_tensor( + cute.recast_ptr(sP.iterator, sdK_epi_layout.inner, dK.element_type), + sdK_epi_layout.outer, + ) + s_epi_dV = cute.make_tensor( + cute.recast_ptr(sdOT.iterator, sdV_epi_layout.inner, dV.element_type), + sdV_epi_layout.outer, + ) + + # Compile-time: stage tile shape and number of stages. + epi_cols_dKV = math.gcd( + 128 // (dV.element_type.width // 8), self.cta_tiler[2] // num_warp_groups + ) + num_epi_stages_dKV = (self.cta_tiler[2] // num_warp_groups) // epi_cols_dKV + total_epi_stages = num_warp_groups * num_epi_stages_dKV + epi_tile_dKV = (self.cta_tiler[1], epi_cols_dKV) + + # Local (M, N) coord tensor for SMEM indexing (no global domain offset + # — cdK/cdV are domain-offset by blk_coord_k * tile_shape_K to match + # the GMEM destination, but the SMEM indexing must be per-CTA-local). + cdV_local = cute.make_identity_tensor((self.cta_tiler[1], self.cta_tiler[2])) + cdK_local = cdV_local + + tiled_t2r_dK = tcgen05.make_tmem_copy(load_op, tdKtdK) + thread_t2r_dK = tiled_t2r_dK.get_slice(dp_idx) + + tTR_cdK = thread_t2r_dK.partition_D(cdK) + tTR_cdK = split_wg(tTR_cdK, num_warp_groups, wg_idx) + tTR_cdK_local = thread_t2r_dK.partition_D(cdK_local) + tTR_cdK_local = split_wg(tTR_cdK_local, num_warp_groups, wg_idx) + tTR_gdK = thread_t2r_dK.partition_D(gdK) + tTR_gdK = split_wg(tTR_gdK, num_warp_groups, wg_idx) + tTR_rdK = cute.make_rmem_tensor(tTR_cdK.shape, self.acc_dtype) + tTR_tdK = thread_t2r_dK.partition_S(tdKtdK) + tTR_tdK = split_wg(tTR_tdK, num_warp_groups, wg_idx) + + mdV_in = cute.make_tensor( + dV.iterator, cute.make_layout((K, self.cta_tiler[2], HB), stride=dV.stride) + ) + offset_mdV = cute.assume(blk_offset[1] * mdV_in.stride[0], divby=64) + mdV = cute.make_tensor(mdV_in.iterator + offset_mdV, mdV_in.layout) + gdV = cute.local_tile(mdV, (self.cta_tiler[1], self.cta_tiler[2]), (None, None, None)) + gdV = gdV[None, None, blk_coord_k, 0, blk_coord_batch] + + cdV = cute.domain_offset( + (blk_coord_k * self.cta_tiler[1], 0), + cute.make_identity_tensor((self.cta_tiler[1], self.cta_tiler[2])), + ) + + tdVtdV = tdVtdV[(None, None), 0, 0] + + tiled_t2r_dV = tcgen05.make_tmem_copy(load_op, tdVtdV) + thread_t2r_dV = tiled_t2r_dV.get_slice(dp_idx) + + tTR_cdV = thread_t2r_dV.partition_D(cdV) + tTR_cdV = split_wg(tTR_cdV, num_warp_groups, wg_idx) + tTR_cdV_local = thread_t2r_dV.partition_D(cdV_local) + tTR_cdV_local = split_wg(tTR_cdV_local, num_warp_groups, wg_idx) + tTR_gdV = thread_t2r_dV.partition_D(gdV) + tTR_gdV = split_wg(tTR_gdV, num_warp_groups, wg_idx) + tTR_rdV = cute.make_rmem_tensor(tTR_cdV.shape, self.acc_dtype) + tTR_tdV = thread_t2r_dV.partition_S(tdVtdV) + tTR_tdV = split_wg(tTR_tdV, num_warp_groups, wg_idx) + + # GMEM destinations for the multi-stage TMA path (gated on not-varlen). + if cutlass.const_expr(not varlen): + mdV_tma_3d = cute.make_tensor( + dV_tma.iterator, + cute.make_layout((K, self.cta_tiler[2], HB), stride=dV_tma.stride), + ) + mdV_tma_cur = mdV_tma_3d[None, None, blk_coord_batch] + gdV_tma = cute.local_tile( + mdV_tma_cur, (self.cta_tiler[1], self.cta_tiler[2]), (blk_coord_k, 0) + ) + gdV_tma_epi = cute.local_tile(gdV_tma, epi_tile_dKV, (0, None)) + + mdK_tma_3d = cute.make_tensor( + dK_tma.iterator, + cute.make_layout((K, self.cta_tiler[2], HB), stride=dK_tma.stride), + ) + mdK_tma_cur = mdK_tma_3d[None, None, blk_coord_batch] + gdK_tma = cute.local_tile( + mdK_tma_cur, (self.cta_tiler[1], self.cta_tiler[2]), (blk_coord_k, 0) + ) + gdK_tma_epi = cute.local_tile(gdK_tma, epi_tile_dKV, (0, None)) + + cta_threads = self.num_compute_warps * self.threads_per_warp + + dkdv_handle = mma_compute_dKdV_consumer.wait_and_advance() + + if blk_coord_k * self.tile_shape_K < problem_shape_k_cur_batch: + cute.copy(tiled_t2r_dV, tTR_tdV, tTR_rdV) + tTR_rdV_cast = cute.make_rmem_tensor(tTR_rdV.shape, dV.element_type) + tTR_rdV_cast.store(tTR_rdV.load().to(dV.element_type)) + + if cutlass.const_expr(not varlen): + # reg -> SMEM via per-element indexed stores using tTR_cdV's + # per-thread (M, N) coords. (M, N) is per-CTA cdV space (M=0..63, + # N=0..255). We map N=(n%epi_cols, n//epi_cols) → (N_within, stage) + # of the 3D s_epi_dV tensor. + for _i in cutlass.range_constexpr(cute.size(tTR_cdV_local, mode=[2])): + for _j in cutlass.range_constexpr(cute.size(tTR_cdV_local[None, 0, _i])): + c = tTR_cdV_local[None, 0, _i][_j] + m_pos = c[0] + n_pos = c[1] + stage_pos = n_pos // epi_cols_dKV + n_within_pos = n_pos % epi_cols_dKV + v = tTR_rdV_cast[None, 0, _i][_j] + s_epi_dV[m_pos, n_within_pos, stage_pos] = v + cute.arch.fence_view_async_shared() + # Inter-WG barrier — both warp-groups must finish their writes + # before the leader warp reads SMEM via TMA. + cute.arch.barrier(barrier_id=5, number_of_threads=cta_threads) + # TMA bulk store, one (64, 64) box per stage. + if leader_warp and wg_idx == 0: + for _stage in cutlass.range_constexpr(total_epi_stages): + sdV_stage = s_epi_dV[None, None, _stage] + gdV_stage = gdV_tma_epi[None, None, _stage] + td_sdV, td_gdV = cpasync.tma_partition( + tma_atom_dV, + 0, + cute.make_layout(1), + cute.group_modes(sdV_stage, 0, 2), + cute.group_modes(gdV_stage, 0, 2), + ) + cute.copy(tma_atom_dV, td_sdV, td_gdV) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + else: + self.store(tTR_gdV, tTR_rdV, tTR_cdV, (K, D)) + + cute.arch.fence_view_async_tmem_load() + dkdv_handle.release() + + dkdv_handle = mma_compute_dKdV_consumer.wait_and_advance() + + if blk_coord_k * self.tile_shape_K < problem_shape_k_cur_batch: + cute.copy(tiled_t2r_dK, tTR_tdK, tTR_rdK) + + for i in cutlass.range(cute.size(tTR_rdK), unroll_full=True): + tTR_rdK[i] = scale_softmax * tTR_rdK[i] + + tTR_rdK_cast = cute.make_rmem_tensor(tTR_rdK.shape, dK.element_type) + tTR_rdK_cast.store(tTR_rdK.load().to(dK.element_type)) + + if cutlass.const_expr(not varlen): + for _i in cutlass.range_constexpr(cute.size(tTR_cdK_local, mode=[2])): + for _j in cutlass.range_constexpr(cute.size(tTR_cdK_local[None, 0, _i])): + c = tTR_cdK_local[None, 0, _i][_j] + m_pos = c[0] + n_pos = c[1] + stage_pos = n_pos // epi_cols_dKV + n_within_pos = n_pos % epi_cols_dKV + v = tTR_rdK_cast[None, 0, _i][_j] + s_epi_dK[m_pos, n_within_pos, stage_pos] = v + cute.arch.fence_view_async_shared() + cute.arch.barrier(barrier_id=6, number_of_threads=cta_threads) + if leader_warp and wg_idx == 0: + for _stage in cutlass.range_constexpr(total_epi_stages): + sdK_stage = s_epi_dK[None, None, _stage] + gdK_stage = gdK_tma_epi[None, None, _stage] + td_sdK, td_gdK = cpasync.tma_partition( + tma_atom_dK, + 0, + cute.make_layout(1), + cute.group_modes(sdK_stage, 0, 2), + cute.group_modes(gdK_stage, 0, 2), + ) + cute.copy(tma_atom_dK, td_sdK, td_gdK) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + else: + self.store(tTR_gdK, tTR_rdK, tTR_cdK, (K, D)) + + cute.arch.fence_view_async_tmem_load() + dkdv_handle.release() + + return mma_compute_dKdV_consumer + + def get_workspace_tensor( + self, + problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]], + workspace: cute.Tensor, + acc_dtype: type[cutlass.Numeric], + varlen: bool, + ) -> tuple[cute.Tensor, cute.Tensor, cute.Tensor]: + """Get workspace tensor.""" + D = problem_shape[2] + H, B = cute.size(problem_shape[3][0]), cute.size(problem_shape[3][1]) + H_r, H_k = problem_shape[3][0] + D = cute.round_up(D, 8) + + # b = 1 for varlen, else batch_size + b = workspace.shape[0] + # s_q_sum for varlen, else s_q_max, already rounded to 8 + S_Q = workspace.shape[1] + + acc_bytes = acc_dtype.width // 8 + sum_OdO_bytes = cute.assume(b * H * S_Q * acc_bytes, divby=acc_bytes) + scaled_lse_bytes = cute.assume(b * H * S_Q * acc_bytes, divby=acc_bytes) + + sum_OdO_iter = workspace.iterator + scaled_lse_iter = sum_OdO_iter + sum_OdO_bytes + + sum_OdO_iter = cute.recast_ptr(sum_OdO_iter, dtype=self.acc_dtype) + scaled_lse_iter = cute.recast_ptr(scaled_lse_iter, dtype=self.acc_dtype) + + sum_OdO = cute.make_tensor( + sum_OdO_iter, + cute.make_layout( + (S_Q, ((H_r, H_k), B)), + stride=(1, ((S_Q, S_Q * H_r), 0 if varlen else S_Q * H)), + ), + ) + scaled_lse = cute.make_tensor( + scaled_lse_iter, + cute.make_layout( + (S_Q, ((H_r, H_k), B)), + stride=(1, ((S_Q, S_Q * H_r), 0 if varlen else S_Q * H)), + ), + ) + + return sum_OdO, scaled_lse + + @staticmethod + def _compute_sum_OdO_grid( + problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]], + block_q: int, + ) -> tuple[Int32, Int32, Int32]: + """Compute grid shape for sum_OdO kernel.""" + return ( + cute.ceil_div(cute.size(problem_shape[0]), block_q), + cute.size(problem_shape[3][0]), # H + cute.size(problem_shape[3][1]), # B + ) + + @staticmethod + def _compute_bwd_grid( + problem_shape: tuple[Int32, Int32, Int32, tuple[tuple[Int32, Int32], Int32]], + block_k: int, + ) -> tuple[Int32, Int32, Int32]: + """Compute grid shape for bwd kernel.""" + K = problem_shape[1] + _, H_K = problem_shape[3][0] + B = problem_shape[3][1] + return (cute.ceil_div(K, block_k), cute.size(H_K), cute.size(B)) + + @staticmethod + def get_workspace_size(s_q: int, d: int, h: int, b: int, acc_dtype: type[cutlass.Numeric]): + """Get workspace size.""" + d = (d + 7) // 8 * 8 # round up to 8 + s_q = (s_q + 7) // 8 * 8 # round up to 8 + workspace_bytes = 0 + # OdO vector + workspace_bytes += acc_dtype.width // 8 + # scaled LSE vector + workspace_bytes += acc_dtype.width // 8 + # FP32 versions of outputs that are churned (start off with Q only) + workspace_bytes += d * acc_dtype.width // 8 + return (b, s_q, h, workspace_bytes) + + def make_and_init_load_mma_K_pipeline(self, load_mma_K_mbar_ptr, cluster_layout_vmnk): + """Create and initialize barrier for load mma Q.""" + load_mma_K_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ) + load_mma_K_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + return pipeline.PipelineTmaUmma.create( + barrier_storage=load_mma_K_mbar_ptr, + num_stages=self.load_mma_K_stage, + producer_group=load_mma_K_producer_group, + consumer_group=load_mma_K_consumer_group, + tx_count=self.tma_copy_K_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + def make_and_init_load_mma_V_pipeline(self, load_mma_V_mbar_ptr, cluster_layout_vmnk): + """Create and initialize barrier for load mma Q.""" + load_mma_V_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ) + load_mma_V_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + return pipeline.PipelineTmaUmma.create( + barrier_storage=load_mma_V_mbar_ptr, + num_stages=self.load_mma_V_stage, + producer_group=load_mma_V_producer_group, + consumer_group=load_mma_V_consumer_group, + tx_count=self.tma_copy_V_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + def make_and_init_load_mma_Q_pipeline(self, load_mma_Q_mbar_ptr, cluster_layout_vmnk): + """Create and initialize barrier for load mma Q.""" + load_mma_Q_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ) + load_mma_Q_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + return pipeline.PipelineTmaUmma.create( + barrier_storage=load_mma_Q_mbar_ptr, + num_stages=self.load_mma_Q_stage, + producer_group=load_mma_Q_producer_group, + consumer_group=load_mma_Q_consumer_group, + tx_count=self.tma_copy_Q_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + def make_and_init_load_mma_QT_pipeline(self, load_mma_QT_mbar_ptr, cluster_layout_vmnk): + """Create and initialize barrier for load mma QT.""" + load_mma_QT_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ) + load_mma_QT_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + return pipeline.PipelineTmaUmma.create( + barrier_storage=load_mma_QT_mbar_ptr, + num_stages=self.load_mma_QT_stage, + producer_group=load_mma_QT_producer_group, + consumer_group=load_mma_QT_consumer_group, + tx_count=self.tma_copy_Q_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + def make_and_init_load_mma_dO_pipeline(self, load_mma_dO_mbar_ptr, cluster_layout_vmnk): + """Create and initialize barrier for load mma dO.""" + load_mma_dO_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.load_warp_id]) + ) + load_mma_dO_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + return pipeline.PipelineTmaUmma.create( + barrier_storage=load_mma_dO_mbar_ptr, + num_stages=self.load_mma_dO_stage, + producer_group=load_mma_dO_producer_group, + consumer_group=load_mma_dO_consumer_group, + tx_count=self.tma_copy_dO_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + def make_and_init_load_compute_LSE_pipeline(self, load_compute_lse_mbar_ptr): + """Create and initialize barrier for load compute lse.""" + load_compute_lse_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp, + # self.threads_per_warp, + ) + load_compute_lse_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * self.num_compute_warps, + # self.threads_per_warp * self.num_compute_warps, + ) + return pipeline.PipelineCpAsync.create( + barrier_storage=load_compute_lse_mbar_ptr, + num_stages=self.load_compute_LSE_stage, + producer_group=load_compute_lse_producer_group, + consumer_group=load_compute_lse_consumer_group, + ) + + def make_and_init_load_compute_sum_OdO_pipeline(self, load_compute_sum_OdO_mbar_ptr): + """Create and initialize barrier for load sum OdO.""" + load_compute_sum_OdO_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp, + # self.threads_per_warp, + ) + load_compute_sum_OdO_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * self.num_compute_warps, + # self.threads_per_warp * self.num_compute_warps, + ) + return pipeline.PipelineCpAsync.create( + barrier_storage=load_compute_sum_OdO_mbar_ptr, + num_stages=self.load_compute_sum_OdO_stage, + producer_group=load_compute_sum_OdO_producer_group, + consumer_group=load_compute_sum_OdO_consumer_group, + ) + + def make_and_init_mma_compute_S_pipeline(self, mma_compute_S_mbar_ptr, cluster_layout_vmnk): + """Create and initialize barrier for mma S.""" + mma_compute_S_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len([self.mma_warp_id]), + ) + mma_compute_S_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0], + ) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=mma_compute_S_mbar_ptr, + num_stages=self.mma_compute_S_stage, + producer_group=mma_compute_S_producer_group, + consumer_group=mma_compute_S_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + # Barrier to between dP = v * dO and consume of dP in compute() + def make_and_init_mma_compute_dP_pipeline(self, mma_compute_dP_mbar_ptr, cluster_layout_vmnk): + """Create and initialize barrier for mma Q.""" + mma_compute_dP_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len([self.mma_warp_id]), + ) + mma_compute_dP_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0], + ) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=mma_compute_dP_mbar_ptr, + num_stages=self.mma_compute_dP_stage, + producer_group=mma_compute_dP_producer_group, + consumer_group=mma_compute_dP_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + def make_and_init_compute_mma_P_pipeline(self, compute_mma_P_mbar_ptr, cluster_layout_vmnk): + """Create and initialize barrier for mma P.""" + compute_mma_P_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0], + # self.num_compute_warps * self.threads_per_warp, + ) + compute_mma_P_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len([self.mma_warp_id]), + ) + return pipeline.PipelineAsyncUmma.create( + barrier_storage=compute_mma_P_mbar_ptr, + num_stages=self.compute_mma_P_stage, + producer_group=compute_mma_P_producer_group, + consumer_group=compute_mma_P_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + def make_and_init_compute_mma_dS_pipeline(self, compute_mma_dS_mbar_ptr, cluster_layout_vmnk): + """Create and initialize barrier for mma dS.""" + + compute_mma_dS_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0], + # self.num_compute_warps * self.threads_per_warp, + ) + compute_mma_dS_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len([self.mma_warp_id]), + ) + + return pipeline.PipelineAsyncUmma.create( + barrier_storage=compute_mma_dS_mbar_ptr, + num_stages=self.compute_mma_dS_stage, + producer_group=compute_mma_dS_producer_group, + consumer_group=compute_mma_dS_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + ) + + def make_and_init_mma_compute_dKdV_pipeline( + self, mma_compute_dKdV_mbar_ptr, cluster_layout_vmnk + ): + """Create and initialize barrier for mma dKdV.""" + mma_compute_dKdV_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len([self.mma_warp_id]), + ) + mma_compute_dKdV_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.num_compute_warps * self.threads_per_warp * cluster_layout_vmnk.shape[0][0], + # self.num_compute_warps * self.threads_per_warp, + ) + return pipeline.PipelineUmmaAsync.create( + barrier_storage=mma_compute_dKdV_mbar_ptr, + num_stages=self.mma_compute_dKdV_stage, + producer_group=mma_compute_dKdV_producer_group, + consumer_group=mma_compute_dKdV_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + ) diff --git a/fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_backward_dqkernel.py b/fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_backward_dqkernel.py new file mode 100644 index 000000000..25d6a91de --- /dev/null +++ b/fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_backward_dqkernel.py @@ -0,0 +1,2317 @@ +# Copyright (c) 2025, Siyu Wang, Shengbin Di, Yuxi Chi, Johnsonms, Linfeng Zheng, Haoyan Huang, Lanbo Li, Yun Zhong, Man Yuan, Minmin Sun, Yong Li, Wei Lin. + +from typing import Type, Tuple, Optional + +import cuda.bindings.driver as cuda + +import math +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +from cutlass.cute.nvgpu import cpasync +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.typing import Int32, Int64, Float32 + +from cutlass.utils import ClcDynamicPersistentTileScheduler +from flash_attn.cute.tile_scheduler import ( + ClcState, + compute_sm100_fmha_grid as compute_grid, + compute_sm100_fmha_grid_clc as compute_grid_clc, + make_sm100_thread_cooperative_group as make_thread_cooperative_group, + Sm100FmhaStaticTileScheduler as FmhaStaticTileScheduler, + Sm100FmhaStaticTileSchedulerParams as FmhaStaticTileSchedulerParams, + Sm100FmhaClcDynamicTileScheduler as FmhaClcDynamicTileScheduler, + Sm100FmhaClcDynamicTileSchedulerParams as FmhaClcDynamicTileSchedulerParams, +) +from flash_attn.cute.mask import ( + Sm100FusedMask as FusedMask, +) +from flash_attn.cute.tile_scheduler import SM100_TMEM_CAPACITY_COLUMNS +import flash_attn.cute.copy_utils as fa_copy_utils + + +class BlackwellFusedMultiHeadAttentionBackwardDQKernel: + def __init__( + self, + acc_dtype: Type[cutlass.Numeric], + mma_tiler: Tuple[int, int, int], + is_causal: bool, + window_size_left: int | None, + window_size_right: int | None, + is_persistent: bool, + split_head: bool, + use_clc_scheduler: bool = False, + ): + self.acc_dtype = acc_dtype + self.mma_tiler = mma_tiler + self.is_causal = is_causal + self.window_size_left = window_size_left + # Keep original behavior (known-good in this repo) + window_size_left = ( + None + if (window_size_left is None or window_size_left < 0) + else cutlass.Int32(window_size_left) + ) + window_size_right = ( + None + if (window_size_right is None or window_size_right < 0) + else cutlass.Int32(window_size_right) + ) + self.window_size_left = None if self.is_causal else window_size_left + self.window_size_right = cutlass.Int32(0) if self.is_causal else window_size_right + self.is_local = (not self.is_causal) and ( + self.window_size_left is not None or self.window_size_right is not None + ) + assert mma_tiler[0] == 128 and mma_tiler[1] == 128, "Only 128x128 tile impl is supported" + assert mma_tiler[2] == 256, "Only 256 is supported for 128x128 tile impl" + self.cta_tiler = ( + mma_tiler[0], + mma_tiler[1], + mma_tiler[2], + ) + self.qk_mma_tiler = ( + 2 * mma_tiler[0], + mma_tiler[1], + min(self.cta_tiler[2], 128) if split_head else self.cta_tiler[2], + ) + self.dov_mma_tiler = self.qk_mma_tiler + self.dsk_mma_tiler = ( + 2 * mma_tiler[0], + min(self.cta_tiler[2], 128) if split_head else self.cta_tiler[2], + mma_tiler[1], + ) + + self.dsk_block_tiler = ( + self.dsk_mma_tiler[0] // 2, + self.dsk_mma_tiler[1], + self.dsk_mma_tiler[2], + ) + self.iterations_qk = self.cta_tiler[2] // self.qk_mma_tiler[2] + self.iterations_dov = self.cta_tiler[2] // self.dov_mma_tiler[2] + self.iterations_dsk = self.cta_tiler[2] // self.dsk_mma_tiler[1] + self.cluster_shape_mn = (2, 1) + self.tmem_warp_shape_mn = (4, 1) + self.is_persistent = is_persistent + self.use_clc_scheduler = use_clc_scheduler + self.use_semantic_trip_range = self.is_causal or self.is_local + + self.compute_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.empty_warp_id = (10, 11) + self.sched_warp_id = self.empty_warp_id[0] if use_clc_scheduler else None + self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + self.num_compute_warps = len(self.compute_warp_ids) + + self.cta_sync_bar_id = 0 + self.tmem_alloc_sync_bar_id = 1 + self.compute_sync_bar_id = 2 + + self.threads_per_warp = 32 + self.threads_per_cta = self.threads_per_warp * len( + ( + *self.compute_warp_ids, # this is to get a round num threads + *self.epilogue_warp_ids, + self.mma_warp_id, + self.load_warp_id, + *self.empty_warp_id, + ) + ) + + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_cta, + ) + + self.tmem_s_offset = 0 + self.tmem_dp_offset = 128 + self.tmem_dq_offset = 256 + + self.num_regs_compute = 256 + self.num_regs_epilogue = 160 + self.num_regs_other = 32 + + self.buffer_align_bytes = 1024 + + def _setup_attributes(self): + self.q_stage = self.iterations_qk + self.k_stage = self.iterations_qk + self.do_stage = self.iterations_dov + self.v_stage = self.iterations_dov + self.kt_stage = 1 + self.qk_acc_stage = 1 + self.dov_acc_stage = 1 + self.dsk_acc_stage = 1 + self.epi_stage = 1 + self.load_compute_LSE_stage = 1 + self.load_compute_sum_OdO_stage = 1 + if cutlass.const_expr(self.use_clc_scheduler): + self.num_clc_stage = 1 + self.num_clc_response_bytes = 16 + + @cute.jit + def __call__( + self, + q_tensor: cute.Tensor, + k_tensor: cute.Tensor, + v_tensor: cute.Tensor, + dq_tensor: cute.Tensor, + do_tensor: cute.Tensor, + lse_tensor: cute.Tensor, + sum_odo_tensor: cute.Tensor, + cum_seqlen_q: Optional[cute.Tensor], + cum_seqlen_k: Optional[cute.Tensor], + scale_softmax: cutlass.Float32, + stream: cuda.CUstream, + ): + varlen = cum_seqlen_q is not None or cum_seqlen_k is not None + # Infer shape metadata from normalized 5D tensors (B, S, H_k, H_r, D), + # similar to the dedicated hd256 forward path. + s_q = q_tensor.shape[1] + s_k = k_tensor.shape[1] + d = q_tensor.shape[4] + h_k = q_tensor.shape[2] + h_r = q_tensor.shape[3] + if cutlass.const_expr(cum_seqlen_q is not None): + b = cum_seqlen_q.shape[0] - 1 + elif cutlass.const_expr(cum_seqlen_k is not None): + b = cum_seqlen_k.shape[0] - 1 + else: + b = q_tensor.shape[0] + # `lse_tensor` / `sum_odo_tensor` are preallocated FP32 buffers (lse_log2 / dpsum) + # whose sequence dimension is padded (rounded up) by the caller. + # Use their leading dimension as the LSE length to ensure correct batch strides. + s_lse = lse_tensor.shape[0] + s_q64 = Int64(s_q) + s_k64 = Int64(s_k) + s_lse64 = Int64(s_lse) + h_r64 = Int64(h_r) + h_k64 = Int64(h_k) + b64 = Int64(b) + # Packed-varlen representation uses batch-dim = 1 and sequence-dim = total_{q,k}. + # Keep the *physical* sequence extent in the tensor layouts so that applying + # `cuseqlen_*` offsets stays within the tensor domain. + s_q_total = q_tensor.shape[1] if cum_seqlen_q is not None else s_q64 + s_k_total = k_tensor.shape[1] if cum_seqlen_k is not None else s_k64 + b_lse = b64 if cum_seqlen_q is None else 1 + stride_b_lse = h_r64 * h_k64 * s_lse64 if cum_seqlen_q is None else 0 + + # (s, d, ((h_r, h_k), b)) + q_layout = cute.make_layout( + (s_q_total, d, ((h_r, h_k), b)), + stride=( + cute.assume(q_tensor.stride[1], divby=64), + q_tensor.stride[4], + ( + (q_tensor.stride[3], q_tensor.stride[2]), + 0 if cum_seqlen_q is not None else cute.assume(q_tensor.stride[0], divby=64), + ), + ), + ) + q = cute.make_tensor(q_tensor.iterator, q_layout) + # (s, d, ((h_r, h_k), b)) + do_layout = cute.make_layout( + (s_q_total, d, ((h_r, h_k), b)), + stride=( + cute.assume(do_tensor.stride[1], divby=64), + do_tensor.stride[4], + ( + (do_tensor.stride[3], do_tensor.stride[2]), + 0 if cum_seqlen_q is not None else cute.assume(do_tensor.stride[0], divby=64), + ), + ), + ) + do = cute.make_tensor(do_tensor.iterator, do_layout) + # (s, d, ((h_r, h_k), b)), 0-stride for h_r to broadcast + k_layout = cute.make_layout( + (s_k_total, d, ((h_r, h_k), b)), + stride=( + cute.assume(k_tensor.stride[1], divby=64), + k_tensor.stride[4], + ( + (0, k_tensor.stride[2]), + 0 if cum_seqlen_k is not None else cute.assume(k_tensor.stride[0], divby=64), + ), + ), + ) + k = cute.make_tensor(k_tensor.iterator, k_layout) + # (d, s, ((h_r, h_k), b)), 0-stride for h_r to broadcast + kt_layout = cute.make_layout( + (d, s_k_total, ((h_r, h_k), b)), + stride=( + k_tensor.stride[4], + cute.assume(k_tensor.stride[1], divby=64), + ( + (0, k_tensor.stride[2]), + 0 if cum_seqlen_k is not None else cute.assume(k_tensor.stride[0], divby=64), + ), + ), + ) + kt = cute.make_tensor(k_tensor.iterator, kt_layout) + # (s, d, ((h_r, h_k), b)), 0-stride for h_r to broadcast + v_layout = cute.make_layout( + (s_k_total, d, ((h_r, h_k), b)), + stride=( + cute.assume(v_tensor.stride[1], divby=64), + v_tensor.stride[4], + ( + (0, v_tensor.stride[2]), + 0 if cum_seqlen_k is not None else cute.assume(v_tensor.stride[0], divby=64), + ), + ), + ) + v = cute.make_tensor(v_tensor.iterator, v_layout) + # (s, ((h_r, h_k), b)) + lse = cute.make_tensor(lse_tensor.iterator, lse_tensor.layout) + # (s, ((h_r, h_k), b)) + sum_odo_layout = cute.make_layout( + (s_lse64, ((h_r, h_k), b_lse)), + stride=(1, ((s_lse64, h_r64 * s_lse64), stride_b_lse)), + ) + sum_odo = cute.make_tensor(sum_odo_tensor.iterator, sum_odo_layout) + # (s, d, ((h_r, h_k), b)) + dq_layout = cute.make_layout( + (s_q_total, d, ((h_r, h_k), b)), + stride=( + cute.assume(dq_tensor.stride[1], divby=64), + dq_tensor.stride[4], + ( + (dq_tensor.stride[3], dq_tensor.stride[2]), + 0 if cum_seqlen_q is not None else cute.assume(dq_tensor.stride[0], divby=64), + ), + ), + ) + dq = cute.make_tensor(dq_tensor.iterator, dq_layout) + + # setup static attributes before smem/grid/tma computation + self.q_dtype = q.element_type + self.k_dtype = k.element_type + self.v_dtype = v.element_type + self.do_dtype = do.element_type + self.dq_dtype = dq.element_type + self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.q_dtype.width + + if cutlass.const_expr(self.use_clc_scheduler): + self.tile_sched_params, grid = compute_grid_clc( + (s_q, dq.shape[1], dq.shape[2]) if cum_seqlen_q is not None else dq.shape, + self.cta_tiler, + (*self.cluster_shape_mn, 1), + ) + else: + self.tile_sched_params, grid = compute_grid( + (s_q, dq.shape[1], dq.shape[2]) if cum_seqlen_q is not None else dq.shape, + self.cta_tiler, + self.is_persistent, + ) + + self.q_major_mode = utils.LayoutEnum.from_tensor(q).mma_major_mode() + self.do_major_mode = utils.LayoutEnum.from_tensor(do).mma_major_mode() + self.k_major_mode = utils.LayoutEnum.from_tensor(k).mma_major_mode() + self.v_major_mode = utils.LayoutEnum.from_tensor(v).mma_major_mode() + self.dq_layout = utils.LayoutEnum.from_tensor(dq) + + if cutlass.const_expr(self.q_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of q is not supported") + if cutlass.const_expr(self.k_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of k is not supported") + if cutlass.const_expr(self.v_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of v is not supported") + if cutlass.const_expr(self.do_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of v is not supported") + + # check type consistency + if cutlass.const_expr(self.q_dtype != self.k_dtype): + raise TypeError(f"Type mismatch: {self.q_dtype} != {self.k_dtype}") + if cutlass.const_expr(self.q_dtype != self.v_dtype): + raise TypeError(f"Type mismatch: {self.q_dtype} != {self.v_dtype}") + if cutlass.const_expr(self.q_dtype != self.do_dtype): + raise TypeError(f"Type mismatch: {self.q_dtype} != {self.do_dtype}") + + self._setup_attributes() + + cta_group = tcgen05.CtaGroup.TWO + # the intermediate tensor p is from tmem & k-major + ds_source = tcgen05.OperandSource.TMEM + ds_major_mode = tcgen05.OperandMajorMode.K + k_trans_major_mode = tcgen05.OperandMajorMode.MN + qk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + self.q_major_mode, + self.k_major_mode, + self.acc_dtype, + cta_group, + self.qk_mma_tiler[:2], + ) + dov_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.do_dtype, + self.do_major_mode, + self.v_major_mode, + self.acc_dtype, + cta_group, + self.dov_mma_tiler[:2], + ) + dsk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + ds_major_mode, + k_trans_major_mode, + self.acc_dtype, + cta_group, + self.dsk_mma_tiler[:2], + ds_source, + ) + + self.cluster_shape_mnk = (*self.cluster_shape_mn, 1) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (qk_tiled_mma.thr_id.shape,), + ) + + self.epi_tile = self.dsk_block_tiler[:2] + + q_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.q_stage, + ) + k_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.qk_mma_tiler, + self.k_dtype, + self.k_stage, + ) + do_smem_layout_staged = sm100_utils.make_smem_layout_a( + dov_tiled_mma, + self.dov_mma_tiler, + self.do_dtype, + self.do_stage, + ) + v_smem_layout_staged = sm100_utils.make_smem_layout_b( + dov_tiled_mma, + self.dov_mma_tiler, + self.v_dtype, + self.v_stage, + ) + ds_tmem_layout_staged = sm100_utils.make_smem_layout_a( + dsk_tiled_mma, + self.dsk_mma_tiler, + self.q_dtype, + self.qk_acc_stage, + ) + ds_tmem_layout = cute.select(ds_tmem_layout_staged, mode=[0, 1, 2]) + kt_smem_layout_staged = sm100_utils.make_smem_layout_b( + dsk_tiled_mma, + self.dsk_mma_tiler, + self.k_dtype, + self.dsk_acc_stage, + ) + + # TMA load for Q + tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) + + q_smem_layout = cute.select(q_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q, tma_tensor_q = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q, + q_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + # TMA load for K + k_smem_layout = cute.select(k_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_k, tma_tensor_k = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + k, + k_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + # TMA load for dO + do_smem_layout = cute.select(do_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_do, tma_tensor_do = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + do, + do_smem_layout, + self.dov_mma_tiler, + dov_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + # TMA load for V + v_smem_layout = cute.select(v_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + v, + v_smem_layout, + self.dov_mma_tiler, + dov_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + # TMA load for KT + kt_smem_layout = cute.select(kt_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_kt, tma_tensor_kt = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + kt, + kt_smem_layout, + self.dsk_mma_tiler, + dsk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + lse_smem_layout = cute.make_layout((self.cta_tiler[0], self.load_compute_LSE_stage)) + sum_odo_smem_layout = cute.make_layout((self.cta_tiler[0], self.load_compute_sum_OdO_stage)) + + q_copy_size = cute.size_in_bytes(self.q_dtype, q_smem_layout) + k_copy_size = cute.size_in_bytes(self.k_dtype, k_smem_layout) + do_copy_size = cute.size_in_bytes(self.do_dtype, do_smem_layout) + v_copy_size = cute.size_in_bytes(self.v_dtype, v_smem_layout) + kt_copy_size = cute.size_in_bytes(self.k_dtype, kt_smem_layout) + self.tma_copy_q_bytes = q_copy_size * cute.size(qk_tiled_mma.thr_id.shape) + self.tma_copy_k_bytes = k_copy_size * cute.size(qk_tiled_mma.thr_id.shape) + self.tma_copy_do_bytes = do_copy_size * cute.size(qk_tiled_mma.thr_id.shape) + self.tma_copy_v_bytes = v_copy_size * cute.size(qk_tiled_mma.thr_id.shape) + self.tma_copy_kt_bytes = kt_copy_size * cute.size(qk_tiled_mma.thr_id.shape) + + # dQ epilogue TMA store. Use a single (M, 64) SMEM stage and rotate it + # across the hd=256 N dimension to stay within the SMEM budget. + tma_store_op = cpasync.CopyBulkTensorTileS2GOp() + epi_cols_dQ = math.gcd(128 // (dq.element_type.width // 8), self.epi_tile[1]) + epi_tile_dQ = (self.epi_tile[0], epi_cols_dQ) + sdQ_epi_layout = sm100_utils.make_smem_layout_epi( + dq.element_type, + self.dq_layout, + epi_tile_dQ, + 1, + ) + tma_atom_dQ, tma_tensor_dQ = cpasync.make_tiled_tma_atom( + tma_store_op, + dq, + cute.select(sdQ_epi_layout, mode=[0, 1]), + epi_tile_dQ, + ) + + @cute.struct + class SharedStorage: + # TMA G2S load barriers: LOAD warp (producer) -> MMA warp (consumer) + load_q_mbar_ptr: cute.struct.MemRange[ + Int64, self.q_stage * 2 + ] # load_q_{producer,consumer} + load_do_mbar_ptr: cute.struct.MemRange[ + Int64, self.do_stage * 2 + ] # load_do_{producer,consumer} + load_k_mbar_ptr: cute.struct.MemRange[ + Int64, self.k_stage * 2 + ] # load_k_{producer,consumer} + load_kt_mbar_ptr: cute.struct.MemRange[ + Int64, self.kt_stage * 2 + ] # load_kt_{producer,consumer} + load_v_mbar_ptr: cute.struct.MemRange[ + Int64, self.v_stage * 2 + ] # load_v_{producer,consumer} + mma_s_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2] + mma_dp_mbar_ptr: cute.struct.MemRange[Int64, self.dov_acc_stage * 2] + mma_dq_mbar_ptr: cute.struct.MemRange[Int64, self.epi_stage * 2] + ds_mma_mbar_ptr: cute.struct.MemRange[Int64, self.dsk_acc_stage * 2] + lse_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.load_compute_LSE_stage * 2] + sum_odo_mbar_ptr: cute.struct.MemRange[ + cutlass.Int64, self.load_compute_sum_OdO_stage * 2 + ] + # A CTA-wide "TMEM lifetime" barrier used to safely deallocate TMEM after all users finish. + tmem_dealloc_mbar_ptr: Int64 + # Tmem holding buffer + tmem_holding_buf: Int32 + # CLC pipeline barriers and response buffer + clc_mbar_ptr: cute.struct.MemRange[Int64, 2] + clc_response: cute.struct.MemRange[Int32, 4] + + self.shared_storage = SharedStorage + + grid = cute.round_up(grid, self.cluster_shape_mnk) + # Launch the kernel synchronously + self.kernel( + qk_tiled_mma, + dov_tiled_mma, + dsk_tiled_mma, + tma_atom_q, + tma_tensor_q, + tma_atom_k, + tma_tensor_k, + tma_atom_v, + tma_tensor_v, + tma_atom_do, + tma_tensor_do, + tma_atom_kt, + tma_tensor_kt, + tma_atom_dQ, + tma_tensor_dQ, + lse, + sum_odo, + dq, + cum_seqlen_q, + cum_seqlen_k, + scale_softmax, + self.window_size_left, + self.window_size_right, + self.cluster_layout_vmnk, + q_smem_layout_staged, + k_smem_layout_staged, + v_smem_layout_staged, + do_smem_layout_staged, + kt_smem_layout_staged, + ds_tmem_layout, + sdQ_epi_layout, + lse_smem_layout, + sum_odo_smem_layout, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + stream=stream, + min_blocks_per_mp=1, + ) + + # GPU device kernel + @cute.kernel + def kernel( + self, + qk_tiled_mma: cute.TiledMma, + dov_tiled_mma: cute.TiledMma, + dsk_tiled_mma: cute.TiledMma, + tma_atom_q: cute.CopyAtom, + mQ_qdl: cute.Tensor, + tma_atom_k: cute.CopyAtom, + mK_kdl: cute.Tensor, + tma_atom_v: cute.CopyAtom, + mV_dkl: cute.Tensor, + tma_atom_do: cute.CopyAtom, + mdO_qdl: cute.Tensor, + tma_atom_kt: cute.CopyAtom, + mK_dkl: cute.Tensor, + tma_atom_dQ: cute.CopyAtom, + mdQ_tma: cute.Tensor, + mLSE: cute.Tensor, + mSum_OdO: cute.Tensor, + mdQ_qdl: cute.Tensor, + cum_seqlen_q: Optional[cute.Tensor], + cum_seqlen_k: Optional[cute.Tensor], + scale_softmax: Float32, + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + cluster_layout_vmnk: cute.Layout, + q_smem_layout_staged: cute.ComposedLayout, + k_smem_layout_staged: cute.ComposedLayout, + v_smem_layout_staged: cute.ComposedLayout, + do_smem_layout_staged: cute.ComposedLayout, + kt_smem_layout_staged: cute.ComposedLayout, + ds_tmem_layout_staged: cute.ComposedLayout, + sdQ_epi_layout: cute.ComposedLayout, + lse_smem_layout: cute.Layout, + sum_odo_smem_layout: cute.Layout, + tile_sched_params: FmhaStaticTileSchedulerParams | FmhaClcDynamicTileSchedulerParams, + ): + # llvm.inline_asm( + # None, + # [], + # '.pragma "global knob CommonIntoMultiBlockLoop=1";', + # "", + # has_side_effects=True, + # is_align_stack=False, + # asm_dialect=llvm.AsmDialect.AD_ATT, + # ) + + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + # + # Prefetch tma desc + # + if warp_idx == self.load_warp_id: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_k) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_v) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_do) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_kt) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_dQ) + + bidx, bidy, bidz = cute.arch.block_idx() + tidx, _, _ = cute.arch.thread_idx() + varlen = cum_seqlen_q is not None or cum_seqlen_k is not None + mma_tile_coord_v = bidx % cute.size(qk_tiled_mma.thr_id.shape) + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + + # Alloc + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + load_q_producer, load_q_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.q_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_q_bytes, + barrier_storage=storage.load_q_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + load_k_producer, load_k_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.k_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_k_bytes, + barrier_storage=storage.load_k_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + load_v_producer, load_v_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.v_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_v_bytes, + barrier_storage=storage.load_v_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + load_do_producer, load_do_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.do_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_do_bytes, + barrier_storage=storage.load_do_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + load_kt_producer, load_kt_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.kt_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_kt_bytes, + barrier_storage=storage.load_kt_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + mma_s_producer, mma_s_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.qk_acc_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group( + len(self.compute_warp_ids) * self.threads_per_warp * self.cluster_shape_mnk[0], + ), + barrier_storage=storage.mma_s_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + mma_dp_producer, mma_dp_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.dov_acc_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group( + len(self.compute_warp_ids) * self.threads_per_warp * self.cluster_shape_mnk[0], + ), + barrier_storage=storage.mma_dp_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + ds_mma_producer, ds_mma_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.dsk_acc_stage, + producer_group=make_thread_cooperative_group( + len(self.compute_warp_ids) * self.threads_per_warp * self.cluster_shape_mnk[0], + ), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + barrier_storage=storage.ds_mma_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + mma_dq_producer, mma_dq_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.epi_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group( + len(self.epilogue_warp_ids) * self.threads_per_warp * self.cluster_shape_mnk[0], + ), + barrier_storage=storage.mma_dq_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + + load_lse_producer, load_lse_consumer = pipeline.PipelineCpAsync.create( + num_stages=self.load_compute_LSE_stage, + producer_group=make_thread_cooperative_group(self.threads_per_warp), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * self.num_compute_warps + ), + barrier_storage=storage.lse_mbar_ptr.data_ptr(), + ).make_participants() + load_sum_odo_producer, load_sum_odo_consumer = pipeline.PipelineCpAsync.create( + num_stages=self.load_compute_sum_OdO_stage, + producer_group=make_thread_cooperative_group(self.threads_per_warp), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * self.num_compute_warps + ), + barrier_storage=storage.sum_odo_mbar_ptr.data_ptr(), + ).make_participants() + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.epilogue_warp_ids[0], + is_two_cta=True, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + tmem.allocate(self.tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + + # Initialize CLC state if using dynamic scheduler + if cutlass.const_expr(self.use_clc_scheduler): + clc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + cluster_size = cute.size(self.cluster_shape_mnk) + num_clc_consumer_threads = self.threads_per_warp * ( + 1 # sched_warp (CTA 0 only) + + cluster_size + * ( + len(self.compute_warp_ids) + + len(self.epilogue_warp_ids) + + 1 # mma_warp + + 1 # load_warp + ) + ) + clc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_clc_consumer_threads + ) + clc_response_ptr = storage.clc_response.data_ptr() + clc = ClcState.create( + hw_scheduler=ClcDynamicPersistentTileScheduler.create( + self.tile_sched_params.clc_hw_params(), + cute.arch.block_idx(), + cute.arch.grid_dim(), + clc_response_ptr, + ), + pipeline=pipeline.PipelineClcFetchAsync.create( + barrier_storage=storage.clc_mbar_ptr.data_ptr(), + num_stages=self.num_clc_stage, + producer_group=clc_pipeline_producer_group, + consumer_group=clc_pipeline_consumer_group, + tx_count=self.num_clc_response_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ), + consumer_state=pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_clc_stage + ), + producer_state=pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_clc_stage + ), + ) + else: + clc = None + clc_response_ptr = None + + # Cluster arrive after barrier init + pipeline.pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + + sQ = smem.allocate_tensor( + element_type=self.q_dtype, + layout=q_smem_layout_staged.outer, + swizzle=q_smem_layout_staged.inner, + byte_alignment=128, + ) + sK = smem.allocate_tensor( + element_type=self.k_dtype, + layout=k_smem_layout_staged.outer, + swizzle=k_smem_layout_staged.inner, + byte_alignment=128, + ) + # K and V now use separate memory since we removed the transform stage + sV = smem.allocate_tensor( + element_type=self.v_dtype, + layout=v_smem_layout_staged.outer, + swizzle=v_smem_layout_staged.inner, + byte_alignment=128, + ) + sdO = smem.allocate_tensor( + element_type=self.do_dtype, + layout=do_smem_layout_staged.outer, + swizzle=do_smem_layout_staged.inner, + byte_alignment=128, + ) + sKT = smem.allocate_tensor( + element_type=self.k_dtype, + layout=kt_smem_layout_staged.outer, + swizzle=kt_smem_layout_staged.inner, + byte_alignment=128, + ) + sLSE = smem.allocate_tensor( + element_type=self.acc_dtype, + layout=lse_smem_layout, + byte_alignment=128, + ) + sSum_OdO = smem.allocate_tensor( + element_type=self.acc_dtype, + layout=sum_odo_smem_layout, + byte_alignment=128, + ) + # Alias the dQ TMA epilogue staging buffer onto sdO. A standalone + # (128, 64) bf16 allocation exceeds B200's per-block SMEM limit for + # this kernel, while sdO is no longer needed by the epilogue warp once + # the dQ accumulator is ready to store. + s_epi_dQ = cute.make_tensor( + cute.recast_ptr(sdO.iterator, sdQ_epi_layout.inner, self.dq_dtype), + sdQ_epi_layout.outer, + ) + qk_thr_mma = qk_tiled_mma.get_slice(mma_tile_coord_v) # default 1sm + dov_thr_mma = dov_tiled_mma.get_slice(mma_tile_coord_v) # default 1sm + dsk_thr_mma = dsk_tiled_mma.get_slice(mma_tile_coord_v) # default 1sm + tSrQ = qk_thr_mma.make_fragment_A(sQ) + tSrK = qk_thr_mma.make_fragment_B(sK) + tdPrdO = dov_thr_mma.make_fragment_A(sdO) + tdPrV = dov_thr_mma.make_fragment_B(sV) + tdQrKT = dsk_thr_mma.make_fragment_B(sKT) + qk_acc_shape = qk_thr_mma.partition_shape_C((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) + tStS = qk_thr_mma.make_fragment_C(cute.append(qk_acc_shape, self.qk_acc_stage)) + dov_acc_shape = dov_thr_mma.partition_shape_C( + (self.dov_mma_tiler[0], self.dov_mma_tiler[1]) + ) + tdPtdP = dov_thr_mma.make_fragment_C(cute.append(dov_acc_shape, self.dov_acc_stage)) + dsk_acc_shape = dsk_thr_mma.partition_shape_C( + (self.dsk_mma_tiler[0], self.dsk_mma_tiler[1]) + ) + tdQtdQ = dsk_thr_mma.make_fragment_C(dsk_acc_shape) + tdQtdQ_layout = cute.append( + tdQtdQ.layout, + cute.make_layout( + self.iterations_dsk, + stride=self.dsk_mma_tiler[1] // self.tmem_warp_shape_mn[1], + ), + ) + tStS = cute.make_tensor(tStS.iterator + self.tmem_s_offset, tStS.layout) + tdPtdP = cute.make_tensor(tdPtdP.iterator + self.tmem_dp_offset, tdPtdP.layout) + tdQtdQ_staged = cute.make_tensor(tdQtdQ.iterator + self.tmem_dq_offset, tdQtdQ_layout) + + # /////////////////////////////////////////////////////////////////////////////// + # EMPTY + # /////////////////////////////////////////////////////////////////////////////// + for _i in cutlass.range_constexpr(len(self.empty_warp_id)): + if warp_idx == self.empty_warp_id[_i]: + cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + + if cutlass.const_expr(self.use_clc_scheduler): + tile_sched = FmhaClcDynamicTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + clc_response_ptr, + clc, + ) + else: + blk_idx = cute.arch.block_idx() + tile_sched = FmhaStaticTileScheduler( + tile_sched_params, blk_idx[0], blk_idx, cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + # Cluster wait + pipeline.pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + # /////////////////////////////////////////////////////////////////////////////// + # LOAD + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.load_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + batch_coord = curr_block_coord[2][1] + seqlen_q = mQ_qdl.shape[0] + seqlen_k = mK_kdl.shape[0] + cuseqlen_q = Int32(0) + cuseqlen_k = Int32(0) + is_valid_q = True + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + is_valid_q = FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.qk_mma_tiler[0], + mma_block_coord[0], + seqlen_q, + ) + if cutlass.const_expr(cum_seqlen_k is not None): + cuseqlen_k = cum_seqlen_k[batch_coord] + seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k + seqlen_kv_loop_start, seqlen_kv_loop_steps = ( + FusedMask.get_trip_start_count_via_block_info( + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + self.is_causal, + self.is_local, + window_size_left, + window_size_right, + ) + ) + is_valid_k = seqlen_kv_loop_steps > 0 + has_work = is_valid_q and is_valid_k + + if has_work: + block_offset = ( + cuseqlen_q, + cuseqlen_k, + Int32(0), + ((Int32(0), Int32(0)), Int32(0)), + ) + mQ_qdl_ = cute.domain_offset(cute.select(block_offset, mode=[0, 2, 3]), mQ_qdl) + mK_kdl_ = cute.domain_offset(cute.select(block_offset, mode=[1, 2, 3]), mK_kdl) + mdO_qdl_ = cute.domain_offset( + cute.select(block_offset, mode=[0, 2, 3]), mdO_qdl + ) + mV_dkl_ = cute.domain_offset(cute.select(block_offset, mode=[1, 2, 3]), mV_dkl) + mK_dkl_ = cute.domain_offset(cute.select(block_offset, mode=[2, 1, 3]), mK_dkl) + block_offset_stats = block_offset + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q_stats = cute.assume( + (cuseqlen_q + batch_coord * self.cta_tiler[0]) + // self.cta_tiler[0] + * self.cta_tiler[0], + divby=self.cta_tiler[0], + ) + block_offset_stats = ( + cuseqlen_q_stats, + block_offset[1], + block_offset[2], + block_offset[3], + ) + LSE = cute.domain_offset(cute.select(block_offset_stats, mode=[0, 3]), mLSE) + sum_OdO = cute.domain_offset( + cute.select(block_offset_stats, mode=[0, 3]), mSum_OdO + ) + + # Local tile partition global tensors + q_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # (bM, bK, loopM, loopK, loopL) + gQ_qdl = cute.flat_divide(mQ_qdl_, cute.select(self.qk_mma_tiler, mode=[0, 2])) + tSgQ_qdl = qk_thr_mma.partition_A(gQ_qdl) + tQsQ, tQgQ_qdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_q, + block_in_cluster_coord_vmnk[2], + q_cta_layout, + cute.group_modes(sQ, 0, 3), + cute.group_modes(tSgQ_qdl, 0, 3), + ) + k_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + gK_kdl = cute.flat_divide(mK_kdl_, cute.select(self.qk_mma_tiler, mode=[1, 2])) + tSgK_kdl = qk_thr_mma.partition_B(gK_kdl) + tKsK, tKgK_kdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_k, + block_in_cluster_coord_vmnk[1], + k_cta_layout, + cute.group_modes(sK, 0, 3), + cute.group_modes(tSgK_kdl, 0, 3), + ) + do_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # (bM, bK, loopM, loopK, loopL) + gdO_qdl = cute.flat_divide( + mdO_qdl_, cute.select(self.dov_mma_tiler, mode=[0, 2]) + ) + tdPgdO_qdl = dov_thr_mma.partition_A(gdO_qdl) + tdOsdO, tdOgdO_qdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_do, + block_in_cluster_coord_vmnk[2], + do_cta_layout, + cute.group_modes(sdO, 0, 3), + cute.group_modes(tdPgdO_qdl, 0, 3), + ) + v_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + gV_dkl = cute.flat_divide(mV_dkl_, cute.select(self.dov_mma_tiler, mode=[1, 2])) + tSgV_dkl = dov_thr_mma.partition_B(gV_dkl) + tVsV, tVgV_dkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_v, + block_in_cluster_coord_vmnk[1], + v_cta_layout, + cute.group_modes(sV, 0, 3), + cute.group_modes(tSgV_dkl, 0, 3), + ) + # kt layout + kt_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + gK_dkl = cute.flat_divide(mK_dkl_, cute.select(self.dsk_mma_tiler, mode=[1, 2])) + tdQgK_dkl = dsk_thr_mma.partition_B(gK_dkl) + tKTsKT, tKgK_dkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_kt, + block_in_cluster_coord_vmnk[1], + kt_cta_layout, + cute.group_modes(sKT, 0, 3), + cute.group_modes(tdQgK_dkl, 0, 3), + ) + # ((atom_v, rest_v), RestK) + tQgQ = tQgQ_qdl[None, mma_block_coord[0], None, mma_block_coord[2]] + # ((atom_v, rest_v), RestK) + tdOgdO = tdOgdO_qdl[None, mma_block_coord[0], None, mma_block_coord[2]] + # ((atom_v, rest_v), RestN, RestK) + tKgK = tKgK_kdl[None, None, None, mma_block_coord[2]] + # ((atom_v, rest_v), RestN, RestK) + tVgV = tVgV_dkl[None, None, None, mma_block_coord[2]] + # ((atom_v, rest_v), RestN, RestK) + tKTgKT = tKgK_dkl[None, None, None, mma_block_coord[2]] + + # LSE + lse_handle = load_lse_producer.acquire_and_advance() + # 32 threads loading 128 values of 32b each + # so 4*32b = 128b + thread_idx = tidx % self.threads_per_warp + async_copy_num_elts = sLSE.shape[0] // self.threads_per_warp + atom_async_copy = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.ALWAYS), + self.acc_dtype, + num_bits_per_copy=self.acc_dtype.width, + ) + sLSE_for_copy = cute.flat_divide(sLSE, (1,)) + LSE_for_copy = cute.flat_divide(LSE, (1,)) + for i in cutlass.range_constexpr(async_copy_num_elts): + LSE_idx = ( + self.cta_tiler[0] * curr_block_coord[0] + + thread_idx * async_copy_num_elts + ) + if cute.elem_less(LSE_idx + i, seqlen_q): + cute.copy( + atom_async_copy, + LSE_for_copy[None, LSE_idx + i, curr_block_coord[2]], + sLSE_for_copy[ + None, + thread_idx * async_copy_num_elts + i, + lse_handle.index, + ], + ) + else: + sLSE_for_copy[ + None, + thread_idx * async_copy_num_elts + i, + lse_handle.index, + ].fill(0.0) + lse_handle.commit() + + sum_odo_handle = load_sum_odo_producer.acquire_and_advance() + sSum_OdO_for_copy = cute.flat_divide(sSum_OdO, (1,)) + sum_OdO_for_copy = cute.flat_divide(sum_OdO, (1,)) + for i in cutlass.range_constexpr(async_copy_num_elts): + sum_OdO_idx = ( + self.cta_tiler[0] * curr_block_coord[0] + + thread_idx * async_copy_num_elts + ) + if cute.elem_less(sum_OdO_idx + i, seqlen_q): + cute.copy( + atom_async_copy, + sum_OdO_for_copy[None, sum_OdO_idx + i, curr_block_coord[2]], + sSum_OdO_for_copy[ + None, + thread_idx * async_copy_num_elts + i, + sum_odo_handle.index, + ], + ) + else: + sSum_OdO_for_copy[ + None, + thread_idx * async_copy_num_elts + i, + sum_odo_handle.index, + ].fill(0.0) + sum_odo_handle.commit() + + # Q + for iter in cutlass.range(self.iterations_qk, unroll=1): + q_handle = load_q_producer.acquire_and_advance() + cute.copy( + tma_atom_q, + tQgQ[None, iter], + tQsQ[None, q_handle.index], + tma_bar_ptr=q_handle.barrier, + ) + # dO + for iter in cutlass.range(self.iterations_dov, unroll=1): + do_handle = load_do_producer.acquire_and_advance() + cute.copy( + tma_atom_do, + tdOgdO[None, iter], + tdOsdO[None, do_handle.index], + tma_bar_ptr=do_handle.barrier, + ) + + kv_coord = seqlen_kv_loop_start + for i in cutlass.range(0, seqlen_kv_loop_steps, 1, unroll=1): + # Ki + for iter in cutlass.range(self.iterations_qk, unroll=1): + k_handle = load_k_producer.acquire_and_advance() + cute.copy( + tma_atom_k, + tKgK[None, kv_coord, iter], + tKsK[None, k_handle.index], + tma_bar_ptr=k_handle.barrier, + ) + # Vi + for iter in cutlass.range(self.iterations_dov, unroll=1): + v_handle = load_v_producer.acquire_and_advance() + cute.copy( + tma_atom_v, + tVgV[None, kv_coord, iter], + tVsV[None, v_handle.index], + tma_bar_ptr=v_handle.barrier, + ) + # KTi + for iter in cutlass.range(self.iterations_dsk, unroll=1): + kt_handle = load_kt_producer.acquire_and_advance() + cute.copy( + tma_atom_kt, + tKTgKT[None, iter, kv_coord], + tKTsKT[None, kt_handle.index], + tma_bar_ptr=kt_handle.barrier, + ) + kv_coord += 1 + + work_tile = tile_sched.advance_to_next_work() + # End of persistent scheduler loop + load_k_producer.tail() + load_v_producer.tail() + load_kt_producer.tail() + load_q_producer.tail() + load_do_producer.tail() + load_lse_producer.tail() + load_sum_odo_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # MMA + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + is_leader_cta = cta_rank_in_cluster % 2 == 0 + + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + seqlen_q = mQ_qdl.shape[0] + seqlen_k = mK_kdl.shape[0] + batch_coord = curr_block_coord[2][1] + is_valid_q = True + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + is_valid_q = FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.qk_mma_tiler[0], + mma_block_coord[0], + seqlen_q, + ) + if cutlass.const_expr(cum_seqlen_k is not None): + cuseqlen_k = cum_seqlen_k[batch_coord] + seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k + seqlen_kv_loop_start, seqlen_kv_loop_steps = ( + FusedMask.get_trip_start_count_via_block_info( + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + self.is_causal, + self.is_local, + window_size_left, + window_size_right, + ) + ) + is_valid_k = seqlen_kv_loop_steps > 0 + has_work = is_valid_q and is_valid_k + + if has_work: + # dq_handle = mma_dq_producer.acquire_and_advance() + load_q_releaser = load_q_consumer.clone() + load_do_releaser = load_do_consumer.clone() + dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + num_innerloop = 8 + + if is_leader_cta: + dq_handle = mma_dq_producer.acquire_and_advance() + if seqlen_kv_loop_steps > 1: + # QK0 + s_handle = mma_s_producer.acquire_and_advance() + tStS_slice = tStS[None, None, None, s_handle.index] + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for iter in cutlass.range(self.iterations_qk, unroll=1): + load_q_consumer.wait_and_advance() + tSrQ_slice = tSrQ[None, None, None, iter] + + k_handle = load_k_consumer.wait_and_advance() + tSrK_trans_slice = tSrK[None, None, None, k_handle.index] + num_kphases = cute.size(tSrQ_slice, mode=[2]) + if cutlass.const_expr(num_kphases % num_innerloop == 0): + num_outer_iter = num_kphases // num_innerloop + for outer_iter in cutlass.range(num_outer_iter, unroll=1): + for kphase_idx in cutlass.range( + num_innerloop, unroll_full=True + ): + kphase_coord = ( + None, + None, + outer_iter * num_innerloop + kphase_idx, + ) + cute.gemm( + qk_tiled_mma, + tStS_slice, + tSrQ_slice[kphase_coord], + tSrK_trans_slice[kphase_coord], + tStS_slice, + ) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + else: + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + qk_tiled_mma, + tStS_slice, + tSrQ_slice[kphase_coord], + tSrK_trans_slice[kphase_coord], + tStS_slice, + ) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + k_handle.release() + cute.arch.fence_view_async_tmem_store() + s_handle.commit() + + # dOV0 + dp_handle = mma_dp_producer.acquire_and_advance() + tdPtdP_slice = tdPtdP[None, None, None, dp_handle.index] + dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for iter in cutlass.range(self.iterations_dov, unroll=1): + load_do_consumer.wait_and_advance() + tdPrdO_slice = tdPrdO[None, None, None, iter] + v_handle = load_v_consumer.wait_and_advance() + tdPrV_trans_slice = tdPrV[None, None, None, v_handle.index] + num_kphases = cute.size(tdPrdO_slice, mode=[2]) + if cutlass.const_expr(num_kphases % num_innerloop == 0): + num_outer_iter = num_kphases // num_innerloop + for outer_iter in cutlass.range(num_outer_iter, unroll=1): + for kphase_idx in cutlass.range( + num_innerloop, unroll_full=True + ): + kphase_coord = ( + None, + None, + outer_iter * num_innerloop + kphase_idx, + ) + cute.gemm( + dov_tiled_mma, + tdPtdP_slice, + tdPrdO_slice[kphase_coord], + tdPrV_trans_slice[kphase_coord], + tdPtdP_slice, + ) + dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + else: + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + dov_tiled_mma, + tdPtdP_slice, + tdPrdO_slice[kphase_coord], + tdPrV_trans_slice[kphase_coord], + tdPtdP_slice, + ) + dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + v_handle.release() + cute.arch.fence_view_async_tmem_store() + dp_handle.commit() + + for i in cutlass.range(1, seqlen_kv_loop_steps - 1, 1, unroll=1): + # QKi + s_handle = mma_s_producer.acquire_and_advance() + + tStS_slice = tStS[None, None, None, s_handle.index] + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for iter in cutlass.range(self.iterations_qk, unroll=1): + tSrQ_slice = tSrQ[None, None, None, iter] + k_handle = load_k_consumer.wait_and_advance() + tSrK_trans_slice = tSrK[None, None, None, k_handle.index] + num_kphases = cute.size(tSrQ_slice, mode=[2]) + if cutlass.const_expr(num_kphases % num_innerloop == 0): + num_outer_iter = num_kphases // num_innerloop + for outer_iter in cutlass.range(num_outer_iter, unroll=1): + for kphase_idx in cutlass.range( + num_innerloop, unroll_full=True + ): + kphase_coord = ( + None, + None, + outer_iter * num_innerloop + kphase_idx, + ) + cute.gemm( + qk_tiled_mma, + tStS_slice, + tSrQ_slice[kphase_coord], + tSrK_trans_slice[kphase_coord], + tStS_slice, + ) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + else: + for kphase_idx in cutlass.range( + num_kphases, unroll_full=True + ): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + qk_tiled_mma, + tStS_slice, + tSrQ_slice[kphase_coord], + tSrK_trans_slice[kphase_coord], + tStS_slice, + ) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + k_handle.release() + s_handle.commit() + + # dSKTi + ds_handle = ds_mma_consumer.wait_and_advance() + dsk_whether_acc = dsk_tiled_mma.get(tcgen05.Field.ACCUMULATE) + for iter in cutlass.range(self.iterations_dsk, unroll=1): + kt_handle = load_kt_consumer.wait_and_advance() + dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, dsk_whether_acc) + tdQtdQ_slice = tdQtdQ_staged[None, None, None, iter] + tdStdS_slice = tdPtdP[None, None, None, ds_handle.index] + tdS = cute.make_tensor( + tdStdS_slice.iterator, ds_tmem_layout_staged.outer + ) + tdQrdS = dsk_thr_mma.make_fragment_A(tdS) + tdQrdS_slice = cute.make_tensor( + cute.recast_ptr(tdStdS_slice.iterator, dtype=self.q_dtype), + tdQrdS.layout, + ) + + tdQrKT_slice = tdQrKT[None, None, None, kt_handle.index] + num_kphases = cute.size(tdQrKT_slice, mode=[2]) + if cutlass.const_expr(num_kphases % num_innerloop == 0): + num_outer_iter = num_kphases // num_innerloop + for outer_iter in cutlass.range(num_outer_iter, unroll=1): + for kphase_idx in cutlass.range( + num_innerloop, unroll_full=True + ): + kphase_coord = ( + None, + None, + outer_iter * num_innerloop + kphase_idx, + ) + cute.gemm( + dsk_tiled_mma, + tdQtdQ_slice, + tdQrdS_slice[kphase_coord], + tdQrKT_slice[kphase_coord], + tdQtdQ_slice, + ) + dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + else: + for kphase_idx in cutlass.range( + num_kphases, unroll_full=True + ): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + dsk_tiled_mma, + tdQtdQ_slice, + tdQrdS_slice[kphase_coord], + tdQrKT_slice[kphase_coord], + tdQtdQ_slice, + ) + dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + kt_handle.release() + ds_handle.release() + + # dOVi + dp_handle = mma_dp_producer.acquire_and_advance() + tdPtdP_slice = tdPtdP[None, None, None, dp_handle.index] + dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for iter in cutlass.range(self.iterations_dov, unroll=1): + tdPrdO_slice = tdPrdO[None, None, None, iter] + v_handle = load_v_consumer.wait_and_advance() + tdPrV_trans_slice = tdPrV[None, None, None, v_handle.index] + num_kphases = cute.size(tdPrdO_slice, mode=[2]) + if cutlass.const_expr(num_kphases % num_innerloop == 0): + num_outer_iter = num_kphases // num_innerloop + for outer_iter in cutlass.range(num_outer_iter, unroll=1): + for kphase_idx in cutlass.range( + num_innerloop, unroll_full=True + ): + kphase_coord = ( + None, + None, + outer_iter * num_innerloop + kphase_idx, + ) + cute.gemm( + dov_tiled_mma, + tdPtdP_slice, + tdPrdO_slice[kphase_coord], + tdPrV_trans_slice[kphase_coord], + tdPtdP_slice, + ) + dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + else: + for kphase_idx in cutlass.range( + num_kphases, unroll_full=True + ): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + dov_tiled_mma, + tdPtdP_slice, + tdPrdO_slice[kphase_coord], + tdPrV_trans_slice[kphase_coord], + tdPtdP_slice, + ) + dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + v_handle.release() + dp_handle.commit() + + # QKend + s_handle = mma_s_producer.acquire_and_advance() + tStS_slice = tStS[None, None, None, s_handle.index] + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for iter in cutlass.range(self.iterations_qk, unroll=1): + tSrQ_slice = tSrQ[None, None, None, iter] + k_handle = load_k_consumer.wait_and_advance() + + tSrK_trans_slice = tSrK[None, None, None, k_handle.index] + + num_kphases = cute.size(tSrQ_slice, mode=[2]) + if cutlass.const_expr(num_kphases % num_innerloop == 0): + num_outer_iter = num_kphases // num_innerloop + for outer_iter in cutlass.range(num_outer_iter, unroll=1): + for kphase_idx in cutlass.range( + num_innerloop, unroll_full=True + ): + kphase_coord = ( + None, + None, + outer_iter * num_innerloop + kphase_idx, + ) + cute.gemm( + qk_tiled_mma, + tStS_slice, + tSrQ_slice[kphase_coord], + tSrK_trans_slice[kphase_coord], + tStS_slice, + ) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + else: + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + qk_tiled_mma, + tStS_slice, + tSrQ_slice[kphase_coord], + tSrK_trans_slice[kphase_coord], + tStS_slice, + ) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + k_handle.release() + load_q_releaser.release() + load_q_releaser.advance() + s_handle.commit() + + # dSKTend - 1 + ds_handle = ds_mma_consumer.wait_and_advance() + dsk_whether_acc = dsk_tiled_mma.get(tcgen05.Field.ACCUMULATE) + for iter in cutlass.range(self.iterations_dsk, unroll=1): + kt_handle = load_kt_consumer.wait_and_advance() + dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, dsk_whether_acc) + tdQtdQ_slice = tdQtdQ_staged[None, None, None, iter] + tdStdS_slice = tdPtdP[None, None, None, ds_handle.index] + tdS = cute.make_tensor( + tdStdS_slice.iterator, ds_tmem_layout_staged.outer + ) + tdQrdS = dsk_thr_mma.make_fragment_A(tdS) + tdQrdS_slice = cute.make_tensor( + cute.recast_ptr(tdStdS_slice.iterator, dtype=self.q_dtype), + tdQrdS.layout, + ) + + tdQrKT_slice = tdQrKT[None, None, None, kt_handle.index] + num_kphases = cute.size(tdQrKT_slice, mode=[2]) + if cutlass.const_expr(num_kphases % num_innerloop == 0): + num_outer_iter = num_kphases // num_innerloop + for outer_iter in cutlass.range(num_outer_iter, unroll=1): + for kphase_idx in cutlass.range( + num_innerloop, unroll_full=True + ): + kphase_coord = ( + None, + None, + outer_iter * num_innerloop + kphase_idx, + ) + cute.gemm( + dsk_tiled_mma, + tdQtdQ_slice, + tdQrdS_slice[kphase_coord], + tdQrKT_slice[kphase_coord], + tdQtdQ_slice, + ) + dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + else: + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + dsk_tiled_mma, + tdQtdQ_slice, + tdQrdS_slice[kphase_coord], + tdQrKT_slice[kphase_coord], + tdQtdQ_slice, + ) + dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + kt_handle.release() + ds_handle.release() + + # dOVend + dp_handle = mma_dp_producer.acquire_and_advance() + tdPtdP_slice = tdPtdP[None, None, None, dp_handle.index] + dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for iter in cutlass.range(self.iterations_dov, unroll=1): + tdPrdO_slice = tdPrdO[None, None, None, iter] + v_handle = load_v_consumer.wait_and_advance() + tdPrV_trans_slice = tdPrV[None, None, None, v_handle.index] + num_kphases = cute.size(tdPrdO_slice, mode=[2]) + if cutlass.const_expr(num_kphases % num_innerloop == 0): + num_outer_iter = num_kphases // num_innerloop + for outer_iter in cutlass.range(num_outer_iter, unroll=1): + for kphase_idx in cutlass.range( + num_innerloop, unroll_full=True + ): + kphase_coord = ( + None, + None, + outer_iter * num_innerloop + kphase_idx, + ) + cute.gemm( + dov_tiled_mma, + tdPtdP_slice, + tdPrdO_slice[kphase_coord], + tdPrV_trans_slice[kphase_coord], + tdPtdP_slice, + ) + dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + else: + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + dov_tiled_mma, + tdPtdP_slice, + tdPrdO_slice[kphase_coord], + tdPrV_trans_slice[kphase_coord], + tdPtdP_slice, + ) + dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + v_handle.release() + load_do_releaser.release() + load_do_releaser.advance() + dp_handle.commit() + # dSKTend + ds_handle = ds_mma_consumer.wait_and_advance() + dsk_whether_acc = dsk_tiled_mma.get(tcgen05.Field.ACCUMULATE) + for iter in cutlass.range(self.iterations_dsk, unroll=1): + kt_handle = load_kt_consumer.wait_and_advance() + dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, dsk_whether_acc) + tdQtdQ_slice = tdQtdQ_staged[None, None, None, iter] + tdStdS_slice = tdPtdP[None, None, None, ds_handle.index] + tdS = cute.make_tensor( + tdStdS_slice.iterator, ds_tmem_layout_staged.outer + ) + tdQrdS = dsk_thr_mma.make_fragment_A(tdS) + tdQrdS_slice = cute.make_tensor( + cute.recast_ptr(tdStdS_slice.iterator, dtype=self.q_dtype), + tdQrdS.layout, + ) + + tdQrKT_slice = tdQrKT[None, None, None, kt_handle.index] + num_kphases = cute.size(tdQrKT_slice, mode=[2]) + if cutlass.const_expr(num_kphases % num_innerloop == 0): + num_outer_iter = num_kphases // num_innerloop + for outer_iter in cutlass.range(num_outer_iter, unroll=1): + for kphase_idx in cutlass.range( + num_innerloop, unroll_full=True + ): + kphase_coord = ( + None, + None, + outer_iter * num_innerloop + kphase_idx, + ) + cute.gemm( + dsk_tiled_mma, + tdQtdQ_slice, + tdQrdS_slice[kphase_coord], + tdQrKT_slice[kphase_coord], + tdQtdQ_slice, + ) + dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + else: + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + dsk_tiled_mma, + tdQtdQ_slice, + tdQrdS_slice[kphase_coord], + tdQrKT_slice[kphase_coord], + tdQtdQ_slice, + ) + dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + kt_handle.release() + ds_handle.release() + else: + # QK0 + s_handle = mma_s_producer.acquire_and_advance() + tStS_slice = tStS[None, None, None, s_handle.index] + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + + for iter in cutlass.range(self.iterations_qk, unroll=1): + load_q_consumer.wait_and_advance() + tSrQ_slice = tSrQ[None, None, None, iter] + k_handle = load_k_consumer.wait_and_advance() + tSrK_trans_slice = tSrK[None, None, None, k_handle.index] + num_kphases = cute.size(tSrQ_slice, mode=[2]) + if cutlass.const_expr(num_kphases % num_innerloop == 0): + num_outer_iter = num_kphases // num_innerloop + for outer_iter in cutlass.range(num_outer_iter, unroll=1): + for kphase_idx in cutlass.range( + num_innerloop, unroll_full=True + ): + kphase_coord = ( + None, + None, + outer_iter * num_innerloop + kphase_idx, + ) + cute.gemm( + qk_tiled_mma, + tStS_slice, + tSrQ_slice[kphase_coord], + tSrK_trans_slice[kphase_coord], + tStS_slice, + ) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + else: + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + qk_tiled_mma, + tStS_slice, + tSrQ_slice[kphase_coord], + tSrK_trans_slice[kphase_coord], + tStS_slice, + ) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + k_handle.release() + load_q_releaser.release() + load_q_releaser.advance() + s_handle.commit() + + # dOV0 + dp_handle = mma_dp_producer.acquire_and_advance() + tdPtdP_slice = tdPtdP[None, None, None, dp_handle.index] + dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for iter in cutlass.range(self.iterations_dov, unroll=1): + load_do_consumer.wait_and_advance() + tdPrdO_slice = tdPrdO[None, None, None, iter] + v_handle = load_v_consumer.wait_and_advance() + tdPrV_trans_slice = tdPrV[None, None, None, v_handle.index] + num_kphases = cute.size(tdPrdO_slice, mode=[2]) + if cutlass.const_expr(num_kphases % num_innerloop == 0): + num_outer_iter = num_kphases // num_innerloop + for outer_iter in cutlass.range(num_outer_iter, unroll=1): + for kphase_idx in cutlass.range( + num_innerloop, unroll_full=True + ): + kphase_coord = ( + None, + None, + outer_iter * num_innerloop + kphase_idx, + ) + cute.gemm( + dov_tiled_mma, + tdPtdP_slice, + tdPrdO_slice[kphase_coord], + tdPrV_trans_slice[kphase_coord], + tdPtdP_slice, + ) + dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + else: + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + dov_tiled_mma, + tdPtdP_slice, + tdPrdO_slice[kphase_coord], + tdPrV_trans_slice[kphase_coord], + tdPtdP_slice, + ) + dov_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + v_handle.release() + load_do_releaser.release() + load_do_releaser.advance() + dp_handle.commit() + + # dSKT0 + ds_handle = ds_mma_consumer.wait_and_advance() + dsk_whether_acc = dsk_tiled_mma.get(tcgen05.Field.ACCUMULATE) + for iter in cutlass.range(self.iterations_dsk, unroll=1): + kt_handle = load_kt_consumer.wait_and_advance() + dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, dsk_whether_acc) + tdQtdQ_slice = tdQtdQ_staged[None, None, None, iter] + tdStdS_slice = tdPtdP[None, None, None, ds_handle.index] + tdS = cute.make_tensor( + tdStdS_slice.iterator, ds_tmem_layout_staged.outer + ) + tdQrdS = dsk_thr_mma.make_fragment_A(tdS) + tdQrdS_slice = cute.make_tensor( + cute.recast_ptr(tdStdS_slice.iterator, dtype=self.q_dtype), + tdQrdS.layout, + ) + + tdQrKT_slice = tdQrKT[None, None, None, kt_handle.index] + num_kphases = cute.size(tdQrKT_slice, mode=[2]) + if cutlass.const_expr(num_kphases % num_innerloop == 0): + num_outer_iter = num_kphases // num_innerloop + for outer_iter in cutlass.range(num_outer_iter, unroll=1): + for kphase_idx in cutlass.range( + num_innerloop, unroll_full=True + ): + kphase_coord = ( + None, + None, + outer_iter * num_innerloop + kphase_idx, + ) + cute.gemm( + dsk_tiled_mma, + tdQtdQ_slice, + tdQrdS_slice[kphase_coord], + tdQrKT_slice[kphase_coord], + tdQtdQ_slice, + ) + dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + else: + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + dsk_tiled_mma, + tdQtdQ_slice, + tdQrdS_slice[kphase_coord], + tdQrKT_slice[kphase_coord], + tdQtdQ_slice, + ) + dsk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + kt_handle.release() + ds_handle.release() + dq_handle.commit() + work_tile = tile_sched.advance_to_next_work() + # End of persistent scheduler loop + mma_s_producer.tail() + mma_dp_producer.tail() + mma_dq_producer.tail() + + # Softmax and dSoftmax warp + if warp_idx >= self.compute_warp_ids[0] and warp_idx <= self.compute_warp_ids[-1]: + # increase register after decreasing + cute.arch.warpgroup_reg_alloc(self.num_regs_compute) + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + batch_coord = curr_block_coord[2][1] + seqlen_q = mQ_qdl.shape[0] + seqlen_k = mK_kdl.shape[0] + cuseqlen_q = Int32(0) + is_valid_q = True + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + is_valid_q = FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.qk_mma_tiler[0], + mma_block_coord[0], + seqlen_q, + ) + if cutlass.const_expr(cum_seqlen_k is not None): + cuseqlen_k = cum_seqlen_k[batch_coord] + seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k + start_count, trip_count = FusedMask.get_trip_start_count_via_block_info( + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + self.is_causal, + self.is_local, + window_size_left, + window_size_right, + ) + is_valid_k = trip_count > 0 + has_work = is_valid_q and is_valid_k + + if has_work: + end_count = start_count + trip_count + if cutlass.const_expr(self.use_semantic_trip_range): + n_block_min_causal_local_mask, n_block_min_before_local_mask = ( + FusedMask.get_trip_mask_bounds_via_block_info( + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + self.is_causal, + self.is_local, + window_size_left, + window_size_right, + ) + ) + + cS_base = cute.make_identity_tensor( + (self.qk_mma_tiler[0], self.qk_mma_tiler[1]) + ) + cS = cute.domain_offset((mma_block_coord[0] * self.qk_mma_tiler[0], 0), cS_base) + tScS = qk_thr_mma.partition_C(cS) + + cdP_base = cute.make_identity_tensor( + (self.dov_mma_tiler[0], self.dov_mma_tiler[1]) + ) + cdP = cute.domain_offset( + (mma_block_coord[0] * self.dov_mma_tiler[0], 0), cdP_base + ) + tdPcdP = dov_thr_mma.partition_C(cdP) + + lse_handle = load_lse_consumer.wait_and_advance() + sum_odo_handle = load_sum_odo_consumer.wait_and_advance() + for step in cutlass.range(start_count, end_count, 1, unroll=1): + cS_iter = cute.domain_offset((0, step * self.qk_mma_tiler[1]), cS) + tScS_iter = qk_thr_mma.partition_C(cS_iter) + + cdP_iter = cute.domain_offset((0, step * self.dov_mma_tiler[1]), cdP) + + tdPcdP_iter = dov_thr_mma.partition_C(cdP_iter) + + # Si, dPi -> dSi + if cutlass.const_expr(self.use_semantic_trip_range): + need_apply_mask = ( + step >= n_block_min_causal_local_mask + or step < n_block_min_before_local_mask + ) + else: + need_apply_mask = step == end_count - 1 + mma_s_consumer, mma_dp_consumer, ds_mma_producer = self.compute_step( + (need_apply_mask, window_size_left, window_size_right), + ( + seqlen_q, + seqlen_k, + scale_softmax, + batch_coord, + curr_block_coord[0], + varlen, + ), + (tStS, tScS_iter, tdPtdP, tdPcdP_iter, sLSE, sSum_OdO), + ( + mma_s_consumer, + mma_dp_consumer, + ds_mma_producer, + lse_handle, + sum_odo_handle, + ), + step, + ) + lse_handle.release() + sum_odo_handle.release() + + work_tile = tile_sched.advance_to_next_work() + ds_mma_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # Epilogue + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx >= self.epilogue_warp_ids[0] and warp_idx <= self.epilogue_warp_ids[-1]: + cute.arch.warpgroup_reg_dealloc(self.num_regs_epilogue) + + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + batch_coord = curr_block_coord[2][1] + # cute.printf("batch_coord={}", batch_coord) + seqlen_q = mQ_qdl.shape[0] + seqlen_k = mK_kdl.shape[0] + cuseqlen_q = Int32(0) + is_valid_q = True + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + is_valid_q = FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.qk_mma_tiler[0], + mma_block_coord[0], + seqlen_q, + ) + if cutlass.const_expr(cum_seqlen_k is not None): + cuseqlen_k = cum_seqlen_k[batch_coord] + seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k + seqlen_kv_loop_start, seqlen_kv_loop_steps = ( + FusedMask.get_trip_start_count_via_block_info( + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + self.is_causal, + self.is_local, + window_size_left, + window_size_right, + ) + ) + is_valid_k = seqlen_kv_loop_steps > 0 + has_work = is_valid_q and is_valid_k + + mdQ_qdl_eff = mdQ_qdl + if cutlass.const_expr(cum_seqlen_q is not None): + block_offset_dQ = (cuseqlen_q,) + (None,) * 2 + mdQ_qdl_eff = cute.domain_offset(block_offset_dQ, mdQ_qdl) + + # (bM, bN, loopM, loopN, loopL) + gdQ_qdl = cute.flat_divide( + mdQ_qdl_eff, cute.select(self.dsk_block_tiler, mode=[0, 1]) + ) + cdQ_qdl = cute.flat_divide( + cute.make_identity_tensor(mdQ_qdl_eff.shape), + cute.select(self.dsk_block_tiler, mode=[0, 1]), + ) + + gdQ_staged = gdQ_qdl[None, None, curr_block_coord[0], None, curr_block_coord[2]] + cdQ_staged = cdQ_qdl[None, None, curr_block_coord[0], None, curr_block_coord[2]] + gdQ_tma_staged = gdQ_staged + + if cutlass.const_expr(not varlen): + gdQ_tma_qdl = cute.flat_divide( + mdQ_tma, cute.select(self.dsk_block_tiler, mode=[0, 1]) + ) + gdQ_tma_staged = gdQ_tma_qdl[ + None, None, curr_block_coord[0], None, curr_block_coord[2] + ] + + if has_work: + # dQ TMEM to GMEM + mma_dq_consumer = self.dQ_epilogue( + seqlen_q, + (mma_dq_consumer, gdQ_staged, cdQ_staged, tdQtdQ_staged), + self.epi_tile, + (tma_atom_dQ, gdQ_tma_staged, s_epi_dQ, varlen), + ) + else: + self.dQ_epilogue_write_zero( + seqlen_q, + gdQ_staged, + cdQ_staged, + ) + + work_tile = tile_sched.advance_to_next_work() + # NOTE: tmem.free() moved to kernel end to enable cluster-wide sync + + # /////////////////////////////////////////////////////////////////////////////// + # Scheduler Warp (only for CLC dynamic scheduler) + # /////////////////////////////////////////////////////////////////////////////// + if cutlass.const_expr(self.use_clc_scheduler): + is_first_cta_in_cluster = cta_rank_in_cluster == 0 + + if warp_idx == self.sched_warp_id and is_first_cta_in_cluster: + cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + while work_tile.is_valid_tile: + tile_sched.prefetch_next_work() + work_tile = tile_sched.advance_to_next_work() + tile_sched.producer_tail() + + # /////////////////////////////////////////////////////////////////////////////// + # Empty warps reg dealloc + # /////////////////////////////////////////////////////////////////////////////// + if cutlass.const_expr(self.use_clc_scheduler): + if warp_idx > self.load_warp_id: + if not (warp_idx == self.sched_warp_id and is_first_cta_in_cluster): + cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + else: + if warp_idx > self.load_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + + # /////////////////////////////////////////////////////////////////////////////// + # Cooperative TMEM Deallocation (2CTA) + # /////////////////////////////////////////////////////////////////////////////// + # All warps (including scheduler) have finished by this point. + # Cluster-wide sync ensures both CTAs reach here before dealloc. + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + return + + @cute.jit + def compute_step( + self, + mask_args: Tuple, + value_args: Tuple, + tensor_args: Tuple, + pipeline_args: Tuple, + step: Int32, + ) -> Tuple[Float32, Float32, pipeline.PipelineConsumer, pipeline.PipelineProducer]: + need_apply_mask, window_size_left, window_size_right = mask_args + seqlen_q, seqlen_k, scale_softmax, batch_coord, block_m_idx, varlen = value_args + tStS, tScS, tdPtdP, tdPcdP, sLSE, sSum_OdO = tensor_args + mma_s_consumer, mma_dp_consumer, ds_mma_producer, lse_handle, sum_odo_handle = pipeline_args + + bidx = block_m_idx + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.compute_warp_ids)) + s_handle = mma_s_consumer.wait_and_advance() + tStS_slice = tStS[(None, None), 0, 0, s_handle.index] + tScS_slice = tScS[(None, None), 0, 0] + tmem_load_atom = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(16)), self.acc_dtype + ) + tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS_slice) + thr_load = tmem_tiled_load.get_slice(thread_idx) + tTMEM_LOADtS = thr_load.partition_S(tStS_slice) + tTMEM_LOADcS = thr_load.partition_D(tScS_slice) + tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.acc_dtype) + cute.copy(tmem_tiled_load, tTMEM_LOADtS, tTMEM_LOADrS) + cute.arch.fence_view_async_tmem_load() + s_handle.release() + if need_apply_mask: + FusedMask.apply_mask_via_causal_local( + tTMEM_LOADrS, + tTMEM_LOADcS, + seqlen_q, + seqlen_k, + self.use_semantic_trip_range, + self.is_causal, + self.is_local, + window_size_left, + window_size_right, + ) + + log2_e = cutlass.Float32(math.log2(math.e)) + softmax_scale_log2_e = scale_softmax * log2_e + tTMEM_STORErP = cute.make_rmem_tensor(tTMEM_LOADrS.shape, self.q_dtype) + for k in cutlass.range(0, cute.size(tTMEM_LOADrS), 2, unroll_full=True): + lse = ( + -sLSE[ + cute.get(tTMEM_LOADcS[k], mode=[0]) - bidx * self.cta_tiler[0], + lse_handle.index, + ], + -sLSE[ + cute.get(tTMEM_LOADcS[k + 1], mode=[0]) - bidx * self.cta_tiler[0], + lse_handle.index, + ], + ) + tTMEM_LOADrS[k], tTMEM_LOADrS[k + 1] = cute.arch.fma_packed_f32x2( + (tTMEM_LOADrS[k], tTMEM_LOADrS[k + 1]), + (softmax_scale_log2_e, softmax_scale_log2_e), + lse, + ) + tTMEM_LOADrS[k] = cute.math.exp2(tTMEM_LOADrS[k], fastmath=True) + tTMEM_LOADrS[k + 1] = cute.math.exp2(tTMEM_LOADrS[k + 1], fastmath=True) + + dp_handle = mma_dp_consumer.wait_and_advance() + tdPtdP_slice = tdPtdP[(None, None), 0, 0, dp_handle.index] + tdPcdP_slice = tdPcdP[(None, None), 0, 0] + thr_load = tmem_tiled_load.get_slice(thread_idx) + tTMEM_LOADtdP = thr_load.partition_S(tdPtdP_slice) + tTMEM_LOADcdP = thr_load.partition_D(tdPcdP_slice) + tTMEM_LOADrdP = cute.make_rmem_tensor(tTMEM_LOADcdP.shape, self.acc_dtype) + cute.copy(tmem_tiled_load, tTMEM_LOADtdP, tTMEM_LOADrdP) + cute.arch.fence_view_async_tmem_load() + dp_handle.release() + tTMEM_STORErdP = cute.make_rmem_tensor(tTMEM_LOADrdP.shape, self.q_dtype) + + for k in cutlass.range(0, cute.size(tTMEM_LOADrdP), 2, unroll_full=True): + dpsum_0 = -sSum_OdO[ + cute.get(tTMEM_LOADcdP[k], mode=[0]) - bidx * self.cta_tiler[0], + sum_odo_handle.index, + ] + dpsum_1 = -sSum_OdO[ + cute.get(tTMEM_LOADcdP[k + 1], mode=[0]) - bidx * self.cta_tiler[0], + sum_odo_handle.index, + ] + if cutlass.const_expr(varlen): + if not cute.elem_less(cute.get(tTMEM_LOADcdP[k], mode=[0]), seqlen_q): + dpsum_0 = 0.0 + if not cute.elem_less(cute.get(tTMEM_LOADcdP[k + 1], mode=[0]), seqlen_q): + dpsum_1 = 0.0 + tTMEM_LOADrdP[k], tTMEM_LOADrdP[k + 1] = cute.arch.add_packed_f32x2( + (tTMEM_LOADrdP[k], tTMEM_LOADrdP[k + 1]), + (dpsum_0, dpsum_1), + ) + tTMEM_LOADrdP[k], tTMEM_LOADrdP[k + 1] = cute.arch.mul_packed_f32x2( + (tTMEM_LOADrdP[k], tTMEM_LOADrdP[k + 1]), (tTMEM_LOADrS[k], tTMEM_LOADrS[k + 1]) + ) + tTMEM_LOADrdP[k], tTMEM_LOADrdP[k + 1] = cute.arch.mul_packed_f32x2( + (tTMEM_LOADrdP[k], tTMEM_LOADrdP[k + 1]), (scale_softmax, scale_softmax) + ) + dp_vec = tTMEM_LOADrdP.load() + tTMEM_STORErdP.store(dp_vec.to(self.q_dtype)) + + ds_handle = ds_mma_producer.acquire_and_advance() + tmem_store_atom = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(32)), self.acc_dtype + ) + tilePlikeFP32 = tdPtdP_slice.shape[1] // Float32.width * self.q_dtype.width + tdPtdP_dS_layout = cute.composition( + tdPtdP_slice.layout, cute.make_layout((tdPtdP_slice.shape[0], tilePlikeFP32)) + ) + tdPtdP_dS = cute.make_tensor(tdPtdP_slice.iterator, tdPtdP_dS_layout) + tdPcdP_dS_layout = cute.composition( + tdPcdP_slice.layout, cute.make_layout((tdPcdP_slice.shape[0], tilePlikeFP32)) + ) + tdPcdP_dS = cute.make_tensor(tdPcdP_slice.iterator, tdPcdP_dS_layout) + tmem_tiled_store = tcgen05.make_tmem_copy(tmem_store_atom, tdPtdP_dS) + + thr_store = tmem_tiled_store.get_slice(thread_idx) + tTMEM_STOREtdS = thr_store.partition_D(tdPtdP_dS) + tTMEM_STOREcdP = thr_store.partition_S(tdPcdP_dS) + tTMEM_STORErdS_ = cute.make_tensor( + cute.recast_ptr(tTMEM_STORErdP.iterator, dtype=self.acc_dtype), + tTMEM_STOREcdP.shape, + ) + cute.copy(tmem_tiled_store, tTMEM_STORErdS_, tTMEM_STOREtdS) + cute.arch.fence_view_async_tmem_store() + ds_handle.commit() + return mma_s_consumer, mma_dp_consumer, ds_mma_producer + + @cute.jit + def dQ_epilogue( + self, + seqlen_q: int, + dq_args: Tuple, + epi_tile: cute.Tile, + tma_args: Tuple, + ) -> Tuple[pipeline.PipelineConsumer, pipeline.PipelineProducer]: + (mma_dq_consumer, gdQ_staged, cdQ_staged, tdQtdQ_staged) = dq_args + tma_atom_dQ, gdQ_tma_staged, s_epi_dQ, varlen = tma_args + dq_handle = mma_dq_consumer.wait_and_advance() + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, bidz = cute.arch.block_idx() + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + cute.arch.fence_view_async_shared() + + epi_cols_dQ = math.gcd(128 // (self.dq_dtype.width // 8), epi_tile[1]) + num_epi_stages_dQ = epi_tile[1] // epi_cols_dQ + epi_tile_dQ = (epi_tile[0], epi_cols_dQ) + leader_warp = (cute.arch.make_warp_uniform(cute.arch.warp_idx()) % 4) == 0 + + for iter in cutlass.range(self.iterations_dsk): + gdQ = gdQ_staged[None, None, iter] + cdQ = cdQ_staged[None, None, iter] + tdQtdQ = tdQtdQ_staged[(None, None), 0, 0, iter] + tdQtdQ_epi = cute.zipped_divide(tdQtdQ, epi_tile) + cdQ_epi = cute.zipped_divide(cdQ, epi_tile) + gdQ_epi = cute.zipped_divide(gdQ, epi_tile) + cdQ_local = cute.make_identity_tensor(epi_tile) + cdQ_local_epi = cute.zipped_divide(cdQ_local, epi_tile) + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.epilogue_warp_ids)) + tmem_copy_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.acc_dtype + ) + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tdQtdQ_epi) + thr_tmem_load = tiled_tmem_load.get_slice(thread_idx) + tTMEM_LOADtdQ = thr_tmem_load.partition_S(tdQtdQ_epi) + tTMEM_LOADgdQ = thr_tmem_load.partition_D(gdQ_epi) + tTMEM_LOADcdQ = thr_tmem_load.partition_D(cdQ_epi) + tTMEM_LOADcdQ_local = thr_tmem_load.partition_D(cdQ_local_epi) + + if cutlass.const_expr(not varlen): + gdQ_tma = gdQ_tma_staged[None, None, iter] + gdQ_tma_epi = cute.local_tile(gdQ_tma, epi_tile_dQ, (0, None)) + sdQ_stage = s_epi_dQ[None, None, 0] + + for stage_k in cutlass.range_constexpr(num_epi_stages_dQ): + for i in cutlass.range(cute.size(tTMEM_LOADtdQ, mode=[1]), unroll_full=True): + tTMEM_LOADtdQ_i = tTMEM_LOADtdQ[None, i, 0] + tTMEM_LOADcdQ_i_local = tTMEM_LOADcdQ_local[None, i, 0] + tTMrdQ = cute.make_rmem_tensor(tTMEM_LOADcdQ_i_local.shape, self.acc_dtype) + cute.copy(tiled_tmem_load, tTMEM_LOADtdQ_i, tTMrdQ) + tSMrdQ = cute.make_rmem_tensor(tTMrdQ.shape, self.q_dtype) + dq_vec = tTMrdQ.load() + tSMrdQ.store(dq_vec.to(self.q_dtype)) + for j in cutlass.range_constexpr(cute.size(tTMEM_LOADcdQ_i_local)): + c = tTMEM_LOADcdQ_i_local[j] + m_pos = c[0] + n_pos = c[1] + if n_pos // epi_cols_dQ == stage_k: + s_epi_dQ[m_pos, n_pos % epi_cols_dQ, 0] = tSMrdQ[j] + + cute.arch.fence_view_async_shared() + cute.arch.barrier(barrier_id=3, number_of_threads=128) + + if leader_warp: + gdQ_stage = gdQ_tma_epi[None, None, stage_k] + td_sdQ, td_gdQ = cpasync.tma_partition( + tma_atom_dQ, + 0, + cute.make_layout(1), + cute.group_modes(sdQ_stage, 0, 2), + cute.group_modes(gdQ_stage, 0, 2), + ) + cute.copy(tma_atom_dQ, td_sdQ, td_gdQ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + # Non-issuing threads must not rotate the SMEM buffer + # until the leader warp's TMA read has completed. + cute.arch.barrier(barrier_id=3, number_of_threads=128) + else: + for i in cutlass.range(cute.size(tTMEM_LOADtdQ, mode=[1]), unroll_full=True): + tTMEM_LOADtdQ_i = tTMEM_LOADtdQ[None, i, 0] + tTMEM_LOADgdQ_i = tTMEM_LOADgdQ[None, i, 0] + tTMEM_LOADcdQ_i = tTMEM_LOADcdQ[None, i, 0] + tTMrdQ = cute.make_rmem_tensor(tTMEM_LOADcdQ[None, 0, i].shape, self.acc_dtype) + cute.copy(tiled_tmem_load, tTMEM_LOADtdQ_i, tTMrdQ) + tSMrdQ = cute.make_rmem_tensor(tTMrdQ.shape, self.q_dtype) + dq_vec = tTMrdQ.load() + tSMrdQ.store(dq_vec.to(self.q_dtype)) + if cute.elem_less(tTMEM_LOADcdQ_i[0][0], seqlen_q): + cute.autovec_copy(tSMrdQ, tTMEM_LOADgdQ_i) + dq_handle.release() + return mma_dq_consumer + + @cute.jit + def dQ_epilogue_write_zero( + self, + seqlen_q, + gdQ_staged, + cdQ_staged, + ): + num_epi_threads = self.threads_per_warp * len(self.epilogue_warp_ids) + tidx = cute.arch.thread_idx()[0] % num_epi_threads + + tiled_copy_r2g = fa_copy_utils.tiled_copy_2d( + self.dq_dtype, cute.size(gdQ_staged.shape[1]), num_epi_threads + ) + + thr_copy_r2g = tiled_copy_r2g.get_slice(tidx) + tdQgdQ_staged = thr_copy_r2g.partition_D(gdQ_staged) + tdQcdQ_staged = thr_copy_r2g.partition_D(cdQ_staged) + + tdQrdQ = cute.make_rmem_tensor_like(tdQgdQ_staged[None, 0, None, 0]) + tdQrdQ.fill(self.dq_dtype(0.0)) + + for iter in cutlass.range(self.iterations_dsk, unroll_full=True): + tdQgdQ = tdQgdQ_staged[None, None, None, iter] + tdQcdQ = tdQcdQ_staged[None, None, None, iter] + for m in cutlass.range(cute.size(tdQgdQ.shape[1]), unroll_full=True): + if cute.elem_less(tdQcdQ[0, m, 0][0], seqlen_q): + cute.copy(tiled_copy_r2g, tdQrdQ, tdQgdQ[None, m, None]) diff --git a/fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_forward.py b/fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_forward.py new file mode 100644 index 000000000..b21fc16c7 --- /dev/null +++ b/fastvideo-kernel/fa4/sm100_hd256_2cta_fmha_forward.py @@ -0,0 +1,1889 @@ +# Copyright (c) 2025, Siyu Wang, Shengbin Di, Yuxi Chi, Johnsonms, Linfeng Zheng, Haoyan Huang, Lanbo Li, Yun Zhong, Man Yuan, Minmin Sun, Yong Li, Wei Lin. + +import math +from typing import Tuple, Optional + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.utils.blackwell_helpers as sm100_utils +from cutlass.cute.typing import Int32, Int64, Float32 + +from cutlass.utils import ClcDynamicPersistentTileScheduler +from flash_attn.cute.tile_scheduler import ( + ClcState, + compute_sm100_fmha_grid as compute_grid, + compute_sm100_fmha_grid_clc as compute_grid_clc, + make_sm100_thread_cooperative_group as make_thread_cooperative_group, + Sm100FmhaStaticTileScheduler as FmhaStaticTileScheduler, + Sm100FmhaStaticTileSchedulerParams as FmhaStaticTileSchedulerParams, + Sm100FmhaClcDynamicTileScheduler as FmhaClcDynamicTileScheduler, + Sm100FmhaClcDynamicTileSchedulerParams as FmhaClcDynamicTileSchedulerParams, +) +from flash_attn.cute.mask import ( + Sm100FusedMask as FusedMask, +) +from flash_attn.cute.tile_scheduler import SM100_TMEM_CAPACITY_COLUMNS +from flash_attn.cute.flash_fwd_sm100 import DescaleTensors, _TUNING_CONFIG +from flash_attn.cute.utils import ex2_emulation_2 +from flash_attn.cute.utils import AuxData + + +class BlackwellFusedMultiHeadAttentionForward: + def __init__( + self, + head_dim: int, + head_dim_v: Optional[int] = None, + qhead_per_kvhead: int = 1, + is_causal: bool = False, + is_local: bool = False, + is_split_kv: bool = False, + pack_gqa: bool = False, + q_subtile_factor: int = 1, + m_block_size: int = 128, + n_block_size: int = 128, + q_stage: int = 2, + is_persistent: bool = True, + score_mod=None, + mask_mod=None, + has_aux_tensors: bool = False, + paged_kv_non_tma: bool = False, + is_varlen_q: bool = False, + use_2cta_instrs: bool = False, + use_clc_scheduler: bool = False, + ): + head_dim_v = head_dim if head_dim_v is None else head_dim_v + assert head_dim == 256 and head_dim_v == 256, ( + "SM100 dedicated kernel only supports (head_dim, head_dim_v) = (256, 256)" + ) + assert score_mod is None, "SM100 forward with head_dim=256 does not support score_mod" + assert mask_mod is None, "SM100 forward with head_dim=256 does not support mask_mod" + assert not has_aux_tensors, "SM100 forward with head_dim=256 does not support aux tensors" + assert not paged_kv_non_tma, ( + "SM100 hd256 2CTA supports TMA paged KV only (page_size must equal tile_n=128)" + ) + assert not pack_gqa, "SM100 forward with head_dim=256 does not support pack_gqa" + assert not is_split_kv, "SM100 forward with head_dim=256 does not support SplitKV" + assert q_subtile_factor == 1, ( + "SM100 forward with head_dim=256 does not support q_subtile_factor" + ) + assert m_block_size == 128 and n_block_size == 128, ( + "SM100 dedicated kernel only supports tile_m=128 and tile_n=128" + ) + # q_stage / persistence / scheduler knobs are accepted for interface parity, + # but this dedicated kernel uses fixed internal settings. + + qk_acc_dtype = cutlass.Float32 + pv_acc_dtype = cutlass.Float32 + mma_tiler = (128, 128, head_dim) + self.qk_acc_dtype = qk_acc_dtype + self.pv_acc_dtype = pv_acc_dtype + self.qhead_per_kvhead = qhead_per_kvhead + self.mma_tiler = mma_tiler + assert mma_tiler[0] == 128 and mma_tiler[1] == 128, "Only 128x128 tile impl is supported" + assert mma_tiler[2] == 256, "Only 256 is supported for 128x128 tile impl" + self.cta_tiler = ( + mma_tiler[0], + mma_tiler[1], + mma_tiler[2], + ) + self.qk_mma_tiler = ( + 2 * mma_tiler[0], + mma_tiler[1], + min(self.cta_tiler[2], 128), + ) + self.pv_mma_tiler = self.qk_mma_tiler + self.pv_block_tiler = ( + self.pv_mma_tiler[0] // 2, + self.pv_mma_tiler[1], + self.pv_mma_tiler[2], + ) + self.iterations_qk = self.cta_tiler[2] // self.qk_mma_tiler[2] + self.iterations_pv = self.cta_tiler[2] // self.pv_mma_tiler[1] + self.cluster_shape_mn = (2, 1) + self.tmem_warp_shape_mn = (4, 1) + # Dedicated hd256 kernel uses fixed scheduling policy. + self.is_persistent = False + self.is_causal = is_causal + self.is_local = is_local + self.use_semantic_trip_range = is_causal or is_local + self.use_clc_scheduler = False + + self.softmax_warp_ids = (0, 1, 2, 3) + self.correction_warp_ids = (4, 5, 6, 7) + self.mma_warp_id = 8 + self.load_warp_id = 9 + self.empty_warp_id = (10, 11) + self.sched_warp_id = self.empty_warp_id[0] if use_clc_scheduler else None + self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + + self.threads_per_warp = 32 + self.threads_per_cta = self.threads_per_warp * len( + ( + *self.softmax_warp_ids, # this is to get a round num threads + *self.correction_warp_ids, + self.mma_warp_id, + self.load_warp_id, + *self.empty_warp_id, + ) + ) + + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_cta, + ) + + self.tmem_s_offset = 0 + self.tmem_o_offset = 256 + self.tmem_p_offset = self.tmem_s_offset + + _tune_key = (True, is_causal, 256, False) # hd256: always 2cta, no sm103 variant + _tune = _TUNING_CONFIG.get(_tune_key, {}) + self.num_regs_softmax = _tune.get("num_regs_softmax", 256) + self.num_regs_correction = _tune.get("num_regs_correction", 160) + self.num_regs_other = 32 # fixed for hd256; not derived from 512 budget like other kernels + self.ex2_emu_freq = _tune.get("ex2_emu_freq", 4) + self.ex2_emu_res = _tune.get("ex2_emu_res", 3) + self.ex2_emu_start_frg = _tune.get("ex2_emu_start_frg", 0) + + self.buffer_align_bytes = 1024 + + def _setup_attributes(self): + self.q_stage = self.iterations_qk + self.kv_stage = 4 + self.qk_acc_stage = 2 + self.mma_corr_stage = 1 + if cutlass.const_expr(self.use_clc_scheduler): + self.num_clc_stage = 1 + self.num_clc_response_bytes = 16 + + @cute.jit + def __call__( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mO: cute.Tensor, + mLSE: Optional[cute.Tensor], + softmax_scale: Float32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + mPageTable: Optional[cute.Tensor] = None, + window_size_left: Int32 | int | None = None, + window_size_right: Int32 | int | None = None, + learnable_sink: Optional[cute.Tensor] = None, + descale_tensors: Optional[DescaleTensors] = None, + blocksparse_tensors: Optional[cute.Tensor] = None, + aux_data: AuxData = AuxData(), + stream: cuda.CUstream = None, + ): + # Keep parity with FlashAttentionForwardSm100.__call__ interface. + # (TODO@wangsiyu) Implement these features. + assert mSeqUsedQ is None and mSeqUsedK is None, ( + "SM100 forward with head_dim=256 does not support seqused_q/seqused_k" + ) + assert learnable_sink is None, ( + "SM100 forward with head_dim=256 does not support learnable_sink" + ) + assert blocksparse_tensors is None, ( + "SM100 forward with head_dim=256 does not support block sparsity" + ) + assert aux_data.tensors is None, ( + "SM100 forward with head_dim=256 does not support aux_tensors" + ) + assert aux_data.scalars is None, ( + "SM100 forward with head_dim=256 does not support aux_scalars" + ) + assert not self.is_local, ( + "SM100 forward with head_dim=256 does not support local attention yet" + ) + assert window_size_left is None and window_size_right is None, ( + "SM100 forward with head_dim=256 does not support runtime window_size overrides" + ) + assert descale_tensors is None, ( + "SM100 forward with head_dim=256 does not support descale_tensors" + ) + + q_tensor, k_tensor, v_tensor, o_tensor = mQ, mK, mV, mO + lse_tensor = mLSE + cum_seqlen_q = mCuSeqlensQ + cum_seqlen_k = mCuSeqlensK + + q_rank = len(mQ.shape) + k_rank = len(mK.shape) + if cutlass.const_expr(cum_seqlen_q is not None): + # Varlen path accepts either legacy 5D tensors or standard 3D tensors. + if cutlass.const_expr(q_rank == 5): + s_q = mQ.shape[1] + h_q = mQ.shape[2] * mQ.shape[3] + d = mQ.shape[4] + elif cutlass.const_expr(q_rank == 3): + s_q = mQ.shape[0] + h_q = mQ.shape[1] + d = mQ.shape[2] + else: + raise RuntimeError(f"hd256 forward varlen expects q rank 3 or 5, got rank {q_rank}") + else: + # Non-varlen path accepts either legacy 5D tensors or standard 4D tensors. + if cutlass.const_expr(q_rank == 5): + s_q = mQ.shape[1] + h_q = mQ.shape[2] * mQ.shape[3] + d = mQ.shape[4] + elif cutlass.const_expr(q_rank == 4): + s_q = mQ.shape[1] + h_q = mQ.shape[2] + d = mQ.shape[3] + else: + raise RuntimeError( + f"hd256 forward non-varlen expects q rank 4 or 5, got rank {q_rank}" + ) + + if cutlass.const_expr(cum_seqlen_k is not None): + if cutlass.const_expr(k_rank == 5): + s_k = mK.shape[1] + h_k = mK.shape[2] + elif cutlass.const_expr(k_rank == 3): + s_k = mK.shape[0] + h_k = mK.shape[1] + else: + raise RuntimeError(f"hd256 forward varlen expects k rank 3 or 5, got rank {k_rank}") + else: + if cutlass.const_expr(k_rank == 5): + s_k = mK.shape[1] + h_k = mK.shape[2] + elif cutlass.const_expr(k_rank == 4): + s_k = mK.shape[1] + h_k = mK.shape[2] + else: + raise RuntimeError( + f"hd256 forward non-varlen expects k rank 4 or 5, got rank {k_rank}" + ) + if cutlass.const_expr(cum_seqlen_q is not None): + b = mCuSeqlensQ.shape[0] - 1 + elif cutlass.const_expr(cum_seqlen_k is not None): + b = mCuSeqlensK.shape[0] - 1 + else: + b = mQ.shape[0] + + scale_softmax = softmax_scale + scale_softmax_log2 = softmax_scale * math.log2(math.exp(1.0)) + scale_output = 1.0 + s_lse = s_q + h_r = h_q // h_k + s_q64 = Int64(s_q) + s_k64 = Int64(s_k) + s_lse64 = Int64(s_lse) + d64 = cute.assume(Int64(d), divby=128) + h_r64 = Int64(h_r) + h_k64 = Int64(h_k) + b64 = Int64(b) + s_q_total = ( + q_tensor.shape[1] + if cum_seqlen_q is not None and q_rank == 5 + else (q_tensor.shape[0] if cum_seqlen_q is not None else s_q64) + ) + s_k_total = ( + k_tensor.shape[1] + if cum_seqlen_k is not None and k_rank == 5 + else (k_tensor.shape[0] if cum_seqlen_k is not None else s_k64) + ) + stride_b_qo = h_r64 * h_k64 * s_q64 * d64 if cum_seqlen_q is None else 0 + stride_b_kv = h_k64 * s_k64 * d64 if cum_seqlen_k is None else 0 + b_lse = b64 if cum_seqlen_q is None else 1 + stride_b_lse = h_r64 * h_k64 * s_lse64 if cum_seqlen_q is None else 0 + + # (s, d, ((h_r, h_k), b)) + q_layout = cute.make_layout( + (s_q_total, d, ((h_r, h_k), b)), + stride=(d64 * h_r64 * h_k64, 1, ((d64, d64 * h_r64), stride_b_qo)), + ) + q = cute.make_tensor(q_tensor.iterator, q_layout) + if cutlass.const_expr(mPageTable is not None): + # Paged: K layout (num_pages, page_size, h_k, d); page_table maps kv_coord→physical page. + num_pages = k_tensor.shape[0] + page_size = k_tensor.shape[1] + page_size64 = Int64(page_size) + max_seqlen_k_paged = Int32(mPageTable.shape[1] * page_size) + k_paged_layout = cute.make_layout( + (page_size, d, h_k, num_pages), + stride=(d64 * h_k64, 1, d64, page_size64 * d64 * h_k64), + ) + k = cute.make_tensor(k_tensor.iterator, k_paged_layout) + v_paged_layout = cute.make_layout( + (d, page_size, h_k, num_pages), + stride=(1, d64 * h_k64, d64, page_size64 * d64 * h_k64), + ) + v = cute.make_tensor(v_tensor.iterator, v_paged_layout) + page_table_layout = cute.make_layout( + (b, mPageTable.shape[1]), + stride=(Int64(mPageTable.shape[1]), 1), + ) + page_table = cute.make_tensor(mPageTable.iterator, page_table_layout) + else: + # (s, d, ((h_r, h_k), b)), 0-stride for h_r to broadcast + k_layout = cute.make_layout( + (s_k_total, d, ((h_r, h_k), b)), + stride=(d64 * h_k64, 1, ((0, d64), stride_b_kv)), + ) + k = cute.make_tensor(k_tensor.iterator, k_layout) + # (d, s, ((h_r, h_k), b)), 0-stride for h_r to broadcast + v_layout = cute.make_layout( + (d, s_k_total, ((h_r, h_k), b)), + stride=(1, d64 * h_k64, ((0, d64), stride_b_kv)), + ) + v = cute.make_tensor(v_tensor.iterator, v_layout) + page_table = None + max_seqlen_k_paged = None + # (s, d, ((h_r, h_k), b)) + o_layout = cute.make_layout( + (s_q_total, d, ((h_r, h_k), b)), + stride=(d64 * h_r64 * h_k64, 1, ((d64, d64 * h_r64), stride_b_qo)), + ) + o = cute.make_tensor(o_tensor.iterator, o_layout) + if cutlass.const_expr(lse_tensor is not None): + # (s, ((h_r, h_k), b)) + lse_layout = cute.make_layout( + (s_lse64, ((h_r, h_k), b_lse)), + stride=(1, ((s_lse64, h_r64 * s_lse64), stride_b_lse)), + ) + lse = cute.make_tensor(lse_tensor.iterator, lse_layout) + else: + lse = None + + # setup static attributes before smem/grid/tma computation + self.q_dtype = q.element_type + self.k_dtype = k.element_type + self.v_dtype = v.element_type + self.o_dtype = o.element_type + self.tilePlikeFP32 = self.qk_mma_tiler[1] // Float32.width * self.q_dtype.width + + if cutlass.const_expr(self.use_clc_scheduler): + self.tile_sched_params, grid = compute_grid_clc( + (s_q, o.shape[1], o.shape[2]) if cum_seqlen_q is not None else o.shape, + self.cta_tiler, + (*self.cluster_shape_mn, 1), + ) + else: + self.tile_sched_params, grid = compute_grid( + (s_q, o.shape[1], o.shape[2]) if cum_seqlen_q is not None else o.shape, + self.cta_tiler, + self.is_persistent, + ) + + self.q_major_mode = utils.LayoutEnum.from_tensor(q).mma_major_mode() + self.k_major_mode = utils.LayoutEnum.from_tensor(k).mma_major_mode() + self.v_major_mode = utils.LayoutEnum.from_tensor(v).mma_major_mode() + self.o_layout = utils.LayoutEnum.from_tensor(o) + + if cutlass.const_expr(self.q_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of q is not supported") + if cutlass.const_expr(self.k_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of k is not supported") + if cutlass.const_expr(self.v_major_mode != tcgen05.OperandMajorMode.MN): + raise RuntimeError("The layout of v is not supported") + + # check type consistency + if cutlass.const_expr(self.q_dtype != self.k_dtype): + raise TypeError(f"Type mismatch: {self.q_dtype} != {self.k_dtype}") + if cutlass.const_expr(self.q_dtype != self.v_dtype): + raise TypeError(f"Type mismatch: {self.q_dtype} != {self.v_dtype}") + self._setup_attributes() + + cta_group = tcgen05.CtaGroup.TWO + # the intermediate tensor p is from tmem & k-major + p_source = tcgen05.OperandSource.TMEM + p_major_mode = tcgen05.OperandMajorMode.K + qk_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.q_dtype, + self.q_major_mode, + self.k_major_mode, + self.qk_acc_dtype, + cta_group, + self.qk_mma_tiler[:2], + ) + pv_tiled_mma = sm100_utils.make_trivial_tiled_mma( + self.v_dtype, + p_major_mode, + self.v_major_mode, + self.pv_acc_dtype, + cta_group, + self.pv_mma_tiler[:2], + p_source, + ) + + self.cluster_shape_mnk = (*self.cluster_shape_mn, 1) + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout(self.cluster_shape_mnk), + (qk_tiled_mma.thr_id.shape,), + ) + + self.epi_tile = self.pv_block_tiler[:2] + + q_smem_layout_staged = sm100_utils.make_smem_layout_a( + qk_tiled_mma, + self.qk_mma_tiler, + self.q_dtype, + self.q_stage, + ) + k_smem_layout_staged = sm100_utils.make_smem_layout_b( + qk_tiled_mma, + self.qk_mma_tiler, + self.k_dtype, + self.kv_stage, + ) + p_tmem_layout_staged = sm100_utils.make_smem_layout_a( + pv_tiled_mma, + self.pv_mma_tiler, + self.q_dtype, + self.qk_acc_stage, + ) + p_tmem_layout = cute.select(p_tmem_layout_staged, mode=[0, 1, 2]) + v_smem_layout_staged = sm100_utils.make_smem_layout_b( + pv_tiled_mma, + self.pv_mma_tiler, + self.v_dtype, + self.kv_stage, + ) + # TMA load for Q + tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp(cta_group) + + q_smem_layout = cute.select(q_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_q, tma_tensor_q = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, + q, + q_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # TMA load for K + k_smem_layout = cute.select(k_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_k, tma_tensor_k = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + k, + k_smem_layout, + self.qk_mma_tiler, + qk_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + # TMA load for V + v_smem_layout = cute.select(v_smem_layout_staged, mode=[0, 1, 2]) + tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, + v, + v_smem_layout, + self.pv_mma_tiler, + pv_tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + q_copy_size = cute.size_in_bytes(self.q_dtype, q_smem_layout) + k_copy_size = cute.size_in_bytes(self.k_dtype, k_smem_layout) + self.tma_copy_q_bytes = q_copy_size * cute.size(qk_tiled_mma.thr_id.shape) + self.tma_copy_kv_bytes = k_copy_size * cute.size(qk_tiled_mma.thr_id.shape) + + @cute.struct + class SharedStorage: + # TMA G2S load barriers: LOAD warp (producer) -> MMA warp (consumer) + load_q_mbar_ptr: cute.struct.MemRange[ + Int64, self.q_stage * 2 + ] # load_q_{producer,consumer} + load_kv_mbar_ptr: cute.struct.MemRange[ + Int64, self.kv_stage * 2 + ] # load_kv_{producer,consumer} + mma_s_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2] + p_mma_mbar_ptr: cute.struct.MemRange[Int64, self.qk_acc_stage * 2] + # Softmax -> Correction signaling barriers (row_max/row_sum vec ready) + s_corr_mbar_ptr: cute.struct.MemRange[ + Int64, self.qk_acc_stage * 2 + ] # s_corr_{producer,consumer} + sum_mbar_ptr: cute.struct.MemRange[Int64, 2] + # MMA -> Correction ownership barriers for O_partial tokens (online rescale/finalize) + mma_corr_mbar_ptr: cute.struct.MemRange[ + Int64, self.mma_corr_stage * 2 + ] # mma_corr_{producer,consumer} + # A CTA-wide "TMEM lifetime" barrier used to safely deallocate TMEM after all users finish. + tmem_dealloc_mbar_ptr: Int64 + # Tmem holding buffer + tmem_holding_buf: Int32 + # CLC pipeline barriers and response buffer + clc_mbar_ptr: cute.struct.MemRange[Int64, 2] + clc_response: cute.struct.MemRange[Int32, 4] + + self.shared_storage = SharedStorage + + grid = cute.round_up(grid, self.cluster_shape_mnk) + # Launch the kernel synchronously + self.kernel( + qk_tiled_mma, + pv_tiled_mma, + tma_atom_q, + tma_tensor_q, + tma_atom_k, + tma_tensor_k, + tma_atom_v, + tma_tensor_v, + o, + cum_seqlen_q, + cum_seqlen_k, + lse, + scale_softmax_log2, + scale_softmax, + scale_output, + page_table, + max_seqlen_k_paged, + window_size_left, + window_size_right, + self.cluster_layout_vmnk, + q_smem_layout_staged, + k_smem_layout_staged, + p_tmem_layout, + v_smem_layout_staged, + self.tile_sched_params, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + stream=stream, + min_blocks_per_mp=1, + ) + + # GPU device kernel + @cute.kernel + def kernel( + self, + qk_tiled_mma: cute.TiledMma, + pv_tiled_mma: cute.TiledMma, + tma_atom_q: cute.CopyAtom, + mQ_qdl: cute.Tensor, + tma_atom_k: cute.CopyAtom, + mK_kdl: cute.Tensor, + tma_atom_v: cute.CopyAtom, + mV_dkl: cute.Tensor, + mO_qdl: cute.Tensor, + cum_seqlen_q: Optional[cute.Tensor], + cum_seqlen_k: Optional[cute.Tensor], + mLSE: Optional[cute.Tensor], + scale_softmax_log2: Float32, + scale_softmax: Float32, + scale_output: Float32, + mPageTable: Optional[cute.Tensor], + max_seqlen_k: Optional[Int32], + window_size_left: Optional[Int32], + window_size_right: Optional[Int32], + cluster_layout_vmnk: cute.Layout, + q_smem_layout_staged: cute.ComposedLayout, + k_smem_layout_staged: cute.ComposedLayout, + p_tmem_layout_staged: cute.ComposedLayout, + v_smem_layout_staged: cute.ComposedLayout, + tile_sched_params: FmhaStaticTileSchedulerParams | FmhaClcDynamicTileSchedulerParams, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + # + # Prefetch tma desc + # + if warp_idx == self.load_warp_id: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_k) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_v) + + bidx, _, _ = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(qk_tiled_mma.thr_id.shape) + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + + # Alloc + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + load_q_producer, load_q_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.q_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_q_bytes, + barrier_storage=storage.load_q_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + load_kv_producer, load_kv_consumer = pipeline.PipelineTmaUmma.create( + num_stages=self.kv_stage, + producer_group=make_thread_cooperative_group(len([self.load_warp_id])), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + tx_count=self.tma_copy_kv_bytes, + barrier_storage=storage.load_kv_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + mma_s_producer, mma_s_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.qk_acc_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group( + len(self.softmax_warp_ids) * self.threads_per_warp * self.cluster_shape_mnk[0], + ), + barrier_storage=storage.mma_s_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + p_mma_producer, p_mma_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.qk_acc_stage, + producer_group=make_thread_cooperative_group( + len(self.softmax_warp_ids) * self.threads_per_warp * self.cluster_shape_mnk[0], + ), + consumer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + barrier_storage=storage.p_mma_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + s_corr_producer, s_corr_consumer = pipeline.PipelineAsync.create( + num_stages=self.qk_acc_stage, + producer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.softmax_warp_ids) + ), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.correction_warp_ids) + ), + barrier_storage=storage.s_corr_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + sum_producer, sum_consumer = pipeline.PipelineAsync.create( + num_stages=1, + producer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.softmax_warp_ids) + ), + consumer_group=make_thread_cooperative_group( + self.threads_per_warp * len(self.correction_warp_ids) + ), + barrier_storage=storage.sum_mbar_ptr.data_ptr(), + defer_sync=True, + ).make_participants() + mma_corr_producer, mma_corr_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.mma_corr_stage, + producer_group=make_thread_cooperative_group(len([self.mma_warp_id])), + consumer_group=make_thread_cooperative_group( + len(self.correction_warp_ids) * self.threads_per_warp * self.cluster_shape_mnk[0], + ), + barrier_storage=storage.mma_corr_mbar_ptr.data_ptr(), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ).make_participants() + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.correction_warp_ids[0], + is_two_cta=True, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr, + ) + tmem.allocate(self.tmem_alloc_cols) + tmem.wait_for_alloc() + tmem_ptr = tmem.retrieve_ptr(self.qk_acc_dtype) + # Initialize CLC state if using dynamic scheduler + if cutlass.const_expr(self.use_clc_scheduler): + clc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + cluster_size = cute.size(self.cluster_shape_mnk) + num_clc_consumer_threads = self.threads_per_warp * ( + 1 # sched_warp (CTA 0 only) + + cluster_size + * ( + len(self.softmax_warp_ids) + + len(self.correction_warp_ids) + + 1 # mma_warp + + 1 # load_warp + ) + ) + clc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_clc_consumer_threads + ) + clc_response_ptr = storage.clc_response.data_ptr() + clc = ClcState.create( + hw_scheduler=ClcDynamicPersistentTileScheduler.create( + self.tile_sched_params.clc_hw_params(), + cute.arch.block_idx(), + cute.arch.grid_dim(), + clc_response_ptr, + ), + pipeline=pipeline.PipelineClcFetchAsync.create( + barrier_storage=storage.clc_mbar_ptr.data_ptr(), + num_stages=self.num_clc_stage, + producer_group=clc_pipeline_producer_group, + consumer_group=clc_pipeline_consumer_group, + tx_count=self.num_clc_response_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ), + consumer_state=pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_clc_stage + ), + producer_state=pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_clc_stage + ), + ) + else: + clc = None + clc_response_ptr = None + + # Cluster arrive after barrier init + pipeline.pipeline_init_arrive(cluster_shape_mn=cluster_layout_vmnk, is_relaxed=True) + + sQ = smem.allocate_tensor( + element_type=self.q_dtype, + layout=q_smem_layout_staged.outer, + swizzle=q_smem_layout_staged.inner, + byte_alignment=128, + ) + sK = smem.allocate_tensor( + element_type=self.k_dtype, + layout=k_smem_layout_staged.outer, + swizzle=k_smem_layout_staged.inner, + byte_alignment=128, + ) + # K and V now use separate memory since we removed the transform stage + sV = smem.allocate_tensor( + element_type=self.v_dtype, + layout=v_smem_layout_staged.outer, + swizzle=v_smem_layout_staged.inner, + byte_alignment=128, + ) + + sSum = smem.allocate_tensor( + element_type=self.qk_acc_dtype, + layout=cute.make_layout(len(self.softmax_warp_ids) * self.threads_per_warp), + byte_alignment=128, + ) + qk_thr_mma = qk_tiled_mma.get_slice(mma_tile_coord_v) # default 1sm + pv_thr_mma = pv_tiled_mma.get_slice(mma_tile_coord_v) # default 1sm + tSrQ = qk_thr_mma.make_fragment_A(sQ) + tSrK = qk_thr_mma.make_fragment_B(sK) + tOrV = pv_thr_mma.make_fragment_B(sV) + qk_acc_shape = qk_thr_mma.partition_shape_C((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) + tStS = qk_thr_mma.make_fragment_C(cute.append(qk_acc_shape, self.qk_acc_stage)) + pv_acc_shape = pv_thr_mma.partition_shape_C((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) + tOtO = pv_thr_mma.make_fragment_C(pv_acc_shape) + tOtO_layout = cute.append( + tOtO.layout, + cute.make_layout( + self.iterations_pv, + stride=self.pv_mma_tiler[1] // self.tmem_warp_shape_mn[1], + ), + ) + tStS = cute.make_tensor(tStS.iterator + self.tmem_s_offset, tStS.layout) + tOtO_staged = cute.make_tensor(tOtO.iterator + self.tmem_o_offset, tOtO_layout) + + # /////////////////////////////////////////////////////////////////////////////// + # EMPTY + # /////////////////////////////////////////////////////////////////////////////// + for _i in cutlass.range_constexpr(len(self.empty_warp_id)): + if warp_idx == self.empty_warp_id[_i]: + cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + + if cutlass.const_expr(self.use_clc_scheduler): + tile_sched = FmhaClcDynamicTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + clc_response_ptr, + clc, + ) + else: + blk_idx = cute.arch.block_idx() + tile_sched = FmhaStaticTileScheduler( + tile_sched_params, blk_idx[0], blk_idx, cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + # Cluster wait + pipeline.pipeline_init_wait(cluster_shape_mn=cluster_layout_vmnk) + + # /////////////////////////////////////////////////////////////////////////////// + # LOAD + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.load_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx # (q_tile_idx, 0, (head_idx, batch_idx)) + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + continue_cond = False + batch_coord = curr_block_coord[2][1] + seqlen_q = mQ_qdl.shape[0] + seqlen_k = ( + mK_kdl.shape[0] if cutlass.const_expr(mPageTable is None) else max_seqlen_k + ) + cuseqlen_q = Int32(0) + cuseqlen_k = Int32(0) + block_offset = ( + Int32(0), + Int32(0), + Int32(0), + ((Int32(0), Int32(0)), Int32(0)), + ) + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + if cutlass.const_expr(cum_seqlen_k is not None): + cuseqlen_k = cum_seqlen_k[batch_coord] + seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k + block_offset = ( + cuseqlen_q, + cuseqlen_k, + Int32(0), + ((Int32(0), Int32(0)), Int32(0)), + ) + continue_cond = not FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.qk_mma_tiler[0], + mma_block_coord[0], + seqlen_q, + ) + if not continue_cond: + mQ_qdl_ = cute.domain_offset(cute.select(block_offset, mode=[0, 2, 3]), mQ_qdl) + # Local tile partition global tensors + q_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape + ) + # (bM, bK, loopM, loopK, loopL) + gQ_qdl = cute.flat_divide(mQ_qdl_, cute.select(self.qk_mma_tiler, mode=[0, 2])) + tSgQ_qdl = qk_thr_mma.partition_A(gQ_qdl) + tQsQ, tQgQ_qdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_q, + block_in_cluster_coord_vmnk[2], + q_cta_layout, + cute.group_modes(sQ, 0, 3), + cute.group_modes(tSgQ_qdl, 0, 3), + ) + kv_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape + ) + if cutlass.const_expr(mPageTable is None): + # Dense path: domain_offset K/V by batch block, select batch via mma_block_coord[2]. + mK_kdl_ = cute.domain_offset( + cute.select(block_offset, mode=[1, 2, 3]), mK_kdl + ) + mV_dkl_ = cute.domain_offset( + cute.select(block_offset, mode=[2, 1, 3]), mV_dkl + ) + gK_kdl = cute.flat_divide( + mK_kdl_, cute.select(self.qk_mma_tiler, mode=[1, 2]) + ) + tSgK_kdl = qk_thr_mma.partition_B(gK_kdl) + tKsK, tKgK_kdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_k, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + cute.group_modes(sK, 0, 3), + cute.group_modes(tSgK_kdl, 0, 3), + ) + gV_dkl = cute.flat_divide( + mV_dkl_, cute.select(self.pv_mma_tiler, mode=[1, 2]) + ) + tSgV_dkl = pv_thr_mma.partition_B(gV_dkl) + tVsV, tVgV_dkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_v, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + cute.group_modes(sV, 0, 3), + cute.group_modes(tSgV_dkl, 0, 3), + ) + # ((atom_v, rest_v), RestN, RestK) + tKgK = tKgK_kdl[None, None, None, mma_block_coord[2]] + tVgV = tVgV_dkl[None, None, None, mma_block_coord[2]] + else: + # Paged path: slice K/V by KV head, keep num_pages dim for page_idx-based TMA. + head_kv_coord = curr_block_coord[2][0] // self.qhead_per_kvhead + mK_kdl_ = mK_kdl[None, None, head_kv_coord, None] + mV_dkl_ = mV_dkl[None, None, head_kv_coord, None] + gK_kdl = cute.flat_divide( + mK_kdl_, cute.select(self.qk_mma_tiler, mode=[1, 2]) + ) + tSgK_kdl = qk_thr_mma.partition_B(gK_kdl) + tKsK, tKgK_kdl = cute.nvgpu.cpasync.tma_partition( + tma_atom_k, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + cute.group_modes(sK, 0, 3), + cute.group_modes(tSgK_kdl, 0, 3), + ) + gV_dkl = cute.flat_divide( + mV_dkl_, cute.select(self.pv_mma_tiler, mode=[1, 2]) + ) + tSgV_dkl = pv_thr_mma.partition_B(gV_dkl) + tVsV, tVgV_dkl = cute.nvgpu.cpasync.tma_partition( + tma_atom_v, + block_in_cluster_coord_vmnk[1], + kv_cta_layout, + cute.group_modes(sV, 0, 3), + cute.group_modes(tSgV_dkl, 0, 3), + ) + tKgK = tKgK_kdl + tVgV = tVgV_dkl + # ((atom_v, rest_v), RestK) + tQgQ = tQgQ_qdl[None, mma_block_coord[0], None, mma_block_coord[2]] + + seqlen_kv_loop_start, seqlen_kv_loop_steps = ( + FusedMask.get_trip_start_count_via_block_info( + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + self.is_causal, + self.is_local, + window_size_left, + window_size_right, + ) + ) + seqlen_kv_loop_end = seqlen_kv_loop_start + seqlen_kv_loop_steps + # Q + for iter in cutlass.range(self.iterations_qk, unroll=1): + q_handle = load_q_producer.acquire_and_advance() + cute.copy( + tma_atom_q, + tQgQ[None, iter], + tQsQ[None, q_handle.index], + tma_bar_ptr=q_handle.barrier, + ) + + # K0 + kv_coord = seqlen_kv_loop_start + k_page_idx = ( + mPageTable[batch_coord, kv_coord] + if cutlass.const_expr(mPageTable is not None) + else None + ) + for iter in cutlass.range(self.iterations_qk, unroll=1): + k_handle = load_kv_producer.acquire_and_advance() + cute.copy( + tma_atom_k, + tKgK[None, kv_coord, iter] + if cutlass.const_expr(mPageTable is None) + else tKgK[None, 0, iter, k_page_idx], + tKsK[None, k_handle.index], + tma_bar_ptr=k_handle.barrier, + ) + kv_coord += 1 + # v_page_idx_prev carries K[i-1]'s page index for use as V[i-1]'s page + # (K and V for the same KV block share the same physical page). + # Also serves as the Vend page index when seqlen_kv_loop_steps == 1. + v_page_idx_prev = ( + k_page_idx if cutlass.const_expr(mPageTable is not None) else None + ) + # Prefetch K1 page after K0 TMA dispatch to hide L2 latency. + if cutlass.const_expr(mPageTable is not None): + if seqlen_kv_loop_steps > 1: + k_page_idx = mPageTable[batch_coord, kv_coord] + + for i in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + # Ki: k_page_idx was prefetched at end of previous iteration + # (or in the prologue for i==1); L2 latency already hidden. + for iter in cutlass.range(self.iterations_qk, unroll=1): + k_handle = load_kv_producer.acquire_and_advance() + cute.copy( + tma_atom_k, + tKgK[None, kv_coord, iter] + if cutlass.const_expr(mPageTable is None) + else tKgK[None, 0, iter, k_page_idx], + tKsK[None, k_handle.index], + tma_bar_ptr=k_handle.barrier, + ) + # Vi-1: reuse v_page_idx_prev (= K[i-1]'s page), no extra GMEM read. + for iter in cutlass.range(self.iterations_pv, unroll=1): + v_handle = load_kv_producer.acquire_and_advance() + cute.copy( + tma_atom_v, + tVgV[None, iter, kv_coord - 1] + if cutlass.const_expr(mPageTable is None) + else tVgV[None, iter, 0, v_page_idx_prev], + tVsV[None, v_handle.index], + tma_bar_ptr=v_handle.barrier, + ) + v_page_idx_prev = ( + k_page_idx if cutlass.const_expr(mPageTable is not None) else None + ) + kv_coord += 1 + # Prefetch next K page while V TMA is in flight. + if cutlass.const_expr(mPageTable is not None): + if kv_coord < seqlen_kv_loop_end: + k_page_idx = mPageTable[batch_coord, kv_coord] + # Vend: reuse v_page_idx_prev (= K[end-1]'s page), no extra GMEM read. + for iter in cutlass.range(self.iterations_pv, unroll=1): + v_handle = load_kv_producer.acquire_and_advance() + cute.copy( + tma_atom_v, + tVgV[None, iter, seqlen_kv_loop_end - 1] + if cutlass.const_expr(mPageTable is None) + else tVgV[None, iter, 0, v_page_idx_prev], + tVsV[None, v_handle.index], + tma_bar_ptr=v_handle.barrier, + ) + + work_tile = tile_sched.advance_to_next_work() + # End of persistent scheduler loop + load_kv_producer.tail() + load_q_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # MMA + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + is_leader_cta = cta_rank_in_cluster % 2 == 0 + + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + continue_cond = False + seqlen_q = mQ_qdl.shape[0] + seqlen_k = ( + mK_kdl.shape[0] if cutlass.const_expr(mPageTable is None) else max_seqlen_k + ) + batch_coord = curr_block_coord[2][1] + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + continue_cond = not FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.qk_mma_tiler[0], + mma_block_coord[0], + seqlen_q, + ) + + if not continue_cond: + if cutlass.const_expr(cum_seqlen_k is not None): + cuseqlen_k = cum_seqlen_k[batch_coord] + seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k + + seqlen_kv_loop_start, seqlen_kv_loop_steps = ( + FusedMask.get_trip_start_count_via_block_info( + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + self.is_causal, + self.is_local, + window_size_left, + window_size_right, + ) + ) + seqlen_kv_loop_end = seqlen_kv_loop_start + seqlen_kv_loop_steps + + load_q_releaser = load_q_consumer.clone() + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + if seqlen_kv_loop_steps > 1: + # QK0 + if is_leader_cta: + s_handle = mma_s_producer.acquire_and_advance() + tStS_slice = tStS[None, None, None, s_handle.index] + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for iter in cutlass.range(self.iterations_qk, unroll=1): + load_q_consumer.wait_and_advance() + tSrQ_slice = tSrQ[None, None, None, iter] + k_handle = load_kv_consumer.wait_and_advance() + tSrK_trans_slice = tSrK[None, None, None, k_handle.index] + num_kphases = cute.size(tSrQ_slice, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + qk_tiled_mma, + tStS_slice, + tSrQ_slice[kphase_coord], + tSrK_trans_slice[kphase_coord], + tStS_slice, + ) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + k_handle.release() + s_handle.commit() + for i in cutlass.range(1, seqlen_kv_loop_steps - 1, 1, unroll=1): + # QKi + if is_leader_cta: + s_handle = mma_s_producer.acquire_and_advance() + tStS_slice = tStS[None, None, None, s_handle.index] + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for iter in cutlass.range(self.iterations_qk, unroll=1): + tSrQ_slice = tSrQ[None, None, None, iter] + k_handle = load_kv_consumer.wait_and_advance() + tSrK_trans_slice = tSrK[None, None, None, k_handle.index] + num_kphases = cute.size(tSrQ_slice, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + qk_tiled_mma, + tStS_slice, + tSrQ_slice[kphase_coord], + tSrK_trans_slice[kphase_coord], + tStS_slice, + ) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + k_handle.release() + s_handle.commit() + + # PVi-1 + p_handle = p_mma_consumer.wait_and_advance() + o_handle = mma_corr_producer.acquire_and_advance() + pv_whether_acc = pv_tiled_mma.get(tcgen05.Field.ACCUMULATE) + for iter in cutlass.range(self.iterations_pv, unroll=1): + v_handle = load_kv_consumer.wait_and_advance() + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, pv_whether_acc) + tOtO_slice = tOtO_staged[None, None, None, iter] + tStS_slice = tStS[None, None, None, p_handle.index] + tP = cute.make_tensor( + tStS_slice.iterator, p_tmem_layout_staged.outer + ) + tOrP = pv_thr_mma.make_fragment_A(tP) + tOrP_slice = cute.make_tensor( + cute.recast_ptr(tStS_slice.iterator, dtype=self.q_dtype), + tOrP.layout, + ) + tOrV_slice = tOrV[None, None, None, v_handle.index] + num_kphases = cute.size(tOrV_slice, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + pv_tiled_mma, + tOtO_slice, + tOrP_slice[kphase_coord], + tOrV_slice[kphase_coord], + tOtO_slice, + ) + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + v_handle.release() + o_handle.commit() + p_handle.release() + if is_leader_cta: + # QKend + s_handle = mma_s_producer.acquire_and_advance() + tStS_slice = tStS[None, None, None, s_handle.index] + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for iter in cutlass.range(self.iterations_qk, unroll=1): + tSrQ_slice = tSrQ[None, None, None, iter] + k_handle = load_kv_consumer.wait_and_advance() + tSrK_trans_slice = tSrK[None, None, None, k_handle.index] + num_kphases = cute.size(tSrQ_slice, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + qk_tiled_mma, + tStS_slice, + tSrQ_slice[kphase_coord], + tSrK_trans_slice[kphase_coord], + tStS_slice, + ) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + k_handle.release() + load_q_releaser.release() + load_q_releaser.advance() + s_handle.commit() + + # PVend-1 + p_handle = p_mma_consumer.wait_and_advance() + o_handle = mma_corr_producer.acquire_and_advance() + pv_whether_acc = pv_tiled_mma.get(tcgen05.Field.ACCUMULATE) + for iter in cutlass.range(self.iterations_pv, unroll=1): + v_handle = load_kv_consumer.wait_and_advance() + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, pv_whether_acc) + tOtO_slice = tOtO_staged[None, None, None, iter] + tStS_slice = tStS[None, None, None, p_handle.index] + tP = cute.make_tensor( + tStS_slice.iterator, p_tmem_layout_staged.outer + ) + tOrP = pv_thr_mma.make_fragment_A(tP) + tOrP_slice = cute.make_tensor( + cute.recast_ptr(tStS_slice.iterator, dtype=self.q_dtype), + tOrP.layout, + ) + tOrV_slice = tOrV[None, None, None, v_handle.index] + num_kphases = cute.size(tOrV_slice, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + pv_tiled_mma, + tOtO_slice, + tOrP_slice[kphase_coord], + tOrV_slice[kphase_coord], + tOtO_slice, + ) + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + v_handle.release() + o_handle.commit() + p_handle.release() + else: + if is_leader_cta: + # QK0 + s_handle = mma_s_producer.acquire_and_advance() + tStS_slice = tStS[None, None, None, s_handle.index] + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, False) + for iter in cutlass.range(self.iterations_qk, unroll=1): + load_q_consumer.wait_and_advance() + tSrQ_slice = tSrQ[None, None, None, iter] + k_handle = load_kv_consumer.wait_and_advance() + tSrK_trans_slice = tSrK[None, None, None, k_handle.index] + num_kphases = cute.size(tSrQ_slice, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + qk_tiled_mma, + tStS_slice, + tSrQ_slice[kphase_coord], + tSrK_trans_slice[kphase_coord], + tStS_slice, + ) + qk_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + k_handle.release() + load_q_releaser.release() + load_q_releaser.advance() + s_handle.commit() + + if is_leader_cta: + # PVend + p_handle = p_mma_consumer.wait_and_advance() + o_handle = mma_corr_producer.acquire_and_advance() + pv_whether_acc = pv_tiled_mma.get(tcgen05.Field.ACCUMULATE) + for iter in cutlass.range(self.iterations_pv, unroll=1): + v_handle = load_kv_consumer.wait_and_advance() + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, pv_whether_acc) + tOtO_slice = tOtO_staged[None, None, None, iter] + tStS_slice = tStS[None, None, None, p_handle.index] + tP = cute.make_tensor(tStS_slice.iterator, p_tmem_layout_staged.outer) + tOrP = pv_thr_mma.make_fragment_A(tP) + tOrP_slice = cute.make_tensor( + cute.recast_ptr(tStS_slice.iterator, dtype=self.q_dtype), + tOrP.layout, + ) + tOrV_slice = tOrV[None, None, None, v_handle.index] + num_kphases = cute.size(tOrV_slice, mode=[2]) + for kphase_idx in cutlass.range(num_kphases, unroll_full=True): + kphase_coord = (None, None, kphase_idx) + cute.gemm( + pv_tiled_mma, + tOtO_slice, + tOrP_slice[kphase_coord], + tOrV_slice[kphase_coord], + tOtO_slice, + ) + pv_tiled_mma.set(tcgen05.Field.ACCUMULATE, True) + v_handle.release() + o_handle.commit() + p_handle.release() + work_tile = tile_sched.advance_to_next_work() + # End of persistent scheduler loop + mma_s_producer.tail() + mma_corr_producer.tail() + + if warp_idx < self.correction_warp_ids[0] and warp_idx >= self.softmax_warp_ids[0]: + # increase register after decreasing + cute.arch.warpgroup_reg_alloc(self.num_regs_softmax) + + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + batch_coord = curr_block_coord[2][1] + continue_cond = False + seqlen_q = mQ_qdl.shape[0] + seqlen_k = ( + mK_kdl.shape[0] if cutlass.const_expr(mPageTable is None) else max_seqlen_k + ) + cuseqlen_q = Int32(0) + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + continue_cond = not FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.qk_mma_tiler[0], + mma_block_coord[0], + seqlen_q, + ) + if not continue_cond: + if cutlass.const_expr(cum_seqlen_k is not None): + cuseqlen_k = cum_seqlen_k[batch_coord] + seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k + + row_max = -Float32.inf + row_max_prev = -Float32.inf + row_sum = 0.0 + + start_count, trip_count = FusedMask.get_trip_start_count_via_block_info( + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + self.is_causal, + self.is_local, + window_size_left, + window_size_right, + ) + end_count = start_count + trip_count + # require at least one softmax iteration for zero trip_count case; + # rely on masking this iteration for correctness + if end_count <= start_count: + start_count = 0 + end_count = 1 + if cutlass.const_expr(self.use_semantic_trip_range): + n_block_min_causal_local_mask, n_block_min_before_local_mask = ( + FusedMask.get_trip_mask_bounds_via_block_info( + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + self.is_causal, + self.is_local, + window_size_left, + window_size_right, + ) + ) + cS_base = cute.make_identity_tensor( + (self.qk_mma_tiler[0], self.qk_mma_tiler[1]) + ) + cS = cute.domain_offset((mma_block_coord[0] * self.qk_mma_tiler[0], 0), cS_base) + tScS = qk_thr_mma.partition_C(cS) + + for step in cutlass.range(start_count, end_count, 1, unroll=1): + cS_iter = cute.domain_offset((0, step * self.qk_mma_tiler[1]), cS) + tScS_iter = qk_thr_mma.partition_C(cS_iter) + if cutlass.const_expr(self.use_semantic_trip_range): + need_apply_mask = ( + step >= n_block_min_causal_local_mask + or step < n_block_min_before_local_mask + or step == end_count - 1 + ) + else: + # Residual path only needs seqlen masking on the last K tile. + need_apply_mask = step == end_count - 1 + # Si -> Pi + ( + row_max, + row_sum, + mma_s_consumer, + p_mma_producer, + s_corr_producer, + ) = self.softmax_step( + (need_apply_mask, window_size_left, window_size_right), + ( + row_max_prev, + row_sum, + seqlen_q, + seqlen_k, + scale_softmax_log2, + ), + (tStS, tScS_iter), + (mma_s_consumer, p_mma_producer, s_corr_producer), + ) + row_max_prev = row_max + sum_producer = self.store_sum_max( + row_max, + mLSE, + row_sum, + sSum, + sum_producer, + curr_block_coord, + seqlen_q, + cum_seqlen_q, + cuseqlen_q, + scale_softmax, + ) + work_tile = tile_sched.advance_to_next_work() + p_mma_producer.tail() + s_corr_producer.tail() + + # /////////////////////////////////////////////////////////////////////////////// + # Correction + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx >= self.correction_warp_ids[0] and warp_idx < self.mma_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_correction) + + while work_tile.is_valid_tile: + curr_block_coord = work_tile.tile_idx + mma_block_coord = ( + curr_block_coord[0] // cute.size(qk_tiled_mma.thr_id.shape), + curr_block_coord[1], + curr_block_coord[2], + ) + batch_coord = curr_block_coord[2][1] + seqlen_q = mQ_qdl.shape[0] + seqlen_k = ( + mK_kdl.shape[0] if cutlass.const_expr(mPageTable is None) else max_seqlen_k + ) + continue_cond = False + cuseqlen_q = Int32(0) + if cutlass.const_expr(cum_seqlen_q is not None): + cuseqlen_q = cum_seqlen_q[batch_coord] + seqlen_q = cum_seqlen_q[batch_coord + 1] - cuseqlen_q + continue_cond = not FmhaStaticTileScheduler.check_valid_work_for_seqlen_q( + self.qk_mma_tiler[0], + mma_block_coord[0], + seqlen_q, + ) + + if not continue_cond: + if cutlass.const_expr(cum_seqlen_k is not None): + cuseqlen_k = cum_seqlen_k[batch_coord] + seqlen_k = cum_seqlen_k[batch_coord + 1] - cuseqlen_k + + mO_qdl_eff = mO_qdl + if cutlass.const_expr(cum_seqlen_q is not None): + block_offset_o = ( + cuseqlen_q, + Int32(0), + Int32(0), + ((Int32(0), Int32(0)), Int32(0)), + ) + mO_qdl_eff = cute.domain_offset( + cute.select(block_offset_o, mode=[0, 2, 3]), mO_qdl + ) + + # (bM, bN, loopM, loopN, loopL) + gO_qdl = cute.flat_divide( + mO_qdl_eff, cute.select(self.pv_block_tiler, mode=[0, 1]) + ) + cO_qdl = cute.flat_divide( + cute.make_identity_tensor(mO_qdl_eff.shape), + cute.select(self.pv_block_tiler, mode=[0, 1]), + ) + + _, seqlen_kv_loop_steps = FusedMask.get_trip_start_count_via_block_info( + mma_block_coord, + self.qk_mma_tiler, + seqlen_q, + seqlen_k, + self.is_causal, + self.is_local, + window_size_left, + window_size_right, + ) + gO_staged = gO_qdl[None, None, curr_block_coord[0], None, curr_block_coord[2]] + cO_staged = cO_qdl[None, None, curr_block_coord[0], None, curr_block_coord[2]] + cS = cute.make_identity_tensor((self.qk_mma_tiler[0], self.qk_mma_tiler[1])) + tScS = qk_thr_mma.partition_C(cS) + + # Empty step as the first step is no need for correction + stats_handle = s_corr_consumer.wait_and_advance() + stats_handle.release() + for step in cutlass.range(1, seqlen_kv_loop_steps, 1, unroll=1): + # Oi-1 -> Oi + mma_corr_consumer, s_corr_consumer = self.correction_rescale( + scale_softmax_log2, + (s_corr_consumer, tStS, tScS), + (mma_corr_consumer, tOtO_staged, cO_staged), + self.epi_tile, + ) + # O_partial -> O_final + mma_corr_consumer, sum_consumer = self.correction_epilog( + (seqlen_q, scale_output), + (sum_consumer, sSum), + (mma_corr_consumer, gO_staged, cO_staged, tOtO_staged), + self.epi_tile, + ) + work_tile = tile_sched.advance_to_next_work() + # NOTE: tmem.free() moved to kernel end to enable cluster-wide sync + + # /////////////////////////////////////////////////////////////////////////////// + # Scheduler Warp (only for CLC dynamic scheduler) + # /////////////////////////////////////////////////////////////////////////////// + if cutlass.const_expr(self.use_clc_scheduler): + is_first_cta_in_cluster = cta_rank_in_cluster == 0 + + if warp_idx == self.sched_warp_id and is_first_cta_in_cluster: + cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + while work_tile.is_valid_tile: + tile_sched.prefetch_next_work() + work_tile = tile_sched.advance_to_next_work() + tile_sched.producer_tail() + + # /////////////////////////////////////////////////////////////////////////////// + # Empty warps reg dealloc + # /////////////////////////////////////////////////////////////////////////////// + if cutlass.const_expr(self.use_clc_scheduler): + if warp_idx > self.load_warp_id: + if not (warp_idx == self.sched_warp_id and is_first_cta_in_cluster): + cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + else: + if warp_idx > self.load_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + + # /////////////////////////////////////////////////////////////////////////////// + # Cooperative TMEM Deallocation (2CTA) + # /////////////////////////////////////////////////////////////////////////////// + # All warps (including scheduler) have finished by this point. + # Cluster-wide sync ensures both CTAs reach here before dealloc. + cute.arch.cluster_arrive() + cute.arch.cluster_wait() + tmem.relinquish_alloc_permit() + tmem.free(tmem_ptr) + + return + + @cute.jit + def softmax_step( + self, + mask_args: Tuple, + value_args: Tuple, + tensor_args: Tuple, + pipeline_args: Tuple, + ) -> Tuple[Float32, Float32, pipeline.PipelineConsumer, pipeline.PipelineProducer]: + need_apply_mask, window_size_left, window_size_right = mask_args + row_max, row_sum, seqlen_q, seqlen_k, scale_softmax_log2 = value_args + tStS, tScS = tensor_args + mma_s_consumer, p_mma_producer, s_corr_producer = pipeline_args + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + s_handle = mma_s_consumer.wait_and_advance() + tStS_slice = tStS[(None, None), 0, 0, s_handle.index] + tScS_slice = tScS[(None, None), 0, 0] + tmem_load_atom = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(32)), self.qk_acc_dtype + ) + tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tStS_slice) + thr_load = tmem_tiled_load.get_slice(thread_idx) + tTMEM_LOADtS = thr_load.partition_S(tStS_slice) + tTMEM_LOADcS = thr_load.partition_D(tScS_slice) + tTMEM_LOADrS = cute.make_rmem_tensor(tTMEM_LOADcS.shape, self.qk_acc_dtype) + cute.copy(tmem_tiled_load, tTMEM_LOADtS, tTMEM_LOADrS) + + cute.arch.fence_view_async_tmem_load() + s_handle.release() + if need_apply_mask: + FusedMask.apply_mask_via_causal_local( + tTMEM_LOADrS, + tTMEM_LOADcS, + seqlen_q, + seqlen_k, + self.use_semantic_trip_range, + self.is_causal, + self.is_local, + window_size_left, + window_size_right, + ) + old_row_max = row_max + row_max = tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0) + row_max_safe = row_max + if row_max == -cutlass.Float32.inf: + row_max_safe = 0.0 + + stats_handle = s_corr_producer.acquire_and_advance() + stats_layout = cute.composition( + tStS_slice.layout, cute.make_layout((tStS_slice.shape[0], 2)) + ) + stats_c_layout = cute.composition( + tScS_slice.layout, cute.make_layout((tScS_slice.shape[0], 2)) + ) + tOtStats = cute.make_tensor(tStS_slice.iterator + self.tilePlikeFP32, stats_layout) + tOcStats = cute.make_tensor(tScS_slice.iterator, stats_c_layout) + tmem_store_stats_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(2)), + self.qk_acc_dtype, + ) + tiled_tmem_store_stats = tcgen05.make_tmem_copy(tmem_store_stats_atom, tOtStats) + thr_tmem_store_stats = tiled_tmem_store_stats.get_slice(thread_idx) + tTMEM_STOREcStats = thr_tmem_store_stats.partition_S(tOcStats) + tTMEM_STORErStats = cute.make_rmem_tensor(tTMEM_STOREcStats.shape, self.qk_acc_dtype) + tTMEM_STORErStats[0] = old_row_max + tTMEM_STORErStats[1] = row_max_safe + tTMEM_STOREtStats = thr_tmem_store_stats.partition_D(tOtStats) + cute.copy(tiled_tmem_store_stats, tTMEM_STORErStats, tTMEM_STOREtStats) + cute.arch.fence_view_async_tmem_store() + stats_handle.commit() + + scale = scale_softmax_log2 + minus_row_max_scale = (0.0 - row_max_safe) * scale + # Acquire P write slot early — overlaps any pipeline stall with exp2 compute + p_handle = p_mma_producer.acquire_and_advance() + # Fragment-based FMA + exp2 + bf16 conversion + # Trades SFU for FMA via polynomial emulation on a fraction of elements + ex2_frg_tile = 32 + ex2_frg_cnt = cute.size(tTMEM_LOADrS) // ex2_frg_tile + tTMEM_LOADrS_ex2 = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(ex2_frg_tile)) + tTMEM_STORErP = cute.make_rmem_tensor(tTMEM_LOADrS.shape, self.q_dtype) + tTMEM_STORErP_ex2 = cute.logical_divide(tTMEM_STORErP, cute.make_layout(ex2_frg_tile)) + for j in cutlass.range_constexpr(ex2_frg_cnt): + for k in cutlass.range_constexpr(0, ex2_frg_tile, 2): + tTMEM_LOADrS_ex2[k, j], tTMEM_LOADrS_ex2[k + 1, j] = cute.arch.fma_packed_f32x2( + (tTMEM_LOADrS_ex2[k, j], tTMEM_LOADrS_ex2[k + 1, j]), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + ) + if cutlass.const_expr(self.ex2_emu_freq == 0): + tTMEM_LOADrS_ex2[k, j] = cute.math.exp2(tTMEM_LOADrS_ex2[k, j], fastmath=True) + tTMEM_LOADrS_ex2[k + 1, j] = cute.math.exp2( + tTMEM_LOADrS_ex2[k + 1, j], fastmath=True + ) + else: + if cutlass.const_expr( + k % self.ex2_emu_freq < self.ex2_emu_freq - self.ex2_emu_res + or j >= ex2_frg_cnt - 1 + or j < self.ex2_emu_start_frg + ): + tTMEM_LOADrS_ex2[k, j] = cute.math.exp2( + tTMEM_LOADrS_ex2[k, j], fastmath=True + ) + tTMEM_LOADrS_ex2[k + 1, j] = cute.math.exp2( + tTMEM_LOADrS_ex2[k + 1, j], fastmath=True + ) + else: + tTMEM_LOADrS_ex2[k, j], tTMEM_LOADrS_ex2[k + 1, j] = ex2_emulation_2( + tTMEM_LOADrS_ex2[k, j], tTMEM_LOADrS_ex2[k + 1, j] + ) + tTMEM_STORErP_ex2[None, j].store(tTMEM_LOADrS_ex2[None, j].load().to(self.q_dtype)) + tmem_store_atom = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(32)), self.qk_acc_dtype + ) + tilePlikeFP32 = tStS_slice.shape[1] // Float32.width * self.q_dtype.width + tStS_P_layout = cute.composition( + tStS_slice.layout, cute.make_layout((tStS_slice.shape[0], tilePlikeFP32)) + ) + tStS_P = cute.make_tensor(tStS_slice.iterator, tStS_P_layout) + tScS_P_layout = cute.composition( + tScS_slice.layout, cute.make_layout((tScS_slice.shape[0], tilePlikeFP32)) + ) + tScS_P = cute.make_tensor(tScS_slice.iterator, tScS_P_layout) + tmem_tiled_store = tcgen05.make_tmem_copy(tmem_store_atom, tStS_P) + thr_store = tmem_tiled_store.get_slice(thread_idx) + tTMEM_STOREtP = thr_store.partition_D(tStS_P) + tTMEM_STOREcS = thr_store.partition_S(tScS_P) + tTMEM_STORErP_ = cute.make_tensor( + cute.recast_ptr(tTMEM_STORErP.iterator, dtype=self.qk_acc_dtype), + tTMEM_STOREcS.shape, + ) + cute.copy(tmem_tiled_store, tTMEM_STORErP_, tTMEM_STOREtP) + cute.arch.fence_view_async_tmem_store() + + p_handle.commit() + acc_scale_ = scale * (old_row_max - row_max_safe) + acc_scale = cute.math.exp2(acc_scale_, fastmath=True) * 0.5 + # TODO: calc row sum with TensorSSA + row_sum *= acc_scale + local_row_sum_0 = (row_sum, row_sum) + local_row_sum_1 = (0.0, 0.0) + local_row_sum_2 = (0.0, 0.0) + local_row_sum_3 = (0.0, 0.0) + reduction_unroll = 4 + frg_tile = cute.size(tTMEM_LOADrS) // reduction_unroll + tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile)) + for j in cutlass.range_constexpr(0, cute.size(tTMEM_LOADrS_frg, mode=[0]), 2): + local_row_sum_0 = cute.arch.add_packed_f32x2( + local_row_sum_0, (tTMEM_LOADrS_frg[j, 0], tTMEM_LOADrS_frg[j + 1, 0]) + ) + local_row_sum_1 = cute.arch.add_packed_f32x2( + local_row_sum_1, (tTMEM_LOADrS_frg[j, 1], tTMEM_LOADrS_frg[j + 1, 1]) + ) + local_row_sum_2 = cute.arch.add_packed_f32x2( + local_row_sum_2, (tTMEM_LOADrS_frg[j, 2], tTMEM_LOADrS_frg[j + 1, 2]) + ) + local_row_sum_3 = cute.arch.add_packed_f32x2( + local_row_sum_3, (tTMEM_LOADrS_frg[j, 3], tTMEM_LOADrS_frg[j + 1, 3]) + ) + local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_1) + local_row_sum_2 = cute.arch.add_packed_f32x2(local_row_sum_2, local_row_sum_3) + local_row_sum_0 = cute.arch.add_packed_f32x2(local_row_sum_0, local_row_sum_2) + row_sum = local_row_sum_0[0] + local_row_sum_0[1] + return row_max, row_sum, mma_s_consumer, p_mma_producer, s_corr_producer + + @cute.jit + def correction_rescale( + self, + scale_softmax_log2: Float32, + stats_args: tuple, + o_args: tuple, + epi_tile: cute.Tile, + ) -> pipeline.PipelineConsumer: + (s_corr_consumer, tStS, tScS) = stats_args + (mma_o_consumer, tOtO_staged, cO_staged) = o_args + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + + stats_handle = s_corr_consumer.wait_and_advance() + tStS_slice = tStS[(None, None), 0, 0, stats_handle.index] + tScS_slice = tScS[(None, None), 0, 0] + stats_layout = cute.composition( + tStS_slice.layout, cute.make_layout((tStS_slice.shape[0], 2)) + ) + stats_c_layout = cute.composition( + tScS_slice.layout, cute.make_layout((tScS_slice.shape[0], 2)) + ) + tOtStats = cute.make_tensor(tStS_slice.iterator + self.tilePlikeFP32, stats_layout) + tOcStats = cute.make_tensor(tScS_slice.iterator, stats_c_layout) + tmem_load_stats_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(2)), + self.qk_acc_dtype, + ) + tiled_tmem_load_stats = tcgen05.make_tmem_copy(tmem_load_stats_atom, tOtStats) + thr_tmem_load_stats = tiled_tmem_load_stats.get_slice(thread_idx) + tTMEM_LOADtStats = thr_tmem_load_stats.partition_S(tOtStats) + tTMEM_LOADcStats = thr_tmem_load_stats.partition_D(tOcStats) + tTMEM_LOADrStats = cute.make_rmem_tensor(tTMEM_LOADcStats.shape, self.qk_acc_dtype) + cute.copy(tiled_tmem_load_stats, tTMEM_LOADtStats, tTMEM_LOADrStats) + + scale = scale_softmax_log2 * (tTMEM_LOADrStats[0] - tTMEM_LOADrStats[1]) + scale = cute.math.exp2(scale, fastmath=True) + stats_handle.release() + o_handle = mma_o_consumer.wait_and_advance() + for iter in cutlass.range(self.iterations_pv, unroll_full=True): + tOtO = tOtO_staged[(None, None), 0, 0, iter] + cO = cO_staged[None, None, iter] + tOtO_epi = cute.zipped_divide(tOtO, epi_tile) + cO_epi = cute.zipped_divide(cO, epi_tile) + tmem_load_atom = cute.make_copy_atom( + tcgen05.Ld32x32bOp(tcgen05.Repetition(16)), + self.pv_acc_dtype, + ) + tmem_tiled_load = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_epi) + thr_load = tmem_tiled_load.get_slice(thread_idx) + tmem_store_atom = cute.make_copy_atom( + tcgen05.St32x32bOp(tcgen05.Repetition(16)), + self.pv_acc_dtype, + ) + tmem_store_atom = tcgen05.make_tmem_copy(tmem_store_atom, tOtO_epi) + thr_store = tmem_store_atom.get_slice(thread_idx) + tTMEM_LOADtO = thr_load.partition_S(tOtO_epi) + tTMEM_LOADcO = thr_load.partition_D(cO_epi) + tTMEM_STOREtO = thr_store.partition_D(tOtO_epi) + tTMrO = cute.make_rmem_tensor_like( + cute.append( + cute.make_layout(tTMEM_LOADcO[None, 0, 0].shape), + cute.make_layout(2, stride=cute.size(tTMEM_LOADcO[None, 0, 0].shape)), + ), + self.pv_acc_dtype, + ) + tTMEM_LOADtO_0 = tTMEM_LOADtO[None, 0, 0] + cute.copy(tmem_tiled_load, tTMEM_LOADtO_0, tTMrO[None, 0]) + iter_num = cute.size(tTMEM_LOADtO, mode=[1]) + for i in cutlass.range(1, iter_num, unroll_full=True): + tTMEM_LOADtO_i = tTMEM_LOADtO[None, i, 0] + cute.copy(tmem_tiled_load, tTMEM_LOADtO_i, tTMrO[None, i % 2]) + for j in cutlass.range(0, cute.size(tTMrO, mode=[0]), 2, unroll_full=True): + tTMrO[j, (i - 1) % 2], tTMrO[j + 1, (i - 1) % 2] = cute.arch.mul_packed_f32x2( + (tTMrO[j, (i - 1) % 2], tTMrO[j + 1, (i - 1) % 2]), + (scale, scale), + ) + tTMEM_STOREtO_prev_i = tTMEM_STOREtO[None, i - 1, 0] + cute.copy(tmem_store_atom, tTMrO[None, (i - 1) % 2], tTMEM_STOREtO_prev_i) + + for j in cutlass.range(0, cute.size(tTMrO, mode=[0]), 2, unroll_full=True): + tTMrO[j, (iter_num - 1) % 2], tTMrO[j + 1, (iter_num - 1) % 2] = ( + cute.arch.mul_packed_f32x2( + ( + tTMrO[j, (iter_num - 1) % 2], + tTMrO[j + 1, (iter_num - 1) % 2], + ), + (scale, scale), + ) + ) + cute.copy( + tmem_store_atom, + tTMrO[None, (iter_num - 1) % 2], + tTMEM_STOREtO[None, iter_num - 1, 0], + ) + cute.arch.fence_view_async_tmem_store() + o_handle.release() + return mma_o_consumer, s_corr_consumer + + @cute.jit + def correction_epilog( + self, + value_args: Tuple, + sum_args: Tuple, + o_args: Tuple, + epi_tile: cute.Tile, + ) -> Tuple[pipeline.PipelineConsumer, pipeline.PipelineProducer]: + (seqlen_q, scale_output) = value_args + (sum_consumer, sSum) = sum_args + (mma_o_consumer, gO_staged, cO_staged, tOtO_staged) = o_args + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + sum_handle = sum_consumer.wait_and_advance() + row_sum = sSum[thread_idx] + cute.arch.fence_view_async_shared() + sum_handle.release() + row_sum_is_zero_or_nan = row_sum == 0.0 or row_sum != row_sum + scale = scale_output / row_sum if not row_sum_is_zero_or_nan else 0.0 + o_handle = mma_o_consumer.wait_and_advance() + for iter in cutlass.range(self.iterations_pv): + gO = gO_staged[None, None, iter] + cO = cO_staged[None, None, iter] + tOtO = tOtO_staged[(None, None), 0, 0, iter] + tOtO_epi = cute.zipped_divide(tOtO, epi_tile) + cO_epi = cute.zipped_divide(cO, epi_tile) + gO_epi = cute.zipped_divide(gO, epi_tile) + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + tmem_copy_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.pv_acc_dtype + ) + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_epi) + thr_tmem_load = tiled_tmem_load.get_slice(thread_idx) + tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_epi) + tTMEM_LOADgO = thr_tmem_load.partition_D(gO_epi) + tTMEM_LOADcO = thr_tmem_load.partition_D(cO_epi) + for i in cutlass.range(cute.size(tTMEM_LOADtO, mode=[1]), unroll_full=True): + tTMEM_LOADtO_i = tTMEM_LOADtO[None, i, 0] + tTMEM_LOADgO_i = tTMEM_LOADgO[None, i, 0] + tTMEM_LOADcO_i = tTMEM_LOADcO[None, i, 0] + tTMrO = cute.make_rmem_tensor(tTMEM_LOADcO[None, 0, i].shape, self.pv_acc_dtype) + cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO) + for j in cutlass.range(0, cute.size(tTMrO), 2, unroll_full=True): + tTMrO[j], tTMrO[j + 1] = cute.arch.mul_packed_f32x2( + (tTMrO[j], tTMrO[j + 1]), + (scale, scale), + ) + tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype) + o_vec = tTMrO.load() + tSMrO.store(o_vec.to(self.o_dtype)) + if cute.elem_less(tTMEM_LOADcO_i[0][0], seqlen_q): + cute.autovec_copy(tSMrO, tTMEM_LOADgO_i) + o_handle.release() + return mma_o_consumer, sum_consumer + + @cute.jit + def store_sum_max( + self, + row_max, + mLSE, + row_sum, + sSum, + sum_producer, + current_block_coord, + seqlen_q, + cum_seqlen_q, + cuseqlen_q, + scale_softmax, + ): + tidx, _, _ = cute.arch.thread_idx() + thread_idx = tidx % (self.threads_per_warp * len(self.softmax_warp_ids)) + sum_handle = sum_producer.acquire_and_advance() + sSum[thread_idx] = row_sum + cute.arch.fence_view_async_shared() + sum_handle.commit() + row_sum_is_zero_or_nan = row_sum == 0.0 or row_sum != row_sum + + if cutlass.const_expr(mLSE is not None): + q_idx = current_block_coord[0] * self.cta_tiler[0] + tidx + hb_idx = ( + (current_block_coord[2][0], Int32(0)) + if cutlass.const_expr(cum_seqlen_q is not None) + else current_block_coord[2] + ) + lse_value = ( + scale_softmax * row_max + cute.math.log(row_sum, fastmath=True) + if not row_sum_is_zero_or_nan + else -Float32.inf + ) + if cute.elem_less(q_idx, seqlen_q): + global_q_idx = ( + q_idx + cuseqlen_q if cutlass.const_expr(cum_seqlen_q is not None) else q_idx + ) + mLSE[global_q_idx, hb_idx] = lse_value + return sum_producer diff --git a/fastvideo-kernel/fa4/sm90_config_search.py b/fastvideo-kernel/fa4/sm90_config_search.py new file mode 100644 index 000000000..6c9584ea3 --- /dev/null +++ b/fastvideo-kernel/fa4/sm90_config_search.py @@ -0,0 +1,402 @@ +"""Search feasible SM90 fwd/bwd attention configs for given (head_dim, head_dim_v). + +Enumerates tile sizes, swap modes, atom layouts, and staging options. +Checks GMMA divisibility, register budget, and shared memory budget. + +Usage: + python flash_attn/cute/sm90_config_search.py --headdim 128 + python flash_attn/cute/sm90_config_search.py --mode fwd --headdim 192-128 + python flash_attn/cute/sm90_config_search.py --mode bwd --headdim 192 --tile-n 64,96 +""" + +import math + +# H100 hardware limits +SMEM_LIMIT = 224 * 1024 # 228 KB minus ~3 KB for LSE, dPsum, mbarriers +REG_LIMITS = {2: 216, 3: 128} # per-WG budget: 2WG=240-24, 3WG=160-32 +THREADS_PER_WG = 128 + + +def _divisors(n): + return [d for d in range(1, n + 1) if n % d == 0] + + +def _acc_regs(M, N, num_wg): + """Accumulator registers per thread per WG.""" + return M * N // (num_wg * THREADS_PER_WG) + + +def _check_mma(M, N, num_wg, atom_layout_m, swap_AB): + """Check MMA feasibility. Returns regs per WG, or None if infeasible. + + GMMA atom M=64. Swap exchanges (M, N) and atom layout. + Requires: M divisible by (atom_layout_m * 64), N by (atom_layout_n * 8). + """ + if swap_AB: + M, N = N, M + atom_layout_m = num_wg // atom_layout_m + atom_layout_n = num_wg // atom_layout_m + if M % (atom_layout_m * 64) != 0 or N % (atom_layout_n * 8) != 0: + return None + return _acc_regs(M, N, num_wg) + + +def _mma_traffic(M_eff, N_eff, K_red, num_wg, wg_n, is_rs=False): + """Total SMEM read traffic for one MMA (all WGs combined). + + num_instr = (M_eff / 64) * wg_n instructions total. + Each reads A(64, K_red) and B(N_eff/wg_n, K_red) from smem (bf16). + """ + num_instr = (M_eff // 64) * wg_n + A_per = 64 * K_red * 2 if not is_rs else 0 + B_per = (N_eff // wg_n) * K_red * 2 + return num_instr * (A_per + B_per) + + +# ============================================================================ +# Backward +# ============================================================================ + + +def _check_bwd_config( + hdim, + hdimv, + tile_m, + tile_n, + num_wg, + SdP_swapAB, + dKV_swapAB, + dQ_swapAB, + AtomLayoutMSdP, + AtomLayoutNdKV, + AtomLayoutMdQ, +): + reg_limit = REG_LIMITS[num_wg] + + # MMA feasibility + regs_SdP = _check_mma(tile_m, tile_n, num_wg, AtomLayoutMSdP, SdP_swapAB) + regs_dK = _check_mma(tile_n, hdim, num_wg, AtomLayoutNdKV, dKV_swapAB) + regs_dV = _check_mma(tile_n, hdimv, num_wg, AtomLayoutNdKV, dKV_swapAB) + regs_dQ = _check_mma(tile_m, hdim, num_wg, AtomLayoutMdQ, dQ_swapAB) + if any(r is None for r in (regs_SdP, regs_dK, regs_dV, regs_dQ)): + return None + + # Peak regs: max(S+dP, dQ) + dK + dV + total_regs = max(2 * regs_SdP, regs_dQ) + regs_dK + regs_dV + if total_regs > reg_limit: + return None + + # SMEM + mma_dkv_is_rs = ( + AtomLayoutMSdP == 1 and AtomLayoutNdKV == num_wg and SdP_swapAB and not dKV_swapAB + ) + Q_stage, PdS_stage = 2, 1 + + for dO_stage in (2, 1): + sQ = tile_m * hdim * 2 * Q_stage + sK = tile_n * hdim * 2 + sV = tile_n * hdimv * 2 + sdO = tile_m * hdimv * 2 * dO_stage + sPdS = tile_m * tile_n * 2 * PdS_stage + sP = sPdS if not mma_dkv_is_rs else 0 + sdQaccum = tile_m * hdim * 4 + smem = sQ + sK + sV + sdO + sP + sPdS + sdQaccum + if smem <= SMEM_LIMIT: + break + else: + return None + + # SMEM traffic + def _swap(a, b, s): + return (b, a) if s else (a, b) + + def _wg_n(al_m, s): + return al_m if s else num_wg // al_m + + M_s, N_s = _swap(tile_m, tile_n, SdP_swapAB) + wn_SdP = _wg_n(AtomLayoutMSdP, SdP_swapAB) + traffic_S = _mma_traffic(M_s, N_s, hdim, num_wg, wn_SdP) + traffic_dP = _mma_traffic(M_s, N_s, hdimv, num_wg, wn_SdP) + + wn_dKV = _wg_n(AtomLayoutNdKV, dKV_swapAB) + M_dv, N_dv = _swap(tile_n, hdimv, dKV_swapAB) + traffic_dV = _mma_traffic(M_dv, N_dv, tile_m, num_wg, wn_dKV, is_rs=mma_dkv_is_rs) + M_dk, N_dk = _swap(tile_n, hdim, dKV_swapAB) + traffic_dK = _mma_traffic(M_dk, N_dk, tile_m, num_wg, wn_dKV, is_rs=mma_dkv_is_rs) + + M_dq, N_dq = _swap(tile_m, hdim, dQ_swapAB) + wn_dQ = _wg_n(AtomLayoutMdQ, dQ_swapAB) + traffic_dQ = _mma_traffic(M_dq, N_dq, tile_n, num_wg, wn_dQ) + + traffic_P_store = tile_m * tile_n * 2 if not mma_dkv_is_rs else 0 + traffic_dS_store = tile_m * tile_n * 2 + traffic_dQ_smem = tile_m * hdim * 4 * 2 # store + TMA load + + smem_traffic = ( + traffic_S + + traffic_dP + + traffic_dV + + traffic_dK + + traffic_dQ + + traffic_P_store + + traffic_dS_store + + traffic_dQ_smem + ) + + return dict( + tile_m=tile_m, + tile_n=tile_n, + num_wg=num_wg, + Q_stage=Q_stage, + dO_stage=dO_stage, + PdS_stage=PdS_stage, + SdP_swapAB=SdP_swapAB, + dKV_swapAB=dKV_swapAB, + dQ_swapAB=dQ_swapAB, + AtomLayoutMSdP=AtomLayoutMSdP, + AtomLayoutNdKV=AtomLayoutNdKV, + AtomLayoutMdQ=AtomLayoutMdQ, + mma_dkv_is_rs=mma_dkv_is_rs, + regs_SdP=regs_SdP, + regs_dK=regs_dK, + regs_dV=regs_dV, + regs_dQ=regs_dQ, + total_regs=total_regs, + reg_limit=reg_limit, + smem_bytes=smem, + smem_kb=smem / 1024, + smem_traffic=smem_traffic, + smem_traffic_kb=smem_traffic / 1024, + smem_traffic_per_block=smem_traffic / (tile_m * tile_n), + ) + + +def find_feasible_bwd_configs( + head_dim, + head_dim_v=None, + tile_m_choices=(64, 80, 96, 112, 128), + tile_n_choices=(64, 80, 96, 112, 128), +): + if head_dim_v is None: + head_dim_v = head_dim + hdim = int(math.ceil(head_dim / 32) * 32) + hdimv = int(math.ceil(head_dim_v / 32) * 32) + + results = [] + for num_wg in (2, 3): + divs = _divisors(num_wg) + for tile_m in tile_m_choices: + for tile_n in tile_n_choices: + for SdP_swap in (False, True): + if (tile_n if SdP_swap else tile_m) % 64 != 0: + continue + for dKV_swap in (False, True): + if not dKV_swap and tile_n % 64 != 0: + continue + if dKV_swap and (hdim % 64 != 0 or hdimv % 64 != 0): + continue + for dQ_swap in (False, True): + if (hdim if dQ_swap else tile_m) % 64 != 0: + continue + for a1 in divs: + for a2 in divs: + for a3 in divs: + cfg = _check_bwd_config( + hdim, + hdimv, + tile_m, + tile_n, + num_wg, + SdP_swap, + dKV_swap, + dQ_swap, + a1, + a2, + a3, + ) + if cfg is not None: + results.append(cfg) + + results.sort(key=lambda c: (-c["tile_n"], -c["tile_m"], c["smem_traffic_per_block"])) + return results + + +def print_bwd_configs(configs, max_results=20): + if not configs: + print("No feasible configs found!") + return + n = min(len(configs), max_results) + print(f"Found {len(configs)} feasible configs (showing top {n}):\n") + hdr = ( + f"{'wg':>2} {'tm':>3} {'tn':>3} " + f"{'SdP':>3} {'dKV':>3} {'dQ':>3} " + f"{'aSdP':>4} {'adKV':>4} {'adQ':>4} " + f"{'Qs':>2} {'dOs':>3} " + f"{'rS':>3} {'rdK':>3} {'rdV':>3} {'rdQ':>3} {'tot':>4}/{'':<3} " + f"{'smem':>5} {'traffic':>7} {'tr/blk':>6}" + ) + print(hdr) + print("-" * len(hdr)) + B = lambda b: "T" if b else "F" + for c in configs[:max_results]: + print( + f"{c['num_wg']:>2} {c['tile_m']:>3} {c['tile_n']:>3} " + f"{B(c['SdP_swapAB']):>3} {B(c['dKV_swapAB']):>3} {B(c['dQ_swapAB']):>3} " + f"{c['AtomLayoutMSdP']:>4} {c['AtomLayoutNdKV']:>4} {c['AtomLayoutMdQ']:>4} " + f"{c['Q_stage']:>2} {c['dO_stage']:>3} " + f"{c['regs_SdP']:>3} {c['regs_dK']:>3} {c['regs_dV']:>3} {c['regs_dQ']:>3} " + f"{c['total_regs']:>4}/{c['reg_limit']:<3} " + f"{c['smem_kb']:>4.0f}K " + f"{c['smem_traffic_kb']:>6.0f}K " + f"{c['smem_traffic_per_block']:>6.1f}" + ) + + +# ============================================================================ +# Forward +# ============================================================================ + + +def _check_fwd_config(hdim, hdimv, tile_n, num_wg, pv_is_rs, overlap_wg): + reg_limit = REG_LIMITS[num_wg] + tile_m = num_wg * 64 + + if tile_n % 8 != 0: + return None + + regs_S = _acc_regs(tile_m, tile_n, num_wg) + regs_O = _acc_regs(tile_m, hdimv, num_wg) + regs_P = regs_S // 2 # bf16 = half of f32 + + if overlap_wg: + total_regs = regs_S + regs_P + regs_O + else: + total_regs = regs_S + regs_O + + if total_regs > reg_limit: + return None + + # SMEM: 1 stage Q, 2 stages K/V, O overlaps Q, sP if not RS + sQ = tile_m * hdim * 2 + sK = tile_n * hdim * 2 * 2 + sV = tile_n * hdimv * 2 * 2 + sO = tile_m * hdimv * 2 + sP = tile_m * tile_n * 2 if not pv_is_rs else 0 + smem = max(sQ, sO) + sK + sV + sP + if smem > SMEM_LIMIT: + return None + + # SMEM traffic: num_instr = num_wg (all WGs in M, wg_n=1) + traffic_S = num_wg * (64 * hdim * 2 + tile_n * hdim * 2) + A_pv = 64 * tile_n * 2 if not pv_is_rs else 0 + traffic_O = num_wg * (A_pv + hdimv * tile_n * 2) + traffic_P_store = tile_m * tile_n * 2 if not pv_is_rs else 0 + smem_traffic = traffic_S + traffic_O + traffic_P_store + + return dict( + tile_m=tile_m, + tile_n=tile_n, + num_wg=num_wg, + pv_is_rs=pv_is_rs, + overlap_wg=overlap_wg, + regs_S=regs_S, + regs_O=regs_O, + regs_P=regs_P, + total_regs=total_regs, + reg_limit=reg_limit, + smem_bytes=smem, + smem_kb=smem / 1024, + smem_traffic=smem_traffic, + smem_traffic_kb=smem_traffic / 1024, + smem_traffic_per_block=smem_traffic / (tile_m * tile_n), + ) + + +def find_feasible_fwd_configs( + head_dim, head_dim_v=None, tile_n_choices=(64, 80, 96, 112, 128, 144, 160, 176, 192) +): + if head_dim_v is None: + head_dim_v = head_dim + hdim = int(math.ceil(head_dim / 32) * 32) + hdimv = int(math.ceil(head_dim_v / 32) * 32) + + results = [] + for num_wg in (2, 3): + for tile_n in tile_n_choices: + for pv_is_rs in (True, False): + for overlap_wg in (True, False): + cfg = _check_fwd_config(hdim, hdimv, tile_n, num_wg, pv_is_rs, overlap_wg) + if cfg is not None: + results.append(cfg) + + results.sort(key=lambda c: (-c["tile_n"], c["smem_traffic_per_block"])) + return results + + +def print_fwd_configs(configs, max_results=20): + if not configs: + print("No feasible configs found!") + return + n = min(len(configs), max_results) + print(f"Found {len(configs)} feasible configs (showing top {n}):\n") + hdr = ( + f"{'wg':>2} {'tm':>3} {'tn':>3} " + f"{'RS':>2} {'olap':>4} " + f"{'rS':>3} {'rP':>3} {'rO':>3} {'tot':>4}/{'':<3} " + f"{'smem':>5} {'traffic':>7} {'tr/blk':>6}" + ) + print(hdr) + print("-" * len(hdr)) + B = lambda b: "T" if b else "F" + for c in configs[:max_results]: + print( + f"{c['num_wg']:>2} {c['tile_m']:>3} {c['tile_n']:>3} " + f"{B(c['pv_is_rs']):>2} {B(c['overlap_wg']):>4} " + f"{c['regs_S']:>3} {c['regs_P']:>3} {c['regs_O']:>3} " + f"{c['total_regs']:>4}/{c['reg_limit']:<3} " + f"{c['smem_kb']:>4.0f}K " + f"{c['smem_traffic_kb']:>6.0f}K " + f"{c['smem_traffic_per_block']:>6.1f}" + ) + + +# ============================================================================ +# CLI +# ============================================================================ + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Search feasible SM90 MMA configs") + parser.add_argument("--mode", choices=["fwd", "bwd", "both"], default="both") + parser.add_argument( + "--headdim", type=str, default="128", help="Head dim, or hdim-hdimv (e.g. 192-128)" + ) + parser.add_argument("--tile-m", type=str, default="64,80,96,112,128", help="Bwd tile_m choices") + parser.add_argument( + "--tile-n", + type=str, + default=None, + help="tile_n choices (default: fwd up to 192, bwd up to 128)", + ) + parser.add_argument("-n", "--num-results", type=int, default=30) + args = parser.parse_args() + + parts = args.headdim.split("-") + hdim = int(parts[0]) + hdimv = int(parts[1]) if len(parts) > 1 else hdim + + TN_FWD = "64,80,96,112,128,144,160,176,192" + TN_BWD = "64,80,96,112,128" + + if args.mode in ("fwd", "both"): + tn = tuple(int(x) for x in (args.tile_n or TN_FWD).split(",")) + print(f"=== FWD configs: hdim={hdim}, hdimv={hdimv} ===\n") + print_fwd_configs(find_feasible_fwd_configs(hdim, hdimv, tn), args.num_results) + print() + + if args.mode in ("bwd", "both"): + tm = tuple(int(x) for x in args.tile_m.split(",")) + tn = tuple(int(x) for x in (args.tile_n or TN_BWD).split(",")) + print(f"=== BWD configs: hdim={hdim}, hdimv={hdimv} ===\n") + print_bwd_configs(find_feasible_bwd_configs(hdim, hdimv, tm, tn), args.num_results) diff --git a/fastvideo-kernel/fa4/softmax.py b/fastvideo-kernel/fa4/softmax.py new file mode 100644 index 000000000..231e469d4 --- /dev/null +++ b/fastvideo-kernel/fa4/softmax.py @@ -0,0 +1,740 @@ +# Copyright (c) 2025, Tri Dao. + +import math +import operator +from typing import Tuple +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Boolean + +from quack import layout_utils +import flash_attn.cute.utils as utils +from quack.cute_dsl_utils import ParamsBase +from flash_attn.cute.seqlen_info import SeqlenInfoQK +from flash_attn.cute.utils import AuxData + + +@cute.jit +def call_score_mod( + score_mod: cutlass.Constexpr, + score, + batch_idx, + head_idx, + q_idx, + kv_idx, + seqlen_info, + aux_data: AuxData, +): + aux_tensors = aux_data.tensors if aux_data.tensors is not None else () + # Compatibility shim for pre-aux_scalars score_mod callables. + if cutlass.const_expr(aux_data.scalars is not None): + return score_mod( + score, + batch_idx, + head_idx, + q_idx=q_idx, + kv_idx=kv_idx, + seqlen_info=seqlen_info, + aux_tensors=aux_tensors, + aux_scalars=aux_data.scalars, + ) + return score_mod( + score, + batch_idx, + head_idx, + q_idx=q_idx, + kv_idx=kv_idx, + seqlen_info=seqlen_info, + aux_tensors=aux_tensors, + ) + + +@cute.jit +def call_score_mod_bwd( + score_mod_bwd: cutlass.Constexpr, + grad, + score, + batch_idx, + head_idx, + q_idx, + kv_idx, + seqlen_info, + aux_data: AuxData, +): + aux_tensors = aux_data.tensors if aux_data.tensors is not None else () + # Compatibility shim for pre-aux_scalars score_mod_bwd callables. + if cutlass.const_expr(aux_data.scalars is not None): + return score_mod_bwd( + grad, + score, + batch_idx, + head_idx, + q_idx=q_idx, + kv_idx=kv_idx, + seqlen_info=seqlen_info, + aux_tensors=aux_tensors, + aux_scalars=aux_data.scalars, + ) + return score_mod_bwd( + grad, + score, + batch_idx, + head_idx, + q_idx=q_idx, + kv_idx=kv_idx, + seqlen_info=seqlen_info, + aux_tensors=aux_tensors, + ) + + +@dataclass +class Softmax(ParamsBase): + scale_log2: Float32 + num_rows: cutlass.Constexpr[int] + row_max: cute.Tensor + row_sum: cute.Tensor + arch: cutlass.Constexpr[int] = 80 + softmax_scale: Float32 | None = None + + @staticmethod + def create( + scale_log2: Float32, + num_rows: cutlass.Constexpr[int], + arch: cutlass.Constexpr[int] = 80, + softmax_scale: Float32 | None = None, + ): + row_max = cute.make_rmem_tensor(num_rows, Float32) + row_sum = cute.make_rmem_tensor(num_rows, Float32) + return Softmax(scale_log2, num_rows, row_max, row_sum, arch, softmax_scale) + + def reset(self) -> None: + self.row_max.fill(-Float32.inf) + self.row_sum.fill(0.0) + + def _compute_row_max( + self, acc_S_row: cute.TensorSSA, init_val: float | Float32 | None = None + ) -> Float32: + return utils.fmax_reduce(acc_S_row, init_val, arch=self.arch) + + def _compute_row_sum( + self, acc_S_row_exp: cute.TensorSSA, init_val: float | Float32 | None = None + ) -> Float32: + return utils.fadd_reduce(acc_S_row_exp, init_val, arch=self.arch) + + @cute.jit + def online_softmax( + self, + acc_S: cute.Tensor, + is_first: cutlass.Constexpr[bool] = False, + check_inf: cutlass.Constexpr[bool] = True, + ) -> cute.Tensor: + """Apply online softmax and return the row_scale to rescale O. + + :param acc_S: acc_S tensor + :type acc_S: cute.Tensor + :param is_first: is first n_block + :type is_first: cutlass.Constexpr + """ + # Change acc_S to M,N layout view. + acc_S_mn = layout_utils.reshape_acc_to_mn(acc_S) + row_scale = cute.make_fragment_like(self.row_max, Float32) + + row_max = self.row_max + row_sum = self.row_sum + scale_log2 = self.scale_log2 + arch = self.arch + + # Each iteration processes one row of acc_S + for r in cutlass.range(cute.size(row_max), unroll_full=True): + acc_S_row = acc_S_mn[r, None].load() # (n_block_size) + + row_max_cur = utils.fmax_reduce( + acc_S_row, + init_val=row_max[r] if cutlass.const_expr(not is_first) else None, + arch=arch, + ) + + row_max_cur = cute.arch.warp_reduction_max(row_max_cur, threads_in_group=4) + # Update row_max before changing row_max_cur to safe value for -inf + row_max_prev = row_max[r] + row_max[r] = row_max_cur + + if cutlass.const_expr(check_inf): + row_max_cur = 0.0 if row_max_cur == -Float32.inf else row_max_cur + + if cutlass.const_expr(is_first): + row_max_cur_scaled = row_max_cur * scale_log2 + acc_S_row_exp = cute.math.exp2( + acc_S_row * scale_log2 - row_max_cur_scaled, fastmath=True + ) + acc_S_row_sum = utils.fadd_reduce(acc_S_row_exp, init_val=None, arch=arch) + row_scale[r] = 1.0 + else: + row_max_cur_scaled = row_max_cur * scale_log2 + acc_S_row_exp = cute.math.exp2( + acc_S_row * scale_log2 - row_max_cur_scaled, fastmath=True + ) + # row_scale[r] = cute.math.exp2(row_max_prev * self.scale_log2 - row_max_cur_scaled) + row_scale[r] = cute.math.exp2( + (row_max_prev - row_max_cur) * scale_log2, fastmath=True + ) + acc_S_row_sum = utils.fadd_reduce( + acc_S_row_exp, init_val=row_sum[r] * row_scale[r], arch=arch + ) + + row_sum[r] = acc_S_row_sum + acc_S_mn[r, None].store(acc_S_row_exp) + + return row_scale + + @cute.jit + def finalize( + self, final_scale: Float32 = 1.0, sink_val: Float32 | cute.Tensor | None = None + ) -> cute.Tensor: + """Finalize the online softmax by computing the scale and logsumexp.""" + if cutlass.const_expr(sink_val is not None and isinstance(sink_val, cute.Tensor)): + assert cute.size(sink_val) == cute.size(self.row_sum) + row_sum = self.row_sum + row_max = self.row_max + scale_log2 = self.scale_log2 + + # quad reduction for row_sum as we didn't do it during each iteration of online softmax + row_sum.store(utils.warp_reduce(row_sum.load(), operator.add, width=4)) + row_scale = cute.make_fragment_like(row_max, Float32) + + for r in cutlass.range(cute.size(row_sum), unroll_full=True): + if cutlass.const_expr(sink_val is not None): + sink_val_cur = sink_val if not isinstance(sink_val, cute.Tensor) else sink_val[r] + LOG2_E = math.log2(math.e) + row_sum[r] += cute.math.exp2( + sink_val_cur * LOG2_E - row_max[r] * scale_log2, fastmath=True + ) + + # if row_sum is zero or nan, set acc_O_mn_row to 1.0 + acc_O_mn_row_is_zero_or_nan = row_sum[r] == 0.0 or row_sum[r] != row_sum[r] + row_scale[r] = ( + cute.arch.rcp_approx(row_sum[r] if not acc_O_mn_row_is_zero_or_nan else 1.0) + ) * final_scale + row_sum_cur = row_sum[r] + LN2 = math.log(2.0) + row_sum[r] = ( + (row_max[r] * scale_log2 + cute.math.log2(row_sum_cur, fastmath=True)) * LN2 + if not acc_O_mn_row_is_zero_or_nan + else -Float32.inf + ) + return row_scale + + @cute.jit + def rescale_O(self, acc_O: cute.Tensor, row_scale: cute.Tensor) -> None: + """Scale each row of acc_O by the given scale tensor. + :param acc_O: input tensor + :type acc_O: cute.Tensor + :param row_scale: row_scale tensor + :type row_scale: cute.Tensor + """ + acc_O_mn = layout_utils.reshape_acc_to_mn(acc_O) + assert cute.size(row_scale) == cute.size(acc_O_mn, mode=[0]) + for r in cutlass.range(cute.size(row_scale), unroll_full=True): + acc_O_mn[r, None].store(acc_O_mn[r, None].load() * row_scale[r]) + + +@dataclass +class SoftmaxSm100(Softmax): + rescale_threshold: cutlass.Constexpr[float] = 0.0 + max_offset: cutlass.Constexpr[int] = 0 + row_pair_stride: cutlass.Constexpr[int] = 0 + + @staticmethod + def create( + scale_log2: Float32, + rescale_threshold: cutlass.Constexpr[float] = 0.0, + softmax_scale: Float32 | None = None, + max_offset: cutlass.Constexpr[int] = 0, + row_pair_stride: cutlass.Constexpr[int] = 0, + ): + num_rows = 1 + arch = 100 + row_max = cute.make_rmem_tensor(num_rows, Float32) + row_sum = cute.make_rmem_tensor(num_rows, Float32) + return SoftmaxSm100( + scale_log2, + num_rows, + row_max, + row_sum, + arch, + softmax_scale, + rescale_threshold=rescale_threshold, + max_offset=max_offset, + row_pair_stride=row_pair_stride, + ) + + @cute.jit + def compute_row_max_local(self, acc_S_row: cute.TensorSSA, is_first: Boolean) -> Float32: + if cutlass.const_expr(is_first): + row_max_new = self._compute_row_max(acc_S_row) + else: + row_max_old = self.row_max[0] + row_max_new = self._compute_row_max(acc_S_row, init_val=row_max_old) + return row_max_new + + @cute.jit + def update_row_max_from_local( + self, + row_max_new: Float32, + is_first: Boolean, + ) -> Tuple[Float32, Float32]: + if cutlass.const_expr(self.row_pair_stride != 0): + row_max_new = utils.fmax( + row_max_new, + cute.arch.shuffle_sync_bfly( + row_max_new, + offset=self.row_pair_stride, + mask=-1, + mask_and_clamp=31, + ), + ) + if cutlass.const_expr(is_first): + row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0 + acc_scale = 0.0 + else: + row_max_old = self.row_max[0] + row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0 + acc_scale_ = (row_max_old - row_max_safe) * self.scale_log2 + acc_scale = cute.math.exp2(acc_scale_) + if cutlass.const_expr(self.rescale_threshold > 0.0): + if acc_scale_ >= -self.rescale_threshold: + row_max_new = row_max_old + row_max_safe = row_max_old + acc_scale = 1.0 + self.row_max[0] = row_max_new + return row_max_safe, acc_scale + + @cute.jit + def update_row_max(self, acc_S_row: cute.TensorSSA, is_first: int) -> Tuple[Float32, Float32]: + if cutlass.const_expr(is_first): + row_max_new = self._compute_row_max(acc_S_row) + if cutlass.const_expr(self.row_pair_stride != 0): + row_max_new = utils.fmax( + row_max_new, + cute.arch.shuffle_sync_bfly( + row_max_new, + offset=self.row_pair_stride, + mask=-1, + mask_and_clamp=31, + ), + ) + row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0 + acc_scale = 0.0 + else: + row_max_old = self.row_max[0] + row_max_new = self._compute_row_max(acc_S_row, init_val=row_max_old) + if cutlass.const_expr(self.row_pair_stride != 0): + row_max_new = utils.fmax( + row_max_new, + cute.arch.shuffle_sync_bfly( + row_max_new, + offset=self.row_pair_stride, + mask=-1, + mask_and_clamp=31, + ), + ) + row_max_safe = row_max_new if row_max_new != -cutlass.Float32.inf else 0.0 + acc_scale_ = (row_max_old - row_max_safe) * self.scale_log2 + acc_scale = cute.math.exp2(acc_scale_, fastmath=True) + if cutlass.const_expr(self.rescale_threshold > 0.0): + if acc_scale_ >= -self.rescale_threshold: + row_max_new = row_max_old + row_max_safe = row_max_old + acc_scale = 1.0 + self.row_max[0] = row_max_new + return row_max_safe, acc_scale + + def update_row_sum( + self, acc_S_row_exp: cute.TensorSSA, row_scale: Float32, is_first: int = False + ) -> None: + if cutlass.const_expr(self.row_pair_stride != 0): + local_sum = self._compute_row_sum(acc_S_row_exp) + row_sum_new = local_sum + cute.arch.shuffle_sync_bfly( + local_sum, + offset=self.row_pair_stride, + mask=-1, + mask_and_clamp=31, + ) + if cutlass.const_expr(not is_first): + row_sum_new = row_sum_new + self.row_sum[0] * row_scale + self.row_sum[0] = row_sum_new + else: + init_val = self.row_sum[0] * row_scale if cutlass.const_expr(not is_first) else None + self.row_sum[0] = self._compute_row_sum(acc_S_row_exp, init_val=init_val) + + @cute.jit + def scale_subtract_rowmax( + self, + acc_S_row: cute.Tensor, + row_max: Float32, + ): + assert cute.size(acc_S_row.shape) % 2 == 0, "acc_S_row must have an even number of elements" + row_max_scaled = row_max * self.scale_log2 + max_offset = Float32(self.max_offset) + bias = max_offset - row_max_scaled + for i in cutlass.range(0, cute.size(acc_S_row.shape), 2, unroll_full=True): + acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2( + (acc_S_row[i], acc_S_row[i + 1]), + (self.scale_log2, self.scale_log2), + (bias, bias), + ) + + @cute.jit + def apply_exp2_convert( + self, + acc_S_row: cute.Tensor, + acc_S_row_converted: cute.Tensor, + ex2_emu_freq: cutlass.Constexpr[int] = 0, + ex2_emu_res: cutlass.Constexpr[int] = 4, + ex2_emu_start_frg: cutlass.Constexpr[int] = 0, + ): + assert cute.size(acc_S_row.shape) % 2 == 0, "acc_S_row must have an even number of elements" + frg_tile = 32 + assert frg_tile % 2 == 0 + frg_cnt = cute.size(acc_S_row) // frg_tile + assert cute.size(acc_S_row) % frg_tile == 0 + acc_S_row_frg = cute.logical_divide(acc_S_row, cute.make_layout(frg_tile)) + acc_S_row_converted_frg = cute.logical_divide( + acc_S_row_converted, cute.make_layout(frg_tile) + ) + for j in cutlass.range_constexpr(frg_cnt): + for k in cutlass.range_constexpr(0, cute.size(acc_S_row_frg, mode=[0]), 2): + # acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True) + # acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True) + if cutlass.const_expr(ex2_emu_freq == 0): + acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True) + acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True) + else: + if cutlass.const_expr( + k % ex2_emu_freq < ex2_emu_freq - ex2_emu_res + or j >= frg_cnt - 1 + or j < ex2_emu_start_frg + ): + acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True) + acc_S_row_frg[k + 1, j] = cute.math.exp2( + acc_S_row_frg[k + 1, j], fastmath=True + ) + else: + # acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] = utils.e2e_asm2(acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j]) + acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] = utils.ex2_emulation_2( + acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] + ) + acc_S_row_converted_frg[None, j].store( + acc_S_row_frg[None, j].load().to(acc_S_row_converted.element_type) + ) + + @cute.jit + def scale_apply_exp2_convert( + self, + acc_S_row: cute.Tensor, + row_max: Float32, + acc_S_row_converted: cute.Tensor, + ): + assert cute.size(acc_S_row.shape) % 2 == 0, "acc_S_row must have an even number of elements" + minus_row_max_scaled = -row_max * self.scale_log2 + for i in cutlass.range_constexpr(0, cute.size(acc_S_row.shape), 2): + acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2( + (acc_S_row[i], acc_S_row[i + 1]), + (self.scale_log2, self.scale_log2), + (minus_row_max_scaled, minus_row_max_scaled), + ) + + # for i in cutlass.range_constexpr(0, cute.size(acc_S_row.shape), 2): + # acc_S_row[i], acc_S_row[i + 1] = cute.arch.fma_packed_f32x2( + # (acc_S_row[i], acc_S_row[i + 1]), + # (self.scale_log2, self.scale_log2), + # (minus_row_max_scaled, minus_row_max_scaled), + # ) + # acc_S_row[i] = cute.math.exp2(acc_S_row[i], fastmath=True) + # acc_S_row[i + 1] = cute.math.exp2(acc_S_row[i + 1], fastmath=True) + + frg_tile = 32 + assert frg_tile % 2 == 0 + frg_cnt = cute.size(acc_S_row) // frg_tile + assert cute.size(acc_S_row) % frg_tile == 0 + acc_S_row_frg = cute.logical_divide(acc_S_row, cute.make_layout(frg_tile)) + acc_S_row_converted_frg = cute.logical_divide( + acc_S_row_converted, cute.make_layout(frg_tile) + ) + for j in cutlass.range_constexpr(frg_cnt): + for k in cutlass.range_constexpr(0, cute.size(acc_S_row_frg, mode=[0]), 2): + # acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j] = ( + # cute.arch.fma_packed_f32x2( + # (acc_S_row_frg[k, j], acc_S_row_frg[k + 1, j]), + # (self.scale_log2, self.scale_log2), + # (minus_row_max_scaled, minus_row_max_scaled), + # ) + # ) + # acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True) + # acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True) + acc_S_row_frg[k, j] = cute.math.exp2(acc_S_row_frg[k, j], fastmath=True) + acc_S_row_frg[k + 1, j] = cute.math.exp2(acc_S_row_frg[k + 1, j], fastmath=True) + acc_S_row_converted_frg[None, j].store( + acc_S_row_frg[None, j].load().to(acc_S_row_converted.element_type) + ) + + +@cute.jit +def floor_if_packed( + q_idx, + qhead_per_kvhead: cutlass.Constexpr[int], +) -> cute.Tensor: + """Convert q_idx to packed format for Pack-GQA.""" + if cutlass.const_expr(qhead_per_kvhead == 1): + return q_idx + return q_idx // qhead_per_kvhead + + +@cute.jit +def apply_score_mod_inner( + score_tensor, + index_tensor, + score_mod: cutlass.Constexpr, + batch_idx, + head_idx, + softmax_scale, + vec_size: cutlass.Constexpr, + qk_acc_dtype: cutlass.Constexpr, + aux_data: AuxData, + fastdiv_mods, + seqlen_info: SeqlenInfoQK, + constant_q_idx: cutlass.Constexpr, + qhead_per_kvhead: cutlass.Constexpr[int] = 1, + transpose_indices: cutlass.Constexpr[bool] = False, +): + """Shared implementation for applying score modification. + + Args: + score_tensor: The scores to modify (acc_S for flash_fwd, tSrS_t2r for sm100) + index_tensor: Index positions (tScS for flash_fwd, tScS_t2r for sm100) + score_mod: The score modification function to apply + batch_idx: Batch index + head_idx: Head index + softmax_scale: Scale to apply + vec_size: Vector size for processing elements + qk_acc_dtype: Data type for accumulator + aux_tensors: Optional aux_tensors for FlexAttention + aux_scalars: Optional runtime scalar captures for FlexAttention + fastdiv_mods: Tuple of (seqlen_q_divmod, seqlen_k_divmod) for wrapping + seqlen_info: Sequence length info + constant_q_idx: If provided, use this constant for all q_idx values + If None, compute q_idx per-element + qhead_per_kvhead_packgqa: Pack-GQA replication factor. Divide q_idx by this + when greater than 1 so score mods see logical heads. + transpose_indices: If True, swap q_idx/kv_idx in index_tensor (for bwd kernel where S is transposed) + """ + # Index positions in the index_tensor tuple + # Forward: index_tensor[...][0] = q_idx, index_tensor[...][1] = kv_idx + # Backward (transposed): index_tensor[...][0] = kv_idx, index_tensor[...][1] = q_idx + if cutlass.const_expr(transpose_indices): + q_idx_pos = cutlass.const_expr(1) + kv_idx_pos = cutlass.const_expr(0) + else: + q_idx_pos = cutlass.const_expr(0) + kv_idx_pos = cutlass.const_expr(1) + + n_vals = cutlass.const_expr(cute.size(score_tensor.shape)) + score_vec = cute.make_rmem_tensor(vec_size, qk_acc_dtype) + kv_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) + + # SSA values for batch (constant across all elements) + batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32).broadcast_to((vec_size,)) + + # Handle q_idx based on whether it's constant + q_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) + + # For Pack-GQA with non-constant q_idx, we need per-element head indices + # since a thread may process multiple query head indices + if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): + head_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) + + for i in cutlass.range(0, n_vals, vec_size, unroll_full=True): + for j in cutlass.range(vec_size, unroll_full=True): + score_vec[j] = score_tensor[i + j] * softmax_scale + + # Extract head offset from packed q_idx for Pack-GQA + if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): + q_idx_packed = index_tensor[i + j][q_idx_pos] + # Building up the logical q_head idx: final_q_head = kv_head * qhead_per_kvhead + (q_physical % qhead_per_kvhead) + q_idx_logical = q_idx_packed // qhead_per_kvhead + head_offset = q_idx_packed - q_idx_logical * qhead_per_kvhead + head_idx_vec[j] = head_idx * qhead_per_kvhead + head_offset + + # If we will do loads we mod, in order to not read OOB + if cutlass.const_expr(aux_data.tensors is not None and fastdiv_mods is not None): + if cutlass.const_expr(constant_q_idx is None): + seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods + q_idx_floored = floor_if_packed( + index_tensor[i + j][q_idx_pos], qhead_per_kvhead + ) + _, q_idx_wrapped = divmod(q_idx_floored, seqlen_q_divmod) + q_idx_vec[j] = q_idx_wrapped + else: + _, seqlen_k_divmod = fastdiv_mods + + _, kv_idx_wrapped = divmod(index_tensor[i + j][kv_idx_pos], seqlen_k_divmod) + kv_idx_vec[j] = kv_idx_wrapped + else: + # No bounds checking - direct indexing + if constant_q_idx is None: + q_idx_vec[j] = floor_if_packed(index_tensor[i + j][q_idx_pos], qhead_per_kvhead) + kv_idx_vec[j] = index_tensor[i + j][kv_idx_pos] + + # Convert to SSA for score_mod call + score_ssa = score_vec.load() + kv_idx_ssa = kv_idx_vec.load() + if cutlass.const_expr(constant_q_idx is None): + q_idx_ssa = q_idx_vec.load() + else: + # NB we do not apply Pack-GQA division here, as constant_q_idx is assumed to already be logical + q_idx_const = constant_q_idx + q_idx_ssa = utils.scalar_to_ssa(q_idx_const, cutlass.Int32).broadcast_to((vec_size,)) + + # Compute head_idx_ssa: per-element for Pack-GQA with non-constant q_idx, constant otherwise + if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): + head_idx_ssa = head_idx_vec.load() + else: + head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32).broadcast_to((vec_size,)) + + post_mod_scores = call_score_mod( + score_mod, + score_ssa, + batch_idx_ssa, + head_idx_ssa, + q_idx_ssa, + kv_idx_ssa, + seqlen_info, + aux_data, + ) + + # Write back modified scores + score_vec.store(post_mod_scores) + for j in cutlass.range(vec_size, unroll_full=True): + score_tensor[i + j] = score_vec[j] + + +@cute.jit +def apply_score_mod_bwd_inner( + grad_tensor, + score_tensor, + index_tensor, + score_mod_bwd: cutlass.Constexpr, + batch_idx, + head_idx, + softmax_scale, + vec_size: cutlass.Constexpr, + qk_acc_dtype: cutlass.Constexpr, + aux_data: AuxData, + fastdiv_mods, + seqlen_info, + constant_q_idx: cutlass.Constexpr, + qhead_per_kvhead: cutlass.Constexpr[int] = 1, + transpose_indices: cutlass.Constexpr[bool] = False, +): + """Apply backward score modification (joint graph). + + Args: + grad_tensor: in/out: dlogits rewritten in-place with d(scaled_scores) + score_tensor: pre-mod scores (unscaled QK tile), scaled by softmax_scale internally + index_tensor: Index positions (same as forward) + score_mod_bwd: The backward score modification function (joint graph) + batch_idx: Batch index + head_idx: Head index + softmax_scale: Scale to apply to score_tensor + vec_size: Vector size for processing elements + qk_acc_dtype: Data type for accumulator + aux_tensors: Optional aux_tensors for FlexAttention + aux_scalars: Optional runtime scalar captures for FlexAttention + fastdiv_mods: Tuple of (seqlen_q_divmod, seqlen_k_divmod) for wrapping + seqlen_info: Sequence length info + constant_q_idx: If provided, use this constant for all q_idx values + qhead_per_kvhead: Pack-GQA replication factor + transpose_indices: If True, swap q_idx/kv_idx in index_tensor + """ + # Index positions in the index_tensor tuple + # Forward: index_tensor[...][0] = q_idx, index_tensor[...][1] = kv_idx + # Backward (transposed): index_tensor[...][0] = kv_idx, index_tensor[...][1] = q_idx + if cutlass.const_expr(transpose_indices): + q_idx_pos = cutlass.const_expr(1) + kv_idx_pos = cutlass.const_expr(0) + else: + q_idx_pos = cutlass.const_expr(0) + kv_idx_pos = cutlass.const_expr(1) + n_vals = cutlass.const_expr(cute.size(grad_tensor.shape)) + grad_vec = cute.make_rmem_tensor(vec_size, qk_acc_dtype) + score_vec = cute.make_rmem_tensor(vec_size, qk_acc_dtype) + kv_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) + batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32).broadcast_to((vec_size,)) + q_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) + + # For Pack-GQA with non-constant q_idx, we need per-element head indices + if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): + head_idx_vec = cute.make_rmem_tensor(vec_size, cutlass.Int32) + + for i in cutlass.range(0, n_vals, vec_size, unroll_full=True): + for j in cutlass.range(vec_size, unroll_full=True): + grad_vec[j] = grad_tensor[i + j] + # Scale score so joint graph sees same value as forward score_mod + score_vec[j] = score_tensor[i + j] * softmax_scale + + if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): + q_idx_packed = index_tensor[i + j][q_idx_pos] + q_idx_logical = q_idx_packed // qhead_per_kvhead + head_offset = q_idx_packed - q_idx_logical * qhead_per_kvhead + head_idx_vec[j] = head_idx * qhead_per_kvhead + head_offset + + if cutlass.const_expr(aux_data.tensors is not None and fastdiv_mods is not None): + if cutlass.const_expr(constant_q_idx is None): + seqlen_q_divmod, seqlen_k_divmod = fastdiv_mods + q_idx_floored = floor_if_packed( + index_tensor[i + j][q_idx_pos], qhead_per_kvhead + ) + _, q_idx_wrapped = divmod(q_idx_floored, seqlen_q_divmod) + q_idx_vec[j] = q_idx_wrapped + else: + _, seqlen_k_divmod = fastdiv_mods + + _, kv_idx_wrapped = divmod(index_tensor[i + j][kv_idx_pos], seqlen_k_divmod) + kv_idx_vec[j] = kv_idx_wrapped + else: + # No bounds checking - direct indexing + if constant_q_idx is None: + q_idx_vec[j] = floor_if_packed(index_tensor[i + j][q_idx_pos], qhead_per_kvhead) + kv_idx_vec[j] = index_tensor[i + j][kv_idx_pos] + + grad_ssa = grad_vec.load() + score_ssa = score_vec.load() + kv_idx_ssa = kv_idx_vec.load() + + if cutlass.const_expr(constant_q_idx is None): + q_idx_ssa = q_idx_vec.load() + else: + q_idx_ssa = utils.scalar_to_ssa(constant_q_idx, cutlass.Int32).broadcast_to((vec_size,)) + + if cutlass.const_expr(qhead_per_kvhead > 1 and constant_q_idx is None): + head_idx_ssa = head_idx_vec.load() + else: + head_idx_ssa = utils.scalar_to_ssa(head_idx, cutlass.Int32).broadcast_to((vec_size,)) + + grad_out_ssa = call_score_mod_bwd( + score_mod_bwd, + grad_ssa, + score_ssa, + batch_idx_ssa, + head_idx_ssa, + q_idx_ssa, + kv_idx_ssa, + seqlen_info, + aux_data, + ) + + grad_vec.store(grad_out_ssa) + for j in cutlass.range(vec_size, unroll_full=True): + grad_tensor[i + j] = grad_vec[j] diff --git a/fastvideo-kernel/fa4/testing.py b/fastvideo-kernel/fa4/testing.py new file mode 100644 index 000000000..e6b2cf20d --- /dev/null +++ b/fastvideo-kernel/fa4/testing.py @@ -0,0 +1,487 @@ +import math +from contextlib import nullcontext +from functools import wraps +from typing import Optional + +import torch +import torch.nn.functional as F +from einops import rearrange, repeat +from torch._guards import active_fake_mode +from torch._subclasses.fake_tensor import FakeTensorMode + + +class IndexFirstAxis(torch.autograd.Function): + @staticmethod + def forward(ctx, input, indices): + ctx.save_for_backward(indices) + assert input.ndim >= 2 + ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:] + second_dim = other_shape.numel() + return torch.gather( + rearrange(input, "b ... -> b (...)"), + 0, + repeat(indices, "z -> z d", d=second_dim), + ).reshape(-1, *other_shape) + + @staticmethod + def backward(ctx, grad_output): + (indices,) = ctx.saved_tensors + assert grad_output.ndim >= 2 + other_shape = grad_output.shape[1:] + grad_output = rearrange(grad_output, "b ... -> b (...)") + grad_input = torch.zeros( + [ctx.first_axis_dim, grad_output.shape[1]], + device=grad_output.device, + dtype=grad_output.dtype, + ) + grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output) + return grad_input.reshape(ctx.first_axis_dim, *other_shape), None + + +index_first_axis = IndexFirstAxis.apply + + +class IndexPutFirstAxis(torch.autograd.Function): + @staticmethod + def forward(ctx, values, indices, first_axis_dim): + ctx.save_for_backward(indices) + assert indices.ndim == 1 + assert values.ndim >= 2 + output = torch.zeros( + first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype + ) + output[indices] = values + return output + + @staticmethod + def backward(ctx, grad_output): + (indices,) = ctx.saved_tensors + grad_values = grad_output[indices] + return grad_values, None, None + + +index_put_first_axis = IndexPutFirstAxis.apply + + +def unpad_input(hidden_states, attention_mask, unused_mask=None): + all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask + seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) + used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + in_fake_mode = active_fake_mode() is not None + if not in_fake_mode: + indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + else: + # torch.nonzero and .item() are not supported in FakeTensorMode + batch_size, seqlen = attention_mask.shape + indices = torch.arange(batch_size * seqlen, device=hidden_states.device) + max_seqlen_in_batch = seqlen + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + return ( + index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices), + indices, + cu_seqlens, + max_seqlen_in_batch, + used_seqlens_in_batch, + ) + + +def pad_input(hidden_states, indices, batch, seqlen): + output = index_put_first_axis(hidden_states, indices, batch * seqlen) + return rearrange(output, "(b s) ... -> b s ...", b=batch) + + +def generate_random_padding_mask( + max_seqlen, batch_size, device, mode="random", zero_lengths=False, min_seqlen=None +): + assert mode in ["full", "random", "third"] + min_seqlen = min_seqlen if min_seqlen is not None else 0 if zero_lengths else 1 + if mode == "full": + lengths = torch.full((batch_size, 1), max_seqlen, device=device, dtype=torch.int32) + elif mode == "random": + lengths = torch.randint( + max(min_seqlen, max_seqlen - 20), + max_seqlen + 1, + (batch_size, 1), + device=device, + ) + else: + lengths = torch.randint( + max(min_seqlen, max_seqlen // 3), + max_seqlen + 1, + (batch_size, 1), + device=device, + ) + + if zero_lengths: + for i in range(batch_size): + if i % 5 == 0: + lengths[i] = 0 + lengths[-1] = 0 + padding_mask = ( + repeat(torch.arange(max_seqlen, device=device), "s -> b s", b=batch_size) < lengths + ) + return padding_mask + + +def generate_qkv( + q, + k, + v, + query_padding_mask=None, + key_padding_mask=None, + qv=None, + kvpacked=False, + qkvpacked=False, + query_unused_mask=None, + key_unused_mask=None, +): + assert not (kvpacked and qkvpacked) + batch_size, seqlen_q, nheads, d = q.shape + d_v = v.shape[-1] + _, seqlen_k, nheads_k, _ = k.shape + assert k.shape == (batch_size, seqlen_k, nheads_k, d) + assert v.shape == (batch_size, seqlen_k, nheads_k, d_v) + if query_unused_mask is not None or key_unused_mask is not None: + assert not kvpacked + assert not qkvpacked + + if query_padding_mask is not None: + q_unpad, indices_q, cu_seqlens_q, max_seqlen_q, seqused_q = unpad_input( + q, query_padding_mask, query_unused_mask + ) + output_pad_fn = lambda output_unpad: pad_input( + output_unpad, indices_q, batch_size, seqlen_q + ) + qv_unpad = rearrange(qv, "b s ... -> (b s) ...")[indices_q] if qv is not None else None + else: + q_unpad = rearrange(q, "b s h d -> (b s) h d") + cu_seqlens_q = torch.arange( + 0, (batch_size + 1) * seqlen_q, step=seqlen_q, dtype=torch.int32, device=q_unpad.device + ) + seqused_q = None + max_seqlen_q = seqlen_q + output_pad_fn = lambda output_unpad: rearrange( + output_unpad, "(b s) h d -> b s h d", b=batch_size + ) + qv_unpad = rearrange(qv, "b s ... -> (b s) ...") if qv is not None else None + + if key_padding_mask is not None: + k_unpad, indices_k, cu_seqlens_k, max_seqlen_k, seqused_k = unpad_input( + k, key_padding_mask, key_unused_mask + ) + v_unpad, *_ = unpad_input(v, key_padding_mask, key_unused_mask) + else: + k_unpad = rearrange(k, "b s h d -> (b s) h d") + v_unpad = rearrange(v, "b s h d -> (b s) h d") + cu_seqlens_k = torch.arange( + 0, (batch_size + 1) * seqlen_k, step=seqlen_k, dtype=torch.int32, device=k_unpad.device + ) + seqused_k = None + max_seqlen_k = seqlen_k + + if qkvpacked: + assert (query_padding_mask == key_padding_mask).all() + assert nheads == nheads_k + qkv_unpad = torch.stack([q_unpad, k_unpad, v_unpad], dim=1) + qkv = torch.stack([q, k, v], dim=2) + if query_padding_mask is not None: + dqkv_pad_fn = lambda dqkv_unpad: pad_input(dqkv_unpad, indices_q, batch_size, seqlen_q) + else: + dqkv_pad_fn = lambda dqkv_unpad: rearrange( + dqkv_unpad, "(b s) t h d -> b s t h d", b=batch_size + ) + return ( + qkv_unpad.detach().requires_grad_(), + cu_seqlens_q, + max_seqlen_q, + qkv.detach().requires_grad_(), + output_pad_fn, + dqkv_pad_fn, + ) + elif kvpacked: + kv_unpad = torch.stack([k_unpad, v_unpad], dim=1) + kv = torch.stack([k, v], dim=2) + dq_pad_fn = output_pad_fn + if key_padding_mask is not None: + dkv_pad_fn = lambda dkv_unpad: pad_input(dkv_unpad, indices_k, batch_size, seqlen_k) + else: + dkv_pad_fn = lambda dkv_unpad: rearrange( + dkv_unpad, "(b s) t h d -> b s t h d", b=batch_size + ) + return ( + q_unpad.detach().requires_grad_(), + kv_unpad.detach().requires_grad_(), + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + q.detach().requires_grad_(), + kv.detach().requires_grad_(), + output_pad_fn, + dq_pad_fn, + dkv_pad_fn, + ) + else: + dq_pad_fn = output_pad_fn + if key_padding_mask is not None: + dk_pad_fn = lambda dk_unpad: pad_input(dk_unpad, indices_k, batch_size, seqlen_k) + else: + dk_pad_fn = lambda dk_unpad: rearrange(dk_unpad, "(b s) h d -> b s h d", b=batch_size) + return ( + q_unpad.detach().requires_grad_(), + k_unpad.detach().requires_grad_(), + v_unpad.detach().requires_grad_(), + qv_unpad.detach() if qv is not None else None, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + q.detach().requires_grad_(), + k.detach().requires_grad_(), + v.detach().requires_grad_(), + qv.detach() if qv is not None else None, + output_pad_fn, + dq_pad_fn, + dk_pad_fn, + ) + + +def construct_local_mask( + seqlen_q, + seqlen_k, + window_size=(None, None), + sink_token_length=0, + query_padding_mask=None, + key_padding_mask=None, + key_leftpad=None, + device=None, +): + row_idx = rearrange(torch.arange(seqlen_q, device=device, dtype=torch.long), "s -> s 1") + col_idx = torch.arange(seqlen_k, device=device, dtype=torch.long) + if key_leftpad is not None: + key_leftpad = rearrange(key_leftpad, "b -> b 1 1 1") + col_idx = repeat(col_idx, "s -> b 1 1 s", b=key_leftpad.shape[0]) + col_idx = torch.where(col_idx >= key_leftpad, col_idx - key_leftpad, 2**32) + sk = ( + seqlen_k + if key_padding_mask is None + else rearrange(key_padding_mask.sum(-1), "b -> b 1 1 1") + ) + sq = ( + seqlen_q + if query_padding_mask is None + else rearrange(query_padding_mask.sum(-1), "b -> b 1 1 1") + ) + if window_size[0] is None: + return col_idx > row_idx + sk - sq + window_size[1] + else: + sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk + if window_size[1] is None: + local_mask_left = col_idx > sk + else: + local_mask_left = col_idx > torch.minimum(row_idx + sk - sq + window_size[1], sk) + return torch.logical_or( + local_mask_left, + torch.logical_and( + col_idx < row_idx + sk - sq - window_size[0], col_idx >= sink_token_length + ), + ) + + +def construct_chunk_mask( + seqlen_q, + seqlen_k, + attention_chunk, + query_padding_mask=None, + key_padding_mask=None, + key_leftpad=None, + device=None, +): + row_idx = rearrange(torch.arange(seqlen_q, device=device, dtype=torch.long), "s -> s 1") + col_idx = torch.arange(seqlen_k, device=device, dtype=torch.long) + if key_leftpad is not None: + key_leftpad = rearrange(key_leftpad, "b -> b 1 1 1") + col_idx = repeat(col_idx, "s -> b 1 1 s", b=key_leftpad.shape[0]) + col_idx = torch.where(col_idx >= key_leftpad, col_idx - key_leftpad, 2**32) + sk = ( + seqlen_k + if key_padding_mask is None + else rearrange(key_padding_mask.sum(-1), "b -> b 1 1 1") + ) + sq = ( + seqlen_q + if query_padding_mask is None + else rearrange(query_padding_mask.sum(-1), "b -> b 1 1 1") + ) + sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk + col_limit_left_chunk = row_idx + sk - sq - (row_idx + sk - sq) % attention_chunk + return torch.logical_or( + col_idx < col_limit_left_chunk, col_idx >= col_limit_left_chunk + attention_chunk + ) + + +def attention_ref( + q, + k, + v, + query_padding_mask=None, + key_padding_mask=None, + key_leftpad=None, + attn_bias=None, + dropout_p=0.0, + dropout_mask=None, + causal=False, + qv=None, + q_descale=None, + k_descale=None, + v_descale=None, + window_size=(None, None), + attention_chunk=0, + sink_token_length=0, + learnable_sink: Optional[torch.Tensor] = None, + softcap=0.0, + upcast=True, + reorder_ops=False, + intermediate_dtype=None, + return_lse=False, + gather_kv_indices=None, +): + assert v is not None + has_qk = q is not None and k is not None + assert has_qk or qv is not None + if causal: + window_size = (window_size[0], 0) + dtype_og = v.dtype + q_shape = q.shape if q is not None else qv.shape + if upcast: + q, k, v, qv = [t.float() if t is not None else None for t in (q, k, v, qv)] + if q_descale is not None: + q_descale = repeat(q_descale, "b h -> b 1 (h g) 1", g=q_shape[2] // v.shape[2]) + q, qv = [(t.float() * q_descale).to(t.dtype) if t is not None else None for t in (q, qv)] + if k_descale is not None: + k = (k.float() * rearrange(k_descale, "b h -> b 1 h 1")).to(dtype=k.dtype) + if v_descale is not None: + v = (v.float() * rearrange(v_descale, "b h -> b 1 h 1")).to(dtype=v.dtype) + seqlen_q, seqlen_k = q_shape[1], v.shape[1] + k, v = [ + repeat(t, "b s h d -> b s (h g) d", g=q_shape[2] // t.shape[2]) if t is not None else None + for t in (k, v) + ] + d = q_shape[-1] # == dv for qv + dv = v.shape[-1] + softmax_scale = 1.0 / math.sqrt(d if qv is None or q is None else d + dv) + if has_qk: + scores = torch.einsum( + "bthd,bshd->bhts", + q if reorder_ops else q * softmax_scale, + k * softmax_scale if reorder_ops else k, + ) + if qv is not None: + qv_scores = torch.einsum( + "bthd,bshd->bhts", + qv if reorder_ops else qv * softmax_scale, + v * softmax_scale if reorder_ops else v, + ) + scores = qv_scores if not has_qk else scores + qv_scores + if softcap > 0: + scores = torch.tanh(scores / softcap) * softcap + if key_padding_mask is not None: + scores.masked_fill_(rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf")) + local_mask = None + if window_size[0] is not None or window_size[1] is not None: + local_mask = construct_local_mask( + seqlen_q, + seqlen_k, + window_size, + sink_token_length, + query_padding_mask, + key_padding_mask, + key_leftpad=key_leftpad, + device=v.device, + ) + if attention_chunk > 0: + chunk_mask = construct_chunk_mask( + seqlen_q, + seqlen_k, + attention_chunk, + query_padding_mask, + key_padding_mask, + key_leftpad=key_leftpad, + device=v.device, + ) + local_mask = ( + torch.logical_or(local_mask, chunk_mask) if local_mask is not None else chunk_mask + ) + if gather_kv_indices is not None: + batch = q_shape[0] + topk_len = gather_kv_indices.shape[2] + if topk_len < seqlen_k: + topk_index_mask = torch.full( + (batch, seqlen_q, seqlen_k), False, device="cuda" + ).scatter_(-1, gather_kv_indices, True) + scores.masked_fill_(rearrange(~topk_index_mask, "b t s -> b 1 t s"), float("-inf")) + if local_mask is not None: + scores.masked_fill_(local_mask, float("-inf")) + if attn_bias is not None: + scores = scores + attn_bias + # After all masks are applied, before softmax: + # scores shape: [b, h, t, s] + lse = torch.logsumexp(scores, dim=-1) # [b, h, t] + if learnable_sink is None: + attention = torch.softmax(scores, dim=-1).to(v.dtype) + else: + scores_fp32 = scores.to(torch.float32) + logits_max = torch.amax(scores_fp32, dim=-1, keepdim=True) + learnable_sink = rearrange(learnable_sink, "h -> h 1 1") + logits_or_sinks_max = torch.maximum(learnable_sink, logits_max) + unnormalized_scores = torch.exp(scores_fp32 - logits_or_sinks_max) + normalizer = unnormalized_scores.sum(dim=-1, keepdim=True) + torch.exp( + learnable_sink - logits_or_sinks_max + ) + # LSE with sink: log(Z) = log(normalizer) + max + lse = (torch.log(normalizer.squeeze(-1)) + logits_or_sinks_max.squeeze(-1)).to(dtype_og) + attention = (unnormalized_scores / normalizer).to(v.dtype) + if query_padding_mask is not None: + attention = attention.masked_fill(rearrange(~query_padding_mask, "b s -> b 1 s 1"), 0.0) + if key_padding_mask is not None: + attention = attention.masked_fill(rearrange(~key_padding_mask, "b s -> b 1 1 s"), 0.0) + if local_mask is not None: + attention = attention.masked_fill(torch.all(local_mask, dim=-1, keepdim=True), 0.0) + dropout_scaling = 1.0 / (1 - dropout_p) + if dropout_mask is not None: + attention_drop = attention.masked_fill(~dropout_mask, 0.0) + else: + attention_drop = attention + if intermediate_dtype is not None: + attention_drop = attention_drop.to(intermediate_dtype).to(attention_drop.dtype) + output = torch.einsum("bhts,bshd->bthd", attention_drop, v * dropout_scaling) + if query_padding_mask is not None: + output.masked_fill_(rearrange(~query_padding_mask, "b s -> b s 1 1"), 0.0) + if return_lse: + return output.to(dtype_og), attention.to(dtype_og), lse.to(dtype_og) + return output.to(dtype=dtype_og), attention.to(dtype=dtype_og) + + +def maybe_fake_tensor_mode(fake: bool = True): + """ + One way to populate/pre-compile cache is to use torch fake tensor mode, + which does not allocate actual GPU tensors but retains tensor shape/dtype + metadata for cute.compile. + """ + + def decorator(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + with FakeTensorMode() if fake else nullcontext(): + return fn(*args, **kwargs) + + return wrapper + + return decorator + + +def is_fake_mode() -> bool: + return active_fake_mode() is not None diff --git a/fastvideo-kernel/fa4/tile_scheduler.py b/fastvideo-kernel/fa4/tile_scheduler.py new file mode 100644 index 000000000..404d22c4c --- /dev/null +++ b/fastvideo-kernel/fa4/tile_scheduler.py @@ -0,0 +1,1650 @@ +# Copyright (c) 2025, Tri Dao, Siyu Wang, Shengbin Di, Yuxi Chi, Johnsonms, Linfeng Zheng, Haoyan Huang, Lanbo Li, Yun Zhong, Man Yuan, Minmin Sun, Yong Li, Wei Lin. + +from enum import IntEnum, auto +from typing import Optional, Tuple, Protocol, runtime_checkable +from dataclasses import dataclass + +try: + from typing import override +except ImportError: # Python < 3.12 + from typing_extensions import override + +import cutlass +from cutlass.pipeline import PipelineClcFetchAsync, PipelineState +from cutlass._mlir import ir +import cutlass.cute as cute +from cutlass import Int32, const_expr +from cutlass.cute import FastDivmodDivisor +from cutlass.utils import ClcDynamicPersistentTileScheduler, ClcDynamicPersistentTileSchedulerParams +from cutlass.cute.typing import Boolean +from cutlass.cutlass_dsl import ( + min as dsl_min, + extract_mlir_values, + new_from_mlir_values, +) +from cutlass.utils.hardware_info import HardwareInfo + +from quack.cute_dsl_utils import ParamsBase + +import flash_attn.cute.utils as utils +from flash_attn.cute.fast_math import clz + + +class SchedulingMode(IntEnum): + NONE = auto() + STATIC = auto() + DYNAMIC = auto() + CLC = auto() + + +@dataclass +class ClcState(ParamsBase): + """Owns the runtime state shared by CLC-capable tile schedulers. + + `FlashAttentionForwardSm100` constructs this state because it owns the CLC + response buffer, mbarrier storage, and launch geometry needed to initialize + the hardware scheduler and async pipeline. Individual tile schedulers then + consume this state and map the returned hardware work tiles into their own + logical `WorkTileInfo` coordinates. + + To add CLC support to a scheduler: + - implement `clc_problem_shape(params)` so the kernel can create the hardware scheduler + - accept `clc: ClcState | None` in `create(...)` / `__init__` + - map `clc.initial_work_tile_info()` and `clc.get_current_work()` into scheduler coordinates + """ + + _hw_scheduler: ClcDynamicPersistentTileScheduler + _pipeline: PipelineClcFetchAsync + _consumer_state: PipelineState + _producer_state: PipelineState + + @staticmethod + def create( + *, + hw_scheduler: ClcDynamicPersistentTileScheduler, + pipeline: PipelineClcFetchAsync, + consumer_state: PipelineState, + producer_state: PipelineState, + ) -> "ClcState": + return ClcState(hw_scheduler, pipeline, consumer_state, producer_state) + + def initial_work_tile_info(self): + return self._hw_scheduler.initial_work_tile_info() + + def get_current_work(self): + return self._hw_scheduler.get_current_work() + + def prefetch_next_work(self, *, loc=None, ip=None): + self._pipeline.producer_acquire(self._producer_state, loc=loc, ip=ip) + mbarrier_addr = self._pipeline.producer_get_barrier(self._producer_state, loc=loc, ip=ip) + self._hw_scheduler.advance_to_next_work(mbarrier_addr, loc=loc, ip=ip) + self._producer_state.advance(loc=loc, ip=ip) + + def consumer_wait(self, *, loc=None, ip=None): + self._pipeline.consumer_wait(self._consumer_state, loc=loc, ip=ip) + + def consumer_release(self, *, loc=None, ip=None): + self._pipeline.consumer_release(self._consumer_state, loc=loc, ip=ip) + self._consumer_state.advance(loc=loc, ip=ip) + + def producer_tail(self, *, loc=None, ip=None): + self._pipeline.producer_tail(self._producer_state, loc=loc, ip=ip) + + +class WorkTileInfo(cutlass.utils.WorkTileInfo): + """Altered WorkTileInfo which includes four axes: (block, head, batch, split)""" + + @override + def __new_from_mlir_values__(self, values: list[ir.Value]) -> "WorkTileInfo": + assert len(values) == 5 + new_tile_idx = cutlass.new_from_mlir_values(self._tile_idx, values[:-1]) + new_is_valid_tile = cutlass.new_from_mlir_values(self._is_valid_tile, [values[-1]]) + return WorkTileInfo(new_tile_idx, new_is_valid_tile) + + +@runtime_checkable +class TileSchedulerProtocol(Protocol): + """Protocol defining the interface all tile schedulers must implement. + + Schedulers are responsible for: + 1. Coordinate mapping: linear tile index -> (m_block, head, batch, split) + 2. Work distribution: how to get the next tile (static grid-stride vs CLC dynamic) + """ + + def get_current_work(self) -> WorkTileInfo: + """Get the current work tile coordinates.""" + ... + + def initial_work_tile_info(self) -> WorkTileInfo: + """Get the initial work tile for this CTA.""" + ... + + def advance_to_next_work(self, *, loc=None, ip=None): + """Consumer-side advance: move to next tile and return it. + + For static schedulers: grid-stride increment + get_current_work. + For CLC schedulers: consumer wait + get_current_work + consumer release + state advance. + """ + ... + + def prefetch_next_work(self, *, loc=None, ip=None) -> None: + """Producer-side prefetch of next work tile (no-op for static schedulers). + + For CLC schedulers: producer acquire + issue CLC query + producer state advance. + Only called by the scheduler warp. + """ + ... + + def producer_tail(self, *, loc=None, ip=None) -> None: + """Producer-side cleanup after the last tile. + + No-op for static schedulers. For CLC schedulers: pipeline producer_tail. + """ + ... + + +@dataclass +class TileSchedulerArguments(ParamsBase): + num_block: Int32 + num_head: Int32 + num_batch: Int32 + num_splits: Int32 + seqlen_k: Int32 + headdim: Int32 + headdim_v: Int32 + total_q: Int32 + tile_shape_mn: cutlass.Constexpr[Tuple[int, int]] + cluster_shape_mn: cutlass.Constexpr[Tuple[int, int]] = (1, 1) + mCuSeqlensQ: Optional[cute.Tensor] = None + mSeqUsedQ: Optional[cute.Tensor] = None + qhead_per_kvhead_packgqa: cutlass.Constexpr[int] = 1 + element_size: cutlass.Constexpr[int] = 2 + is_persistent: cutlass.Constexpr[bool] = False + lpt: cutlass.Constexpr[bool] = False + is_split_kv: cutlass.Constexpr[bool] = False + head_swizzle: cutlass.Constexpr[bool] = False + use_cluster_idx: cutlass.Constexpr[bool] = False + + +class SingleTileScheduler: + @dataclass + class Params(ParamsBase): + num_block: Int32 + num_head: Int32 + num_batch: Int32 + num_splits: Int32 + num_splits_divmod: FastDivmodDivisor + is_split_kv: cutlass.Constexpr[bool] = False + cluster_shape_mn: cutlass.Constexpr[Tuple[int, int]] = (1, 1) + use_cluster_idx: cutlass.Constexpr[bool] = False + + @staticmethod + def create( + args: TileSchedulerArguments, *, loc=None, ip=None + ) -> "SingleTileScheduler.Params": + return SingleTileScheduler.Params( + args.num_block, + args.num_head, + args.num_batch, + args.num_splits, + FastDivmodDivisor(args.num_splits), + args.is_split_kv, + args.cluster_shape_mn, + args.use_cluster_idx, + ) + + def __init__(self, params: Params, blk_coord: cute.Coord, *, loc=None, ip=None): + self.params = params + self._blk_coord = blk_coord + self._is_first_block = True + self._loc = loc + self._ip = ip + + @staticmethod + def to_underlying_arguments( + args: TileSchedulerArguments, + *, + scheduling_mode: SchedulingMode = SchedulingMode.STATIC, + loc=None, + ip=None, + ) -> Params: + assert scheduling_mode == SchedulingMode.STATIC, ( + f"SingleTileScheduler only supports STATIC, got {scheduling_mode!r}" + ) + return SingleTileScheduler.Params.create(args, loc=loc, ip=ip) + + @staticmethod + def create( + params: Params, clc: ClcState | None = None, *, loc=None, ip=None + ) -> "SingleTileScheduler": + if const_expr(cute.size(params.cluster_shape_mn) == 1 or not params.use_cluster_idx): + blk_coord = cute.arch.block_idx() + else: + blk_coord = cute.arch.cluster_idx() + return SingleTileScheduler(params, blk_coord, loc=loc, ip=ip) + + # called by host + @staticmethod + def get_grid_shape( + params: Params, + *, + loc=None, + ip=None, + ) -> Tuple[Int32, Int32, Int32]: + # TODO: this hard-codes the fact that we only use cluster = (1, 1) or (2, 1) + assert params.cluster_shape_mn[1] == 1, "Only cluster_shape_mn[1] == 1 is supported" + if const_expr(params.use_cluster_idx): + # Grid must have num_block * cluster_m physical blocks so that there are num_block clusters + grid_x = params.num_block * params.cluster_shape_mn[0] + else: + grid_x = cute.round_up(params.num_block, params.cluster_shape_mn[0]) + return ( + grid_x, + params.num_head * params.num_splits, + params.num_batch, + ) + + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + block_idx, head_idx, batch_idx = self._blk_coord + if const_expr(self.params.is_split_kv): + head_idx, split_idx = divmod(head_idx, self.params.num_splits_divmod) + else: + split_idx = Int32(0) + return WorkTileInfo( + (block_idx, head_idx, batch_idx, split_idx), + self._is_first_block, + ) + + def initial_work_tile_info(self, *, loc=None, ip=None): + return self.get_current_work(loc=loc, ip=ip) + + def prefetch_next_work(self, *, loc=None, ip=None): + pass + + def advance_to_next_work(self, *, loc=None, ip=None): + self._is_first_block = False + return self.get_current_work() + + def producer_tail(self, *, loc=None, ip=None): + pass + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in [self.params, self._blk_coord]: + obj_values = cutlass.extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + for obj, n_items in zip([self.params, self._blk_coord], self._values_pos): + obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return SingleTileScheduler(*(tuple(obj_list)), loc=self._loc) + + +class StaticPersistentTileScheduler: + @dataclass + class Params(ParamsBase): + num_block_cluster_divmod: FastDivmodDivisor + num_head_divmod: FastDivmodDivisor + total_blocks_cluster: Int32 + cluster_shape_m: cutlass.Constexpr[int] = 1 + + @staticmethod + def create( + args: TileSchedulerArguments, *, loc=None, ip=None + ) -> "StaticPersistentTileScheduler.Params": + num_block_cluster = cute.ceil_div(args.num_block, cute.size(args.cluster_shape_mn)) + total_blocks_cluster = num_block_cluster * args.num_head * args.num_batch + return StaticPersistentTileScheduler.Params( + FastDivmodDivisor(num_block_cluster), + FastDivmodDivisor(args.num_head), + total_blocks_cluster, + cluster_shape_m=args.cluster_shape_mn[0], + ) + + def __init__(self, params: Params, tile_idx: Int32, *, loc=None, ip=None): + self.params = params + self._tile_idx = tile_idx + self._loc = loc + self._ip = ip + + @staticmethod + def to_underlying_arguments( + args: TileSchedulerArguments, + *, + scheduling_mode: SchedulingMode = SchedulingMode.STATIC, + loc=None, + ip=None, + ) -> Params: + assert scheduling_mode == SchedulingMode.STATIC, ( + f"StaticPersistentTileScheduler only supports STATIC, got {scheduling_mode!r}" + ) + return StaticPersistentTileScheduler.Params.create(args, loc=loc, ip=ip) + + @staticmethod + def create( + params: Params, clc: ClcState | None = None, *, loc=None, ip=None + ) -> "StaticPersistentTileScheduler": + if const_expr(cute.size(params.cluster_shape_m) == 1): + tile_idx = cute.arch.block_idx()[0] + else: + tile_idx = cute.arch.cluster_idx()[0] + return StaticPersistentTileScheduler(params, tile_idx, loc=loc, ip=ip) + + @staticmethod + def get_grid_shape( + params: Params, + *, + loc=None, + ip=None, + ) -> Tuple[Int32, Int32, Int32]: + hardware_info = cutlass.utils.HardwareInfo() + sm_count = hardware_info.get_device_multiprocessor_count() + max_ctas = (sm_count // params.cluster_shape_m) * params.cluster_shape_m + grid_x = cutlass.min(max_ctas, params.total_blocks_cluster * params.cluster_shape_m) + return (grid_x, Int32(1), Int32(1)) + + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + hn_idx, block_idx = divmod(self._tile_idx, self.params.num_block_cluster_divmod) + batch_idx, head_idx = divmod(hn_idx, self.params.num_head_divmod) + is_valid = self._tile_idx < self.params.total_blocks_cluster + return WorkTileInfo( + (Int32(block_idx), Int32(head_idx), Int32(batch_idx), Int32(0)), is_valid + ) + + def initial_work_tile_info(self, *, loc=None, ip=None): + return self.get_current_work(loc=loc, ip=ip) + + def prefetch_next_work(self, *, loc=None, ip=None): + pass + + def advance_to_next_work(self, *, loc=None, ip=None): + if const_expr(self.params.cluster_shape_m == 1): + self._tile_idx += cute.arch.grid_dim()[0] + else: + self._tile_idx += cute.arch.cluster_dim()[0] + return self.get_current_work() + + def producer_tail(self, *, loc=None, ip=None): + pass + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in [self.params, self._tile_idx]: + obj_values = cutlass.extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + for obj, n_items in zip( + [self.params, self._tile_idx], + self._values_pos, + ): + obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return StaticPersistentTileScheduler(*(tuple(obj_list)), loc=self._loc) + + +class SingleTileLPTScheduler: + @dataclass + class Params(ParamsBase): + total_blocks: Int32 + num_splits: Int32 + num_block: Int32 + num_head: Int32 + num_batch: Int32 + l2_minor: Int32 + num_head_divmod: FastDivmodDivisor + l2_minor_divmod: FastDivmodDivisor + l2_major_divmod: FastDivmodDivisor + l2_minor_residual_divmod: FastDivmodDivisor + num_hb_quotient: Int32 + num_splits_divmod: FastDivmodDivisor + is_split_kv: cutlass.Constexpr[bool] = False + cluster_shape_m: cutlass.Constexpr[int] = 1 + scheduling_mode: cutlass.Constexpr[SchedulingMode] = SchedulingMode.STATIC + lpt: cutlass.Constexpr[bool] = True + use_cluster_idx: cutlass.Constexpr[bool] = True + + @staticmethod + @cute.jit + def create( + args: TileSchedulerArguments, + *, + scheduling_mode: SchedulingMode = SchedulingMode.STATIC, + loc=None, + ip=None, + ) -> "SingleTileLPTScheduler.Params": + assert scheduling_mode in (SchedulingMode.STATIC, SchedulingMode.CLC), ( + f"Only STATIC and CLC are supported, got {scheduling_mode!r}" + ) + # int64: this product overflows int32 once seqlen_k * (headdim + + # headdim_v) * element_size > 2**31 (seqlen_k > ~4M for hdim-128 bf16). + size_one_kv_head = ( + cutlass.Int64(args.seqlen_k) * (args.headdim + args.headdim_v) * args.element_size + ) + size_one_head = size_one_kv_head + size_l2 = 50 * 1024 * 1024 # 40 MB for K & V + # Swizzle is the size of each "section". Round swizzle to a power of 2 + # Need to be careful about the case where only one head will fit + # swizzle is how many heads can fit in L2 + # Seems faster if swizzle is a power of 2 + log2_floor = lambda n: 31 - clz(n) + swizzle = ( + 1 if size_l2 < size_one_head else (1 << log2_floor(Int32(size_l2 // size_one_head))) + ) + # If we're in the last section (called residual), we don't want to divide by + # swizzle. Instead we want to divide by the remainder. + num_hb_quotient = (args.num_head * args.num_batch) // swizzle + num_hb_remainder = (args.num_head * args.num_batch) % swizzle + return SingleTileLPTScheduler.Params( + total_blocks=args.num_block * args.num_head * args.num_batch, + num_block=args.num_block, + num_head=args.num_head, + num_batch=args.num_batch, + l2_minor=Int32(swizzle), + num_head_divmod=FastDivmodDivisor(args.num_head), + l2_minor_divmod=FastDivmodDivisor(swizzle), + l2_major_divmod=FastDivmodDivisor(swizzle * args.num_block), + l2_minor_residual_divmod=FastDivmodDivisor(max(num_hb_remainder, 1)), + num_hb_quotient=Int32(num_hb_quotient), + num_splits=args.num_splits, + num_splits_divmod=FastDivmodDivisor(args.num_splits), + is_split_kv=args.is_split_kv, + cluster_shape_m=args.cluster_shape_mn[0], + scheduling_mode=scheduling_mode, + lpt=args.lpt, + use_cluster_idx=args.use_cluster_idx, + ) + + def __init__( + self, + params: Params, + tile_idx: Int32, + split_idx: Int32, + clc: ClcState | None = None, + *, + loc=None, + ip=None, + ): + self.params = params + self._tile_idx = tile_idx + self._split_idx = split_idx + self.clc = clc + self._loc = loc + self._ip = ip + + @staticmethod + def to_underlying_arguments( + args: TileSchedulerArguments, + *, + scheduling_mode: SchedulingMode = SchedulingMode.STATIC, + loc=None, + ip=None, + ) -> Params: + return SingleTileLPTScheduler.Params.create( + args, scheduling_mode=scheduling_mode, loc=loc, ip=ip + ) + + @staticmethod + def _clc_grid_shape(params: Params): + num_batch_splits = ( + params.num_batch * params.num_splits + if const_expr(params.is_split_kv) + else params.num_batch + ) + if const_expr(params.use_cluster_idx): + # Grid must have num_block * cluster_m physical blocks so that there are num_block clusters + grid_x = params.num_block * params.cluster_shape_m + else: + grid_x = cute.round_up(params.num_block, params.cluster_shape_m) + return ( + grid_x, + params.num_head, + num_batch_splits, + ) + + @staticmethod + @cute.jit + def clc_problem_shape(params: Params): + return ClcDynamicPersistentTileSchedulerParams( + problem_shape_ntile_mnl=SingleTileLPTScheduler._clc_grid_shape(params), + cluster_shape_mnk=(params.cluster_shape_m, 1, 1), + ) + + @staticmethod + @cute.jit + def create( + params: Params, clc: ClcState | None = None, *, loc=None, ip=None + ) -> "SingleTileLPTScheduler": + if const_expr(params.scheduling_mode == SchedulingMode.CLC): + return SingleTileLPTScheduler( + params, cute.arch.block_idx()[0], Int32(0), clc, loc=loc, ip=ip + ) + tile_idx, split_idx, _ = cute.arch.block_idx() + return SingleTileLPTScheduler(params, tile_idx, split_idx, loc=loc, ip=ip) + + @staticmethod + def get_grid_shape( + params: Params, + *, + loc=None, + ip=None, + ) -> Tuple[Int32, Int32, Int32]: + if const_expr(params.scheduling_mode == SchedulingMode.CLC): + return SingleTileLPTScheduler._clc_grid_shape(params) + return (params.total_blocks, params.num_splits, Int32(1)) + + @cute.jit + def clc_work_to_coords(self, work) -> WorkTileInfo: + """Convert CLC response (block, head, batch_split) to WorkTileInfo. + + CLC returns raw grid coordinates — no L2 swizzle (hardware decides order). + We only apply cluster division, optional LPT block reversal, and split_kv unpacking. + """ + block_idx = work.tile_idx[0] + if const_expr(self.params.cluster_shape_m > 1): + block_idx = block_idx // self.params.cluster_shape_m + if const_expr(self.params.lpt): + # Longest-processing-time-first: reverse block order + if const_expr(self.params.cluster_shape_m > 1 and not self.params.use_cluster_idx): + num_block = self.params.num_block // self.params.cluster_shape_m + else: + num_block = self.params.num_block + block_idx = num_block - 1 - block_idx + split_idx = Int32(0) + if const_expr(self.params.is_split_kv): + batch_idx, split_idx = divmod(work.tile_idx[2], self.params.num_splits_divmod) + else: + batch_idx = work.tile_idx[2] + if const_expr(self.params.cluster_shape_m > 1 and not self.params.use_cluster_idx): + bidx_in_cluster = cute.arch.block_in_cluster_idx() + block_idx = block_idx * self.params.cluster_shape_m + bidx_in_cluster[0] + return WorkTileInfo( + (Int32(block_idx), Int32(work.tile_idx[1]), Int32(batch_idx), Int32(split_idx)), + work.is_valid_tile, + ) + + @cute.jit + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + work = self.clc.get_current_work() + self._tile_idx = work.tile_idx[0] + return self.clc_work_to_coords(work) + # Static path: L2-swizzled coordinate mapping + params = self.params + # Implement LPT scheduling coordinate calculation + bidhb, l2_mod = divmod(self._tile_idx, params.l2_major_divmod) + # If we're in the last section (called residual), we don't want to divide by + # swizzle. Instead we want to divide by the remainder. + block, bidhb_residual = 0, 0 + if bidhb < params.num_hb_quotient: + block, bidhb_residual = divmod(l2_mod, params.l2_minor_divmod) + else: + block, bidhb_residual = divmod(l2_mod, params.l2_minor_residual_divmod) + bidhb_actual = bidhb * params.l2_minor + bidhb_residual + batch_idx, head_idx = divmod(bidhb_actual, params.num_head_divmod) + # Longest-processing-time-first + if const_expr(params.lpt): + block = params.num_block - 1 - block + is_valid = self._tile_idx < params.total_blocks + return WorkTileInfo( + (Int32(block), Int32(head_idx), Int32(batch_idx), Int32(self._split_idx)), is_valid + ) + + @cute.jit + def initial_work_tile_info(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + work = self.clc.initial_work_tile_info() + self._tile_idx = work.tile_idx[0] + return self.clc_work_to_coords(work) + return self.get_current_work(loc=loc, ip=ip) + + def prefetch_next_work(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + self.clc.prefetch_next_work(loc=loc, ip=ip) + + def advance_to_next_work(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + self.clc.consumer_wait(loc=loc, ip=ip) + work = self.get_current_work() + self.clc.consumer_release(loc=loc, ip=ip) + return work + # Single tile scheduler - set to invalid tile_idx to indicate no more work + self._tile_idx = self.params.total_blocks + return self.get_current_work() + + def producer_tail(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + self.clc.producer_tail(loc=loc, ip=ip) + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + objs = [self.params, self._tile_idx, self._split_idx] + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + objs += [self.clc] + for obj in objs: + obj_values = cutlass.extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + objs = [self.params, self._tile_idx, self._split_idx] + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + objs += [self.clc] + for obj, n_items in zip(objs, self._values_pos): + obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return self.__class__(*obj_list, loc=self._loc) + + +class SingleTileLPTBwdScheduler: + @dataclass + class Params(ParamsBase): + total_blocks: Int32 + num_block: Int32 + l2_minor: Int32 + num_head_divmod: FastDivmodDivisor + l2_minor_divmod: FastDivmodDivisor + l2_major_divmod: FastDivmodDivisor + l2_minor_residual_divmod: FastDivmodDivisor + num_hb_quotient: Int32 + cluster_shape_mn: cutlass.Constexpr[Tuple[int, int]] = (1, 1) + spt: cutlass.Constexpr[bool] = True + + @staticmethod + @cute.jit + def create( + args: TileSchedulerArguments, *, loc=None, ip=None + ) -> "SingleTileLPTBwdScheduler.Params": + size_l2 = 50 * 1024 * 1024 + # int64: these products overflow int32 at large seqlen_k (> ~4M for + # hdim-128 bf16; the dqaccum *4 term wraps even sooner). + size_one_qdo_head = ( + cutlass.Int64(args.seqlen_k) * (args.headdim + args.headdim_v) * args.element_size + ) + size_one_dqaccum_head = cutlass.Int64(args.seqlen_k) * (args.headdim) * 4 + # size_one_dqaccum_head = 0 + size_one_head = size_one_qdo_head + size_one_dqaccum_head + log2_floor = lambda n: 31 - clz(n) + swizzle = ( + 1 if size_l2 < size_one_head else (1 << log2_floor(Int32(size_l2 // size_one_head))) + ) + # swizzle = 8 + # If we're in the last section (called residual), we don't want to divide by + # swizzle. Instead we want to divide by the remainder. + num_hb_quotient = (args.num_head * args.num_batch) // swizzle + num_hb_remainder = (args.num_head * args.num_batch) % swizzle + num_block = cute.ceil_div(args.num_block, args.cluster_shape_mn[0]) + return SingleTileLPTBwdScheduler.Params( + total_blocks=(num_block * args.cluster_shape_mn[0]) + * args.num_head + * args.num_batch, + num_block=num_block, + l2_minor=Int32(swizzle), + num_head_divmod=FastDivmodDivisor(args.num_head), + l2_minor_divmod=FastDivmodDivisor(swizzle), + l2_major_divmod=FastDivmodDivisor(swizzle * num_block), + l2_minor_residual_divmod=FastDivmodDivisor( + max(num_hb_remainder, 1) + ), # don't divide by 0 + num_hb_quotient=Int32(num_hb_quotient), + cluster_shape_mn=args.cluster_shape_mn, + spt=args.lpt, + ) + + def __init__(self, params: Params, tile_idx: Int32, *, loc=None, ip=None): + self.params = params + self._tile_idx = tile_idx + self._loc = loc + self._ip = ip + + @staticmethod + def to_underlying_arguments( + args: TileSchedulerArguments, + *, + scheduling_mode: SchedulingMode = SchedulingMode.STATIC, + loc=None, + ip=None, + ) -> Params: + assert scheduling_mode == SchedulingMode.STATIC, ( + f"SingleTileLPTBwdScheduler only supports STATIC, got {scheduling_mode!r}" + ) + return SingleTileLPTBwdScheduler.Params.create(args, loc=loc, ip=ip) + + @staticmethod + @cute.jit + def create(params: Params, *, loc=None, ip=None) -> "SingleTileLPTBwdScheduler": + tile_idx = cute.arch.block_idx()[0] + return SingleTileLPTBwdScheduler(params, tile_idx, loc=loc, ip=ip) + + # called by host + @staticmethod + def get_grid_shape( + params: Params, + *, + loc=None, + ip=None, + ) -> Tuple[Int32, Int32, Int32]: + return (params.total_blocks, Int32(1), Int32(1)) + + @cute.jit + def get_current_work(self, *, loc=None, ip=None) -> cutlass.utils.WorkTileInfo: + cluster_idx = self._tile_idx // self.params.cluster_shape_mn[0] + params = self.params + # Implement LPT scheduling coordinate calculation + bidhb, l2_mod = divmod(cluster_idx, params.l2_major_divmod) + # If we're in the last section (called residual), we don't want to divide by + # swizzle. Instead we want to divide by the remainder. + block, bidhb_residual = 0, 0 + if bidhb < params.num_hb_quotient: + block, bidhb_residual = divmod(l2_mod, params.l2_minor_divmod) + else: + block, bidhb_residual = divmod(l2_mod, params.l2_minor_residual_divmod) + bidhb_actual = bidhb * params.l2_minor + bidhb_residual + batch_idx, head_idx = divmod(bidhb_actual, params.num_head_divmod) + if cutlass.const_expr(params.spt): + block = params.num_block - 1 - block + if cutlass.const_expr(params.cluster_shape_mn[0] > 1): + bidx_in_cluster = cute.arch.block_in_cluster_idx() + block = block * params.cluster_shape_mn[0] + bidx_in_cluster[0] + is_valid = self._tile_idx < params.total_blocks + return WorkTileInfo((Int32(block), Int32(head_idx), Int32(batch_idx), Int32(0)), is_valid) + + def initial_work_tile_info(self, *, loc=None, ip=None): + return self.get_current_work(loc=loc, ip=ip) + + def prefetch_next_work(self, *, loc=None, ip=None): + pass + + def advance_to_next_work(self, *, loc=None, ip=None): + # Single tile scheduler - set to invalid tile_idx to indicate no more work + self._tile_idx = self.params.total_blocks + return self.get_current_work() + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in [self.params, self._tile_idx]: + obj_values = cutlass.extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + for obj, n_items in zip([self.params, self._tile_idx], self._values_pos): + obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return self.__class__(*(tuple(obj_list)), loc=self._loc) + + +class SingleTileVarlenScheduler: + @dataclass + class Params(ParamsBase): + num_head: Int32 + num_batch: Int32 + total_q: Int32 + num_splits: Int32 + max_kvblock_in_l2: Int32 + tile_shape_mn: cutlass.Constexpr[Tuple[int, int]] + mCuSeqlensQ: Optional[cute.Tensor] = None + mSeqUsedQ: Optional[cute.Tensor] = None + qhead_per_kvhead_packgqa: cutlass.Constexpr[int] = 1 + lpt: cutlass.Constexpr[bool] = False + is_split_kv: cutlass.Constexpr[bool] = False + head_swizzle: cutlass.Constexpr[bool] = False + cluster_shape_m: cutlass.Constexpr[int] = 1 + use_cluster_idx: cutlass.Constexpr[bool] = False + scheduling_mode: cutlass.Constexpr[SchedulingMode] = SchedulingMode.STATIC + + @staticmethod + @cute.jit + def create( + args: TileSchedulerArguments, + *, + scheduling_mode: SchedulingMode = SchedulingMode.STATIC, + loc=None, + ip=None, + ) -> "SingleTileVarlenScheduler.Params": + assert scheduling_mode in (SchedulingMode.STATIC, SchedulingMode.CLC), ( + f"Only STATIC and CLC are supported, got {scheduling_mode!r}" + ) + size_l2 = 50 * 1024 * 1024 # 50 MB for K & V + # if backward, this is qdo block size + kv_block_size = ( + (args.headdim + args.headdim_v) * args.element_size * args.tile_shape_mn[1] + ) + # if backward, add dqaccum block size to calculate swizzle + if args.head_swizzle: + kv_block_size += args.headdim * 4 * args.tile_shape_mn[1] + max_kvblock_in_l2 = size_l2 // kv_block_size + assert args.mCuSeqlensQ is not None or args.mSeqUsedQ is not None, ( + "At least one of mCuSeqlensQ or mSeqUsedQ must be provided" + ) + assert args.cluster_shape_mn[1] == 1, "Only cluster_shape_mn[1] == 1 is supported" + return SingleTileVarlenScheduler.Params( + num_head=args.num_head, + num_batch=args.num_batch, + total_q=args.total_q, + num_splits=args.num_splits, + max_kvblock_in_l2=max_kvblock_in_l2, + tile_shape_mn=args.tile_shape_mn, + mCuSeqlensQ=args.mCuSeqlensQ, + mSeqUsedQ=args.mSeqUsedQ, + qhead_per_kvhead_packgqa=args.qhead_per_kvhead_packgqa, + lpt=args.lpt, + is_split_kv=args.is_split_kv, + head_swizzle=args.head_swizzle, + cluster_shape_m=args.cluster_shape_mn[0], + use_cluster_idx=args.use_cluster_idx, + scheduling_mode=scheduling_mode, + ) + + def __init__( + self, + params: Params, + tile_idx: Int32, + split_idx: Int32, + clc: ClcState | None = None, + *, + loc=None, + ip=None, + ): + self.params = params + self._tile_idx = tile_idx + self._split_idx = split_idx + self._is_first_block = True + self.clc = clc + self._loc = loc + self._ip = ip + + @staticmethod + def to_underlying_arguments( + args: TileSchedulerArguments, + *, + scheduling_mode: SchedulingMode = SchedulingMode.STATIC, + loc=None, + ip=None, + ) -> Params: + return SingleTileVarlenScheduler.Params.create( + args, scheduling_mode=scheduling_mode, loc=loc, ip=ip + ) + + @staticmethod + @cute.jit + def clc_problem_shape(params: Params): + return ClcDynamicPersistentTileSchedulerParams( + problem_shape_ntile_mnl=SingleTileVarlenScheduler.get_grid_shape(params), + cluster_shape_mnk=(1, 1, 1), + ) + + @staticmethod + @cute.jit + def create( + params: Params, clc: ClcState | None = None, *, loc=None, ip=None + ) -> "SingleTileVarlenScheduler": + if const_expr(params.scheduling_mode == SchedulingMode.CLC): + block_idx = cute.arch.block_idx() + split_idx = Int32(0) + if const_expr(params.is_split_kv): + split_idx = block_idx[1] + return SingleTileVarlenScheduler( + params, + block_idx[0], + split_idx, + clc, + loc=loc, + ip=ip, + ) + tile_idx, split_idx, _ = cute.arch.block_idx() + return SingleTileVarlenScheduler(params, tile_idx, split_idx, loc=loc, ip=ip) + + # called by host + @staticmethod + def get_grid_shape( + params: Params, + *, + loc=None, + ip=None, + ) -> Tuple[Int32, Int32, Int32]: + total_blocks_max = ( + params.total_q + + params.num_batch * (params.cluster_shape_m * params.tile_shape_mn[0] - 1) + ) // params.tile_shape_mn[0] + # Round down to nearest multiple of cluster since odd excess is always padding. + total_blocks_max = total_blocks_max // params.cluster_shape_m * params.cluster_shape_m + return (total_blocks_max * params.num_head, params.num_splits, Int32(1)) + + @cute.jit + def _get_num_m_blocks(self, lane: Int32, bidb_start: Int32) -> Int32: + params = self.params + batch_idx = lane + bidb_start + if cutlass.const_expr(params.mSeqUsedQ is not None): + seqlen = Int32(0) + if batch_idx < params.num_batch: + seqlen = params.mSeqUsedQ[batch_idx] + else: + assert params.mCuSeqlensQ is not None + cur_cu_seqlen = Int32(0) + if batch_idx <= params.num_batch: + cur_cu_seqlen = params.mCuSeqlensQ[batch_idx] + next_cu_seqlen = cute.arch.shuffle_sync_down(cur_cu_seqlen, offset=1) + seqlen = next_cu_seqlen - cur_cu_seqlen + if cutlass.const_expr(params.qhead_per_kvhead_packgqa > 1): + seqlen *= params.qhead_per_kvhead_packgqa + return ( + cute.ceil_div(cute.ceil_div(seqlen, params.tile_shape_mn[0]), params.cluster_shape_m) + if batch_idx < params.num_batch and lane < cute.arch.WARP_SIZE - 1 + else Int32(0) + ) + + @cute.jit + def _varlen_coord_map(self) -> WorkTileInfo: + """Map self._tile_idx to (block, head, batch) via warp-level prefix sums.""" + params = self.params + lane_idx = cute.arch.lane_idx() + num_m_blocks = self._get_num_m_blocks(lane_idx, bidb_start=0) + num_m_blocks_cumulative = utils.warp_prefix_sum(num_m_blocks, lane_idx) + # Total number of blocks for the next 31 batches + m_blocks_in_group = cute.arch.shuffle_sync(num_m_blocks_cumulative, cute.arch.WARP_SIZE - 1) + # Same for all lanes + group_end_tile = m_blocks_in_group * params.num_head + # if cute.arch.thread_idx()[0] == 128 + 31: cute.printf("SingleTileVarlenScheduler: tile_idx=%d, group_end_tile = %d, num_m_blocks=%d, num_m_blocks_cumulative = %d, m_blocks_in_group = %d", self._tile_idx, group_end_tile, num_m_blocks, num_m_blocks_cumulative, m_blocks_in_group) + block, head_idx, batch_idx = Int32(0), Int32(0), Int32(0) + next_tile_idx = self._tile_idx // params.cluster_shape_m + while group_end_tile <= next_tile_idx: + batch_idx += cute.arch.WARP_SIZE - 1 + if batch_idx >= params.num_batch: + batch_idx = Int32(params.num_batch) + group_end_tile = next_tile_idx + 1 + else: + num_m_blocks = self._get_num_m_blocks(lane_idx, bidb_start=batch_idx) + num_m_blocks_cumulative = utils.warp_prefix_sum(num_m_blocks, lane_idx) + m_blocks_in_group = cute.arch.shuffle_sync( + num_m_blocks_cumulative, cute.arch.WARP_SIZE - 1 + ) + group_end_tile += m_blocks_in_group * params.num_head + is_valid = False + if batch_idx >= params.num_batch: + block, head_idx, batch_idx = Int32(0), Int32(0), Int32(params.num_batch) + else: + group_start_tile = group_end_tile - m_blocks_in_group * params.num_head + # if cute.arch.thread_idx()[0] == 128 + 31: cute.printf("SingleTileVarlenScheduler: tile_idx=%d, group_end_tile = %d, num_m_blocks=%d, batch_idx = %d", self._tile_idx, group_end_tile, num_m_blocks, batch_idx) + # The next problem to process is the first one that does not have ending tile position + # that is greater than or equal to tile index. + batch_idx_in_group = cute.arch.popc( + cute.arch.vote_ballot_sync( + group_start_tile + num_m_blocks_cumulative * params.num_head <= next_tile_idx + ) + ) + batch_idx += batch_idx_in_group + num_m_blocks_prev_lane = ( + 0 + if batch_idx_in_group == 0 + else cute.arch.shuffle_sync(num_m_blocks_cumulative, batch_idx_in_group - 1) + ) + num_m_blocks = cute.arch.shuffle_sync(num_m_blocks, batch_idx_in_group) + mh_block = next_tile_idx - group_start_tile - num_m_blocks_prev_lane * params.num_head + if cutlass.const_expr(params.lpt or params.head_swizzle): + # This is a version of the SingleTileLPTScheduler, complicated by the fact that + # the seqlen can vary per batch. + # TODO: is there any case where num_m_blocks is 0? + # TODO: by right we should read the seqlen_kv but we're assuming seqlen_q == seqlen_k here + num_n_blocks = ( + num_m_blocks + * params.tile_shape_mn[0] + * params.cluster_shape_m + // params.qhead_per_kvhead_packgqa + // params.tile_shape_mn[1] + ) + # nheads_in_l2 = min(max(self.max_kvblock_in_l2 // num_n_blocks, 1), self.num_head) + # Seems faster to have this be a power of 2 + nheads_in_l2 = ( + 16 + if num_n_blocks * 16 <= params.max_kvblock_in_l2 + else ( + 8 + if num_n_blocks * 8 <= params.max_kvblock_in_l2 + else ( + 4 + if num_n_blocks * 4 <= params.max_kvblock_in_l2 + else (2 if num_n_blocks * 2 <= params.max_kvblock_in_l2 else 1) + ) + ) + ) + nheads_in_l2 = min(nheads_in_l2, params.num_head) + mh_in_l2 = nheads_in_l2 * num_m_blocks + section_idx = mh_block // mh_in_l2 + l2_mod = mh_block - section_idx * mh_in_l2 + # Deal with tail section + nheads_in_this_section = ( + nheads_in_l2 + if nheads_in_l2 * (section_idx + 1) <= params.num_head + else params.num_head - section_idx * nheads_in_l2 + ) + block = l2_mod // nheads_in_this_section + head_idx_residual = l2_mod - block * nheads_in_this_section + head_idx = section_idx * nheads_in_l2 + head_idx_residual + if cutlass.const_expr(params.lpt): + block = num_m_blocks - 1 - block + else: + head_idx = mh_block // num_m_blocks + block = mh_block - head_idx * num_m_blocks + is_valid = self._is_first_block and batch_idx < params.num_batch + if cutlass.const_expr(params.cluster_shape_m > 1 and not params.use_cluster_idx): + bidx_in_cluster = cute.arch.block_in_cluster_idx() + block = block * params.cluster_shape_m + bidx_in_cluster[0] + # if cute.arch.thread_idx()[0] == 128: cute.printf("SingleTileVarlenScheduler: tile_idx=%d, batch_idx=%d, head_idx=%d, block=%d, is_valid = %d", self._tile_idx, batch_idx, head_idx, block, is_valid) + split_idx = self._split_idx if const_expr(params.is_split_kv) else Int32(0) + return WorkTileInfo((Int32(block), Int32(head_idx), Int32(batch_idx), split_idx), is_valid) + + @cute.jit + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + clc_work = self.clc.get_current_work() + # Default to grid_dim (one past last valid flat index) so _varlen_coord_map + # returns is_valid=False when CLC is exhausted. CLC tile_idx is garbage when + # invalid, so we can't trust it. Local-then-assign avoids CuTe DSL structural + # mismatch on self inside the runtime if. + new_tile_idx = cute.arch.grid_dim()[0] + new_split_idx = Int32(0) + if clc_work.is_valid_tile: + new_tile_idx = clc_work.tile_idx[0] + if const_expr(self.params.is_split_kv): + new_split_idx = clc_work.tile_idx[1] + self._tile_idx = new_tile_idx + self._split_idx = new_split_idx + return self._varlen_coord_map() + + @cute.jit + def initial_work_tile_info(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + clc_work = self.clc.initial_work_tile_info() + # See get_current_work for why grid_dim and local-then-assign. + new_tile_idx = cute.arch.grid_dim()[0] + new_split_idx = Int32(0) + if clc_work.is_valid_tile: + new_tile_idx = clc_work.tile_idx[0] + if const_expr(self.params.is_split_kv): + new_split_idx = clc_work.tile_idx[1] + self._tile_idx = new_tile_idx + self._split_idx = new_split_idx + return self._varlen_coord_map() + + def prefetch_next_work(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + self.clc.prefetch_next_work(loc=loc, ip=ip) + + def advance_to_next_work(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + self.clc.consumer_wait(loc=loc, ip=ip) + work = self.get_current_work() + self.clc.consumer_release(loc=loc, ip=ip) + return work + self._is_first_block = False + return self.get_current_work() + + def producer_tail(self, *, loc=None, ip=None): + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + self.clc.producer_tail(loc=loc, ip=ip) + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + objs = [self.params, self._tile_idx, self._split_idx] + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + objs += [self.clc] + for obj in objs: + obj_values = cutlass.extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + objs = [self.params, self._tile_idx, self._split_idx] + if const_expr(self.params.scheduling_mode == SchedulingMode.CLC): + objs += [self.clc] + for obj, n_items in zip(objs, self._values_pos): + obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return self.__class__(*obj_list, loc=self._loc) + + +# ----------------------------------------------------------------------------- +# SM100 FMHA-specific schedulers (kept separate from generic schedulers). +# ----------------------------------------------------------------------------- + + +class Sm100FmhaStaticTileSchedulerParams: + """A class to represent parameters for the FMHA (Fused Multi-Head Attention) static tile scheduler. + + This class holds the configuration parameters needed to initialize and configure + the tile scheduler for FMHA operations. + + :ivar is_persistent: Whether to use persistent kernel mode. + :type is_persistent: bool + :ivar problem_shape_mbh: Problem shape in (M, B, H) format. + :type problem_shape_mbh: cute.Shape + """ + + def __init__( + self, + is_persistent: bool, + problem_shape_mbh: cute.Shape, + *, + loc=None, + ip=None, + ): + """ + Initializes the Sm100FmhaStaticTileSchedulerParams with the given parameters. + + :param is_persistent: Whether to use persistent kernel mode. + :type is_persistent: bool + :param problem_shape_mbh: Problem shape in (M, B, H) format. + :type problem_shape_mbh: cute.Shape + """ + self.is_persistent = is_persistent + self.problem_shape_mbh = problem_shape_mbh + self._loc = loc + self._ip = ip + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in [self.problem_shape_mbh]: + obj_values = extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + for obj, n_items in zip([self.problem_shape_mbh], self._values_pos): + obj_list.append(new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return Sm100FmhaStaticTileSchedulerParams( + self.is_persistent, *(tuple(obj_list)), loc=self._loc + ) + + +class Sm100FmhaStaticTileScheduler: + """A static tile scheduler for FMHA (Fused Multi-Head Attention) operations. + + This class manages the scheduling of work tiles for FMHA kernels, supporting + both persistent and non-persistent kernel modes. It tracks the current work + position and advances through the problem space efficiently. + + :ivar _params: Scheduler parameters. + :type _params: Sm100FmhaStaticTileSchedulerParams + :ivar _blk_coord: Block coordinates. + :type _blk_coord: cute.Coord + :ivar _grid_shape: Grid shape for the kernel. + :type _grid_shape: cute.Shape + :ivar _is_persistent: Whether to use persistent kernel mode. + :type _is_persistent: bool + :ivar _current_work_linear_idx: Current linear work index. + :type _current_work_linear_idx: Int32 + :ivar _problem_shape_mbh: Problem shape in (M, B, H) format. + :type _problem_shape_mbh: cute.Layout + :ivar _num_blocks: Number of blocks in the problem. + :type _num_blocks: Int32 + :ivar _is_first_block: Whether this is the first block. + :type _is_first_block: bool + :ivar num_persistent_sm: Number of persistent SMs. + :type num_persistent_sm: Int32 + """ + + def __init__( + self, + params: Sm100FmhaStaticTileSchedulerParams, + current_work_linear_idx: Int32, + blk_coord: cute.Coord, + grid_shape: cute.Shape, + *, + loc=None, + ip=None, + ): + """ + Initializes the Sm100FmhaStaticTileScheduler with the given parameters. + + :param params: Scheduler parameters. + :type params: Sm100FmhaStaticTileSchedulerParams + :param current_work_linear_idx: Current linear work index. + :type current_work_linear_idx: Int32 + :param blk_coord: Block coordinates. + :type blk_coord: cute.Coord + :param grid_shape: Grid shape for the kernel. + :type grid_shape: cute.Shape + """ + self._params = params + self._blk_coord = blk_coord + self._grid_shape = grid_shape + self._is_persistent = params.is_persistent + self._current_work_linear_idx = current_work_linear_idx + self._problem_shape_mbh = cute.make_layout(params.problem_shape_mbh, loc=loc, ip=ip) + self._num_blocks = cute.size(self._problem_shape_mbh, loc=loc, ip=ip) + self._is_first_block = True + self.num_persistent_sm = cute.size(grid_shape, loc=loc, ip=ip) + self._loc = loc + self._ip = ip + + # called by host + @staticmethod + def get_grid_shape( + params: Sm100FmhaStaticTileSchedulerParams, + *, + loc=None, + ip=None, + ) -> cute.Shape: + """ + Determine the grid shape for the FMHA kernel. + + For persistent kernels, the grid shape is limited by the number of SMs + (Streaming Multiprocessors) available on the device. For non-persistent + kernels, the grid shape matches the problem shape. + + :param params: Scheduler parameters. + :type params: Sm100FmhaStaticTileSchedulerParams + + :return: Grid shape as (M, B, H) tuple. + :rtype: cute.Shape + """ + if params.is_persistent: + hardware_info = HardwareInfo() + sm_count = hardware_info.get_device_multiprocessor_count() + return ( + dsl_min(sm_count, cute.size(params.problem_shape_mbh, loc=loc, ip=ip)), + 1, + 1, + ) + else: + return params.problem_shape_mbh + + @staticmethod + def check_valid_work_for_seqlen_q( + q_tiler: int, + current_idx: Int32, + seqlen_q: Int32, + ) -> Boolean: + """ + Check if the current work index is valid for the given query sequence length. + + This method verifies that the current work tile index multiplied by the + query tiler size is within the bounds of the query sequence length. + + :param q_tiler: Query tiler size. + :type q_tiler: int + :param current_idx: Current work index. + :type current_idx: Int32 + :param seqlen_q: Query sequence length. + :type seqlen_q: Int32 + + :return: True if the work is valid, False otherwise. + :rtype: Boolean + """ + return current_idx * q_tiler < seqlen_q + + def get_current_work(self, *, loc=None, ip=None) -> cutlass.utils.WorkTileInfo: + """ + Get information about the current work tile. + + Determines if the current work is valid and computes the tile coordinates + based on whether the kernel is persistent or non-persistent. + + :return: WorkTileInfo containing tile coordinates and validity flag. + :rtype: WorkTileInfo + """ + is_valid = ( + self._current_work_linear_idx < self._num_blocks + if self._is_persistent + else self._is_first_block + ) + + blk_coord = (0, 0, 0) + if self._is_persistent: + blk_coord = self._problem_shape_mbh.get_hier_coord( + self._current_work_linear_idx, loc=loc, ip=ip + ) + else: + blk_coord = self._blk_coord + + # cur_tile_coord is (mid, 0, (bid, hid)) + cur_tile_coord = ( + blk_coord[0], + 0, + (blk_coord[1], blk_coord[2]), + ) + + return cutlass.utils.WorkTileInfo(cur_tile_coord, is_valid) + + def initial_work_tile_info(self, *, loc=None, ip=None): + """ + Get the initial work tile information. + + :return: Initial WorkTileInfo. + :rtype: WorkTileInfo + """ + return self.get_current_work(loc=loc, ip=ip) + + def advance_to_next_work(self, *, advance_count=1, loc=None, ip=None): + """ + Advance to the next work tile and return it. + + For persistent kernels, advances by the number of persistent SMs. + For non-persistent kernels, marks that the first block has been processed. + """ + if self._is_persistent: + self._current_work_linear_idx += advance_count * self.num_persistent_sm + self._is_first_block = False + return self.get_current_work() + + def prefetch_next_work(self, *, loc=None, ip=None): + """No-op for static scheduler.""" + pass + + def producer_tail(self, *, loc=None, ip=None): + """No-op for static scheduler.""" + pass + + def __extract_mlir_values__(self): + values = extract_mlir_values(self._params) + values.extend(extract_mlir_values(self._current_work_linear_idx)) + values.extend(extract_mlir_values(self._blk_coord)) + values.extend(extract_mlir_values(self._grid_shape)) + return values + + def __new_from_mlir_values__(self, values): + assert len(values) == 10 + new_params = new_from_mlir_values(self._params, values[0:3]) + new_current_work_linear_idx = new_from_mlir_values( + self._current_work_linear_idx, [values[3]] + ) + new_blk_coord = new_from_mlir_values(self._blk_coord, values[4:7]) + new_grid_shape = new_from_mlir_values(self._grid_shape, values[7:]) + return Sm100FmhaStaticTileScheduler( + new_params, new_current_work_linear_idx, new_blk_coord, new_grid_shape + ) + + +def compute_sm100_fmha_grid( + o_shape: cute.Shape, + cta_tiler: Tuple[int, int, int], + is_persistent: bool, +) -> Tuple[Sm100FmhaStaticTileSchedulerParams, Tuple[int, int, int]]: + """Compute grid parameters for FMHA (static scheduler). + + The output tensor o has shape (s, d, ((h_r, h_k), b)). + """ + tile_sched_params = Sm100FmhaStaticTileSchedulerParams( + is_persistent, + ( + cute.ceil_div(cute.size(o_shape[0]), cta_tiler[0]), + cute.size(o_shape[2][0]), + cute.size(o_shape[2][1]), + ), + ) + grid = Sm100FmhaStaticTileScheduler.get_grid_shape(tile_sched_params) + return tile_sched_params, grid + + +############################################################################## +# Fmha CLC dynamic tile scheduler +############################################################################## + + +class Sm100FmhaClcDynamicTileSchedulerParams: + """Parameters for FMHA CLC dynamic persistent tile scheduler. + + This class manages the layout of tiles for CLC (Cluster Launch Control) + based dynamic scheduling, adapted for FMHA's (M, B, H) problem shape. + + :ivar problem_shape_mbh: Problem shape in (M, B, H) format. + :type problem_shape_mbh: cute.Shape + :ivar cluster_shape_mnk: Cluster shape in (M, N, K) format. + :type cluster_shape_mnk: cute.Shape + """ + + def __init__( + self, + problem_shape_mbh: cute.Shape, + cluster_shape_mnk: cute.Shape, + *, + loc=None, + ip=None, + ): + self.problem_shape_mbh = problem_shape_mbh + self._cluster_shape_mnk = cluster_shape_mnk + self.cluster_shape_mn = cluster_shape_mnk[:2] + self._loc = loc + self._ip = ip + + # FMHA uses linear indexing over (M, B, H), convert to (M, N, L) style + # For FMHA: M dim is tile count along sequence, N=1, L=(B*H) + self.problem_shape_ntile_mnl = ( + problem_shape_mbh[0], # M tiles + 1, # N tiles (always 1 for FMHA) + problem_shape_mbh[1] * problem_shape_mbh[2], # L = B * H + ) + + # Create layout for cluster-to-tile mapping + self.problem_layout_ncluster_mnl = cute.make_layout( + cute.ceil_div(self.problem_shape_ntile_mnl, cluster_shape_mnk[:2]), + loc=loc, + ip=ip, + ) + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in [ + self.problem_shape_mbh, + self._cluster_shape_mnk, + ]: + obj_values = extract_mlir_values(obj) + values += obj_values + self._values_pos.append(len(obj_values)) + return values + + def __new_from_mlir_values__(self, values): + obj_list = [] + values_copy = list(values) + for obj, n_items in zip( + [self.problem_shape_mbh, self._cluster_shape_mnk], + self._values_pos, + ): + obj_list.append(new_from_mlir_values(obj, values_copy[:n_items])) + values_copy = values_copy[n_items:] + return Sm100FmhaClcDynamicTileSchedulerParams(*(tuple(obj_list)), loc=self._loc) + + def get_grid_shape(self, *, loc=None, ip=None) -> Tuple[int, int, int]: + """Compute grid shape aligned with cluster shape.""" + return cute.round_up(self.problem_shape_ntile_mnl, self._cluster_shape_mnk) + + def clc_hw_params(self) -> ClcDynamicPersistentTileSchedulerParams: + """Return params for the upstream CLC hardware scheduler.""" + return ClcDynamicPersistentTileSchedulerParams( + problem_shape_ntile_mnl=self.problem_shape_ntile_mnl, + cluster_shape_mnk=self._cluster_shape_mnk, + ) + + +class Sm100FmhaClcDynamicTileScheduler: + """CLC dynamic persistent tile scheduler for FMHA. + + This scheduler uses Blackwell's Cluster Launch Control hardware mechanism + for dynamic tile distribution, providing automatic load balancing. + Adapted for FMHA's (M, B, H) problem shape. + """ + + def __init__( + self, + params: Sm100FmhaClcDynamicTileSchedulerParams, + cta_id_in_cluster: cute.Coord, + num_tiles_executed: Int32, + clc_response_ptr: cute.Pointer, + block_idx: Tuple, + clc: ClcState = None, + *, + loc=None, + ip=None, + ): + self.params = params + self.cta_id_in_cluster = cta_id_in_cluster + self._num_tiles_executed = num_tiles_executed + self._clc_response_ptr = clc_response_ptr + self._block_idx = block_idx + self.clc = clc + self._loc = loc + self._ip = ip + + def __extract_mlir_values__(self): + values = extract_mlir_values(self.cta_id_in_cluster) + values.extend(extract_mlir_values(self._num_tiles_executed)) + values.extend(extract_mlir_values(self._clc_response_ptr)) + values.extend(extract_mlir_values(self._block_idx)) + if self.clc is not None: + values.extend(extract_mlir_values(self.clc)) + return values + + def __new_from_mlir_values__(self, values): + new_cta_id_in_cluster = new_from_mlir_values(self.cta_id_in_cluster, values[0:3]) + new_num_tiles_executed = new_from_mlir_values(self._num_tiles_executed, [values[3]]) + new_clc_response_ptr = new_from_mlir_values(self._clc_response_ptr, [values[4]]) + new_block_idx = new_from_mlir_values(self._block_idx, values[5:8]) + new_clc = None + if self.clc is not None: + new_clc = new_from_mlir_values(self.clc, values[8:]) + return Sm100FmhaClcDynamicTileScheduler( + self.params, + new_cta_id_in_cluster, + new_num_tiles_executed, + new_clc_response_ptr, + new_block_idx, + new_clc, + ) + + @staticmethod + def create( + params: Sm100FmhaClcDynamicTileSchedulerParams, + block_idx: Tuple, + grid_dim: Tuple, + clc_response_ptr: cute.Pointer, + clc: ClcState = None, + *, + loc=None, + ip=None, + ): + """Create a CLC dynamic tile scheduler instance.""" + bidx, bidy, bidz = block_idx + + # CTA id in cluster + cta_id_in_cluster = ( + Int32(bidx % params.cluster_shape_mn[0]), + Int32(bidy % params.cluster_shape_mn[1]), + Int32(0), + ) + + num_tiles_executed = Int32(0) + + return Sm100FmhaClcDynamicTileScheduler( + params, + cta_id_in_cluster, + num_tiles_executed, + clc_response_ptr, + block_idx, + clc, + ) + + @staticmethod + def get_grid_shape( + params: Sm100FmhaClcDynamicTileSchedulerParams, + *, + loc=None, + ip=None, + ) -> Tuple[int, int, int]: + """Get grid shape for kernel launch.""" + return params.get_grid_shape(loc=loc, ip=ip) + + def work_tile_info_from_clc_response(self, result_addr: cute.Pointer, *, loc=None, ip=None): + """Parse CLC response and convert to FMHA tile coordinates.""" + m_idx, n_idx, l_idx, vld = cute.arch.clc_response(result_addr, loc=loc, ip=ip) + cute.arch.fence_proxy("async.shared", space="cta") + + # CLC returns first CTA coordinates: m_idx=x, l_idx=z + # l_idx is the L (batch) dimension; decode to (bid, hid) + hid = l_idx % self.params.problem_shape_mbh[2] + bid = l_idx // self.params.problem_shape_mbh[2] + + cta_idx_in_cluster, cta_idy_in_cluster, _ = self.cta_id_in_cluster + cur_tile_coord = ( + m_idx + cta_idx_in_cluster, # M dimension + 0, # N always 0 for FMHA + (bid, hid), # (B, H) packed + ) + + return cutlass.utils.WorkTileInfo(cur_tile_coord, vld) + + def get_current_work(self, *, loc=None, ip=None): + """Get current work tile from CLC response.""" + return self.work_tile_info_from_clc_response(self._clc_response_ptr, loc=loc, ip=ip) + + def initial_work_tile_info(self, *, loc=None, ip=None): + """Get initial work tile based on block index.""" + bidx, bidy, bidz = self._block_idx + # bidz is the L (batch) dimension; decode to (bid, hid) + hid = bidz % self.params.problem_shape_mbh[2] + bid = bidz // self.params.problem_shape_mbh[2] + return cutlass.utils.WorkTileInfo((bidx, 0, (bid, hid)), True) + + def advance_to_next_work(self, *, loc=None, ip=None): + """Consumer-side advance: wait for next tile, read coordinates, release.""" + self.clc.consumer_wait(loc=loc, ip=ip) + work = self.get_current_work(loc=loc, ip=ip) + self.clc.consumer_release(loc=loc, ip=ip) + self._num_tiles_executed += Int32(1) + return work + + def prefetch_next_work(self, *, loc=None, ip=None): + """Producer-side: issue CLC query for next tile.""" + self.clc.prefetch_next_work(loc=loc, ip=ip) + + def producer_tail(self, *, loc=None, ip=None): + """Producer-side cleanup after last tile.""" + self.clc.producer_tail(loc=loc, ip=ip) + + @property + def num_tiles_executed(self) -> Int32: + return self._num_tiles_executed + + +def compute_sm100_fmha_grid_clc( + o_shape: cute.Shape, + cta_tiler: Tuple[int, int, int], + cluster_shape_mnk: Tuple[int, int, int], +) -> Tuple[Sm100FmhaClcDynamicTileSchedulerParams, Tuple[int, int, int]]: + """Compute grid parameters for FMHA with CLC dynamic scheduling.""" + problem_shape_mbh = ( + cute.ceil_div(cute.size(o_shape[0]), cta_tiler[0]), + cute.size(o_shape[2][0]), + cute.size(o_shape[2][1]), + ) + tile_sched_params = Sm100FmhaClcDynamicTileSchedulerParams(problem_shape_mbh, cluster_shape_mnk) + grid = Sm100FmhaClcDynamicTileScheduler.get_grid_shape(tile_sched_params) + return tile_sched_params, grid + + +############################################################################## +# Fused Mask +############################################################################## + + +def make_sm100_thread_cooperative_group(size: int): + return cutlass.pipeline.CooperativeGroup(cutlass.pipeline.Agent.Thread, size) + + +SM100_TMEM_CAPACITY_COLUMNS = 512 diff --git a/fastvideo-kernel/fa4/topk_gather_kv.py b/fastvideo-kernel/fa4/topk_gather_kv.py new file mode 100644 index 000000000..680af3614 --- /dev/null +++ b/fastvideo-kernel/fa4/topk_gather_kv.py @@ -0,0 +1,278 @@ +from typing import Type, Optional +from dataclasses import dataclass +import operator + +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +from cutlass.cute.nvgpu import cpasync +from cutlass import Int32, Uint32, const_expr, Boolean + +from flash_attn.cute import utils +from flash_attn.cute.utils import warp_reduce +from quack.cute_dsl_utils import ParamsBase + +import math + + +@dataclass +class CpasyncGatherKVManager(ParamsBase): + mIndexTopk: cute.Tensor + sBitmask: Optional[cute.Tensor] + + cta_rank_in_cluster: Int32 + thread_idx: Int32 + warp_idx: Int32 + + topk_length: Int32 + seqlen_k_limit: Int32 + tile_n: Int32 + num_threads: cutlass.Constexpr[Int32] + hdim: cutlass.Constexpr[Int32] + hdim_v: cutlass.Constexpr[Int32] + num_hdimv_splits: cutlass.Constexpr[Int32] + cta_group_size: cutlass.Constexpr[Int32] + + gmem_threads_per_row: cutlass.Constexpr[Int32] + topk_indices_per_thread: Int32 + async_copy_elems: Int32 + + gmem_tiled_copy_KV: cute.TiledCopy + gmem_thr_copy_KV: cute.TiledCopy + + rTopk: cute.Tensor + rTopkHalf: cute.Tensor + # for bitmask + rTopk_NonInterleaved: cute.Tensor + + pipeline_bitmask: Optional[pipeline.PipelineAsync] + cpasync_barrier: Optional[pipeline.NamedBarrier] + + disable_bitmask: cutlass.Constexpr[Boolean] + + @staticmethod + def create( + mIndexTopk: cute.Tensor, + cta_rank_in_cluster: Int32, + thread_idx: Int32, + warp_idx: Int32, + topk_length: Int32, + seqlen_k_limit: Int32, + tile_n: cutlass.Constexpr[Int32], + hdim: cutlass.Constexpr[Int32], + hdim_v: cutlass.Constexpr[Int32], + num_hdimv_splits: cutlass.Constexpr[Int32], + num_threads: cutlass.Constexpr[Int32], + dtype: Type[cutlass.Numeric], + cta_group_size: cutlass.Constexpr[Int32], + cpasync_barrier: Optional[pipeline.NamedBarrier] = None, + disable_bitmask: cutlass.Constexpr[Boolean] = False, + sBitmask: Optional[cute.Tensor] = None, + pipeline_bitmask: Optional[pipeline.PipelineAsync] = None, + ): + assert tile_n % num_threads == 0 + assert num_threads == 128 + assert hdim % 64 == 0 + assert (hdim_v // num_hdimv_splits // cta_group_size) % 64 == 0 + assert num_threads % cute.arch.WARP_SIZE == 0 + universal_copy_bits = 128 + async_copy_elems = universal_copy_bits // dtype.width + dtype_bytes = dtype.width // 8 + # assumes hdim is never part of transposed operand + gmem_k_block_size = math.gcd( + hdim, + hdim_v // num_hdimv_splits // cta_group_size, + 128 // dtype_bytes, + ) + assert gmem_k_block_size % async_copy_elems == 0 + gmem_threads_per_row = gmem_k_block_size // async_copy_elems + assert cute.arch.WARP_SIZE % gmem_threads_per_row == 0 + atom_async_copy = cute.make_copy_atom( + cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + dtype, + num_bits_per_copy=universal_copy_bits, + ) + thr_layout = cute.make_ordered_layout( + (num_threads // gmem_threads_per_row, gmem_threads_per_row), + order=(1, 0), + ) + val_layout = cute.make_layout((1, async_copy_elems)) + gmem_tiled_copy_KV = cute.make_tiled_copy_tv(atom_async_copy, thr_layout, val_layout) + gmem_thr_copy_KV = gmem_tiled_copy_KV.get_slice(thread_idx) + topk_indices_per_thread = tile_n // num_threads + + rTopk = cute.make_rmem_tensor((topk_indices_per_thread,), Int32) + rTopkHalf = cute.make_rmem_tensor((topk_indices_per_thread,), Int32) + rTopk_NonInterleaved = cute.make_rmem_tensor((topk_indices_per_thread,), Int32) + + return CpasyncGatherKVManager( + mIndexTopk, + sBitmask, + cta_rank_in_cluster, + thread_idx, + warp_idx, + topk_length, + seqlen_k_limit, + tile_n, + num_threads, + hdim, + hdim_v, + num_hdimv_splits, + cta_group_size, + gmem_threads_per_row, + topk_indices_per_thread, + async_copy_elems, + gmem_tiled_copy_KV, + gmem_thr_copy_KV, + rTopk, + rTopkHalf, + rTopk_NonInterleaved, + pipeline_bitmask, + cpasync_barrier, + disable_bitmask, + ) + + @cute.jit + def load_index_topk( + self, + n_block: Int32, + transpose: bool, + ): + entries_per_thread = self.topk_indices_per_thread + rTopk = self.rTopk if const_expr(transpose) else self.rTopkHalf + + for i in cutlass.range_constexpr(entries_per_thread): + row = ( + i * self.num_threads + + (self.thread_idx % self.gmem_threads_per_row) + * (self.num_threads // self.gmem_threads_per_row) + + (self.thread_idx // self.gmem_threads_per_row) + ) + # need this if not offset in load_X + # if const_expr(not transpose): + # row += self.cta_rank_in_cluster * (self.tile_n//self.cta_group_size) + # row = row % self.tile_n + row_idx = n_block * self.tile_n + row + rTopk[i] = self.mIndexTopk[row_idx] + + if const_expr(not transpose and not self.disable_bitmask): + row_non_interleaved = i * self.num_threads + self.thread_idx + row_idx_non_interleaved = n_block * self.tile_n + row_non_interleaved + self.rTopk_NonInterleaved[0] = self.mIndexTopk[row_idx_non_interleaved] + + @cute.jit + def compute_bitmask( + self, + producer_state_bitmask, + ): + assert self.pipeline_bitmask is not None, "pipeline_bitmask not provided" + assert self.cpasync_barrier is not None, "cpasync barrier not provided" + + lane_idx = cute.arch.lane_idx() + assert cute.size(self.rTopk_NonInterleaved) == 1 + bitmask = Uint32(0) + + # Step 1. Construct per-thread bitmask + topk_idx = self.rTopk_NonInterleaved[0] + is_valid = topk_idx >= 0 and topk_idx < self.seqlen_k_limit + if is_valid: + bitmask = Uint32(1 << lane_idx) + + # Step 2. Warp shuffle bitwise OR = add since indices are exclusive. + bitmask = warp_reduce(bitmask, operator.add) + + self.pipeline_bitmask.producer_acquire(producer_state_bitmask) + # store to smem and sync threads + if lane_idx == 0: + self.sBitmask[self.warp_idx, producer_state_bitmask.index] = bitmask + self.cpasync_barrier.arrive_and_wait() + + self.pipeline_bitmask.producer_commit(producer_state_bitmask) + producer_state_bitmask.advance() + return producer_state_bitmask + + @cute.jit + def compute_X_ptr( + self, + mX: cute.Tensor, + transpose: bool, + d_offset: int = 0, + ): + entries_per_thread = self.topk_indices_per_thread + tPrXPtr = cute.make_rmem_tensor((entries_per_thread,), cutlass.Int64) + tPrRowValid = cute.make_rmem_tensor((entries_per_thread,), cutlass.Int32) + rTopk = self.rTopk if const_expr(transpose) else self.rTopkHalf + + for i in cutlass.range_constexpr(entries_per_thread): + topk_idx = rTopk[i] + if const_expr(not self.disable_bitmask): + row_valid = topk_idx >= 0 and topk_idx < self.seqlen_k_limit + tPrRowValid[i] = row_valid + if const_expr(not transpose): + tPrXPtr[i] = utils.elem_pointer(mX, (topk_idx, d_offset)).toint() + else: + tPrXPtr[i] = utils.elem_pointer(mX, (d_offset, topk_idx)).toint() + + return tPrXPtr, tPrRowValid + + @cute.jit + def load_X( + self, + mX: cute.Tensor, + sX: cute.Tensor, + transpose: bool, + K_or_V: str, + d_offset: int = 0, + ): + assert K_or_V in ("K", "V") + cta_tile_n = self.tile_n if const_expr(transpose) else self.tile_n // self.cta_group_size + head_dim = self.hdim if const_expr(K_or_V == "K") else self.hdim_v // self.num_hdimv_splits + if const_expr(transpose): + head_dim = head_dim // self.cta_group_size + order = (1, 0) if const_expr(transpose) else (0, 1) + + sX_nd_layout = cute.make_ordered_layout((cta_tile_n, head_dim), order=order) + sX_nd = cute.composition(sX, sX_nd_layout) + + cX = cute.make_identity_tensor((cta_tile_n, head_dim)) + tXsX = self.gmem_thr_copy_KV.partition_D(sX_nd) + tXcX = self.gmem_thr_copy_KV.partition_S(cX) + + tPrXPtr, tPrRowValid = self.compute_X_ptr(mX, transpose, d_offset) + + if const_expr(not transpose): + offset = self.cta_rank_in_cluster * (self.gmem_threads_per_row // self.cta_group_size) + else: + offset = 0 + + for m in cutlass.range_constexpr(cute.size(tXsX, mode=[1])): + if const_expr(not self.disable_bitmask): + row_valid = utils.shuffle_sync( + tPrRowValid[m // self.gmem_threads_per_row], + (m + offset) % self.gmem_threads_per_row, + width=self.gmem_threads_per_row, + ) + should_load = cute.make_fragment_like(tXsX[(0, None), m, 0], Boolean) + should_load.fill(Boolean(row_valid)) + x_ptr_i64 = utils.shuffle_sync( + tPrXPtr[m // self.gmem_threads_per_row], + (m + offset) % self.gmem_threads_per_row, + width=self.gmem_threads_per_row, + ) + x_gmem_ptr = cute.make_ptr( + mX.element_type, x_ptr_i64, cute.AddressSpace.gmem, assumed_align=16 + ) + mX_cur = cute.make_tensor(x_gmem_ptr, cute.make_layout((head_dim,))) + mX_cur_copy = cute.tiled_divide(mX_cur, (self.async_copy_elems,)) + + for k in cutlass.range_constexpr(cute.size(tXsX, mode=[2])): + ki = tXcX[0, 0, k][1] // self.async_copy_elems + mX_cur_copy_ki = mX_cur_copy[None, ki] + tXsX_k = tXsX[None, m, k] + mX_cur_copy_ki = cute.make_tensor(mX_cur_copy_ki.iterator, tXsX_k.layout) + cute.copy( + self.gmem_tiled_copy_KV, + mX_cur_copy_ki, + tXsX_k, + pred=should_load if const_expr(not self.disable_bitmask) else None, + ) diff --git a/fastvideo-kernel/fa4/utils.py b/fastvideo-kernel/fa4/utils.py new file mode 100644 index 000000000..9ac6ac97b --- /dev/null +++ b/fastvideo-kernel/fa4/utils.py @@ -0,0 +1,957 @@ +# Copyright (c) 2025, Tri Dao. + +import math +import hashlib +import inspect +import os +from functools import partial +from typing import Type, Callable, Optional, Tuple, overload, NamedTuple + +import cutlass +import cutlass.cute as cute + +from cutlass import Float32, Int32, const_expr +from cutlass.cute import FastDivmodDivisor +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass._mlir.dialects import nvvm, llvm +from cutlass.cute.runtime import from_dlpack + + +import quack.activation + +_MIXER_ATTRS = ("__vec_size__",) + + +class AuxData(NamedTuple): + tensors: tuple | list | None = None + scalars: tuple | None = None + + +# Obtained from sollya: +# fpminimax(exp(x * log(2.0)), 1, [|1,24...|],[0;1],relative); +POLY_EX2 = { + 0: (1.0), + 1: ( + 1.0, + 0.922497093677520751953125, + ), + 2: ( + 1.0, + 0.6657850742340087890625, + 0.330107033252716064453125, + ), + 3: ( + 1.0, + 0.695146143436431884765625, + 0.227564394474029541015625, + 0.077119089663028717041015625, + ), + 4: ( + 1.0, + 0.693042695522308349609375, + 0.2412912547588348388671875, + 5.2225358784198760986328125e-2, + 1.3434938155114650726318359375e-2, + ), + 5: ( + 1.0, + 0.693151414394378662109375, + 0.24016360938549041748046875, + 5.5802188813686370849609375e-2, + 9.01452265679836273193359375e-3, + 1.86810153536498546600341796875e-3, + ), +} + +_fa_clc_enabled: bool = os.environ.get("FA_CLC", "0") == "1" +_fa_disable_2cta_enabled: bool = os.environ.get("FA_DISABLE_2CTA", "0") == "1" + + +def _is_cuda_12() -> bool: + """Check if the CUDA toolkit version is 12.x. + + 2CTA forward non-causal has a codegen regression on CUDA 12 that causes + ~18% slowdown compared to 1CTA. This is fixed in CUDA 13.x. + """ + try: + import torch + + cuda_version = torch.version.cuda + if cuda_version is not None: + major = cuda_version.split(".")[0] + return int(major) == 12 + except Exception: + pass + return False + + +_fa_disable_2cta_cuda12: bool = _is_cuda_12() + + +def _get_use_clc_scheduler_default() -> bool: + return _fa_clc_enabled + + +def _get_disable_2cta_default(is_fwd: bool = False) -> bool: + if is_fwd: + return _fa_disable_2cta_enabled or _fa_disable_2cta_cuda12 + else: + return _fa_disable_2cta_enabled + + +def _compute_base_hash(func: Callable) -> str: + """Compute hash from source code or bytecode and closure values.""" + try: + data = inspect.getsource(func).encode() + except (OSError, TypeError): + if hasattr(func, "__code__") and func.__code__ is not None: + data = func.__code__.co_code + else: + data = repr(func).encode() + + hasher = hashlib.sha256(data) + + if hasattr(func, "__closure__") and func.__closure__ is not None: + for cell in func.__closure__: + hasher.update(repr(cell.cell_contents).encode()) + + return hasher.hexdigest() + + +def hash_callable( + func: Callable, mixer_attrs: Tuple[str] = _MIXER_ATTRS, set_cute_hash: bool = True +) -> str: + """Hash a callable based on the source code or bytecode and closure values. + Fast-path: if the callable (or its __wrapped__ base) has a ``__cute_hash__`` + attribute, that value is returned immediately as the base hash, then + metadata dunders are mixed in to produce the final dict-key hash. + set_cute_hash: whether or not to set func.__cute_hash__ + """ + # Resolve base hash + if hasattr(func, "__cute_hash__"): + base_hash = func.__cute_hash__ + else: + # Unwrap decorated functions (e.g., cute.jit wrappers). + base_func = getattr(func, "__wrapped__", func) + + if hasattr(base_func, "__cute_hash__"): + base_hash = base_func.__cute_hash__ + else: + base_hash = _compute_base_hash(base_func) + + if set_cute_hash: + base_func.__cute_hash__ = base_hash + + # Mix in mutable metadata dunders + mixer_values = tuple(getattr(func, attr, None) for attr in mixer_attrs) + + if all(v is None for v in mixer_values): + return base_hash + + hasher = hashlib.sha256(base_hash.encode()) + + for attr, val in zip(mixer_attrs, mixer_values): + hasher.update(f"{attr}={val!r}".encode()) + + return hasher.hexdigest() + + +def create_softcap_scoremod(softcap_val): + @cute.jit + def scoremod_premask_fn( + acc_S_SSA, batch_idx, head_idx, q_idx, kv_idx, seqlen_info, aux_tensors + ): + scores = acc_S_SSA / softcap_val + return softcap_val * cute.math.tanh(scores, fastmath=True) + + return scoremod_premask_fn + + +def create_softcap_scoremod_bwd(softcap_val): + @cute.jit + def scoremod_bwd_fn( + grad_out_SSA, score_SSA, batch_idx, head_idx, q_idx, kv_idx, seqlen_info, aux_tensors + ): + scores = score_SSA / softcap_val + tanh_scores = cute.math.tanh(scores, fastmath=True) + return grad_out_SSA * (1.0 - tanh_scores * tanh_scores) + + return scoremod_bwd_fn + + +LOG2_E = math.log2(math.e) + + +def compute_softmax_scale_log2(softmax_scale, score_mod): + """Compute softmax_scale_log2 and adjusted softmax_scale based on whether score_mod is used. + + When score_mod is None, fold the log2(e) factor into softmax_scale_log2 and set softmax_scale + to None. When score_mod is present, keep softmax_scale separate so it can be applied before + the score_mod, and set softmax_scale_log2 to just the change-of-base constant. + + Returns (softmax_scale_log2, softmax_scale). + """ + if const_expr(score_mod is None): + return softmax_scale * LOG2_E, None + else: + return LOG2_E, softmax_scale + + +def compute_fastdiv_mods(mQ, mK, qhead_per_kvhead, pack_gqa, aux_tensors, mPageTable=None): + """Compute FastDivmodDivisor pairs for aux_tensors index computation. + + Returns a (seqlen_q_divmod, seqlen_k_divmod) tuple, or None if aux_tensors is None. + """ + if const_expr(aux_tensors is None): + return None + seqlen_q = cute.size(mQ.shape[0]) // (qhead_per_kvhead if const_expr(pack_gqa) else 1) + seqlen_k = ( + cute.size(mK.shape[0]) + if const_expr(mPageTable is None) + else mK.shape[0] * mPageTable.shape[1] + ) + return (FastDivmodDivisor(seqlen_q), FastDivmodDivisor(seqlen_k)) + + +def convert_from_dlpack(x, leading_dim, alignment=16, divisibility=1) -> cute.Tensor: + return ( + from_dlpack(x, assumed_align=alignment) + .mark_layout_dynamic(leading_dim=leading_dim) + .mark_compact_shape_dynamic( + mode=leading_dim, stride_order=x.dim_order(), divisibility=divisibility + ) + ) + + +def convert_from_dlpack_compact_dynamic( + x, + *, + dynamic_modes: tuple[int, ...], + alignment: int = 16, + stride_order=None, + divisibility: int = 1, + enable_tvm_ffi: bool = False, +) -> cute.Tensor: + """Convert via DLPack and mark selected compact dimensions as dynamic.""" + if isinstance(dynamic_modes, int): + dynamic_modes = (dynamic_modes,) + if stride_order is None: + stride_order = x.dim_order() + t = ( + from_dlpack(x, assumed_align=alignment, enable_tvm_ffi=True) + if enable_tvm_ffi + else from_dlpack(x, assumed_align=alignment) + ) + for mode in dynamic_modes: + t = t.mark_compact_shape_dynamic( + mode=mode, + stride_order=stride_order, + divisibility=divisibility, + ) + return t + + +def convert_from_dlpack_leading_static( + x, leading_dim, alignment=16, static_modes=None, stride_order=None +) -> cute.Tensor: + if stride_order is None: + stride_order = x.dim_order() + x_ = from_dlpack(x, assumed_align=alignment) + for i in range(x.ndim): + if i != leading_dim and (static_modes is None or i not in static_modes): + x_ = x_.mark_compact_shape_dynamic(mode=i, stride_order=stride_order) + return x_ + + +def make_tiled_copy_A( + copy_atom: cute.CopyAtom, tiled_mma: cute.TiledMma, swapAB: cutlass.Constexpr[bool] = False +) -> cute.TiledCopy: + if const_expr(swapAB): + return cute.make_tiled_copy_B(copy_atom, tiled_mma) + else: + return cute.make_tiled_copy_A(copy_atom, tiled_mma) + + +def make_tiled_copy_B( + copy_atom: cute.CopyAtom, tiled_mma: cute.TiledMma, swapAB: cutlass.Constexpr[bool] = False +) -> cute.TiledCopy: + if const_expr(swapAB): + return cute.make_tiled_copy_A(copy_atom, tiled_mma) + else: + return cute.make_tiled_copy_B(copy_atom, tiled_mma) + + +def mma_make_fragment_A( + smem: cute.Tensor, thr_mma: cute.ThrMma, swapAB: cutlass.Constexpr[bool] = False +) -> cute.Tensor: + if const_expr(swapAB): + return mma_make_fragment_B(smem, thr_mma) + else: + return thr_mma.make_fragment_A(thr_mma.partition_A(smem)) + + +def mma_make_fragment_B( + smem: cute.Tensor, thr_mma: cute.ThrMma, swapAB: cutlass.Constexpr[bool] = False +) -> cute.Tensor: + if const_expr(swapAB): + return mma_make_fragment_A(smem, thr_mma) + else: + return thr_mma.make_fragment_B(thr_mma.partition_B(smem)) + + +def get_smem_store_atom( + arch: cutlass.Constexpr[int], element_type: Type[cute.Numeric], transpose: bool = False +) -> cute.CopyAtom: + if const_expr(arch < 90 or element_type.width != 16): + return cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + element_type, + num_bits_per_copy=2 * element_type.width, + ) + else: + return cute.make_copy_atom( + cute.nvgpu.warp.StMatrix8x8x16bOp(transpose=transpose, num_matrices=4), + element_type, + ) + + +@cute.jit +def warp_reduce( + val: cute.TensorSSA | cute.Numeric, + op: Callable, + width: cutlass.Constexpr[int] = cute.arch.WARP_SIZE, +) -> cute.TensorSSA | cute.Numeric: + if const_expr(isinstance(val, cute.TensorSSA)): + res = cute.make_rmem_tensor(val.shape, val.dtype) + res.store(val) + for i in cutlass.range_constexpr(cute.size(val.shape)): + res[i] = warp_reduce(res[i], op, width) + return res.load() + else: + for i in cutlass.range_constexpr(int(math.log2(width))): + val = op(val, cute.arch.shuffle_sync_bfly(val, offset=1 << i)) + return val + + +@dsl_user_op +def smid(*, loc=None, ip=None) -> Int32: + return Int32( + llvm.inline_asm( + T.i32(), + [], + "mov.u32 $0, %smid;", + "=r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def fmax( + a: float | Float32, b: float | Float32, c: float | Float32 | None = None, *, loc=None, ip=None +) -> Float32: + return Float32( + nvvm.fmax( + Float32(a).ir_value(loc=loc, ip=ip), + Float32(b).ir_value(loc=loc, ip=ip), + c=Float32(c).ir_value(loc=loc, ip=ip) if c is not None else None, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def fmax_reduce( + x: cute.TensorSSA, init_val: float | Float32 | None = None, arch: cutlass.Constexpr[int] = 80 +) -> Float32: + if const_expr(arch < 100 or cute.size(x.shape) % 8 != 0): + # if const_expr(init_val is None): + # init_val = -cutlass.Float32.if + # return x.reduce(cute.ReductionOp.MAX, init_val, 0) + res = cute.make_rmem_tensor(x.shape, Float32) + res.store(x) + # local_max = [res[0], res[1]] + # for i in cutlass.range_constexpr(2, cute.size(x.shape), 2): + # local_max[0] = fmax(local_max[0], res[i + 0]) + # local_max[1] = fmax(local_max[1], res[i + 1]) + # local_max[0] = fmax(local_max[0], local_max[1]) + # return local_max[0] if const_expr(init_val is None) else fmax(local_max[0], init_val) + local_max = [res[0], res[1], res[2], res[3]] + for i in cutlass.range_constexpr(4, cute.size(x.shape), 4): + local_max[0] = fmax(local_max[0], res[i + 0]) + local_max[1] = fmax(local_max[1], res[i + 1]) + local_max[2] = fmax(local_max[2], res[i + 2]) + local_max[3] = fmax(local_max[3], res[i + 3]) + local_max[0] = fmax(local_max[0], local_max[1]) + local_max[2] = fmax(local_max[2], local_max[3]) + local_max[0] = fmax(local_max[0], local_max[2]) + return local_max[0] if const_expr(init_val is None) else fmax(local_max[0], init_val) + else: + # [2025-06-15] x.reduce only seems to use 50% 3-input max and 50% 2-input max + # We instead force the 3-input max. + res = cute.make_rmem_tensor(x.shape, Float32) + res.store(x) + local_max_0 = ( + fmax(init_val, res[0], res[1]) + if const_expr(init_val is not None) + else fmax(res[0], res[1]) + ) + local_max = [ + local_max_0, + fmax(res[2], res[3]), + fmax(res[4], res[5]), + fmax(res[6], res[7]), + ] + for i in cutlass.range_constexpr(8, cute.size(x.shape), 8): + local_max[0] = fmax(local_max[0], res[i], res[i + 1]) + local_max[1] = fmax(local_max[1], res[i + 2], res[i + 3]) + local_max[2] = fmax(local_max[2], res[i + 4], res[i + 5]) + local_max[3] = fmax(local_max[3], res[i + 6], res[i + 7]) + local_max[0] = fmax(local_max[0], local_max[1]) + return fmax(local_max[0], local_max[2], local_max[3]) + + +@cute.jit +def fadd_reduce( + x: cute.TensorSSA, init_val: float | Float32 | None = None, arch: cutlass.Constexpr[int] = 80 +) -> Float32: + if const_expr(arch < 100 or cute.size(x.shape) % 8 != 0): + if const_expr(init_val is None): + init_val = Float32.zero + return x.reduce(cute.ReductionOp.ADD, init_val, 0) + # res = cute.make_rmem_tensor(x.shape, Float32) + # res.store(x) + # local_sum = [res[0], res[1], res[2], res[3]] + # for i in cutlass.range_constexpr(4, cute.size(x.shape), 4): + # local_sum[0] += res[i + 0] + # local_sum[1] += res[i + 1] + # local_sum[2] += res[i + 2] + # local_sum[3] += res[i + 3] + # local_sum[0] += local_sum[1] + # local_sum[2] += local_sum[3] + # local_sum[0] += local_sum[2] + # return local_sum[0] if const_expr(init_val is None) else local_sum[0] + init_val + else: + res = cute.make_rmem_tensor(x.shape, Float32) + res.store(x) + local_sum_0 = ( + cute.arch.add_packed_f32x2((init_val, 0.0), (res[0], res[1])) + # cute.arch.add_packed_f32x2((init_val / 2, init_val / 2), (res[0], res[1])) + if const_expr(init_val is not None) + else (res[0], res[1]) + ) + local_sum = [local_sum_0, (res[2], res[3]), (res[4], res[5]), (res[6], res[7])] + for i in cutlass.range_constexpr(8, cute.size(x.shape), 8): + local_sum[0] = cute.arch.add_packed_f32x2(local_sum[0], (res[i + 0], res[i + 1])) + local_sum[1] = cute.arch.add_packed_f32x2(local_sum[1], (res[i + 2], res[i + 3])) + local_sum[2] = cute.arch.add_packed_f32x2(local_sum[2], (res[i + 4], res[i + 5])) + local_sum[3] = cute.arch.add_packed_f32x2(local_sum[3], (res[i + 6], res[i + 7])) + local_sum[0] = cute.arch.add_packed_f32x2(local_sum[0], local_sum[1]) + local_sum[2] = cute.arch.add_packed_f32x2(local_sum[2], local_sum[3]) + local_sum[0] = cute.arch.add_packed_f32x2(local_sum[0], local_sum[2]) + return local_sum[0][0] + local_sum[0][1] + + +@dsl_user_op +def atomic_add_fp32(a: float | Float32, gmem_ptr: cute.Pointer, *, loc=None, ip=None) -> None: + # gmem_ptr_i64 = gmem_ptr.toint(loc=loc, ip=ip).ir_value() + # # cache_hint = cutlass.Int64(0x12F0000000000000) + # llvm.inline_asm( + # None, + # [gmem_ptr_i64, Float32(a).ir_value(loc=loc, ip=ip)], + # # [gmem_ptr_i64, Float32(a).ir_value(loc=loc, ip=ip), cache_hint.ir_value()], + # "red.global.add.f32 [$0], $1;", + # # "red.global.add.L2::cache_hint.f32 [$0], $1, 0x12F0000000000000;", + # # "red.global.add.L2::cache_hint.f32 [$0], $1, $2;", + # "l,f", + # # "l,f,l", + # has_side_effects=True, + # is_align_stack=False, + # asm_dialect=llvm.AsmDialect.AD_ATT, + # ) + nvvm.atomicrmw( + res=T.f32(), op=nvvm.AtomicOpKind.FADD, ptr=gmem_ptr.llvm_ptr, a=Float32(a).ir_value() + ) + + +@dsl_user_op +def elem_pointer(x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None) -> cute.Pointer: + return x.iterator + cute.crd2idx(coord, x.layout, loc=loc, ip=ip) + + +@cute.jit +def predicate_k(tAcA: cute.Tensor, limit: cutlass.Int32) -> cute.Tensor: + # Only compute predicates for the "k" dimension. For the mn dimension, we will use "if" + tApA = cute.make_rmem_tensor( + cute.make_layout( + (cute.size(tAcA, mode=[0, 1]), cute.size(tAcA, mode=[1]), cute.size(tAcA, mode=[2])), + stride=(cute.size(tAcA, mode=[2]), 0, 1), + ), + cutlass.Boolean, + ) + for rest_v in cutlass.range_constexpr(tApA.shape[0]): + for rest_k in cutlass.range_constexpr(tApA.shape[2]): + tApA[rest_v, 0, rest_k] = cute.elem_less(tAcA[(0, rest_v), 0, rest_k][1], limit) + return tApA + + +def canonical_warp_group_idx(sync: bool = True) -> cutlass.Int32: + warp_group_idx = cute.arch.thread_idx()[0] // 128 + if const_expr(sync): + warp_group_idx = cute.arch.make_warp_uniform(warp_group_idx) + return warp_group_idx + + +# @dsl_user_op +# def warp_vote_any_lt(a: float | Float32, b: float | Float32, *, loc=None, ip=None) -> cutlass.Boolean: +# mask = cutlass.Int32(-1) +# return cutlass.Boolean( +# llvm.inline_asm( +# T.i32(), +# [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip), mask.ir_value(loc=loc, ip=ip)], +# ".pred p1, p2;\n" +# "setp.lt.f32 p1, $1, $2;\n" +# "vote.sync.any.pred p2, p1, $3;\n" +# "selp.u32 $0, 1, 0, p2;", +# # "selp.u32 $0, 1, 0, p1;", +# "=r,f,f,r", +# has_side_effects=False, +# is_align_stack=False, +# asm_dialect=llvm.AsmDialect.AD_ATT, +# ) +# ) + + +@cute.jit +def shuffle_sync( + value: cute.Numeric, + offset: cute.typing.Int, + width: cutlass.Constexpr[int] = cute.arch.WARP_SIZE, +) -> cute.Numeric: + assert value.width % 32 == 0, "value type must be a multiple of 32 bits" + # 1 -> 0b11111, 2 -> 0b11110, 4 -> 0b11100, 8 -> 0b11000, 16 -> 0b10000, 32 -> 0b00000 + mask = cute.arch.WARP_SIZE - width + clamp = cute.arch.WARP_SIZE - 1 + mask_and_clamp = mask << 8 | clamp + # important: need stride 1 and not 0 for recast_tensor to work + val = cute.make_rmem_tensor(cute.make_layout((1,), stride=(1,)), type(value)) + val[0] = value + val_i32 = cute.recast_tensor(val, cutlass.Int32) + for i in cutlass.range_constexpr(cute.size(val_i32)): + val_i32[i] = cute.arch.shuffle_sync(val_i32[i], offset, mask_and_clamp=mask_and_clamp) + return val[0] + + +@dsl_user_op +def shl_u32(val: cutlass.Uint32, shift: cutlass.Uint32, *, loc=None, ip=None) -> cutlass.Uint32: + """ + Left-shift val by shift bits using PTX shl.b32 (sign-agnostic). + + Named ``shl_u32`` (not ``shl_b32``) because python type annotations + distinguish signed/unsigned. + + PTX semantics (§9.7.8.8): "Shift amounts greater than the register width N + are clamped to N." So ``shl.b32 d, a, 32`` is well-defined and yields 0. + + This differs from C/C++ and LLVM IR, where shifting by >= the type width is + undefined behavior. CuTeDSL compiles through MLIR -> LLVM IR, so a plain + Python-level ``Uint32(x) << Uint32(n)`` inherits LLVM's UB: the optimizer + may treat the result as poison and eliminate dependent code. Inline PTX + bypasses the LLVM IR shift entirely — the instruction is emitted verbatim + into PTX where clamping makes it safe for all shift amounts. + """ + return cutlass.Uint32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Uint32(val).ir_value(loc=loc, ip=ip), + cutlass.Uint32(shift).ir_value(loc=loc, ip=ip), + ], + "shl.b32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def shr_u32(val: cutlass.Uint32, shift: cutlass.Uint32, *, loc=None, ip=None) -> cutlass.Uint32: + """ + Unsigned right-shift val by shift bits using PTX shr.u32 (zero-fills). + + See ``shl_u32`` docstring for why inline PTX is used instead of plain + CuTeDSL shift operators (LLVM shift-by-type-width UB). + """ + return cutlass.Uint32( + llvm.inline_asm( + T.i32(), + [ + cutlass.Uint32(val).ir_value(loc=loc, ip=ip), + cutlass.Uint32(shift).ir_value(loc=loc, ip=ip), + ], + "shr.u32 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@cute.jit +def warp_prefix_sum(val: cutlass.Int32, lane: Optional[cutlass.Int32] = None) -> cutlass.Int32: + if const_expr(lane is None): + lane = cute.arch.lane_idx() + # if cute.arch.thread_idx()[0] >= 128 and cute.arch.thread_idx()[0] < 128 + 32 and cute.arch.block_idx()[0] == 0: cute.printf("tidx = %d, val = %d", cute.arch.thread_idx()[0] % 32, val) + for i in cutlass.range_constexpr(int(math.log2(cute.arch.WARP_SIZE))): + offset = 1 << i + # Very important that we set mask_and_clamp to 0 + partial_sum = cute.arch.shuffle_sync_up(val, offset=offset, mask_and_clamp=0) + if lane >= offset: + val += partial_sum + # if cute.arch.thread_idx()[0] >= 128 and cute.arch.thread_idx()[0] < 128 + 32 and cute.arch.block_idx()[0] == 0: cute.printf("tidx = %d, partial_sum = %d, val = %d", cute.arch.thread_idx()[0] % 32, partial_sum, val) + return val + + +@dsl_user_op +def cvt_f16x2_f32( + a: float | Float32, b: float | Float32, to_dtype: Type, *, loc=None, ip=None +) -> cutlass.Int32: + assert to_dtype in [cutlass.BFloat16, cutlass.Float16], "to_dtype must be BFloat16 or Float16" + return cutlass.Int32( + llvm.inline_asm( + T.i32(), + [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)], + f"cvt.rn.{'bf16x2' if to_dtype is cutlass.BFloat16 else 'f16x2'}.f32 $0, $2, $1;", + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@overload +def cvt_f16(src: cute.Tensor, dst: cute.Tensor) -> None: ... + + +@overload +def cvt_f16(src: cute.Tensor, dtype: Type[cute.Numeric]) -> cute.Tensor: ... + + +@cute.jit +def cvt_f16(src: cute.Tensor, dst_or_dtype): + """Convert Float32 tensor to Float16/BFloat16. + + Args: + src: Source tensor with Float32 element type + dst_or_dtype: Either a destination tensor or a dtype (Float16/BFloat16) + + Returns: + None if dst is a tensor, or a new tensor if dtype is provided + """ + if const_expr(isinstance(dst_or_dtype, type)): + # dtype variant: create new tensor and call the tensor variant + dtype = dst_or_dtype + dst = cute.make_rmem_tensor(src.shape, dtype) + cvt_f16(src, dst) + return dst + else: + # tensor variant: write to dst + dst = dst_or_dtype + assert cute.size(dst.shape) == cute.size(src.shape), "dst and src must have the same size" + assert cute.size(src.shape) % 2 == 0, "src must have an even number of elements" + assert dst.element_type in [cutlass.BFloat16, cutlass.Float16], ( + "dst must be BFloat16 or Float16" + ) + assert src.element_type is Float32, "src must be Float32" + dst_i32 = cute.recast_tensor(dst, cutlass.Int32) + assert cute.size(dst_i32.shape) * 2 == cute.size(src.shape) + for i in cutlass.range_constexpr(cute.size(dst_i32)): + dst_i32[i] = cvt_f16x2_f32(src[2 * i], src[2 * i + 1], dst.element_type) + + +@dsl_user_op +@cute.jit +def evaluate_polynomial(x: Float32, poly: Tuple[Float32, ...], *, loc=None, ip=None) -> Float32: + deg = len(poly) - 1 + out = poly[deg] + for i in cutlass.range_constexpr(deg - 1, -1, -1): + out = out * x + poly[i] + return out + + +@dsl_user_op +@cute.jit +def evaluate_polynomial_2( + x: Float32, y: Float32, poly: Tuple[Float32, ...], *, loc=None, ip=None +) -> Tuple[Float32, Float32]: + deg = len(poly) - 1 + out = (poly[deg], poly[deg]) + for i in cutlass.range_constexpr(deg - 1, -1, -1): + out = cute.arch.fma_packed_f32x2(out, (x, y), (poly[i], poly[i])) + return out + + +@dsl_user_op +def add_round_down(x: float | Float32, y: float | Float32, *, loc=None, ip=None) -> Float32: + # There's probably a way to call llvm or nvvm to do this instead of ptx + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [Float32(x).ir_value(loc=loc, ip=ip), Float32(y).ir_value(loc=loc, ip=ip)], + "add.rm.ftz.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def combine_int_frac_ex2(x_rounded: Float32, frac_ex2: Float32, *, loc=None, ip=None) -> Float32: + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [ + Float32(x_rounded).ir_value(loc=loc, ip=ip), + Float32(frac_ex2).ir_value(loc=loc, ip=ip), + ], + "{\n\t" + ".reg .s32 x_rounded_i, frac_ex_i, x_rounded_e, out_i;\n\t" + "mov.b32 x_rounded_i, $1;\n\t" + "mov.b32 frac_ex_i, $2;\n\t" + "shl.b32 x_rounded_e, x_rounded_i, 23;\n\t" + # add.u32 generates IMAD instruction and add.s32 generates LEA instruction + # IMAD uses the FMA pipeline and LEA uses the ALU pipeline, afaik + "add.s32 out_i, x_rounded_e, frac_ex_i;\n\t" + "mov.b32 $0, out_i;\n\t" + "}\n", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def ex2_emulation(x: Float32, *, poly_degree: int = 3, loc=None, ip=None) -> Float32: + assert poly_degree in POLY_EX2, f"Polynomial degree {poly_degree} not supported" + # We assume x <= 127.0 + fp32_round_int = float(2**23 + 2**22) + x_clamped = cute.arch.fmax(x, -127.0) + # We want to round down here, so that the fractional part is in [0, 1) + x_rounded = add_round_down(x_clamped, fp32_round_int, loc=loc, ip=ip) + # The integer floor of x is now in the last 8 bits of x_rounded + # We assume the next 2 ops round to nearest even. The rounding mode is important. + x_rounded_back = x_rounded - fp32_round_int + x_frac = x_clamped - x_rounded_back + x_frac_ex2 = evaluate_polynomial(x_frac, POLY_EX2[poly_degree], loc=loc, ip=ip) + return combine_int_frac_ex2(x_rounded, x_frac_ex2, loc=loc, ip=ip) + + +# TODO: check that the ex2_emulation_2 produces the same SASS as the ptx version +@dsl_user_op +def ex2_emulation_2( + x: Float32, y: Float32, *, poly_degree: int = 3, loc=None, ip=None +) -> Tuple[Float32, Float32]: + # We assume x <= 127.0 and y <= 127.0 + fp32_round_int = float(2**23 + 2**22) + xy_clamped = (cute.arch.fmax(x, -127.0), cute.arch.fmax(y, -127.0)) + # We want to round down here, so that the fractional part is in [0, 1) + xy_rounded = cute.arch.add_packed_f32x2(xy_clamped, (fp32_round_int, fp32_round_int), rnd="rm") + # The integer floor of x & y are now in the last 8 bits of xy_rounded + # We want the next 2 ops to round to nearest even. The rounding mode is important. + xy_rounded_back = quack.activation.sub_packed_f32x2( + xy_rounded, (fp32_round_int, fp32_round_int) + ) + xy_frac = quack.activation.sub_packed_f32x2(xy_clamped, xy_rounded_back) + xy_frac_ex2 = evaluate_polynomial_2(*xy_frac, POLY_EX2[poly_degree], loc=loc, ip=ip) + x_out = combine_int_frac_ex2(xy_rounded[0], xy_frac_ex2[0], loc=loc, ip=ip) + y_out = combine_int_frac_ex2(xy_rounded[1], xy_frac_ex2[1], loc=loc, ip=ip) + return x_out, y_out + + +@dsl_user_op +def e2e_asm2(x: Float32, y: Float32, *, loc=None, ip=None) -> Tuple[Float32, Float32]: + out_f32x2 = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [Float32(x).ir_value(loc=loc, ip=ip), Float32(y, loc=loc, ip=ip).ir_value()], + "{\n\t" + ".reg .f32 f1, f2, f3, f4, f5, f6, f7;\n\t" + ".reg .b64 l1, l2, l3, l4, l5, l6, l7, l8, l9, l10;\n\t" + ".reg .s32 r1, r2, r3, r4, r5, r6, r7, r8;\n\t" + "max.ftz.f32 f1, $2, 0fC2FE0000;\n\t" + "max.ftz.f32 f2, $3, 0fC2FE0000;\n\t" + "mov.b64 l1, {f1, f2};\n\t" + "mov.f32 f3, 0f4B400000;\n\t" + "mov.b64 l2, {f3, f3};\n\t" + "add.rm.ftz.f32x2 l7, l1, l2;\n\t" + "sub.rn.ftz.f32x2 l8, l7, l2;\n\t" + "sub.rn.ftz.f32x2 l9, l1, l8;\n\t" + "mov.f32 f7, 0f3D9DF09D;\n\t" + "mov.b64 l6, {f7, f7};\n\t" + "mov.f32 f6, 0f3E6906A4;\n\t" + "mov.b64 l5, {f6, f6};\n\t" + "mov.f32 f5, 0f3F31F519;\n\t" + "mov.b64 l4, {f5, f5};\n\t" + "mov.f32 f4, 0f3F800000;\n\t" + "mov.b64 l3, {f4, f4};\n\t" + "fma.rn.ftz.f32x2 l10, l9, l6, l5;\n\t" + "fma.rn.ftz.f32x2 l10, l10, l9, l4;\n\t" + "fma.rn.ftz.f32x2 l10, l10, l9, l3;\n\t" + "mov.b64 {r1, r2}, l7;\n\t" + "mov.b64 {r3, r4}, l10;\n\t" + "shl.b32 r5, r1, 23;\n\t" + "add.s32 r7, r5, r3;\n\t" + "shl.b32 r6, r2, 23;\n\t" + "add.s32 r8, r6, r4;\n\t" + "mov.b32 $0, r7;\n\t" + "mov.b32 $1, r8;\n\t" + "}\n", + "=r,=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + out0 = Float32(llvm.extractvalue(T.f32(), out_f32x2, [0], loc=loc, ip=ip)) + out1 = Float32(llvm.extractvalue(T.f32(), out_f32x2, [1], loc=loc, ip=ip)) + return out0, out1 + + +@dsl_user_op +def domain_offset_aligned( + coord: cute.Coord, tensor: cute.Tensor, *, loc=None, ip=None +) -> cute.Tensor: + assert isinstance(tensor.iterator, cute.Pointer) + # We assume that applying the offset does not change the pointer alignment + new_ptr = cute.make_ptr( + tensor.element_type, + elem_pointer(tensor, coord).toint(), + tensor.memspace, + assumed_align=tensor.iterator.alignment, + ) + return cute.make_tensor(new_ptr, tensor.layout) + + +@dsl_user_op +def warp_reduction( + val: cute.Numeric, op: Callable, *, threads_in_group: int = 32, loc=None, ip=None +) -> cute.Numeric: + """Warp-wide reduction helper for a custom binary op.""" + offset = threads_in_group // 2 + while offset > 0: + val = op( + val, + cute.arch.shuffle_sync_bfly( + val, offset=offset, mask=-1, mask_and_clamp=31, loc=loc, ip=ip + ), + ) + offset //= 2 + return val + + +warp_reduction_max = partial( + warp_reduction, op=lambda x, y: fmax(x, y) if isinstance(x, Float32) else max(x, y) +) +warp_reduction_sum = partial(warp_reduction, op=lambda x, y: x + y) # noqa: FURB118 + + +@dsl_user_op +def make_cotiled_copy( + atom: cute.CopyAtom, atom_layout_tv: cute.Layout, data_layout: cute.Layout, *, loc=None, ip=None +) -> cute.TiledCopy: + """Compatibility wrapper for deprecated CuTeDSL `make_cotiled_copy`.""" + assert cute.is_static(atom_layout_tv.type), "atom_layout_tv must be static" + assert cute.is_static(data_layout.type), "data_layout must be static" + + inv_layout_ = cute.left_inverse(data_layout, loc=loc, ip=ip) + inv_data_layout = cute.make_layout( + (inv_layout_.shape, (1)), stride=(inv_layout_.stride, (0)), loc=loc, ip=ip + ) + layout_tv_data = cute.composition(inv_data_layout, atom_layout_tv, loc=loc, ip=ip) + + atom_layout_v_to_check = cute.coalesce( + cute.make_layout(atom_layout_tv.shape[1], stride=atom_layout_tv.stride[1], loc=loc, ip=ip), + loc=loc, + ip=ip, + ) + data_layout_v_to_check = cute.coalesce( + cute.composition( + data_layout, + cute.make_layout( + layout_tv_data.shape[1], stride=layout_tv_data.stride[1], loc=loc, ip=ip + ), + loc=loc, + ip=ip, + ), + loc=loc, + ip=ip, + ) + assert data_layout_v_to_check == atom_layout_v_to_check, ( + "the memory pointed to by atom_layout_tv does not exist in the data_layout." + ) + + flat_data_shape = cute.product_each(data_layout.shape, loc=loc, ip=ip) + tiler = tuple( + cute.filter( + cute.composition( + cute.make_layout( + flat_data_shape, + stride=tuple(0 if j != i else 1 for j in range(cute.rank(flat_data_shape))), + loc=loc, + ip=ip, + ), + layout_tv_data, + loc=loc, + ip=ip, + ), + loc=loc, + ip=ip, + ) + for i in range(cute.rank(flat_data_shape)) + ) + tile2data = cute.composition( + cute.make_layout(flat_data_shape, loc=loc, ip=ip), tiler, loc=loc, ip=ip + ) + layout_tv = cute.composition( + cute.left_inverse(tile2data, loc=loc, ip=ip), layout_tv_data, loc=loc, ip=ip + ) + return cute.make_tiled_copy(atom, layout_tv, tiler, loc=loc, ip=ip) + + +@cute.jit +def scalar_to_ssa(a: cute.Numeric, dtype) -> cute.TensorSSA: + """Convert a scalar to a cute TensorSSA of shape (1,) and given dtype""" + vec = cute.make_rmem_tensor(1, dtype) + vec[0] = a + return vec.load() + + +def ssa_to_scalar(val): + """Could inline but nice for reflecting the above api""" + return val[0] + + +@cute.jit +def get_batch_from_cu_tensor(idx: Int32, cu_tensor: cute.Tensor) -> Int32: + """Binary search to determine batch from packed index in a cumulative tensor""" + batch_size = cute.size(cu_tensor) - 1 + lo = Int32(0) + hi = batch_size + + while lo < hi: + mid = (lo + hi) // 2 + if cu_tensor[mid + 1] <= idx: + lo = mid + 1 + else: + hi = mid + + return lo diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py index 34c56d6aa..ca5732583 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py @@ -6,9 +6,13 @@ edge ("route A"), and requires no optional dependencies. The FA4 CuTe block-sparse fastpath (intended for Blackwell sm_100+) is *opt-in* via ``FASTVIDEO_VSA_CUTEDSL=1``. It routes to -:mod:`fastvideo_kernel.block_sparse_attn_cute_fwd`, which natively operates -on 128-token KV blocks (this wrapper expands the logical 256-block map / -sizes into that physical 128-block representation). The CuTe kernel +:mod:`fastvideo_kernel.block_sparse_attn_cute_fwd`. By default this wrapper +expands logical KV256 blocks to the historical physical KV128 tiles. Set +``FASTVIDEO_VSA_FA4_BLOCK_SHAPE=128x64`` for the fine-grained FA4 schedule. +The production ``64x64`` selection pairs adjacent Q64 children into physical +Q128 tiles: their KV lists are identical because this wrapper derives them +from one original Q256 metadata row. Direct Q64 callers with arbitrary masks +continue to use the native Q64/KV64 path. The CuTe kernel (``flash_attn.cute`` with block-sparsity) is an optional dependency, imported lazily only when this fastpath is selected. @@ -30,7 +34,18 @@ from .block_sparse_attn import block_sparse_attn_triton, _force_triton # FA4 CuTe build (``flash_attn.cute``) and make it a hard dependency of the # default Triton path. -_KV_BLOCK_PHYS = 128 # FA4 CuTe BSA forward uses 128-token KV blocks. +_LOGICAL_BLOCK_SIZE = 256 +_FA4_BLOCK_SHAPES = { + # The public name describes the original VSA mask. FA4 internally splits + # its logical KV256 edge into two physical KV128 tiles. + "256x256": (256, 128), + "128x64": (128, 64), + # This pairing is safe only at this original-Q256 wrapper boundary: all + # four Q64 children inherit exactly the same KV list from their parent. + # Keep arbitrary/distinct Q64 maps on block_sparse_attn_cute_fwd's native + # Q64/KV64 path, where they are not coalesced. + "64x64": (128, 64), +} _KV_BLOCK_TRITON = 64 # Existing Triton path uses 64-token KV blocks. @@ -49,28 +64,49 @@ def _resolve_backend() -> str: return "triton" -def _expand_mask_and_sizes_256_to_128( +def _resolve_fa4_block_shape() -> Tuple[int, int]: + requested = os.environ.get("FASTVIDEO_VSA_FA4_BLOCK_SHAPE", "256x256").lower() + try: + return _FA4_BLOCK_SHAPES[requested] + except KeyError as exc: + choices = ", ".join(_FA4_BLOCK_SHAPES) + raise ValueError( + f"unsupported FASTVIDEO_VSA_FA4_BLOCK_SHAPE={requested!r}; choose one of {choices}" + ) from exc + + +def _expand_mask_and_sizes_256_for_fa4( logical_mask_256: torch.Tensor, logical_kv_sizes_256: torch.Tensor, + q_block_size: int, + kv_block_size: int, ) -> Tuple[torch.Tensor, torch.Tensor]: - """Expand a [B, H, Qb256, KVb256] map to [B, H, Qb256, KVb128]. + """Expand a logical Q256/KV256 map to one supported physical FA4 shape. - Each logical 256-token KV block splits into two physical 128-token - children. Each child inherits the logical mask edge; its valid-token - count is the logical count clamped into the child's window. + Every child inherits its parent edge. KV valid-token counts are clamped + into each child's window, preserving partial and empty tail blocks. """ - expanded_mask = logical_mask_256.repeat_interleave(2, dim=3) + if _LOGICAL_BLOCK_SIZE % q_block_size or _LOGICAL_BLOCK_SIZE % kv_block_size: + raise ValueError( + f"FA4 QxKV blocks must divide {_LOGICAL_BLOCK_SIZE}, got " + f"{q_block_size}x{kv_block_size}" + ) + q_factor = _LOGICAL_BLOCK_SIZE // q_block_size + kv_factor = _LOGICAL_BLOCK_SIZE // kv_block_size + expanded_mask = logical_mask_256.repeat_interleave(q_factor, dim=2) + expanded_mask = expanded_mask.repeat_interleave(kv_factor, dim=3) sizes_i32 = logical_kv_sizes_256.to(torch.int32) - child0 = torch.clamp(sizes_i32, min=0, max=_KV_BLOCK_PHYS) - child1 = torch.clamp(sizes_i32 - _KV_BLOCK_PHYS, min=0, max=_KV_BLOCK_PHYS) - expanded_sizes = torch.empty( - (sizes_i32.numel() * 2, ), + offsets = torch.arange( + kv_factor, dtype=torch.int32, device=sizes_i32.device, - ) - expanded_sizes[0::2] = child0 - expanded_sizes[1::2] = child1 + ) * kv_block_size + expanded_sizes = torch.clamp( + sizes_i32[:, None] - offsets[None, :], + min=0, + max=kv_block_size, + ).reshape(-1) return expanded_mask, expanded_sizes @@ -126,9 +162,15 @@ def block_sparse_attn_256( if _resolve_backend() == "triton": return _triton_via_route_a(q, k, v, logical_block_map_256, logical_variable_block_sizes_256) - mask_128, sizes_128 = _expand_mask_and_sizes_256_to_128(logical_block_map_256, logical_variable_block_sizes_256) + q_block_size, kv_block_size = _resolve_fa4_block_shape() + physical_mask, physical_sizes = _expand_mask_and_sizes_256_for_fa4( + logical_block_map_256, + logical_variable_block_sizes_256, + q_block_size, + kv_block_size, + ) from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd - return block_sparse_attn_cute_fwd(q, k, v, mask_128, sizes_128) + return block_sparse_attn_cute_fwd(q, k, v, physical_mask, physical_sizes) def block_sparse_attn_256_bshd( @@ -156,6 +198,12 @@ def block_sparse_attn_256_bshd( ) return out_bhsd.transpose(1, 2).contiguous(), aux - mask_128, sizes_128 = _expand_mask_and_sizes_256_to_128(logical_block_map_256, logical_variable_block_sizes_256) + q_block_size, kv_block_size = _resolve_fa4_block_shape() + physical_mask, physical_sizes = _expand_mask_and_sizes_256_for_fa4( + logical_block_map_256, + logical_variable_block_sizes_256, + q_block_size, + kv_block_size, + ) from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd_bshd - return block_sparse_attn_cute_fwd_bshd(q, k, v, mask_128, sizes_128) + return block_sparse_attn_cute_fwd_bshd(q, k, v, physical_mask, physical_sizes) diff --git a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py index 498e5ee6d..9713dd9bd 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py +++ b/fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py @@ -9,10 +9,11 @@ The BSHD variant is preferred from VSA-256 callers to avoid layout round-trips on the hot path. The FA4 CuTe block-sparse kernel (``flash_attn.cute`` with -``block_sparsity``) is an *optional* dependency: it is imported lazily and -only exercised when the VSA-256 CuTe fastpath is explicitly selected -(``FASTVIDEO_VSA_CUTEDSL=1``). The default VSA-256 path is Triton and does -not require it. Also needs ``nvidia-cutlass-dsl`` and ``quack-kernels``. +``block_sparsity``) is an *optional* dependency: it is imported lazily. The +legacy VSA-256 route selects it with ``FASTVIDEO_VSA_CUTEDSL=1``; direct +fine-grained callers can use Q128/KV64 and, on SM100, Q64/KV64. The default +VSA-256 path is Triton and does not require FA4. The CuTe path also needs +``nvidia-cutlass-dsl`` and ``quack-kernels``. """ from __future__ import annotations @@ -22,11 +23,11 @@ from typing import Tuple import torch -_FA4_IMPORT_HINT = ("VSA-256 CuTe fastpath requires a FlashAttention-4 CuTe build that " +_FA4_IMPORT_HINT = ("The CuTe block-sparse fastpath requires a FlashAttention-4 CuTe build that " "provides `flash_attn.cute` with block-sparsity support (plus " "`nvidia-cutlass-dsl` and `quack-kernels`). This is an optional " "dependency; the default VSA-256 path is Triton. Install the FA4 CuTe " - "build and set FASTVIDEO_VSA_CUTEDSL=1 to enable the CuTe fastpath.") + "build; set FASTVIDEO_VSA_CUTEDSL=1 when selecting the VSA-256 route.") def _load_fa4_cute(): @@ -44,7 +45,7 @@ def _load_fa4_cute(): return BlockSparseTensorsTorch, _flash_attn_fwd -# Q-side tile size; kv_block_size comes from the caller's VSA logical KV block. +# Default Q-side tile size; the Q64/KV64 specialization uses a 64-row tile. _M_BLOCK_SIZE_DEFAULT = 128 @@ -72,6 +73,18 @@ def _choose_q_sparse_block_size(q_len: int, m_block_size: int = _M_BLOCK_SIZE_DE return m_block_size +def _choose_m_block_size(q_block_size: int, kv_block_size: int) -> int: + """Choose the physical FA4 Q tile without changing sparse-map semantics.""" + if q_block_size == 64: + if kv_block_size != 64: + raise ValueError("the FA4 Q64 specialization requires KV64 blocks") + major, _minor = torch.cuda.get_device_capability() + if major != 10: + raise RuntimeError("the FA4 Q64/KV64 specialization currently requires SM100") + return 64 + return _M_BLOCK_SIZE_DEFAULT + + def _aggregate_q_block_map( block_map: torch.Tensor, q_sparse_block_size: int, @@ -146,11 +159,18 @@ def _cute_forward( ) -> Tuple[torch.Tensor, torch.Tensor]: """Internal: FA4 CuTe BSA fwd with BSHD inputs.""" BlockSparseTensorsTorch, _flash_attn_fwd = _load_fa4_cute() - q_sparse_candidate = _choose_q_sparse_block_size(q_bshd.shape[1]) - q_sparse_block_size = max( - q_block_size, - ((q_sparse_candidate + q_block_size - 1) // q_block_size) * q_block_size, - ) + m_block_size = _choose_m_block_size(q_block_size, kv_block_size) + if q_block_size == m_block_size: + # Q128 and Q64 maps may select different KV blocks in adjacent rows. + # Preserve their declared granularity instead of unioning them into the + # historical two-stage Q256 schedule. + q_sparse_block_size = q_block_size + else: + q_sparse_candidate = _choose_q_sparse_block_size(q_bshd.shape[1], m_block_size) + q_sparse_block_size = max( + q_block_size, + ((q_sparse_candidate + q_block_size - 1) // q_block_size) * q_block_size, + ) sparse_map = _aggregate_q_block_map( block_map, q_sparse_block_size=q_sparse_block_size, @@ -177,7 +197,7 @@ def _cute_forward( q_bshd, k_bshd, v_bshd, - tile_mn=(_M_BLOCK_SIZE_DEFAULT, kv_block_size), + tile_mn=(m_block_size, kv_block_size), mask_mod=_build_vbs_mask_mod(kv_block_size), block_sparse_tensors=sparse_tensors, aux_tensors=[variable_block_sizes], diff --git a/fastvideo-kernel/tests/test_fa4_vsa_block_shapes.py b/fastvideo-kernel/tests/test_fa4_vsa_block_shapes.py new file mode 100644 index 000000000..b5dfcf157 --- /dev/null +++ b/fastvideo-kernel/tests/test_fa4_vsa_block_shapes.py @@ -0,0 +1,526 @@ +"""Correctness coverage for FastVideo's fine-grained FA4 sparse tiles.""" + +from __future__ import annotations + +import math + +import pytest +import torch + + +@pytest.fixture(autouse=True) +def _require_fa4_sm100(monkeypatch): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + major, _minor = torch.cuda.get_device_capability() + if major not in (10, 11): + pytest.skip("fine-grained FA4 VSA requires SM100/SM110") + pytest.importorskip( + "flash_attn.cute.block_sparsity", + reason="optional FastVideo FA4 source package is not installed", + ) + monkeypatch.setenv("FASTVIDEO_FA4_VSA_DUAL_STREAM", "1") + monkeypatch.setenv("FASTVIDEO_FA4_VSA_SP_DOUBLE_BUFFER", "1") + + +def _selected_block_map( + heads: int, + q_blocks: int, + kv_blocks: int, + selected_blocks: int, + device: torch.device, +) -> torch.Tensor: + block_map = torch.zeros( + 1, + heads, + q_blocks, + kv_blocks, + dtype=torch.bool, + device=device, + ) + for head in range(heads): + for q_block in range(q_blocks): + start = (3 * head + 5 * q_block) % kv_blocks + indices = [(start + 2 * offset) % kv_blocks for offset in range(selected_blocks)] + block_map[0, head, q_block, indices] = True + return block_map + + +def _ordered_sparse_tensors( + block_map: torch.Tensor, + q_block_size: int, + kv_block_size: int, + masked_blocks: int, +): + from flash_attn.cute.block_sparsity import BlockSparseTensorsTorch + + batch, heads, q_blocks, kv_blocks = block_map.shape + selected_blocks = int(block_map.sum(dim=-1).min().item()) + assert bool((block_map.sum(dim=-1) == selected_blocks).all()) + if not 0 <= masked_blocks <= selected_blocks: + raise ValueError("masked_blocks must be within the selected-block count") + + selected = torch.arange(kv_blocks, dtype=torch.int32, device=block_map.device) + selected = selected.view(1, 1, 1, kv_blocks).expand(batch, heads, q_blocks, kv_blocks) + selected = selected.masked_select(block_map).view(batch, heads, q_blocks, selected_blocks) + selected = selected.sort(dim=-1).values + + mask_count = masked_blocks + full_count = selected_blocks - masked_blocks + mask_idx = torch.zeros( + batch, + heads, + q_blocks, + max(mask_count, 1), + dtype=torch.int32, + device=block_map.device, + ) + full_idx = torch.zeros( + batch, + heads, + q_blocks, + max(full_count, 1), + dtype=torch.int32, + device=block_map.device, + ) + if mask_count: + mask_idx[..., :mask_count] = selected[..., :mask_count] + if full_count: + full_idx[..., :full_count] = selected[..., mask_count:] + + return BlockSparseTensorsTorch( + mask_block_cnt=torch.full( + (batch, heads, q_blocks), + mask_count, + dtype=torch.int32, + device=block_map.device, + ), + mask_block_idx=mask_idx, + full_block_cnt=torch.full( + (batch, heads, q_blocks), + full_count, + dtype=torch.int32, + device=block_map.device, + ), + full_block_idx=full_idx, + block_size=(q_block_size, kv_block_size), + ) + + +def _torch_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_map: torch.Tensor, + q_block_size: int, + kv_block_size: int, + variable_block_sizes: torch.Tensor | None = None, +) -> torch.Tensor: + token_mask = block_map.repeat_interleave(q_block_size, dim=2) + token_mask = token_mask.repeat_interleave(kv_block_size, dim=3) + if variable_block_sizes is not None: + kv_positions = torch.arange(k.shape[1], device=k.device) + kv_blocks = kv_positions // kv_block_size + kv_offsets = kv_positions % kv_block_size + valid_tokens = kv_offsets < variable_block_sizes[kv_blocks] + token_mask = token_mask & valid_tokens.view(1, 1, 1, -1) + q_heads = q.permute(0, 2, 1, 3).float() + k_heads = k.permute(0, 2, 1, 3).float() + v_heads = v.permute(0, 2, 1, 3).float() + scores = torch.matmul(q_heads, k_heads.transpose(-2, -1)) / math.sqrt(q.shape[-1]) + scores.masked_fill_(~token_mask, float("-inf")) + out = torch.matmul(torch.softmax(scores, dim=-1), v_heads) + return out.permute(0, 2, 1, 3).contiguous() + + +@pytest.mark.cuda +@pytest.mark.parametrize( + ( + "q_block_size", + "kv_block_size", + "selected_blocks", + "masked_blocks", + "sp_double_buffer", + ), + [ + pytest.param(128, 64, 1, 0, True, id="q128_kv64_one_stream_empty"), + pytest.param(128, 64, 2, 1, True, id="q128_kv64_even_mixed"), + pytest.param(128, 64, 3, 1, True, id="q128_kv64_odd_mixed"), + pytest.param(128, 64, 3, 0, True, id="q128_kv64_odd_full_only"), + pytest.param(128, 64, 3, 1, False, id="q128_kv64_legacy_optout"), + pytest.param(64, 64, 3, 1, True, id="q64_kv64_odd_mixed"), + ], +) +def test_fa4_vsa_fine_grained_forward( + q_block_size: int, + kv_block_size: int, + selected_blocks: int, + masked_blocks: int, + sp_double_buffer: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from flash_attn.cute.interface import _flash_attn_fwd + + monkeypatch.setenv( + "FASTVIDEO_FA4_VSA_SP_DOUBLE_BUFFER", + "1" if sp_double_buffer else "0", + ) + torch.manual_seed(2026 + q_block_size + selected_blocks) + device = torch.device("cuda") + batch, heads, head_dim = 1, 8, 128 + q_blocks, kv_blocks = 4, 7 + q_len = q_blocks * q_block_size + kv_len = kv_blocks * kv_block_size + q = torch.randn(batch, q_len, heads, head_dim, dtype=torch.bfloat16, device=device) + k = torch.randn(batch, kv_len, heads, head_dim, dtype=torch.bfloat16, device=device) + v = torch.randn_like(k) + block_map = _selected_block_map(heads, q_blocks, kv_blocks, selected_blocks, device) + sparse_tensors = _ordered_sparse_tensors( + block_map, + q_block_size, + kv_block_size, + masked_blocks, + ) + + out = _flash_attn_fwd( + q, + k, + v, + tile_mn=(q_block_size, kv_block_size), + block_sparse_tensors=sparse_tensors, + causal=False, + return_lse=True, + )[0] + ref = _torch_reference(q, k, v, block_map, q_block_size, kv_block_size) + + assert bool(torch.isfinite(out).all()) + error = (out.float() - ref).abs() + avg_abs = float(error.mean()) + max_rel = float(error.max() / (ref.abs().mean() + 1e-6)) + print( + f"[fa4 {q_block_size}x{kv_block_size} selected={selected_blocks} " + f"masked={masked_blocks} spdb={sp_double_buffer}] " + f"avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}" + ) + assert avg_abs < 1e-3 + assert max_rel < 0.2 + + +@pytest.mark.cuda +@pytest.mark.parametrize("q_block_size", [128, 64], ids=["q128_kv64", "q64_kv64"]) +def test_fa4_vsa_dual_stream_persistent_phase_transitions(q_block_size: int) -> None: + """Exercise even, odd, empty, and one-block tiles in one persistent launch.""" + if q_block_size == 64 and torch.cuda.get_device_capability()[0] != 10: + pytest.skip("the Q64/KV64 specialization currently requires SM100") + + from flash_attn.cute.block_sparsity import BlockSparseTensorsTorch + from flash_attn.cute.interface import _flash_attn_fwd + + torch.manual_seed(4554 + q_block_size) + device = torch.device("cuda") + kv_block_size = 64 + counts = [2, 3, 0, 1, 4] + masked_counts = [1, 1, 0, 1, 2] + batch, heads, head_dim, kv_blocks = 1, 4, 128, 8 + q_blocks = len(counts) + q = torch.randn( + batch, + q_blocks * q_block_size, + heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + k = torch.randn( + batch, + kv_blocks * kv_block_size, + heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + v = torch.randn_like(k) + + mask_count = torch.empty(batch, heads, q_blocks, dtype=torch.int32, device=device) + full_count = torch.empty_like(mask_count) + mask_idx = torch.zeros( + batch, + heads, + q_blocks, + max(masked_counts), + dtype=torch.int32, + device=device, + ) + full_idx = torch.zeros( + batch, + heads, + q_blocks, + max(count - masked for count, masked in zip(counts, masked_counts)), + dtype=torch.int32, + device=device, + ) + block_map = torch.zeros( + batch, + heads, + q_blocks, + kv_blocks, + dtype=torch.bool, + device=device, + ) + for head in range(heads): + for q_block, (count, masked) in enumerate(zip(counts, masked_counts)): + chosen = sorted({(head + 3 * q_block + 2 * offset) % kv_blocks for offset in range(count)}) + assert len(chosen) == count + mask_count[0, head, q_block] = masked + full_count[0, head, q_block] = count - masked + if masked: + mask_idx[0, head, q_block, :masked] = torch.tensor( + chosen[:masked], + dtype=torch.int32, + device=device, + ) + if count > masked: + full_idx[0, head, q_block, :count - masked] = torch.tensor( + chosen[masked:], + dtype=torch.int32, + device=device, + ) + if count: + block_map[0, head, q_block, chosen] = True + + sparse_tensors = BlockSparseTensorsTorch( + mask_block_cnt=mask_count, + mask_block_idx=mask_idx, + full_block_cnt=full_count, + full_block_idx=full_idx, + block_size=(q_block_size, kv_block_size), + ) + out, lse = _flash_attn_fwd( + q, + k, + v, + tile_mn=(q_block_size, kv_block_size), + block_sparse_tensors=sparse_tensors, + causal=False, + return_lse=True, + )[:2] + ref = _torch_reference(q, k, v, block_map, q_block_size, kv_block_size) + + nonempty_rows = torch.tensor( + [count > 0 for count in counts], + dtype=torch.bool, + device=device, + ).repeat_interleave(q_block_size) + error = (out[:, nonempty_rows].float() - ref[:, nonempty_rows]).abs() + empty_out = out[:, ~nonempty_rows].float() + empty_lse = lse[:, :, ~nonempty_rows] + avg_abs = float(error.mean()) + max_abs = float(error.max()) + print( + f"[fa4 persistent {q_block_size}x{kv_block_size}] " + f"avg_abs={avg_abs:.6e}, max_abs={max_abs:.6e}" + ) + assert avg_abs < 1e-3 + assert max_abs < 1e-2 + assert float(empty_out.abs().max()) == 0.0 + assert bool(torch.isneginf(empty_lse).all()) + + +@pytest.mark.cuda +def test_fa4_vsa_q128_kv64_bshd_wrapper_variable_blocks() -> None: + """Cover partial and empty KV blocks through FastVideo's VBS mask mod.""" + from fastvideo_kernel.block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd_bshd + + torch.manual_seed(5128) + device = torch.device("cuda") + batch, heads, head_dim = 1, 8, 128 + q_block_size, kv_block_size = 128, 64 + q_blocks, kv_blocks, selected_blocks = 4, 7, 3 + q = torch.randn( + batch, + q_blocks * q_block_size, + heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + k = torch.randn( + batch, + kv_blocks * kv_block_size, + heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + v = torch.randn_like(k) + block_map = _selected_block_map(heads, q_blocks, kv_blocks, selected_blocks, device) + variable_block_sizes = torch.tensor( + [64, 23, 64, 7, 0, 51, 64], + dtype=torch.int32, + device=device, + ) + + out, lse = block_sparse_attn_cute_fwd_bshd( + q, + k, + v, + block_map, + variable_block_sizes, + ) + ref = _torch_reference( + q, + k, + v, + block_map, + q_block_size, + kv_block_size, + variable_block_sizes, + ) + + assert bool(torch.isfinite(out).all()) + assert bool(torch.isfinite(lse).all()) + error = (out.float() - ref).abs() + avg_abs = float(error.mean()) + max_rel = float(error.max() / (ref.abs().mean() + 1e-6)) + print(f"[fa4 wrapper Q128/KV64 VBS] avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}") + assert avg_abs < 1e-3 + assert max_rel < 0.2 + + +@pytest.mark.cuda +@pytest.mark.parametrize( + ("q_block_size", "kv_block_size"), + [ + pytest.param(128, 64, id="q128_kv64"), + pytest.param(64, 64, id="q64_kv64"), + ], +) +def test_fa4_vsa_fine_grained_bshd_wrapper( + q_block_size: int, + kv_block_size: int, +) -> None: + """Exercise the FastVideo BSHD entrypoint used by the H3 backend.""" + if q_block_size == 64 and torch.cuda.get_device_capability()[0] != 10: + pytest.skip("the Q64/KV64 specialization currently requires SM100") + + from fastvideo_kernel.block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd_bshd + + torch.manual_seed(4048 + q_block_size) + device = torch.device("cuda") + batch, heads, head_dim = 1, 8, 128 + q_blocks, kv_blocks, selected_blocks = 4, 7, 3 + q = torch.randn( + batch, + q_blocks * q_block_size, + heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + k = torch.randn( + batch, + kv_blocks * kv_block_size, + heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + v = torch.randn_like(k) + block_map = _selected_block_map(heads, q_blocks, kv_blocks, selected_blocks, device) + variable_block_sizes = torch.full( + (kv_blocks,), + kv_block_size, + dtype=torch.int32, + device=device, + ) + + out, lse = block_sparse_attn_cute_fwd_bshd( + q, + k, + v, + block_map, + variable_block_sizes, + ) + ref = _torch_reference(q, k, v, block_map, q_block_size, kv_block_size) + + assert bool(torch.isfinite(out).all()) + assert bool(torch.isfinite(lse).all()) + error = (out.float() - ref).abs() + avg_abs = float(error.mean()) + max_rel = float(error.max() / (ref.abs().mean() + 1e-6)) + print( + f"[fa4 wrapper {q_block_size}x{kv_block_size}] " + f"avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}" + ) + assert avg_abs < 1e-3 + assert max_rel < 0.2 + + +@pytest.mark.cuda +@pytest.mark.parametrize("block_shape", ["256x256", "128x64", "64x64"]) +def test_vsa256_bshd_fa4_block_shape_routes( + block_shape: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Route one logical Q256/KV256 mask through every FA4 specialization.""" + if block_shape == "64x64" and torch.cuda.get_device_capability()[0] != 10: + pytest.skip("the Q64/KV64 specialization currently requires SM100") + + from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_256_bshd + + monkeypatch.setenv("FASTVIDEO_VSA_CUTEDSL", "1") + monkeypatch.setenv("FASTVIDEO_VSA_FA4_BLOCK_SHAPE", block_shape) + torch.manual_seed(6256) + device = torch.device("cuda") + batch, heads, head_dim = 1, 8, 128 + logical_block_size = 256 + q_blocks, kv_blocks, selected_blocks = 4, 7, 3 + q = torch.randn( + batch, + q_blocks * logical_block_size, + heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + k = torch.randn( + batch, + kv_blocks * logical_block_size, + heads, + head_dim, + dtype=torch.bfloat16, + device=device, + ) + v = torch.randn_like(k) + block_map = _selected_block_map(heads, q_blocks, kv_blocks, selected_blocks, device) + variable_block_sizes = torch.tensor( + [256, 137, 255, 64, 201, 1, 192], + dtype=torch.int32, + device=device, + ) + + out, lse = block_sparse_attn_256_bshd( + q, + k, + v, + block_map, + variable_block_sizes, + ) + ref = _torch_reference( + q, + k, + v, + block_map, + logical_block_size, + logical_block_size, + variable_block_sizes, + ) + + assert bool(torch.isfinite(out).all()) + assert bool(torch.isfinite(lse).all()) + error = (out.float() - ref).abs() + avg_abs = float(error.mean()) + max_rel = float(error.max() / (ref.abs().mean() + 1e-6)) + print(f"[VSA256 route {block_shape}] avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}") + assert avg_abs < 1e-3 + assert max_rel < 0.2 diff --git a/handoff.md b/handoff.md new file mode 100644 index 000000000..10c7759ad --- /dev/null +++ b/handoff.md @@ -0,0 +1,170 @@ +# FA4 VSA block-64 handoff + +## Status + +This branch vendors FlashAttention-4's CuTe source as a first-class FastVideo +fork, adds the issue-#4554 benchmark, and implements forward-only SM100 sparse +attention paths for Q128/KV64 and Q64/KV64. + +The requested single-Q-block/two-KV-workstream design is implemented. The two +workstreams split the canonical sparse traversal by ordinal parity, maintain +independent online-softmax state, and merge their FP32 output accumulators in +the correction epilogue. Q128/KV64 additionally uses four score/probability +slots to overlap successive blocks. + +The measured result is not fully zero-loss: + +- Native Q128/KV64 reaches about 90% of raw FA4 Q256/KV256 throughput. +- Native Q64/KV64 reaches about 55% of raw FA4 Q256/KV256 throughput. +- A newly added production-only VSA256 adapter reaches 80-86% of the public + Q256 wrapper. It is selected with + `FASTVIDEO_VSA_FA4_BLOCK_SHAPE=64x64`, but physically runs Q128/KV64 by + pairing adjacent Q64 children whose sparse lists are identical at the + original Q256 metadata boundary. It must not be presented as native + Q64/KV64 performance. + +## Source and provenance + +- Fork: `fastvideo-kernel/fa4/` +- Upstream base: Dao-AILab/flash-attention commit + `82d6441eec5d4dfec120153db2c0145ae855a083` +- Q64 primitives adapted from upstream commit + `526c18d25bcbc7fc7d6740ab3c7c84ed2d42cb0b` +- Full refresh instructions: `fastvideo-kernel/fa4/UPSTREAM.md` +- Root `pyproject.toml` resolves `flash-attn-4` to the editable local fork. + +The fork retains the upstream `flash-attn-4` distribution name and +`flash_attn.cute` import namespace. + +## Main implementation + +- `fastvideo-kernel/fa4/interface.py` + - Strict compile-time gates for the dual-stream and double-buffer paths. + - `FASTVIDEO_FA4_VSA_DUAL_STREAM=0` retains the upstream-style single-stream + fallback. + - `FASTVIDEO_FA4_VSA_SP_DOUBLE_BUFFER=0` retains the earlier dual-stream + schedule for rollback and comparison. +- `fastvideo-kernel/fa4/block_sparse_utils.py` + - Canonical ordinal-parity sparse traversal shared by load, MMA, and + softmax code. + - Correct handling for masked-only, full-only, mixed, odd-length, and empty + streams. +- `fastvideo-kernel/fa4/flash_fwd_sm100.py` + - One Q block, two KV workstreams, four S/P slots for Q128/KV64, independent + phase tracking, and stable FP32 merge. + - Native M64 score/output layouts and row-pair softmax support. +- `fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_256.py` + - Public VSA256 shape selection via `FASTVIDEO_VSA_FA4_BLOCK_SHAPE`. + - The `64x64` option is the newly added coalescing adapter described above. +- `fastvideo-kernel/python/fastvideo_kernel/block_sparse_attn_cute_fwd.py` + - Direct Q128/KV64 and native Q64/KV64 BSHD entrypoints. + +## Benchmark + +Harness: + +```bash +fastvideo-kernel/benchmarks/bench_vsa_blackwell.py +``` + +It derives from FlashInfer issue #4554 and immutable gist revision +`e15ac9066f23ef3690e33e1cc1fdac45b4b9099f`. The default `exact256` mask mode +uses identical selected token pairs across block shapes, checks every output +against FP32 masked SDPA, and reports sparse-aware algorithmic TFLOP/s plus MFU. + +Four-GPU command: + +```bash +PYTHONPATH="$PWD/fastvideo-kernel/fa4:$PWD/fastvideo-kernel/python" \ +CUTE_DSL_ENABLE_TVM_FFI=1 \ +FASTVIDEO_VSA_CUTEDSL=1 \ +FASTVIDEO_FA4_VSA_DUAL_STREAM=1 \ +FASTVIDEO_FA4_VSA_SP_DOUBLE_BUFFER=1 \ +uv run --no-sync torchrun --standalone --nproc-per-node=4 \ + fastvideo-kernel/benchmarks/bench_vsa_blackwell.py \ + --seq_lens 32768 --sparsities dense 90 --mask_mode exact256 \ + --arms cutedsl256 fa4_wrapper fa4 \ + --block_shapes 256x256 128x64 64x64 \ + --warmup 5 --rep 20 --out /tmp/fa4_final_tray.json +``` + +Environment used for the final run: + +- 4x NVIDIA GB200, SM100 +- PyTorch `2.12.0+cu130` +- CUDA 13.0, driver 580.159.04 +- BF16, B=1, H=12, D=128, S=32768 +- MFU denominator: 2.5 PFLOP/s dense BF16 per GPU, 10 PFLOP/s per tray + +### Four-GPU aggregate results + +| Raw schedule | Dense PFLOP/s | Dense MFU | 90% sparse PFLOP/s | 90% sparse MFU | Relative to raw Q256 | +|---|---:|---:|---:|---:|---:| +| Q256/KV256 | 6.353 | 63.53% | 5.852 | 58.52% | 100% | +| Q128/KV64 | 5.764 | 57.64% | 5.234 | 52.34% | 90.7% dense / 89.5% sparse | +| Native Q64/KV64 | 3.408 | 34.08% | 3.215 | 32.15% | 53.6% dense / 54.9% sparse | + +| Public VSA256 route | Dense PFLOP/s | Dense MFU | 90% sparse PFLOP/s | 90% sparse MFU | Relative to public Q256 | +|---|---:|---:|---:|---:|---:| +| Default Q256 | 6.185 | 61.85% | 4.494 | 44.94% | 100% | +| Added `64x64` coalescing adapter | 5.317 | 53.17% | 3.584 | 35.84% | 86.0% dense / 79.7% sparse | + +Fixed-base NCU showed that score/probability double buffering improves the +Q128/KV64 end-to-end wrapper by 8.74% dense and 4.35% at 90% sparsity over +the retained legacy dual-stream schedule. + +## Validation + +Final GB200 runs: + +```bash +CUDA_VISIBLE_DEVICES=3 \ +PYTHONPATH="$PWD/fastvideo-kernel/fa4:$PWD/fastvideo-kernel/python" \ +CUTE_DSL_ENABLE_TVM_FFI=1 \ +FASTVIDEO_VSA_CUTEDSL=1 \ +FASTVIDEO_FA4_VSA_DUAL_STREAM=1 \ +FASTVIDEO_FA4_VSA_SP_DOUBLE_BUFFER=1 \ +uv run --no-sync pytest -vs \ + fastvideo-kernel/tests/test_fa4_vsa_block_shapes.py +``` + +Result: 14 passed. Coverage includes selected counts 0-5, odd/even and +one-stream-empty cases, persistent phase transitions, masked/full/mixed +metadata, variable block sizes, native Q64, Q128, public VSA256 routes, and +the legacy double-buffer opt-out. + +Existing compatibility tests: + +```bash +uv run --no-sync pytest -vs \ + fastvideo-kernel/tests/test_vsa256_forward.py \ + fastvideo-kernel/tests/test_vsa256_forward_cross.py \ + fastvideo-kernel/tests/test_vsa256_forward_vbs.py +``` + +Result: 4 passed. + +Other completed gates: + +- `pre-commit run --files ...`: passed for all scoped files. +- `python -m py_compile`: passed for the fork, wrappers, tests, and benchmark. +- FA4 sdist and wheel builds: passed. +- `twine check`: passed for both distributions. +- Isolated wheel import: resolved `flash_attn.cute` and the expected local + version. +- `git diff --check`: passed. + +## Known limitations and follow-up + +- These optimized specializations are forward-only. +- Native Q128/KV64 and Q64/KV64 do not meet the original zero-loss/30%-loss + targets; keep the native and coalesced-adapter numbers separate. +- The existing kernel CI lane uses H100. The SM100/SM110 tests skip there, so + the four-GB200 local run is currently the only hardware regression gate. +- The workspace's untracked `uv.lock` predates the local source selection and + was deliberately not overwritten. `uv lock --dry-run --offline` resolves + only the expected `flash-attn-4` source/version update. +- Rejected performance experiments include a second UMMA issuer, paired wait + batching, cooperative 2-CTA execution, cluster multicast, CTA-wide stats + barriers, and correction wait reordering. All either regressed or failed to + improve the fixed-base profiles and were not retained. diff --git a/pyproject.toml b/pyproject.toml index d7630ba34..0bfb1f55a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -122,9 +122,9 @@ fastvideo-kernel = [ # uv reads UV_TORCH_BACKEND / --torch-backend (not a pyproject setting). The # supported backends (cu126, cu130) both provide torch 2.12.0. imagebind = { git = "https://github.com/facebookresearch/ImageBind.git", rev = "53680b02d7e37b19b124fa37bae4b6c98c38f5be" } -# FA4 cute, pinned to a cutlass-4.5-compatible revision. torch.compile support -# comes from FastVideo's own custom_op wrappers. -flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "82d6441eec5d4dfec120153db2c0145ae855a083", subdirectory = "flash_attn/cute" } +# FastVideo's FA4 CuTe source fork. Its upstream commit and refresh procedure +# are recorded in fastvideo-kernel/fa4/UPSTREAM.md. +flash-attn-4 = { path = "fastvideo-kernel/fa4", editable = true } [project.optional-dependencies] @@ -189,9 +189,8 @@ streaming = [ dreamverse = [ "uvicorn[standard]>=0.41.0", "cerebras-cloud-sdk", - # PyPI forbids direct URL deps in published metadata; pin upstream cute via - # [tool.uv.sources] above (same pattern as imagebind) so the wheel stays - # publishable while uv workspace installs still get it. + # Keep the published dependency name portable; [tool.uv.sources] above + # selects FastVideo's in-tree CuTe fork for workspace installs. "flash-attn-4", "flashinfer-python", "openai>=1.40",