diff --git a/CHANGELOG.md b/CHANGELOG.md index 54cc6f84..1eb5a29a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,20 @@ # CUTLASS 4.x +## [4.3.1](https://github.com/NVIDIA/cutlass/releases/tag/v4.3.1) (2025-11-26) + +### CuTe DSL +* New features + - Added Blackwell SM103 support + - Multiple dependent DSOs in the wheel have been merged into one single DSO +* Bug fixing and improvements + - Fixed device reset issue with tvm-ffi + - Fixed tvm-ffi export compiled function + +### CUTLASS C++ +* Support blockscaled variant of ragged contiguous grouped gemm with the new simplified MoE API in [example 92](https://github.com/NVIDIA/cutlass/tree/main/examples/92_blackwell_moe_gemm/). + - The new example works for all microscaling types. + ## [4.3.0](https://github.com/NVIDIA/cutlass/releases/tag/v4.3.0) (2025-11-21) ### CuTe DSL diff --git a/README.md b/README.md index 5470dd49..6e4c30e6 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,9 @@ ![ALT](./media/images/gemm-hierarchy-with-epilogue-no-labels.png "Complete CUDA GEMM decomposition") # Overview -# CUTLASS 4.3.0 +# CUTLASS 4.3.1 -_CUTLASS 4.3.0 - Nov 2025_ +_CUTLASS 4.3.1 - Nov 2025_ CUTLASS is a collection of abstractions for implementing high-performance matrix-matrix multiplication (GEMM) and related computations at all levels and scales within CUDA. It incorporates strategies for @@ -51,6 +51,8 @@ To get started quickly - please refer : - Added fake tensor and stream to decouple compile jit function with "from_dlpack" flow. Now we no longer require users to have real tensor when compile jit function. - Added FastDivmodDivisor with Python operator overloads, new APIs, Cute dialect integration, and optimized static tile scheduler performance for faster index mapping. - Added l2 cache evict priority for tma related ops. Users could do fine-grain l2 cache control. + - Added Blackwell SM103 support. + - Multiple dependent DSOs in the wheel have been merged into one single DSO. * Debuggability improvements: - Supported source location tracking for DSL APIs (Allow tools like ``nsight`` profiling to correlate perf metrics with Python source code) - Supported dumping PTX and CUBIN code: [Hello World Example](https://github.com/NVIDIA/cutlass/blob/main/examples/python/CuTeDSL/notebooks/hello_world.ipynb) @@ -95,6 +97,8 @@ To get started quickly - please refer : - Fixed TensorSSA.__getitem__ indexing to match CuTe's indexing convention - Fixed an issue with cutlass.max and cutlass.min - Fixed an issue with mark_compact_shape_dynamic + - Fixed device reset issue with tvm-ffi + - Fixed tvm-ffi export compiled function ## CUTLASS C++ * Further enhance Blackwell SM100 Attention kernels in [example 77](https://github.com/NVIDIA/cutlass/tree/main/examples/77_blackwell_fmha/). diff --git a/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_blockscaled_rcgrouped.cu b/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_blockscaled_rcgrouped.cu index 5fb925bc..d945c2f0 100644 --- a/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_blockscaled_rcgrouped.cu +++ b/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_blockscaled_rcgrouped.cu @@ -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 @@ -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(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(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::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 -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(¤t_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; } diff --git a/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_rcgrouped.cu b/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_rcgrouped.cu index 6d9f3907..7be13268 100644 --- a/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_rcgrouped.cu +++ b/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_rcgrouped.cu @@ -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(problem_sizes_host.size()); + m = get<0>(problem_sizes_host.at(0)); + k = get<2>(problem_sizes_host.at(0)); return true; } diff --git a/examples/python/CuTeDSL/ampere/hstu_attention.py b/examples/python/CuTeDSL/ampere/hstu_attention.py index 1c5aeb5f..78e91eb4 100644 --- a/examples/python/CuTeDSL/ampere/hstu_attention.py +++ b/examples/python/CuTeDSL/ampere/hstu_attention.py @@ -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, diff --git a/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py b/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py index a5be849d..b5dfe8ee 100644 --- a/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py +++ b/examples/python/CuTeDSL/blackwell/blockwise_gemm/blockwise_gemm.py @@ -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( diff --git a/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py b/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py index 43b29d06..d49802c1 100644 --- a/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py +++ b/examples/python/CuTeDSL/blackwell/blockwise_gemm/contiguous_grouped_gemm.py @@ -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( diff --git a/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py b/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py index 2457ea04..1d74304b 100644 --- a/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py +++ b/examples/python/CuTeDSL/blackwell/blockwise_gemm/masked_grouped_gemm.py @@ -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( diff --git a/examples/python/CuTeDSL/blackwell/fmha.py b/examples/python/CuTeDSL/blackwell/fmha.py index d1e9c927..d262f23d 100644 --- a/examples/python/CuTeDSL/blackwell/fmha.py +++ b/examples/python/CuTeDSL/blackwell/fmha.py @@ -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], diff --git a/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py new file mode 100644 index 00000000..cd254568 --- /dev/null +++ b/examples/python/CuTeDSL/blackwell/mixed_input_fmha/mixed_input_fmha_decode.py @@ -0,0 +1,2015 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. 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. + +# 3. 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. + +import argparse +import enum +import math +import time +from typing import Type, Tuple +from functools import partial + +import torch +import torch.nn.functional as F +from torch.nn.functional import scaled_dot_product_attention +from torch.nn.attention import SDPBackend, sdpa_kernel + +import cuda.bindings.driver as cuda + +import cutlass +import cutlass.cute as cute +import cutlass.cute.nvgpu.tcgen05 as tcgen05 +import cutlass.utils as utils +import cutlass.pipeline as pipeline +import cutlass.torch as cutlass_torch +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.cute.testing as testing +from cutlass.cute.runtime import from_dlpack +from cutlass.cute.typing import Int32, Int64, Float32, Pointer, AddressSpace + +from cutlass._mlir.dialects import nvvm, llvm + +# Kernel invariants +mma_modes = (0, 1, 2) +mma_dice = (None, None, None) # (MMA, #MMA_M, #MMA_K) +cpy_dice = (None, *mma_dice) # (CPY, #CPY_MMA, #CPY_M, #CPY_K) +warp_threads = 32 +warpgroup_warps = 4 +warpgroup_threads = 128 + +# Math helpers +log2_e = math.log2(math.e) # change exponential base +use_tensor_ssa_math = False # experimental +fadd2 = partial(cute.arch.add_packed_f32x2, ftz=False, rnd=nvvm.FPRoundingMode.RN) +fmul2 = partial(cute.arch.mul_packed_f32x2, ftz=False, rnd=nvvm.FPRoundingMode.RN) +ffma2 = partial(cute.arch.fma_packed_f32x2, ftz=False, rnd=nvvm.FPRoundingMode.RN) +exp2 = partial(cute.math.exp2, fastmath=True) + +class MixedInputFusedMultiHeadAttentionDecode: + def __init__( + self, + headdim, + block_scaledim, # headdim per scale factor; scale factor shape is (batches, heads_k, seqlen, headdim / block_scaledim) + grouped_head_tile, # GQA packing tile size, can be less than group size + dual_convert = False, # Dual warpgroups pingponging on convert stages + deterministic = False, # If True, cluster reduction is disabled + ): + self.headdim = headdim + self.grouped_head_tile = grouped_head_tile + self.block_scaledim = block_scaledim + self.scaledim = headdim // block_scaledim + self.dual_convert = dual_convert + self.deterministic = deterministic + + assert headdim % block_scaledim == 0 + assert grouped_head_tile % 8 == 0 and 0 < grouped_head_tile <= 32 + + warpgroup_id = 0 + + self.softmax_warpgroup_id = warpgroup_id + warpgroup_id += 1 + + self.cvt_warpgroup_ids = (warpgroup_id, (warpgroup_id+1) if dual_convert else None) + cvt_warpgroups = 2 if dual_convert else 1 + warpgroup_id += cvt_warpgroups + + # Why 2 MMA+TMA warps when not MMA bound? + # Less register pressure per warp promotes concise SASS + # hides MMA 'switching' latency that gets exposed with less concise SASS + # We would have 2 leftover warps if we do warpgroup reg realloc + # and less register pressure gives more realloc flexibility + self.mma_kq_warp_id = warpgroup_id * warpgroup_warps + 0 + self.mma_vp_warp_id = warpgroup_id * warpgroup_warps + 1 + self.tma_kv_warp_id = warpgroup_id * warpgroup_warps + 2 + self.tma_qo_warp_id = warpgroup_id * warpgroup_warps + 3 + self.mma_tma_warpgroup_id = warpgroup_id + warpgroup_id += 1 + + self.threads_per_cta = warpgroup_id * warpgroup_threads + + self.use_reg_reconfig = grouped_head_tile > 16 + max_regs_per_wg_thread = (64*1024 // warpgroup_threads) # 64K regs per SM + self.mma_tma_regs = 72 + self.cvt_regs = 112 + self.softmax_regs = min(256, max_regs_per_wg_thread - self.mma_tma_regs + - self.cvt_regs * cvt_warpgroups) + assert (self.mma_tma_regs + self.softmax_regs + + self.cvt_regs * cvt_warpgroups) <= max_regs_per_wg_thread + + self.bs_stages = 2 + self.sp_stages = 2 + self.o_stages = 1 + + def can_implement(self, problem_shape, kv_splits, kv_cluster_dim, q_dtype, kv_dtype, o_dtype, acc_dtype): + b, h_q, h_k, s_k, d = problem_shape + + if h_q % h_k != 0: + raise ValueError(f"heads_q({h_q}) must be a multiple of heads_k({h_k})") + + if kv_splits % kv_cluster_dim != 0: + raise ValueError(f"kv_splits({kv_splits}) must be a multiple of kv_cluster_dim({kv_cluster_dim})") + + if self.deterministic and kv_cluster_dim != 1: + raise ValueError(f"kv_cluster_dim({kv_cluster_dim}) must be 1 for determinism") + + align_scale_bits = 128 # TMA requirement + if self.scaledim * q_dtype.width < align_scale_bits: + align_seq = align_scale_bits // (self.scaledim * q_dtype.width) + if s_k % align_seq != 0: + raise ValueError(f"seqlen({s_k}) must be a multiple of {align_seq}") + + @cute.jit + def __call__( + self, + problem_shape: Tuple[Int32, Int32, Int32, Int32, Int32], # b, h_q, h_k, s_k, d + kv_splits: Int32, # threadblocks per sequence + kv_cluster_dim: Int32, # threadblocks per partial buffer (atomic reduction) + q_iter: cute.Pointer, + k_iter: cute.Pointer, + v_iter: cute.Pointer, + k_scale_iter: cute.Pointer, + v_scale_iter: cute.Pointer, + o_iter: cute.Pointer, + m_iter: cute.Pointer, # colmax_s, must be -inf initialized + l_iter: cute.Pointer, # logsumexp + o_partial_iter: cute.Pointer, # partial O per kv cluster, must be zero initialized if nondeterminism is enabled + m_partial_iter: cute.Pointer, # partial colmax_s per kv cluster + l_partial_iter: cute.Pointer, # partial colsum_p per kv cluster, must be zero initialized if nondeterminism is enabled + scale_qs: Float32, + scale_o: Float32, + stream: cuda.CUstream, + ): + ############################## + # TiledMma creation + ############################## + mma_dtype = q_iter.dtype + acc_dtype = o_partial_iter.dtype + assert acc_dtype is Float32 # don't support other acc types for now + + # Block tile sets the granularity at which threadblocks consume work + blk_tile_s = 128 + blk_tile_h = self.grouped_head_tile + blk_tile_d = self.headdim + blk_tile_shd = (blk_tile_s, blk_tile_h, blk_tile_d) + + # MMA tile sets the granularity at which TMAs + MMAs are issued + mma_tile_m = 128 + mma_tile_n = self.grouped_head_tile + mma_tile_k = 128 if self.headdim % 128 == 0 else 64 + mma_tile_mnk = (mma_tile_m, mma_tile_n, mma_tile_k) + assert self.headdim % mma_tile_k == 0 + + # GEMM1: (S_K, H_R, D, (H_K, B)) + tiled_mma_kq = sm100_utils.make_trivial_tiled_mma( + mma_dtype, + tcgen05.OperandMajorMode.K, # K + tcgen05.OperandMajorMode.K, # Q + acc_dtype, + tcgen05.CtaGroup.ONE, + mma_tile_mnk[:2], + tcgen05.OperandSource.TMEM, # converted K in tmem + ) + + # GEMM2: (D, H_R, S_K, (H_K, B)) + tiled_mma_vp = sm100_utils.make_trivial_tiled_mma(# + mma_dtype, + tcgen05.OperandMajorMode.K, # V + tcgen05.OperandMajorMode.MN, # P + acc_dtype, + tcgen05.CtaGroup.ONE, + mma_tile_mnk[:2], + tcgen05.OperandSource.TMEM, # converted V in tmem + ) + + # Calculate Q stages + self.q_stages = blk_tile_d // mma_tile_k + + # Heuristics to avoid power throttling + max_cvt_stages = 4 if self.grouped_head_tile == 32 and mma_tile_k == 128 else 8 + max_kv_stages = 8 if mma_tile_k == 128 else 14 + + # Calculate KV tmem stages + tmem_alloc_cols = mma_tile_n * self.sp_stages + tmem_alloc_cols += mma_tile_n * self.o_stages * (blk_tile_d // mma_tile_m) + tmem_capacity = 512 + cvt_stage_cols = mma_tile_k * mma_dtype.width // 32 + self.cvt_stages = (tmem_capacity - tmem_alloc_cols) // cvt_stage_cols + self.cvt_stages = min(self.cvt_stages, max_cvt_stages) + + tmem_alloc_cols += self.cvt_stages * cvt_stage_cols + self.tmem_alloc_cols = 2 ** math.ceil(math.log2(tmem_alloc_cols)) # Tmem alloc must be PO2 + + print(f"\tcvt stages: {self.cvt_stages}") + + # Calculate KV smem stages + self.mbarrier_reserved_bytes = 768 + smem_alloc_bits = self.mbarrier_reserved_bytes*8 + smem_alloc_bits += mma_tile_n * 2 * acc_dtype.width # colmax + cluster colmax + smem_alloc_bits += self.scaledim * blk_tile_s * self.bs_stages * mma_dtype.width # block scale + smem_alloc_bits += mma_tile_n * warpgroup_warps * acc_dtype.width # colsum + smem_alloc_bits += mma_tile_n * mma_tile_k * self.q_stages * mma_dtype.width # Q + smem_alloc_bits += mma_tile_m * mma_tile_n * self.sp_stages * mma_dtype.width # P + + smem_capacity = utils.get_smem_capacity_in_bytes("sm_100") + self.kv_stages = (smem_capacity*8 - smem_alloc_bits) // (mma_tile_m * mma_tile_k * k_iter.dtype.width) + self.kv_stages = min(self.kv_stages, max_kv_stages) + + print(f"\tkv stages: {self.kv_stages}") + + ############################## + # TMA creation + ############################## + b, h_q, h_k, s_k, d = problem_shape + h_r = h_q // h_k + kv_clusters = kv_splits // kv_cluster_dim + + q = cute.make_tensor(q_iter, + cute.make_ordered_layout( + shape=(h_r, d, (h_k, b)), + order=( 1, 0, ( 2, 3)) + ) + ) + + k = cute.make_tensor(k_iter, + cute.make_ordered_layout( + shape=(s_k, d, (h_k, b)), + order=( 1, 0, ( 2, 3)) + ) + ) + assert k_iter.dtype is not q_iter.dtype + + v = cute.make_tensor(v_iter, + cute.make_ordered_layout( + shape=(d, s_k, (h_k, b)), + order=(0, 1, ( 2, 3)) + ) + ) + assert v_iter.dtype is k_iter.dtype + + o_partial = cute.make_tensor(o_partial_iter, + cute.make_ordered_layout( + shape=(d, h_r, (h_k, b), kv_clusters), + order=(0, 1, ( 2, 3), 4) + ) + ) + + m = cute.make_tensor(m_iter, + cute.make_ordered_layout( + shape=(h_r, (h_k, b)), + order=( 0, ( 1, 2)), + ) + ) + assert m_iter.dtype is acc_dtype + + m_partial = cute.make_tensor(m_partial_iter, + cute.make_ordered_layout( + shape=(h_r, (h_k, b), kv_clusters), + order=( 0, ( 1, 2), 3), + ) + ) + assert m_partial_iter.dtype is acc_dtype + + l_partial = cute.make_tensor(l_partial_iter, + cute.make_ordered_layout( + shape=(h_r, (h_k, b), kv_clusters), + order=( 0, ( 1, 2), 3), + ) + ) + assert l_partial_iter.dtype is acc_dtype + + align_scale_bits = 128 # TMA requirement + if cutlass.const_expr(self.scaledim * mma_dtype.width >= align_scale_bits): + scale_layout = cute.make_ordered_layout( + shape=(self.scaledim, s_k, (h_k, b)), + order=(0, 1, (2, 3)) + ) + else: + align_seq = align_scale_bits // (self.scaledim * mma_dtype.width) + s_ks = (align_seq, s_k // align_seq) + scale_layout = cute.make_ordered_layout( + shape=(self.scaledim, s_ks, (h_k, b)), + order=(0, (1, 2), (3, 4)) + ) + + k_scale = cute.make_tensor(k_scale_iter, scale_layout) + assert k_scale_iter.dtype is mma_dtype + + v_scale = cute.make_tensor(v_scale_iter, scale_layout) + assert v_scale_iter.dtype is mma_dtype + + + # (MMA, MMA_M/N, MMA_K, Stages) + smem_layout_q = sm100_utils.make_smem_layout_b(tiled_mma_kq, mma_tile_mnk, q_iter.dtype, self.q_stages) + smem_layout_k = sm100_utils.make_smem_layout_a(tiled_mma_kq, mma_tile_mnk, k_iter.dtype, self.kv_stages) + smem_layout_v = sm100_utils.make_smem_layout_a(tiled_mma_vp, mma_tile_mnk, v_iter.dtype, self.kv_stages, + is_k_major=False) # V is always headdim-major (GEMM2 M-major) in gmem+smem + + smem_layout_bs = cute.make_layout((self.scaledim, blk_tile_s, self.bs_stages)) + + smem_layout_atom_o = tcgen05.make_smem_layout_atom(tcgen05.mma.SmemLayoutAtomKind.MN_SW128, o_partial_iter.dtype) + smem_layout_o = cute.tile_to_shape(smem_layout_atom_o, (blk_tile_d, blk_tile_h), order=(1,0)) + smem_layout_o = cute.flat_divide(smem_layout_o, (mma_tile_m, mma_tile_n)) + + tma_load_op = cute.nvgpu.cpasync.CopyBulkTensorTileG2SOp() + tma_store_op = (cute.nvgpu.cpasync.CopyBulkTensorTileS2GOp() if self.deterministic + else cute.nvgpu.cpasync.CopyReduceBulkTensorTileS2GOp()) + + tma_atom_q, tma_tensor_q = cute.nvgpu.make_tiled_tma_atom_B( + tma_load_op, q, cute.select(smem_layout_q, mma_modes), mma_tile_mnk, tiled_mma_kq + ) + tma_atom_k, tma_tensor_k = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, k, cute.select(smem_layout_k, mma_modes), mma_tile_mnk, tiled_mma_kq + ) + tma_atom_v, tma_tensor_v = cute.nvgpu.make_tiled_tma_atom_A( + tma_load_op, v, cute.select(smem_layout_v, mma_modes), mma_tile_mnk, tiled_mma_vp + ) + tma_atom_ks, tma_tensor_ks = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_load_op, k_scale, cute.select(smem_layout_bs, mode=[0,1]), smem_layout_bs.shape[:2] + ) + tma_atom_vs, tma_tensor_vs = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_load_op, v_scale, cute.select(smem_layout_bs, mode=[0,1]), smem_layout_bs.shape[:2] + ) + tma_atom_o, tma_tensor_o = cute.nvgpu.cpasync.make_tiled_tma_atom( + tma_store_op, o_partial, cute.select(smem_layout_o, mode=[0,1]), mma_tile_mnk[:2] + ) + + # K scale and V scale will have the same TMA tensor (coord tensor) + # only difference is base ptr which is stored in copy atom + tma_tensor_bs = tma_tensor_ks + + ############################## + # Decode Kernel launch + ############################## + scale_qs_log2_e = scale_qs * log2_e + + n_tiles = cute.ceil_div(h_r, blk_tile_h) + l_tiles = b * h_k + grid = (kv_splits, n_tiles, l_tiles) + + self.decode( + blk_tile_shd, mma_tile_mnk, tiled_mma_kq, tiled_mma_vp, + q_iter.dtype, smem_layout_q, tma_atom_q, tma_tensor_q, + k_iter.dtype, smem_layout_k, tma_atom_k, tma_tensor_k, + v_iter.dtype, smem_layout_v, tma_atom_v, tma_tensor_v, + smem_layout_bs, tma_atom_ks, tma_atom_vs, tma_tensor_bs, + o_partial_iter.dtype, smem_layout_o, tma_atom_o, tma_tensor_o, + m, m_partial, l_partial, + scale_qs, scale_qs_log2_e, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=[kv_cluster_dim, 1, 1], + stream=stream, + min_blocks_per_mp=1, + ) + + ############################## + # Reduction Kernel launch + ############################## + o = cute.make_tensor(o_iter, cute.make_layout((d, h_q, b))) + m = cute.make_tensor(m_iter, cute.make_layout((h_q, b))) + l = cute.make_tensor(l_iter, cute.make_layout((h_q, b))) + + o_partial = cute.make_tensor(o_partial_iter, cute.make_layout((d, h_q, b, kv_clusters))) + m_partial = cute.make_tensor(m_partial_iter, cute.make_layout((h_q, b, kv_clusters))) + l_partial = cute.make_tensor(l_partial_iter, cute.make_layout((h_q, b, kv_clusters))) + + d_per_blk = 128 + d_blks = cute.ceil_div(d, d_per_blk) + + reduction( + o, m, l, + o_partial, m_partial, l_partial, + scale_o + ).launch( + grid=[d_blks, h_q, b], + block=[d_per_blk, 1, 1], + cluster=[1, 1, 1], + stream=stream, + ) + + @cute.kernel + def decode( + self, + # MMA + blk_tile_shd: cute.Tile, + mma_tile_mnk: cute.Tile, + tiled_mma_kq: cute.TiledMma, + tiled_mma_vp: cute.TiledMma, + # Q + q_dtype: Type[cutlass.Numeric], + smem_layout_q: cute.ComposedLayout, + tma_atom_q: cute.CopyAtom, + mQ: cute.Tensor, + # K + k_dtype: Type[cutlass.Numeric], + smem_layout_k: cute.ComposedLayout, + tma_atom_k: cute.CopyAtom, + mK: cute.Tensor, + # V + v_dtype: Type[cutlass.Numeric], + smem_layout_v: cute.ComposedLayout, + tma_atom_v: cute.CopyAtom, + mV: cute.Tensor, + # K/V block scale + smem_layout_bs: cute.Layout, + tma_atom_ks: cute.CopyAtom, + tma_atom_vs: cute.CopyAtom, + mBS: cute.Tensor, + # O + o_dtype: Type[cutlass.Numeric], + smem_layout_o: cute.ComposedLayout, + tma_atom_o: cute.CopyAtom, + mO: cute.Tensor, + # Rest + mM: cute.Tensor, + mM_partial: cute.Tensor, + mL_partial: cute.Tensor, + scale_qs: Float32, + scale_qs_log2_e: Float32, + ): + # Read special registers + kv_splits, tiles_hr, tiles_hb = cute.arch.grid_dim() + kv_split_idx, coord_hr, coord_hb = cute.arch.block_idx() + kv_split_in_cluster, _, _ = cute.arch.block_in_cluster_idx() + kv_cluster_dim, _, _ = cute.arch.block_in_cluster_dim() + kv_cluster_idx, _, _ = cute.arch.cluster_idx() + tidx, _, _ = cute.arch.thread_idx() + lane_idx = cute.arch.lane_idx() + warp_idx = cute.arch.make_warp_uniform(tidx // warp_threads) + warpgroup_idx = cute.arch.make_warp_uniform(tidx // warpgroup_threads) + warpgroup_tidx = tidx % warpgroup_threads + warpgroup_widx = warp_idx % warpgroup_warps + init_warp = 0 + + # No multicast + mcast_coord = 0 + mcast_layout = cute.make_layout((1,1,1,1)) # vmnk + + # Alias types + mma_dtype = q_dtype + acc_dtype = o_dtype + + # Shapes for MMA tile indexing (Read TMA partition for example) + blk_tile_s, blk_tile_h, blk_tile_d = blk_tile_shd + mma_tile_m, mma_tile_n, mma_tile_k = mma_tile_mnk + tiles_dm, tiles_sk = cute.ceil_div((blk_tile_d, blk_tile_s), (mma_tile_m, mma_tile_k)) + tiles_dk, tiles_sm = cute.ceil_div((blk_tile_d, blk_tile_s), (mma_tile_k, mma_tile_m)) + tiles_s = cute.ceil_div(mK.shape[0], blk_tile_s) + iters_s = cute.ceil_div(tiles_s - kv_split_idx, kv_splits) + prefetch_iters = self.sp_stages - 1 + if iters_s < prefetch_iters: + prefetch_iters = iters_s + assert tiles_sm == 1 + + # Runtime checks + exit_early = kv_split_idx >= tiles_s + do_cluster_reduction = not self.deterministic and kv_cluster_dim > 1 + lane_store_max = mma_tile_n == warp_threads or lane_idx < mma_tile_n + + # Smem alloc helper + svector_align = 16 + stensor_align = 128 + smem = utils.SmemAllocator() + + ############################## + # Prefetch TMA descriptor + ############################## + if warp_idx == init_warp and not exit_early: + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_q) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_k) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_v) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_ks) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_vs) + cute.nvgpu.cpasync.prefetch_descriptor(tma_atom_o) + init_warp += 1 + + ############################## + # Tmem Allocation + ############################## + tmem_ptr_smem_ptr = smem.allocate_array(Int32) + if warp_idx == init_warp and not exit_early: + cute.arch.alloc_tmem(self.tmem_alloc_cols, tmem_ptr_smem_ptr) + init_warp += 1 + + ############################## + # Pipeline Allocation + Init + ############################## + # Allocate Mbarriers + q_pipeline_ptr = smem.allocate_array(Int64, self.q_stages * 2) + kv_pipeline_ptr = smem.allocate_array(Int64, self.kv_stages * 2) + bs_pipeline_ptr = smem.allocate_array(Int64, self.bs_stages * 2) + cvt_pipeline_ptr = smem.allocate_array(Int64, self.cvt_stages * 2) + s_pipeline_ptr = smem.allocate_array(Int64, self.sp_stages * 2) + p_pipeline_ptr = smem.allocate_array(Int64, self.sp_stages * 2) + o_pipeline_ptr = smem.allocate_array(Int64, self.o_stages * 2) + m_cluster_full_ptr = smem.allocate_array(Int64) # signal cluster colmax is in split 0 smem + + assert smem._allocated_bytes <= self.mbarrier_reserved_bytes + + # Declare named barriers + softmax_nbar_id = 1 + mma_kq_nbar_id = 2 + mma_vp_nbar_id = 3 + + # Initialize cluster colmax + mbar (even if this split exits early) + sM_layout = cute.make_layout(shape=(mma_tile_m, mma_tile_n), stride=(0, 1)) + sM_cluster = smem.allocate_tensor(acc_dtype, sM_layout, svector_align) + if do_cluster_reduction: + if warp_idx == init_warp and kv_split_in_cluster == 0: + if lane_store_max: + sM_cluster[(0,lane_idx)] = -Float32.inf + cute.arch.fence_acq_rel_cluster() + init_warp += 1 + + if warp_idx == init_warp: + # split 0 waits for one arrive per colmax elt per split in cluster + # other splits in cluster wait for one arrive from split 0 + arrive_count = (kv_cluster_dim * mma_tile_n) if kv_split_in_cluster == 0 else 1 + cute.arch.mbarrier_init(m_cluster_full_ptr, arrive_count) + cute.arch.mbarrier_init_fence() + init_warp += 1 + + cute.arch.cluster_arrive_relaxed() + + # Setup up thread cooperatives + elect_one_cooperative = pipeline.CooperativeGroup(pipeline.Agent.Thread) + warpgroup_cooperative = pipeline.CooperativeGroup(pipeline.Agent.Thread, warpgroup_threads) + dual_warpgroup_cooperative = pipeline.CooperativeGroup(pipeline.Agent.Thread, warpgroup_threads * 2) + mma_group = elect_one_cooperative + tma_group = elect_one_cooperative + tma_qo_group = elect_one_cooperative + cvt_group = warpgroup_cooperative + softmax_group = warpgroup_cooperative + + # Initialize pipelines + q_producer, q_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.q_stages, + producer_group=tma_group, + consumer_group=softmax_group, # Reuse Q consumer mbarriers to sync O store + tx_count=cute.size_in_bytes(q_dtype, cute.select(smem_layout_q, mma_modes)), + barrier_storage=q_pipeline_ptr, + tidx=mcast_coord, cta_layout_vmnk=mcast_layout, + defer_sync=True, + ).make_participants() + kv_producer, kv_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.kv_stages, + producer_group=tma_group, consumer_group=cvt_group, + tx_count=cute.size_in_bytes(k_dtype, cute.select(smem_layout_k, mma_modes)), + barrier_storage=kv_pipeline_ptr, + tidx=mcast_coord, cta_layout_vmnk=mcast_layout, + defer_sync=True, + ).make_participants() + bs_producer, bs_consumer = pipeline.PipelineTmaAsync.create( + num_stages=self.bs_stages, + producer_group=tma_qo_group, + consumer_group=(dual_warpgroup_cooperative if self.dual_convert else cvt_group), + tx_count=cute.size_in_bytes(mma_dtype, cute.select(smem_layout_bs, mode=[0,1])), + barrier_storage=bs_pipeline_ptr, + tidx=mcast_coord, cta_layout_vmnk=mcast_layout, + defer_sync=True, + ).make_participants() + cvt_producer, cvt_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.cvt_stages, + producer_group=cvt_group, consumer_group=mma_group, + barrier_storage=cvt_pipeline_ptr, + defer_sync=True, + ).make_participants() + s_producer, s_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.sp_stages, + producer_group=mma_group, consumer_group=softmax_group, + barrier_storage=s_pipeline_ptr, + defer_sync=True, + ).make_participants() + p_producer, p_consumer = pipeline.PipelineAsyncUmma.create( + num_stages=self.sp_stages, + producer_group=softmax_group, consumer_group=mma_group, + barrier_storage=p_pipeline_ptr, + defer_sync=True, + ).make_participants() + o_producer, o_consumer = pipeline.PipelineUmmaAsync.create( + num_stages=self.o_stages, + producer_group=mma_group, consumer_group=softmax_group, + barrier_storage=o_pipeline_ptr, + defer_sync=True, + ).make_participants() + + # Ensure visibility of cluster mbarrier + colmax init + if do_cluster_reduction: + cute.arch.cluster_wait() + + # Ensure visibility of local mbarrier inits and tmem alloc + cute.arch.sync_threads() + + ############################## + # MMA Partition + Allocate + ############################## + # Threadblock slice + thrblk_mma_kq = tiled_mma_kq.get_slice(0) + thrblk_mma_vp = tiled_mma_vp.get_slice(0) + + # M - colmax + sM = smem.allocate_tensor(acc_dtype, sM_layout, svector_align) + tCsM = thrblk_mma_kq.partition_C(sM) + tCsM_cluster = thrblk_mma_kq.partition_C(sM_cluster) + + # L - colsum + sL_layout = cute.make_layout(shape=(mma_tile_m, mma_tile_n, warpgroup_warps), stride=(0, 1, mma_tile_n)) + sL = smem.allocate_tensor(acc_dtype, sL_layout, svector_align) + tCsL = thrblk_mma_kq.partition_C(sL) + + # BS - block scale + sBS = smem.allocate_tensor(mma_dtype, smem_layout_bs, stensor_align) # (SCALE, TILE_S, bs_stages) + + # Q + tBsQ = smem.allocate_tensor(q_dtype, smem_layout_q.outer, stensor_align, smem_layout_q.inner) # (MMA, #MMA_N, #MMA_K, q_stages) + + # K + tAsK = smem.allocate_tensor(k_dtype, smem_layout_k.outer, stensor_align, smem_layout_k.inner) # (MMA, #MMA_M, #MMA_K, kv_stages) + tAtK_cvt_shape = tiled_mma_kq.partition_shape_A((mma_tile_m, mma_tile_k, self.cvt_stages)) # (MMA, #MMA_M, #MMA_K, cvt_stages) + tAtK_cvt = thrblk_mma_kq.make_fragment_A(tAtK_cvt_shape) + + # V + tAsV_iterator = cute.recast_ptr(tAsK.iterator, smem_layout_v.inner, dtype=v_dtype) # KV share input buffers + tAsV = cute.make_tensor(tAsV_iterator, smem_layout_v.outer) # (MMA, #MMA_M, #MMA_K, kv_stages) + tAtV_cvt_shape = tiled_mma_vp.partition_shape_A((mma_tile_m, mma_tile_k, self.cvt_stages)) # (MMA, #MMA_M, #MMA_K, cvt_stages) + tAtV_cvt = thrblk_mma_vp.make_fragment_A(tAtV_cvt_shape) + + # S + tCtS_shape = tiled_mma_kq.partition_shape_C((mma_tile_m, mma_tile_n, self.sp_stages)) + tCtS = thrblk_mma_kq.make_fragment_C(tCtS_shape) # (MMA_MN, #MMA_M=1, #MMA_N=1, sp_stages) + + # P - Treat MN C tile of BMM0 as NM B tile of BMM1 + # (MMA_NK, #MMA_N, #MMA_K=MMA_TILE_M/MMA_K, sp_stages) + mma_tile_nm = (None, mma_tile_n, mma_tile_m) + tBsP_nm_layout = sm100_utils.make_smem_layout_b(tiled_mma_vp, mma_tile_nm, mma_dtype, self.sp_stages) + tBsP_nm = smem.allocate_tensor(mma_dtype, tBsP_nm_layout.outer, stensor_align, tBsP_nm_layout.inner) + + # Tile for NK B tile iteration + # (MMA_NK, #MMA_N, #MMA_K=MMA_TILE_K/MMA_K, #TILES_SK=MMA_TILE_M/MMA_TILE_K, sp_stages) + tBsP_nk_tile = thrblk_mma_vp.partition_shape_B((mma_tile_n, mma_tile_k)) + tBsP_nk = cute.local_tile(tBsP_nm, tBsP_nk_tile, (0, 0, None, None)) + + # Reshape NM B tile of BMM1 to become MN C tile of BMM0 + # (MMA_NK, #MMA_N, #MMA_K=MMA_TILE_M/MMA_K, sp_stages) -> + # (MMA_MN, #MMA_M, #MMA_N, sp_stages) + tCsP_tile = cute.make_ordered_layout(tCtS_shape, order=((2, 0), 3, 1, 4)) + tCsP = cute.composition(tBsP_nm, tCsP_tile) + + # O + sO_iterator = cute.recast_ptr(tBsQ.iterator, smem_layout_o.inner, dtype=o_dtype) # Reuse QKV smem for O TMA store + sO_mma = cute.make_tensor(sO_iterator, smem_layout_o.outer) # (MMA_TILE_M, MMA_TILE_N, #TILE_DM, #TILE_HN) + tCsO = thrblk_mma_vp.partition_C(sO_mma) # (MMA, #MMA_M, #MMA_N, #TILE_DM, #TILE_HN) + tCtO = thrblk_mma_vp.make_fragment_C(tCsO.shape) + + # Tmem tensor allocation + tmem_ptr = cute.arch.retrieve_tmem_ptr(Int32, 16, tmem_ptr_smem_ptr) + tmem_offset = 0 + + tAtK_cvt = cute.make_tensor(cute.recast_ptr(tmem_ptr + tmem_offset, dtype=mma_dtype), tAtK_cvt.layout) + tAtV_cvt = cute.make_tensor(cute.recast_ptr(tmem_ptr + tmem_offset, dtype=mma_dtype), tAtV_cvt.layout) + tmem_offset += tcgen05.find_tmem_tensor_col_offset(tAtK_cvt) + + tCtS = cute.make_tensor(cute.recast_ptr(tmem_ptr + tmem_offset, dtype=acc_dtype), tCtS.layout) + tmem_offset += tcgen05.find_tmem_tensor_col_offset(tCtS) + + tCtO = cute.make_tensor(cute.recast_ptr(tmem_ptr + tmem_offset, dtype=acc_dtype), tCtO.layout) + tmem_offset += tcgen05.find_tmem_tensor_col_offset(tCtO) + + print(f"\t{tmem_offset} tmem cols used, {self.tmem_alloc_cols} tmem cols allocated") + assert tmem_offset <= self.tmem_alloc_cols + + ############################## + # Exit early + ############################## + if exit_early: + if do_cluster_reduction and tidx == 0: + waiting_split_in_cluster = 0 + cute.arch.mbarrier_arrive( + m_cluster_full_ptr, + waiting_split_in_cluster, + arrive_count=mma_tile_n + ) + + ############################## + # TMA KV Dispatch + ############################## + elif warp_idx == self.tma_kv_warp_id: + # Free registers + if cutlass.const_expr(self.use_reg_reconfig): + cute.arch.warpgroup_reg_dealloc(self.mma_tma_regs) + + # Apply block tiler and slice + gK = cute.local_tile(mK, tiler=(blk_tile_s, blk_tile_d), coord=(None, 0, coord_hb)) # (TILE_S, TILE_D, #TILE_S) + gV = cute.local_tile(mV, tiler=(blk_tile_d, blk_tile_s), coord=(0, None, coord_hb)) # (TILE_D, TILE_S, #TILE_S) + + # Apply MMA tiler and MMA partition + gK_mma = cute.flat_divide(gK, (mma_tile_m, mma_tile_k)) # (MMA_TILE_M, MMA_TILE_K, #TILE_SM, #TILE_DK, #TILE_S) + gV_mma = cute.flat_divide(gV, (mma_tile_m, mma_tile_k)) # (MMA_TILE_M, MMA_TILE_K, #TILE_DM, #TILE_SK, #TILE_S) + tAgK = thrblk_mma_kq.partition_A(gK_mma) # (MMA, #MMA_M, #MMA_K, #TILE_SM, #TILE_DK, #TILE_S) + tAgV = thrblk_mma_vp.partition_A(gV_mma) # (MMA, #MMA_M, #MMA_K, #TILE_DM, #TILE_SK, #TILE_S) + + # #TILE_SM=TILE_S/MMA_TILE_M, #TILE_HN=TILE_H/MMA_TILE_N, #TILE_DK=TILE_D/MMA_TILE_K + # #TILE_DM=TILE_D/MMA_TILE_M, #TILE_HN=TILE_H/MMA_TILE_N, #TILE_SK=TILE_S/MMA_TILE_K + # + # Example with TILE_S=MMA_TILE_M=128, TILE_H=MMA_TILE_N=8, MMA_TILE_K=64, TILE_D=512 + # BMM1: MMA=128x8x16, #MMA_M=1, #MMA_N=1, #MMA_K=4, #TILE_SM=1, #TILE_HN=1, #TILE_DK=8, #TILE_S=S/128 + # BMM2: MMA=128x8x16, #MMA_M=1, #MMA_N=1, #MMA_K=4, #TILE_DM=4, #TILE_HN=1, #TILE_SK=2, #TILE_S=S/128 + + # TMA partition + # (MMA, #MMA_M, #MMA_K, Rest...) -> (TMA, Rest...) + tGSsK, tGSgK = cute.nvgpu.cpasync.tma_partition( + tma_atom_k, mcast_coord, mcast_layout, + smem_tensor=cute.group_modes(tAsK, 0, 3), + gmem_tensor=cute.group_modes(tAgK, 0, 3)) + + tGSsV, tGSgV = cute.nvgpu.cpasync.tma_partition( + tma_atom_v, mcast_coord, mcast_layout, + smem_tensor=cute.group_modes(tAsV, 0, 3), + gmem_tensor=cute.group_modes(tAgV, 0, 3)) + + # + # Sequence loop + # + prefetch_tiles = prefetch_iters * kv_splits + for s in cutlass.range(kv_split_idx, prefetch_tiles + tiles_s, kv_splits): + # Load K + if s < tiles_s: + tGSgK_s = tGSgK[None, None, None, s] + for dk in cutlass.range_constexpr(tiles_dk): + k_handle = kv_producer.acquire_and_advance() + cute.copy( + tma_atom_k, + tGSgK_s[None, 0, dk], + tGSsK[None, k_handle.index], + tma_bar_ptr=k_handle.barrier) + + # Load V + if s >= prefetch_tiles: + tGSgV_s = tGSgV[None, None, None, s-prefetch_tiles] + for sk in cutlass.range_constexpr(tiles_sk): + for dm in cutlass.range_constexpr(tiles_dm): + v_handle = kv_producer.acquire_and_advance() + cute.copy( + tma_atom_v, + tGSgV_s[None, dm, sk], + tGSsV[None, v_handle.index], + tma_bar_ptr=v_handle.barrier) + + ############################## + # TMA QO Dispatch + ############################## + elif warp_idx == self.tma_qo_warp_id: + # Free registers + if cutlass.const_expr(self.use_reg_reconfig): + cute.arch.warpgroup_reg_dealloc(self.mma_tma_regs) + + # Apply block tiler and slice + gQ = cute.local_tile(mQ, tiler=(blk_tile_h, blk_tile_d), coord=(coord_hr, 0, coord_hb)) # (TILE_H, TILE_D) + gO = cute.local_tile(mO, tiler=(blk_tile_d, blk_tile_h), coord=(0, coord_hr, coord_hb, kv_cluster_idx)) # (TILE_D, TILE_H) + gBS = cute.local_tile(mBS, tiler=(self.scaledim, blk_tile_s), coord=(0, None, coord_hb)) # (SCALE, TILE_S, #TILE_S) + + # Apply MMA tiler and MMA partition + gQ_mma = cute.flat_divide(gQ, (mma_tile_n, mma_tile_k)) # (MMA_TILE_N, MMA_TILE_K, #TILE_HN, #TILE_DK) + gO_mma = cute.flat_divide(gO, (mma_tile_m, mma_tile_n)) # (MMA_TILE_M, MMA_TILE_N, #TILE_DM, #TILE_HN) + tBgQ = thrblk_mma_kq.partition_B(gQ_mma) # (MMA, #MMA_N, #MMA_K, #TILE_HN, #TILE_DK) + + # TMA partition + tGSsQ, tGSgQ = cute.nvgpu.cpasync.tma_partition( + tma_atom_q, mcast_coord, mcast_layout, + smem_tensor=cute.group_modes(tBsQ, 0, 3), + gmem_tensor=cute.group_modes(tBgQ, 0, 3)) + + tSGsO, tSGgO = cute.nvgpu.cpasync.tma_partition( + tma_atom_o, mcast_coord, mcast_layout, + smem_tensor=cute.group_modes(sO_mma, 0, 2), + gmem_tensor=cute.group_modes(gO_mma, 0, 2)) + + tGSsBS, tGSgBS = cute.nvgpu.cpasync.tma_partition( + tma_atom_ks, mcast_coord, mcast_layout, + smem_tensor=cute.group_modes(sBS, 0, 2), + gmem_tensor=cute.group_modes(gBS, 0, 2)) + + # Load Q + for dk in cutlass.range_constexpr(tiles_dk): + q_handle = q_producer.acquire_and_advance() + cute.copy( + tma_atom_q, + tGSgQ[None, 0, dk], # stages_q == tiles_dk by construction + tGSsQ[None, dk], + tma_bar_ptr=q_handle.barrier) + + # Sequence Loop + prefetch_tiles = prefetch_iters * kv_splits + for s in cutlass.range(kv_split_idx, prefetch_tiles + tiles_s, kv_splits): + if s < tiles_s: + bs_handle = bs_producer.acquire_and_advance() + cute.copy( + tma_atom_ks, + tGSgBS[None, s], + tGSsBS[None, bs_handle.index], + tma_bar_ptr=bs_handle.barrier) + + if s >= prefetch_tiles: + bs_handle = bs_producer.acquire_and_advance() + cute.copy( + tma_atom_vs, + tGSgBS[None, s-prefetch_tiles], + tGSsBS[None, bs_handle.index], + tma_bar_ptr=bs_handle.barrier) + + # Store O + for dm in cutlass.range_constexpr(tiles_dm): + q_producer.acquire_and_advance() # Reuse Q load barriers to sync O store + cute.copy( + tma_atom_o, + tSGsO[None, dm, 0], + tSGgO[None, dm, 0]) + + ############################## + # Convert Dispatch + ############################## + elif warpgroup_idx in self.cvt_warpgroup_ids: + # Free registers + if cutlass.const_expr(self.use_reg_reconfig): + cute.arch.warpgroup_reg_dealloc(self.cvt_regs) + + # Initialize for dual convert if necessary + convert_warpgroups = 1 + convert_phase = 0 + if cutlass.const_expr(self.dual_convert): + assert tiles_dk % 2 == 0 + assert (tiles_dm * tiles_sk) % 2 == 0 + convert_warpgroups = 2 + convert_phase = warpgroup_idx % convert_warpgroups + if convert_phase == 1: + kv_consumer.advance() + cvt_producer.advance() + + # Construct tiled copy and partition K + tmem_op_width = 32 + tmem_op_repeat = tcgen05.Repetition(mma_tile_k * mma_dtype.width // tmem_op_width) + tmem_store_atom_k = cute.make_copy_atom(tcgen05.St32x32bOp(tmem_op_repeat), mma_dtype) + tmem_store_k = tcgen05.make_tmem_copy(tmem_store_atom_k, tAtK_cvt[*mma_dice, 0]) + thr_store_k = tmem_store_k.get_slice(warpgroup_tidx) # tmem copy is always 128 threads + + tKsK = thr_store_k.partition_S(tAsK) + tKtK_cvt = thr_store_k.partition_D(tAtK_cvt) + + # Construct tiled copy and partition V + tmem_op_width = 128 + tmem_op_repeat = tcgen05.Repetition(mma_tile_k * mma_dtype.width // tmem_op_width) + tmem_store_atom_v = cute.make_copy_atom(tcgen05.St16x128bOp(tmem_op_repeat), mma_dtype) + tmem_store_v = tcgen05.make_tmem_copy(tmem_store_atom_v, tAtV_cvt[*mma_dice, 0]) + thr_store_v = tmem_store_v.get_slice(warpgroup_tidx) + + smem_load_atom_v = cute.make_copy_atom(cute.nvgpu.warp.LdMatrix16x16x8bOp(num_matrices=2), v_dtype) + smem_load_v = cute.make_tiled_copy_S(smem_load_atom_v, tmem_store_v) + thr_load_v = smem_load_v.get_slice(warpgroup_tidx) + + tVsV = thr_load_v.partition_S(tAsV) + tVrV_shape = smem_load_v.get_slice(0).partition_D(tAsV).shape[:-1] + tVrV_cvt_shape = thr_store_v.get_slice(0).partition_S(tAtV_cvt).shape[:-1] + tVtV_cvt = thr_store_v.partition_D(tAtV_cvt) + + # Partition KS - K block scale + sKS_layout = cute.make_layout( + shape=(blk_tile_s, (self.block_scaledim, self.scaledim), self.bs_stages), + stride=(self.scaledim, (0, 1), blk_tile_s * self.scaledim)) + sKS = cute.make_tensor(sBS.iterator, sKS_layout) # (TILE_S, TILE_D, bs_stages) + # (MMA_TILE_M, MMA_TILE_K, (#TILE_SM, #TILE_DK), bs_stages) + sKS_mma = cute.group_modes(cute.flat_divide(sKS, (mma_tile_m, mma_tile_k)), 2, 4) + tAsKS = thrblk_mma_kq.partition_A(sKS_mma) # (MMA, #MMA_M, #MMA_K, (#TILE_SM, #TILE_DK), bs_stages) + tKsKS = thr_store_k.partition_S(tAsKS) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE, bs_stages) + tKrKS = cute.make_rmem_tensor_like(tKsKS[*cpy_dice, None, 0]) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE) + + # Partition VS - V block scale + sVS_layout = cute.make_layout( + shape=((self.block_scaledim, self.scaledim), blk_tile_s, self.bs_stages), + stride=((0, 1), self.scaledim, blk_tile_s * self.scaledim)) + sVS = cute.make_tensor(sBS.iterator, sVS_layout) # (TILE_D, TILE_S, bs_stages) + # (MMA_TILE_M, MMA_TILE_K, (#TILE_DM, #TILE_SK), bs_stages) + sVS_mma = cute.group_modes(cute.flat_divide(sVS, (mma_tile_m, mma_tile_k)), 2, 4) + tAsVS = thrblk_mma_vp.partition_A(sVS_mma) # (MMA, #MMA_M, #MMA_K, (#TILE_DM, #TILE_SK), bs_stages) + tVsVS = thr_load_v.partition_D(tAsVS) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE, bs_stages) + tVrVS = cute.make_rmem_tensor_like(tVsVS[*cpy_dice, None, 0]) # (CPY, CPY_MMA, CPY_M, CPY_K, #TILE) + + # + # Sequence loop + # + for s in cutlass.range(prefetch_iters + iters_s): + if s < iters_s: + # Load K scale + bs_handle = bs_consumer.wait_and_advance() + cute.autovec_copy(tKsKS[*cpy_dice, None, bs_handle.index], tKrKS) + cute.arch.fence_view_async_shared() + bs_handle.release() + + # Convert and scale K + for dk in cutlass.range(tiles_dk // convert_warpgroups, unroll=2): + tKrK = cute.make_rmem_tensor(tKsK.shape[:-1], k_dtype) + tKrK_cvt = cute.make_rmem_tensor(tKsK.shape[:-1], mma_dtype) + + kv_handle = kv_consumer.wait_and_advance() + cute.autovec_copy(tKsK[*cpy_dice, kv_handle.index], tKrK) + cute.arch.fence_view_async_shared() + kv_handle.release() + + coord_dk = dk * convert_warpgroups + convert_phase + scale_k = tKrKS[*cpy_dice, coord_dk].load() + tKrK_cvt.store(tKrK.load().to(mma_dtype) * scale_k) + + cvt_handle = cvt_producer.acquire_and_advance() + cute.copy(thr_store_k, tKrK_cvt, tKtK_cvt[*cpy_dice, cvt_handle.index]) + cute.arch.fence_view_async_tmem_store() + cvt_handle.commit() + + # Advance again for dual warpgroups + if cutlass.const_expr(self.dual_convert): + kv_consumer.advance() + cvt_producer.advance() + + if s >= prefetch_iters: + # Load V scale + bs_handle = bs_consumer.wait_and_advance() + cute.autovec_copy(tVsVS[*cpy_dice, None, bs_handle.index], tVrVS) + cute.arch.fence_view_async_shared() + bs_handle.release() + + # Convert and scale V + for dmsk in cutlass.range(tiles_dm * tiles_sk // convert_warpgroups, unroll=2): + tVrV = cute.make_rmem_tensor(tVrV_shape, v_dtype) + tVrV_cvt = cute.make_rmem_tensor(tVrV_cvt_shape, mma_dtype) + + kv_handle = kv_consumer.wait_and_advance() + cute.copy(thr_load_v, tVsV[*cpy_dice, kv_handle.index], tVrV) + cute.arch.fence_view_async_shared() + kv_handle.release() + + coord_dmsk = dmsk * convert_warpgroups + convert_phase + scale_v = tVrVS[*cpy_dice, coord_dmsk].load() + tVrV_cvt.store(tVrV.load().to(mma_dtype) * scale_v) + + cvt_handle = cvt_producer.acquire_and_advance() + cute.copy(thr_store_v, tVrV_cvt, tVtV_cvt[*cpy_dice, cvt_handle.index]) + cute.arch.fence_view_async_tmem_store() + cvt_handle.commit() + + # Advance again for dual warpgroups + if cutlass.const_expr(self.dual_convert): + kv_consumer.advance() + cvt_producer.advance() + + ############################## + # MMA KQ Dispatch + ############################## + elif warp_idx == self.mma_kq_warp_id: + # Free registers + if cutlass.const_expr(self.use_reg_reconfig): + cute.arch.warpgroup_reg_dealloc(self.mma_tma_regs) + + # Setup mma descriptors + tBsQ_desc = thrblk_mma_kq.make_fragment_B(tBsQ) + + # Wait for Q + for dk in cutlass.range_constexpr(tiles_dk): + q_consumer.wait_and_advance() + + # Sequence loop + s_token = True # Producer always acquires first + for s in cutlass.range(iters_s): + # BMM1 + k_token = cvt_consumer.try_wait() + tiled_mma_kq.set(tcgen05.Field.ACCUMULATE, False) + s_handle = s_producer.acquire_and_advance(s_token) + for dk in cutlass.range_constexpr(tiles_dk): + is_last_iter = dk == tiles_dk - 1 + k_handle = cvt_consumer.wait_and_advance(k_token) + # Signal BMM2 to start + if is_last_iter: + cute.arch.barrier_arrive(barrier_id=mma_kq_nbar_id, number_of_threads=64) + for mma_k in cutlass.range_constexpr(tAtK_cvt.shape[2]): + cute.gemm( + tiled_mma_kq, + tCtS[*mma_dice, s_handle.index], + tAtK_cvt[None, None, mma_k, k_handle.index], + tBsQ_desc[None, None, mma_k, dk], + tCtS[*mma_dice, s_handle.index], + ) + if dk == 0 and mma_k == 0: + tiled_mma_kq.set(tcgen05.Field.ACCUMULATE, True) + k_handle.release() + if not is_last_iter: + k_token = cvt_consumer.try_wait() + s_handle.commit() + + # Advance and wait for BMM 2 + if s > 0: + for _ in cutlass.range_constexpr(tiles_dm * tiles_sk): + cvt_consumer.advance() + cute.arch.barrier(barrier_id=mma_vp_nbar_id, number_of_threads=64) + s_token = s_producer.try_acquire() + + + ############################## + # MMA VP Dispatch + ############################## + elif warp_idx == self.mma_vp_warp_id: + # Free registers + if cutlass.const_expr(self.use_reg_reconfig): + cute.arch.warpgroup_reg_dealloc(self.mma_tma_regs) + + # Setup mma descriptors + tiled_mma_vp.set(tcgen05.Field.ACCUMULATE, True) + tBsP_desc = thrblk_mma_vp.make_fragment_B(tBsP_nk) + + for _ in cutlass.range_constexpr(tiles_dk): + cvt_consumer.advance() + cute.arch.barrier(barrier_id=mma_kq_nbar_id, number_of_threads=64) + + # Sequence loop + p_token = False + o_token = True # Producer always acquires first + for s in cutlass.range(iters_s): + # Advance and wait for BMM1 + if s < iters_s - 1: + for _ in cutlass.range_constexpr(tiles_dk): + cvt_consumer.advance() + cute.arch.barrier(barrier_id=mma_kq_nbar_id, number_of_threads=64) + p_token = p_consumer.try_wait() + + # BMM2 + v_token = cvt_consumer.try_wait() + p_handle = p_consumer.wait_and_advance(p_token) + o_handle = o_producer.acquire_and_advance(o_token) + for sk in cutlass.range_constexpr(tiles_sk): + for dm in cutlass.range_constexpr(tiles_dm): + is_last_iter = sk == tiles_sk - 1 and dm == tiles_dm - 1 + v_handle = cvt_consumer.wait_and_advance(v_token) + # Signal BMM1 to start + if is_last_iter: + cute.arch.barrier_arrive(barrier_id=mma_vp_nbar_id, number_of_threads=64) + for mma_k in cutlass.range_constexpr(tAtV_cvt.shape[2]): + cute.gemm( + tiled_mma_vp, + tCtO[*mma_dice, dm, 0], + tAtV_cvt[None, None, mma_k, v_handle.index], + tBsP_desc[None, None, mma_k, sk, p_handle.index], + tCtO[*mma_dice, dm, 0], + ) + v_handle.release() + if not is_last_iter: + v_token = cvt_consumer.try_wait() + p_handle.release() + o_handle.commit() + o_token = o_producer.try_acquire() + + # Wait for signal to dealloc tmem, then dealloc + o_producer.tail() + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.dealloc_tmem(tmem_ptr, self.tmem_alloc_cols) + + ############################## + # Softmax + Correction Dispatch + ############################## + elif warpgroup_idx == self.softmax_warpgroup_id: + # Alloc registers + if cutlass.const_expr(self.use_reg_reconfig): + cute.arch.warpgroup_reg_alloc(self.softmax_regs) + + # Construct tiled copies + tmem_op_width = 32 + tmem_op_repeat = tcgen05.Repetition(mma_tile_n * acc_dtype.width // tmem_op_width) + tmem_load_atom_s = cute.make_copy_atom(tcgen05.Ld32x32bOp(tmem_op_repeat), acc_dtype) + tmem_load_s = tcgen05.make_tmem_copy(tmem_load_atom_s, tCtS[*mma_dice, 0]) + thr_load_s = tmem_load_s.get_slice(warpgroup_tidx) + + tmem_store_atom_o = cute.make_copy_atom(tcgen05.St32x32bOp(tmem_op_repeat), o_dtype) + tmem_store_o = tcgen05.make_tmem_copy(tmem_store_atom_o, tCtO[*mma_dice, 0, 0]) + thr_store_o = tmem_store_o.get_slice(warpgroup_tidx) + + # Partition S and P + tStS = thr_load_s.partition_S(tCtS) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, stages_sp) + tSsP = thr_load_s.partition_D(tCsP) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, stages_sp) + + # Partition O + tStO = thr_load_s.partition_S(tCtO) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, #TILE_DM, #TILE_HN) + tSsO = thr_load_s.partition_D(tCsO) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, #TILE_DM, #TILE_HN) + tSrO = cute.make_rmem_tensor_like(tSsO) + + # Partition colmax and initialize in RF + tSsM = thr_load_s.partition_D(tCsM) # (CPY, #CPY_MMA, #CPY_M, #CPY_N) + tSsM_cluster = thr_load_s.partition_D(tCsM_cluster) + tSrM_prev = cute.make_rmem_tensor_like(tSsM) + tSrM_prev.fill(-Float32.inf) + + # Partition colsum and initialize in RF + # Each thread maintains a local colsum in RF, smem reduction happens after loop + tSsL = thr_load_s.partition_D(tCsL) # (CPY, #CPY_MMA, #CPY_M, #CPY_N, WARPS) + tSrL = cute.make_rmem_tensor_like(tSsL[*cpy_dice, 0]) + tSrL.fill(Float32(0)) + + assert warp_threads >= cute.size(tSsM) + + # get gmem colmax + colsum to store to + inbound_hr = coord_hr * blk_tile_h + lane_idx < mM.shape[0] + gM = cute.local_tile(mM, tiler=(mma_tile_n, 1), coord=(coord_hr, coord_hb)) # (TILE_H) = (MMA_TILE_N) + gM_partial = cute.local_tile(mM_partial, tiler=(mma_tile_n, 1), coord=(coord_hr, coord_hb, kv_cluster_idx)) + gL_partial = cute.local_tile(mL_partial, tiler=(mma_tile_n, 1), coord=(coord_hr, coord_hb, kv_cluster_idx)) + + # Initialize O + tSrO.fill(Float32(0)) + cute.copy(thr_store_o, tSrO, tStO) + + # Initialize colsum and colmax in smem and wait + if warpgroup_widx == 0 and lane_store_max: + tSsM[lane_idx] = -Float32.inf + if warpgroup_widx == 1 and lane_store_max: + tSsL[lane_idx] = Float32(0) + cute.arch.barrier(barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads) + + # + # Sequence loop + # + for s in cutlass.range(iters_s): + # Load S from tmem + s_handle = s_consumer.wait_and_advance() + tSrS = cute.make_rmem_tensor(tSsP.shape[:-1], acc_dtype) + cute.copy(tmem_load_s, tStS[*cpy_dice, s_handle.index], tSrS) + cute.arch.fence_view_async_tmem_load() + s_handle.release() + + # Reduce colmax in warp RF + tSrM = cute.make_rmem_tensor_like(tSsM) + tSrM_lane = Float32(0) # Avoid dynamic register indexing + for i in cutlass.range_constexpr(cute.size(tSrS)): + tSrM[i] = warp_fmax(tSrS[i]) + if i == lane_idx: + tSrM_lane = tSrM[i] + + # Reduce colmax in smem + if lane_store_max: + smem_fmax(tSsM.iterator + tSsM.layout(lane_idx), tSrM_lane) + + # Wait for colmax then load + cute.arch.barrier(barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads) + cute.autovec_copy(tSsM, tSrM) + + # Compute online softmax + tSrP = cute.make_rmem_tensor(tSsP.shape[:-1], mma_dtype) + if cutlass.const_expr(use_tensor_ssa_math): + tSrP_f32 = exp2(scale_qs_log2_e * (tSrS.load() - tSrM.load())) + tSrP.store(tSrP_f32.to(mma_dtype)) # convert + else: + tSrP_f32 = cute.make_rmem_tensor(tSrS.shape, acc_dtype) + for i in cutlass.range_constexpr(0, cute.size(tSrS), 2): + p_f32x2 = fadd2((tSrS[i], tSrS[i+1]), (-tSrM[i], -tSrM[i+1])) + p_f32x2 = fmul2(p_f32x2, (scale_qs_log2_e, scale_qs_log2_e)) + tSrP_f32[i] = exp2(p_f32x2[0]) + tSrP_f32[i+1] = exp2(p_f32x2[1]) + tSrP.store(tSrP_f32.load().to(mma_dtype)) + + # Store P to smem + p_handle = p_producer.acquire_and_advance() + cute.autovec_copy(tSrP, tSsP[*cpy_dice, p_handle.index]) + cute.arch.fence_view_async_shared() + p_handle.commit() + + # Compute correction and correct colsum + if cutlass.const_expr(use_tensor_ssa_math): + correction = exp2(scale_qs_log2_e * (tSrM_prev.load() - tSrM.load())) + tSrL.store(tSrL.load() * correction + tSrP_f32) + else: + correction = cute.make_rmem_tensor_like(tSrM) + for i in cutlass.range_constexpr(0, cute.size(tSrM), 2): + c_f32x2 = fadd2((tSrM_prev[i], tSrM_prev[i+1]), (-tSrM[i], -tSrM[i+1])) + c_f32x2 = fmul2(c_f32x2, (scale_qs_log2_e, scale_qs_log2_e)) + c_f32x2 = (exp2(c_f32x2[0]), exp2(c_f32x2[1])) + correction[i] = c_f32x2[0] + correction[i+1] = c_f32x2[1] + l_f32x2 = ffma2(c_f32x2, (tSrL[i], tSrL[i+1]), (tSrP_f32[i], tSrP_f32[i+1])) + tSrL[i] = l_f32x2[0] + tSrL[i+1] = l_f32x2[1] + + # Correct O + if s > 0: + # Wait for O + o_handle = o_consumer.wait_and_advance() + + # Apply correction + for dm in cutlass.range_constexpr(tiles_dm): + tSrO_dm = cute.make_rmem_tensor_like(tSsO[*cpy_dice, 0, 0]) + cute.copy(thr_load_s, tStO[*cpy_dice, dm, 0], tSrO_dm) + + for i in cutlass.range_constexpr(0, cute.size(tSrO_dm), 2): + o_f32x2 = fmul2((tSrO_dm[i], tSrO_dm[i+1]), (correction[i], correction[i+1])) + tSrO_dm[i] = o_f32x2[0] + tSrO_dm[i+1] = o_f32x2[1] + + cute.copy(thr_store_o, tSrO_dm, tStO[*cpy_dice, dm, 0]) + + # Notify MMA + cute.arch.fence_view_async_tmem_store() + o_handle.release() + + # Update colmax + tSrM_prev.store(tSrM.load()) + + # + # Softmax Epilogue + # + + # Reduce colsum in warp RF + tSrL_lane = Float32(0.0) + for i in cutlass.range_constexpr(cute.size(tSrL)): + tSrL[i] = cute.arch.warp_reduction_sum(tSrL[i]) + if i == lane_idx: + tSrL_lane = tSrL[i] + + # Store partial colsum in smem + if lane_store_max: + tSsL[*cpy_dice, warpgroup_widx][lane_idx] = tSrL_lane + + # Reduce cluster colmax and correct O + if do_cluster_reduction: + # Reduce cluster colmax + if warpgroup_widx == 0: + if lane_store_max: + dsmem_fmax( + sM_cluster.iterator + sM_layout((0, lane_idx)), + sM[(0, lane_idx)], + m_cluster_full_ptr + ) + + # split 0 waits for cluster colmax to finish reduction + # other splits wait for split 0 to notify cluster colmax is ready + cute.arch.mbarrier_wait(m_cluster_full_ptr, phase=0) + + if warpgroup_widx == 0: + if kv_split_in_cluster == 0: + # notify other splits that cluster colmax is ready in split 0 smem + if lane_idx > 0 and lane_idx < kv_cluster_dim: + waiting_split_in_cluster = lane_idx + cute.arch.mbarrier_arrive( + m_cluster_full_ptr, waiting_split_in_cluster, arrive_count=1 + ) + else: + # other splits copy cluster colmax into local smem + if lane_store_max: + sM_cluster[0, lane_idx] = dsmem_load( + sM_cluster.iterator + sM_layout((0, lane_idx)) + ) + + # warpgroup waits for cluster colmax to load into local smem + cute.arch.barrier(barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads) + + if warpgroup_widx == 0 and lane_store_max and inbound_hr: + # Load colsum and colmax + sL_lane_wg = sL[0, lane_idx, None] + sL_lane = sL_lane_wg[0] + sL_lane_wg[1] + sL_lane_wg[2] + sL_lane_wg[3] + sM_prev_lane = sM[0, lane_idx] + sM_lane = sM_cluster[0, lane_idx] + + # Correct colsum and scale colmax + correction = exp2(scale_qs_log2_e * (sM_prev_lane - sM_lane)) + sL_lane = sL_lane * correction + sM_lane = sM_lane * scale_qs + + # Store colsum and colmax + gmem_fadd(gL_partial.iterator + gL_partial.layout(lane_idx), sL_lane) + gmem_fmax(gM.iterator + gM.layout(lane_idx), sM_lane) + if kv_split_in_cluster == 0: + gM_partial[lane_idx] = sM_lane + + # Load cluster colmax + tSrM = cute.make_rmem_tensor_like(tSsM) + cute.autovec_copy(tSsM_cluster, tSrM) + + # Wait and load O + o_handle = o_consumer.wait_and_advance() + cute.copy(thr_load_s, tStO, tSrO) + cute.arch.fence_view_async_tmem_load() + o_handle.release() # Final release signals tmem dealloc + + # Apply cluster correction + if cutlass.const_expr(use_tensor_ssa_math): + correction = exp2(scale_qs_log2_e * (tSrM_prev.load() - tSrM.load())) + tSrO.store(tSrO.load() * correction) + else: + correction = cute.make_rmem_tensor_like(tSsM) + for i in cutlass.range_constexpr(0, cute.size(tSrM), 2): + c_f32x2 = fadd2((tSrM_prev[i], tSrM_prev[i+1]), (-tSrM[i], -tSrM[i+1])) + c_f32x2 = fmul2(c_f32x2, (scale_qs_log2_e, scale_qs_log2_e)) + correction[i] = exp2(c_f32x2[0]) + correction[i+1] = exp2(c_f32x2[1]) + + for i in cutlass.range_constexpr(0, cute.size(tSrO), cute.size(correction)): + for j in cutlass.range_constexpr(0, cute.size(correction), 2): + o_f32x2 = fmul2((tSrO[i+j], tSrO[i+j+1]), (correction[j], correction[j+1])) + tSrO[i+j] = o_f32x2[0] + tSrO[i+j+1] = o_f32x2[1] + + # Wait and load O without cluster correction + else: + # Wait for colsum + cute.arch.barrier(barrier_id=softmax_nbar_id, number_of_threads=warpgroup_threads) + + if warpgroup_widx == 0 and lane_store_max and inbound_hr: + # Load colsum and colmax + sL_lane_wg = sL[0, lane_idx, None] + sL_lane = sL_lane_wg[0] + sL_lane_wg[1] + sL_lane_wg[2] + sL_lane_wg[3] + sM_lane = sM[0, lane_idx] + + # Scale colmax + sM_lane = sM_lane * scale_qs + + # Store colsum and colmax + gL_partial[lane_idx] = sL_lane + gM_partial[lane_idx] = sM_lane + gmem_fmax(gM.iterator + gM.layout(lane_idx), sM_lane) + + o_handle = o_consumer.wait_and_advance() + cute.copy(thr_load_s, tStO, tSrO) + cute.arch.fence_view_async_tmem_load() + o_handle.release() # Final release signals tmem dealloc + + # Store O to smem + for dm in cutlass.range_constexpr(tiles_dm): + tOrO_dm = tSrO[*cpy_dice, dm, 0] + tOsO_dm = tSsO[*cpy_dice, dm, 0] + + cute.autovec_copy(tOrO_dm, tOsO_dm) + cute.arch.fence_view_async_shared() + + # Reuse Q consumer barriers to notify O TMA store + q_consumer.release() + q_consumer.advance() + + + # Ensure split 0 doesn't exit before all splits read cluster colmax + if do_cluster_reduction: + cute.arch.cluster_arrive_relaxed() + if kv_split_in_cluster == 0 and warp_idx == self.tma_qo_warp_id: + cute.arch.cluster_wait() + + return + + @cute.kernel + def reduction( + o : cute.Tensor, + m : cute.Tensor, + l : cute.Tensor, + o_partial : cute.Tensor, + m_partial : cute.Tensor, + l_partial : cute.Tensor, + scale_o : Float32, + ): + d_blk_idx, coord_h, coord_b = cute.arch.block_idx() + d_per_blk, _, _ = cute.arch.block_dim() + d_idx, _, _ = cute.arch.thread_idx() + coord_d = d_blk_idx * d_per_blk + d_idx + + o_dhb = Float32(0) + m_hb = m[coord_h, coord_b] + l_hb = Float32(0) + + o_partial_dhb = o_partial[coord_d, coord_h, coord_b, None] + m_partial_hb = m_partial[coord_h, coord_b, None] + l_partial_hb = l_partial[coord_h, coord_b, None] + + for partial_idx in cutlass.range(o_partial.shape[-1]): + correction = exp2(log2_e * (m_partial_hb[partial_idx] - m_hb)) + o_dhb += correction * o_partial_dhb[partial_idx] + l_hb += correction * l_partial_hb[partial_idx] + + if coord_d < o.shape[0]: + o[coord_d, coord_h, coord_b] = o.element_type(scale_o * (o_dhb / l_hb)) + + if d_blk_idx == 0 and d_idx == 0: + l[coord_h, coord_b] = m_hb + cute.math.log(l_hb, fastmath=False) + + return + + @cute.jit + def _mapa(ptr : Pointer, cta_rank_in_cluster : Int32 = 0): + llvm_ptr = ptr.llvm_ptr + return nvvm.mapa_shared_cluster( + llvm_ptr.type, + llvm_ptr, + Int32(cta_rank_in_cluster).ir_value(), + ) + + @cute.jit + def dsmem_load(val_ptr : Pointer): + val_llvm_ptr = _mapa(val_ptr, 0) + + ret = llvm.inline_asm( + Float32.mlir_type, + [val_llvm_ptr], + "ld.relaxed.cta.shared::cluster.f32 $0, [$1];", + "=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + return Float32(ret) + + @cute.jit + def warp_fmax(val : Float32): + ret = llvm.inline_asm( + Float32.mlir_type, + [val.ir_value()], + "redux.sync.max.NaN.f32 $0, $1, 0xffffffff;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + return Float32(ret) + + @cute.jit + def smem_fmax(ptr : Pointer, val : Float32): + # https://stackoverflow.com/a/72461459 + # Works with canonical NaN which warp_redux_fmax should return + llvm.inline_asm( + None, + [ptr.llvm_ptr, val.ir_value()], + """{\n\t + .reg .pred p;\n\t + setp.lt.s32 p, $1, 0x0; + @p red.relaxed.shared::cta.min.u32 [$0], $1;\n\t + @!p red.relaxed.shared::cta.max.s32 [$0], $1;\n\t + }\n\t""", + "r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + @cute.jit + def dsmem_fmax(val_ptr : Pointer, val : Float32, mbar_ptr : Pointer): + expect_tx_bytes = Int32(Float32.width // 8) + val_llvm_ptr = _mapa(val_ptr, 0) + mbar_llvm_ptr = _mapa(mbar_ptr, 0) + + nvvm.mbarrier_txn( + mbar_llvm_ptr, + expect_tx_bytes.ir_value(), + kind=nvvm.MBarrierTxnKind.ARRIVE_EXPECT_TX, + space=nvvm.MBarrierSpaceKind.CLUSTER, + ) + + llvm.inline_asm( + None, + [val_llvm_ptr, val.ir_value(), mbar_llvm_ptr], + """{\n\t + .reg .pred p;\n\t + setp.lt.s32 p, $1, 0x0; + @p red.async.relaxed.cluster.shared::cluster.mbarrier::complete_tx::bytes.min.u32 [$0], $1, [$2];\n\t + @!p red.async.relaxed.cluster.shared::cluster.mbarrier::complete_tx::bytes.max.s32 [$0], $1, [$2];\n\t + }\n\t""", + "r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + @cute.jit + def gmem_fmax(ptr : Pointer, val : Float32): + llvm.inline_asm( + None, + [ptr.llvm_ptr, val.ir_value()], + """{\n\t + .reg .pred p;\n\t + setp.lt.s32 p, $1, 0x0; + @p red.relaxed.gpu.global.min.u32 [$0], $1;\n\t + @!p red.relaxed.gpu.global.max.s32 [$0], $1;\n\t + }\n\t""", + "l,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + @cute.jit + def smem_fadd(ptr : Pointer, val : Float32): + # Expensive + llvm.inline_asm( + None, + [ptr.llvm_ptr, val.ir_value()], + "red.relaxed.shared::cta.add.f32 [$0], $1;", + "r,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + @cute.jit + def gmem_fadd(ptr : Pointer, val : Float32): + llvm.inline_asm( + None, + [ptr.llvm_ptr, val.ir_value()], + "red.relaxed.cluster.global.add.f32 [$0], $1;", + "l,f", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +def run( + batches : int, + seqlen : int, + heads_q : int, + heads_k : int, + headdim : int, + block_scaledim: int, + kv_splits: int, + kv_cluster_dim: int, + q_dtype: Type[cutlass.Numeric], + kv_dtype: Type[cutlass.Numeric], + o_dtype: Type[cutlass.Numeric], + acc_dtype: Type[cutlass.Numeric], + tolerance: float, + scale_q: float, + scale_o: float, + scale_s: float, + warmup_iterations: int, + iterations: int, + skip_ref_check: bool, + use_cold_l2: bool = False, + **kwargs, +): + print(f"Running Blackwell SM100 Mixed Input FMHA Decode test with:") + print(f"\tbatches: {batches}, seqlen: {seqlen}") + print(f"\theads_q: {heads_q}, heads_k: {heads_k}") + print(f"\theaddim: {headdim}, block_scaledim: {block_scaledim}") + print(f"\tkv_splits: {kv_splits}, kv_cluster_dim: {kv_cluster_dim}") + print(f"\tq_dtype: {q_dtype}") + print(f"\tkv_dtype: {kv_dtype}") + print(f"\to_dtype: {o_dtype}") + print(f"\tacc_dtype: {acc_dtype}") + print(f"\ttolerance: {tolerance}") + print(f"\tscale_q: {scale_q}") + print(f"\tscale_o: {scale_o}") + print(f"\tscale_s: {scale_s}") + print(f"\twarmup_iterations: {warmup_iterations}") + print(f"\titerations: {iterations}") + print(f"\tskip_ref_check: {skip_ref_check}") + print(f"\tuse_cold_l2: {use_cold_l2}") + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + # + # Config Kernel + # + grouped_heads = heads_q // heads_k + fmha = MixedInputFusedMultiHeadAttentionDecode( + headdim=headdim, + block_scaledim=block_scaledim, + grouped_head_tile=min(cute.round_up(grouped_heads, 8), 32), + dual_convert=(headdim > 128), + deterministic=(kv_cluster_dim == 1), + ) + + if scale_s == 0.0: # default to 1/sqrt(d) + scale_s = 1.0 / math.sqrt(headdim) + scale_qs = scale_q * scale_s + + if kv_splits == 0: + hardware_info = cutlass.utils.HardwareInfo() + sm_count = hardware_info.get_device_multiprocessor_count() + sm_count = 148 if sm_count <= 0 else sm_count + grid_yz = batches * heads_k * math.ceil(grouped_heads / fmha.grouped_head_tile) + kv_splits = sm_count // grid_yz # 1 wave + kv_splits = max(1, kv_splits) + if sm_count == 148 and grid_yz == 32: + kv_splits = 9 # 2 waves + print(f"\tauto kv_splits: {kv_splits}") + + seqlen_q = 1 + seqlen_k = seqlen + kv_clusters = kv_splits // kv_cluster_dim + + problem_shape = (batches, heads_q, heads_k, seqlen_k, headdim) + + fmha.can_implement(problem_shape, kv_splits, kv_cluster_dim, q_dtype, kv_dtype, o_dtype, acc_dtype) + + # + # Allocate Tensors + # + torch.manual_seed(1111) + def create_tensor(shape, dtype, init = True): + init_type = cutlass.torch.TensorInitType.RANDOM + init_config = cutlass.torch.RandomInitConfig(min_val=-2, max_val=2) + + if init is False or init is None: + init_type = cutlass.torch.TensorInitType.SKIP + init_config = None + elif isinstance(init, int) or isinstance(init, float): + init_type = cutlass.torch.TensorInitType.SCALAR + init_config = cutlass.torch.ScalarInitConfig(value=init) + elif isinstance(init, tuple) or isinstance(init, list): + if len(init) == 2: + init_type = cutlass.torch.TensorInitType.RANDOM + init_config = cutlass.torch.RandomInitConfig(min_val=init[0], max_val=init[1]) + if len(init) == 3: + init_type = cutlass.torch.TensorInitType.GAUSSIAN + init_config = cutlass.torch.RandomInitConfig(mean=init[0], std=init[1], scale=init[2]) + + f32_torch_tensor = cutlass_torch.create_and_permute_torch_tensor( + shape, + torch.float32, + permute_order=None, + init_type=init_type, + init_config=init_config, + ) + + _, torch_tensor = cutlass_torch.cute_tensor_like( + f32_torch_tensor, + dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + + # Create dtype cute tensor with offset (gpu) + cute_tensor = from_dlpack(torch_tensor, assumed_align=16) + cute_tensor.element_type = dtype + + return ( + f32_torch_tensor, + cute_tensor, + torch_tensor, + ) + + qo_shape = (kv_clusters, batches, heads_q, seqlen_q, headdim) + kv_shape = ( batches, heads_k, seqlen_k, headdim) + scale_shape = ( batches, heads_k, seqlen_k, fmha.scaledim) + + q_ref, q_cute, q_torch = create_tensor(qo_shape[1:], q_dtype, init=[-8,7]) + k_ref, k_cute, k_torch = create_tensor(kv_shape, kv_dtype, init=[-8,7]) + v_ref, v_cute, v_torch = create_tensor(kv_shape, kv_dtype, init=[-8,7]) + k_scale_ref, k_scale_cute, k_scale_torch = create_tensor(scale_shape, q_dtype, init=[-2, 2]) + v_scale_ref, v_scale_cute, v_scale_torch = create_tensor(scale_shape, q_dtype, init=[-2, 2]) + _, o_cute, o_torch = create_tensor(qo_shape[1:], o_dtype, init=False) + _, m_cute, m_torch = create_tensor(qo_shape[1:-1], acc_dtype, init=-math.inf) + _, l_cute, l_torch = create_tensor(qo_shape[1:-1], acc_dtype, init=False) + _, o_partial_cute, o_partial_torch = create_tensor(qo_shape, acc_dtype, init=0) + _, m_partial_cute, m_partial_torch = create_tensor(qo_shape[:-1], acc_dtype, init=-math.inf) + _, l_partial_cute, l_partial_torch = create_tensor(qo_shape[:-1], acc_dtype, init=0) + + # + # Compile + # + current_stream = cutlass_torch.default_stream() + compiled_fmha = cute.compile( + fmha, + problem_shape, + kv_splits, + kv_cluster_dim, + q_cute.iterator, + k_cute.iterator, + v_cute.iterator, + k_scale_cute.iterator, + v_scale_cute.iterator, + o_cute.iterator, + m_cute.iterator, + l_cute.iterator, + o_partial_cute.iterator, + m_partial_cute.iterator, + l_partial_cute.iterator, + scale_qs, + scale_o, + current_stream, + options=f"--opt-level 2", + ) + print("Finished Compiling") + + # + # Refcheck + # + def run_torch_fmha(q_ref, k_ref, v_ref, k_scale_ref, v_scale_ref, scale_qs=1.0, scale_o=1.0): + for i in range(0, headdim // block_scaledim): + j = i * block_scaledim + k_ref[..., j:j+block_scaledim] *= k_scale_ref[..., i:i+1] + v_ref[..., j:j+block_scaledim] *= v_scale_ref[..., i:i+1] + + with sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.MATH], set_priority=True): + o_ref = scaled_dot_product_attention( + q_ref, + k_ref, + v_ref, + attn_mask=None, + dropout_p=0.0, + scale=scale_qs, + is_causal=False, + enable_gqa=(heads_q != heads_k), + ) + + return o_ref * scale_o + + if not skip_ref_check: + # Execute kernel once for reference checking + print("Running...") + compiled_fmha( + problem_shape, + kv_splits, + kv_cluster_dim, + q_cute.iterator, + k_cute.iterator, + v_cute.iterator, + k_scale_cute.iterator, + v_scale_cute.iterator, + o_cute.iterator, + m_cute.iterator, + l_cute.iterator, + o_partial_cute.iterator, + m_partial_cute.iterator, + l_partial_cute.iterator, + scale_qs, + scale_o, + current_stream, + ) + print("Verifying results...") + o_ref = run_torch_fmha(q_ref, k_ref, v_ref, k_scale_ref, v_scale_ref, scale_qs, scale_o) + torch.testing.assert_close(o_ref, o_torch.float().cpu(), atol=tolerance, rtol=1e-05) + + def generate_tensors(): + _, q_cute, _ = create_tensor(qo_shape[1:], q_dtype, init=[-8,7]) + _, k_cute, _ = create_tensor(kv_shape, kv_dtype, init=[-8,7]) + _, v_cute, _ = create_tensor(kv_shape, kv_dtype, init=[-8,7]) + _, k_scale_cute, _ = create_tensor(scale_shape, q_dtype, init=[-8,7]) + _, v_scale_cute, _ = create_tensor(scale_shape, q_dtype, init=[-8,7]) + _, o_cute, _ = create_tensor(qo_shape[1:], o_dtype, init=False) + _, m_cute, _ = create_tensor(qo_shape[1:-1], acc_dtype, init=-math.inf) + _, l_cute, _ = create_tensor(qo_shape[1:-1], acc_dtype, init=False) + _, o_partial_cute, _ = create_tensor(qo_shape, acc_dtype, init=0) + _, m_partial_cute, _ = create_tensor(qo_shape[:-1], acc_dtype, init=-math.inf) + _, l_partial_cute, _ = create_tensor(qo_shape[:-1], acc_dtype, init=0) + return testing.JitArguments( + problem_shape, + kv_splits, + kv_cluster_dim, + q_cute.iterator, + k_cute.iterator, + v_cute.iterator, + k_scale_cute.iterator, + v_scale_cute.iterator, + o_cute.iterator, + m_cute.iterator, + l_cute.iterator, + o_partial_cute.iterator, + m_partial_cute.iterator, + l_partial_cute.iterator, + scale_qs, + scale_o, + current_stream, + ) + + # + # Profile + # + workspace_count = 1 + if use_cold_l2: + q_torch_effective = q_torch.values() if q_torch.is_nested else q_torch + k_torch_effective = k_torch.values() if k_torch.is_nested else k_torch + v_torch_effective = v_torch.values() if v_torch.is_nested else v_torch + k_scale_torch_effective = k_scale_torch.values() if k_scale_torch.is_nested else k_scale_torch + v_scale_torch_effective = v_scale_torch.values() if v_scale_torch.is_nested else v_scale_torch + o_torch_effective = o_torch.values() if o_torch.is_nested else o_torch + m_torch_effective = m_torch.values() if m_torch.is_nested else m_torch + l_torch_effective = l_torch.values() if l_torch.is_nested else l_torch + o_partial_torch_effective = o_partial_torch.values() if o_partial_torch.is_nested else o_partial_torch + m_partial_torch_effective = m_partial_torch.values() if m_partial_torch.is_nested else m_partial_torch + l_partial_torch_effective = l_partial_torch.values() if l_partial_torch.is_nested else l_partial_torch + one_workspace_bytes = ( + q_torch_effective.numel() * q_torch_effective.element_size() + + k_torch_effective.numel() * k_torch_effective.element_size() + + v_torch_effective.numel() * v_torch_effective.element_size() + + k_scale_torch_effective.numel() * k_scale_torch_effective.element_size() + + v_scale_torch_effective.numel() * v_scale_torch_effective.element_size() + + o_torch_effective.numel() * o_torch_effective.element_size() + + m_torch_effective.numel() * m_torch_effective.element_size() + + l_torch_effective.numel() * l_torch_effective.element_size() + + o_partial_torch_effective.numel() * o_partial_torch_effective.element_size() + + m_partial_torch_effective.numel() * m_partial_torch_effective.element_size() + + l_partial_torch_effective.numel() * l_partial_torch_effective.element_size() + ) + workspace_count = testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = 0.0 + if iterations > 0: + exec_time = testing.benchmark( + compiled_fmha, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + ) + print(exec_time) + + return exec_time # Return execution time in microseconds + +if __name__ == "__main__": + + def parse_comma_separated_ints(s: str): + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) + + parser = argparse.ArgumentParser(description="Example of FMHA on Blackwell.") + + parser.add_argument( + "--batches","--batch","--b", + type=int, + default=1, + help="batch size" + ) + + parser.add_argument( + "--seqlen","--seqlen_k","--seq","--s", + type=int, + default=1024, + help="key/value sequence length" + ) + + parser.add_argument( + "--heads_q","--h_q", + type=int, + default=32, + help="query heads" + ) + + parser.add_argument( + "--heads_k","--h_k", + type=int, + default=4, + help="key/value heads" + ) + + parser.add_argument( + "--headdim","--d", + type=int, + default=512, + help="head dimmension" + ) + + parser.add_argument( + "--block_scaledim","--bs", + type=int, + default=512, + help="headdim per scale factor" + ) + + parser.add_argument( + "--kv_splits","--splits", + type=int, + default=0, + help="threadblocks per sequence", + ) + + parser.add_argument( + "--kv_cluster_dim","--cluster", + type=int, + default=1, + help="threadblocks per partial buffer", + ) + + parser.add_argument( + "--q_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + help="query data type", + ) + + parser.add_argument( + "--kv_dtype", + type=cutlass.dtype, + default=cutlass.Int8, + help="key/value data type", + ) + + parser.add_argument( + "--o_dtype", + type=cutlass.dtype, + default=cutlass.BFloat16, + help="output data type", + ) + + parser.add_argument( + "--acc_dtype", + type=cutlass.dtype, + default=Float32, + help="accumulator/reduction data type", + ) + + parser.add_argument( + "--tolerance", type=float, default=1e-01, help="Tolerance for validation" + ) + + parser.add_argument( + "--scale_q", + type=float, + default=1.0, + help="query scale factor", + ) + + parser.add_argument( + "--scale_o", + type=float, + default=1.0, + help="output scale factor", + ) + + parser.add_argument( + "--scale_s","--scale", + type=float, + default=0.0, + help="score (Q*K) scale factor; if zero, defaults to 1/sqrt(D)", + ) + + parser.add_argument( + "--warmup_iterations", + type=int, + default=0, + help="Number of iterations for warmup", + ) + + parser.add_argument( + "--iterations", + type=int, + default=0, + help="Number of iterations after warmup", + ) + + parser.add_argument( + "--skip_ref_check", + action="store_true", + help="Skip reference check", + ) + + parser.add_argument( + "--use_cold_l2", + action="store_true", + default=False, + help="Use circular buffer tensor sets to ensure L2 cold cache", + ) + + args = parser.parse_args() + + run( + args.batches, + args.seqlen, + args.heads_q, + args.heads_k, + args.headdim, + args.block_scaledim, + args.kv_splits, + args.kv_cluster_dim, + args.q_dtype, + args.kv_dtype, + args.o_dtype, + args.acc_dtype, + args.tolerance, + args.scale_q, + args.scale_o, + args.scale_s, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + ) + + print("PASS") diff --git a/examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py b/examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py index b1d7a97c..92b40da8 100644 --- a/examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py +++ b/examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py @@ -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") diff --git a/include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized_rcggemm.hpp b/include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized_rcggemm.hpp index 2b8d7179..a133c8a8 100644 --- a/include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized_rcggemm.hpp +++ b/include/cutlass/gemm/collective/sm100_blockscaled_mma_array_warpspecialized_rcggemm.hpp @@ -774,7 +774,8 @@ struct CollectiveMma< if constexpr (IsTensorMapUpdateAsync) { return ret; - } else { + } + else { // Fetch a copy of tensormaps for the CTA from Params auto input_tensormaps = tensormaps_init(params, shared_tensormaps, sm_count, sm_idx); return cute::tuple_cat(ret, cute::make_tuple(input_tensormaps)); diff --git a/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst b/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst index ad70f2a2..42d0fdea 100644 --- a/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst +++ b/media/docs/pythonDSL/cute_dsl_general/compile_with_tvm_ffi.rst @@ -35,13 +35,35 @@ There are two ways to enable TVM FFI in |DSL|: a_cute = cute.runtime.from_dlpack(a_torch, enable_tvm_ffi=True).mark_layout_dynamic() b_cute = cute.runtime.from_dlpack(b_torch, enable_tvm_ffi=True).mark_layout_dynamic() - compiled_add = cute.compile(add, a_cute, b_cute, options="--enable-tvm-ffi") + compiled_add = cute.compile(add, a_torch, b_torch, options="--enable-tvm-ffi") Note that the object returned by ``cute.compile`` is a Python function specific to TVM FFI. 2. Alternatively, you can enable TVM FFI globally by setting the environment variable ``CUTE_DSL_ENABLE_TVM_FFI=1``. Please note that this setting will apply to all JIT compilations within the environment. + +Minimizing Host Overhead +------------------------ + +Eager kernel invocation overhead on the CPU host can sometimes become a bottleneck +for latency-sensitive applications. TVM FFI can help greatly reduce this overhead. +To maximize performance benefits, we recommend setting up your workflow as follows +(detailed instructions are provided in subsequent sections): + +- **Compile the kernel with TVM FFI enabled.** +- **Declare shape constraints using fake tensors** and reuse the compiled function + throughout your execution. +- **Pass PyTorch tensors directly** to the compiled function to avoid explicit DLPack conversion. +- **Use the environment stream flag** to implicitly synchronize with the current PyTorch stream. +- **Rely on compiled argument validation** instead of Python-side attribute validation, + as TVM FFI functions perform fast compiled checks. + +Following these steps can significantly reduce the host-side overhead of eager kernel execution. +The sections below provide detailed examples and explanations for each step. +You may find it helpful to refer back to this summary after you review the implementation details. + + Fake tensor for compilation --------------------------- @@ -92,6 +114,18 @@ The fake tensor is a placeholder that mimics the interface of a real tensor but It is used in compilation or testing scenarios where only shape/type/layout information is needed. All attempts to access or mutate data will raise errors. +Note on Stride Order +~~~~~~~~~~~~~~~~~~~~ + +Note that CuTe's convention is to write the stride order for dimensions from left to right, +where a lower order number means higher priority. In the context of the ``make_fake_compact_tensor`` API, +for shape ``(2, 3, 4)`` and stride order ``(0, 1, 2)``, the stride is ``(1, 2, 6)``. +This is commonly known as column-major order. If you want to create a fake tensor with compact row-major order, +you should explicitly pass in ``stride_order=tuple(reversed(range(len(shape))))`` +to ``make_fake_compact_tensor``. Alternatively, you can always precisely control the +stride via the ``stride`` argument in the ``make_fake_tensor`` API. + + ``cute.Tensor`` adapter for TVM FFI ----------------------------------- @@ -173,6 +207,9 @@ The following example demonstrates this approach; the function accepts ``torch.c print("result of b_torch after compiled_add_one(a_torch, b_torch, torch_stream)") print(b_torch) +Using Environment Stream +~~~~~~~~~~~~~~~~~~~~~~~~ + The second option is to rely on the environment-stream flag. Pass ``use_tvm_ffi_env_stream=True`` to ``make_fake_stream`` to mark the argument as an environment stream so it no longer has to be provided explicitly. @@ -202,8 +239,54 @@ before each call. The example below shows this flow: print("result of b_torch after compiled_add_one(a_torch, b_torch)") print(b_torch) -Using the environment-stream flag both speeds up calls and simplifies integration +Using the environment stream flag both speeds up calls and simplifies integration with frameworks such as PyTorch, since no explicit stream parameter is required. +We recommend using the environment stream flag to both simplify framework integration +and minimize host-side calling overhead. + +Working with Tuples +------------------- + +TVM FFI functions can also accept tuples as arguments. Tuples can be recursively +composed of the types that are supported by TVM FFI. The example below shows how to use tuples as arguments: + +.. code-block:: python + + import torch + from cutlass import cute + + @cute.kernel + def device_add_one(a: cute.Tensor, b: cute.Tensor, c: cute.Float32): + threads_per_block = 128 + cta_x_, _, _ = cute.arch.block_idx() + tid_x, _, _ = cute.arch.thread_idx() + tid = cta_x_ * threads_per_block + tid_x + if tid < a.shape[0]: + b[tid] = a[tid] + c + + @cute.jit + def add_one_with_tuple(a: Tuple[cute.Tensor, cute.Tensor, cute.Float32]): + n = a[0].shape[0] + threads_per_block = 128 + blocks = (n + threads_per_block - 1) // threads_per_block + device_add_one(a[0], a[1], a[2]).launch(grid=(blocks, 1, 1), block=(threads_per_block, 1, 1)) + + def example_add_one_with_tuple(): + n = cute.sym_int() + a_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,)) + b_cute = cute.runtime.make_fake_compact_tensor(cute.Float32, (n,)) + compiled_add_one = cute.compile( + add_one_with_tuple, (a_cute, b_cute, cute.Float32(4)), + options="--enable-tvm-ffi" + ) + a_torch = torch.arange(10, dtype=torch.float32, device="cuda") + b_torch = torch.empty(10, dtype=torch.float32, device="cuda") + compiled_add_one((a_torch, b_torch, 5)) + print("result of b_torch after compiled_add_one((a_torch, b_torch, 5))") + print(b_torch) + + example_add_one_with_tuple() + Supported types --------------- @@ -230,6 +313,8 @@ The TVM FFI function supports the following |DSL|-specific types as arguments: - ``tvm_ffi.Shape`` or python tuple of ints. * - CUDA stream ``cuda.CUstream`` - A stream class that implements the CUDA stream protocol (e.g. ``torch.cuda.Stream``, ``cuda.CUstream``). + * - Tuple of types (e.g. ``Tuple[cute.Tensor, cute.Tensor, cutlass.Int32]``) + - Python tuple of corresponding call-time types. Error handling -------------- diff --git a/media/docs/pythonDSL/cute_dsl_general/framework_integration.rst b/media/docs/pythonDSL/cute_dsl_general/framework_integration.rst index 7d223ebb..2fdda5fa 100644 --- a/media/docs/pythonDSL/cute_dsl_general/framework_integration.rst +++ b/media/docs/pythonDSL/cute_dsl_general/framework_integration.rst @@ -415,24 +415,36 @@ The following example demonstrates how to use ``mark_compact_shape_dynamic`` to ) # Layout in DLTensorWrapper has int32 overflow risk. Please set use_32bit_stride to False. +Leveraging TVM FFI for Faster PyTorch Interop +--------------------------------------------- + +The latest version of |DSL| supports TVM FFI to improve interoperability with PyTorch +and other machine learning frameworks. Using TVM FFI provides the following features: + +- Faster JIT function invocation. +- Direct acceptance of ``torch.Tensor`` objects as function arguments. +- Enhanced error handling and kernel validation. +- Seamless integration with multiple programming languages. + +For more details, see :doc:`compile_with_tvm_ffi`. Bypass the DLPack Protocol -------------------------- -In certain scenarios, users may wish to bypass the DLPack protocol and invoke the JIT function directly. -This can be accomplished by creating a lightweight JIT wrapper around the existing JIT function, +In certain scenarios, users may wish to bypass the DLPack protocol and invoke the JIT function directly. +This can be accomplished by creating a lightweight JIT wrapper around the existing JIT function, utilizing ``cute.ptr`` and ``cute.make_tensor`` to pass pointers and construct tensors directly. Typical use cases for bypassing DLPack include: 1. Users want to call the JIT function directly to avoid the overhead introduced by the DLPack protocol. -2. DLPack canonicalizes the stride of shape-1 dimensions to 1, which may result in incorrect alignment +2. DLPack canonicalizes the stride of shape-1 dimensions to 1, which may result in incorrect alignment propagation and affect memory access or performance. 3. DLPack may lack support for some narrow data types. The following example illustrates how to bypass the DLPack protocol when invoking a JIT function. -Assume we have a pre-defined ``TensorOpGemm`` kernel whose JIT interface expects three -arguments of type ``cute.Tensor``. To enable direct invocation without DLPack, we first define a JIT wrapper -function that accepts ``cute.Pointer`` types as parameters. Within this wrapper, we use ``cute.make_tensor`` +Assume we have a pre-defined ``TensorOpGemm`` kernel whose JIT interface expects three +arguments of type ``cute.Tensor``. To enable direct invocation without DLPack, we first define a JIT wrapper +function that accepts ``cute.Pointer`` types as parameters. Within this wrapper, we use ``cute.make_tensor`` to construct tensors from the provided pointers, and then call the ``TensorOpGemm`` kernel as usual. .. code-block:: python @@ -459,7 +471,7 @@ to construct tensors from the provided pointers, and then call the ``TensorOpGem mA = cute.make_tensor(a_ptr, layout=a_layout) mB = cute.make_tensor(b_ptr, layout=b_layout) mC = cute.make_tensor(c_ptr, layout=c_layout) - + # TensorOpGemm is a pre-defined kernel from our example tensor_op_gemm = TensorOpGemm( a_ptr.value_type, c_ptr.value_type, cutlass.Float32, (2, 2, 1) @@ -467,9 +479,9 @@ to construct tensors from the provided pointers, and then call the ``TensorOpGem tensor_op_gemm(mA, mB, mC) -To pass a PyTorch tensor to this new JIT wrapper, we retrieve the raw pointer from the PyTorch tensor +To pass a PyTorch tensor to this new JIT wrapper, we retrieve the raw pointer from the PyTorch tensor and create a ``cute.Pointer`` instance using ``cute.make_ptr``. -This approach allows us to bypass the DLPack protocol entirely, avoiding its overhead and potential +This approach allows us to bypass the DLPack protocol entirely, avoiding its overhead and potential issues with shape-1 dimension handling. .. code-block:: python @@ -483,7 +495,7 @@ issues with shape-1 dimension handling. c = torch.randn( n, m, l, dtype=torch.float16, device="cuda" ).permute(1, 2, 0) - + # from cutlass.cute.runtime import make_ptr a_ptr = make_ptr( cutlass.Float16, a.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 @@ -495,3 +507,5 @@ issues with shape-1 dimension handling. cutlass.Float16, c.data_ptr(), cute.AddressSpace.gmem, assumed_align=32 ) tensor_op_gemm_wrapper(a_ptr, b_ptr, c_ptr, m, n, k, l) + + diff --git a/python/CuTeDSL/cutlass/__init__.py b/python/CuTeDSL/cutlass/__init__.py index 38bd7626..e42a8e3a 100644 --- a/python/CuTeDSL/cutlass/__init__.py +++ b/python/CuTeDSL/cutlass/__init__.py @@ -9,6 +9,10 @@ # and related documentation outside the scope permitted by the EULA # is strictly prohibited. +from ._mlir._mlir_libs import _cutlass_ir + +_cutlass_ir.populate(_cutlass_ir) + from .cutlass_dsl import ( Constexpr, dsl_user_op, diff --git a/python/CuTeDSL/cutlass/base_dsl/arch.py b/python/CuTeDSL/cutlass/base_dsl/arch.py index 86876172..60d2060d 100644 --- a/python/CuTeDSL/cutlass/base_dsl/arch.py +++ b/python/CuTeDSL/cutlass/base_dsl/arch.py @@ -32,6 +32,9 @@ class Arch(Enum): sm_101 = (10, 1, "") sm_101a = (10, 1, "a") sm_101f = (10, 1, "f") + sm_103 = (10, 3, "") + sm_103a = (10, 3, "a") + sm_103f = (10, 3, "f") sm_110 = (11, 0, "") sm_110a = (11, 0, "a") sm_110f = (11, 0, "f") @@ -80,6 +83,9 @@ class Arch(Enum): Arch.sm_101, Arch.sm_101a, Arch.sm_101f, + Arch.sm_103, + Arch.sm_103a, + Arch.sm_103f, Arch.sm_110, Arch.sm_110a, Arch.sm_110f, diff --git a/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py b/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py index fdb09b9e..a81fd55f 100644 --- a/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py +++ b/python/CuTeDSL/cutlass/base_dsl/ast_helpers.py @@ -529,6 +529,7 @@ def _get_self_module(): return inspect.getmodule(_get_self_module) +@lru_cache(maxsize=16) def cf_symbol_check(symbol): """ Check if the symbol is control flow symbol from current module. diff --git a/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py b/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py index 16dc74ca..81083348 100644 --- a/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py +++ b/python/CuTeDSL/cutlass/base_dsl/ast_preprocessor.py @@ -1344,6 +1344,25 @@ class DSLPreprocessor(ast.NodeTransformer): ), node, ) + elif func.id == "super" and node.args == [] and node.keywords == []: + # If it's a Python3 argument free super(), rewrite to old style super with args + # So if this call is under dynamic control flow, it still works. + return ast.copy_location( + ast.Call( + func=func, + args=node.args + + [ + ast.Attribute( + value=ast.Name(id="self", ctx=ast.Load()), + attr="__class__", + ctx=ast.Load(), + ), + ast.Name(id="self", ctx=ast.Load()), + ], + keywords=node.keywords, + ), + node, + ) elif isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name): def create_downcast_call(arg): @@ -1853,21 +1872,41 @@ class DSLPreprocessor(ast.NodeTransformer): # Handle elif case elif_node = node.orelse[0] nested_if_name = elif_region_name - # Recursion for nested elif - nested_if = self.create_if_function( - nested_if_name, elif_node, write_args, full_write_args_count - ) - else_block = ast.FunctionDef( - name=else_block_name, - args=func_then_else_arguments, - body=[ - nested_if, - ast.Return( - value=ast.Name(id=nested_if_name, ctx=ast.Load()) - ), - ], - decorator_list=[], - ) + # AST cannot distinguish between the following two cases: + # elif pred: + # and + # else: + # if pred: + # And under both cases, the `pred` can be a const_expr, so we need to handle it here. + if self.is_node_constexpr(elif_node): + self.generic_visit(elif_node) + check = self._insert_cf_symbol_check(elif_node.test.func) + else_block = ast.FunctionDef( + name=else_block_name, + args=func_then_else_arguments, + body=[ + check, + elif_node, + ast.Return(value=return_list), + ], + decorator_list=[], + ) + else: + # Recursion for nested elif + nested_if = self.create_if_function( + nested_if_name, elif_node, write_args, full_write_args_count + ) + else_block = ast.FunctionDef( + name=else_block_name, + args=func_then_else_arguments, + body=[ + nested_if, + ast.Return( + value=ast.Name(id=nested_if_name, ctx=ast.Load()) + ), + ], + decorator_list=[], + ) else: else_body = [] for stmt in node.orelse: diff --git a/python/CuTeDSL/cutlass/base_dsl/compiler.py b/python/CuTeDSL/cutlass/base_dsl/compiler.py index 3b47c578..3124519b 100644 --- a/python/CuTeDSL/cutlass/base_dsl/compiler.py +++ b/python/CuTeDSL/cutlass/base_dsl/compiler.py @@ -356,7 +356,11 @@ class GPUArch(StringCompileOption): class EnableTVMFFI(EmptyCompileOption): - option_name = "enable-tvm-ffi" + pass + + +class DumpDir(EmptyCompileOption): + option_name = "dump-dir" class CompileOptions: @@ -380,6 +384,7 @@ class CompileOptions: GPUArch: GPUArch(""), LinkLibraries: LinkLibraries(""), EnableTVMFFI: EnableTVMFFI(False), + DumpDir: DumpDir(""), } if options is not None: @@ -416,19 +421,24 @@ class CompileOptions: if self.options[GPUArch].value == "" else self.options[GPUArch].value ) + dump_dir = ( + envar.dump_dir + if self.options[DumpDir].value == "" + else self.options[DumpDir].value + ) if self.options[KeepPTX].value: self.options[KeepPTX].dump_path = os.path.join( - envar.dump_dir, f"{function_name}" + dump_dir, f"{function_name}" ) self.options[KeepPTX].full_ptx_path = os.path.join( - envar.dump_dir, f"{function_name}.{arch}.ptx" + dump_dir, f"{function_name}.{arch}.ptx" ) if self.options[KeepCUBIN].value: self.options[KeepCUBIN].dump_path = os.path.join( - envar.dump_dir, f"{function_name}" + dump_dir, f"{function_name}" ) self.options[KeepCUBIN].full_cubin_path = os.path.join( - envar.dump_dir, f"{function_name}.{arch}.cubin" + dump_dir, f"{function_name}.{arch}.cubin" ) @property @@ -504,6 +514,7 @@ def _parse_compile_options_from_str(options: str) -> CompileOptions: "keep_ptx": KeepPTX, "gpu_arch": GPUArch, "enable_tvm_ffi": EnableTVMFFI, + "dump_dir": DumpDir, } return mapping[option_str] @@ -520,6 +531,7 @@ def _parse_compile_options_from_str(options: str) -> CompileOptions: parser.add_argument("--ptxas-options", type=str, default="") parser.add_argument("--gpu-arch", type=str, default="") parser.add_argument("--enable-tvm-ffi", action="store_true", default=False) + parser.add_argument("--dump-dir", type=str, default="") compile_options = CompileOptions() try: # Use shlex to properly handle options with spaces @@ -545,7 +557,7 @@ class CompileCallable: def __init__(self, options=None): def preprocess_options(option): if type(option) is type and issubclass( - option, (BooleanCompileOption, BooleanBasedFileDumpOption) + option, (BooleanCompileOption, BooleanBasedFileDumpOption, EnableTVMFFI) ): # Automatically creates a True instance of the option return option(True) diff --git a/python/CuTeDSL/cutlass/base_dsl/dsl.py b/python/CuTeDSL/cutlass/base_dsl/dsl.py index e9e8662f..e19b0888 100644 --- a/python/CuTeDSL/cutlass/base_dsl/dsl.py +++ b/python/CuTeDSL/cutlass/base_dsl/dsl.py @@ -53,7 +53,7 @@ from .runtime.jit_arg_adapters import is_argument_constexpr, JitArgAdapterRegist from .ast_preprocessor import DSLPreprocessor from .common import * -from .typing import get_c_pointers, get_mlir_types, Integer, arg_compatible_with_tvm_ffi +from .typing import get_c_pointers, get_mlir_types, Integer from .arch import Arch # ============================================================================= @@ -810,10 +810,6 @@ class BaseDSL: if is_host: if self.envar.enable_tvm_ffi: - if not arg_compatible_with_tvm_ffi(arg): - raise DSLRuntimeError( - f"Argument #{i + 1} ({arg_name}) is not a TVM FFI argument." - ) jit_exec_arg.extend([arg]) else: jit_exec_arg.extend(get_c_pointers(arg)) @@ -1218,6 +1214,8 @@ class BaseDSL: no_cache, func_type=JitCompiledFunction, *, + full_args=None, + full_kwargs=None, dynamic_args=None, dynamic_kwargs=None, original_function_name=None, @@ -1388,6 +1386,8 @@ class BaseDSL: pipeline, args_spec, no_cache, + full_args=args, + full_kwargs=kwargs, dynamic_args=dynamic_args, dynamic_kwargs=dynamic_kwargs, original_function_name=original_function_name, @@ -1400,7 +1400,7 @@ class BaseDSL: module_hash, ) jit_function = self.jit_cache[module_hash] - + finally: self.post_compilation_cleanup() diff --git a/python/CuTeDSL/cutlass/base_dsl/env_manager.py b/python/CuTeDSL/cutlass/base_dsl/env_manager.py index 94873e94..cd0aa79b 100644 --- a/python/CuTeDSL/cutlass/base_dsl/env_manager.py +++ b/python/CuTeDSL/cutlass/base_dsl/env_manager.py @@ -235,27 +235,20 @@ def get_prefix_dsl_libs(prefix: str): return prefix_libs_existing def get_libs_cand(start): - target_libs = { - "mlir_c_runner_utils", - "mlir_runner_utils", - "mlir_cuda_runtime", + target_cuda_dialect_libs = { + "cuda_dialect_runtime", } lib_folder_guesses = [ "lib", ] - libs_cand = find_libs_in_ancestors(start, target_libs, lib_folder_guesses) - - optional_libs_cand = find_libs_in_ancestors( - start, {"cuda_dialect_runtime"}, lib_folder_guesses - ) - - if libs_cand: - dsl_libs = ":".join(libs_cand) - if optional_libs_cand: - dsl_libs += ":" + ":".join(optional_libs_cand) - return dsl_libs - + for target_libs in [ + target_cuda_dialect_libs, + ]: + libs_cand = find_libs_in_ancestors(start, target_libs, lib_folder_guesses) + if libs_cand: + dsl_libs = ":".join(libs_cand) + return dsl_libs return None # find from install folder diff --git a/python/CuTeDSL/cutlass/base_dsl/jit_executor.py b/python/CuTeDSL/cutlass/base_dsl/jit_executor.py index 0982953f..f36ad862 100644 --- a/python/CuTeDSL/cutlass/base_dsl/jit_executor.py +++ b/python/CuTeDSL/cutlass/base_dsl/jit_executor.py @@ -216,6 +216,69 @@ class ExecutionArgs: return exe_args, adapted_args + def get_rectified_args_from_original_args(self, full_args, full_kwargs): + """ + This function is used to rectify the original arguments to the runtime + arguments that matched the original args_spec. + + :param full_args: The original full arguments to filter. + :param full_kwargs: The original full keyword arguments to filter. + :return: The filtered arguments and keyword arguments. + """ + arg_spec = self.original_args_spec + + if arg_spec.defaults: + defaults_start_idx = len(arg_spec.args) - len(arg_spec.defaults) + else: + defaults_start_idx = len(arg_spec.args) + + runtime_args = [] + # Filter arguments and maintain their properties + for i, arg_name in enumerate(arg_spec.args): + arg_type = arg_spec.annotations.get(arg_name, None) + + # Skip compile-time arguments + if is_arg_spec_constexpr(arg_type, arg_name, i, self.function_name): + continue + + # Keep corresponding default if it exists + if i >= defaults_start_idx: + default_idx = i - defaults_start_idx + runtime_args.append(arg_spec.defaults[default_idx]) + else: + runtime_args.append(full_args[i]) + + # Filter keyword-only arguments + runtime_kwargs = {} + if arg_spec.kwonlyargs: + for i, kwarg in enumerate(arg_spec.kwonlyargs): + arg_type = arg_spec.annotations.get(kwarg, None) + + # Skip compile-time arguments + if is_arg_spec_constexpr(arg_type, kwarg, i, self.function_name): + continue + + # Keep runtime keyword-only arguments + if kwarg in full_kwargs: + runtime_kwargs[kwarg] = full_kwargs[kwarg] + elif arg_spec.kwonlydefaults and kwarg in arg_spec.kwonlydefaults: + runtime_kwargs[kwarg] = arg_spec.kwonlydefaults[kwarg] + + if (len(runtime_args) != len(self.args_spec.args) or + len(runtime_kwargs) != len(self.args_spec.kwonlyargs)): + raise DSLRuntimeError( + "input args/kwargs length does not match runtime function signature!", + context={ + "input args length": len(runtime_args), + "input kwargs length": len(runtime_kwargs), + "function signature args length": len(self.args_spec.args), + "function signature kwonlyargs length": len(self.args_spec.kwonlyargs), + }, + ) + + return runtime_args + list(runtime_kwargs.values()) + + def filter_runtime_arg_spec(self, arg_spec: inspect.FullArgSpec): runtime_args = [] runtime_annotations = {} @@ -562,6 +625,13 @@ class JitCompiledFunction: kernel_modules[sym] = CudaModuleAndKernel(sym, cubin_module, kernel, attrs) return list(kernel_modules.values()) + def _validate_engine(self): + if self.engine is None: + raise DSLRuntimeError( + "The compiled function does not have a valid execution engine.", + suggestion="For cross-compilation, please use `cute.export.export_to_c` to serialize the compiled function and load/execute it on target device.", + ) + def to(self, device=None) -> JitExecutor: """Returns an executable function bound to the given device. @@ -573,6 +643,7 @@ class JitCompiledFunction: :return: A callable executor function. :rtype: JitExecutor """ + self._validate_engine() with self._executor_lock: # We need to ensure that the modules are loaded if not already if self.jit_module is None: @@ -621,4 +692,4 @@ class JitCompiledFunction: # object alive as it hold a reference to self. proxy_self = weakref.proxy(self) self._default_executor = proxy_self.to(None) - self._default_executor.run_compiled_program(exe_args) + return self._default_executor.run_compiled_program(exe_args) diff --git a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/call_provider.py b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/call_provider.py index 03512528..c9ad4d20 100644 --- a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/call_provider.py +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/call_provider.py @@ -20,6 +20,18 @@ from ..._mlir.dialects import llvm from .tvm_ffi_builder import CallContext, CallProvider, TVMFFIBuilder +def _flatten_tuple_params(params: list[spec.Param]) -> list[spec.Param]: + """Recursively flatten TupleParam into list of params.""" + flattened = [] + for param in params: + if isinstance(param, spec.TupleParam): + # Recursively flatten nested tuples + flattened.extend(_flatten_tuple_params(param.params)) + else: + flattened.append(param) + return flattened + + class NopCallProvider(CallProvider): """No-op call provider for testing purposes.""" @@ -53,7 +65,11 @@ class DynamicParamPackCallProvider(CallProvider, TVMFFIBuilder): """ def __init__( - self, target_func: str, include_num_args: bool = False, struct_call: bool = False, + self, + target_func: str, + include_num_args: bool = False, + struct_call: bool = False, + flatten_tuple_params: bool = True, ) -> None: import tvm_ffi @@ -61,8 +77,12 @@ class DynamicParamPackCallProvider(CallProvider, TVMFFIBuilder): self.target_func = target_func self.include_num_args = include_num_args self.struct_call = struct_call + self.flatten_tuple_params = flatten_tuple_params self.float4x2_dtype = tvm_ffi.dtype("float4_e2m1fnx2") + if not self.flatten_tuple_params: + raise RuntimeError("flatten_tuple_params=False is not supported yet") + def get_callee_struct_for_param_tensor( self, param: spec.Tensor, @@ -157,8 +177,15 @@ class DynamicParamPackCallProvider(CallProvider, TVMFFIBuilder): self, current_block: ir.Block, context: CallContext ) -> list[tuple[ir.Type, ir.Value]]: """Pack a parameter to a struct.""" + # Flatten TupleParam into list of params if enabled + if self.flatten_tuple_params: + flattened_params = _flatten_tuple_params(context.params) + else: + flattened_params = context.params + + # Pack each parameter packed_params = [] - for param in context.params: + for param in flattened_params: if isinstance(param, spec.Tensor): packed_params.append( self.pack_param_tensor(current_block, context, param) @@ -177,6 +204,9 @@ class DynamicParamPackCallProvider(CallProvider, TVMFFIBuilder): packed_params.append( self.pack_param_var(current_block, context, param.var) ) + elif isinstance(param, spec.ConstNone): + # const none is not packed + continue else: raise NotImplementedError(f"Unsupported parameter type: {type(param)}") return packed_params diff --git a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py index 943ce262..36a3aadc 100644 --- a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/mlir_builder.py @@ -302,7 +302,10 @@ class MLIRBuilder(MLIRTypeBuilder): cond: ir.Value, true_block: ir.Block, false_block: ir.Block, - branch_weights=None + *, + branch_weights=None, + true_dest_operands: Sequence[ir.Value] = (), + false_dest_operands: Sequence[ir.Value] = (), ) -> None: """Create a conditional branch. @@ -318,6 +321,10 @@ class MLIRBuilder(MLIRTypeBuilder): Optional branch weights [true_weight, false_weight] for optimization hints. Higher values indicate higher probability. For example, (99, 1) indicates the true branch is much more likely than the false branch. + true_dest_operands : Sequence[ir.Value] + Operands to pass to the true destination block. + false_dest_operands : Sequence[ir.Value] + Operands to pass to the false destination block. """ if branch_weights is not None: # Branch weights should be a tuple/list of two integers [true_weight, false_weight] @@ -325,8 +332,8 @@ class MLIRBuilder(MLIRTypeBuilder): raise ValueError("branch_weights must have exactly 2 elements") llvm.cond_br( cond, - true_dest_operands=[], - false_dest_operands=[], + true_dest_operands=true_dest_operands, + false_dest_operands=false_dest_operands, true_dest=true_block, false_dest=false_block, branch_weights=ir.DenseI32ArrayAttr.get(list(branch_weights)) @@ -334,8 +341,8 @@ class MLIRBuilder(MLIRTypeBuilder): else: llvm.cond_br( cond, - true_dest_operands=[], - false_dest_operands=[], + true_dest_operands=true_dest_operands, + false_dest_operands=false_dest_operands, true_dest=true_block, false_dest=false_block, ) @@ -401,6 +408,17 @@ class MLIRBuilder(MLIRTypeBuilder): ) func_op.attributes["llvm.linkage"] = ir.StringAttr.get("external") + def create_alloca(self, entry_block: ir.Block, alloca_type: ir.Type, array_size: int) -> ir.Value: + """Create an alloca operation.""" + with ir.InsertionPoint(entry_block.operations[0]): + # declare the struct type + alloca = llvm.alloca( + res=self.ptr_type, + elem_type=alloca_type, + array_size=self.i32(array_size), + ) + return alloca + def pack_values_to_alloca( self, current_block: ir.Block, @@ -423,14 +441,11 @@ class MLIRBuilder(MLIRTypeBuilder): tuple[ir.Type, ir.Value] The struct type and the alloca. """ - with ir.InsertionPoint(entry_block.operations[0]): - # declare the struct type - struct_type = self.struct_type(fields=[value.type for value in values]) - alloca = llvm.alloca( - res=self.ptr_type, - elem_type=struct_type, - array_size=self.i32(1), - ) + # Declare the struct type from the values + struct_type = self.struct_type(fields=[value.type for value in values]) + + # Create alloca using the helper method + alloca = self.create_alloca(entry_block, struct_type, array_size=1) with ir.InsertionPoint(current_block): for index, value in enumerate(values): diff --git a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/spec.py b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/spec.py index 62f7698b..be66771d 100644 --- a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/spec.py +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/spec.py @@ -136,7 +136,10 @@ class Shape(Param): name: str shape: list[Union[int, Var]] - def __init__(self, name: str, shape: list[Union[int, Var]]) -> None: + def __init__( + self, + name: str, shape: list[Union[int, Var]], + ) -> None: """Initialize a Shape parameter. Parameters @@ -146,6 +149,9 @@ class Shape(Param): shape : list[int | Var] The shape of the parameter. + unpack_shape: bool + Whether to unpack the shape into list of arguments when calling + the call provider function. """ self.name = name self.shape = shape @@ -301,6 +307,106 @@ class DataPointer(Param): self.address_space = address_space +class ConstNone(Param): + """ConstNone parameter. + + Parameters + ---------- + name : str + The parameter name. + """ + name: str + + def __init__(self, name: str) -> None: + """Initialize a ConstExpr parameter. + + Parameters + ---------- + name : str + The parameter name. + """ + self.name = name + + +class TupleParam(Param): + """Tuple parameter. + + Parameters + ---------- + name : str + The parameter name. + """ + name: str + params: list[Param] + + def __init__(self, name: str, params: list[Param]) -> None: + """Initialize a TupleParam parameter. + + Parameters + ---------- + name : str + The parameter name. + params : list[Param] + The parameters of the tuple. + """ + self.name = name + self.params = params + + +def format_param_type(param: Param) -> str: + """Format a parameter type as a string, recursively handling nested types. + + Parameters + ---------- + param : Param + The parameter to format. + + Returns + ------- + str + The formatted type string. + + Raises + ------ + TypeError + If an unsupported parameter type is encountered. + """ + if isinstance(param, Var): + return str(param.dtype) + elif isinstance(param, Tensor): + # Format tensor shape + shape_strs = [] + for dim in param.shape: + if isinstance(dim, Var): + shape_strs.append(dim.name) + else: + shape_strs.append(str(dim)) + shape_str = "[" + ", ".join(shape_strs) + "]" + return f"Tensor({shape_str}, {param.dtype})" + elif isinstance(param, Shape): + # Format shape parameter + shape_strs = [] + for dim in param.shape: + if isinstance(dim, Var): + shape_strs.append(dim.name) + else: + shape_strs.append(str(dim)) + shape_str = "[" + ", ".join(shape_strs) + "]" + return f"Shape({shape_str})" + elif isinstance(param, Stream): + return "Stream" + elif isinstance(param, DataPointer): + return "DataPointer" + elif isinstance(param, ConstNone): + return "None" + elif isinstance(param, TupleParam): + # Recursively format tuple elements + element_types = [format_param_type(p) for p in param.params] + return f"Tuple[{', '.join(element_types)}]" + else: + raise TypeError(f"Unsupported parameter type: {type(param)}") + + def signature(name: str, params: list[Param]) -> str: """Generate a function signature string from name and parameters. @@ -325,39 +431,13 @@ def signature(name: str, params: list[Param]) -> str: param_strs = [] for param in params: - if isinstance(param, Var): - param_str = f"{param.name}: {param.dtype}" - elif isinstance(param, Tensor): - # Format tensor shape - shape_strs = [] - for dim in param.shape: - if isinstance(dim, Var): - shape_strs.append(dim.name) - else: - shape_strs.append(str(dim)) - shape_str = "[" + ", ".join(shape_strs) + "]" - param_str = f"{param.name}: Tensor({shape_str}, {param.dtype})" - elif isinstance(param, Shape): - # Format shape parameter - shape_strs = [] - for dim in param.shape: - if isinstance(dim, Var): - shape_strs.append(dim.name) - else: - shape_strs.append(str(dim)) - shape_str = "[" + ", ".join(shape_strs) + "]" - param_str = f"{param.name}: Shape({shape_str})" - elif isinstance(param, Stream): - param_str = f"{param.name}: Stream" - elif isinstance(param, DataPointer): - param_str = f"{param.name}: DataPointer" - elif isinstance(param, EnvStream): + if isinstance(param, EnvStream): # env stream is not part of the FFI function signature # continue to skip append continue - else: - raise TypeError(f"Unsupported parameter type: {type(param)}") + param_type = format_param_type(param) + param_str = f"{param.name}: {param_type}" param_strs.append(param_str) return f"{name}({', '.join(param_strs)})" diff --git a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/tvm_ffi_builder.py b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/tvm_ffi_builder.py index 0fe414f9..7353210c 100644 --- a/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/tvm_ffi_builder.py +++ b/python/CuTeDSL/cutlass/base_dsl/tvm_ffi_builder/tvm_ffi_builder.py @@ -27,6 +27,64 @@ from .mlir_builder import MLIRBuilder from dataclasses import dataclass +@dataclass +class ArgContext: + """Context information for parameter decoding error messages. + + :ivar param_name: The name of the parameter. + :vartype param_name: str + :ivar arg_index: The index of the argument in the function call. + :vartype arg_index: int + :ivar tuple_indices: List of tuple indices for nested tuple access (e.g., [0, 1] for tuple[0][1]). + :vartype tuple_indices: list[int] + """ + param_name: str + arg_index: int + tuple_indices: list[int] + + def get(self) -> list[str]: + """Get the context as a list of strings for error messages. + + :returns: Context strings like ["on argument ", "#0"] or ["on my_tuple[0][1] in argument ", "#0"]. + :rtype: list[str] + """ + if not self.tuple_indices: + # Top-level argument: "on argument #0" + return ["on argument ", f"#{self.arg_index}"] + else: + # Nested tuple element: "on my_tuple[0][1] in argument #0" + indices_str = "".join(f"[{i}]" for i in self.tuple_indices) + return [f"on {self.param_name}{indices_str} in argument ", f"#{self.arg_index}"] + + def get_field_name(self, field_suffix: str) -> str: + """Get the field name with tuple indices for shape/stride access. + + :param field_suffix: The field suffix (e.g., ".shape", ".strides"). + :type field_suffix: str + :returns: Field name like "my_param.shape" or "my_tuple[0][1].shape". + :rtype: str + """ + if not self.tuple_indices: + return f"{self.param_name}{field_suffix}" + else: + indices_str = "".join(f"[{i}]" for i in self.tuple_indices) + return f"{self.param_name}{indices_str}{field_suffix}" + + def get_element_context(self, element_index: int) -> "ArgContext": + """Create a nested context for a tuple element. + + :param element_index: The index within the tuple. + :type element_index: int + :returns: New context for the nested element. + :rtype: ArgContext + """ + return ArgContext( + param_name=self.param_name, + arg_index=self.arg_index, + tuple_indices=self.tuple_indices + [element_index], + ) + + @dataclass class CallContext: """Call context that contains the information of the call.""" @@ -119,7 +177,7 @@ class TVMFFIBuilder(MLIRBuilder): super().__init__() # this is a number we can tune to minimize the register size # it is 6 by default to minimize the register size - self.set_raised_from_cstr_parts_max_num_parts = 6 + self.set_raised_from_cstr_parts_max_num_parts = 8 self.set_raised_from_cstr_parts_cache: dict[int, str] = {} self.tvm_ffi_any_type = self.struct_type( name="TVMFFIAny", @@ -488,26 +546,28 @@ class TVMFFIBuilder(MLIRBuilder): ) -> ir.Value: """Downcast i64 to lower bits.""" overflow_flags = llvm.IntegerOverflowFlags.none + if (hasattr(tvm_ffi._dtype.DataTypeCode, "BOOL") and + target_dtype.type_code == tvm_ffi._dtype.DataTypeCode.BOOL): + # LLVM use i1 (boolean) for boolean + return llvm.icmp(llvm.ICmpPredicate.ne, v_int64, self.i64(0)) if target_dtype.bits == 64: - result = v_int64 - elif target_dtype.bits == 32: - result = llvm.trunc( + return v_int64 + if target_dtype.bits == 32: + return llvm.trunc( res=self.i32_type, arg=v_int64, overflow_flags=overflow_flags ) - elif target_dtype.bits == 16: - result = llvm.trunc( + if target_dtype.bits == 16: + return llvm.trunc( res=self.i16_type, arg=v_int64, overflow_flags=overflow_flags ) - elif target_dtype.bits == 8: - result = llvm.trunc( + if target_dtype.bits == 8: + return llvm.trunc( res=self.i8_type, arg=v_int64, overflow_flags=overflow_flags ) - elif target_dtype.bits == 1: + if target_dtype.bits == 1: # For i1 (boolean), convert i64 to boolean by checking if non-zero - result = llvm.icmp(llvm.ICmpPredicate.ne, v_int64, self.i64(0)) - else: - raise ValueError(f"Unsupported Var dtype: {target_dtype}") - return result + return llvm.icmp(llvm.ICmpPredicate.ne, v_int64, self.i64(0)) + raise ValueError(f"Unsupported Var dtype: {target_dtype}") def is_contiguous( self, @@ -767,8 +827,40 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): self.matched_var_binding = {} self.matched_var_source = {} + def find_or_declare_extern_func( + self, name: str, params: Sequence[ir.Type], ret: ir.Type + ) -> None: + """Find an existing extern function or declare it if it doesn't exist. + + This method checks if a function with the given name already exists in the module. + If it does, the method returns without doing anything. Otherwise, it declares + the function as an external function. + + Parameters + ---------- + name : str + The name of the extern function. + params : Sequence[ir.Type] + The parameter types of the function. + ret : ir.Type + The return type of the function. + """ + # Check if the function already exists + existing_func = self.find_func_in_module(self.module, name) + if existing_func is not None: + # Function already declared, nothing to do + return + + # Function doesn't exist, declare it + self.declare_extern_func(name, params, ret) + def decode_param_int( - self, current_block: ir.Block, param: spec.Var, args: ir.Value, arg_index: int + self, + current_block: ir.Block, + param: spec.Var, + args: ir.Value, + arg_index: int, + arg_context: ArgContext, ) -> ir.Block: """Decode the integer parameter at the given index.""" # read the type index @@ -785,8 +877,8 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): lambda: is_int_or_bool, "TypeError", [ - "Mismatched type on argument ", - f"#{arg_index}", + "Mismatched type ", + *arg_context.get(), self._fn_call_context, ", expected int", ], @@ -801,14 +893,19 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): param, v_int64, [ - "value on argument ", - f"#{arg_index}", + "value ", + *arg_context.get(), self._fn_call_context, ], ) def decode_param_float( - self, current_block: ir.Block, param: spec.Var, args: ir.Value, arg_index: int + self, + current_block: ir.Block, + param: spec.Var, + args: ir.Value, + arg_index: int, + arg_context: ArgContext, ) -> ir.Block: """Decode the float parameter at the given index.""" # read the type index @@ -891,8 +988,8 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): self.raise_error_and_return( "TypeError", [ - "Mismatched type on argument ", - f"#{arg_index}", + "Mismatched type ", + *arg_context.get(), self._fn_call_context, ", expected float", ], @@ -912,6 +1009,7 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): param: spec.Var, args: ir.Value, arg_index: int, + arg_context: ArgContext, *, allow_int_as_ptr: bool = False, address_space: Optional[int] = None, @@ -941,8 +1039,8 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): lambda: is_opaque_ptr_or_nullptr, "TypeError", [ - "Mismatched type on argument ", - f"#{arg_index}", + "Mismatched type ", + *arg_context.get(), self._fn_call_context, expect_message, ], @@ -959,6 +1057,37 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): return current_block + def decode_param_const_none( + self, + current_block: ir.Block, + param: spec.ConstNone, + args: ir.Value, + arg_index: int, + arg_context: ArgContext, + ) -> ir.Block: + """Decode the opaque handle parameter at the given index.""" + # read the type index + with ir.InsertionPoint(current_block): + type_index: ir.Value = self.load_ffi_any_array_item_type_index(args, arg_index) + # Check if type is a nullptr + is_nullptr = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFINone)) + + expect_message = ", expected None" + + # Break error message into reusable parts for better string deduplication + current_block = self.check_condition( + current_block, + lambda: is_nullptr, + "TypeError", + [ + "Mismatched type ", + *arg_context.get(), + self._fn_call_context, + expect_message, + ], + ) + return current_block + def check_int_value_dtype_bound( self, current_block: ir.Block, @@ -1105,9 +1234,16 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): *error_msg_context, f", expected to be {var}" ] + + def check_value_mismatch() -> ir.Value: + cond = self.equal(value, expected_value) + if skip_check_predicate is not None: + cond = self.or_(skip_check_predicate, cond) + return cond + return self.check_condition( current_block, - lambda: self.equal(value, expected_value), + check_value_mismatch, error_kind, error_msg_mismatch ) @@ -1117,17 +1253,18 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): current_block: ir.Block, var: Union[spec.Var, int], value: ir.Value, - field: str, - arg_index: int, + field_suffix: str, shape_index: int, + arg_context: ArgContext, *, skip_check_predicate: Optional[ir.Value] = None, ) -> ir.Block: """Load the shape value from the argument or match the shape value from the parameter.""" + field_name = arg_context.get_field_name(field_suffix) error_msg = [ - field, - f"[{shape_index}] on argument ", - f"#{arg_index}", + field_name, + f"[{shape_index}] ", + *arg_context.get(), self._fn_call_context, ] return self.set_or_check_matched_var_binding( @@ -1135,7 +1272,11 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): ) def decode_param_shape_from_ffi_array( - self, current_block: ir.Block, param: spec.Shape, arg_index: int, array_cell: ir.Value + self, + current_block: ir.Block, + param: spec.Shape, + arg_context: ArgContext, + array_cell: ir.Value, ) -> tuple[ir.Block, list[ir.Value]]: """Decode the shape parameter from the TVMFFIArrayCell.""" with ir.InsertionPoint(current_block): @@ -1148,8 +1289,8 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): lambda: self.equal(array_size, self.i64(len(param.shape))), "ValueError", [ - "Mismatched Shape on argument ", - f"#{arg_index}", + "Mismatched Shape ", + *arg_context.get(), self._fn_call_context, f", expected shape size={len(param.shape)}", ], @@ -1157,35 +1298,49 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): # Load and validate each element of the array load_shapes = [] - for i in range(len(param.shape)): - with ir.InsertionPoint(current_block): - type_index: ir.Value = self.load_ffi_any_array_item_type_index(array_data, i) + + def validate_and_load_shape_element( + block: ir.Block, index: int + ) -> tuple[ir.Block, ir.Value]: + """Validate and load a single shape element from the array.""" + with ir.InsertionPoint(block): + type_index: ir.Value = self.load_ffi_any_array_item_type_index(array_data, index) # Check if type is int or bool (both use v_int64, bool can be converted to int) - is_int = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIInt)) + is_int_val = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIInt)) # Check that the element is an integer - current_block = self.check_condition( - current_block, - lambda: is_int, + field_name = arg_context.get_field_name("") + block = self.check_condition( + block, + lambda: is_int_val, "TypeError", [ - f"Invalid shape element type ", - f"{param.name}[{i}]", - f" on argument ", - f"#{arg_index}", + "Invalid shape element type ", + f"{field_name}[{index}]", + " ", + *arg_context.get(), self._fn_call_context, - f", expected int", + ", expected int", ], ) - with ir.InsertionPoint(current_block): - v_int64: ir.Value = self.load_ffi_any_array_item_v_int64(array_data, i) - load_shapes.append(v_int64) + with ir.InsertionPoint(block): + v_int64: ir.Value = self.load_ffi_any_array_item_v_int64(array_data, index) + + return block, v_int64 + + for i in range(len(param.shape)): + current_block, v_int64 = validate_and_load_shape_element(current_block, i) + load_shapes.append(v_int64) return (current_block, load_shapes) def decode_param_shape_from_ffi_shape( - self, current_block: ir.Block, param: spec.Shape, arg_index: int, shape_cell: ir.Value + self, + current_block: ir.Block, + param: spec.Shape, + arg_context: ArgContext, + shape_cell: ir.Value, ) -> tuple[ir.Block, list[ir.Value]]: """Decode the shape parameter from the TVMFFIShapeCell.""" with ir.InsertionPoint(current_block): @@ -1198,8 +1353,8 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): lambda: self.equal(shape_size, self.i64(len(param.shape))), "ValueError", [ - "Mismatched Shape on argument ", - f"#{arg_index}", + "Mismatched Shape ", + *arg_context.get(), self._fn_call_context, f", expected shape size={len(param.shape)}", ], @@ -1211,7 +1366,12 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): return (current_block, load_shapes) def decode_param_shape( - self, current_block: ir.Block, param: spec.Shape, args: ir.Value, arg_index: int + self, + current_block: ir.Block, + param: spec.Shape, + args: ir.Value, + arg_index: int, + arg_context: ArgContext, ) -> ir.Block: """Decode the shape parameter at the given index.""" with ir.InsertionPoint(current_block): @@ -1238,7 +1398,7 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): self.load_ffi_any_array_item_v_ptr(args, arg_index) ) ffi_shape_block, load_shapes = self.decode_param_shape_from_ffi_shape( - ffi_shape_block, param, arg_index, shape_cell + ffi_shape_block, param, arg_context, shape_cell ) with ir.InsertionPoint(ffi_shape_block): self.br(subsequent_block, args=load_shapes) @@ -1256,7 +1416,7 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): self.load_ffi_any_array_item_v_ptr(args, arg_index) ) ffi_array_block, load_shapes = self.decode_param_shape_from_ffi_array( - ffi_array_block, param, arg_index, array_cell_ptr + ffi_array_block, param, arg_context, array_cell_ptr ) with ir.InsertionPoint(ffi_array_block): self.br(subsequent_block, args=load_shapes) @@ -1267,8 +1427,8 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): self.raise_error_and_return( "TypeError", [ - "Mismatched type on argument ", - f"#{arg_index}", + "Mismatched type ", + *arg_context.get(), self._fn_call_context, ", expected ffi.Shape or ffi.Array", ], @@ -1279,7 +1439,7 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): shape_values = list(subsequent_block.arguments) for i, dim in enumerate(param.shape): subsequent_block = self.set_or_check_matched_var_binding_from_shape( - subsequent_block, dim, shape_values[i], f"{param.name}", arg_index, i + subsequent_block, dim, shape_values[i], "", i, arg_context ) return subsequent_block @@ -1290,6 +1450,7 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): param: spec.Tensor, args: ir.Value, arg_index: int, + arg_context: ArgContext, ) -> tuple[ir.Block, ir.Value]: """Decode tensor step0: check index and find out DLTensor*.""" # read the type index @@ -1333,8 +1494,8 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): self.raise_error_and_return( "TypeError", [ - "Mismatched type on argument ", - f"#{arg_index}", + "Mismatched type ", + *arg_context.get(), self._fn_call_context, ", expected Tensor", ], @@ -1352,10 +1513,11 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): param: spec.Tensor, args: ir.Value, arg_index: int, + arg_context: ArgContext, ) -> ir.Block: """Decode the tensor parameter at the given index.""" current_block, dl_tensor_ptr = self.decode_param_tensor_dltensor_ptr( - current_block, param, args, arg_index + current_block, param, args, arg_index, arg_context ) with ir.InsertionPoint(current_block): data = self.load_dltensor_data_ptr(dl_tensor_ptr) @@ -1381,8 +1543,8 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): check_alignment, "ValueError", [ - "Misaligned Tensor data on argument ", - f"#{arg_index}", + "Misaligned Tensor data ", + *arg_context.get(), self._fn_call_context, f", expected data alignment={param.data_alignment} bytes", ], @@ -1401,8 +1563,8 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): lambda: self.equal(ndim, self.i32(expected_ndim)), "ValueError", [ - "Mismatched Tensor on argument ", - f"#{arg_index}", + "Mismatched Tensor ", + *arg_context.get(), self._fn_call_context, f", expected ndim={expected_ndim}", ], @@ -1414,8 +1576,8 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): lambda: self.equal(device_type, self.i32(param.dlpack_device_type)), "ValueError", [ - "Mismatched Tensor on argument ", - f"#{arg_index}", + "Mismatched Tensor ", + *arg_context.get(), self._fn_call_context, f", expected device_type={param.device_type_name}", ], @@ -1437,8 +1599,8 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): dtype_equal, "ValueError", [ - "Mismatched Tensor on argument ", - f"#{arg_index}", + "Mismatched Tensor ", + *arg_context.get(), self._fn_call_context, f", expected dtype={param.dtype}", ], @@ -1450,8 +1612,8 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): lambda: self.equal(byte_offset, self.i64(0)), "ValueError", [ - "Mismatched Tensor on argument ", - f"#{arg_index}", + "Mismatched Tensor ", + *arg_context.get(), self._fn_call_context, ", expected byte_offset=0", ], @@ -1474,9 +1636,9 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): current_block, param.shape[index], load_shapes[index], - f"{param.name}.shape", - arg_index, + ".shape", index, + arg_context, ) if param.strides is not None: @@ -1490,9 +1652,9 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): current_block, param.strides[index], load_strides[index], - f"{param.name}.strides", - arg_index, + ".strides", index, + arg_context, skip_check_predicate=skip_check_predicate, ) else: @@ -1502,8 +1664,8 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): lambda: self.is_contiguous(param.shape, load_shapes, load_strides), "ValueError", [ - "Mismatched Tensor on argument ", - f"#{arg_index}", + "Mismatched Tensor ", + *arg_context.get(), self._fn_call_context, ", expected contiguous", ], @@ -1516,11 +1678,12 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): param: spec.Stream, args: ir.Value, arg_index: int, + arg_context: ArgContext, ) -> ir.Block: """Decode the stream parameter at the given index.""" # stream is decoded as opaque handle return self.decode_param_opaque_handle( - current_block, param.var, args, arg_index + current_block, param.var, args, arg_index, arg_context ) def decode_param_data_pointer( @@ -1529,6 +1692,7 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): param: spec.DataPointer, args: ir.Value, arg_index: int, + arg_context: ArgContext, ) -> ir.Block: """Decode the data pointer parameter at the given index.""" # data pointer is decoded as opaque handle @@ -1537,6 +1701,7 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): param.var, args, arg_index, + arg_context, allow_int_as_ptr=True, address_space=param.address_space, ) @@ -1578,35 +1743,117 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): expected_num_args += 1 return expected_num_args - def decode_param( # noqa: PLR0911 - self, current_block: ir.Block, param: spec.Param, args: ir.Value, arg_index: int + def decode_param_tuple( + self, + current_block: ir.Block, + param: spec.TupleParam, + args: ir.Value, + arg_index: int, + arg_context: ArgContext, ) -> ir.Block: - """Decode the parameter at the given index.""" + """Decode the tuple parameter at the given index.""" + # Check if type is kTVMFFIArray + with ir.InsertionPoint(current_block): + type_index: ir.Value = self.load_ffi_any_array_item_type_index(args, arg_index) + is_ffi_array = self.equal(type_index, self.i32(TVMFFITypeIndex.kTVMFFIArray)) + + # Check that the type is an array + current_block = self.check_condition( + current_block, + lambda: is_ffi_array, + "TypeError", + [ + "Mismatched type ", + *arg_context.get(), + self._fn_call_context, + ", expected ffi.Array for tuple", + ], + ) + + # Load the array cell + with ir.InsertionPoint(current_block): + array_cell_ptr = self.get_object_cell_ptr( + self.load_ffi_any_array_item_v_ptr(args, arg_index) + ) + array_data = self.load_array_cell_data_ptr(array_cell_ptr) + array_size = self.load_array_cell_size_as_i64(array_cell_ptr) + + # Check that the array size matches the expected tuple size + current_block = self.check_condition( + current_block, + lambda: self.equal(array_size, self.i64(len(param.params))), + "ValueError", + [ + "Mismatched tuple size ", + *arg_context.get(), + self._fn_call_context, + f", expected tuple size={len(param.params)}", + ], + ) + + # Recursively decode each element of the tuple + for i, tuple_param in enumerate(param.params): + # Create nested context for the tuple element + nested_context = arg_context.get_element_context(i) + current_block = self.decode_param(current_block, tuple_param, array_data, i, nested_context) + + return current_block + + def decode_param( # noqa: PLR0911 + self, + current_block: ir.Block, + param: spec.Param, + args: ir.Value, + arg_index: int, + arg_context: ArgContext, + ) -> ir.Block: + """Decode the parameter at the given index. + + Parameters + ---------- + current_block : ir.Block + The current IR block. + param : spec.Param + The parameter specification to decode. + args : ir.Value + The FFI arguments array. + arg_index : int + The index in the args array. + arg_context : ArgContext + Context information for error messages. + """ if isinstance(param, spec.Var): if param.dtype.type_code == tvm_ffi._dtype.DataTypeCode.INT: - return self.decode_param_int(current_block, param, args, arg_index) + return self.decode_param_int(current_block, param, args, arg_index, arg_context) elif param.dtype.type_code == tvm_ffi._dtype.DataTypeCode.UINT: # UINT uses the same logic as INT since both are stored in v_int64 - return self.decode_param_int(current_block, param, args, arg_index) + return self.decode_param_int(current_block, param, args, arg_index, arg_context) + elif (hasattr(tvm_ffi._dtype.DataTypeCode, "BOOL") and + param.dtype.type_code == tvm_ffi._dtype.DataTypeCode.BOOL): + return self.decode_param_int(current_block, param, args, arg_index, arg_context) elif param.dtype.type_code == tvm_ffi._dtype.DataTypeCode.FLOAT: - return self.decode_param_float(current_block, param, args, arg_index) + return self.decode_param_float(current_block, param, args, arg_index, arg_context) elif param.dtype.type_code == tvm_ffi._dtype.DataTypeCode.HANDLE: return self.decode_param_opaque_handle( - current_block, param, args, arg_index + current_block, param, args, arg_index, arg_context ) else: raise ValueError(f"Unsupported parameter type: {param.dtype.type_code}") elif isinstance(param, spec.Shape): - return self.decode_param_shape(current_block, param, args, arg_index) + return self.decode_param_shape(current_block, param, args, arg_index, arg_context) elif isinstance(param, spec.Tensor): - return self.decode_param_tensor(current_block, param, args, arg_index) + return self.decode_param_tensor(current_block, param, args, arg_index, arg_context) elif isinstance(param, spec.Stream): - return self.decode_param_stream(current_block, param, args, arg_index) + return self.decode_param_stream(current_block, param, args, arg_index, arg_context) elif isinstance(param, spec.EnvStream): # decode of env stream is deferred after we go through all parameters return current_block elif isinstance(param, spec.DataPointer): - return self.decode_param_data_pointer(current_block, param, args, arg_index) + return self.decode_param_data_pointer(current_block, param, args, arg_index, arg_context) + elif isinstance(param, spec.ConstNone): + return self.decode_param_const_none(current_block, param, args, arg_index, arg_context) + elif isinstance(param, spec.TupleParam): + return self.decode_param_tuple(current_block, param, args, arg_index, arg_context) else: raise ValueError(f"Unsupported parameter type: {type(param)}") @@ -1652,20 +1899,20 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): with ir.InsertionPoint(self.module.body): # type: ignore[union-attr] # void TVMFFIErrorSetRaisedFromCStr( # const char* error_kind, const char* message); - self.declare_extern_func( + self.find_or_declare_extern_func( "TVMFFIErrorSetRaisedFromCStr", [self.ptr_type, self.ptr_type], self.void_type, ) # void TVMFFIErrorSetRaisedFromCStrParts( # const char* error_kind, const char* messages, int32_t num_parts); - self.declare_extern_func( + self.find_or_declare_extern_func( "TVMFFIErrorSetRaisedFromCStrParts", [self.ptr_type, self.ptr_type, self.i32_type], self.void_type, ) # void* TVMFFIEnvGetStream(int32_t device_type, int32_t device_id); - self.declare_extern_func( + self.find_or_declare_extern_func( "TVMFFIEnvGetStream", [self.i32_type, self.i32_type], self.ptr_type, @@ -1697,7 +1944,12 @@ class TVMFFIFunctionBuilder(TVMFFIBuilder): # decode parameters to populate the matched var binding for arg_index, param in enumerate(params_list): - current_block = self.decode_param(current_block, param, args, arg_index) + arg_context = ArgContext( + param_name=param.name, + arg_index=arg_index, + tuple_indices=[], + ) + current_block = self.decode_param(current_block, param, args, arg_index, arg_context) with ir.InsertionPoint(current_block): env_stream = self.find_env_stream(params_list) diff --git a/python/CuTeDSL/cutlass/base_dsl/typing.py b/python/CuTeDSL/cutlass/base_dsl/typing.py index a9db8b29..73051206 100644 --- a/python/CuTeDSL/cutlass/base_dsl/typing.py +++ b/python/CuTeDSL/cutlass/base_dsl/typing.py @@ -239,19 +239,6 @@ def get_c_pointers(obj): return [] -def arg_compatible_with_tvm_ffi(arg): - """ - Given the `arg`, check if it is a compatible argument for TVM FFI - """ - import tvm_ffi - - return ( - hasattr(arg, "__tvm_ffi_object__") - or isinstance(arg, (int, float, bool)) - or isinstance(arg, tvm_ffi.Shape) - ) - - def get_mlir_types(obj): """ Given the `obj`, recursively go through it to extract all contained MLIR types diff --git a/python/CuTeDSL/cutlass/cute/__init__.py b/python/CuTeDSL/cutlass/cute/__init__.py index db42bc7d..1a5cb619 100644 --- a/python/CuTeDSL/cutlass/cute/__init__.py +++ b/python/CuTeDSL/cutlass/cute/__init__.py @@ -205,6 +205,7 @@ KeepCUBIN = _dsl.KeepCUBIN KeepPTX = _dsl.KeepPTX GPUArch = _dsl.GPUArch LinkLibraries = _dsl.LinkLibraries +EnableTVMFFI = _dsl.EnableTVMFFI # attach the TVM FFI ABI interface postprocessor to the DSL from . import _tvm_ffi_args_spec_converter diff --git a/python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py b/python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py index 642f0495..d0499ab9 100644 --- a/python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py +++ b/python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py @@ -42,7 +42,7 @@ from .typing import ( ) import cuda.bindings.driver as cuda -from typing import List, Dict, Any, Optional +from typing import List, Dict, Any, Optional, Tuple, get_origin, get_args import inspect NumericToTVMFFIDtype = { @@ -91,6 +91,11 @@ def _get_llvm_address_space_from_memspace( return 1 return None +def _is_gpu_memspace( + memspace: _cute_ir.AddressSpace, +) -> bool: + return memspace != _cute_ir.AddressSpace.generic + class SymIntId: def __init__(self, sym_int: SymInt): @@ -103,118 +108,187 @@ class SymIntId: return self.sym_int is other.sym_int -def _tvm_ffi_args_spec_converter( - function_name: str, - args_spec: inspect.FullArgSpec, - dynamic_args: List[Any], - dynamic_kwargs: Dict[str, Any], -): - """Convert cute algebra args to tvm ffi spec params. - This function converts the cute arguments specs to tvm ffi spec params. - """ - exec_args = ExecutionArgs(args_spec, function_name) - rectified_args = exec_args.get_rectified_args(dynamic_args, dynamic_kwargs) - arg_names = exec_args.args_spec.args + exec_args.args_spec.kwonlyargs +class ConverterContext: + """Context for managing variable allocation during TVM FFI args conversion.""" - params = [] - num_dyn_shape_vars = 0 - num_dyn_stride_vars = 0 - sym_int_id_mapping = {} + def __init__(self): + self.num_dyn_shape_vars = 0 + self.num_dyn_stride_vars = 0 + self.sym_int_id_mapping = {} - - def alloc_shape_name(): - nonlocal num_dyn_shape_vars - name = f"n{num_dyn_shape_vars}" - num_dyn_shape_vars += 1 + def alloc_shape_name(self) -> str: + """Allocate a new dynamic shape variable name.""" + name = f"n{self.num_dyn_shape_vars}" + self.num_dyn_shape_vars += 1 return name - def alloc_stride_name(): - nonlocal num_dyn_stride_vars - name = f"s{num_dyn_stride_vars}" - num_dyn_stride_vars += 1 + def alloc_stride_name(self) -> str: + """Allocate a new dynamic stride variable name.""" + name = f"s{self.num_dyn_stride_vars}" + self.num_dyn_stride_vars += 1 return name - def alloc_or_reuse_symint_var(value, name_alloc_func): - nonlocal sym_int_id_mapping + def alloc_or_reuse_symint_var(self, value: SymInt, name_alloc_func): + """Allocate or reuse a symbolic integer variable.""" sym_int_id = SymIntId(value) - if sym_int_id in sym_int_id_mapping: - return sym_int_id_mapping[sym_int_id] + if sym_int_id in self.sym_int_id_mapping: + return self.sym_int_id_mapping[sym_int_id] name = name_alloc_func() if value.width == 32: dtype = NumericToTVMFFIDtype[Int32] else: dtype = NumericToTVMFFIDtype[Int64] var = spec.Var(name, dtype, divisibility=value.divisibility) - sym_int_id_mapping[sym_int_id] = var + self.sym_int_id_mapping[sym_int_id] = var return var - for arg, arg_name in zip(rectified_args, arg_names): - arg_type = args_spec.annotations.get(arg_name, None) - if isinstance(arg, Numeric) and arg.dtype in AcceptableNumericTypesForScalar: - params.append(spec.Var(arg_name, NumericToTVMFFIDtype[arg.dtype])) - elif is_cute_algebra_type(arg_type): - shape = [] - for i in range(len(arg)): - if isinstance(arg[i], int): - shape.append(arg[i]) - elif isinstance(arg[i], SymInt): - shape.append(alloc_or_reuse_symint_var(arg[i], alloc_shape_name)) - else: - shape.append(spec.Var(alloc_shape_name(), NumericToTVMFFIDtype[arg[i].dtype])) - params.append(spec.Shape(arg_name, shape)) - elif isinstance(arg, Tensor): - shapes = [] - for i, dyn_mask in enumerate(arg.dynamic_shapes_mask): - if not dyn_mask: - shapes.append(arg.shape[i]) - elif isinstance(arg.shape[i], SymInt): - shapes.append(alloc_or_reuse_symint_var(arg.shape[i], alloc_shape_name)) - else: - shapes.append(spec.Var(alloc_shape_name(), NumericToTVMFFIDtype[Int32])) - strides = [] - for i, dyn_mask in enumerate(arg.dynamic_strides_mask): - if not dyn_mask: - strides.append(arg.stride[i]) - elif isinstance(arg.stride[i], SymInt): - strides.append(alloc_or_reuse_symint_var(arg.stride[i], alloc_stride_name)) - else: - if hasattr(arg, "_use_32bit_stride") and arg._use_32bit_stride: - dtype = NumericToTVMFFIDtype[Int32] - else: - dtype = NumericToTVMFFIDtype[Int64] - strides.append(spec.Var(alloc_stride_name(), dtype)) +def _convert_single_arg( + arg, + arg_name: str, + arg_type, + ctx: ConverterContext +) -> spec.Param: + """Convert a single argument to a spec.Param. + Parameters + ---------- + arg : Any + The argument value to convert. + arg_name : str + The name of the argument. + arg_type : type + The type annotation of the argument. + ctx : ConverterContext + The converter context for managing variable allocation. + + Returns + ------- + spec.Param + The converted parameter specification. + """ + if arg is None: + return spec.ConstNone(arg_name) + elif (isinstance(arg, Numeric) and arg.dtype in AcceptableNumericTypesForScalar): + return spec.Var(arg_name, NumericToTVMFFIDtype[arg.dtype]) + elif arg_type in AcceptableNumericTypesForScalar: + return spec.Var(arg_name, NumericToTVMFFIDtype[arg_type]) + elif is_cute_algebra_type(arg_type): + shape = [] + for i in range(len(arg)): + if isinstance(arg[i], int): + shape.append(arg[i]) + elif isinstance(arg[i], SymInt): + shape.append(ctx.alloc_or_reuse_symint_var(arg[i], ctx.alloc_shape_name)) + else: + shape.append(spec.Var(ctx.alloc_shape_name(), NumericToTVMFFIDtype[arg[i].dtype])) + return spec.Shape(arg_name, shape) + elif isinstance(arg, Tensor): + shapes = [] + for i, dyn_mask in enumerate(arg.dynamic_shapes_mask): + if not dyn_mask: + shapes.append(arg.shape[i]) + elif isinstance(arg.shape[i], SymInt): + shapes.append(ctx.alloc_or_reuse_symint_var(arg.shape[i], ctx.alloc_shape_name)) + else: + shapes.append(spec.Var(ctx.alloc_shape_name(), NumericToTVMFFIDtype[Int32])) + strides = [] + + for i, dyn_mask in enumerate(arg.dynamic_strides_mask): + if not dyn_mask: + strides.append(arg.stride[i]) + elif isinstance(arg.stride[i], SymInt): + strides.append(ctx.alloc_or_reuse_symint_var(arg.stride[i], ctx.alloc_stride_name)) + else: + if hasattr(arg, "_use_32bit_stride") and arg._use_32bit_stride: + dtype = NumericToTVMFFIDtype[Int32] + else: + dtype = NumericToTVMFFIDtype[Int64] + strides.append(spec.Var(ctx.alloc_stride_name(), dtype)) + if hasattr(arg, "_tvm_ffi_tensor"): + tvm_ffi_tensor = arg._tvm_ffi_tensor + dtype = tvm_ffi_tensor.dtype + tvm_ffi_cute_tensor = spec.Tensor( + arg_name, + shapes, + arg._tvm_ffi_tensor.dtype, + strides=strides, + data_alignment=arg._assumed_align, + device_type=tvm_ffi_tensor.device.type + ) + else: + # for FakeTensor, strictly follow the shape and stride from the cute tensor + device_type = "cuda" if _is_gpu_memspace(arg.memspace) else "cpu" tvm_ffi_cute_tensor = spec.Tensor( arg_name, shapes, NumericToTVMFFIDtype[arg.element_type], strides=strides, data_alignment=arg._assumed_align, + device_type=device_type, ) if arg.element_type == Float4E2M1FN: tvm_ffi_cute_tensor = spec.create_map_tensor_dtype_f4x2_to_f4_spec( tvm_ffi_cute_tensor ) - params.append(tvm_ffi_cute_tensor) - elif isinstance(arg, Pointer): - address_space = None - if hasattr(arg, "memspace"): - address_space = _get_llvm_address_space_from_memspace(arg.memspace) - params.append(spec.DataPointer(arg_name, address_space=address_space)) - elif isinstance(arg, _FakeStream): - if arg.use_tvm_ffi_env_stream: - params.append(spec.EnvStream(arg_name)) - else: - params.append(spec.Stream(arg_name)) - elif isinstance(arg, cuda.CUstream): - params.append(spec.Stream(arg_name)) + return tvm_ffi_cute_tensor + elif isinstance(arg, Pointer) or arg_type == Pointer: + address_space = None + if hasattr(arg, "memspace"): + address_space = _get_llvm_address_space_from_memspace(arg.memspace) + return spec.DataPointer(arg_name, address_space=address_space) + elif isinstance(arg, _FakeStream): + if arg.use_tvm_ffi_env_stream: + return spec.EnvStream(arg_name) else: - raise DSLRuntimeError(f"Unsupported argument type: {type(arg)}") - # The following code can obtain signature of the function - # that maybe useful for future debugging and usecases. - # signature = spec.signature(function_name, params) + return spec.Stream(arg_name) + elif isinstance(arg, cuda.CUstream): + return spec.Stream(arg_name) + elif arg_type is not None and get_origin(arg_type) is tuple: + # Handle Tuple[X, Y, ...] type annotations + tuple_element_types = get_args(arg_type) + if not isinstance(arg, (tuple, list)): + raise DSLRuntimeError(f"Expected tuple for argument {arg_name}, got {type(arg)}") + if len(arg) != len(tuple_element_types): + raise DSLRuntimeError( + f"Tuple length mismatch for argument {arg_name}: " + f"expected {len(tuple_element_types)}, got {len(arg)}" + ) + + # Recursively convert each tuple element + tuple_params = [] + for i, (elem, elem_type) in enumerate(zip(arg, tuple_element_types)): + elem_name = f"{arg_name}[{i}]" + elem_param = _convert_single_arg(elem, elem_name, elem_type, ctx) + tuple_params.append(elem_param) + + return spec.TupleParam(arg_name, tuple_params) + else: + raise DSLRuntimeError(f"Unsupported argument type: {type(arg)}") + + +def _tvm_ffi_args_spec_converter( + function_name: str, + args_spec: inspect.FullArgSpec, + full_args: List[Any], + full_kwargs: Dict[str, Any], +): + """Convert cute algebra args to tvm ffi spec params. + + This function converts the cute arguments specs to tvm ffi spec params. + """ + exec_args = ExecutionArgs(args_spec, function_name) + rectified_args = exec_args.get_rectified_args_from_original_args(full_args, full_kwargs) + arg_names = exec_args.args_spec.args + exec_args.args_spec.kwonlyargs + params = [] + ctx = ConverterContext() + + for arg, arg_name in zip(rectified_args, arg_names): + arg_type = args_spec.annotations.get(arg_name, None) + param = _convert_single_arg(arg, arg_name, arg_type, ctx) + params.append(param) + return params diff --git a/python/CuTeDSL/cutlass/cute/arch/mbar.py b/python/CuTeDSL/cutlass/cute/arch/mbar.py index dad7e380..33363733 100644 --- a/python/CuTeDSL/cutlass/cute/arch/mbar.py +++ b/python/CuTeDSL/cutlass/cute/arch/mbar.py @@ -209,6 +209,7 @@ def mbarrier_conditional_try_wait( def mbarrier_arrive( mbar_ptr: Pointer, peer_cta_rank_in_cluster: Optional[Int] = None, + arrive_count: Int = 1, *, loc=None, ip=None, @@ -239,7 +240,7 @@ def mbarrier_arrive( nvvm.mbarrier_txn( mbar_llvm_ptr, - Int32(1).ir_value(loc=loc, ip=ip), + Int32(arrive_count).ir_value(loc=loc, ip=ip), kind=nvvm.MBarrierTxnKind.ARRIVE, space=space, loc=loc, diff --git a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py index 2a116cb0..57cafc26 100644 --- a/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py +++ b/python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py @@ -201,13 +201,13 @@ class MmaOp(Tcgen05MmaOp): if self.cta_group == CtaGroup.ONE: if m not in [64, 128]: raise OpError(self, f"expects the M-mode to be 64 or 128, but got {m}") - if m == 64: - if (n < 8) or (n > 256) or (n % 8 != 0): + if self.b_dtype.width == 8 and self.b_major_mode == OperandMajorMode.MN: + if (n < 16) or (n > 256) or (n % 16 != 0): raise OpError( self, - f"expects the N-mode to satisfy 8 <= N <= 256 and N % 8 == 0, but got {n}", + f"expects the N-mode to satisfy 16 <= N <= 256 and N % 16 == 0, but got {n}", ) - elif m == 128: + else: if (n < 8) or (n > 256) or (n % 8 != 0): raise OpError( self, @@ -216,11 +216,18 @@ class MmaOp(Tcgen05MmaOp): else: if m not in [128, 256]: raise OpError(self, f"expects the M-mode to be 128 or 256, but got {m}") - if (n < 16) or (n > 256) or (n % 16 != 0): - raise OpError( - self, - f"expects the N-mode to satisfy 16 <= N <= 256 and N % 16 == 0, but got {n}", - ) + if self.b_dtype.width == 8 and self.b_major_mode == OperandMajorMode.MN: + if (n < 32) or (n > 256) or (n % 32 != 0): + raise OpError( + self, + f"expects the N-mode to satisfy 32 <= N <= 256 and N % 32 == 0, but got {n}", + ) + else: + if (n < 16) or (n > 256) or (n % 16 != 0): + raise OpError( + self, + f"expects the N-mode to satisfy 16 <= N <= 256 and N % 16 == 0, but got {n}", + ) def __str__(self) -> str: return ( @@ -302,6 +309,7 @@ class BlockScaledMmaOp(Tcgen05MmaOp): admissible_archs = [ Arch.sm_100a, + Arch.sm_103a, ] def __post_init__(self) -> None: diff --git a/python/CuTeDSL/cutlass/cute/runtime.py b/python/CuTeDSL/cutlass/cute/runtime.py index ab04f27e..19815b5c 100644 --- a/python/CuTeDSL/cutlass/cute/runtime.py +++ b/python/CuTeDSL/cutlass/cute/runtime.py @@ -10,7 +10,6 @@ # is strictly prohibited. import ctypes -import os import sys from pathlib import Path from functools import lru_cache @@ -20,6 +19,7 @@ from typing import Union, Optional, Type, List # MLIR modules imports from cutlass._mlir import ir +from cutlass.base_dsl.env_manager import get_prefix_dsl_libs import cutlass._mlir.dialects.cute as _cute_ir import cutlass._mlir.dialects.cuda as _cuda_dialect @@ -144,6 +144,7 @@ class _Tensor(Tensor): self._tvm_ffi_tensor = tvm_ffi.from_dlpack(tensor) self._dlpack_data = self._tvm_ffi_tensor.__dlpack__() + self._dltensor_wrapper = None self._assumed_align = assumed_align self._is_dynamic = False @@ -387,7 +388,16 @@ class _Tensor(Tensor): return CoreTensor(values[0].value, self._dtype) def __tvm_ffi_object__(self): - return self._tvm_ffi_tensor + try: + return self._tvm_ffi_tensor + except AttributeError: + raise DSLRuntimeError( + ( + "runtime._Tensor is not a TVM-FFI tensor. " + "Enable TVM-FFI with `from_dlpack(..., enable_tvm_ffi=True)` " + "or `CUTE_DSL_ENABLE_TVM_FFI=1`." + ) + ) def _get_cute_type_str(inp): @@ -411,7 +421,8 @@ class _FakeCompactTensor(Tensor): self._dtype = dtype self._shape = shape self._stride_order = stride_order or tuple(range(len(shape))) - self._memspace = memspace or AddressSpace.gmem + # cannot use memspace or AddressSpace.gmem because AddressSpace.generic is 0 + self._memspace = memspace if memspace is not None else AddressSpace.gmem self._assumed_align = assumed_align or -(-dtype.width // 8) self._use_32bit_stride = use_32bit_stride @@ -510,7 +521,8 @@ class _FakeTensor(Tensor): self._dtype = dtype self._shape = shape self._stride = stride - self._memspace = memspace or AddressSpace.generic + # cannot use memspace or AddressSpace.generic because AddressSpace.generic is 0 + self._memspace = memspace if memspace is not None else AddressSpace.gmem self._assumed_align = assumed_align if assumed_align is None: # use the bytes width of the element dtype. The alignment is at least one byte align. @@ -605,7 +617,7 @@ def make_fake_compact_tensor( :param shape: Shape of the tensor. :type shape: tuple[int, ...] :param stride_order: Order in which strides (memory layout) are assigned to the tensor dimensions. - If None, the default layout is row-major. Otherwise, it should be a permutation of the dimension indices. + If None, the default layout is col-major. Otherwise, it should be a permutation of the dimension indices. :type stride_order: tuple[int, ...], optional :param memspace: Memory space where the fake tensor resides. Optional. :type memspace: str, optional @@ -644,6 +656,7 @@ def make_fake_compact_tensor( use_32bit_stride=use_32bit_stride, ) + def make_fake_tensor(dtype, shape, stride, *, memspace=None, assumed_align=None): """ Create a fake tensor with the specified element type, shape, and stride. @@ -859,21 +872,22 @@ def find_runtime_libraries(*, enable_tvm_ffi: bool = True) -> List[str]: """ def _get_cuda_dialect_runtime_path(): - libs = os.environ.get("CUTE_DSL_LIBS") - if libs: - sep = ";" if sys.platform.startswith("win32") else ":" - for path in libs.split(sep): - if path.endswith("libcuda_dialect_runtime.so"): - return path - try: - # find package library from wheel package - pkg_base = Path(__file__).resolve().parent.parent - lib_path = pkg_base / "lib" / "libcuda_dialect_runtime.so" - if lib_path.is_file(): - return str(lib_path) - except OSError: + libs = get_prefix_dsl_libs("CUTE_DSL") + if libs is None: return None + # check if the separator is ; for windows + if sys.platform.startswith("win32") and ";" in libs: + libs = libs.split(";") + else: + libs = libs.split(":") + + for path in libs: + if path.endswith("libcuda_dialect_runtime.so"): + return path + + return None + libs = [] cuda_dialect_runtime_path = _get_cuda_dialect_runtime_path() if cuda_dialect_runtime_path: diff --git a/python/CuTeDSL/cutlass/cute/testing.py b/python/CuTeDSL/cutlass/cute/testing.py index 6f3d7e17..766dff88 100644 --- a/python/CuTeDSL/cutlass/cute/testing.py +++ b/python/CuTeDSL/cutlass/cute/testing.py @@ -20,6 +20,7 @@ from typing import Type, Union, Callable, Optional, Dict, List, Any import cuda.bindings.driver as cuda_driver import cuda.bindings.runtime as cuda_runtime +import cutlass import cutlass.base_dsl.jit_executor from cutlass.cutlass_dsl import Constexpr, CuTeDSL, T, dsl_user_op @@ -233,6 +234,7 @@ def convert(src: cute.Tensor, dst: cute.Tensor): src.shape[leading_mode] % elem_per_copy == 0 and dst.shape[leading_mode] % elem_per_copy == 0 ) + _convert(src, dst, leading_mode, elem_per_copy) diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py b/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py index 10dc52d0..806992d3 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/__init__.py @@ -52,6 +52,7 @@ from ..base_dsl.compiler import ( KeepPTX, GPUArch, LinkLibraries, + EnableTVMFFI, ) from ..base_dsl.runtime.jit_arg_adapters import * diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py b/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py index 08321ee5..24405c86 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cuda_jit_executor.py @@ -21,7 +21,12 @@ import cuda.bindings.runtime as cuda_runtime import cuda.bindings.driver as cuda_driver # Local modules imports -from ..base_dsl.jit_executor import JitExecutor, ExecutionArgs, JitFunctionArtifacts +from ..base_dsl.jit_executor import ( + JitExecutor, + JitCompiledFunction, + ExecutionArgs, + JitFunctionArtifacts, +) from ..base_dsl.utils.logger import log from ..base_dsl.common import DSLCudaRuntimeError, DSLRuntimeError from ..base_dsl.typing import Int32 @@ -57,7 +62,7 @@ class CudaDialectJitModule: self.unload() -class CudaDialectJitCompiledFunction: +class CudaDialectJitCompiledFunction(JitCompiledFunction): """Holds a compiled function and its module.""" def __init__( @@ -103,21 +108,6 @@ class CudaDialectJitCompiledFunction: self._executor_lock = threading.RLock() self._default_executor = None - @property - def __ptx__(self): - """Returns the PTX code of the JIT-compiled function.""" - return self.artifacts.PTX if self.artifacts is not None else None - - @property - def __cubin__(self): - """Returns the CUBIN data of the JIT-compiled function.""" - return self.artifacts.CUBIN if self.artifacts is not None else None - - @property - def __mlir__(self): - """Returns the MLIR code of the JIT-compiled function.""" - return self.artifacts.MLIR if self.artifacts is not None else None - @functools.cached_property def num_devices(self): """Returns the number of CUDA devices available.""" @@ -285,6 +275,7 @@ class CudaDialectJitCompiledFunction: :return: A callable executor function. :rtype: JitExecutor """ + super()._validate_engine() with self._executor_lock: # We need to ensure that the modules are loaded if not already if self.jit_module is None or self.jit_module.is_unloaded(): @@ -297,37 +288,3 @@ class CudaDialectJitCompiledFunction: ) return JitExecutor(self.jit_module, None, self.jit_time_profiling) - - def generate_execution_args(self, *args, **kwargs): - return self.args_spec.generate_execution_args(args, kwargs) - - def set_dynamic_args(self, dynamic_args, dynamic_kwargs): - """Sets the dynamic argument information required for export to c code generation.""" - self.dynamic_args = dynamic_args - self.dynamic_kwargs = dynamic_kwargs - - def __call__(self, *args, **kwargs): - """Executes the jit-compiled function under the currently active CUDA context. - - Calling this method multiple devices is not allowed and will result in unexpected - CUDA errors. If you need to call the kernel on multiple devices use `to` - to return a per-device function. - """ - exe_args, adapted_args = self.generate_execution_args(*args, **kwargs) - return self.run_compiled_program(exe_args) - - def run_compiled_program(self, exe_args): - """Executes the jit-compiled function under the currently active CUDA context. - - Calling this method multiple devices is not allowed and will result in unexpected - CUDA errors. If you need to call the kernel on multiple devices use `to` - to return a per-device function. - """ - with self._executor_lock: - if self._default_executor is None: - log().debug("Creating default executor.") - # We use a weak reference here so that this instance does not keep this - # object alive as it hold a reference to self. - proxy_self = weakref.proxy(self) - self._default_executor = proxy_self.to(None) - return self._default_executor.run_compiled_program(exe_args) diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cuda_stream_adapter.py b/python/CuTeDSL/cutlass/cutlass_dsl/cuda_stream_adapter.py index 05d3d6b4..2a1ac5cb 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/cuda_stream_adapter.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cuda_stream_adapter.py @@ -42,3 +42,7 @@ class CudaDialectStreamAdapter: def __get_mlir_types__(self): return [cuda.StreamType.get()] + + def __cuda_stream__(self): + # support cuda stream protocol + return (0, int(self._arg)) diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py index cf8162e6..477b3152 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py @@ -32,6 +32,7 @@ import pkgutil from dataclasses import is_dataclass, fields from math import ceil from itertools import chain +from pathlib import Path from collections.abc import Sequence import builtins import ctypes @@ -334,8 +335,15 @@ class CutlassBaseDSL(BaseDSL): ) try: # update the version hash of the cutlass shared library + giant_dso_name = str( + next( + (Path(dsl_path) / "_mlir" / "_mlir_libs").glob( + "_cutlass_ir.cpython*" + ) + ).name + ) with open( - os.path.join(dsl_path, "_mlir/_mlir_libs/libCutlassIRPythonCAPI.so"), + os.path.join(dsl_path, f"_mlir/_mlir_libs/{giant_dso_name}"), "rb", ) as f: while True: @@ -345,7 +353,7 @@ class CutlassBaseDSL(BaseDSL): version_hash.update(chunk) except Exception: raise DSLRuntimeError( - "Failed to read the shared library file libCutlassIRPythonCAPI.so." + f"Failed to read the shared library file {giant_dso_name}." "The file may not exist or may not be readable." "Please re-install the package." ) @@ -386,6 +394,8 @@ class CutlassBaseDSL(BaseDSL): args_spec, no_cache, *, + full_args=None, + full_kwargs=None, dynamic_args=None, dynamic_kwargs=None, original_function_name=None, @@ -399,6 +409,8 @@ class CutlassBaseDSL(BaseDSL): :param pipeline: The pipeline to use for compilation. :param args_spec: The args spec to use for compilation. :param no_cache: Whether to cache the result. + :param full_args: The full arguments to use for compilation. + :param full_kwargs: The full keyword arguments to use for compilation. :param dynamic_args: The dynamic arguments to use for compilation. :param dynamic_kwargs: The dynamic keyword arguments to use for compilation. :param original_function_name: The name of the original function without mangling. @@ -415,7 +427,7 @@ class CutlassBaseDSL(BaseDSL): assert self._tvm_ffi_args_spec_converter is not None tvm_ffi_spec_params = self._tvm_ffi_args_spec_converter( - function_name, args_spec, dynamic_args, dynamic_kwargs + function_name, args_spec, full_args, full_kwargs ) tvm_ffi_provider = TVMFFICuteCallProvider(function_name) @@ -445,6 +457,8 @@ class CutlassBaseDSL(BaseDSL): args_spec, no_cache, TVMFFIJitCompiledFunction, + full_args=full_args, + full_kwargs=full_kwargs, dynamic_args=dynamic_args, dynamic_kwargs=dynamic_kwargs, ) @@ -457,6 +471,8 @@ class CutlassBaseDSL(BaseDSL): args_spec, no_cache, CudaDialectJitCompiledFunction, + full_args=full_args, + full_kwargs=full_kwargs, dynamic_args=dynamic_args, dynamic_kwargs=dynamic_kwargs, original_function_name=original_function_name, diff --git a/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py b/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py index 0a430543..4636615c 100644 --- a/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py +++ b/python/CuTeDSL/cutlass/cutlass_dsl/tvm_ffi_provider.py @@ -17,18 +17,24 @@ from cutlass.base_dsl.tvm_ffi_builder import ( ) from cutlass._mlir import ir from cutlass._mlir.dialects import llvm -from cutlass._mlir._mlir_libs import _execution_engine_extra +from cutlass._mlir._mlir_libs._cutlass_ir import _execution_engine_extra from cutlass.cutlass_dsl.cuda_jit_executor import CudaDialectJitCompiledFunction from cutlass.base_dsl.common import DSLRuntimeError +from typing import Optional import tvm_ffi class TVMFFICuteCallProvider(DynamicParamPackCallProvider): """Cute call provider that uses cute call convention.""" + cuda_device_index: Optional[ir.Value] + cuda_error_handle_block: Optional[ir.Block] + def __init__(self, target_func: str): super().__init__(target_func, struct_call=True) self.cuda_global_state_symbol = f"__{target_func}_cuda_state" + self.cuda_device_index = None + self.cuda_error_handle_block = None def get_callee_struct_for_param_tensor( self, @@ -41,7 +47,10 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): ) -> ir.Type: """Routine used to override the tensor passing struct convention""" with ir.InsertionPoint(current_block): - data_type = self.gpu_ptr_type + if param.dlpack_device_type == tvm_ffi.DLDeviceType.kDLCPU: + data_type = self.ptr_type + else: + data_type = self.gpu_ptr_type strides_type = ( self.struct_type(fields=[x.type for x in strides]) if len(strides) != 1 @@ -79,22 +88,27 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): def declare_extern_funcs(self, current_block: ir.Block, context: CallContext): """Append the error handling function to the current block.""" with ir.InsertionPoint(context.module.body): - self.declare_extern_func( + context.builder.find_or_declare_extern_func( "cuda_dialect_get_error_name", [self.i32_type], self.ptr_type, ) - self.declare_extern_func( + context.builder.find_or_declare_extern_func( + "_cudaGetDevice", + [self.ptr_type], + self.i32_type, + ) + context.builder.find_or_declare_extern_func( "_cudaSetDevice", [self.i32_type], self.i32_type, ) - self.declare_extern_func( + context.builder.find_or_declare_extern_func( "cuda_dialect_init_library_once", [self.ptr_type, self.ptr_type, self.ptr_type, self.ptr_type], self.i32_type, ) - self.declare_extern_func( + context.builder.find_or_declare_extern_func( "cuda_dialect_unload_library_once", [self.ptr_type], self.void_type, @@ -116,7 +130,7 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): self.cuda_global_state_symbol, self.ptr_type ) cuda_init_ptr = self.address_of("cuda_init", self.ptr_type) - cuda_load_ptr = self.address_of("cuda_load", self.ptr_type) + cuda_load_to_device_ptr = self.address_of("cuda_load_to_device", self.ptr_type) set_error_ptr = self.address_of( "TVMFFIErrorSetRaisedFromCStr", self.ptr_type ) @@ -127,7 +141,7 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): callee_operands=[ cuda_global_state_ptr, cuda_init_ptr, - cuda_load_ptr, + cuda_load_to_device_ptr, set_error_ptr, ], op_bundle_sizes=[], @@ -207,34 +221,70 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): def check_cuda_error( self, code: ir.Value, current_block: ir.Block, context: CallContext ): - """Check if the CUDA error is raised and return the error string if so.""" + """Check if the CUDA error is raised and return the error string if so. + + Uses a shared error handling block to avoid code duplication. The error code + is passed as a block argument to the shared error handler. + """ + assert self.cuda_error_handle_block is not None with ir.InsertionPoint(current_block): # check if the call is successful - error_block = current_block.create_after() - success_block = error_block.create_after() - # Check if call is successful (non-zero return code) + success_block = current_block.create_after() + # Check if call is successful (zero return code means success) self.cond_br( cond=self.equal(code, self.i32(0)), true_block=success_block, - false_block=error_block, + false_block=self.cuda_error_handle_block, branch_weights=self.BRANCH_WEIGHTS_LIKELY, + false_dest_operands=[code], # Pass error code to shared error block + ) + return success_block + + def set_cuda_device_if_mismatch( + self, + current_block: ir.Block, + context: CallContext, + current_device: Optional[ir.Value], + target_device: Optional[ir.Value], + ) -> ir.Block: + """Set the CUDA device index if it differs from the target device. + """ + # If either device is None, no switching needed + if current_device is None: + assert target_device is None + return current_block + + with ir.InsertionPoint(current_block): + # Check if devices are different + devices_differ = self.not_equal(current_device, target_device) + + # Create blocks for conditional device switching + switch_device_block = current_block.create_after() + continuation_block = switch_device_block.create_after() + # For this specific case, avoid branch weights for now + # mainly to avoid too drastic reordering of the code + self.cond_br( + cond=devices_differ, + true_block=switch_device_block, + false_block=continuation_block ) - # Error block: raise error and return - with ir.InsertionPoint(error_block): - error_str = llvm.call( - result=self.ptr_type, - callee="cuda_dialect_get_error_name", - callee_operands=[code], + # Switch device block: call cudaSetDevice + with ir.InsertionPoint(switch_device_block): + result = llvm.call( + result=self.i32_type, + callee="_cudaSetDevice", + callee_operands=[target_device], op_bundle_sizes=[], op_bundle_operands=[], ) - # Raise error and return -1 - context.builder.raise_error_and_return( - error_kind="RuntimeError", - error_message_parts=["CUDA Error: ", error_str], - ) - return success_block + + # Check for errors and branch to continuation + switch_device_block = self.check_cuda_error(result, switch_device_block, context) + with ir.InsertionPoint(switch_device_block): + self.br(continuation_block) + + return continuation_block def generate_llvm_call( self, @@ -243,6 +293,36 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): context: CallContext, ) -> ir.Block: """Generate the LLVM call operation and check if the call is successful.""" + old_cuda_device_index: Optional[ir.Value] = None + + # If we need to manage CUDA device context + if self.cuda_device_index is not None: + # Create an alloca in the entry block to store the current device index + device_index_alloca = context.builder.create_alloca( + context.entry_block, self.i32_type, array_size=1 + ) + + # Get the current device + with ir.InsertionPoint(current_block): + get_device_result = llvm.call( + result=self.i32_type, + callee="_cudaGetDevice", + callee_operands=[device_index_alloca], + op_bundle_sizes=[], + op_bundle_operands=[], + ) + current_block = self.check_cuda_error(get_device_result, current_block, context) + + # Load the current device index from the alloca + with ir.InsertionPoint(current_block): + old_cuda_device_index = llvm.load(self.i32_type, device_index_alloca) + + # Switch to target device if different + current_block = self.set_cuda_device_if_mismatch( + current_block, context, old_cuda_device_index, self.cuda_device_index + ) + + # Execute the main call with ir.InsertionPoint(current_block): result = llvm.call( result=self.i32_type, @@ -251,42 +331,72 @@ class TVMFFICuteCallProvider(DynamicParamPackCallProvider): op_bundle_sizes=[], op_bundle_operands=[], ) - return self.check_cuda_error(result, current_block, context) - def insert_set_cuda_device(self, current_block: ir.Block, context: CallContext): - """Call the _cudaSetDevice function if we can find device id from tensor parameters.""" + # Restore the original device BEFORE checking for errors + # This ensures device is restored even if the main call failed + current_block = self.set_cuda_device_if_mismatch( + current_block, context, self.cuda_device_index, old_cuda_device_index + ) - def find_cuda_device_index_from_params(): - for param in context.params: - if ( - isinstance(param, spec.Tensor) - and param.dlpack_device_type != tvm_ffi.DLDeviceType.kDLCPU - ): - return context.matched_var_binding[param.device_id] - return None + # Now check for errors from the main call + current_block = self.check_cuda_error(result, current_block, context) - cuda_device_index = find_cuda_device_index_from_params() + return current_block - if cuda_device_index is None: - return current_block - with ir.InsertionPoint(current_block): - result = llvm.call( - result=self.i32_type, - callee="_cudaSetDevice", - callee_operands=[cuda_device_index], + def find_cuda_device_index_from_params(self, context: CallContext): + """Find the CUDA device index from tensor parameters.""" + for param in context.params: + if ( + isinstance(param, spec.Tensor) + and param.dlpack_device_type != tvm_ffi.DLDeviceType.kDLCPU + ): + return context.matched_var_binding[param.device_id] + return None + + def create_shared_cuda_error_block( + self, + current_block: ir.Block, + context: CallContext + ) -> ir.Block: + """Create a shared error handling block for all CUDA errors. + """ + # Create the shared error block after the current block (setup phase) + # This block will be branched to from multiple error checking sites + # It accepts the error code as a block argument + error_block = current_block.create_after() + error_code = error_block.add_argument(self.i32_type, ir.Location.unknown()) + + # Populate the error block + with ir.InsertionPoint(error_block): + error_str = llvm.call( + result=self.ptr_type, + callee="cuda_dialect_get_error_name", + callee_operands=[error_code], op_bundle_sizes=[], op_bundle_operands=[], ) + # Raise error and return -1 + context.builder.raise_error_and_return( + error_kind="RuntimeError", + error_message_parts=["CUDA Error: ", error_str], + ) - return self.check_cuda_error(result, current_block, context) + return error_block def __call__(self, current_block: ir.Block, context: CallContext) -> ir.Block: current_block = self.declare_extern_funcs(current_block, context) current_block = self.insert_lazy_init_cuda(current_block, context) current_block = self.append_unload_to_global_dtors(current_block, context) - current_block = self.insert_set_cuda_device(current_block, context) + # Create shared CUDA error handling block after the setup blocks + # This reduces code duplication - all CUDA errors branch to this single block + self.cuda_error_handle_block = self.create_shared_cuda_error_block(current_block, context) + # setup device index, will be set around the call to the target function + self.cuda_device_index = self.find_cuda_device_index_from_params(context) current_block = super().__call__(current_block, context) + self.cuda_device_index = None + self.cuda_error_handle_block = None + # reset the device index and error block return current_block @@ -316,12 +426,15 @@ class TVMFFIJitCompiledFunction(tvm_ffi.Function, CudaDialectJitCompiledFunction if self.__chandle__() != 0: raise DSLRuntimeError("TVM FFI function is already initialized") # get the MLIR function pointer from the execution engine - tvm_ffi_function_ptr = self.engine.raw_lookup("__tvm_ffi_" + self.function_name) - tvm_ffi_function = tvm_ffi.Function.__from_mlir_packed_safe_call__( - tvm_ffi_function_ptr - ) - # move the handle from the tvm_ffi.Function to the current instance - self.__move_handle_from__(tvm_ffi_function) + if self.engine is not None: + tvm_ffi_function_ptr = self.engine.raw_lookup( + "__tvm_ffi_" + self.function_name + ) + tvm_ffi_function = tvm_ffi.Function.__from_mlir_packed_safe_call__( + tvm_ffi_function_ptr + ) + # move the handle from the tvm_ffi.Function to the current instance + self.__move_handle_from__(tvm_ffi_function) def to(self, device=None): """TVM FFI function itself is already support all devices.""" diff --git a/python/CuTeDSL/cutlass/torch.py b/python/CuTeDSL/cutlass/torch.py index ac3fe2e2..bbc8f378 100644 --- a/python/CuTeDSL/cutlass/torch.py +++ b/python/CuTeDSL/cutlass/torch.py @@ -165,6 +165,16 @@ def create_and_permute_torch_tensor( return dtype_torch_tensor +def get_leading_dim(torch_tensor: torch.Tensor) -> int: + """ + Get the leading dimension of a torch tensor + """ + for i, stride in enumerate(torch_tensor.stride()): + if stride == 1: + return i + return None + + def convert_cute_tensor( f32_torch_tensor: "torch.Tensor", cute_tensor: Tensor, @@ -189,8 +199,10 @@ def convert_cute_tensor( }: fp32_cute_tensor = from_dlpack(f32_torch_tensor) if is_dynamic_layout: + # note: dim_order to not always maps to leading dimension, + # so we need to get the leading dimension from the torch tensor strides fp32_cute_tensor = fp32_cute_tensor.mark_layout_dynamic( - f32_torch_tensor.dim_order()[-1] + leading_dim=get_leading_dim(f32_torch_tensor) ) # Copy and convert from f32 cute tensor to dtype cute tensor cute.testing.convert(fp32_cute_tensor, cute_tensor) @@ -297,11 +309,9 @@ def cute_tensor_like( # create cute tensor using the device buffer cute_tensor = from_dlpack(torch_tensor, assumed_align=assumed_align) cute_tensor.element_type = cutlass_dtype + if is_dynamic_layout: - for i, stride in enumerate(torch_tensor.stride()): - if stride == 1: - leading_dim = i - break + leading_dim = get_leading_dim(torch_tensor) cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim) # initialize the cute tensor data @@ -316,5 +326,4 @@ def cute_tensor_like( ) else: torch_tensor.copy_(data_ref.to(dtype=torch_dtype)) - return cute_tensor, torch_tensor diff --git a/python/CuTeDSL/cutlass/utils/blackwell_helpers.py b/python/CuTeDSL/cutlass/utils/blackwell_helpers.py index 3c5c97b1..17b002e5 100644 --- a/python/CuTeDSL/cutlass/utils/blackwell_helpers.py +++ b/python/CuTeDSL/cutlass/utils/blackwell_helpers.py @@ -664,6 +664,7 @@ def make_smem_layout_a( a_dtype: Type[Numeric], num_stages: int, *, + is_k_major=None, loc=None, ip=None, ) -> Union[cute.Layout, cute.ComposedLayout]: @@ -687,7 +688,8 @@ def make_smem_layout_a( :rtype: Union[cute.Layout, cute.ComposedLayout] """ - is_k_major = tiled_mma.op.a_major_mode == OperandMajorMode.K + is_k_major = (tiled_mma.op.a_major_mode == OperandMajorMode.K) if is_k_major is None else is_k_major + a_major_mode = OperandMajorMode.K if is_k_major else OperandMajorMode.MN a_smem_shape = tiled_mma.partition_shape_A( cute.dice(mma_tiler_mnk, (1, None, 1), loc=loc, ip=ip), loc=loc, ip=ip ) @@ -696,7 +698,7 @@ def make_smem_layout_a( cute.size(a_smem_shape[0][1], loc=loc, ip=ip) * a_smem_shape[2], ) smem_layout_atom_kind = get_smem_layout_atom_ab( - tiled_mma.op.a_major_mode, a_dtype, a_smem_shape_mn_k, loc=loc, ip=ip + a_major_mode, a_dtype, a_smem_shape_mn_k, loc=loc, ip=ip ) a_smem_layout_atom = make_smem_layout_atom( smem_layout_atom_kind, a_dtype, loc=loc, ip=ip @@ -716,6 +718,7 @@ def make_smem_layout_b( b_dtype: Type[Numeric], num_stages: int, *, + is_k_major=None, loc=None, ip=None, ) -> Union[cute.Layout, cute.ComposedLayout]: @@ -739,7 +742,8 @@ def make_smem_layout_b( :rtype: Union[cute.Layout, cute.ComposedLayout] """ - is_k_major = tiled_mma.op.b_major_mode == OperandMajorMode.K + is_k_major = (tiled_mma.op.b_major_mode == OperandMajorMode.K) if is_k_major is None else is_k_major + b_major_mode = OperandMajorMode.K if is_k_major else OperandMajorMode.MN b_smem_shape = tiled_mma.partition_shape_B( cute.dice(mma_tiler_mnk, (None, 1, 1), loc=loc, ip=ip), loc=loc, ip=ip ) @@ -749,7 +753,7 @@ def make_smem_layout_b( ) smem_layout_atom_kind = get_smem_layout_atom_ab( - tiled_mma.op.b_major_mode, b_dtype, b_smem_shape_nk, loc=loc, ip=ip + b_major_mode, b_dtype, b_smem_shape_nk, loc=loc, ip=ip ) b_smem_layout_atom = make_smem_layout_atom( smem_layout_atom_kind, b_dtype, loc=loc, ip=ip diff --git a/python/CuTeDSL/cutlass/utils/hardware_info.py b/python/CuTeDSL/cutlass/utils/hardware_info.py index a80d2471..2605d0bb 100644 --- a/python/CuTeDSL/cutlass/utils/hardware_info.py +++ b/python/CuTeDSL/cutlass/utils/hardware_info.py @@ -8,12 +8,10 @@ # Any use, reproduction, disclosure, or distribution of this software # and related documentation outside the scope permitted by the EULA # is strictly prohibited. - -from cuda.bindings import driver, nvrtc, runtime -from cutlass.cutlass_dsl.cuda_jit_executor import CudaDialectJitModule +from cuda.bindings import driver, runtime from cutlass.base_dsl.common import DSLRuntimeError - -import cutlass.cute as cute +from cutlass import cute +import tempfile """ This class is used to get the hardware info of given GPU device. @@ -53,8 +51,8 @@ class HardwareInfo: f"Cluster size must be between 1 and 32, {cluster_size} is not supported" ) - self._get_device_function(self.device) - + # must do get kernel after set device so runtime context is set correctly + self.kernel = self._get_device_function() max_shared_memory_per_block = self._checkCudaErrors( driver.cuDeviceGetAttribute( driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, @@ -152,8 +150,6 @@ class HardwareInfo: if isinstance(error, driver.CUresult): err, name = driver.cuGetErrorName(error) return name if err == driver.CUresult.CUDA_SUCCESS else "" - elif isinstance(error, nvrtc.nvrtcResult): - return nvrtc.nvrtcGetErrorString(error)[1] else: raise RuntimeError("Unknown error type: {}".format(error)) @@ -175,14 +171,20 @@ class HardwareInfo: ) # get a empty kernel to compute occupancy - def _get_device_function(self, device) -> driver.CUfunction: - self.compiled_kernel = cute.compile(self._host_function).to(device) - assert isinstance(self.compiled_kernel.jit_module, CudaDialectJitModule) - err, kernels = runtime.cudaLibraryEnumerateKernels( - 1, self.compiled_kernel.jit_module.cuda_library[0] - ) - if err is not runtime.cudaError_t.cudaSuccess: - raise DSLRuntimeError(f"Failed to enumerate kernels: {err}") - self.kernel = kernels[0] - self.kernel = self._checkCudaErrors(driver.cuKernelGetFunction(self.kernel)) - return self.kernel + def _get_device_function(self) -> driver.CUfunction: + """ + Get a device function by compiling a dummy kernel using cuteDSL pipeline. + """ + # Create a temporary directory for dumping artifacts + with tempfile.TemporaryDirectory() as temp_dir: + # keep-cubin will keep the cubin in the artifacts + compiled_func = cute.compile(self._host_function, options=f"--dump-dir={temp_dir} --keep-cubin") + # Get the CUBIN from artifacts + cubin_data = compiled_func.artifacts.CUBIN + cuda_library = self._checkCudaErrors( + driver.cuLibraryLoadData(cubin_data, None, None, 0, None, None, 0) + ) + # Enumerate kernels from the library + kernels = self._checkCudaErrors(driver.cuLibraryEnumerateKernels(1, cuda_library)) + # Get the function from the kernel + return self._checkCudaErrors(driver.cuKernelGetFunction(kernels[0])) diff --git a/python/cutlass_cppgen/__init__.py b/python/cutlass_cppgen/__init__.py index 2491da63..9bdd259c 100644 --- a/python/cutlass_cppgen/__init__.py +++ b/python/cutlass_cppgen/__init__.py @@ -133,7 +133,7 @@ def get_option_registry(): this._option_registry = OptionRegistry(device_cc()) return this._option_registry -this.__version__ = '4.3.0' +this.__version__ = '4.2.1' from cutlass_cppgen.backend import create_memory_pool from cutlass_cppgen.emit.pytorch import pytorch diff --git a/python/setup_cutlass.py b/python/setup_cutlass.py index 8b53d8f4..acc0c46e 100644 --- a/python/setup_cutlass.py +++ b/python/setup_cutlass.py @@ -51,7 +51,7 @@ setup_pycute.perform_setup() setup( name='cutlass_cppgen', - version='4.3.0', + version='4.2.0', description='CUTLASS Pythonic Interface', package_dir={'': '.'}, packages=[ diff --git a/python/setup_pycute.py b/python/setup_pycute.py index 37ba01cb..0bad050f 100644 --- a/python/setup_pycute.py +++ b/python/setup_pycute.py @@ -36,7 +36,7 @@ from setuptools import setup def perform_setup(): setup( name='pycute', - version='4.3.0', + version='4.2.1', description='Python implementation of CuTe', packages=['pycute'], )