diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 7d16d7a17..e67385c43 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -552,6 +552,17 @@ def _smem_desc_params( # cta_tile_k_bytes, and an MN-major tile is MN/group_elems groups of # group_elems * cta_tile_k_bytes — both give cta_tile_k_bytes per MN element. a_smem_m_step_bytes = (cta_smem_m // cfg.num_mma_m) * cfg.cta_tile_k_bytes + a_mcast_slices, b_mcast_slices, ab_empty_full_mask = _mcast_slice_plan( + chain.matmul.a_major, + chain.matmul.b_major, + cfg.cgrp_size_m, + cfg.cgrp_size_n, + cta_group, + cta_smem_m, + cta_smem_n, + per_cta_a=cta_group == 2 and bool(chain.mainloop_a_ops), + per_cta_b=cta_group == 2 and bool(chain.mainloop_b_ops), + ) lines = [ f"# Tile config: {cfg.name}", @@ -560,7 +571,7 @@ def _smem_desc_params( # Template `cta_tile_mnk` = per-CTA SMEM/TMA box dims (B's N halved under # 2-CTA MMA), NOT the logical per-CTA tile from TileConfig. f"cta_tile_mnk = {cfg.cta_smem_tile_mnk(elem_bytes, cta_group)}", - f"epi_tile_mn = {cfg.epi_tile_mn}", + f"epi_tile_mn = {(cfg.epi_tile_mn[0], _epi_n(cfg, cta_group, out_dt))}", f"threads_per_cta = {cfg.threads_per_cta}", f"cluster_shape_mnk = {cfg.cluster_shape}", f"matmul_batch = {chain.matmul.batch}", @@ -570,9 +581,12 @@ def _smem_desc_params( f"b_is_n_major = {chain.matmul.b_major == 'n'}", f"mma_a_major = {1 if chain.matmul.a_major == 'm' else 0}", f"mma_b_major = {1 if chain.matmul.b_major == 'n' else 0}", - f"ab_stages = {cfg.max_ab_stages(cta_group)}", + f"ab_stages = {cfg.max_ab_stages(cta_group, moe=chain.has_moe)}", f"multicast_a = {cfg.multicast_a}", f"multicast_b = {cfg.multicast_b(cta_group)}", + f"a_mcast_slices = {a_mcast_slices}", + f"b_mcast_slices = {b_mcast_slices}", + f"ab_empty_full_mask = {ab_empty_full_mask}", f"ab_smem_swizzle = cutlass.experimental.primitives.Tcgen05SmemSwizzle.{smem_swizzle_name}", f"ab_smem_swizzle_bytes = {smem_swizzle_bytes}", f"a_smem_desc_leading_byte_offset = {a_lbo}", @@ -606,6 +620,7 @@ def _smem_desc_params( f"cd_fake_n_div = {2 if out_dt == 'fp4_e2m1' else 1}", # M-major TMA-store C-descriptor inner-M box = 128 B swizzle span / elem_bytes. f"cd_mmajor_atom_m = {128 // DTYPE_BYTES[out_dt]}", + *_epi_swizzle_lines(cfg, cta_group, out_dt), ] # Persistent kernel always: double-TMEM + L2 N-super-block swizzle. # (acc_stages is emitted below, once the TMEM budget is known.) @@ -648,9 +663,8 @@ def _smem_desc_params( # Per-CTA SMEM B-tile N is halved under 2-CTA MMA (the pair splits B's N). smem_n = cfg.cta_tile_n // cta_group - per_stage = (chain.num_a_operands * cfg.cta_tile_m + chain.num_b_operands * smem_n) * cfg.cta_tile_k_bytes + 2 * 8 - fixed = 2 * acc_stages * 8 + 8 - avail = _sm_smem_ab_budget_bytes(cfg.pipeline) - fixed + per_stage = (chain.num_a_operands * cfg.cta_tile_m + chain.num_b_operands * smem_n) * cfg.cta_tile_k_bytes + avail = _sm_smem_ab_budget_bytes(cfg.pipeline, moe=chain.has_moe) ab_stages_mg = min(avail // per_stage, _AB_STAGES_CAP) if ab_stages_mg < 1: raise NotImplementedError( @@ -710,7 +724,7 @@ def _smem_desc_params( # Final ab_stages override: account for the TMA-D SMEM buffer (fixed, when # TMA-store is active) AND a mixed-input mainloop's narrow LOAD buffer # (per-stage). Otherwise leave the plain max. - smem_d_bytes = _smem_d_bytes(cfg, chain) if use_tma else 0 + smem_d_bytes = _smem_d_bytes(cfg, chain, cta_group) if use_tma else 0 cast_extra_per_stage = 0 if chain.has_mainloop_fusion and (chain.mainloop_a_cast or chain.mainloop_b_cast): smem_n = cfg.cta_tile_n // cta_group @@ -724,6 +738,7 @@ def _smem_desc_params( cta_group, extra_smem_bytes=smem_d_bytes, extra_per_stage_bytes=cast_extra_per_stage, + moe=chain.has_moe, ) 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)) @@ -893,6 +908,32 @@ def _grid_num_clusters(cfg: TileConfig, device=None) -> int: return max_active_clusters(cfg.cgrp_size_m * cfg.cgrp_size_n, device) +_SMEM_SWIZZLE_ATOM_ROWS = 8 + + +def _mcast_slice_plan( + a_major: str, + b_major: str, + cluster_m: int, + cluster_n: int, + cta_group: int, + cta_smem_m: int, + cta_smem_n: int, + *, + per_cta_a: bool = False, + per_cta_b: bool = False, +) -> tuple[int, int, bool]: + """(a_slices, b_slices, needs_full_empty_mask) for the cluster's TMA multicast.""" + a_group = 1 if per_cta_a else cluster_n + b_group = 1 if per_cta_b else cluster_m // cta_group + atom = _SMEM_SWIZZLE_ATOM_ROWS + a_slices = a_group if (a_major == "k" and a_group > 1 and cta_smem_m % (a_group * atom) == 0) else 1 + b_slices = b_group if (b_major == "k" and b_group > 1 and cta_smem_n % (b_group * atom) == 0) else 1 + a_closed = cluster_n == 1 or a_slices > 1 + b_closed = cluster_m // cta_group == 1 or b_slices > 1 + return a_slices, b_slices, not (a_closed and b_closed) + + 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. @@ -1031,6 +1072,9 @@ def _render_block_scale_tile_constants( cta_m = cfg.cta_tile_m cta_n = cfg.cta_tile_n + bs_a_mcast_slices, bs_b_mcast_slices, bs_ab_empty_full_mask = _mcast_slice_plan( + chain.matmul.a_major, chain.matmul.b_major, cfg.cgrp_size_m, cfg.cgrp_size_n, cta_group, cta_m, cta_n // cta_group + ) # MMA K-instruction width (sm100 → 32 bytes): fp4 → 64 elems, fp8 → 32. mma_inst_k_bytes = cfg.mma_inst_k_bytes mma_inst_k_elems = mma_inst_k_bytes * 8 // data_elem_bits @@ -1081,24 +1125,38 @@ def _bs_smem_desc_params(is_mn_major, mn_extent, name): sf_k4 = sf_k // 4 # 4 SF-K per utccp atom nb_m = cta_m // 128 nb_n = cta_n // 128 + mma_nb_m = cfg.mma_inst_m // 128 + mma_nb_n = cfg.mma_inst_n // 128 + sfa_nb_m = mma_nb_m * cfg.num_mma_m _REGISTERS_PER_ATOM = 4 # cols per 128×4 utccp atom scales_per_inst = mma_inst_k_elems // bs.block_size word_scales = max(_REGISTERS_PER_ATOM, scales_per_inst) # cols per block-word word_atoms = -(-word_scales // _REGISTERS_PER_ATOM) # atoms copied per word (ceil: a partial atom still costs one) insts_per_word = max(_REGISTERS_PER_ATOM // scales_per_inst, 1) - num_sf_words = max(num_kblocks // insts_per_word, 1) # utccp refreshes / k-tile _REGISTERS_PER_BLOCK = word_scales # SF word width per block - # One SF word per 128-row / 128-column block, packed back to back. SFB must - # stay contiguous (one instruction's SFB read extent grows with n_dim, so a - # single scale_b reads every N block as one span); SFA blocks each get their - # own scale_a at the same stride. - sfa_tmem_cols = nb_m * _REGISTERS_PER_BLOCK - sfb_tmem_cols = nb_n * _REGISTERS_PER_BLOCK # fixed SF word width (SFB) if is_sm103: + # A whole K-tile of SF is TMEM-resident, so a region is sf_k wide. num_sf_words = sf_k4 - sfa_tmem_cols = nb_m * sf_k - sfb_tmem_cols = nb_n * sf_k + sfa_tmem_cols = sfa_nb_m * sf_k + sfb_tmem_cols = mma_nb_n * sf_k + else: + # One SF word per 128-row / 128-column block, packed back to back and + # refreshed per word. SFB must stay contiguous (one instruction's SFB read + # extent grows with n_dim, so a single scale_b reads every N block as one + # span); SFA is one block per M sub-block, so its M instructions stay + # independent. + num_sf_words = max(num_kblocks // insts_per_word, 1) # utccp refreshes / k-tile + if num_sf_words * insts_per_word != num_kblocks: + raise NotImplementedError( + f"block-scale {cfg.name!r}: {num_sf_words} SF word(s) x {insts_per_word} " + f"instruction(s) per word covers {num_sf_words * insts_per_word} K-blocks, " + f"but the K-tile has {num_kblocks}. The MMA would read scale bytes that were " + f"never staged, or skip the trailing K-blocks. Reachable only if the " + f"cta_tile_k_bytes == 128 requirement is relaxed." + ) + sfa_tmem_cols = sfa_nb_m * _REGISTERS_PER_BLOCK + sfb_tmem_cols = mma_nb_n * _REGISTERS_PER_BLOCK # utccp SMEM-source offsets (16-byte units). One 128×4 atom = 512 B = 32; # consecutive K-atoms 1 atom apart; each M/N-block of 128 rows is sf_k4 atoms # further along the SF SMEM tile. @@ -1121,9 +1179,6 @@ def _bs_smem_desc_params(is_mn_major, mn_extent, name): # double-TMEM pipelining at a single mbar. total_tmem = _tmem_cols_for_arch() - def _align16(x: int) -> int: - return (x + 15) & ~15 - # --- TMEM acc-stage + overlap (per-GEMM budget; arch- & count-agnostic) --- # SF = one fixed word PER DISTINCT OPERAND (shared A → one SFA word). # Each GEMM's acc gets its OWN region; the 2 tile-stage buffers overlap @@ -1144,7 +1199,29 @@ def _align16(x: int) -> int: acc_cols_per_stage = num_mma_m * epi_cols_per_mma_m na, nb = chain.num_a_operands, chain.num_b_operands sf_total_cols = na * sfa_tmem_cols + nb * sfb_tmem_cols - per_gemm = (total_tmem - sf_total_cols) // num_gemms + # Columns each instruction reads from its scale base -- ISA opUTCHMMA, + # "Load A/B scale factors" (sf{a,b}_tmem_cols). + if is_sm103: + # The resident K-tile means the scale pointer advances per k-block, so a + # read window can reach past the operand's own region. + sf_ids = [scales_per_inst * j % 4 for j in range(num_kblocks)] + wide_vec = bs.block_size == 16 + sfb_extra = 4 if cta_n <= 128 else 8 + sfa_off = [scales_per_inst * j // 4 * 4 * sfa_nb_m for j in range(num_kblocks)] + sfb_off = [scales_per_inst * j // 4 * 4 * mma_nb_n for j in range(num_kblocks)] + extra = [sfb_extra if (wide_vec or sf_ids[j] >= 2) else 0 for j in range(num_kblocks)] + sfa_spans = [(sfa_off[j], 4 if (not wide_vec and sf_ids[j] < 2) else 8) for j in range(num_kblocks)] + sfb_spans = [(sfb_off[j], 2 * ((cta_n + 63) // 64) + extra[j]) for j in range(num_kblocks)] + sf_reserved_cols = max( + sf_total_cols, + (na - 1) * sfa_tmem_cols + max(off + cols for off, cols in sfa_spans), + na * sfa_tmem_cols + (nb - 1) * sfb_tmem_cols + max(off + cols for off, cols in sfb_spans), + ) + else: + # Every SF word is re-utccp'd into the same columns and the scale base is fixed, + # so an instruction reads exactly its own operand's region. + sf_reserved_cols = sf_total_cols + per_gemm = (total_tmem - sf_reserved_cols) // num_gemms if per_gemm < acc_cols_per_stage: raise NotImplementedError( f"block-scale {cfg.name!r}: per-GEMM TMEM budget {per_gemm} < one acc " @@ -1156,12 +1233,10 @@ def _align16(x: int) -> int: acc_stages = 2 # full per-GEMM double-buffer else: acc_stages = 1 - # the overlap drain order assumes one contiguous M block per stage - if num_mma_m == 1: - gran = 32 # epilogue TMEM-load drain unit (cols) - ov = ((2 * acc_cols_per_stage - per_gemm + gran - 1) // gran) * gran - if ov < acc_cols_per_stage: # else no room → plain 1-stage - acc_overlap_cols = ov + gran = _epi_n(cfg, cta_group, chain.output_dtype) # epilogue TMEM-load drain unit (cols) + ov = ((2 * acc_cols_per_stage - per_gemm + gran - 1) // gran) * gran + if ov < acc_cols_per_stage: # else no room -> plain 1-stage + acc_overlap_cols = ov use_acc_overlap = acc_overlap_cols > 0 # within-GEMM per-stage stride + per-GEMM region size: acc_stage_stride = (acc_cols_per_stage - acc_overlap_cols) if use_acc_overlap else acc_cols_per_stage @@ -1171,10 +1246,10 @@ def _align16(x: int) -> int: acc_gemm_stride = 2 * acc_cols_per_stage - acc_overlap_cols else: acc_gemm_stride = acc_cols_per_stage - acc_overlap_subtiles = acc_overlap_cols // 32 + acc_overlap_subtiles = acc_overlap_cols // _epi_n(cfg, cta_group, chain.output_dtype) acc_region_cols = acc_cols_per_stage # per-stage stride WITHIN a GEMM - sf_region_base = _align16(num_gemms * acc_gemm_stride) + sf_region_base = num_gemms * acc_gemm_stride # Per-distinct-operand SF word col bases (single-GEMM → length-1 lists). sfa_col_bases = [sf_region_base + i * sfa_tmem_cols for i in range(na)] sfb_col_bases = [sf_region_base + na * sfa_tmem_cols + j * sfb_tmem_cols for j in range(nb)] @@ -1188,51 +1263,28 @@ def _align16(x: int) -> int: f"block-scale {cfg.name!r}: the accumulator + SF regions need {used_cols} TMEM columns but only {num_tmem_alloc_cols} are allocated" ) - mma_m = cfg.mma_inst_m * cta_group - half_m = cta_group == 2 and mma_m == 128 # unreachable while block-scale pins cta_m=128 - omma_k = mma_inst_k_elems if is_fp4 else 0 - sf_ids = [scales_per_inst * j % 4 for j in range(num_kblocks)] - sfb_extra = 4 if cta_n <= 128 else 8 - sfa_off = [scales_per_inst * j // 4 * 4 * nb_m for j in range(num_kblocks)] - sfb_off = [scales_per_inst * j // 4 * 4 * nb_n for j in range(num_kblocks)] - if omma_k in (96, 128): - # 96 -> 3X (block 32) / 6X (block 16); 128 -> 4X (block 32) / 8X (block 16). - wide_vec = bs.block_size == 16 # the 6X / 8X arm - if omma_k == 128: # extra-enhanced: the extra term is 8X-only - extra = [sfb_extra if wide_vec else 0] * num_kblocks - else: # enhanced: 6X, or 3X with sfb_id >= 2 - extra = [sfb_extra if (wide_vec or sf_ids[j] >= 2) else 0 for j in range(num_kblocks)] - sfa_spans = [(sfa_off[j], 4 if (not wide_vec and omma_k == 96 and sf_ids[j] < 2) else 8) for j in range(num_kblocks)] - sfb_spans = [(sfb_off[j], 2 * ((cta_n + 63) // 64) + extra[j]) for j in range(num_kblocks)] - else: - sfa_spans = [(0, 2 if half_m else 4)] - sfb_spans = [(0, 2 * ((cta_n + 127) // 128) if half_m else 2 * ((cta_n + 63) // 64))] - for _label, _bases, _spans in (("SFA", sfa_col_bases, sfa_spans), ("SFB", sfb_col_bases, sfb_spans)): - _end = max(b + off + cols for b in _bases for off, cols in _spans) - if _end > num_tmem_alloc_cols: - raise NotImplementedError( - f"block-scale {cfg.name!r}: the hardware {_label} TMEM span reaches " - f"column {_end} but only {num_tmem_alloc_cols} are allocated" - + ("; an SFB overrun is NOT reported by the hardware, so this check is the only guard" if _label == "SFB" else "") - ) + if is_sm103: + for _label, _bases, _spans in (("SFA", sfa_col_bases, sfa_spans), ("SFB", sfb_col_bases, sfb_spans)): + _end = max(b + off + cols for b in _bases for off, cols in _spans) + if _end > num_tmem_alloc_cols: + raise NotImplementedError( + f"block-scale {cfg.name!r}: the hardware {_label} TMEM span reaches " + f"column {_end} but only {num_tmem_alloc_cols} are allocated (the MMA " + f"would fault OOR_ADDR)" + ) # --- AB SMEM pipeline depth ---------------------------------------------- - # Per-stage SMEM = (packed data + SF) per DISTINCT operand + 2 mbar. - per_stage = na * (sA_packed_elems + sfa_smem_bytes) + nb * (sB_packed_elems + sfb_smem_bytes) + 2 * 8 from .tile_config import _sm_smem_ab_budget_bytes, _AB_STAGES_CAP - fixed = 2 * acc_stages * 8 + 8 # TMA-store stages output through a fixed SMEM-D buffer; reserve it before # sizing the AB pipeline (else SMEM overflows the cap). - if use_tma_store_epi: - fixed += _smem_d_bytes(cfg, chain) - ab_stages = max(1, min((_sm_smem_ab_budget_bytes(cfg.pipeline) - fixed) // per_stage, _AB_STAGES_CAP)) + ab_budget = _sm_smem_ab_budget_bytes(cfg.pipeline, moe=chain.has_moe) - (_smem_d_bytes(cfg, chain, cta_group) if use_tma_store_epi else 0) if is_sm103: # CUTLASS-style sm103 pipeline: an AB stage is ONE 128-B-K chunk (a # third of the 384-B K-tile), and SF rides its OWN ring (own warp, # own mbars) at 12-SF-per-row group granularity — 4 groups per K-tile # at VS16, 2 at VS32 (both 12 SFs/row). Data-only AB stages + a fixed - # SF ring replace the combined per-K-tile stage above. + # SF ring, instead of one combined per-K-tile stage. sf_stages = 6 a_chunk_bytes = cta_m * 128 b_chunk_bytes = (cta_n // cta_group) * 128 @@ -1240,13 +1292,17 @@ def _align16(x: int) -> int: # SFB is loaded FULL per CTA (the pair MMA reads each CTA's own TMEM # SFB across the whole pair-N range) — same convention as sm100. sfb_group_bytes = cta_n * 12 - sf_ring_bytes = sf_stages * (na * sfa_group_bytes + nb * sfb_group_bytes + 2 * 8) - per_ab_stage = na * a_chunk_bytes + nb * b_chunk_bytes + 2 * 8 - ab_stages = min((_sm_smem_ab_budget_bytes(cfg.pipeline) - fixed - sf_ring_bytes) // per_ab_stage, _AB_STAGES_CAP) + sf_ring_bytes = sf_stages * (na * sfa_group_bytes + nb * sfb_group_bytes) + per_ab_stage = na * a_chunk_bytes + nb * b_chunk_bytes + ab_stages = min((ab_budget - sf_ring_bytes) // per_ab_stage, _AB_STAGES_CAP) if ab_stages < 3: raise NotImplementedError( f"block-scale {cfg.name!r}: only {ab_stages} 128-B AB chunk " f"stages fit in SMEM — the sm103 pipeline needs >= 3 (one " f"K-tile in flight)" ) + else: + # One stage covers a whole K-tile: packed data + SF per DISTINCT operand. + per_stage = na * (sA_packed_elems + sfa_smem_bytes) + nb * (sB_packed_elems + sfb_smem_bytes) + ab_stages = max(1, min(ab_budget // per_stage, _AB_STAGES_CAP)) out_dt = chain.output_dtype vec_bytes_epi = _epi_vec_bytes(chain, cfg, cta_group) @@ -1278,7 +1334,7 @@ def _align16(x: int) -> int: f"cgrp_tile_mnk = ({cta_m * cfg.cgrp_size_m}, {cta_n * cfg.cgrp_size_n}, {cta_k_elems})", f"cgrp_tile_m = {cta_m * cfg.cgrp_size_m}", f"cgrp_tile_n = {cta_n * cfg.cgrp_size_n}", - f"epi_tile_mn = {cfg.epi_tile_mn}", + f"epi_tile_mn = {(cfg.epi_tile_mn[0], _epi_n(cfg, cta_group, out_dt))}", f"threads_per_cta = 256", f"cluster_shape_mnk = {cfg.cluster_shape}", f"matmul_a_batch = {chain.matmul.a_batch}", @@ -1311,6 +1367,9 @@ def _align16(x: int) -> int: # plain load when the runtime pattern names a single peer. f"multicast_a = {cfg.multicast_a}", f"multicast_b = {cfg.multicast_b(cta_group)}", + f"a_mcast_slices = {bs_a_mcast_slices}", + f"b_mcast_slices = {bs_b_mcast_slices}", + f"ab_empty_full_mask = {bs_ab_empty_full_mask}", "", f"# packed data SMEM", # ab_dtype is the width BOTH operands are sized by (sA/sB bytes, B's TMA @@ -1357,6 +1416,7 @@ def _align16(x: int) -> int: f"cd_fake_n_div = {2 if out_dt == 'fp4_e2m1' else 1}", # M-major TMA-store C-descriptor inner-M box = 128 B swizzle span / elem_bytes. f"cd_mmajor_atom_m = {128 // DTYPE_BYTES[out_dt]}", + *_epi_swizzle_lines(cfg, cta_group, out_dt), "", f"# block-scale MMA", f"mma_block_scale_kind = nvvm.MMABlockScaleKind.{bs.mma_block_scale_kind}", @@ -1382,6 +1442,7 @@ def _align16(x: int) -> int: f"num_blocks_n = {nb_n}", f"registers_per_block = {_REGISTERS_PER_BLOCK}", f"epi_cols_per_mma_m = {epi_cols_per_mma_m}", + f"mma_c_dtype = {DTYPE_TO_CUTLASS[chain.matmul.accum_dtype]}", # Byte step from one MMA M sub-block to the next inside the SMEM tile. # sm103 stages ONE 128-B K chunk per AB stage, not the whole K-tile, so # its per-M-row width is the chunk's, not cta_tile_k_bytes. @@ -1423,8 +1484,8 @@ def _align16(x: int) -> int: f"mma_next_chunk_by_j = {tuple((kstep * j + kstep - 1) // 128 for j in range(num_kblocks))}", f"mma_phase16_by_j = {tuple((kstep * j) % 128 // 16 for j in range(num_kblocks))}", f"sf_id_by_j = {tuple(spi * j % 4 for j in range(num_kblocks))}", - f"sfa_mma_col_off_by_j = {tuple(spi * j // 4 * 4 * nb_m for j in range(num_kblocks))}", - f"sfb_mma_col_off_by_j = {tuple(spi * j // 4 * 4 * nb_n for j in range(num_kblocks))}", + f"sfa_mma_col_off_by_j = {tuple(spi * j // 4 * 4 * sfa_nb_m for j in range(num_kblocks))}", + f"sfb_mma_col_off_by_j = {tuple(spi * j // 4 * 4 * mma_nb_n for j in range(num_kblocks))}", ] if is_sm107: # SM 10.7 block-scale MMA: K = 64 bytes per instruction (2x sm100), so @@ -2650,12 +2711,45 @@ def _check_input_alignment(chain: FusionChain) -> None: _EPI_SMEM_STAGES = 2 -def _smem_d_bytes(cfg, chain) -> int: - """SMEM-D buffer bytes for the TMA-store epilogue: `_EPI_SMEM_STAGES` slots - of one epilogue subtile (`epi_tile_mn` = one MMA-M block × 32) + a 16-byte - alignment pad. With num_mma_m > 1 the M blocks reuse the same slots.""" - elem_bytes = DTYPE_BYTES[chain.output_dtype] - return _EPI_SMEM_STAGES * cfg.epi_tile_mn[0] * cfg.epi_tile_mn[1] * elem_bytes + 16 +_EPI_SWIZZLE_BY_ROW_BYTES = {32: (1, "s32b"), 64: (2, "s64b"), 128: (3, "s128b")} +_EPI_ROW_BYTES_MAX = 128 # widest TMA store swizzle +_EPI_N_BASE = 32 # drain width when the epilogue is already hidden behind the MMA +_EPI_N_MAX = 64 # per-lane fp32 registers the drain can hold + + +def _epi_n(cfg, cta_group: int, out_dt: str) -> int: + cols = _epi_tile_cols(cfg, cta_group) + cap = _EPI_N_BASE + if cfg.num_mma_m > 1 and 2 * cfg.num_mma_m * cols > _tmem_cols_for_arch(): + cap = _EPI_N_MAX + n = min(_EPI_ROW_BYTES_MAX // DTYPE_BYTES[out_dt], cap, cols) + return 1 << (n.bit_length() - 1) + + +def _epi_swizzle_lines(cfg, cta_group: int, out_dt: str) -> list[str]: + epi_n = _epi_n(cfg, cta_group, out_dt) + row_bytes = epi_n * DTYPE_BYTES[out_dt] + if row_bytes not in _EPI_SWIZZLE_BY_ROW_BYTES: + raise NotImplementedError( + f"{cfg.name!r}: epilogue subtile row is {row_bytes} bytes " + f"(epi_n={epi_n} x {DTYPE_BYTES[out_dt]}B {out_dt}); " + f"the TMA store swizzle only spans {sorted(_EPI_SWIZZLE_BY_ROW_BYTES)}" + ) + b, tma = _EPI_SWIZZLE_BY_ROW_BYTES[row_bytes] + return [ + f"epi_n = {epi_n}", + f"epi_smem_swizzle = cutlass.Swizzle({b}, 4, 3)", + f"epi_tma_swizzle = _tma.TensorMapSwizzle.{tma}", + ] + + +def _smem_d_bytes(cfg, chain, cta_group: int) -> int: + """SMEM-D buffer bytes for the TMA-store epilogue: `_EPI_SMEM_STAGES` slots of + one epilogue subtile (one MMA-M block x epi_n) + a 16-byte alignment pad. With + num_mma_m > 1 the M blocks reuse the same slots. epi_n MUST be the same value + the kernel renders, or the reserve under-counts and the launch is rejected.""" + out_dt = chain.output_dtype + return _EPI_SMEM_STAGES * cfg.epi_tile_mn[0] * _epi_n(cfg, cta_group, out_dt) * DTYPE_BYTES[out_dt] + 16 def _use_tma_store_epi(chain, cfg, vec_bytes_epi: int, cta_group: int) -> bool: @@ -2666,7 +2760,7 @@ def _use_tma_store_epi(chain, cfg, vec_bytes_epi: int, cta_group: int) -> bool: SMEM source aligned to the descriptor swizzle (else undeclarable). - mma_inst_m == 128: only the 128-rows-per-MMA-block thread→row layout is wired (an M=64 MMA block drains through the packed lane<16 layout). - - out dtype ∈ {bf16, fp16}: matches the hard-coded s64b 32-col swizzle. + - out dtype ∈ {bf16, fp16}: the drain widens to epi_n only for 2-byte output. - M-major output: 16B-aligned M (16x256b TMEM-load + stmatrix.trans + tma_store). """ if chain.has_moe: @@ -2698,10 +2792,11 @@ def _use_tma_store_epi(chain, cfg, vec_bytes_epi: int, cta_group: int) -> bool: return False if chain.output_dtype not in ("bf16", "fp16"): return False - # Fixed 32-col subtile: an N-tile that is not a whole number of subtiles - # would TMA-store a 32-wide box past the tile edge into the neighbouring - # tile (TMA clamps only at the GLOBAL extent) — fall back to STG. - if cfg.cta_tile_n % 32 != 0: + # An N-tile that is not a whole number of subtiles would TMA-store an + # epi_n-wide box past the tile edge into the neighbouring tile (TMA clamps + # only at the GLOBAL extent) — fall back to STG. The N-major arm's span walk + # would halve to fit, but the M-major arm drains a fixed epi_n. + if _epi_tile_cols(cfg, cta_group) % _epi_n(cfg, cta_group, chain.output_dtype) != 0: return False # Under cta_group=2 each CTA holds cta_tile_n//2 cols, so cta_tile_n<64 # (per-CTA n<32) would split a subtile across CTAs — unsupported by the diff --git a/python/cudnn/gemm/frost/kernel_templates/_tile_helpers.py b/python/cudnn/gemm/frost/kernel_templates/_tile_helpers.py index 4496c7b26..97a1d15dd 100644 --- a/python/cudnn/gemm/frost/kernel_templates/_tile_helpers.py +++ b/python/cudnn/gemm/frost/kernel_templates/_tile_helpers.py @@ -34,12 +34,14 @@ def l2_swizzle_tile(raw_m, raw_n, nt_m, nt_n, swizzle_w): return log_m, log_n -def epi_subtile_spans(cols): - """Power-of-two column spans the epilogue drains a tile in (host-side).""" +def epi_subtile_spans(cols, epi_n=32): + """Power-of-two column spans the epilogue drains a tile in (host-side). + Starts at ``epi_n`` and halves to fit the remainder, so any 8-multiple N is + covered whatever the widest span is.""" spans = [] off = 0 while off < cols: - w = 32 + w = epi_n while w > cols - off: w //= 2 spans.append((off, w)) 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 910f7cbba..bcb6fe68b 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 @@ -24,6 +24,7 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( + epi_subtile_spans as _epi_subtile_spans, l2_swizzle_tile as _l2_swizzle_tile, tcgen05_alloc as _tcgen05_alloc, tcgen05_dealloc as _tcgen05_dealloc, @@ -41,6 +42,9 @@ # @@INJECT_TILE_CONSTANTS@@ +if use_acc_overlap and any(_w != epi_n for _, _w in _epi_subtile_spans(epi_cols_per_mma_m, epi_n)): + raise NotImplementedError(f"{__name__}: acc overlap reverses subtiles by index, which needs a uniform drain width") + CLC_SCHED_STAGES = 2 @@ -195,11 +199,15 @@ def _kernel( a_part_arrive = cutlass.Int16(a_pattern) << m_rank b_part_arrive = cutlass.Int16(b_pattern) << (n_rank * cluster_m) - ab_empty_arrive_mask = a_part_arrive | b_part_arrive + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = a_part_arrive | b_part_arrive _smem_sys_reserved = cutlass.Array(cutlass.Int8, 1024, space=cutlass.AddressSpace.smem, alignment=1) ab_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) + sf_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) ab_empty_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) acc_empty_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) acc_full_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) @@ -269,13 +277,18 @@ def _kernel( ) # @@TMA_STORE_ONLY:END@@ - ab_empty_count = cluster_m + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size + else: + ab_empty_count = cluster_m + cluster_n - 1 num_consumer_warps_per_cta = 7 clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: for i in range(ab_stages): if elect_one: nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) + if elect_one: + nvvm.mbarrier_init(sf_full_mbar_ptr.subview(i), 1) if elect_one: nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) for i in range(acc_stages): @@ -300,9 +313,10 @@ def _kernel( sA_bytes = sA_elems * (ab_dtype.width // 8) 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) + ab_only_copy_bytes = num_a_operands * sA_bytes + num_b_operands * sB_bytes + sf_only_copy_bytes = num_a_operands * sfa_smem_bytes + num_b_operands * sfb_smem_bytes - cols_per_acc_stage = cta_tile_mnk[1] + epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m tmem_alloc_bar_count = (num_epilogue_warps + 1) * 32 # @@INJECT_TAP_PTRS@@ @@ -377,8 +391,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) tile_m = init_tile_m @@ -410,7 +423,9 @@ def _kernel( coord_k = k_tile_idx * cta_tile_mnk[2] coord_sf_k = k_tile_idx * sf_tma_box_k if elect_one: - nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) + nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), ab_only_copy_bytes) + if elect_one: + nvvm.mbarrier_arrive_expect_tx(sf_full_mbar_ptr.subview(stage), sf_only_copy_bytes) for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) @@ -419,6 +434,93 @@ def _kernel( tma_sfa_desc = tma_sfa_descs[_ai] sfa_m_block = coord_m_per_cta // 128 if cutlass.const_expr(multicast_a): + if n_rank == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFA_stage, + tma_sfa_desc.get_ptr(), + (0, coord_sf_k, sfa_m_block, tile_l_a), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + else: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFA_stage, + tma_sfa_desc.get_ptr(), + (0, coord_sf_k, sfa_m_block, tile_l_a), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + + for _bj in cutlass.range_constexpr(num_b_operands): + sB_stage = smem_b_list[_bj].subview(sB_elems * stage) + tma_b_desc = tma_b_descs[_bj] + sSFB_stage = smem_sfb_list[_bj].subview(sfb_smem_bytes * stage) + tma_sfb_desc = tma_sfb_descs[_bj] + sfb_n_block = coord_n_per_cta // 128 + if cutlass.const_expr(multicast_b): + if m_rank == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFB_stage, + tma_sfb_desc.get_ptr(), + (0, coord_sf_k, sfb_n_block, tile_l_b), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + else: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFB_stage, + tma_sfb_desc.get_ptr(), + (0, coord_sf_k, sfb_n_block, tile_l_b), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + + for _ai in cutlass.range_constexpr(num_a_operands): + sA_stage = smem_a_list[_ai].subview(sA_elems * stage) + tma_a_desc = tma_a_descs[_ai] + sSFA_stage = smem_sfa_list[_ai].subview(sfa_smem_bytes * stage) + tma_sfa_desc = tma_sfa_descs[_ai] + sfa_m_block = coord_m_per_cta // 128 + if cutlass.const_expr(a_mcast_slices > 1): + _a_rows = cta_tile_mnk[0] // a_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(n_rank * _a_rows * ab_packed_per_row), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + n_rank * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + else: + _a_per_cta = a_mcast_slices // cluster_n + for _asl in cutlass.range(_a_per_cta): + _a_idx = n_rank * _a_per_cta + _asl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(_a_idx * _a_rows * ab_packed_per_row), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + _a_idx * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + elif cutlass.const_expr(multicast_a): if n_rank == 0: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): @@ -475,37 +577,40 @@ def _kernel( multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_1, ) - if cutlass.const_expr(multicast_a): - if n_rank == 0: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFA_stage, - tma_sfa_desc.get_ptr(), - (0, coord_sf_k, sfa_m_block, tile_l_a), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_a, - group=nvvm.CTAGroup.CTA_1, - ) - else: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFA_stage, - tma_sfa_desc.get_ptr(), - (0, coord_sf_k, sfa_m_block, tile_l_a), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_a, - group=nvvm.CTAGroup.CTA_1, - ) - for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) tma_b_desc = tma_b_descs[_bj] sSFB_stage = smem_sfb_list[_bj].subview(sfb_smem_bytes * stage) tma_sfb_desc = tma_sfb_descs[_bj] sfb_n_block = coord_n_per_cta // 128 - if cutlass.const_expr(multicast_b): + if cutlass.const_expr(b_mcast_slices > 1): + _b_rows = cta_tile_mnk[1] // b_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(m_rank * _b_rows * ab_packed_per_row), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + m_rank * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + else: + _b_per_cta = b_mcast_slices // cluster_m + for _bsl in cutlass.range(_b_per_cta): + _b_idx = m_rank * _b_per_cta + _bsl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(_b_idx * _b_rows * ab_packed_per_row), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + _b_idx * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + elif cutlass.const_expr(multicast_b): if m_rank == 0: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): @@ -562,30 +667,6 @@ def _kernel( multicast_mask=tma_mcast_mask_b, group=nvvm.CTAGroup.CTA_1, ) - if cutlass.const_expr(multicast_b): - if m_rank == 0: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFB_stage, - tma_sfb_desc.get_ptr(), - (0, coord_sf_k, sfb_n_block, tile_l_b), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_1, - ) - else: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFB_stage, - tma_sfb_desc.get_ptr(), - (0, coord_sf_k, sfb_n_block, tile_l_b), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_1, - ) - ab_iter += 1 consumer_stage = tile_iter % CLC_SCHED_STAGES @@ -618,14 +699,14 @@ def _kernel( tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -667,7 +748,7 @@ def _kernel( sfa_scale_ptrs = [nvvm.make_tmem_ptr(b, cutlass.Float32) for b in sfa_tmem_bases] sfb_scale_ptrs = [nvvm.make_tmem_ptr(b, cutlass.Float32) for b in sfb_tmem_bases] sfa_dst_ptrs = [ - [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block, cutlass.Float32) for m in range(num_blocks_m)] for i in range(num_a_operands) + [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block, cutlass.Float32) for m in range(num_mma_m)] for i in range(num_a_operands) ] sfb_dst_ptrs = [ [nvvm.make_tmem_ptr(sfb_tmem_bases[j] + m * registers_per_block, cutlass.Float32) for m in range(num_blocks_n)] for j in range(num_b_operands) @@ -708,9 +789,6 @@ def _kernel( if stage == 0 and ab_iter != 0: ab_full_phase_bit = ab_full_phase_bit ^ 1 - while not nvvm.mbarrier_try_wait_parity(ab_full_mbar_ptr.subview(stage), ab_full_phase_bit, time_limit=10_000_000): - pass - desc_a_bases = [ cutlass.experimental.primitives.Tcgen05SmemDesc.build( start_address=smem_a_list[i].subview(sA_elems * stage), @@ -748,56 +826,61 @@ def _kernel( for j in range(num_b_operands) ] - for atom_r in cutlass.range(num_sf_atoms, unroll_full=True): - for _ai in cutlass.range_constexpr(num_a_operands): - for _m in cutlass.range_constexpr(num_blocks_m): - if elect_one: - nvvm.tcgen05_cp( - s2t_shape, - sfa_dst_ptrs[_ai][_m], - desc_sfa_bases[_ai] + (sf_atom_desc_stride * atom_r + sf_block_desc_stride * _m), - group=nvvm.CTAGroup.CTA_1, - multicast=s2t_multicast, - ) + while not nvvm.mbarrier_try_wait_parity(sf_full_mbar_ptr.subview(stage), ab_full_phase_bit, time_limit=10_000_000): + pass + + for sf_word in cutlass.range_constexpr(num_sf_atoms): for _bj in cutlass.range_constexpr(num_b_operands): - for _m in cutlass.range_constexpr(num_blocks_n): + for block_n in cutlass.range_constexpr(num_blocks_n): if elect_one: nvvm.tcgen05_cp( s2t_shape, - sfb_dst_ptrs[_bj][_m], - desc_sfb_bases[_bj] + (sf_atom_desc_stride * atom_r + sf_block_desc_stride * _m), + sfb_dst_ptrs[_bj][block_n], + desc_sfb_bases[_bj] + (sf_atom_desc_stride * sf_word + sf_block_desc_stride * block_n), group=nvvm.CTAGroup.CTA_1, multicast=s2t_multicast, ) - for j in cutlass.range_constexpr(sf_insts_per_atom): - k_block_idx = atom_r * sf_insts_per_atom + j - idesc_k = idesc_by_j[j] - for g in cutlass.range_constexpr(num_gemms): - _ai = gemm_a_idx[g] - _bj = gemm_b_idx[g] - desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) - desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - for mi in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(sf_word == 0): + while not nvvm.mbarrier_try_wait_parity(ab_full_mbar_ptr.subview(stage), ab_full_phase_bit, time_limit=10_000_000): + pass + for mma_k_in_word in cutlass.range_constexpr(sf_insts_per_atom): + mma_k = sf_word * sf_insts_per_atom + mma_k_in_word + idesc_k = idesc_by_j[mma_k_in_word] + for gemm_i in cutlass.range_constexpr(num_gemms): + _ai = gemm_a_idx[gemm_i] + _bj = gemm_b_idx[gemm_i] + desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * mma_k) + desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * mma_k) + for mma_m in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(mma_k_in_word == 0 and _ai not in gemm_a_idx[:gemm_i]): + if elect_one: + nvvm.tcgen05_cp( + s2t_shape, + sfa_dst_ptrs[_ai][mma_m], + desc_sfa_bases[_ai] + (sf_atom_desc_stride * sf_word + sf_block_desc_stride * mma_m), + group=nvvm.CTAGroup.CTA_1, + multicast=s2t_multicast, + ) # The M sub-block offset is a whole SMEM swizzle atom, so the # descriptor's swizzle phase is preserved. B and its SF are # shared; A's SF word block follows the M block. - desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) + desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mma_m) if elect_one: _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_1, - acc_tmem_ptrs[g][mi], + acc_tmem_ptrs[gemm_i][mma_m], desc_a, desc_b, idesc_k, enable_input_d=scale_d, - scale_a=sfa_dst_ptrs[_ai][mi], + scale_a=sfa_dst_ptrs[_ai][mma_m], scale_b=sfb_scale_ptrs[_bj], scale_vec_size=scale_vec_size, - b_collector_op=_b_collector_op(mi), + b_collector_op=_b_collector_op(mma_m), ) # Every accumulator sees scale_d=False on exactly the first - # k_block of the tile, so the flip stays outside mi. + # k_block of the tile, so the flip stays outside mma_m. scale_d = cutlass.Boolean(True) if elect_one: @@ -833,22 +916,20 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity(acc_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass - if cutlass.const_expr(use_acc_overlap): - while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): - pass + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity(acc_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + if cutlass.const_expr(use_acc_overlap): + while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): + pass nvvm.bar_warp_sync(0xFFFFFFFF) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) @@ -875,21 +956,14 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_epi = cutlass.Int32(0) - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - row_id_with_warp_offset = base_row_id - else: - row_id_with_warp_offset = base_row_id + warp_idx * 32 - - # One M block's accumulator columns are contiguous. - subtile_cnt = cute.ceil_div(epi_cols_per_mma_m, 32) - t2r_inst_repx = epi_tile_mn[1] - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - shape = nvvm.Tcgen05LdStShape.SHAPE_16X32BX2 - ld_half_off = 0 - else: - shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B - ld_half_off = None + # @@EPILOGUE_SETUP:BEGIN@@ + row_id_with_warp_offset = base_row_id + warp_idx * 32 + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) + subtile_cnt = len(epi_spans) + shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) @@ -897,7 +971,8 @@ def _kernel( while is_valid != 0: 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] + # @@EPILOGUE_DRAIN:BEGIN@@ + coord_n_c = 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: @@ -912,58 +987,64 @@ def _kernel( else: acc_buf_parity = cutlass.Int32(0) acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # One pass per MMA-M block over its own column region. + for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * mma_inst_shape_mnk[0] - mi_col_base = acc_base_col + mi * epi_cols_per_mma_m + if cutlass.const_expr(use_acc_overlap and num_mma_m > 1): + _mi = mi + (1 - acc_buf_parity) * (num_mma_m - 1 - 2 * mi) + else: + _mi = mi + coord_m = coord_m_tile + _mi * epi_rows_per_mma_m + mi_col_base = acc_base_col + _mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * acc_gemm_stride) for g in range(num_gemms)] - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - row = coord_m + warp_idx * 16 + lane - row_active = lane < 16 - else: - row = coord_m + tidx - row_active = True + row = coord_m + tidx + row_active = True # @@INJECT_AUX_VIEWS@@ - for subtile_idx in cutlass.range(subtile_cnt, unroll_full=True): + for subtile_idx in cutlass.range_constexpr(subtile_cnt): if cutlass.const_expr(use_acc_overlap): _sub = subtile_idx + (1 - acc_buf_parity) * (subtile_cnt - 1 - 2 * subtile_idx) - subtile_col_offset = _sub * 32 + subtile_col_offset = _sub * epi_n + subtile_w = epi_n else: - subtile_col_offset = subtile_idx * 32 + subtile_col_offset, subtile_w = epi_spans[subtile_idx] if cutlass.const_expr(not (use_tma_store_epi and cd_out_is_m_major)): c_rmem_vecs = [] for g in cutlass.range_constexpr(num_gemms): - tmem = cutlass.inttoptr( - tmem_col_addr_gemms[g] + subtile_col_offset, - 6, - cutlass.Float32, - ) - c_rmem_vecs.append(nvvm.tcgen05_ld(shape, tmem, num=t2r_inst_repx, offset=ld_half_off)) + subtile_tmem_addr = tmem_col_addr_gemms[g] + subtile_col_offset + tmem = cutlass.inttoptr(subtile_tmem_addr, 6, mma_c_dtype) + _cv = nvvm.tcgen05_ld(shape, tmem, num=subtile_w) + c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - if use_acc_overlap and (not cd_out_is_m_major) and mi == num_mma_m - 1 and subtile_idx == acc_overlap_subtiles - 1: + if cutlass.const_expr(((not use_acc_overlap) or cd_out_is_m_major) and not (use_tma_store_epi and cd_out_is_m_major)): + if cutlass.const_expr(mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1): + nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) + nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) + if elect_one: + nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) + + if use_acc_overlap and (not cd_out_is_m_major) and mi * subtile_cnt + subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) - col = coord_n + subtile_col_offset + col = coord_n_c + subtile_col_offset # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = (epi_stage_idx + 1) % EPI_SMEM_STAGES smem_subtile_ptr = smem_d_ptr.subview(epi_stage_idx * epi_subtile_elems) - smem_thr_ptr = smem_subtile_ptr.subview(tidx * t2r_inst_repx) + smem_thr_ptr = smem_subtile_ptr.subview(tidx * subtile_w) if cutlass.const_expr(cd_out_is_m_major): ld_col = mi_col_base + subtile_col_offset for _h in cutlass.range(2, unroll_full=True): ld_row = base_row_id + warp_idx * 32 + _h * 16 ld_addr = (ld_row << 16) | ld_col - ld_tmem = cutlass.inttoptr(ld_addr, 6, cutlass.Float32) - _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=4) + ld_tmem = cutlass.inttoptr(ld_addr, 6, mma_c_dtype) + _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=epi_n // 8) vec_f32 = _lv col_j = col linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 @@ -971,7 +1052,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ _i32 = vec_out.bitcast(cutlass.Int32) - for _blk in cutlass.range_constexpr(2): + for _blk in cutlass.range_constexpr(epi_n // 16): _regs = [_i32[_blk * 4 + _j] for _j in range(4)] _n_full = (lane % 8) + 8 * (lane // 16) + 16 * _blk _m_base = warp_idx * 32 + _h * 16 + 8 * ((lane // 8) % 2) @@ -996,7 +1077,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ - smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=cutlass.Swizzle(2, 4, 3)) + smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=epi_smem_swizzle) cute.arch.fence_view_async_shared() nvvm.barrier_cta_sync( @@ -1032,7 +1113,7 @@ def _kernel( # @@STG_ONLY:BEGIN@@ if row_active and row < M: - for j in cutlass.range_constexpr(t2r_inst_repx // vsize): + for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize if col_j + vsize <= N: vec_f32 = c_rmem_vec[j * vsize : (j + 1) * vsize] @@ -1042,12 +1123,14 @@ def _kernel( # @@INJECT_EPILOGUE@@ # @@STG_ONLY:END@@ - if cutlass.const_expr((not use_acc_overlap) or cd_out_is_m_major): + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + if cutlass.const_expr(use_tma_store_epi and cd_out_is_m_major): nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) + # @@EPILOGUE_DRAIN:END@@ consumer_stage = tile_iter % CLC_SCHED_STAGES if consumer_stage == 0 and tile_iter != 0: clc_full_phase_epi = clc_full_phase_epi ^ 1 @@ -1155,7 +1238,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1209,7 +1292,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1260,7 +1343,7 @@ def _host( out_stride_l_0 * cd_dtype.width // 128, ], box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], - swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + swizzle=(epi_tma_swizzle if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), ) tma_c_desc_list = [tma_c_desc] # @@TMA_STORE_ONLY:END@@ 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 0c9382ff8..eee797f4f 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 @@ -26,6 +26,7 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( + epi_subtile_spans as _epi_subtile_spans, l2_swizzle_tile as _l2_swizzle_tile, tcgen05_alloc as _tcgen05_alloc, tcgen05_dealloc as _tcgen05_dealloc, @@ -42,6 +43,9 @@ # @@INJECT_TILE_CONSTANTS@@ +if use_acc_overlap and any(_w != epi_n for _, _w in _epi_subtile_spans(epi_cols_per_mma_m, epi_n)): + raise NotImplementedError(f"{__name__}: acc overlap reverses subtiles by index, which needs a uniform drain width") + # Programmatic Dependent Launch (PDL, sm_90+). USE_PDL = True @@ -194,11 +198,15 @@ def _kernel( a_part_arrive = cutlass.Int16(a_pattern) << m_rank b_part_arrive = cutlass.Int16(b_pattern) << (n_rank * cluster_m) - ab_empty_arrive_mask = a_part_arrive | b_part_arrive + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = a_part_arrive | b_part_arrive _smem_sys_reserved = cutlass.Array(cutlass.Int8, 1024, space=cutlass.AddressSpace.smem, alignment=1) ab_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) + sf_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) ab_empty_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) acc_empty_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) acc_full_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) @@ -255,11 +263,16 @@ def _kernel( ) # @@TMA_STORE_ONLY:END@@ - ab_empty_count = cluster_m + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size + else: + ab_empty_count = cluster_m + cluster_n - 1 if warp_idx == 0: for i in range(ab_stages): if elect_one: nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) + if elect_one: + nvvm.mbarrier_init(sf_full_mbar_ptr.subview(i), 1) if elect_one: nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) for i in range(acc_stages): @@ -281,9 +294,10 @@ def _kernel( sA_bytes = sA_elems * (ab_dtype.width // 8) 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) + ab_only_copy_bytes = num_a_operands * sA_bytes + num_b_operands * sB_bytes + sf_only_copy_bytes = num_a_operands * sfa_smem_bytes + num_b_operands * sfb_smem_bytes - cols_per_acc_stage = cta_tile_mnk[1] + epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m tmem_alloc_bar_count = (num_epilogue_warps + 1) * 32 # @@INJECT_TAP_PTRS@@ @@ -306,8 +320,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) tile_m = init_tile_m @@ -338,7 +351,9 @@ def _kernel( coord_k = k_tile_idx * cta_tile_mnk[2] coord_sf_k = k_tile_idx * sf_tma_box_k if elect_one: - nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) + nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), ab_only_copy_bytes) + if elect_one: + nvvm.mbarrier_arrive_expect_tx(sf_full_mbar_ptr.subview(stage), sf_only_copy_bytes) for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) @@ -347,6 +362,93 @@ def _kernel( tma_sfa_desc = tma_sfa_descs[_ai] sfa_m_block = coord_m_per_cta // 128 if cutlass.const_expr(multicast_a): + if n_rank == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFA_stage, + tma_sfa_desc.get_ptr(), + (0, coord_sf_k, sfa_m_block, tile_l_a), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + else: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFA_stage, + tma_sfa_desc.get_ptr(), + (0, coord_sf_k, sfa_m_block, tile_l_a), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + + for _bj in cutlass.range_constexpr(num_b_operands): + sB_stage = smem_b_list[_bj].subview(sB_elems * stage) + tma_b_desc = tma_b_descs[_bj] + sSFB_stage = smem_sfb_list[_bj].subview(sfb_smem_bytes * stage) + tma_sfb_desc = tma_sfb_descs[_bj] + sfb_n_block = coord_n_per_cta // 128 + if cutlass.const_expr(multicast_b): + if m_rank == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFB_stage, + tma_sfb_desc.get_ptr(), + (0, coord_sf_k, sfb_n_block, tile_l_b), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + else: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFB_stage, + tma_sfb_desc.get_ptr(), + (0, coord_sf_k, sfb_n_block, tile_l_b), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + + for _ai in cutlass.range_constexpr(num_a_operands): + sA_stage = smem_a_list[_ai].subview(sA_elems * stage) + tma_a_desc = tma_a_descs[_ai] + sSFA_stage = smem_sfa_list[_ai].subview(sfa_smem_bytes * stage) + tma_sfa_desc = tma_sfa_descs[_ai] + sfa_m_block = coord_m_per_cta // 128 + if cutlass.const_expr(a_mcast_slices > 1): + _a_rows = cta_tile_mnk[0] // a_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(n_rank * _a_rows * ab_packed_per_row), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + n_rank * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + else: + _a_per_cta = a_mcast_slices // cluster_n + for _asl in cutlass.range(_a_per_cta): + _a_idx = n_rank * _a_per_cta + _asl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(_a_idx * _a_rows * ab_packed_per_row), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + _a_idx * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + elif cutlass.const_expr(multicast_a): if n_rank == 0: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): @@ -403,37 +505,40 @@ def _kernel( multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_1, ) - if cutlass.const_expr(multicast_a): - if n_rank == 0: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFA_stage, - tma_sfa_desc.get_ptr(), - (0, coord_sf_k, sfa_m_block, tile_l_a), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_a, - group=nvvm.CTAGroup.CTA_1, - ) - else: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFA_stage, - tma_sfa_desc.get_ptr(), - (0, coord_sf_k, sfa_m_block, tile_l_a), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_a, - group=nvvm.CTAGroup.CTA_1, - ) - for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) tma_b_desc = tma_b_descs[_bj] sSFB_stage = smem_sfb_list[_bj].subview(sfb_smem_bytes * stage) tma_sfb_desc = tma_sfb_descs[_bj] sfb_n_block = coord_n_per_cta // 128 - if cutlass.const_expr(multicast_b): + if cutlass.const_expr(b_mcast_slices > 1): + _b_rows = cta_tile_mnk[1] // b_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(m_rank * _b_rows * ab_packed_per_row), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + m_rank * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + else: + _b_per_cta = b_mcast_slices // cluster_m + for _bsl in cutlass.range(_b_per_cta): + _b_idx = m_rank * _b_per_cta + _bsl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(_b_idx * _b_rows * ab_packed_per_row), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + _b_idx * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + elif cutlass.const_expr(multicast_b): if m_rank == 0: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): @@ -490,30 +595,6 @@ def _kernel( multicast_mask=tma_mcast_mask_b, group=nvvm.CTAGroup.CTA_1, ) - if cutlass.const_expr(multicast_b): - if m_rank == 0: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFB_stage, - tma_sfb_desc.get_ptr(), - (0, coord_sf_k, sfb_n_block, tile_l_b), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_1, - ) - else: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFB_stage, - tma_sfb_desc.get_ptr(), - (0, coord_sf_k, sfb_n_block, tile_l_b), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_1, - ) - ab_iter += 1 is_valid = cutlass.Int32(0) @@ -523,14 +604,14 @@ def _kernel( tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -573,7 +654,7 @@ def _kernel( sfa_scale_ptrs = [nvvm.make_tmem_ptr(b, cutlass.Float32) for b in sfa_tmem_bases] sfb_scale_ptrs = [nvvm.make_tmem_ptr(b, cutlass.Float32) for b in sfb_tmem_bases] sfa_dst_ptrs = [ - [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block, cutlass.Float32) for m in range(num_blocks_m)] for i in range(num_a_operands) + [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block, cutlass.Float32) for m in range(num_mma_m)] for i in range(num_a_operands) ] sfb_dst_ptrs = [ [nvvm.make_tmem_ptr(sfb_tmem_bases[j] + m * registers_per_block, cutlass.Float32) for m in range(num_blocks_n)] for j in range(num_b_operands) @@ -614,9 +695,6 @@ def _kernel( if stage == 0 and ab_iter != 0: ab_full_phase_bit = ab_full_phase_bit ^ 1 - while not nvvm.mbarrier_try_wait_parity(ab_full_mbar_ptr.subview(stage), ab_full_phase_bit, time_limit=10_000_000): - pass - desc_a_bases = [ cutlass.experimental.primitives.Tcgen05SmemDesc.build( start_address=smem_a_list[i].subview(sA_elems * stage), @@ -654,56 +732,61 @@ def _kernel( for j in range(num_b_operands) ] - for atom_r in cutlass.range(num_sf_atoms, unroll_full=True): - for _ai in cutlass.range_constexpr(num_a_operands): - for _m in cutlass.range_constexpr(num_blocks_m): - if elect_one: - nvvm.tcgen05_cp( - s2t_shape, - sfa_dst_ptrs[_ai][_m], - desc_sfa_bases[_ai] + (sf_atom_desc_stride * atom_r + sf_block_desc_stride * _m), - group=nvvm.CTAGroup.CTA_1, - multicast=s2t_multicast, - ) + while not nvvm.mbarrier_try_wait_parity(sf_full_mbar_ptr.subview(stage), ab_full_phase_bit, time_limit=10_000_000): + pass + + for sf_word in cutlass.range_constexpr(num_sf_atoms): for _bj in cutlass.range_constexpr(num_b_operands): - for _m in cutlass.range_constexpr(num_blocks_n): + for block_n in cutlass.range_constexpr(num_blocks_n): if elect_one: nvvm.tcgen05_cp( s2t_shape, - sfb_dst_ptrs[_bj][_m], - desc_sfb_bases[_bj] + (sf_atom_desc_stride * atom_r + sf_block_desc_stride * _m), + sfb_dst_ptrs[_bj][block_n], + desc_sfb_bases[_bj] + (sf_atom_desc_stride * sf_word + sf_block_desc_stride * block_n), group=nvvm.CTAGroup.CTA_1, multicast=s2t_multicast, ) - for j in cutlass.range_constexpr(sf_insts_per_atom): - k_block_idx = atom_r * sf_insts_per_atom + j - idesc_k = idesc_by_j[j] - for g in cutlass.range_constexpr(num_gemms): - _ai = gemm_a_idx[g] - _bj = gemm_b_idx[g] - desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) - desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - for mi in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(sf_word == 0): + while not nvvm.mbarrier_try_wait_parity(ab_full_mbar_ptr.subview(stage), ab_full_phase_bit, time_limit=10_000_000): + pass + for mma_k_in_word in cutlass.range_constexpr(sf_insts_per_atom): + mma_k = sf_word * sf_insts_per_atom + mma_k_in_word + idesc_k = idesc_by_j[mma_k_in_word] + for gemm_i in cutlass.range_constexpr(num_gemms): + _ai = gemm_a_idx[gemm_i] + _bj = gemm_b_idx[gemm_i] + desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * mma_k) + desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * mma_k) + for mma_m in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(mma_k_in_word == 0 and _ai not in gemm_a_idx[:gemm_i]): + if elect_one: + nvvm.tcgen05_cp( + s2t_shape, + sfa_dst_ptrs[_ai][mma_m], + desc_sfa_bases[_ai] + (sf_atom_desc_stride * sf_word + sf_block_desc_stride * mma_m), + group=nvvm.CTAGroup.CTA_1, + multicast=s2t_multicast, + ) # The M sub-block offset is a whole SMEM swizzle atom, so the # descriptor's swizzle phase is preserved. B and its SF are # shared; A's SF word block follows the M block. - desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) + desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mma_m) if elect_one: _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_1, - acc_tmem_ptrs[g][mi], + acc_tmem_ptrs[gemm_i][mma_m], desc_a, desc_b, idesc_k, enable_input_d=scale_d, - scale_a=sfa_dst_ptrs[_ai][mi], + scale_a=sfa_dst_ptrs[_ai][mma_m], scale_b=sfb_scale_ptrs[_bj], scale_vec_size=scale_vec_size, - b_collector_op=_b_collector_op(mi), + b_collector_op=_b_collector_op(mma_m), ) # Every accumulator sees scale_d=False on exactly the first - # k_block of the tile, so the flip stays outside mi. + # k_block of the tile, so the flip stays outside mma_m. scale_d = cutlass.Boolean(True) if elect_one: @@ -724,21 +807,19 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) - if elect_one: - while not nvvm.mbarrier_try_wait_parity( - acc_empty_mbar_ptr.subview(acc_stage), - acc_empty_phase_bit ^ 1, - time_limit=10_000_000, - ): + while not nvvm.mbarrier_try_wait_parity( + acc_empty_mbar_ptr.subview(acc_stage), + acc_empty_phase_bit ^ 1, + time_limit=10_000_000, + ): + pass + if cutlass.const_expr(use_acc_overlap): + while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): pass - if cutlass.const_expr(use_acc_overlap): - while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): - pass nvvm.bar_warp_sync(0xFFFFFFFF) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) @@ -766,21 +847,14 @@ def _kernel( tile_l = init_tile_l is_valid = cutlass.Int32(1) - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - row_id_with_warp_offset = base_row_id - else: - row_id_with_warp_offset = base_row_id + warp_idx * 32 - - # One M block's accumulator columns are contiguous. - subtile_cnt = cute.ceil_div(epi_cols_per_mma_m, 32) - t2r_inst_repx = epi_tile_mn[1] - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - shape = nvvm.Tcgen05LdStShape.SHAPE_16X32BX2 - ld_half_off = 0 - else: - shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B - ld_half_off = None + # @@EPILOGUE_SETUP:BEGIN@@ + row_id_with_warp_offset = base_row_id + warp_idx * 32 + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) + subtile_cnt = len(epi_spans) + shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) @@ -788,7 +862,8 @@ def _kernel( while is_valid != 0: 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] + # @@EPILOGUE_DRAIN:BEGIN@@ + coord_n_c = 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: @@ -803,59 +878,64 @@ def _kernel( else: acc_buf_parity = cutlass.Int32(0) acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # One pass per MMA-M block over its own column region. + for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * mma_inst_shape_mnk[0] - mi_col_base = acc_base_col + mi * epi_cols_per_mma_m + if cutlass.const_expr(use_acc_overlap and num_mma_m > 1): + _mi = mi + (1 - acc_buf_parity) * (num_mma_m - 1 - 2 * mi) + else: + _mi = mi + coord_m = coord_m_tile + _mi * epi_rows_per_mma_m + mi_col_base = acc_base_col + _mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * acc_gemm_stride) for g in range(num_gemms)] - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - row = coord_m + warp_idx * 16 + lane - row_active = lane < 16 - else: - row = coord_m + tidx - row_active = True + row = coord_m + tidx + row_active = True # @@INJECT_AUX_VIEWS@@ - for subtile_idx in cutlass.range(subtile_cnt, unroll_full=True): + for subtile_idx in cutlass.range_constexpr(subtile_cnt): if cutlass.const_expr(use_acc_overlap): _sub = subtile_idx + (1 - acc_buf_parity) * (subtile_cnt - 1 - 2 * subtile_idx) - subtile_col_offset = _sub * 32 + subtile_col_offset = _sub * epi_n + subtile_w = epi_n else: - subtile_col_offset = subtile_idx * 32 - + subtile_col_offset, subtile_w = epi_spans[subtile_idx] if cutlass.const_expr(not (use_tma_store_epi and cd_out_is_m_major)): c_rmem_vecs = [] for g in cutlass.range_constexpr(num_gemms): - tmem = cutlass.inttoptr( - tmem_col_addr_gemms[g] + subtile_col_offset, - 6, - cutlass.Float32, - ) - c_rmem_vecs.append(nvvm.tcgen05_ld(shape, tmem, num=t2r_inst_repx, offset=ld_half_off)) + subtile_tmem_addr = tmem_col_addr_gemms[g] + subtile_col_offset + tmem = cutlass.inttoptr(subtile_tmem_addr, 6, mma_c_dtype) + _cv = nvvm.tcgen05_ld(shape, tmem, num=subtile_w) + c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - if use_acc_overlap and (not cd_out_is_m_major) and mi == num_mma_m - 1 and subtile_idx == acc_overlap_subtiles - 1: + if cutlass.const_expr(((not use_acc_overlap) or cd_out_is_m_major) and not (use_tma_store_epi and cd_out_is_m_major)): + if cutlass.const_expr(mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1): + nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) + nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) + if elect_one: + nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) + + if use_acc_overlap and (not cd_out_is_m_major) and mi * subtile_cnt + subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) - col = coord_n + subtile_col_offset + col = coord_n_c + subtile_col_offset # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = (epi_stage_idx + 1) % EPI_SMEM_STAGES smem_subtile_ptr = smem_d_ptr.subview(epi_stage_idx * epi_subtile_elems) - smem_thr_ptr = smem_subtile_ptr.subview(tidx * t2r_inst_repx) + smem_thr_ptr = smem_subtile_ptr.subview(tidx * subtile_w) if cutlass.const_expr(cd_out_is_m_major): ld_col = mi_col_base + subtile_col_offset for _h in cutlass.range(2, unroll_full=True): ld_row = base_row_id + warp_idx * 32 + _h * 16 ld_addr = (ld_row << 16) | ld_col - ld_tmem = cutlass.inttoptr(ld_addr, 6, cutlass.Float32) - _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=4) + ld_tmem = cutlass.inttoptr(ld_addr, 6, mma_c_dtype) + _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=epi_n // 8) vec_f32 = _lv col_j = col linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 @@ -863,7 +943,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ _i32 = vec_out.bitcast(cutlass.Int32) - for _blk in cutlass.range_constexpr(2): + for _blk in cutlass.range_constexpr(epi_n // 16): _regs = [_i32[_blk * 4 + _j] for _j in range(4)] _n_full = (lane % 8) + 8 * (lane // 16) + 16 * _blk _m_base = warp_idx * 32 + _h * 16 + 8 * ((lane // 8) % 2) @@ -888,7 +968,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ - smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=cutlass.Swizzle(2, 4, 3)) + smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=epi_smem_swizzle) cute.arch.fence_view_async_shared() nvvm.barrier_cta_sync( @@ -924,7 +1004,7 @@ def _kernel( # @@STG_ONLY:BEGIN@@ if row_active and row < M: - for j in cutlass.range_constexpr(t2r_inst_repx // vsize): + for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize if col_j + vsize <= N: vec_f32 = c_rmem_vec[j * vsize : (j + 1) * vsize] @@ -934,12 +1014,14 @@ def _kernel( # @@INJECT_EPILOGUE@@ # @@STG_ONLY:END@@ - if cutlass.const_expr((not use_acc_overlap) or cd_out_is_m_major): + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + if cutlass.const_expr(use_tma_store_epi and cd_out_is_m_major): nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) + # @@EPILOGUE_DRAIN:END@@ is_valid = cutlass.Int32(0) tile_iter += 1 @@ -1023,7 +1105,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1077,7 +1159,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1129,7 +1211,7 @@ def _host( out_stride_l_0 * cd_dtype.width // 128, ], box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], - swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + swizzle=(epi_tma_swizzle if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), ) tma_c_desc_list = [tma_c_desc] # @@TMA_STORE_ONLY:END@@ 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 3bea40b8f..14b89cb8e 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 @@ -23,6 +23,7 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( + epi_subtile_spans as _epi_subtile_spans, l2_swizzle_tile as _l2_swizzle_tile, tcgen05_alloc as _tcgen05_alloc, tcgen05_dealloc as _tcgen05_dealloc, @@ -40,6 +41,9 @@ # @@INJECT_TILE_CONSTANTS@@ +if use_acc_overlap and any(_w != epi_n for _, _w in _epi_subtile_spans(epi_cols_per_mma_m, epi_n)): + raise NotImplementedError(f"{__name__}: acc overlap reverses subtiles by index, which needs a uniform drain width") + # Scheduler ring depth. CLC_SCHED_STAGES = 2 @@ -197,6 +201,7 @@ def _kernel( _smem_sys_reserved = cutlass.Array(cutlass.Int8, 1024, space=cutlass.AddressSpace.smem, alignment=1) ab_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) + sf_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) ab_empty_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) acc_empty_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) acc_full_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) @@ -268,7 +273,10 @@ def _kernel( acc_empty_count = num_epilogue_warps * 2 cta_group = 2 - ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size // cta_group + else: + ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 num_consumer_warps_per_cta = 7 clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: @@ -281,6 +289,8 @@ def _kernel( for i in range(ab_stages): if elect_one: nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) + if elect_one: + nvvm.mbarrier_init(sf_full_mbar_ptr.subview(i), 1) if elect_one: nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) for i in range(acc_stages): @@ -298,7 +308,8 @@ def _kernel( sA_bytes = sA_elems * (ab_dtype.width // 8) 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 + ab_only_copy_bytes = (num_a_operands * sA_bytes + num_b_operands * sB_bytes) * 2 + sf_only_copy_bytes = (num_a_operands * sfa_smem_bytes + num_b_operands * sfb_smem_bytes) * 2 # Per-CTA logical tile — the cluster cancels out, so these stay compile-time # constants even when the cluster shape is only known at runtime. @@ -308,11 +319,6 @@ def _kernel( # 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 - if cutlass.const_expr(epi_rows_per_mma_m == 64): - # cluster-MMA m=128: the pair also splits N, so each CTA drains N/2. - cols_per_acc_stage = pair_n_size // 2 - else: - cols_per_acc_stage = pair_n_size tmem_alloc_bar_count = (num_epilogue_warps + 1) * 32 nvvm.barrier_cluster_wait() @@ -390,8 +396,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) tile_m = init_tile_m @@ -427,7 +432,9 @@ def _kernel( if is_pair_leader: if elect_one: - nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) + nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), ab_only_copy_bytes) + if elect_one: + nvvm.mbarrier_arrive_expect_tx(sf_full_mbar_ptr.subview(stage), sf_only_copy_bytes) for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) @@ -436,6 +443,92 @@ def _kernel( tma_sfa_desc = tma_sfa_descs[_ai] sfa_m_block = coord_m_per_cta // 128 if cutlass.const_expr(multicast_a): + if n_rank == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFA_stage, + tma_sfa_desc.get_ptr(), + (0, coord_sf_k, sfa_m_block, tile_l_a), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + else: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFA_stage, + tma_sfa_desc.get_ptr(), + (0, coord_sf_k, sfa_m_block, tile_l_a), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + + for _bj in cutlass.range_constexpr(num_b_operands): + sB_stage = smem_b_list[_bj].subview(sB_elems * stage) + tma_b_desc = tma_b_descs[_bj] + sSFB_stage = smem_sfb_list[_bj].subview(sfb_smem_bytes * stage) + tma_sfb_desc = tma_sfb_descs[_bj] + if cutlass.const_expr(multicast_b): + if pair_m_idx == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFB_stage, + tma_sfb_desc.get_ptr(), + (0, coord_sf_k, sfb_n_block, tile_l_b), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + else: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFB_stage, + tma_sfb_desc.get_ptr(), + (0, coord_sf_k, sfb_n_block, tile_l_b), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + + for _ai in cutlass.range_constexpr(num_a_operands): + sA_stage = smem_a_list[_ai].subview(sA_elems * stage) + tma_a_desc = tma_a_descs[_ai] + sSFA_stage = smem_sfa_list[_ai].subview(sfa_smem_bytes * stage) + tma_sfa_desc = tma_sfa_descs[_ai] + sfa_m_block = coord_m_per_cta // 128 + if cutlass.const_expr(a_mcast_slices > 1): + _a_rows = cta_tile_mnk[0] // a_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(n_rank * _a_rows * ab_packed_per_row), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + n_rank * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + else: + _a_per_cta = a_mcast_slices // cluster_n + for _asl in cutlass.range(_a_per_cta): + _a_idx = n_rank * _a_per_cta + _asl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(_a_idx * _a_rows * ab_packed_per_row), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + _a_idx * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + elif cutlass.const_expr(multicast_a): if n_rank == 0: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): @@ -492,36 +585,39 @@ def _kernel( multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_2, ) - if cutlass.const_expr(multicast_a): - if n_rank == 0: + for _bj in cutlass.range_constexpr(num_b_operands): + sB_stage = smem_b_list[_bj].subview(sB_elems * stage) + tma_b_desc = tma_b_descs[_bj] + sSFB_stage = smem_sfb_list[_bj].subview(sfb_smem_bytes * stage) + tma_sfb_desc = tma_sfb_descs[_bj] + if cutlass.const_expr(b_mcast_slices > 1): + _b_rows = cta_tile_mnk[1] // b_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFA_stage, - tma_sfa_desc.get_ptr(), - (0, coord_sf_k, sfa_m_block, tile_l_a), + sB_stage.subview(pair_m_idx * _b_rows * ab_packed_per_row), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + pair_m_idx * _b_rows, tile_l_b), ab_full_mbar_ptr.subview(stage), [], - multicast_mask=tma_mcast_mask_a, + multicast_mask=tma_mcast_mask_b, group=nvvm.CTAGroup.CTA_2, ) - else: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFA_stage, - tma_sfa_desc.get_ptr(), - (0, coord_sf_k, sfa_m_block, tile_l_a), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_a, - group=nvvm.CTAGroup.CTA_2, - ) - - for _bj in cutlass.range_constexpr(num_b_operands): - sB_stage = smem_b_list[_bj].subview(sB_elems * stage) - tma_b_desc = tma_b_descs[_bj] - sSFB_stage = smem_sfb_list[_bj].subview(sfb_smem_bytes * stage) - tma_sfb_desc = tma_sfb_descs[_bj] - if cutlass.const_expr(multicast_b): + else: + _b_per_cta = b_mcast_slices // (cluster_m // cta_group) + for _bsl in cutlass.range(_b_per_cta): + _b_idx = pair_m_idx * _b_per_cta + _bsl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(_b_idx * _b_rows * ab_packed_per_row), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + _b_idx * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + elif cutlass.const_expr(multicast_b): if pair_m_idx == 0: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): @@ -578,30 +674,6 @@ def _kernel( multicast_mask=tma_mcast_mask_b, group=nvvm.CTAGroup.CTA_2, ) - if cutlass.const_expr(multicast_b): - if pair_m_idx == 0: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFB_stage, - tma_sfb_desc.get_ptr(), - (0, coord_sf_k, sfb_n_block, tile_l_b), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_2, - ) - else: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFB_stage, - tma_sfb_desc.get_ptr(), - (0, coord_sf_k, sfb_n_block, tile_l_b), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_2, - ) - ab_iter += 1 consumer_stage = tile_iter % CLC_SCHED_STAGES @@ -634,14 +706,14 @@ def _kernel( tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 pair_mask = cutlass.Int16(3) << pair_leader_rank a_arrive_pattern = a_mcast_pattern @@ -652,7 +724,10 @@ def _kernel( a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) - ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) _tcgen05_alloc( @@ -697,7 +772,7 @@ def _kernel( sfa_scale_ptrs = [nvvm.make_tmem_ptr(b, cutlass.Float32) for b in sfa_tmem_bases] sfb_scale_ptrs = [nvvm.make_tmem_ptr(b, cutlass.Float32) for b in sfb_tmem_bases] sfa_dst_ptrs = [ - [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block, cutlass.Float32) for m in range(num_blocks_m)] for i in range(num_a_operands) + [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block, cutlass.Float32) for m in range(num_mma_m)] for i in range(num_a_operands) ] sfb_dst_ptrs = [ [nvvm.make_tmem_ptr(sfb_tmem_bases[j] + m * registers_per_block, cutlass.Float32) for m in range(num_blocks_n)] for j in range(num_b_operands) @@ -738,13 +813,6 @@ def _kernel( if stage == 0 and ab_iter != 0: ab_full_phase_bit = ab_full_phase_bit ^ 1 - while not nvvm.mbarrier_try_wait_parity( - ab_full_mbar_ptr.subview(stage), - ab_full_phase_bit, - time_limit=10_000_000, - ): - pass - desc_a_bases = [ cutlass.experimental.primitives.Tcgen05SmemDesc.build( start_address=smem_a_list[i].subview(sA_elems * stage), @@ -782,56 +850,69 @@ def _kernel( for j in range(num_b_operands) ] - for atom_r in cutlass.range(num_sf_atoms, unroll_full=True): - for _ai in cutlass.range_constexpr(num_a_operands): - for _m in cutlass.range_constexpr(num_blocks_m): - if elect_one: - nvvm.tcgen05_cp( - s2t_shape, - sfa_dst_ptrs[_ai][_m], - desc_sfa_bases[_ai] + (sf_atom_desc_stride * atom_r + sf_block_desc_stride * _m), - group=nvvm.CTAGroup.CTA_2, - multicast=s2t_multicast, - ) + while not nvvm.mbarrier_try_wait_parity( + sf_full_mbar_ptr.subview(stage), + ab_full_phase_bit, + time_limit=10_000_000, + ): + pass + + for sf_word in cutlass.range_constexpr(num_sf_atoms): for _bj in cutlass.range_constexpr(num_b_operands): - for _m in cutlass.range_constexpr(num_blocks_n): + for block_n in cutlass.range_constexpr(num_blocks_n): if elect_one: nvvm.tcgen05_cp( s2t_shape, - sfb_dst_ptrs[_bj][_m], - desc_sfb_bases[_bj] + (sf_atom_desc_stride * atom_r + sf_block_desc_stride * _m), + sfb_dst_ptrs[_bj][block_n], + desc_sfb_bases[_bj] + (sf_atom_desc_stride * sf_word + sf_block_desc_stride * block_n), group=nvvm.CTAGroup.CTA_2, multicast=s2t_multicast, ) - for j in cutlass.range_constexpr(sf_insts_per_atom): - k_block_idx = atom_r * sf_insts_per_atom + j - idesc_k = idesc_by_j[j] - for g in cutlass.range_constexpr(num_gemms): - _ai = gemm_a_idx[g] - _bj = gemm_b_idx[g] - desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) - desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - for mi in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(sf_word == 0): + while not nvvm.mbarrier_try_wait_parity( + ab_full_mbar_ptr.subview(stage), + ab_full_phase_bit, + time_limit=10_000_000, + ): + pass + for mma_k_in_word in cutlass.range_constexpr(sf_insts_per_atom): + mma_k = sf_word * sf_insts_per_atom + mma_k_in_word + idesc_k = idesc_by_j[mma_k_in_word] + for gemm_i in cutlass.range_constexpr(num_gemms): + _ai = gemm_a_idx[gemm_i] + _bj = gemm_b_idx[gemm_i] + desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * mma_k) + desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * mma_k) + for mma_m in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(mma_k_in_word == 0 and _ai not in gemm_a_idx[:gemm_i]): + if elect_one: + nvvm.tcgen05_cp( + s2t_shape, + sfa_dst_ptrs[_ai][mma_m], + desc_sfa_bases[_ai] + (sf_atom_desc_stride * sf_word + sf_block_desc_stride * mma_m), + group=nvvm.CTAGroup.CTA_2, + multicast=s2t_multicast, + ) # The M sub-block offset is a whole SMEM swizzle atom, so # the descriptor's swizzle phase is preserved. B and its SF # are shared; A's SF word block follows the M block. - desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) + desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mma_m) if elect_one: _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_2, - acc_tmem_ptrs[g][mi], + acc_tmem_ptrs[gemm_i][mma_m], desc_a, desc_b, idesc_k, enable_input_d=scale_d, - scale_a=sfa_dst_ptrs[_ai][mi], + scale_a=sfa_dst_ptrs[_ai][mma_m], scale_b=sfb_scale_ptrs[_bj], scale_vec_size=scale_vec_size, - b_collector_op=_b_collector_op(mi), + b_collector_op=_b_collector_op(mma_m), ) # Every accumulator sees scale_d=False on exactly the first - # k_block of the tile, so the flip stays outside mi. + # k_block of the tile, so the flip stays outside mma_m. scale_d = cutlass.Boolean(True) if elect_one: @@ -868,25 +949,21 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_empty_mbar_ptr.subview(tail_stage), - tail_phase, - time_limit=10_000_000, - ): - pass - nvvm.bar_warp_sync(0xFFFFFFFF) - + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity( + acc_empty_mbar_ptr.subview(tail_stage), + tail_phase, + time_limit=10_000_000, + ): + pass nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): @@ -924,8 +1001,7 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) @@ -959,25 +1035,22 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_epi = cutlass.Int32(0) + # @@EPILOGUE_SETUP:BEGIN@@ row_id_with_warp_offset = base_row_id + warp_idx * 32 - if cutlass.const_expr(cols_per_acc_stage >= 32): - t2r_inst_repx = 32 - subtile_cnt = cols_per_acc_stage // 32 - else: - t2r_inst_repx = cols_per_acc_stage - subtile_cnt = 1 + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) + subtile_cnt = len(epi_spans) shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 - + # @@EPILOGUE_SETUP:END@@ # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) # @@TMA_STORE_ONLY:END@@ while is_valid != 0: coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + # @@EPILOGUE_DRAIN:BEGIN@@ 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 acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: @@ -992,62 +1065,72 @@ def _kernel( else: acc_buf_parity = cutlass.Int32(0) acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # The 2-CTA epilogue drains its own half of the instruction's M, - # epi_rows_per_mma_m rows at a time, so a CTA tile of num_mma_m blocks - # drains in num_mma_m passes over its own column region. + for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * epi_rows_per_mma_m - mi_col_base = acc_base_col + mi * epi_cols_per_mma_m + if cutlass.const_expr(use_acc_overlap and num_mma_m > 1): + _mi = mi + (1 - acc_buf_parity) * (num_mma_m - 1 - 2 * mi) + else: + _mi = mi + coord_m = coord_m_tile + _mi * epi_rows_per_mma_m + mi_col_base = acc_base_col + _mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * acc_gemm_stride) for g in range(num_gemms)] - if cutlass.const_expr(epi_rows_per_mma_m == 64): - row = coord_m + (warp_idx % 2) * 32 + lane - row_active = True - else: - row = coord_m + tidx - row_active = True + row = coord_m + tidx + row_active = True # @@INJECT_AUX_VIEWS@@ - for subtile_idx in cutlass.range(subtile_cnt, unroll_full=True): + for subtile_idx in cutlass.range_constexpr(subtile_cnt): if cutlass.const_expr(use_acc_overlap): _sub = subtile_idx + (1 - acc_buf_parity) * (subtile_cnt - 1 - 2 * subtile_idx) - subtile_col_offset = _sub * t2r_inst_repx + subtile_col_offset = _sub * epi_n + subtile_w = epi_n else: - subtile_col_offset = subtile_idx * t2r_inst_repx - + subtile_col_offset, subtile_w = epi_spans[subtile_idx] if cutlass.const_expr(not (use_tma_store_epi and cd_out_is_m_major)): c_rmem_vecs = [] for g in cutlass.range_constexpr(num_gemms): - tmem = cutlass.inttoptr( - tmem_col_addr_gemms[g] + subtile_col_offset, - 6, - cutlass.Float32, - ) - c_rmem_vecs.append(nvvm.tcgen05_ld(shape, tmem, num=t2r_inst_repx)) + subtile_tmem_addr = tmem_col_addr_gemms[g] + subtile_col_offset + tmem = cutlass.inttoptr(subtile_tmem_addr, 6, mma_c_dtype) + _cv = nvvm.tcgen05_ld(shape, tmem, num=subtile_w) + c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - if use_acc_overlap and (not cd_out_is_m_major) and mi == num_mma_m - 1 and subtile_idx == acc_overlap_subtiles - 1: + if cutlass.const_expr(((not use_acc_overlap) or cd_out_is_m_major) and not (use_tma_store_epi and cd_out_is_m_major)): + if cutlass.const_expr(mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1): + nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) + nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) + if elect_one: + nvvm.mbarrier_arrive( + nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), + scope=nvvm.MemScope.CLUSTER, + relaxed=True, + ) + + if use_acc_overlap and (not cd_out_is_m_major) and mi * subtile_cnt + subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: - mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) - nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) + nvvm.mbarrier_arrive( + nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), + scope=nvvm.MemScope.CLUSTER, + relaxed=True, + ) col = coord_n_c + subtile_col_offset # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = (epi_stage_idx + 1) % EPI_SMEM_STAGES smem_subtile_ptr = smem_d_ptr.subview(epi_stage_idx * epi_subtile_elems) - smem_thr_ptr = smem_subtile_ptr.subview(tidx * t2r_inst_repx) + smem_thr_ptr = smem_subtile_ptr.subview(tidx * subtile_w) if cutlass.const_expr(cd_out_is_m_major): ld_col = mi_col_base + subtile_col_offset for _h in cutlass.range(2, unroll_full=True): ld_row = base_row_id + warp_idx * 32 + _h * 16 ld_addr = (ld_row << 16) | ld_col - ld_tmem = cutlass.inttoptr(ld_addr, 6, cutlass.Float32) - _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=4) + ld_tmem = cutlass.inttoptr(ld_addr, 6, mma_c_dtype) + _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=epi_n // 8) vec_f32 = _lv col_j = col linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 @@ -1055,7 +1138,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ _i32 = vec_out.bitcast(cutlass.Int32) - for _blk in cutlass.range_constexpr(2): + for _blk in cutlass.range_constexpr(epi_n // 16): _regs = [_i32[_blk * 4 + _j] for _j in range(4)] _n_full = (lane % 8) + 8 * (lane // 16) + 16 * _blk _m_base = warp_idx * 32 + _h * 16 + 8 * ((lane // 8) % 2) @@ -1080,7 +1163,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ - smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=cutlass.Swizzle(2, 4, 3)) + smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=epi_smem_swizzle) cute.arch.fence_view_async_shared() nvvm.barrier_cta_sync( @@ -1116,7 +1199,7 @@ def _kernel( # @@STG_ONLY:BEGIN@@ if row_active and row < M: - for j in cutlass.range_constexpr(t2r_inst_repx // vsize): + for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize if col_j + vsize <= N: vec_f32 = c_rmem_vec[j * vsize : (j + 1) * vsize] @@ -1126,13 +1209,15 @@ def _kernel( # @@INJECT_EPILOGUE@@ # @@STG_ONLY:END@@ - if cutlass.const_expr((not use_acc_overlap) or cd_out_is_m_major): + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + if cutlass.const_expr(use_tma_store_epi and cd_out_is_m_major): nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) + # @@EPILOGUE_DRAIN:END@@ consumer_stage = tile_iter % CLC_SCHED_STAGES if consumer_stage == 0 and tile_iter != 0: clc_full_phase_epi = clc_full_phase_epi ^ 1 @@ -1239,7 +1324,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1293,7 +1378,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1345,7 +1430,7 @@ def _host( out_stride_l_0 * cd_dtype.width // 128, ], box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], - swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + swizzle=(epi_tma_swizzle if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), ) tma_c_desc_list = [tma_c_desc] # @@TMA_STORE_ONLY:END@@ 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 ea1f16861..db5e22bd6 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 @@ -24,6 +24,7 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( + epi_subtile_spans as _epi_subtile_spans, l2_swizzle_tile as _l2_swizzle_tile, tcgen05_alloc as _tcgen05_alloc, tcgen05_dealloc as _tcgen05_dealloc, @@ -40,6 +41,9 @@ # @@INJECT_TILE_CONSTANTS@@ +if use_acc_overlap and any(_w != epi_n for _, _w in _epi_subtile_spans(epi_cols_per_mma_m, epi_n)): + raise NotImplementedError(f"{__name__}: acc overlap reverses subtiles by index, which needs a uniform drain width") + # Programmatic Dependent Launch (PDL, sm_90+). USE_PDL = True @@ -194,6 +198,7 @@ def _kernel( _smem_sys_reserved = cutlass.Array(cutlass.Int8, 1024, space=cutlass.AddressSpace.smem, alignment=1) ab_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) + sf_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) ab_empty_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) acc_empty_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) acc_full_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) @@ -251,7 +256,10 @@ def _kernel( acc_empty_count = num_epilogue_warps * 2 cta_group = 2 - ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size // cta_group + else: + ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 if warp_idx == 0: if cutlass.const_expr(use_acc_overlap): if elect_one: @@ -262,6 +270,8 @@ def _kernel( for i in range(ab_stages): if elect_one: nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) + if elect_one: + nvvm.mbarrier_init(sf_full_mbar_ptr.subview(i), 1) if elect_one: nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) for i in range(acc_stages): @@ -274,7 +284,8 @@ def _kernel( sA_bytes = sA_elems * (ab_dtype.width // 8) 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 + ab_only_copy_bytes = (num_a_operands * sA_bytes + num_b_operands * sB_bytes) * 2 + sf_only_copy_bytes = (num_a_operands * sfa_smem_bytes + num_b_operands * sfb_smem_bytes) * 2 # Per-CTA logical tile — the cluster cancels out, so these stay compile-time # constants even when the cluster shape is only known at runtime. @@ -284,11 +295,6 @@ def _kernel( # 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 - if cutlass.const_expr(epi_rows_per_mma_m == 64): - # cluster-MMA m=128: the pair also splits N, so each CTA drains N/2. - cols_per_acc_stage = pair_n_size // 2 - else: - cols_per_acc_stage = pair_n_size tmem_alloc_bar_count = (num_epilogue_warps + 1) * 32 nvvm.barrier_cluster_wait() @@ -314,8 +320,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) tile_m = init_tile_m @@ -350,7 +355,9 @@ def _kernel( if is_pair_leader: if elect_one: - nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) + nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), ab_only_copy_bytes) + if elect_one: + nvvm.mbarrier_arrive_expect_tx(sf_full_mbar_ptr.subview(stage), sf_only_copy_bytes) for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) @@ -359,6 +366,92 @@ def _kernel( tma_sfa_desc = tma_sfa_descs[_ai] sfa_m_block = coord_m_per_cta // 128 if cutlass.const_expr(multicast_a): + if n_rank == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFA_stage, + tma_sfa_desc.get_ptr(), + (0, coord_sf_k, sfa_m_block, tile_l_a), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + else: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFA_stage, + tma_sfa_desc.get_ptr(), + (0, coord_sf_k, sfa_m_block, tile_l_a), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + + for _bj in cutlass.range_constexpr(num_b_operands): + sB_stage = smem_b_list[_bj].subview(sB_elems * stage) + tma_b_desc = tma_b_descs[_bj] + sSFB_stage = smem_sfb_list[_bj].subview(sfb_smem_bytes * stage) + tma_sfb_desc = tma_sfb_descs[_bj] + if cutlass.const_expr(multicast_b): + if pair_m_idx == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFB_stage, + tma_sfb_desc.get_ptr(), + (0, coord_sf_k, sfb_n_block, tile_l_b), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + else: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFB_stage, + tma_sfb_desc.get_ptr(), + (0, coord_sf_k, sfb_n_block, tile_l_b), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + + for _ai in cutlass.range_constexpr(num_a_operands): + sA_stage = smem_a_list[_ai].subview(sA_elems * stage) + tma_a_desc = tma_a_descs[_ai] + sSFA_stage = smem_sfa_list[_ai].subview(sfa_smem_bytes * stage) + tma_sfa_desc = tma_sfa_descs[_ai] + sfa_m_block = coord_m_per_cta // 128 + if cutlass.const_expr(a_mcast_slices > 1): + _a_rows = cta_tile_mnk[0] // a_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(n_rank * _a_rows * ab_packed_per_row), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + n_rank * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + else: + _a_per_cta = a_mcast_slices // cluster_n + for _asl in cutlass.range(_a_per_cta): + _a_idx = n_rank * _a_per_cta + _asl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(_a_idx * _a_rows * ab_packed_per_row), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + _a_idx * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + elif cutlass.const_expr(multicast_a): if n_rank == 0: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): @@ -415,36 +508,39 @@ def _kernel( multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_2, ) - if cutlass.const_expr(multicast_a): - if n_rank == 0: + for _bj in cutlass.range_constexpr(num_b_operands): + sB_stage = smem_b_list[_bj].subview(sB_elems * stage) + tma_b_desc = tma_b_descs[_bj] + sSFB_stage = smem_sfb_list[_bj].subview(sfb_smem_bytes * stage) + tma_sfb_desc = tma_sfb_descs[_bj] + if cutlass.const_expr(b_mcast_slices > 1): + _b_rows = cta_tile_mnk[1] // b_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFA_stage, - tma_sfa_desc.get_ptr(), - (0, coord_sf_k, sfa_m_block, tile_l_a), + sB_stage.subview(pair_m_idx * _b_rows * ab_packed_per_row), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + pair_m_idx * _b_rows, tile_l_b), ab_full_mbar_ptr.subview(stage), [], - multicast_mask=tma_mcast_mask_a, + multicast_mask=tma_mcast_mask_b, group=nvvm.CTAGroup.CTA_2, ) - else: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFA_stage, - tma_sfa_desc.get_ptr(), - (0, coord_sf_k, sfa_m_block, tile_l_a), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_a, - group=nvvm.CTAGroup.CTA_2, - ) - - for _bj in cutlass.range_constexpr(num_b_operands): - sB_stage = smem_b_list[_bj].subview(sB_elems * stage) - tma_b_desc = tma_b_descs[_bj] - sSFB_stage = smem_sfb_list[_bj].subview(sfb_smem_bytes * stage) - tma_sfb_desc = tma_sfb_descs[_bj] - if cutlass.const_expr(multicast_b): + else: + _b_per_cta = b_mcast_slices // (cluster_m // cta_group) + for _bsl in cutlass.range(_b_per_cta): + _b_idx = pair_m_idx * _b_per_cta + _bsl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(_b_idx * _b_rows * ab_packed_per_row), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + _b_idx * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + elif cutlass.const_expr(multicast_b): if pair_m_idx == 0: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): @@ -501,30 +597,6 @@ def _kernel( multicast_mask=tma_mcast_mask_b, group=nvvm.CTAGroup.CTA_2, ) - if cutlass.const_expr(multicast_b): - if pair_m_idx == 0: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFB_stage, - tma_sfb_desc.get_ptr(), - (0, coord_sf_k, sfb_n_block, tile_l_b), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_2, - ) - else: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFB_stage, - tma_sfb_desc.get_ptr(), - (0, coord_sf_k, sfb_n_block, tile_l_b), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_2, - ) - ab_iter += 1 is_valid = cutlass.Int32(0) @@ -534,14 +606,14 @@ def _kernel( tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 pair_mask = cutlass.Int16(3) << pair_leader_rank a_arrive_pattern = a_mcast_pattern @@ -552,7 +624,10 @@ def _kernel( a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) - ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) _tcgen05_alloc( @@ -596,7 +671,7 @@ def _kernel( sfa_scale_ptrs = [nvvm.make_tmem_ptr(b, cutlass.Float32) for b in sfa_tmem_bases] sfb_scale_ptrs = [nvvm.make_tmem_ptr(b, cutlass.Float32) for b in sfb_tmem_bases] sfa_dst_ptrs = [ - [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block, cutlass.Float32) for m in range(num_blocks_m)] for i in range(num_a_operands) + [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block, cutlass.Float32) for m in range(num_mma_m)] for i in range(num_a_operands) ] sfb_dst_ptrs = [ [nvvm.make_tmem_ptr(sfb_tmem_bases[j] + m * registers_per_block, cutlass.Float32) for m in range(num_blocks_n)] for j in range(num_b_operands) @@ -637,13 +712,6 @@ def _kernel( if stage == 0 and ab_iter != 0: ab_full_phase_bit = ab_full_phase_bit ^ 1 - while not nvvm.mbarrier_try_wait_parity( - ab_full_mbar_ptr.subview(stage), - ab_full_phase_bit, - time_limit=10_000_000, - ): - pass - desc_a_bases = [ cutlass.experimental.primitives.Tcgen05SmemDesc.build( start_address=smem_a_list[i].subview(sA_elems * stage), @@ -682,56 +750,69 @@ def _kernel( for j in range(num_b_operands) ] - for atom_r in cutlass.range(num_sf_atoms, unroll_full=True): - for _ai in cutlass.range_constexpr(num_a_operands): - for _m in cutlass.range_constexpr(num_blocks_m): - if elect_one: - nvvm.tcgen05_cp( - s2t_shape, - sfa_dst_ptrs[_ai][_m], - desc_sfa_bases[_ai] + (sf_atom_desc_stride * atom_r + sf_block_desc_stride * _m), - group=nvvm.CTAGroup.CTA_2, - multicast=s2t_multicast, - ) + while not nvvm.mbarrier_try_wait_parity( + sf_full_mbar_ptr.subview(stage), + ab_full_phase_bit, + time_limit=10_000_000, + ): + pass + + for sf_word in cutlass.range_constexpr(num_sf_atoms): for _bj in cutlass.range_constexpr(num_b_operands): - for _m in cutlass.range_constexpr(num_blocks_n): + for block_n in cutlass.range_constexpr(num_blocks_n): if elect_one: nvvm.tcgen05_cp( s2t_shape, - sfb_dst_ptrs[_bj][_m], - desc_sfb_bases[_bj] + (sf_atom_desc_stride * atom_r + sf_block_desc_stride * _m), + sfb_dst_ptrs[_bj][block_n], + desc_sfb_bases[_bj] + (sf_atom_desc_stride * sf_word + sf_block_desc_stride * block_n), group=nvvm.CTAGroup.CTA_2, multicast=s2t_multicast, ) - for j in cutlass.range_constexpr(sf_insts_per_atom): - k_block_idx = atom_r * sf_insts_per_atom + j - idesc_k = idesc_by_j[j] - for g in cutlass.range_constexpr(num_gemms): - _ai = gemm_a_idx[g] - _bj = gemm_b_idx[g] - desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) - desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - for mi in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(sf_word == 0): + while not nvvm.mbarrier_try_wait_parity( + ab_full_mbar_ptr.subview(stage), + ab_full_phase_bit, + time_limit=10_000_000, + ): + pass + for mma_k_in_word in cutlass.range_constexpr(sf_insts_per_atom): + mma_k = sf_word * sf_insts_per_atom + mma_k_in_word + idesc_k = idesc_by_j[mma_k_in_word] + for gemm_i in cutlass.range_constexpr(num_gemms): + _ai = gemm_a_idx[gemm_i] + _bj = gemm_b_idx[gemm_i] + desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * mma_k) + desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * mma_k) + for mma_m in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(mma_k_in_word == 0 and _ai not in gemm_a_idx[:gemm_i]): + if elect_one: + nvvm.tcgen05_cp( + s2t_shape, + sfa_dst_ptrs[_ai][mma_m], + desc_sfa_bases[_ai] + (sf_atom_desc_stride * sf_word + sf_block_desc_stride * mma_m), + group=nvvm.CTAGroup.CTA_2, + multicast=s2t_multicast, + ) # The M sub-block offset is a whole SMEM swizzle atom, so # the descriptor's swizzle phase is preserved. B and its SF # are shared; A's SF word block follows the M block. - desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) + desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mma_m) if elect_one: _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_2, - acc_tmem_ptrs[g][mi], + acc_tmem_ptrs[gemm_i][mma_m], desc_a, desc_b, idesc_k, enable_input_d=scale_d, - scale_a=sfa_dst_ptrs[_ai][mi], + scale_a=sfa_dst_ptrs[_ai][mma_m], scale_b=sfb_scale_ptrs[_bj], scale_vec_size=scale_vec_size, - b_collector_op=_b_collector_op(mi), + b_collector_op=_b_collector_op(mma_m), ) # Every accumulator sees scale_d=False on exactly the first - # k_block of the tile, so the flip stays outside mi. + # k_block of the tile, so the flip stays outside mma_m. scale_d = cutlass.Boolean(True) if elect_one: @@ -753,25 +834,21 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_empty_mbar_ptr.subview(tail_stage), - tail_phase, - time_limit=10_000_000, - ): - pass - nvvm.bar_warp_sync(0xFFFFFFFF) - + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity( + acc_empty_mbar_ptr.subview(tail_stage), + tail_phase, + time_limit=10_000_000, + ): + pass nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): @@ -787,8 +864,7 @@ def _kernel( ) else: if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) @@ -821,24 +897,22 @@ def _kernel( tile_l = init_tile_l is_valid = cutlass.Int32(1) + # @@EPILOGUE_SETUP:BEGIN@@ row_id_with_warp_offset = base_row_id + warp_idx * 32 - if cutlass.const_expr(cols_per_acc_stage >= 32): - t2r_inst_repx = 32 - subtile_cnt = cols_per_acc_stage // 32 - else: - t2r_inst_repx = cols_per_acc_stage - subtile_cnt = 1 + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) + subtile_cnt = len(epi_spans) shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) # @@TMA_STORE_ONLY:END@@ while is_valid != 0: coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + # @@EPILOGUE_DRAIN:BEGIN@@ 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 acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: @@ -853,62 +927,72 @@ def _kernel( else: acc_buf_parity = cutlass.Int32(0) acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # The 2-CTA epilogue drains its own half of the instruction's M, - # epi_rows_per_mma_m rows at a time, so a CTA tile of num_mma_m blocks - # drains in num_mma_m passes over its own column region. + for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * epi_rows_per_mma_m - mi_col_base = acc_base_col + mi * epi_cols_per_mma_m + if cutlass.const_expr(use_acc_overlap and num_mma_m > 1): + _mi = mi + (1 - acc_buf_parity) * (num_mma_m - 1 - 2 * mi) + else: + _mi = mi + coord_m = coord_m_tile + _mi * epi_rows_per_mma_m + mi_col_base = acc_base_col + _mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * acc_gemm_stride) for g in range(num_gemms)] - if cutlass.const_expr(epi_rows_per_mma_m == 64): - row = coord_m + (warp_idx % 2) * 32 + lane - row_active = True - else: - row = coord_m + tidx - row_active = True + row = coord_m + tidx + row_active = True # @@INJECT_AUX_VIEWS@@ - for subtile_idx in cutlass.range(subtile_cnt, unroll_full=True): + for subtile_idx in cutlass.range_constexpr(subtile_cnt): if cutlass.const_expr(use_acc_overlap): _sub = subtile_idx + (1 - acc_buf_parity) * (subtile_cnt - 1 - 2 * subtile_idx) - subtile_col_offset = _sub * t2r_inst_repx + subtile_col_offset = _sub * epi_n + subtile_w = epi_n else: - subtile_col_offset = subtile_idx * t2r_inst_repx - + subtile_col_offset, subtile_w = epi_spans[subtile_idx] if cutlass.const_expr(not (use_tma_store_epi and cd_out_is_m_major)): c_rmem_vecs = [] for g in cutlass.range_constexpr(num_gemms): - tmem = cutlass.inttoptr( - tmem_col_addr_gemms[g] + subtile_col_offset, - 6, - cutlass.Float32, - ) - c_rmem_vecs.append(nvvm.tcgen05_ld(shape, tmem, num=t2r_inst_repx)) + subtile_tmem_addr = tmem_col_addr_gemms[g] + subtile_col_offset + tmem = cutlass.inttoptr(subtile_tmem_addr, 6, mma_c_dtype) + _cv = nvvm.tcgen05_ld(shape, tmem, num=subtile_w) + c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - if use_acc_overlap and (not cd_out_is_m_major) and mi == num_mma_m - 1 and subtile_idx == acc_overlap_subtiles - 1: + if cutlass.const_expr(((not use_acc_overlap) or cd_out_is_m_major) and not (use_tma_store_epi and cd_out_is_m_major)): + if cutlass.const_expr(mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1): + nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) + nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) + if elect_one: + nvvm.mbarrier_arrive( + nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), + scope=nvvm.MemScope.CLUSTER, + relaxed=True, + ) + + if use_acc_overlap and (not cd_out_is_m_major) and mi * subtile_cnt + subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: - mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) - nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) + nvvm.mbarrier_arrive( + nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), + scope=nvvm.MemScope.CLUSTER, + relaxed=True, + ) col = coord_n_c + subtile_col_offset # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = (epi_stage_idx + 1) % EPI_SMEM_STAGES smem_subtile_ptr = smem_d_ptr.subview(epi_stage_idx * epi_subtile_elems) - smem_thr_ptr = smem_subtile_ptr.subview(tidx * t2r_inst_repx) + smem_thr_ptr = smem_subtile_ptr.subview(tidx * subtile_w) if cutlass.const_expr(cd_out_is_m_major): ld_col = mi_col_base + subtile_col_offset for _h in cutlass.range(2, unroll_full=True): ld_row = base_row_id + warp_idx * 32 + _h * 16 ld_addr = (ld_row << 16) | ld_col - ld_tmem = cutlass.inttoptr(ld_addr, 6, cutlass.Float32) - _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=4) + ld_tmem = cutlass.inttoptr(ld_addr, 6, mma_c_dtype) + _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=epi_n // 8) vec_f32 = _lv col_j = col linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 @@ -916,7 +1000,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ _i32 = vec_out.bitcast(cutlass.Int32) - for _blk in cutlass.range_constexpr(2): + for _blk in cutlass.range_constexpr(epi_n // 16): _regs = [_i32[_blk * 4 + _j] for _j in range(4)] _n_full = (lane % 8) + 8 * (lane // 16) + 16 * _blk _m_base = warp_idx * 32 + _h * 16 + 8 * ((lane // 8) % 2) @@ -941,7 +1025,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ - smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=cutlass.Swizzle(2, 4, 3)) + smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=epi_smem_swizzle) cute.arch.fence_view_async_shared() nvvm.barrier_cta_sync( @@ -977,7 +1061,7 @@ def _kernel( # @@STG_ONLY:BEGIN@@ if row_active and row < M: - for j in cutlass.range_constexpr(t2r_inst_repx // vsize): + for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize if col_j + vsize <= N: vec_f32 = c_rmem_vec[j * vsize : (j + 1) * vsize] @@ -987,13 +1071,15 @@ def _kernel( # @@INJECT_EPILOGUE@@ # @@STG_ONLY:END@@ - if cutlass.const_expr((not use_acc_overlap) or cd_out_is_m_major): + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + if cutlass.const_expr(use_tma_store_epi and cd_out_is_m_major): nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) + # @@EPILOGUE_DRAIN:END@@ is_valid = cutlass.Int32(0) tile_iter += 1 @@ -1077,7 +1163,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1131,7 +1217,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1183,7 +1269,7 @@ def _host( out_stride_l_0 * cd_dtype.width // 128, ], box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], - swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + swizzle=(epi_tma_swizzle if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), ) tma_c_desc_list = [tma_c_desc] # @@TMA_STORE_ONLY:END@@ 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 db413c039..9748adf23 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_1ctamma.py @@ -181,7 +181,10 @@ def _kernel( a_part_arrive = cutlass.Int16(a_pattern) << m_rank b_part_arrive = cutlass.Int16(b_pattern) << (n_rank * cluster_m) - ab_empty_arrive_mask = a_part_arrive | b_part_arrive + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = a_part_arrive | b_part_arrive _smem_sys_reserved = cutlass.Array(cutlass.Int8, 1024, space=cutlass.AddressSpace.smem, alignment=1) @@ -237,7 +240,10 @@ def _kernel( ) # @@TMA_STORE_ONLY:END@@ - ab_empty_count = cluster_m + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size + else: + ab_empty_count = cluster_m + cluster_n - 1 num_consumer_warps_per_cta = 7 clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: @@ -285,6 +291,7 @@ def _kernel( # [g*cols_per_acc_stage + mi*epi_cols_per_mma_m + ni*mma_inst_shape_mnk[1], +N) # all at TMEM lane base 0. The N blocks tile one M block's column range, so # the epilogue drains a whole M block as one contiguous span. + epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m epi_cols_per_mma_m = cta_tile_mnk[1] cols_per_acc_stage = num_mma_m * epi_cols_per_mma_m acc_region_cols = num_gemms * cols_per_acc_stage @@ -363,8 +370,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) tile_m = init_tile_m @@ -400,7 +406,34 @@ def _kernel( for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) tma_a_desc = tma_a_descs[_ai] - if cutlass.const_expr(multicast_a): + if cutlass.const_expr(a_mcast_slices > 1): + _a_rows = cta_tile_mnk[0] // a_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(n_rank * _a_rows * cta_tile_mnk[2]), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + n_rank * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + else: + _a_per_cta = a_mcast_slices // cluster_n + for _asl in cutlass.range(_a_per_cta): + _a_idx = n_rank * _a_per_cta + _asl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(_a_idx * _a_rows * cta_tile_mnk[2]), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + _a_idx * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + elif cutlass.const_expr(multicast_a): if n_rank == 0: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): @@ -461,7 +494,34 @@ def _kernel( for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) tma_b_desc = tma_b_descs[_bj] - if cutlass.const_expr(multicast_b): + if cutlass.const_expr(b_mcast_slices > 1): + _b_rows = cta_tile_mnk[1] // b_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(m_rank * _b_rows * cta_tile_mnk[2]), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + m_rank * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + else: + _b_per_cta = b_mcast_slices // cluster_m + for _bsl in cutlass.range(_b_per_cta): + _b_idx = m_rank * _b_per_cta + _bsl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(_b_idx * _b_rows * cta_tile_mnk[2]), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + _b_idx * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + elif cutlass.const_expr(multicast_b): if m_rank == 0: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): @@ -551,14 +611,14 @@ def _kernel( tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -683,21 +743,18 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity(acc_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity(acc_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass - nvvm.bar_warp_sync(0xFFFFFFFF) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) _tcgen05_dealloc( @@ -725,18 +782,13 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_epi = cutlass.Int32(0) - # The drain layout follows the MMA INSTRUCTION's M (the TMEM row - # organisation), not the CTA tile's: at hardware M=64 the accumulator - # occupies data paths 0-15 of each sub-partition, so the 32-data-path - # LDTM layout cannot be used to read it. + # @@EPILOGUE_SETUP:BEGIN@@ if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): row_id_with_warp_offset = base_row_id else: row_id_with_warp_offset = base_row_id + warp_idx * 32 - # One M block's accumulator columns are contiguous (the N-direction MMAs - # tile them), so one span list drains all of them. - epi_spans = _epi_subtile_spans(epi_cols_per_mma_m) + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) subtile_cnt = len(epi_spans) if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): shape = nvvm.Tcgen05LdStShape.SHAPE_16X32BX2 @@ -745,6 +797,7 @@ def _kernel( shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B ld_half_off = None lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) @@ -752,7 +805,8 @@ def _kernel( while is_valid != 0: 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] + # @@EPILOGUE_DRAIN:BEGIN@@ + coord_n_c = 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: @@ -763,11 +817,8 @@ def _kernel( acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # One pass per MMA-M block: the 4 epilogue warps cover mma_inst M - # rows at a time, so a CTA tile of num_mma_m blocks drains in - # num_mma_m passes over its own column region. for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * mma_inst_shape_mnk[0] + coord_m = coord_m_tile + mi * epi_rows_per_mma_m mi_col_base = acc_base_col + mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * cols_per_acc_stage) for g in range(num_gemms)] @@ -796,15 +847,13 @@ def _kernel( c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - # Exactly one acc_empty arrive per epilogue warp per tile — - # the mbar counts warps, not M blocks. - if (not use_tma_store_epi) and mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1: + if (not (use_tma_store_epi and cd_out_is_m_major)) and mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) - col = coord_n + subtile_col_offset + col = coord_n_c + subtile_col_offset # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = (epi_stage_idx + 1) % EPI_SMEM_STAGES @@ -817,7 +866,7 @@ def _kernel( ld_row = base_row_id + warp_idx * 32 + _h * 16 ld_addr = (ld_row << 16) | ld_col ld_tmem = cutlass.inttoptr(ld_addr, 6, mma_c_dtype) - _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=4) + _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=epi_n // 8) if cutlass.const_expr(acc_widen_to_fp32): _accf = _lv.to(cutlass.Float32) _lv = _accf + cutlass.full_like(_accf, 0.0) @@ -828,7 +877,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ _i32 = vec_out.bitcast(cutlass.Int32) - for _blk in cutlass.range_constexpr(2): + for _blk in cutlass.range_constexpr(epi_n // 16): _regs = [_i32[_blk * 4 + _j] for _j in range(4)] _n_full = (lane % 8) + 8 * (lane // 16) + 16 * _blk _m_base = warp_idx * 32 + _h * 16 + 8 * ((lane // 8) % 2) @@ -853,7 +902,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ - smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=cutlass.Swizzle(2, 4, 3)) + smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=epi_smem_swizzle) cute.arch.fence_view_async_shared() nvvm.barrier_cta_sync( @@ -899,12 +948,14 @@ def _kernel( # @@INJECT_EPILOGUE@@ # @@STG_ONLY:END@@ - if cutlass.const_expr(use_tma_store_epi): + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + if cutlass.const_expr(use_tma_store_epi and cd_out_is_m_major): nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) + # @@EPILOGUE_DRAIN:END@@ consumer_stage = tile_iter % CLC_SCHED_STAGES if consumer_stage == 0 and tile_iter != 0: clc_full_phase_epi = clc_full_phase_epi ^ 1 @@ -1016,7 +1067,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, ) ) @@ -1047,7 +1098,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, ) ) @@ -1077,7 +1128,7 @@ def _host( out_stride_l_0 * cd_dtype.width // 128, ], box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], - swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + swizzle=(epi_tma_swizzle if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), ) tma_c_desc_list = [tma_c_desc] # @@TMA_STORE_ONLY:END@@ 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 ac0d6876d..bc34ba115 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 @@ -178,7 +178,10 @@ def _kernel( a_part_arrive = cutlass.Int16(a_pattern) << m_rank b_part_arrive = cutlass.Int16(b_pattern) << (n_rank * cluster_m) - ab_empty_arrive_mask = a_part_arrive | b_part_arrive + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = a_part_arrive | b_part_arrive _smem_sys_reserved = cutlass.Array(cutlass.Int8, 1024, space=cutlass.AddressSpace.smem, alignment=1) @@ -220,7 +223,10 @@ def _kernel( ) # @@TMA_STORE_ONLY:END@@ - ab_empty_count = cluster_m + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size + else: + ab_empty_count = cluster_m + cluster_n - 1 if warp_idx == 0: for i in range(ab_stages): if elect_one: @@ -259,6 +265,7 @@ def _kernel( # [g*cols_per_acc_stage + mi*epi_cols_per_mma_m + ni*mma_inst_shape_mnk[1], +N) # all at TMEM lane base 0. The N blocks tile one M block's column range, so # the epilogue drains a whole M block as one contiguous span. + epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m epi_cols_per_mma_m = cta_tile_mnk[1] cols_per_acc_stage = num_mma_m * epi_cols_per_mma_m acc_region_cols = num_gemms * cols_per_acc_stage @@ -284,8 +291,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) tile_m = init_tile_m @@ -320,7 +326,34 @@ def _kernel( for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) tma_a_desc = tma_a_descs[_ai] - if cutlass.const_expr(multicast_a): + if cutlass.const_expr(a_mcast_slices > 1): + _a_rows = cta_tile_mnk[0] // a_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(n_rank * _a_rows * cta_tile_mnk[2]), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + n_rank * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + else: + _a_per_cta = a_mcast_slices // cluster_n + for _asl in cutlass.range(_a_per_cta): + _a_idx = n_rank * _a_per_cta + _asl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(_a_idx * _a_rows * cta_tile_mnk[2]), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + _a_idx * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + elif cutlass.const_expr(multicast_a): if n_rank == 0: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): @@ -381,7 +414,34 @@ def _kernel( for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) tma_b_desc = tma_b_descs[_bj] - if cutlass.const_expr(multicast_b): + if cutlass.const_expr(b_mcast_slices > 1): + _b_rows = cta_tile_mnk[1] // b_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(m_rank * _b_rows * cta_tile_mnk[2]), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + m_rank * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + else: + _b_per_cta = b_mcast_slices // cluster_m + for _bsl in cutlass.range(_b_per_cta): + _b_idx = m_rank * _b_per_cta + _bsl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(_b_idx * _b_rows * cta_tile_mnk[2]), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + _b_idx * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + elif cutlass.const_expr(multicast_b): if m_rank == 0: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): @@ -448,14 +508,14 @@ def _kernel( tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -564,18 +624,15 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") - if elect_one: - while not nvvm.mbarrier_try_wait_parity( - acc_empty_mbar_ptr.subview(acc_stage), - acc_empty_phase_bit ^ 1, - time_limit=10_000_000, - ): - pass + while not nvvm.mbarrier_try_wait_parity( + acc_empty_mbar_ptr.subview(acc_stage), + acc_empty_phase_bit ^ 1, + time_limit=10_000_000, + ): + pass - nvvm.bar_warp_sync(0xFFFFFFFF) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) _tcgen05_dealloc( @@ -600,18 +657,13 @@ def _kernel( tile_l = init_tile_l is_valid = cutlass.Int32(1) - # The drain layout follows the MMA INSTRUCTION's M (the TMEM row - # organisation), not the CTA tile's: at hardware M=64 the accumulator - # occupies data paths 0-15 of each sub-partition, so the 32-data-path - # LDTM layout cannot be used to read it. + # @@EPILOGUE_SETUP:BEGIN@@ if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): row_id_with_warp_offset = base_row_id else: row_id_with_warp_offset = base_row_id + warp_idx * 32 - # One M block's accumulator columns are contiguous (the N-direction MMAs - # tile them), so one span list drains all of them. - epi_spans = _epi_subtile_spans(epi_cols_per_mma_m) + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) subtile_cnt = len(epi_spans) if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): shape = nvvm.Tcgen05LdStShape.SHAPE_16X32BX2 @@ -620,6 +672,7 @@ def _kernel( shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B ld_half_off = None lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) @@ -627,7 +680,8 @@ def _kernel( while is_valid != 0: 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] + # @@EPILOGUE_DRAIN:BEGIN@@ + coord_n_c = 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: @@ -638,11 +692,8 @@ def _kernel( acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # One pass per MMA-M block: the 4 epilogue warps cover mma_inst M - # rows at a time, so a CTA tile of num_mma_m blocks drains in - # num_mma_m passes over its own column region. for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * mma_inst_shape_mnk[0] + coord_m = coord_m_tile + mi * epi_rows_per_mma_m mi_col_base = acc_base_col + mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * cols_per_acc_stage) for g in range(num_gemms)] @@ -657,7 +708,6 @@ def _kernel( for subtile_idx in cutlass.range_constexpr(subtile_cnt): subtile_col_offset, subtile_w = epi_spans[subtile_idx] - if cutlass.const_expr(not (use_tma_store_epi and cd_out_is_m_major)): c_rmem_vecs = [] for g in cutlass.range_constexpr(num_gemms): @@ -672,15 +722,13 @@ def _kernel( c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - # Exactly one acc_empty arrive per epilogue warp per tile — - # the mbar counts warps, not M blocks. - if (not use_tma_store_epi) and mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1: + if (not (use_tma_store_epi and cd_out_is_m_major)) and mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) - col = coord_n + subtile_col_offset + col = coord_n_c + subtile_col_offset # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = (epi_stage_idx + 1) % EPI_SMEM_STAGES @@ -693,7 +741,7 @@ def _kernel( ld_row = base_row_id + warp_idx * 32 + _h * 16 ld_addr = (ld_row << 16) | ld_col ld_tmem = cutlass.inttoptr(ld_addr, 6, mma_c_dtype) - _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=4) + _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=epi_n // 8) if cutlass.const_expr(acc_widen_to_fp32): _accf = _lv.to(cutlass.Float32) _lv = _accf + cutlass.full_like(_accf, 0.0) @@ -704,7 +752,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ _i32 = vec_out.bitcast(cutlass.Int32) - for _blk in cutlass.range_constexpr(2): + for _blk in cutlass.range_constexpr(epi_n // 16): _regs = [_i32[_blk * 4 + _j] for _j in range(4)] _n_full = (lane % 8) + 8 * (lane // 16) + 16 * _blk _m_base = warp_idx * 32 + _h * 16 + 8 * ((lane // 8) % 2) @@ -729,7 +777,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ - smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=cutlass.Swizzle(2, 4, 3)) + smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=epi_smem_swizzle) cute.arch.fence_view_async_shared() nvvm.barrier_cta_sync( @@ -775,12 +823,14 @@ def _kernel( # @@INJECT_EPILOGUE@@ # @@STG_ONLY:END@@ - if cutlass.const_expr(use_tma_store_epi): + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + if cutlass.const_expr(use_tma_store_epi and cd_out_is_m_major): nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) + # @@EPILOGUE_DRAIN:END@@ is_valid = cutlass.Int32(0) tile_iter += 1 @@ -867,7 +917,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, ) ) @@ -898,7 +948,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, ) ) @@ -928,7 +978,7 @@ def _host( out_stride_l_0 * cd_dtype.width // 128, ], box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], - swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + swizzle=(epi_tma_swizzle if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), ) tma_c_desc_list = [tma_c_desc] # @@TMA_STORE_ONLY:END@@ 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 892ac691b..20147be8c 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_matmul_2ctamma.py @@ -236,7 +236,10 @@ def _kernel( acc_empty_count = num_epilogue_warps * 2 cta_group = 2 - ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size // cta_group + else: + ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 num_consumer_warps_per_cta = 7 clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: @@ -371,8 +374,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) tile_m = init_tile_m @@ -409,7 +411,34 @@ def _kernel( for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) tma_a_desc = tma_a_descs[_ai] - if cutlass.const_expr(multicast_a): + if cutlass.const_expr(a_mcast_slices > 1): + _a_rows = cta_tile_mnk[0] // a_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(n_rank * _a_rows * cta_tile_mnk[2]), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + n_rank * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + else: + _a_per_cta = a_mcast_slices // cluster_n + for _asl in cutlass.range(_a_per_cta): + _a_idx = n_rank * _a_per_cta + _asl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(_a_idx * _a_rows * cta_tile_mnk[2]), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + _a_idx * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + elif cutlass.const_expr(multicast_a): if n_rank == 0: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): @@ -470,7 +499,34 @@ def _kernel( for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) tma_b_desc = tma_b_descs[_bj] - if cutlass.const_expr(multicast_b): + if cutlass.const_expr(b_mcast_slices > 1): + _b_rows = cta_tile_mnk[1] // b_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(pair_m_idx * _b_rows * cta_tile_mnk[2]), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + pair_m_idx * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + else: + _b_per_cta = b_mcast_slices // (cluster_m // cta_group) + for _bsl in cutlass.range(_b_per_cta): + _b_idx = pair_m_idx * _b_per_cta + _bsl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(_b_idx * _b_rows * cta_tile_mnk[2]), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + _b_idx * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + elif cutlass.const_expr(multicast_b): if pair_m_idx == 0: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): @@ -560,14 +616,14 @@ def _kernel( tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 pair_mask = cutlass.Int16(3) << pair_leader_rank a_arrive_pattern = a_mcast_pattern @@ -578,7 +634,10 @@ def _kernel( a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) - ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) _tcgen05_alloc( @@ -710,25 +769,21 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_empty_mbar_ptr.subview(tail_stage), - tail_phase, - time_limit=10_000_000, - ): - pass - nvvm.bar_warp_sync(0xFFFFFFFF) - + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity( + acc_empty_mbar_ptr.subview(tail_stage), + tail_phase, + time_limit=10_000_000, + ): + pass nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): @@ -765,8 +820,7 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) @@ -798,19 +852,21 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_epi = cutlass.Int32(0) + # @@EPILOGUE_SETUP:BEGIN@@ row_id_with_warp_offset = base_row_id + warp_idx * 32 - # One M block's accumulator columns are contiguous, so one span list - # drains all of them. - epi_spans = _epi_subtile_spans(epi_cols_per_mma_m) + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) subtile_cnt = len(epi_spans) shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) # @@TMA_STORE_ONLY:END@@ while is_valid != 0: coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + # @@EPILOGUE_DRAIN:BEGIN@@ 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 @@ -824,9 +880,6 @@ def _kernel( acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # One pass per MMA-M block: the 4 epilogue warps cover this CTA's - # epi_rows_per_mma_m rows at a time, so a CTA tile of num_mma_m blocks - # drains in num_mma_m passes over its own column region. for mi in cutlass.range_constexpr(num_mma_m): coord_m = coord_m_tile + mi * epi_rows_per_mma_m mi_col_base = acc_base_col + mi * epi_cols_per_mma_m @@ -857,9 +910,7 @@ def _kernel( c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - # Exactly one acc_empty arrive per epilogue warp per tile — - # the mbar counts warps, not M blocks. - if (not use_tma_store_epi) and mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1: + if (not (use_tma_store_epi and cd_out_is_m_major)) and mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: @@ -882,7 +933,7 @@ def _kernel( ld_row = base_row_id + warp_idx * 32 + _h * 16 ld_addr = (ld_row << 16) | ld_col ld_tmem = cutlass.inttoptr(ld_addr, 6, mma_c_dtype) - _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=4) + _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=epi_n // 8) if cutlass.const_expr(acc_widen_to_fp32): _accf = _lv.to(cutlass.Float32) _lv = _accf + cutlass.full_like(_accf, 0.0) @@ -893,7 +944,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ _i32 = vec_out.bitcast(cutlass.Int32) - for _blk in cutlass.range_constexpr(2): + for _blk in cutlass.range_constexpr(epi_n // 16): _regs = [_i32[_blk * 4 + _j] for _j in range(4)] _n_full = (lane % 8) + 8 * (lane // 16) + 16 * _blk _m_base = warp_idx * 32 + _h * 16 + 8 * ((lane // 8) % 2) @@ -918,7 +969,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ - smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=cutlass.Swizzle(2, 4, 3)) + smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=epi_smem_swizzle) cute.arch.fence_view_async_shared() nvvm.barrier_cta_sync( @@ -964,13 +1015,15 @@ def _kernel( # @@INJECT_EPILOGUE@@ # @@STG_ONLY:END@@ - if cutlass.const_expr(use_tma_store_epi): + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + if cutlass.const_expr(use_tma_store_epi and cd_out_is_m_major): nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) + # @@EPILOGUE_DRAIN:END@@ consumer_stage = tile_iter % CLC_SCHED_STAGES if consumer_stage == 0 and tile_iter != 0: clc_full_phase_epi = clc_full_phase_epi ^ 1 @@ -1081,7 +1134,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cgrp_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cgrp_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, ) ) @@ -1112,7 +1165,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cgrp_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cgrp_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, ) ) @@ -1142,7 +1195,7 @@ def _host( out_stride_l_0 * cd_dtype.width // 128, ], box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], - swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + swizzle=(epi_tma_swizzle if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), ) tma_c_desc_list = [tma_c_desc] # @@TMA_STORE_ONLY:END@@ 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 6a4a8facb..015538e47 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 @@ -218,7 +218,10 @@ def _kernel( acc_empty_count = num_epilogue_warps * 2 cta_group = 2 - ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size // cta_group + else: + ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 if warp_idx == 0: if elect_one: nvvm.mbarrier_init(tmem_dealloc_mbar_ptr, 32) @@ -293,8 +296,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) tile_m = init_tile_m @@ -330,7 +332,34 @@ def _kernel( for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) tma_a_desc = tma_a_descs[_ai] - if cutlass.const_expr(multicast_a): + if cutlass.const_expr(a_mcast_slices > 1): + _a_rows = cta_tile_mnk[0] // a_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(n_rank * _a_rows * cta_tile_mnk[2]), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + n_rank * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + else: + _a_per_cta = a_mcast_slices // cluster_n + for _asl in cutlass.range(_a_per_cta): + _a_idx = n_rank * _a_per_cta + _asl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(_a_idx * _a_rows * cta_tile_mnk[2]), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + _a_idx * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + elif cutlass.const_expr(multicast_a): if n_rank == 0: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): @@ -391,7 +420,34 @@ def _kernel( for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) tma_b_desc = tma_b_descs[_bj] - if cutlass.const_expr(multicast_b): + if cutlass.const_expr(b_mcast_slices > 1): + _b_rows = cta_tile_mnk[1] // b_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(pair_m_idx * _b_rows * cta_tile_mnk[2]), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + pair_m_idx * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + else: + _b_per_cta = b_mcast_slices // (cluster_m // cta_group) + for _bsl in cutlass.range(_b_per_cta): + _b_idx = pair_m_idx * _b_per_cta + _bsl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(_b_idx * _b_rows * cta_tile_mnk[2]), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + _b_idx * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + elif cutlass.const_expr(multicast_b): if pair_m_idx == 0: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): @@ -458,14 +514,14 @@ def _kernel( tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 pair_mask = cutlass.Int16(3) << pair_leader_rank a_arrive_pattern = a_mcast_pattern @@ -476,7 +532,10 @@ def _kernel( a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) - ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) _tcgen05_alloc( @@ -592,25 +651,21 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_empty_mbar_ptr.subview(tail_stage), - tail_phase, - time_limit=10_000_000, - ): - pass - nvvm.bar_warp_sync(0xFFFFFFFF) - + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity( + acc_empty_mbar_ptr.subview(tail_stage), + tail_phase, + time_limit=10_000_000, + ): + pass nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): @@ -625,8 +680,7 @@ def _kernel( ) else: if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) @@ -658,13 +712,14 @@ def _kernel( tile_l = init_tile_l is_valid = cutlass.Int32(1) + # @@EPILOGUE_SETUP:BEGIN@@ row_id_with_warp_offset = base_row_id + warp_idx * 32 - # One M block's accumulator columns are contiguous, so one span list - # drains all of them. - epi_spans = _epi_subtile_spans(epi_cols_per_mma_m) + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) subtile_cnt = len(epi_spans) shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) @@ -672,6 +727,7 @@ def _kernel( while is_valid != 0: coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + # @@EPILOGUE_DRAIN:BEGIN@@ 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 @@ -685,9 +741,6 @@ def _kernel( acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # One pass per MMA-M block: the 4 epilogue warps cover this CTA's - # epi_rows_per_mma_m rows at a time, so a CTA tile of num_mma_m blocks - # drains in num_mma_m passes over its own column region. for mi in cutlass.range_constexpr(num_mma_m): coord_m = coord_m_tile + mi * epi_rows_per_mma_m mi_col_base = acc_base_col + mi * epi_cols_per_mma_m @@ -704,7 +757,6 @@ def _kernel( for subtile_idx in cutlass.range_constexpr(subtile_cnt): subtile_col_offset, subtile_w = epi_spans[subtile_idx] - if cutlass.const_expr(not (use_tma_store_epi and cd_out_is_m_major)): c_rmem_vecs = [] for g in cutlass.range_constexpr(num_gemms): @@ -719,9 +771,7 @@ def _kernel( c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - # Exactly one acc_empty arrive per epilogue warp per tile — - # the mbar counts warps, not M blocks. - if (not use_tma_store_epi) and mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1: + if (not (use_tma_store_epi and cd_out_is_m_major)) and mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: @@ -744,7 +794,7 @@ def _kernel( ld_row = base_row_id + warp_idx * 32 + _h * 16 ld_addr = (ld_row << 16) | ld_col ld_tmem = cutlass.inttoptr(ld_addr, 6, mma_c_dtype) - _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=4) + _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=epi_n // 8) if cutlass.const_expr(acc_widen_to_fp32): _accf = _lv.to(cutlass.Float32) _lv = _accf + cutlass.full_like(_accf, 0.0) @@ -755,7 +805,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ _i32 = vec_out.bitcast(cutlass.Int32) - for _blk in cutlass.range_constexpr(2): + for _blk in cutlass.range_constexpr(epi_n // 16): _regs = [_i32[_blk * 4 + _j] for _j in range(4)] _n_full = (lane % 8) + 8 * (lane // 16) + 16 * _blk _m_base = warp_idx * 32 + _h * 16 + 8 * ((lane // 8) % 2) @@ -780,7 +830,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ - smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=cutlass.Swizzle(2, 4, 3)) + smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=epi_smem_swizzle) cute.arch.fence_view_async_shared() nvvm.barrier_cta_sync( @@ -826,13 +876,15 @@ def _kernel( # @@INJECT_EPILOGUE@@ # @@STG_ONLY:END@@ - if cutlass.const_expr(use_tma_store_epi): + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + if cutlass.const_expr(use_tma_store_epi and cd_out_is_m_major): nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) + # @@EPILOGUE_DRAIN:END@@ is_valid = cutlass.Int32(0) tile_iter += 1 @@ -919,7 +971,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cgrp_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cgrp_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, ) ) @@ -950,7 +1002,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cgrp_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cgrp_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, ) ) @@ -979,7 +1031,7 @@ def _host( out_stride_l_0 * cd_dtype.width // 128, ], box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], - swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + swizzle=(epi_tma_swizzle if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), ) tma_c_desc_list = [tma_c_desc] # @@TMA_STORE_ONLY:END@@ 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 3d55f8a99..088fc86a8 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 @@ -180,7 +180,10 @@ def _kernel( a_part_arrive = cutlass.Int16(a_pattern) << m_rank b_part_arrive = cutlass.Int16(b_pattern) << (n_rank * cluster_m) - ab_empty_arrive_mask = a_part_arrive | b_part_arrive + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = a_part_arrive | b_part_arrive _smem_sys_reserved = cutlass.Array(cutlass.Int8, 1024, space=cutlass.AddressSpace.smem, alignment=1) @@ -236,7 +239,10 @@ def _kernel( ) # @@TMA_STORE_ONLY:END@@ - ab_empty_count = cluster_m + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size + else: + ab_empty_count = cluster_m + cluster_n - 1 num_consumer_warps_per_cta = 7 + num_mainloop_warps clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: @@ -277,7 +283,6 @@ def _kernel( sB_tma_bytes = sB_elems * (ab_load_b_dtype.width // 8) else: sB_tma_bytes = sB_bytes - num_tma_copy_bytes = sA_tma_bytes + sB_tma_bytes idesc = cutlass.experimental.primitives.Tcgen05InstrDesc.build( a_dtype=mma_a_dtype, @@ -292,6 +297,7 @@ def _kernel( # TMEM accumulator layout, per acc stage: # M block mi, N block ni -> columns # [mi*epi_cols_per_mma_m + ni*mma_inst_shape_mnk[1], +N), all at lane base 0. + epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m epi_cols_per_mma_m = cta_tile_mnk[1] cols_per_acc_stage = num_mma_m * epi_cols_per_mma_m acc_region_cols = num_gemms * cols_per_acc_stage @@ -370,8 +376,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) tile_m = init_tile_m @@ -416,7 +421,34 @@ def _kernel( if elect_one: nvvm.mbarrier_arrive_expect_tx(b_full_mbar_ptr.subview(stage), sB_tma_bytes) - if cutlass.const_expr(multicast_a): + if cutlass.const_expr(a_mcast_slices > 1): + _a_rows = cta_tile_mnk[0] // a_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_tma_dst.subview(n_rank * _a_rows * cta_tile_mnk[2]), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + n_rank * _a_rows, tile_l_a), + a_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + else: + _a_per_cta = a_mcast_slices // cluster_n + for _asl in cutlass.range(_a_per_cta): + _a_idx = n_rank * _a_per_cta + _asl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_tma_dst.subview(_a_idx * _a_rows * cta_tile_mnk[2]), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + _a_idx * _a_rows, tile_l_a), + a_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + elif cutlass.const_expr(multicast_a): if n_rank == 0: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): @@ -474,7 +506,34 @@ def _kernel( group=nvvm.CTAGroup.CTA_1, ) - if cutlass.const_expr(multicast_b): + if cutlass.const_expr(b_mcast_slices > 1): + _b_rows = cta_tile_mnk[1] // b_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_tma_dst.subview(m_rank * _b_rows * cta_tile_mnk[2]), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + m_rank * _b_rows, tile_l_b), + b_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + else: + _b_per_cta = b_mcast_slices // cluster_m + for _bsl in cutlass.range(_b_per_cta): + _b_idx = m_rank * _b_per_cta + _bsl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_tma_dst.subview(_b_idx * _b_rows * cta_tile_mnk[2]), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + _b_idx * _b_rows, tile_l_b), + b_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + elif cutlass.const_expr(multicast_b): if m_rank == 0: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): @@ -564,14 +623,14 @@ def _kernel( tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -709,21 +768,18 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity(acc_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity(acc_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass - nvvm.bar_warp_sync(0xFFFFFFFF) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) _tcgen05_dealloc( @@ -948,18 +1004,13 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_epi = cutlass.Int32(0) - # The drain layout follows the MMA INSTRUCTION's M (the TMEM row - # organisation), not the CTA tile's: at hardware M=64 the accumulator - # occupies data paths 0-15 of each sub-partition, so the 32-data-path - # LDTM layout cannot be used to read it. + # @@EPILOGUE_SETUP:BEGIN@@ if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): row_id_with_warp_offset = base_row_id else: row_id_with_warp_offset = base_row_id + warp_idx * 32 - # One M block's accumulator columns are contiguous (the N-direction MMAs - # tile them), so one span list drains all of them. - epi_spans = _epi_subtile_spans(epi_cols_per_mma_m) + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) subtile_cnt = len(epi_spans) if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): shape = nvvm.Tcgen05LdStShape.SHAPE_16X32BX2 @@ -968,6 +1019,7 @@ def _kernel( shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B ld_half_off = None lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) @@ -975,7 +1027,8 @@ def _kernel( while is_valid != 0: 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] + # @@EPILOGUE_DRAIN:BEGIN@@ + coord_n_c = 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: @@ -986,11 +1039,8 @@ def _kernel( acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # One pass per MMA-M block: the 4 epilogue warps cover mma_inst M - # rows at a time, so a CTA tile of num_mma_m blocks drains in - # num_mma_m passes over its own column region. for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * mma_inst_shape_mnk[0] + coord_m = coord_m_tile + mi * epi_rows_per_mma_m mi_col_base = acc_base_col + mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * cols_per_acc_stage) for g in range(num_gemms)] @@ -1019,15 +1069,13 @@ def _kernel( c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - # Exactly one acc_empty arrive per epilogue warp per tile — - # the mbar counts warps, not M blocks. - if (not use_tma_store_epi) and mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1: + if (not (use_tma_store_epi and cd_out_is_m_major)) and mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) - col = coord_n + subtile_col_offset + col = coord_n_c + subtile_col_offset # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = (epi_stage_idx + 1) % EPI_SMEM_STAGES @@ -1040,7 +1088,7 @@ def _kernel( ld_row = base_row_id + warp_idx * 32 + _h * 16 ld_addr = (ld_row << 16) | ld_col ld_tmem = cutlass.inttoptr(ld_addr, 6, mma_c_dtype) - _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=4) + _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=epi_n // 8) if cutlass.const_expr(acc_widen_to_fp32): _accf = _lv.to(cutlass.Float32) _lv = _accf + cutlass.full_like(_accf, 0.0) @@ -1051,7 +1099,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ _i32 = vec_out.bitcast(cutlass.Int32) - for _blk in cutlass.range_constexpr(2): + for _blk in cutlass.range_constexpr(epi_n // 16): _regs = [_i32[_blk * 4 + _j] for _j in range(4)] _n_full = (lane % 8) + 8 * (lane // 16) + 16 * _blk _m_base = warp_idx * 32 + _h * 16 + 8 * ((lane // 8) % 2) @@ -1076,7 +1124,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ - smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=cutlass.Swizzle(2, 4, 3)) + smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=epi_smem_swizzle) cute.arch.fence_view_async_shared() nvvm.barrier_cta_sync( @@ -1122,12 +1170,14 @@ def _kernel( # @@INJECT_EPILOGUE@@ # @@STG_ONLY:END@@ - if cutlass.const_expr(use_tma_store_epi): + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + if cutlass.const_expr(use_tma_store_epi and cd_out_is_m_major): nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) + # @@EPILOGUE_DRAIN:END@@ consumer_stage = tile_iter % CLC_SCHED_STAGES if consumer_stage == 0 and tile_iter != 0: clc_full_phase_epi = clc_full_phase_epi ^ 1 @@ -1211,7 +1261,7 @@ def _host( a.iterator, cute.make_layout((m, k_sym, a_batch), stride=(a_stride_m, 1, a_stride_l)), ) - a_box_dims = (cta_tile_mnk[0], cta_tile_mnk[2], 1) + a_box_dims = (cta_tile_mnk[0] // a_mcast_slices, cta_tile_mnk[2], 1) a_stride_order = (1, 0, 2) if cutlass.const_expr(b_is_n_major): b_tma_tensor = cute.make_tensor( @@ -1225,7 +1275,7 @@ def _host( b.iterator, cute.make_layout((n, k_sym, b_batch), stride=(b_stride_n, 1, b_stride_l)), ) - b_box_dims = (cta_tile_mnk[1], cta_tile_mnk[2], 1) + b_box_dims = (cta_tile_mnk[1] // b_mcast_slices, cta_tile_mnk[2], 1) b_stride_order = (1, 0, 2) if cutlass.const_expr(mainloop_a_cast): a_desc_dtype = ab_load_a_dtype @@ -1281,7 +1331,7 @@ def _host( out_stride_l_0 * cd_dtype.width // 128, ], box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], - swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + swizzle=(epi_tma_swizzle if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), ) tma_c_desc_list = [tma_c_desc] # @@TMA_STORE_ONLY:END@@ 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 190ee4d49..207b7dc18 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 @@ -238,7 +238,10 @@ def _kernel( acc_empty_count = num_epilogue_warps * 2 cta_group = 2 - ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size // cta_group + else: + ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 num_consumer_warps_per_cta = 7 + num_mainloop_warps clc_empty_count = num_consumer_warps_per_cta * cluster_size mainloop_full_count = num_mainloop_warps * 2 @@ -277,7 +280,6 @@ def _kernel( sB_tma_bytes = sB_elems * (ab_load_b_dtype.width // 8) else: sB_tma_bytes = sB_bytes - num_tma_copy_bytes = (sA_tma_bytes + sB_tma_bytes) * 2 idesc = cutlass.experimental.primitives.Tcgen05InstrDesc.build( a_dtype=mma_a_dtype, @@ -387,8 +389,7 @@ def _kernel( nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) tile_m = init_tile_m @@ -472,6 +473,33 @@ def _kernel( multicast_mask=a_self_mask, group=nvvm.CTAGroup.CTA_1, ) + elif cutlass.const_expr(a_mcast_slices > 1): + _a_rows = cta_tile_mnk[0] // a_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_tma_dst.subview(n_rank * _a_rows * cta_tile_mnk[2]), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + n_rank * _a_rows, tile_l_a), + a_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + else: + _a_per_cta = a_mcast_slices // cluster_n + for _asl in cutlass.range(_a_per_cta): + _a_idx = n_rank * _a_per_cta + _asl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_tma_dst.subview(_a_idx * _a_rows * cta_tile_mnk[2]), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + _a_idx * _a_rows, tile_l_a), + a_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) elif cutlass.const_expr(multicast_a): if n_rank == 0: if cutlass.const_expr(a_is_m_major): @@ -559,6 +587,33 @@ def _kernel( multicast_mask=b_self_mask, group=nvvm.CTAGroup.CTA_1, ) + elif cutlass.const_expr(b_mcast_slices > 1): + _b_rows = cta_tile_mnk[1] // b_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_tma_dst.subview(pair_m_idx * _b_rows * cta_tile_mnk[2]), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + pair_m_idx * _b_rows, tile_l_b), + b_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + else: + _b_per_cta = b_mcast_slices // (cluster_m // cta_group) + for _bsl in cutlass.range(_b_per_cta): + _b_idx = pair_m_idx * _b_per_cta + _bsl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_tma_dst.subview(_b_idx * _b_rows * cta_tile_mnk[2]), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + _b_idx * _b_rows, tile_l_b), + b_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) elif cutlass.const_expr(multicast_b): if pair_m_idx == 0: if cutlass.const_expr(b_is_n_major): @@ -649,14 +704,14 @@ def _kernel( tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 pair_mask = cutlass.Int16(3) << pair_leader_rank a_arrive_pattern = a_mcast_pattern @@ -667,7 +722,10 @@ def _kernel( a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) - ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) _tcgen05_alloc( @@ -808,25 +866,21 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_empty_mbar_ptr.subview(tail_stage), - tail_phase, - time_limit=10_000_000, - ): - pass - nvvm.bar_warp_sync(0xFFFFFFFF) - + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity( + acc_empty_mbar_ptr.subview(tail_stage), + tail_phase, + time_limit=10_000_000, + ): + pass nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): @@ -863,8 +917,7 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) @@ -1091,19 +1144,21 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_epi = cutlass.Int32(0) + # @@EPILOGUE_SETUP:BEGIN@@ row_id_with_warp_offset = base_row_id + warp_idx * 32 - # One M block's accumulator columns are contiguous, so one span list - # drains all of them. - epi_spans = _epi_subtile_spans(epi_cols_per_mma_m) + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) subtile_cnt = len(epi_spans) shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) # @@TMA_STORE_ONLY:END@@ while is_valid != 0: coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + # @@EPILOGUE_DRAIN:BEGIN@@ 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 @@ -1117,9 +1172,6 @@ def _kernel( acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # One pass per MMA-M block: the 4 epilogue warps cover this CTA's - # epi_rows_per_mma_m rows at a time, so a CTA tile of num_mma_m blocks - # drains in num_mma_m passes over its own column region. for mi in cutlass.range_constexpr(num_mma_m): coord_m = coord_m_tile + mi * epi_rows_per_mma_m mi_col_base = acc_base_col + mi * epi_cols_per_mma_m @@ -1150,9 +1202,7 @@ def _kernel( c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - # Exactly one acc_empty arrive per epilogue warp per tile — - # the mbar counts warps, not M blocks. - if (not use_tma_store_epi) and mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1: + if (not (use_tma_store_epi and cd_out_is_m_major)) and mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: @@ -1175,7 +1225,7 @@ def _kernel( ld_row = base_row_id + warp_idx * 32 + _h * 16 ld_addr = (ld_row << 16) | ld_col ld_tmem = cutlass.inttoptr(ld_addr, 6, mma_c_dtype) - _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=4) + _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=epi_n // 8) if cutlass.const_expr(acc_widen_to_fp32): _accf = _lv.to(cutlass.Float32) _lv = _accf + cutlass.full_like(_accf, 0.0) @@ -1186,7 +1236,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ _i32 = vec_out.bitcast(cutlass.Int32) - for _blk in cutlass.range_constexpr(2): + for _blk in cutlass.range_constexpr(epi_n // 16): _regs = [_i32[_blk * 4 + _j] for _j in range(4)] _n_full = (lane % 8) + 8 * (lane // 16) + 16 * _blk _m_base = warp_idx * 32 + _h * 16 + 8 * ((lane // 8) % 2) @@ -1211,7 +1261,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ - smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=cutlass.Swizzle(2, 4, 3)) + smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=epi_smem_swizzle) cute.arch.fence_view_async_shared() nvvm.barrier_cta_sync( @@ -1257,13 +1307,15 @@ def _kernel( # @@INJECT_EPILOGUE@@ # @@STG_ONLY:END@@ - if cutlass.const_expr(use_tma_store_epi): + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + if cutlass.const_expr(use_tma_store_epi and cd_out_is_m_major): nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) + # @@EPILOGUE_DRAIN:END@@ consumer_stage = tile_iter % CLC_SCHED_STAGES if consumer_stage == 0 and tile_iter != 0: clc_full_phase_epi = clc_full_phase_epi ^ 1 @@ -1347,7 +1399,7 @@ def _host( a.iterator, cute.make_layout((m, k_sym, a_batch), stride=(a_stride_m, 1, a_stride_l)), ) - a_box_dims = (cta_tile_mnk[0], cgrp_tile_mnk[2], 1) + a_box_dims = (cta_tile_mnk[0] // a_mcast_slices, cgrp_tile_mnk[2], 1) a_stride_order = (1, 0, 2) if cutlass.const_expr(b_is_n_major): b_tma_tensor = cute.make_tensor( @@ -1361,7 +1413,7 @@ def _host( b.iterator, cute.make_layout((n, k_sym, b_batch), stride=(b_stride_n, 1, b_stride_l)), ) - b_box_dims = (cta_tile_mnk[1], cgrp_tile_mnk[2], 1) + b_box_dims = (cta_tile_mnk[1] // b_mcast_slices, cgrp_tile_mnk[2], 1) b_stride_order = (1, 0, 2) if cutlass.const_expr(mainloop_a_cast): a_desc_dtype = ab_load_a_dtype @@ -1417,7 +1469,7 @@ def _host( out_stride_l_0 * cd_dtype.width // 128, ], box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], - swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + swizzle=(epi_tma_swizzle if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), ) tma_c_desc_list = [tma_c_desc] # @@TMA_STORE_ONLY:END@@ diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_1ctamma.py index c59ccd773..8d55f2b28 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_1ctamma.py @@ -21,6 +21,7 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( copy_tensormap_to_workspace as _copy_tensormap_to_workspace, + epi_subtile_spans as _epi_subtile_spans, fence_tensormap_acquire as _fence_tensormap_acquire, fence_tensormap_release as _fence_tensormap_release, moe_swizzle_tile as _moe_swizzle_tile, @@ -43,6 +44,9 @@ # it to the per-CTA GMEM workspace the TMA reads. # @@INJECT_TILE_CONSTANTS@@ +if use_acc_overlap and any(_w != epi_n for _, _w in _epi_subtile_spans(epi_cols_per_mma_m, epi_n)): + raise NotImplementedError(f"{__name__}: acc overlap reverses subtiles by index, which needs a uniform drain width") + # Per-CTA scheduler ring (replaces CLC): 2 stages, 8 int32 slot words. SCHED_STAGES = 2 @@ -159,11 +163,15 @@ def _kernel( a_part_arrive = cutlass.Int16(a_pattern) << m_rank b_part_arrive = cutlass.Int16(b_pattern) << (n_rank * cluster_m) - ab_empty_arrive_mask = a_part_arrive | b_part_arrive + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = a_part_arrive | b_part_arrive _smem_sys_reserved = cutlass.Array(cutlass.Int8, 1024, space=cutlass.AddressSpace.smem, alignment=1) ab_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) + sf_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) ab_empty_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) acc_empty_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) acc_full_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) @@ -228,12 +236,17 @@ def _kernel( for _ in range(num_b_operands) ] - ab_empty_count = cluster_m + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size + else: + ab_empty_count = cluster_m + cluster_n - 1 sched_empty_count = 1 + 1 + num_epilogue_warps if warp_idx == 0: for i in range(ab_stages): if elect_one: nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) + if elect_one: + nvvm.mbarrier_init(sf_full_mbar_ptr.subview(i), 1) if elect_one: nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) for i in range(acc_stages): @@ -260,9 +273,10 @@ def _kernel( sA_bytes = sA_elems * (ab_dtype.width // 8) 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) + ab_only_copy_bytes = num_a_operands * sA_bytes + num_b_operands * sB_bytes + sf_only_copy_bytes = num_a_operands * sfa_smem_bytes + num_b_operands * sfb_smem_bytes - cols_per_acc_stage = cta_tile_mnk[1] + epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m tmem_alloc_bar_count = (num_epilogue_warps + 1) * 32 # @@INJECT_TAP_PTRS@@ @@ -493,8 +507,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) sched_stage = cutlass.Int32(0) @@ -576,33 +589,60 @@ def _kernel( coord_k = k_tile_idx * cta_tile_mnk[2] coord_sf_k = k_tile_idx * sf_tma_box_k if elect_one: - nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) + nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), ab_only_copy_bytes) + if elect_one: + nvvm.mbarrier_arrive_expect_tx(sf_full_mbar_ptr.subview(stage), sf_only_copy_bytes) a_issue = (not multicast_a) or (n_rank == 0) + if cutlass.const_expr(a_mcast_slices > 1): + a_data_issue = True + _a_off = n_rank * (cta_tile_mnk[0] // a_mcast_slices) + else: + a_data_issue = a_issue + _a_off = 0 b_issue = (not multicast_b) or (m_rank == 0) + if cutlass.const_expr(b_mcast_slices > 1): + b_data_issue = True + _b_off = m_rank * (cta_tile_mnk[1] // b_mcast_slices) + else: + b_data_issue = b_issue + _b_off = 0 if a_issue: for _ai in cutlass.range_constexpr(num_a_operands): if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - smem_a_list[_ai].subview(sA_elems * stage), - a_desc_tma_ptr_list[_ai], - (coord_k, coord_m_group, cutlass.Int32(0)), - ab_full_mbar_ptr.subview(stage), + smem_sfa_list[_ai].subview(sfa_smem_bytes * stage), + tma_sfa_descs[_ai].get_ptr(), + (0, coord_sf_k, sfa_m_block, cutlass.Int32(0)), + sf_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_1, ) + if b_issue: + for _bj in cutlass.range_constexpr(num_b_operands): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + smem_sfb_list[_bj].subview(sfb_smem_bytes * stage), + tma_sfb_descs[_bj].get_ptr(), + (0, coord_sf_k, sfb_n_block, coord_expert), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + if a_data_issue: for _ai in cutlass.range_constexpr(num_a_operands): if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - smem_sfa_list[_ai].subview(sfa_smem_bytes * stage), - tma_sfa_descs[_ai].get_ptr(), - (0, coord_sf_k, sfa_m_block, cutlass.Int32(0)), + smem_a_list[_ai].subview(sA_elems * stage + _a_off * ab_packed_per_row), + a_desc_tma_ptr_list[_ai], + (coord_k, coord_m_group + _a_off, cutlass.Int32(0)), ab_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_1, ) - if b_issue: + if b_data_issue: for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) if cutlass.const_expr(b_is_n_major): @@ -624,39 +664,28 @@ def _kernel( else: if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - sB_stage, + sB_stage.subview(_b_off * ab_packed_per_row), tma_b_descs[_bj].get_ptr(), - (coord_k, coord_n_per_cta, coord_expert), + (coord_k, coord_n_per_cta + _b_off, coord_expert), ab_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_b, group=nvvm.CTAGroup.CTA_1, ) - for _bj in cutlass.range_constexpr(num_b_operands): - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - smem_sfb_list[_bj].subview(sfb_smem_bytes * stage), - tma_sfb_descs[_bj].get_ptr(), - (0, coord_sf_k, sfb_n_block, coord_expert), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_1, - ) ab_iter += 1 tail_stage = ab_iter % ab_stages tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -701,7 +730,7 @@ def _kernel( sfa_scale_ptrs = [nvvm.make_tmem_ptr(b, cutlass.Float32) for b in sfa_tmem_bases] sfb_scale_ptrs = [nvvm.make_tmem_ptr(b, cutlass.Float32) for b in sfb_tmem_bases] sfa_dst_ptrs = [ - [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block, cutlass.Float32) for m in range(num_blocks_m)] for i in range(num_a_operands) + [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block, cutlass.Float32) for m in range(num_mma_m)] for i in range(num_a_operands) ] sfb_dst_ptrs = [ [nvvm.make_tmem_ptr(sfb_tmem_bases[j] + m * registers_per_block, cutlass.Float32) for m in range(num_blocks_n)] for j in range(num_b_operands) @@ -757,13 +786,6 @@ def _kernel( if stage == 0 and ab_iter != 0: ab_full_phase_bit = ab_full_phase_bit ^ 1 - while not nvvm.mbarrier_try_wait_parity( - ab_full_mbar_ptr.subview(stage), - ab_full_phase_bit, - time_limit=10_000_000, - ): - pass - desc_a_bases = [ cutlass.experimental.primitives.Tcgen05SmemDesc.build( start_address=smem_a_list[i].subview(sA_elems * stage), @@ -801,56 +823,69 @@ def _kernel( for j in range(num_b_operands) ] - for atom_r in cutlass.range_constexpr(num_sf_atoms): - for _ai in cutlass.range_constexpr(num_a_operands): - for _m in cutlass.range_constexpr(num_blocks_m): - if elect_one: - nvvm.tcgen05_cp( - s2t_shape, - sfa_dst_ptrs[_ai][_m], - desc_sfa_bases[_ai] + (sf_atom_desc_stride * atom_r + sf_block_desc_stride * _m), - group=nvvm.CTAGroup.CTA_1, - multicast=s2t_multicast, - ) + while not nvvm.mbarrier_try_wait_parity( + sf_full_mbar_ptr.subview(stage), + ab_full_phase_bit, + time_limit=10_000_000, + ): + pass + + for sf_word in cutlass.range_constexpr(num_sf_atoms): for _bj in cutlass.range_constexpr(num_b_operands): - for _m in cutlass.range_constexpr(num_blocks_n): + for block_n in cutlass.range_constexpr(num_blocks_n): if elect_one: nvvm.tcgen05_cp( s2t_shape, - sfb_dst_ptrs[_bj][_m], - desc_sfb_bases[_bj] + (sf_atom_desc_stride * atom_r + sf_block_desc_stride * _m), + sfb_dst_ptrs[_bj][block_n], + desc_sfb_bases[_bj] + (sf_atom_desc_stride * sf_word + sf_block_desc_stride * block_n), group=nvvm.CTAGroup.CTA_1, multicast=s2t_multicast, ) - for j in cutlass.range_constexpr(sf_insts_per_atom): - k_block_idx = atom_r * sf_insts_per_atom + j - idesc_k = idesc_by_j[j] - for g in cutlass.range_constexpr(num_gemms): - _ai = gemm_a_idx[g] - _bj = gemm_b_idx[g] - desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) - desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - for mi in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(sf_word == 0): + while not nvvm.mbarrier_try_wait_parity( + ab_full_mbar_ptr.subview(stage), + ab_full_phase_bit, + time_limit=10_000_000, + ): + pass + for mma_k_in_word in cutlass.range_constexpr(sf_insts_per_atom): + mma_k = sf_word * sf_insts_per_atom + mma_k_in_word + idesc_k = idesc_by_j[mma_k_in_word] + for gemm_i in cutlass.range_constexpr(num_gemms): + _ai = gemm_a_idx[gemm_i] + _bj = gemm_b_idx[gemm_i] + desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * mma_k) + desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * mma_k) + for mma_m in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(mma_k_in_word == 0 and _ai not in gemm_a_idx[:gemm_i]): + if elect_one: + nvvm.tcgen05_cp( + s2t_shape, + sfa_dst_ptrs[_ai][mma_m], + desc_sfa_bases[_ai] + (sf_atom_desc_stride * sf_word + sf_block_desc_stride * mma_m), + group=nvvm.CTAGroup.CTA_1, + multicast=s2t_multicast, + ) # The M sub-block offset is a whole SMEM swizzle atom, so # the descriptor's swizzle phase is preserved. B and its SF # are shared; A's SF word block follows the M block. - desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) + desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mma_m) if elect_one: _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_1, - acc_tmem_ptrs[g][mi], + acc_tmem_ptrs[gemm_i][mma_m], desc_a, desc_b, idesc_k, enable_input_d=scale_d, - scale_a=sfa_dst_ptrs[_ai][mi], + scale_a=sfa_dst_ptrs[_ai][mma_m], scale_b=sfb_scale_ptrs[_bj], scale_vec_size=scale_vec_size, - b_collector_op=_b_collector_op(mi), + b_collector_op=_b_collector_op(mma_m), ) # Every accumulator sees scale_d=False on exactly the first - # k_block of the tile, so the flip stays outside mi. + # k_block of the tile, so the flip stays outside mma_m. scale_d = cutlass.Boolean(True) if elect_one: @@ -869,28 +904,26 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) if tile_iter != 0: tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_empty_mbar_ptr.subview(tail_stage), - tail_phase, - time_limit=10_000_000, - ): - pass - if cutlass.const_expr(use_acc_overlap): - while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): - pass + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity( + acc_empty_mbar_ptr.subview(tail_stage), + tail_phase, + time_limit=10_000_000, + ): + pass + if cutlass.const_expr(use_acc_overlap): + while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): + pass nvvm.bar_warp_sync(0xFFFFFFFF) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) @@ -917,21 +950,14 @@ def _kernel( sched_stage = cutlass.Int32(0) sched_full_phase = cutlass.Int32(0) - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - row_id_with_warp_offset = base_row_id - else: - row_id_with_warp_offset = base_row_id + warp_idx * 32 - - # One M block's accumulator columns are contiguous. - subtile_cnt = cute.ceil_div(epi_cols_per_mma_m, 32) - t2r_inst_repx = epi_tile_mn[1] - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - shape = nvvm.Tcgen05LdStShape.SHAPE_16X32BX2 - ld_half_off = 0 - else: - shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B - ld_half_off = None + # @@EPILOGUE_SETUP:BEGIN@@ + row_id_with_warp_offset = base_row_id + warp_idx * 32 + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) + subtile_cnt = len(epi_spans) + shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ while is_valid != 0: while not nvvm.mbarrier_try_wait_parity( @@ -956,17 +982,14 @@ def _kernel( if is_valid != 0: coord_m_tile = group_begin + 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] + # @@EPILOGUE_DRAIN:BEGIN@@ + coord_n_c = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: acc_full_phase_bit = acc_full_phase_bit ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_full_mbar_ptr.subview(acc_stage), - acc_full_phase_bit, - time_limit=10_000_000, - ): + while not nvvm.mbarrier_try_wait_parity(acc_full_mbar_ptr.subview(acc_stage), acc_full_phase_bit, time_limit=10_000_000): pass if cutlass.const_expr(use_acc_overlap): @@ -975,48 +998,53 @@ def _kernel( else: acc_buf_parity = cutlass.Int32(0) acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # One pass per MMA-M block over its own column region. + for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * mma_inst_shape_mnk[0] - mi_col_base = acc_base_col + mi * epi_cols_per_mma_m + if cutlass.const_expr(use_acc_overlap and num_mma_m > 1): + _mi = mi + (1 - acc_buf_parity) * (num_mma_m - 1 - 2 * mi) + else: + _mi = mi + coord_m = coord_m_tile + _mi * epi_rows_per_mma_m + mi_col_base = acc_base_col + _mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * acc_gemm_stride) for g in range(num_gemms)] - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - row = coord_m + warp_idx * 16 + lane - row_active = lane < 16 - else: - row = coord_m + tidx - row_active = True + row = coord_m + tidx + row_active = True # @@INJECT_AUX_VIEWS@@ for subtile_idx in cutlass.range_constexpr(subtile_cnt): if cutlass.const_expr(use_acc_overlap): _sub = subtile_idx + (1 - acc_buf_parity) * (subtile_cnt - 1 - 2 * subtile_idx) - subtile_col_offset = _sub * 32 + subtile_col_offset = _sub * epi_n + subtile_w = epi_n else: - subtile_col_offset = subtile_idx * 32 + subtile_col_offset, subtile_w = epi_spans[subtile_idx] c_rmem_vecs = [] for g in cutlass.range_constexpr(num_gemms): - tmem = cutlass.inttoptr( - tmem_col_addr_gemms[g] + subtile_col_offset, - 6, - cutlass.Float32, - ) - c_rmem_vecs.append(nvvm.tcgen05_ld(shape, tmem, num=t2r_inst_repx, offset=ld_half_off)) + subtile_tmem_addr = tmem_col_addr_gemms[g] + subtile_col_offset + tmem = cutlass.inttoptr(subtile_tmem_addr, 6, mma_c_dtype) + _cv = nvvm.tcgen05_ld(shape, tmem, num=subtile_w) + c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - if cutlass.const_expr(use_acc_overlap and mi == num_mma_m - 1 and subtile_idx == acc_overlap_subtiles - 1): + if cutlass.const_expr(not use_acc_overlap): + if cutlass.const_expr(mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1): + nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) + nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) + if elect_one: + nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) + + if use_acc_overlap and mi * subtile_cnt + subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) - col = coord_n + subtile_col_offset + col = coord_n_c + subtile_col_offset - # @@STG_ONLY:BEGIN@@ if row_active and row < group_end: - for j in cutlass.range_constexpr(t2r_inst_repx // vsize): + for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize if col_j + vsize <= N: vec_f32 = c_rmem_vec[j * vsize : (j + 1) * vsize] @@ -1024,13 +1052,9 @@ def _kernel( # @@INJECT_STG_VEC_BINDINGS@@ # @@INJECT_EPILOGUE@@ - # @@STG_ONLY:END@@ - if cutlass.const_expr(not use_acc_overlap): - nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) - nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if elect_one: - nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + # @@EPILOGUE_DRAIN:END@@ tile_iter += 1 if cutlass.const_expr(use_acc_overlap): @@ -1094,7 +1118,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1127,7 +1151,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py index cd3074233..8cb12554a 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py @@ -26,6 +26,7 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( copy_tensormap_to_workspace as _copy_tensormap_to_workspace, + epi_subtile_spans as _epi_subtile_spans, fence_tensormap_acquire as _fence_tensormap_acquire, fence_tensormap_release as _fence_tensormap_release, moe_swizzle_tile as _moe_swizzle_tile, @@ -45,6 +46,9 @@ # @@INJECT_TILE_CONSTANTS@@ +if use_acc_overlap and any(_w != epi_n for _, _w in _epi_subtile_spans(epi_cols_per_mma_m, epi_n)): + raise NotImplementedError(f"{__name__}: acc overlap reverses subtiles by index, which needs a uniform drain width") + SCHED_STAGES = 2 SCHED_SLOT_WORDS = 8 @@ -168,6 +172,7 @@ def _kernel( _smem_sys_reserved = cutlass.Array(cutlass.Int8, 1024, space=cutlass.AddressSpace.smem, alignment=1) ab_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) + sf_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) ab_empty_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) acc_empty_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) acc_full_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) @@ -233,7 +238,10 @@ def _kernel( acc_empty_count = num_epilogue_warps * 2 cta_group = 2 - ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size // cta_group + else: + ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 sched_empty_count = 1 + 1 + num_epilogue_warps if warp_idx == 0: if cutlass.const_expr(use_acc_overlap): @@ -245,6 +253,8 @@ def _kernel( for i in range(ab_stages): if elect_one: nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) + if elect_one: + nvvm.mbarrier_init(sf_full_mbar_ptr.subview(i), 1) if elect_one: nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) for i in range(acc_stages): @@ -262,17 +272,13 @@ def _kernel( sA_bytes = sA_elems * (ab_dtype.width // 8) 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 + ab_only_copy_bytes = (num_a_operands * sA_bytes + num_b_operands * sB_bytes) * 2 + sf_only_copy_bytes = (num_a_operands * sfa_smem_bytes + num_b_operands * sfb_smem_bytes) * 2 pair_n_size = cgrp_tile_mnk[1] // cluster_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 - if cutlass.const_expr(epi_rows_per_mma_m == 64): - # cluster-MMA m=128: the pair also splits N, so each CTA drains N/2. - cols_per_acc_stage = pair_n_size // 2 - else: - cols_per_acc_stage = pair_n_size tmem_alloc_bar_count = (num_epilogue_warps + 1) * 32 nvvm.barrier_cluster_wait() @@ -506,8 +512,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) sched_stage = cutlass.Int32(0) @@ -593,33 +598,61 @@ def _kernel( if is_pair_leader: if elect_one: - nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) + nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), ab_only_copy_bytes) + if elect_one: + nvvm.mbarrier_arrive_expect_tx(sf_full_mbar_ptr.subview(stage), sf_only_copy_bytes) a_issue = (not multicast_a) or (n_rank == 0) + if cutlass.const_expr(a_mcast_slices > 1): + a_data_issue = True + _a_off = n_rank * (cta_tile_mnk[0] // a_mcast_slices) + else: + a_data_issue = a_issue + _a_off = 0 b_issue = (not multicast_b) or (pair_m_idx == 0) + if cutlass.const_expr(b_mcast_slices > 1): + b_data_issue = True + _b_off = pair_m_idx * (cta_tile_mnk[1] // b_mcast_slices) + else: + b_data_issue = b_issue + _b_off = 0 if a_issue: for _ai in cutlass.range_constexpr(num_a_operands): if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - smem_a_list[_ai].subview(sA_elems * stage), - a_desc_tma_ptr_list[_ai], - (coord_k, coord_m_group, cutlass.Int32(0)), - ab_full_mbar_ptr.subview(stage), + smem_sfa_list[_ai].subview(sfa_smem_bytes * stage), + tma_sfa_descs[_ai].get_ptr(), + (0, coord_sf_k, sfa_m_block, cutlass.Int32(0)), + sf_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_2, ) + if b_issue: + for _bj in cutlass.range_constexpr(num_b_operands): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + smem_sfb_list[_bj].subview(sfb_smem_bytes * stage), + tma_sfb_descs[_bj].get_ptr(), + (0, coord_sf_k, sfb_n_block, coord_expert), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + + if a_data_issue: for _ai in cutlass.range_constexpr(num_a_operands): if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - smem_sfa_list[_ai].subview(sfa_smem_bytes * stage), - tma_sfa_descs[_ai].get_ptr(), - (0, coord_sf_k, sfa_m_block, cutlass.Int32(0)), + smem_a_list[_ai].subview(sA_elems * stage + _a_off * ab_packed_per_row), + a_desc_tma_ptr_list[_ai], + (coord_k, coord_m_group + _a_off, cutlass.Int32(0)), ab_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_2, ) - if b_issue: + if b_data_issue: for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) if cutlass.const_expr(b_is_n_major): @@ -641,40 +674,28 @@ def _kernel( else: if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - sB_stage, + sB_stage.subview(_b_off * ab_packed_per_row), tma_b_descs[_bj].get_ptr(), - (coord_k, coord_n_per_cta, coord_expert), + (coord_k, coord_n_per_cta + _b_off, coord_expert), ab_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_b, group=nvvm.CTAGroup.CTA_2, ) - for _bj in cutlass.range_constexpr(num_b_operands): - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - smem_sfb_list[_bj].subview(sfb_smem_bytes * stage), - tma_sfb_descs[_bj].get_ptr(), - (0, coord_sf_k, sfb_n_block, coord_expert), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_2, - ) - ab_iter += 1 tail_stage = ab_iter % ab_stages tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 pair_mask = cutlass.Int16(3) << pair_leader_rank a_arrive_pattern = 0 @@ -686,7 +707,10 @@ def _kernel( a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) - ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) _tcgen05_alloc( @@ -730,7 +754,7 @@ def _kernel( sfa_scale_ptrs = [nvvm.make_tmem_ptr(b, cutlass.Float32) for b in sfa_tmem_bases] sfb_scale_ptrs = [nvvm.make_tmem_ptr(b, cutlass.Float32) for b in sfb_tmem_bases] sfa_dst_ptrs = [ - [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block, cutlass.Float32) for m in range(num_blocks_m)] for i in range(num_a_operands) + [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block, cutlass.Float32) for m in range(num_mma_m)] for i in range(num_a_operands) ] sfb_dst_ptrs = [ [nvvm.make_tmem_ptr(sfb_tmem_bases[j] + m * registers_per_block, cutlass.Float32) for m in range(num_blocks_n)] for j in range(num_b_operands) @@ -786,13 +810,6 @@ def _kernel( if stage == 0 and ab_iter != 0: ab_full_phase_bit = ab_full_phase_bit ^ 1 - while not nvvm.mbarrier_try_wait_parity( - ab_full_mbar_ptr.subview(stage), - ab_full_phase_bit, - time_limit=10_000_000, - ): - pass - desc_a_bases = [ cutlass.experimental.primitives.Tcgen05SmemDesc.build( start_address=smem_a_list[i].subview(sA_elems * stage), @@ -830,56 +847,69 @@ def _kernel( for j in range(num_b_operands) ] - for atom_r in cutlass.range(num_sf_atoms, unroll_full=True): - for _ai in cutlass.range_constexpr(num_a_operands): - for _m in cutlass.range_constexpr(num_blocks_m): - if elect_one: - nvvm.tcgen05_cp( - s2t_shape, - sfa_dst_ptrs[_ai][_m], - desc_sfa_bases[_ai] + (sf_atom_desc_stride * atom_r + sf_block_desc_stride * _m), - group=nvvm.CTAGroup.CTA_2, - multicast=s2t_multicast, - ) + while not nvvm.mbarrier_try_wait_parity( + sf_full_mbar_ptr.subview(stage), + ab_full_phase_bit, + time_limit=10_000_000, + ): + pass + + for sf_word in cutlass.range_constexpr(num_sf_atoms): for _bj in cutlass.range_constexpr(num_b_operands): - for _m in cutlass.range_constexpr(num_blocks_n): + for block_n in cutlass.range_constexpr(num_blocks_n): if elect_one: nvvm.tcgen05_cp( s2t_shape, - sfb_dst_ptrs[_bj][_m], - desc_sfb_bases[_bj] + (sf_atom_desc_stride * atom_r + sf_block_desc_stride * _m), + sfb_dst_ptrs[_bj][block_n], + desc_sfb_bases[_bj] + (sf_atom_desc_stride * sf_word + sf_block_desc_stride * block_n), group=nvvm.CTAGroup.CTA_2, multicast=s2t_multicast, ) - for j in cutlass.range_constexpr(sf_insts_per_atom): - k_block_idx = atom_r * sf_insts_per_atom + j - idesc_k = idesc_by_j[j] - for g in cutlass.range_constexpr(num_gemms): - _ai = gemm_a_idx[g] - _bj = gemm_b_idx[g] - desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) - desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - for mi in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(sf_word == 0): + while not nvvm.mbarrier_try_wait_parity( + ab_full_mbar_ptr.subview(stage), + ab_full_phase_bit, + time_limit=10_000_000, + ): + pass + for mma_k_in_word in cutlass.range_constexpr(sf_insts_per_atom): + mma_k = sf_word * sf_insts_per_atom + mma_k_in_word + idesc_k = idesc_by_j[mma_k_in_word] + for gemm_i in cutlass.range_constexpr(num_gemms): + _ai = gemm_a_idx[gemm_i] + _bj = gemm_b_idx[gemm_i] + desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * mma_k) + desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * mma_k) + for mma_m in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(mma_k_in_word == 0 and _ai not in gemm_a_idx[:gemm_i]): + if elect_one: + nvvm.tcgen05_cp( + s2t_shape, + sfa_dst_ptrs[_ai][mma_m], + desc_sfa_bases[_ai] + (sf_atom_desc_stride * sf_word + sf_block_desc_stride * mma_m), + group=nvvm.CTAGroup.CTA_2, + multicast=s2t_multicast, + ) # The M sub-block offset is a whole SMEM swizzle atom, so # the descriptor's swizzle phase is preserved. B and its SF # are shared; A's SF word block follows the M block. - desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) + desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mma_m) if elect_one: _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_2, - acc_tmem_ptrs[g][mi], + acc_tmem_ptrs[gemm_i][mma_m], desc_a, desc_b, idesc_k, enable_input_d=scale_d, - scale_a=sfa_dst_ptrs[_ai][mi], + scale_a=sfa_dst_ptrs[_ai][mma_m], scale_b=sfb_scale_ptrs[_bj], scale_vec_size=scale_vec_size, - b_collector_op=_b_collector_op(mi), + b_collector_op=_b_collector_op(mma_m), ) # Every accumulator sees scale_d=False on exactly the first - # k_block of the tile, so the flip stays outside mi. + # k_block of the tile, so the flip stays outside mma_m. scale_d = cutlass.Boolean(True) if elect_one: @@ -899,26 +929,22 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") if tile_iter != 0: tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_empty_mbar_ptr.subview(tail_stage), - tail_phase, - time_limit=10_000_000, - ): - pass - nvvm.bar_warp_sync(0xFFFFFFFF) - + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity( + acc_empty_mbar_ptr.subview(tail_stage), + tail_phase, + time_limit=10_000_000, + ): + pass nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): @@ -952,8 +978,7 @@ def _kernel( sched_full_phase = sched_full_phase ^ 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) @@ -983,15 +1008,14 @@ def _kernel( sched_stage = cutlass.Int32(0) sched_full_phase = cutlass.Int32(0) + # @@EPILOGUE_SETUP:BEGIN@@ row_id_with_warp_offset = base_row_id + warp_idx * 32 - if cutlass.const_expr(cols_per_acc_stage >= 32): - t2r_inst_repx = 32 - subtile_cnt = cols_per_acc_stage // 32 - else: - t2r_inst_repx = cols_per_acc_stage - subtile_cnt = 1 + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) + subtile_cnt = len(epi_spans) shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ while is_valid != 0: while not nvvm.mbarrier_try_wait_parity( @@ -1016,19 +1040,14 @@ def _kernel( if is_valid != 0: coord_m_tile = group_begin + tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] + # @@EPILOGUE_DRAIN:BEGIN@@ coord_n_c = tile_n * cgrp_tile_mnk[1] + 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 acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: acc_full_phase_bit = acc_full_phase_bit ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_full_mbar_ptr.subview(acc_stage), - acc_full_phase_bit, - time_limit=10_000_000, - ): + while not nvvm.mbarrier_try_wait_parity(acc_full_mbar_ptr.subview(acc_stage), acc_full_phase_bit, time_limit=10_000_000): pass if cutlass.const_expr(use_acc_overlap): @@ -1037,52 +1056,61 @@ def _kernel( else: acc_buf_parity = cutlass.Int32(0) acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # The 2-CTA epilogue drains its own half of the instruction's M, - # epi_rows_per_mma_m rows at a time, so a CTA tile of num_mma_m blocks - # drains in num_mma_m passes over its own column region. + for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * epi_rows_per_mma_m - mi_col_base = acc_base_col + mi * epi_cols_per_mma_m + if cutlass.const_expr(use_acc_overlap and num_mma_m > 1): + _mi = mi + (1 - acc_buf_parity) * (num_mma_m - 1 - 2 * mi) + else: + _mi = mi + coord_m = coord_m_tile + _mi * epi_rows_per_mma_m + mi_col_base = acc_base_col + _mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * acc_gemm_stride) for g in range(num_gemms)] - if cutlass.const_expr(epi_rows_per_mma_m == 64): - row = coord_m + (warp_idx % 2) * 32 + lane - row_active = True - else: - row = coord_m + tidx - row_active = True + row = coord_m + tidx + row_active = True # @@INJECT_AUX_VIEWS@@ - for subtile_idx in cutlass.range(subtile_cnt, unroll_full=True): + for subtile_idx in cutlass.range_constexpr(subtile_cnt): if cutlass.const_expr(use_acc_overlap): _sub = subtile_idx + (1 - acc_buf_parity) * (subtile_cnt - 1 - 2 * subtile_idx) - subtile_col_offset = _sub * t2r_inst_repx + subtile_col_offset = _sub * epi_n + subtile_w = epi_n else: - subtile_col_offset = subtile_idx * t2r_inst_repx + subtile_col_offset, subtile_w = epi_spans[subtile_idx] c_rmem_vecs = [] for g in cutlass.range_constexpr(num_gemms): - tmem = cutlass.inttoptr( - tmem_col_addr_gemms[g] + subtile_col_offset, - 6, - cutlass.Float32, - ) - c_rmem_vecs.append(nvvm.tcgen05_ld(shape, tmem, num=t2r_inst_repx)) + subtile_tmem_addr = tmem_col_addr_gemms[g] + subtile_col_offset + tmem = cutlass.inttoptr(subtile_tmem_addr, 6, mma_c_dtype) + _cv = nvvm.tcgen05_ld(shape, tmem, num=subtile_w) + c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - if use_acc_overlap and mi == num_mma_m - 1 and subtile_idx == acc_overlap_subtiles - 1: + if cutlass.const_expr(not use_acc_overlap): + if cutlass.const_expr(mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1): + nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) + nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) + if elect_one: + nvvm.mbarrier_arrive( + nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), + scope=nvvm.MemScope.CLUSTER, + relaxed=True, + ) + + if use_acc_overlap and mi * subtile_cnt + subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: - mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) - nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) + nvvm.mbarrier_arrive( + nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), + scope=nvvm.MemScope.CLUSTER, + relaxed=True, + ) col = coord_n_c + subtile_col_offset - # @@STG_ONLY:BEGIN@@ - if row_active and row < group_end: - for j in cutlass.range_constexpr(t2r_inst_repx // vsize): + for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize if col_j + vsize <= N: vec_f32 = c_rmem_vec[j * vsize : (j + 1) * vsize] @@ -1091,14 +1119,8 @@ def _kernel( # @@INJECT_EPILOGUE@@ - # @@STG_ONLY:END@@ - - if cutlass.const_expr(not use_acc_overlap): - nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) - nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if elect_one: - mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) - nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + # @@EPILOGUE_DRAIN:END@@ tile_iter += 1 if cutlass.const_expr(use_acc_overlap): @@ -1164,7 +1186,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1197,7 +1219,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_1ctamma.py index 9208d47f1..bf9784ed7 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_1ctamma.py @@ -152,7 +152,10 @@ def _kernel( a_part_arrive = cutlass.Int16(a_pattern) << m_rank b_part_arrive = cutlass.Int16(b_pattern) << (n_rank * cluster_m) - ab_empty_arrive_mask = a_part_arrive | b_part_arrive + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = a_part_arrive | b_part_arrive _smem_sys_reserved = cutlass.Array(cutlass.Int8, 1024, space=cutlass.AddressSpace.smem, alignment=1) @@ -202,7 +205,10 @@ def _kernel( for _ in range(num_a_operands) ] - ab_empty_count = cluster_m + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size + else: + ab_empty_count = cluster_m + cluster_n - 1 sched_empty_count = 1 + 1 + num_epilogue_warps if warp_idx == 0: for i in range(ab_stages): @@ -247,6 +253,7 @@ def _kernel( epi_cols_per_mma_m = cta_tile_mnk[1] cols_per_acc_stage = num_mma_m * epi_cols_per_mma_m acc_region_cols = num_gemms * cols_per_acc_stage + epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m tmem_alloc_bar_count = (num_epilogue_warps + 1) * 32 # @@INJECT_TAP_PTRS@@ @@ -375,8 +382,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) sched_stage = cutlass.Int32(0) @@ -457,21 +463,33 @@ def _kernel( nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) a_issue = (not multicast_a) or (n_rank == 0) - if a_issue: + if cutlass.const_expr(a_mcast_slices > 1): + a_data_issue = True + _a_off = n_rank * (cta_tile_mnk[0] // a_mcast_slices) + else: + a_data_issue = a_issue + _a_off = 0 + if a_data_issue: for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - sA_stage, + sA_stage.subview(_a_off * cta_tile_mnk[2]), a_desc_tma_ptr_list[_ai], - (coord_k, coord_m_group, cutlass.Int32(0)), + (coord_k, coord_m_group + _a_off, cutlass.Int32(0)), ab_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_1, ) b_issue = (not multicast_b) or (m_rank == 0) - if b_issue: + if cutlass.const_expr(b_mcast_slices > 1): + b_data_issue = True + _b_off = m_rank * (cta_tile_mnk[1] // b_mcast_slices) + else: + b_data_issue = b_issue + _b_off = 0 + if b_data_issue: for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) if cutlass.const_expr(b_is_n_major): @@ -493,9 +511,9 @@ def _kernel( else: if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - sB_stage, + sB_stage.subview(_b_off * cta_tile_mnk[2]), tma_b_descs[_bj].get_ptr(), - (coord_k, coord_n_per_cta, coord_expert), + (coord_k, coord_n_per_cta + _b_off, coord_expert), ab_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_b, @@ -507,14 +525,14 @@ def _kernel( tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -642,26 +660,23 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") if tile_iter != 0: tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_empty_mbar_ptr.subview(tail_stage), - tail_phase, - time_limit=10_000_000, - ): - pass + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity( + acc_empty_mbar_ptr.subview(tail_stage), + tail_phase, + time_limit=10_000_000, + ): + pass - nvvm.bar_warp_sync(0xFFFFFFFF) nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) _tcgen05_dealloc( @@ -688,14 +703,13 @@ def _kernel( # The drain layout follows the MMA INSTRUCTION's M: at hardware M=64 the # accumulator occupies data paths 0-15 of each sub-partition, so the # 32-data-path LDTM layout cannot be used to read it. + # @@EPILOGUE_SETUP:BEGIN@@ if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): row_id_with_warp_offset = base_row_id else: row_id_with_warp_offset = base_row_id + warp_idx * 32 - # One M block's accumulator columns are contiguous, so one span list - # drains all of them. - epi_spans = _epi_subtile_spans(epi_cols_per_mma_m) + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) subtile_cnt = len(epi_spans) if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): shape = nvvm.Tcgen05LdStShape.SHAPE_16X32BX2 @@ -704,6 +718,7 @@ def _kernel( shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B ld_half_off = None lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ while not nvvm.mbarrier_try_wait_parity(sched_full_mbar_ptr.subview(sched_stage), sched_full_phase, time_limit=10_000_000): pass @@ -723,7 +738,8 @@ def _kernel( while is_valid != 0: coord_m_tile = group_begin + 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] + # @@EPILOGUE_DRAIN:BEGIN@@ + coord_n_c = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: @@ -734,11 +750,8 @@ def _kernel( acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # One pass per MMA-M block: the 4 epilogue warps cover this CTA's rows for - # one block at a time, so a CTA tile of num_mma_m blocks drains in - # num_mma_m passes over its own column region. for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * mma_inst_shape_mnk[0] + coord_m = coord_m_tile + mi * epi_rows_per_mma_m mi_col_base = acc_base_col + mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * cols_per_acc_stage) for g in range(num_gemms)] @@ -758,23 +771,22 @@ def _kernel( subtile_tmem_addr = tmem_col_addr_gemms[g] + subtile_col_offset tmem = cutlass.inttoptr(subtile_tmem_addr, 6, mma_c_dtype) _cv = nvvm.tcgen05_ld(shape, tmem, num=subtile_w, offset=ld_half_off) + # INT8 int32 accumulate → widen to fp32 (skipped for int32 output). if cutlass.const_expr(acc_widen_to_fp32): _accf = _cv.to(cutlass.Float32) + # `+ 0.0` forces a fresh fp32 register so int32->fp32 isn't folded into an invalid int32->fp8 cast. _cv = _accf + cutlass.full_like(_accf, 0.0) c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - # Exactly one acc_empty arrive per epilogue warp per tile — - # the mbar counts warps, not M blocks. if mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) - col = coord_n + subtile_col_offset + col = coord_n_c + subtile_col_offset - # @@STG_ONLY:BEGIN@@ if row_active and row < group_end: for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize @@ -784,8 +796,9 @@ def _kernel( # @@INJECT_STG_VEC_BINDINGS@@ # @@INJECT_EPILOGUE@@ - # @@STG_ONLY:END@@ + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + # @@EPILOGUE_DRAIN:END@@ tile_iter += 1 while not nvvm.mbarrier_try_wait_parity( @@ -863,7 +876,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, ) ) @@ -894,7 +907,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, ) ) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_2ctamma.py index 6d858a125..9bb75b7cc 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm100_moe_grouped_matmul_fwd_2ctamma.py @@ -212,7 +212,10 @@ def _kernel( acc_empty_count = num_epilogue_warps * 2 cta_group = 2 - ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size // cta_group + else: + ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 num_consumer_warps_per_cta = 1 + 1 + num_epilogue_warps sched_empty_count = num_consumer_warps_per_cta if warp_idx == 0: @@ -392,8 +395,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) sched_stage = cutlass.Int32(0) @@ -476,21 +478,33 @@ def _kernel( if elect_one: nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) a_issue = (not multicast_a) or (n_rank == 0) - if a_issue: + if cutlass.const_expr(a_mcast_slices > 1): + a_data_issue = True + _a_off = n_rank * (cta_tile_mnk[0] // a_mcast_slices) + else: + a_data_issue = a_issue + _a_off = 0 + if a_data_issue: for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - sA_stage, + sA_stage.subview(_a_off * cta_tile_mnk[2]), a_desc_tma_ptr_list[_ai], - (coord_k, coord_m_group, cutlass.Int32(0)), + (coord_k, coord_m_group + _a_off, cutlass.Int32(0)), ab_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_2, ) b_issue = (not multicast_b) or (pair_m_idx == 0) - if b_issue: + if cutlass.const_expr(b_mcast_slices > 1): + b_data_issue = True + _b_off = pair_m_idx * (cta_tile_mnk[1] // b_mcast_slices) + else: + b_data_issue = b_issue + _b_off = 0 + if b_data_issue: for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) if cutlass.const_expr(b_is_n_major): @@ -512,9 +526,9 @@ def _kernel( else: if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - sB_stage, + sB_stage.subview(_b_off * cta_tile_mnk[2]), tma_b_descs[_bj].get_ptr(), - (coord_k, coord_n_per_cta, coord_expert), + (coord_k, coord_n_per_cta + _b_off, coord_expert), ab_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_b, @@ -526,14 +540,14 @@ def _kernel( tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 pair_mask = cutlass.Int16(3) << pair_leader_rank a_arrive_pattern = 0 @@ -545,7 +559,10 @@ def _kernel( a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) - ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) _tcgen05_alloc( @@ -675,26 +692,22 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") if tile_iter != 0: tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_empty_mbar_ptr.subview(tail_stage), - tail_phase, - time_limit=10_000_000, - ): - pass - nvvm.bar_warp_sync(0xFFFFFFFF) - + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity( + acc_empty_mbar_ptr.subview(tail_stage), + tail_phase, + time_limit=10_000_000, + ): + pass nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): @@ -727,8 +740,7 @@ def _kernel( sched_full_phase = sched_full_phase ^ 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) @@ -757,13 +769,14 @@ def _kernel( sched_stage = cutlass.Int32(0) sched_full_phase = cutlass.Int32(0) + # @@EPILOGUE_SETUP:BEGIN@@ row_id_with_warp_offset = base_row_id + warp_idx * 32 - # One M block's accumulator columns are contiguous, so one span list - # drains all of them. - epi_spans = _epi_subtile_spans(epi_cols_per_mma_m) + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) subtile_cnt = len(epi_spans) shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ while is_valid != 0: while not nvvm.mbarrier_try_wait_parity( @@ -788,6 +801,7 @@ def _kernel( if is_valid != 0: coord_m_tile = group_begin + tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] + # @@EPILOGUE_DRAIN:BEGIN@@ coord_n_c = tile_n * cgrp_tile_mnk[1] + 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 @@ -796,18 +810,11 @@ def _kernel( if acc_stage == 0 and tile_iter != 0: acc_full_phase_bit = acc_full_phase_bit ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_full_mbar_ptr.subview(acc_stage), - acc_full_phase_bit, - time_limit=10_000_000, - ): + while not nvvm.mbarrier_try_wait_parity(acc_full_mbar_ptr.subview(acc_stage), acc_full_phase_bit, time_limit=10_000_000): pass acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # One pass per MMA-M block: the 4 epilogue warps cover this CTA's rows for - # one block at a time, so a CTA tile of num_mma_m blocks drains in - # num_mma_m passes over its own column region. for mi in cutlass.range_constexpr(num_mma_m): coord_m = coord_m_tile + mi * epi_rows_per_mma_m mi_col_base = acc_base_col + mi * epi_cols_per_mma_m @@ -829,14 +836,14 @@ def _kernel( subtile_tmem_addr = tmem_col_addr_gemms[g] + subtile_col_offset tmem = cutlass.inttoptr(subtile_tmem_addr, 6, mma_c_dtype) _cv = nvvm.tcgen05_ld(shape, tmem, num=subtile_w) + # INT8 int32 accumulate → widen to fp32 (skipped for int32 output). if cutlass.const_expr(acc_widen_to_fp32): _accf = _cv.to(cutlass.Float32) + # `+ 0.0` forces a fresh fp32 register so int32->fp32 isn't folded into an invalid int32->fp8 cast. _cv = _accf + cutlass.full_like(_accf, 0.0) c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - # Exactly one acc_empty arrive per epilogue warp per tile — - # the mbar counts warps, not M blocks. if mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) @@ -848,6 +855,7 @@ def _kernel( ) col = coord_n_c + subtile_col_offset + if row_active and row < group_end: for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize @@ -857,6 +865,9 @@ def _kernel( # @@INJECT_STG_VEC_BINDINGS@@ # @@INJECT_EPILOGUE@@ + + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + # @@EPILOGUE_DRAIN:END@@ tile_iter += 1 if warp_idx == unused_warp_id: @@ -916,7 +927,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cgrp_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cgrp_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, ) ) @@ -947,7 +958,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cgrp_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cgrp_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, ) ) 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 b09a85c34..b0b9f820f 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 @@ -45,6 +45,7 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( + epi_subtile_spans as _epi_subtile_spans, l2_swizzle_tile as _l2_swizzle_tile, tcgen05_alloc as _tcgen05_alloc, tcgen05_dealloc as _tcgen05_dealloc, @@ -62,6 +63,9 @@ # @@INJECT_TILE_CONSTANTS@@ +if use_acc_overlap and any(_w != epi_n for _, _w in _epi_subtile_spans(epi_cols_per_mma_m, epi_n)): + raise NotImplementedError(f"{__name__}: acc overlap reverses subtiles by index, which needs a uniform drain width") + CLC_SCHED_STAGES = 2 @@ -260,7 +264,10 @@ def _kernel( a_part_arrive = cutlass.Int16(a_pattern) << m_rank b_part_arrive = cutlass.Int16(b_pattern) << (n_rank * cluster_m) - ab_empty_arrive_mask = a_part_arrive | b_part_arrive + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = a_part_arrive | b_part_arrive _smem_sys_reserved = cutlass.Array(cutlass.Int8, 1024, space=cutlass.AddressSpace.smem, alignment=1) @@ -336,7 +343,10 @@ def _kernel( ) # @@TMA_STORE_ONLY:END@@ - ab_empty_count = cluster_m + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size + else: + ab_empty_count = cluster_m + cluster_n - 1 num_consumer_warps_per_cta = 8 clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: @@ -375,7 +385,7 @@ def _kernel( num_tma_ab_chunk_bytes = num_a_operands * a_chunk_bytes + num_b_operands * b_chunk_bytes num_tma_sf_group_bytes = num_a_operands * sfa_group_bytes + num_b_operands * sfb_group_bytes - cols_per_acc_stage = cta_tile_mnk[1] + epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m tmem_alloc_bar_count = (num_epilogue_warps + 1) * 32 # @@INJECT_TAP_PTRS@@ @@ -450,8 +460,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_stage_cur = cutlass.Int32(0) # incremental ring walk — no div in the loop tile_m = init_tile_m @@ -486,7 +495,35 @@ def _kernel( for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(a_chunk_packed_elems * stage) tma_a_desc = tma_a_descs[_ai] - if cutlass.const_expr(multicast_a): + if cutlass.const_expr(a_mcast_slices > 1): + _a_rows = cta_tile_mnk[0] // a_mcast_slices + _a_row_elems = a_chunk_packed_elems // cta_tile_mnk[0] + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(n_rank * _a_rows * _a_row_elems), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + n_rank * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + else: + _a_per_cta = a_mcast_slices // cluster_n + for _asl in cutlass.range(_a_per_cta): + _a_idx = n_rank * _a_per_cta + _asl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(_a_idx * _a_rows * _a_row_elems), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + _a_idx * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + elif cutlass.const_expr(multicast_a): if n_rank == 0: if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -513,7 +550,35 @@ def _kernel( for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(b_chunk_packed_elems * stage) tma_b_desc = tma_b_descs[_bj] - if cutlass.const_expr(multicast_b): + if cutlass.const_expr(b_mcast_slices > 1): + _b_rows = cta_tile_mnk[1] // b_mcast_slices + _b_row_elems = b_chunk_packed_elems // cta_tile_mnk[1] + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(m_rank * _b_rows * _b_row_elems), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + m_rank * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + else: + _b_per_cta = b_mcast_slices // cluster_m + for _bsl in cutlass.range(_b_per_cta): + _b_idx = m_rank * _b_per_cta + _bsl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(_b_idx * _b_rows * _b_row_elems), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + _b_idx * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + elif cutlass.const_expr(multicast_b): if m_rank == 0: if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -571,20 +636,19 @@ def _kernel( # incremental advance already flipped the phase on wrap. tail_stage = ab_stage_cur tail_phase = ab_empty_phase_bit - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 if warp_idx == sf_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") sf_empty_phase_bit = cutlass.Int32(1) sf_stage_cur = cutlass.Int32(0) # incremental ring walk — no div in the loop tile_m = init_tile_m @@ -704,14 +768,14 @@ def _kernel( tail_stage = sf_stage_cur tail_phase = sf_empty_phase_bit - for _ in range(sf_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == sf_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(sf_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(sf_stages): + while not nvvm.mbarrier_try_wait_parity(sf_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == sf_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 if warp_idx == mma_warp_id: nvvm.setmaxregister(mma_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -1034,22 +1098,20 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity(acc_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass - if cutlass.const_expr(use_acc_overlap): - while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): - pass + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity(acc_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + if cutlass.const_expr(use_acc_overlap): + while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): + pass nvvm.bar_warp_sync(0xFFFFFFFF) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) @@ -1076,21 +1138,14 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_epi = cutlass.Int32(0) - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - row_id_with_warp_offset = base_row_id - else: - row_id_with_warp_offset = base_row_id + warp_idx * 32 - - # One M block's accumulator columns are contiguous. - subtile_cnt = cute.ceil_div(epi_cols_per_mma_m, 32) - t2r_inst_repx = epi_tile_mn[1] - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - shape = nvvm.Tcgen05LdStShape.SHAPE_16X32BX2 - ld_half_off = 0 - else: - shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B - ld_half_off = None + # @@EPILOGUE_SETUP:BEGIN@@ + row_id_with_warp_offset = base_row_id + warp_idx * 32 + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) + subtile_cnt = len(epi_spans) + shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) @@ -1098,7 +1153,8 @@ def _kernel( while is_valid != 0: 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] + # @@EPILOGUE_DRAIN:BEGIN@@ + coord_n_c = 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: @@ -1113,58 +1169,64 @@ def _kernel( else: acc_buf_parity = cutlass.Int32(0) acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # One pass per MMA-M block over its own column region. + for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * mma_inst_shape_mnk[0] - mi_col_base = acc_base_col + mi * epi_cols_per_mma_m + if cutlass.const_expr(use_acc_overlap and num_mma_m > 1): + _mi = mi + (1 - acc_buf_parity) * (num_mma_m - 1 - 2 * mi) + else: + _mi = mi + coord_m = coord_m_tile + _mi * epi_rows_per_mma_m + mi_col_base = acc_base_col + _mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * acc_gemm_stride) for g in range(num_gemms)] - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - row = coord_m + warp_idx * 16 + lane - row_active = lane < 16 - else: - row = coord_m + tidx - row_active = True + row = coord_m + tidx + row_active = True # @@INJECT_AUX_VIEWS@@ - for subtile_idx in cutlass.range(subtile_cnt, unroll_full=True): + for subtile_idx in cutlass.range_constexpr(subtile_cnt): if cutlass.const_expr(use_acc_overlap): _sub = subtile_idx + (1 - acc_buf_parity) * (subtile_cnt - 1 - 2 * subtile_idx) - subtile_col_offset = _sub * 32 + subtile_col_offset = _sub * epi_n + subtile_w = epi_n else: - subtile_col_offset = subtile_idx * 32 + subtile_col_offset, subtile_w = epi_spans[subtile_idx] if cutlass.const_expr(not (use_tma_store_epi and cd_out_is_m_major)): c_rmem_vecs = [] for g in cutlass.range_constexpr(num_gemms): - tmem = cutlass.inttoptr( - tmem_col_addr_gemms[g] + subtile_col_offset, - 6, - cutlass.Float32, - ) - c_rmem_vecs.append(nvvm.tcgen05_ld(shape, tmem, num=t2r_inst_repx, offset=ld_half_off)) + subtile_tmem_addr = tmem_col_addr_gemms[g] + subtile_col_offset + tmem = cutlass.inttoptr(subtile_tmem_addr, 6, mma_c_dtype) + _cv = nvvm.tcgen05_ld(shape, tmem, num=subtile_w) + c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - if use_acc_overlap and (not cd_out_is_m_major) and mi == num_mma_m - 1 and subtile_idx == acc_overlap_subtiles - 1: + if cutlass.const_expr(((not use_acc_overlap) or cd_out_is_m_major) and not (use_tma_store_epi and cd_out_is_m_major)): + if cutlass.const_expr(mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1): + nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) + nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) + if elect_one: + nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) + + if use_acc_overlap and (not cd_out_is_m_major) and mi * subtile_cnt + subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) - col = coord_n + subtile_col_offset + col = coord_n_c + subtile_col_offset # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = (epi_stage_idx + 1) % EPI_SMEM_STAGES smem_subtile_ptr = smem_d_ptr.subview(epi_stage_idx * epi_subtile_elems) - smem_thr_ptr = smem_subtile_ptr.subview(tidx * t2r_inst_repx) + smem_thr_ptr = smem_subtile_ptr.subview(tidx * subtile_w) if cutlass.const_expr(cd_out_is_m_major): ld_col = mi_col_base + subtile_col_offset for _h in cutlass.range(2, unroll_full=True): ld_row = base_row_id + warp_idx * 32 + _h * 16 ld_addr = (ld_row << 16) | ld_col - ld_tmem = cutlass.inttoptr(ld_addr, 6, cutlass.Float32) - _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=4) + ld_tmem = cutlass.inttoptr(ld_addr, 6, mma_c_dtype) + _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=epi_n // 8) vec_f32 = _lv col_j = col linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 @@ -1172,7 +1234,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ _i32 = vec_out.bitcast(cutlass.Int32) - for _blk in cutlass.range_constexpr(2): + for _blk in cutlass.range_constexpr(epi_n // 16): _regs = [_i32[_blk * 4 + _j] for _j in range(4)] _n_full = (lane % 8) + 8 * (lane // 16) + 16 * _blk _m_base = warp_idx * 32 + _h * 16 + 8 * ((lane // 8) % 2) @@ -1197,7 +1259,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ - smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=cutlass.Swizzle(2, 4, 3)) + smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=epi_smem_swizzle) cute.arch.fence_view_async_shared() nvvm.barrier_cta_sync( @@ -1233,7 +1295,7 @@ def _kernel( # @@STG_ONLY:BEGIN@@ if row_active and row < M: - for j in cutlass.range_constexpr(t2r_inst_repx // vsize): + for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize if col_j + vsize <= N: vec_f32 = c_rmem_vec[j * vsize : (j + 1) * vsize] @@ -1243,12 +1305,14 @@ def _kernel( # @@INJECT_EPILOGUE@@ # @@STG_ONLY:END@@ - if cutlass.const_expr((not use_acc_overlap) or cd_out_is_m_major): + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + if cutlass.const_expr(use_tma_store_epi and cd_out_is_m_major): nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) + # @@EPILOGUE_DRAIN:END@@ consumer_stage = tile_iter % CLC_SCHED_STAGES if consumer_stage == 0 and tile_iter != 0: clc_full_phase_epi = clc_full_phase_epi ^ 1 @@ -1339,7 +1403,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[ab_tma_box_k_elems, cta_tile_mnk[0], 1], + box_dims=[ab_tma_box_k_elems, cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1377,7 +1441,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[ab_tma_box_k_elems, cta_tile_mnk[1], 1], + box_dims=[ab_tma_box_k_elems, cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1428,7 +1492,7 @@ def _host( out_stride_l_0 * cd_dtype.width // 128, ], box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], - swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + swizzle=(epi_tma_swizzle if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), ) tma_c_desc_list = [tma_c_desc] # @@TMA_STORE_ONLY:END@@ 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 353f592b6..89f0df0b9 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 @@ -45,6 +45,7 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( + epi_subtile_spans as _epi_subtile_spans, l2_swizzle_tile as _l2_swizzle_tile, tcgen05_alloc as _tcgen05_alloc, tcgen05_dealloc as _tcgen05_dealloc, @@ -62,6 +63,9 @@ # @@INJECT_TILE_CONSTANTS@@ +if use_acc_overlap and any(_w != epi_n for _, _w in _epi_subtile_spans(epi_cols_per_mma_m, epi_n)): + raise NotImplementedError(f"{__name__}: acc overlap reverses subtiles by index, which needs a uniform drain width") + CLC_SCHED_STAGES = 2 @@ -335,7 +339,10 @@ def _kernel( acc_empty_count = num_epilogue_warps * 2 cta_group = 2 - ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size // cta_group + else: + ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 num_consumer_warps_per_cta = 8 clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: @@ -383,11 +390,6 @@ def _kernel( # 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 - if cutlass.const_expr(epi_rows_per_mma_m == 64): - # cluster-MMA m=128: the pair also splits N, so each CTA drains N/2. - cols_per_acc_stage = pair_n_size // 2 - else: - cols_per_acc_stage = pair_n_size tmem_alloc_bar_count = (num_epilogue_warps + 1) * 32 nvvm.barrier_cluster_wait() @@ -465,8 +467,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_stage_cur = cutlass.Int32(0) # incremental ring walk — no div in the loop tile_m = init_tile_m @@ -503,7 +504,35 @@ def _kernel( for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(a_chunk_packed_elems * stage) tma_a_desc = tma_a_descs[_ai] - if cutlass.const_expr(multicast_a): + if cutlass.const_expr(a_mcast_slices > 1): + _a_rows = cta_tile_mnk[0] // a_mcast_slices + _a_row_elems = a_chunk_packed_elems // cta_tile_mnk[0] + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(n_rank * _a_rows * _a_row_elems), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + n_rank * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + else: + _a_per_cta = a_mcast_slices // cluster_n + for _asl in cutlass.range(_a_per_cta): + _a_idx = n_rank * _a_per_cta + _asl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(_a_idx * _a_rows * _a_row_elems), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + _a_idx * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + elif cutlass.const_expr(multicast_a): if n_rank == 0: if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -530,7 +559,35 @@ def _kernel( for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(b_chunk_packed_elems * stage) tma_b_desc = tma_b_descs[_bj] - if cutlass.const_expr(multicast_b): + if cutlass.const_expr(b_mcast_slices > 1): + _b_rows = cta_tile_mnk[1] // b_mcast_slices + _b_row_elems = b_chunk_packed_elems // cta_tile_mnk[1] + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(pair_m_idx * _b_rows * _b_row_elems), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + pair_m_idx * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + else: + _b_per_cta = b_mcast_slices // (cluster_m // cta_group) + for _bsl in cutlass.range(_b_per_cta): + _b_idx = pair_m_idx * _b_per_cta + _bsl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(_b_idx * _b_rows * _b_row_elems), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + _b_idx * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + elif cutlass.const_expr(multicast_b): if pair_m_idx == 0: if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( @@ -588,20 +645,19 @@ def _kernel( # incremental advance already flipped the phase on wrap. tail_stage = ab_stage_cur tail_phase = ab_empty_phase_bit - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 if warp_idx == sf_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") sf_empty_phase_bit = cutlass.Int32(1) sf_stage_cur = cutlass.Int32(0) # incremental ring walk — no div in the loop tile_m = init_tile_m @@ -724,14 +780,14 @@ def _kernel( tail_stage = sf_stage_cur tail_phase = sf_empty_phase_bit - for _ in range(sf_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == sf_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(sf_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(sf_stages): + while not nvvm.mbarrier_try_wait_parity(sf_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == sf_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 if warp_idx == mma_warp_id: nvvm.setmaxregister(mma_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -755,7 +811,10 @@ def _kernel( a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) - ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) peer_cta_rank = cta_rank_in_cluster ^ 1 if is_pair_leader: mma_slot_stage = cutlass.Int32(0) # ring pos of the tile's first chunk @@ -1069,19 +1128,17 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity(acc_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity(acc_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass nvvm.bar_warp_sync(0xFFFFFFFF) # Asymmetric pair rendezvous on the TMEM free: leader waits its own @@ -1124,8 +1181,7 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) @@ -1157,25 +1213,22 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_epi = cutlass.Int32(0) + # @@EPILOGUE_SETUP:BEGIN@@ row_id_with_warp_offset = base_row_id + warp_idx * 32 - if cutlass.const_expr(cols_per_acc_stage >= 32): - t2r_inst_repx = 32 - subtile_cnt = cols_per_acc_stage // 32 - else: - t2r_inst_repx = cols_per_acc_stage - subtile_cnt = 1 + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) + subtile_cnt = len(epi_spans) shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 - + # @@EPILOGUE_SETUP:END@@ # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) # @@TMA_STORE_ONLY:END@@ while is_valid != 0: coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + # @@EPILOGUE_DRAIN:BEGIN@@ 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 acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: @@ -1190,61 +1243,72 @@ def _kernel( else: acc_buf_parity = cutlass.Int32(0) acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # The 2-CTA epilogue drains its own half of the instruction's M, - # epi_rows_per_mma_m rows at a time, so a CTA tile of num_mma_m blocks - # drains in num_mma_m passes over its own column region. + for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * epi_rows_per_mma_m - mi_col_base = acc_base_col + mi * epi_cols_per_mma_m + if cutlass.const_expr(use_acc_overlap and num_mma_m > 1): + _mi = mi + (1 - acc_buf_parity) * (num_mma_m - 1 - 2 * mi) + else: + _mi = mi + coord_m = coord_m_tile + _mi * epi_rows_per_mma_m + mi_col_base = acc_base_col + _mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * acc_gemm_stride) for g in range(num_gemms)] - if cutlass.const_expr(epi_rows_per_mma_m == 64): - row = coord_m + (warp_idx % 2) * 32 + lane - row_active = True - else: - row = coord_m + tidx - row_active = True + row = coord_m + tidx + row_active = True # @@INJECT_AUX_VIEWS@@ - for subtile_idx in cutlass.range(subtile_cnt, unroll_full=True): + for subtile_idx in cutlass.range_constexpr(subtile_cnt): if cutlass.const_expr(use_acc_overlap): _sub = subtile_idx + (1 - acc_buf_parity) * (subtile_cnt - 1 - 2 * subtile_idx) - subtile_col_offset = _sub * t2r_inst_repx + subtile_col_offset = _sub * epi_n + subtile_w = epi_n else: - subtile_col_offset = subtile_idx * t2r_inst_repx + subtile_col_offset, subtile_w = epi_spans[subtile_idx] if cutlass.const_expr(not (use_tma_store_epi and cd_out_is_m_major)): c_rmem_vecs = [] for g in cutlass.range_constexpr(num_gemms): - tmem = cutlass.inttoptr( - tmem_col_addr_gemms[g] + subtile_col_offset, - 6, - cutlass.Float32, - ) - c_rmem_vecs.append(nvvm.tcgen05_ld(shape, tmem, num=t2r_inst_repx)) + subtile_tmem_addr = tmem_col_addr_gemms[g] + subtile_col_offset + tmem = cutlass.inttoptr(subtile_tmem_addr, 6, mma_c_dtype) + _cv = nvvm.tcgen05_ld(shape, tmem, num=subtile_w) + c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - if use_acc_overlap and (not cd_out_is_m_major) and mi == num_mma_m - 1 and subtile_idx == acc_overlap_subtiles - 1: + if cutlass.const_expr(((not use_acc_overlap) or cd_out_is_m_major) and not (use_tma_store_epi and cd_out_is_m_major)): + if cutlass.const_expr(mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1): + nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) + nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) + if elect_one: + nvvm.mbarrier_arrive( + nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), + scope=nvvm.MemScope.CLUSTER, + relaxed=True, + ) + + if use_acc_overlap and (not cd_out_is_m_major) and mi * subtile_cnt + subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: - mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) - nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) + nvvm.mbarrier_arrive( + nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), + scope=nvvm.MemScope.CLUSTER, + relaxed=True, + ) col = coord_n_c + subtile_col_offset # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = (epi_stage_idx + 1) % EPI_SMEM_STAGES smem_subtile_ptr = smem_d_ptr.subview(epi_stage_idx * epi_subtile_elems) - smem_thr_ptr = smem_subtile_ptr.subview(tidx * t2r_inst_repx) + smem_thr_ptr = smem_subtile_ptr.subview(tidx * subtile_w) if cutlass.const_expr(cd_out_is_m_major): ld_col = mi_col_base + subtile_col_offset for _h in cutlass.range(2, unroll_full=True): ld_row = base_row_id + warp_idx * 32 + _h * 16 ld_addr = (ld_row << 16) | ld_col - ld_tmem = cutlass.inttoptr(ld_addr, 6, cutlass.Float32) - _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=4) + ld_tmem = cutlass.inttoptr(ld_addr, 6, mma_c_dtype) + _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=epi_n // 8) vec_f32 = _lv col_j = col linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 @@ -1252,7 +1316,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ _i32 = vec_out.bitcast(cutlass.Int32) - for _blk in cutlass.range_constexpr(2): + for _blk in cutlass.range_constexpr(epi_n // 16): _regs = [_i32[_blk * 4 + _j] for _j in range(4)] _n_full = (lane % 8) + 8 * (lane // 16) + 16 * _blk _m_base = warp_idx * 32 + _h * 16 + 8 * ((lane // 8) % 2) @@ -1277,7 +1341,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ - smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=cutlass.Swizzle(2, 4, 3)) + smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=epi_smem_swizzle) cute.arch.fence_view_async_shared() nvvm.barrier_cta_sync( @@ -1313,7 +1377,7 @@ def _kernel( # @@STG_ONLY:BEGIN@@ if row_active and row < M: - for j in cutlass.range_constexpr(t2r_inst_repx // vsize): + for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize if col_j + vsize <= N: vec_f32 = c_rmem_vec[j * vsize : (j + 1) * vsize] @@ -1323,13 +1387,15 @@ def _kernel( # @@INJECT_EPILOGUE@@ # @@STG_ONLY:END@@ - if cutlass.const_expr((not use_acc_overlap) or cd_out_is_m_major): + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + if cutlass.const_expr(use_tma_store_epi and cd_out_is_m_major): nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) + # @@EPILOGUE_DRAIN:END@@ consumer_stage = tile_iter % CLC_SCHED_STAGES if consumer_stage == 0 and tile_iter != 0: clc_full_phase_epi = clc_full_phase_epi ^ 1 @@ -1420,7 +1486,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[ab_tma_box_k_elems, cta_tile_mnk[0], 1], + box_dims=[ab_tma_box_k_elems, cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1458,7 +1524,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[ab_tma_box_k_elems, cta_tile_mnk[1], 1], + box_dims=[ab_tma_box_k_elems, cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1509,7 +1575,7 @@ def _host( out_stride_l_0 * cd_dtype.width // 128, ], box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], - swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + swizzle=(epi_tma_swizzle if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), ) tma_c_desc_list = [tma_c_desc] # @@TMA_STORE_ONLY:END@@ 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 c31269e4a..4b8619a67 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 @@ -42,6 +42,7 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( + epi_subtile_spans as _epi_subtile_spans, l2_swizzle_tile as _l2_swizzle_tile, tcgen05_alloc as _tcgen05_alloc, tcgen05_dealloc as _tcgen05_dealloc, @@ -59,6 +60,9 @@ # @@INJECT_TILE_CONSTANTS@@ +if use_acc_overlap and any(_w != epi_n for _, _w in _epi_subtile_spans(epi_cols_per_mma_m, epi_n)): + raise NotImplementedError(f"{__name__}: acc overlap reverses subtiles by index, which needs a uniform drain width") + CLC_SCHED_STAGES = 2 @@ -212,11 +216,15 @@ def _kernel( a_part_arrive = cutlass.Int16(a_pattern) << m_rank b_part_arrive = cutlass.Int16(b_pattern) << (n_rank * cluster_m) - ab_empty_arrive_mask = a_part_arrive | b_part_arrive + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = a_part_arrive | b_part_arrive _smem_sys_reserved = cutlass.Array(cutlass.Int8, 1024, space=cutlass.AddressSpace.smem, alignment=1) ab_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) + sf_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) ab_empty_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) acc_empty_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) acc_full_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) @@ -286,13 +294,18 @@ def _kernel( ) # @@TMA_STORE_ONLY:END@@ - ab_empty_count = cluster_m + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size + else: + ab_empty_count = cluster_m + cluster_n - 1 num_consumer_warps_per_cta = 7 clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: for i in range(ab_stages): if elect_one: nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) + if elect_one: + nvvm.mbarrier_init(sf_full_mbar_ptr.subview(i), 1) if elect_one: nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) for i in range(acc_stages): @@ -317,9 +330,10 @@ def _kernel( sA_bytes = sA_elems * (ab_dtype.width // 8) 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) + ab_only_copy_bytes = num_a_operands * sA_bytes + num_b_operands * sB_bytes + sf_only_copy_bytes = num_a_operands * sfa_smem_bytes + num_b_operands * sfb_smem_bytes - cols_per_acc_stage = cta_tile_mnk[1] + epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m tmem_alloc_bar_count = (num_epilogue_warps + 1) * 32 # @@INJECT_TAP_PTRS@@ @@ -394,8 +408,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) tile_m = init_tile_m @@ -427,7 +440,9 @@ def _kernel( coord_k = k_tile_idx * cta_tile_mnk[2] coord_sf_k = k_tile_idx * sf_tma_box_k if elect_one: - nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) + nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), ab_only_copy_bytes) + if elect_one: + nvvm.mbarrier_arrive_expect_tx(sf_full_mbar_ptr.subview(stage), sf_only_copy_bytes) for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) @@ -436,6 +451,93 @@ def _kernel( tma_sfa_desc = tma_sfa_descs[_ai] sfa_m_block = coord_m_per_cta // 128 if cutlass.const_expr(multicast_a): + if n_rank == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFA_stage, + tma_sfa_desc.get_ptr(), + (0, coord_sf_k, sfa_m_block, tile_l_a), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + else: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFA_stage, + tma_sfa_desc.get_ptr(), + (0, coord_sf_k, sfa_m_block, tile_l_a), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + + for _bj in cutlass.range_constexpr(num_b_operands): + sB_stage = smem_b_list[_bj].subview(sB_elems * stage) + tma_b_desc = tma_b_descs[_bj] + sSFB_stage = smem_sfb_list[_bj].subview(sfb_smem_bytes * stage) + tma_sfb_desc = tma_sfb_descs[_bj] + sfb_n_block = coord_n_per_cta // 128 + if cutlass.const_expr(multicast_b): + if m_rank == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFB_stage, + tma_sfb_desc.get_ptr(), + (0, coord_sf_k, sfb_n_block, tile_l_b), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + else: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFB_stage, + tma_sfb_desc.get_ptr(), + (0, coord_sf_k, sfb_n_block, tile_l_b), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + + for _ai in cutlass.range_constexpr(num_a_operands): + sA_stage = smem_a_list[_ai].subview(sA_elems * stage) + tma_a_desc = tma_a_descs[_ai] + sSFA_stage = smem_sfa_list[_ai].subview(sfa_smem_bytes * stage) + tma_sfa_desc = tma_sfa_descs[_ai] + sfa_m_block = coord_m_per_cta // 128 + if cutlass.const_expr(a_mcast_slices > 1): + _a_rows = cta_tile_mnk[0] // a_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(n_rank * _a_rows * ab_packed_per_row), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + n_rank * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + else: + _a_per_cta = a_mcast_slices // cluster_n + for _asl in cutlass.range(_a_per_cta): + _a_idx = n_rank * _a_per_cta + _asl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(_a_idx * _a_rows * ab_packed_per_row), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + _a_idx * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_1, + ) + elif cutlass.const_expr(multicast_a): if n_rank == 0: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): @@ -492,37 +594,40 @@ def _kernel( multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_1, ) - if cutlass.const_expr(multicast_a): - if n_rank == 0: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFA_stage, - tma_sfa_desc.get_ptr(), - (0, coord_sf_k, sfa_m_block, tile_l_a), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_a, - group=nvvm.CTAGroup.CTA_1, - ) - else: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFA_stage, - tma_sfa_desc.get_ptr(), - (0, coord_sf_k, sfa_m_block, tile_l_a), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_a, - group=nvvm.CTAGroup.CTA_1, - ) - for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) tma_b_desc = tma_b_descs[_bj] sSFB_stage = smem_sfb_list[_bj].subview(sfb_smem_bytes * stage) tma_sfb_desc = tma_sfb_descs[_bj] sfb_n_block = coord_n_per_cta // 128 - if cutlass.const_expr(multicast_b): + if cutlass.const_expr(b_mcast_slices > 1): + _b_rows = cta_tile_mnk[1] // b_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(m_rank * _b_rows * ab_packed_per_row), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + m_rank * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + else: + _b_per_cta = b_mcast_slices // cluster_m + for _bsl in cutlass.range(_b_per_cta): + _b_idx = m_rank * _b_per_cta + _bsl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(_b_idx * _b_rows * ab_packed_per_row), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + _b_idx * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + elif cutlass.const_expr(multicast_b): if m_rank == 0: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): @@ -579,30 +684,6 @@ def _kernel( multicast_mask=tma_mcast_mask_b, group=nvvm.CTAGroup.CTA_1, ) - if cutlass.const_expr(multicast_b): - if m_rank == 0: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFB_stage, - tma_sfb_desc.get_ptr(), - (0, coord_sf_k, sfb_n_block, tile_l_b), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_1, - ) - else: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFB_stage, - tma_sfb_desc.get_ptr(), - (0, coord_sf_k, sfb_n_block, tile_l_b), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_1, - ) - ab_iter += 1 consumer_stage = tile_iter % CLC_SCHED_STAGES @@ -635,14 +716,14 @@ def _kernel( tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -713,7 +794,7 @@ def _kernel( sfa_dst_ptrs = [ [ [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block + a * registers_per_atom, cutlass.Float32) for a in range(word_atoms)] - for m in range(num_blocks_m) + for m in range(num_mma_m) ] for i in range(num_a_operands) ] @@ -760,9 +841,6 @@ def _kernel( if stage == 0 and ab_iter != 0: ab_full_phase_bit = ab_full_phase_bit ^ 1 - while not nvvm.mbarrier_try_wait_parity(ab_full_mbar_ptr.subview(stage), ab_full_phase_bit, time_limit=10_000_000): - pass - desc_a_bases = [ cutlass.experimental.primitives.Tcgen05SmemDesc.build( start_address=smem_a_list[i].subview(sA_elems * stage), @@ -802,58 +880,63 @@ def _kernel( # One SF word per group of MMAs, refreshed right before they read # it. A word spans word_atoms consecutive K-atoms in SMEM. - for atom_r in cutlass.range(num_sf_atoms, unroll_full=True): - for _ai in cutlass.range_constexpr(num_a_operands): - for _m in cutlass.range_constexpr(num_blocks_m): - for _a in cutlass.range_constexpr(word_atoms): - if elect_one: - nvvm.tcgen05_cp( - s2t_shape, - sfa_dst_ptrs[_ai][_m][_a], - desc_sfa_bases[_ai] + (sf_atom_desc_stride * (atom_r * word_atoms + _a) + sf_block_desc_stride * _m), - group=nvvm.CTAGroup.CTA_1, - multicast=s2t_multicast, - ) + while not nvvm.mbarrier_try_wait_parity(sf_full_mbar_ptr.subview(stage), ab_full_phase_bit, time_limit=10_000_000): + pass + + for sf_word in cutlass.range_constexpr(num_sf_atoms): for _bj in cutlass.range_constexpr(num_b_operands): - for _m in cutlass.range_constexpr(num_blocks_n): + for block_n in cutlass.range_constexpr(num_blocks_n): for _a in cutlass.range_constexpr(word_atoms): if elect_one: nvvm.tcgen05_cp( s2t_shape, - sfb_dst_ptrs[_bj][_m][_a], - desc_sfb_bases[_bj] + (sf_atom_desc_stride * (atom_r * word_atoms + _a) + sf_block_desc_stride * _m), + sfb_dst_ptrs[_bj][block_n][_a], + desc_sfb_bases[_bj] + (sf_atom_desc_stride * (sf_word * word_atoms + _a) + sf_block_desc_stride * block_n), group=nvvm.CTAGroup.CTA_1, multicast=s2t_multicast, ) - for j in cutlass.range_constexpr(sf_insts_per_atom): - k_block_idx = atom_r * sf_insts_per_atom + j - idesc_k = idesc_by_j[j] - for g in cutlass.range_constexpr(num_gemms): - _ai = gemm_a_idx[g] - _bj = gemm_b_idx[g] - desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) - desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - for mi in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(sf_word == 0): + while not nvvm.mbarrier_try_wait_parity(ab_full_mbar_ptr.subview(stage), ab_full_phase_bit, time_limit=10_000_000): + pass + for mma_k_in_word in cutlass.range_constexpr(sf_insts_per_atom): + mma_k = sf_word * sf_insts_per_atom + mma_k_in_word + idesc_k = idesc_by_j[mma_k_in_word] + for gemm_i in cutlass.range_constexpr(num_gemms): + _ai = gemm_a_idx[gemm_i] + _bj = gemm_b_idx[gemm_i] + desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * mma_k) + desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * mma_k) + for mma_m in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(mma_k_in_word == 0 and _ai not in gemm_a_idx[:gemm_i]): + for _a in cutlass.range_constexpr(word_atoms): + if elect_one: + nvvm.tcgen05_cp( + s2t_shape, + sfa_dst_ptrs[_ai][mma_m][_a], + desc_sfa_bases[_ai] + (sf_atom_desc_stride * (sf_word * word_atoms + _a) + sf_block_desc_stride * mma_m), + group=nvvm.CTAGroup.CTA_1, + multicast=s2t_multicast, + ) # The M sub-block offset is a whole SMEM swizzle atom, so the # descriptor's swizzle phase is preserved. B and its SF are # shared; A's SF word block follows the M block. - desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) + desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mma_m) if elect_one: _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_1, - acc_tmem_ptrs[g][mi], + acc_tmem_ptrs[gemm_i][mma_m], desc_a, desc_b, idesc_k, enable_input_d=scale_d, - scale_a=sfa_dst_ptrs[_ai][mi][0], + scale_a=sfa_dst_ptrs[_ai][mma_m][0], scale_b=sfb_scale_ptrs[_bj], scale_vec_size=scale_vec_size, - b_collector_op=_b_collector_op(mi), + b_collector_op=_b_collector_op(mma_m), ) # Every accumulator sees scale_d=False on exactly the first - # k_block of the tile, so the flip stays outside mi. + # k_block of the tile, so the flip stays outside mma_m. scale_d = cutlass.Boolean(True) if elect_one: @@ -889,22 +972,20 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity(acc_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass - if cutlass.const_expr(use_acc_overlap): - while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): - pass + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity(acc_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + if cutlass.const_expr(use_acc_overlap): + while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): + pass nvvm.bar_warp_sync(0xFFFFFFFF) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) @@ -931,21 +1012,14 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_epi = cutlass.Int32(0) - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - row_id_with_warp_offset = base_row_id - else: - row_id_with_warp_offset = base_row_id + warp_idx * 32 - - # One M block's accumulator columns are contiguous. - subtile_cnt = cute.ceil_div(epi_cols_per_mma_m, 32) - t2r_inst_repx = epi_tile_mn[1] - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - shape = nvvm.Tcgen05LdStShape.SHAPE_16X32BX2 - ld_half_off = 0 - else: - shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B - ld_half_off = None + # @@EPILOGUE_SETUP:BEGIN@@ + row_id_with_warp_offset = base_row_id + warp_idx * 32 + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) + subtile_cnt = len(epi_spans) + shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) @@ -953,7 +1027,8 @@ def _kernel( while is_valid != 0: 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] + # @@EPILOGUE_DRAIN:BEGIN@@ + coord_n_c = 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: @@ -968,58 +1043,64 @@ def _kernel( else: acc_buf_parity = cutlass.Int32(0) acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # One pass per MMA-M block over its own column region. + for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * mma_inst_shape_mnk[0] - mi_col_base = acc_base_col + mi * epi_cols_per_mma_m + if cutlass.const_expr(use_acc_overlap and num_mma_m > 1): + _mi = mi + (1 - acc_buf_parity) * (num_mma_m - 1 - 2 * mi) + else: + _mi = mi + coord_m = coord_m_tile + _mi * epi_rows_per_mma_m + mi_col_base = acc_base_col + _mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * acc_gemm_stride) for g in range(num_gemms)] - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - row = coord_m + warp_idx * 16 + lane - row_active = lane < 16 - else: - row = coord_m + tidx - row_active = True + row = coord_m + tidx + row_active = True # @@INJECT_AUX_VIEWS@@ - for subtile_idx in cutlass.range(subtile_cnt, unroll_full=True): + for subtile_idx in cutlass.range_constexpr(subtile_cnt): if cutlass.const_expr(use_acc_overlap): _sub = subtile_idx + (1 - acc_buf_parity) * (subtile_cnt - 1 - 2 * subtile_idx) - subtile_col_offset = _sub * 32 + subtile_col_offset = _sub * epi_n + subtile_w = epi_n else: - subtile_col_offset = subtile_idx * 32 + subtile_col_offset, subtile_w = epi_spans[subtile_idx] if cutlass.const_expr(not (use_tma_store_epi and cd_out_is_m_major)): c_rmem_vecs = [] for g in cutlass.range_constexpr(num_gemms): - tmem = cutlass.inttoptr( - tmem_col_addr_gemms[g] + subtile_col_offset, - 6, - cutlass.Float32, - ) - c_rmem_vecs.append(nvvm.tcgen05_ld(shape, tmem, num=t2r_inst_repx, offset=ld_half_off)) + subtile_tmem_addr = tmem_col_addr_gemms[g] + subtile_col_offset + tmem = cutlass.inttoptr(subtile_tmem_addr, 6, mma_c_dtype) + _cv = nvvm.tcgen05_ld(shape, tmem, num=subtile_w) + c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - if use_acc_overlap and (not cd_out_is_m_major) and mi == num_mma_m - 1 and subtile_idx == acc_overlap_subtiles - 1: + if cutlass.const_expr(((not use_acc_overlap) or cd_out_is_m_major) and not (use_tma_store_epi and cd_out_is_m_major)): + if cutlass.const_expr(mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1): + nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) + nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) + if elect_one: + nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) + + if use_acc_overlap and (not cd_out_is_m_major) and mi * subtile_cnt + subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) - col = coord_n + subtile_col_offset + col = coord_n_c + subtile_col_offset # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = (epi_stage_idx + 1) % EPI_SMEM_STAGES smem_subtile_ptr = smem_d_ptr.subview(epi_stage_idx * epi_subtile_elems) - smem_thr_ptr = smem_subtile_ptr.subview(tidx * t2r_inst_repx) + smem_thr_ptr = smem_subtile_ptr.subview(tidx * subtile_w) if cutlass.const_expr(cd_out_is_m_major): ld_col = mi_col_base + subtile_col_offset for _h in cutlass.range(2, unroll_full=True): ld_row = base_row_id + warp_idx * 32 + _h * 16 ld_addr = (ld_row << 16) | ld_col - ld_tmem = cutlass.inttoptr(ld_addr, 6, cutlass.Float32) - _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=4) + ld_tmem = cutlass.inttoptr(ld_addr, 6, mma_c_dtype) + _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=epi_n // 8) vec_f32 = _lv col_j = col linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 @@ -1027,7 +1108,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ _i32 = vec_out.bitcast(cutlass.Int32) - for _blk in cutlass.range_constexpr(2): + for _blk in cutlass.range_constexpr(epi_n // 16): _regs = [_i32[_blk * 4 + _j] for _j in range(4)] _n_full = (lane % 8) + 8 * (lane // 16) + 16 * _blk _m_base = warp_idx * 32 + _h * 16 + 8 * ((lane // 8) % 2) @@ -1052,7 +1133,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ - smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=cutlass.Swizzle(2, 4, 3)) + smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=epi_smem_swizzle) cute.arch.fence_view_async_shared() nvvm.barrier_cta_sync( @@ -1088,7 +1169,7 @@ def _kernel( # @@STG_ONLY:BEGIN@@ if row_active and row < M: - for j in cutlass.range_constexpr(t2r_inst_repx // vsize): + for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize if col_j + vsize <= N: vec_f32 = c_rmem_vec[j * vsize : (j + 1) * vsize] @@ -1098,12 +1179,14 @@ def _kernel( # @@INJECT_EPILOGUE@@ # @@STG_ONLY:END@@ - if cutlass.const_expr((not use_acc_overlap) or cd_out_is_m_major): + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + if cutlass.const_expr(use_tma_store_epi and cd_out_is_m_major): nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) + # @@EPILOGUE_DRAIN:END@@ consumer_stage = tile_iter % CLC_SCHED_STAGES if consumer_stage == 0 and tile_iter != 0: clc_full_phase_epi = clc_full_phase_epi ^ 1 @@ -1211,7 +1294,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1265,7 +1348,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1316,7 +1399,7 @@ def _host( out_stride_l_0 * cd_dtype.width // 128, ], box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], - swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + swizzle=(epi_tma_swizzle if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), ) tma_c_desc_list = [tma_c_desc] # @@TMA_STORE_ONLY:END@@ 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 e1b9cf461..360f86d85 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 @@ -41,6 +41,7 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( + epi_subtile_spans as _epi_subtile_spans, l2_swizzle_tile as _l2_swizzle_tile, tcgen05_alloc as _tcgen05_alloc, tcgen05_dealloc as _tcgen05_dealloc, @@ -58,6 +59,9 @@ # @@INJECT_TILE_CONSTANTS@@ +if use_acc_overlap and any(_w != epi_n for _, _w in _epi_subtile_spans(epi_cols_per_mma_m, epi_n)): + raise NotImplementedError(f"{__name__}: acc overlap reverses subtiles by index, which needs a uniform drain width") + # Scheduler ring depth. CLC_SCHED_STAGES = 2 @@ -214,6 +218,7 @@ def _kernel( _smem_sys_reserved = cutlass.Array(cutlass.Int8, 1024, space=cutlass.AddressSpace.smem, alignment=1) ab_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) + sf_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) ab_empty_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) acc_empty_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) acc_full_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) @@ -285,7 +290,10 @@ def _kernel( acc_empty_count = num_epilogue_warps * 2 cta_group = 2 - ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size // cta_group + else: + ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 num_consumer_warps_per_cta = 7 clc_empty_count = num_consumer_warps_per_cta * cluster_size if warp_idx == 0: @@ -298,6 +306,8 @@ def _kernel( for i in range(ab_stages): if elect_one: nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) + if elect_one: + nvvm.mbarrier_init(sf_full_mbar_ptr.subview(i), 1) if elect_one: nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) for i in range(acc_stages): @@ -315,7 +325,8 @@ def _kernel( sA_bytes = sA_elems * (ab_dtype.width // 8) 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 + ab_only_copy_bytes = (num_a_operands * sA_bytes + num_b_operands * sB_bytes) * 2 + sf_only_copy_bytes = (num_a_operands * sfa_smem_bytes + num_b_operands * sfb_smem_bytes) * 2 # Per-CTA logical tile — the cluster cancels out, so these stay compile-time # constants even when the cluster shape is only known at runtime. @@ -325,11 +336,6 @@ def _kernel( # 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 - if cutlass.const_expr(epi_rows_per_mma_m == 64): - # cluster-MMA m=128: the pair also splits N, so each CTA drains N/2. - cols_per_acc_stage = pair_n_size // 2 - else: - cols_per_acc_stage = pair_n_size tmem_alloc_bar_count = (num_epilogue_warps + 1) * 32 nvvm.barrier_cluster_wait() @@ -407,8 +413,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) tile_m = init_tile_m @@ -444,7 +449,9 @@ def _kernel( if is_pair_leader: if elect_one: - nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) + nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), ab_only_copy_bytes) + if elect_one: + nvvm.mbarrier_arrive_expect_tx(sf_full_mbar_ptr.subview(stage), sf_only_copy_bytes) for _ai in cutlass.range_constexpr(num_a_operands): sA_stage = smem_a_list[_ai].subview(sA_elems * stage) @@ -453,6 +460,92 @@ def _kernel( tma_sfa_desc = tma_sfa_descs[_ai] sfa_m_block = coord_m_per_cta // 128 if cutlass.const_expr(multicast_a): + if n_rank == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFA_stage, + tma_sfa_desc.get_ptr(), + (0, coord_sf_k, sfa_m_block, tile_l_a), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + else: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFA_stage, + tma_sfa_desc.get_ptr(), + (0, coord_sf_k, sfa_m_block, tile_l_a), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + + for _bj in cutlass.range_constexpr(num_b_operands): + sB_stage = smem_b_list[_bj].subview(sB_elems * stage) + tma_b_desc = tma_b_descs[_bj] + sSFB_stage = smem_sfb_list[_bj].subview(sfb_smem_bytes * stage) + tma_sfb_desc = tma_sfb_descs[_bj] + if cutlass.const_expr(multicast_b): + if pair_m_idx == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFB_stage, + tma_sfb_desc.get_ptr(), + (0, coord_sf_k, sfb_n_block, tile_l_b), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + else: + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sSFB_stage, + tma_sfb_desc.get_ptr(), + (0, coord_sf_k, sfb_n_block, tile_l_b), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + + for _ai in cutlass.range_constexpr(num_a_operands): + sA_stage = smem_a_list[_ai].subview(sA_elems * stage) + tma_a_desc = tma_a_descs[_ai] + sSFA_stage = smem_sfa_list[_ai].subview(sfa_smem_bytes * stage) + tma_sfa_desc = tma_sfa_descs[_ai] + sfa_m_block = coord_m_per_cta // 128 + if cutlass.const_expr(a_mcast_slices > 1): + _a_rows = cta_tile_mnk[0] // a_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(n_rank * _a_rows * ab_packed_per_row), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + n_rank * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + else: + _a_per_cta = a_mcast_slices // cluster_n + for _asl in cutlass.range(_a_per_cta): + _a_idx = n_rank * _a_per_cta + _asl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sA_stage.subview(_a_idx * _a_rows * ab_packed_per_row), + tma_a_desc.get_ptr(), + (coord_k, coord_m_per_cta + _a_idx * _a_rows, tile_l_a), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_a, + group=nvvm.CTAGroup.CTA_2, + ) + elif cutlass.const_expr(multicast_a): if n_rank == 0: if cutlass.const_expr(a_is_m_major): for m_group in cutlass.range_constexpr(cta_tile_mnk[0] // a_tma_group_elems): @@ -509,36 +602,39 @@ def _kernel( multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_2, ) - if cutlass.const_expr(multicast_a): - if n_rank == 0: + for _bj in cutlass.range_constexpr(num_b_operands): + sB_stage = smem_b_list[_bj].subview(sB_elems * stage) + tma_b_desc = tma_b_descs[_bj] + sSFB_stage = smem_sfb_list[_bj].subview(sfb_smem_bytes * stage) + tma_sfb_desc = tma_sfb_descs[_bj] + if cutlass.const_expr(b_mcast_slices > 1): + _b_rows = cta_tile_mnk[1] // b_mcast_slices + if cutlass.const_expr(fallback_cluster_shape_mnk is None): if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFA_stage, - tma_sfa_desc.get_ptr(), - (0, coord_sf_k, sfa_m_block, tile_l_a), + sB_stage.subview(pair_m_idx * _b_rows * ab_packed_per_row), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + pair_m_idx * _b_rows, tile_l_b), ab_full_mbar_ptr.subview(stage), [], - multicast_mask=tma_mcast_mask_a, + multicast_mask=tma_mcast_mask_b, group=nvvm.CTAGroup.CTA_2, ) - else: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFA_stage, - tma_sfa_desc.get_ptr(), - (0, coord_sf_k, sfa_m_block, tile_l_a), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_a, - group=nvvm.CTAGroup.CTA_2, - ) - - for _bj in cutlass.range_constexpr(num_b_operands): - sB_stage = smem_b_list[_bj].subview(sB_elems * stage) - tma_b_desc = tma_b_descs[_bj] - sSFB_stage = smem_sfb_list[_bj].subview(sfb_smem_bytes * stage) - tma_sfb_desc = tma_sfb_descs[_bj] - if cutlass.const_expr(multicast_b): + else: + _b_per_cta = b_mcast_slices // (cluster_m // cta_group) + for _bsl in cutlass.range(_b_per_cta): + _b_idx = pair_m_idx * _b_per_cta + _bsl + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + sB_stage.subview(_b_idx * _b_rows * ab_packed_per_row), + tma_b_desc.get_ptr(), + (coord_k, coord_n_per_cta + _b_idx * _b_rows, tile_l_b), + ab_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + elif cutlass.const_expr(multicast_b): if pair_m_idx == 0: if cutlass.const_expr(b_is_n_major): for n_group in cutlass.range_constexpr(cta_tile_mnk[1] // b_tma_group_elems): @@ -595,30 +691,6 @@ def _kernel( multicast_mask=tma_mcast_mask_b, group=nvvm.CTAGroup.CTA_2, ) - if cutlass.const_expr(multicast_b): - if pair_m_idx == 0: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFB_stage, - tma_sfb_desc.get_ptr(), - (0, coord_sf_k, sfb_n_block, tile_l_b), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_2, - ) - else: - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - sSFB_stage, - tma_sfb_desc.get_ptr(), - (0, coord_sf_k, sfb_n_block, tile_l_b), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_2, - ) - ab_iter += 1 consumer_stage = tile_iter % CLC_SCHED_STAGES @@ -651,14 +723,14 @@ def _kernel( tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 pair_mask = cutlass.Int16(3) << pair_leader_rank a_arrive_pattern = a_mcast_pattern @@ -669,7 +741,10 @@ def _kernel( a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) - ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) _tcgen05_alloc( @@ -744,7 +819,7 @@ def _kernel( sfa_dst_ptrs = [ [ [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block + a * registers_per_atom, cutlass.Float32) for a in range(word_atoms)] - for m in range(num_blocks_m) + for m in range(num_mma_m) ] for i in range(num_a_operands) ] @@ -791,13 +866,6 @@ def _kernel( if stage == 0 and ab_iter != 0: ab_full_phase_bit = ab_full_phase_bit ^ 1 - while not nvvm.mbarrier_try_wait_parity( - ab_full_mbar_ptr.subview(stage), - ab_full_phase_bit, - time_limit=10_000_000, - ): - pass - desc_a_bases = [ cutlass.experimental.primitives.Tcgen05SmemDesc.build( start_address=smem_a_list[i].subview(sA_elems * stage), @@ -837,58 +905,71 @@ def _kernel( # One SF word per group of MMAs, refreshed right before they # read it. A word spans word_atoms consecutive K-atoms in SMEM. - for atom_r in cutlass.range(num_sf_atoms, unroll_full=True): - for _ai in cutlass.range_constexpr(num_a_operands): - for _m in cutlass.range_constexpr(num_blocks_m): - for _a in cutlass.range_constexpr(word_atoms): - if elect_one: - nvvm.tcgen05_cp( - s2t_shape, - sfa_dst_ptrs[_ai][_m][_a], - desc_sfa_bases[_ai] + (sf_atom_desc_stride * (atom_r * word_atoms + _a) + sf_block_desc_stride * _m), - group=nvvm.CTAGroup.CTA_2, - multicast=s2t_multicast, - ) + while not nvvm.mbarrier_try_wait_parity( + sf_full_mbar_ptr.subview(stage), + ab_full_phase_bit, + time_limit=10_000_000, + ): + pass + + for sf_word in cutlass.range_constexpr(num_sf_atoms): for _bj in cutlass.range_constexpr(num_b_operands): - for _m in cutlass.range_constexpr(num_blocks_n): + for block_n in cutlass.range_constexpr(num_blocks_n): for _a in cutlass.range_constexpr(word_atoms): if elect_one: nvvm.tcgen05_cp( s2t_shape, - sfb_dst_ptrs[_bj][_m][_a], - desc_sfb_bases[_bj] + (sf_atom_desc_stride * (atom_r * word_atoms + _a) + sf_block_desc_stride * _m), + sfb_dst_ptrs[_bj][block_n][_a], + desc_sfb_bases[_bj] + (sf_atom_desc_stride * (sf_word * word_atoms + _a) + sf_block_desc_stride * block_n), group=nvvm.CTAGroup.CTA_2, multicast=s2t_multicast, ) - for j in cutlass.range_constexpr(sf_insts_per_atom): - k_block_idx = atom_r * sf_insts_per_atom + j - idesc_k = idesc_by_j[j] - for g in cutlass.range_constexpr(num_gemms): - _ai = gemm_a_idx[g] - _bj = gemm_b_idx[g] - desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) - desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - for mi in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(sf_word == 0): + while not nvvm.mbarrier_try_wait_parity( + ab_full_mbar_ptr.subview(stage), + ab_full_phase_bit, + time_limit=10_000_000, + ): + pass + for mma_k_in_word in cutlass.range_constexpr(sf_insts_per_atom): + mma_k = sf_word * sf_insts_per_atom + mma_k_in_word + idesc_k = idesc_by_j[mma_k_in_word] + for gemm_i in cutlass.range_constexpr(num_gemms): + _ai = gemm_a_idx[gemm_i] + _bj = gemm_b_idx[gemm_i] + desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * mma_k) + desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * mma_k) + for mma_m in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(mma_k_in_word == 0 and _ai not in gemm_a_idx[:gemm_i]): + for _a in cutlass.range_constexpr(word_atoms): + if elect_one: + nvvm.tcgen05_cp( + s2t_shape, + sfa_dst_ptrs[_ai][mma_m][_a], + desc_sfa_bases[_ai] + (sf_atom_desc_stride * (sf_word * word_atoms + _a) + sf_block_desc_stride * mma_m), + group=nvvm.CTAGroup.CTA_2, + multicast=s2t_multicast, + ) # The M sub-block offset is a whole SMEM swizzle atom, so # the descriptor's swizzle phase is preserved. B and its SF # are shared; A's SF word block follows the M block. - desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) + desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mma_m) if elect_one: _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_2, - acc_tmem_ptrs[g][mi], + acc_tmem_ptrs[gemm_i][mma_m], desc_a, desc_b, idesc_k, enable_input_d=scale_d, - scale_a=sfa_dst_ptrs[_ai][mi][0], + scale_a=sfa_dst_ptrs[_ai][mma_m][0], scale_b=sfb_scale_ptrs[_bj], scale_vec_size=scale_vec_size, - b_collector_op=_b_collector_op(mi), + b_collector_op=_b_collector_op(mma_m), ) # Every accumulator sees scale_d=False on exactly the first - # k_block of the tile, so the flip stays outside mi. + # k_block of the tile, so the flip stays outside mma_m. scale_d = cutlass.Boolean(True) if elect_one: @@ -925,25 +1006,21 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_empty_mbar_ptr.subview(tail_stage), - tail_phase, - time_limit=10_000_000, - ): - pass - nvvm.bar_warp_sync(0xFFFFFFFF) - + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity( + acc_empty_mbar_ptr.subview(tail_stage), + tail_phase, + time_limit=10_000_000, + ): + pass nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): @@ -981,8 +1058,7 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) @@ -1016,25 +1092,22 @@ def _kernel( is_valid = cutlass.Int32(1) clc_full_phase_epi = cutlass.Int32(0) + # @@EPILOGUE_SETUP:BEGIN@@ row_id_with_warp_offset = base_row_id + warp_idx * 32 - if cutlass.const_expr(cols_per_acc_stage >= 32): - t2r_inst_repx = 32 - subtile_cnt = cols_per_acc_stage // 32 - else: - t2r_inst_repx = cols_per_acc_stage - subtile_cnt = 1 + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) + subtile_cnt = len(epi_spans) shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 - + # @@EPILOGUE_SETUP:END@@ # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) # @@TMA_STORE_ONLY:END@@ while is_valid != 0: coord_m_tile = tile_m * cgrp_tile_m_cur + m_rank * cta_tile_mnk[0] + # @@EPILOGUE_DRAIN:BEGIN@@ 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 acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: @@ -1049,62 +1122,72 @@ def _kernel( else: acc_buf_parity = cutlass.Int32(0) acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # The 2-CTA epilogue drains its own half of the instruction's M, - # epi_rows_per_mma_m rows at a time, so a CTA tile of num_mma_m blocks - # drains in num_mma_m passes over its own column region. + for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * epi_rows_per_mma_m - mi_col_base = acc_base_col + mi * epi_cols_per_mma_m + if cutlass.const_expr(use_acc_overlap and num_mma_m > 1): + _mi = mi + (1 - acc_buf_parity) * (num_mma_m - 1 - 2 * mi) + else: + _mi = mi + coord_m = coord_m_tile + _mi * epi_rows_per_mma_m + mi_col_base = acc_base_col + _mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * acc_gemm_stride) for g in range(num_gemms)] - if cutlass.const_expr(epi_rows_per_mma_m == 64): - row = coord_m + (warp_idx % 2) * 32 + lane - row_active = True - else: - row = coord_m + tidx - row_active = True + row = coord_m + tidx + row_active = True # @@INJECT_AUX_VIEWS@@ - for subtile_idx in cutlass.range(subtile_cnt, unroll_full=True): + for subtile_idx in cutlass.range_constexpr(subtile_cnt): if cutlass.const_expr(use_acc_overlap): _sub = subtile_idx + (1 - acc_buf_parity) * (subtile_cnt - 1 - 2 * subtile_idx) - subtile_col_offset = _sub * t2r_inst_repx + subtile_col_offset = _sub * epi_n + subtile_w = epi_n else: - subtile_col_offset = subtile_idx * t2r_inst_repx - + subtile_col_offset, subtile_w = epi_spans[subtile_idx] if cutlass.const_expr(not (use_tma_store_epi and cd_out_is_m_major)): c_rmem_vecs = [] for g in cutlass.range_constexpr(num_gemms): - tmem = cutlass.inttoptr( - tmem_col_addr_gemms[g] + subtile_col_offset, - 6, - cutlass.Float32, - ) - c_rmem_vecs.append(nvvm.tcgen05_ld(shape, tmem, num=t2r_inst_repx)) + subtile_tmem_addr = tmem_col_addr_gemms[g] + subtile_col_offset + tmem = cutlass.inttoptr(subtile_tmem_addr, 6, mma_c_dtype) + _cv = nvvm.tcgen05_ld(shape, tmem, num=subtile_w) + c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - if use_acc_overlap and (not cd_out_is_m_major) and mi == num_mma_m - 1 and subtile_idx == acc_overlap_subtiles - 1: + if cutlass.const_expr(((not use_acc_overlap) or cd_out_is_m_major) and not (use_tma_store_epi and cd_out_is_m_major)): + if cutlass.const_expr(mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1): + nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) + nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) + if elect_one: + nvvm.mbarrier_arrive( + nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), + scope=nvvm.MemScope.CLUSTER, + relaxed=True, + ) + + if use_acc_overlap and (not cd_out_is_m_major) and mi * subtile_cnt + subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: - mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) - nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) + nvvm.mbarrier_arrive( + nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), + scope=nvvm.MemScope.CLUSTER, + relaxed=True, + ) col = coord_n_c + subtile_col_offset # @@TMA_STORE_ONLY:BEGIN@@ epi_stage_idx = (epi_stage_idx + 1) % EPI_SMEM_STAGES smem_subtile_ptr = smem_d_ptr.subview(epi_stage_idx * epi_subtile_elems) - smem_thr_ptr = smem_subtile_ptr.subview(tidx * t2r_inst_repx) + smem_thr_ptr = smem_subtile_ptr.subview(tidx * subtile_w) if cutlass.const_expr(cd_out_is_m_major): ld_col = mi_col_base + subtile_col_offset for _h in cutlass.range(2, unroll_full=True): ld_row = base_row_id + warp_idx * 32 + _h * 16 ld_addr = (ld_row << 16) | ld_col - ld_tmem = cutlass.inttoptr(ld_addr, 6, cutlass.Float32) - _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=4) + ld_tmem = cutlass.inttoptr(ld_addr, 6, mma_c_dtype) + _lv = nvvm.tcgen05_ld(nvvm.Tcgen05LdStShape.SHAPE_16X256B, ld_tmem, num=epi_n // 8) vec_f32 = _lv col_j = col linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 @@ -1112,7 +1195,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ _i32 = vec_out.bitcast(cutlass.Int32) - for _blk in cutlass.range_constexpr(2): + for _blk in cutlass.range_constexpr(epi_n // 16): _regs = [_i32[_blk * 4 + _j] for _j in range(4)] _n_full = (lane % 8) + 8 * (lane // 16) + 16 * _blk _m_base = warp_idx * 32 + _h * 16 + 8 * ((lane // 8) % 2) @@ -1137,7 +1220,7 @@ def _kernel( # @@INJECT_EPILOGUE@@ - smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=cutlass.Swizzle(2, 4, 3)) + smem_thr_ptr.data_ptr().store_swizzled(vec_out, alignment=64, swizzle=epi_smem_swizzle) cute.arch.fence_view_async_shared() nvvm.barrier_cta_sync( @@ -1173,7 +1256,7 @@ def _kernel( # @@STG_ONLY:BEGIN@@ if row_active and row < M: - for j in cutlass.range_constexpr(t2r_inst_repx // vsize): + for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize if col_j + vsize <= N: vec_f32 = c_rmem_vec[j * vsize : (j + 1) * vsize] @@ -1183,13 +1266,15 @@ def _kernel( # @@INJECT_EPILOGUE@@ # @@STG_ONLY:END@@ - if cutlass.const_expr((not use_acc_overlap) or cd_out_is_m_major): + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + if cutlass.const_expr(use_tma_store_epi and cd_out_is_m_major): nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) + # @@EPILOGUE_DRAIN:END@@ consumer_stage = tile_iter % CLC_SCHED_STAGES if consumer_stage == 0 and tile_iter != 0: clc_full_phase_epi = clc_full_phase_epi ^ 1 @@ -1296,7 +1381,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1350,7 +1435,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1402,7 +1487,7 @@ def _host( out_stride_l_0 * cd_dtype.width // 128, ], box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], - swizzle=(_tma.TensorMapSwizzle.s64b if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), + swizzle=(epi_tma_swizzle if cutlass.const_expr(use_tma_store_epi) else _tma.TensorMapSwizzle.none), ) tma_c_desc_list = [tma_c_desc] # @@TMA_STORE_ONLY:END@@ diff --git a/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_1ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_1ctamma.py index 33ee8800f..fa7409421 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_1ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_1ctamma.py @@ -40,6 +40,7 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( copy_tensormap_to_workspace as _copy_tensormap_to_workspace, + epi_subtile_spans as _epi_subtile_spans, fence_tensormap_acquire as _fence_tensormap_acquire, fence_tensormap_release as _fence_tensormap_release, moe_swizzle_tile as _moe_swizzle_tile, @@ -62,6 +63,9 @@ # it to the per-CTA GMEM workspace the TMA reads. # @@INJECT_TILE_CONSTANTS@@ +if use_acc_overlap and any(_w != epi_n for _, _w in _epi_subtile_spans(epi_cols_per_mma_m, epi_n)): + raise NotImplementedError(f"{__name__}: acc overlap reverses subtiles by index, which needs a uniform drain width") + # Per-CTA scheduler ring (replaces CLC): 2 stages, 8 int32 slot words. SCHED_STAGES = 2 @@ -177,11 +181,15 @@ def _kernel( a_part_arrive = cutlass.Int16(a_pattern) << m_rank b_part_arrive = cutlass.Int16(b_pattern) << (n_rank * cluster_m) - ab_empty_arrive_mask = a_part_arrive | b_part_arrive + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = a_part_arrive | b_part_arrive _smem_sys_reserved = cutlass.Array(cutlass.Int8, 1024, space=cutlass.AddressSpace.smem, alignment=1) ab_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) + sf_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) ab_empty_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) acc_empty_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) acc_full_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) @@ -246,12 +254,17 @@ def _kernel( for _ in range(num_b_operands) ] - ab_empty_count = cluster_m + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size + else: + ab_empty_count = cluster_m + cluster_n - 1 sched_empty_count = 1 + 1 + num_epilogue_warps if warp_idx == 0: for i in range(ab_stages): if elect_one: nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) + if elect_one: + nvvm.mbarrier_init(sf_full_mbar_ptr.subview(i), 1) if elect_one: nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) for i in range(acc_stages): @@ -278,9 +291,10 @@ def _kernel( sA_bytes = sA_elems * (ab_dtype.width // 8) 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) + ab_only_copy_bytes = num_a_operands * sA_bytes + num_b_operands * sB_bytes + sf_only_copy_bytes = num_a_operands * sfa_smem_bytes + num_b_operands * sfb_smem_bytes - cols_per_acc_stage = cta_tile_mnk[1] + epi_rows_per_mma_m = cta_tile_mnk[0] // num_mma_m tmem_alloc_bar_count = (num_epilogue_warps + 1) * 32 # @@INJECT_TAP_PTRS@@ @@ -511,8 +525,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) sched_stage = cutlass.Int32(0) @@ -594,33 +607,60 @@ def _kernel( coord_k = k_tile_idx * cta_tile_mnk[2] coord_sf_k = k_tile_idx * sf_tma_box_k if elect_one: - nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) + nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), ab_only_copy_bytes) + if elect_one: + nvvm.mbarrier_arrive_expect_tx(sf_full_mbar_ptr.subview(stage), sf_only_copy_bytes) a_issue = (not multicast_a) or (n_rank == 0) + if cutlass.const_expr(a_mcast_slices > 1): + a_data_issue = True + _a_off = n_rank * (cta_tile_mnk[0] // a_mcast_slices) + else: + a_data_issue = a_issue + _a_off = 0 b_issue = (not multicast_b) or (m_rank == 0) + if cutlass.const_expr(b_mcast_slices > 1): + b_data_issue = True + _b_off = m_rank * (cta_tile_mnk[1] // b_mcast_slices) + else: + b_data_issue = b_issue + _b_off = 0 if a_issue: for _ai in cutlass.range_constexpr(num_a_operands): if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - smem_a_list[_ai].subview(sA_elems * stage), - a_desc_tma_ptr_list[_ai], - (coord_k, coord_m_group, cutlass.Int32(0)), - ab_full_mbar_ptr.subview(stage), + smem_sfa_list[_ai].subview(sfa_smem_bytes * stage), + tma_sfa_descs[_ai].get_ptr(), + (0, coord_sf_k, sfa_m_block, cutlass.Int32(0)), + sf_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_1, ) + if b_issue: + for _bj in cutlass.range_constexpr(num_b_operands): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + smem_sfb_list[_bj].subview(sfb_smem_bytes * stage), + tma_sfb_descs[_bj].get_ptr(), + (0, coord_sf_k, sfb_n_block, coord_expert), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_1, + ) + if a_data_issue: for _ai in cutlass.range_constexpr(num_a_operands): if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - smem_sfa_list[_ai].subview(sfa_smem_bytes * stage), - tma_sfa_descs[_ai].get_ptr(), - (0, coord_sf_k, sfa_m_block, cutlass.Int32(0)), + smem_a_list[_ai].subview(sA_elems * stage + _a_off * ab_packed_per_row), + a_desc_tma_ptr_list[_ai], + (coord_k, coord_m_group + _a_off, cutlass.Int32(0)), ab_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_1, ) - if b_issue: + if b_data_issue: for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) if cutlass.const_expr(b_is_n_major): @@ -642,39 +682,28 @@ def _kernel( else: if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - sB_stage, + sB_stage.subview(_b_off * ab_packed_per_row), tma_b_descs[_bj].get_ptr(), - (coord_k, coord_n_per_cta, coord_expert), + (coord_k, coord_n_per_cta + _b_off, coord_expert), ab_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_b, group=nvvm.CTAGroup.CTA_1, ) - for _bj in cutlass.range_constexpr(num_b_operands): - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - smem_sfb_list[_bj].subview(sfb_smem_bytes * stage), - tma_sfb_descs[_bj].get_ptr(), - (0, coord_sf_k, sfb_n_block, coord_expert), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_1, - ) ab_iter += 1 tail_stage = ab_iter % ab_stages tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) @@ -748,7 +777,7 @@ def _kernel( sfa_dst_ptrs = [ [ [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block + a * registers_per_atom, cutlass.Float32) for a in range(word_atoms)] - for m in range(num_blocks_m) + for m in range(num_mma_m) ] for i in range(num_a_operands) ] @@ -810,13 +839,6 @@ def _kernel( if stage == 0 and ab_iter != 0: ab_full_phase_bit = ab_full_phase_bit ^ 1 - while not nvvm.mbarrier_try_wait_parity( - ab_full_mbar_ptr.subview(stage), - ab_full_phase_bit, - time_limit=10_000_000, - ): - pass - desc_a_bases = [ cutlass.experimental.primitives.Tcgen05SmemDesc.build( start_address=smem_a_list[i].subview(sA_elems * stage), @@ -856,58 +878,71 @@ def _kernel( # One SF word per group of MMAs, refreshed right before they # read it. A word spans word_atoms consecutive K-atoms in SMEM. - for atom_r in cutlass.range_constexpr(num_sf_atoms): - for _ai in cutlass.range_constexpr(num_a_operands): - for _m in cutlass.range_constexpr(num_blocks_m): - for _a in cutlass.range_constexpr(word_atoms): - if elect_one: - nvvm.tcgen05_cp( - s2t_shape, - sfa_dst_ptrs[_ai][_m][_a], - desc_sfa_bases[_ai] + (sf_atom_desc_stride * (atom_r * word_atoms + _a) + sf_block_desc_stride * _m), - group=nvvm.CTAGroup.CTA_1, - multicast=s2t_multicast, - ) + while not nvvm.mbarrier_try_wait_parity( + sf_full_mbar_ptr.subview(stage), + ab_full_phase_bit, + time_limit=10_000_000, + ): + pass + + for sf_word in cutlass.range_constexpr(num_sf_atoms): for _bj in cutlass.range_constexpr(num_b_operands): - for _m in cutlass.range_constexpr(num_blocks_n): + for block_n in cutlass.range_constexpr(num_blocks_n): for _a in cutlass.range_constexpr(word_atoms): if elect_one: nvvm.tcgen05_cp( s2t_shape, - sfb_dst_ptrs[_bj][_m][_a], - desc_sfb_bases[_bj] + (sf_atom_desc_stride * (atom_r * word_atoms + _a) + sf_block_desc_stride * _m), + sfb_dst_ptrs[_bj][block_n][_a], + desc_sfb_bases[_bj] + (sf_atom_desc_stride * (sf_word * word_atoms + _a) + sf_block_desc_stride * block_n), group=nvvm.CTAGroup.CTA_1, multicast=s2t_multicast, ) - for j in cutlass.range_constexpr(sf_insts_per_atom): - k_block_idx = atom_r * sf_insts_per_atom + j - idesc_k = idesc_by_j[j] - for g in cutlass.range_constexpr(num_gemms): - _ai = gemm_a_idx[g] - _bj = gemm_b_idx[g] - desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) - desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - for mi in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(sf_word == 0): + while not nvvm.mbarrier_try_wait_parity( + ab_full_mbar_ptr.subview(stage), + ab_full_phase_bit, + time_limit=10_000_000, + ): + pass + for mma_k_in_word in cutlass.range_constexpr(sf_insts_per_atom): + mma_k = sf_word * sf_insts_per_atom + mma_k_in_word + idesc_k = idesc_by_j[mma_k_in_word] + for gemm_i in cutlass.range_constexpr(num_gemms): + _ai = gemm_a_idx[gemm_i] + _bj = gemm_b_idx[gemm_i] + desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * mma_k) + desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * mma_k) + for mma_m in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(mma_k_in_word == 0 and _ai not in gemm_a_idx[:gemm_i]): + for _a in cutlass.range_constexpr(word_atoms): + if elect_one: + nvvm.tcgen05_cp( + s2t_shape, + sfa_dst_ptrs[_ai][mma_m][_a], + desc_sfa_bases[_ai] + (sf_atom_desc_stride * (sf_word * word_atoms + _a) + sf_block_desc_stride * mma_m), + group=nvvm.CTAGroup.CTA_1, + multicast=s2t_multicast, + ) # The M sub-block offset is a whole SMEM swizzle atom, so # the descriptor's swizzle phase is preserved. B and its SF # are shared; A's SF word block follows the M block. - desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) + desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mma_m) if elect_one: _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_1, - acc_tmem_ptrs[g][mi], + acc_tmem_ptrs[gemm_i][mma_m], desc_a, desc_b, idesc_k, enable_input_d=scale_d, - scale_a=sfa_dst_ptrs[_ai][mi][0], + scale_a=sfa_dst_ptrs[_ai][mma_m][0], scale_b=sfb_scale_ptrs[_bj], scale_vec_size=scale_vec_size, - b_collector_op=_b_collector_op(mi), + b_collector_op=_b_collector_op(mma_m), ) # Every accumulator sees scale_d=False on exactly the first - # k_block of the tile, so the flip stays outside mi. + # k_block of the tile, so the flip stays outside mma_m. scale_d = cutlass.Boolean(True) if elect_one: @@ -926,28 +961,26 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_1) if tile_iter != 0: tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_empty_mbar_ptr.subview(tail_stage), - tail_phase, - time_limit=10_000_000, - ): - pass - if cutlass.const_expr(use_acc_overlap): - while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): - pass + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity( + acc_empty_mbar_ptr.subview(tail_stage), + tail_phase, + time_limit=10_000_000, + ): + pass + if cutlass.const_expr(use_acc_overlap): + while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): + pass nvvm.bar_warp_sync(0xFFFFFFFF) alloc_ptr = cutlass.inttoptr(tmem_raw_addr, 6, cutlass.Int32) @@ -974,21 +1007,14 @@ def _kernel( sched_stage = cutlass.Int32(0) sched_full_phase = cutlass.Int32(0) - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - row_id_with_warp_offset = base_row_id - else: - row_id_with_warp_offset = base_row_id + warp_idx * 32 - - # One M block's accumulator columns are contiguous. - subtile_cnt = cute.ceil_div(epi_cols_per_mma_m, 32) - t2r_inst_repx = epi_tile_mn[1] - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - shape = nvvm.Tcgen05LdStShape.SHAPE_16X32BX2 - ld_half_off = 0 - else: - shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B - ld_half_off = None + # @@EPILOGUE_SETUP:BEGIN@@ + row_id_with_warp_offset = base_row_id + warp_idx * 32 + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) + subtile_cnt = len(epi_spans) + shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ while is_valid != 0: while not nvvm.mbarrier_try_wait_parity( @@ -1013,17 +1039,14 @@ def _kernel( if is_valid != 0: coord_m_tile = group_begin + 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] + # @@EPILOGUE_DRAIN:BEGIN@@ + coord_n_c = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: acc_full_phase_bit = acc_full_phase_bit ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_full_mbar_ptr.subview(acc_stage), - acc_full_phase_bit, - time_limit=10_000_000, - ): + while not nvvm.mbarrier_try_wait_parity(acc_full_mbar_ptr.subview(acc_stage), acc_full_phase_bit, time_limit=10_000_000): pass if cutlass.const_expr(use_acc_overlap): @@ -1032,48 +1055,53 @@ def _kernel( else: acc_buf_parity = cutlass.Int32(0) acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # One pass per MMA-M block over its own column region. + for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * mma_inst_shape_mnk[0] - mi_col_base = acc_base_col + mi * epi_cols_per_mma_m + if cutlass.const_expr(use_acc_overlap and num_mma_m > 1): + _mi = mi + (1 - acc_buf_parity) * (num_mma_m - 1 - 2 * mi) + else: + _mi = mi + coord_m = coord_m_tile + _mi * epi_rows_per_mma_m + mi_col_base = acc_base_col + _mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * acc_gemm_stride) for g in range(num_gemms)] - if cutlass.const_expr(mma_inst_shape_mnk[0] == 64): - row = coord_m + warp_idx * 16 + lane - row_active = lane < 16 - else: - row = coord_m + tidx - row_active = True + row = coord_m + tidx + row_active = True # @@INJECT_AUX_VIEWS@@ for subtile_idx in cutlass.range_constexpr(subtile_cnt): if cutlass.const_expr(use_acc_overlap): _sub = subtile_idx + (1 - acc_buf_parity) * (subtile_cnt - 1 - 2 * subtile_idx) - subtile_col_offset = _sub * 32 + subtile_col_offset = _sub * epi_n + subtile_w = epi_n else: - subtile_col_offset = subtile_idx * 32 + subtile_col_offset, subtile_w = epi_spans[subtile_idx] c_rmem_vecs = [] for g in cutlass.range_constexpr(num_gemms): - tmem = cutlass.inttoptr( - tmem_col_addr_gemms[g] + subtile_col_offset, - 6, - cutlass.Float32, - ) - c_rmem_vecs.append(nvvm.tcgen05_ld(shape, tmem, num=t2r_inst_repx, offset=ld_half_off)) + subtile_tmem_addr = tmem_col_addr_gemms[g] + subtile_col_offset + tmem = cutlass.inttoptr(subtile_tmem_addr, 6, mma_c_dtype) + _cv = nvvm.tcgen05_ld(shape, tmem, num=subtile_w) + c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - if cutlass.const_expr(use_acc_overlap and mi == num_mma_m - 1 and subtile_idx == acc_overlap_subtiles - 1): + if cutlass.const_expr(not use_acc_overlap): + if cutlass.const_expr(mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1): + nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) + nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) + if elect_one: + nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) + + if use_acc_overlap and mi * subtile_cnt + subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) - col = coord_n + subtile_col_offset + col = coord_n_c + subtile_col_offset - # @@STG_ONLY:BEGIN@@ if row_active and row < group_end: - for j in cutlass.range_constexpr(t2r_inst_repx // vsize): + for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize if col_j + vsize <= N: vec_f32 = c_rmem_vec[j * vsize : (j + 1) * vsize] @@ -1081,13 +1109,9 @@ def _kernel( # @@INJECT_STG_VEC_BINDINGS@@ # @@INJECT_EPILOGUE@@ - # @@STG_ONLY:END@@ - if cutlass.const_expr(not use_acc_overlap): - nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) - nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if elect_one: - nvvm.mbarrier_arrive(acc_empty_mbar_ptr.subview(acc_stage)) + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + # @@EPILOGUE_DRAIN:END@@ tile_iter += 1 if cutlass.const_expr(use_acc_overlap): @@ -1151,7 +1175,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1184,7 +1208,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) diff --git a/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_2ctamma.py b/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_2ctamma.py index b91e571cf..4a33805f3 100644 --- a/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_2ctamma.py +++ b/python/cudnn/gemm/frost/kernel_templates/sm107_moe_grouped_block_scale_matmul_fwd_2ctamma.py @@ -45,6 +45,7 @@ import cutlass.experimental.primitives as nvvm from cudnn.gemm.frost.kernel_templates._tile_helpers import ( copy_tensormap_to_workspace as _copy_tensormap_to_workspace, + epi_subtile_spans as _epi_subtile_spans, fence_tensormap_acquire as _fence_tensormap_acquire, fence_tensormap_release as _fence_tensormap_release, moe_swizzle_tile as _moe_swizzle_tile, @@ -64,6 +65,9 @@ # @@INJECT_TILE_CONSTANTS@@ +if use_acc_overlap and any(_w != epi_n for _, _w in _epi_subtile_spans(epi_cols_per_mma_m, epi_n)): + raise NotImplementedError(f"{__name__}: acc overlap reverses subtiles by index, which needs a uniform drain width") + SCHED_STAGES = 2 SCHED_SLOT_WORDS = 8 @@ -186,6 +190,7 @@ def _kernel( _smem_sys_reserved = cutlass.Array(cutlass.Int8, 1024, space=cutlass.AddressSpace.smem, alignment=1) ab_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) + sf_full_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) ab_empty_mbar_ptr = cutlass.Array(cutlass.Int64, ab_stages, space=cutlass.AddressSpace.smem) acc_empty_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) acc_full_mbar_ptr = cutlass.Array(cutlass.Int64, acc_stages, space=cutlass.AddressSpace.smem) @@ -251,7 +256,10 @@ def _kernel( acc_empty_count = num_epilogue_warps * 2 cta_group = 2 - ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_count = cluster_size // cta_group + else: + ab_empty_count = (cluster_m // cta_group) + cluster_n - 1 sched_empty_count = 1 + 1 + num_epilogue_warps if warp_idx == 0: if cutlass.const_expr(use_acc_overlap): @@ -263,6 +271,8 @@ def _kernel( for i in range(ab_stages): if elect_one: nvvm.mbarrier_init(ab_full_mbar_ptr.subview(i), 1) + if elect_one: + nvvm.mbarrier_init(sf_full_mbar_ptr.subview(i), 1) if elect_one: nvvm.mbarrier_init(ab_empty_mbar_ptr.subview(i), ab_empty_count) for i in range(acc_stages): @@ -280,17 +290,13 @@ def _kernel( sA_bytes = sA_elems * (ab_dtype.width // 8) 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 + ab_only_copy_bytes = (num_a_operands * sA_bytes + num_b_operands * sB_bytes) * 2 + sf_only_copy_bytes = (num_a_operands * sfa_smem_bytes + num_b_operands * sfb_smem_bytes) * 2 pair_n_size = cgrp_tile_mnk[1] // cluster_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 - if cutlass.const_expr(epi_rows_per_mma_m == 64): - # cluster-MMA m=128: the pair also splits N, so each CTA drains N/2. - cols_per_acc_stage = pair_n_size // 2 - else: - cols_per_acc_stage = pair_n_size tmem_alloc_bar_count = (num_epilogue_warps + 1) * 32 nvvm.barrier_cluster_wait() @@ -524,8 +530,7 @@ def _kernel( if warp_idx == tma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("wait") + nvvm.griddepcontrol("wait") ab_empty_phase_bit = cutlass.Int32(1) ab_iter = cutlass.Int32(0) sched_stage = cutlass.Int32(0) @@ -611,33 +616,61 @@ def _kernel( if is_pair_leader: if elect_one: - nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), num_tma_copy_bytes) + nvvm.mbarrier_arrive_expect_tx(ab_full_mbar_ptr.subview(stage), ab_only_copy_bytes) + if elect_one: + nvvm.mbarrier_arrive_expect_tx(sf_full_mbar_ptr.subview(stage), sf_only_copy_bytes) a_issue = (not multicast_a) or (n_rank == 0) + if cutlass.const_expr(a_mcast_slices > 1): + a_data_issue = True + _a_off = n_rank * (cta_tile_mnk[0] // a_mcast_slices) + else: + a_data_issue = a_issue + _a_off = 0 b_issue = (not multicast_b) or (pair_m_idx == 0) + if cutlass.const_expr(b_mcast_slices > 1): + b_data_issue = True + _b_off = pair_m_idx * (cta_tile_mnk[1] // b_mcast_slices) + else: + b_data_issue = b_issue + _b_off = 0 if a_issue: for _ai in cutlass.range_constexpr(num_a_operands): if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - smem_a_list[_ai].subview(sA_elems * stage), - a_desc_tma_ptr_list[_ai], - (coord_k, coord_m_group, cutlass.Int32(0)), - ab_full_mbar_ptr.subview(stage), + smem_sfa_list[_ai].subview(sfa_smem_bytes * stage), + tma_sfa_descs[_ai].get_ptr(), + (0, coord_sf_k, sfa_m_block, cutlass.Int32(0)), + sf_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_2, ) + if b_issue: + for _bj in cutlass.range_constexpr(num_b_operands): + if elect_one: + nvvm.cp_async_bulk_tensor_shared_cluster_global( + smem_sfb_list[_bj].subview(sfb_smem_bytes * stage), + tma_sfb_descs[_bj].get_ptr(), + (0, coord_sf_k, sfb_n_block, coord_expert), + sf_full_mbar_ptr.subview(stage), + [], + multicast_mask=tma_mcast_mask_b, + group=nvvm.CTAGroup.CTA_2, + ) + + if a_data_issue: for _ai in cutlass.range_constexpr(num_a_operands): if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - smem_sfa_list[_ai].subview(sfa_smem_bytes * stage), - tma_sfa_descs[_ai].get_ptr(), - (0, coord_sf_k, sfa_m_block, cutlass.Int32(0)), + smem_a_list[_ai].subview(sA_elems * stage + _a_off * ab_packed_per_row), + a_desc_tma_ptr_list[_ai], + (coord_k, coord_m_group + _a_off, cutlass.Int32(0)), ab_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_a, group=nvvm.CTAGroup.CTA_2, ) - if b_issue: + if b_data_issue: for _bj in cutlass.range_constexpr(num_b_operands): sB_stage = smem_b_list[_bj].subview(sB_elems * stage) if cutlass.const_expr(b_is_n_major): @@ -659,40 +692,28 @@ def _kernel( else: if elect_one: nvvm.cp_async_bulk_tensor_shared_cluster_global( - sB_stage, + sB_stage.subview(_b_off * ab_packed_per_row), tma_b_descs[_bj].get_ptr(), - (coord_k, coord_n_per_cta, coord_expert), + (coord_k, coord_n_per_cta + _b_off, coord_expert), ab_full_mbar_ptr.subview(stage), [], multicast_mask=tma_mcast_mask_b, group=nvvm.CTAGroup.CTA_2, ) - for _bj in cutlass.range_constexpr(num_b_operands): - if elect_one: - nvvm.cp_async_bulk_tensor_shared_cluster_global( - smem_sfb_list[_bj].subview(sfb_smem_bytes * stage), - tma_sfb_descs[_bj].get_ptr(), - (0, coord_sf_k, sfb_n_block, coord_expert), - ab_full_mbar_ptr.subview(stage), - [], - multicast_mask=tma_mcast_mask_b, - group=nvvm.CTAGroup.CTA_2, - ) - ab_iter += 1 tail_stage = ab_iter % ab_stages tail_phase = ab_empty_phase_bit if tail_stage == 0 and ab_iter != 0: tail_phase = tail_phase ^ 1 - for _ in range(ab_stages - 1): - tail_stage = tail_stage + 1 - if tail_stage == ab_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - if elect_one: - while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): - pass + if cutlass.const_expr(cluster_shape_mnk[0] * cluster_shape_mnk[1] > 1): + for _ in range(ab_stages): + while not nvvm.mbarrier_try_wait_parity(ab_empty_mbar_ptr.subview(tail_stage), tail_phase, time_limit=10_000_000): + pass + tail_stage = tail_stage + 1 + if tail_stage == ab_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 pair_mask = cutlass.Int16(3) << pair_leader_rank a_arrive_pattern = 0 @@ -704,7 +725,10 @@ def _kernel( a_part = a_arrive_pattern << m_rank a_part = a_part | (a_part << 1) b_part = b_arrive_pattern << (n_rank * cluster_m) - ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) + if cutlass.const_expr(ab_empty_full_mask): + ab_empty_arrive_mask = cutlass.Int16((1 << cluster_size) - 1) + else: + ab_empty_arrive_mask = cutlass.Int16(a_part | b_part) if warp_idx == mma_warp_id: nvvm.setmaxregister(prod_reg_count, nvvm.SetMaxRegisterAction.DECREASE) _tcgen05_alloc( @@ -777,7 +801,7 @@ def _kernel( sfa_dst_ptrs = [ [ [nvvm.make_tmem_ptr(sfa_tmem_bases[i] + m * registers_per_block + a * registers_per_atom, cutlass.Float32) for a in range(word_atoms)] - for m in range(num_blocks_m) + for m in range(num_mma_m) ] for i in range(num_a_operands) ] @@ -839,13 +863,6 @@ def _kernel( if stage == 0 and ab_iter != 0: ab_full_phase_bit = ab_full_phase_bit ^ 1 - while not nvvm.mbarrier_try_wait_parity( - ab_full_mbar_ptr.subview(stage), - ab_full_phase_bit, - time_limit=10_000_000, - ): - pass - desc_a_bases = [ cutlass.experimental.primitives.Tcgen05SmemDesc.build( start_address=smem_a_list[i].subview(sA_elems * stage), @@ -885,58 +902,72 @@ def _kernel( # One SF word per group of MMAs, refreshed right before they # read it. A word spans word_atoms consecutive K-atoms in SMEM. - for atom_r in cutlass.range(num_sf_atoms, unroll_full=True): - for _ai in cutlass.range_constexpr(num_a_operands): - for _m in cutlass.range_constexpr(num_blocks_m): - for _a in cutlass.range_constexpr(word_atoms): - if elect_one: - nvvm.tcgen05_cp( - s2t_shape, - sfa_dst_ptrs[_ai][_m][_a], - desc_sfa_bases[_ai] + (sf_atom_desc_stride * (atom_r * word_atoms + _a) + sf_block_desc_stride * _m), - group=nvvm.CTAGroup.CTA_2, - multicast=s2t_multicast, - ) + while not nvvm.mbarrier_try_wait_parity( + sf_full_mbar_ptr.subview(stage), + ab_full_phase_bit, + time_limit=10_000_000, + ): + pass + + for sf_word in cutlass.range_constexpr(num_sf_atoms): for _bj in cutlass.range_constexpr(num_b_operands): - for _m in cutlass.range_constexpr(num_blocks_n): + for block_n in cutlass.range_constexpr(num_blocks_n): for _a in cutlass.range_constexpr(word_atoms): if elect_one: nvvm.tcgen05_cp( s2t_shape, - sfb_dst_ptrs[_bj][_m][_a], - desc_sfb_bases[_bj] + (sf_atom_desc_stride * (atom_r * word_atoms + _a) + sf_block_desc_stride * _m), + sfb_dst_ptrs[_bj][block_n][_a], + desc_sfb_bases[_bj] + (sf_atom_desc_stride * (sf_word * word_atoms + _a) + sf_block_desc_stride * block_n), group=nvvm.CTAGroup.CTA_2, multicast=s2t_multicast, ) - for j in cutlass.range_constexpr(sf_insts_per_atom): - k_block_idx = atom_r * sf_insts_per_atom + j - idesc_k = idesc_by_j[j] - for g in cutlass.range_constexpr(num_gemms): - _ai = gemm_a_idx[g] - _bj = gemm_b_idx[g] - desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * k_block_idx) - desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * k_block_idx) - for mi in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(sf_word == 0): + while not nvvm.mbarrier_try_wait_parity( + ab_full_mbar_ptr.subview(stage), + ab_full_phase_bit, + time_limit=10_000_000, + ): + pass + for mma_k_in_word in cutlass.range_constexpr(sf_insts_per_atom): + mma_k = sf_word * sf_insts_per_atom + mma_k_in_word + idesc_k = idesc_by_j[mma_k_in_word] + for gemm_i in cutlass.range_constexpr(num_gemms): + _ai = gemm_a_idx[gemm_i] + _bj = gemm_b_idx[gemm_i] + desc_a_k = desc_a_bases[_ai].advance_start_address(a_smem_k_step_bytes * mma_k) + desc_b = desc_b_bases[_bj].advance_start_address(b_smem_k_step_bytes * mma_k) + for mma_m in cutlass.range_constexpr(num_mma_m): + if cutlass.const_expr(mma_k_in_word == 0 and _ai not in gemm_a_idx[:gemm_i]): + for _a in cutlass.range_constexpr(word_atoms): + if elect_one: + nvvm.tcgen05_cp( + s2t_shape, + sfa_dst_ptrs[_ai][mma_m][_a], + desc_sfa_bases[_ai] + + (sf_atom_desc_stride * (sf_word * word_atoms + _a) + sf_block_desc_stride * mma_m), + group=nvvm.CTAGroup.CTA_2, + multicast=s2t_multicast, + ) # The M sub-block offset is a whole SMEM swizzle atom, so # the descriptor's swizzle phase is preserved. B and its SF # are shared; A's SF word block follows the M block. - desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mi) + desc_a = desc_a_k.advance_start_address(a_smem_m_step_bytes * mma_m) if elect_one: _tcgen05_mma_block_scale( mma_block_scale_kind, nvvm.CTAGroup.CTA_2, - acc_tmem_ptrs[g][mi], + acc_tmem_ptrs[gemm_i][mma_m], desc_a, desc_b, idesc_k, enable_input_d=scale_d, - scale_a=sfa_dst_ptrs[_ai][mi][0], + scale_a=sfa_dst_ptrs[_ai][mma_m][0], scale_b=sfb_scale_ptrs[_bj], scale_vec_size=scale_vec_size, - b_collector_op=_b_collector_op(mi), + b_collector_op=_b_collector_op(mma_m), ) # Every accumulator sees scale_d=False on exactly the first - # k_block of the tile, so the flip stays outside mi. + # k_block of the tile, so the flip stays outside mma_m. scale_d = cutlass.Boolean(True) if elect_one: @@ -956,26 +987,22 @@ def _kernel( tile_iter += 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") if tile_iter != 0: tail_stage = acc_stage tail_phase = acc_empty_phase_bit - if elect_one: - for _ in range(acc_stages): - tail_stage = tail_stage + 1 - if tail_stage == acc_stages: - tail_stage = cutlass.Int32(0) - tail_phase = tail_phase ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_empty_mbar_ptr.subview(tail_stage), - tail_phase, - time_limit=10_000_000, - ): - pass - nvvm.bar_warp_sync(0xFFFFFFFF) - + for _ in range(acc_stages): + tail_stage = tail_stage + 1 + if tail_stage == acc_stages: + tail_stage = cutlass.Int32(0) + tail_phase = tail_phase ^ 1 + while not nvvm.mbarrier_try_wait_parity( + acc_empty_mbar_ptr.subview(tail_stage), + tail_phase, + time_limit=10_000_000, + ): + pass nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) while not nvvm.mbarrier_try_wait_parity(tmem_dealloc_mbar_ptr, 0, time_limit=10_000_000): @@ -1009,8 +1036,7 @@ def _kernel( sched_full_phase = sched_full_phase ^ 1 if cutlass.const_expr(USE_PDL): - if elect_one: - nvvm.griddepcontrol("launch_dependents") + nvvm.griddepcontrol("launch_dependents") nvvm.tcgen05_relinquish_alloc_permit(group=nvvm.CTAGroup.CTA_2) peer_mbar = nvvm.mapa(tmem_dealloc_mbar_ptr, peer_cta_rank) @@ -1040,15 +1066,14 @@ def _kernel( sched_stage = cutlass.Int32(0) sched_full_phase = cutlass.Int32(0) + # @@EPILOGUE_SETUP:BEGIN@@ row_id_with_warp_offset = base_row_id + warp_idx * 32 - if cutlass.const_expr(cols_per_acc_stage >= 32): - t2r_inst_repx = 32 - subtile_cnt = cols_per_acc_stage // 32 - else: - t2r_inst_repx = cols_per_acc_stage - subtile_cnt = 1 + + epi_spans = _epi_subtile_spans(epi_cols_per_mma_m, epi_n) + subtile_cnt = len(epi_spans) shape = nvvm.Tcgen05LdStShape.SHAPE_32X32B lane = tidx % 32 + # @@EPILOGUE_SETUP:END@@ while is_valid != 0: while not nvvm.mbarrier_try_wait_parity( @@ -1073,19 +1098,14 @@ def _kernel( if is_valid != 0: coord_m_tile = group_begin + tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] + # @@EPILOGUE_DRAIN:BEGIN@@ coord_n_c = tile_n * cgrp_tile_mnk[1] + 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 acc_stage = tile_iter % acc_stages if acc_stage == 0 and tile_iter != 0: acc_full_phase_bit = acc_full_phase_bit ^ 1 - while not nvvm.mbarrier_try_wait_parity( - acc_full_mbar_ptr.subview(acc_stage), - acc_full_phase_bit, - time_limit=10_000_000, - ): + while not nvvm.mbarrier_try_wait_parity(acc_full_mbar_ptr.subview(acc_stage), acc_full_phase_bit, time_limit=10_000_000): pass if cutlass.const_expr(use_acc_overlap): @@ -1094,52 +1114,61 @@ def _kernel( else: acc_buf_parity = cutlass.Int32(0) acc_base_col = base_col_id_root + acc_stage * acc_region_cols - # The 2-CTA epilogue drains its own half of the instruction's M, - # epi_rows_per_mma_m rows at a time, so a CTA tile of num_mma_m blocks - # drains in num_mma_m passes over its own column region. + for mi in cutlass.range_constexpr(num_mma_m): - coord_m = coord_m_tile + mi * epi_rows_per_mma_m - mi_col_base = acc_base_col + mi * epi_cols_per_mma_m + if cutlass.const_expr(use_acc_overlap and num_mma_m > 1): + _mi = mi + (1 - acc_buf_parity) * (num_mma_m - 1 - 2 * mi) + else: + _mi = mi + coord_m = coord_m_tile + _mi * epi_rows_per_mma_m + mi_col_base = acc_base_col + _mi * epi_cols_per_mma_m tmem_col_addr_gemms = [(row_id_with_warp_offset << 16) | (mi_col_base + g * acc_gemm_stride) for g in range(num_gemms)] - if cutlass.const_expr(epi_rows_per_mma_m == 64): - row = coord_m + (warp_idx % 2) * 32 + lane - row_active = True - else: - row = coord_m + tidx - row_active = True + row = coord_m + tidx + row_active = True # @@INJECT_AUX_VIEWS@@ - for subtile_idx in cutlass.range(subtile_cnt, unroll_full=True): + for subtile_idx in cutlass.range_constexpr(subtile_cnt): if cutlass.const_expr(use_acc_overlap): _sub = subtile_idx + (1 - acc_buf_parity) * (subtile_cnt - 1 - 2 * subtile_idx) - subtile_col_offset = _sub * t2r_inst_repx + subtile_col_offset = _sub * epi_n + subtile_w = epi_n else: - subtile_col_offset = subtile_idx * t2r_inst_repx + subtile_col_offset, subtile_w = epi_spans[subtile_idx] c_rmem_vecs = [] for g in cutlass.range_constexpr(num_gemms): - tmem = cutlass.inttoptr( - tmem_col_addr_gemms[g] + subtile_col_offset, - 6, - cutlass.Float32, - ) - c_rmem_vecs.append(nvvm.tcgen05_ld(shape, tmem, num=t2r_inst_repx)) + subtile_tmem_addr = tmem_col_addr_gemms[g] + subtile_col_offset + tmem = cutlass.inttoptr(subtile_tmem_addr, 6, mma_c_dtype) + _cv = nvvm.tcgen05_ld(shape, tmem, num=subtile_w) + c_rmem_vecs.append(_cv) c_rmem_vec = c_rmem_vecs[0] - if use_acc_overlap and mi == num_mma_m - 1 and subtile_idx == acc_overlap_subtiles - 1: + if cutlass.const_expr(not use_acc_overlap): + if cutlass.const_expr(mi == num_mma_m - 1 and subtile_idx == subtile_cnt - 1): + nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) + nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) + if elect_one: + nvvm.mbarrier_arrive( + nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), + scope=nvvm.MemScope.CLUSTER, + relaxed=True, + ) + + if use_acc_overlap and mi * subtile_cnt + subtile_idx == acc_overlap_subtiles - 1: nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) if elect_one: - mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) - nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) + nvvm.mbarrier_arrive( + nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank), + scope=nvvm.MemScope.CLUSTER, + relaxed=True, + ) col = coord_n_c + subtile_col_offset - # @@STG_ONLY:BEGIN@@ - if row_active and row < group_end: - for j in cutlass.range_constexpr(t2r_inst_repx // vsize): + for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize if col_j + vsize <= N: vec_f32 = c_rmem_vec[j * vsize : (j + 1) * vsize] @@ -1148,14 +1177,8 @@ def _kernel( # @@INJECT_EPILOGUE@@ - # @@STG_ONLY:END@@ - - if cutlass.const_expr(not use_acc_overlap): - nvvm.tcgen05_wait(kind=nvvm.Tcgen05Wait.LOAD) - nvvm.tcgen05_fence(nvvm.Tcgen05Fence.BEFORE_THREAD_SYNC) - if elect_one: - mbar_pair_ptr = nvvm.mapa(acc_empty_mbar_ptr.subview(acc_stage), pair_leader_rank) - nvvm.mbarrier_arrive(mbar_pair_ptr, scope=nvvm.MemScope.CLUSTER, relaxed=True) + # The M-major TMA path loads its accumulator inside the store loop, so its release cannot move up. + # @@EPILOGUE_DRAIN:END@@ tile_iter += 1 if cutlass.const_expr(use_acc_overlap): @@ -1221,7 +1244,7 @@ def _host( a_stride_m * ab_dtype.width // 128, a_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[0], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[0] // a_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) @@ -1254,7 +1277,7 @@ def _host( b_stride_n * ab_dtype.width // 128, b_stride_l * ab_dtype.width // 128, ], - box_dims=[cta_tile_mnk[2], cta_tile_mnk[1], 1], + box_dims=[cta_tile_mnk[2], cta_tile_mnk[1] // b_mcast_slices, 1], swizzle=ab_tma_swizzle, tma_format=ab_tma_format, ) diff --git a/python/cudnn/gemm/frost/tile_config.py b/python/cudnn/gemm/frost/tile_config.py index b67e5dc6e..9de6d6157 100644 --- a/python/cudnn/gemm/frost/tile_config.py +++ b/python/cudnn/gemm/frost/tile_config.py @@ -40,13 +40,17 @@ def _sm_smem_budget_bytes(device=None) -> int: return _sm_smem_budget_bytes_of(resolve_device(device)) -# Per-CTA SMEM held back off the top when sizing the ab/acc pipeline, keyed by -# the kernel template's pipeline: the CLC ring, smem barriers, the TMEM base-address +# Per-CTA SMEM held back off the top when sizing the ab pipeline: the scheduler ring, +# every smem barrier, the TMEM base address and — on the MoE templates — the per-CTA TMA +# tensormap scratch. The ab pipeline itself only counts the operand tensors, so these +# fragments are budgeted once here instead of being modelled stage by stage. _SMEM_FIXED_RESERVE_BY_PIPELINE = {"sm100": 2048, "sm103": 2048, "sm107": 2048} +_SMEM_FIXED_RESERVE_MOE_BY_PIPELINE = {"sm100": 4096, "sm103": 4096, "sm107": 4096} -def _sm_smem_ab_budget_bytes(pipeline: str, device=None) -> int: - return _sm_smem_budget_bytes(device) - _pipeline_fact(_SMEM_FIXED_RESERVE_BY_PIPELINE, pipeline, "SMEM fixed reserve") +def _sm_smem_ab_budget_bytes(pipeline: str, device=None, *, moe: bool = False) -> int: + table = _SMEM_FIXED_RESERVE_MOE_BY_PIPELINE if moe else _SMEM_FIXED_RESERVE_BY_PIPELINE + return _sm_smem_budget_bytes(device) - _pipeline_fact(table, pipeline, "SMEM fixed reserve") _L2_RETENTION_DIVISOR = 3 @@ -106,16 +110,15 @@ def smem_max_ab_stages( cta_tile_k_bytes: int, *, cta_group: int = 1, - acc_stages: int = 2, extra_smem_bytes: int = 0, extra_per_stage_bytes: int = 0, pipeline: str, + moe: bool = False, device=None, ) -> int: smem_b_n = cta_tile_n // cta_group - per_stage = (cta_tile_m + smem_b_n) * cta_tile_k_bytes + extra_per_stage_bytes + 2 * 8 - fixed = 2 * acc_stages * 8 + 8 - avail = _sm_smem_ab_budget_bytes(pipeline, device) - fixed - extra_smem_bytes + per_stage = (cta_tile_m + smem_b_n) * cta_tile_k_bytes + extra_per_stage_bytes + avail = _sm_smem_ab_budget_bytes(pipeline, device, moe=moe) - extra_smem_bytes if avail < per_stage: raise ValueError( f"tile ({cta_tile_m},{cta_tile_n},K={cta_tile_k_bytes}B) " @@ -307,6 +310,7 @@ def max_ab_stages( *, extra_smem_bytes: int = 0, extra_per_stage_bytes: int = 0, + moe: bool = False, ) -> int: """Largest SMEM pipeline depth under ``cta_group`` (2-CTA MMA halves B's SMEM N, so it fits more stages).""" @@ -315,10 +319,10 @@ def max_ab_stages( self.cta_tile_n, self.cta_tile_k_bytes, cta_group=cta_group, - acc_stages=self.acc_stages, extra_smem_bytes=extra_smem_bytes, extra_per_stage_bytes=extra_per_stage_bytes, pipeline=self.pipeline, + moe=moe, ) # -- multicast model ----------------------------------------------------- diff --git a/test/python/gemm/frost/test_block_scale_matmul.py b/test/python/gemm/frost/test_block_scale_matmul.py index f626d5b8c..0877a02f8 100644 --- a/test/python/gemm/frost/test_block_scale_matmul.py +++ b/test/python/gemm/frost/test_block_scale_matmul.py @@ -1293,10 +1293,9 @@ def test_fp4_rejects_non_k_major(): _CFG_256 = "CONFIG_sm103_128x256x384_128x256x48_cluster1x1" # The kernel compiles on any Blackwell-family GPU (the K=96 mode is an idesc -# bit, not a mnemonic); it RUNS only on sm103. requires_sm103 = pytest.mark.skipif( - _SM != 103, - reason="sm103 block-scale kernels run only on an SM 103 GPU, have " + ("none" if _SM is None else f"sm_{_SM}"), + _SM is None or not (103 <= _SM < 110), + reason="sm103 block-scale kernels run only on 103 <= SM < 110, have " + ("none" if _SM is None else f"sm_{_SM}"), ) @@ -1931,16 +1930,16 @@ def test_render_sm107_tile_constants(_pretend_sm107, combo, cta_n, omma, k_mode, assert got["idesc_a_dtype"] == idesc_dtype -@pytest.mark.parametrize("pipeline,mma_k,runs_on_512", [("sm100", 32, True), ("sm107", 64, False)]) -def test_tmem_columns_follow_the_arch_not_the_pipeline(pipeline, mma_k, runs_on_512, monkeypatch): +@pytest.mark.parametrize("pipeline,mma_k", [("sm100", 32), ("sm107", 64)]) +def test_tmem_columns_follow_the_arch_not_the_pipeline(pipeline, mma_k, monkeypatch): """TMEM size is a property of the GPU, so an sm100-pipeline kernel gets SM 10.7's 576 columns just like an sm107 one — and past 512 the alloc has to ask for the exclusive mode. The extra columns are what lets a 256-wide N tile double-buffer its accumulator instead of overlapping the two. - Only the sm100 config is checked on a 512-column part: PIPELINE_ARCH_RANGES - confines sm107 kernels to 107..109, which always have 576, and this tile's - SFB span needs 520 of them.""" + The 512-column arm is hypothetical for sm107 (PIPELINE_ARCH_RANGES confines + those kernels to 107..109, which always have 576), but both pipelines size + the SF region the same way, so both render the overlap fallback.""" chain = analyze(_bs_chain()) cfg = by_name(f"CONFIG_{pipeline}_128x256x128_128x256x{mma_k}_cluster1x1") @@ -1954,10 +1953,6 @@ def render(): assert got["acc_stages"] == "2" and got["use_acc_overlap"] == "False" monkeypatch.setattr(C, "_current_arch", lambda: 100) - if not runs_on_512: - with pytest.raises(NotImplementedError, match="TMEM span reaches column 520"): - render() - return got = render() assert got["num_tmem_alloc_cols"] == "512" and got["tmem_alloc_exclusive"] == "False" assert got["acc_stages"] == "1" and got["use_acc_overlap"] == "True" diff --git a/test/python/gemm/frost/test_matmul_epilogue_fusion.py b/test/python/gemm/frost/test_matmul_epilogue_fusion.py index e7b63013d..4728612f4 100644 --- a/test/python/gemm/frost/test_matmul_epilogue_fusion.py +++ b/test/python/gemm/frost/test_matmul_epilogue_fusion.py @@ -2526,16 +2526,21 @@ def _pw_aux_order(compiled, aux_bufs): return order +def _prod_dim1(s): + """Product along dim 1 on the CPU -- torch's device-side `prod` NVRTC-JITs its kernel.""" + return s.cpu().prod(dim=1, keepdim=True).to(s.device) + + _RED_CASES = { "avg_full": (cudnn.reduction_mode.AVG, (1, 1, 1), lambda s: s.mean().view(1, 1, 1), 0.0), "avg_row": (cudnn.reduction_mode.AVG, (1, _PW_M, 1), lambda s: s.mean(dim=1, keepdim=True).view(1, _PW_M, 1), 0.0), "avg_col": (cudnn.reduction_mode.AVG, (1, 1, _PW_N), lambda s: s.mean(dim=0, keepdim=True).view(1, 1, _PW_N), 0.0), "norm1_row": (cudnn.reduction_mode.NORM1, (1, _PW_M, 1), lambda s: s.abs().sum(dim=1, keepdim=True).view(1, _PW_M, 1), 0.0), - "mul_row": (cudnn.reduction_mode.MUL, (1, _PW_M, 1), lambda s: s.prod(dim=1, keepdim=True).view(1, _PW_M, 1), 1.0), + "mul_row": (cudnn.reduction_mode.MUL, (1, _PW_M, 1), lambda s: _prod_dim1(s).view(1, _PW_M, 1), 1.0), "mul_no_zeros_row": ( cudnn.reduction_mode.MUL_NO_ZEROS, (1, _PW_M, 1), - lambda s: torch.where(s == 0, torch.ones_like(s), s).prod(dim=1, keepdim=True).view(1, _PW_M, 1), + lambda s: _prod_dim1(torch.where(s == 0, torch.ones_like(s), s)).view(1, _PW_M, 1), 1.0, ), } diff --git a/test/python/gemm/frost/test_moe_grouped_matmul_fwd.py b/test/python/gemm/frost/test_moe_grouped_matmul_fwd.py index bcba678c0..ff88d7de7 100644 --- a/test/python/gemm/frost/test_moe_grouped_matmul_fwd.py +++ b/test/python/gemm/frost/test_moe_grouped_matmul_fwd.py @@ -50,6 +50,8 @@ def _vp_moe(compiled, token, weight, fto, output): ("CONFIG_sm100_256x256x128_128x256x32_cluster2x1", 2), # num_mma_m=2 on the pair ("CONFIG_sm100_256x128x128_128x128x32_cluster1x1", 1), # num_mma_m=2 ("CONFIG_sm100_128x128x128_64x128x32_cluster1x1", 1), # num_mma_m=2 at mma_inst_m=64 + ("CONFIG_sm100_128x256x128_128x256x32_cluster2x2", 1), + ("CONFIG_sm100_128x256x128_128x256x32_cluster4x2", 2), ] # The plain-e2e test additionally covers N-tiles that are not a multiple of 32 diff --git a/test/python/gemm/frost/test_template_epilogue_parity.py b/test/python/gemm/frost/test_template_epilogue_parity.py new file mode 100644 index 000000000..65bb122ef --- /dev/null +++ b/test/python/gemm/frost/test_template_epilogue_parity.py @@ -0,0 +1,208 @@ +"""The epilogue drain is ONE logic across all 20 kernel templates. + +Source-level only -- no GPU, no render, no JIT. See CLAUDE.md +"The epilogue is ONE logic across all 20 templates". +""" + +import ast +import difflib +import pathlib +import re +import textwrap + +import pytest + +import cudnn.gemm.frost + +pytestmark = pytest.mark.L0 + +_MARKER = re.compile(r"^[ \t]*# *@@EPILOGUE_(SETUP|DRAIN):(BEGIN|END)@@[ \t]*$") + +# Groups the region must be identical within. A template must appear in exactly +# one group per region; adding a template makes the completeness test fail until +# its group is declared here. +_PLAIN_1 = [ + "sm100_matmul_1ctamma.py", + "sm100_matmul_1ctamma_static.py", + "sm100_matmul_mainloop_1ctamma.py", +] +_PLAIN_2 = [ + "sm100_matmul_2ctamma.py", + "sm100_matmul_2ctamma_static.py", + "sm100_matmul_mainloop_2ctamma.py", +] +_BS_1 = [ + "sm100_block_scale_matmul_1ctamma.py", + "sm100_block_scale_matmul_1ctamma_static.py", + "sm103_block_scale_matmul_1ctamma.py", + "sm107_block_scale_matmul_1ctamma.py", +] +_BS_2 = [ + "sm100_block_scale_matmul_2ctamma.py", + "sm100_block_scale_matmul_2ctamma_static.py", + "sm103_block_scale_matmul_2ctamma.py", + "sm107_block_scale_matmul_2ctamma.py", +] +_MOE_PLAIN_1 = ["sm100_moe_grouped_matmul_fwd_1ctamma.py"] +_MOE_PLAIN_2 = ["sm100_moe_grouped_matmul_fwd_2ctamma.py"] +_MOE_BS_1 = [ + "sm100_moe_grouped_block_scale_matmul_fwd_1ctamma.py", + "sm107_moe_grouped_block_scale_matmul_fwd_1ctamma.py", +] +_MOE_BS_2 = [ + "sm100_moe_grouped_block_scale_matmul_fwd_2ctamma.py", + "sm107_moe_grouped_block_scale_matmul_fwd_2ctamma.py", +] + +# SETUP (LDTM shape + row base + span list) depends only on cta_group and on +# whether the pipeline is block-scaled, so MoE joins its family. +_SETUP_GROUPS = { + "1ctamma": _PLAIN_1 + _MOE_PLAIN_1, + "2ctamma": _PLAIN_2 + _MOE_PLAIN_2, + "block_scale_1ctamma": _BS_1 + _MOE_BS_1, + "block_scale_2ctamma": _BS_2 + _MOE_BS_2, +} + +# DRAIN additionally splits on MoE: no TMA-store half, no mixed CGA, and the +# store is bounded by the routed group rather than by M. +_DRAIN_GROUPS = { + "1ctamma": _PLAIN_1, + "2ctamma": _PLAIN_2, + "block_scale_1ctamma": _BS_1, + "block_scale_2ctamma": _BS_2, + "moe_1ctamma": _MOE_PLAIN_1, + "moe_2ctamma": _MOE_PLAIN_2, + "moe_block_scale_1ctamma": _MOE_BS_1, + "moe_block_scale_2ctamma": _MOE_BS_2, +} + +_BLOCK_SCALE = set(_BS_1 + _BS_2 + _MOE_BS_1 + _MOE_BS_2) + + +def _template_dir(): + # kernel_templates has no __init__.py (it is exec'd per render), so go + # through the package that does. + return pathlib.Path(cudnn.gemm.frost.__file__).parent / "kernel_templates" + + +def _templates(): + return sorted(p for p in _template_dir().glob("sm*.py")) + + +def _region(path, name): + """The marked region's text, dedented (MoE nests one level deeper).""" + src = path.read_text().splitlines(keepends=True) + begin = end = None + for i, line in enumerate(src): + m = _MARKER.match(line.rstrip("\n")) + if not m or m.group(1) != name: + continue + if m.group(2) == "BEGIN": + assert begin is None, f"{path.name}: duplicate {name}:BEGIN" + begin = i + else: + assert begin is not None, f"{path.name}: {name}:END before BEGIN" + assert end is None, f"{path.name}: duplicate {name}:END" + end = i + assert begin is not None and end is not None, f"{path.name}: missing @@EPILOGUE_{name}@@ markers" + return textwrap.dedent("".join(src[begin + 1 : end])) + + +def _diff(name_a, text_a, name_b, text_b): + return "\n".join(difflib.unified_diff(text_a.splitlines(), text_b.splitlines(), fromfile=name_a, tofile=name_b, lineterm="")) + + +@pytest.mark.parametrize("region,groups", [("SETUP", _SETUP_GROUPS), ("DRAIN", _DRAIN_GROUPS)]) +def test_every_template_is_assigned_to_exactly_one_group(region, groups): + declared = [f for names in groups.values() for f in names] + assert len(declared) == len(set(declared)), f"{region}: a template is in two groups" + on_disk = {p.name for p in _templates()} + assert set(declared) == on_disk, ( + f"{region} group table is out of date -- declare the group of every template.\n" + f" missing from the table: {sorted(on_disk - set(declared))}\n" + f" no longer on disk: {sorted(set(declared) - on_disk)}" + ) + + +@pytest.mark.parametrize( + "region,group", + [("SETUP", g) for g in _SETUP_GROUPS] + [("DRAIN", g) for g in _DRAIN_GROUPS], +) +def test_the_region_is_identical_within_its_group(region, group): + names = (_SETUP_GROUPS if region == "SETUP" else _DRAIN_GROUPS)[group] + d = _template_dir() + ref_name = names[0] + ref = _region(d / ref_name, region) + assert ref.strip(), f"{ref_name}: empty {region} region" + for name in names[1:]: + got = _region(d / name, region) + assert got == ref, ( + f"{region} region of {name} has drifted from its group '{group}'.\n" + f"A new epilogue feature lands in EVERY template of the group, in the same shape.\n" + _diff(ref_name, ref, name, got) + ) + + +def test_the_packed_ldtm_arm_keys_on_the_hardware_mma_m(): + """foot-gun #18: a 2-CTA cluster-MMA m=128 tile also has + epi_rows_per_mma_m == 64, but must NOT take the packed 16x32bx2 path -- + keying the LDTM shape on it is a silent miscompute, not a fault.""" + offenders = [] + for path in _templates(): + for i, line in enumerate(path.read_text().splitlines(), 1): + if "SHAPE_16X32BX2" not in line: + continue + window = path.read_text().splitlines()[max(0, i - 3) : i] + if not any("mma_inst_shape_mnk[0] == 64" in w for w in window): + offenders.append(f"{path.name}:{i}") + assert not offenders, "packed LDTM shape not guarded by the HARDWARE MMA M:\n " + "\n ".join(offenders) + + +def test_block_scale_templates_carry_no_m64_path(): + """tile_config.validate_block_scale_config rejects mma_inst_m % 128 != 0, so + the hardware-M=64 arms are provably dead there and are deleted, not gated.""" + offenders = [f"{p.name}: {tok}" for p in _templates() if p.name in _BLOCK_SCALE for tok in ("SHAPE_16X32BX2", "ld_half_off") if tok in p.read_text()] + assert not offenders, "block-scale forbids mma_inst_m=64 -- drop the dead arm:\n " + "\n ".join(offenders) + + +def test_the_retired_fixed_width_drain_name_is_gone(): + """The drain width comes from `_epi_subtile_spans`; `t2r_inst_repx` was the + fixed-width spelling it replaced.""" + offenders = [p.name for p in _templates() if "t2r_inst_repx" in p.read_text()] + assert not offenders, f"use epi_spans / subtile_w, not t2r_inst_repx: {offenders}" + + +def test_cols_per_acc_stage_has_one_meaning(): + """It briefly named two different quantities: `num_mma_m * epi_cols_per_mma_m` + in the plain pipeline and `epi_cols_per_mma_m` in block-scale.""" + bad_def, bad_use = [], [] + for path in _templates(): + src = path.read_text() + if "cols_per_acc_stage" not in src: + continue + if path.name in _BLOCK_SCALE: + bad_use.append(path.name) + continue + if "cols_per_acc_stage = num_mma_m * epi_cols_per_mma_m" not in src: + bad_def.append(path.name) + assert not bad_use, f"block-scale means epi_cols_per_mma_m -- say so: {bad_use}" + assert not bad_def, f"cols_per_acc_stage must be num_mma_m * epi_cols_per_mma_m: {bad_def}" + + +def test_the_overlap_arm_never_indexes_the_span_list(): + """`epi_spans` is a Python list, so it takes a COMPILE-TIME index; under + acc overlap the subtile index depends on `tile_iter % 2` and is a runtime + value. That arm computes the column offset arithmetically instead.""" + offenders = [] + for path in _templates(): + src = path.read_text() + if "_sub = subtile_idx" not in src: + continue + for i, line in enumerate(src.splitlines(), 1): + if "epi_spans[_sub]" in line: + offenders.append(f"{path.name}:{i}") + assert not offenders, "runtime index into a Python list:\n " + "\n ".join(offenders) + + +def test_the_markers_do_not_break_the_template_parse(): + for path in _templates(): + ast.parse(path.read_text())