fa4 cleanup (#19727)
This commit is contained in:
Rain Jiang
@@ -59,6 +59,7 @@ dependencies = [
|
||||
"scipy",
|
||||
"sentencepiece",
|
||||
"setproctitle",
|
||||
"sgl-fa4==4.0.3",
|
||||
"sgl-kernel==0.3.21",
|
||||
"soundfile==0.13.1",
|
||||
"tiktoken",
|
||||
|
||||
@@ -1,4 +0,0 @@
|
||||
[flake8]
|
||||
max-line-length = 100
|
||||
# W503: line break before binary operator
|
||||
ignore = E731, E741, F841, W503
|
||||
@@ -1,5 +0,0 @@
|
||||
Tri Dao, tri@tridao.me
|
||||
Jay Shah
|
||||
Ted Zadouri
|
||||
Markus Hoehnerbach
|
||||
Vijay Thakkar
|
||||
@@ -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.
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
@@ -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 = {}
|
||||
@@ -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
|
||||
@@ -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())
|
||||
@@ -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
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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],
|
||||
)
|
||||
@@ -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
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -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],
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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]
|
||||
@@ -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<B,M,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
|
||||
@@ -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()
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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),
|
||||
)
|
||||
@@ -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
|
||||
]
|
||||
@@ -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)
|
||||
@@ -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]
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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<b,m,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<b,m,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]
|
||||
@@ -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)]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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(
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user