leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

v4.3.1 update. (#2817)

Browse Source
This commit is contained in:
Junkai-Wu
2025-11-27 22:49:30 +08:00
committed by GitHub
parent 2052fd3885
commit 1de3a576cc
44 changed files with 3316 additions and 510 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
examples/python/CuTeDSL/ampere/hstu_attention.py
View File

@@ -330,7 +330,10 @@ class HSTUAttentionForwardAmpere(object):
if cutlass.const_expr(self._is_causal):
n_block = (
cute.ceil_div((m_block + 1) * self._m_block_size, self._n_block_size)
cute.ceil_div(
min((m_block + 1) * self._m_block_size, mK.shape[1]),
self._n_block_size,
)
- 1
) # for causal case, only process the first n_block tiles
else:
@@ -652,7 +655,7 @@ class HSTUAttentionForwardAmpere(object):
# m residue handling for RAB
for m in cutlass.range_constexpr(cute.size(tRABcRAB.shape[1])):
if cute.elem_less(
tRABcRAB[0, m, 0, n_block][1], mRAB.layout.shape[2]
tRABcRAB[0, m, 0, n_block_idx - 1][1], mRAB.layout.shape[2]
):
cute.copy(
gmem_tiled_copy_QKV,

16
examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py
View File

@@ -1977,33 +1977,21 @@ class BlockwiseGemmKernel:
tcgen05.copy.Ld16x256bOp(tcgen05.copy.Repetition(8)),
self.acc_dtype,
)
elif cutlass.const_expr(self.mma_tiler[0] == 128):
else:
tmem_load_atom = cute.make_copy_atom(
tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)),
self.acc_dtype,
)
else:
# default: 16dp
tmem_load_atom = cute.make_copy_atom(
tcgen05.copy.Ld16x256bOp(tcgen05.copy.Repetition(1)),
self.acc_dtype,
)
if cutlass.const_expr(self.mma_tiler[0] == 64):
tmem_store_atom = cute.make_copy_atom(
tcgen05.copy.St16x256bOp(tcgen05.copy.Repetition(8)),
self.acc_dtype,
)
elif cutlass.const_expr(self.mma_tiler[0] == 128):
else:
tmem_store_atom = cute.make_copy_atom(
tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)),
self.acc_dtype,
)
else:
# default: 16dp
tmem_store_atom = cute.make_copy_atom(
tcgen05.copy.St16x256bOp(tcgen05.copy.Repetition(1)),
self.acc_dtype,
)
tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0, None)], epi_tile)
tAcc_final_epi = cute.flat_divide(

16
examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py
View File

@@ -2010,33 +2010,21 @@ class BlockwiseContiguousGroupedGemmKernel:
tcgen05.copy.Ld16x256bOp(tcgen05.copy.Repetition(8)),
self.acc_dtype,
)
elif cutlass.const_expr(self.mma_tiler[0] == 128):
else:
tmem_load_atom = cute.make_copy_atom(
tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)),
self.acc_dtype,
)
else:
# default: 16dp
tmem_load_atom = cute.make_copy_atom(
tcgen05.copy.Ld16x256bOp(tcgen05.copy.Repetition(1)),
self.acc_dtype,
)
if cutlass.const_expr(self.mma_tiler[0] == 64):
tmem_store_atom = cute.make_copy_atom(
tcgen05.copy.St16x256bOp(tcgen05.copy.Repetition(8)),
self.acc_dtype,
)
elif cutlass.const_expr(self.mma_tiler[0] == 128):
else:
tmem_store_atom = cute.make_copy_atom(
tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)),
self.acc_dtype,
)
else:
# default: 16dp
tmem_store_atom = cute.make_copy_atom(
tcgen05.copy.St16x256bOp(tcgen05.copy.Repetition(1)),
self.acc_dtype,
)
tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0, None)], epi_tile)
tAcc_final_epi = cute.flat_divide(

16
examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py
View File

@@ -2010,33 +2010,21 @@ class BlockwiseMaskedGroupedGemmKernel:
tcgen05.copy.Ld16x256bOp(tcgen05.copy.Repetition(8)),
self.acc_dtype,
)
elif cutlass.const_expr(self.mma_tiler[0] == 128):
else:
tmem_load_atom = cute.make_copy_atom(
tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)),
self.acc_dtype,
)
else:
# default: 16dp
tmem_load_atom = cute.make_copy_atom(
tcgen05.copy.Ld16x256bOp(tcgen05.copy.Repetition(1)),
self.acc_dtype,
)
if cutlass.const_expr(self.mma_tiler[0] == 64):
tmem_store_atom = cute.make_copy_atom(
tcgen05.copy.St16x256bOp(tcgen05.copy.Repetition(8)),
self.acc_dtype,
)
elif cutlass.const_expr(self.mma_tiler[0] == 128):
else:
tmem_store_atom = cute.make_copy_atom(
tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)),
self.acc_dtype,
)
else:
# default: 16dp
tmem_store_atom = cute.make_copy_atom(
tcgen05.copy.St16x256bOp(tcgen05.copy.Repetition(1)),
self.acc_dtype,
)
tAcc_epi = cute.flat_divide(tAcc[((None, None), 0, 0, None)], epi_tile)
tAcc_final_epi = cute.flat_divide(

2
examples/python/CuTeDSL/blackwell/fmha.py
View File

@@ -247,7 +247,7 @@ class BlackwellFusedMultiHeadAttentionForward:
k_iter: cute.Pointer,
v_iter: cute.Pointer,
o_iter: cute.Pointer,
problem_size: Tuple[Int32, Int32, Int32, Int32, Int32, Int32],
problem_size: Tuple[Int32, Int32, Int32, Int32, Int32, Int32, Int32],
cum_seqlen_q: Optional[cute.Tensor],
cum_seqlen_k: Optional[cute.Tensor],
lse_iter: Optional[cute.Pointer],

2015
examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py Normal file
View File

File diff suppressed because it is too large Load Diff

1
examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py
View File

@@ -77,6 +77,7 @@ def main():
# compile the kernel with "--enable-tvm-ffi" option
compiled_add_one = cute.compile(add_one, a_cute, b_cute, options="--enable-tvm-ffi")
os.makedirs("./build", exist_ok=True)
object_file_path = "./build/add_one.o"
lib_path = "./build/add_one.so"
compiled_add_one.export_to_c(object_file_path, function_name="add_one")