From e416068a814afb954febd96c00edf04ecbc4cee7 Mon Sep 17 00:00:00 2001 From: Yanqin Zhai Date: Sat, 15 Aug 2026 09:43:05 -0700 Subject: [PATCH 1/3] initial-commit-for-mixed-cga --- python/cudnn/gemm/frost/compiler.py | 118 +++++++++++++++++- .../sm107_block_scale_matmul_1ctamma.py | 75 +++++++---- .../sm107_block_scale_matmul_2ctamma.py | 98 ++++++++++----- test/python/gemm/frost/gemm_test_utils.py | 1 + .../gemm/frost/test_block_scale_matmul.py | 89 +++++++++++++ 5 files changed, 324 insertions(+), 57 deletions(-) diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 95864c67f..24eef22e9 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -890,12 +890,110 @@ def _grid_num_clusters(cfg: TileConfig, device=None) -> int: return max_active_clusters(cfg.cgrp_size_m * cfg.cgrp_size_n, device) +def _cluster_mcast_patterns(cluster_m: int, cluster_n: int, cta_group: int) -> tuple[int, int]: + """(A, B) multicast bit patterns for a cluster shape, at CTA rank 0. + + A is shared along N (one bit per n_rank, stride cluster_m); B along the M + pairs (one bit per MMA pair, stride cta_group). The kernel shifts each by + its own rank. + """ + a_pattern = 0 + for n_idx in range(cluster_n): + a_pattern |= 1 << (n_idx * cluster_m) + b_pattern = 0 + for pair_idx in range(cluster_m // cta_group): + b_pattern |= 1 << (pair_idx * cta_group) + return a_pattern, b_pattern + + +# Preferred-cluster substitution (CU_LAUNCH_ATTRIBUTE_PREFERRED_CLUSTER_DIMENSION) +# is a property of the PART, like the B collector above — but unlike it, every +# part from SM 10.0 up can do it, so this is a FLOOR, not a range: there is no +# known ceiling to encode. (No driver check rides along: the attribute and +# Blackwell support shipped together, so a part this new implies a driver that +# knows it.) +_FLEX_CGA_MIN_ARCH = 100 + + +def _flex_cga_supported(arch: int | None = None) -> bool: + """Whether this GPU can group blocks into a preferred cluster and fall back + to a smaller one where the preferred does not fit — SM 10.0 and up.""" + a = _current_arch() if arch is None else arch + return a is not None and a >= _FLEX_CGA_MIN_ARCH + + +def min_fallback_cluster(cta_group: int) -> tuple[int, int]: + """The smallest cluster this MMA mode can fall back to: one CTA for a 1-CTA + MMA, and the 2-CTA pair for a 2-CTA one (the pair must stay inside a single + cluster). Smaller is better — it lets the device place a fallback cluster on + any leftover SM.""" + return (cta_group, 1) + + +@functools.lru_cache(maxsize=None) +def _template_reads_fallback_cluster(template_file: str) -> bool: + """Whether a template implements flexible CGA — i.e. whether it consumes the + ``fallback_cluster_shape_mnk`` constant. Read off the source rather than a + hand-kept capability flag, so it cannot drift from what the template does.""" + return "fallback_cluster_shape_mnk" in (_TEMPLATE_DIR / template_file).read_text() + + +def _flex_cga_fallback(cfg: TileConfig, cta_group: int, template_file: str) -> tuple[int, int] | None: + """The fallback cluster to attach to this launch, or ``None`` for a plain + fixed-cluster launch — byte-for-byte the pre-flexible-CGA behavior. + + Nothing here is a caller knob: the shape is `min_fallback_cluster(cta_group)` + and the rest is facts. It is OFF when the GPU cannot substitute clusters, + when the template has not been ported (its cluster constants are baked to the + preferred shape, so a CTA landing in a smaller cluster would wait on arrivals + that never come — a hang, not a lost optimization), when the config's cluster + is ALREADY the minimum, and when the config pins the N-super-block walk (that + rasterization is not invariant across the two cluster shapes). + """ + if os.environ.get("CUDNN_FROST_DISABLE_FLEX_CGA"): + return None + if not _flex_cga_supported(): + return None + if not _template_reads_fallback_cluster(template_file): + return None + fallback = min_fallback_cluster(cta_group) + if fallback == (cfg.cgrp_size_m, cfg.cgrp_size_n) or cfg.tile_swizzle_n > 1: + return None + return fallback + + +def _flex_cga_constants(cfg: TileConfig, cta_group: int, fallback_cluster: tuple[int, int] | None) -> list[str]: + """Constants for the kernel's runtime cluster-shape select. + + The config's own cluster is the PREFERRED (wide) shape; ``fallback_cluster`` + is the smaller one the device substitutes when a preferred cluster does not + fit. Everything else the kernel needs follows arithmetically from the cluster + dims it reads at runtime — only the multicast bit patterns are loop-built, so + both are precomputed here. + """ + a_pref, b_pref = _cluster_mcast_patterns(cfg.cgrp_size_m, cfg.cgrp_size_n, cta_group) + if fallback_cluster is None: + a_fb, b_fb = a_pref, b_pref + shape = "None" + else: + a_fb, b_fb = _cluster_mcast_patterns(fallback_cluster[0], fallback_cluster[1], cta_group) + shape = f"({fallback_cluster[0]}, {fallback_cluster[1]}, 1)" + return [ + f"fallback_cluster_shape_mnk = {shape}", + f"flex_a_pattern_pref = {a_pref}", + f"flex_b_pattern_pref = {b_pref}", + f"flex_a_pattern_fb = {a_fb}", + f"flex_b_pattern_fb = {b_fb}", + ] + + def _render_block_scale_tile_constants( cfg: TileConfig, chain: FusionChain, cta_group: int, *, use_tma_store_epi: bool = False, + fallback_cluster: tuple[int, int] | None = None, ) -> str: """Emit module-level constants for the block-scale matmul template. @@ -1197,8 +1295,13 @@ def _align16(x: int) -> int: f"acc_gemm_stride = {acc_gemm_stride}", f"sfa_col_bases = {tuple(sfa_col_bases)}", f"sfb_col_bases = {tuple(sfb_col_bases)}", - f"tile_swizzle_n = {cfg.tile_swizzle_n}", + # Flexible CGA pins the walk to the identity map: the super-block + # rasterization is not invariant across the two cluster shapes. + f"tile_swizzle_n = {1 if fallback_cluster is not None else cfg.tile_swizzle_n}", f"swizzle_l2_budget_bytes = {_l2_swizzle_budget_bytes()}", + # Read off the PREFERRED cluster; a fallback cluster is a divisor of it, + # so these flags dominate and the multicast code path degenerates to the + # plain load when the runtime pattern names a single peer. f"multicast_a = {cfg.multicast_a}", f"multicast_b = {cfg.multicast_b(cta_group)}", "", @@ -1338,6 +1441,7 @@ def _align16(x: int) -> int: # routed-group base offset = group_begin rows × this. NOT ab_dtype.width # (that is the packed Float4E2M1FNx2 8-bit type). lines.append(f"ab_data_elem_bits = {data_elem_bits}") + lines.extend(_flex_cga_constants(cfg, cta_group, fallback_cluster)) lines.extend(_quant_device_imports(chain)) return "\n".join(lines) @@ -1569,6 +1673,7 @@ def _render_block_scale_template( config: TileConfig, cta_group: int, scheduler: str, + fallback_cluster: tuple[int, int] | None = None, ) -> str: """Render the block-scale matmul template. Picks TMA-store when _use_tma_store_epi allows, else STG; SF TMA descriptors are hardcoded in the @@ -1588,7 +1693,7 @@ def _render_block_scale_template( host_aux_pass = _aux_call_block(aux_tensors) compile_aux_fakes = _aux_fake_block(aux_tensors, dynamic_strides=True, align_reqs=_aux_align_reqs(chain, vec_bytes=vec_bytes_epi)) compile_aux_pass = _aux_call_block(aux_tensors, prefix="fake_") - tile_constants = _render_block_scale_tile_constants(config, chain, cta_group, use_tma_store_epi=use_tma) + tile_constants = _render_block_scale_tile_constants(config, chain, cta_group, use_tma_store_epi=use_tma, fallback_cluster=fallback_cluster) if snippets.tap_constants: tile_constants += "\n" + "\n".join(snippets.tap_constants) @@ -2793,6 +2898,12 @@ def jit_from_cudnn_graph( `tile_config.CATALOG`. Execution strategy: ``cta_group`` ∈ {1, 2} and ``scheduler`` ∈ {"clc", "static"} pick the template (mainloop auto-detected). ``force_stg_epi=True`` skips the TMA-store path even when its gate accepts. + + Flexible CGA needs no argument and no caller change: where the GPU and the + template both support it, the launch carries ``config``'s cluster as the + PREFERRED shape plus the smallest fallback the MMA mode allows, so the device + fills the SMs a wide fixed cluster leaves idle (:func:`_flex_cga_fallback`). + Everywhere else the launch is the plain fixed cluster it always was. """ chain, binding = analyze_with_binding(graph) _dtype_reason = dtype_arch_reject(chain, _current_arch()) @@ -3322,6 +3433,7 @@ def _jit_block_scale( _arch_reason = _tmpl.active_reject(config) if _arch_reason is not None: raise NotImplementedError(_arch_reason) + fallback_cluster = _flex_cga_fallback(config, cta_group, _tmpl.file) _compute_output_vec_bytes(chain) # eager: rejects bad output alignment vec_bytes_epi = _epi_vec_bytes(chain, config, cta_group) _check_block_quant_supported(chain, vec_bytes_epi, config, cta_group) @@ -3332,7 +3444,7 @@ def _jit_block_scale( output_elem_bytes=DTYPE_BYTES[chain.output_dtype], use_tma_store=use_tma, ) - src = _render_block_scale_template(chain, snippets, config, cta_group, scheduler) + src = _render_block_scale_template(chain, snippets, config, cta_group, scheduler, fallback_cluster=fallback_cluster) mod = _import_kernel(src) digest = hashlib.sha256(src.encode("utf-8")).hexdigest()[:16] return CompiledFusedGemm( diff --git a/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_1ctamma.py index 7a58499fe..8f568f131 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_1ctamma.py @@ -147,8 +147,24 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - cluster_m = cluster_shape_mnk[0] - cluster_n = cluster_shape_mnk[1] + # Flexible CGA: the launch carries a preferred (wide) cluster plus a smaller + # fallback one, and the device picks per cluster — a CTA can only tell which + # by reading the hardware cluster dims. Everything cluster-shaped below then + # follows from those, so the two kinds share one body; only A's multicast bit + # pattern is loop-built and comes in precomputed per shape. + a_mcast_pattern = flex_a_pattern_pref + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + cluster_m = cluster_shape_mnk[0] + cluster_n = cluster_shape_mnk[1] + else: + cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() + cluster_m = cdim_x + cluster_n = cdim_y + a_mcast_pattern = cutlass.Int32(flex_a_pattern_pref) + # Bitwise, not `or`: both operands are runtime Booleans (this is the form + # cutlass.cute.experimental.is_preferred_cluster uses). + if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): + a_mcast_pattern = cutlass.Int32(flex_a_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -157,8 +173,6 @@ def _kernel( is_cluster_leader_cta = cta_rank_in_cluster == 0 - full_cluster_mask = cutlass.Int16((1 << cluster_size) - 1) - if warp_idx == mma_warp_id: for _i in cutlass.range_constexpr(num_a_operands): nvvm.prefetch_tensormap(tma_a_descs[_i].get_ptr()) @@ -181,10 +195,11 @@ def _kernel( ) init_tile_l = bidz - a_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_pattern = a_pattern | (1 << (n_idx * cluster_m)) - b_pattern = (1 << cluster_m) - 1 + a_pattern = a_mcast_pattern + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + b_pattern = (1 << cluster_m) - 1 + else: + b_pattern = (cutlass.Int32(1) << cluster_m) - 1 if cutlass.const_expr(multicast_a): tma_mcast_mask_a = cutlass.Int16(a_pattern) << m_rank @@ -314,9 +329,12 @@ def _kernel( M = m N = n - num_tile_m = cute.ceil_div(M, cgrp_tile_mnk[0]) - num_tile_n = cute.ceil_div(N, cgrp_tile_mnk[1]) num_k_tiles = cute.ceil_div(k, cta_tile_mnk[2]) + # The tile this cluster owns spans its OWN cluster shape; both shapes walk + # the grid as the identity map (tile == blockIdx), so they tile the problem + # identically and every output tile is still covered exactly once. + cgrp_tile_m_cur = cta_tile_mnk[0] * cluster_m + cgrp_tile_n_cur = cta_tile_mnk[1] * cluster_n if warp_idx == scheduler_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -387,8 +405,8 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_tma = cutlass.Int32(0) while is_valid != 0: - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_per_cta = tile_n * cgrp_tile_n_cur + n_rank * cta_tile_mnk[1] if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -934,8 +952,8 @@ def _kernel( # @@TMA_STORE_ONLY:END@@ while is_valid != 0: - coord_m_tile = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] + coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n = tile_n * cgrp_tile_n_cur + n_rank * cta_tile_mnk[1] acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: @@ -1312,7 +1330,7 @@ def _host( grid_x = num_tile_m_host * cluster_m grid_y = num_tile_n_host * cluster_n grid_shape = (grid_x, grid_y, batch) - _kernel( + launch = _kernel( problem_size[0], problem_size[1], problem_size[2], @@ -1323,13 +1341,28 @@ def _host( # @@TMA_STORE_ONLY:BEGIN@@ # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ - ).launch( - grid=grid_shape, - block=(threads_per_cta, 1, 1), - cluster=cluster_shape_mnk, - use_pdl=USE_PDL, - stream=stream, ) + # Flexible CGA: `cluster` is the preferred (wide) shape and + # `fallback_cluster` the regular one the device groups blocks into when a + # preferred cluster does not fit. The grid is already a multiple of the + # preferred shape, which the driver requires. + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) + else: + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + fallback_cluster=fallback_cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) @lru_cache(maxsize=None) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_2ctamma.py index 4c20d844b..f9340e7a2 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_2ctamma.py @@ -147,8 +147,27 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - cluster_m = cluster_shape_mnk[0] - cluster_n = cluster_shape_mnk[1] + # Flexible CGA: the launch carries a preferred (wide) cluster plus a smaller + # fallback one, and the device picks per cluster — a CTA can only tell which + # by reading the hardware cluster dims. Everything cluster-shaped below then + # follows from those, so the two kinds share one body; only the multicast bit + # patterns are loop-built and come in precomputed per shape. + a_mcast_pattern = flex_a_pattern_pref + b_mcast_pattern = flex_b_pattern_pref + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + cluster_m = cluster_shape_mnk[0] + cluster_n = cluster_shape_mnk[1] + else: + cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() + cluster_m = cdim_x + cluster_n = cdim_y + a_mcast_pattern = cutlass.Int32(flex_a_pattern_pref) + b_mcast_pattern = cutlass.Int32(flex_b_pattern_pref) + # Bitwise, not `or`: both operands are runtime Booleans (this is the form + # cutlass.cute.experimental.is_preferred_cluster uses). + if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): + a_mcast_pattern = cutlass.Int32(flex_a_pattern_fb) + b_mcast_pattern = cutlass.Int32(flex_b_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -183,19 +202,12 @@ def _kernel( ) init_tile_l = bidz - a_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_pattern = a_pattern | (1 << (n_idx * cluster_m)) - b_pattern = 0 - for pm_idx in cutlass.range_constexpr(cluster_m // 2): - b_pattern = b_pattern | (1 << (pm_idx * 2)) - if cutlass.const_expr(multicast_a): - tma_mcast_mask_a = cutlass.Int16(a_pattern << m_rank) + tma_mcast_mask_a = cutlass.Int16(a_mcast_pattern << m_rank) else: tma_mcast_mask_a = cutlass.Int16(1 << cta_rank_in_cluster) if cutlass.const_expr(multicast_b): - tma_mcast_mask_b = cutlass.Int16((b_pattern << pair_member) << (n_rank * cluster_m)) + tma_mcast_mask_b = cutlass.Int16((b_mcast_pattern << pair_member) << (n_rank * cluster_m)) else: tma_mcast_mask_b = cutlass.Int16(1 << cta_rank_in_cluster) @@ -305,7 +317,11 @@ def _kernel( sB_bytes = sB_elems * (ab_dtype.width // 8) num_tma_copy_bytes = (num_a_operands * (sA_bytes + sfa_smem_bytes) + num_b_operands * (sB_bytes + sfb_smem_bytes)) * 2 - pair_n_size = cgrp_tile_mnk[1] // cluster_n + # Per-CTA logical tile — the cluster cancels out, so these stay compile-time + # constants even when the cluster shape is only known at runtime. + logical_cta_tile_m = cgrp_tile_mnk[0] // cluster_shape_mnk[0] + logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_shape_mnk[1] + pair_n_size = logical_cta_tile_n # Per-CTA output rows one MMA-M block covers. The pair splits M, so this is # the per-CTA mma_inst_m — half the instruction's hardware M. epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m @@ -326,9 +342,12 @@ def _kernel( M = m N = n - num_tile_m = cute.ceil_div(M, cgrp_tile_mnk[0]) - num_tile_n = cute.ceil_div(N, cgrp_tile_mnk[1]) num_k_tiles = cute.ceil_div(k, cgrp_tile_mnk[2]) + # The tile this cluster owns spans its OWN cluster shape; both shapes walk + # the grid as the identity map (tile == blockIdx), so they tile the problem + # identically and every output tile is still covered exactly once. + cgrp_tile_m_cur = logical_cta_tile_m * cluster_m + cgrp_tile_n_cur = logical_cta_tile_n * cluster_n if warp_idx == scheduler_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -399,9 +418,8 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_tma = cutlass.Int32(0) while is_valid != 0: - logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_n - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_per_cta = tile_n * cgrp_tile_n_cur + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -421,7 +439,7 @@ def _kernel( coord_k = k_tile_idx * cgrp_tile_mnk[2] coord_sf_k = k_tile_idx * sf_tma_box_k - coord_n_pair = tile_n * cgrp_tile_mnk[1] + n_rank * logical_cta_tile_n + coord_n_pair = tile_n * cgrp_tile_n_cur + n_rank * logical_cta_tile_n sfb_n_block = coord_n_pair // 128 if is_pair_leader: @@ -643,12 +661,11 @@ def _kernel( pass pair_mask = cutlass.Int16(3) << pair_leader_rank - a_arrive_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_arrive_pattern = a_arrive_pattern | (1 << (n_idx * cluster_m)) - b_arrive_pattern = 0 - for m_idx in cutlass.range_constexpr(cluster_m): - b_arrive_pattern = b_arrive_pattern | (1 << m_idx) + a_arrive_pattern = a_mcast_pattern + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + b_arrive_pattern = (1 << cluster_m) - 1 + else: + b_arrive_pattern = (cutlass.Int32(1) << cluster_m) - 1 a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) @@ -1014,8 +1031,8 @@ def _kernel( # @@TMA_STORE_ONLY:END@@ while is_valid != 0: - coord_m_tile = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_c = tile_n * cgrp_tile_mnk[1] + n_rank * pair_n_size + coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_c = tile_n * cgrp_tile_n_cur + n_rank * pair_n_size if cutlass.const_expr(epi_rows_per_mma_m == 64): coord_n_c = coord_n_c + (warp_idx // 2) * cols_per_acc_stage @@ -1399,7 +1416,7 @@ def _host( grid_x = num_tile_m_host * cluster_m grid_y = num_tile_n_host * cluster_n grid_shape = (grid_x, grid_y, batch) - _kernel( + launch = _kernel( problem_size[0], problem_size[1], problem_size[2], @@ -1410,13 +1427,28 @@ def _host( # @@TMA_STORE_ONLY:BEGIN@@ # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ - ).launch( - grid=grid_shape, - block=(threads_per_cta, 1, 1), - cluster=cluster_shape_mnk, - use_pdl=USE_PDL, - stream=stream, ) + # Flexible CGA: `cluster` is the preferred (wide) shape and + # `fallback_cluster` the regular one the device groups blocks into when a + # preferred cluster does not fit. The grid is already a multiple of the + # preferred shape, which the driver requires. + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) + else: + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + fallback_cluster=fallback_cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) @lru_cache(maxsize=None) diff --git a/test/python/gemm/frost/gemm_test_utils.py b/test/python/gemm/frost/gemm_test_utils.py index ebe53420d..3916e2cfc 100644 --- a/test/python/gemm/frost/gemm_test_utils.py +++ b/test/python/gemm/frost/gemm_test_utils.py @@ -76,6 +76,7 @@ def __init__(self, graph, config=None, cta_group=2, scheduler="clc", force_stg_e self.binding = self._compiled.binding self.block_scale = self.chain.has_block_scale self.aux_names = [t.name for t in self.chain.aux_tensors] + self.generated_path = self._compiled.generated_path def __call__(self, variant_pack): return self._compiled(variant_pack) diff --git a/test/python/gemm/frost/test_block_scale_matmul.py b/test/python/gemm/frost/test_block_scale_matmul.py index 3ba949fd3..494d0323d 100644 --- a/test/python/gemm/frost/test_block_scale_matmul.py +++ b/test/python/gemm/frost/test_block_scale_matmul.py @@ -13,6 +13,8 @@ import cudnn import cudnn.gemm.frost # noqa: F401 (installs recorder) +import dataclasses + import pytest import torch @@ -2012,6 +2014,93 @@ def test_sm107_block_scale_matmul_shapes_and_clusters(combo, config_name, M, N, _run_bs_numeric(combo, config_name, M, N, K) +@requires_sm107 +@pytest.mark.parametrize("combo", ["nvfp4", "mxfp8"]) +@pytest.mark.parametrize( + "config_name", + [ + "CONFIG_sm107_128x128x128_128x128x64_cluster4x1_2ctamma", + "CONFIG_sm107_128x128x128_128x128x64_cluster4x2_2ctamma", + "CONFIG_sm107_256x256x128_128x256x64_cluster2x4_2ctamma", + "CONFIG_sm107_128x128x128_128x128x64_cluster4x1_1ctamma", + "CONFIG_sm107_128x128x128_128x128x64_cluster1x4_1ctamma", + "CONFIG_sm107_128x128x128_128x128x64_cluster2x2_1ctamma", + ], +) +def test_sm107_block_scale_flexible_cga(combo, config_name): + """Flexible CGA rides along with no caller change: any config whose cluster is + wider than the MMA mode's minimum launches it as the PREFERRED shape plus that + minimum as the fallback. The tile decomposition is the identity map for either + cluster shape, so both kinds cover the problem exactly once; only the multicast + masks, mbarrier arrival counts and rank math follow the shape the CTA actually + landed in. M is large enough that the grid outruns what the preferred clusters + hold resident, which is when the device substitutes the fallback shape.""" + cta_group = 2 if config_name.endswith("_2ctamma") else 1 + cfg = by_name(config_name.rsplit("_", 1)[0]) + assert C._flex_cga_fallback(cfg, cta_group, f"sm107_block_scale_matmul_{cta_group}ctamma.py") == (cta_group, 1) + _run_bs_numeric(combo, config_name, 1920, 1920, 512) + + +@requires_sm107 +@pytest.mark.parametrize("cta_group", [1, 2]) +def test_flex_cga_fallback_is_the_mma_mode_minimum(cta_group): + """The fallback shape is derived, never passed: one CTA for a 1-CTA MMA, the + pair for a 2-CTA one — and a config already AT that minimum has nothing to + fall back to, so it launches as a plain fixed cluster.""" + tmpl = f"sm107_block_scale_matmul_{cta_group}ctamma.py" + assert C.min_fallback_cluster(cta_group) == (cta_group, 1) + wide = by_name("CONFIG_sm107_128x128x128_128x128x64_cluster4x2") + assert C._flex_cga_fallback(wide, cta_group, tmpl) == (cta_group, 1) + minimal = by_name(f"CONFIG_sm107_128x128x128_128x128x64_cluster{cta_group}x1") + assert C._flex_cga_fallback(minimal, cta_group, tmpl) is None + + +@requires_sm107 +def test_flex_cga_is_off_where_it_cannot_be_honored(monkeypatch): + """Every gate is a fact, not a knob: the GPU's ability to substitute clusters, + whether the template consumes the fallback constant at all (an unported one + would hang — its cluster constants are baked to the preferred shape), and + whether the config pins the N-super-block walk (not invariant across the two + cluster shapes).""" + wide = by_name("CONFIG_sm107_128x128x128_128x128x64_cluster4x2") + sm107_tmpl = "sm107_block_scale_matmul_2ctamma.py" + assert C._flex_cga_fallback(wide, 2, sm107_tmpl) == (2, 1) + + # Template that never reads the constant -> no fallback attached. + assert not C._template_reads_fallback_cluster("sm100_block_scale_matmul_2ctamma.py") + assert C._flex_cga_fallback(wide, 2, "sm100_block_scale_matmul_2ctamma.py") is None + + # Substitution is a floor, not a range: every part from SM 10.0 up can do it. + assert C._flex_cga_supported(100) and C._flex_cga_supported(110) + # A pre-Blackwell part -> plain fixed cluster, as before. + monkeypatch.setattr(C, "_current_arch", lambda: 90) + assert not C._flex_cga_supported() + assert C._flex_cga_fallback(wide, 2, sm107_tmpl) is None + monkeypatch.undo() + + # A pinned N-super-block walk -> skipped rather than silently mis-tiled. + pinned = dataclasses.replace(wide, tile_swizzle_n=8) + assert C._flex_cga_fallback(pinned, 2, sm107_tmpl) is None + + # The escape hatch for A/B measurement. + monkeypatch.setenv("CUDNN_FROST_DISABLE_FLEX_CGA", "1") + assert C._flex_cga_fallback(wide, 2, sm107_tmpl) is None + + +@requires_sm107 +def test_flex_cga_leaves_unported_templates_byte_identical(): + """An sm100 template on this GPU gets no fallback in its launch — so its + generated kernel, including the L2 swizzle knob the flexible path pins, is + exactly what it was before flexible CGA existed.""" + sm100_cfg = "CONFIG_sm100_128x128x128_128x128x32_cluster4x2_2ctamma" + _run_bs_numeric("nvfp4", sm100_cfg, 512, 512, 512) + g = _build_nvfp4_graph(256, 256, 512) + compiled = _plan(g, **_kw(sm100_cfg)) + src = compiled.generated_path.read_text() + assert "fallback_cluster=" not in src + assert "fallback_cluster_shape_mnk = None" in src + + @requires_sm107 @pytest.mark.parametrize( "config_name", From b2369f15f40a0f3fdec2a30987f9498176aee14c Mon Sep 17 00:00:00 2001 From: Yanqin Zhai Date: Sun, 16 Aug 2026 14:31:02 -0700 Subject: [PATCH 2/3] enable_mixed_cga_for_all --- python/cudnn/gemm/frost/compiler.py | 46 ++++---- python/cudnn/gemm/frost/kernel_registry.py | 29 ++--- .../sm100_block_scale_matmul_1ctamma.py | 75 +++++++++---- ...sm100_block_scale_matmul_1ctamma_static.py | 75 +++++++++---- .../sm100_block_scale_matmul_2ctamma.py | 98 +++++++++++------ ...sm100_block_scale_matmul_2ctamma_static.py | 98 +++++++++++------ .../kernel_templates/sm100_matmul_1ctamma.py | 76 +++++++++---- .../sm100_matmul_1ctamma_static.py | 76 +++++++++---- .../kernel_templates/sm100_matmul_2ctamma.py | 97 +++++++++++------ .../sm100_matmul_2ctamma_static.py | 97 +++++++++++------ .../sm100_matmul_mainloop_1ctamma.py | 76 +++++++++---- .../sm100_matmul_mainloop_2ctamma.py | 97 +++++++++++------ .../sm103_block_scale_matmul_1ctamma.py | 79 ++++++++++---- .../sm103_block_scale_matmul_2ctamma.py | 103 ++++++++++++------ .../sm107_block_scale_matmul_1ctamma.py | 10 +- .../sm107_block_scale_matmul_2ctamma.py | 16 +-- .../gemm/frost/test_block_scale_matmul.py | 53 +++++---- 17 files changed, 793 insertions(+), 408 deletions(-) diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 24eef22e9..117bd302b 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -484,6 +484,8 @@ def _render_tile_constants( chain: FusionChain, cta_group: int, use_tma: bool = True, + *, + fallback_cluster: tuple[int, int] | None = None, ) -> str: """Emit module-level tile + dtype constants for the config/chain, appended below the template's defaults (last assignment wins). TileConfig geometry is @@ -607,7 +609,7 @@ def _smem_desc_params( ] # Persistent kernel always: double-TMEM + L2 N-super-block swizzle. # (acc_stages is emitted below, once the TMEM budget is known.) - lines.append(f"tile_swizzle_n = {cfg.tile_swizzle_n}") + lines.append(f"tile_swizzle_n = {1 if fallback_cluster is not None else cfg.tile_swizzle_n}") lines.append(f"swizzle_l2_budget_bytes = {_l2_swizzle_budget_bytes()}") # Multi-GEMM (parallel matmuls sharing the epilogue). Always emitted; # single-GEMM = (1, 1, 1). gemm_a_idx[g]/gemm_b_idx[g] pick GEMM g's operand @@ -725,6 +727,7 @@ def _smem_desc_params( ) lines.append(f"ab_stages = {new_ab} # SMEM-D {smem_d_bytes}B fixed" f" + cast LOAD {cast_extra_per_stage}B/stage") lines.extend(_quant_device_imports(chain)) + lines.extend(_mixed_cga_constants(cfg, cta_group, fallback_cluster)) return "\n".join(lines) @@ -912,14 +915,14 @@ def _cluster_mcast_patterns(cluster_m: int, cluster_n: int, cta_group: int) -> t # known ceiling to encode. (No driver check rides along: the attribute and # Blackwell support shipped together, so a part this new implies a driver that # knows it.) -_FLEX_CGA_MIN_ARCH = 100 +_MIXED_CGA_MIN_ARCH = 100 -def _flex_cga_supported(arch: int | None = None) -> bool: +def _mixed_cga_supported(arch: int | None = None) -> bool: """Whether this GPU can group blocks into a preferred cluster and fall back to a smaller one where the preferred does not fit — SM 10.0 and up.""" a = _current_arch() if arch is None else arch - return a is not None and a >= _FLEX_CGA_MIN_ARCH + return a is not None and a >= _MIXED_CGA_MIN_ARCH def min_fallback_cluster(cta_group: int) -> tuple[int, int]: @@ -932,15 +935,15 @@ def min_fallback_cluster(cta_group: int) -> tuple[int, int]: @functools.lru_cache(maxsize=None) def _template_reads_fallback_cluster(template_file: str) -> bool: - """Whether a template implements flexible CGA — i.e. whether it consumes the + """Whether a template implements mixed CGA — i.e. whether it consumes the ``fallback_cluster_shape_mnk`` constant. Read off the source rather than a hand-kept capability flag, so it cannot drift from what the template does.""" return "fallback_cluster_shape_mnk" in (_TEMPLATE_DIR / template_file).read_text() -def _flex_cga_fallback(cfg: TileConfig, cta_group: int, template_file: str) -> tuple[int, int] | None: +def _mixed_cga_fallback(cfg: TileConfig, cta_group: int, template_file: str) -> tuple[int, int] | None: """The fallback cluster to attach to this launch, or ``None`` for a plain - fixed-cluster launch — byte-for-byte the pre-flexible-CGA behavior. + fixed-cluster launch — byte-for-byte the pre-mixed-CGA behavior. Nothing here is a caller knob: the shape is `min_fallback_cluster(cta_group)` and the rest is facts. It is OFF when the GPU cannot substitute clusters, @@ -950,9 +953,9 @@ def _flex_cga_fallback(cfg: TileConfig, cta_group: int, template_file: str) -> t is ALREADY the minimum, and when the config pins the N-super-block walk (that rasterization is not invariant across the two cluster shapes). """ - if os.environ.get("CUDNN_FROST_DISABLE_FLEX_CGA"): + if os.environ.get("CUDNN_FROST_DISABLE_MIXED_CGA"): return None - if not _flex_cga_supported(): + if not _mixed_cga_supported(): return None if not _template_reads_fallback_cluster(template_file): return None @@ -962,7 +965,7 @@ def _flex_cga_fallback(cfg: TileConfig, cta_group: int, template_file: str) -> t return fallback -def _flex_cga_constants(cfg: TileConfig, cta_group: int, fallback_cluster: tuple[int, int] | None) -> list[str]: +def _mixed_cga_constants(cfg: TileConfig, cta_group: int, fallback_cluster: tuple[int, int] | None) -> list[str]: """Constants for the kernel's runtime cluster-shape select. The config's own cluster is the PREFERRED (wide) shape; ``fallback_cluster`` @@ -980,10 +983,10 @@ def _flex_cga_constants(cfg: TileConfig, cta_group: int, fallback_cluster: tuple shape = f"({fallback_cluster[0]}, {fallback_cluster[1]}, 1)" return [ f"fallback_cluster_shape_mnk = {shape}", - f"flex_a_pattern_pref = {a_pref}", - f"flex_b_pattern_pref = {b_pref}", - f"flex_a_pattern_fb = {a_fb}", - f"flex_b_pattern_fb = {b_fb}", + f"mixed_a_pattern_pref = {a_pref}", + f"mixed_b_pattern_pref = {b_pref}", + f"mixed_a_pattern_fb = {a_fb}", + f"mixed_b_pattern_fb = {b_fb}", ] @@ -1295,7 +1298,7 @@ def _align16(x: int) -> int: f"acc_gemm_stride = {acc_gemm_stride}", f"sfa_col_bases = {tuple(sfa_col_bases)}", f"sfb_col_bases = {tuple(sfb_col_bases)}", - # Flexible CGA pins the walk to the identity map: the super-block + # Mixed CGA pins the walk to the identity map: the super-block # rasterization is not invariant across the two cluster shapes. f"tile_swizzle_n = {1 if fallback_cluster is not None else cfg.tile_swizzle_n}", f"swizzle_l2_budget_bytes = {_l2_swizzle_budget_bytes()}", @@ -1441,7 +1444,7 @@ def _align16(x: int) -> int: # routed-group base offset = group_begin rows × this. NOT ab_dtype.width # (that is the packed Float4E2M1FNx2 8-bit type). lines.append(f"ab_data_elem_bits = {data_elem_bits}") - lines.extend(_flex_cga_constants(cfg, cta_group, fallback_cluster)) + lines.extend(_mixed_cga_constants(cfg, cta_group, fallback_cluster)) lines.extend(_quant_device_imports(chain)) return "\n".join(lines) @@ -1512,6 +1515,7 @@ def _render_template( from .kernel_registry import select_template tmpl = select_template(chain, config, cta_group, scheduler) + fallback_cluster = _mixed_cga_fallback(config, cta_group, tmpl.file) template_path = _TEMPLATE_DIR / tmpl.file src = template_path.read_text() # Strip the unused epilogue path FIRST so its @@INJECT_EPILOGUE@@ marker @@ -1566,7 +1570,7 @@ def _render_template( align_reqs=_aux_align_reqs(chain, vec_bytes=vec_bytes_epi), ) compile_aux_pass = _aux_call_block(aux_tensors, prefix="fake_") - tile_constants = _render_tile_constants(config, chain, cta_group, use_tma) + tile_constants = _render_tile_constants(config, chain, cta_group, use_tma, fallback_cluster=fallback_cluster) if snippets.tap_constants: tile_constants += "\n" + "\n".join(snippets.tap_constants) # Multi-output tap plumbing. Empty lists → markers expand to nothing (kernel @@ -2899,10 +2903,10 @@ def jit_from_cudnn_graph( ``scheduler`` ∈ {"clc", "static"} pick the template (mainloop auto-detected). ``force_stg_epi=True`` skips the TMA-store path even when its gate accepts. - Flexible CGA needs no argument and no caller change: where the GPU and the + Mixed CGA needs no argument and no caller change: where the GPU and the template both support it, the launch carries ``config``'s cluster as the PREFERRED shape plus the smallest fallback the MMA mode allows, so the device - fills the SMs a wide fixed cluster leaves idle (:func:`_flex_cga_fallback`). + fills the SMs a wide fixed cluster leaves idle (:func:`_mixed_cga_fallback`). Everywhere else the launch is the plain fixed cluster it always was. """ chain, binding = analyze_with_binding(graph) @@ -3433,7 +3437,7 @@ def _jit_block_scale( _arch_reason = _tmpl.active_reject(config) if _arch_reason is not None: raise NotImplementedError(_arch_reason) - fallback_cluster = _flex_cga_fallback(config, cta_group, _tmpl.file) + fallback_cluster = _mixed_cga_fallback(config, cta_group, _tmpl.file) _compute_output_vec_bytes(chain) # eager: rejects bad output alignment vec_bytes_epi = _epi_vec_bytes(chain, config, cta_group) _check_block_quant_supported(chain, vec_bytes_epi, config, cta_group) @@ -3783,7 +3787,7 @@ def _jit_moe_block_scale( output_elem_bytes=DTYPE_BYTES[chain.output_dtype], use_tma_store=(not _FORCE_STG_EPI) and _use_tma_store_epi(chain, config, vec_bytes_epi, cta_group), ) - src = _render_block_scale_template(chain, snippets, config, cta_group, scheduler) + src = _render_block_scale_template(chain, snippets, config, cta_group, scheduler, fallback_cluster=_mixed_cga_fallback(config, cta_group, _tmpl.file)) mod = _import_kernel(src) digest = hashlib.sha256(src.encode("utf-8")).hexdigest()[:16] cluster_m, cluster_n = config.cgrp_size_m, config.cgrp_size_n diff --git a/python/cudnn/gemm/frost/kernel_registry.py b/python/cudnn/gemm/frost/kernel_registry.py index 82df5f1a2..56cb68d15 100644 --- a/python/cudnn/gemm/frost/kernel_registry.py +++ b/python/cudnn/gemm/frost/kernel_registry.py @@ -270,7 +270,7 @@ class KernelTemplate: graph_type: GraphType # the single graph type this template supports mainloop: bool # mainloop-fusion variant (transform A/B before MMA) # Multi-GEMM support (templates without it reject multi-GEMM chains). - supports_multi_gemm: bool = False + supports_multi_gemm: bool = True # A CTA tile spanning several MMA instructions along M (num_mma_m > 1). supports_multi_mma_m: bool = True @@ -433,7 +433,7 @@ def _mm( static: bool, mainloop: bool = False, graph_type: GraphType = GraphType.MATMUL, - supports_multi_gemm: bool = False, + supports_multi_gemm: bool = True, supports_multi_mma_m: bool = True, ) -> KernelTemplate: pipeline = _pipeline_from_file(file) @@ -456,19 +456,17 @@ def _mm( TEMPLATES: tuple[KernelTemplate, ...] = ( # plain matmul - _mm("sm100_matmul_1ctamma.py", cta_group=1, static=False, supports_multi_gemm=True), + _mm("sm100_matmul_1ctamma.py", cta_group=1, static=False), _mm( "sm100_matmul_1ctamma_static.py", cta_group=1, static=True, - supports_multi_gemm=True, ), - _mm("sm100_matmul_2ctamma.py", cta_group=2, static=False, supports_multi_gemm=True), + _mm("sm100_matmul_2ctamma.py", cta_group=2, static=False), _mm( "sm100_matmul_2ctamma_static.py", cta_group=2, static=True, - supports_multi_gemm=True, ), # block-scaled matmul _mm( @@ -476,28 +474,24 @@ def _mm( cta_group=1, static=False, graph_type=GraphType.BLOCK_SCALE_MATMUL, - supports_multi_gemm=True, ), _mm( "sm100_block_scale_matmul_1ctamma_static.py", cta_group=1, static=True, graph_type=GraphType.BLOCK_SCALE_MATMUL, - supports_multi_gemm=True, ), _mm( "sm100_block_scale_matmul_2ctamma.py", cta_group=2, static=False, graph_type=GraphType.BLOCK_SCALE_MATMUL, - supports_multi_gemm=True, ), _mm( "sm100_block_scale_matmul_2ctamma_static.py", cta_group=2, static=True, graph_type=GraphType.BLOCK_SCALE_MATMUL, - supports_multi_gemm=True, ), # sm103 block-scaled matmul: fp4-only (nvfp4/mxfp4), K=48B UTCOMMA # (K-tile 384 B, 8 MMAs over 3× 128-B chunks via circular SMEM descs). @@ -524,18 +518,19 @@ def _mm( cta_group=1, static=False, graph_type=GraphType.BLOCK_SCALE_MATMUL, - supports_multi_gemm=True, ), _mm( "sm107_block_scale_matmul_2ctamma.py", cta_group=2, static=False, graph_type=GraphType.BLOCK_SCALE_MATMUL, - supports_multi_gemm=True, ), # mainloop-fusion matmul (CLC only — no static / block-scale variant yet) - _mm("sm100_matmul_mainloop_1ctamma.py", cta_group=1, static=False, mainloop=True), - _mm("sm100_matmul_mainloop_2ctamma.py", cta_group=2, static=False, mainloop=True), + # The mainloop templates have no per-GEMM operand indexing (no gemm_a_idx / + # gemm_b_idx in the MMA warp), so a second GEMM's accumulator would never be + # written. The analyzer also only detects mainloop at len(matmuls) == 1. + _mm("sm100_matmul_mainloop_1ctamma.py", cta_group=1, static=False, mainloop=True, supports_multi_gemm=False), + _mm("sm100_matmul_mainloop_2ctamma.py", cta_group=2, static=False, mainloop=True, supports_multi_gemm=False), # MoE grouped matmul fwd (own grouped persistent scheduler; static_sched # irrelevant, registered False so default scheduler="clc" selects). _mm( @@ -543,14 +538,12 @@ def _mm( cta_group=1, static=False, graph_type=GraphType.MOE, - supports_multi_gemm=True, ), _mm( "sm100_moe_grouped_matmul_fwd_2ctamma.py", cta_group=2, static=False, graph_type=GraphType.MOE, - supports_multi_gemm=True, ), # MoE grouped matmul with block-scaled (FP4/FP8 + SF) inputs. _mm( @@ -558,28 +551,24 @@ def _mm( cta_group=1, static=False, graph_type=GraphType.MOE_BLOCK_SCALE, - supports_multi_gemm=True, ), _mm( "sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py", cta_group=2, static=False, graph_type=GraphType.MOE_BLOCK_SCALE, - supports_multi_gemm=True, ), _mm( "sm107_moe_grouped_block_scale_matmul_fwd_1ctamma.py", cta_group=1, static=False, graph_type=GraphType.MOE_BLOCK_SCALE, - supports_multi_gemm=True, ), _mm( "sm107_moe_grouped_block_scale_matmul_fwd_2ctamma.py", cta_group=2, static=False, graph_type=GraphType.MOE_BLOCK_SCALE, - supports_multi_gemm=True, ), ) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma.py index ee7c5c44a..910f7cbba 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma.py @@ -130,8 +130,24 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - cluster_m = cluster_shape_mnk[0] - cluster_n = cluster_shape_mnk[1] + # Mixed CGA: the launch carries a preferred (wide) cluster plus a smaller + # fallback one, and the device picks per cluster — a CTA can only tell which + # by reading the hardware cluster dims. Everything cluster-shaped below then + # follows from those, so the two kinds share one body; only the multicast bit + # pattern is loop-built and comes in precomputed per shape. + a_mcast_pattern = mixed_a_pattern_pref + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + cluster_m = cluster_shape_mnk[0] + cluster_n = cluster_shape_mnk[1] + else: + cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() + cluster_m = cdim_x + cluster_n = cdim_y + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_pref) + # Bitwise, not `or`: both operands are runtime Booleans (this is the form + # cutlass.cute.experimental.is_preferred_cluster uses). + if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -140,8 +156,6 @@ def _kernel( is_cluster_leader_cta = cta_rank_in_cluster == 0 - full_cluster_mask = cutlass.Int16((1 << cluster_size) - 1) - if warp_idx == mma_warp_id: for _i in cutlass.range_constexpr(num_a_operands): nvvm.prefetch_tensormap(tma_a_descs[_i].get_ptr()) @@ -164,10 +178,11 @@ def _kernel( ) init_tile_l = bidz - a_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_pattern = a_pattern | (1 << (n_idx * cluster_m)) - b_pattern = (1 << cluster_m) - 1 + a_pattern = a_mcast_pattern + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + b_pattern = (1 << cluster_m) - 1 + else: + b_pattern = (cutlass.Int32(1) << cluster_m) - 1 if cutlass.const_expr(multicast_a): tma_mcast_mask_a = cutlass.Int16(a_pattern) << m_rank @@ -297,9 +312,12 @@ def _kernel( M = m N = n - num_tile_m = cute.ceil_div(M, cgrp_tile_mnk[0]) - num_tile_n = cute.ceil_div(N, cgrp_tile_mnk[1]) num_k_tiles = cute.ceil_div(k, cta_tile_mnk[2]) + # The tile this cluster owns spans its OWN cluster shape; both shapes walk + # the grid as the identity map (tile == blockIdx), so they tile the problem + # identically and every output tile is still covered exactly once. + cgrp_tile_m_cur = cta_tile_mnk[0] * cluster_m + cgrp_tile_n_cur = cta_tile_mnk[1] * cluster_n if warp_idx == scheduler_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -370,8 +388,8 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_tma = cutlass.Int32(0) while is_valid != 0: - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_per_cta = tile_n * cgrp_tile_n_cur + n_rank * cta_tile_mnk[1] if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -878,8 +896,8 @@ def _kernel( # @@TMA_STORE_ONLY:END@@ while is_valid != 0: - coord_m_tile = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] + coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n = tile_n * cgrp_tile_n_cur + n_rank * cta_tile_mnk[1] acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: @@ -1256,7 +1274,7 @@ def _host( grid_x = num_tile_m_host * cluster_m grid_y = num_tile_n_host * cluster_n grid_shape = (grid_x, grid_y, batch) - _kernel( + launch = _kernel( problem_size[0], problem_size[1], problem_size[2], @@ -1267,13 +1285,28 @@ def _host( # @@TMA_STORE_ONLY:BEGIN@@ # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ - ).launch( - grid=grid_shape, - block=(threads_per_cta, 1, 1), - cluster=cluster_shape_mnk, - use_pdl=USE_PDL, - stream=stream, ) + # Mixed CGA: `cluster` is the preferred (wide) shape and `fallback_cluster` + # the regular one the device groups blocks into when a preferred cluster does + # not fit. The grid is already a multiple of the preferred shape, which the + # driver requires. + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) + else: + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + fallback_cluster=fallback_cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) @lru_cache(maxsize=None) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma_static.py b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma_static.py index 640cbdfc8..0c9382ff8 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma_static.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_1ctamma_static.py @@ -129,8 +129,24 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - cluster_m = cluster_shape_mnk[0] - cluster_n = cluster_shape_mnk[1] + # Mixed CGA: the launch carries a preferred (wide) cluster plus a smaller + # fallback one, and the device picks per cluster — a CTA can only tell which + # by reading the hardware cluster dims. Everything cluster-shaped below then + # follows from those, so the two kinds share one body; only the multicast bit + # pattern is loop-built and comes in precomputed per shape. + a_mcast_pattern = mixed_a_pattern_pref + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + cluster_m = cluster_shape_mnk[0] + cluster_n = cluster_shape_mnk[1] + else: + cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() + cluster_m = cdim_x + cluster_n = cdim_y + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_pref) + # Bitwise, not `or`: both operands are runtime Booleans (this is the form + # cutlass.cute.experimental.is_preferred_cluster uses). + if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -139,8 +155,6 @@ def _kernel( is_cluster_leader_cta = cta_rank_in_cluster == 0 - full_cluster_mask = cutlass.Int16((1 << cluster_size) - 1) - if warp_idx == mma_warp_id: for _i in cutlass.range_constexpr(num_a_operands): nvvm.prefetch_tensormap(tma_a_descs[_i].get_ptr()) @@ -163,10 +177,11 @@ def _kernel( ) init_tile_l = bidz - a_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_pattern = a_pattern | (1 << (n_idx * cluster_m)) - b_pattern = (1 << cluster_m) - 1 + a_pattern = a_mcast_pattern + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + b_pattern = (1 << cluster_m) - 1 + else: + b_pattern = (cutlass.Int32(1) << cluster_m) - 1 if cutlass.const_expr(multicast_a): tma_mcast_mask_a = cutlass.Int16(a_pattern) << m_rank @@ -278,9 +293,12 @@ def _kernel( M = m N = n - num_tile_m = cute.ceil_div(M, cgrp_tile_mnk[0]) - num_tile_n = cute.ceil_div(N, cgrp_tile_mnk[1]) num_k_tiles = cute.ceil_div(k, cta_tile_mnk[2]) + # The tile this cluster owns spans its OWN cluster shape; both shapes walk + # the grid as the identity map (tile == blockIdx), so they tile the problem + # identically and every output tile is still covered exactly once. + cgrp_tile_m_cur = cta_tile_mnk[0] * cluster_m + cgrp_tile_n_cur = cta_tile_mnk[1] * cluster_n if warp_idx == scheduler_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -298,8 +316,8 @@ def _kernel( tile_iter = cutlass.Int32(0) is_valid = cutlass.Int32(1) while is_valid != 0: - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_per_cta = tile_n * cgrp_tile_n_cur + n_rank * cta_tile_mnk[1] if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -769,8 +787,8 @@ def _kernel( # @@TMA_STORE_ONLY:END@@ while is_valid != 0: - coord_m_tile = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] + coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n = tile_n * cgrp_tile_n_cur + n_rank * cta_tile_mnk[1] acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: @@ -1125,7 +1143,7 @@ def _host( grid_x = num_tile_m_host * cluster_m grid_y = num_tile_n_host * cluster_n grid_shape = (grid_x, grid_y, batch) - _kernel( + launch = _kernel( problem_size[0], problem_size[1], problem_size[2], @@ -1136,13 +1154,28 @@ def _host( # @@TMA_STORE_ONLY:BEGIN@@ # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ - ).launch( - grid=grid_shape, - block=(threads_per_cta, 1, 1), - cluster=cluster_shape_mnk, - use_pdl=USE_PDL, - stream=stream, ) + # Mixed CGA: `cluster` is the preferred (wide) shape and `fallback_cluster` + # the regular one the device groups blocks into when a preferred cluster does + # not fit. The grid is already a multiple of the preferred shape, which the + # driver requires. + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) + else: + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + fallback_cluster=fallback_cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) @lru_cache(maxsize=None) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma.py index 1f33dd1c0..3bea40b8f 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma.py @@ -130,8 +130,27 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - cluster_m = cluster_shape_mnk[0] - cluster_n = cluster_shape_mnk[1] + # Mixed CGA: the launch carries a preferred (wide) cluster plus a smaller + # fallback one, and the device picks per cluster — a CTA can only tell which + # by reading the hardware cluster dims. Everything cluster-shaped below then + # follows from those, so the two kinds share one body; only the multicast bit + # patterns are loop-built and come in precomputed per shape. + a_mcast_pattern = mixed_a_pattern_pref + b_mcast_pattern = mixed_b_pattern_pref + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + cluster_m = cluster_shape_mnk[0] + cluster_n = cluster_shape_mnk[1] + else: + cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() + cluster_m = cdim_x + cluster_n = cdim_y + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_pref) + b_mcast_pattern = cutlass.Int32(mixed_b_pattern_pref) + # Bitwise, not `or`: both operands are runtime Booleans (this is the form + # cutlass.cute.experimental.is_preferred_cluster uses). + if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_fb) + b_mcast_pattern = cutlass.Int32(mixed_b_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -166,19 +185,12 @@ def _kernel( ) init_tile_l = bidz - a_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_pattern = a_pattern | (1 << (n_idx * cluster_m)) - b_pattern = 0 - for pm_idx in cutlass.range_constexpr(cluster_m // 2): - b_pattern = b_pattern | (1 << (pm_idx * 2)) - if cutlass.const_expr(multicast_a): - tma_mcast_mask_a = cutlass.Int16(a_pattern << m_rank) + tma_mcast_mask_a = cutlass.Int16(a_mcast_pattern << m_rank) else: tma_mcast_mask_a = cutlass.Int16(1 << cta_rank_in_cluster) if cutlass.const_expr(multicast_b): - tma_mcast_mask_b = cutlass.Int16((b_pattern << pair_member) << (n_rank * cluster_m)) + tma_mcast_mask_b = cutlass.Int16((b_mcast_pattern << pair_member) << (n_rank * cluster_m)) else: tma_mcast_mask_b = cutlass.Int16(1 << cta_rank_in_cluster) @@ -288,7 +300,11 @@ def _kernel( sB_bytes = sB_elems * (ab_dtype.width // 8) num_tma_copy_bytes = (num_a_operands * (sA_bytes + sfa_smem_bytes) + num_b_operands * (sB_bytes + sfb_smem_bytes)) * 2 - pair_n_size = cgrp_tile_mnk[1] // cluster_n + # Per-CTA logical tile — the cluster cancels out, so these stay compile-time + # constants even when the cluster shape is only known at runtime. + logical_cta_tile_m = cgrp_tile_mnk[0] // cluster_shape_mnk[0] + logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_shape_mnk[1] + pair_n_size = logical_cta_tile_n # Per-CTA output rows one MMA-M block covers. The pair splits M, so this is # the per-CTA mma_inst_m — half the instruction's hardware M. epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m @@ -309,9 +325,12 @@ def _kernel( M = m N = n - num_tile_m = cute.ceil_div(M, cgrp_tile_mnk[0]) - num_tile_n = cute.ceil_div(N, cgrp_tile_mnk[1]) num_k_tiles = cute.ceil_div(k, cgrp_tile_mnk[2]) + # The tile this cluster owns spans its OWN cluster shape; both shapes walk + # the grid as the identity map (tile == blockIdx), so they tile the problem + # identically and every output tile is still covered exactly once. + cgrp_tile_m_cur = logical_cta_tile_m * cluster_m + cgrp_tile_n_cur = logical_cta_tile_n * cluster_n if warp_idx == scheduler_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -382,9 +401,8 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_tma = cutlass.Int32(0) while is_valid != 0: - logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_n - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_per_cta = tile_n * cgrp_tile_n_cur + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -404,7 +422,7 @@ def _kernel( coord_k = k_tile_idx * cgrp_tile_mnk[2] coord_sf_k = k_tile_idx * sf_tma_box_k - coord_n_pair = tile_n * cgrp_tile_mnk[1] + n_rank * logical_cta_tile_n + coord_n_pair = tile_n * cgrp_tile_n_cur + n_rank * logical_cta_tile_n sfb_n_block = coord_n_pair // 128 if is_pair_leader: @@ -626,12 +644,11 @@ def _kernel( pass pair_mask = cutlass.Int16(3) << pair_leader_rank - a_arrive_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_arrive_pattern = a_arrive_pattern | (1 << (n_idx * cluster_m)) - b_arrive_pattern = 0 - for m_idx in cutlass.range_constexpr(cluster_m): - b_arrive_pattern = b_arrive_pattern | (1 << m_idx) + a_arrive_pattern = a_mcast_pattern + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + b_arrive_pattern = (1 << cluster_m) - 1 + else: + b_arrive_pattern = (cutlass.Int32(1) << cluster_m) - 1 a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) @@ -957,8 +974,8 @@ def _kernel( # @@TMA_STORE_ONLY:END@@ while is_valid != 0: - coord_m_tile = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_c = tile_n * cgrp_tile_mnk[1] + n_rank * pair_n_size + coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_c = tile_n * cgrp_tile_n_cur + n_rank * pair_n_size if cutlass.const_expr(epi_rows_per_mma_m == 64): coord_n_c = coord_n_c + (warp_idx // 2) * cols_per_acc_stage @@ -1342,7 +1359,7 @@ def _host( grid_x = num_tile_m_host * cluster_m grid_y = num_tile_n_host * cluster_n grid_shape = (grid_x, grid_y, batch) - _kernel( + launch = _kernel( problem_size[0], problem_size[1], problem_size[2], @@ -1353,13 +1370,28 @@ def _host( # @@TMA_STORE_ONLY:BEGIN@@ # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ - ).launch( - grid=grid_shape, - block=(threads_per_cta, 1, 1), - cluster=cluster_shape_mnk, - use_pdl=USE_PDL, - stream=stream, ) + # Mixed CGA: `cluster` is the preferred (wide) shape and `fallback_cluster` + # the regular one the device groups blocks into when a preferred cluster does + # not fit. The grid is already a multiple of the preferred shape, which the + # driver requires. + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) + else: + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + fallback_cluster=fallback_cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) @lru_cache(maxsize=None) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma_static.py b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma_static.py index 8d9b50a5b..ea1f16861 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma_static.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_block_scale_matmul_2ctamma_static.py @@ -127,8 +127,27 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - cluster_m = cluster_shape_mnk[0] - cluster_n = cluster_shape_mnk[1] + # Mixed CGA: the launch carries a preferred (wide) cluster plus a smaller + # fallback one, and the device picks per cluster — a CTA can only tell which + # by reading the hardware cluster dims. Everything cluster-shaped below then + # follows from those, so the two kinds share one body; only the multicast bit + # patterns are loop-built and come in precomputed per shape. + a_mcast_pattern = mixed_a_pattern_pref + b_mcast_pattern = mixed_b_pattern_pref + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + cluster_m = cluster_shape_mnk[0] + cluster_n = cluster_shape_mnk[1] + else: + cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() + cluster_m = cdim_x + cluster_n = cdim_y + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_pref) + b_mcast_pattern = cutlass.Int32(mixed_b_pattern_pref) + # Bitwise, not `or`: both operands are runtime Booleans (this is the form + # cutlass.cute.experimental.is_preferred_cluster uses). + if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_fb) + b_mcast_pattern = cutlass.Int32(mixed_b_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -163,19 +182,12 @@ def _kernel( ) init_tile_l = bidz - a_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_pattern = a_pattern | (1 << (n_idx * cluster_m)) - b_pattern = 0 - for pm_idx in cutlass.range_constexpr(cluster_m // 2): - b_pattern = b_pattern | (1 << (pm_idx * 2)) - if cutlass.const_expr(multicast_a): - tma_mcast_mask_a = cutlass.Int16(a_pattern << m_rank) + tma_mcast_mask_a = cutlass.Int16(a_mcast_pattern << m_rank) else: tma_mcast_mask_a = cutlass.Int16(1 << cta_rank_in_cluster) if cutlass.const_expr(multicast_b): - tma_mcast_mask_b = cutlass.Int16((b_pattern << pair_member) << (n_rank * cluster_m)) + tma_mcast_mask_b = cutlass.Int16((b_mcast_pattern << pair_member) << (n_rank * cluster_m)) else: tma_mcast_mask_b = cutlass.Int16(1 << cta_rank_in_cluster) @@ -264,7 +276,11 @@ def _kernel( sB_bytes = sB_elems * (ab_dtype.width // 8) num_tma_copy_bytes = (num_a_operands * (sA_bytes + sfa_smem_bytes) + num_b_operands * (sB_bytes + sfb_smem_bytes)) * 2 - pair_n_size = cgrp_tile_mnk[1] // cluster_n + # Per-CTA logical tile — the cluster cancels out, so these stay compile-time + # constants even when the cluster shape is only known at runtime. + logical_cta_tile_m = cgrp_tile_mnk[0] // cluster_shape_mnk[0] + logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_shape_mnk[1] + pair_n_size = logical_cta_tile_n # Per-CTA output rows one MMA-M block covers. The pair splits M, so this is # the per-CTA mma_inst_m — half the instruction's hardware M. epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m @@ -285,9 +301,12 @@ def _kernel( M = m N = n - num_tile_m = cute.ceil_div(M, cgrp_tile_mnk[0]) - num_tile_n = cute.ceil_div(N, cgrp_tile_mnk[1]) num_k_tiles = cute.ceil_div(k, cgrp_tile_mnk[2]) + # The tile this cluster owns spans its OWN cluster shape; both shapes walk + # the grid as the identity map (tile == blockIdx), so they tile the problem + # identically and every output tile is still covered exactly once. + cgrp_tile_m_cur = logical_cta_tile_m * cluster_m + cgrp_tile_n_cur = logical_cta_tile_n * cluster_n if warp_idx == scheduler_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -305,9 +324,8 @@ def _kernel( tile_iter = cutlass.Int32(0) is_valid = cutlass.Int32(1) while is_valid != 0: - logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_n - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_per_cta = tile_n * cgrp_tile_n_cur + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -327,7 +345,7 @@ def _kernel( coord_k = k_tile_idx * cgrp_tile_mnk[2] coord_sf_k = k_tile_idx * sf_tma_box_k - coord_n_pair = tile_n * cgrp_tile_mnk[1] + n_rank * logical_cta_tile_n + coord_n_pair = tile_n * cgrp_tile_n_cur + n_rank * logical_cta_tile_n sfb_n_block = coord_n_pair // 128 if is_pair_leader: @@ -526,12 +544,11 @@ def _kernel( pass pair_mask = cutlass.Int16(3) << pair_leader_rank - a_arrive_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_arrive_pattern = a_arrive_pattern | (1 << (n_idx * cluster_m)) - b_arrive_pattern = 0 - for m_idx in cutlass.range_constexpr(cluster_m): - b_arrive_pattern = b_arrive_pattern | (1 << m_idx) + a_arrive_pattern = a_mcast_pattern + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + b_arrive_pattern = (1 << cluster_m) - 1 + else: + b_arrive_pattern = (cutlass.Int32(1) << cluster_m) - 1 a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) @@ -818,8 +835,8 @@ def _kernel( # @@TMA_STORE_ONLY:END@@ while is_valid != 0: - coord_m_tile = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_c = tile_n * cgrp_tile_mnk[1] + n_rank * pair_n_size + coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_c = tile_n * cgrp_tile_n_cur + n_rank * pair_n_size if cutlass.const_expr(epi_rows_per_mma_m == 64): coord_n_c = coord_n_c + (warp_idx // 2) * cols_per_acc_stage @@ -1180,7 +1197,7 @@ def _host( grid_x = num_tile_m_host * cluster_m grid_y = num_tile_n_host * cluster_n grid_shape = (grid_x, grid_y, batch) - _kernel( + launch = _kernel( problem_size[0], problem_size[1], problem_size[2], @@ -1191,13 +1208,28 @@ def _host( # @@TMA_STORE_ONLY:BEGIN@@ # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ - ).launch( - grid=grid_shape, - block=(threads_per_cta, 1, 1), - cluster=cluster_shape_mnk, - use_pdl=USE_PDL, - stream=stream, ) + # Mixed CGA: `cluster` is the preferred (wide) shape and `fallback_cluster` + # the regular one the device groups blocks into when a preferred cluster does + # not fit. The grid is already a multiple of the preferred shape, which the + # driver requires. + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) + else: + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + fallback_cluster=fallback_cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) @lru_cache(maxsize=None) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py index 835373875..db413c039 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py @@ -118,8 +118,24 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - cluster_m = cluster_shape_mnk[0] - cluster_n = cluster_shape_mnk[1] + # Mixed CGA: the launch carries a preferred (wide) cluster plus a smaller + # fallback one, and the device picks per cluster — a CTA can only tell which + # by reading the hardware cluster dims. Everything cluster-shaped below then + # follows from those, so the two kinds share one body; only the multicast bit + # pattern is loop-built and comes in precomputed per shape. + a_mcast_pattern = mixed_a_pattern_pref + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + cluster_m = cluster_shape_mnk[0] + cluster_n = cluster_shape_mnk[1] + else: + cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() + cluster_m = cdim_x + cluster_n = cdim_y + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_pref) + # Bitwise, not `or`: both operands are runtime Booleans (this is the form + # cutlass.cute.experimental.is_preferred_cluster uses). + if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -128,8 +144,6 @@ def _kernel( is_cluster_leader_cta = cta_rank_in_cluster == 0 - full_cluster_mask = cutlass.Int16((1 << cluster_size) - 1) - if warp_idx == mma_warp_id: for _i in cutlass.range_constexpr(num_a_operands): nvvm.prefetch_tensormap(tma_a_descs[_i].get_ptr()) @@ -150,10 +164,11 @@ def _kernel( ) init_tile_l = bidz - a_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_pattern = a_pattern | (1 << (n_idx * cluster_m)) - b_pattern = (1 << cluster_m) - 1 + a_pattern = a_mcast_pattern + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + b_pattern = (1 << cluster_m) - 1 + else: + b_pattern = (cutlass.Int32(1) << cluster_m) - 1 if cutlass.const_expr(multicast_a): tma_mcast_mask_a = cutlass.Int16(a_pattern) << m_rank @@ -282,10 +297,12 @@ def _kernel( M = m N = n - num_tile_m = cute.ceil_div(M, cgrp_tile_mnk[0]) - num_tile_n = cute.ceil_div(N, cgrp_tile_mnk[1]) - total_tiles = num_tile_m * num_tile_n num_k_tiles = cute.ceil_div(k, cta_tile_mnk[2]) + # The tile this cluster owns spans its OWN cluster shape; both shapes walk + # the grid as the identity map (tile == blockIdx), so they tile the problem + # identically and every output tile is still covered exactly once. + cgrp_tile_m_cur = cta_tile_mnk[0] * cluster_m + cgrp_tile_n_cur = cta_tile_mnk[1] * cluster_n num_k_blocks = cta_tile_mnk[2] // mma_inst_shape_mnk[2] if warp_idx == scheduler_warp_id: @@ -357,8 +374,8 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_tma = cutlass.Int32(0) while is_valid != 0: - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_per_cta = tile_n * cgrp_tile_n_cur + n_rank * cta_tile_mnk[1] if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -734,8 +751,8 @@ def _kernel( # @@TMA_STORE_ONLY:END@@ while is_valid != 0: - coord_m_tile = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] + coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n = tile_n * cgrp_tile_n_cur + n_rank * cta_tile_mnk[1] acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: @@ -1074,7 +1091,7 @@ def _host( grid_x = num_tile_m_host * cluster_m grid_y = num_tile_n_host * cluster_n grid_shape = (grid_x, grid_y, batch) - _kernel( + launch = _kernel( problem_size[0], problem_size[1], problem_size[2], @@ -1085,13 +1102,28 @@ def _host( # @@TMA_STORE_ONLY:BEGIN@@ # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ - ).launch( - grid=grid_shape, - block=(threads_per_cta, 1, 1), - cluster=cluster_shape_mnk, - use_pdl=USE_PDL, - stream=stream, ) + # Mixed CGA: `cluster` is the preferred (wide) shape and `fallback_cluster` + # the regular one the device groups blocks into when a preferred cluster does + # not fit. The grid is already a multiple of the preferred shape, which the + # driver requires. + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) + else: + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + fallback_cluster=fallback_cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) @lru_cache(maxsize=None) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma_static.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma_static.py index d32b5c6ae..ac0d6876d 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma_static.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma_static.py @@ -115,8 +115,24 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - cluster_m = cluster_shape_mnk[0] - cluster_n = cluster_shape_mnk[1] + # Mixed CGA: the launch carries a preferred (wide) cluster plus a smaller + # fallback one, and the device picks per cluster — a CTA can only tell which + # by reading the hardware cluster dims. Everything cluster-shaped below then + # follows from those, so the two kinds share one body; only the multicast bit + # pattern is loop-built and comes in precomputed per shape. + a_mcast_pattern = mixed_a_pattern_pref + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + cluster_m = cluster_shape_mnk[0] + cluster_n = cluster_shape_mnk[1] + else: + cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() + cluster_m = cdim_x + cluster_n = cdim_y + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_pref) + # Bitwise, not `or`: both operands are runtime Booleans (this is the form + # cutlass.cute.experimental.is_preferred_cluster uses). + if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -125,8 +141,6 @@ def _kernel( is_cluster_leader_cta = cta_rank_in_cluster == 0 - full_cluster_mask = cutlass.Int16((1 << cluster_size) - 1) - if warp_idx == mma_warp_id: for _i in cutlass.range_constexpr(num_a_operands): nvvm.prefetch_tensormap(tma_a_descs[_i].get_ptr()) @@ -147,10 +161,11 @@ def _kernel( ) init_tile_l = bidz - a_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_pattern = a_pattern | (1 << (n_idx * cluster_m)) - b_pattern = (1 << cluster_m) - 1 + a_pattern = a_mcast_pattern + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + b_pattern = (1 << cluster_m) - 1 + else: + b_pattern = (cutlass.Int32(1) << cluster_m) - 1 if cutlass.const_expr(multicast_a): tma_mcast_mask_a = cutlass.Int16(a_pattern) << m_rank @@ -255,10 +270,12 @@ def _kernel( vsize = (VEC_BYTES * 8) // cd_dtype.width M = m N = n - num_tile_m = cute.ceil_div(M, cgrp_tile_mnk[0]) - num_tile_n = cute.ceil_div(N, cgrp_tile_mnk[1]) - total_tiles = num_tile_m * num_tile_n num_k_tiles = cute.ceil_div(k, cta_tile_mnk[2]) + # The tile this cluster owns spans its OWN cluster shape; both shapes walk + # the grid as the identity map (tile == blockIdx), so they tile the problem + # identically and every output tile is still covered exactly once. + cgrp_tile_m_cur = cta_tile_mnk[0] * cluster_m + cgrp_tile_n_cur = cta_tile_mnk[1] * cluster_n num_k_blocks = cta_tile_mnk[2] // mma_inst_shape_mnk[2] if warp_idx == scheduler_warp_id: @@ -277,8 +294,8 @@ def _kernel( tile_iter = cutlass.Int32(0) is_valid = cutlass.Int32(1) while is_valid != 0: - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_per_cta = tile_n * cgrp_tile_n_cur + n_rank * cta_tile_mnk[1] if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -609,8 +626,8 @@ def _kernel( # @@TMA_STORE_ONLY:END@@ while is_valid != 0: - coord_m_tile = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] + coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n = tile_n * cgrp_tile_n_cur + n_rank * cta_tile_mnk[1] acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: @@ -925,7 +942,7 @@ def _host( grid_x = num_tile_m_host * cluster_m grid_y = num_tile_n_host * cluster_n grid_shape = (grid_x, grid_y, batch) - _kernel( + launch = _kernel( problem_size[0], problem_size[1], problem_size[2], @@ -936,13 +953,28 @@ def _host( # @@TMA_STORE_ONLY:BEGIN@@ # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ - ).launch( - grid=grid_shape, - block=(threads_per_cta, 1, 1), - cluster=cluster_shape_mnk, - use_pdl=USE_PDL, - stream=stream, ) + # Mixed CGA: `cluster` is the preferred (wide) shape and `fallback_cluster` + # the regular one the device groups blocks into when a preferred cluster does + # not fit. The grid is already a multiple of the preferred shape, which the + # driver requires. + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) + else: + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + fallback_cluster=fallback_cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) @lru_cache(maxsize=None) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py index 65ad6b0eb..892ac691b 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py @@ -116,8 +116,27 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - cluster_m = cluster_shape_mnk[0] - cluster_n = cluster_shape_mnk[1] + # Mixed CGA: the launch carries a preferred (wide) cluster plus a smaller + # fallback one, and the device picks per cluster — a CTA can only tell which + # by reading the hardware cluster dims. Everything cluster-shaped below then + # follows from those, so the two kinds share one body; only the multicast bit + # patterns are loop-built and come in precomputed per shape. + a_mcast_pattern = mixed_a_pattern_pref + b_mcast_pattern = mixed_b_pattern_pref + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + cluster_m = cluster_shape_mnk[0] + cluster_n = cluster_shape_mnk[1] + else: + cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() + cluster_m = cdim_x + cluster_n = cdim_y + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_pref) + b_mcast_pattern = cutlass.Int32(mixed_b_pattern_pref) + # Bitwise, not `or`: both operands are runtime Booleans (this is the form + # cutlass.cute.experimental.is_preferred_cluster uses). + if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_fb) + b_mcast_pattern = cutlass.Int32(mixed_b_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -150,19 +169,12 @@ def _kernel( ) init_tile_l = bidz - a_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_pattern = a_pattern | (1 << (n_idx * cluster_m)) - b_pattern = 0 - for pm_idx in cutlass.range_constexpr(cluster_m // 2): - b_pattern = b_pattern | (1 << (pm_idx * 2)) - if cutlass.const_expr(multicast_a): - tma_mcast_mask_a = cutlass.Int16(a_pattern << m_rank) + tma_mcast_mask_a = cutlass.Int16(a_mcast_pattern << m_rank) else: tma_mcast_mask_a = cutlass.Int16(1 << cta_rank_in_cluster) if cutlass.const_expr(multicast_b): - tma_mcast_mask_b = cutlass.Int16((b_pattern << pair_member) << (n_rank * cluster_m)) + tma_mcast_mask_b = cutlass.Int16((b_mcast_pattern << pair_member) << (n_rank * cluster_m)) else: tma_mcast_mask_b = cutlass.Int16(1 << cta_rank_in_cluster) @@ -263,7 +275,11 @@ def _kernel( ) # TMEM: per-pair (CTA_2 group), double-buffered along N. - pair_n_size = cgrp_tile_mnk[1] // cluster_n + # Per-CTA logical tile — the cluster cancels out, so these stay compile-time + # constants even when the cluster shape is only known at runtime. + logical_cta_tile_m = cgrp_tile_mnk[0] // cluster_shape_mnk[0] + logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_shape_mnk[1] + pair_n_size = logical_cta_tile_n # Per-CTA output rows one MMA-M block covers. The pair splits M, so this is # the per-CTA mma_inst_m — half the instruction's hardware M. epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m @@ -289,10 +305,12 @@ def _kernel( vsize = (VEC_BYTES * 8) // cd_dtype.width M = m N = n - num_tile_m = cute.ceil_div(M, cgrp_tile_mnk[0]) - num_tile_n = cute.ceil_div(N, cgrp_tile_mnk[1]) - total_tiles = num_tile_m * num_tile_n num_k_tiles = cute.ceil_div(k, cgrp_tile_mnk[2]) + # The tile this cluster owns spans its OWN cluster shape; both shapes walk + # the grid as the identity map (tile == blockIdx), so they tile the problem + # identically and every output tile is still covered exactly once. + cgrp_tile_m_cur = logical_cta_tile_m * cluster_m + cgrp_tile_n_cur = logical_cta_tile_n * cluster_n num_k_blocks = cgrp_tile_mnk[2] // mma_inst_shape_mnk[2] if warp_idx == scheduler_warp_id: @@ -364,9 +382,8 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_tma = cutlass.Int32(0) while is_valid != 0: - logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_n - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_per_cta = tile_n * cgrp_tile_n_cur + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -553,12 +570,11 @@ def _kernel( pass pair_mask = cutlass.Int16(3) << pair_leader_rank - a_arrive_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_arrive_pattern = a_arrive_pattern | (1 << (n_idx * cluster_m)) - b_arrive_pattern = 0 - for m_idx in cutlass.range_constexpr(cluster_m): - b_arrive_pattern = b_arrive_pattern | (1 << m_idx) + a_arrive_pattern = a_mcast_pattern + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + b_arrive_pattern = (1 << cluster_m) - 1 + else: + b_arrive_pattern = (cutlass.Int32(1) << cluster_m) - 1 a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) @@ -794,8 +810,8 @@ def _kernel( # @@TMA_STORE_ONLY:END@@ while is_valid != 0: - coord_m_tile = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_c = tile_n * cgrp_tile_mnk[1] + n_rank * pair_n_size + coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_c = tile_n * cgrp_tile_n_cur + n_rank * pair_n_size if cutlass.const_expr(epi_rows_per_mma_m == 64): coord_n_c = coord_n_c + (warp_idx // 2) * epi_cols_per_mma_m @@ -1140,7 +1156,7 @@ def _host( grid_x = num_tile_m_host * cluster_m grid_y = num_tile_n_host * cluster_n grid_shape = (grid_x, grid_y, batch) - _kernel( + launch = _kernel( problem_size[0], problem_size[1], problem_size[2], @@ -1151,13 +1167,28 @@ def _host( # @@TMA_STORE_ONLY:BEGIN@@ # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ - ).launch( - grid=grid_shape, - block=(threads_per_cta, 1, 1), - cluster=cluster_shape_mnk, - use_pdl=USE_PDL, - stream=stream, ) + # Mixed CGA: `cluster` is the preferred (wide) shape and `fallback_cluster` + # the regular one the device groups blocks into when a preferred cluster does + # not fit. The grid is already a multiple of the preferred shape, which the + # driver requires. + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) + else: + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + fallback_cluster=fallback_cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) @lru_cache(maxsize=None) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma_static.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma_static.py index 987cd5181..6a4a8facb 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma_static.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma_static.py @@ -113,8 +113,27 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - cluster_m = cluster_shape_mnk[0] - cluster_n = cluster_shape_mnk[1] + # Mixed CGA: the launch carries a preferred (wide) cluster plus a smaller + # fallback one, and the device picks per cluster — a CTA can only tell which + # by reading the hardware cluster dims. Everything cluster-shaped below then + # follows from those, so the two kinds share one body; only the multicast bit + # patterns are loop-built and come in precomputed per shape. + a_mcast_pattern = mixed_a_pattern_pref + b_mcast_pattern = mixed_b_pattern_pref + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + cluster_m = cluster_shape_mnk[0] + cluster_n = cluster_shape_mnk[1] + else: + cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() + cluster_m = cdim_x + cluster_n = cdim_y + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_pref) + b_mcast_pattern = cutlass.Int32(mixed_b_pattern_pref) + # Bitwise, not `or`: both operands are runtime Booleans (this is the form + # cutlass.cute.experimental.is_preferred_cluster uses). + if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_fb) + b_mcast_pattern = cutlass.Int32(mixed_b_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -147,19 +166,12 @@ def _kernel( ) init_tile_l = bidz - a_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_pattern = a_pattern | (1 << (n_idx * cluster_m)) - b_pattern = 0 - for pm_idx in cutlass.range_constexpr(cluster_m // 2): - b_pattern = b_pattern | (1 << (pm_idx * 2)) - if cutlass.const_expr(multicast_a): - tma_mcast_mask_a = cutlass.Int16(a_pattern << m_rank) + tma_mcast_mask_a = cutlass.Int16(a_mcast_pattern << m_rank) else: tma_mcast_mask_a = cutlass.Int16(1 << cta_rank_in_cluster) if cutlass.const_expr(multicast_b): - tma_mcast_mask_b = cutlass.Int16((b_pattern << pair_member) << (n_rank * cluster_m)) + tma_mcast_mask_b = cutlass.Int16((b_mcast_pattern << pair_member) << (n_rank * cluster_m)) else: tma_mcast_mask_b = cutlass.Int16(1 << cta_rank_in_cluster) @@ -237,7 +249,11 @@ def _kernel( b_major=mma_b_major, ) - pair_n_size = cgrp_tile_mnk[1] // cluster_n + # Per-CTA logical tile — the cluster cancels out, so these stay compile-time + # constants even when the cluster shape is only known at runtime. + logical_cta_tile_m = cgrp_tile_mnk[0] // cluster_shape_mnk[0] + logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_shape_mnk[1] + pair_n_size = logical_cta_tile_n # Per-CTA output rows one MMA-M block covers. The pair splits M, so this is # the per-CTA mma_inst_m — half the instruction's hardware M. epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m @@ -263,10 +279,12 @@ def _kernel( vsize = (VEC_BYTES * 8) // cd_dtype.width M = m N = n - num_tile_m = cute.ceil_div(M, cgrp_tile_mnk[0]) - num_tile_n = cute.ceil_div(N, cgrp_tile_mnk[1]) - total_tiles = num_tile_m * num_tile_n num_k_tiles = cute.ceil_div(k, cgrp_tile_mnk[2]) + # The tile this cluster owns spans its OWN cluster shape; both shapes walk + # the grid as the identity map (tile == blockIdx), so they tile the problem + # identically and every output tile is still covered exactly once. + cgrp_tile_m_cur = logical_cta_tile_m * cluster_m + cgrp_tile_n_cur = logical_cta_tile_n * cluster_n num_k_blocks = cgrp_tile_mnk[2] // mma_inst_shape_mnk[2] if warp_idx == scheduler_warp_id: @@ -285,9 +303,8 @@ def _kernel( tile_iter = cutlass.Int32(0) is_valid = cutlass.Int32(1) while is_valid != 0: - logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_n - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_per_cta = tile_n * cgrp_tile_n_cur + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -451,12 +468,11 @@ def _kernel( pass pair_mask = cutlass.Int16(3) << pair_leader_rank - a_arrive_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_arrive_pattern = a_arrive_pattern | (1 << (n_idx * cluster_m)) - b_arrive_pattern = 0 - for m_idx in cutlass.range_constexpr(cluster_m): - b_arrive_pattern = b_arrive_pattern | (1 << m_idx) + a_arrive_pattern = a_mcast_pattern + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + b_arrive_pattern = (1 << cluster_m) - 1 + else: + b_arrive_pattern = (cutlass.Int32(1) << cluster_m) - 1 a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) @@ -655,8 +671,8 @@ def _kernel( # @@TMA_STORE_ONLY:END@@ while is_valid != 0: - coord_m_tile = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_c = tile_n * cgrp_tile_mnk[1] + n_rank * pair_n_size + coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_c = tile_n * cgrp_tile_n_cur + n_rank * pair_n_size if cutlass.const_expr(epi_rows_per_mma_m == 64): coord_n_c = coord_n_c + (warp_idx // 2) * epi_cols_per_mma_m @@ -977,7 +993,7 @@ def _host( grid_x = num_tile_m_host * cluster_m grid_y = num_tile_n_host * cluster_n grid_shape = (grid_x, grid_y, batch) - _kernel( + launch = _kernel( problem_size[0], problem_size[1], problem_size[2], @@ -988,13 +1004,28 @@ def _host( # @@TMA_STORE_ONLY:BEGIN@@ # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ - ).launch( - grid=grid_shape, - block=(threads_per_cta, 1, 1), - cluster=cluster_shape_mnk, - use_pdl=USE_PDL, - stream=stream, ) + # Mixed CGA: `cluster` is the preferred (wide) shape and `fallback_cluster` + # the regular one the device groups blocks into when a preferred cluster does + # not fit. The grid is already a multiple of the preferred shape, which the + # driver requires. + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) + else: + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + fallback_cluster=fallback_cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) @lru_cache(maxsize=None) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_1ctamma.py index 96f1dc1c5..3d55f8a99 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_1ctamma.py @@ -119,8 +119,24 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - cluster_m = cluster_shape_mnk[0] - cluster_n = cluster_shape_mnk[1] + # Mixed CGA: the launch carries a preferred (wide) cluster plus a smaller + # fallback one, and the device picks per cluster — a CTA can only tell which + # by reading the hardware cluster dims. Everything cluster-shaped below then + # follows from those, so the two kinds share one body; only the multicast bit + # pattern is loop-built and comes in precomputed per shape. + a_mcast_pattern = mixed_a_pattern_pref + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + cluster_m = cluster_shape_mnk[0] + cluster_n = cluster_shape_mnk[1] + else: + cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() + cluster_m = cdim_x + cluster_n = cdim_y + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_pref) + # Bitwise, not `or`: both operands are runtime Booleans (this is the form + # cutlass.cute.experimental.is_preferred_cluster uses). + if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -129,8 +145,6 @@ def _kernel( is_cluster_leader_cta = cta_rank_in_cluster == 0 - full_cluster_mask = cutlass.Int16((1 << cluster_size) - 1) - if warp_idx == mma_warp_id: nvvm.prefetch_tensormap(tma_a_desc.get_ptr()) nvvm.prefetch_tensormap(tma_b_desc.get_ptr()) @@ -149,10 +163,11 @@ def _kernel( ) init_tile_l = bidz - a_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_pattern = a_pattern | (1 << (n_idx * cluster_m)) - b_pattern = (1 << cluster_m) - 1 + a_pattern = a_mcast_pattern + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + b_pattern = (1 << cluster_m) - 1 + else: + b_pattern = (cutlass.Int32(1) << cluster_m) - 1 if cutlass.const_expr(multicast_a): tma_mcast_mask_a = cutlass.Int16(a_pattern) << m_rank @@ -289,10 +304,12 @@ def _kernel( M = m N = n - num_tile_m = cute.ceil_div(M, cgrp_tile_mnk[0]) - num_tile_n = cute.ceil_div(N, cgrp_tile_mnk[1]) - total_tiles = num_tile_m * num_tile_n num_k_tiles = cute.ceil_div(k, cta_tile_mnk[2]) + # The tile this cluster owns spans its OWN cluster shape; both shapes walk + # the grid as the identity map (tile == blockIdx), so they tile the problem + # identically and every output tile is still covered exactly once. + cgrp_tile_m_cur = cta_tile_mnk[0] * cluster_m + cgrp_tile_n_cur = cta_tile_mnk[1] * cluster_n num_k_blocks = cta_tile_mnk[2] // mma_inst_shape_mnk[2] if warp_idx == scheduler_warp_id: @@ -364,8 +381,8 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_tma = cutlass.Int32(0) while is_valid != 0: - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_per_cta = tile_n * cgrp_tile_n_cur + n_rank * cta_tile_mnk[1] if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -957,8 +974,8 @@ def _kernel( # @@TMA_STORE_ONLY:END@@ while is_valid != 0: - coord_m_tile = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] + coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n = tile_n * cgrp_tile_n_cur + n_rank * cta_tile_mnk[1] acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: @@ -1278,7 +1295,7 @@ def _host( grid_x = num_tile_m_host * cluster_m grid_y = num_tile_n_host * cluster_n grid_shape = (grid_x, grid_y, batch) - _kernel( + launch = _kernel( problem_size[0], problem_size[1], problem_size[2], @@ -1289,13 +1306,28 @@ def _host( # @@TMA_STORE_ONLY:BEGIN@@ # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ - ).launch( - grid=grid_shape, - block=(threads_per_cta, 1, 1), - cluster=cluster_shape_mnk, - use_pdl=USE_PDL, - stream=stream, ) + # Mixed CGA: `cluster` is the preferred (wide) shape and `fallback_cluster` + # the regular one the device groups blocks into when a preferred cluster does + # not fit. The grid is already a multiple of the preferred shape, which the + # driver requires. + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) + else: + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + fallback_cluster=fallback_cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) @lru_cache(maxsize=None) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_2ctamma.py index f89a9f569..190ee4d49 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_mainloop_2ctamma.py @@ -121,8 +121,27 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - cluster_m = cluster_shape_mnk[0] - cluster_n = cluster_shape_mnk[1] + # Mixed CGA: the launch carries a preferred (wide) cluster plus a smaller + # fallback one, and the device picks per cluster — a CTA can only tell which + # by reading the hardware cluster dims. Everything cluster-shaped below then + # follows from those, so the two kinds share one body; only the multicast bit + # patterns are loop-built and come in precomputed per shape. + a_mcast_pattern = mixed_a_pattern_pref + b_mcast_pattern = mixed_b_pattern_pref + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + cluster_m = cluster_shape_mnk[0] + cluster_n = cluster_shape_mnk[1] + else: + cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() + cluster_m = cdim_x + cluster_n = cdim_y + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_pref) + b_mcast_pattern = cutlass.Int32(mixed_b_pattern_pref) + # Bitwise, not `or`: both operands are runtime Booleans (this is the form + # cutlass.cute.experimental.is_preferred_cluster uses). + if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_fb) + b_mcast_pattern = cutlass.Int32(mixed_b_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -153,19 +172,12 @@ def _kernel( ) init_tile_l = bidz - a_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_pattern = a_pattern | (1 << (n_idx * cluster_m)) - b_pattern = 0 - for pm_idx in cutlass.range_constexpr(cluster_m // 2): - b_pattern = b_pattern | (1 << (pm_idx * 2)) - if cutlass.const_expr(multicast_a): - tma_mcast_mask_a = cutlass.Int16(a_pattern << m_rank) + tma_mcast_mask_a = cutlass.Int16(a_mcast_pattern << m_rank) else: tma_mcast_mask_a = cutlass.Int16(1 << cta_rank_in_cluster) if cutlass.const_expr(multicast_b): - tma_mcast_mask_b = cutlass.Int16((b_pattern << pair_member) << (n_rank * cluster_m)) + tma_mcast_mask_b = cutlass.Int16((b_mcast_pattern << pair_member) << (n_rank * cluster_m)) else: tma_mcast_mask_b = cutlass.Int16(1 << cta_rank_in_cluster) @@ -277,7 +289,11 @@ def _kernel( b_major=mma_b_major, ) - pair_n_size = cgrp_tile_mnk[1] // cluster_n + # Per-CTA logical tile — the cluster cancels out, so these stay compile-time + # constants even when the cluster shape is only known at runtime. + logical_cta_tile_m = cgrp_tile_mnk[0] // cluster_shape_mnk[0] + logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_shape_mnk[1] + pair_n_size = logical_cta_tile_n # Per-CTA output rows one MMA-M block covers. The pair splits M, so this is # the per-CTA mma_inst_m — half the instruction's hardware M. epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m @@ -304,10 +320,12 @@ def _kernel( M = m N = n - num_tile_m = cute.ceil_div(M, cgrp_tile_mnk[0]) - num_tile_n = cute.ceil_div(N, cgrp_tile_mnk[1]) - total_tiles = num_tile_m * num_tile_n num_k_tiles = cute.ceil_div(k, cgrp_tile_mnk[2]) + # The tile this cluster owns spans its OWN cluster shape; both shapes walk + # the grid as the identity map (tile == blockIdx), so they tile the problem + # identically and every output tile is still covered exactly once. + cgrp_tile_m_cur = logical_cta_tile_m * cluster_m + cgrp_tile_n_cur = logical_cta_tile_n * cluster_n num_k_blocks = cgrp_tile_mnk[2] // mma_inst_shape_mnk[2] if warp_idx == scheduler_warp_id: @@ -380,9 +398,8 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_tma = cutlass.Int32(0) while is_valid != 0: - logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_n - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_per_cta = tile_n * cgrp_tile_n_cur + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -642,12 +659,11 @@ def _kernel( pass pair_mask = cutlass.Int16(3) << pair_leader_rank - a_arrive_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_arrive_pattern = a_arrive_pattern | (1 << (n_idx * cluster_m)) - b_arrive_pattern = 0 - for m_idx in cutlass.range_constexpr(cluster_m): - b_arrive_pattern = b_arrive_pattern | (1 << m_idx) + a_arrive_pattern = a_mcast_pattern + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + b_arrive_pattern = (1 << cluster_m) - 1 + else: + b_arrive_pattern = (cutlass.Int32(1) << cluster_m) - 1 a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) @@ -1087,8 +1103,8 @@ def _kernel( # @@TMA_STORE_ONLY:END@@ while is_valid != 0: - coord_m_tile = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_c = tile_n * cgrp_tile_mnk[1] + n_rank * pair_n_size + coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_c = tile_n * cgrp_tile_n_cur + n_rank * pair_n_size if cutlass.const_expr(epi_rows_per_mma_m == 64): coord_n_c = coord_n_c + (warp_idx // 2) * epi_cols_per_mma_m @@ -1415,7 +1431,7 @@ def _host( grid_x = num_tile_m_host * cluster_m grid_y = num_tile_n_host * cluster_n grid_shape = (grid_x, grid_y, batch) - _kernel( + launch = _kernel( problem_size[0], problem_size[1], problem_size[2], @@ -1426,13 +1442,28 @@ def _host( # @@TMA_STORE_ONLY:BEGIN@@ # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ - ).launch( - grid=grid_shape, - block=(threads_per_cta, 1, 1), - cluster=cluster_shape_mnk, - use_pdl=USE_PDL, - stream=stream, ) + # Mixed CGA: `cluster` is the preferred (wide) shape and `fallback_cluster` + # the regular one the device groups blocks into when a preferred cluster does + # not fit. The grid is already a multiple of the preferred shape, which the + # driver requires. + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) + else: + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + fallback_cluster=fallback_cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) @lru_cache(maxsize=None) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_1ctamma.py index 22dc5a5ae..b09a85c34 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_1ctamma.py @@ -195,8 +195,24 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - cluster_m = cluster_shape_mnk[0] - cluster_n = cluster_shape_mnk[1] + # Mixed CGA: the launch carries a preferred (wide) cluster plus a smaller + # fallback one, and the device picks per cluster — a CTA can only tell which + # by reading the hardware cluster dims. Everything cluster-shaped below then + # follows from those, so the two kinds share one body; only the multicast bit + # pattern is loop-built and comes in precomputed per shape. + a_mcast_pattern = mixed_a_pattern_pref + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + cluster_m = cluster_shape_mnk[0] + cluster_n = cluster_shape_mnk[1] + else: + cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() + cluster_m = cdim_x + cluster_n = cdim_y + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_pref) + # Bitwise, not `or`: both operands are runtime Booleans (this is the form + # cutlass.cute.experimental.is_preferred_cluster uses). + if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -205,8 +221,6 @@ def _kernel( is_cluster_leader_cta = cta_rank_in_cluster == 0 - full_cluster_mask = cutlass.Int16((1 << cluster_size) - 1) - if warp_idx == mma_warp_id: for _i in cutlass.range_constexpr(num_a_operands): nvvm.prefetch_tensormap(tma_a_descs[_i].get_ptr()) @@ -229,10 +243,11 @@ def _kernel( ) init_tile_l = bidz - a_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_pattern = a_pattern | (1 << (n_idx * cluster_m)) - b_pattern = (1 << cluster_m) - 1 + a_pattern = a_mcast_pattern + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + b_pattern = (1 << cluster_m) - 1 + else: + b_pattern = (cutlass.Int32(1) << cluster_m) - 1 if cutlass.const_expr(multicast_a): tma_mcast_mask_a = cutlass.Int16(a_pattern) << m_rank @@ -370,9 +385,12 @@ def _kernel( M = m N = n - num_tile_m = cute.ceil_div(M, cgrp_tile_mnk[0]) - num_tile_n = cute.ceil_div(N, cgrp_tile_mnk[1]) num_k_tiles = cute.ceil_div(k, cta_tile_mnk[2]) + # The tile this cluster owns spans its OWN cluster shape; both shapes walk + # the grid as the identity map (tile == blockIdx), so they tile the problem + # identically and every output tile is still covered exactly once. + cgrp_tile_m_cur = cta_tile_mnk[0] * cluster_m + cgrp_tile_n_cur = cta_tile_mnk[1] * cluster_n if warp_idx == scheduler_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -443,8 +461,8 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_tma = cutlass.Int32(0) while is_valid != 0: - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_per_cta = tile_n * cgrp_tile_n_cur + n_rank * cta_tile_mnk[1] if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -576,8 +594,8 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_sf = cutlass.Int32(0) while is_valid != 0: - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_per_cta = tile_n * cgrp_tile_n_cur + n_rank * cta_tile_mnk[1] if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -1079,8 +1097,8 @@ def _kernel( # @@TMA_STORE_ONLY:END@@ while is_valid != 0: - coord_m_tile = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] + coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n = tile_n * cgrp_tile_n_cur + n_rank * cta_tile_mnk[1] acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: @@ -1424,7 +1442,7 @@ def _host( grid_x = num_tile_m_host * cluster_m grid_y = num_tile_n_host * cluster_n grid_shape = (grid_x, grid_y, batch) - _kernel( + launch = _kernel( problem_size[0], problem_size[1], problem_size[2], @@ -1435,13 +1453,28 @@ def _host( # @@TMA_STORE_ONLY:BEGIN@@ # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ - ).launch( - grid=grid_shape, - block=(threads_per_cta, 1, 1), - cluster=cluster_shape_mnk, - use_pdl=USE_PDL, - stream=stream, ) + # Mixed CGA: `cluster` is the preferred (wide) shape and `fallback_cluster` + # the regular one the device groups blocks into when a preferred cluster does + # not fit. The grid is already a multiple of the preferred shape, which the + # driver requires. + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) + else: + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + fallback_cluster=fallback_cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) @lru_cache(maxsize=None) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_2ctamma.py index 03ae73089..353f592b6 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm103_block_scale_matmul_2ctamma.py @@ -195,8 +195,27 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - cluster_m = cluster_shape_mnk[0] - cluster_n = cluster_shape_mnk[1] + # Mixed CGA: the launch carries a preferred (wide) cluster plus a smaller + # fallback one, and the device picks per cluster — a CTA can only tell which + # by reading the hardware cluster dims. Everything cluster-shaped below then + # follows from those, so the two kinds share one body; only the multicast bit + # patterns are loop-built and come in precomputed per shape. + a_mcast_pattern = mixed_a_pattern_pref + b_mcast_pattern = mixed_b_pattern_pref + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + cluster_m = cluster_shape_mnk[0] + cluster_n = cluster_shape_mnk[1] + else: + cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() + cluster_m = cdim_x + cluster_n = cdim_y + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_pref) + b_mcast_pattern = cutlass.Int32(mixed_b_pattern_pref) + # Bitwise, not `or`: both operands are runtime Booleans (this is the form + # cutlass.cute.experimental.is_preferred_cluster uses). + if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_fb) + b_mcast_pattern = cutlass.Int32(mixed_b_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -231,19 +250,12 @@ def _kernel( ) init_tile_l = bidz - a_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_pattern = a_pattern | (1 << (n_idx * cluster_m)) - b_pattern = 0 - for pm_idx in cutlass.range_constexpr(cluster_m // 2): - b_pattern = b_pattern | (1 << (pm_idx * 2)) - if cutlass.const_expr(multicast_a): - tma_mcast_mask_a = cutlass.Int16(a_pattern << m_rank) + tma_mcast_mask_a = cutlass.Int16(a_mcast_pattern << m_rank) else: tma_mcast_mask_a = cutlass.Int16(1 << cta_rank_in_cluster) if cutlass.const_expr(multicast_b): - tma_mcast_mask_b = cutlass.Int16((b_pattern << pair_member) << (n_rank * cluster_m)) + tma_mcast_mask_b = cutlass.Int16((b_mcast_pattern << pair_member) << (n_rank * cluster_m)) else: tma_mcast_mask_b = cutlass.Int16(1 << cta_rank_in_cluster) @@ -363,7 +375,11 @@ def _kernel( num_tma_ab_chunk_bytes = (num_a_operands * a_chunk_bytes + num_b_operands * b_chunk_bytes) * 2 num_tma_sf_group_bytes = (num_a_operands * sfa_group_bytes + num_b_operands * sfb_group_bytes) * 2 - pair_n_size = cgrp_tile_mnk[1] // cluster_n + # Per-CTA logical tile — the cluster cancels out, so these stay compile-time + # constants even when the cluster shape is only known at runtime. + logical_cta_tile_m = cgrp_tile_mnk[0] // cluster_shape_mnk[0] + logical_cta_tile_n_c = cgrp_tile_mnk[1] // cluster_shape_mnk[1] + pair_n_size = logical_cta_tile_n_c # Per-CTA output rows one MMA-M block covers. The pair splits M, so this is # the per-CTA mma_inst_m — half the instruction's hardware M. epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m @@ -384,9 +400,12 @@ def _kernel( M = m N = n - num_tile_m = cute.ceil_div(M, cgrp_tile_mnk[0]) - num_tile_n = cute.ceil_div(N, cgrp_tile_mnk[1]) num_k_tiles = cute.ceil_div(k, cta_tile_mnk[2]) + # The tile this cluster owns spans its OWN cluster shape; both shapes walk + # the grid as the identity map (tile == blockIdx), so they tile the problem + # identically and every output tile is still covered exactly once. + cgrp_tile_m_cur = logical_cta_tile_m * cluster_m + cgrp_tile_n_cur = logical_cta_tile_n_c * cluster_n if warp_idx == scheduler_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -457,9 +476,9 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_tma = cutlass.Int32(0) while is_valid != 0: - logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_n - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] + logical_cta_tile_n = logical_cta_tile_n_c + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_per_cta = tile_n * cgrp_tile_n_cur + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -592,10 +611,10 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_sf = cutlass.Int32(0) while is_valid != 0: - logical_cta_tile_n = cgrp_tile_mnk[1] // cluster_n - coord_m_per_cta = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] + logical_cta_tile_n = logical_cta_tile_n_c + coord_m_per_cta = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] # SFB covers the FULL pair-N range per CTA (no pair_member offset). - coord_n_pair = tile_n * cgrp_tile_mnk[1] + n_rank * logical_cta_tile_n + coord_n_pair = tile_n * cgrp_tile_n_cur + n_rank * logical_cta_tile_n if cutlass.const_expr(matmul_a_batch == 1): tile_l_a = cutlass.Int32(0) else: @@ -728,12 +747,11 @@ def _kernel( base_col_id_root = tmem_raw_addr & 0xFFFF base_row_id = tmem_raw_addr >> 16 pair_mask = cutlass.Int16(3) << pair_leader_rank - a_arrive_pattern = 0 - for n_idx in cutlass.range_constexpr(cluster_n): - a_arrive_pattern = a_arrive_pattern | (1 << (n_idx * cluster_m)) - b_arrive_pattern = 0 - for m_idx in cutlass.range_constexpr(cluster_m): - b_arrive_pattern = b_arrive_pattern | (1 << m_idx) + a_arrive_pattern = a_mcast_pattern + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + b_arrive_pattern = (1 << cluster_m) - 1 + else: + b_arrive_pattern = (cutlass.Int32(1) << cluster_m) - 1 a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) @@ -1154,8 +1172,8 @@ def _kernel( # @@TMA_STORE_ONLY:END@@ while is_valid != 0: - coord_m_tile = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] - coord_n_c = tile_n * cgrp_tile_mnk[1] + n_rank * pair_n_size + coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + coord_n_c = tile_n * cgrp_tile_n_cur + n_rank * pair_n_size if cutlass.const_expr(epi_rows_per_mma_m == 64): coord_n_c = coord_n_c + (warp_idx // 2) * cols_per_acc_stage @@ -1505,7 +1523,7 @@ def _host( grid_x = num_tile_m_host * cluster_m grid_y = num_tile_n_host * cluster_n grid_shape = (grid_x, grid_y, batch) - _kernel( + launch = _kernel( problem_size[0], problem_size[1], problem_size[2], @@ -1516,13 +1534,28 @@ def _host( # @@TMA_STORE_ONLY:BEGIN@@ # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ - ).launch( - grid=grid_shape, - block=(threads_per_cta, 1, 1), - cluster=cluster_shape_mnk, - use_pdl=USE_PDL, - stream=stream, ) + # Mixed CGA: `cluster` is the preferred (wide) shape and `fallback_cluster` + # the regular one the device groups blocks into when a preferred cluster does + # not fit. The grid is already a multiple of the preferred shape, which the + # driver requires. + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) + else: + launch.launch( + grid=grid_shape, + block=(threads_per_cta, 1, 1), + cluster=cluster_shape_mnk, + fallback_cluster=fallback_cluster_shape_mnk, + use_pdl=USE_PDL, + stream=stream, + ) @lru_cache(maxsize=None) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_1ctamma.py index 8f568f131..c31269e4a 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_1ctamma.py @@ -147,12 +147,12 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - # Flexible CGA: the launch carries a preferred (wide) cluster plus a smaller + # Mixed CGA: the launch carries a preferred (wide) cluster plus a smaller # fallback one, and the device picks per cluster — a CTA can only tell which # by reading the hardware cluster dims. Everything cluster-shaped below then # follows from those, so the two kinds share one body; only A's multicast bit # pattern is loop-built and comes in precomputed per shape. - a_mcast_pattern = flex_a_pattern_pref + a_mcast_pattern = mixed_a_pattern_pref if cutlass.const_expr(fallback_cluster_shape_mnk is None): cluster_m = cluster_shape_mnk[0] cluster_n = cluster_shape_mnk[1] @@ -160,11 +160,11 @@ def _kernel( cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() cluster_m = cdim_x cluster_n = cdim_y - a_mcast_pattern = cutlass.Int32(flex_a_pattern_pref) + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_pref) # Bitwise, not `or`: both operands are runtime Booleans (this is the form # cutlass.cute.experimental.is_preferred_cluster uses). if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): - a_mcast_pattern = cutlass.Int32(flex_a_pattern_fb) + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -1342,7 +1342,7 @@ def _host( # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ ) - # Flexible CGA: `cluster` is the preferred (wide) shape and + # Mixed CGA: `cluster` is the preferred (wide) shape and # `fallback_cluster` the regular one the device groups blocks into when a # preferred cluster does not fit. The grid is already a multiple of the # preferred shape, which the driver requires. diff --git a/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_2ctamma.py index f9340e7a2..e1b9cf461 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_2ctamma.py @@ -147,13 +147,13 @@ def _kernel( gridx = cute.arch.grid_dim()[0] gridy = cute.arch.grid_dim()[1] - # Flexible CGA: the launch carries a preferred (wide) cluster plus a smaller + # Mixed CGA: the launch carries a preferred (wide) cluster plus a smaller # fallback one, and the device picks per cluster — a CTA can only tell which # by reading the hardware cluster dims. Everything cluster-shaped below then # follows from those, so the two kinds share one body; only the multicast bit # patterns are loop-built and come in precomputed per shape. - a_mcast_pattern = flex_a_pattern_pref - b_mcast_pattern = flex_b_pattern_pref + a_mcast_pattern = mixed_a_pattern_pref + b_mcast_pattern = mixed_b_pattern_pref if cutlass.const_expr(fallback_cluster_shape_mnk is None): cluster_m = cluster_shape_mnk[0] cluster_n = cluster_shape_mnk[1] @@ -161,13 +161,13 @@ def _kernel( cdim_x, cdim_y, _cdim_z = cute.arch.block_in_cluster_dim() cluster_m = cdim_x cluster_n = cdim_y - a_mcast_pattern = cutlass.Int32(flex_a_pattern_pref) - b_mcast_pattern = cutlass.Int32(flex_b_pattern_pref) + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_pref) + b_mcast_pattern = cutlass.Int32(mixed_b_pattern_pref) # Bitwise, not `or`: both operands are runtime Booleans (this is the form # cutlass.cute.experimental.is_preferred_cluster uses). if (cdim_x != cluster_shape_mnk[0]) | (cdim_y != cluster_shape_mnk[1]): - a_mcast_pattern = cutlass.Int32(flex_a_pattern_fb) - b_mcast_pattern = cutlass.Int32(flex_b_pattern_fb) + a_mcast_pattern = cutlass.Int32(mixed_a_pattern_fb) + b_mcast_pattern = cutlass.Int32(mixed_b_pattern_fb) cluster_size = cluster_m * cluster_n * cluster_shape_mnk[2] cta_rank_in_cluster = cute.arch.block_idx_in_cluster() @@ -1428,7 +1428,7 @@ def _host( # @@INJECT_HOST_TMA_C_PASS@@ # @@TMA_STORE_ONLY:END@@ ) - # Flexible CGA: `cluster` is the preferred (wide) shape and + # Mixed CGA: `cluster` is the preferred (wide) shape and # `fallback_cluster` the regular one the device groups blocks into when a # preferred cluster does not fit. The grid is already a multiple of the # preferred shape, which the driver requires. diff --git a/test/python/gemm/frost/test_block_scale_matmul.py b/test/python/gemm/frost/test_block_scale_matmul.py index 494d0323d..adde16a89 100644 --- a/test/python/gemm/frost/test_block_scale_matmul.py +++ b/test/python/gemm/frost/test_block_scale_matmul.py @@ -2028,7 +2028,7 @@ def test_sm107_block_scale_matmul_shapes_and_clusters(combo, config_name, M, N, ], ) def test_sm107_block_scale_flexible_cga(combo, config_name): - """Flexible CGA rides along with no caller change: any config whose cluster is + """Mixed CGA rides along with no caller change: any config whose cluster is wider than the MMA mode's minimum launches it as the PREFERRED shape plus that minimum as the fallback. The tile decomposition is the identity map for either cluster shape, so both kinds cover the problem exactly once; only the multicast @@ -2037,26 +2037,26 @@ def test_sm107_block_scale_flexible_cga(combo, config_name): hold resident, which is when the device substitutes the fallback shape.""" cta_group = 2 if config_name.endswith("_2ctamma") else 1 cfg = by_name(config_name.rsplit("_", 1)[0]) - assert C._flex_cga_fallback(cfg, cta_group, f"sm107_block_scale_matmul_{cta_group}ctamma.py") == (cta_group, 1) + assert C._mixed_cga_fallback(cfg, cta_group, f"sm107_block_scale_matmul_{cta_group}ctamma.py") == (cta_group, 1) _run_bs_numeric(combo, config_name, 1920, 1920, 512) @requires_sm107 @pytest.mark.parametrize("cta_group", [1, 2]) -def test_flex_cga_fallback_is_the_mma_mode_minimum(cta_group): +def test_mixed_cga_fallback_is_the_mma_mode_minimum(cta_group): """The fallback shape is derived, never passed: one CTA for a 1-CTA MMA, the pair for a 2-CTA one — and a config already AT that minimum has nothing to fall back to, so it launches as a plain fixed cluster.""" tmpl = f"sm107_block_scale_matmul_{cta_group}ctamma.py" assert C.min_fallback_cluster(cta_group) == (cta_group, 1) wide = by_name("CONFIG_sm107_128x128x128_128x128x64_cluster4x2") - assert C._flex_cga_fallback(wide, cta_group, tmpl) == (cta_group, 1) + assert C._mixed_cga_fallback(wide, cta_group, tmpl) == (cta_group, 1) minimal = by_name(f"CONFIG_sm107_128x128x128_128x128x64_cluster{cta_group}x1") - assert C._flex_cga_fallback(minimal, cta_group, tmpl) is None + assert C._mixed_cga_fallback(minimal, cta_group, tmpl) is None @requires_sm107 -def test_flex_cga_is_off_where_it_cannot_be_honored(monkeypatch): +def test_mixed_cga_is_off_where_it_cannot_be_honored(monkeypatch): """Every gate is a fact, not a knob: the GPU's ability to substitute clusters, whether the template consumes the fallback constant at all (an unported one would hang — its cluster constants are baked to the preferred shape), and @@ -2064,40 +2064,47 @@ def test_flex_cga_is_off_where_it_cannot_be_honored(monkeypatch): cluster shapes).""" wide = by_name("CONFIG_sm107_128x128x128_128x128x64_cluster4x2") sm107_tmpl = "sm107_block_scale_matmul_2ctamma.py" - assert C._flex_cga_fallback(wide, 2, sm107_tmpl) == (2, 1) + assert C._mixed_cga_fallback(wide, 2, sm107_tmpl) == (2, 1) - # Template that never reads the constant -> no fallback attached. - assert not C._template_reads_fallback_cluster("sm100_block_scale_matmul_2ctamma.py") - assert C._flex_cga_fallback(wide, 2, "sm100_block_scale_matmul_2ctamma.py") is None + # Template that never reads the constant -> no fallback attached. The MoE + # ones stay that way: their fixed-grid persistent scheduler strides by a + # host-computed cluster count, which mixed clusters invalidate. + moe_tmpl = "sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py" + assert not C._template_reads_fallback_cluster(moe_tmpl) + assert C._mixed_cga_fallback(wide, 2, moe_tmpl) is None # Substitution is a floor, not a range: every part from SM 10.0 up can do it. - assert C._flex_cga_supported(100) and C._flex_cga_supported(110) + assert C._mixed_cga_supported(100) and C._mixed_cga_supported(110) # A pre-Blackwell part -> plain fixed cluster, as before. monkeypatch.setattr(C, "_current_arch", lambda: 90) - assert not C._flex_cga_supported() - assert C._flex_cga_fallback(wide, 2, sm107_tmpl) is None + assert not C._mixed_cga_supported() + assert C._mixed_cga_fallback(wide, 2, sm107_tmpl) is None monkeypatch.undo() # A pinned N-super-block walk -> skipped rather than silently mis-tiled. pinned = dataclasses.replace(wide, tile_swizzle_n=8) - assert C._flex_cga_fallback(pinned, 2, sm107_tmpl) is None + assert C._mixed_cga_fallback(pinned, 2, sm107_tmpl) is None # The escape hatch for A/B measurement. - monkeypatch.setenv("CUDNN_FROST_DISABLE_FLEX_CGA", "1") - assert C._flex_cga_fallback(wide, 2, sm107_tmpl) is None + monkeypatch.setenv("CUDNN_FROST_DISABLE_MIXED_CGA", "1") + assert C._mixed_cga_fallback(wide, 2, sm107_tmpl) is None @requires_sm107 -def test_flex_cga_leaves_unported_templates_byte_identical(): - """An sm100 template on this GPU gets no fallback in its launch — so its - generated kernel, including the L2 swizzle knob the flexible path pins, is - exactly what it was before flexible CGA existed.""" +def test_mixed_cga_ported_templates_attach_a_fallback(): + """A ported template on a wide-cluster config launches with both shapes and + still computes the same result; an unported one renders exactly as it did + before mixed CGA existed.""" sm100_cfg = "CONFIG_sm100_128x128x128_128x128x32_cluster4x2_2ctamma" _run_bs_numeric("nvfp4", sm100_cfg, 512, 512, 512) g = _build_nvfp4_graph(256, 256, 512) - compiled = _plan(g, **_kw(sm100_cfg)) - src = compiled.generated_path.read_text() - assert "fallback_cluster=" not in src + src = _plan(g, **_kw(sm100_cfg)).generated_path.read_text() + assert "fallback_cluster=fallback_cluster_shape_mnk" in src + assert "fallback_cluster_shape_mnk = (2, 1, 1)" in src + + # Already-minimal cluster -> nothing to fall back to, plain fixed launch. + minimal_cfg = "CONFIG_sm100_128x128x128_128x128x32_cluster2x1_2ctamma" + src = _plan(_build_nvfp4_graph(256, 256, 512), **_kw(minimal_cfg)).generated_path.read_text() assert "fallback_cluster_shape_mnk = None" in src From 72ca3f46f27a30c7d3a88c1cb13c25946ca32ac2 Mon Sep 17 00:00:00 2001 From: Yanqin Zhai Date: Mon, 17 Aug 2026 13:35:20 -0700 Subject: [PATCH 3/3] address-code-rabbit-comments --- python/cudnn/gemm/frost/compiler.py | 8 ++++++-- python/cudnn/gemm/frost/kernel_registry.py | 3 +++ test/python/gemm/frost/test_block_scale_matmul.py | 5 ++++- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 117bd302b..7d16d7a17 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -950,8 +950,10 @@ def _mixed_cga_fallback(cfg: TileConfig, cta_group: int, template_file: str) -> when the template has not been ported (its cluster constants are baked to the preferred shape, so a CTA landing in a smaller cluster would wait on arrivals that never come — a hang, not a lost optimization), when the config's cluster - is ALREADY the minimum, and when the config pins the N-super-block walk (that - rasterization is not invariant across the two cluster shapes). + is ALREADY the minimum, when the preferred cluster is not an integer multiple + of it per dim (what the driver requires of the pair), and when the config pins + the N-super-block walk (that rasterization is not invariant across the two + cluster shapes). """ if os.environ.get("CUDNN_FROST_DISABLE_MIXED_CGA"): return None @@ -962,6 +964,8 @@ def _mixed_cga_fallback(cfg: TileConfig, cta_group: int, template_file: str) -> fallback = min_fallback_cluster(cta_group) if fallback == (cfg.cgrp_size_m, cfg.cgrp_size_n) or cfg.tile_swizzle_n > 1: return None + if cfg.cgrp_size_m % fallback[0] or cfg.cgrp_size_n % fallback[1]: + return None return fallback diff --git a/python/cudnn/gemm/frost/kernel_registry.py b/python/cudnn/gemm/frost/kernel_registry.py index 56cb68d15..4ffaf8b7a 100644 --- a/python/cudnn/gemm/frost/kernel_registry.py +++ b/python/cudnn/gemm/frost/kernel_registry.py @@ -499,11 +499,13 @@ def _mm( # unwritten SMEM in K, period 192 B) and the ab_stages budget under-counts, # so cta_tile_m=256 also overruns the SMEM cap. Both are silent-wrong / # launch-fail, hence the gate. See CLAUDE.md for what was ruled out. + # Multi-GEMM has never been validated on this pipeline either — same gate. _mm( "sm103_block_scale_matmul_1ctamma.py", cta_group=1, static=False, graph_type=GraphType.BLOCK_SCALE_MATMUL, + supports_multi_gemm=False, supports_multi_mma_m=False, ), _mm( @@ -511,6 +513,7 @@ def _mm( cta_group=2, static=False, graph_type=GraphType.BLOCK_SCALE_MATMUL, + supports_multi_gemm=False, supports_multi_mma_m=False, ), _mm( diff --git a/test/python/gemm/frost/test_block_scale_matmul.py b/test/python/gemm/frost/test_block_scale_matmul.py index adde16a89..bc614d956 100644 --- a/test/python/gemm/frost/test_block_scale_matmul.py +++ b/test/python/gemm/frost/test_block_scale_matmul.py @@ -2027,7 +2027,7 @@ def test_sm107_block_scale_matmul_shapes_and_clusters(combo, config_name, M, N, "CONFIG_sm107_128x128x128_128x128x64_cluster2x2_1ctamma", ], ) -def test_sm107_block_scale_flexible_cga(combo, config_name): +def test_sm107_block_scale_mixed_cga(combo, config_name): """Mixed CGA rides along with no caller change: any config whose cluster is wider than the MMA mode's minimum launches it as the PREFERRED shape plus that minimum as the fallback. The tile decomposition is the identity map for either @@ -2101,6 +2101,9 @@ def test_mixed_cga_ported_templates_attach_a_fallback(): src = _plan(g, **_kw(sm100_cfg)).generated_path.read_text() assert "fallback_cluster=fallback_cluster_shape_mnk" in src assert "fallback_cluster_shape_mnk = (2, 1, 1)" in src + # What makes the tile walk the identity map for BOTH shapes: the renderer + # pins the N-super-block width, so _auto_swizzle_w const-folds to 1. + assert "tile_swizzle_n = 1" in src # Already-minimal cluster -> nothing to fall back to, plain fixed launch. minimal_cfg = "CONFIG_sm100_128x128x128_128x128x32_cluster2x1_2ctamma"