diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 95864c67f..7d16d7a17 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) @@ -890,12 +893,114 @@ 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.) +_MIXED_CGA_MIN_ARCH = 100 + + +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 >= _MIXED_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 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 _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-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, + 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, 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 + if not _mixed_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 + if cfg.cgrp_size_m % fallback[0] or cfg.cgrp_size_n % fallback[1]: + return None + return fallback + + +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`` + 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"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}", + ] + + 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 +1302,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}", + # 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()}", + # 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 +1448,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(_mixed_cga_constants(cfg, cta_group, fallback_cluster)) lines.extend(_quant_device_imports(chain)) return "\n".join(lines) @@ -1408,6 +1519,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 @@ -1462,7 +1574,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 @@ -1569,6 +1681,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 +1701,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 +2906,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. + + 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:`_mixed_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 +3441,7 @@ def _jit_block_scale( _arch_reason = _tmpl.active_reject(config) if _arch_reason is not None: raise NotImplementedError(_arch_reason) + 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) @@ -3332,7 +3452,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( @@ -3671,7 +3791,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..4ffaf8b7a 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). @@ -505,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( @@ -517,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( @@ -524,18 +521,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 +541,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 +554,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 7a58499fe..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,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] + # 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 = 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() @@ -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, ) + # 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_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm107_block_scale_matmul_2ctamma.py index 4c20d844b..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,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] + # 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() @@ -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, ) + # 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/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..bc614d956 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,103 @@ 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_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 + 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._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_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._mixed_cga_fallback(wide, cta_group, tmpl) == (cta_group, 1) + minimal = by_name(f"CONFIG_sm107_128x128x128_128x128x64_cluster{cta_group}x1") + assert C._mixed_cga_fallback(minimal, cta_group, tmpl) is None + + +@requires_sm107 +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 + 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._mixed_cga_fallback(wide, 2, sm107_tmpl) == (2, 1) + + # 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._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._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._mixed_cga_fallback(pinned, 2, sm107_tmpl) is None + + # The escape hatch for A/B measurement. + monkeypatch.setenv("CUDNN_FROST_DISABLE_MIXED_CGA", "1") + assert C._mixed_cga_fallback(wide, 2, sm107_tmpl) is None + + +@requires_sm107 +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) + 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" + src = _plan(_build_nvfp4_graph(256, 256, 512), **_kw(minimal_cfg)).generated_path.read_text() + assert "fallback_cluster_shape_mnk = None" in src + + @requires_sm107 @pytest.mark.parametrize( "config_name",