import functools from torch._inductor.ir import get_stride_order from torch._inductor.runtime.cutedsl_cache import disk_cache_get, disk_cache_set from torch._inductor.runtime.runtime_utils import ceildiv from cutlass.utils import TensorMapUpdateMode {{gen_defines()}} # ---- Import GroupedGemm implementation, copied on PyTorch build from Cutlass repository: cutlass/examples/python/CuTeDSL/blackwell/grouped_gemm.py ---- from torch._inductor.kernel.vendored_templates.cutedsl.kernels.cutedsl_grouped_gemm import ( GroupedGemmKernel, ) # Caching is split into two levels: # # 1. _prep_fn_cache — caches the compiled executor for # build_group_ptrs_from_bases(). Depends only on tensor shapes, strides, # and dtypes of A/B/C, so it can be reused across different group # partitionings (`offs`). # # 2. _gemm_fn_cache — caches the compiled Grouped GEMM executor. Its key # extends the prep key with (max_active_clusters, total_num_clusters) # because different `offs` tensors change per-group problem sizes and # thus total_num_clusters, which alters the grid shape and persistent # scheduler configuration. Kernels compiled for one grid cannot be # safely reused for another. # # Compile-time constexprs (TILE_M, TILE_N, CLUSTER_M/N, USE_2_CTA, etc.) # are baked into each generated file, so they appear only in # _KERNEL_CONFIG_KEY (used by the disk cache) and not in the per-call # runtime keys. # # The @lru_cache on get_hardware_info() avoids redundant MLIR # recompilation overhead from hw.get_max_active_clusters(), which depends # only on GPU type. _prep_fn_cache = {} _gemm_fn_cache = {} _KERNEL_CONFIG_KEY = ( TILE_M, TILE_N, CLUSTER_M, CLUSTER_N, int(USE_2_CTA), str(ACC_DTYPE), str(TENSORMAP_UPDATE_MODE), ) def _to_cute_tensor(t, assumed_align=16): return from_dlpack(t.detach(), assumed_align=assumed_align, enable_tvm_ffi=True).mark_layout_dynamic() _TORCH_TO_CUTLASS_DTYPE = { torch.float16: cutlass.Float16, torch.bfloat16: cutlass.BFloat16, torch.float32: cutlass.Float32, torch.float64: cutlass.Float64, torch.int8: cutlass.Int8, torch.int32: cutlass.Int32, torch.int64: cutlass.Int64, } def _to_fake_cute_tensor(t, assumed_align=16): dyn_shape = tuple(map(int, t.shape)) stride_order = tuple(get_stride_order(t.stride())) return cute.runtime.make_fake_compact_tensor( _TORCH_TO_CUTLASS_DTYPE[t.dtype], dyn_shape, stride_order=stride_order, assumed_align=assumed_align, ) @functools.lru_cache def get_hardware_info(): hw = cutlass.utils.HardwareInfo() sm_count = hw.get_max_active_clusters(1) max_active_clusters = hw.get_max_active_clusters(CLUSTER_M * CLUSTER_N) return (sm_count, max_active_clusters) def get_prep_cache_key(input_a, input_b, output): return ( tuple(input_a.shape), tuple(input_a.stride()), input_a.dtype, tuple(input_b.shape), tuple(input_b.stride()), input_b.dtype, tuple(output.shape), tuple(output.stride()), output.dtype, ) def get_gemm_cache_key(prep_cache_key, max_active_clusters, total_num_clusters): return ( prep_cache_key, max_active_clusters, total_num_clusters, ) def compute_total_num_clusters(problem_sizes_mnkl, cluster_tile_shape_mn): total_num_clusters = 0 for m, n, _, _ in problem_sizes_mnkl: num_clusters_mn = tuple( ceildiv(x, y) for x, y in zip((m, n), cluster_tile_shape_mn) ) total_num_clusters += functools.reduce(lambda x, y: x * y, num_clusters_mn) return total_num_clusters def compute_cluster_tile_shape(mma_tiler_mn, cluster_shape_mn, use_2cta_instrs): cta_tile_shape_mn = list(mma_tiler_mn) if use_2cta_instrs: cta_tile_shape_mn[0] = cta_tile_shape_mn[0] // 2 return tuple(x * y for x, y in zip(cta_tile_shape_mn, cluster_shape_mn)) @cute.kernel def build_group_ptrs_from_bases_kernel( base_A_u64: cutlass.Int64, # device addr of input_a (bytes) base_B_u64: cutlass.Int64, # device addr of input_b (bytes) base_C_u64: cutlass.Int64, # device addr of Output (bytes) offs: cute.Tensor, # [G], cutlass.Int32/64 cumulative K: cutlass.Constexpr, N: cutlass.Constexpr, sizeof_element: cutlass.Int32, # bytes # -------- STRIDES (in ELEMENTS) -------- stride_A_m_elems: cutlass.Constexpr, # A.stride(0) stride_A_k_elems: cutlass.Constexpr, # A.stride(1) stride_B0_elems: cutlass.Constexpr, # B.stride(0) stride_Bk_elems: cutlass.Constexpr, # B.stride(1) stride_Bn_elems: cutlass.Constexpr, # B.stride(2) stride_C_m_elems: cutlass.Constexpr, # C.stride(0) stride_C_n_elems: cutlass.Constexpr, # C.stride(1) # -------- OUTPUTS -------- out_ptrs: cute.Tensor, # [G,3] cutlass.Int64: (A_ptr, B_ptr, C_ptr) out_problem: cute.Tensor, # [G,4] cutlass.Int32: (m_g, n, k, 1) out_strides_abc: cute.Tensor, # [G,3,2] cutlass.Int32 [[A_m,A_k],[B_n,B_k],[C_m,C_n]] ): tidx, _, _ = cute.arch.thread_idx() g = tidx m_beg_i32 = 0 if g > 0: m_beg_i32 = offs[g - 1] m_end_i32 = offs[g] m_g_i32 = m_end_i32 - m_beg_i32 a_byte_off = ( cutlass.Int64(m_beg_i32) * stride_A_m_elems * cutlass.Int64(sizeof_element) ) c_byte_off = ( cutlass.Int64(m_beg_i32) * stride_C_m_elems * cutlass.Int64(sizeof_element) ) b_byte_off = cutlass.Int64(g) * stride_B0_elems * cutlass.Int64(sizeof_element) # ---- pointers ---- out_ptrs[g, 0] = base_A_u64 + a_byte_off out_ptrs[g, 1] = base_B_u64 + b_byte_off out_ptrs[g, 2] = base_C_u64 + c_byte_off # ---- (m, n, k, 1) ---- out_problem[g, 0] = m_g_i32 out_problem[g, 1] = N out_problem[g, 2] = K out_problem[g, 3] = cutlass.Int32(1) # ---- strides ---- out_strides_abc[g, 0, 0] = cutlass.Int32(stride_A_m_elems) out_strides_abc[g, 0, 1] = cutlass.Int32(stride_A_k_elems) out_strides_abc[g, 1, 0] = cutlass.Int32(stride_Bn_elems) out_strides_abc[g, 1, 1] = cutlass.Int32(stride_Bk_elems) out_strides_abc[g, 2, 0] = cutlass.Int32(stride_C_m_elems) out_strides_abc[g, 2, 1] = cutlass.Int32(stride_C_n_elems) @cute.jit def launch_build_group_ptrs_from_bases( base_A_u64: cutlass.Int64, base_B_u64: cutlass.Int64, base_C_u64: cutlass.Int64, offs: cute.Tensor, G: cutlass.Constexpr, K: cutlass.Constexpr, N: cutlass.Constexpr, sizeof_element: cutlass.Constexpr, stride_A_m_elems: cutlass.Constexpr, stride_A_k_elems: cutlass.Constexpr, stride_B0_elems: cutlass.Constexpr, stride_Bk_elems: cutlass.Constexpr, stride_Bn_elems: cutlass.Constexpr, stride_C_m_elems: cutlass.Constexpr, stride_C_n_elems: cutlass.Constexpr, out_ptrs: cute.Tensor, # [G,3] cutlass.Int64 out_problem: cute.Tensor, # [G,4] cutlass.Int32 out_strides_abc: cute.Tensor, # [3,2] cutlass.Int32 stream: cuda.CUstream, ): build_group_ptrs_from_bases_kernel( base_A_u64, base_B_u64, base_C_u64, offs, K, N, sizeof_element, stride_A_m_elems, stride_A_k_elems, stride_B0_elems, stride_Bk_elems, stride_Bn_elems, stride_C_m_elems, stride_C_n_elems, out_ptrs, out_problem, out_strides_abc, ).launch(grid=(1, 1, 1), block=(G, 1, 1), stream=stream) {{def_kernel("input_a", "input_b", "input_a_offs")}} stream = cuda.CUstream(stream) input_b = input_b.transpose(1, 2) sumM, K = input_a.shape G, N, Kb = input_b.shape dev = input_a.device dev_idx = dev.index if dev.index is not None else 0 base_A_u64 = int(input_a.data_ptr()) base_B_u64 = int(input_b.data_ptr()) base_C_u64 = int({{get_output()}}.data_ptr()) ptrs_t = torch.empty((G, 3), device=dev, dtype=torch.int64) probs_t = torch.empty((G, 4), device=dev, dtype=torch.int32) strides_t = torch.empty((G, 3, 2), device=dev, dtype=torch.int32) prep_cache_key = get_prep_cache_key(input_a, input_b, {{get_output()}}) prep_executor = disk_cache_get(_prep_fn_cache, __file__, _KERNEL_CONFIG_KEY, prep_cache_key, dev_idx) if prep_executor is None: sizeof_element = int(input_a.element_size()) sA_m, sA_k = map(int, input_a.stride()) sB_0, sB_n, sB_k = map(int, input_b.stride()) sC_m, sC_n = map(int, {{get_output()}}.stride()) prep_executor = cute.compile( launch_build_group_ptrs_from_bases, base_A_u64=base_A_u64, base_B_u64=base_B_u64, base_C_u64=base_C_u64, offs=_to_fake_cute_tensor(input_a_offs, assumed_align=4), G=int(G), K=int(K), N=int(N), sizeof_element=sizeof_element, stride_A_m_elems=sA_m, stride_A_k_elems=sA_k, stride_B0_elems=sB_0, stride_Bk_elems=sB_k, stride_Bn_elems=sB_n, stride_C_m_elems=sC_m, stride_C_n_elems=sC_n, out_ptrs=_to_fake_cute_tensor(ptrs_t), out_problem=_to_fake_cute_tensor(probs_t, assumed_align=4), out_strides_abc=_to_fake_cute_tensor(strides_t, assumed_align=4), stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), options="--enable-tvm-ffi", ) disk_cache_set(_prep_fn_cache, __file__, _KERNEL_CONFIG_KEY, prep_cache_key, prep_executor, dev_idx) prep_executor( base_A_u64, base_B_u64, base_C_u64, input_a_offs, ptrs_t, probs_t, strides_t, ) # --- Tensormap workspace per SM --- num_tensormap_buffers, max_active_clusters = get_hardware_info() tensormap_shape = ( num_tensormap_buffers, GroupedGemmKernel.num_tensormaps, GroupedGemmKernel.bytes_per_tensormap // 8, ) tensormap_workspace_t = torch.empty(tensormap_shape, device=dev, dtype=torch.int64) # --- Total clusters --- cluster_tile_shape_mn = compute_cluster_tile_shape( (TILE_M, TILE_N), (CLUSTER_M, CLUSTER_N), bool(USE_2_CTA) ) total_num_clusters = int(compute_total_num_clusters(probs_t, cluster_tile_shape_mn)) gemm_cache_key = get_gemm_cache_key( prep_cache_key, max_active_clusters, total_num_clusters ) gemm_executor = disk_cache_get(_gemm_fn_cache, __file__, _KERNEL_CONFIG_KEY, gemm_cache_key, dev_idx) if gemm_executor is None: grouped_gemm = GroupedGemmKernel( acc_dtype=ACC_DTYPE, use_2cta_instrs=USE_2_CTA, mma_tiler_mn=(TILE_M, TILE_N), cluster_shape_mn=(CLUSTER_M, CLUSTER_N), tensormap_update_mode=TENSORMAP_UPDATE_MODE, ) a_unsq = input_a.unsqueeze(-1) b_unsq = input_b[0].unsqueeze(-1) c_unsq = {{get_output()}}.unsqueeze(-1) gemm_executor = cute.compile( grouped_gemm, _to_fake_cute_tensor(a_unsq), _to_fake_cute_tensor(b_unsq), _to_fake_cute_tensor(c_unsq), G, _to_fake_cute_tensor(probs_t, assumed_align=4), _to_fake_cute_tensor(strides_t, assumed_align=4), _to_fake_cute_tensor(ptrs_t), total_num_clusters, _to_fake_cute_tensor(tensormap_workspace_t), max_active_clusters, cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), options="--enable-tvm-ffi", ) disk_cache_set(_gemm_fn_cache, __file__, _KERNEL_CONFIG_KEY, gemm_cache_key, gemm_executor, dev_idx) gemm_executor( input_a.unsqueeze(-1), input_b[0].unsqueeze(-1), {{get_output()}}.unsqueeze(-1), probs_t, strides_t, ptrs_t, tensormap_workspace_t, ) def {{kernel_name}}_precompile(precompile_shapes, precompile_strides, precompile_dtypes, device_index=0, device_capability=None, hw_info=None): """Compile CuTe DSL kernels using fake tensors (no CUDA tensor allocation). Called by the subprocess worker to warm the disk cache. Uses fake compact CuTe tensors for cute.compile(), avoiding torch.cuda.* calls so this can run in forked subprocess workers where PyTorch CUDA is unavailable. """ import logging _precompile_log = logging.getLogger(__name__) dtype_a = getattr(torch, precompile_dtypes["input_a"]) dtype_b = getattr(torch, precompile_dtypes["input_b"]) dtype_out = getattr(torch, precompile_dtypes["output"]) shape_a = tuple(precompile_shapes["input_a"]) shape_b = tuple(precompile_shapes["input_b"]) shape_out = tuple(precompile_shapes["output"]) shape_offs = tuple(precompile_shapes["input_a_offs"]) stride_a = tuple(precompile_strides["input_a"]) stride_b = tuple(precompile_strides["input_b"]) stride_out = tuple(precompile_strides["output"]) sumM, K = shape_a G = shape_offs[0] N = shape_b[2] # _main transposes input_b: (G, K, N) -> (G, N, K) trans_shape_b = (shape_b[0], shape_b[2], shape_b[1]) trans_stride_b = (stride_b[0], stride_b[2], stride_b[1]) # Cache key must match get_prep_cache_key() which uses post-transpose tensors prep_cache_key = ( shape_a, stride_a, dtype_a, trans_shape_b, trans_stride_b, dtype_b, shape_out, stride_out, dtype_out, ) # --- Build CPU proxy tensors from metadata --- # Subprocess workers cannot allocate CUDA tensors, so we create CPU # tensors with matching shape/stride/dtype for _to_fake_cute_tensor. def _cpu_proxy(shape, stride, dtype): return torch.empty_strided(shape, stride, dtype=dtype, device="cpu") # _main transposes then unsqueezes: a.unsqueeze(-1), b[0].unsqueeze(-1), c.unsqueeze(-1) # Replicate those stride transformations on CPU proxies. proxy_a = _cpu_proxy(shape_a, stride_a, dtype_a).unsqueeze(-1) proxy_b_transposed = _cpu_proxy( trans_shape_b, trans_stride_b, dtype_b, ) proxy_b_unsq = proxy_b_transposed[0].unsqueeze(-1) proxy_out = _cpu_proxy(shape_out, stride_out, dtype_out).unsqueeze(-1) proxy_offs = _cpu_proxy(shape_offs, (1,), torch.int32) proxy_ptrs = _cpu_proxy((G, 3), (3, 1), torch.int64) proxy_probs = _cpu_proxy((G, 4), (4, 1), torch.int32) proxy_strides = _cpu_proxy((G, 3, 2), (6, 2, 1), torch.int32) # --- Compile prep_executor --- sizeof_element = torch.tensor(0, dtype=dtype_a).element_size() sA_m, sA_k = stride_a sB_0, sB_n, sB_k = trans_stride_b sC_m, sC_n = stride_out prep_executor = cute.compile( launch_build_group_ptrs_from_bases, base_A_u64=0, base_B_u64=0, base_C_u64=0, offs=_to_fake_cute_tensor(proxy_offs, assumed_align=4), G=int(G), K=int(K), N=int(N), sizeof_element=sizeof_element, stride_A_m_elems=sA_m, stride_A_k_elems=sA_k, stride_B0_elems=sB_0, stride_Bk_elems=sB_k, stride_Bn_elems=sB_n, stride_C_m_elems=sC_m, stride_C_n_elems=sC_n, out_ptrs=_to_fake_cute_tensor(proxy_ptrs), out_problem=_to_fake_cute_tensor(proxy_probs, assumed_align=4), out_strides_abc=_to_fake_cute_tensor(proxy_strides, assumed_align=4), stream=cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), options="--enable-tvm-ffi", ) disk_cache_set( _prep_fn_cache, __file__, _KERNEL_CONFIG_KEY, prep_cache_key, prep_executor, device_index, device_capability=device_capability, ) # --- Compile GEMM executor --- if hw_info is None: _precompile_log.debug( "hw_info not provided in precompile metadata; " "skipping GEMM executor precompile (will compile lazily at runtime)" ) return sm_count, max_active_clusters = hw_info tensormap_shape = ( sm_count, GroupedGemmKernel.num_tensormaps, GroupedGemmKernel.bytes_per_tensormap // 8, ) proxy_tensormap = _cpu_proxy(tensormap_shape, ( tensormap_shape[1] * tensormap_shape[2], tensormap_shape[2], 1, ), torch.int64) # Compute per-group problem sizes analytically from offsets offs_data = [sumM * (i + 1) // G for i in range(G)] problem_sizes = [] for i in range(G): m_beg = 0 if i == 0 else offs_data[i - 1] problem_sizes.append((offs_data[i] - m_beg, N, K, 1)) cluster_tile_shape_mn = compute_cluster_tile_shape( (TILE_M, TILE_N), (CLUSTER_M, CLUSTER_N), bool(USE_2_CTA), ) total_num_clusters = int(compute_total_num_clusters( problem_sizes, cluster_tile_shape_mn, )) gemm_cache_key = get_gemm_cache_key( prep_cache_key, max_active_clusters, total_num_clusters, ) grouped_gemm = GroupedGemmKernel( acc_dtype=ACC_DTYPE, use_2cta_instrs=USE_2_CTA, mma_tiler_mn=(TILE_M, TILE_N), cluster_shape_mn=(CLUSTER_M, CLUSTER_N), tensormap_update_mode=TENSORMAP_UPDATE_MODE, ) gemm_executor = cute.compile( grouped_gemm, _to_fake_cute_tensor(proxy_a), _to_fake_cute_tensor(proxy_b_unsq), _to_fake_cute_tensor(proxy_out), G, _to_fake_cute_tensor(proxy_probs, assumed_align=4), _to_fake_cute_tensor(proxy_strides, assumed_align=4), _to_fake_cute_tensor(proxy_ptrs), total_num_clusters, _to_fake_cute_tensor(proxy_tensormap), max_active_clusters, cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), options="--enable-tvm-ffi", ) disk_cache_set( _gemm_fn_cache, __file__, _KERNEL_CONFIG_KEY, gemm_cache_key, gemm_executor, device_index, device_capability=device_capability, )