From 3992a023e65ed68cb203fc90164a6962378f187d Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Sat, 24 Jan 2026 15:25:03 +0800 Subject: [PATCH] Move fa4 from sgl-kernel to jit kernel (#17353) --- .github/workflows/pr-test.yml | 5 + .pre-commit-config.yaml | 2 +- python/sglang/jit_kernel/__init__.py | 0 .../jit_kernel/flash_attention/cute/.flake8 | 4 + .../jit_kernel/flash_attention/cute/AUTHORS | 5 + .../jit_kernel/flash_attention/cute/LICENSE | 29 + .../jit_kernel/flash_attention/cute/README.md | 0 .../flash_attention/cute/__init__.py | 21 + .../flash_attention/cute/ampere_helpers.py | 103 + .../flash_attention/cute/barrier.py | 71 + .../flash_attention/cute/benchmark.py | 268 ++ .../flash_attention/cute/blackwell_helpers.py | 753 +++++ .../flash_attention/cute/block_info.py | 108 + .../cute/block_sparse_utils.py | 1451 ++++++++ .../flash_attention/cute/block_sparsity.py | 250 ++ .../cute/compute_block_sparsity.py | 377 +++ .../flash_attention/cute/copy_utils.py | 340 ++ .../flash_attention/cute/cute_dsl_utils.py | 144 + .../flash_attention/cute/fast_math.py | 21 + .../flash_attention/cute/flash_bwd.py | 1262 +++++++ .../cute/flash_bwd_postprocess.py | 463 +++ .../cute/flash_bwd_preprocess.py | 365 ++ .../flash_attention/cute/flash_bwd_sm100.py | 2950 +++++++++++++++++ .../flash_attention/cute/flash_bwd_sm90.py | 1706 ++++++++++ .../flash_attention/cute/flash_fwd.py | 2485 ++++++++++++++ .../flash_attention/cute/flash_fwd_combine.py | 704 ++++ .../flash_attention/cute/flash_fwd_sm100.py | 2740 +++++++++++++++ .../flash_attention/cute/hopper_helpers.py | 101 + .../flash_attention/cute/interface.py | 1758 ++++++++++ .../jit_kernel/flash_attention/cute/mask.py | 651 ++++ .../flash_attention/cute/mma_sm100_desc.py | 291 ++ .../flash_attention/cute/named_barrier.py | 31 + .../flash_attention/cute/pack_gqa.py | 164 + .../flash_attention/cute/paged_kv.py | 214 ++ .../flash_attention/cute/pipeline.py | 272 ++ .../flash_attention/cute/pyproject.toml | 56 + .../flash_attention/cute/seqlen_info.py | 138 + .../flash_attention/cute/softmax.py | 582 ++++ .../flash_attention/cute/testing.py | 423 +++ .../flash_attention/cute/tile_scheduler.py | 719 ++++ .../jit_kernel/flash_attention/cute/utils.py | 859 +++++ .../sglang/jit_kernel/flash_attention_v4.py | 170 + .../tests/test_flash_attention_4.py | 1504 +++++++++ .../layers/attention/backends/flash_attn.py | 12 +- .../attention/flashattention_backend.py | 33 +- 45 files changed, 24594 insertions(+), 11 deletions(-) create mode 100644 python/sglang/jit_kernel/__init__.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/.flake8 create mode 100644 python/sglang/jit_kernel/flash_attention/cute/AUTHORS create mode 100644 python/sglang/jit_kernel/flash_attention/cute/LICENSE create mode 100644 python/sglang/jit_kernel/flash_attention/cute/README.md create mode 100644 python/sglang/jit_kernel/flash_attention/cute/__init__.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/ampere_helpers.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/barrier.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/benchmark.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/blackwell_helpers.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/block_info.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/block_sparse_utils.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/block_sparsity.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/compute_block_sparsity.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/copy_utils.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/cute_dsl_utils.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/fast_math.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/flash_bwd.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/flash_bwd_postprocess.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/flash_bwd_preprocess.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/flash_bwd_sm100.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/flash_bwd_sm90.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/flash_fwd.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/flash_fwd_combine.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/flash_fwd_sm100.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/hopper_helpers.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/interface.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/mask.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/mma_sm100_desc.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/named_barrier.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/pack_gqa.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/paged_kv.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/pipeline.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/pyproject.toml create mode 100644 python/sglang/jit_kernel/flash_attention/cute/seqlen_info.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/softmax.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/testing.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/tile_scheduler.py create mode 100644 python/sglang/jit_kernel/flash_attention/cute/utils.py create mode 100644 python/sglang/jit_kernel/flash_attention_v4.py create mode 100644 python/sglang/jit_kernel/tests/test_flash_attention_4.py diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index c6da837ef..e7bcf9c06 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -1241,6 +1241,11 @@ jobs: fi IS_BLACKWELL=1 python3 run_suite.py --hw cuda --suite stage-b-test-4-gpu-b200 $CONTINUE_ON_ERROR_FLAG + - name: Run FA4 jit_kernel tests (SM100+) + timeout-minutes: 5 + run: | + IS_BLACKWELL=1 python3 -m pytest -q python/sglang/jit_kernel/tests/test_flash_attention_4.py + stage-c-test-large-4-gpu: needs: [check-changes, call-gate, wait-for-stage-b, sgl-kernel-build-wheels] if: | diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5cb362dfd..7abe48029 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,5 +1,5 @@ default_stages: [pre-commit, pre-push, manual] -exclude: ^python/sglang/multimodal_gen/csrc +exclude: ^(python/sglang/multimodal_gen/csrc|python/sglang/jit_kernel/flash_attention/cute) repos: - repo: https://github.com/pre-commit/pre-commit-hooks diff --git a/python/sglang/jit_kernel/__init__.py b/python/sglang/jit_kernel/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/jit_kernel/flash_attention/cute/.flake8 b/python/sglang/jit_kernel/flash_attention/cute/.flake8 new file mode 100644 index 000000000..bae5b85c0 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/.flake8 @@ -0,0 +1,4 @@ +[flake8] +max-line-length = 100 +# W503: line break before binary operator +ignore = E731, E741, F841, W503 diff --git a/python/sglang/jit_kernel/flash_attention/cute/AUTHORS b/python/sglang/jit_kernel/flash_attention/cute/AUTHORS new file mode 100644 index 000000000..bc3991c67 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/AUTHORS @@ -0,0 +1,5 @@ +Tri Dao, tri@tridao.me +Jay Shah +Ted Zadouri +Markus Hoehnerbach +Vijay Thakkar \ No newline at end of file diff --git a/python/sglang/jit_kernel/flash_attention/cute/LICENSE b/python/sglang/jit_kernel/flash_attention/cute/LICENSE new file mode 100644 index 000000000..5860e4b33 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/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/python/sglang/jit_kernel/flash_attention/cute/README.md b/python/sglang/jit_kernel/flash_attention/cute/README.md new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/jit_kernel/flash_attention/cute/__init__.py b/python/sglang/jit_kernel/flash_attention/cute/__init__.py new file mode 100644 index 000000000..7706f746a --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/__init__.py @@ -0,0 +1,21 @@ +"""Flash Attention CUTE (CUDA Template Engine) implementation.""" + +__version__ = "0.1.0" + +import cutlass.cute as cute + +from .interface import ( + flash_attn_func, + flash_attn_varlen_func, +) + +from .cute_dsl_utils import cute_compile_patched + +# Patch cute.compile to optionally dump SASS +cute.compile = cute_compile_patched + + +__all__ = [ + "flash_attn_func", + "flash_attn_varlen_func", +] diff --git a/python/sglang/jit_kernel/flash_attention/cute/ampere_helpers.py b/python/sglang/jit_kernel/flash_attention/cute/ampere_helpers.py new file mode 100644 index 000000000..e3072d8ce --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/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/python/sglang/jit_kernel/flash_attention/cute/barrier.py b/python/sglang/jit_kernel/flash_attention/cute/barrier.py new file mode 100644 index 000000000..c999b1801 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/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/python/sglang/jit_kernel/flash_attention/cute/benchmark.py b/python/sglang/jit_kernel/flash_attention/cute/benchmark.py new file mode 100644 index 000000000..9a7820e7b --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/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/python/sglang/jit_kernel/flash_attention/cute/blackwell_helpers.py b/python/sglang/jit_kernel/flash_attention/cute/blackwell_helpers.py new file mode 100644 index 000000000..e39ac5eeb --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/blackwell_helpers.py @@ -0,0 +1,753 @@ +# 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 sglang.jit_kernel.flash_attention.cute.mma_sm100_desc as sm100_desc +from .utils import parse_swizzle_from_pointer + + +@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, +) -> 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_constexpr(cute.size(tCrA.shape[2])): + 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, + **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, **kwargs + ) + + +@cute.jit +def gemm( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + zero_init: bool | Boolean = False, +) -> cute.TiledMma: + for k in cutlass.range_constexpr(cute.size(tCrA.shape[2])): + tiled_mma.set(tcgen05.Field.ACCUMULATE, not zero_init or k != 0) + cute.gemm(tiled_mma, acc, tCrA[None, None, k], tCrB[None, None, k], acc) + return tiled_mma + + +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)) + if const_expr(not is_ts): + sA_swizzle = parse_swizzle_from_pointer(sA.iterator) + 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 = parse_swizzle_from_pointer(sB.iterator) + 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::f16 [$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::f16 [$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)) + if const_expr(not is_ts): + sA_swizzle = parse_swizzle_from_pointer(sA.iterator) + 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 = parse_swizzle_from_pointer(sB.iterator) + 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::f16 [$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::f16 [$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::f16 [$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::f16 [$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, + zero_init: bool | Boolean = False, + # sA_offset: Int32 = 0, + # acc_offset: Int32 = 0, + tA_addr: Optional[Int32] = None, +) -> 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)) + if const_expr(not is_ts): + sA_swizzle = parse_swizzle_from_pointer(sA.iterator) + 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 = parse_swizzle_from_pointer(sB.iterator) + 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::1.kind::f16 [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::1.kind::f16 [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" + 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::1.kind::f16 [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::1.kind::f16 [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 cute.size(tCrA.shape[2]) // 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"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 + {hex(offset_a[k])}], smem_desc_b, idesc, 1;\n\t" + ) + for k in range(cute.size(tCrA.shape[2]) // 4 * 3, 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)) + 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::f16 [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::f16 [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::f16 [$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::f16 [$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, + ) diff --git a/python/sglang/jit_kernel/flash_attention/cute/block_info.py b/python/sglang/jit_kernel/flash_attention/cute/block_info.py new file mode 100644 index 000000000..a5a2544a7 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/block_info.py @@ -0,0 +1,108 @@ +# 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 .seqlen_info import SeqlenInfoQK + + +@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: cutlass.Int32 = 0, + num_splits: cutlass.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 = ( + cutlass.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_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)) diff --git a/python/sglang/jit_kernel/flash_attention/cute/block_sparse_utils.py b/python/sglang/jit_kernel/flash_attention/cute/block_sparse_utils.py new file mode 100644 index 000000000..15a606df0 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/block_sparse_utils.py @@ -0,0 +1,1451 @@ +""" +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 +from functools import partial +import math +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, const_expr + +# Import data structures from block_sparsity +from .block_sparsity import BlockSparseTensors +import sglang.jit_kernel.flash_attention.cute.utils as utils +import sglang.jit_kernel.flash_attention.cute.copy_utils as copy_utils +from .named_barrier import NamedBarrierBwd + + +@cute.jit +def load_block_list( + block_indices: cute.Tensor, + block_count, + load_q_with_first: cutlass.Constexpr, + first_block_preloaded: cutlass.Constexpr, + kv_producer_state, + load_Q, + load_K, + load_V, + pipeline_k, + pipeline_v, + use_tma_q: cutlass.Constexpr, + tma_q_bytes: cutlass.Constexpr, + intra_wg_overlap: cutlass.Constexpr, +): + """Iterate over the sparse blocks and load K, V (and Q) 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. + + 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): + # Peel first iteration: the first block may need to load Q alongside K, + # Parameters are already Constexpr, so no need to wrap in const_expr() + n_block_first = block_indices[block_count - 1] + extra_tx = tma_q_bytes if const_expr(load_q_with_first) and const_expr(use_tma_q) else 0 + pipeline_k.producer_acquire(kv_producer_state, extra_tx_count=extra_tx) + + if const_expr(load_q_with_first and use_tma_q): + load_Q(tma_bar_ptr=pipeline_k.producer_get_barrier(kv_producer_state)) + + load_K(src_idx=n_block_first, producer_state=kv_producer_state) + pipeline_v.producer_acquire(kv_producer_state) + load_V(src_idx=n_block_first, producer_state=kv_producer_state) + kv_producer_state.advance() + + for offset in cutlass.range(1, 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): + extra_tx = ( + tma_q_bytes if const_expr(load_q_with_first) and const_expr(use_tma_q) else 0 + ) + pipeline_k.producer_acquire(kv_producer_state, extra_tx_count=extra_tx) + + if const_expr(load_q_with_first and use_tma_q): + load_Q(tma_bar_ptr=pipeline_k.producer_get_barrier(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], +): + """Map packed m_block indices to block-sparse tensor indices.""" + if const_expr(qhead_per_kvhead != 1): + return m_block // qhead_per_kvhead + return m_block + + +@cute.jit +def produce_block_sparse_loads( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + m_block, + kv_producer_state, + load_Q, + load_K, + load_V, + pipeline_k, + pipeline_v, + use_tma_q: cutlass.Constexpr, + tma_q_bytes: cutlass.Constexpr, + intra_wg_overlap: cutlass.Constexpr, + qhead_per_kvhead: cutlass.Constexpr[int] = 1, +): + """Iterate over the mask and full block lists for a single tile. + + 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. + """ + + mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx = blocksparse_tensors + + m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead) + + curr_mask_block_cnt = mask_block_cnt[batch_idx, head_idx, m_block_sparse] + curr_mask_block_idx = mask_block_idx[batch_idx, head_idx, m_block_sparse, None] + + if const_expr(full_block_cnt is not None): + curr_full_block_cnt = full_block_cnt[batch_idx, head_idx, m_block_sparse] + curr_full_block_idx = full_block_idx[batch_idx, head_idx, m_block_sparse, None] + else: + curr_full_block_cnt = Int32(0) + curr_full_block_idx = None + + 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 Q+K load. + kv_producer_state = load_block_list( + curr_full_block_idx, + curr_full_block_cnt, + load_q_with_first=True, + first_block_preloaded=False, + kv_producer_state=kv_producer_state, + load_Q=load_Q, + load_K=load_K, + load_V=load_V, + pipeline_k=pipeline_k, + pipeline_v=pipeline_v, + use_tma_q=use_tma_q, + tma_q_bytes=tma_q_bytes, + 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: load Q together with the first masked K so consumers can + # start immediately. When overlap is disabled this fully drains the list. + kv_producer_state = load_block_list( + curr_mask_block_idx, + curr_mask_block_cnt, + load_q_with_first=True, + first_block_preloaded=False, + kv_producer_state=kv_producer_state, + load_Q=load_Q, + load_K=load_K, + load_V=load_V, + pipeline_k=pipeline_k, + pipeline_v=pipeline_v, + use_tma_q=use_tma_q, + tma_q_bytes=tma_q_bytes, + 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, + load_q_with_first=False, + first_block_preloaded=True, + kv_producer_state=kv_producer_state, + load_Q=load_Q, + load_K=load_K, + load_V=load_V, + pipeline_k=pipeline_k, + pipeline_v=pipeline_v, + use_tma_q=use_tma_q, + tma_q_bytes=tma_q_bytes, + 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 (skipping the Q + # reload because the masked list already issued it). + kv_producer_state = load_block_list( + curr_full_block_idx, + curr_full_block_cnt, + load_q_with_first=False, + first_block_preloaded=False, + kv_producer_state=kv_producer_state, + load_Q=load_Q, + load_K=load_K, + load_V=load_V, + pipeline_k=pipeline_k, + pipeline_v=pipeline_v, + use_tma_q=use_tma_q, + tma_q_bytes=tma_q_bytes, + 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, + 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, +): + """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. + """ + + mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx = blocksparse_tensors + + m_block_sparse = sparse_tensor_m_block(m_block, qhead_per_kvhead) + + curr_mask_block_cnt = mask_block_cnt[batch_idx, head_idx, m_block_sparse] + curr_mask_block_idx = mask_block_idx[batch_idx, head_idx, m_block_sparse, None] + curr_full_block_cnt = full_block_cnt[batch_idx, head_idx, m_block_sparse] + curr_full_block_idx = full_block_idx[batch_idx, head_idx, m_block_sparse, None] + + 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, + 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, + 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, + 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, + 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, + 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 load_block_list_sm100( + block_indices: cute.Tensor, + block_count, + load_q_with_first: cutlass.Constexpr, + m_block, + 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).""" + if block_count > 0: + # First iteration: load Q alongside K if requested + n_block_first = block_indices[block_count - 1] + + if const_expr(load_q_with_first): + # SM100 loads Q0 and optionally Q1 + load_Q(block=q_stage * m_block + 0, stage=0) + if const_expr(q_stage == 2): + load_Q(block=q_stage * m_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_count - 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, + kv_producer_state, + load_Q, + load_K, + load_V, + pipeline_kv, + q_stage: cutlass.Constexpr, + q_producer_phase: Int32, + qhead_per_kvhead: 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 + """ + # NB: Compute unpacked index for sparse tensor access + if const_expr(qhead_per_kvhead != 1): + m_block_sparse = m_block // qhead_per_kvhead + else: + m_block_sparse = m_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_sparse] + curr_mask_block_idx = mask_block_idx[batch_idx, head_idx, m_block_sparse, None] + + if const_expr(full_block_cnt is not None): + curr_full_block_cnt = full_block_cnt[batch_idx, head_idx, m_block_sparse] + curr_full_block_idx = full_block_idx[batch_idx, head_idx, m_block_sparse, None] + else: + curr_full_block_cnt = Int32(0) + curr_full_block_idx = None + + mask_empty = curr_mask_block_cnt == 0 + full_empty = curr_full_block_cnt == 0 + + 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, + curr_full_block_cnt, + load_q_with_first=True, + m_block=m_block, + 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, + curr_mask_block_cnt, + load_q_with_first=True, + m_block=m_block, + 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, + curr_full_block_cnt, + load_q_with_first=False, + m_block=m_block, + 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, + qhead_per_kvhead: cutlass.Constexpr, +): + # NB: Convert packed m_block to unpacked for sparse tensor indexing + if const_expr(qhead_per_kvhead != 1): + m_block_sparse = m_block // qhead_per_kvhead + else: + m_block_sparse = m_block + + mask_block_cnt, mask_block_idx, full_block_cnt, full_block_idx = blocksparse_tensors + if const_expr(full_block_cnt is not None): + return ( + mask_block_cnt[batch_idx, head_idx, m_block_sparse] + + full_block_cnt[batch_idx, head_idx, m_block_sparse] + ) + else: + return mask_block_cnt[batch_idx, head_idx, m_block_sparse] + + +@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, + m_block: Int32, + head_idx: Int32, + batch_idx: Int32, + split_idx: Int32, + sScale: cute.Tensor, + stats: list, + correction_epilogue: Callable, + thr_mma_pv: cute.core.ThrMma, + tOtOs: tuple[cute.Tensor], + sO: cute.Tensor, + mbar_ptr, + mbar_softmax_corr_full_offset: Int32, + mbar_softmax_corr_empty_offset: Int32, + mbar_P_full_O_rescaled_offset: Int32, + mbar_P_full_2_offset: Int32, + mbar_corr_epi_full_offset: Int32, + mbar_corr_epi_empty_offset: Int32, + softmax_corr_consumer_phase: Int32, + o_corr_consumer_phase: Int32, + corr_epi_producer_phase: Int32, + softmax_scale_log2: Float32, + mO_cur: Optional[cute.Tensor] = None, + gO: Optional[cute.Tensor] = None, + gmem_tiled_copy_O: Optional[cute.TiledCopy] = None, +): + """Handle the block-sparse case where a tile is fully masked: + * zero staged results + * seed stats + * satisfy the usual barrier protocol so downstream warps continue to make progress. + """ + LOG2_E = Float32(math.log2(math.e)) + + 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 = Float32(1.0) + else: + row_sum_value = row_sum_value + utils.exp2f( + sink_val * LOG2_E - row_max_value * softmax_scale_log2 + ) + if tidx < m_block_size: + scale_row_idx = tidx + stage * m_block_size + sScale[scale_row_idx] = row_sum_value + if const_expr(mLSE is not None or learnable_sink is not None): + sScale[scale_row_idx + m_block_size * 2] = row_max_value + 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) + + cute.arch.mbarrier_wait( + mbar_ptr + mbar_softmax_corr_full_offset + stage, + softmax_corr_consumer_phase, + ) + cute.arch.mbarrier_arrive(mbar_ptr + mbar_softmax_corr_empty_offset + stage) + + if const_expr(gmem_tiled_copy_O is None): + cute.arch.mbarrier_wait( + mbar_ptr + mbar_corr_epi_empty_offset + stage, + corr_epi_producer_phase, + ) + correction_epilogue( + thr_mma_pv, + tOtOs[stage], + tidx, + stage, + m_block, + seqlen.seqlen_q, + Float32(0.0), # zero scale ensures empty tile writes zeros into staged outputs + sO[None, None, stage], + mO_cur, + gO, + gmem_tiled_copy_O, + ) + if const_expr(gmem_tiled_copy_O is None): + cute.arch.mbarrier_arrive(mbar_ptr + mbar_corr_epi_full_offset + stage) + cute.arch.mbarrier_arrive(mbar_ptr + mbar_P_full_O_rescaled_offset + stage) + cute.arch.mbarrier_arrive(mbar_ptr + mbar_P_full_2_offset + stage) + + softmax_corr_consumer_phase ^= 1 + o_corr_consumer_phase ^= 1 + corr_epi_producer_phase ^= 1 + + return ( + softmax_corr_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, + 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, + mbar_ptr, + mbar_softmax_corr_full_offset: Int32, + mbar_softmax_corr_empty_offset: Int32, + mbar_P_full_O_rescaled_offset: Int32, + mbar_P_full_2_offset: Int32, + q_stage: cutlass.Constexpr, + stage_idx: Int32, + check_m_boundary: bool, + qhead_per_kvhead: cutlass.Constexpr, +): + # Convert packed m_block to unpacked for sparse tensor indexing + if const_expr(qhead_per_kvhead != 1): + m_block_sparse = m_block // qhead_per_kvhead + else: + m_block_sparse = m_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_sparse] + curr_mask_block_idx = mask_block_idx[batch_idx, head_idx, m_block_sparse, None] + + if const_expr(full_block_cnt is not None): + curr_full_block_cnt = full_block_cnt[batch_idx, head_idx, m_block_sparse] + curr_full_block_idx = full_block_idx[batch_idx, head_idx, m_block_sparse, None] + else: + curr_full_block_cnt = Int32(0) + curr_full_block_idx = None + + total_block_cnt = curr_mask_block_cnt + curr_full_block_cnt + + if total_block_cnt == 0: + cute.arch.mbarrier_arrive(mbar_ptr + mbar_softmax_corr_full_offset + stage_idx) + cute.arch.mbarrier_arrive(mbar_ptr + mbar_P_full_O_rescaled_offset + stage_idx) + cute.arch.mbarrier_arrive(mbar_ptr + mbar_P_full_2_offset + stage_idx) + cute.arch.mbarrier_arrive(mbar_ptr + mbar_softmax_corr_empty_offset + stage_idx) + else: + if curr_mask_block_cnt > 0: + mask_n_block = curr_mask_block_idx[curr_mask_block_cnt - 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, curr_mask_block_cnt): + mask_n_block = curr_mask_block_idx[curr_mask_block_cnt - 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 curr_full_block_cnt > 0: + full_n_block = curr_full_block_idx[curr_full_block_cnt - 1] + if curr_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=False, check_q_boundary=check_m_boundary + ), + ) + for i in cutlass.range(1, curr_full_block_cnt): + full_n_block = curr_full_block_idx[curr_full_block_cnt - 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, + ) + + +# ============================================================================= +# 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, + 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 * 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 + 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, 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, + 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, + 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 * 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], + 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 // subtile_factor + subtile_offset = iter_idx % 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 * 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) + with cute.arch.elect_one(): + 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) + with cute.arch.elect_one(): + 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, + 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 * subtile_factor, unroll=1): + sparse_idx = iter_idx // subtile_factor + subtile_offset = iter_idx % subtile_factor + m_block = curr_q_idx[sparse_idx] * 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 * subtile_factor, unroll=1): + sparse_idx = iter_idx // subtile_factor + subtile_offset = iter_idx % subtile_factor + m_block = curr_full_idx[sparse_idx] * 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, + softmax_scale, + seqlen, + subtile_factor: cutlass.Constexpr, + m_block_max: int, + aux_tensors=None, + 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_tensors=aux_tensors, + 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_tensors=aux_tensors, + fastdiv_mods=fastdiv_mods, + ) + + for iter_idx in cutlass.range(curr_q_cnt * subtile_factor, unroll=1): + sparse_idx = iter_idx // subtile_factor + subtile_offset = iter_idx % subtile_factor + m_block = curr_q_idx[sparse_idx] * 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, + dKV_accumulate=dKV_accumulate, + thr_mma_SdP=thr_mma_SdP, + batch_idx=batch_idx, + head_idx=head_idx, + n_block=n_block, + softmax_scale=softmax_scale, + seqlen=seqlen, + aux_tensors=aux_tensors, + fastdiv_mods=fastdiv_mods, + ) + dKV_accumulate = True + + if const_expr(full_cnt is not None): + for iter_idx in cutlass.range(curr_full_cnt * subtile_factor, unroll=1): + sparse_idx = iter_idx // subtile_factor + subtile_offset = iter_idx % subtile_factor + m_block = curr_full_idx[sparse_idx] * 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, + dKV_accumulate=dKV_accumulate, + thr_mma_SdP=thr_mma_SdP, + batch_idx=batch_idx, + head_idx=head_idx, + n_block=n_block, + softmax_scale=softmax_scale, + seqlen=seqlen, + aux_tensors=aux_tensors, + fastdiv_mods=fastdiv_mods, + ) + 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_mma_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_mma_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() + for warp_group_idx in cutlass.range_constexpr(num_mma_warp_groups): + cute.arch.cp_async_bulk_wait_group(num_mma_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, + ) + + +@cute.jit +def dQaccum_store_block_sparse_bwd_sm90( + blocksparse_tensors: BlockSparseTensors, + batch_idx, + head_idx, + n_block, + sdQaccum: cute.Tensor, + gdQaccum: cute.Tensor, + subtile_factor: cutlass.Constexpr, + m_block_max: int, + num_mma_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 * subtile_factor, unroll=1): + sparse_idx = iter_idx // subtile_factor + subtile_offset = iter_idx % subtile_factor + m_block = curr_q_idx[sparse_idx] * subtile_factor + subtile_offset + + if m_block < m_block_max: + _store_one_dQaccum_sm90( + m_block, + sdQaccum, + gdQaccum, + num_mma_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 * subtile_factor, unroll=1): + sparse_idx = iter_idx // subtile_factor + subtile_offset = iter_idx % subtile_factor + m_block = curr_full_idx[sparse_idx] * subtile_factor + subtile_offset + + if m_block < m_block_max: + _store_one_dQaccum_sm90( + m_block, + sdQaccum, + gdQaccum, + num_mma_warp_groups, + num_threads_per_warp_group, + tma_copy_bytes_dQ, + ) diff --git a/python/sglang/jit_kernel/flash_attention/cute/block_sparsity.py b/python/sglang/jit_kernel/flash_attention/cute/block_sparsity.py new file mode 100644 index 000000000..99aa69083 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/block_sparsity.py @@ -0,0 +1,250 @@ +""" +Block-sparsity utilities for FlexAttention +""" + +from typing import Callable, NamedTuple, Tuple + +import cutlass.cute as cute +import torch + +from .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 + full_block_idx: cute.Tensor | None + + def __new_from_mlir_values__(self, values): + if len(values) == 2: + values = (*values, None, None) + return BlockSparseTensors(*values) + + +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 + + +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, int, int], + expected_index_shape: Tuple[int, int, int, 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 + ) + expanded_idx = _expand_sparsity_tensor( + idx, expected_index_shape, f"{name}_block_idx", context, hint + ) + return expanded_cnt, expanded_idx + + +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 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, + 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 subtile_factor * m_block_size. + """ + sparse_block_size_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, int, int], + expected_index_shape: Tuple[int, int, int, 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") + + return BlockSparseTensorsTorch( + mask_block_cnt=mask_cnt, + mask_block_idx=mask_idx, + full_block_cnt=full_cnt, + full_block_idx=full_idx, + ) + + +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, + ): + if tensor is not None: + patterns.append(get_broadcast_dims(tensor)) + else: + patterns.append(None) + return tuple(patterns) + + +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, + mask_block_idx, + full_block_cnt, + full_block_idx, + ) = tensors + + ( + 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 (mask_block_cnt, 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 (full_block_cnt, full_block_idx) + ] + + return BlockSparseTensors( + mask_block_cnt_tensor, + mask_block_idx_tensor, + full_block_cnt_tensor, + full_block_idx_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/python/sglang/jit_kernel/flash_attention/cute/compute_block_sparsity.py b/python/sglang/jit_kernel/flash_attention/cute/compute_block_sparsity.py new file mode 100644 index 000000000..062ecb288 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/compute_block_sparsity.py @@ -0,0 +1,377 @@ +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 .block_sparsity import ( + BlockSparseTensors, + BlockSparseTensorsTorch, + to_cute_block_sparse_tensors, +) +from .utils import hash_callable, scalar_to_ssa, ssa_to_scalar +from .seqlen_info import SeqlenInfoQK + + +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 + - varlen support + - 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, + aux_tensors: Optional[list] = None, + ): + self.mask_cnt, self.mask_idx, self.full_cnt, self.full_idx = blocksparse_tensors + + if const_expr(self.compute_full_blocks): + assert self.full_cnt is not None and self.full_idx is not None, ( + "full block tensors must be provided when computing full blocks" + ) + + batch_size, num_heads, num_m_blocks, num_n_blocks = self.mask_idx.shape + # launch 1 CTA per m block + grid = [num_m_blocks, num_heads, batch_size] + + 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 + + self.kernel( + self.mask_cnt, + self.mask_idx, + self.full_cnt, + self.full_idx, + num_n_blocks, + seqlen_q, + seqlen_k, + aux_tensors, + ).launch(grid=grid, block=[num_threads, 1, 1]) + + @cute.kernel + def kernel( + self, + mask_cnt: cute.Tensor, + mask_idx: cute.Tensor, + full_cnt: cute.Tensor, + full_idx: cute.Tensor, + num_n_blocks: Int32, + seqlen_q: Int32, + seqlen_k: Int32, + aux_tensors: Optional[list] = None, + ): + tidx, _, _ = cute.arch.thread_idx() + warp_idx = cute.arch.warp_idx() + lane_id = cute.arch.lane_idx() + m_block, head_idx, batch_idx = cute.arch.block_idx() + + ssa = partial(scalar_to_ssa, dtype=Int32) + + seqlen = SeqlenInfoQK.create( + batch_idx, + seqlen_q, + seqlen_k, + mCuSeqlensQ=None, + mCuSeqlensK=None, + mSeqUsedQ=None, + mSeqUsedK=None, + ) + + @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)) + ) + + num_mask_blocks = Int32(0) + num_full_blocks = Int32(0) + + for n_block in cutlass.range(num_n_blocks, unroll_full=True): + m_base = m_block * self.tile_mn[0] + n_base = n_block * self.tile_mn[1] + + if const_expr(self.use_fast_sampling): + # Fast path: 5-point sampling (4 corners + center) + # Clamps OOB indices to nearest in bounds. + thread_result = Boolean(False) + thread_is_valid = Boolean(False) + q_idx = Int32(0) + kv_idx = Int32(0) + + if tidx == 0: + # Top-left corner (0, 0); always in bounds + q_idx = m_base + kv_idx = n_base + elif tidx == 1: + # Top-right corner + q_idx = m_base + kv_idx = cutlass.min(n_base + self.tile_mn[1] - 1, seqlen_k - 1) + elif tidx == 2: + # Bottom-left corner + q_idx = cutlass.min(m_base + self.tile_mn[0] - 1, seqlen_q - 1) + kv_idx = n_base + elif tidx == 3: + # Bottom-right corner + q_idx = cutlass.min(m_base + self.tile_mn[0] - 1, seqlen_q - 1) + kv_idx = cutlass.min(n_base + self.tile_mn[1] - 1, seqlen_k - 1) + elif tidx == 4: + # Center point + q_idx = m_base + (cutlass.min(seqlen_q - m_base, self.tile_mn[0])) // 2 + kv_idx = n_base + (cutlass.min(seqlen_k - n_base, self.tile_mn[1])) // 2 + else: + thread_is_valid = Boolean(False) + + # Check bounds and determine if this thread has a valid index pair + if tidx < 5 and q_idx < seqlen_q and kv_idx < seqlen_k: + thread_is_valid = Boolean(True) + q_idx_ssa = ssa(q_idx) + kv_idx_ssa = ssa(kv_idx) + thread_result = ssa_to_scalar( + self.mask_mod( + ssa(batch_idx), + ssa(head_idx), + q_idx_ssa, + kv_idx_ssa, + seqlen, + aux_tensors, + ) + ) + else: + thread_is_valid = Boolean(False) + + # Use vote_any_sync to see if any valid thread found unmasked or masked + # Only count results from threads that checked valid indices + 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: check all elements in the block + # Track if this thread's row has any masked or unmasked elements + thread_has_unmasked = Boolean(False) + thread_has_masked = Boolean(False) + thread_is_valid = Boolean(False) + + # Each thread handles 1 row + q_idx = m_base + tidx + kv_idx = Int32(0) + if tidx < self.tile_mn[0] and q_idx < seqlen_q: + thread_is_valid = Boolean(True) + q_idx_ssa = ssa(q_idx) + + # Loop over all columns in this row + for c in cutlass.range(self.tile_mn[1], unroll_full=True): + kv_idx = n_base + c + kv_idx_ssa = ssa(kv_idx) + + # Only check elements within valid sequence bounds + if kv_idx < seqlen_k: + # Direct scalar call + mask_val = ssa_to_scalar( + self.mask_mod( + ssa(batch_idx), + ssa(head_idx), + q_idx_ssa, + kv_idx_ssa, + seqlen, + aux_tensors, + ) + ) + + # Update tracking flags + if mask_val: + thread_has_unmasked = Boolean(True) + else: + thread_has_masked = Boolean(True) + + # Block-level reduction to combine results across all threads + # Only count votes from threads that checked valid indices + warp_has_unmasked_mask = cute.arch.vote_any_sync( + thread_has_unmasked & thread_is_valid + ) + warp_has_masked_mask = cute.arch.vote_any_sync(thread_has_masked & thread_is_valid) + + # lane 0 writes the ballot mask to shared memory + lane_id = tidx % 32 + if lane_id == 0: + # Store as Int8 + reduction_buffer[warp_idx, 0] = Int8(1) if warp_has_unmasked_mask else Int8(0) + reduction_buffer[warp_idx, 1] = Int8(1) if warp_has_masked_mask else Int8(0) + + cute.arch.sync_threads() + + # Thread 0 ORs all warp results together + has_unmasked = Boolean(False) + has_masked = Boolean(False) + if tidx == 0: + for w in cutlass.range(self.num_warps): + if reduction_buffer[w, 0]: + has_unmasked = Boolean(True) + if reduction_buffer[w, 1]: + has_masked = Boolean(True) + + # 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: + mask_idx[batch_idx, head_idx, m_block, num_mask_blocks] = n_block + num_mask_blocks += 1 + elif is_full and const_expr(self.compute_full_blocks): + full_idx[batch_idx, head_idx, m_block, num_full_blocks] = n_block + num_full_blocks += 1 + + # Only thread 0 writes back the counts + if tidx == 0: + 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], # list[cute.Tensor] + device, + compute_full_blocks: bool = True, + use_fast_sampling: bool = False, +) -> Tuple[BlockSparseTensors, 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. + 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: + A tuple of `BlockSparseTensors` and `BlockSparseTensorsTorch`. + """ + # Check if mask_mod is marked as suitable for 5-point fast 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 + + 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, + ) + + mask_mod_hash = hash_callable(mask_mod) + blocksparse_tensors = to_cute_block_sparse_tensors( + blocksparse_tensors_torch, enable_tvm_ffi=True + ) + + compile_key = ( + tile_m, + tile_n, + mask_mod_hash, + compute_full_blocks, + aux_tensors is not None, + use_fast_sampling, + ) + if compile_key not in compute_block_sparsity.compile_cache: + 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, aux_tensors, options="--enable-tvm-ffi" + ) + + compute_block_sparsity.compile_cache[compile_key]( + blocksparse_tensors_torch, + seqlen_q, + seqlen_k, + aux_tensors, + ) + + return blocksparse_tensors, blocksparse_tensors_torch + + +compute_block_sparsity.compile_cache = {} diff --git a/python/sglang/jit_kernel/flash_attention/cute/copy_utils.py b/python/sglang/jit_kernel/flash_attention/cute/copy_utils.py new file mode 100644 index 000000000..cfdcbdb80 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/copy_utils.py @@ -0,0 +1,340 @@ +# 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_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/python/sglang/jit_kernel/flash_attention/cute/cute_dsl_utils.py b/python/sglang/jit_kernel/flash_attention/cute/cute_dsl_utils.py new file mode 100644 index 000000000..14723872b --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/cute_dsl_utils.py @@ -0,0 +1,144 @@ +# Copyright (c) 2025, Tri Dao. + +import os +import pathlib +from typing import Tuple +from functools import partial, lru_cache +from dataclasses import dataclass, fields + +import torch + +try: + from triton.tools.disasm import extract +except ImportError: + extract = None + +import cutlass +import cutlass.cute as cute +from cutlass.base_dsl.typing import JitArgument +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, +} + + +@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) + + +@dataclass +class ParamsBase: + def __extract_mlir_values__(self): + all_fields = [getattr(self, field.name) for field in fields(self)] + non_constexpr_fields = [f for f in all_fields if not isinstance(f, StaticTypes)] + values, self._values_pos = [], [] + for obj in non_constexpr_fields: + 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): + all_fields = {field.name: getattr(self, field.name) for field in fields(self)} + constexpr_fields = {n: f for n, f in all_fields.items() if isinstance(f, StaticTypes)} + non_constexpr_fields = { + n: f for n, f in all_fields.items() if not isinstance(f, StaticTypes) + } + for (name, field), n_items in zip(non_constexpr_fields.items(), self._values_pos): + non_constexpr_fields[name] = cutlass.new_from_mlir_values(field, values[:n_items]) + values = values[n_items:] + return self.__class__(**non_constexpr_fields, **constexpr_fields) + + +@dataclass +class ArgumentsBase(JitArgument): + def __c_pointers__(self): + all_fields = [getattr(self, field.name) for field in fields(self)] + non_constexpr_fields = [f for f in all_fields if not isinstance(f, StaticTypes)] + c_ptrs = [] + for obj in non_constexpr_fields: + if hasattr(obj, "__c_pointers__"): + c_ptrs.extend(obj.__c_pointers__()) + return c_ptrs + + def __get_mlir_types__(self): + all_fields = [getattr(self, field.name) for field in fields(self)] + non_constexpr_fields = [f for f in all_fields if not isinstance(f, StaticTypes)] + types, self._values_pos = [], [] + for obj in non_constexpr_fields: + if hasattr(obj, "__get_mlir_types__"): + obj_types = obj.__get_mlir_types__() + types.extend(obj_types) + self._values_pos.append(len(obj_types)) + else: + self._values_pos.append(0) + return types + + def __new_from_mlir_values__(self, values): + all_fields = {field.name: getattr(self, field.name) for field in fields(self)} + constexpr_fields = {n: f for n, f in all_fields.items() if isinstance(f, StaticTypes)} + non_constexpr_fields = { + n: f for n, f in all_fields.items() if not isinstance(f, StaticTypes) + } + for (name, field), n_items in zip(non_constexpr_fields.items(), self._values_pos): + non_constexpr_fields[name] = cutlass.new_from_mlir_values(field, values[:n_items]) + values = values[n_items:] + return self.__class__(**non_constexpr_fields, **constexpr_fields) + + +def load_cubin_module_data_patched(cubin_data, filepath): + pathlib.Path(filepath).write_bytes(cubin_data) + return load_cubin_module_data_og(cubin_data) + + +def cute_compile_patched(*args, **kwargs): + """A patched version of cute.compile that dump the SASS to a file if CUTE_CUBIN_PATH is set.""" + cubin_path = os.getenv("CUTE_CUBIN_PATH", None) + if cubin_path is not None: + cutlass.base_dsl.runtime.cuda.load_cubin_module_data = partial( + load_cubin_module_data_patched, filepath=cubin_path + ) + output = cute_compile_og(*args, **kwargs) + if cubin_path is not None: + cutlass.base_dsl.runtime.cuda.load_cubin_module_data = load_cubin_module_data_og + if extract is not None: + sass = extract(cubin_path, None) + pathlib.Path(cubin_path).with_suffix(".annotated.sass").write_text(sass) + return output + + +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.""" + 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 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()) diff --git a/python/sglang/jit_kernel/flash_attention/cute/fast_math.py b/python/sglang/jit_kernel/flash_attention/cute/fast_math.py new file mode 100644 index 000000000..c56ea89e7 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/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/python/sglang/jit_kernel/flash_attention/cute/flash_bwd.py b/python/sglang/jit_kernel/flash_attention/cute/flash_bwd.py new file mode 100644 index 000000000..1ac795bcf --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/flash_bwd.py @@ -0,0 +1,1262 @@ +# 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 Float32, Int32 +import cutlass.utils as utils_basic + +import sglang.jit_kernel.flash_attention.cute.ampere_helpers as sm80_utils +import sglang.jit_kernel.flash_attention.cute.utils as utils +from .mask import AttentionMask +from .seqlen_info import SeqlenInfoQK +from .tile_scheduler import ParamsBase, SingleTileScheduler, SingleTileVarlenScheduler, TileSchedulerArguments + + +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, + ): + """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 16 as k_block_size + 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 + + @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(not mQ_type in [cutlass.Float16, cutlass.BFloat16]): + raise TypeError("Only Float16 or BFloat16 is supported") + if cutlass.const_expr(not mLSE_type in [cutlass.Float32]): + raise TypeError("LSE tensor must be Float32") + if cutlass.const_expr(not mdPsum_type in [cutlass.Float32]): + raise TypeError("dPsum tensor must be Float32") + if cutlass.const_expr(not mdQaccum_type 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, + stream: cuda.CUstream, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + softcap: Float32 | float | None = None, + window_size_left: Int32 | int | None = None, + window_size_right: Int32 | int | None = None, + mdQ_semaphore: Optional[cute.Tensor] = None, + ): + assert mdQ_semaphore is None, "semaphore not supported yet" + # 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))) + # Assume all strides are divisible by 128 bits except the last stride + # Skip cute.assume() for stride=0 (broadcast dims from expand() are Python ints) + new_stride = lambda t: (*(cute.assume(s, divby=128 // t.element_type.width) if s != 0 else s for s in t.stride[:-1]), t.stride[-1]) + mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV = [cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) if t is not None else None 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 * math.log2(math.e) + 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, + ).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], + ): + # 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) + + 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.offset_q + batch_idx * self.m_block_size + 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 = [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_fragment(acc_shape_dK, cutlass.Float32) + acc_dV = cute.make_fragment(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 = utils.make_acc_tensor_mn_view(thr_mma_sdp.partition_C(sLSEMma))[LSEslice] + tSsdPsumMma = utils.make_acc_tensor_mn_view(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_log2=softmax_scale_log2, + ) + + # /////////////////////////////////////////////////////////////////////////////// + # 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.seqlen_q, seqlen.seqlen_k) + mask_fn = partial( + mask.apply_mask, n_block=n_block, thr_mma=thr_mma_sdp, + 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_log2: cutlass.Float32, + 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_fragment(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, + ) + 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 + ) + if cutlass.const_expr(mask_fn is not None): + mask_fn(acc_S, m_block=m_block) + acc_S_mn = utils.make_acc_tensor_mn_view(acc_S) + 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(utils.exp2f(acc_S_mn[r, None].load() * softmax_scale_log2 - tLSErLSE[r])) + # 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_fragment(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 = utils.make_acc_tensor_mn_view(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): + acc_dP_mn[r, None].store(acc_S_mn[r, None].load() * (acc_dP_mn[r, None].load() - tLSErdPsum[r])) + # 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 = cute.make_tensor(rP.iterator, utils.convert_layout_acc_frgA(rP.layout)) + 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_fragment(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 = cute.make_tensor(rdS.iterator, utils.convert_layout_acc_frgA(rdS.layout)) + 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/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_postprocess.py b/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_postprocess.py new file mode 100644 index 000000000..3850a80ff --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_postprocess.py @@ -0,0 +1,463 @@ +# 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, Literal + +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 + +import sglang.jit_kernel.flash_attention.cute.utils as utils +import sglang.jit_kernel.flash_attention.cute.copy_utils as copy_utils +import sglang.jit_kernel.flash_attention.cute.ampere_helpers as sm80_utils +import sglang.jit_kernel.flash_attention.cute.hopper_helpers as sm90_utils +from .seqlen_info import SeqlenInfoQK +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +from .tile_scheduler import ( + ParamsBase, + SingleTileScheduler, + SingleTileVarlenScheduler, + TileSchedulerArguments, +) + + +class FlashAttentionBackwardPostprocess: + def __init__( + self, + dtype: Type[cutlass.Numeric], + head_dim: int, + arch: Literal[80, 90, 100], + tile_m: int = 128, + num_threads: int = 256, + AtomLayoutMdQ: int = 1, + dQ_swapAB: bool = False, + ): + """ + :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 in [80, 90, 100], ( + "Only Ampere (80), Hopper (90), and Blackwell (100) 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 + + @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 == 80): + 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 == 90): + num_mma_warp_groups = self.num_threads // 128 + atom_layout_dQ = (self.AtomLayoutMdQ, num_mma_warp_groups // 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 in [80, 90]): + 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 == 80) else 4 + if const_expr(self.arch == 80): + 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 == 90): + num_threads_per_warp_group = 128 + num_mma_warp_groups = 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_mma_warp_groups)), # thr_layout + cute.make_layout(128 // Float32.width), # val_layout + ) + self.sdQaccum_layout = cute.make_layout( + (self.tile_m * self.tile_hdim // num_mma_warp_groups, num_mma_warp_groups) + ) + 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) + ) + + self.gmem_tiled_copy_dQ = copy_utils.tiled_copy_2d( + self.dtype, self.tile_hdim, self.num_threads + ) + # /////////////////////////////////////////////////////////////////////////////// + # 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 == 80): + 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 == 90): + self.sdQ_layout = sm90_utils.make_smem_layout( + self.dtype, LayoutEnum.ROW_MAJOR, (self.tile_m, self.tile_hdim) + ) + 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], + stream: cuda.CUstream, + ): + # 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") + + # Assume all strides are divisible by 128 bits except the last stride + new_stride = lambda t: ( + *(cute.assume(s, divby=128 // t.element_type.width) for s in t.stride[:-1]), + t.stride[-1], + ) + mdQaccum, mdQ = [ + cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(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 in [80, 90]): + 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 = 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, + ) + 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.offset_q + batch_idx * self.tile_m + if cutlass.const_expr(self.arch >= 90): + padded_offset_q = padded_offset_q // self.tile_m * self.tile_m + 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) + + # 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 in [80, 90]): + acc_shape = tiled_mma.partition_shape_C( + tile_shape if const_expr(not dQ_swapAB) else tile_shape[::-1] + ) + acc = cute.make_fragment(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_fragment(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 in [80, 90]): + 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 in [80, 90]): + 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/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_preprocess.py b/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_preprocess.py new file mode 100644 index 000000000..b677b140b --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_preprocess.py @@ -0,0 +1,365 @@ +# 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. +import math +import operator +from typing import Callable, Type, Optional, Literal + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass import Float32 + +import sglang.jit_kernel.flash_attention.cute.utils as utils +import sglang.jit_kernel.flash_attention.cute.copy_utils as copy_utils +from .seqlen_info import SeqlenInfoQK +from .tile_scheduler import ( + ParamsBase, + SingleTileScheduler, + SingleTileVarlenScheduler, + TileSchedulerArguments, +) + + +class FlashAttentionBackwardPreprocess: + def __init__( + self, + dtype: Type[cutlass.Numeric], + head_dim: int, + arch: Literal[80, 90, 100], + m_block_size: int = 128, + num_threads: int = 128, + ): + """ + 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 num_threads: number of threads + :type num_threads: int + """ + self.dtype = dtype + self.m_block_size = m_block_size + self.arch = arch + # 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.check_hdim_oob = head_dim != self.head_dim_padded + self.num_threads = num_threads + + @staticmethod + def can_implement(dtype, head_dim, m_block_size, 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 m_block_size: m block size + :type m_block_size: 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 < m_block_size: # 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_padded % 128 == 0 + else ( + 64 + if self.head_dim_padded % 64 == 0 + else (32 if self.head_dim_padded % 32 == 0 else 16) + ) + ) + self.gmem_tiled_copy_O = copy_utils.tiled_copy_2d( + self.dtype, gmem_k_block_size, self.num_threads + ) + universal_copy_bits = 128 + num_copy_elems_dQaccum = universal_copy_bits // Float32.width + assert ( + self.m_block_size * 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, + mdO: cute.Tensor, + mdPsum: cute.Tensor, + mLSE: Optional[cute.Tensor], + mLSElog2: Optional[cute.Tensor], + mdQaccum: Optional[cute.Tensor], + mCuSeqlensQ: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + stream: cuda.CUstream, + ): + # Get the data type and check if it is fp16 or bf16 + if cutlass.const_expr(not (mO.element_type == mdO.element_type)): + raise TypeError("All tensors must have the same data type") + if cutlass.const_expr(mO.element_type not in [cutlass.Float16, cutlass.BFloat16]): + raise TypeError("Only Float16 or BFloat16 is supported") + if cutlass.const_expr(mdPsum.element_type not in [Float32]): + raise TypeError("dPsum tensor must be Float32") + if cutlass.const_expr(mdQaccum is not None): + if cutlass.const_expr(mdQaccum.element_type not in [Float32]): + raise TypeError("dQaccum tensor must be Float32") + if cutlass.const_expr(mLSE is not None): + assert mLSElog2 is not None, "If mLSE is provided, mLSElog2 must also be provided" + if cutlass.const_expr(mLSE.element_type not in [Float32]): + raise TypeError("LSE tensor must be Float32") + if cutlass.const_expr(mLSElog2.element_type not in [Float32]): + raise TypeError("LSElog2 tensor must be Float32") + + # Assume all strides are divisible by 128 bits except the last stride + new_stride = lambda t: ( + *(cute.assume(s, divby=128 // t.element_type.width) for s in t.stride[:-1]), + t.stride[-1], + ) + mO, mdO, mdQaccum = [ + cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) + if t is not None + else None + for t in (mO, mdO, mdQaccum) + ] + + self._setup_attributes() + + if cutlass.const_expr(mCuSeqlensQ is not None): + TileScheduler = SingleTileVarlenScheduler + num_head = mO.shape[1] + num_batch = mCuSeqlensQ.shape[0] - 1 + else: + TileScheduler = SingleTileScheduler + num_head = mO.shape[2] + num_batch = mO.shape[0] + + tile_sched_args = TileSchedulerArguments( + num_block=cute.ceil_div(mO.shape[1], self.m_block_size), + num_head=num_head, + num_batch=num_batch, + num_splits=1, + seqlen_k=0, + headdim=0, + headdim_v=mO.shape[2], + total_q=mO.shape[0], + tile_shape_mn=(self.m_block_size, 1), + mCuSeqlensQ=mCuSeqlensQ, + mSeqUsedQ=mSeqUsedQ, + ) + + tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args) + grid_dim = TileScheduler.get_grid_shape(tile_sched_params) + + self.kernel( + mO, + mdO, + mdPsum, + mLSE, + mLSElog2, + mdQaccum, + mCuSeqlensQ, + mSeqUsedQ, + 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, + ) + + @cute.kernel + def kernel( + self, + mO: cute.Tensor, + mdO: cute.Tensor, + mdPsum: cute.Tensor, + mLSE: Optional[cute.Tensor], + mLSElog2: Optional[cute.Tensor], + mdQaccum: Optional[cute.Tensor], + mCuSeqlensQ: Optional[cute.Tensor], + mSeqUsedQ: Optional[cute.Tensor], + 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 + + if work_tile.is_valid_tile: + # /////////////////////////////////////////////////////////////////////////////// + # Get the appropriate tiles for this thread block. + # /////////////////////////////////////////////////////////////////////////////// + seqlen = SeqlenInfoQK.create( + batch_idx, + mO.shape[1], + 0, + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=None, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=None, + ) + + if cutlass.const_expr(not seqlen.has_cu_seqlens_q): + mO_cur = mO[batch_idx, None, head_idx, None] + mdO_cur = mdO[batch_idx, None, head_idx, None] + mdPsum_cur = mdPsum[batch_idx, head_idx, None] + headdim_v = mO.shape[3] + else: + mO_cur = cute.domain_offset((seqlen.offset_q, 0), mO[None, head_idx, None]) + mdO_cur = cute.domain_offset((seqlen.offset_q, 0), mdO[None, head_idx, None]) + + padded_offset_q = seqlen.offset_q + batch_idx * self.m_block_size + if cutlass.const_expr(self.arch >= 90): + padded_offset_q = padded_offset_q // self.m_block_size * self.m_block_size + mdPsum_cur = cute.domain_offset((padded_offset_q,), mdPsum[head_idx, None]) + headdim_v = mO.shape[2] + + blkOdO_shape = (self.m_block_size, self.head_dim_padded) + # (m_block_size, head_dim) + gO = cute.local_tile(mO_cur, blkOdO_shape, (m_block, 0)) + gdO = cute.local_tile(mdO_cur, blkOdO_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) + + # /////////////////////////////////////////////////////////////////////////////// + # Predicate: Mark indices that need to copy when problem_shape isn't a multiple + # of tile_shape + # /////////////////////////////////////////////////////////////////////////////// + # Construct identity layout for KV + cO = cute.make_identity_tensor((self.m_block_size, self.head_dim_padded)) + tOcO = gmem_thr_copy_O.partition_S(cO) + t0OcO = gmem_thr_copy_O.get_slice(0).partition_S(cO) + tOpO = utils.predicate_k(tOcO, limit=headdim_v) + tOpdO = utils.predicate_k(tOcO, limit=headdim_v) + + seqlen_q = seqlen.seqlen_q + seqlen_q_rounded = cute.round_up(seqlen_q, self.m_block_size) + + if cutlass.const_expr(mLSE is not None): + if cutlass.const_expr(not seqlen.has_cu_seqlens_q): + mLSE_cur = mLSE[batch_idx, head_idx, None] + else: + mLSE_cur = cute.domain_offset((seqlen.offset_q,), mLSE[head_idx, None]) + + gLSE = cute.local_tile(mLSE_cur, (self.m_block_size,), (m_block,)) + lse = Float32.inf + if tidx < seqlen_q - m_block * self.m_block_size: + lse = gLSE[tidx] + + tOrO = cute.make_fragment_like(tOgO) + tOrdO = cute.make_fragment_like(tOgdO) + assert cute.size(tOgO, mode=[0]) == cute.size(tOgdO, mode=[0]) + assert cute.size(tOgO, mode=[1]) == cute.size(tOgdO, mode=[1]) + assert cute.size(tOgO, mode=[2]) == cute.size(tOgdO, mode=[2]) + 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 + # (seqlen_q - m_block * kBlockM). This is because the entries of t0OcO are known at compile time. + if t0OcO[0, m, 0][0] < seqlen_q - m_block * self.m_block_size - tOcO[0][0]: + cute.copy( + gmem_thr_copy_O, + tOgO[None, m, None], + tOrO[None, m, None], + pred=tOpO[None, m, None] + if cutlass.const_expr(self.check_hdim_oob) + else None, + ) + cute.copy( + gmem_thr_copy_O, + tOgdO[None, m, None], + tOrdO[None, m, None], + pred=tOpdO[None, m, None] + if cutlass.const_expr(self.check_hdim_oob) + else None, + ) + # Sum across the "k" dimension + dpsum = (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 + dpsum = utils.warp_reduce(dpsum, operator.add, width=threads_per_row) + dP_sum = cute.make_fragment(cute.size(tOrO, mode=[1]), Float32) + dP_sum.store(dpsum) + + # Write dPsum from rmem -> gmem + gdPsum = cute.local_tile(mdPsum_cur, (self.m_block_size,), (m_block,)) + # Only the thread corresponding to column 0 writes out the dPsum to gmem + if tOcO[0, 0, 0][1] == 0: + for m in cutlass.range(cute.size(dP_sum), unroll_full=True): + row = tOcO[0, m, 0][0] + gdPsum[row] = dP_sum[m] if row < seqlen_q - m_block * self.m_block_size else 0.0 + + # Clear dQaccum + if cutlass.const_expr(mdQaccum is not None): + if cutlass.const_expr(not seqlen.has_cu_seqlens_q): + mdQaccum_cur = mdQaccum[batch_idx, head_idx, None] + else: + mdQaccum_cur = cute.domain_offset( + (padded_offset_q * self.head_dim_padded,), mdQaccum[head_idx, None] + ) + + # HACK: Compiler doesn't seem to recognize that padding + # by padded_offset_q * self.head_dim_padded 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) + + blkdQaccum_shape = (self.m_block_size * 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_fragment_like(tdQgdQaccum) + zero.fill(0.0) + cute.copy(gmem_tiled_copy_dQaccum, zero, tdQgdQaccum) + + if cutlass.const_expr(mLSE is not None): + if cutlass.const_expr(not seqlen.has_cu_seqlens_q): + mLSElog2_cur = mLSElog2[batch_idx, head_idx, None] + else: + mLSElog2_cur = cute.domain_offset((padded_offset_q,), mLSElog2[head_idx, None]) + + gLSElog2 = cute.local_tile(mLSElog2_cur, (self.m_block_size,), (m_block,)) + LOG2_E = math.log2(math.e) + if tidx < seqlen_q_rounded - m_block * self.m_block_size: + gLSElog2[tidx] = lse * LOG2_E if lse != -Float32.inf else 0.0 diff --git a/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_sm100.py b/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_sm100.py new file mode 100644 index 000000000..272655e54 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_sm100.py @@ -0,0 +1,2950 @@ +# 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, 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, PipelineConsumer + +import sglang.jit_kernel.flash_attention.cute.utils as utils +import sglang.jit_kernel.flash_attention.cute.copy_utils as copy_utils +import sglang.jit_kernel.flash_attention.cute.pipeline as pipeline +from .blackwell_helpers import gemm_w_idx, gemm_ptx_w_idx # noqa +from .mask import AttentionMask +from .seqlen_info import SeqlenInfoQK +from .block_info import BlockInfo +from .tile_scheduler import ( + TileSchedulerArguments, + SingleTileScheduler, + SingleTileLPTBwdScheduler, # noqa + SingleTileVarlenScheduler, + ParamsBase, +) + +import sglang.jit_kernel.flash_attention.cute.barrier as barrier +from .named_barrier import NamedBarrierBwdSm100 +from .softmax import apply_score_mod_inner, apply_score_mod_bwd_inner +from .block_sparsity import BlockSparseTensors +from .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, + cluster_size: int = 1, + 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, + 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 + assert head_dim == head_dim_v, "head_dim and head_dim_v must be the same for now" + self.tile_hdimv = int(math.ceil(head_dim_v / hdim_multiple_of) * hdim_multiple_of) + assert self.tile_hdim == self.tile_hdimv, ( + "tile_hdim and tile_hdimv must be the same for now" + ) + 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 + + # CTA tiler + self.cta_tiler = (tile_n, tile_m, self.tile_hdim) + # S = K @ Q.T + self.mma_tiler_kq = (tile_n, tile_m, self.tile_hdim) + # dP = V @ dO.T + self.mma_tiler_vdo = (tile_n, tile_m, self.tile_hdimv) + # dV = P.T @ dO + self.mma_tiler_pdo = (tile_n, self.tile_hdimv, tile_m) + # dK = dS.T @ Q (N, M) (M, D) + self.mma_tiler_dsq = (tile_n, self.tile_hdimv, tile_m) + # dQ = dS @ K + self.mma_tiler_dsk = (tile_m, self.tile_hdimv, tile_n) + + 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 + + # 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.subtile_factor = 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 + self.use_smem_dS_for_mma_dK = self.deterministic and self.is_causal + + 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.epi_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.epi_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 + SM100_TMEM_CAPACITY_COLUMNS = 512 + self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + + # 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 + + 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_dP_offset # overlap with dP + 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 = 152 + self.num_regs_compute = 136 + else: + self.num_regs_reduce = 136 + self.num_regs_compute = 144 + self.num_regs_other = 96 - 8 + self.num_regs_empty = 24 + assert self.num_regs_reduce + self.num_regs_compute * 2 + self.num_regs_other <= 512 + + self.buffer_align_bytes = 1024 + + def _setup_attributes(self): + self.Q_stage = 2 + self.dO_stage = 1 + # LSE_stage = Q_stage and dPsum_stage = dO_stage + # self.sdKVaccum_stage = 2 + # number of tma reduce adds per dQacc mma + self.dQ_reduce_ncol = 32 + self.sdQaccum_stage = 64 // self.dQ_reduce_ncol + assert self.tile_hdim % self.dQ_reduce_ncol == 0 + self.dQaccum_reduce_stage = self.tile_hdim // self.dQ_reduce_ncol + self.cluster_reduce_dQ = False and cute.size(self.cluster_shape_mn) > 1 + # number of tma reduce adds for dKacc and dVacc epilogue + self.dK_reduce_ncol = 32 + + def _get_tiled_mma(self): + cta_group = tcgen05.CtaGroup.ONE + # S = 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, + cta_group, + self.mma_tiler_kq[:2], + ) + # dP = 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, + cta_group, + self.mma_tiler_vdo[:2], + ) + # dV += P @ 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, + 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, + 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, + 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 = 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 = 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 @ 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.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.sdKV_epi_tile = ( + self.tile_n, + min(128 // (self.dk_dtype.width // 8), self.tile_hdim // 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.sdKV_epi_tile[1]) + self.sdKV_flat_epi_tile = self.tile_n * (self.tile_hdim // 2) // self.num_epi_stages + # TODO: dK and dV could have different shapes + if const_expr(not self.dKV_postprocess): + self.sdKV_layout = sm100_utils_basic.make_smem_layout_epi( + self.dk_dtype, + LayoutEnum.ROW_MAJOR, + self.sdKV_epi_tile, + 2, # num compute wgs + ) + else: + self.sdKV_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, + stream: cuda.CUstream, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + softcap: Float32 | float | None = 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_tensors: Optional[list] = None, + # Block-sparse tensors (Q direction - for iterating m_blocks per n_block): + blocksparse_tensors: Optional[BlockSparseTensors] = 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.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" + + # Assume all strides are divisible by 128 bits except the last stride + new_stride = lambda t: ( + *(cute.assume(s, divby=128 // t.element_type.width) for s in t.stride[:-1]), + t.stride[-1], + ) + ( + mdQaccum, + mdK, + mdV, + ) = [ + cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) + if t is not None + else None + 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 = [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 = [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 = [ + 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 = [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 = utils.select(mdO, mode=dO_transpose) + + # (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 = 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 = [ + 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() + + cta_group = tcgen05.CtaGroup.ONE + + 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.sdKV_layout, mode=[0, 1]), + self.sdKV_epi_tile, + 1, # no mcast + ) + tma_atom_dV, mdV_tma_tensor = cpasync.make_tiled_tma_atom( + tma_copy_op_dKV, + mdV, + cute.select(self.sdKV_layout, mode=[0, 1]), + self.sdKV_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(cta_group) + tma_load_op_multicast = cpasync.CopyBulkTensorTileG2SMulticastOp(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( + # tma_load_op if const_expr(self.cluster_shape_mnk[0] == 1) else tma_load_op_multicast, + 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, + ) + 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( + # tma_load_op if const_expr(self.cluster_shape_mnk[0] == 1) else tma_load_op_multicast, + 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, + ) + + 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, 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 + + # TileScheduler = SingleTileScheduler + if const_expr(self.is_varlen_k): + TileScheduler = SingleTileVarlenScheduler + elif const_expr(self.deterministic): + TileScheduler = SingleTileLPTBwdScheduler + else: + TileScheduler = SingleTileScheduler + # reads n_blocks right-to-left + 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.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) + # cute.printf("grid_dim = {}", grid_dim) + + # 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.sdKV_layout), + ) + sdO_alloc_bytes = max( + cute.size_in_bytes(self.dv_dtype, self.sdKV_layout), + cute.size_in_bytes(self.do_dtype, self.sdO_layout), + ) + # Sanity check that layouts fit in allocation + sdV_bytes = cute.size_in_bytes(self.dv_dtype, self.sdKV_layout) + sdK_bytes = cute.size_in_bytes(self.dk_dtype, self.sdKV_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" + + @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 * 1] + dP_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * 1] + dS_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * 1] + dKV_mbar_ptr: cute.struct.MemRange[cutlass.Int64, 2 * 2] + 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: cute.struct.MemRange[cutlass.Int64, 1] + + # Smem tensors + + # sQ is reused for sdK which in the non-MHA case needs float32 + 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 is reused for sdV which in the non-MHA case needs float32 + 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_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_block_sparsity or aux_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_K, + tma_tensor_V, + mLSE, + mdPsum, + tma_tensor_dO, + mdV, + mdK, + mdQaccum, + mdV_tma_tensor, + mdK_tma_tensor, + mdQ_semaphore, + mdK_semaphore, + mdV_semaphore, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + tma_atom_Q, + tma_atom_K, + tma_atom_V, + tma_atom_dO, + tma_atom_dV, + tma_atom_dK, + self.sQ_layout, + self.sQt_layout, + self.sK_layout, + self.sV_layout, + self.sLSE_layout, + self.sdPsum_layout, + self.sdO_layout, + self.sdOt_layout, + self.sdSt_layout, + self.sdS_layout, + self.sKt_layout, + self.sdQaccum_layout, + self.sdKV_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_tensors, + 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, + ) + + @cute.kernel + def kernel( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mLSE: cute.Tensor, + mdPsum: cute.Tensor, + mdO: 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_K: cute.CopyAtom, + tma_atom_V: cute.CopyAtom, + tma_atom_dO: 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, + 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, + sKt_layout: cute.ComposedLayout, + sdQaccum_layout: cute.Layout, + sdKV_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_tensors: Optional[list] = None, + fastdiv_mods=(None, None), + blocksparse_tensors: Optional[BlockSparseTensors] = None, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + # Prefetch tma descriptor + if warp_idx == self.load_warp_id: + with cute.arch.elect_one(): + cpasync.prefetch_descriptor(tma_atom_Q) + cpasync.prefetch_descriptor(tma_atom_K) + cpasync.prefetch_descriptor(tma_atom_V) + 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) + + tmem_dealloc_mbar_ptr = storage.tmem_dealloc_mbar_ptr.data_ptr() + 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 warp_idx == 1: + cute.arch.mbarrier_init( + tmem_dealloc_mbar_ptr, cute.arch.WARP_SIZE * len(self.compute_warp_ids) + ) + 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) + + # UMMA producers and AsyncThread consumers + pipeline_producer_group_MMA_AsyncThread = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + # Only 1 thread per warp will signal + pipeline_consumer_group_MMA_AsyncThread = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, len(self.compute_warp_ids) + ) + 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(), + ) + 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(), + ) + 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(), + ) + pipeline_consumer_group_MMA_AsyncThread_dQ = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, + len(self.reduce_warp_ids), + ) # 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(), + ) + + # 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) + ) # 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(), + ) + + # 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) * self.num_mcast_ctas_b + 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, + # init_wait=False, + ) + 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, + # init_wait=False, + ) + 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, + init_wait=False, + ) + 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, + init_wait=True, + ) + + sQ = storage.sQ.get_tensor(sQ_layout.outer, swizzle=sQ_layout.inner, dtype=self.q_dtype) + 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) + 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) + sdO = storage.sdO.get_tensor( + sdO_layout.outer, swizzle=sdO_layout.inner, dtype=self.do_dtype + ) + 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(not self.dKV_postprocess): + sdV = storage.sdO.get_tensor( + sdKV_layout.outer, swizzle=sdKV_layout.inner, dtype=self.dv_dtype + ) + sdK = storage.sQ.get_tensor( + sdKV_layout.outer, swizzle=sdKV_layout.inner, dtype=self.dk_dtype + ) + else: + sdV = storage.sdO.get_tensor(sdKV_layout, dtype=self.dv_dtype) + sdK = storage.sQ.get_tensor(sdKV_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(0) + 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(0) + 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(0) + 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(0) + 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(0) + 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, + ) + TileSchedulerCls = partial(self.tile_scheduler_cls.create, tile_sched_params) + + AttentionMaskCls = partial( + AttentionMask, + self.tile_m, + self.tile_n, + swap_AB=True, + window_size_left=window_size_left, + window_size_right=window_size_right, + ) + + # EMPTY + # (15) + if warp_idx == self.empty_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_empty) + + # EPI + # (14) + if warp_idx == self.epi_warp_id: + # currently no-op, could use for tma store/reduce + cute.arch.warpgroup_reg_dealloc(self.num_regs_empty) + + # LOAD + # (13) + if warp_idx == self.load_warp_id: + cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + self.load( + thr_mma_S, + thr_mma_dP, + thr_mma_dV, + mQ, + mK, + mV, + mLSE, + mdPsum, + mdO, + sQ, + sK, + sV, + sLSE, + sdPsum, + sdO, + tma_atom_Q, + tma_atom_K, + tma_atom_V, + tma_atom_dO, + pipeline_Q, + 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.warpgroup_reg_dealloc(self.num_regs_other) + + # Alloc tmem buffer + tmem_alloc_cols = Int32(self.tmem_alloc_cols) + cute.arch.alloc_tmem(tmem_alloc_cols, storage.tmem_holding_buf) + cute.arch.sync_warp() + + self.mma( + tiled_mma_S, + tiled_mma_dP, + tiled_mma_dV, + tiled_mma_dK, + tiled_mma_dQ, + sQ, + sQt, + sK, + sV, + sdO, + sdOt, + sdSt, + sdS, + sKt, + tP, + tdS, + tStS, + tdPtdP, + tdVtdV, + tdKtdK, + tdQtdQ, + pipeline_Q.make_consumer(), + pipeline_dO, + pipeline_S_P, + pipeline_dS, + pipeline_dKV, + pipeline_dP, + pipeline_dQ, + block_info, + SeqlenInfoCls, + TileSchedulerCls, + blocksparse_tensors, + ) + cute.arch.relinquish_tmem_alloc_permit() + tmem_ptr = cute.arch.retrieve_tmem_ptr( + Float32, alignment=16, ptr_to_buffer_holding_addr=storage.tmem_holding_buf + ) + + cute.arch.mbarrier_wait(tmem_dealloc_mbar_ptr, 0) + tmem_alloc_cols = Int32(self.tmem_alloc_cols) + cute.arch.dealloc_tmem(tmem_ptr, tmem_alloc_cols, is_two_cta=False) + + # 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.warpgroup_reg_alloc(self.num_regs_compute) # 8 warps + self.compute_loop( + thr_mma_S, + thr_mma_dP, + thr_mma_dV, + thr_mma_dK, + tStS, + sLSE, + sdPsum, + tdVtdV, + tdKtdK, + mdV, + mdK, + sdS, + tdPtdP, + pipeline_LSE, + pipeline_dPsum, + pipeline_S_P, + pipeline_dS, + pipeline_dKV, + pipeline_dP, + 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_tensors, + fastdiv_mods, + blocksparse_tensors, + ) + cute.arch.mbarrier_arrive(tmem_dealloc_mbar_ptr) + + # Reduce + # (0, 1, 2, 3) - dQ + if warp_idx >= self.reduce_warp_ids[0] and warp_idx <= self.reduce_warp_ids[-1]: + cute.arch.warpgroup_reg_alloc(self.num_regs_reduce) + self.dQacc_reduce( + mdQaccum, + sdQaccum, + thr_mma_dQ, + tdQtdQ, + pipeline_dQ, + block_info, + SeqlenInfoCls, + TileSchedulerCls, + mdQ_semaphore, + blocksparse_tensors, + ) + + return + + @cute.jit + def load( + self, + thr_mma_S: cute.core.ThrMma, + thr_mma_dP: cute.core.ThrMma, + thr_mma_dV: cute.core.ThrMma, + mQ: cute.Tensor, + mK: cute.Tensor, + mV: cute.Tensor, + mLSE: cute.Tensor, + mdPsum: cute.Tensor, + mdO: cute.Tensor, + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + sLSE: cute.Tensor, + sdPsum: cute.Tensor, + sdO: cute.Tensor, + tma_atom_Q: cute.CopyAtom, + tma_atom_K: cute.CopyAtom, + tma_atom_V: cute.CopyAtom, + tma_atom_dO: cute.CopyAtom, + pipeline_Q: 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_dO_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 + 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 + ] + + gK = cute.local_tile(mK_cur, cute.select(self.mma_tiler_kq, mode=[0, 2]), (n_block, 0)) + tSgK = thr_mma_S.partition_A(gK) + gV = cute.local_tile(mV_cur, cute.select(self.mma_tiler_vdo, mode=[0, 2]), (n_block, 0)) + tdPgV = thr_mma_dP.partition_A(gV) + 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) + + load_K, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_K, 0, cute.make_layout(1), tSgK, sK, single_stage=True + ) + load_V, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_V, + 0, + cute.make_layout(1), + tdPgV, + sV, + 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) + 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=tdPgdO, + dst_tensor=sdO, + mcast_mask=q_do_mcast_mask, + ) + load_dO = copy_utils.tma_producer_copy_fn(load_dO, pipeline_dO) + 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, + subtile_factor=self.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, + subtile_factor=self.subtile_factor, + m_block_max=m_block_max, + ) + ) + else: + first_m_block = m_block_min + + # First iteration: load K together w Q & LSE, then V together w dO & dPsum + if const_expr(should_load_Q): + 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) + 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"] + ) + 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) + 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, 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() + + # Dense path: iterate from m_block_min+1 to m_block_max + for m_block in cutlass.range(m_block_min + 1, m_block_max, unroll=1): + if const_expr(should_load_Q): + 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) + 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) + load_dO(m_block, 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], + 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(should_load_Q): + pipeline_Q.producer_tail( + producer_state_Q_LSE.clone() + ) # will hang if we don't clone + pipeline_LSE.producer_tail(producer_state_Q_LSE) + 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, + sV: cute.Tensor, + sdO: cute.Tensor, + sdOt: cute.Tensor, + sdSt: cute.Tensor, + sdS: cute.Tensor, + sKt: cute.Tensor, + tP: cute.Tensor, + tdS: cute.Tensor, + tStS: cute.Tensor, + tdPtdP: cute.Tensor, + tdVtdV: cute.Tensor, + tdKtdK: cute.Tensor, + tdQtdQ: cute.Tensor, + pipeline_Q_consumer: PipelineConsumer, + 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, + 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 @ dO.T + tdPrV = tiled_mma_dP.make_fragment_A(sV) + tdPrdOt = tiled_mma_dP.make_fragment_B(sdOt) + # dK = dS.T @ Q + if const_expr(self.use_smem_dS_for_mma_dK): + tdKrdS = tiled_mma_dK.make_fragment_A(sdSt) + else: + tdKrdS = tiled_mma_dK.make_fragment_A(tdS) + 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 + ) + # 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, + ) + # 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, + ) + mma_dsk_fn = partial(gemm_w_idx, tiled_mma_dQ, tdQtdQ, tdQrdS, tdQrK, zero_init=True) + # 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): + 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, + ) + + 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 + consumer_state_dS = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Consumer, 1 + ) + # producer_state_dKV = cutlass.pipeline.make_pipeline_state( + # cutlass.pipeline.PipelineUserType.Producer, 2 + # ) + producer_phase_dKV = Int32(1) + cta_group = pipeline_S_P.cta_group + + 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, + subtile_factor=self.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 process_tile: + accumulate_dK = False + # ----------------------------------------------------------- + ###### Prologue + # ----------------------------------------------------------- + # 1. S = Q0 @ K.T + # 2. dP = V @ dO.T + # 3. dV = P @ dO + # 1) S = Q0 @ K.T + 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) + # Don't release Q yet + pipeline_S_P.sync_object_full.arrive(0, pipeline_S_P.producer_mask, cta_group) + + # 2) dP = V @ dO.T + pipeline_dO.consumer_wait(consumer_state_dO) + pipeline_dP.sync_object_empty.wait(0, producer_phase_acc) + # dQ uses the same tmem as dP + pipeline_dQ.sync_object_empty.wait(0, producer_phase_acc) + mma_dov_fn(B_idx=consumer_state_dO.index) + # Don't release dO yet + pipeline_dP.sync_object_full.arrive(0, pipeline_dP.producer_mask, cta_group) + + producer_phase_acc ^= 1 + # 3) dV = P.T @ dO + # wait for P to be ready, which uses the same tmem as S + 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 @ dO.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 + ) + for _ in cutlass.range(main_loop_iters, unroll=1): + # 1) S = K @ Q_i + handle_Q_next = pipeline_Q_consumer.wait_and_advance() + # Don't need to wait for S, as P must have been ready ealier, i.e., S is ready + 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-3) + # Do dK = dS.T @ Q, then dQ = dS @ K if dS in tmem for first mma + # Otherwise, reverse order + pipeline_dS.consumer_wait(consumer_state_dS) + + if const_expr(self.use_smem_dS_for_mma_dK): + mma_dsk_fn() + pipeline_dQ.sync_object_full.arrive(0, pipeline_dQ.producer_mask, cta_group) + mma_dsq_fn(B_idx=handle_Q.index, zero_init=not accumulate_dK) + accumulate_dK = True + handle_Q.release() + else: + mma_dsq_fn(B_idx=handle_Q.index, zero_init=not accumulate_dK) + accumulate_dK = True + handle_Q.release() + mma_dsk_fn() + pipeline_dQ.sync_object_full.arrive(0, pipeline_dQ.producer_mask, cta_group) + + # dP uses the same tmem as dQ + # However, if dS is ready, then dP must have been ready, + # so we don't need this wait before mma_dsk_fn() + # pipeline_dP.sync_object_empty.wait(0, producer_phase_acc) + + pipeline_dS.consumer_release(consumer_state_dS) + consumer_state_dS.advance() + + # 4) dP = V @ dO.T + pipeline_dO.consumer_wait(consumer_state_dO) + # dQ uses the same tmem as dP + 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 + # 5) dV += P @ dO + # wait for P to be ready, which uses the same tmem as S + 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) + + # ----------------------------------------------------------- + ###### Remaining 2 + # ----------------------------------------------------------- + # 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.producer_commit(producer_state_dKV) + pipeline_dKV.sync_object_full.arrive(1, pipeline_dKV.producer_mask, cta_group) + # producer_state_dKV.advance() + producer_phase_dKV ^= 1 + + # 2) dQ = dS @ K + # dS is done, so dP must have been ready, we don't need to wait + mma_dsk_fn() + pipeline_dQ.sync_object_full.arrive(0, pipeline_dQ.producer_mask, cta_group) + # Wait until dQ is done before releasing Q, since K and Q0 uses the same mbarrier + 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_tensors=None, + 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) + cS = cute.make_identity_tensor((self.tile_n, self.tile_m)) + cS = cute.domain_offset((n_block * self.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_tensors, + 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_tensors=None, + 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_tensors, + 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.core.ThrMma, + thr_mma_dP: cute.core.ThrMma, + thr_mma_dV: cute.core.ThrMma, + thr_mma_dK: cute.core.ThrMma, + tStS: cute.Tensor, + sLSE: cute.Tensor, + sdPsum: cute.Tensor, + tdVtdV: cute.Tensor, + tdKtdK: cute.Tensor, + mdV: cute.Tensor, + mdK: cute.Tensor, + sdS: cute.Tensor, + tdPtdP: cute.Tensor, + pipeline_LSE: PipelineAsync, + pipeline_dPsum: PipelineAsync, + pipeline_S_P: PipelineAsync, + pipeline_dS: PipelineAsync, + pipeline_dKV: PipelineAsync, + pipeline_dP: PipelineAsync, + 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_tensors: Optional[list] = None, + 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 = utils.transpose_view(sLSE_2D) + sdPsum_2D = 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.mma_tiler_kq[0] // 32 * self.v_dtype.width # 64 for tile_n = 128 + # 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)) + + 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_layout = sm100_utils_basic.make_smem_layout_epi( + self.ds_dtype, LayoutEnum.ROW_MAJOR, (self.tile_n, self.tile_m), 1 + ).outer # ((8,16), (64,2), (1, 1)) + sdS_layout = cute.slice_(sdS_layout, (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) + + 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 = cutlass.pipeline.make_pipeline_state( + 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) + # TODO: condition mask_seqlen + mask_fn = partial( + mask.apply_mask_sm100_transposed, + tScS_t2r=tScS_t2r, + t0ScS_t2r=t0ScS_t2r, + n_block=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_tensors=aux_tensors, + fastdiv_mods=fastdiv_mods, + ) + + # prefetch_LSE = not self.is_causal + prefetch_LSE = False + + # some tiles might be empty due to block sparsity + if const_expr(self.use_block_sparsity): + ( + 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, + subtile_factor=self.subtile_factor, + m_block_max=m_block_max, + ) + process_tile = loop_count > Int32(0) + else: + 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 + + # 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): + 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, + subtile_factor=self.subtile_factor, + m_block_max=m_block_max, + ) + m_block_oob = m_block >= m_block_max + else: + m_block = m_block_min + iter_idx + m_block_oob = False + is_full_block = False + # Prefetch 1 stage of LSE + pipeline_LSE.consumer_wait(consumer_state_LSE) + tSrLSE_s2r = cute.make_fragment(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_fragment(tScS_t2r.shape, Float32) + cute.copy(thr_copy_t2r, tStS_t2r, tSrS_t2r) + 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_tensors, + 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_fragment(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] = utils.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() + self.compute_sync_barrier.arrive_and_wait() + + 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) + pipeline_LSE.consumer_release(consumer_state_LSE) + # consumer_state_S_P_dP.advance() + 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) + 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_fragment(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] = utils.sub_packed_f32x2( + (tdPrdP_cur[2 * v], tdPrdP_cur[2 * v + 1]), dPsum_pair + ) + tdPrdP_cur[2 * v], tdPrdP_cur[2 * v + 1] = utils.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] + cS_bwd = cute.make_identity_tensor((self.tile_n, self.tile_m)) + cS_bwd = cute.domain_offset( + (n_block * self.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_tensors, + 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) + cute.autovec_copy(tdPrdS_cvt, tRS_sdS[None, stage]) + if const_expr(not self.use_smem_dS_for_mma_dK): + tdPrdS_r2t_f32 = cute.recast_tensor(tdPrdS_cvt, Float32) + cute.copy(thr_copy_r2t, tdPrdS_r2t_f32, tdPtdS_r2t[None, stage, 0, 0]) + + if const_expr(not self.use_smem_dS_for_mma_dK): + cute.arch.fence_view_async_tmem_store() + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, space=cute.arch.SharedSpace.shared_cta + ) + self.compute_sync_barrier.arrive_and_wait() + + # with cute.arch.elect_one(): + # The mma warp no longer waits for dP (it waits for dS), so we don't have to arrive + # pipeline_dP.sync_object_empty.arrive(0, pipeline_dP.consumer_mask) + pipeline_dPsum.consumer_release(consumer_state_dPsum) + consumer_state_dPsum.advance() + 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, + ) + #### 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, + ) + # 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: + # like other epis, currently assumes hdim == hdimv + gmem_tiled_copy_zero_dKV = copy_utils.tiled_copy_2d( + self.dk_dtype, + self.tile_hdim, + 128, # num_threads + ) + gmem_thr_copy_zero_dKV = gmem_tiled_copy_zero_dKV.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, (self.tile_n, self.tile_hdim), (n_block, 0)) + gdV = cute.local_tile(mdV_cur, (self.tile_n, self.tile_hdimv), (n_block, 0)) + tdKgdK = gmem_thr_copy_zero_dKV.partition_D(gdK) + tdVgdV = gmem_thr_copy_zero_dKV.partition_D(gdV) + assert tdKgdK.shape[2] == 1 + assert tdVgdV.shape[2] == 1 + cdKV = cute.make_identity_tensor((self.tile_n, self.tile_hdim)) + tdKVcdKV = gmem_thr_copy_zero_dKV.partition_D(cdKV) + 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 = tdKVcdKV[0, i, 0][0] + if row_idx < seqlen.seqlen_k - self.tile_n * n_block: + cute.copy(gmem_tiled_copy_zero_dKV, zero, tdKgdK[None, i, 0]) + else: + for i in cutlass.range_constexpr(tdVgdV.shape[1]): + row_idx = tdKVcdKV[0, i, 0][0] + if row_idx < seqlen.seqlen_k - self.tile_n * n_block: + cute.copy(gmem_tiled_copy_zero_dKV, zero, tdVgdV[None, i, 0]) + + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + + @cute.jit + def dQacc_reduce( + self, + mdQaccum: cute.Tensor, + sdQaccum: cute.Tensor, + thr_mma_dQ: cute.core.ThrMma, + tdQtdQ: cute.Tensor, + pipeline_dQ: PipelineAsync, + 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 + # TMEM -> RMEM + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(self.dQ_reduce_ncol)), 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 + assert cute.size(tdQrdQ_t2r_shape, mode=[1]) == self.dQaccum_reduce_stage, ( + "dQaccum reduce stage mismatch" + ) + + 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 + 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] + ) + 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): + mdQ_semaphore_cur = mdQ_semaphore[None, None, head_idx, batch_idx] + + delay_semaphore_release = self.is_causal + n_block_global_max = cute.ceil_div(seqlen.seqlen_k, self.tile_n) + + # some tiles might be empty due to block sparsity + if const_expr(self.use_block_sparsity): + ( + 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, + subtile_factor=self.subtile_factor, + m_block_max=m_block_max, + ) + process_tile = loop_count > Int32(0) + else: + 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 + + # 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): + 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, + subtile_factor=self.subtile_factor, + m_block_max=m_block_max, + ) + if m_block_max > 0: + m_block = cutlass.min(m_block, m_block_max - 1) + else: + m_block = m_block_min + iter_idx + pipeline_dQ.consumer_wait(dQ_consumer_state) + # TMEM -> RMEM + tdQrdQ_t2r = cute.make_fragment(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() + + gdQaccum_cur = gdQaccum[None, None, m_block] + + for stage in cutlass.range_constexpr(cute.size(tdQrdQ_t2r, mode=[1])): # 4 + smem_idx = dQ_tma_store_producer_state.index + tdQsdQ_r2s = tdQsdQ[None, None, smem_idx] + tdQrdQ_r2s = cute.make_tensor( + tdQrdQ_t2r[None, stage, None, None].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_proxy( + cute.arch.ProxyKind.async_shared, space=cute.arch.SharedSpace.shared_cta + ) + # semaphore acquire + if const_expr(self.deterministic and stage == 0): + if const_expr(self.spt): + if const_expr( + self.is_causal or block_info.window_size_right is not None + ): + n_idx_right = ( + (m_block + 1) * self.tile_m + seqlen.seqlen_k - seqlen.seqlen_q + ) + if const_expr(block_info.window_size_right is not None): + n_idx_right += block_info.window_size_right + n_block_max_for_m_block = min( + n_block_global_max, + cute.ceil_div(n_idx_right, self.tile_n), + ) + else: + n_block_max_for_m_block = n_block_global_max + lock_value = n_block_max_for_m_block - 1 - n_block + else: + lock_value = n_block + barrier.wait_eq( + mdQ_semaphore_cur[(m_block, None)].iterator, tidx, 0, lock_value + ) + self.reduce_sync_barrier.arrive_and_wait() + # Copy from shared memory to global memory + if is_tma_warp: + with cute.arch.elect_one(): + copy_utils.cpasync_reduce_bulk_add_f32( + sdQaccum[None, smem_idx].iterator, + gdQaccum_cur[None, stage].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) + self.reduce_sync_barrier.arrive_and_wait() + dQ_tma_store_producer_state.advance() + # Directly add to gmem, much slower + # tdQgdQ = thr_copy_dQaccum_r2s.partition_D(gdQaccum[None, stage, m_block]) + # assert cute.size(tdQrdQ_r2s) == cute.size(tdQgdQ) + # for i in cutlass.range(cute.size(tdQrdQ_r2s) // 4, unroll_full=True): + # copy_utils.atomic_add_fp32x4( + # tdQrdQ_r2s[4 * i], + # tdQrdQ_r2s[4 * i + 1], + # tdQrdQ_r2s[4 * i + 2], + # tdQrdQ_r2s[4 * i + 3], + # utils.elem_pointer(tdQgdQ, 4 * i), + # ) + # semaphore release for prior m_block + 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, 0, 1 + ) + + # semaphore release + # NOTE: arrive_inc calls red_release which issues membar + if const_expr(self.deterministic and not delay_semaphore_release): + if is_tma_warp: + cute.arch.cp_async_bulk_wait_group(0, read=read_flag) + self.reduce_sync_barrier.arrive_and_wait() + barrier.arrive_inc(mdQ_semaphore_cur[m_block, None].iterator, tidx, 0, 1) + + if const_expr(not self.is_local) or m_block_min < m_block_max: + 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, 0, 1 + ) + + 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, tidx, 0, 1) + + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + + @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.core.ThrMma, + thr_mma_dK: cute.core.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_fragment(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_fragment(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.tile_n, self.tile_hdimv), (None, 0)) + gdV_tile = gdV[None, None, n_block] + + 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_fragment(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_fragment(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.tile_n, self.tile_hdimv), (None, 0)) + gdK_tile = gdK[None, None, n_block] + + 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) + consumer_state_dKV.advance() + 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.core.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], + ) -> cutlass.pipeline.PipelineState: + # assumes mma_tiler_pdo = mma_tiler_dsq = (tile_n, head_dim) + # head_dim = head_dim_v, dk_dtype = dv_dtype + 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 + + 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, self.tile_hdim), (n_block, 0) + ) # (tile_n, hdim) + gdKV = self.split_wg(gdKV_p, wg_idx, num_wg) # (tile_n, hdim / 2) + gdKV_epi = cute.local_tile( + gdKV, self.sdKV_epi_tile, (0, None) + ) # (tile_n, 64, epi_stage = (hdim / 2) / 64) + else: + 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 * self.tile_hdim,), mdKV[None, head_idx_kv] + ) + gdKV_p = cute.local_tile( + mdKV_cur, (self.tile_n * self.tile_hdim,), (n_block,) + ) # (tile_n * hdim) + gdKV = cute.logical_divide(gdKV_p, (self.tile_n * self.tile_hdim // num_wg,))[ + ((None, wg_idx),) + ] # (tile_n * hdim / 2) + gdKV_epi = cute.flat_divide( + gdKV, (self.sdKV_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): + 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]) + assert num_epi_stages == self.num_epi_stages, "Epi stage calculation is wrong" + else: + num_epi_stages = self.num_epi_stages + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), 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((self.tile_n, self.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_fragment(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] = utils.mul_packed_f32x2( + (tdKVrdKV_t2r[2 * i], tdKVrdKV_t2r[2 * i + 1]), (scale, scale) + ) + tdKVrdKV = cute.make_fragment(tdKVrdKV_t2r.shape, self.dv_dtype) # (32 columns) + tdKVrdKV.store(tdKVrdKV_t2r.load().to(self.dv_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_proxy( + cute.arch.ProxyKind.async_shared, space=cute.arch.SharedSpace.shared_cta + ) + 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_proxy( + cute.arch.ProxyKind.async_shared, space=cute.arch.SharedSpace.shared_cta + ) + 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/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_sm90.py b/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_sm90.py new file mode 100644 index 000000000..01b43531b --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_sm90.py @@ -0,0 +1,1706 @@ +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.arch import ProxyKind, SharedSpace +from cutlass.cute import FastDivmodDivisor +from cutlass import Float32, Int32, Boolean, const_expr +from cutlass.utils import LayoutEnum + +import sglang.jit_kernel.flash_attention.cute.hopper_helpers as sm90_utils +import sglang.jit_kernel.flash_attention.cute.utils as utils +import sglang.jit_kernel.flash_attention.cute.copy_utils as copy_utils +from .hopper_helpers import gemm_zero_init, gemm_w_idx +from .mask import AttentionMask +from .seqlen_info import SeqlenInfoQK +from .block_info import BlockInfo +import sglang.jit_kernel.flash_attention.cute.pipeline as pipeline +from .tile_scheduler import TileSchedulerArguments, SingleTileScheduler, ParamsBase +from .named_barrier import NamedBarrierFwd, NamedBarrierBwd +from .softmax import apply_score_mod_inner, apply_score_mod_bwd_inner +from .block_sparsity import BlockSparseTensors +from .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, +) + + +def mma_partition_fragment_AB( + thr_mma: cute.core.ThrMma, sA: Optional[cute.Tensor], sB: Optional[cute.Tensor], swap_AB: bool +): + if const_expr(not swap_AB): + return ( + thr_mma.make_fragment_A(thr_mma.partition_A(sA)) if sA is not None else None, + thr_mma.make_fragment_B(thr_mma.partition_B(sB)) if sB is not None else None, + ) + else: + return ( + thr_mma.make_fragment_B(thr_mma.partition_B(sA)) if sA is not None else None, + thr_mma.make_fragment_A(thr_mma.partition_A(sB)) if sB is not None else None, + ) + + +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, + 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, + subtile_factor: cutlass.Constexpr[int] = 1, + ): + 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 = False + 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_mma_warp_groups = (self.num_threads // 128) - 1 + self.mma_dkv_is_rs = ( + AtomLayoutMSdP == 1 + and AtomLayoutNdKV == self.num_mma_warp_groups + and SdP_swapAB + and not dKV_swapAB + ) + self.V_in_regs = V_in_regs + if qhead_per_kvhead > 1: + assert self.same_hdim_kv, "GQA backward requires head_dim == head_dim_v" + assert self.num_mma_warp_groups == 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. + # TODO: impl these for hdim 64 + self.shuffle_LSE = self.SdP_swapAB and self.tile_hdim <= 64 + self.shuffle_dPsum = self.SdP_swapAB and self.tile_hdim <= 64 + + 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.subtile_factor = 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 + + @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): + self.sQ_layout, self.sK_layout, self.sV_layout, self.sdO_layout, self.sPdS_layout = [ + sm90_utils.make_smem_layout(self.dtype, LayoutEnum.ROW_MAJOR, shape, stage) + for shape, stage in [ + ((self.tile_m, self.tile_hdim), self.Q_stage), + ((self.tile_n, self.tile_hdim), None), + ((self.tile_n, self.tile_hdimv), None), + ((self.tile_m, self.tile_hdimv), self.dO_stage), + ((self.tile_m, self.tile_n), self.PdS_stage), + ] + ] + self.sdQaccum_layout = cute.make_layout( + (self.tile_m * self.tile_hdim // self.num_mma_warp_groups, self.num_mma_warp_groups) + ) + # 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_mma_warp_groups)), + cute.make_layout(128 // Float32.width), # val_layout + ) + # dKVaccum for GQA epilogue - reuses sV+sK memory recast as f32 + self.sdKVaccum_layout = cute.make_layout( + (self.tile_n * self.tile_hdim // self.num_mma_warp_groups, self.num_mma_warp_groups) + ) + # dKVaccum R->S (same pattern as dQaccum but sized for tile_n) + self.r2s_tiled_copy_dKVaccum = 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_mma_warp_groups)), + cute.make_layout(128 // Float32.width), + ) + + def _get_tiled_mma(self): + # S = Q @ K.T, dP = dO @ V.T + atom_layout_SdP = (self.AtomLayoutMSdP, self.num_mma_warp_groups // self.AtomLayoutMSdP) + 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=(atom_layout_SdP if not self.SdP_swapAB else atom_layout_SdP[::-1]) + + (1,), + tiler_mn=tiler_mn_SdP if not self.SdP_swapAB else tiler_mn_SdP[::-1], + ) + # dV = P.T @ dO, dK = dS.T @ Q + atom_layout_dKV = (self.AtomLayoutNdKV, self.num_mma_warp_groups // self.AtomLayoutNdKV) + 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=(atom_layout_dKV if not self.dKV_swapAB else atom_layout_dKV[::-1]) + + (1,), + tiler_mn=tiler_mn_d if not self.dKV_swapAB else tiler_mn_d[::-1], + 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 + atom_layout_dQ = (self.AtomLayoutMdQ, self.num_mma_warp_groups // self.AtomLayoutMdQ) + 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=(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], + ) + return tiled_mma_SdP, tiled_mma_dK, tiled_mma_dV, tiled_mma_dQ + + def _get_shared_storage_cls(self): + sQ_alignment = sK_alignment = sV_alighment = sdQaccum_alignment = sdO_alignment = 1024 + + sQ_struct, sK_struct, sV_struct, sdO_struct, sdQaccum_struct = [ + cute.struct.Align[cute.struct.MemRange[type, cute.cosize(layout)], alignment] + for (layout, type, alignment) in [ + (self.sQ_layout, self.dtype, sQ_alignment), + (self.sK_layout, self.dtype, sK_alignment), + (self.sV_layout, self.dtype, sV_alighment), + (self.sdO_layout, self.dtype, sdO_alignment), + (self.sdQaccum_layout, Float32, sdQaccum_alignment), + ] + ] + + 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, + stream: cuda.CUstream, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + softcap: Float32 | float | None = 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_tensors: Optional[list] = None, + blocksparse_tensors: Optional[BlockSparseTensors] = None, + ): + assert mdQ_semaphore is None and mdK_semaphore is None and mdV_semaphore is None, ( + "determinism not supported yet for Sm90" + ) + + 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) + ) + ) + + # Assume all strides are divisible by 128 bits except the last stride + # Skip cute.assume() for stride=0 (broadcast dims from expand() are Python ints) + new_stride = lambda t: ( + *( + cute.assume(s, divby=128 // t.element_type.width) if s != 0 else s + for s in t.stride[:-1] + ), + t.stride[-1], + ) + mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV = [ + cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) + if t is not None + else None + for t in (mQ, mK, mV, mdO, mLSE, mdPsum, mdQaccum, mdK, mdV) + ] + + layout_transpose = [1, 3, 2, 0] # (b, s, n, h) --> (s, h, n, b) + mQ, mK, mV, mdO = [utils.select(t, layout_transpose) for t in (mQ, mK, mV, mdO)] + if const_expr(self.qhead_per_kvhead == 1): + mdK, mdV = [utils.select(t, layout_transpose) for t in (mdK, mdV)] + else: + accum_transpose = [2, 1, 0] # (b, n, s*h) -> (s*h, n, b) + mdK, mdV = [utils.select(t, accum_transpose) for t in (mdK, mdV)] + LSE_dPsum_dQaccum_transpose = [2, 1, 0] # (b, n, s) -> (s, n, b) + mLSE, mdPsum, mdQaccum = [ + 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() + + 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 + + self.num_mma_regs = 240 + self.num_producer_regs = 24 + # self.num_mma_regs = 232 + # self.num_producer_regs = 40 + + 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_mma_warp_groups + ) + 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): + tma_atom_dK, tma_tensor_dK = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + mdK, + 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, + 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 + + TileScheduler = SingleTileScheduler + tile_sched_args = TileSchedulerArguments( + cute.ceil_div(cute.size(mK.shape[0]), self.tile_n), + cute.size(mQ.shape[2]), + cute.size(mQ.shape[3]), + 1, # num_splits + cute.size(mK.shape[0]), + mQ.shape[1], + mV.shape[1], + total_q=cute.size(mQ.shape[0]) * cute.size(mQ.shape[3]), + tile_shape_mn=(self.tile_m, self.tile_n), + mCuSeqlensQ=None, + mSeqUsedQ=None, + qhead_per_kvhead_packgqa=1, + element_size=self.dtype.width // 8, + is_persistent=False, + lpt=False, + ) + + 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_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) + + 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, + self.sQ_layout, + self.sK_layout, + self.sV_layout, + self.sPdS_layout, + self.sdO_layout, + self.sdQaccum_layout, + self.sdKVaccum_layout, + self.r2s_tiled_copy_dQaccum, + self.r2s_tiled_copy_dKVaccum, + tiled_mma_SdP, + tiled_mma_dK, + tiled_mma_dV, + tiled_mma_dQ, + softmax_scale_log2, + softmax_scale, + tile_sched_params, + TileScheduler, + SharedStorage, + aux_tensors, + fastdiv_mods, + blocksparse_tensors, + qhead_per_kvhead_divmod, + ).launch( + grid=grid_dim, + block=[self.num_threads, 1, 1], + smem=SharedStorage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + + @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, + sQ_layout: cute.ComposedLayout, + sK_layout: cute.ComposedLayout, + sV_layout: cute.ComposedLayout, + sPdS_layout: cute.ComposedLayout, + sdO_layout: cute.ComposedLayout, + sdQaccum_layout: cute.Layout, + sdKVaccum_layout: cute.Layout, + r2s_tiled_copy_dQaccum: cute.TiledCopy, + r2s_tiled_copy_dKVaccum: 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_tensors: Optional[list] = None, + fastdiv_mods=(None, None), + blocksparse_tensors: Optional[BlockSparseTensors] = None, + qhead_per_kvhead_divmod: Optional[FastDivmodDivisor] = None, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + # prefetch TMA descriptors + if warp_idx == 0: + cpasync.prefetch_descriptor(tma_atom_Q) + cpasync.prefetch_descriptor(tma_atom_K) + cpasync.prefetch_descriptor(tma_atom_V) + cpasync.prefetch_descriptor(tma_atom_dO) + + 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 + None, + None, + qhead_per_kvhead_packgqa=1, + ) + SeqlenInfoCls = partial( + SeqlenInfoQK.create, + seqlen_q_static=mQ.shape[0], + seqlen_k_static=mK.shape[0], + mCuSeqlensQ=None, + mCuSeqlensK=None, + mSeqUsedQ=None, + mSeqUsedK=None, + ) + AttentionMaskCls = partial( + AttentionMask, + self.tile_m, + self.tile_n, + window_size_left=None, + window_size_right=None, + swap_AB=self.SdP_swapAB, + ) + TileSchedulerCls = partial(TileScheduler.create, tile_sched_params) + + if warp_idx < 4: + cute.arch.warpgroup_reg_dealloc(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: + for warp_group_idx in cutlass.range(self.num_mma_warp_groups): + 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, + ) + self.dQaccum_store( + mdQaccum, + sdQaccum, + block_info, + TileSchedulerCls, + SeqlenInfoCls, + blocksparse_tensors, + ) + else: + cute.arch.warpgroup_reg_alloc(self.num_mma_regs) + tidx, _, _ = cute.arch.thread_idx() + tidx = tidx - 128 + self.mma( + tiled_mma_SdP, + tiled_mma_dK, + tiled_mma_dV, + tiled_mma_dQ, + mdK, + mdV, + 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, + r2s_tiled_copy_dKVaccum, + sdKVaccum_layout, + softmax_scale_log2, + softmax_scale, + block_info, + SeqlenInfoCls, + AttentionMaskCls, + TileSchedulerCls, + aux_tensors, + fastdiv_mods, + blocksparse_tensors, + qhead_per_kvhead_divmod, + ) + + @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 = mK[None, None, head_idx_kv, batch_idx] + gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (n_block, 0)) + mV_cur = mV[None, None, head_idx_kv, batch_idx] + gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (n_block, 0)) + + mQ_cur = mQ[None, None, head_idx, batch_idx] + gQ = cute.local_tile(mQ_cur, (self.tile_m, self.tile_hdim), (None, 0)) + mdO_cur = mdO[None, None, head_idx, batch_idx] + gdO = cute.local_tile(mdO_cur, (self.tile_m, self.tile_hdimv), (None, 0)) + mLSE_cur = mLSE[None, head_idx, batch_idx] + gLSE = cute.local_tile(mLSE_cur, (self.tile_m,), (None,)) + mdPsum_cur = mdPsum[None, head_idx, batch_idx] + 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) 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, + subtile_factor=self.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) + with cute.arch.elect_one(): + 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) + with cute.arch.elect_one(): + 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) + with cute.arch.elect_one(): + 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) + with cute.arch.elect_one(): + 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), + subtile_factor=self.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.core.ThrMma, + batch_idx, + head_idx, + m_block, + n_block, + softmax_scale, + seqlen_info: SeqlenInfoQK, + aux_tensors=None, + 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_tensors, + 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.core.ThrMma, + batch_idx, + head_idx, + m_block, + n_block, + softmax_scale, + seqlen_info: SeqlenInfoQK, + aux_tensors=None, + 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_tensors, + 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, + 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, + r2s_tiled_copy_dKVaccum: cute.TiledCopy, + sdKVaccum_layout: cute.Layout, + softmax_scale_log2: Float32, + softmax_scale: Float32, + block_info: BlockInfo, + SeqlenInfoCls: Callable, + AttentionMaskCls: Callable, + TileSchedulerCls: Callable, + aux_tensors: Optional[list] = None, + fastdiv_mods=(None, None), + blocksparse_tensors: Optional[BlockSparseTensors] = None, + qhead_per_kvhead_divmod: Optional[FastDivmodDivisor] = None, + ): + warp_group_idx = cute.arch.make_warp_uniform(tidx // self.num_threads_per_warp_group) + warp_group_thread_layout = cute.make_layout( + self.num_mma_warp_groups, 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 = tiled_mma_dQ.get_slice(warp_group_thread_layout(warp_group_idx)) + # S = Q @ K.T + tSrQ, tSrK = mma_partition_fragment_AB(wg_mma_SdP, sQ, sK, self.SdP_swapAB) + # dP = dO @ V.T + tdPrdO, tdPrV = mma_partition_fragment_AB(wg_mma_SdP, sdO, sV, self.SdP_swapAB) + # dV += P.T @ dO + sPt = utils.transpose_view(sP) if sP is not None else None + sdOt = utils.transpose_view(sdO) + tdVrPt, tdVrdOt = mma_partition_fragment_AB(wg_mma_dV, sPt, sdOt, self.dKV_swapAB) + # dK += dS.T @ Q + sdSt = utils.transpose_view(sdS) + sQt = utils.transpose_view(sQ) + tdKrdSt, tdKrQt = mma_partition_fragment_AB(wg_mma_dK, sdSt, sQt, self.dKV_swapAB) + # dQ = dS @ K + sKt = utils.transpose_view(sK) + tdQrdS, tdQrKt = mma_partition_fragment_AB(wg_mma_dQ, sdS, sKt, self.dQ_swapAB) + + # Smem copy atom tiling + smem_copy_atom_PdS = utils.get_smem_store_atom( + self.arch, self.dtype, transpose=self.SdP_swapAB + ) + smem_thr_copy_PdS = cute.make_tiled_copy_C(smem_copy_atom_PdS, tiled_mma_SdP).get_slice( + tidx + ) + tPsP = None + if const_expr(sP is not None): + tPsP = smem_thr_copy_PdS.partition_D(sP if const_expr(not self.SdP_swapAB) else sPt) + tdSsdS = smem_thr_copy_PdS.partition_D(sdS if const_expr(not self.SdP_swapAB) else sdSt) + + sLSE_mma = 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_mma = 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): + sLSE_mma = utils.transpose_view(sLSE_mma) + sdPsum_mma = utils.transpose_view(sdPsum_mma) + LSEslice = (None, 0, None) if const_expr(not self.SdP_swapAB) else (0, None, None) + tLSEsLSE = utils.make_acc_tensor_mn_view(thr_mma_SdP.partition_C(sLSE_mma))[LSEslice] + tLSEsdPsum = utils.make_acc_tensor_mn_view(thr_mma_SdP.partition_C(sdPsum_mma))[LSEslice] + + smem_thr_copy_dQaccum = r2s_tiled_copy_dQaccum.get_slice(tidx) + tdQsdQaccum = smem_thr_copy_dQaccum.partition_D(sdQaccum) + + dV_shape = (self.tile_n, self.tile_hdimv) + acc_dV = cute.make_fragment( + tiled_mma_dV.partition_shape_C(dV_shape if not self.dKV_swapAB else dV_shape[::-1]), + Float32, + ) + dK_shape = (self.tile_n, self.tile_hdim) + acc_dK = cute.make_fragment( + tiled_mma_dK.partition_shape_C(dK_shape if not self.dKV_swapAB else dK_shape[::-1]), + Float32, + ) + + mma_qk_fn = partial( + gemm_zero_init, + tiled_mma_SdP, + (self.tile_m, self.tile_n), + tSrQ, + tSrK, + swap_AB=self.SdP_swapAB, + ) + mma_dov_fn = partial( + gemm_zero_init, + tiled_mma_SdP, + (self.tile_m, self.tile_n), + tdPrdO, + tdPrV, + swap_AB=self.SdP_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 + ) + mma_dsq_fn = partial( + gemm_w_idx, tiled_mma_dK, acc_dK, tdKrdSt, tdKrQt, swap_AB=self.dKV_swapAB + ) + else: + assert not self.dKV_swapAB + mma_pdo_fn = partial(gemm_w_idx, tiled_mma_dV, acc_dV, tCrB=tdVrdOt) + mma_dsq_fn = partial(gemm_w_idx, tiled_mma_dK, acc_dK, tCrB=tdKrQt) + mma_dsk_fn = partial( + gemm_zero_init, + tiled_mma_dQ, + (self.tile_m, self.tile_hdim), + tdQrdS, + tdQrKt, + swap_AB=self.dQ_swapAB, + ) + + 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, + pipeline_Q=pipeline_Q, + pipeline_dO=pipeline_dO, + tLSEsLSE=tLSEsLSE, + tLSEsdPsum=tLSEsdPsum, + tPsP=tPsP, + tdSsdS=tdSsdS, + tdQsdQaccum=tdQsdQaccum, + smem_thr_copy_PdS=smem_thr_copy_PdS, + smem_thr_copy_dQaccum=smem_thr_copy_dQaccum, + softmax_scale_log2=softmax_scale_log2, + # acc_dV=acc_dV, + # acc_dK=acc_dK, + ) + + 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) + 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) 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, + subtile_factor=self.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_tensors=aux_tensors, + 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, + dKV_accumulate=dKV_accumulate, + thr_mma_SdP=thr_mma_SdP, + batch_idx=batch_idx, + head_idx=head_idx, + n_block=n_block, + softmax_scale=softmax_scale, + seqlen=seqlen, + aux_tensors=aux_tensors, + fastdiv_mods=fastdiv_mods, + ) + 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, + softmax_scale=softmax_scale, + seqlen=seqlen, + subtile_factor=self.subtile_factor, + m_block_max=m_block_max, + aux_tensors=aux_tensors, + 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, + r2s_tiled_copy_dKVaccum, + sdKVaccum_layout, + tidx, + n_block, + head_idx, + batch_idx, + qhead_per_kvhead_divmod, + ) + else: + # Block sparsity: KV tile with zero Q blocks produces no dK/dV; write zeros. + if const_expr(self.use_block_sparsity): + 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, + r2s_tiled_copy_dKVaccum, + sdKVaccum_layout, + tidx, + n_block, + head_idx, + batch_idx, + qhead_per_kvhead_divmod, + ) + + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + + @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, + pipeline_Q: cutlass.pipeline.PipelineAsync, + pipeline_dO: cutlass.pipeline.PipelineAsync, + tLSEsLSE: cute.Tensor, + tLSEsdPsum: cute.Tensor, + tPsP: Optional[cute.Tensor], + tdSsdS: Optional[cute.Tensor], + tdQsdQaccum: cute.Tensor, + smem_thr_copy_PdS: cute.TiledCopy, + smem_thr_copy_dQaccum: cute.TiledCopy, + softmax_scale_log2: Float32, + mask_fn: Optional[Callable] = None, + dKV_accumulate: Boolean = True, + thr_mma_SdP: Optional[cute.core.ThrMma] = None, + batch_idx: Int32 = 0, + head_idx: Int32 = 0, + n_block: Int32 = 0, + softmax_scale: Float32 = 1.0, + seqlen: Optional[SeqlenInfoQK] = None, + aux_tensors: Optional[list] = None, + fastdiv_mods=(None, None), + ): + consumer_state_dO_cur = ( + consumer_state_dO if const_expr(self.Q_stage == self.dO_stage) else consumer_state_Q + ) + 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) + 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): + self.apply_score_mod( + acc_S, + thr_mma_SdP, + batch_idx, + head_idx, + m_block, + n_block, + softmax_scale, + seqlen, + aux_tensors, + fastdiv_mods, + ) + + # (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 = utils.make_acc_tensor_mn_view(acc_S, transpose=self.SdP_swapAB) + for r in cutlass.range_constexpr(cute.size(acc_S_mn, mode=[0])): + 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 - tLSErLSE[r], fastmath=True + ) + tLSErdPsum = copy_utils.load_s2r(tLSEsdPsum[None, smem_idx_dO]) + + # Convert P from f32 -> f16 + tdVrP = utils.cvt_f16(utils.make_acc_tensor_frgA_view(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): + cute.arch.barrier( + barrier_id=int(NamedBarrierBwd.PdS), number_of_threads=self.num_mma_threads + ) + tPrP = smem_thr_copy_PdS.retile(tdVrP) + cute.copy(smem_thr_copy_PdS, tPrP, tPsP[None, None, None, smem_idx_PdS]) + + # (4) [Pointwise 2] dS = P*(dP-dPsum) + warpgroup.wait_group(0) + acc_dP_mn = utils.make_acc_tensor_mn_view(acc_dP, transpose=self.SdP_swapAB) + for r in cutlass.range_constexpr(cute.size(acc_dP_mn, mode=[0])): + 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] - tLSErdPsum[r]) + + if const_expr(self.score_mod_bwd is not None): + self.apply_score_mod_bwd( + acc_dP, + acc_S_pre, + thr_mma_SdP, + batch_idx, + head_idx, + m_block, + n_block, + softmax_scale, + seqlen, + aux_tensors, + fastdiv_mods, + ) + + # Convert dS from f32 -> f16 + tdKrdS = utils.cvt_f16(utils.make_acc_tensor_frgA_view(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_proxy(ProxyKind.async_shared, space=SharedSpace.shared_cta) + cute.arch.barrier( + barrier_id=int(NamedBarrierBwd.PdS), number_of_threads=self.num_mma_threads + ) + + # R2S for dS + tdSrdS = smem_thr_copy_PdS.retile(tdKrdS) + cute.copy(smem_thr_copy_PdS, tdSrdS, tdSsdS[None, None, None, 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_proxy(ProxyKind.async_shared, space=SharedSpace.shared_cta) + cute.arch.barrier( + barrier_id=int(NamedBarrierBwd.PdS), number_of_threads=self.num_mma_threads + ) + # (6) [GEMM 4] dQ = dS @ K + acc_dQ = mma_dsk_fn(A_idx=smem_idx_PdS, wg_wait=1) + # if cute.arch.thread_idx()[0] == 128: cute.print_tensor(acc_dV) + 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) + # if cute.arch.thread_idx()[0] == 128: cute.print_tensor(acc_dQ) + + 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_proxy(ProxyKind.async_shared, space=SharedSpace.shared_cta) + 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) + # if cute.arch.thread_idx()[0] == 128: cute.print_tensor(acc_dK) + pipeline_Q.consumer_release(consumer_state_Q) + # if cute.arch.thread_idx()[0] % 32 == 0: cute.printf("tidx = {}, m_block = {}, after pipeline_Q consumer release", cute.arch.thread_idx()[0], m_block) + + 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, + r2s_tiled_copy_dKVaccum: cute.TiledCopy, + sdKVaccum_layout: cute.Layout, + tidx: Int32, + n_block: Int32, + head_idx: Int32, + batch_idx: Int32, + qhead_per_kvhead_divmod: Optional[FastDivmodDivisor] = None, + ): + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + + if const_expr(self.qhead_per_kvhead == 1): + rdV = cute.make_fragment_like(acc_dV, self.dtype) + rdV.store(acc_dV.load().to(self.dtype)) + rdK = utils.cvt_f16(acc_dK, self.dtype) + + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_mma_threads + ) + + smem_copy_atom_dKV = cute.make_copy_atom( + cute.nvgpu.warp.StMatrix8x8x16bOp(transpose=self.dKV_swapAB, num_matrices=4), + self.dtype, + ) + smem_thr_copy_dK = cute.make_tiled_copy_C(smem_copy_atom_dKV, tiled_mma_dK).get_slice( + tidx + ) + smem_thr_copy_dV = cute.make_tiled_copy_C(smem_copy_atom_dKV, tiled_mma_dV).get_slice( + tidx + ) + mdV_cur = mdV[None, None, head_idx, batch_idx] + mdK_cur = mdK[None, None, head_idx, batch_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 + ) + + taccdVrdV = smem_thr_copy_dV.retile(rdV) + sdV = sV if const_expr(not self.dKV_swapAB) else utils.transpose_view(sV) + taccdVsdV = smem_thr_copy_dV.partition_D(sdV) + cute.copy(smem_copy_atom_dKV, taccdVrdV, taccdVsdV) + cute.arch.fence_proxy(ProxyKind.async_shared, space=SharedSpace.shared_cta) + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_mma_threads + ) + if warp_idx == 4: + store_dV() + taccdKrdK = smem_thr_copy_dK.retile(rdK) + sdK = sK if const_expr(not self.dKV_swapAB) else utils.transpose_view(sK) + taccdKsdK = smem_thr_copy_dK.partition_D(sdK) + cute.copy(smem_copy_atom_dKV, taccdKrdK, taccdKsdK) + cute.arch.fence_proxy(ProxyKind.async_shared, space=SharedSpace.shared_cta) + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_mma_threads + ) + if warp_idx == 4: + store_dK() + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + else: + head_idx_kv = head_idx // qhead_per_kvhead_divmod + + mdKaccum_cur = mdK[None, head_idx_kv, batch_idx] + gdKaccum_ = cute.local_tile(mdKaccum_cur, (self.tile_n * self.tile_hdim,), (n_block,)) + gdKaccum = cute.flat_divide( + gdKaccum_, (self.tile_n * self.tile_hdim // self.num_mma_warp_groups,) + ) + + mdVaccum_cur = mdV[None, head_idx_kv, batch_idx] + gdVaccum_ = cute.local_tile(mdVaccum_cur, (self.tile_n * self.tile_hdimv,), (n_block,)) + gdVaccum = cute.flat_divide( + gdVaccum_, (self.tile_n * self.tile_hdimv // self.num_mma_warp_groups,) + ) + + sdKVaccum = cute.make_tensor( + cute.recast_ptr(sV.iterator, dtype=Float32), + sdKVaccum_layout, + ) + + smem_thr_copy_dKVaccum = r2s_tiled_copy_dKVaccum.get_slice(tidx) + tdKsdKVaccum = smem_thr_copy_dKVaccum.partition_D(sdKVaccum) + + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_mma_threads + ) + + tdKrdKaccum_flat = cute.make_tensor( + acc_dK.iterator, cute.make_layout(tdKsdKVaccum.shape) + ) + cute.autovec_copy(tdKrdKaccum_flat, tdKsdKVaccum) + cute.arch.fence_proxy(ProxyKind.async_shared, space=SharedSpace.shared_cta) + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_mma_threads + ) + + if warp_idx == 4: + with cute.arch.elect_one(): + for wg_idx in cutlass.range_constexpr(self.num_mma_warp_groups): + copy_utils.cpasync_reduce_bulk_add_f32( + sdKVaccum[None, wg_idx].iterator, + gdKaccum[None, wg_idx].iterator, + self.tma_copy_bytes["dKacc"] // self.num_mma_warp_groups, + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_mma_threads + ) + + tdVrdVaccum_flat = cute.make_tensor( + acc_dV.iterator, cute.make_layout(tdKsdKVaccum.shape) + ) + cute.autovec_copy(tdVrdVaccum_flat, tdKsdKVaccum) + cute.arch.fence_proxy(ProxyKind.async_shared, space=SharedSpace.shared_cta) + cute.arch.barrier( + barrier_id=int(NamedBarrierFwd.Epilogue), number_of_threads=self.num_mma_threads + ) + + if warp_idx == 4: + with cute.arch.elect_one(): + for wg_idx in cutlass.range_constexpr(self.num_mma_warp_groups): + copy_utils.cpasync_reduce_bulk_add_f32( + sdKVaccum[None, wg_idx].iterator, + gdVaccum[None, wg_idx].iterator, + self.tma_copy_bytes["dVacc"] // self.num_mma_warp_groups, + ) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + + @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, + ): + 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) + mdQaccum_cur = mdQaccum[None, head_idx, batch_idx] + gdQaccum_ = cute.local_tile(mdQaccum_cur, (self.tile_m * self.tile_hdim,), (None,)) + # (M * K / WG, WG, _) + gdQaccum = cute.flat_divide( + gdQaccum_, (self.tile_m * self.tile_hdim // self.num_mma_warp_groups,) + ) + 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) 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, + subtile_factor=self.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 + + for warp_group_idx in cutlass.range_constexpr(self.num_mma_warp_groups): + 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() + for warp_group_idx in cutlass.range_constexpr(self.num_mma_warp_groups): + cute.arch.cp_async_bulk_wait_group( + self.num_mma_warp_groups - 1 - warp_group_idx, read=True + ) + 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, + ) + else: + dQaccum_store_block_sparse_bwd_sm90( + blocksparse_tensors, + batch_idx, + head_idx, + n_block, + sdQaccum, + gdQaccum, + subtile_factor=self.subtile_factor, + m_block_max=m_block_max, + num_mma_warp_groups=self.num_mma_warp_groups, + num_threads_per_warp_group=self.num_threads_per_warp_group, + tma_copy_bytes_dQ=self.tma_copy_bytes["dQ"], + ) + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() diff --git a/python/sglang/jit_kernel/flash_attention/cute/flash_fwd.py b/python/sglang/jit_kernel/flash_attention/cute/flash_fwd.py new file mode 100644 index 000000000..91b7b7fae --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/flash_fwd.py @@ -0,0 +1,2485 @@ +# 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, List +from functools import partial + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +from cutlass import Constexpr, Float32, Int32, const_expr, Boolean +from cutlass.cute.nvgpu import cpasync, warp, warpgroup +from cutlass.cute.arch import ProxyKind, SharedSpace +import cutlass.utils as utils_basic +from cutlass.utils import LayoutEnum +import cutlass.utils.hopper_helpers as sm90_utils_basic + +from quack import copy_utils as quack_copy_utils + +import sglang.jit_kernel.flash_attention.cute.ampere_helpers as sm80_utils +import sglang.jit_kernel.flash_attention.cute.hopper_helpers as sm90_utils +import sglang.jit_kernel.flash_attention.cute.utils as utils +import sglang.jit_kernel.flash_attention.cute.copy_utils as copy_utils +from .mask import AttentionMask +from .softmax import Softmax, apply_score_mod_inner +from .seqlen_info import SeqlenInfoQK +from .block_info import BlockInfo +from .block_sparsity import BlockSparseTensors +from .block_sparse_utils import ( + produce_block_sparse_loads, + consume_block_sparse_loads, +) +import sglang.jit_kernel.flash_attention.cute.pipeline as pipeline +from .pack_gqa import PackGQA +from .named_barrier import NamedBarrierFwd +from .tile_scheduler import ( + TileSchedulerArguments, + SingleTileScheduler, + SingleTileLPTScheduler, + SingleTileVarlenScheduler, + ParamsBase, +) +from cutlass.cute import FastDivmodDivisor + + +class FlashAttentionForwardBase: + arch: int = 80 + + 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, + ): + """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_in_regs = Q_in_regs + self.score_mod = score_mod + self.mask_mod = mask_mod + self.qk_acc_dtype = Float32 + if const_expr(has_aux_tensors): + self.vec_size: cutlass.Constexpr = 1 + else: + self.vec_size: cutlass.Constexpr = 2 + + @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, + stream: cuda.CUstream, + ): + """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, 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 = quack_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): + if const_expr(not seqlen.has_cu_seqlens_q): + 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) + mLSE_cur = cute.domain_offset((offset,), mLSE[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 = utils.make_acc_tensor_mn_view(thr_mma.partition_C(gLSE_expanded)) + assert cute.size(taccOgLSE, mode=[0]) == cute.size(lse) + taccOcO = utils.make_acc_tensor_mn_view(thr_mma.partition_C(cO)) + t0accOcO = utils.make_acc_tensor_mn_view(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_constexpr(cute.size(taccOgLSE.shape[1])): + 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) + + if const_expr(not seqlen.has_cu_seqlens_q): + mO_cur = mO[None, None, head_idx, batch_idx] + else: + offset = seqlen.offset_q if const_expr(not self.pack_gqa) else (0, seqlen.offset_q) + mO_cur = cute.domain_offset((offset, 0), mO[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_proxy(ProxyKind.async_shared, space=SharedSpace.shared_cta) + 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], + stream: cuda.CUstream, + softmax_scale: Optional[Float32] = None, + window_size_left: Optional[Int32] = None, + window_size_right: Optional[Int32] = None, + learnable_sink: Optional[cute.Tensor] = None, + aux_tensors=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)) + ) + 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 >= 90 + self._setup_attributes() + SharedStorage = self._get_shared_storage_cls() + # Assume all strides are divisible by 128 bits except the last stride + # Skip cute.assume() for stride=0 (broadcast dims from expand() are Python ints) + new_stride = lambda t: ( + *( + cute.assume(s, divby=128 // t.element_type.width) + if s != 0 + else s + for s in t.stride[:-1] + ), + t.stride[-1], + ) + mQ, mK, mV, mO = [ + cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(t))) + for t in (mQ, mK, mV, mO) + ] + mQ, mK, mV, mO = [ + cute.make_tensor(t.iterator, cute.select(t.layout, mode=[1, 3, 2, 0])) + for t in (mQ, mK, mV, mO) + ] + mLSE = cute.make_tensor(mLSE.iterator, cute.select(mLSE.layout, mode=[2, 1, 0])) + # grid_dim: (m_block, num_head, batch_size) + grid_dim = ( + cute.ceil_div(mQ.shape[0], self.tile_m), + cute.size(mQ.shape[2]), + cute.size(mQ.shape[3]), + ) + LOG2_E = math.log2(math.e) + if const_expr(self.score_mod is None): + softmax_scale_log2 = Float32(softmax_scale * LOG2_E) + softmax_scale = None + else: + # NB: If a user passes in a score mod, we want to apply the score-mod in the sm_scaled qk + # But in the original base 10. We hijack softmax_scale_log2 to just be the change of base + # and correctly apply the softmax_scale prior to score_mod in the softmax step + softmax_scale_log2 = Float32(LOG2_E) + softmax_scale = Float32(softmax_scale) + + fastdiv_mods = None + if const_expr(aux_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.kernel( + mQ, + mK, + mV, + mO, + mLSE, + 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, + aux_tensors, + 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], + 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, + aux_tensors=None, + fastdiv_mods=None, + ): + # Thread index, block index + tidx, _, _ = cute.arch.thread_idx() + m_block, num_head, batch_size = cute.arch.block_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(seqlen_q_static=mQ.shape[0], seqlen_k_static=mK.shape[0]) + n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block) + # TODO: return early if n_block_max == 0 + # if self.is_causal: + # if n_block_max <= 0: + # return + n_block = n_block_max - 1 + + # /////////////////////////////////////////////////////////////////////////////// + # 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) + gQ = cute.local_tile(mQ[None, None, num_head, batch_size], blkQ_shape, (m_block, 0)) + num_head_kv = num_head // self.qhead_per_kvhead + gK = cute.local_tile(mK[None, None, num_head_kv, batch_size], blkK_shape, (None, 0)) + gV = cute.local_tile(mV[None, None, num_head_kv, batch_size], 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 = 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_fragment(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_tensors=aux_tensors, + 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.seqlen_q, + seqlen.seqlen_k, + window_size_left, + window_size_right, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + mask_fn = partial( + mask.apply_mask, + m_block=m_block, + thr_mma=thr_mma_qk, + mask_causal=self.is_causal, + mask_local=self.is_local, + 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, + check_inf=True, + mask_fn=partial(mask_fn, 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, + check_inf=True, + mask_fn=partial(mask_fn, mask_seqlen=False), + ) + 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, check_inf=True + ) + 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_tensors=None, + 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_fragment(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, + seqlen, + softmax_scale=softmax.softmax_scale, + aux_tensors=aux_tensors, + 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 = cute.make_tensor(rP.iterator, utils.convert_layout_acc_frgA(rP.layout)) + 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() + + +class FlashAttentionForwardSm90(FlashAttentionForwardBase): + arch = 90 + + def __init__( + self, + *args, + intra_wg_overlap: bool = True, + mma_pv_is_rs: bool = True, + **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 + + 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), # Might need (1, 2, 1) for hdim 512 + 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, + ) + tiled_mma_pv_rs = 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, + ) + return tiled_mma_qk, tiled_mma_pv, tiled_mma_pv_rs + + def _get_shared_storage_cls(self): + # If we use cp.async to load Q, we want sQ to align to 1024 bytes + 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 for Q, 1 for O, self.num_stages*2 for K, self.num_stages*2 for V, + mbar_ptr_QO_struct = cute.struct.MemRange[cutlass.Int64, 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: mbar_ptr_QO_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: mbar_ptr_QO_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, + stream: cuda.CUstream, + 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_tensors: Optional[list] = 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) + ) + ) + + # Assume all strides are divisible by 128 bits except the last stride + # Skip cute.assume() for stride=0 (broadcast dims from expand() are Python ints) + new_stride = lambda t: ( + *( + cute.assume(s, divby=128 // t.element_type.width) + if s != 0 + else s + for s in t.stride[:-1] + ), + t.stride[-1], + ) + + mQ, mK, mV, mO = [ + cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(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 = [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 = [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 = utils.select(mLSE, LSE_layout_transpose) if const_expr(mLSE is not None) else None + + tiled_mma_qk, tiled_mma_pv, tiled_mma_pv_rs = self._get_tiled_mma() + self.num_mma_threads = tiled_mma_qk.size + self.num_threads_per_warp_group = 128 + self.num_mma_warp_groups = self.num_mma_threads // self.num_threads_per_warp_group + self.num_threads = self.num_threads_per_warp_group * (self.num_mma_warp_groups + 1) + self.num_producer_threads = 32 + self.num_Q_load_threads = self.num_mma_threads # If not TMA_Q, MMA threads load Q + self.num_epilogue_threads = self.num_mma_threads + self.num_mma_regs = ( + 256 + if self.num_mma_warp_groups == 1 + else (240 if self.num_mma_warp_groups == 2 else 160) + ) + self.num_producer_regs = ( + 56 if self.num_mma_warp_groups == 1 else (24 if self.num_mma_warp_groups == 2 else 32) + ) + # self.num_mma_regs = 232 + # self.num_producer_regs = 40 + self.use_block_sparsity = cutlass.const_expr(blocksparse_tensors is not None) + + self.use_scheduler_barrier = ( + (self.num_mma_warp_groups >= 2 and self.tile_hdim <= 128) + if const_expr(self.intra_wg_overlap) + else (self.num_mma_warp_groups == 2) + ) + self.use_tma_Q = self.arch >= 90 and not ( + self.pack_gqa and self.tile_m % self.qhead_per_kvhead != 0 + ) + self.use_tma_O = ( + self.arch >= 90 and mCuSeqlensQ is None and mSeqUsedQ is None and not self.pack_gqa + ) + # TODO: rescale_O_before_gemm + 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.dtype, LayoutEnum.ROW_MAJOR, (self.tile_m, self.tile_n) + ) + + SharedStorage = self._get_shared_storage_cls() + + if const_expr(self.pack_gqa): + shape_Q_packed = ( + (self.qhead_per_kvhead, mQ.shape[0]), + mQ.shape[1], + mK.shape[2], + *mQ.shape[3:], + ) + stride_Q_packed = ( + (mQ.stride[2], mQ.stride[0]), + mQ.stride[1], + mQ.stride[2] * self.qhead_per_kvhead, + *mQ.stride[3:], + ) + mQ = cute.make_tensor( + mQ.iterator, cute.make_layout(shape_Q_packed, stride=stride_Q_packed) + ) + shape_O_packed = ( + (self.qhead_per_kvhead, mO.shape[0]), + mK.shape[1], + mK.shape[2], + *mO.shape[3:], + ) + stride_O_packed = ( + (mO.stride[2], mO.stride[0]), + mO.stride[1], + mO.stride[2] * self.qhead_per_kvhead, + *mO.stride[3:], + ) + mO = cute.make_tensor( + mO.iterator, cute.make_layout(shape_O_packed, stride=stride_O_packed) + ) + if const_expr(mLSE is not None): + shape_LSE_packed = ( + (self.qhead_per_kvhead, mLSE.shape[0]), + mK.shape[2], + *mLSE.shape[2:], + ) + stride_LSE_packed = ( + (mLSE.stride[1], mLSE.stride[0]), + mLSE.stride[1] * self.qhead_per_kvhead, + *mLSE.stride[2:], + ) + mLSE = cute.make_tensor( + mLSE.iterator, cute.make_layout(shape_LSE_packed, stride=stride_LSE_packed) + ) + + # 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), + ] + } + tma_atom_Q, tma_tensor_Q = None, None + if const_expr(self.use_tma_Q): + tma_atom_Q, tma_tensor_Q = cpasync.make_tiled_tma_atom( + gmem_tiled_copy_Q, + mQ, + self.sQ_layout, + (self.tile_m, self.tile_hdim), # No mcast + ) + 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): + tma_atom_O, tma_tensor_O = cpasync.make_tiled_tma_atom( + gmem_tiled_copy_O, + mO, + 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]), + 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) + LOG2_E = math.log2(math.e) + if const_expr(self.score_mod is None): + softmax_scale_log2 = softmax_scale * LOG2_E + softmax_scale = None + else: + # NB: If a user passes in a score mod, we want to apply the score-mod in the sm_scaled qk + # But in the original base 10. We hijack softmax_scale_log2 to just be the change of base + # and correctly apply the softmax_scale prior to score_mod in the softmax step + softmax_scale_log2 = LOG2_E + softmax_scale = softmax_scale + 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_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]) + if const_expr(mPageTable is None) + else mK.shape[0] * mPageTable.shape[1] + ) + seqlen_q_divmod = FastDivmodDivisor(seqlen_q) + seqlen_k_divmod = FastDivmodDivisor(seqlen_k) + fastdiv_mods = (seqlen_q_divmod, seqlen_k_divmod) + + self.kernel( + tma_tensor_Q if const_expr(self.use_tma_Q) else mQ, + tma_tensor_K, + tma_tensor_V, + tma_tensor_O if const_expr(self.use_tma_O) else mO, + mLSE, + mCuSeqlensQ, + mCuSeqlensK, + mSeqUsedQ, + mSeqUsedK, + 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, + tiled_mma_pv_rs, + tile_sched_params, + TileScheduler, + SharedStorage, + aux_tensors, + 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], + 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, + tiled_mma_pv_rs: cute.TiledMma, + tile_sched_params: ParamsBase, + TileScheduler: cutlass.Constexpr[Callable], + SharedStorage: cutlass.Constexpr[Callable], + aux_tensors=Optional[list[cute.Tensor]], + 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 init + mbar_ptr_Q = storage.mbar_ptr.data_ptr() + if warp_idx == 1: + # if tidx < 2: + # # barrierO num threads should be self.num_mma_threads + # cute.arch.mbarrier_init(mbar_ptr_Q + tidx, 1 if tidx == 0 else self.num_mma_threads) + if const_expr(not self.use_tma_Q): + cute.arch.mbarrier_init(mbar_ptr_Q, self.num_Q_load_threads) + # cute.arch.mbarrier_init(mbar_ptr_Q + 1, self.num_mma_threads) + # We rely on pipeline_k and pipeline_v to initialize the mbarrier fence and sync + pipeline_kv_producer_group = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread + ) + pipeline_kv_consumer_group = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, self.num_mma_threads // cute.arch.WARP_SIZE + ) + pipeline_k = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.mbar_ptr_K.data_ptr(), + num_stages=self.num_stages, + producer_group=pipeline_kv_producer_group, + consumer_group=pipeline_kv_consumer_group, + tx_count=self.tma_copy_bytes["K"], + defer_sync=True, + ) + pipeline_v = pipeline.PipelineTmaAsync.create( + barrier_storage=storage.mbar_ptr_V.data_ptr(), + num_stages=self.num_stages, + producer_group=pipeline_kv_producer_group, + consumer_group=pipeline_kv_consumer_group, + tx_count=self.tma_copy_bytes["V"], + defer_sync=False + ) + + # /////////////////////////////////////////////////////////////////////////////// + # 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 = 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], + mCuSeqlensQ=mCuSeqlensQ, + mCuSeqlensK=mCuSeqlensK, + mSeqUsedQ=mSeqUsedQ, + mSeqUsedK=mSeqUsedK, + ) + 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) + + if warp_idx < 4: # Producer + cute.arch.warpgroup_reg_dealloc(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, + mbar_ptr_Q, + blocksparse_tensors, + block_info, + SeqlenInfoCls, + TileSchedulerCls, + ) + + else: # Consumer + cute.arch.warpgroup_reg_alloc(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, + tiled_mma_pv_rs, + mQ, + mO, + mLSE, + sQ, + sK, + sVt, + sP, + sO, + learnable_sink, + pipeline_k, + pipeline_v, + mbar_ptr_Q, + gmem_tiled_copy_Q, + gmem_tiled_copy_O, + tma_atom_O, + tidx, + softmax_scale_log2, + softmax_scale, + block_info, + SeqlenInfoCls, + AttentionMaskCls, + TileSchedulerCls, + blocksparse_tensors, + aux_tensors, + 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: cute.CopyAtom, + tma_atom_K: cute.CopyAtom, + tma_atom_V: cute.CopyAtom, + pipeline_k: cutlass.pipeline.PipelineAsync, + pipeline_v: cutlass.pipeline.PipelineAsync, + mbar_ptr_Q: cutlass.Pointer, + 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 + if warp_idx_in_wg == 0: + q_producer_phase = Int32(1) + kv_producer_state = pipeline.make_pipeline_state( + cutlass.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 + ) + 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)) + 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 + ) + # TODO: mcast + # TODO check warp_idx if we have 128 producer threads + load_K, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_K, 0, cute.make_layout(1), gK, sK + ) + load_K = copy_utils.tma_producer_copy_fn(load_K, pipeline_k) + load_V, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_V, 0, cute.make_layout(1), gV, sV + ) + load_V = copy_utils.tma_producer_copy_fn(load_V, pipeline_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) + # 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) + # First iteration: load both Q & K with the same mbarrier + n_block = n_block_max - 1 + pipeline_k.producer_acquire( + kv_producer_state, + extra_tx_count=self.tma_copy_bytes["Q"] + if const_expr(self.use_tma_Q) + else 0, + ) + if const_expr(self.use_tma_Q): + load_Q(tma_bar_ptr=pipeline_k.producer_get_barrier(kv_producer_state)) + load_K(src_idx=n_block, producer_state=kv_producer_state) + + if const_expr(not self.intra_wg_overlap): + pipeline_v.producer_acquire(kv_producer_state) + load_V(src_idx=n_block, producer_state=kv_producer_state) + 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 + 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: + 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 + 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) + n_block = n_block_min + pipeline_v.producer_acquire(kv_producer_state) + load_V(src_idx=n_block, producer_state=kv_producer_state) + kv_producer_state.advance() + else: + kv_producer_state = produce_block_sparse_loads( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + kv_producer_state, + load_Q, + load_K, + load_V, + pipeline_k, + pipeline_v, + self.use_tma_Q, + self.tma_copy_bytes["Q"], + self.intra_wg_overlap, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + + tile_scheduler.prefetch_next_work() + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + # End of persistent scheduler loop + + @cute.jit + def mma( + self, + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + tiled_mma_pv_rs: cute.TiledMma, + # softmax: Softmax, + # acc_O: cute.Tensor, + mQ: cute.Tensor, + 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: cutlass.pipeline.PipelineAsync, + pipeline_v: cutlass.pipeline.PipelineAsync, + mbar_ptr_Q: cutlass.Pointer, + gmem_tiled_copy_Q: cute.TiledCopy, + 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_tensors: Optional[list], + fastdiv_mods=None, + ): + warp_group_idx = cute.arch.make_warp_uniform(tidx // self.num_threads_per_warp_group) + warp_group_thread_layout = cute.make_layout( + self.num_mma_warp_groups, 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 = tiled_mma_qk.make_fragment_A(wg_mma_qk.partition_A(sQ)) + tSrK = tiled_mma_qk.make_fragment_B(wg_mma_qk.partition_B(sK)) + if const_expr(self.mma_pv_is_rs): + acc_S_shape = tiled_mma_qk.partition_shape_C((self.tile_m, self.tile_n)) + tOrP = cute.make_fragment( + utils.convert_layout_acc_frgA(cute.make_layout(acc_S_shape)), self.dtype + ) + else: + tOrP = tiled_mma_pv.make_fragment_A(wg_mma_pv.partition_A(sP)) + tOrVt = tiled_mma_pv.make_fragment_B(wg_mma_pv.partition_B(sVt)) + + # /////////////////////////////////////////////////////////////////////////////// + # Smem copy atom tiling + # /////////////////////////////////////////////////////////////////////////////// + smem_copy_atom_P = utils.get_smem_store_atom(self.arch, 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_pi) if const_expr(sP_pi is not None) else None + tPsP = smem_thr_copy_P.partition_D(sP) if const_expr(sP is not None) else None + # if cute.arch.thread_idx()[0] == 0: + # cute.printf(sP_pi.layout, sP_pi.iterator) + # cute.printf(sP.layout, sP.iterator) + # cute.printf(tPsP.layout, tPsP.iterator) + + self.mma_init() + + acc_shape_O = tiled_mma_pv.partition_shape_C((self.tile_m, self.tile_hdimv)) + acc_O = cute.make_fragment(acc_shape_O, Float32) + smem_copy_params = SimpleNamespace(smem_thr_copy_P=smem_thr_copy_P, tPsP=tPsP) + + mma_qk_fn = partial( + sm90_utils.gemm_zero_init, tiled_mma_qk, (self.tile_m, self.tile_n), tSrQ, tSrK + ) + mma_pv_fn = partial(sm90_utils.gemm_w_idx, tiled_mma_pv, acc_O, tOrP, tOrVt) + + 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, + tiled_mma_pv_rs=tiled_mma_pv_rs, + pipeline_k=pipeline_k, + pipeline_v=pipeline_v, + acc_O=acc_O, + tOrP=tOrP, + smem_copy_params=smem_copy_params, + check_inf=True, + ) + + q_consumer_phase = Int32(0) + kv_consumer_state = pipeline.make_pipeline_state( + cutlass.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, + ) + + 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, + softmax=softmax, + ) + process_last_half_block = partial( + self.last_half_block_overlap, + pipeline_v=pipeline_v, + mma_pv_fn=mma_pv_fn, + ) + 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_tensors=aux_tensors, + 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_tensors=aux_tensors, + fastdiv_mods=fastdiv_mods, + ) + mma_one_n_block = partial( + mma_one_n_block_all, + seqlen=seqlen, + softmax=softmax, + score_mod_fn=score_mod_fn, + ) + # Load Q if not TMA_Q + 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 + ) + mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx] + # gmem_thr_copy_Q = gmem_tiled_copy_Q.get_slice(tidx) + # gQ = cute.local_tile(mQ_cur, (self.tile_m, self.tile_hdim), (m_block, 0)) + # self.load_Q(gmem_thr_copy_Q, gQ, sQ, m_block, seqlen=seqlen.seqlen_q, + # headdim=mQ.shape[1]) + pack_gqa.load_Q(mQ_cur, sQ, gmem_tiled_copy_Q, tidx, m_block, seqlen.seqlen_q) + cute.arch.cp_async_mbarrier_arrive_noinc(mbar_ptr_Q) + + n_block_min, n_block_max = block_info.get_n_block_min_max(seqlen, m_block) + if const_expr(not self.use_tma_Q): + cute.arch.mbarrier_wait(mbar_ptr_Q, phase=q_consumer_phase) + q_consumer_phase ^= 1 + # 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, + ) + # Need to initialize tOrO in the case of RescaleOBeforeGemm where we will scale tOrO even in the 1st iter + # acc_O.fill(0.0) + 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 + # 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, + ) + + # Handle empty case (when no blocks to process) + if not processed_any: + softmax.reset() + acc_O.fill(0.0) + + 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_fragment_like(softmax.row_max, Float32) + cS = cute.make_identity_tensor((self.tile_m, self.tile_n)) + tScS_mn = utils.make_acc_tensor_mn_view(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, + 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) + + softmax.online_softmax(acc_S, is_first=is_first_block) + + tOrP_acc = cute.make_tensor(acc_S.iterator, utils.convert_layout_acc_frgA(acc_S.layout)) + tOrP_cur = ( + tOrP if const_expr(self.mma_pv_is_rs) else cute.make_fragment_like(tOrP_acc, self.dtype) + ) + tOrP_cur.store(tOrP_acc.load().to(self.dtype)) + + # if pv gemm not rs + 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_proxy( + cute.arch.ProxyKind.async_shared, space=cute.arch.SharedSpace.shared_cta + ) + cute.arch.sync_warp() + + return kv_consumer_state + + @cute.jit + def last_half_block_overlap( + self, + kv_consumer_state, + pipeline_v, + mma_pv_fn: Callable, + zero_init: bool, + ): + """Processes the final PV GEMM when using intra-warpgroup-overlap""" + + 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: cutlass.pipeline.PipelineState | pipeline.PipelineStateSimple, + n_block: Int32, + mma_qk_fn: Callable, + mma_pv_fn: Callable, + tiled_mma_pv_rs: cute.TiledMma, + pipeline_k: cutlass.pipeline.PipelineAsync, + pipeline_v: cutlass.pipeline.PipelineAsync, + acc_O: cute.Tensor, + tOrP: cute.Tensor, + smem_copy_params: SimpleNamespace, + softmax: Softmax, + seqlen: SeqlenInfoQK, + 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(utils.make_acc_tensor_mn_view(acc_S)) + tOrP_acc = cute.make_tensor(acc_S.iterator, utils.convert_layout_acc_frgA(acc_S.layout)) + tOrP_cur = ( + tOrP if const_expr(self.mma_pv_is_rs) else cute.make_fragment_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_proxy(ProxyKind.async_shared, space=SharedSpace.shared_cta) + 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: cutlass.pipeline.PipelineState | pipeline.PipelineStateSimple, + n_block: Int32, + mma_qk_fn: Callable, + mma_pv_fn: Callable, + tiled_mma_pv_rs: cute.TiledMma, + pipeline_k: cutlass.pipeline.PipelineAsync, + pipeline_v: cutlass.pipeline.PipelineAsync, + acc_O: cute.Tensor, + tOrP: cute.Tensor, + smem_copy_params: SimpleNamespace, + softmax: Softmax, + seqlen: SeqlenInfoQK, + 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) + 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(utils.make_acc_tensor_mn_view(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 = cute.make_tensor(acc_S.iterator, utils.convert_layout_acc_frgA(acc_S.layout)) + tOrP_cur = ( + tOrP if const_expr(self.mma_pv_is_rs) else cute.make_fragment_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) + 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_proxy(ProxyKind.async_shared, space=SharedSpace.shared_cta) + 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_tensors: Optional[list] = None, + 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.vec_size, + self.qk_acc_dtype, + aux_tensors, + 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_mma_warp_groups in [2, 3] + cur_wg = utils.canonical_warp_group_idx(sync=False) - 1 + if const_expr(self.num_mma_warp_groups == 2): + next_wg = 1 - cur_wg + else: + t = cur_wg + 1 + next_wg = t % self.num_mma_warp_groups + cute.arch.barrier_arrive( + barrier_id=int(NamedBarrierFwd.WarpSchedulerWG1) + next_wg, + number_of_threads=2 * self.num_threads_per_warp_group, + ) + diff --git a/python/sglang/jit_kernel/flash_attention/cute/flash_fwd_combine.py b/python/sglang/jit_kernel/flash_attention/cute/flash_fwd_combine.py new file mode 100644 index 000000000..b6f4acf51 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/flash_fwd_combine.py @@ -0,0 +1,704 @@ +# 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 +import operator +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, const_expr + +import sglang.jit_kernel.flash_attention.cute.utils as utils +from .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, + m_block_size: 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 m_block_size: 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.m_block_size = m_block_size + 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, + m_block_size, + 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 m_block_size % 8 != 0: + return False + max_splits = 1 << log_max_splits + if max_splits > 256: + return False + if (m_block_size * 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.m_block_size % 128 == 0 + else ( + 64 + if self.m_block_size % 64 == 0 + else ( + 32 + if self.m_block_size % 32 == 0 + else (16 if self.m_block_size % 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.m_block_size), (0, 1) + ) + + # O partial shared memory layout (simple layout for pipeline stages) + self.smem_layout_o = cute.make_ordered_layout( + (self.m_block_size, 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, + semaphore_to_reset: Optional[cute.Tensor] = None, + 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)" + ) + + # Assume all strides are divisible by 128 bits except the last stride + new_stride = lambda t: ( + *(cute.assume(s, divby=128 // t.element_type.width) for s in t.stride[:-1]), + t.stride[-1], + ) + mO_partial, mO = [ + cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(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.m_block_size], 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.m_block_size), + 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, + 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], + 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, batch_idx = cute.arch.block_idx() + + # /////////////////////////////////////////////////////////////////////////////// + # 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.m_block_size,)) + sO = storage.sO.get_tensor(smem_layout_o) + + # Handle semaphore reset + 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 batch_idx == cute.arch.grid_dim()[2] - 1 + ): + semaphore_to_reset[0] = 0 + + # Get number of splits + num_splits = ( + num_splits_dynamic_ptr[batch_idx] + 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, + ) + 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.m_block_size < max_idx + ): + # =============================== + # Step 1: Load LSE_partial from gmem to shared memory + # =============================== + + if const_expr(cu_seqlens is None): + # mLSE_partial_cur = mLSE_partial[None, None, None, batch_idx] + mLSE_partial_cur = utils.coord_offset_i64(mLSE_partial, batch_idx, dim=3) + else: + # mLSE_partial_cur = cute.domain_offset((offset, 0, 0), mLSE_partial) + mLSE_partial_cur = utils.domain_offset_i64((offset, 0, 0), mLSE_partial) + 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.m_block_size)) + 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.m_block_size + 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.m_block_size, self.k_block_size)) + tOcO = gmem_thr_copy_O_partial.partition_D(cO) + tOsO_partial = gmem_thr_copy_O_partial.partition_D(sO) + if const_expr(cu_seqlens is None): + # mO_partial_cur = mO_partial[None, None, None, None, batch_idx] + mO_partial_cur = utils.coord_offset_i64(mO_partial, batch_idx, dim=4) + else: + # mO_partial_cur = cute.domain_offset((offset, 0, 0, 0), mO_partial) + mO_partial_cur = utils.domain_offset_i64((offset, 0, 0, 0), mO_partial) + + # Precompute these values to avoid recomputing them in the loop + num_rows = const_expr(cute.size(tOcO, mode=[1])) + tOmidx = cute.make_fragment(num_rows, cutlass.Int32) + tOhidx = cute.make_fragment(num_rows, cutlass.Int32) + tOrOptr = cute.make_fragment(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.m_block_size + 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_i64( + mO_partial_cur, (tOmidx[m], k_block * self.k_block_size, 0, tOhidx[m]) + ).toint() + if idx >= max_idx: + tOhidx[m] = -1 + + tOpO = cute.make_fragment(cute.size(tOcO, [2]), cutlass.Boolean) + if const_expr(not self.is_even_k): + 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_fragment_like(ts2rsLSE) + cute.copy(s2r_tiled_copy_LSE, ts2rsLSE, ts2rrLSE) + + # =============================== + # Step 4: Compute final LSE along split dimension + # =============================== + + lse_sum = cute.make_fragment(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_fragment(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 = utils.warp_reduce( + ts2rrLSE[None, None, m] + .load() + .reduce(cute.ReductionOp.MAX, init_val=-Float32.inf, reduction_profile=0), + op=cute.arch.fmax, + width=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] = utils.warp_reduce(max_valid_idx, max, width=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 = utils.exp2f(ts2rrLSE[0, s, m] * LOG2_E - (lse_max_cur * LOG2_E)) + lse_sum_cur += scale + ts2rrLSE[0, s, m] = scale # Store scale for later use + lse_sum_cur = utils.warp_reduce(lse_sum_cur, operator.add, width=threads_per_col) + lse_sum[m] = utils.logf(lse_sum_cur) + 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.m_block_size: + 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] + mLSE_cur = utils.coord_offset_i64(mLSE, batch_idx, dim=2) + else: + # mLSE_cur = cute.domain_offset((offset, 0), mLSE) + mLSE_cur = utils.domain_offset_i64((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.m_block_size + 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])): + thr_max_valid_split = max(thr_max_valid_split, sMaxValidSplit[tOcO[0, m, 0][0]]) + + tOrO_partial = cute.make_fragment_like(tOsO_partial[None, None, None, 0]) + tOrO = cute.make_fragment_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_fragment(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_fragment_like(tOrO, self.dtype) + rO.store(tOrO.load().to(self.dtype)) + if const_expr(cu_seqlens is None): + # mO_cur = mO[None, None, None, batch_idx] + mO_cur = utils.coord_offset_i64(mO, batch_idx, dim=3) + else: + # mO_cur = cute.domain_offset((offset, 0, 0), mO) + mO_cur = utils.domain_offset_i64((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: 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(self.is_even_k) or tOpO[k]: + cute.copy( + gmem_tiled_copy_O_partial, + # mO_partial_cur_copy[None, k_idx, split], + utils.coord_offset_i64(mO_partial_cur_copy, split, dim=2)[None, k_idx], + tOsO_partial_cur[None, m, k], + ) diff --git a/python/sglang/jit_kernel/flash_attention/cute/flash_fwd_sm100.py b/python/sglang/jit_kernel/flash_attention/cute/flash_fwd_sm100.py new file mode 100644 index 000000000..b01b1b33f --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/flash_fwd_sm100.py @@ -0,0 +1,2740 @@ +# 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 enum +import math +from typing import Type, Tuple, Callable, Optional, Literal +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 +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +import cutlass.utils.blackwell_helpers as sm100_utils_basic + +from .paged_kv import PagedKVManager +import sglang.jit_kernel.flash_attention.cute.utils as utils +import sglang.jit_kernel.flash_attention.cute.copy_utils as copy_utils +import sglang.jit_kernel.flash_attention.cute.pipeline as pipeline +from .mask import AttentionMask +from .softmax import SoftmaxSm100, apply_score_mod_inner +from .seqlen_info import SeqlenInfoQK +from .block_info import BlockInfo +from .block_sparsity import BlockSparseTensors +from .block_sparse_utils import ( + get_total_block_count, + produce_block_sparse_loads_sm100, + softmax_block_sparse_sm100, + handle_block_sparse_empty_tile_correction_sm100, +) +from .pack_gqa import PackGQA +import sglang.jit_kernel.flash_attention.cute.mma_sm100_desc as sm100_desc +import sglang.jit_kernel.flash_attention.cute.blackwell_helpers as sm100_utils +from cutlass.cute import FastDivmodDivisor +from .tile_scheduler import ( + TileSchedulerArguments, + SingleTileScheduler, + StaticPersistentTileScheduler, + SingleTileLPTScheduler, + SingleTileVarlenScheduler, + ParamsBase, +) + + +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 FlashAttentionForwardSm100: + arch = 100 + + 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, + m_block_size: int = 128, + n_block_size: int = 128, + q_stage: cutlass.Constexpr[int] = 2, + 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, + ): + 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.q_stage = q_stage + assert self.q_stage in [1, 2] + + # 2 Q tile per CTA + self.cta_tiler = (self.q_stage * m_block_size, n_block_size, self.head_dim_padded) + self.mma_tiler_qk = (m_block_size, n_block_size, self.head_dim_padded) + self.mma_tiler_pv = (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 = (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.use_correction_warps_for_epi = is_varlen_q + self.qhead_per_kvhead = qhead_per_kvhead + self.is_split_kv = is_split_kv + self.pack_gqa = pack_gqa + if pack_gqa: + assert m_block_size % self.qhead_per_kvhead == 0, ( + "For PackGQA, m_block_size must be divisible by qhead_per_kvhead" + ) + 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 + if cutlass.const_expr(has_aux_tensors): + self.vec_size: cutlass.Constexpr = 1 + else: + self.vec_size: cutlass.Constexpr = 2 + # 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) + 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" + ) + + 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,) + SM100_TMEM_CAPACITY_COLUMNS = 512 + self.tmem_alloc_cols = SM100_TMEM_CAPACITY_COLUMNS + + 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, + ) + ) + + if self.q_stage == 1: + if not self.use_tma_KV: + 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 + elif self.is_varlen_q: # fallback + self.epilogue_warp_ids = (13, 14) + + self.tmem_s_offset = [0, self.n_block_size] # e.g., 0, 128 + self.tmem_o_offset = [ + self.tmem_s_offset[-1] + self.n_block_size + i * self.head_dim_v_padded + for i in range(self.q_stage) + ] # e.g., 256, 384 + self.tmem_total = self.tmem_o_offset[-1] + self.head_dim_v_padded + assert self.tmem_total <= SM100_TMEM_CAPACITY_COLUMNS + 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(2) + ] # 0, 128 + + # vec buffer for row_max & row_sum + self.tmem_vec_offset = self.tmem_s_offset + + 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: + # self.num_regs_softmax = 192 if self.is_causal or self.is_local else 184 + self.num_regs_softmax = 200 if not paged_kv_non_tma else 184 + # self.num_regs_softmax = 176 + # self.num_regs_correction = 96 + # self.num_regs_correction = 80 + # self.num_regs_correction = 64 if self.is_causal or self.is_local else 80 + self.num_regs_correction = 64 + # self.num_regs_other = 32 + # self.num_regs_other = 64 + # self.num_regs_other = 80 + self.num_regs_other = 48 if not paged_kv_non_tma else 80 + # self.num_regs_other = 96 if self.is_causal or self.is_local else 80 + # self.num_regs_other = 64 if self.is_causal or self.is_local else 80 + self.num_regs_empty = 24 + + 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 + """ + + self.kv_stage = 4 if self.q_dtype.width == 8 or self.q_stage == 1 else 3 + self.acc_stage = 1 + # For hdim 192,128, 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.m_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, + stream: cuda.CUstream, + 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_tensors: Optional[list] = 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 + # Assume all strides are divisible by 128 bits except the last stride + new_stride = lambda t: ( + *(cute.assume(s, divby=128 // t.element_type.width) for s in t.stride[:-1]), + t.stride[-1], + ) + mQ, mK, mV, mO = [ + cute.make_tensor(t.iterator, cute.make_layout(t.shape, stride=new_stride(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)) + + self.q_major_mode = cutlass.utils.LayoutEnum.from_tensor(mQ).mma_major_mode() + self.k_major_mode = cutlass.utils.LayoutEnum.from_tensor(mK).mma_major_mode() + self.v_major_mode = cutlass.utils.LayoutEnum.from_tensor(mV).mma_major_mode() + self.o_layout = cutlass.utils.LayoutEnum.from_tensor(mO) + + if const_expr(self.q_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of mQ is not supported") + if const_expr(self.k_major_mode != tcgen05.OperandMajorMode.K): + raise RuntimeError("The layout of mK is not supported") + if const_expr(self.v_major_mode != tcgen05.OperandMajorMode.MN): + raise RuntimeError("The layout of mV is not supported") + + # 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}") + self._setup_attributes() + self.use_tma_O = self.arch >= 90 and mCuSeqlensQ is None and mSeqUsedQ is None + # This can be tuned + self.e2e_freq = 16 + if const_expr( + self.head_dim_padded > 64 and not self.is_causal and not self.is_local and self.pack_gqa + ): + self.e2e_freq = 32 if mCuSeqlensQ is not None or mSeqUsedQ is not None else 10 + + cta_group = tcgen05.CtaGroup.ONE + # the intermediate tensor p is from tmem & mK-major + p_source = tcgen05.OperandSource.TMEM + p_major_mode = tcgen05.OperandMajorMode.K + tiled_mma_qk = sm100_utils_basic.make_trivial_tiled_mma( + self.q_dtype, + self.q_major_mode, + self.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, + self.v_major_mode, + self.pv_acc_dtype, + cta_group, + self.mma_tiler_pv[: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), + (tiled_mma_qk.thr_id.shape,), + ) + + self.epi_tile = self.mma_tiler_pv[:2] + + 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.acc_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): + shape_Q_packed = ( + (self.qhead_per_kvhead, mQ.shape[0]), + mQ.shape[1], + mK.shape[2], + *mQ.shape[3:], + ) + stride_Q_packed = ( + (mQ.stride[2], mQ.stride[0]), + mQ.stride[1], + mQ.stride[2] * self.qhead_per_kvhead, + *mQ.stride[3:], + ) + mQ = cute.make_tensor( + mQ.iterator, cute.make_layout(shape_Q_packed, stride=stride_Q_packed) + ) + shape_O_packed = ( + (self.qhead_per_kvhead, mO.shape[0]), + mO.shape[1], + mK.shape[2], + *mO.shape[3:], + ) + stride_O_packed = ( + (mO.stride[2], mO.stride[0]), + mO.stride[1], + mO.stride[2] * self.qhead_per_kvhead, + *mO.stride[3:], + ) + mO = cute.make_tensor( + mO.iterator, cute.make_layout(shape_O_packed, stride=stride_O_packed) + ) + if const_expr(mLSE is not None): + shape_LSE_packed = ( + (self.qhead_per_kvhead, mLSE.shape[0]), + mK.shape[2], + *mLSE.shape[2:], + ) + stride_LSE_packed = ( + (mLSE.stride[1], mLSE.stride[0]), + mLSE.stride[1] * self.qhead_per_kvhead, + *mLSE.stride[2:], + ) + mLSE = cute.make_tensor( + mLSE.iterator, cute.make_layout(shape_LSE_packed, stride=stride_LSE_packed) + ) + + 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), + ] + } + + # TMA load for Q + tma_load_op = cpasync.CopyBulkTensorTileG2SOp(cta_group) + tma_store_op = cpasync.CopyBulkTensorTileS2GOp() + + 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, + self.cluster_layout_vmnk.shape, + ) + + 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, + self.cluster_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, + self.cluster_layout_vmnk.shape, + ) + else: + tma_atom_K = None + tma_atom_V = None + + o_cta_v_layout = cute.composition(cute.make_identity_layout(mO.shape), self.epi_tile) + + 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]), + o_cta_v_layout, + ) + 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) + + if const_expr(mCuSeqlensQ is not None or mSeqUsedQ is not None): + TileScheduler = SingleTileVarlenScheduler + else: + if const_expr(self.is_causal or self.is_local): + TileScheduler = SingleTileLPTScheduler + else: + TileScheduler = ( + SingleTileScheduler + if const_expr(not self.is_persistent) + else StaticPersistentTileScheduler + ) + tile_sched_args = TileSchedulerArguments( + cute.ceil_div(cute.size(mQ.shape[0]), self.cta_tiler[0]), + 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, + ) + tile_sched_params = TileScheduler.to_underlying_arguments(tile_sched_args) + self.tile_scheduler_cls = TileScheduler + grid_dim = TileScheduler.get_grid_shape(tile_sched_params) + + self.mbar_load_q_full_offset = 0 + self.mbar_load_q_empty_offset = self.mbar_load_q_full_offset + self.q_stage + self.mbar_load_kv_full_offset = self.mbar_load_q_empty_offset + self.q_stage + self.mbar_load_kv_empty_offset = self.mbar_load_kv_full_offset + self.kv_stage + self.mbar_P_full_O_rescaled_offset = self.mbar_load_kv_empty_offset + self.kv_stage + self.mbar_S_full_offset = self.mbar_P_full_O_rescaled_offset + self.q_stage + self.mbar_O_full_offset = self.mbar_S_full_offset + self.q_stage + self.mbar_softmax_corr_full_offset = self.mbar_O_full_offset + self.q_stage + self.mbar_softmax_corr_empty_offset = self.mbar_softmax_corr_full_offset + self.q_stage + self.mbar_corr_epi_full_offset = self.mbar_softmax_corr_empty_offset + self.q_stage + self.mbar_corr_epi_empty_offset = self.mbar_corr_epi_full_offset + self.q_stage + self.mbar_s0_s1_sequence_offset = self.mbar_corr_epi_empty_offset + self.q_stage + self.mbar_tmem_dealloc_offset = self.mbar_s0_s1_sequence_offset + 8 + self.mbar_P_full_2_offset = self.mbar_tmem_dealloc_offset + 1 + self.mbar_total = self.mbar_P_full_2_offset + self.q_stage + + 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) + ) + + @cute.struct + class SharedStorage: + # m_barriers for pipelines + mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.mbar_total] + # Tmem holding buffer + tmem_holding_buf: Int32 + # Smem tensors + # store row max and row sum + sScale: cute.struct.MemRange[Float32, self.q_stage * self.m_block_size * 2] + 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, + ] + + self.shared_storage = SharedStorage + + LOG2_E = math.log2(math.e) + if const_expr(self.score_mod is None): + softmax_scale_log2 = softmax_scale * LOG2_E + softmax_scale = None + else: + # NB: If a users passes in a score mod, we want to apply the score-mod in the sm_scaled qk + # But in the original base 10. We hijack softmax_scale_log2 to just be the change of base + # and correctly apply the softmax_scale prior to score_mod in the softmax step + softmax_scale_log2 = LOG2_E + softmax_scale = softmax_scale + + 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 cutlass.const_expr(aux_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]) + if const_expr(mPageTable is None) + else mK.shape[0] * mPageTable.shape[1] + ) + seqlen_q_divmod = FastDivmodDivisor(seqlen_q) + seqlen_k_divmod = FastDivmodDivisor(seqlen_k) + fastdiv_mods = (seqlen_q_divmod, seqlen_k_divmod) + + 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") + + # 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, + blocksparse_tensors, + sQ_layout, + sK_layout, + tP_layout, + sV_layout, + sO_layout, + gmem_tiled_copy_O, + tiled_mma_qk, + tiled_mma_pv, + tile_sched_params, + num_splits, + aux_tensors, + fastdiv_mods, + head_divmod, + ).launch( + grid=grid_dim, + block=[self.threads_per_cta, 1, 1], + cluster=self.cluster_shape_mnk, + smem=self.shared_storage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=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: 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], + 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_O: Optional[cute.TiledCopy], + tiled_mma_qk: cute.TiledMma, + tiled_mma_pv: cute.TiledMma, + tile_sched_params: ParamsBase, + num_splits: Int32, + aux_tensors: Optional[list] = None, + 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: + cpasync.prefetch_descriptor(tma_atom_Q) + if const_expr(tma_atom_K is not None): + cpasync.prefetch_descriptor(tma_atom_K) + if const_expr(tma_atom_V is not None): + cpasync.prefetch_descriptor(tma_atom_V) + if const_expr(tma_atom_O is not None): + cpasync.prefetch_descriptor(tma_atom_O) + + # Alloc + smem = cutlass.utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + mbar_ptr = storage.mbar_ptr.data_ptr() + # Use the first N warps to initialize barriers + if warp_idx == 1: + # Init "full" barrier with number of producers, "empty" barrier with number of consumers + for i in cutlass.range_constexpr(self.q_stage): + cute.arch.mbarrier_init( + mbar_ptr + self.mbar_load_q_full_offset + i, 1 + ) + cute.arch.mbarrier_init( + mbar_ptr + self.mbar_load_q_empty_offset + i, len([self.mma_warp_id]) + ) + if warp_idx == 2: + for i in cutlass.range_constexpr(self.q_stage): + cute.arch.mbarrier_init( + mbar_ptr + self.mbar_softmax_corr_empty_offset + i, cute.arch.WARP_SIZE * 4 + ) + cute.arch.mbarrier_init( + mbar_ptr + self.mbar_softmax_corr_full_offset + i, cute.arch.WARP_SIZE * 4 + ) + if warp_idx == 3: + if const_expr(self.s0_s1_barrier): + for i in cutlass.range_constexpr(8): + cute.arch.mbarrier_init( + mbar_ptr + self.mbar_s0_s1_sequence_offset + i, cute.arch.WARP_SIZE + ) + if const_expr(not self.use_correction_warps_for_epi) and warp_idx == 4: + for i in cutlass.range_constexpr(self.q_stage): + cute.arch.mbarrier_init( + mbar_ptr + self.mbar_corr_epi_full_offset + i, + cute.arch.WARP_SIZE * len(self.correction_warp_ids), + ) + cute.arch.mbarrier_init( + mbar_ptr + self.mbar_corr_epi_empty_offset + i, + cute.arch.WARP_SIZE * len(self.epilogue_warp_ids), + ) + if warp_idx == 5: + for i in cutlass.range_constexpr(self.q_stage): + cute.arch.mbarrier_init( + mbar_ptr + self.mbar_P_full_O_rescaled_offset + i, + cute.arch.WARP_SIZE + * (len(self.softmax0_warp_ids) + len(self.correction_warp_ids)), + ) + cute.arch.mbarrier_init( + mbar_ptr + self.mbar_S_full_offset + i, len([self.mma_warp_id]) + ) + cute.arch.mbarrier_init( + mbar_ptr + self.mbar_O_full_offset + i, len([self.mma_warp_id]) + ) + if warp_idx == 6: + for i in cutlass.range_constexpr(self.q_stage): + cute.arch.mbarrier_init( + mbar_ptr + self.mbar_P_full_2_offset + i, + cute.arch.WARP_SIZE * len(self.softmax0_warp_ids), + ) + if warp_idx == 7: + cute.arch.mbarrier_init( + mbar_ptr + self.mbar_tmem_dealloc_offset, + cute.arch.WARP_SIZE + * len( + ( + *self.softmax0_warp_ids, + *self.softmax1_warp_ids, + *self.correction_warp_ids, + ) + ), + ) + # Relying on pipeline_kv constructor to call mbarrier_init_fence and sync + pipeline_kv = self.make_and_init_load_kv_pipeline(mbar_ptr + self.mbar_load_kv_full_offset) + + # 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.q_stage * self.m_block_size * 2)) + + thr_mma_qk = tiled_mma_qk.get_slice(0) # default 1SM + thr_mma_pv = tiled_mma_pv.get_slice(0) # default 1SM + + qk_acc_shape = thr_mma_qk.partition_shape_C(self.mma_tiler_qk[:2]) + tStS_fake = thr_mma_qk.make_fragment_C(qk_acc_shape) + # 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) + tStS = cute.make_tensor(tmem_ptr, tStS_fake.layout) + + pv_acc_shape = thr_mma_pv.partition_shape_C(self.mma_tiler_pv[:2]) + tOtO = thr_mma_pv.make_fragment_C(pv_acc_shape) + + tStSs = tuple( + cute.make_tensor(tStS.iterator + self.tmem_s_offset[stage], tStS.layout) + for stage in range(self.q_stage) + ) + tOtOs = tuple( + cute.make_tensor(tOtO.iterator + self.tmem_o_offset[stage], tOtO.layout) + for stage in range(self.q_stage) + ) + + tP = cute.make_tensor(tStS.iterator, tP_layout.outer) + tOrP = thr_mma_pv.make_fragment_A(tP)[None, None, None, 0] + + tOrPs = [ + cute.make_tensor( + tOrP.iterator + + self.qk_acc_dtype.width // self.q_dtype.width * self.tmem_p_offset[stage], + tOrP.layout, + ) + for stage in range(self.q_stage) + ] + + 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, + 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, + ) + AttentionMaskCls = 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, + ) + TileSchedulerCls = partial(self.tile_scheduler_cls.create, tile_sched_params) + + # /////////////////////////////////////////////////////////////////////////////// + # EMPTY + # /////////////////////////////////////////////////////////////////////////////// + for i in cutlass.range_constexpr(len(self.empty_warp_ids)): + if warp_idx == self.empty_warp_ids[i]: + cute.arch.warpgroup_reg_dealloc(self.num_regs_empty) + + # /////////////////////////////////////////////////////////////////////////////// + # LOAD + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx >= self.load_warp_ids[0] and warp_idx <= self.load_warp_ids[-1]: + cute.arch.warpgroup_reg_dealloc(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, + pipeline_kv, + mbar_ptr, + block_info, + num_splits, + SeqlenInfoCls, + TileSchedulerCls, + blocksparse_tensors, + ) + + # /////////////////////////////////////////////////////////////////////////////// + # MMA + # /////////////////////////////////////////////////////////////////////////////// + if warp_idx == self.mma_warp_id: + # if warp_idx == self.mma_warp_id or warp_idx == self.empty_warp_ids: + cute.arch.warpgroup_reg_dealloc(self.num_regs_other) + # Alloc tmem buffer + tmem_alloc_cols = Int32(self.tmem_alloc_cols) + if warp_idx == self.mma_warp_id: + cute.arch.alloc_tmem(tmem_alloc_cols, storage.tmem_holding_buf) + cute.arch.sync_warp() + + self.mma( + tiled_mma_qk, + tiled_mma_pv, + sQ, + sK, + sV, + tStSs, + tOtOs, + tOrPs, + pipeline_kv, + mbar_ptr, + block_info, + num_splits, + SeqlenInfoCls, + TileSchedulerCls, + blocksparse_tensors, + ) + + # if warp_idx == self.mma_warp_id: + # dealloc tmem buffer + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.mbarrier_wait(mbar_ptr + self.mbar_tmem_dealloc_offset, 0) + tmem_alloc_cols = Int32(self.tmem_alloc_cols) + # Retrieving tmem ptr and make acc + tmem_ptr = cute.arch.retrieve_tmem_ptr( + Float32, + alignment=16, + ptr_to_buffer_holding_addr=storage.tmem_holding_buf, + ) + cute.arch.dealloc_tmem(tmem_ptr, tmem_alloc_cols) + + # /////////////////////////////////////////////////////////////////////////////// + # 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.warpgroup_reg_dealloc(self.num_regs_other) + self.epilogue_s2g( + mO, + sO, + gmem_tiled_copy_O, + tma_atom_O, + mbar_ptr, + block_info, + num_splits, + SeqlenInfoCls, + TileSchedulerCls, + ) + + # /////////////////////////////////////////////////////////////////////////////// + # Softmax + # /////////////////////////////////////////////////////////////////////////////// + if ( + (const_expr(self.q_stage == 2) and warp_idx <= self.softmax1_warp_ids[-1]) or + (const_expr(self.q_stage == 1) and warp_idx <= self.softmax0_warp_ids[-1]) + ): + # increase register after decreasing + cute.arch.warpgroup_reg_alloc(self.num_regs_softmax) + softmax_loop = partial( + self.softmax_loop, + softmax_scale_log2=softmax_scale_log2, + softmax_scale=softmax_scale, + thr_mma_qk=thr_mma_qk, + sScale=sScale, + mLSE=mLSE, + learnable_sink=learnable_sink, + mbar_ptr=mbar_ptr, + block_info=block_info, + num_splits=num_splits, + SeqlenInfoCls=SeqlenInfoCls, + AttentionMaskCls=AttentionMaskCls, + TileSchedulerCls=TileSchedulerCls, + aux_tensors=aux_tensors, + fastdiv_mods=fastdiv_mods, + head_divmod=head_divmod, + blocksparse_tensors=blocksparse_tensors, + ) + + if const_expr(not self.s0_s1_barrier): + stage = Int32(0 if const_expr(self.q_stage == 1) or warp_idx < self.softmax1_warp_ids[0] else 1) + softmax_loop( + stage=stage, + tStSi=cute.make_tensor( + tStS.iterator + + (self.tmem_s_offset[0] if stage == 0 else self.tmem_s_offset[1]), + tStS.layout, + ), + ) + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_tmem_dealloc_offset) + 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]: + tStSi = cute.make_tensor(tStS.iterator + self.tmem_s_offset[0], tStS.layout) + softmax_loop(stage=0, tStSi=tStSi) + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_tmem_dealloc_offset) + if warp_idx < self.correction_warp_ids[0] and warp_idx >= self.softmax1_warp_ids[0]: + tStSi = cute.make_tensor(tStS.iterator + self.tmem_s_offset[1], tStS.layout) + softmax_loop(stage=1, tStSi=tStSi) + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_tmem_dealloc_offset) + + # /////////////////////////////////////////////////////////////////////////////// + # 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) + self.correction_loop( + thr_mma_qk, + thr_mma_pv, + tStS, + tOtOs, + sScale, + mO, + mLSE, + sO, + learnable_sink, + gmem_tiled_copy_O, + tma_atom_O, + mbar_ptr, + softmax_scale_log2, + block_info, + num_splits, + SeqlenInfoCls, + TileSchedulerCls, + blocksparse_tensors, + ) + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_tmem_dealloc_offset) + + return + + @cute.jit + def load( + self, + thr_mma_qk: cute.core.ThrMma, + thr_mma_pv: cute.core.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: cute.CopyAtom, + tma_atom_K: Optional[cute.CopyAtom], + tma_atom_V: Optional[cute.CopyAtom], + pipeline_kv: cutlass.pipeline.PipelineAsync, + mbar_ptr: cute.Pointer, + block_info: BlockInfo, + num_splits: Int32, + SeqlenInfoCls: Callable, + TileSchedulerCls: Callable, + blocksparse_tensors: Optional[BlockSparseTensors], + ): + num_load_threads = len(self.load_warp_ids) * cute.arch.WARP_SIZE + tidx = cute.arch.thread_idx()[0] % num_load_threads + q_producer_phase = Int32(1) + kv_producer_state = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Producer, self.kv_stage + ) + tile_scheduler = TileSchedulerCls() + 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] + gQ = cute.local_tile(mQ_cur, cute.select(self.mma_tiler_qk, mode=[0, 2]), (None, 0)) + + 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) + ) + tSgQ = thr_mma_qk.partition_A(gQ) + tSgK = thr_mma_qk.partition_B(gK) + tOgV = thr_mma_pv.partition_B(gV) + load_Q_fn, _, _ = copy_utils.tma_get_copy_fn( + tma_atom_Q, 0, cute.make_layout(1), tSgQ, sQ + ) + + 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_Q = partial( + self.load_Q, + load_Q_fn, + mbar_ptr + self.mbar_load_q_full_offset, + mbar_ptr + self.mbar_load_q_empty_offset, + phase=q_producer_phase, + ) + # We have to use mbarrier directly in the load for KV instead of replying on + # pipeline_kv, because we could have different number of TMA bytes for K and V + load_K = partial( + self.load_KV, + tma_atom_K, + tKgK, + tKsK, + paged_kv_manager, + sK, + mbar_ptr + self.mbar_load_kv_full_offset, + mbar_ptr + self.mbar_load_kv_empty_offset, + K_or_V="K", + ) + load_V = partial( + self.load_KV, + tma_atom_V, + tVgV, + tVsV, + paged_kv_manager, + sV, + mbar_ptr + self.mbar_load_kv_full_offset, + mbar_ptr + self.mbar_load_kv_empty_offset, + 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: + if const_expr(self.use_tma_KV) or tidx < cute.arch.WARP_SIZE: + load_Q(block=self.q_stage * m_block + 0, stage=0) # Q0 + 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) + load_K(block=n_block_max - 1, producer_state=kv_producer_state, page_idx=page_idx) # K0 + kv_producer_state.advance() + if const_expr(self.q_stage == 2) and (const_expr(self.use_tma_KV) or tidx < cute.arch.WARP_SIZE): + load_Q(block=self.q_stage * m_block + 1, stage=1) # Q1 + q_producer_phase ^= 1 + 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) + 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: + kv_producer_state, q_producer_phase = produce_block_sparse_loads_sm100( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + 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, + ) + + + tile_scheduler.prefetch_next_work() + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + # End of persistent scheduler loop + + @cute.jit + def mma( + self, + tiled_mma_qk: cute.core.ThrMma, + tiled_mma_pv: cute.core.ThrMma, + sQ: cute.Tensor, + sK: cute.Tensor, + sV: cute.Tensor, + tStSs: Tuple[cute.Tensor, cute.Tensor], + tOtOs: tuple[cute.Tensor], + tOrPs: Tuple[cute.Tensor, cute.Tensor], + pipeline_kv: cutlass.pipeline.PipelineAsync, + mbar_ptr: cute.Pointer, + block_info: BlockInfo, + num_splits: Int32, + SeqlenInfoCls: Callable, + TileSchedulerCls: Callable, + blocksparse_tensors: Optional[BlockSparseTensors], + ): + 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 + + gemm_Si = [ + partial( + sm100_utils.gemm_ptx_partial, + qk_mma_op, + self.tmem_s_offset[stage], + tSrQs[stage], + sA=sQ[None, None, None, stage], + zero_init=True, + ) + for stage in range(self.q_stage) + ] + gemm_Pi = [ + partial( + sm100_utils.gemm_ptx_partial, + pv_mma_op, + self.tmem_o_offset[stage], + tOrPs[stage], + sA=None, + ) + for stage in range(self.q_stage) + ] + + mma_q_consumer_phase = Int32(0) + mma_kv_consumer_state = cutlass.pipeline.make_pipeline_state( + cutlass.pipeline.PipelineUserType.Consumer, self.kv_stage + ) + P_full_O_rescaled_phase = Int32(0) + + tile_scheduler = TileSchedulerCls() + 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, self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1) + 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: + for stage in cutlass.range_constexpr(self.q_stage): + # GEMM_QK00 (Q0 * K0 -> S0) or GEMM_QK01 (Q1 * K0 -> S1) + # 1. wait for Q0 / Q1 + cute.arch.mbarrier_wait( + mbar_ptr + self.mbar_load_q_full_offset + stage, mma_q_consumer_phase + ) + # 2. wait for K0 + if const_expr(stage == 0): + pipeline_kv.consumer_wait(mma_kv_consumer_state) + tSrKi = tSrK[None, None, None, mma_kv_consumer_state.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 + # tiled_mma_qk = sm100_utils.gemm(tiled_mma_qk, tStSs[stage], tSrQs[stage], tSrKi, zero_init=True) + sK_cur = sK[None, None, None, mma_kv_consumer_state.index] + if const_expr(self.uneven_kv_smem): + sK_cur = self.offset_kv_smem( + sK_cur, mma_kv_consumer_state.index, mma_kv_consumer_state.phase + ) + gemm_Si[stage](tCrB=tSrKi, sB=sK_cur) + # 4. release S0 / S1 + with cute.arch.elect_one(): + tcgen05.commit(mbar_ptr + self.mbar_S_full_offset + 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.q_stage): + # 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 has finished. + cute.arch.mbarrier_wait( + mbar_ptr + self.mbar_P_full_O_rescaled_offset + 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, + zero_init=not O_should_accumulate, + mbar_ptr=mbar_ptr + self.mbar_P_full_2_offset + stage, + mbar_phase=P_full_O_rescaled_phase, + ) + # 4. release accumulated O0_partial / O1_partial + # Don't need to signal O_full to the correction warps anymore 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. + # with cute.arch.elect_one(): + # tcgen05.commit(mbar_ptr + self.mbar_O_full_offset + stage) + # 5. release V(i-1) + if const_expr(stage == self.q_stage - 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. + # tiled_mma_qk = sm100_utils.gemm(tiled_mma_qk, tStSs[stage], tSrQs[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) + # 3. release S0 + with cute.arch.elect_one(): + tcgen05.commit(mbar_ptr + self.mbar_S_full_offset + 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 + with cute.arch.elect_one(): + for stage in cutlass.range_constexpr(self.q_stage): + tcgen05.commit(mbar_ptr + self.mbar_load_q_empty_offset + 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.q_stage): + # 2. acquire corrected Oi_partial and Pi + cute.arch.mbarrier_wait( + mbar_ptr + self.mbar_P_full_O_rescaled_offset + 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, + zero_init=not O_should_accumulate, + mbar_ptr=mbar_ptr + self.mbar_P_full_2_offset + stage, + 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 compute + # the row sum of the current tile. It does not guarantee that the 1st tile + # of the next work tile has been computed yet. + with cute.arch.elect_one(): + tcgen05.commit(mbar_ptr + self.mbar_O_full_offset + 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 + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + # End of persistent scheduler loop + + + # for both softmax0 and softmax1 warp group + @cute.jit + def softmax_loop( + self, + stage: int | Int32, + softmax_scale_log2: Float32, + softmax_scale: Float32, + thr_mma_qk: cute.core.ThrMma, + tStSi: cute.Tensor, + sScale: cute.Tensor, + mLSE: Optional[cute.Tensor], + learnable_sink: Optional[cute.Tensor], + mbar_ptr: cute.Pointer, + block_info: BlockInfo, + num_splits: Int32, + SeqlenInfoCls: Callable, + AttentionMaskCls: Callable, + TileSchedulerCls: Callable, + aux_tensors: Optional[list] = None, + fastdiv_mods=(None, None), + head_divmod=None, + blocksparse_tensors: Optional[BlockSparseTensors] = 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)) + ) + + tStScale = cute.composition(tStSi, cute.make_layout((self.m_block_size, 1))) + tScS = thr_mma_qk.partition_C(cute.make_identity_tensor(self.mma_tiler_qk[:2])) + tScScale = cute.composition(tScS, cute.make_layout((self.m_block_size, 1))) + + tilePlikeFP32 = self.mma_tiler_qk[1] // 32 * self.v_dtype.width + tStP_layout = cute.composition( + tStSi.layout, cute.make_layout((self.m_block_size, tilePlikeFP32)) + ) + tStP = cute.make_tensor(tStSi.iterator + self.tmem_s_to_p_offset, tStP_layout) + + tmem_load_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), + Float32, + ) + thr_tmem_load = tcgen05.make_tmem_copy(tmem_load_atom, tStSi).get_slice(tidx) + tStS_t2r = thr_tmem_load.partition_S(tStSi) + + tmem_store_scale_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(1)), + Float32, + ) + thr_tmem_store_scale = tcgen05.make_tmem_copy(tmem_store_scale_atom, tStScale).get_slice( + tidx + ) + + tStScale_r2t = thr_tmem_store_scale.partition_D(tStScale) + tmem_store_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(16)), + Float32, + ) + thr_tmem_store = tcgen05.make_tmem_copy(tmem_store_atom, tStP).get_slice(tidx) + tStP_r2t = thr_tmem_store.partition_D(tStP) + + mma_si_consumer_phase = Int32(0) + si_corr_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 + mbar_s0_s1_sequence_offset = self.mbar_s0_s1_sequence_offset + warp_idx_in_wg + + tile_scheduler = TileSchedulerCls() + 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) + + mask = AttentionMaskCls(seqlen) + shared_mask_kwargs = dict( + m_block=self.q_stage * m_block + stage, + 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_tensors=aux_tensors, + ) + + # 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, + **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, + **shared_mask_kwargs, + ) + else: + mask_fn_none = None + + softmax = SoftmaxSm100.create( + softmax_scale_log2, + rescale_threshold=8.0 if const_expr(self.q_dtype.width == 16) else 0.0, + softmax_scale=softmax_scale, + ) + softmax.reset() + + if const_expr(self.use_block_sparsity): + tile_block_count = get_total_block_count(blocksparse_tensors, batch_idx, head_idx, m_block, self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1) + 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, + mbar_ptr=mbar_ptr, + mbar_s0_s1_sequence_offset=mbar_s0_s1_sequence_offset, + thr_mma_qk=thr_mma_qk, + thr_tmem_load=thr_tmem_load, + thr_tmem_store=thr_tmem_store, + thr_tmem_store_scale=thr_tmem_store_scale, + tStS_t2r=tStS_t2r, + tStScale_r2t=tStScale_r2t, + tStP_r2t=tStP_r2t, + sScale=sScale, + stage=stage, + batch_idx=batch_idx, + head_idx=head_idx, + m_block=self.q_stage * m_block + stage, + seqlen=seqlen, + aux_tensors=aux_tensors, + fastdiv_mods=fastdiv_mods, + head_divmod=head_divmod, + ) + + if has_work: + # Softmax acts as the producer: wait until correction signals the stage is empty + cute.arch.mbarrier_wait( + mbar_ptr + self.mbar_softmax_corr_empty_offset + stage, si_corr_producer_phase + ) + si_corr_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 + stage + 1) * self.m_block_size + check_m_boundary = m_tile_end > seqlen.seqlen_q + else: + check_m_boundary = False + ( + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + empty_tile, + ) = softmax_block_sparse_sm100( + blocksparse_tensors, + batch_idx, + head_idx, + m_block, + softmax_step, + mask_fn, + mask_fn_none, + mma_si_consumer_phase, + si_corr_producer_phase, + s0_s1_sequence_phase, + mbar_ptr, + self.mbar_softmax_corr_full_offset, + self.mbar_softmax_corr_empty_offset, + self.mbar_P_full_O_rescaled_offset, + self.mbar_P_full_2_offset, + self.q_stage, + Int32(stage), + check_m_boundary, + self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1, + ) + if not empty_tile: + sScale[tidx + stage * self.m_block_size] = softmax.row_sum[0] + if const_expr(mLSE is not None or learnable_sink is not None): + sScale[ + tidx + stage * self.m_block_size + self.m_block_size * 2 + ] = 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]) + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_softmax_corr_full_offset + stage) + # 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, si_corr_producer_phase, s0_s1_sequence_phase = softmax_step( + mma_si_consumer_phase, + si_corr_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, 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), + ) + ) + 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, 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), + ) + 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, + ) + # 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, 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), + ) + ) + # Now that we no longer already have the 1st iteration, need mask_seqlen=True here + + # Dense path always writes scale / signals + sScale[tidx + stage * self.m_block_size] = softmax.row_sum[0] + if const_expr(mLSE is not None or learnable_sink is not None): + sScale[ + tidx + stage * self.m_block_size + self.m_block_size * 2 + ] = softmax.row_max[0] + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_softmax_corr_full_offset + stage) + + # # 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 + utils.log2f(softmax.row_sum[0])) * 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 + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + # End of persistent scheduler loop + + @cute.jit + def softmax_step( + self, + mma_si_consumer_phase: Int32, + si_corr_producer_phase: Int32, + s0_s1_sequence_phase: Int32, + n_block: Int32, + softmax: SoftmaxSm100, + mbar_ptr: cute.Pointer, + mbar_s0_s1_sequence_offset: Int32, + thr_mma_qk: cute.core.ThrMma, + thr_tmem_load: cute.CopyAtom, + thr_tmem_store: cute.CopyAtom, + thr_tmem_store_scale: cute.CopyAtom, + tStS_t2r: cute.Tensor, + tStScale_r2t: cute.Tensor, + tStP_r2t: cute.Tensor, + sScale: cute.Tensor, + stage: int | Int32, + batch_idx: Int32, + head_idx: Int32, + m_block: Int32, + seqlen, + aux_tensors: Optional[list] = None, + 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 + """ + 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])) + tScScale = cute.composition(tScS, cute.make_layout((self.m_block_size, 1))) + tScP = cute.composition(tScS, cute.make_layout((self.m_block_size, tilePlikeFP32))) + + # Wait for Si + cute.arch.mbarrier_wait(mbar_ptr + self.mbar_S_full_offset + stage, mma_si_consumer_phase) + tSrS_t2r = cute.make_fragment(thr_tmem_load.partition_D(tScS).shape, self.qk_acc_dtype) + cute.copy(thr_tmem_load, tStS_t2r, tSrS_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_tensors, + 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): + # tSrScale_r2t = cute.make_fragment(thr_tmem_store_scale.partition_S(tScScale).shape, Float32) + # tSrScale_r2t[0] = acc_scale + # cute.copy(thr_tmem_store_scale, tSrScale_r2t, tStScale_r2t) + # cute.arch.fence_view_async_tmem_store() + thread_idx = thr_tmem_load.thr_idx + sScale[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 + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_softmax_corr_full_offset + stage) + + # if thread_idx == 0 and stage == 0: cute.print_tensor(tSrS_t2r) + # print(tSrS_t2r) + softmax.scale_subtract_rowmax(tSrS_t2r, row_max) + # Sequence barrier wait + if const_expr(self.s0_s1_barrier): + cute.arch.mbarrier_wait( + mbar_ptr + mbar_s0_s1_sequence_offset + stage * 4, s0_s1_sequence_phase + ) + tSrP_r2t_f32 = cute.make_fragment(thr_tmem_store.partition_S(tScP).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, + e2e=mask_fn is None and self.head_dim_padded <= 128, + e2e_freq=self.e2e_freq, + ) + # Sequence barrier arrive + if const_expr(self.s0_s1_barrier): + cute.arch.mbarrier_arrive(mbar_ptr + mbar_s0_s1_sequence_offset + (1 - stage) * 4) + # print(tSrP_r2t_f32, tStP_r2t) + # cute.copy(thr_tmem_store, tSrP_r2t_f32, tStP_r2t) + for i in cutlass.range_constexpr(cute.size(tStP_r2t.shape[2]) // 4 * 3): + cute.copy(thr_tmem_store, tSrP_r2t_f32[None, None, i], tStP_r2t[None, None, i]) + cute.arch.fence_view_async_tmem_store() + # Notify mma warp that P is ready + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_P_full_O_rescaled_offset + stage) + for i in cutlass.range_constexpr( + cute.size(tStP_r2t.shape[2]) // 4 * 3, cute.size(tStP_r2t.shape[2]) + ): + cute.copy(thr_tmem_store, tSrP_r2t_f32[None, None, i], tStP_r2t[None, None, i]) + cute.arch.fence_view_async_tmem_store() + # Notify mma warp that the 2nd half of P is ready + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_P_full_2_offset + stage) + cute.arch.mbarrier_wait( + mbar_ptr + self.mbar_softmax_corr_empty_offset + stage, si_corr_producer_phase + ) + softmax.update_row_sum(tSrS_t2r.load(), acc_scale, is_first) + # acc_scale = cute.arch.exp2(acc_scale_) + return mma_si_consumer_phase ^ 1, si_corr_producer_phase ^ 1, s0_s1_sequence_phase ^ 1 + + @cute.jit + def correction_loop( + self, + thr_mma_qk: cute.core.ThrMma, + thr_mma_pv: cute.core.ThrMma, + tStS: cute.Tensor, + tOtOs: tuple[cute.Tensor], + sScale: cute.Tensor, + mO: cute.Tensor, + mLSE: cute.Tensor, + sO: cute.Tensor, + learnable_sink: Optional[cute.Tensor], + gmem_tiled_copy_O: cute.TiledCopy, + tma_atom_O: cute.CopyAtom, + mbar_ptr: cute.Pointer, + softmax_scale_log2: Float32, + block_info: BlockInfo, + num_splits: Int32, + SeqlenInfoCls: Callable, + TileSchedulerCls: Callable, + blocksparse_tensors: Optional[BlockSparseTensors] = None, + ): + tidx = cute.arch.thread_idx()[0] % (cute.arch.WARP_SIZE * len(self.correction_warp_ids)) + tScS = thr_mma_qk.partition_C(cute.make_identity_tensor(self.mma_tiler_qk[:2])) + tStScale_layout = cute.composition(tStS.layout, cute.make_layout((self.m_block_size, 1))) + tStScales = tuple( + cute.make_tensor(tStS.iterator + self.tmem_vec_offset[stage], tStScale_layout) + for stage in range(self.q_stage) + ) + tScScale = cute.composition(tScS, cute.make_layout((self.m_block_size, 1))) + tmem_load_v_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(1)), + self.qk_acc_dtype, + ) + thr_tmem_load_vec = tcgen05.make_tmem_copy(tmem_load_v_atom, tStScales[0]).get_slice(tidx) + + tStScales_t2r = [thr_tmem_load_vec.partition_S(tStScales[stage]) for stage in range(self.q_stage)] + tSrScale_t2r_shape = thr_tmem_load_vec.partition_D(tScScale).shape + + # First iter: no correction is required + for stage in cutlass.range_constexpr(self.q_stage): + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_P_full_O_rescaled_offset + stage) + + softmax_corr_consumer_phase = Int32(0) + o_corr_consumer_phase = Int32(0) + corr_epi_producer_phase = Int32(1) + + tile_scheduler = TileSchedulerCls() + 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) + + 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 = cute.local_tile(mO_cur, (self.m_block_size, self.head_dim_v_padded), (None, 0)) + + # 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, self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1) + 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 + cute.arch.mbarrier_wait( + mbar_ptr + self.mbar_softmax_corr_full_offset + 0, softmax_corr_consumer_phase + ) + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_softmax_corr_empty_offset + 0) + if const_expr(self.q_stage == 2): + cute.arch.mbarrier_wait( + mbar_ptr + self.mbar_softmax_corr_full_offset + 1, softmax_corr_consumer_phase + ) + softmax_corr_consumer_phase ^= 1 + + tSrScale_t2r = cute.make_fragment(tSrScale_t2r_shape, Float32) + for i in cutlass.range(total_block_count - 1, unroll=1): + for stage in cutlass.range_constexpr(self.q_stage): + # wait for S0 / S1 + cute.arch.mbarrier_wait( + mbar_ptr + self.mbar_softmax_corr_full_offset + stage, + softmax_corr_consumer_phase, + ) + # cute.copy(tiled_tmem_load_vec, tStScales_t2r[stage], tSrScale_t2r) + # cute.arch.fence_view_async_tmem_load() + # scale = tSrScale_t2r[0] + scale = sScale[tidx + stage * self.m_block_size] + should_rescale = cute.arch.vote_ballot_sync(scale < 1.0) != 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. + # cute.arch.mbarrier_wait(mbar_ptr + self.mbar_O_full_offset + stage, o_corr_consumer_phase) + if should_rescale: + self.correction_rescale( + thr_mma_pv, tOtOs[stage], tidx, scale + ) + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_P_full_O_rescaled_offset + stage) + if const_expr(self.q_stage == 2): + cute.arch.mbarrier_arrive( + mbar_ptr + self.mbar_softmax_corr_empty_offset + (1 - stage) + ) + else: + cute.arch.mbarrier_arrive( + mbar_ptr + self.mbar_softmax_corr_empty_offset + stage + ) + softmax_corr_consumer_phase ^= 1 + # o_corr_consumer_phase ^= 1 + if const_expr(self.q_stage == 2): + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_softmax_corr_empty_offset + 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 = ( + (self.q_stage * m_block + stage) * self.m_block_size + 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): + cute.arch.mbarrier_wait( + mbar_ptr + self.mbar_softmax_corr_full_offset + stage, + softmax_corr_consumer_phase, + ) + # 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[tidx + stage * self.m_block_size] + if const_expr(mLSE is not None or learnable_sink is not None): + row_max = sScale[tidx + stage * self.m_block_size + self.m_block_size * 2] + else: + row_max = None + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_softmax_corr_empty_offset + 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) + row_sum = Float32(1.0) + else: + row_sum += utils.exp2f( + sink_val * LOG2_E - row_max * softmax_scale_log2 + ) + 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) + cute.arch.mbarrier_wait( + mbar_ptr + self.mbar_O_full_offset + stage, o_corr_consumer_phase + ) + if const_expr(not self.use_correction_warps_for_epi): + cute.arch.mbarrier_wait( + mbar_ptr + self.mbar_corr_epi_empty_offset + stage, corr_epi_producer_phase + ) + self.correction_epilogue( + thr_mma_pv, + tOtOs[stage], + tidx, + stage, + m_block, + seqlen.seqlen_q, + scale, + sO[None, None, stage], + mO_cur, + gO, + gmem_tiled_copy_O, + ) + if const_expr(not self.use_correction_warps_for_epi): + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_corr_epi_full_offset + stage) + # Signal for the next work tile that O buffers in tmem are already read, so + # mma warp can write to them + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_P_full_O_rescaled_offset + stage) + # if tidx == 0: cute.printf("Correction final scale for stage %d: %f\n", stage, scale) + + o_corr_consumer_phase ^= 1 + softmax_corr_consumer_phase ^= 1 + corr_epi_producer_phase ^= 1 + else: + # WARNING: we need some code before the const_expr, see https://github.com/NVIDIA/cutlass/issues/2781 + if const_expr(self.use_correction_warps_for_epi): + gmem_tiled_copy_O_for_empty_tile = gmem_tiled_copy_O + else: + gmem_tiled_copy_O_for_empty_tile = None + if const_expr(self.use_block_sparsity): + ( + softmax_corr_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, + sScale, + stats, + self.correction_epilogue, + thr_mma_pv, + tOtOs, + sO, + mbar_ptr, + self.mbar_softmax_corr_full_offset, + self.mbar_softmax_corr_empty_offset, + self.mbar_P_full_O_rescaled_offset, + self.mbar_P_full_2_offset, + self.mbar_corr_epi_full_offset, + self.mbar_corr_epi_empty_offset, + softmax_corr_consumer_phase, + o_corr_consumer_phase, + corr_epi_producer_phase, + softmax_scale_log2, + 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): + gLSE = cute.local_tile( + mLSE_cur, (self.m_block_size,), (self.q_stage * m_block + stage,) + ) + 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 + utils.log2f(row_sum)) * 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 tidx < seqlen_q - (self.q_stage * m_block + stage) * self.m_block_size: + # This actually just works with PackGQA too + gLSE[tidx] = lse + + # Advance to next tile + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + # End of persistent scheduler loop + + @cute.jit + def correction_rescale( + self, + thr_mma: cute.core.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_fragment((tOrO_t2r_shape, frg_count), self.pv_acc_dtype) + for i in cutlass.range_constexpr(frg_count): + tOrO_frg = cute.make_fragment(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] = utils.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.core.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, + ): + """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.core.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 = 32 * 8 // self.o_dtype.width + tOsO = thr_mma.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=False, + ) + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_i[(None, None), 0]).get_slice( + tidx + ) + thr_tmem_load = tiled_tmem_load.get_slice(tidx) + 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 = thr_tmem_load.partition_D(tOsO_i[(None, None), None]) + tOcO_t2r = thr_tmem_load.partition_D(tOcO_i[(None, None), None]) + for i in cutlass.range_constexpr(self.head_dim_v_padded // corr_tile_size): + tOtO_t2r_i = tOtO_t2r[None, 0, 0, i] + tOsO_r2s_i = tOsO_s2r[None, 0, 0, i] + tOrO_frg = cute.make_fragment(tOcO_t2r[None, 0, 0, i].shape, self.pv_acc_dtype) + cute.copy(tiled_tmem_load, tOtO_t2r_i, tOrO_frg) + for j in cutlass.range_constexpr(0, cute.size(tOrO_frg), 2): + tOrO_frg[j], tOrO_frg[j + 1] = utils.mul_packed_f32x2( + (tOrO_frg[j], tOrO_frg[j + 1]), + (scale, scale), + ) + tOrO_frg_cvt = cute.make_fragment(tOrO_frg.shape, self.o_dtype) + tOrO_frg_cvt.store(tOrO_frg.load().to(self.o_dtype)) + cute.copy(tiled_smem_store, tOrO_frg_cvt, tOsO_r2s_i) + # fence view async shared + cute.arch.fence_proxy( + cute.arch.ProxyKind.async_shared, + space=cute.arch.SharedSpace.shared_cta, + ) + + 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(NamedBarrierFwd.Epilogue), + number_of_threads=len(self.epilogue_warp_ids) * cute.arch.WARP_SIZE) + gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx) + tOsO = gmem_thr_copy_O.partition_S(sO) + cO = cute.make_identity_tensor((self.m_block_size, self.head_dim_v_padded)) + 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_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): + for rest_m in cutlass.range_constexpr(cute.size(tOrO.shape[1])): + if ( + t0OcO[0, rest_m, 0][0] + < seqlen_q + - (self.q_stage * m_block + stage) * self.m_block_size + - tOcO[0][0] + ): + cute.copy( + gmem_tiled_copy_O, + tOrO[None, rest_m, None], + tOgO[None, rest_m, None, self.q_stage * m_block + stage], + 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, + self.q_stage * m_block + stage, + 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], + mbar_ptr: cute.Pointer, + block_info: BlockInfo, + num_splits: int, + SeqlenInfoCls: Callable, + TileSchedulerCls: Callable, + ): + epi_consumer_phase = Int32(0) + tile_scheduler = TileSchedulerCls() + 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) + + if const_expr(not self.is_split_kv) or n_block_min < n_block_max: + 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 = cute.local_tile(mO_cur, (self.m_block_size, self.head_dim_v_padded), (None, 0)) + 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_constexpr(self.q_stage): + # wait from corr, issue tma store on smem + # 1. wait for O0 / O1 final + cute.arch.mbarrier_wait( + mbar_ptr + self.mbar_corr_epi_full_offset + stage, epi_consumer_phase + ) + # 2. copy O0 / O1 to gmem + store_O(src_idx=stage, dst_idx=self.q_stage * m_block + 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 + if const_expr(self.q_stage == 2): + cute.arch.cp_async_bulk_wait_group(1 - stage, read=True) + else: + cute.arch.cp_async_bulk_wait_group(0, read=True) + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_corr_epi_empty_offset + stage) + else: + tidx = cute.arch.thread_idx()[0] % ( + cute.arch.WARP_SIZE * len(self.epilogue_warp_ids) + ) + gmem_thr_copy_O = gmem_tiled_copy_O.get_slice(tidx) + tOsO = gmem_thr_copy_O.partition_S(sO) + cO = cute.make_identity_tensor((self.m_block_size, self.head_dim_v_padded)) + 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]) + pack_gqa = PackGQA( + self.m_block_size, + self.head_dim_v_padded, + self.check_hdim_v_oob, + self.qhead_per_kvhead, + ) + for stage in cutlass.range_constexpr(self.q_stage): + # wait from corr, issue tma store on smem + # 1. wait for O0 / O1 final + cute.arch.mbarrier_wait( + mbar_ptr + self.mbar_corr_epi_full_offset + stage, epi_consumer_phase + ) + # 2. copy O0 / O1 to gmem + # load acc O from smem to rmem for wider vectorization + tOrO = cute.make_fragment_like(tOsO[None, None, None, 0], self.o_dtype) + cute.autovec_copy(tOsO[None, None, None, stage], tOrO) + # copy acc O from rmem to gmem + if const_expr(not self.pack_gqa): + for rest_m in cutlass.range_constexpr(cute.size(tOrO.shape[1])): + if ( + t0OcO[0, rest_m, 0][0] + < seqlen.seqlen_q + - (self.q_stage * m_block + stage) * self.m_block_size + - tOcO[0][0] + ): + cute.copy( + gmem_tiled_copy_O, + tOrO[None, rest_m, None], + tOgO[None, rest_m, None, self.q_stage * m_block + stage], + 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, + self.q_stage * m_block + stage, + seqlen.seqlen_q, + ) + cute.arch.mbarrier_arrive(mbar_ptr + self.mbar_corr_epi_empty_offset + stage) + + epi_consumer_phase ^= 1 + + # Advance to next tile + tile_scheduler.advance_to_next_work() + work_tile = tile_scheduler.get_current_work() + + def load_Q( + self, + load_Q_fn: Callable, + mbar_full_ptr: cute.Pointer, + mbar_empty_ptr: cute.Pointer, + block: Int32, + stage: int, + phase: Int32, + ): + cute.arch.mbarrier_wait(mbar_empty_ptr + stage, phase) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(mbar_full_ptr + stage, self.tma_copy_bytes["Q"]) + load_Q_fn(src_idx=block, dst_idx=stage, tma_bar_ptr=mbar_full_ptr + 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, + mbar_full_ptr: cute.Pointer, + mbar_empty_ptr: cute.Pointer, + block: Int32, + producer_state: cutlass.pipeline.PipelineState, + K_or_V: Literal["K", "V"], + page_idx: Optional[Int32] = None, + ): + assert K_or_V in ("K", "V") + stage, phase = producer_state.index, producer_state.phase + cute.arch.mbarrier_wait(mbar_empty_ptr + stage, phase) + 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: + cute.arch.mbarrier_wait(mbar_empty_ptr + 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 + ) + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx( + mbar_full_ptr + stage, self.tma_copy_bytes[K_or_V], + ) + 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=mbar_full_ptr + stage) + else: + assert paged_kv_manager is not None + paged_kv_manager.load_KV(block, sX[None, None, None, stage], K_or_V) + cute.arch.cp_async_commit_group() + cute.arch.cp_async_mbarrier_arrive_noinc(mbar_full_ptr + 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) + return cute.make_tensor(sX.iterator + offset, sX.layout) + else: + return sX + + def make_and_init_load_kv_pipeline(self, load_kv_mbar_ptr): + load_kv_consumer_group = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, len([self.mma_warp_id]) + ) + if self.use_tma_KV: + load_kv_producer_group = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, len(self.load_warp_ids) + ) + return cutlass.pipeline.PipelineTmaUmma.create( + barrier_storage=load_kv_mbar_ptr, + num_stages=self.kv_stage, + producer_group=load_kv_producer_group, + consumer_group=load_kv_consumer_group, + tx_count=self.tma_copy_bytes["K"], + ) + else: + load_kv_producer_group = cutlass.pipeline.CooperativeGroup( + cutlass.pipeline.Agent.Thread, len(self.load_warp_ids) * cute.arch.WARP_SIZE + ) + return cutlass.pipeline.PipelineAsyncUmma.create( + num_stages=self.kv_stage, + producer_group=load_kv_producer_group, + consumer_group=load_kv_consumer_group, + barrier_storage=load_kv_mbar_ptr, + ) + + # @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(NamedBarrierFwd.WarpSchedulerWG1), number_of_threads=2 * 128, + # ) + + # def warp_scheduler_barrier_sync(self): + # cute.arch.barrier( + # barrier_id=int(NamedBarrierFwd.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(NamedBarrierFwd.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_tensors=None, + 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_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_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.vec_size, + self.qk_acc_dtype, + aux_tensors, + 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/python/sglang/jit_kernel/flash_attention/cute/hopper_helpers.py b/python/sglang/jit_kernel/flash_attention/cute/hopper_helpers.py new file mode 100644 index 000000000..c6a1c3019 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/hopper_helpers.py @@ -0,0 +1,101 @@ +# Copyright (c) 2025, Tri Dao. +from typing import Type, Union, Optional +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Float32, Boolean, const_expr +from cutlass.cute.nvgpu import warpgroup +from cutlass.cutlass_dsl import Numeric, dsl_user_op +from cutlass.utils import LayoutEnum +import cutlass.utils.hopper_helpers as sm90_utils_og + + +@cute.jit +def gemm( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + zero_init: cutlass.Constexpr[bool] = False, + wg_wait: cutlass.Constexpr[int] = 0, + # A_in_regs: cutlass.Constexpr[bool] = False, + swap_AB: cutlass.Constexpr[bool] = False, +) -> None: + if const_expr(swap_AB): + gemm(tiled_mma, acc, tCrB, tCrA, zero_init=zero_init, wg_wait=wg_wait, swap_AB=False) + else: + warpgroup.fence() + # We make a new mma_atom since we'll be modifying its attribute (accumulate). + # Otherwise the compiler complains "operand #0 does not dominate this use" + mma_atom = cute.make_mma_atom(tiled_mma.op) + mma_atom.set(warpgroup.Field.ACCUMULATE, not zero_init) + for k in cutlass.range_constexpr(cute.size(tCrA.shape[2])): + cute.gemm(mma_atom, acc, tCrA[None, None, k], tCrB[None, None, k], acc) + mma_atom.set(warpgroup.Field.ACCUMULATE, True) + warpgroup.commit_group() + if const_expr(wg_wait >= 0): + warpgroup.wait_group(wg_wait) + + +def gemm_zero_init( + tiled_mma: cute.TiledMma, + shape: cute.Shape, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + A_idx: Optional[Int32] = None, + B_idx: Optional[Int32] = None, + wg_wait: int = -1, + swap_AB: bool = False, +) -> cute.Tensor: + if const_expr(swap_AB): + return gemm_zero_init( + tiled_mma, shape[::-1], tCrB, tCrA, B_idx, A_idx, wg_wait, swap_AB=False + ) + else: + acc = cute.make_fragment(tiled_mma.partition_shape_C(shape), Float32) + 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] + gemm(tiled_mma, acc, rA, rB, zero_init=True, wg_wait=wg_wait) + return acc + + +def gemm_w_idx( + tiled_mma: cute.TiledMma, + acc: cute.Tensor, + tCrA: cute.Tensor, + tCrB: cute.Tensor, + zero_init: Boolean, + A_idx: Optional[Int32] = None, + B_idx: Optional[Int32] = None, + wg_wait: int = -1, + swap_AB: bool = False, +) -> None: + if const_expr(swap_AB): + gemm_w_idx(tiled_mma, acc, tCrB, tCrA, zero_init, B_idx, A_idx, wg_wait, 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] + gemm(tiled_mma, acc, rA, rB, zero_init=zero_init, wg_wait=wg_wait) + + +@dsl_user_op +def make_smem_layout( + dtype: Type[Numeric], + layout: LayoutEnum, + shape: cute.Shape, + stage: Optional[int] = None, + *, + loc=None, + ip=None, +) -> Union[cute.Layout, cute.ComposedLayout]: + major_mode_size = shape[1] if layout.is_n_major_c() else shape[0] + smem_layout_atom = warpgroup.make_smem_layout_atom( + sm90_utils_og.get_smem_layout_atom(layout, dtype, major_mode_size), + dtype, + ) + order = (1, 0, 2) if const_expr(layout.is_m_major_c()) else (0, 1, 2) + smem_layout_staged = cute.tile_to_shape( + smem_layout_atom, + cute.append(shape, stage) if const_expr(stage is not None) else shape, + order=order if const_expr(stage is not None) else order[:2], + ) + return smem_layout_staged diff --git a/python/sglang/jit_kernel/flash_attention/cute/interface.py b/python/sglang/jit_kernel/flash_attention/cute/interface.py new file mode 100644 index 000000000..244ddc967 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/interface.py @@ -0,0 +1,1758 @@ +# 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. + +# Supported features: +# - BF16 & FP16 dtype +# - noncausal & causal attention +# - MHA, GQA, MQA +# - hdim 64, 96, 128. +# - (hdim_qk, hdim_v) = (192, 128) for Blackwell (i.e. DeepSeek shape) +# - varlen +# - sliding window +# - bwd pass for Ampere (will also run on Hopper/Blackwell, but will be slow) + +# Features not supported yet: +# - split (i.e. FlashDecoding) +# - tuned block sizes +# - paged KV +# - append KV to existing KV cache +# - FP8 +# - bwd pass optimized for Hopper/Blackwell + +import math +from functools import lru_cache +from typing import Optional, Tuple, Callable + +import torch + + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute + +import sglang.jit_kernel.flash_attention.cute.utils as utils +from .cute_dsl_utils import to_cute_tensor +from .flash_fwd import FlashAttentionForwardSm90 +from .flash_fwd_sm100 import FlashAttentionForwardSm100 +from .flash_bwd_preprocess import FlashAttentionBackwardPreprocess +from .flash_bwd import FlashAttentionBackwardSm80 +from .flash_bwd_sm90 import FlashAttentionBackwardSm90 +from .flash_bwd_sm100 import FlashAttentionBackwardSm100 +from .flash_bwd_postprocess import FlashAttentionBackwardPostprocess +from .flash_fwd_combine import FlashAttentionForwardCombine + +from .block_sparsity import ( + BlockSparseTensorsTorch, + to_cute_block_sparse_tensors, + normalize_block_sparse_tensors, + get_block_sparse_expected_shapes, + get_block_sparse_expected_shapes_bwd, + get_block_sparse_broadcast_pattern, +) + +@lru_cache(maxsize=None) +def _get_device_capability(): + """Cached device capability check.""" + return torch.cuda.get_device_capability()[0] + +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}" + 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, +} + + +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 + + # 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 _flash_attn_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + 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, + 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, + # m_block_size: int = 128, + # n_block_size: int = 64, + # num_threads: int = 128, + m_block_size: int = 128, + n_block_size: int = 128, + num_threads: int = 384, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + _compute_capability: 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, +) -> Tuple[torch.Tensor, 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 + 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. + """ + q, k, v = [maybe_contiguous(t) for t in (q, k, v)] + 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 = k.shape[:2] + seqlen_k = num_pages * page_size + else: + num_pages, page_size = None, None + seqlen_k = k.shape[-3] + num_head_kv = k.shape[-2] + head_dim_v = v.shape[-1] + if cu_seqlens_k is None: + if page_table 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 == (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.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 q.dtype in [torch.float16, torch.bfloat16], "inputs must be float16 or bfloat16" + assert q.dtype == k.dtype == v.dtype, "inputs must have the same 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" + + assert all( + t is None or t.is_cuda + for t in ( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + page_table, + learnable_sink, + ) + ), "inputs must be on CUDA device" + assert num_head % num_head_kv == 0, "num_head must be divisible by num_head_kv" + assert head_dim <= 256, "head_dim must be less than or equal to 256" + alignment = 16 // q.element_size() + assert head_dim % alignment == 0, f"head_dim must be divisible by {alignment}" + assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {alignment}" + if softmax_scale is None: + softmax_scale = 1.0 / math.sqrt(head_dim) + 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 + + out_torch_dtype = q.dtype + device = q.device + q_batch_seqlen_shape = (batch_size, seqlen_q) if cu_seqlens_q is None else (total_q,) + lse_shape = (batch_size, num_head, seqlen_q) if cu_seqlens_q is None else (num_head, total_q) + requires_grad = q.requires_grad or k.requires_grad or v.requires_grad + + 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) + + dtype = torch2cute_dtype_map[q.dtype] + compute_capability = ( + _get_device_capability() + if _compute_capability is None + else _compute_capability + ) + + assert compute_capability in [9, 10, 11], "Unsupported compute capability. Supported: 9.x, 10.x, 11.x" + + use_block_sparsity = block_sparse_tensors is not None + + if mask_mod is None: + if causal: + window_size_right = 0 + local = window_size_left is not None or window_size_right is not 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: + causal, local = False, False + + current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + if compute_capability == 9: # TODO: tune block size according to hdim. + if head_dim == head_dim_v == 128 and not causal and not local and not use_block_sparsity: + n_block_size = 192 + + if compute_capability in [10, 11]: + if ( + pack_gqa + and (128 % qhead_per_kvhead != 0) + ): + pack_gqa = False + # TODO: fix GQA + SplitKV + non-varlen + if pack_gqa and num_splits != 1 and cu_seqlens_q is None: + pack_gqa = False + + 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 + seqlen_q_packgqa = max_seqlen_q * qhead_per_kvhead + if compute_capability == 10: + q_stage = 2 if seqlen_q_packgqa > m_block_size else 1 + else: + q_stage = 1 + + if num_splits < 1: + m_block_size_effective = q_stage * m_block_size + seqlen_k_loaded = max_seqlen_k if not local else max(0, min(max_seqlen_k, window_size_right + window_size_left + 1 + m_block_size)) + num_n_blocks = (seqlen_k_loaded + n_block_size - 1) // n_block_size + 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_splits = num_splits_heuristic( + total_mblocks, + torch.cuda.get_device_properties(device).multi_processor_count, + num_n_blocks, + 128, + ) + + 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) + + # 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 + + 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) + + 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 + ) + + if mask_mod is not None: + if is_varlen: + raise NotImplementedError( + "mask_mod with aux_tensors is not yet supported for varlen sequences. This will be fixed in a future PR." + ) + + if use_block_sparsity: + if is_varlen: + raise NotImplementedError( + "Block sparsity is not yet supported for varlen sequences. This will be fixed in a future PR." + ) + # NB: pack_gqa requires block sparse head dim == 1 (broadcasted) + if pack_gqa and block_sparse_tensors.mask_block_cnt.shape[1] != 1: + pack_gqa = False + if is_split_kv: + raise NotImplementedError( + "Block sparsity is not yet supported with SplitKV. TODO: partition sparse block lists per split." + ) + + # See get_broadcast_dims for why this is needed in compile key + block_sparse_broadcast_pattern = None + normalized_block_sparse_tensors = None + if block_sparse_tensors is not None: + if seqlen_q is None: + raise ValueError("Block sparsity requires fixed-length sequences (seqlen_q must be known).") + expected_count_shape, expected_index_shape = get_block_sparse_expected_shapes( + batch_size, num_head, seqlen_q, seqlen_k, + m_block_size, n_block_size, q_stage, + ) + normalized_block_sparse_tensors = normalize_block_sparse_tensors( + block_sparse_tensors, + expected_count_shape=expected_count_shape, + expected_index_shape=expected_index_shape, + ) + block_sparse_broadcast_pattern = get_block_sparse_broadcast_pattern( + normalized_block_sparse_tensors + ) + + 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, + len(aux_tensors) if aux_tensors is not None else 0, + 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, + m_block_size, + n_block_size, + q_stage, + num_threads, + is_split_kv, + pack_gqa, + compute_capability, + page_size not in [None, 128], # paged KV non-TMA + ) + 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) + elif lse is not None: + lse_tensor = to_cute_tensor(lse, assumed_align=4) + else: + lse_tensor = 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 + if aux_tensors is not None: + cute_aux_tensors = [to_cute_tensor(buf, assumed_align=None, fully_dynamic=True) for buf in aux_tensors] + + if compute_capability == 9: + assert page_table is None, "paged KV not supported on SM 9.0" + assert not is_split_kv, "SplitKV not supported on SM 9.0" + # fa_fwd = FlashAttentionForwardSm80( + fa_fwd = FlashAttentionForwardSm90( + dtype, + head_dim, + head_dim_v, + qhead_per_kvhead, + is_causal=causal, + is_local=local, + pack_gqa=pack_gqa, + tile_m=m_block_size, + tile_n=n_block_size, + # num_stages=1, + num_stages=2, + num_threads=num_threads, + Q_in_regs=False, + intra_wg_overlap=True, + mma_pv_is_rs=True, + mask_mod=mask_mod, + score_mod=score_mod, + has_aux_tensors=aux_tensors is not None, + ) + elif compute_capability in [10, 11]: + fa_fwd = FlashAttentionForwardSm100( + 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=m_block_size, + n_block_size=n_block_size, + 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, 128], + is_varlen_q=cu_seqlens_q is not None + or seqused_q is not None, + ) + else: + raise ValueError( + f"Unsupported compute capability: {compute_capability}. Supported: 9.x, 10.x, 11.x" + ) + # TODO: check @can_implement + _flash_attn_fwd.compile_cache[compile_key] = cute.compile( + fa_fwd, + q_tensor, + k_tensor, + v_tensor, + o_tensor, + lse_tensor, + softmax_scale, + current_stream, + 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, + sparse_tensors, + cute_aux_tensors, + options="--enable-tvm-ffi", + ) + + _flash_attn_fwd.compile_cache[compile_key]( + q, + k, + v, + out if not is_split_kv else out_partial, + lse_partial if is_split_kv else lse, + softmax_scale, + current_stream, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + page_table, + window_size_left, + window_size_right, + learnable_sink, + normalized_block_sparse_tensors, + aux_tensors, + ) + 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 + + +_flash_attn_fwd.compile_cache = {} + + +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, + block_sparse_tensors: Optional[BlockSparseTensorsTorch] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + compute_capability = _get_device_capability() + assert compute_capability in [9, 10, 11], "Unsupported compute capability. Supported: 9.x, 10.x, 11.x" + + if compute_capability == 9: + m_block_size = 80 if not causal else 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 = not causal + AtomLayoutMSdP = 1 + AtomLayoutNdKV = 2 + AtomLayoutMdQ = 1 + cluster_size = 1 + assert window_size_left is None and window_size_right is None, "local not supported yet on 9.x" + else: + m_block_size = 128 + n_block_size = 128 + dQ_swapAB = False + dKV_swapAB = False + AtomLayoutMdQ = 1 + AtomLayoutNdKV = 1 + # TODO: support cluster size 2 + cluster_size = 1 + 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) + ] + 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 + 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] + head_dim_v = v.shape[-1] + + if causal: + window_size_right = 0 + local = window_size_left is not None or window_size_right is not None + if local: + if window_size_left is None and window_size_right == 0: + causal, local = True, False + window_size_right = None + else: + causal, local = False, True + + use_block_sparsity = block_sparse_tensors is not None + + # SM90 block-sparse backward: tile_m=64 is the GCD between a m_block_size that fits, + # the base block_m of 128 from forward, and block-sparse size for subtiling. + if compute_capability == 9 and use_block_sparsity: + m_block_size = 64 + # dQ_swapAB tuning: use False when m_block_size=64 (same as causal case) + dQ_swapAB = False + + # NB: this could be derived from the block_sparse_tensors but for now we hardcode it to 2 + subtile_factor = 2 + sparse_block_size_q = subtile_factor * m_block_size + + 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 + + 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" + 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" + assert head_dim <= 256, "head_dim must be less than or equal to 256" + alignment = 16 // q.element_size() + assert head_dim % alignment == 0, f"head_dim must be divisible by {alignment}" + assert head_dim_v % alignment == 0, f"head_dim_v must be divisible by {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 compute_capability not in [10, 11]: + assert deterministic is False, "bwd deterministic only supported for sm100/sm110 for now" + + if score_mod is not None: + assert score_mod_bwd is not None, "score_mod_bwd is required when score_mod is provided" + assert softcap == 0.0, "softcap and score_mod are mutually exclusive (different log2 scaling)" + assert cu_seqlens_q is None and cu_seqlens_k is None, ( + "varlen + score_mod not supported in bwd yet" + ) + + 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 = 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 = 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) + + dKV_postprocess = qhead_per_kvhead > 1 + if dKV_postprocess: + head_dim_v_rounded = (head_dim_v + 32 - 1) // 32 * 32 + if cu_seqlens_k is None: + 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 + 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: + total_k_rounded_padded = ( + (total_k + cu_seqlens_k.shape[0] * n_block_size - 1) // n_block_size * n_block_size + ) + num_n_blocks = total_k_rounded_padded // n_block_size + if cluster_size == 2 and num_n_blocks % cluster_size != 0: + total_k_rounded_padded = total_k_rounded_padded + n_block_size + 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 = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + if deterministic: + dQ_semaphore = torch.zeros(batch_size, num_head, seqlen_q_rounded // m_block_size, 1, dtype=torch.int32, device="cuda") + 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="cuda") + dV_semaphore = torch.zeros(batch_size, num_head_kv, seqlen_k_rounded // n_block_size, 2, dtype=torch.int32, device="cuda") + else: + dK_semaphore = None + dV_semaphore = None + + # Preprocess kernel: compute (o * dout).sum(dim=-1), lse * log2_e, and zero out dq_accum. + compile_key_pre = ( + compute_capability, + dtype, + head_dim_v, + m_block_size, + num_threads, + cu_seqlens_q is None, + seqused_q is None, + ) + if compile_key_pre not in _flash_attn_bwd.compile_cache_pre: + o_tensor, do_tensor = [to_cute_tensor(t) for t in (out, dout)] + dq_accum_tensor, dpsum_tensor, lse_log2_tensor = [ + to_cute_tensor(t) for t in (dq_accum, dpsum, lse_log2) + ] + lse_tensor = to_cute_tensor(lse, assumed_align=4) + cu_seqlens_q_tensor, seqused_q_tensor = [ + to_cute_tensor(t, assumed_align=4) if t is not None else None + for t in (cu_seqlens_q, seqused_q) + ] + arch = compute_capability * 10 + fa_bwd_pre = FlashAttentionBackwardPreprocess( + dtype, + head_dim_v, + arch, + m_block_size, + num_threads=num_threads, + ) + # TODO: check @can_implement + _flash_attn_bwd.compile_cache_pre[compile_key_pre] = cute.compile( + fa_bwd_pre, + o_tensor, + do_tensor, + dpsum_tensor, + lse_tensor, + lse_log2_tensor, + dq_accum_tensor, + cu_seqlens_q_tensor, + seqused_q_tensor, + current_stream, + options="--enable-tvm-ffi", + ) + _flash_attn_bwd.compile_cache_pre[compile_key_pre]( + out, + dout, + dpsum, + lse, + lse_log2, + dq_accum, + cu_seqlens_q, + seqused_q, + current_stream, + ) + + # NB num_threads application for 3 kernels + # There are pre, main, post processing kernels, currenlty num_threads is only actually + # used for the pre proc, and then we hard code to 384 for the main and post proc, and we do + # before cache key gen + 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 + 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: + 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, subtile_factor, + ) + normalized_block_sparse_tensors = normalize_block_sparse_tensors( + block_sparse_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, and optionally full_q_cnt/full_q_idx). " + f"Regenerate the backward BlockMask with BLOCK_SIZE=({sparse_block_size_q}, {n_block_size}) " + f"(sparse_block_size_q={sparse_block_size_q})." + ), + ) + block_sparse_broadcast_pattern = get_block_sparse_broadcast_pattern( + normalized_block_sparse_tensors + ) + + if compute_capability == 9: + compile_key = ( + compute_capability, + dtype, + head_dim, + head_dim_v, + qhead_per_kvhead, + causal, + softcap != 0.0, + 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, + 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, + use_block_sparsity, + block_sparse_broadcast_pattern, + ) + else: + compile_key = ( + compute_capability, + dtype, + head_dim, + head_dim_v, + qhead_per_kvhead, + causal, + window_size_left is not None, + window_size_right is not None, + softcap != 0.0, + m_block_size, + n_block_size, + num_threads, + pack_gqa, + cluster_size, + deterministic, + score_mod_hash, + score_mod_bwd_hash, + mask_mod_hash, + num_aux_tensors, + 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, + ) + 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) + ] + dq_accum_tensor, dpsum_tensor, lse_log2_tensor = [ + to_cute_tensor(t) for t in (dq_accum, dpsum, lse_log2) + ] + 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) + ] + fa_bwd_sm80 = FlashAttentionBackwardSm80( + 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, + ) + if compute_capability == 9: + fa_bwd_obj = FlashAttentionBackwardSm90( + dtype, + head_dim, + head_dim_v, + qhead_per_kvhead, + causal, + m_block_size, + n_block_size, + num_stages_Q, + num_stages_dO, + num_stages_PdS, + SdP_swapAB, + dKV_swapAB, + dQ_swapAB, + AtomLayoutMSdP, + AtomLayoutNdKV, + AtomLayoutMdQ, + 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, + subtile_factor=subtile_factor, + ) + 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, + # cluster_size=1, + deterministic=deterministic, + score_mod=score_mod, + score_mod_bwd=score_mod_bwd, + mask_mod=mask_mod, + has_aux_tensors=aux_tensors is not None, + subtile_factor=subtile_factor, + ) + + # Block sparse tensors for backward use Q-direction indexing (transposed from forward). + # sparse_block_size_q = subtile_factor * tile_m matches BlockMask granularity. + sparse_tensors_compile = None + if normalized_block_sparse_tensors is not None: + sparse_tensors_compile = to_cute_block_sparse_tensors(normalized_block_sparse_tensors) + + # 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, + current_stream, + cu_seqlens_q_tensor, + cu_seqlens_k_tensor, + seqused_q_tensor, + seqused_k_tensor, + None, # softcap - not yet supported in backward + window_size_left, + window_size_right, + dQ_semaphore_tensor, + dK_semaphore_tensor, + dV_semaphore_tensor, + cute_aux_tensors, + sparse_tensors_compile, + options="--enable-tvm-ffi", + ) + _flash_attn_bwd.compile_cache[compile_key]( + q, + k, + v, + dout, + lse_log2, + dpsum, + dq_accum, + dk if not dKV_postprocess else dk_accum, + dv if not dKV_postprocess else dv_accum, + softmax_scale, + current_stream, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + None, # softcap - not yet supported in backward + window_size_left, + window_size_right, + dQ_semaphore, + dK_semaphore, + dV_semaphore, + aux_tensors, + normalized_block_sparse_tensors, + ) + + num_threads = 256 if compute_capability == 9 else 128 + arch = compute_capability * 10 + # Postprocess kernel: convert dq_accum from float32 to dq in bf16/fp16 + compile_key_post = ( + compute_capability, + dtype, + head_dim, + m_block_size, + num_threads, + AtomLayoutMdQ, + dQ_swapAB, + cu_seqlens_q is None, + seqused_q is None, + ) + if compile_key_post not in _flash_attn_bwd.compile_cache_post: + dq_accum_tensor = to_cute_tensor(dq_accum) + dq_tensor = to_cute_tensor(dq) + cu_seqlens_q_tensor, seqused_q_tensor = [ + to_cute_tensor(t, assumed_align=4) if t is not None else None + for t in (cu_seqlens_q, seqused_q) + ] + fa_bwd_post = FlashAttentionBackwardPostprocess( + dtype, head_dim, arch, m_block_size, num_threads, AtomLayoutMdQ, dQ_swapAB + ) + # TODO: check @can_implement + _flash_attn_bwd.compile_cache_post[compile_key_post] = cute.compile( + fa_bwd_post, + dq_accum_tensor, + dq_tensor, + softmax_scale, + cu_seqlens_q_tensor, + seqused_q_tensor, + current_stream, + options="--enable-tvm-ffi", + ) + _flash_attn_bwd.compile_cache_post[compile_key_post]( + dq_accum, + dq, + softmax_scale, + cu_seqlens_q, + seqused_q, + current_stream, + ) + + if dKV_postprocess: + # Postprocess kernel: convert dk_accum & dv_accum from float32 to bf16/fp16 + compile_key_post = ( + compute_capability, + dtype, + head_dim, + n_block_size, + num_threads, + AtomLayoutNdKV, + dKV_swapAB, + cu_seqlens_k is None, + seqused_k is None, + ) + if compile_key_post not in _flash_attn_bwd.compile_cache_post: + dk_accum_tensor = to_cute_tensor(dk_accum) + dk_tensor = to_cute_tensor(dk) + cu_seqlens_k_tensor, seqused_k_tensor = [ + to_cute_tensor(t, assumed_align=4) if t is not None else None + for t in (cu_seqlens_k, seqused_k) + ] + arch = compute_capability * 10 + fa_bwd_post = FlashAttentionBackwardPostprocess( + dtype, head_dim, arch, n_block_size, num_threads, AtomLayoutNdKV, dKV_swapAB + ) + # TODO: check @can_implement + _flash_attn_bwd.compile_cache_post[compile_key_post] = cute.compile( + fa_bwd_post, + dk_accum_tensor, + dk_tensor, + softmax_scale, + cu_seqlens_k_tensor, + seqused_k_tensor, + current_stream, + options="--enable-tvm-ffi", + ) + _flash_attn_bwd.compile_cache_post[compile_key_post]( + dk_accum, + dk, + softmax_scale, + cu_seqlens_k, + seqused_k, + current_stream, + ) + compile_key_post = ( + compute_capability, + dtype, + head_dim_v, + n_block_size, + num_threads, + AtomLayoutNdKV, + dKV_swapAB, + cu_seqlens_k is None, + seqused_k is None, + ) + if compile_key_post not in _flash_attn_bwd.compile_cache_post: + dv_accum_tensor = to_cute_tensor(dv_accum) + dv_tensor = to_cute_tensor(dv) + cu_seqlens_k_tensor, seqused_k_tensor = [ + to_cute_tensor(t, assumed_align=4) if t is not None else None + for t in (cu_seqlens_k, seqused_k) + ] + arch = compute_capability * 10 + fa_bwd_post = FlashAttentionBackwardPostprocess( + dtype, head_dim_v, arch, n_block_size, num_threads, AtomLayoutNdKV, dKV_swapAB + ) + # TODO: check @can_implement + _flash_attn_bwd.compile_cache_post[compile_key_post] = cute.compile( + fa_bwd_post, + dv_accum_tensor, + dv_tensor, + cutlass.Float32(1.0), + cu_seqlens_k_tensor, + seqused_k_tensor, + current_stream, + options="--enable-tvm-ffi", + ) + _flash_attn_bwd.compile_cache_post[compile_key_post]( + dv_accum, + dv, + 1.0, + cu_seqlens_k, + seqused_k, + current_stream, + ) + + return dq, dk, dv + + +_flash_attn_bwd.compile_cache_pre = {} +_flash_attn_bwd.compile_cache = {} +_flash_attn_bwd.compile_cache_post = {} + + +class FlashAttnFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + 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, + mask_mod: Optional[Callable] = None, + full_block_cnt: Optional[torch.Tensor] = None, + full_block_idx: Optional[torch.Tensor] = None, + mask_block_cnt: Optional[torch.Tensor] = None, + mask_block_idx: Optional[torch.Tensor] = None, + ): + # Only create block sparse tensors if at least one block sparse parameter is provided + block_sparse_tensors = None + if any(t is not None for t in [full_block_cnt, full_block_idx, mask_block_cnt, mask_block_idx]): + block_sparse_tensors = BlockSparseTensorsTorch( + full_block_cnt=full_block_cnt, + full_block_idx=full_block_idx, + mask_block_cnt=mask_block_cnt, + mask_block_idx=mask_block_idx, + ) + out, lse = _flash_attn_fwd( + q, + k, + v, + 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, + mask_mod=mask_mod, + block_sparse_tensors=block_sparse_tensors + ) + ctx.save_for_backward(q, k, v, out, lse) + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.softcap = softcap + ctx.deterministic = deterministic + return out, lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, lse = ctx.saved_tensors + 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, + ) + return dq, dk, dv, *((None,) * 20) # Extra Nones is fine + + +class FlashAttnVarlenFunc(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor], + cu_seqlens_k: Optional[torch.Tensor], + seqused_q: Optional[torch.Tensor] = None, + seqused_k: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_k: Optional[int] = 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, + aux_tensors: Optional[list] = None, + ): + out, lse = _flash_attn_fwd( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_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, + aux_tensors=aux_tensors, + ) + ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k) + 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 + return out, lse + + @staticmethod + def backward(ctx, dout, *args): + q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k, seqused_q, seqused_k = ctx.saved_tensors + assert ctx.softcap == 0.0 + 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, + ) + + return dq, dk, dv, *((None,) * 20) + + +def flash_attn_func( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + 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, + mask_mod: Optional[Callable] = None, + full_block_cnt: Optional[torch.Tensor] = None, + full_block_idx: Optional[torch.Tensor] = None, + mask_block_cnt: Optional[torch.Tensor] = None, + mask_block_idx: Optional[torch.Tensor] = None, +): + return FlashAttnFunc.apply( + q, + k, + v, + softmax_scale, + causal, + window_size, + learnable_sink, + softcap, + num_splits, + pack_gqa, + deterministic, + mask_mod, + full_block_cnt, + full_block_idx, + mask_block_cnt, + mask_block_idx, + ) + + +def flash_attn_varlen_func( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + 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, + seqused_q: Optional[torch.Tensor] = None, + seqused_k: 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, + aux_tensors: Optional[list] = None, +): + return FlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + seqused_q, + seqused_k, + max_seqlen_q, + max_seqlen_k, + page_table, + softmax_scale, + causal, + window_size, + learnable_sink, + softcap, + num_splits, + pack_gqa, + deterministic, + score_mod, + aux_tensors, + ) + + +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, + 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 + """ + # Input validation + assert out_partial.dim() in [4, 5], "out_partial must have 4 or 5 dimensions" + assert lse_partial.dim() in [3, 4], "lse_partial must have 3 or 4 dimensions" + assert out_partial.dtype in [torch.float16, torch.bfloat16, torch.float32], ( + "out_partial must be fp16, bf16, or fp32" + ) + assert lse_partial.dtype == torch.float32, "lse_partial must be fp32" + assert out_partial.is_cuda and lse_partial.is_cuda, "tensors must be on CUDA device" + assert out_partial.stride(-1) == 1, "out_partial must be contiguous in the last dimension" + assert lse_partial.stride(-2) == 1, "lse_partial must be contiguous in the seqlen dimension" + assert lse_partial.shape == out_partial.shape[:-1] + + # Determine if this is variable length based on dimensions + is_varlen = out_partial.dim() == 4 + + # Validate output tensor shapes and types + assert out.shape == out_partial.shape[1:], "out shape mismatch" + if lse is not None: + assert lse.shape == lse_partial.shape[1:], "lse shape mismatch" + assert lse.dtype == torch.float32, "lse must be fp32" + + # 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: + assert t.dtype == torch.int32, f"{name} must be int32" + 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). + m_block_size = 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 m_block_size == 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) + + current_stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + + # 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, + m_block_size, + k_block_size, + log_max_splits, + cu_seqlens is not None, + seqused is not None, + lse is not None, + ) + + if compile_key not in _flash_attn_fwd_combine.compile_cache: + out_partial_tensor = to_cute_tensor( + out_partial, leading_dim=4 if not is_varlen else 3 + ) + lse_partial_tensor = to_cute_tensor( + lse_partial, assumed_align=4, leading_dim=lse_partial.ndim - 2 + ) + out_tensor = to_cute_tensor(out, leading_dim=3 if not is_varlen else 2) + lse_tensor = ( + to_cute_tensor(lse, assumed_align=4, leading_dim=lse.ndim - 2) + if lse is not None + else None + ) + + optional_tensors = [ + to_cute_tensor(t, assumed_align=4, leading_dim=0) + if t is not None + else None + for t in (cu_seqlens, seqused, num_splits_dynamic_ptr, semaphore_to_reset) + ] + cu_seqlens_tensor, seqused_tensor, num_splits_dynamic_tensor, semaphore_tensor = ( + optional_tensors + ) + fa_combine = FlashAttentionForwardCombine( + dtype=dtype, + dtype_partial=dtype_partial, + head_dim=head_dim, + m_block_size=m_block_size, + k_block_size=k_block_size, + log_max_splits=log_max_splits, + ) + + # Check if implementation is supported + if not fa_combine.can_implement( + dtype, + dtype_partial, + head_dim, + m_block_size, + k_block_size, + log_max_splits, + num_threads=256, + ): + raise RuntimeError( + "FlashAttention combine kernel cannot be implemented with given parameters" + ) + + _flash_attn_fwd_combine.compile_cache[compile_key] = cute.compile( + fa_combine, + out_partial_tensor, + lse_partial_tensor, + out_tensor, + lse_tensor, + cu_seqlens_tensor, + seqused_tensor, + num_splits_dynamic_tensor, + semaphore_tensor, + current_stream, + options="--enable-tvm-ffi", + ) + _flash_attn_fwd_combine.compile_cache[compile_key]( + out_partial, + lse_partial, + out, + lse, + cu_seqlens, + seqused, + num_splits_dynamic_ptr, + semaphore_to_reset, + current_stream, + ) + + +_flash_attn_fwd_combine.compile_cache = {} + + +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, + 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 + 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" + assert lse_partial.dim() in [3, 4], "lse_partial must have 3 or 4 dimensions" + assert out_partial.dtype == torch.float32, "out_partial must be fp32 (from accumulation)" + assert lse_partial.dtype == torch.float32, "lse_partial must be fp32" + + # 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 + assert lse_partial.shape == (num_splits, total_q, num_heads), ( + "lse_partial shape mismatch for varlen" + ) + 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 + assert lse_partial.shape == (num_splits, batch_size, seqlen, num_heads), ( + "lse_partial shape mismatch" + ) + + # 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).transpose( + 0, 1 + ) + else: + lse = torch.empty( + batch_size, num_heads, seqlen, dtype=torch.float32, device=device + ).transpose(1, 2) + else: + lse = None + + _flash_attn_fwd_combine( + out_partial, + lse_partial, + out, + lse, + cu_seqlens, + seqused, + ) + return out, lse diff --git a/python/sglang/jit_kernel/flash_attention/cute/mask.py b/python/sglang/jit_kernel/flash_attention/cute/mask.py new file mode 100644 index 000000000..170598244 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/mask.py @@ -0,0 +1,651 @@ +# Copyright (c) 2025, Tri Dao. + +from typing import Optional, Callable +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, const_expr + +import sglang.jit_kernel.flash_attention.cute.utils as utils +from .seqlen_info import SeqlenInfoQK + + +@cute.jit +def mask_r2p(X: cute.Tensor, col_limit: Int32, arch: int = 90, rank1: bool = False) -> None: + # Bit manipulation, compiles down to the R2P instruction + # For sm100: we know that tScS_t2r[i][1] == i, for the particular tmem copy atom we're using. + # For sm90: instead of comparing limit to 0, 1, 8, 9, 16, 17, ..., + # we compare a transformed version of limit to 0, 1, 2, 3, 4, 5, ... + if const_expr(arch == 90): + col_limit_transformed = col_limit // 8 * 2 + min(col_limit % 8, 2) + else: + col_limit_transformed = col_limit + ncol = const_expr(cute.size(X.shape[cute.rank(X) - 1]) if not rank1 else cute.size(X.shape)) + # Ideally we'd move by 32 instead of 24, but mask >> i isn't correct for i == 31 + for s in cutlass.range_constexpr(cute.ceil_div(ncol, 24)): + # Don't need to clamp to 32 since the shr.u32 instruction does that already + col_limit_right_s = max(col_limit_transformed - s * 24, 0) + # 0 -> 0b00...00, 1 -> 0b00...01, ..., 31 -> 0b01...11, 32 -> 0b11...11 + mask = (1 << col_limit_right_s) - 1 + # This needs to be range_constexpr, o/w the compiler can't generate the R2P instruction + for i in cutlass.range_constexpr(min(24, ncol - s * 24)): + in_bound = cutlass.Boolean(mask & (1 << i)) + c = s * 24 + i + if const_expr(rank1): + X[c] = X[c] if in_bound else -Float32.inf + # This is the equivalent of: + # X[s * 24 + i] = X[s * 24 + i] if col_limit_right_s <= i 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 mask_r2p_transposed(X: cute.Tensor, row_limit_top: Int32, num_rep: int) -> None: + # Bit manipulation, compiles down to the R2P instruction + # For sm100: we know that tScS_t2r[i][0] has the form 0, 1, ..., 31, 64, ..., 127 + # or 0, 1, ..., 15, 32, ..., 47, 64, ... + # We compare a transformed version of limit to 0, 1, 2, 3, 4, 5, ... + # Here we hardcode for the case of 2 warp groups. + num_wg = 2 + row_limit_top_transformed = row_limit_top // (num_rep * num_wg) * num_rep + min( + row_limit_top % (num_rep * num_wg), num_rep + ) + ncol = cute.size(X.shape) + # Ideally we'd move by 32 instead of 24, but mask >> i isn't correct for i == 31 + for s in cutlass.range_constexpr(cute.ceil_div(ncol, 24)): + row_limit_top_s = max(row_limit_top_transformed - s * 24, 0) + # 0 -> 0b00...00, 1 -> 0b00...01, ..., 31 -> 0b01...11, 32 -> 0b11...11 + mask = (1 << row_limit_top_s) - 1 + # This needs to be range_constexpr, o/w the compiler can't generate the R2P instruction + for i in cutlass.range_constexpr(min(24, ncol - s * 24)): + out_bound = cutlass.Boolean(mask & (1 << i)) + c = s * 24 + i + X[c] = -Float32.inf if out_bound else X[c] + # tidx = cute.arch.thread_idx()[0] % 256 + # if tidx == 128: + # cute.printf("tidx = {}, s = {}, i = {}, row_limit_top = {}, row_limit_top_s = {}, mask = {}, out_bound = {}", tidx, s, i, row_limit_top, row_limit_top_s, mask, out_bound) + + +@cute.jit +def mask_r2p_dual_bound( + X: cute.Tensor, + col_limit_left: Int32, # Inclusive lower bound + col_limit_right: Int32, # Exclusive upper bound +) -> None: + """ + Dual-bound masking using two bitmasks for SM100, following mask_r2p. + Masks elements where: NOT (col_limit_left <= col < col_limit_right) + + Uses bit manipulation to create a range mask: + mask_right = (1 << right) - 1 -> bits (right-1)..0 are 1 + mask_left = (1 << left) - 1 -> bits (left-1)..0 are 1 + mask_range = mask_range = mask_right & ~ mask_left -> bits (right-1)..left are 1 + """ + ncol = const_expr(cute.size(X.shape)) + + for s in cutlass.range_constexpr(cute.ceil_div(ncol, 24)): + right_s = max(col_limit_right - s * 24, 0) + left_s = max(col_limit_left - s * 24, 0) + + # otherwise cute dsl complains about python int too large to convert into c long + right_s = min(right_s, 24) + left_s = min(left_s, 24) + + # bits (right-1)..left are 1 + mask_right = (1 << right_s) - 1 + mask_left = (1 << left_s) - 1 + mask_range = mask_right & ~mask_left + + # This needs to be range_constexpr, o/w the compiler can't generate the R2P instruction + for i in cutlass.range_constexpr(min(24, ncol - s * 24)): + in_bound = cutlass.Boolean(mask_range & (1 << i)) + c = s * 24 + 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_tensors: Optional[list] = None, + fastdiv_mods=(None, None), + ) -> None: + assert not (mask_causal and mask_local), "mask_causal and mask_local cannot be both True" + acc_S_mn = utils.make_acc_tensor_mn_view(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 = utils.make_acc_tensor_mn_view(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 = utils.make_acc_tensor_mn_view( + 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): + # The compiler now choses not to use R2P + r2p = const_expr(False and 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: + mask_r2p(acc_S_mn, seqlenk_col_limit, arch=90) + + 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_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 = mask_mod( + batch_idx_ssa, + head_idx_ssa, + q_idx_ssa, + kv_idx_ssa, + self.seqlen_info, + aux_tensors, + ) + 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: + mask_r2p(acc_S_mn[r, None], col_limit_right, arch=90, 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 + ) + 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 cute.arch.thread_idx()[0] == 128: cute.printf("n_block = {}, r = {}, row_idx = {}, causal_row_offset = {}, col_limit_right = {}, col_limit_left = {}", n_block, r, row_idx, causal_row_offset, col_limit_right, col_limit_left) + # traverse column index. + for c in cutlass.range(cute.size(tScS_mn.shape[1]), unroll_full=True): + col_idx = t0ScS_mn[0, c][1] + # only consider the column index, so the row index sets to 0. + if col_idx >= col_limit_right or col_idx < col_limit_left: + acc_S_mn[r, c] = -Float32.inf + 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 + else col0 - causal_row_offset - self.window_size_right + ) + # TODO: do we need col_limit_sink? + row_limit_bot = col0 - causal_row_offset + self.window_size_left + 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_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_tensors: Optional[list] = None, + fastdiv_mods=(None, None), + head_divmod=None, + check_q_boundary: bool = False, + ) -> 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_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 + r2p = True + if 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(acc_S, seqlenk_col_limit, arch=100, rank1=True) + + elif const_expr(not mask_causal and not mask_local and mask_mod is not None): + # Block sparse case w/ mask_mod + 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_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_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 = mask_mod( + batch_idx_ssa, + head_idx_ssa, + mask_row_ssa, + kv_idx_ssa, + self.seqlen_info, + aux_tensors, + ) + 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] + + else: # Causal or local + causal_row_offset = 1 + 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 + 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(acc_S, col_limit_right, arch=100, rank1=True) + else: + 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 + ) + 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: + # XOR-based R2P dual bound masking + mask_r2p_dual_bound(acc_S, col_limit_left, col_limit_right) + + @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_tensors: Optional[list] = None, + 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 conventio: + - 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" + 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_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 = mask_mod( + batch_idx_ssa, + head_idx_ssa, + q_idx_ssa, + kv_idx_ssa, + self.seqlen_info, + aux_tensors, + ) + 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 + mask_r2p_transposed(acc_S, row_limit_top, num_rep) + 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 + 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] diff --git a/python/sglang/jit_kernel/flash_attention/cute/mma_sm100_desc.py b/python/sglang/jit_kernel/flash_attention/cute/mma_sm100_desc.py new file mode 100644 index 000000000..16336c346 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/mma_sm100_desc.py @@ -0,0 +1,291 @@ +# 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.FloatE4M3FN: + return MXF8F6F4Format.E4M3 + if cutlass_type is cutlass.FloatE5M2: + 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: + # No idea what the right way to get B, M, S is – so we're just parsing it from the __str__ + # Swizzle string has the form "S" + swz_str = str(swizzle) + inside = swz_str[swz_str.index("<") + 1 : swz_str.index(">")] # '3,4,3' + B, M, S = [int(x) for x in inside.split(",")] # [3, 4, 3] + + 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 diff --git a/python/sglang/jit_kernel/flash_attention/cute/named_barrier.py b/python/sglang/jit_kernel/flash_attention/cute/named_barrier.py new file mode 100644 index 000000000..777c44079 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/named_barrier.py @@ -0,0 +1,31 @@ +# 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 NamedBarrierBwd(enum.IntEnum): + Epilogue = enum.auto() + WarpSchedulerWG1 = enum.auto() + WarpSchedulerWG2 = enum.auto() + WarpSchedulerWG3 = enum.auto() + PdS = enum.auto() + dQFullWG0 = enum.auto() + dQFullWG1 = enum.auto() + dQEmptyWG0 = enum.auto() + dQEmptyWG1 = enum.auto() + + +class NamedBarrierBwdSm100(enum.IntEnum): + EpilogueWG1 = enum.auto() + EpilogueWG2 = enum.auto() + Compute = enum.auto() + dQaccReduce = enum.auto() diff --git a/python/sglang/jit_kernel/flash_attention/cute/pack_gqa.py b/python/sglang/jit_kernel/flash_attention/cute/pack_gqa.py new file mode 100644 index 000000000..02c89df6e --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/pack_gqa.py @@ -0,0 +1,164 @@ +# Copyright (c) 2025, Tri Dao. + + +import cutlass +import cutlass.cute as cute + +import sglang.jit_kernel.flash_attention.cute.utils as utils + + +class PackGQA: + def __init__( + self, + m_block_size: cutlass.Constexpr[int], + head_dim_padded: cutlass.Constexpr[int], + check_hdim_oob: cutlass.Constexpr[bool], + qhead_per_kvhead: cutlass.Constexpr[bool], + ): + self.m_block_size = m_block_size + self.head_dim_padded = head_dim_padded + self.check_hdim_oob = check_hdim_oob + self.qhead_per_kvhead = qhead_per_kvhead + + @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_fragment(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 = utils.make_acc_tensor_mn_view(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/python/sglang/jit_kernel/flash_attention/cute/paged_kv.py b/python/sglang/jit_kernel/flash_attention/cute/paged_kv.py new file mode 100644 index 000000000..d45824222 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/paged_kv.py @@ -0,0 +1,214 @@ +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 + +import sglang.jit_kernel.flash_attention.cute.utils as utils +from .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] + + 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], + ): + 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) + tVpV = utils.predicate_k(tVcV, limit=mV_paged.shape[0]) + + 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, + 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): + tPrXPtr = cute.make_rmem_tensor((self.page_entry_per_thread,), cutlass.Int64) + for i in cutlass.range(self.page_entry_per_thread, unroll=1): + page = self.tPrPage[i] + page_offset = self.tPrPageOffset[i] + if const_expr(K_or_V == "K"): + tPrXPtr[i] = utils.elem_pointer(self.mK_paged, (page_offset, 0, page)).toint() + else: + tPrXPtr[i] = utils.elem_pointer(self.mV_paged, (0, page_offset, page)).toint() + return tPrXPtr + + @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) + + # Finesse sX layout to be (M, N). + 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"): + # Need to transpose V + sX_pi = cute.make_tensor(sX_pi.iterator, cute.select(sX_pi.layout, mode=[1, 0])) + + 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,)) + + 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, + ) diff --git a/python/sglang/jit_kernel/flash_attention/cute/pipeline.py b/python/sglang/jit_kernel/flash_attention/cute/pipeline.py new file mode 100644 index 000000000..54981bca1 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/pipeline.py @@ -0,0 +1,272 @@ +# Copyright (c) 2025, Tri Dao. + +# import math +from typing import Optional +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cutlass import Boolean, Int32, const_expr +from cutlass.cutlass_dsl import if_generate +from cutlass.pipeline import PipelineAsync, PipelineState, Agent, CooperativeGroup +from cutlass.pipeline import PipelineUserType, PipelineOp +from cutlass.pipeline import PipelineTmaAsync as PipelineTmaAsyncOg +from cutlass.pipeline import PipelineTmaUmma as PipelineTmaUmmaOg + + +# We deviate from cute-dsl implementation to use cute.arch.cluster_arrive_relaxed +def pipeline_init_wait(cta_layout_vmnk: Optional[cute.Layout] = None): + """ + Fences the mbarrier init and syncs the threadblock or cluster + """ + cute.arch.mbarrier_init_fence() + + if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1: + # If not using clusters, sync the threadblock + _sync(Agent.ThreadBlock) + else: + # If using clusters, sync the cluster + _sync(Agent.ThreadBlockCluster) + + +def _sync(group: Agent): + """ + Syncs all threads within an agent. + """ + if group is Agent.Thread: + raise NotImplementedError("Error: Not supported.") + elif group is Agent.ThreadBlock: + cute.arch.sync_threads() + elif group is Agent.ThreadBlockCluster: + cute.arch.cluster_arrive_relaxed() + cute.arch.cluster_wait() + else: + assert False, ( + "Error: No explicit sync instruction exists. Please use barriers (named / mbarrier) instead." + ) + + +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): + # assert stages < 2**16 + # self._log_stages = int(math.log2(stages)) + # assert 1 << self._log_stages == stages, "Number of stages must be a power of 2." + 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 1 << self._log_stages + return self._stages + + @property + def index(self) -> Int32: + # return self._phase_index & 0xFFFF + # return self._phase_index & ((1 << self._log_stages) - 1) + if const_expr(self._stages == 1): + return Int32(0) + else: + return self._phase_index % self._stages + + @property + def phase(self) -> Int32: + # return self._phase_index >> 16 + # 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. + # return (self._phase_index >> self._log_stages) % 2 + # return self._phase_index >> self._log_stages + 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 then_body(phase_index): + # # XOR the phase bit and set the index to 0 + # return (phase_index & 0xFFFF0000) ^ (1 << 16) + + # def else_body(phase_index): + # return phase_index + + # self._phase_index = if_generate( + # (self._phase_index & 0xFFFF) == self.stages, + # then_body, + # else_body, + # [self._phase_index], + # [Int32], + # ) + + 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(1 << 16)) + 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." + + +@dataclass(frozen=True) +class PipelineTmaAsync(PipelineTmaAsyncOg): + """ + Override producer_acquire to take in extra_tx_count parameter. + """ + + @staticmethod + def create(*args, **kwargs): + obj = PipelineTmaAsyncOg.create(*args, **kwargs) + # Can't assign to __class__ directly since the dataclass is frozen + # obj.__class__ = PipelineTmaAsync + object.__setattr__(obj, "__class__", PipelineTmaAsync) + return obj + + def producer_acquire( + self, + state: PipelineState, + try_acquire_token: Optional[Boolean] = None, + extra_tx_count: int = 0, + ): + """ + 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), + ) + if const_expr(extra_tx_count == 0): + self.sync_object_full.arrive(state.index, self.producer_mask) + else: + tx_count = self.sync_object_full.tx_count + extra_tx_count + self.sync_object_full.arrive_and_expect_tx(state.index, tx_count) + + +@dataclass(frozen=True) +class PipelineTmaUmma(PipelineTmaUmmaOg): + @staticmethod + def create( + *, + num_stages: int, + producer_group: CooperativeGroup, + consumer_group: CooperativeGroup, + tx_count: int, + barrier_storage: cute.Pointer = None, + cta_layout_vmnk: Optional[cute.Layout] = None, + mcast_mode_mn: tuple[int, int] = (1, 1), + init_wait: cutlass.Constexpr[bool] = True, + ): + """ + This helper function computes any necessary attributes and returns an instance of PipelineTmaUmma. + :param barrier_storage: Pointer to the smem address for this pipeline's mbarriers + :type barrier_storage: cute.Pointer + :param num_stages: Number of buffer stages for this pipeline + :type num_stages: Int32 + :param producer_group: `CooperativeGroup` for the producer agent + :type producer_group: CooperativeGroup + :param consumer_group: `CooperativeGroup` for the consumer agent + :type consumer_group: CooperativeGroup + :param tx_count: Number of bytes expected to be written to the transaction barrier for one stage + :type tx_count: int + :param cta_layout_vmnk: Layout of the cluster shape + :type cta_layout_vmnk: cute.Layout | None + :param mcast_mode_mn: Tuple of two integers, specifying whether mcast is enabled for the m and n modes. At least one of the two integers must be 1. + :type mcast_mode_mn: tuple[int, int] + """ + if not isinstance(barrier_storage, cute.Pointer): + raise ValueError( + f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}" + ) + + producer_type = PipelineOp.TmaLoad + consumer_type = PipelineOp.TCGen05Mma + + producer = (producer_type, producer_group) + consumer = (consumer_type, consumer_group) + + sync_object_full = PipelineAsync._make_sync_object( + barrier_storage.align(min_align=8), num_stages, producer, tx_count + ) + sync_object_empty = PipelineAsync._make_sync_object( + barrier_storage.align(min_align=8) + num_stages, num_stages, consumer + ) + + if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1: + # No mcast mask if not using clusters + producer_mask = None + # All threadblocks are leaders if not using clusters + is_leader_cta = True + else: + producer_mask = PipelineTmaUmma._compute_mcast_arrival_mask( + cta_layout_vmnk, mcast_mode_mn + ) + is_leader_cta = PipelineTmaUmma._compute_is_leader_cta(cta_layout_vmnk) + + cta_group = ( + cute.nvgpu.tcgen05.CtaGroup.ONE + if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1 + else cute.nvgpu.tcgen05.CtaGroup.TWO + ) + + consumer_mask = producer_mask + + if const_expr(init_wait): + pipeline_init_wait(cta_layout_vmnk) + + return PipelineTmaUmma( + sync_object_full, + sync_object_empty, + num_stages, + producer_mask, + consumer_mask, + is_leader_cta, + cta_group, + ) + + def producer_acquire( + self, + state: PipelineState, + try_acquire_token: Optional[Boolean] = None, + extra_tx_count: int = 0, + ): + """ + 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), + ) + if const_expr(extra_tx_count == 0): + if_generate( + self.is_leader_cta, + lambda: self.sync_object_full.arrive(state.index, self.producer_mask), + ) + 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), + ) diff --git a/python/sglang/jit_kernel/flash_attention/cute/pyproject.toml b/python/sglang/jit_kernel/flash_attention/cute/pyproject.toml new file mode 100644 index 000000000..1503556c1 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/pyproject.toml @@ -0,0 +1,56 @@ +[build-system] +requires = ["setuptools"] +build-backend = "setuptools.build_meta" + +[project] +name = "flash-attn-cute" +version = "0.1.0" +description = "Flash Attention CUTE (CUDA Template Engine) implementation" +readme = "README.md" +requires-python = ">=3.10" +license = {text = "BSD 3-Clause License"} +authors = [ + {name = "Tri Dao"}, +] +classifiers = [ + "Development Status :: 3 - Alpha", + "License :: OSI Approved :: BSD License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", +] + +dependencies = [ + "nvidia-cutlass-dsl>=4.3.5,<4.4.0", + "torch", + "einops", + "typing_extensions", + "apache-tvm-ffi>=0.1.5,<0.2", + "torch-c-dlpack-ext", + "quack-kernels==0.2.4", +] + +[project.optional-dependencies] +dev = [ + "pytest", + "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.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 +] diff --git a/python/sglang/jit_kernel/flash_attention/cute/seqlen_info.py b/python/sglang/jit_kernel/flash_attention/cute/seqlen_info.py new file mode 100644 index 000000000..6d8c6feb2 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/seqlen_info.py @@ -0,0 +1,138 @@ +from typing import Optional +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, const_expr + +""" +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: cutlass.Int32 + seqlen: cutlass.Int32 + + @staticmethod + def create( + batch_idx: cutlass.Int32, + seqlen_static: cutlass.Int32, + cu_seqlens: Optional[cute.Tensor] = None, + seqused: Optional[cute.Tensor] = None, + ): + offset = 0 if const_expr(cu_seqlens is None) else cu_seqlens[batch_idx] + 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, seqlen) + + +@dataclass(frozen=True) +class SeqlenInfoQK: + offset_q: cutlass.Int32 + offset_k: cutlass.Int32 + padded_offset_q: cutlass.Int32 + padded_offset_k: cutlass.Int32 + seqlen_q: cutlass.Int32 + seqlen_k: cutlass.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: cutlass.Int32, + seqlen_q_static: cutlass.Int32, + seqlen_k_static: cutlass.Int32, + mCuSeqlensQ: Optional[cute.Tensor] = None, + mCuSeqlensK: Optional[cute.Tensor] = None, + mSeqUsedQ: Optional[cute.Tensor] = None, + mSeqUsedK: Optional[cute.Tensor] = None, + tile_m: cutlass.Constexpr[cutlass.Int32] = 128, + tile_n: cutlass.Constexpr[cutlass.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 (offset_q + batch_idx * tile_m) // tile_m * tile_m + ) + padded_offset_k = ( + 0 + if const_expr(mCuSeqlensK is None) + else (offset_k + batch_idx * tile_n) // tile_n * 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 + ) + has_cu_seqlens_q: int = mCuSeqlensQ is not None + has_cu_seqlens_k: int = mCuSeqlensK is not None + has_seqused_q: int = mSeqUsedQ is not None + has_seqused_k: int = mSeqUsedK is not None + return SeqlenInfoQK( + offset_q, + offset_k, + padded_offset_q, + padded_offset_k, + seqlen_q, + seqlen_k, + has_cu_seqlens_q, + has_cu_seqlens_k, + has_seqused_q, + has_seqused_k, + ) + + def offset_batch_Q( + self, + mQ: cute.Tensor, + batch_idx: Int32, + dim: int, + padded: cutlass.Constexpr[bool] = False, + ) -> cute.Tensor: + """Seqlen must be the first dimension of mQ""" + 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 = offset_q if const_expr(cute.rank(mQ.shape[0]) == 1) else (0, offset_q) + idx = (offset,) + (0,) * (cute.rank(mQ) - 1) + return cute.domain_offset(idx, mQ) + + def offset_batch_K( + self, + mK: cute.Tensor, + batch_idx: Int32, + dim: int, + padded: cutlass.Constexpr[bool] = False, + ) -> cute.Tensor: + """Seqlen must be the first dimension of mK""" + 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 + idx = (offset_k,) + (0,) * (cute.rank(mK) - 1) + return cute.domain_offset(idx, mK) diff --git a/python/sglang/jit_kernel/flash_attention/cute/softmax.py b/python/sglang/jit_kernel/flash_attention/cute/softmax.py new file mode 100644 index 000000000..378932a96 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/softmax.py @@ -0,0 +1,582 @@ +# 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 + +import sglang.jit_kernel.flash_attention.cute.utils as utils +from .cute_dsl_utils import ParamsBase +from .seqlen_info import SeqlenInfoQK + + +@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 = utils.make_acc_tensor_mn_view(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 = utils.warp_reduce(row_max_cur, cute.arch.fmax, width=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 = utils.exp2f(acc_S_row * scale_log2 - row_max_cur_scaled) + + 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 = utils.exp2f(acc_S_row * scale_log2 - row_max_cur_scaled) + # row_scale[r] = utils.exp2f(row_max_prev * self.scale_log2 - row_max_cur_scaled) + row_scale[r] = utils.exp2f((row_max_prev - row_max_cur) * scale_log2) + + 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] += utils.exp2f(sink_val_cur * LOG2_E - row_max[r] * scale_log2) + + # 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 + utils.log2f(row_sum_cur)) * 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 = utils.make_acc_tensor_mn_view(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 + + @staticmethod + def create( + scale_log2: Float32, + rescale_threshold: cutlass.Constexpr[float] = 0.0, + softmax_scale: Float32 | None = None, + ): + 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, + ) + + @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) + 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) + 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 = utils.exp2f(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 + + def update_row_sum( + self, acc_S_row_exp: cute.TensorSSA, row_scale: Float32, is_first: int = False + ) -> None: + 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=self.row_sum[0] * row_scale) + self.row_sum[0] = self._compute_row_sum(acc_S_row_exp, init_val=init_val) + # tmp = self._compute_row_sum(acc_S_row_exp) + # self.row_sum[0] = self.row_sum[0] * row_scale + tmp + + @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 + 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] = utils.fma_packed_f32x2( + (acc_S_row[i], acc_S_row[i + 1]), + (self.scale_log2, self.scale_log2), + (-row_max_scaled, -row_max_scaled), + ) + + @cute.jit + def apply_exp2_convert( + self, + acc_S_row: cute.Tensor, + acc_S_row_converted: cute.Tensor, + e2e: cutlass.Constexpr[bool] = False, + e2e_freq: cutlass.Constexpr[int] = 16, + e2e_res: cutlass.Constexpr[int] = 4, + e2e_frg_limit: cutlass.Constexpr[int] = 1, + ): + 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] = utils.exp2f(acc_S_row_frg[k, j]) + # acc_S_row_frg[k + 1, j] = utils.exp2f(acc_S_row_frg[k + 1, j]) + if cutlass.const_expr(not e2e): + acc_S_row_frg[k, j] = cute.arch.exp2(acc_S_row_frg[k, j]) + acc_S_row_frg[k + 1, j] = cute.arch.exp2(acc_S_row_frg[k + 1, j]) + else: + if cutlass.const_expr( + k % e2e_freq < e2e_freq - e2e_res or j >= frg_cnt - e2e_frg_limit + ): + acc_S_row_frg[k, j] = cute.arch.exp2(acc_S_row_frg[k, j]) + acc_S_row_frg[k + 1, j] = cute.arch.exp2(acc_S_row_frg[k + 1, j]) + 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] = utils.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] = utils.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.arch.exp2(acc_S_row[i]) + # acc_S_row[i + 1] = cute.arch.exp2(acc_S_row[i + 1]) + + 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] = ( + # utils.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] = utils.exp2f(acc_S_row_frg[k, j]) + # acc_S_row_frg[k + 1, j] = utils.exp2f(acc_S_row_frg[k + 1, j]) + acc_S_row_frg[k, j] = cute.arch.exp2(acc_S_row_frg[k, j]) + acc_S_row_frg[k + 1, j] = cute.arch.exp2(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 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_tensors, + 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 + 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 my 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_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,)) + + aux_args = [] + if cutlass.const_expr(aux_tensors is not None): + aux_args = aux_tensors + + post_mod_scores = score_mod( + score_ssa, + batch_idx_ssa, + head_idx_ssa, + q_idx=q_idx_ssa, + kv_idx=kv_idx_ssa, + seqlen_info=seqlen_info, + aux_tensors=aux_args, + ) + + # 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_tensors, + 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 + 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_fragment(vec_size, qk_acc_dtype) + score_vec = cute.make_fragment(vec_size, qk_acc_dtype) + kv_idx_vec = cute.make_fragment(vec_size, cutlass.Int32) + batch_idx_ssa = utils.scalar_to_ssa(batch_idx, cutlass.Int32).broadcast_to((vec_size,)) + q_idx_vec = cute.make_fragment(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_fragment(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_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,)) + + aux_args = [] + if cutlass.const_expr(aux_tensors is not None): + aux_args = aux_tensors + + grad_out_ssa = score_mod_bwd( + grad_ssa, + score_ssa, + batch_idx_ssa, + head_idx_ssa, + q_idx=q_idx_ssa, + kv_idx=kv_idx_ssa, + seqlen_info=seqlen_info, + aux_tensors=aux_args, + ) + + 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/python/sglang/jit_kernel/flash_attention/cute/testing.py b/python/sglang/jit_kernel/flash_attention/cute/testing.py new file mode 100644 index 000000000..2897e64fc --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/testing.py @@ -0,0 +1,423 @@ +import math +from typing import Optional + +import torch +import torch.nn.functional as F +from einops import rearrange, repeat + + +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) + indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + 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): + assert mode in ["full", "random", "third"] + if mode == "full": + lengths = torch.full((batch_size, 1), max_seqlen, device=device, dtype=torch.int32) + elif mode == "random": + lengths = torch.randint( + max(0 if zero_lengths else 1, max_seqlen - 20), + max_seqlen + 1, + (batch_size, 1), + device=device, + ) + else: + lengths = torch.randint( + max(0 if zero_lengths else 1, 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, +): + if causal: + window_size = (window_size[0], 0) + dtype_og = q.dtype + if upcast: + q, k, v = q.float(), k.float(), v.float() + qv = qv.float() if qv is not None else None + if q_descale is not None: + q_descale = repeat(q_descale, "b h -> b 1 (h g) 1", g=q.shape[2] // k.shape[2]) + q = (q.float() * q_descale).to(q.dtype) + qv = (qv.float() * q_descale).to(qv.dtype) if qv is not None else None + 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], k.shape[1] + k = repeat(k, "b s h d -> b s (h g) d", g=q.shape[2] // k.shape[2]) + v = repeat(v, "b s h d -> b s (h g) d", g=q.shape[2] // v.shape[2]) + d = q.shape[-1] + dv = v.shape[-1] + softmax_scale = 1.0 / math.sqrt(d if qv is None else d + dv) + if not reorder_ops: + scores = torch.einsum("bthd,bshd->bhts", q * softmax_scale, k) + else: + scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale) + if qv is not None: + scores = scores + torch.einsum("bthd,bshd->bhts", qv * softmax_scale, v) + 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=q.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=q.device, + ) + local_mask = ( + torch.logical_or(local_mask, chunk_mask) if local_mask is not None else chunk_mask + ) + if local_mask is not None: + scores.masked_fill_(local_mask, float("-inf")) + if attn_bias is not None: + scores = scores + attn_bias + 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 + ) + 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) + return output.to(dtype=dtype_og), attention.to(dtype=dtype_og) diff --git a/python/sglang/jit_kernel/flash_attention/cute/tile_scheduler.py b/python/sglang/jit_kernel/flash_attention/cute/tile_scheduler.py new file mode 100644 index 000000000..c11f18ca9 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/tile_scheduler.py @@ -0,0 +1,719 @@ +# Copyright (c) 2025, Tri Dao. + +from typing import Optional, Tuple +from dataclasses import dataclass, fields + +try: + from typing import override +except ImportError: # Python < 3.12 + from typing_extensions import override + +import cutlass +from cutlass._mlir import ir +import cutlass.cute as cute +from cutlass import Int32, const_expr + +import sglang.jit_kernel.flash_attention.cute.utils as utils +from .fast_math import clz +from cutlass.cute import FastDivmodDivisor + + +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) + + +@dataclass +class ParamsBase: + def __extract_mlir_values__(self): + all_fields = [getattr(self, field.name) for field in fields(self)] + non_constexpr_fields = [f for f in all_fields if not isinstance(f, cutlass.Constexpr)] + values, self._values_pos = [], [] + for obj in non_constexpr_fields: + 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): + all_fields = {field.name: getattr(self, field.name) for field in fields(self)} + constexpr_fields = {n: f for n, f in all_fields.items() if isinstance(f, cutlass.Constexpr)} + non_constexpr_fields = { + n: f for n, f in all_fields.items() if not isinstance(f, cutlass.Constexpr) + } + for (name, field), n_items in zip(non_constexpr_fields.items(), self._values_pos): + non_constexpr_fields[name] = cutlass.new_from_mlir_values(field, values[:n_items]) + values = values[n_items:] + return self.__class__(**non_constexpr_fields, **constexpr_fields) + + +@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 + + +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) + + @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, + ) + + 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, *, loc=None, ip=None) -> Params: + return SingleTileScheduler.Params.create(args, loc=loc, ip=ip) + + @staticmethod + def create(params: Params, *, loc=None, ip=None) -> "SingleTileScheduler": + blk_coord = cute.arch.block_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" + return ( + cute.round_up(params.num_block, params.cluster_shape_mn[0]), + 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 + + 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_divmod: FastDivmodDivisor + num_head_divmod: FastDivmodDivisor + total_blocks: Int32 + + @staticmethod + def create( + args: TileSchedulerArguments, *, loc=None, ip=None + ) -> "StaticPersistentTileScheduler.Params": + total_blocks = args.num_block * args.num_head * args.num_batch + return StaticPersistentTileScheduler.Params( + FastDivmodDivisor(args.num_block), FastDivmodDivisor(args.num_head), total_blocks + ) + + 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, *, loc=None, ip=None) -> Params: + return StaticPersistentTileScheduler.Params.create(args, loc=loc, ip=ip) + + @staticmethod + def create(params: Params, *, loc=None, ip=None) -> "StaticPersistentTileScheduler": + tile_idx = cute.arch.block_idx()[0] + return StaticPersistentTileScheduler(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]: + hardware_info = cutlass.utils.HardwareInfo() + sm_count = hardware_info.get_device_multiprocessor_count() + return (cutlass.min(sm_count, params.total_blocks), Int32(1), Int32(1)) + + # @cute.jit + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + hn_idx, block_idx = divmod(self._tile_idx, self.params.num_block_divmod) + batch_idx, head_idx = divmod(hn_idx, self.params.num_head_divmod) + is_valid = self._tile_idx < self.params.total_blocks + # if cute.arch.thread_idx()[0] == 0: + # cute.printf("TileScheduler: tile_idx=%d, hn_idx=%d, block_idx=%d, batch_idx=%d, head_idx=%d, is_valid=%d", self._tile_idx, hn_idx, block_idx, batch_idx, head_idx, is_valid) + 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): + self._tile_idx += cute.arch.grid_dim()[0] + + 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 + l2_minor: Int32 + num_block_divmod: FastDivmodDivisor + num_head_divmod: FastDivmodDivisor + l2_minor_divmod: FastDivmodDivisor + l2_major_divmod: FastDivmodDivisor + l2_minor_residual_divmod: FastDivmodDivisor + num_hb_quotient: Int32 + is_split_kv: cutlass.Constexpr[bool] = False + + @staticmethod + @cute.jit + def create( + args: TileSchedulerArguments, *, loc=None, ip=None + ) -> "SingleTileLPTScheduler.Params": + # cute.printf(args.num_block, args.num_head, args.num_batch, args.seqlen_k, args.headdim, args.headdim_v, args.total_q, args.tile_shape_mn, args.qhead_per_kvhead_packgqa, args.element_size) + size_one_kv_head = 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 + # swizzle = 1 if size_l2 < size_one_head else (size_l2 // size_one_head) + # Seems faster if swizzle if a power of 2 + log2_floor = lambda n: 31 - clz(n) + swizzle = 1 if size_l2 < size_one_head else (1 << log2_floor(size_l2 // size_one_head)) + # swizzle = 1 if size_l2 < size_one_head else (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, + l2_minor=Int32(swizzle), + num_block_divmod=FastDivmodDivisor(args.num_block), + 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) + ), # don't divide by 0 + num_hb_quotient=Int32(num_hb_quotient), + num_splits=args.num_splits, + is_split_kv=args.is_split_kv, + ) + + def __init__(self, params: Params, tile_idx: Int32, split_idx: Int32, *, loc=None, ip=None): + self.params = params + self._tile_idx = tile_idx + self._split_idx = split_idx + self._loc = loc + self._ip = ip + + @staticmethod + def to_underlying_arguments(args: TileSchedulerArguments, *, loc=None, ip=None) -> Params: + return SingleTileLPTScheduler.Params.create(args, loc=loc, ip=ip) + + @staticmethod + @cute.jit + def create(params: Params, *, loc=None, ip=None) -> "SingleTileLPTScheduler": + tile_idx, split_idx, _ = cute.arch.block_idx() + return SingleTileLPTScheduler(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]: + return (params.total_blocks, params.num_splits, Int32(1)) + + @cute.jit + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + 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 + 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 + ) + + 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 + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in [self.params, self._tile_idx, self._split_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._split_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 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 + size_one_qdo_head = args.seqlen_k * (args.headdim + args.headdim_v) * args.element_size + # size_one_dqaccum_head = 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(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, *, loc=None, ip=None) -> Params: + 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) + is_valid = self._tile_idx < params.total_blocks + bidx_in_cluster = cute.arch.block_in_cluster_idx() + block = block * params.cluster_shape_mn[0] + bidx_in_cluster[0] + if cutlass.const_expr(params.spt): + block = params.num_block - 1 - block + 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 + + 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 + + @staticmethod + @cute.jit + def create( + args: TileSchedulerArguments, *, loc=None, ip=None + ) -> "SingleTileVarlenScheduler.Params": + size_l2 = 50 * 1024 * 1024 # 50 MB for K & V + max_kvblock_in_l2 = size_l2 // ( + (args.headdim + args.headdim_v) * args.element_size * args.tile_shape_mn[1] + ) + assert args.mCuSeqlensQ is not None or args.mSeqUsedQ is not None, ( + "At least one of mCuSeqlensQ or mSeqUsedQ must be provided" + ) + 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, + ) + + def __init__(self, params: Params, tile_idx: Int32, split_idx: Int32, *, loc=None, ip=None): + self.params = params + self._tile_idx = tile_idx + self._split_idx = split_idx + self._is_first_block = True + self._loc = loc + self._ip = ip + + @staticmethod + def to_underlying_arguments(args: TileSchedulerArguments, *, loc=None, ip=None) -> Params: + return SingleTileVarlenScheduler.Params.create(args, loc=loc, ip=ip) + + @staticmethod + def create(params: Params, *, loc=None, ip=None) -> "SingleTileVarlenScheduler": + 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.tile_shape_mn[0] - 1) + ) // params.tile_shape_mn[0] + 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(seqlen, params.tile_shape_mn[0]) + if batch_idx < params.num_batch and lane < cute.arch.WARP_SIZE - 1 + else Int32(0) + ) + + @cute.jit + def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo: + 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 + 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.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 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) + + 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._is_first_block = False + + def __extract_mlir_values__(self): + values, self._values_pos = [], [] + for obj in [self.params, self._tile_idx, self._split_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._split_idx], + self._values_pos, + ): + obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items])) + values = values[n_items:] + return SingleTileVarlenScheduler(*(tuple(obj_list)), loc=self._loc) diff --git a/python/sglang/jit_kernel/flash_attention/cute/utils.py b/python/sglang/jit_kernel/flash_attention/cute/utils.py new file mode 100644 index 000000000..f31d85c5d --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention/cute/utils.py @@ -0,0 +1,859 @@ +# Copyright (c) 2025, Tri Dao. + +import math +import hashlib +import inspect +import re +from typing import Type, Callable, Optional, Tuple, overload +from functools import partial + +import cutlass +import cutlass.cute as cute + +from cutlass import Float32, const_expr +from cutlass.cutlass_dsl import T, dsl_user_op +from cutlass._mlir.dialects import nvvm, llvm +from cutlass.cute.runtime import from_dlpack + + +# cute.arch.{fma,mul,add}_packed_f32x2 uses RZ rounding mode by default +fma_packed_f32x2 = partial(cute.arch.fma_packed_f32x2, rnd=nvvm.RoundingModeKind.RN) +mul_packed_f32x2 = partial(cute.arch.mul_packed_f32x2, rnd=nvvm.RoundingModeKind.RN) +add_packed_f32x2 = partial(cute.arch.add_packed_f32x2, rnd=nvvm.RoundingModeKind.RN) +sub_packed_f32x2 = partial( + cute.arch.calc_packed_f32x2_op, + src_c=None, + calc_func=nvvm.sub_packed_f32x2, + rnd=nvvm.RoundingModeKind.RN, +) + + +def hash_callable(func: Callable, set_cute_hash=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. Code-generation backends such + as Inductor can set this attribute to avoid expensive runtime hashing. + + set_cute_hash: whether or not to set func.__cute_hash__ if not present + """ + if hasattr(func, "__cute_hash__"): + return func.__cute_hash__ + + # Unwrap decorated functions (e.g., cute.jit wrappers). + if hasattr(func, "__wrapped__"): + base_func = func.__wrapped__ + if hasattr(base_func, "__cute_hash__"): + return base_func.__cute_hash__ + func = base_func + + 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 idx, cell in enumerate(func.__closure__): + cell_value = cell.cell_contents + hasher.update(repr(cell_value).encode()) + + hash = hasher.hexdigest() + + if set_cute_hash: + func.__cute_hash__ = hash + + return hash + + +def create_softcap_scoremod(softcap_val): + inv_softcap = 1.0 / softcap_val + + @cute.jit + def scoremod_premask_fn(acc_S_SSA, batch_idx, head_idx, q_idx, kv_idx, aux_tensors): + scores = acc_S_SSA * inv_softcap + return scores * cute.math.tanh(scores, fastmath=True) + + return scoremod_premask_fn + + +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_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.core.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.core.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_fragment(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 + + +def convert_layout_acc_mn(acc_layout: cute.Layout, transpose: bool = False) -> cute.Layout: + """ + For Sm80, convert ((2, 2), MMA_M, MMA_N, ...) to ((2, MMA_M), (2, MMA_N), ...). + For Sm90, convert ((2, 2, V), MMA_M, MMA_N, ...) to ((2, MMA_M), (2, V, MMA_N), ...). + """ + acc_layout_col_major = cute.make_layout(acc_layout.shape) + shape = ( + (acc_layout_col_major.shape[0][1], acc_layout_col_major.shape[1]), # MMA_M + ( + acc_layout_col_major.shape[0][0], + *acc_layout_col_major.shape[0][2:], + acc_layout_col_major.shape[2], + ), # MMA_N + *acc_layout_col_major.shape[3:], + ) + stride = ( + (acc_layout_col_major.stride[0][1], acc_layout_col_major.stride[1]), # MMA_M + ( + acc_layout_col_major.stride[0][0], + *acc_layout_col_major.stride[0][2:], + acc_layout_col_major.stride[2], + ), # MMA_N + *acc_layout_col_major.stride[3:], + ) + if const_expr(transpose): + shape = (shape[1], shape[0], *shape[2:]) + stride = (stride[1], stride[0], *stride[2:]) + acc_layout_mn = cute.make_layout(shape, stride=stride) + return cute.composition(acc_layout, acc_layout_mn) + + +def make_acc_tensor_mn_view(acc: cute.Tensor, transpose: bool = False) -> cute.Tensor: + return cute.make_tensor(acc.iterator, convert_layout_acc_mn(acc.layout, transpose=transpose)) + + +@cute.jit +def convert_layout_acc_frgA(acc_layout: cute.Layout) -> cute.Layout: + # For back to back gemm, convert layout of acc0 to gemm 1 accept layout. + # For Sm80, as the mma instruction shape is 16x8x16, we need to convert from (4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2) + # For Sm90, FP16/BF16, convert acc_layout from ((2, 2, N / 8), MMA_M, MMA_N) to ((2, 2, 2), MMA_M, (N / 16, MMA_N)) + # TODO: Sm90 FP8 + if const_expr(cute.rank(acc_layout.shape[0]) == 3): # Sm90 + l = cute.logical_divide( + acc_layout, ((None, None, 2), None, None) + ) # ((2, 2, (2, N / 16)), MMA_M, MMA_N) + rA_mma_view = cute.make_layout( + ( + (l.shape[0][0], l.shape[0][1], l.shape[0][2][0]), + l.shape[1], + (l.shape[0][2][1], l.shape[2]), + ), + stride=( + (l.stride[0][0], l.stride[0][1], l.stride[0][2][0]), + l.stride[1], + (l.stride[0][2][1], l.stride[2]), + ), + ) + else: # Sm80 + # (4, MMA_M, MMA_N) -> (4, MMA_M, (2, MMA_N / 2)) + l = cute.logical_divide(acc_layout, (None, None, 2)) + rA_mma_view = cute.make_layout( + ( + (l.shape[0], l.shape[2][0]), + l.shape[1], + l.shape[2][1], + ), + stride=( + (l.stride[0], l.stride[2][0]), + l.stride[1], + l.stride[2][1], + ), + ) + return rA_mma_view + + +def make_acc_tensor_frgA_view(acc: cute.Tensor) -> cute.Tensor: + return cute.make_tensor(acc.iterator, convert_layout_acc_frgA(acc.layout)) + + +def select(a: cute.Tensor, mode: list[int]) -> cute.Tensor: + return cute.make_tensor(a.iterator, cute.select(a.layout, mode)) + + +def transpose_view(a: cute.Tensor) -> cute.Tensor: + """Transpose the first two dimensions of a tensor on smem.""" + shape = (a.shape[1], a.shape[0], *a.shape[2:]) + order = (1, 0, *range(2, cute.rank(a))) + return cute.composition(a, cute.make_ordered_layout(shape, order=order)) + # stride = (a.layout.stride[1], a.layout.stride[0], *a.layout.stride[2:]) + # return cute.make_tensor(a.iterator, cute.make_layout(shape, stride=stride)) + + +def parse_swizzle_from_pointer(ptr: cute.Pointer) -> cute.Swizzle: + """Extract swizzle parameters from a pointer's swizzle_type. + + The swizzle_type string has the form '!cute.swizzle<"S">' where + b, m, s are the swizzle parameters (bits, base, shift). + + Returns: + A cute.Swizzle object constructed from the extracted parameters + + Raises: + ValueError: If the swizzle_type string cannot be parsed + """ + # Ideally there should be a better API to get swizzle parameters, but we'll just parse + # the string here. + swizzle_str = str(ptr.type.swizzle_type) + # Extract the inner part "S" + match = re.search(r"S<(\d+),(\d+),(\d+)>", swizzle_str) + if match: + b, m, s = int(match.group(1)), int(match.group(2)), int(match.group(3)) + return cute.make_swizzle(b, m, s) + else: + raise ValueError(f"Could not parse swizzle_type: {swizzle_str}") + + +@cute.jit +def exp2f(x: cute.TensorSSA | Float32) -> cute.TensorSSA | Float32: + """exp2f calculation for both vector and scalar. + :param x: input value + :type x: cute.TensorSSA or Float32 + :return: exp2 value + :rtype: cute.TensorSSA or Float32 + """ + if const_expr(isinstance(x, cute.TensorSSA)): + res = cute.make_fragment(x.shape, Float32) + res.store(x) + for i in cutlass.range_constexpr(cute.size(x.shape)): + res[i] = cute.arch.exp2(res[i]) + return res.load() + else: + return cute.arch.exp2(x) + + +@dsl_user_op +def log2f(a: float | Float32, *, loc=None, ip=None) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(a).ir_value(loc=loc, ip=ip)], + "lg2.approx.ftz.f32 $0, $1;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + ) + + +@dsl_user_op +def logf(a: float | Float32, *, loc=None, ip=None) -> Float32: + return log2f(a, loc=loc, ip=ip) * math.log(2.0) + + +@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( + T.f32(), + 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_fragment(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_fragment(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_fragment(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_fragment(x.shape, Float32) + res.store(x) + local_sum_0 = ( + add_packed_f32x2((init_val, 0.0), (res[0], res[1])) + # 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] = add_packed_f32x2(local_sum[0], (res[i + 0], res[i + 1])) + local_sum[1] = add_packed_f32x2(local_sum[1], (res[i + 2], res[i + 3])) + local_sum[2] = add_packed_f32x2(local_sum[2], (res[i + 4], res[i + 5])) + local_sum[3] = add_packed_f32x2(local_sum[3], (res[i + 6], res[i + 7])) + local_sum[0] = add_packed_f32x2(local_sum[0], local_sum[1]) + local_sum[2] = add_packed_f32x2(local_sum[2], local_sum[3]) + local_sum[0] = 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) + + +@dsl_user_op +def elem_pointer_i64(x: cute.Tensor, coord: cute.Coord, *, loc=None, ip=None) -> cute.Pointer: + flat_coord_i64 = tuple(cutlass.Int64(c) for c in cute.flatten(coord)) + flat_stride = cute.flatten_to_tuple(x.stride) + assert len(flat_coord_i64) == len(flat_stride), ( + "Coordinate and stride must have the same length" + ) + offset = sum(c * s for c, s in zip(flat_coord_i64, flat_stride)) + # HACK: we assume that applying the offset does not change the pointer alignment + byte_offset = offset * x.element_type.width // 8 + return cute.make_ptr( + x.element_type, + x.iterator.toint() + byte_offset, + x.memspace, + assumed_align=x.iterator.alignment, + ) + + +@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_fragment( + 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 shr_u32(val: cutlass.Uint32, shift: cutlass.Uint32, *, loc=None, ip=None) -> cutlass.Uint32: + 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.s32 $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_fragment(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 = 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, *, loc=None, ip=None) -> Float32: + # We assume x <= 127.0 + poly_ex2_deg3 = ( + 1.0, + 0.695146143436431884765625, + 0.227564394474029541015625, + 0.077119089663028717041015625, + ) + 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_deg3, 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, *, loc=None, ip=None) -> Tuple[Float32, Float32]: + # We assume x <= 127.0 and y <= 127.0 + poly_ex2_deg3 = ( + 1.0, + 0.695146143436431884765625, + 0.227564394474029541015625, + 0.077119089663028717041015625, + ) + 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=nvvm.RoundingModeKind.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 = sub_packed_f32x2(xy_rounded, (fp32_round_int, fp32_round_int)) + xy_frac = sub_packed_f32x2(xy_clamped, xy_rounded_back) + xy_frac_ex2 = evaluate_polynomial_2(*xy_frac, poly_ex2_deg3, 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 domain_offset_i64(coord: cute.Coord, tensor: cute.Tensor, *, loc=None, ip=None) -> cute.Tensor: + flat_coord_i64 = tuple(cutlass.Int64(c) for c in cute.flatten(coord)) + flat_stride = cute.flatten_to_tuple(tensor.stride) + assert len(flat_coord_i64) == len(flat_stride), ( + "Coordinate and stride must have the same length" + ) + offset = sum(c * s for c, s in zip(flat_coord_i64, flat_stride)) + assert isinstance(tensor.iterator, cute.Pointer) + # HACK: we assume that applying the offset does not change the pointer alignment + new_ptr = cute.make_ptr( + tensor.element_type, + tensor.iterator.toint() + offset * tensor.element_type.width // 8, + tensor.memspace, + assumed_align=tensor.iterator.max_alignment, + ) + return cute.make_tensor(new_ptr, tensor.layout) + + +@dsl_user_op +def coord_offset_i64( + tensor: cute.Tensor, idx: cute.typing.Int, dim: int, *, loc=None, ip=None +) -> cute.Tensor: + offset = cutlass.Int64(idx) * cute.size(tensor.stride[dim]) + assert isinstance(tensor.iterator, cute.Pointer) + # HACK: we assume that applying the offset does not change the pointer alignment + new_ptr = cute.make_ptr( + tensor.element_type, + tensor.iterator.toint() + offset * tensor.element_type.width // 8, + tensor.memspace, + assumed_align=tensor.iterator.max_alignment, + ) + new_layout = cute.slice_( + tensor.layout, (*[None] * dim, 0, *[None] * (cute.rank(tensor) - dim - 1)) + ) + return cute.make_tensor(new_ptr, new_layout) + + +@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_fragment(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] diff --git a/python/sglang/jit_kernel/flash_attention_v4.py b/python/sglang/jit_kernel/flash_attention_v4.py new file mode 100644 index 000000000..6817a1d54 --- /dev/null +++ b/python/sglang/jit_kernel/flash_attention_v4.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +from typing import Callable, Optional, Tuple, Union + +import torch + +try: + from sglang.jit_kernel.flash_attention.cute import ( + flash_attn_varlen_func as _flash_attn_varlen_func, + ) +except Exception as _e: # pragma: no cover + _flash_attn_varlen_func = None + _flash_attn_import_error = _e +else: + _flash_attn_import_error = None + + +def _maybe_contiguous(x: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + return x.contiguous() if x is not None and x.stride(-1) != 1 else x + + +def flash_attn_varlen_func( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + 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, + page_table: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + softcap: Optional[float] = None, + window_size: Tuple[Optional[int], Optional[int]] = (-1, -1), + learnable_sink: Optional[torch.Tensor] = None, + sinks: Optional[torch.Tensor] = None, + num_splits: int = 1, + pack_gqa: Optional[bool] = None, + score_mod: Optional[Callable] = None, + aux_tensors: Optional[list] = None, + return_softmax_lse: bool = False, + **_: object, +): + if _flash_attn_varlen_func is None: # pragma: no cover + raise ImportError( + "Vendored FlashAttention CUTE is not available (cannot import " + "sglang.jit_kernel.flash_attention.cute). Please check your source tree." + ) from _flash_attn_import_error + + q, k, v = [_maybe_contiguous(t) for t in (q, k, v)] + cu_seqlens_q, cu_seqlens_k = [ + _maybe_contiguous(t) for t in (cu_seqlens_q, cu_seqlens_k) + ] + seqused_q, seqused_k = [_maybe_contiguous(t) for t in (seqused_q, seqused_k)] + page_table = _maybe_contiguous(page_table) + + if learnable_sink is None and sinks is not None: + learnable_sink = sinks + + if window_size == (-1, -1): + window_size = (None, None) + + result = _flash_attn_varlen_func( + q=q, + k=k, + v=v, + 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, + page_table=page_table, + softmax_scale=softmax_scale, + causal=causal, + softcap=softcap, + window_size=window_size, + learnable_sink=learnable_sink, + num_splits=num_splits, + pack_gqa=pack_gqa, + score_mod=score_mod, + aux_tensors=aux_tensors, + ) + + if return_softmax_lse: + return result + if isinstance(result, tuple): + return result[0] + return result + + +def flash_attn_with_kvcache( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k: Optional[torch.Tensor] = None, + v: Optional[torch.Tensor] = None, + qv: Optional[torch.Tensor] = None, + rotary_cos: Optional[torch.Tensor] = None, + rotary_sin: Optional[torch.Tensor] = None, + cache_seqlens: Optional[Union[int, torch.Tensor]] = None, + cache_batch_idx: Optional[torch.Tensor] = None, + cache_leftpad: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_k_new: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + rotary_seqlens: Optional[torch.Tensor] = None, + q_descale: Optional[torch.Tensor] = None, + k_descale: Optional[torch.Tensor] = None, + v_descale: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + causal: bool = False, + window_size: Tuple[int, int] = (-1, -1), + attention_chunk: Optional[int] = None, + softcap: float = 0.0, + rotary_interleaved: bool = True, + scheduler_metadata=None, + num_splits: int = 0, + pack_gqa: Optional[bool] = None, + sm_margin: int = 0, + sinks: Optional[torch.Tensor] = None, + score_mod: Optional[Callable] = None, + aux_tensors: Optional[list] = None, + return_softmax_lse: bool = False, + **_: object, +): + if k is not None or v is not None or qv is not None: + raise NotImplementedError("FA4 does not support updating KV cache in-place.") + if rotary_cos is not None or rotary_sin is not None or rotary_seqlens is not None: + raise NotImplementedError("FA4 path does not support rotary embedding.") + if cache_batch_idx is not None or cache_leftpad is not None: + raise NotImplementedError( + "FA4 path does not support non-consecutive batch indices or left padding." + ) + if q_descale is not None or k_descale is not None or v_descale is not None: + raise NotImplementedError("FA4 path does not support descale.") + + if isinstance(cache_seqlens, int): + cache_seqlens = torch.full( + (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device + ) + + result = flash_attn_varlen_func( + q=q, + k=k_cache, + v=v_cache, + cu_seqlens_q=cu_seqlens_q, + seqused_k=cache_seqlens, + max_seqlen_q=max_seqlen_q, + page_table=page_table, + softmax_scale=softmax_scale, + causal=causal, + softcap=softcap if softcap != 0.0 else None, + window_size=window_size, + num_splits=num_splits if num_splits != 0 else 1, + pack_gqa=pack_gqa, + learnable_sink=sinks, + score_mod=score_mod, + aux_tensors=aux_tensors, + return_softmax_lse=True, + ) + + if return_softmax_lse: + return result + if isinstance(result, tuple): + return result[0] + return result diff --git a/python/sglang/jit_kernel/tests/test_flash_attention_4.py b/python/sglang/jit_kernel/tests/test_flash_attention_4.py new file mode 100644 index 000000000..fe19c9175 --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_flash_attention_4.py @@ -0,0 +1,1504 @@ +# Adapted from https://github.com/Dao-AILab/flash-attention/blob/8ecf128f683266735ba68e3c106ff67a2611886e/tests/cute/test_flash_attn.py + +# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. + +import itertools +import math + +import pytest +import torch +import torch.nn.functional as F +from einops import rearrange, repeat + +from sglang.jit_kernel.flash_attention_v4 import flash_attn_varlen_func + +# Skip this test on Hopper machine +skip_condition = torch.cuda.get_device_capability() < (10, 0) + + +def apply_rotary_emb( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + seqlen_offsets: torch.Tensor | int | None = 0, + interleaved: bool = False, +) -> torch.Tensor: + rotary_dim = cos.shape[-1] * 2 + x_rot = x[..., :rotary_dim] + x_pass = x[..., rotary_dim:] + + cos = cos.to(dtype=x.dtype) + sin = sin.to(dtype=x.dtype) + + if x_rot.dim() < 2: + raise ValueError(f"apply_rotary_emb expects x.dim() >= 2, got {x_rot.dim()}") + + b = x_rot.shape[0] + s = x_rot.shape[1] + + if seqlen_offsets is None: + seqlen_offsets = 0 + + if isinstance(seqlen_offsets, int): + positions = ( + torch.arange(s, device=x_rot.device, dtype=torch.long) + seqlen_offsets + ) + cos_s = cos.index_select(0, positions) + sin_s = sin.index_select(0, positions) + cos_s = cos_s.unsqueeze(0).expand(b, -1, -1) + sin_s = sin_s.unsqueeze(0).expand(b, -1, -1) + else: + if seqlen_offsets.dim() != 1 or seqlen_offsets.shape[0] != b: + raise ValueError( + "apply_rotary_emb expects seqlen_offsets to be int or shape [batch]" + ) + positions = torch.arange(s, device=x_rot.device, dtype=torch.long).unsqueeze( + 0 + ) + seqlen_offsets.to(dtype=torch.long).unsqueeze(1) + cos_s = cos.index_select(0, positions.reshape(-1)).view(b, s, -1) + sin_s = sin.index_select(0, positions.reshape(-1)).view(b, s, -1) + + x_rot = x_rot.reshape(b, s, -1, rotary_dim) + cos_s = cos_s.unsqueeze(2) + sin_s = sin_s.unsqueeze(2) + + if interleaved: + x1 = x_rot[..., ::2] + x2 = x_rot[..., 1::2] + o1 = x1 * cos_s - x2 * sin_s + o2 = x2 * cos_s + x1 * sin_s + x_rot = torch.stack((o1, o2), dim=-1).flatten(-2) + else: + x1, x2 = torch.chunk(x_rot, 2, dim=-1) + o1 = x1 * cos_s - x2 * sin_s + o2 = x2 * cos_s + x1 * sin_s + x_rot = torch.cat((o1, o2), dim=-1) + + x_rot = x_rot.reshape_as(x[..., :rotary_dim]) + return torch.cat((x_rot, x_pass), dim=-1) + + +def unpad_input(hidden_states, attention_mask, unused_mask=None): + """ + Arguments: + hidden_states: (batch, seqlen, ...) + attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid. + unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused. + Return: + hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask. + indices: (total_nnz), the indices of masked tokens from the flattened input sequence. + cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states. + max_seqlen_in_batch: int + seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask. + """ + 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) + indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the + # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim + # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to + # index with integer indices. + return ( + 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): + """ + Arguments: + hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask. + indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence. + batch: int, batch size for the padded sequence. + seqlen: int, maximum sequence length for the padded sequence. + Return: + hidden_states: (batch, seqlen, ...) + """ + dim = hidden_states.shape[1:] + output = torch.zeros( + (batch * seqlen), *dim, device=hidden_states.device, dtype=hidden_states.dtype + ) + output[indices] = hidden_states + return rearrange(output, "(b s) ... -> b s ...", b=batch) + + +def generate_random_padding_mask( + max_seqlen, batch_size, device, mode="random", zero_lengths=False +): + assert mode in ["full", "random", "third"] + if mode == "full": + lengths = torch.full( + (batch_size, 1), max_seqlen, device=device, dtype=torch.int32 + ) + elif mode == "random": + lengths = torch.randint( + max(0 if zero_lengths else 1, max_seqlen - 20), + max_seqlen + 1, + (batch_size, 1), + device=device, + ) + elif mode == "third": + lengths = torch.randint( + max_seqlen // 3, max_seqlen + 1, (batch_size, 1), device=device + ) + else: + # This should never happen due to the assertion above, but for linter + lengths = torch.full( + (batch_size, 1), max_seqlen, device=device, dtype=torch.int32 + ) + + if zero_lengths: + # Generate zero-lengths every 5 batches and the last batch. + 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, +): + """ + Arguments: + q: (batch_size, seqlen_q, nheads, d) + k: (batch_size, seqlen_k, nheads_k, d) + v: (batch_size, seqlen_k, nheads_k, d_v) + query_padding_mask: (batch_size, seqlen), bool + key_padding_mask: (batch_size, seqlen), bool + """ + 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, *rest = 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 + return torch.logical_or( + col_idx > torch.minimum(row_idx + sk - sq + window_size[1], sk), + 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 + # Subtract remainder instead of divide and then multiply to take care of negative values + 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=None, + softcap=0.0, + upcast=True, + reorder_ops=False, + intermediate_dtype=None, +): + """ + Arguments: + q: (batch_size, seqlen_q, nheads, head_dim) + k: (batch_size, seqlen_k, nheads, head_dim) + v: (batch_size, seqlen_k, nheads, head_dim_v) + qv: (batch_size, seqlen_q, nheads, head_dim_v) + query_padding_mask: (batch_size, seqlen_q) + key_padding_mask: (batch_size, seqlen_k) + attn_bias: broadcastable to (batch_size, nheads, seqlen_q, seqlen_k) + dropout_p: float + dropout_mask: (batch_size, nheads, seqlen_q, seqlen_k) + causal: whether to apply causal masking + upcast: whether to cast all inputs to fp32, do all computation in fp32, then cast + output back to fp16/bf16. + reorder_ops: whether to change the order of operations (scaling k instead of scaling k, etc.) + without changing the math. This is to estimate the numerical error from operation + reordering. + Output: + output: (batch_size, seqlen_q, nheads, head_dim_v) + attention: (batch_size, nheads, seqlen_q, seqlen_k), softmax after dropout + """ + if causal: + window_size = (window_size[0], 0) + dtype_og = q.dtype + if upcast: + q, k, v = q.float(), k.float(), v.float() + qv = qv.float() if qv is not None else None + if q_descale is not None: + q_descale = repeat(q_descale, "b h -> b 1 (h g) 1", g=q.shape[2] // k.shape[2]) + q = (q.float() * q_descale).to(q.dtype) + qv = (qv.float() * q_descale).to(qv.dtype) if qv is not None else None + 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], k.shape[1] + k = repeat(k, "b s h d -> b s (h g) d", g=q.shape[2] // k.shape[2]) + v = repeat(v, "b s h d -> b s (h g) d", g=q.shape[2] // v.shape[2]) + d = q.shape[-1] + dv = v.shape[-1] + softmax_scale = 1.0 / math.sqrt(d if qv is None else d + dv) + if not reorder_ops: + scores = torch.einsum("bthd,bshd->bhts", q * softmax_scale, k) + else: + scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale) + if qv is not None: + scores = scores + torch.einsum("bthd,bshd->bhts", qv * softmax_scale, v) + 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=q.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=q.device, + ) + local_mask = ( + torch.logical_or(local_mask, chunk_mask) + if local_mask is not None + else chunk_mask + ) + if local_mask is not None: + scores.masked_fill_(local_mask, float("-inf")) + if attn_bias is not None: + scores = scores + attn_bias + 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 + ) + attention = (unnormalized_scores / normalizer).to(v.dtype) + # We want to mask here so that the attention matrix doesn't have any NaNs + # Otherwise we'll get NaN in dV + if query_padding_mask is not None: + attention = attention.masked_fill( + rearrange(~query_padding_mask, "b s -> b 1 s 1"), 0.0 + ) + # Without this we might get NaN in dv + if key_padding_mask is not None: + attention = attention.masked_fill( + rearrange(~key_padding_mask, "b s -> b 1 1 s"), 0.0 + ) + # Some rows might be completely masked out so we fill them with zero instead of NaN + 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) + # attention_drop = attention.masked_fill(~dropout_mask, 0.0) * dropout_scaling + # output = torch.einsum('bhts,bshd->bthd', attention_drop , v) + 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) + return output.to(dtype=dtype_og), attention.to(dtype=dtype_og) + + +@pytest.mark.skipif( + skip_condition, reason="FA4 Requires compute capability of 10 or above." +) +# @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("mha_type", ["mha", "mqa", "gqa"]) +# @pytest.mark.parametrize("mha_type", ["mqa"]) +@pytest.mark.parametrize("has_learnable_sink", [False, True]) +# @pytest.mark.parametrize("has_learnable_sink", [False]) +# @pytest.mark.parametrize("has_qv", [False, True]) +@pytest.mark.parametrize("has_qv", [False]) +# @pytest.mark.parametrize("deterministic", [False, True]) +@pytest.mark.parametrize("deterministic", [False]) +# @pytest.mark.parametrize("softcap", [0.0, 15.0]) +@pytest.mark.parametrize("softcap", [0.0]) +# @pytest.mark.parametrize("local", [False, True]) +@pytest.mark.parametrize("local", [False]) +@pytest.mark.parametrize("causal", [False, True]) +# @pytest.mark.parametrize("causal", [False]) +# @pytest.mark.parametrize("add_unused_qkv", [False, True]) +@pytest.mark.parametrize("add_unused_qkv", [False]) +# @pytest.mark.parametrize("d", [32, 64, 96, 128, 160, 192, 224, 256]) +# @pytest.mark.parametrize('d', [32, 40, 64, 80, 96, 128, 160, 192, 256]) +# @pytest.mark.parametrize('d', [32, 64, 96, 128, 160, 192]) +# @pytest.mark.parametrize('d', [56, 80]) +# @pytest.mark.parametrize('d', [32, 40, 64, 80, 96, 128]) +# @pytest.mark.parametrize("d", [64, 96, 128]) +@pytest.mark.parametrize("d", [64, 128]) +# @pytest.mark.parametrize("d", [192]) +@pytest.mark.parametrize( + "seqlen_q,seqlen_k", + [ + # (1, 1), + # (1, 3), + # (2, 1), + (511, 1), + (3, 513), + (64, 128), + (128, 128), + (256, 256), + # (113, 203), + # (128, 217), + # (113, 211), + # (108, 256), + # (256, 512), + (307, 256), + (640, 128), + (512, 256), + (1024, 1024), + (1023, 1024), + (1024, 1023), + (2048, 2048), + ], +) +def test_flash_attn_varlen_output( + seqlen_q, + seqlen_k, + d, + add_unused_qkv, + causal, + local, + softcap, + deterministic, + has_qv, + has_learnable_sink, + mha_type, + dtype, +): + if ( + causal or local + ): # Right now we only support causal attention with seqlen_k == seqlen_q + seqlen_k = seqlen_q + device = "cuda" + # set seed + torch.random.manual_seed(seqlen_q + seqlen_k + d + int(causal) * 2 + int(local)) + batch_size = 49 if seqlen_q <= 1024 else 7 + nheads = 6 + # batch_size = 1 + # nheads = 1 + nheads_kv = nheads if mha_type == "mha" else (3 if mha_type == "gqa" else 1) + dtype_ref = torch.bfloat16 if dtype == torch.float8_e4m3fn else dtype + # dv_vals = [128, d] if d > 128 and d <= 192 else ([256, 512, d] if d <= 64 else [d]) + dv_vals = [128] if d == 192 else ([d] if d != 128 else [64, d]) + if dtype == torch.float8_e4m3fn: + dv_vals = [d] + # attention_chunk_vals = [torch.randint(1, seqlen_k * 2, (1,)).item(), 0] if seqlen_q <= seqlen_k else [0] + attention_chunk_vals = [0] + for dv, attention_chunk in itertools.product(dv_vals, attention_chunk_vals): + q_ref = torch.randn( + batch_size, seqlen_q, nheads, d, device=device, dtype=dtype_ref + ) + if softcap > 0.0: + # Ensure the values of qk are at least within softcap range. + q_ref = (q_ref * softcap / 4).detach().requires_grad_() + q_ref = q_ref.to(dtype).to(dtype_ref).requires_grad_() + k_ref = ( + torch.randn( + batch_size, seqlen_k, nheads_kv, d, device=device, dtype=dtype_ref + ) + .to(dtype) + .to(dtype_ref) + .requires_grad_() + ) + v_ref = ( + torch.randn( + batch_size, seqlen_k, nheads_kv, dv, device=device, dtype=dtype_ref + ) + .to(dtype) + .to(dtype_ref) + .requires_grad_() + ) + if has_qv: + qv_ref = ( + torch.randn( + batch_size, seqlen_q, nheads, dv, device=device, dtype=dtype_ref + ) + .to(dtype) + .to(dtype_ref) + ) + else: + qv_ref = None + # Put window_size after QKV randn so that window_size changes from test to test + window_size = ( + (None, None) if not local else torch.randint(0, seqlen_k, (2,)).tolist() + ) + if has_learnable_sink: + learnable_sink = torch.randn(nheads, dtype=torch.bfloat16, device=device) + else: + learnable_sink = None + if dtype == torch.float8_e4m3fn: + q_descale, k_descale, v_descale = [ + torch.rand(batch_size, nheads_kv, device=device, dtype=torch.float32) + * 2 + for _ in range(3) + ] + else: + q_descale, k_descale, v_descale = None, None, None + q, k, v = [x.detach().requires_grad_() for x in (q_ref, k_ref, v_ref)] + qv = qv_ref.detach() if has_qv else None + query_padding_mask = generate_random_padding_mask( + seqlen_q, batch_size, device, mode="random", zero_lengths=False + ) + # TODO: test zero_lengths + key_padding_mask = generate_random_padding_mask( + # seqlen_k, batch_size, device, mode="random", zero_lengths=True + seqlen_k, + batch_size, + device, + mode="random", + zero_lengths=False, + ) + + def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): + if add_unused: + another_mask = generate_random_padding_mask(max_seq_len, bs, device) + attn_mask = torch.logical_and(padding_mask, another_mask) + unused_mask = torch.logical_xor( + torch.logical_or(padding_mask, another_mask), attn_mask + ) + else: + attn_mask = padding_mask + unused_mask = None + return attn_mask, unused_mask + + query_padding_mask, query_unused_mask = _gen_unused_masks( + query_padding_mask, add_unused_qkv, seqlen_q, batch_size, q.device + ) + # query_padding_mask[:] = True + # query_unused_mask = None + key_padding_mask, key_unused_mask = _gen_unused_masks( + key_padding_mask, add_unused_qkv, seqlen_k, batch_size, k.device + ) + + if causal or local: + key_padding_mask = query_padding_mask + + result = generate_qkv( + q, + k, + v, + query_padding_mask, + key_padding_mask, + qv=qv, + kvpacked=False, + query_unused_mask=query_unused_mask, + key_unused_mask=key_unused_mask, + ) + ( + q_unpad, # 0 + k_unpad, # 1 + v_unpad, # 2 + qv_unpad, # 3 + cu_seqlens_q, # 4 + cu_seqlens_k, # 5 + seqused_q, # 6 + seqused_k, # 7 + max_seqlen_q, # 8 + max_seqlen_k, # 9 + q, # 10 + k, # 11 + v, # 12 + qv, # 13 + output_pad_fn, # 14 + dq_pad_fn, # 15 + dk_pad_fn, # 16 + ) = result + q_unpad, k_unpad, v_unpad = [ + x.detach().to(dtype).requires_grad_() for x in (q_unpad, k_unpad, v_unpad) + ] + out_ref, attn_ref = attention_ref( + q_ref, + k_ref, + v_ref, + query_padding_mask, + key_padding_mask, + causal=causal, + qv=qv_ref, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + window_size=window_size, + attention_chunk=attention_chunk, + learnable_sink=learnable_sink, + softcap=softcap, + ) + out_pt, attn_pt = attention_ref( + q_ref, + k_ref, + v_ref, + query_padding_mask, + key_padding_mask, + causal=causal, + qv=qv_ref, + q_descale=q_descale, + k_descale=k_descale, + v_descale=v_descale, + window_size=window_size, + attention_chunk=attention_chunk, + learnable_sink=learnable_sink, + softcap=softcap, + upcast=False, + reorder_ops=True, + intermediate_dtype=dtype if dtype == torch.float8_e4m3fn else None, + ) + + print(f"Pytorch max diff: {(out_pt - out_ref).abs().max().item()}") + print(f"Pytorch mean diff: {(out_pt - out_ref).abs().mean().item()}") + + if query_unused_mask is not None: + q_zero_masking = rearrange(query_unused_mask, "b s -> b s 1 1") + + # Numerical error if we just do any arithmetic on out_ref + fwd_atol = 2 * (out_ref + 0.3 - 0.3 - out_ref).abs().max().item() + rtol = 2 if softcap == 0.0 else 3 + + pack_gqa_vals = [False, True, None] + # num_splits_vals = [1, 3] + num_splits_vals = [1] + for pack_gqa, num_splits in itertools.product(pack_gqa_vals, num_splits_vals): + out_unpad, lse = flash_attn_varlen_func( + q_unpad, + k_unpad, + v_unpad, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + # max_seqlen_q and max_seqlen_k not needed for FA4 + seqused_q=seqused_q, + seqused_k=seqused_k, + causal=causal, + window_size=window_size, + softcap=softcap, + sinks=learnable_sink, # FA4 uses learnable_sink, not sinks + pack_gqa=pack_gqa, + return_softmax_lse=True, + ) + out = output_pad_fn(out_unpad) + if query_unused_mask is not None: + out.masked_fill_(q_zero_masking, 0.0) + print(f"Output max diff: {(out - out_ref).abs().max().item()}") + print(f"Output mean diff: {(out - out_ref).abs().mean().item()}") + # if not causal: + # print(f"LSE max diff: {(lse - lse_ref).abs().max().item()}") + # breakpoint() + + # Check that FlashAttention's numerical error is at most 3x the numerical error + # of a Pytorch implementation. + assert (out - out_ref).abs().max().item() <= rtol * ( + out_pt - out_ref + ).abs().max().item() + fwd_atol + + if ( + dtype != torch.float8_e4m3fn + and not has_qv + and not dv > 256 + and not attention_chunk != 0 + and dv == d + and not has_learnable_sink + and False + ): + g_unpad = torch.randn_like(out_unpad) + do_o = ((g_unpad.float() * out_unpad.float()).sum(-1)).transpose(-1, -2) + # import flash_attn_3_cuda + # dq_unpad, dk_unpad, dv_unpad, softmax_d, dq_accum, lse_log2 = flash_attn_3_cuda.bwd_varlen( + # g_unpad, + # q_unpad, + # k_unpad, + # v_unpad, + # out_unpad, + # lse, + # None, + # None, + # None, + # cu_seqlens_q, + # cu_seqlens_k, + # None, None, + # max_seqlen_q, + # max_seqlen_k, + # d ** (-0.5), + # causal, + # window_size[0], window_size[1], + # softcap, + # deterministic, + # 0, # sm_margin + # ) + dq_unpad, dk_unpad, dv_unpad = torch.autograd.grad( + out_unpad, (q_unpad, k_unpad, v_unpad), g_unpad + ) + dq = dq_pad_fn(dq_unpad) + dk = dk_pad_fn(dk_unpad) + dv = dk_pad_fn(dv_unpad) + if key_unused_mask is not None: + k_zero_masking = rearrange(key_unused_mask, "b s -> b s 1 1") + dk.masked_fill_(k_zero_masking, 0.0) + dv.masked_fill_(k_zero_masking, 0.0) + if query_unused_mask is not None: + dq.masked_fill_(q_zero_masking, 0.0) + # print(f"dO_O max diff: {(softmax_d - do_o).abs().max().item()}") + # assert (softmax_d - do_o).abs().max().item() <= 1e-5 + # assert dq_accum.abs().max().item() == 0.0 + g = output_pad_fn(g_unpad) + + # qk = torch.einsum('bthd,bshd->bhts', q / (d ** 0.5), k).float() + # qk = torch.masked_fill(qk, rearrange(~key_padding_mask, "b s -> b 1 1 s"), float("-inf")) + # dS = torch.einsum('bthd,bshd->bhts', g.float(), v.float()) + # P = torch.softmax(qk, -1) + # dP = P * (dS - (g.float() * out.float()).sum(-1).transpose(1, 2).unsqueeze(-1)) + # dQ = torch.einsum('bhts,bshd->bthd', dP, k.float()) + # dV = torch.einsum('bhts,bthd->bshd', P, g.float()) + # dK = torch.einsum('bhts,bthd->bshd', dP, q.float()) + + # dq, dk, dv = torch.autograd.grad(out, (q, k, v), g) + dq_ref, dk_ref, dv_ref = torch.autograd.grad( + out_ref, (q_ref, k_ref, v_ref), g + ) + dq_pt, dk_pt, dv_pt = torch.autograd.grad(out_pt, (q_ref, k_ref, v_ref), g) + print(f"dQ max diff: {(dq - dq_ref).abs().max().item()}") + print(f"dK max diff: {(dk - dk_ref).abs().max().item()}") + print(f"dV max diff: {(dv - dv_ref).abs().max().item()}") + print(f"dQ mean diff: {(dq - dq_ref).abs().mean().item()}") + print(f"dK mean diff: {(dk - dk_ref).abs().mean().item()}") + print(f"dV mean diff: {(dv - dv_ref).abs().mean().item()}") + print(f"dQ Pytorch max diff: {(dq_pt - dq_ref).abs().max().item()}") + print(f"dK Pytorch max diff: {(dk_pt - dk_ref).abs().max().item()}") + print(f"dV Pytorch max diff: {(dv_pt - dv_ref).abs().max().item()}") + print(f"dQ Pytorch mean diff: {(dq_pt - dq_ref).abs().mean().item()}") + print(f"dK Pytorch mean diff: {(dk_pt - dk_ref).abs().mean().item()}") + print(f"dV Pytorch mean diff: {(dv_pt - dv_ref).abs().mean().item()}") + # breakpoint() + dq_atol = 2 * (dq_ref + 0.3 - 0.3 - dq_ref).abs().max().item() + ( + 0 if softcap == 0 else 3e-4 + ) + assert (dq - dq_ref).abs().max().item() <= rtol * ( + dq_pt - dq_ref + ).abs().max().item() + dq_atol + dk_atol = 2 * (dk_ref + 0.3 - 0.3 - dk_ref).abs().max().item() + ( + 0 if softcap == 0 else 3e-4 + ) + assert (dk - dk_ref).abs().max().item() <= rtol * ( + dk_pt - dk_ref + ).abs().max().item() + dk_atol + dv_atol = 2 * (dv_ref + 0.3 - 0.3 - dv_ref).abs().max().item() + ( + 0 if softcap == 0 else 3e-4 + ) + assert (dv - dv_ref).abs().max().item() <= rtol * ( + dv_pt - dv_ref + ).abs().max().item() + dv_atol + + +@pytest.mark.skipif( + skip_condition, reason="FA4 Requires compute capability of 10 or above." +) +# @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +# @pytest.mark.parametrize("dtype", [torch.float8_e4m3fn]) +@pytest.mark.parametrize("mha_type", ["mha", "mqa", "gqa"]) +# @pytest.mark.parametrize("mha_type", ["mha"]) +@pytest.mark.parametrize("has_learnable_sink", [False, True]) +# @pytest.mark.parametrize("has_learnable_sink", [False]) +# @pytest.mark.parametrize("new_kv", [False, True]) +@pytest.mark.parametrize("new_kv", [False]) +# @pytest.mark.parametrize("local", [False, True]) +@pytest.mark.parametrize("local", [False]) +# @pytest.mark.parametrize("causal", [False, True]) +@pytest.mark.parametrize("causal", [True]) +# @pytest.mark.parametrize("seqlen_new_eq_seqlen_q", [True, False]) +@pytest.mark.parametrize("seqlen_new_eq_seqlen_q", [False]) +# @pytest.mark.parametrize("has_rotary_seqlens", [False, True]) +@pytest.mark.parametrize("has_rotary_seqlens", [False]) +# @pytest.mark.parametrize("rotary_interleaved", [False, True]) +@pytest.mark.parametrize("rotary_interleaved", [True]) +# @pytest.mark.parametrize("rotary_fraction", [0.0, 0.5, 1.0]) +@pytest.mark.parametrize("rotary_fraction", [0.0]) +# @pytest.mark.parametrize("page_size", [None] + ([1, 4, 128])) +# @pytest.mark.parametrize("page_size", [None, 128]) +@pytest.mark.parametrize("page_size", [128]) +# @pytest.mark.parametrize("has_leftpad", [False, True]) +@pytest.mark.parametrize("has_leftpad", [False]) +# @pytest.mark.parametrize("has_batch_idx", [False, True]) +@pytest.mark.parametrize("has_batch_idx", [False]) +# @pytest.mark.parametrize("varlen_q", [False, True]) +@pytest.mark.parametrize("varlen_q", [False]) +# @pytest.mark.parametrize("d", [32, 59, 64, 80, 128, 256]) +# @pytest.mark.parametrize("d", [32, 64, 96, 128, 160, 192, 224, 256]) +# @pytest.mark.parametrize('d', [32, 40, 64, 80, 96, 128, 160, 192]) +# @pytest.mark.parametrize('d', [56, 80]) +# @pytest.mark.parametrize("d", [128]) +@pytest.mark.parametrize("d", [64]) +# @pytest.mark.parametrize("d", [192]) +@pytest.mark.parametrize( + "seqlen_q,seqlen_k", + [ + (1, 128), + (1, 339), + (3, 1024), + (64, 800), + (64, 256), + (3, 799), + (64, 2048), + (16, 20000), + # # (1, 128 * 1024), + # # (16, 128 * 1024), + # (128, 128), + # (256, 512), # To test appending KV with more than 1 block + # (2048, 3577), # Enough tile to test persistent scheduler + ], +) +# @pytest.mark.parametrize('seqlen_q,seqlen_k', [(256, 128)]) +def test_flash_attn_kvcache( + seqlen_q, + seqlen_k, + d, + varlen_q, + has_batch_idx, + has_leftpad, + page_size, + rotary_fraction, + rotary_interleaved, + has_rotary_seqlens, + seqlen_new_eq_seqlen_q, + causal, + local, + new_kv, + has_learnable_sink, + mha_type, + dtype, +): + if page_size is not None and seqlen_k % page_size != 0: + pytest.skip() + if seqlen_q > seqlen_k and new_kv: + pytest.skip() + if not new_kv and rotary_fraction > 0.0: + pytest.skip() + if rotary_fraction == 0.0 and has_rotary_seqlens: + pytest.skip() + device = "cuda" + # set seed + torch.random.manual_seed(0) + batch_size = 5 + # batch_size = 1 + batch_size_cache = batch_size if not has_batch_idx else batch_size * 2 + nheads = 6 + # nheads = 1 + # rotary_dim must be a multiple of 16, and must be <= d + rotary_dim = math.floor(int(rotary_fraction * d) / 16) * 16 + nheads_k = nheads if mha_type == "mha" else (1 if mha_type == "mqa" else 3) + assert nheads % nheads_k == 0 + dtype_ref = torch.bfloat16 if dtype == torch.float8_e4m3fn else dtype + # dv_vals = [128, d] if d > 128 and d <= 192 else ([256, 512, d] if d <= 64 else [d]) + dv_vals = [d] + if dtype == torch.float8_e4m3fn: + dv_vals = [d] + # attention_chunk_vals = [torch.randint(1, seqlen_k * 2, (1,)).item(), 0] if (causal or local) else [0] + attention_chunk_vals = [0] + for dv, attention_chunk in itertools.product(dv_vals, attention_chunk_vals): + # has_qv = d == 64 and dv >= 256 + has_qv = False + q = ( + torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype_ref) + .to(dtype) + .to(dtype_ref) + ) + if has_qv: + qv = ( + torch.randn( + batch_size, seqlen_q, nheads, dv, device=device, dtype=dtype_ref + ) + .to(dtype) + .to(dtype_ref) + ) + else: + qv = None + if varlen_q: + query_padding_mask = generate_random_padding_mask( + seqlen_q, batch_size, device, mode="random" + ) + q_unpad, indices_q, cu_seqlens_q, max_seqlen_q, *rest = unpad_input( + q, query_padding_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 has_qv else None + ) + else: + query_padding_mask = None + q_unpad = q + qv_unpad = qv + cu_seqlens_q, max_seqlen_q = None, None + # Put window_size after QKV randn so that window_size changes from test to test + window_size = ( + (None, None) if not local else torch.randint(0, seqlen_k, (2,)).tolist() + ) + if has_learnable_sink: + learnable_sink = torch.randn(nheads, dtype=torch.bfloat16, device=device) + else: + learnable_sink = None + + seqlen_new = ( + seqlen_q + if seqlen_new_eq_seqlen_q + else torch.randint(1, seqlen_q + 1, (1,)).item() + ) + cu_seqlens_k_new = None + key_new_padding_mask = None + if new_kv: + k = ( + torch.randn( + batch_size, seqlen_new, nheads_k, d, device=device, dtype=dtype_ref + ) + .to(dtype) + .to(dtype_ref) + ) + v = ( + torch.randn( + batch_size, seqlen_new, nheads_k, dv, device=device, dtype=dtype_ref + ) + .to(dtype) + .to(dtype_ref) + ) + if varlen_q: # k & v are also varlen + key_new_padding_mask = generate_random_padding_mask( + seqlen_new, batch_size, device, mode="random" + ) + k_unpad, indices_k, cu_seqlens_k_new, *rest = unpad_input( + k, key_new_padding_mask + ) + v_unpad, *rest = unpad_input(v, key_new_padding_mask) + else: + k_unpad, v_unpad = k, v + else: + k, v, k_unpad, v_unpad = None, None, None, None + if page_size is None: + k_cache = ( + torch.randn( + batch_size_cache, + seqlen_k, + nheads_k, + d, + device=device, + dtype=dtype_ref, + ) + .to(dtype) + .to(dtype_ref) + ) + v_cache = ( + torch.randn( + batch_size_cache, + seqlen_k, + nheads_k, + dv, + device=device, + dtype=dtype_ref, + ) + .to(dtype) + .to(dtype_ref) + ) + page_table = None + num_blocks = None + else: + ( + k_cache, + v_cache, + page_table, + k_cache_paged, + v_cache_paged, + num_blocks, + ) = _generate_block_kvcache( + seqlen_k, + page_size, + batch_size_cache, + nheads_k, + d, + dv, + device, + dtype, + dtype_ref, + ) + cache_seqlens = torch.randint( + 0 if new_kv else 1, + # If we don't use seqlen_q in the case of causal and rotary, cos/sin won't be long enough + ( + ( + seqlen_k + - (seqlen_q if (causal or local) and rotary_dim > 1 else seqlen_new) + + 1 + ) + if new_kv + else (seqlen_k + 1) + ), + (batch_size,), + dtype=torch.int32, + device=device, + ) + if has_leftpad: + cache_leftpad = torch.cat( + [ + ( + torch.randint( + 0, + cache_seqlens[i].item(), + (1,), + dtype=torch.int32, + device=device, + ) + if cache_seqlens[i].item() > 0 + else torch.zeros(1, dtype=torch.int32, device=device) + ) + for i in range(batch_size) + ] + ) + else: + cache_leftpad = None + if has_batch_idx: + cache_batch_idx = torch.randperm( + batch_size_cache, dtype=torch.int32, device=device + )[:batch_size] + else: + cache_batch_idx = None + arange = rearrange(torch.arange(seqlen_k, device=device), "s -> 1 s") + cache_seqlens_expanded = rearrange(cache_seqlens, "b -> b 1") + if not new_kv: + key_padding_mask = arange < cache_seqlens_expanded + else: + k_new_seqlens = ( + key_new_padding_mask.sum(-1, keepdims=True) if varlen_q else seqlen_new + ) + key_padding_mask = arange < cache_seqlens_expanded + k_new_seqlens + if has_leftpad: + key_padding_mask = torch.logical_and( + key_padding_mask, + arange >= cache_leftpad.unsqueeze(-1).expand(-1, seqlen_k), + ) + # cache_seqlens = torch.tensor([64], dtype=torch.int32, device=device) + rotary_seqlens = cache_seqlens if not has_rotary_seqlens else cache_seqlens // 2 + if rotary_dim > 0: + angle = ( + torch.rand( + seqlen_k if page_size is None else num_blocks * page_size, + rotary_dim // 2, + device=device, + ) + * 2 + * math.pi + ) + cos = torch.cos(angle).to(dtype=dtype_ref).to(dtype).to(dtype_ref) + sin = torch.sin(angle).to(dtype=dtype_ref).to(dtype).to(dtype_ref) + if causal or local: + q_ro = apply_rotary_emb( + q, + cos, + sin, + seqlen_offsets=rotary_seqlens, + interleaved=rotary_interleaved, + ) + else: + q_ro = rearrange( + apply_rotary_emb( + rearrange(q, "b s h d -> b 1 (s h) d"), + cos, + sin, + seqlen_offsets=rotary_seqlens, + interleaved=rotary_interleaved, + ), + "b 1 (s h) d -> b s h d", + s=seqlen_q, + ) + # q_ro = q + k_ro = apply_rotary_emb( + k, + cos, + sin, + seqlen_offsets=rotary_seqlens, + interleaved=rotary_interleaved, + ) + else: + cos, sin = None, None + q_ro, k_ro = q, k + # k_cache[:, 64:] = -1 + k_cache_ref = ( + k_cache if not has_batch_idx else k_cache[cache_batch_idx] + ).clone() + v_cache_ref = ( + v_cache if not has_batch_idx else v_cache[cache_batch_idx] + ).clone() + if new_kv: + update_mask = torch.logical_and( + cache_seqlens_expanded <= arange, + arange < cache_seqlens_expanded + k_new_seqlens, + ) + k_to_update = rearrange(k_ro, "b s ... -> (b s) ...") + v_to_update = rearrange(v, "b s ... -> (b s) ...") + if varlen_q: + k_to_update = k_to_update[indices_k] + v_to_update = v_to_update[indices_k] + k_cache_ref[update_mask] = k_to_update + v_cache_ref[update_mask] = v_to_update + k_cache_rep = repeat( + k_cache_ref, "b s h d -> b s (h g) d", g=nheads // nheads_k + ) + v_cache_rep = repeat( + v_cache_ref, "b s h d -> b s (h g) d", g=nheads // nheads_k + ) + out_ref, _ = attention_ref( + q_ro, + k_cache_rep, + v_cache_rep, + query_padding_mask, + key_padding_mask, + causal=causal, + qv=qv, + window_size=window_size, + learnable_sink=learnable_sink, + attention_chunk=attention_chunk, + key_leftpad=cache_leftpad, + ) + out_pt, _ = attention_ref( + q_ro, + k_cache_rep, + v_cache_rep, + query_padding_mask, + key_padding_mask, + causal=causal, + qv=qv, + window_size=window_size, + learnable_sink=learnable_sink, + attention_chunk=attention_chunk, + upcast=False, + reorder_ops=True, + key_leftpad=cache_leftpad, + intermediate_dtype=dtype if dtype == torch.float8_e4m3fn else None, + ) + q = q.to(dtype) + q_unpad = q_unpad.to(dtype) if varlen_q else None + k_cache = k_cache.to(dtype) + v_cache = v_cache.to(dtype) + k_cache_paged = k_cache_paged.to(dtype) if page_size is not None else None + v_cache_paged = v_cache_paged.to(dtype) if page_size is not None else None + k = k.to(dtype) if k is not None else None + v = v.to(dtype) if v is not None else None + k_unpad = k_unpad.to(dtype) if k_unpad is not None else None + v_unpad = v_unpad.to(dtype) if v_unpad is not None else None + qv = qv.to(dtype) if qv is not None else None + qv_unpad = qv_unpad.to(dtype) if (varlen_q and qv is not None) else None + cos = cos.to(dtype) if cos is not None else None + sin = sin.to(dtype) if sin is not None else None + k_cache_saved = k_cache.clone() if page_size is None else k_cache_paged.clone() + v_cache_saved = v_cache.clone() if page_size is None else v_cache_paged.clone() + # num_splits_vals = [1, 0] + num_splits_vals = [1] + # precompute_metadata_vals = [False, True] + precompute_metadata_vals = [False] + for num_splits, precompute_metadata in itertools.product( + num_splits_vals, precompute_metadata_vals + ): + # if precompute_metadata: + # scheduler_metadata = get_scheduler_metadata( + # batch_size, max_seqlen_q if varlen_q else seqlen_q, seqlen_k, nheads, nheads_k, d, + # cache_seqlens, q.dtype, headdim_v=dv, cu_seqlens_q=cu_seqlens_q, + # cu_seqlens_k_new=cu_seqlens_k_new, cache_leftpad=cache_leftpad, + # max_seqlen_k_new=seqlen_new, page_size=page_size, + # causal=causal, window_size=window_size, attention_chunk=attention_chunk, + # num_splits=num_splits + # ) + # else: + # scheduler_metadata = None + scheduler_metadata = None + # Repeat to test metadata reuse + for _ in range(1 if not precompute_metadata else 2): + if page_size is None: + k_cache.copy_(k_cache_saved) + v_cache.copy_(v_cache_saved) + else: + k_cache_paged.copy_(k_cache_saved) + v_cache_paged.copy_(v_cache_saved) + # For FA4, use flash_attn_varlen_func directly instead of flash_attn_with_kvcache + # This matches the pattern from the original FA4 test + out, lse = flash_attn_varlen_func( + q if not varlen_q else q_unpad, + k_cache if page_size is None else k_cache_paged, + v_cache if page_size is None else v_cache_paged, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=None, # FA4 doesn't use cu_seqlens_k for KV cache + # max_seqlen_q and max_seqlen_k not needed for FA4 + seqused_k=cache_seqlens, # Use cache_seqlens as seqused_k + page_table=page_table, + causal=causal, + window_size=window_size, + sinks=learnable_sink, # FA4 uses learnable_sink, not sinks + softcap=0.0, + pack_gqa=None, + return_softmax_lse=True, + ) + if varlen_q: + out = output_pad_fn(out) + # out = flash_attn_with_kvcache( + # q, k_cache, v_cache, cache_seqlens=cache_seqlens, causal=causal, window_size=window_size + # ) + # out = flash_attn_with_kvcache(q, k_cache, v_cache, causal=causal, window_size=window_size) + # qk = torch.einsum("bqhd,bkhd->bhqk", q, k_cache_ref) + # m = qk.amax(-1, keepdim=True) + # s_tmp = torch.exp((qk - m) / math.sqrt(d)) + # o1 = torch.einsum('bhst,bthd->bshd', s_tmp, v_cache_ref) + # lse_ref = torch.logsumexp(qk / math.sqrt(d), -1) + # probs = torch.softmax(qk, dim=-1) + print(f"Output max diff: {(out - out_ref).abs().max().item()}") + print(f"Output mean diff: {(out - out_ref).abs().mean().item()}") + print(f"Pytorch max diff: {(out_pt - out_ref).abs().max().item()}") + print(f"Pytorch mean diff: {(out_pt - out_ref).abs().mean().item()}") + # breakpoint() + + # Check that FlashAttention's numerical error is at most twice the numerical error + # of a Pytorch implementation. + if new_kv: + if page_size is None: + k_cache_select = ( + k_cache.to(dtype_ref) + if not has_batch_idx + else k_cache.to(dtype_ref)[cache_batch_idx] + ) + v_cache_select = ( + v_cache.to(dtype_ref) + if not has_batch_idx + else v_cache.to(dtype_ref)[cache_batch_idx] + ) + else: + k_cache_select = rearrange( + k_cache_paged.to(dtype_ref)[ + ( + page_table + if not has_batch_idx + else page_table[cache_batch_idx] + ).flatten() + ], + "(b nblocks) block_size ... -> b (nblocks block_size) ...", + b=batch_size, + )[:, :seqlen_k].to(dtype_ref) + v_cache_select = rearrange( + v_cache_paged.to(dtype_ref)[ + ( + page_table + if not has_batch_idx + else page_table[cache_batch_idx] + ).flatten() + ], + "(b nblocks) block_size ... -> b (nblocks block_size) ...", + b=batch_size, + )[:, :seqlen_k].to(dtype_ref) + k_cache_ref = k_cache_ref.to(dtype).to(dtype_ref) + v_cache_ref = v_cache_ref.to(dtype).to(dtype_ref) + if dtype is not torch.float8_e4m3fn: + assert torch.equal(v_cache_select, v_cache_ref) + else: + assert torch.allclose( + v_cache_select, v_cache_ref, rtol=1e-3, atol=1e-3 + ) + # breakpoint() + # if rotary_dim == 0 and dtype is not torch.float8_e4m3fn: + if rotary_dim == 0: + assert torch.equal(k_cache_select, k_cache_ref) + else: + # if not torch.allclose(k_cache_select, k_cache_ref, rtol=1e-3, atol=1e-3): + # breakpoint() + if dtype is not torch.float8_e4m3fn: + assert torch.allclose( + k_cache_select, k_cache_ref, rtol=1e-3, atol=1e-3 + ) + else: + assert torch.allclose( + k_cache_select, k_cache_ref, rtol=1e-1, atol=1e-1 + ) + mult = 4 if dtype == torch.float8_e4m3fn else 2 + assert (out - out_ref).abs().max().item() <= mult * ( + out_pt - out_ref + ).abs().max().item() + 1e-5 + mult_mean = 3 if dtype == torch.float8_e4m3fn else 1.5 + assert (out - out_ref).abs().mean().item() <= mult_mean * ( + out_pt - out_ref + ).abs().mean().item() + + +def _generate_block_kvcache( + seqlen_k, page_size, batch_size, nheads_k, d, dv, device, dtype, dtype_ref +): + num_blocks = math.ceil(seqlen_k / page_size) * batch_size * 3 + k_cache_paged = ( + torch.randn(num_blocks, page_size, nheads_k, d, device=device, dtype=dtype_ref) + .to(dtype) + .to(dtype_ref) + ) + v_cache_paged = ( + torch.randn(num_blocks, page_size, nheads_k, dv, device=device, dtype=dtype_ref) + .to(dtype) + .to(dtype_ref) + ) + page_table = rearrange( + torch.randperm(num_blocks, dtype=torch.int32, device=device), + "(b nblocks) -> b nblocks", + b=batch_size, + ) + k_cache = rearrange( + k_cache_paged[page_table.flatten()], + "(b nblocks) block_size ... -> b (nblocks block_size) ...", + b=batch_size, + )[:, :seqlen_k] + v_cache = rearrange( + v_cache_paged[page_table.flatten()], + "(b nblocks) block_size ... -> b (nblocks block_size) ...", + b=batch_size, + )[:, :seqlen_k] + return k_cache, v_cache, page_table, k_cache_paged, v_cache_paged, num_blocks + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py index af84009e1..4fbe4f0c7 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py @@ -16,9 +16,15 @@ from sglang.multimodal_gen.runtime.platforms import ( try: from sgl_kernel.flash_attn import flash_attn_varlen_func - # flash_attn 3 no longer have a different API, see following commit: - # https://github.com/Dao-AILab/flash-attention/commit/ed209409acedbb2379f870bbd03abce31a7a51b7 - flash_attn_func = flash_attn_varlen_func + from sglang.jit_kernel.flash_attention_v4 import ( + flash_attn_varlen_func as flash_attn_varlen_func_fa4, + ) + + def flash_attn_func(*args, ver: int = 3, **kwargs): + if ver == 4: + return flash_attn_varlen_func_fa4(*args, **kwargs) + return flash_attn_varlen_func(*args, ver=ver, **kwargs) + except ImportError as e: raise e diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 367dff9b7..d03b0d4bd 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -22,7 +22,18 @@ if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner from sgl_kernel import merge_state_v2 -from sgl_kernel.flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache +from sgl_kernel.flash_attn import flash_attn_varlen_func as flash_attn_varlen_func_fa3 +from sgl_kernel.flash_attn import flash_attn_with_kvcache as flash_attn_with_kvcache_fa3 + +flash_attn_varlen_func = flash_attn_varlen_func_fa3 +flash_attn_with_kvcache = flash_attn_with_kvcache_fa3 + +from sglang.jit_kernel.flash_attention_v4 import ( + flash_attn_varlen_func as flash_attn_varlen_func_fa4, +) +from sglang.jit_kernel.flash_attention_v4 import ( + flash_attn_with_kvcache as flash_attn_with_kvcache_fa4, +) @dataclass @@ -803,10 +814,21 @@ class FlashAttentionBackend(AttentionBackend): and not is_swa_layer ) - # For fa3 interface version compatibility, we put new fields into conditional keyword args + flash_attn_varlen_func_base = flash_attn_varlen_func_fa3 + flash_attn_with_kvcache_base = flash_attn_with_kvcache_fa3 + + flash_attn_varlen_func = ( + flash_attn_varlen_func_fa4 + if self.fa_impl_ver == 4 + else flash_attn_varlen_func_base + ) + flash_attn_with_kvcache = ( + flash_attn_with_kvcache_fa4 + if self.fa_impl_ver == 4 + else flash_attn_with_kvcache_base + ) + kwargs = {} - if self.fa_impl_ver != 3: - kwargs["ver"] = self.fa_impl_ver if sinks is not None: kwargs["sinks"] = sinks @@ -1110,10 +1132,7 @@ class FlashAttentionBackend(AttentionBackend): if layer.is_cross_attention or layer.attn_type == AttentionType.ENCODER_ONLY: causal = False - # For fa3 interface version compatibility, we put new fields into conditional keyword args kwargs = {} - if self.fa_impl_ver != 3: - kwargs["ver"] = self.fa_impl_ver if sinks is not None: kwargs["sinks"] = sinks