diff --git a/python/pyproject.toml b/python/pyproject.toml index e0c6e8863..fa5c705fa 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -59,6 +59,7 @@ dependencies = [ "scipy", "sentencepiece", "setproctitle", + "sgl-fa4==4.0.3", "sgl-kernel==0.3.21", "soundfile==0.13.1", "tiktoken", diff --git a/python/sglang/jit_kernel/flash_attention/cute/.flake8 b/python/sglang/jit_kernel/flash_attention/cute/.flake8 deleted file mode 100644 index bae5b85c0..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/.flake8 +++ /dev/null @@ -1,4 +0,0 @@ -[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 deleted file mode 100644 index bc3991c67..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/AUTHORS +++ /dev/null @@ -1,5 +0,0 @@ -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 deleted file mode 100644 index 5860e4b33..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/LICENSE +++ /dev/null @@ -1,29 +0,0 @@ -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 deleted file mode 100644 index e69de29bb..000000000 diff --git a/python/sglang/jit_kernel/flash_attention/cute/__init__.py b/python/sglang/jit_kernel/flash_attention/cute/__init__.py deleted file mode 100644 index 7706f746a..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -"""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 deleted file mode 100644 index e3072d8ce..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/ampere_helpers.py +++ /dev/null @@ -1,103 +0,0 @@ -# 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 deleted file mode 100644 index c999b1801..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/barrier.py +++ /dev/null @@ -1,71 +0,0 @@ -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 deleted file mode 100644 index 9a7820e7b..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/benchmark.py +++ /dev/null @@ -1,268 +0,0 @@ -# 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 deleted file mode 100644 index e39ac5eeb..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/blackwell_helpers.py +++ /dev/null @@ -1,753 +0,0 @@ -# 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 deleted file mode 100644 index a5a2544a7..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/block_info.py +++ /dev/null @@ -1,108 +0,0 @@ -# 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 deleted file mode 100644 index 15a606df0..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/block_sparse_utils.py +++ /dev/null @@ -1,1451 +0,0 @@ -""" -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 deleted file mode 100644 index 99aa69083..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/block_sparsity.py +++ /dev/null @@ -1,250 +0,0 @@ -""" -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 deleted file mode 100644 index 062ecb288..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/compute_block_sparsity.py +++ /dev/null @@ -1,377 +0,0 @@ -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 deleted file mode 100644 index cfdcbdb80..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/copy_utils.py +++ /dev/null @@ -1,340 +0,0 @@ -# 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 deleted file mode 100644 index 34ae4594d..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/cute_dsl_utils.py +++ /dev/null @@ -1,146 +0,0 @@ -# Copyright (c) 2025, Tri Dao. - -import os -import pathlib -from typing import Tuple -from functools import partial -from dataclasses import dataclass, fields - -import torch - -from sglang.jit_kernel.utils import cache_once - -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, -} - - -@cache_once -def get_max_active_clusters(cluster_size): - return cutlass.utils.HardwareInfo().get_max_active_clusters(cluster_size=cluster_size) - - -@cache_once -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 deleted file mode 100644 index c56ea89e7..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/fast_math.py +++ /dev/null @@ -1,21 +0,0 @@ -# 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 deleted file mode 100644 index 1ac795bcf..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/flash_bwd.py +++ /dev/null @@ -1,1262 +0,0 @@ -# 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 deleted file mode 100644 index 3850a80ff..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_postprocess.py +++ /dev/null @@ -1,463 +0,0 @@ -# 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 deleted file mode 100644 index b677b140b..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_preprocess.py +++ /dev/null @@ -1,365 +0,0 @@ -# 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 deleted file mode 100644 index 272655e54..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_sm100.py +++ /dev/null @@ -1,2950 +0,0 @@ -# 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 deleted file mode 100644 index 01b43531b..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/flash_bwd_sm90.py +++ /dev/null @@ -1,1706 +0,0 @@ -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 deleted file mode 100644 index 19e300bb9..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/flash_fwd.py +++ /dev/null @@ -1,2513 +0,0 @@ -# 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, - page_size: Optional[int] = None, - ): - """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 - self.page_size = page_size - 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.tiles_per_page = self.page_size // self.tile_n if cutlass.const_expr(self.page_size is not None) else None - 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, - mPageTable, - tma_atom_Q, - tma_atom_K, - tma_atom_V, - tma_atom_O, - softmax_scale_log2, - softmax_scale, - window_size_left, - window_size_right, - learnable_sink, - blocksparse_tensors, - self.sQ_layout, - self.sK_layout, - self.sV_layout, - self.sO_layout, - self.sP_layout, - self.gmem_tiled_copy_Q, - self.gmem_tiled_copy_K, - self.gmem_tiled_copy_V, - self.gmem_tiled_copy_O, - tiled_mma_qk, - tiled_mma_pv, - 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], - mPageTable: Optional[cute.Tensor], - tma_atom_Q: Optional[cute.CopyAtom], - tma_atom_K: Optional[cute.CopyAtom], - tma_atom_V: Optional[cute.CopyAtom], - tma_atom_O: Optional[cute.CopyAtom], - softmax_scale_log2: Float32, - softmax_scale: Optional[Float32], - window_size_left: Optional[Int32], - window_size_right: Optional[Int32], - learnable_sink: Optional[cute.Tensor], - blocksparse_tensors: Optional[BlockSparseTensors], - sQ_layout: cute.ComposedLayout, - sK_layout: cute.ComposedLayout, - sV_layout: cute.ComposedLayout, - sO_layout: cute.ComposedLayout, - sP_layout: cute.ComposedLayout | None, - gmem_tiled_copy_Q: cute.TiledCopy, - gmem_tiled_copy_K: cute.TiledCopy, - gmem_tiled_copy_V: cute.TiledCopy, - gmem_tiled_copy_O: cute.TiledCopy, - tiled_mma_qk: cute.TiledMma, - tiled_mma_pv: cute.TiledMma, - 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] 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.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, - mPageTable, - 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, - mPageTable: Optional[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 - 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((seqlen.offset_k, 0), mV[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)) - else: - mK_cur, mV_cur = [t[None, None, head_idx_kv, None] for t in (mK, mV)] - gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (None, 0, None)) - gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (None, 0, None)) - gK = cute.group_modes(gK, 2, 4) - gV = cute.group_modes(gV, 2, 4) - 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 = self.get_n_block(batch_idx, n_block_max - 1, mPageTable) - 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 = self.get_n_block(batch_idx, n_block_max - 1 - i - 1, mPageTable) - 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 = self.get_n_block(batch_idx, n_block_max - i - 1, mPageTable) - n_block = self.get_n_block(batch_idx, n_block_max - i - 2, mPageTable) - 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 = self.get_n_block(batch_idx, n_block_min, mPageTable) - 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 get_n_block( - self, - batch_idx: int, - n_block: int, - mPageTable: Optional[cute.Tensor], - ): - if cutlass.const_expr(mPageTable is not None): - page_idx = mPageTable[batch_idx, n_block // self.tiles_per_page] - residue = n_block % self.tiles_per_page - n_block = page_idx * self.tiles_per_page + residue - return n_block - - @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 deleted file mode 100644 index b6f4acf51..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/flash_fwd_combine.py +++ /dev/null @@ -1,704 +0,0 @@ -# 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 deleted file mode 100644 index b01b1b33f..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/flash_fwd_sm100.py +++ /dev/null @@ -1,2740 +0,0 @@ -# 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 deleted file mode 100644 index c6a1c3019..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/hopper_helpers.py +++ /dev/null @@ -1,101 +0,0 @@ -# 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 deleted file mode 100644 index 141228da4..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/interface.py +++ /dev/null @@ -1,1762 +0,0 @@ -# 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 typing import Optional, Tuple, Callable - -import torch - - -from sglang.jit_kernel.utils import cache_once - - -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, -) - -@cache_once -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, page_size, - 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_size is None or page_size % n_block_size == 0, f"Only page_size values that are multiples of {n_block_size} are supported for paged KV 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, - page_size=page_size, - ) - 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, - return_lse=True, - ) - 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 deleted file mode 100644 index 170598244..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/mask.py +++ /dev/null @@ -1,651 +0,0 @@ -# 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 deleted file mode 100644 index 16336c346..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/mma_sm100_desc.py +++ /dev/null @@ -1,291 +0,0 @@ -# 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 deleted file mode 100644 index 777c44079..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/named_barrier.py +++ /dev/null @@ -1,31 +0,0 @@ -# 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 deleted file mode 100644 index 02c89df6e..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/pack_gqa.py +++ /dev/null @@ -1,164 +0,0 @@ -# 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 deleted file mode 100644 index d45824222..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/paged_kv.py +++ /dev/null @@ -1,214 +0,0 @@ -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 deleted file mode 100644 index 54981bca1..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/pipeline.py +++ /dev/null @@ -1,272 +0,0 @@ -# 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 deleted file mode 100644 index 1503556c1..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/pyproject.toml +++ /dev/null @@ -1,56 +0,0 @@ -[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 deleted file mode 100644 index 6d8c6feb2..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/seqlen_info.py +++ /dev/null @@ -1,138 +0,0 @@ -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 deleted file mode 100644 index 378932a96..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/softmax.py +++ /dev/null @@ -1,582 +0,0 @@ -# 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 deleted file mode 100644 index 2897e64fc..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/testing.py +++ /dev/null @@ -1,423 +0,0 @@ -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 deleted file mode 100644 index c11f18ca9..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/tile_scheduler.py +++ /dev/null @@ -1,719 +0,0 @@ -# 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 deleted file mode 100644 index f31d85c5d..000000000 --- a/python/sglang/jit_kernel/flash_attention/cute/utils.py +++ /dev/null @@ -1,859 +0,0 @@ -# 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 index 6817a1d54..8958ae1c0 100644 --- a/python/sglang/jit_kernel/flash_attention_v4.py +++ b/python/sglang/jit_kernel/flash_attention_v4.py @@ -5,9 +5,7 @@ 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, - ) + from sgl_fa4.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 @@ -46,7 +44,7 @@ def flash_attn_varlen_func( 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." + "sgl_fa4.cute). Please check your source tree." ) from _flash_attn_import_error q, k, v = [_maybe_contiguous(t) for t in (q, k, v)] 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 4fbe4f0c7..72100c65b 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 @@ -23,7 +23,7 @@ try: 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) + return flash_attn_varlen_func(*args, **kwargs) except ImportError as e: raise e diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 56c2be07e..31ad7c010 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -42,8 +42,6 @@ from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.utils import is_cuda, is_hip -# from sgl_kernel.flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache - if TYPE_CHECKING: from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.model_executor.model_runner import ModelRunner @@ -1783,8 +1781,6 @@ class NativeSparseAttnBackend( ) # Use FA3 for SM90 (Hopper/H200) - fa_version = 3 - return flash_attn_varlen_func( q=q, k=k, @@ -1795,7 +1791,6 @@ class NativeSparseAttnBackend( max_seqlen_k=max_seqlen_k, softmax_scale=layer.scaling, causal=causal, - ver=fa_version, ) def _forward_tilelang( diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index a8b25ecf6..4abe5b2c4 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -35,7 +35,22 @@ _is_hip = is_hip() if _is_cuda: from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache - from sgl_kernel.flash_attn import flash_attn_varlen_func + + try: + from sgl_kernel.flash_attn import 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, **kwargs) + + except ImportError as e: + raise e + if _is_npu: import torch_npu @@ -391,7 +406,7 @@ class VisionFlash3Attention(nn.Module): """ if envs.SGLANG_VIT_ENABLE_CUDA_GRAPH.get(): max_seqlen = cu_seqlens[1] - output = flash_attn_varlen_func( + output = flash_attn_func( q, k, v, @@ -406,7 +421,7 @@ class VisionFlash3Attention(nn.Module): seq_lens = cu_seqlens[1:] - cu_seqlens[:-1] max_seqlen = seq_lens.max().item() - output = flash_attn_varlen_func( + output = flash_attn_func( q, k, v, @@ -457,7 +472,7 @@ class VisionFlash4Attention(nn.Module): seq_lens = cu_seqlens[1:] - cu_seqlens[:-1] max_seqlen = seq_lens.max().item() - output = flash_attn_varlen_func( + output = flash_attn_func( q, k, v, diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index ec9f10ef0..53688ec86 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -570,44 +570,3 @@ install(DIRECTORY "${repo-triton_SOURCE_DIR}/python/triton_kernels/triton_kernel DESTINATION "triton_kernels" PATTERN ".git*" EXCLUDE PATTERN "__pycache__" EXCLUDE) - -# ============================ Extra Install: FA4 ============================= # -# TODO: find a better install condition. -if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_SM100A) - - set(FLASH_ATTN_CUTE_SRC "${repo-flash-attention_SOURCE_DIR}/flash_attn/cute") - set(FLASH_ATTN_CUTE_DST "${CMAKE_CURRENT_BINARY_DIR}/flash_attn_origin/cute") - - file(MAKE_DIRECTORY "${FLASH_ATTN_CUTE_DST}") - - file(COPY "${FLASH_ATTN_CUTE_SRC}/" - DESTINATION "${FLASH_ATTN_CUTE_DST}" - PATTERN ".git*" EXCLUDE - PATTERN "__pycache__" EXCLUDE) - - file(GLOB_RECURSE FLASH_ATTN_CUTE_DST_PY - "${FLASH_ATTN_CUTE_DST}/*.py") - - foreach(FILE_PATH IN LISTS FLASH_ATTN_CUTE_DST_PY) - file(READ "${FILE_PATH}" FILE_CONTENT) - - set(MODIFIED_CONTENT "${FILE_CONTENT}") - - # The main goal is to avoid using "flash_attn" so that other libraries (such as transformers) do not mistakenly assume that "flash_attn" is already installed. - - string(REPLACE "flash_attn.cute" - "flash_attn_origin.cute" - MODIFIED_CONTENT "${MODIFIED_CONTENT}") - - if (NOT FILE_CONTENT STREQUAL MODIFIED_CONTENT) - file(WRITE "${FILE_PATH}" "${MODIFIED_CONTENT}") - message(STATUS " - [FA4 Patch] Patched: ${FILE_PATH}") - endif() - endforeach() - - install(DIRECTORY "${FLASH_ATTN_CUTE_DST}/" - DESTINATION "flash_attn_origin/cute" - PATTERN ".git*" EXCLUDE - PATTERN "__pycache__" EXCLUDE) - -endif() diff --git a/sgl-kernel/python/sgl_kernel/flash_attn.py b/sgl-kernel/python/sgl_kernel/flash_attn.py index 12afdb790..4281a5e6b 100644 --- a/sgl-kernel/python/sgl_kernel/flash_attn.py +++ b/sgl-kernel/python/sgl_kernel/flash_attn.py @@ -10,11 +10,6 @@ except: "Can not import FA3 in sgl_kernel. Please check your installation." ) -try: - from ._fa4_interface import flash_attn_varlen_func as flash_attn_varlen_func_v4 -except ImportError: - flash_attn_varlen_func_v4 = None - @lru_cache(maxsize=1) def is_fa3_supported(device=None) -> bool: @@ -160,45 +155,6 @@ def flash_attn_with_kvcache( logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax normalization factor). """ - if ver == 4: - assert ( - flash_attn_varlen_func_v4 is not None - ), "FA4 is not available, please check your installation." - # Using `(-1, -1)` as no sliding window causes correctness issues for FA4. - assert ( - k is None and v is None - ), "FA4 does not support updating KV cache in-place." - assert ( - rotary_cos is None and rotary_sin is None and rotary_seqlens is None - ), "FA4 does not support rotary embedding." - assert ( - cache_batch_idx is None and cache_leftpad is None - ), "FA4 does not support non-consecutive batch indices or left padding." - assert ( - q_descale is None and k_descale is None and v_descale is None - ), "FA4 does not support descale." - - if window_size == (-1, -1): - window_size = (None, None) - - return flash_attn_varlen_func_v4( - q=q, - k=k_cache, - v=v_cache, - cu_seqlens_q=cu_seqlens_q, - seqused_k=cache_seqlens, - softmax_scale=softmax_scale, - causal=causal, - window_size=window_size, - softcap=softcap, - num_splits=num_splits, - pack_gqa=pack_gqa, - return_softmax_lse=return_softmax_lse, - learnable_sink=sinks, - page_table=page_table, - score_mod=score_mod, - aux_tensors=aux_tensors, - ) assert k_cache.stride(-1) == 1, "k_cache must have contiguous last dimension" assert v_cache.stride(-1) == 1, "v_cache must have contiguous last dimension" @@ -298,32 +254,6 @@ def flash_attn_varlen_func( aux_tensors=None, ver=3, ): - if ver == 4: - assert ( - flash_attn_varlen_func_v4 is not None - ), "FA4 is not available, please check your installation." - # Using `(-1, -1)` as no sliding window causes correctness issues for FA4. - if window_size == (-1, -1): - window_size = (None, None) - return flash_attn_varlen_func_v4( - q, - k, - v, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - seqused_q=seqused_q, - seqused_k=seqused_k, - page_table=page_table, - softmax_scale=softmax_scale, - causal=causal, - window_size=window_size, - softcap=softcap, - pack_gqa=pack_gqa, - learnable_sink=sinks, - return_softmax_lse=return_softmax_lse, - score_mod=score_mod, - aux_tensors=aux_tensors, - ) if not is_fa3_supported(): raise NotImplementedError( diff --git a/sgl-kernel/tests/test_flash_attention_4.py b/sgl-kernel/tests/test_flash_attention_4.py deleted file mode 100644 index 45a5c9cf9..000000000 --- a/sgl-kernel/tests/test_flash_attention_4.py +++ /dev/null @@ -1,1458 +0,0 @@ -# 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 -from functools import partial - -import pytest -import torch -import torch.nn.functional as F -from einops import rearrange, repeat - -try: - from flash_attn.layers.rotary import apply_rotary_emb -except ImportError: - apply_rotary_emb = None - -from sgl_kernel.flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache -from sgl_kernel.testing.rotary_embedding import _apply_rotary_emb as apply_rotary_emb - -# from utils import is_hopper # Not used in this test - -# Force sgl_kernel.flash_attn wrappers to use FA4 (Cute-DSL) implementations. -# The wrappers accept a superset of args; for FA4, extra args are ignored. -flash_attn_varlen_func = partial(flash_attn_varlen_func, ver=4) -flash_attn_with_kvcache = partial(flash_attn_with_kvcache, ver=4) - -# Skip this test on Hopper machine -skip_condition = torch.cuda.get_device_capability() < (10, 0) - - -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, - ver=4, # Use FA4 - ) - 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, - ver=4, # Use FA4 - ) - 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__])