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

v4.3.1 update. (#2817)

Browse Source
This commit is contained in:
Junkai-Wu
2025-11-27 09:49:30 -05: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
examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_blockscaled_rcgrouped.cu
+12 -7
View File
@@ -54,6 +54,7 @@
1 256x128x256
2 256x256x256 and so on
Note that one must keep m and k consistent across groups in the benchmark file.
*/
#include <iostream>
@@ -334,7 +335,6 @@ struct Options {
if (!benchmark_path.empty()) {
if (!benchmark_problems()) {
problem_sizes_host.clear();
tokens_per_expert_host.clear();
return;
}
}
@@ -378,6 +378,7 @@ struct Options {
problem_sizes_host.push_back({m, n, k});
tokens_per_expert_host.push_back(n);
}
groups = static_cast<int>(problem_sizes_host.size());
}
/// Load a benchmark
@@ -409,6 +410,7 @@ struct Options {
problem_sizes_host.push_back({extent.m(), extent.n(), extent.k()});
tokens_per_expert_host.push_back(extent.n());
}
groups = static_cast<int>(problem_sizes_host.size());
m = get<0>(problem_sizes_host.at(0));
k = get<2>(problem_sizes_host.at(0));
@@ -631,8 +633,7 @@ void initialize(const Options &options) {
/// Populates a Gemm::Arguments structure from the given commandline options
template <typename Gemm>
typename Gemm::Arguments args_from_options(Options &options)
{
typename Gemm::Arguments args_from_options(Options &options) {
cutlass::KernelHardwareInfo hw_info;
// Change device_id to another value if you are running on a machine with multiple GPUs and wish
// to use a GPU other than that with device ID 0.
@@ -749,14 +750,14 @@ bool verify(const Options &options) {
block_D.at(i).sync_host();
// Check if output from CUTLASS kernel and reference kernel are equal or not
passed &= cutlass::reference::host::TensorEquals(block_ref_D.at(i).host_view(), block_D.at(i).host_view());
}
return passed;
}
/// Execute a given example GEMM computation
template <typename Gemm>
int run(Options &options, bool host_problem_shapes_available = true)
int run(Options &options)
{
std::cout << " Problem Sizes, Alpha, Beta " << std::endl;
for (int32_t i = 0; i < options.groups; ++i) {
@@ -836,7 +837,9 @@ int main(int argc, char const **args) {
// CUTLASS must be compiled with CUDA 12.8 Toolkit to run this example
if (__CUDACC_VER_MAJOR__ < 12 ||
((__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 8))) {
((__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 8)
)
) {
std::cerr << "This example requires CUDA 12.8 or newer.\n";
// Returning zero so this test passes on older Toolkits. Its actions are no-op.
return 0;
@@ -847,7 +850,9 @@ int main(int argc, char const **args) {
CUDA_CHECK(cudaGetDevice(&current_device_id));
CUDA_CHECK(cudaGetDeviceProperties(&props, current_device_id));
cudaError_t error = cudaGetDeviceProperties(&props, 0);
if (props.major != 10 || (props.minor != 0 && props.minor != 1 && props.minor != 3)) {
if (props.major != 10 || (props.minor != 0 && props.minor != 1 && props.minor != 3
)
) {
std::cerr << "This example requires a GPU with compute capability 100a|f, 101a|f, or 103a|f)." << std::endl;
return 0;
}
examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_rcgrouped.cu
+4
View File
@@ -392,6 +392,7 @@ struct Options {
}
problem_sizes_host.push_back({m, n, k});
tokens_per_expert_host.push_back(n);
}
}
@@ -423,8 +424,11 @@ struct Options {
extent.at(i) = std::atoi(tokens.at(i).c_str());
}
problem_sizes_host.push_back({extent.m(), extent.n(), extent.k()});
tokens_per_expert_host.push_back(extent.n());
}
groups = static_cast<int>(problem_sizes_host.size());
m = get<0>(problem_sizes_host.at(0));
k = get<2>(problem_sizes_host.at(0));
return true;
}
examples/python/CuTeDSL/ampere/hstu_attention.py
+5 -2
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,
examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py
+2 -14
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(
examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py
+2 -14
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(
examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py
+2 -14
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(
examples/python/CuTeDSL/blackwell/fmha.py
+1 -1
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],
examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py
+2015
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py
+1
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")