diff --git a/python/cudnn/gemm/frost/compiler.py b/python/cudnn/gemm/frost/compiler.py index 77c793ef5..5c3f9a05b 100644 --- a/python/cudnn/gemm/frost/compiler.py +++ b/python/cudnn/gemm/frost/compiler.py @@ -1854,6 +1854,13 @@ def _render_block_scale_template( moe_host_ma_pass = ",\n".join([f"a_{i}" for i in range(na)] + [f"_a_stride_sets[{i}][0]" for i in range(na)]) if moe_host_ma_pass: moe_host_ma_pass += "," + moe_kernel_msfa_params = ",\n".join(f"mSFA_{i}: cute.Tensor" for i in range(na)) + if moe_kernel_msfa_params: + moe_kernel_msfa_params += "," + moe_msfa_list = "mSFA_list = [" + ", ".join(f"mSFA_{i}" for i in range(na)) + "]" + moe_host_msfa_pass = ",\n".join(f"_sfa_operands[{i}]" for i in range(na)) + if moe_host_msfa_pass: + moe_host_msfa_pass += "," replacements = { "INJECT_TILE_CONSTANTS": tile_constants, @@ -1907,6 +1914,9 @@ def _render_block_scale_template( "INJECT_MOE_KERNEL_MA_PARAMS": moe_kernel_ma_params, "INJECT_MOE_MA_LIST": moe_ma_list, "INJECT_MOE_HOST_MA_PASS": moe_host_ma_pass, + "INJECT_MOE_KERNEL_MSFA_PARAMS": moe_kernel_msfa_params, + "INJECT_MOE_MSFA_LIST": moe_msfa_list, + "INJECT_MOE_HOST_MSFA_PASS": moe_host_msfa_pass, } ) @@ -2759,12 +2769,9 @@ def _use_tma_store_epi(chain, cfg, vec_bytes_epi: int, cta_group: int) -> bool: - 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}: the drain widens to epi_n only for 2-byte output. - - M-major output: 16B-aligned M (16x256b TMEM-load + stmatrix.trans + tma_store). + - M-major output: 16B-aligned M (16x256b TMEM-load + stmatrix.trans + tma_store), + and not MoE: the six MoE templates carry only the N-major TMA store. """ - if chain.has_moe: - # MoE scatters output rows by routed group; the TMA-store path writes - # contiguous tiles with no group offset → STG-only. - return False if chain.is_multi_gemm: # No multi-accumulator hook in the TMA-store path → STG only. return False @@ -2802,6 +2809,8 @@ def _use_tma_store_epi(chain, cfg, vec_bytes_epi: int, cta_group: int) -> bool: if cta_group == 2 and cfg.cta_tile_n < 64: return False if chain.out_major == "m": + if chain.has_moe: + return False m_align = 16 // DTYPE_BYTES[chain.output_dtype] return chain.matmul.M % m_align == 0 return True @@ -3211,18 +3220,22 @@ class CompiledMoeGemm: # The epilogue chunk width (bytes) the kernel was RENDERED with (tile- # clamped); drives the runtime output/aux alignment requirements. vec_bytes_epi: "int | None" = None + # Tensormap slots each CTA patches: one per distinct A operand, plus the + # output descriptor when the TMA-store epilogue is on (re-dimensioned per + # routed group so the hardware clips the ragged tail). + _desc_slots_per_cta: int = 0 accepts_stream: ClassVar[bool] = True # stream-aware dispatch (see CompiledFusedGemm) @property def workspace_bytes(self) -> int: - """Per-CTA A-descriptor scratch: one 128-byte tensormap slot per CTA per - distinct A operand. The persistent grid is shape-independent, so this is + """Per-CTA tensormap scratch: one 128-byte slot per CTA per patched + descriptor. The persistent grid is shape-independent, so this is constant for the plan — which is why override-shape needs no re-query.""" - return self._grid_ctas * self.chain.num_a_operands * _MOE_DESC_SLOT_BYTES + return self._grid_ctas * self._desc_slots_per_cta * _MOE_DESC_SLOT_BYTES def _make_workspace(self, n_slots, caller=None): - """The per-CTA A-descriptor GMEM workspace (16 int64/slot, 128-byte - aligned). ``n_slots`` = grid_ctas * num_a_operands. Carved from the + """The per-CTA tensormap GMEM workspace (16 int64/slot, 128-byte + aligned). ``n_slots`` = grid_ctas * _desc_slots_per_cta. Carved from the CALLER's buffer when execute() supplied one; otherwise from one this plan owns (the direct jit_from_cudnn_graph path passes no workspace).""" if caller is None: @@ -3297,8 +3310,8 @@ def _launch_single(self, token, weight, first_token_offset, output, snke, worksp (_wrap_raw_tensor(ci) if (spec.is_reduction or spec.is_quant_scale) else _maybe_wrap_layout(ci, _LEADING_DIM_C)) for spec, ci in zip(outputs_spec, c_perms) ] - # A-descriptor workspace: one 128-byte tensormap slot per CTA. - workspace = self._make_workspace(self._grid_ctas * self.chain.num_a_operands, workspace) + # Tensormap workspace: one 128-byte slot per CTA per patched descriptor. + workspace = self._make_workspace(self._grid_ctas * self._desc_slots_per_cta, workspace) return self._launchable( problem_size, first_token_offset, @@ -3418,8 +3431,8 @@ def _call_multi_gemm(self, gemm_pairs, first_token_offset, output, snke, *aux, w f"per-group aux {ref.name!r} must be rank-3 with leading dim " f"{num_groups} (the first_token_offset length); got shape {tuple(t.shape)}" ) aux = tuple(_maybe_wrap_layout(_reshape_aux_to_fake(t, ref), _LEADING_DIM_AUX) for ref, t in zip(chain.aux_tensors, aux)) - # Workspace: one 128-B A descriptor per distinct A operand per CTA. - workspace = self._make_workspace(self._grid_ctas * na, workspace) + # Workspace: one 128-B tensormap slot per patched descriptor per CTA. + workspace = self._make_workspace(self._grid_ctas * self._desc_slots_per_cta, workspace) return self._launchable( problem_size, first_token_offset, @@ -3459,22 +3472,16 @@ def _jit_moe( _check_dtype_config_compat(chain, config, cta_group) _check_input_alignment(chain) _compute_output_vec_bytes(chain) - global _FORCE_STG_EPI - prev_force = _FORCE_STG_EPI - _FORCE_STG_EPI = True # MoE epilogue is STG-only - try: - vec_bytes_epi = _epi_vec_bytes(chain, config, cta_group) - _check_block_quant_supported(chain, vec_bytes_epi, config, cta_group) - use_tma = (not _FORCE_STG_EPI) and _use_tma_store_epi(chain, config, vec_bytes_epi, cta_group) - snippets = generate( - chain, - vec_bytes_epi=vec_bytes_epi, - output_elem_bytes=DTYPE_BYTES[chain.output_dtype], - use_tma_store=use_tma, - ) - src = _render_template(chain, snippets, config, cta_group) - finally: - _FORCE_STG_EPI = prev_force + vec_bytes_epi = _epi_vec_bytes(chain, config, cta_group) + _check_block_quant_supported(chain, vec_bytes_epi, config, cta_group) + use_tma = (not _FORCE_STG_EPI) and _use_tma_store_epi(chain, config, vec_bytes_epi, cta_group) + snippets = generate( + chain, + vec_bytes_epi=vec_bytes_epi, + output_elem_bytes=DTYPE_BYTES[chain.output_dtype], + use_tma_store=use_tma, + ) + src = _render_template(chain, snippets, config, cta_group) mod = _import_kernel(src) digest = hashlib.sha256(src.encode("utf-8")).hexdigest()[:16] cluster_m, cluster_n = config.cgrp_size_m, config.cgrp_size_n @@ -3489,6 +3496,7 @@ def _jit_moe( aux_names=[aux.name for aux in chain.aux_tensors], binding=binding, vec_bytes_epi=vec_bytes_epi, + _desc_slots_per_cta=chain.num_a_operands + (1 if use_tma else 0), ) @@ -3587,6 +3595,10 @@ class CompiledMoeBlockScaleGemm: # The epilogue chunk width (bytes) the kernel was RENDERED with (tile- # clamped); drives the runtime output/aux alignment requirements. vec_bytes_epi: "int | None" = None + # Tensormap slots each CTA patches: one per distinct A operand, one per SFA + # (its base carries start_sf_block_m and its m extent bounds the group), plus + # the output descriptor when the TMA-store epilogue re-dimensions it per group. + _desc_slots_per_cta: int = 0 accepts_stream: ClassVar[bool] = True # stream-aware dispatch (see CompiledFusedGemm) @property @@ -3594,7 +3606,7 @@ def workspace_bytes(self) -> int: """Per-CTA A-descriptor scratch: one 128-byte tensormap slot per CTA per distinct A operand. The persistent grid is shape-independent, so this is constant for the plan — which is why override-shape needs no re-query.""" - return self._grid_ctas * self.chain.num_a_operands * _MOE_DESC_SLOT_BYTES + return self._grid_ctas * self._desc_slots_per_cta * _MOE_DESC_SLOT_BYTES def _make_workspace(self, n_slots, caller=None): """The per-CTA A-descriptor GMEM workspace (16 int64/slot, 128-byte @@ -3682,7 +3694,7 @@ def _launch_single(self, token, weight, sfa, sfb, first_token_offset, output, sn ] msfa = _maybe_wrap_layout(sfa.permute(1, 2, 0), _LEADING_DIM_AUX) msfb = _maybe_wrap_layout(sfb.permute(1, 2, 0), _LEADING_DIM_AUX) - workspace = self._make_workspace(self._grid_ctas * self.chain.num_a_operands, workspace) + workspace = self._make_workspace(self._grid_ctas * self._desc_slots_per_cta, workspace) return self._launchable( problem_size, first_token_offset, @@ -3817,7 +3829,7 @@ def _call_multi_gemm(self, gemm_pairs, first_token_offset, output, snke, *aux, w f"per-group aux {ref.name!r} must be rank-3 with leading dim " f"{num_groups} (the first_token_offset length); got shape {tuple(t.shape)}" ) aux = tuple(_maybe_wrap_layout(_reshape_aux_to_fake(t, ref), _LEADING_DIM_AUX) for ref, t in zip(chain.aux_tensors, aux)) - workspace = self._make_workspace(self._grid_ctas * na, workspace) + workspace = self._make_workspace(self._grid_ctas * self._desc_slots_per_cta, workspace) return self._launchable( problem_size, first_token_offset, @@ -3870,12 +3882,13 @@ def _jit_moe_block_scale( raise NotImplementedError(_arch_reason) _compute_output_vec_bytes(chain) vec_bytes_epi = _epi_vec_bytes(chain, config, cta_group) + use_tma = (not _FORCE_STG_EPI) and _use_tma_store_epi(chain, config, vec_bytes_epi, cta_group) _check_block_quant_supported(chain, vec_bytes_epi, config, cta_group) snippets = generate( chain, vec_bytes_epi=vec_bytes_epi, output_elem_bytes=DTYPE_BYTES[chain.output_dtype], - use_tma_store=(not _FORCE_STG_EPI) and _use_tma_store_epi(chain, config, vec_bytes_epi, cta_group), + use_tma_store=use_tma, ) src = _render_block_scale_template(chain, snippets, config, cta_group, fallback_cluster=_mixed_cga_fallback(config, cta_group, _tmpl.file)) mod = _import_kernel(src) @@ -3891,4 +3904,5 @@ def _jit_moe_block_scale( _grid_ctas=grid_ctas, binding=binding, vec_bytes_epi=vec_bytes_epi, + _desc_slots_per_cta=chain.num_a_operands * 2 + (1 if use_tma else 0), ) diff --git a/python/cudnn/gemm/frost/kernel_templates/_tile_helpers.py b/python/cudnn/gemm/frost/kernel_templates/_tile_helpers.py index a5a473e6f..8464d2476 100644 --- a/python/cudnn/gemm/frost/kernel_templates/_tile_helpers.py +++ b/python/cudnn/gemm/frost/kernel_templates/_tile_helpers.py @@ -79,6 +79,16 @@ def replace_tensormap_global_dim_1(desc_ptr, new_dim) -> None: ) +@cute.jit +def replace_tensormap_global_dim_2(desc_ptr, new_dim) -> None: + nvvm.tensormap_replace( + nvvm.TensormapField.GLOBAL_DIM, + desc_ptr, + new_value=cutlass.Int32(new_dim), + ord=2, + ) + + @cute.jit def replace_tensormap_global_address(desc_ptr, new_address) -> None: nvvm.tensormap_replace( 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 56441776a..285eb5cea 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 @@ -27,6 +27,7 @@ moe_swizzle_tile as _moe_swizzle_tile, replace_tensormap_global_address as _replace_tensormap_global_address, replace_tensormap_global_dim_1 as _replace_tensormap_global_dim_1, + replace_tensormap_global_dim_2 as _replace_tensormap_global_dim_2, tcgen05_alloc as _tcgen05_alloc, tcgen05_dealloc as _tcgen05_dealloc, tcgen05_mma_block_scale as _tcgen05_mma_block_scale, @@ -44,6 +45,10 @@ # it to the per-CTA GMEM workspace the TMA reads. # @@INJECT_TILE_CONSTANTS@@ +# Tensormap workspace slots per CTA: the A operands, plus the output descriptor +# when the TMA-store epilogue re-dimensions it per routed group. +moe_desc_slots = num_a_operands * 2 + (1 if use_tma_store_epi else 0) + 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") @@ -103,12 +108,21 @@ def _kernel( a_tma_workspace: cute.Tensor, # @@INJECT_KERNEL_AB_DESC_PARAMS@@ # @@INJECT_MOE_KERNEL_MA_PARAMS@@ + # @@INJECT_MOE_KERNEL_MSFA_PARAMS@@ # @@INJECT_KERNEL_TAP_PARAMS@@ # @@INJECT_KERNEL_REDUCTION_STRIDE_PARAMS@@ # @@INJECT_KERNEL_AUX_PARAMS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_KERNEL_TMA_C_PARAMS@@ + # @@TMA_STORE_ONLY:END@@ ) -> None: # @@INJECT_AB_DESC_LISTS@@ # @@INJECT_MOE_MA_LIST@@ + # @@INJECT_MOE_MSFA_LIST@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_TMA_C_LISTS@@ + tma_c_desc = tma_c_descs[0] + # @@TMA_STORE_ONLY:END@@ mma_warp_id = 4 tma_warp_id = 5 @@ -145,6 +159,10 @@ def _kernel( nvvm.prefetch_tensormap(tma_b_descs[_j].get_ptr()) nvvm.prefetch_tensormap(tma_sfb_descs[_j].get_ptr()) + # @@TMA_STORE_ONLY:BEGIN@@ + nvvm.prefetch_tensormap(tma_c_desc.get_ptr()) + # @@TMA_STORE_ONLY:END@@ + cluster_linear_init = bidx // cluster_m a_pattern = 0 @@ -196,6 +214,32 @@ def _kernel( ) for _ in range(num_a_operands) ] + tma_sfa_desc_smem_list = [ + cutlass.Array( + cutlass.Int64, + TENSOR_MAP_QWORDS, + space=cutlass.AddressSpace.smem, + alignment=128, + ) + for _ in range(num_a_operands) + ] + + # @@TMA_STORE_ONLY:BEGIN@@ + # One epilogue subtile = one MMA-M block x 32 cols; the M blocks reuse it. + epi_subtile_elems = epi_tile_mn[0] * epi_tile_mn[1] + smem_d_ptr = cutlass.Array( + cd_dtype, + epi_subtile_elems * EPI_SMEM_STAGES, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + tma_c_desc_smem = cutlass.Array( + cutlass.Int64, + TENSOR_MAP_QWORDS, + space=cutlass.AddressSpace.smem, + alignment=128, + ) + # @@TMA_STORE_ONLY:END@@ sA_elems = sA_packed_elems sB_elems = sB_packed_elems @@ -516,7 +560,7 @@ def _kernel( lane = tidx % 32 block_linear = bidx + bidy * gridx - cta_desc_base_list = [a_tma_workspace.iterator.raw_ptr() + (block_linear * num_a_operands + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands)] + cta_desc_base_list = [a_tma_workspace.iterator.raw_ptr() + (block_linear * moe_desc_slots + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands)] a_desc_tma_ptr_list = [ cute.make_ptr( cutlass.Int64, @@ -525,10 +569,24 @@ def _kernel( ) for _ai in range(num_a_operands) ] + sfa_desc_base_list = [ + a_tma_workspace.iterator.raw_ptr() + (block_linear * moe_desc_slots + num_a_operands + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands) + ] + sfa_desc_tma_ptr_list = [ + cute.make_ptr( + cutlass.Int64, + sfa_desc_base_list[_ai].toint(), + mem_space=cute.AddressSpace.generic, + ) + for _ai in range(num_a_operands) + ] + sfa_block_bytes = 512 * (((k // block_size) + 3) // 4) previous_group_begin = cutlass.Int32(-1) if elect_one: for _ai in cutlass.range_constexpr(num_a_operands): _copy_tensormap_to_workspace(tma_a_descs[_ai].get_ptr(), tma_a_desc_smem_list[_ai]) + for _ai in cutlass.range_constexpr(num_a_operands): + _copy_tensormap_to_workspace(tma_sfa_descs[_ai].get_ptr(), tma_sfa_desc_smem_list[_ai]) nvvm.bar_warp_sync(0xFFFFFFFF) while is_valid != 0: @@ -556,7 +614,7 @@ def _kernel( if is_valid != 0: coord_m_group = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] - sfa_m_block = start_sf_block_m + coord_m_group // 128 + sfa_m_block = coord_m_group // 128 sfb_n_block = coord_n_per_cta // 128 if group_begin != previous_group_begin: @@ -573,6 +631,18 @@ def _kernel( (cta_desc_base_list[_ai] + lane).store((tma_a_desc_smem_list[_ai].subview(lane)).load()) nvvm.bar_warp_sync(0xFFFFFFFF) _fence_tensormap_release() + for _ai in cutlass.range_constexpr(num_a_operands): + _fence_tensormap_acquire(sfa_desc_tma_ptr_list[_ai]) + for _ai in cutlass.range_constexpr(num_a_operands): + if elect_one: + sfa_base = mSFA_list[_ai].iterator.raw_ptr().toint() + start_sf_block_m * sfa_block_bytes + _replace_tensormap_global_address(tma_sfa_desc_smem_list[_ai], sfa_base) + _replace_tensormap_global_dim_2(tma_sfa_desc_smem_list[_ai], cute.ceil_div(group_end - group_begin, 128)) + nvvm.bar_warp_sync(0xFFFFFFFF) + if lane < TENSOR_MAP_QWORDS: + (sfa_desc_base_list[_ai] + lane).store((tma_sfa_desc_smem_list[_ai].subview(lane)).load()) + nvvm.bar_warp_sync(0xFFFFFFFF) + _fence_tensormap_release() for k_tile_idx in range(num_k_tiles): stage = ab_iter % ab_stages @@ -611,7 +681,7 @@ def _kernel( 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(), + sfa_desc_tma_ptr_list[_ai], (0, coord_sf_k, sfa_m_block, cutlass.Int32(0)), sf_full_mbar_ptr.subview(stage), [], @@ -965,6 +1035,24 @@ def _kernel( lane = tidx % 32 # @@EPILOGUE_SETUP:END@@ + # @@TMA_STORE_ONLY:BEGIN@@ + epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) + # The routed output is a single (1, S, N) tensor, so the batch coord is fixed. + tile_l = cutlass.Int32(0) + epi_block_linear = bidx + bidy * gridx + d_desc_base = a_tma_workspace.iterator.raw_ptr() + (epi_block_linear * moe_desc_slots + num_a_operands * 2) * TENSOR_MAP_QWORDS + d_desc_tma_ptr = cute.make_ptr( + cutlass.Int64, + d_desc_base.toint(), + mem_space=cute.AddressSpace.generic, + ) + previous_group_end = cutlass.Int32(-1) + if warp_idx == 0: + if elect_one: + _copy_tensormap_to_workspace(tma_c_desc.get_ptr(), tma_c_desc_smem) + nvvm.bar_warp_sync(0xFFFFFFFF) + # @@TMA_STORE_ONLY:END@@ + while is_valid != 0: while not nvvm.mbarrier_try_wait_parity( sched_full_mbar_ptr.subview(sched_stage), @@ -988,6 +1076,22 @@ def _kernel( if is_valid != 0: coord_m_tile = group_begin + tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] + # @@TMA_STORE_ONLY:BEGIN@@ + # Re-dimension D to this group's last row so the hardware clips the + # ragged tail; the base stays put, so the store coords are global. + if warp_idx == 0: + if group_end != previous_group_end: + previous_group_end = group_end + nvvm.cp_async_bulk_wait_group(0, read=True) + _fence_tensormap_acquire(d_desc_tma_ptr) + if elect_one: + _replace_tensormap_global_dim_1(tma_c_desc_smem, group_end) + nvvm.bar_warp_sync(0xFFFFFFFF) + if lane < TENSOR_MAP_QWORDS: + (d_desc_base + lane).store((tma_c_desc_smem.subview(lane)).load()) + nvvm.bar_warp_sync(0xFFFFFFFF) + _fence_tensormap_release() + # @@TMA_STORE_ONLY:END@@ # @@EPILOGUE_DRAIN:BEGIN@@ coord_n_c = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] @@ -1049,6 +1153,43 @@ def _kernel( 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 * subtile_w) + + vec_f32 = c_rmem_vec + col_j = col + linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 + + # @@INJECT_EPILOGUE@@ + + 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( + barrier_id=EPI_SYNC_BAR_ID, + thread_count=num_epilogue_warps * 32, + ) + + if warp_idx == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_global_shared_cta( + d_desc_tma_ptr, + smem_subtile_ptr, + (col, coord_m, tile_l), + ) + if elect_one: + nvvm.cp_async_bulk_commit_group() + nvvm.cp_async_bulk_wait_group(EPI_SMEM_STAGES - 1, read=True) + + nvvm.barrier_cta_sync( + barrier_id=EPI_SYNC_BAR_ID, + thread_count=num_epilogue_warps * 32, + ) + # @@TMA_STORE_ONLY:END@@ + + # @@STG_ONLY:BEGIN@@ if row_active and row < group_end: for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize @@ -1058,6 +1199,7 @@ 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@@ @@ -1081,6 +1223,9 @@ def _host( # @@INJECT_HOST_AB_PARAMS@@ # @@INJECT_HOST_TAP_PARAMS@@ # @@INJECT_HOST_AUX_PARAMS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_PARAMS@@ + # @@TMA_STORE_ONLY:END@@ stream: _cuda.CUstream, ) -> None: # @@INJECT_HOST_AB_LISTS@@ @@ -1112,6 +1257,23 @@ def _host( _stride_idx += 3 # @@INJECT_HOST_REDUCTION_STRIDES@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_LISTS@@ + c = _tma_c_outputs[0] + tma_c_desc = _tma.create_tensor_map_tiled( + global_address=c.iterator.toint(), + dtype=cd_tma_dtype, + global_dims=[n, m, 1], + global_strides=[ + out_stride_m_0 * cd_dtype.width // 128, + out_stride_l_0 * cd_dtype.width // 128, + ], + box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], + swizzle=epi_tma_swizzle, + ) + tma_c_desc_list = [tma_c_desc] + # @@TMA_STORE_ONLY:END@@ + tma_a_desc_list = [] for _a_idx, _a_op in enumerate(_a_operands): a_stride_m, a_stride_k, a_stride_l = _a_stride_sets[_a_idx] @@ -1225,9 +1387,13 @@ def _host( a_tma_workspace, # @@INJECT_HOST_KERNEL_DESC_PASS@@ # @@INJECT_MOE_HOST_MA_PASS@@ + # @@INJECT_MOE_HOST_MSFA_PASS@@ # @@INJECT_HOST_TAP_PASS@@ # @@INJECT_HOST_REDUCTION_STRIDE_PASS@@ # @@INJECT_HOST_AUX_PASS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_PASS@@ + # @@TMA_STORE_ONLY:END@@ ).launch( grid=grid_shape, block=(threads_per_cta, 1, 1), @@ -1298,7 +1464,7 @@ def _make_fake_sfb(): grid_ctas = grid_num_clusters * cluster_m * cluster_n fake_a_tma_workspace = make_fake_compact_tensor( cutlass.Int64, - (grid_ctas * num_a_operands * 16,), + (grid_ctas * moe_desc_slots * 16,), stride_order=(0,), assumed_align=128, ) @@ -1317,6 +1483,18 @@ def _sym_operand_strides(is_mn_major: bool) -> tuple: # @@INJECT_COMPILE_REDUCTION_STRIDE_DECLS@@ # @@INJECT_COMPILE_AB_FAKES@@ # @@INJECT_COMPILE_TAP_FAKES@@ + + # @@TMA_STORE_ONLY:BEGIN@@ + def _make_fake_c(): + return make_fake_compact_tensor( + cd_dtype, + (sym_m, sym_n // cd_fake_n_div, 1), + stride_order=(1, 0, 2), + assumed_align=16, + ) + + # @@INJECT_COMPILE_TMA_C_FAKES@@ + # @@TMA_STORE_ONLY:END@@ problem_size = ( sym_m, sym_n, @@ -1337,6 +1515,9 @@ def _sym_operand_strides(is_mn_major: bool) -> tuple: # @@INJECT_COMPILE_AB_PASS@@ # @@INJECT_COMPILE_TAP_PASS@@ # @@INJECT_COMPILE_AUX_PASS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_COMPILE_TMA_C_PASS@@ + # @@TMA_STORE_ONLY:END@@ stream=_fake_stream, options=frost_compile_options, ) 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 3f621acbb..4756029cb 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 @@ -32,6 +32,7 @@ moe_swizzle_tile as _moe_swizzle_tile, replace_tensormap_global_address as _replace_tensormap_global_address, replace_tensormap_global_dim_1 as _replace_tensormap_global_dim_1, + replace_tensormap_global_dim_2 as _replace_tensormap_global_dim_2, tcgen05_alloc as _tcgen05_alloc, tcgen05_dealloc as _tcgen05_dealloc, tcgen05_mma_block_scale as _tcgen05_mma_block_scale, @@ -46,6 +47,10 @@ # @@INJECT_TILE_CONSTANTS@@ +# Tensormap workspace slots per CTA: the A operands, plus the output descriptor +# when the TMA-store epilogue re-dimensions it per routed group. +moe_desc_slots = num_a_operands * 2 + (1 if use_tma_store_epi else 0) + 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") @@ -104,13 +109,22 @@ def _kernel( a_tma_workspace: cute.Tensor, # @@INJECT_KERNEL_AB_DESC_PARAMS@@ # @@INJECT_MOE_KERNEL_MA_PARAMS@@ + # @@INJECT_MOE_KERNEL_MSFA_PARAMS@@ # @@INJECT_KERNEL_TAP_PARAMS@@ # @@INJECT_KERNEL_REDUCTION_STRIDE_PARAMS@@ # @@INJECT_KERNEL_AUX_PARAMS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_KERNEL_TMA_C_PARAMS@@ + # @@TMA_STORE_ONLY:END@@ ) -> None: # @@INJECT_AB_DESC_LISTS@@ # @@INJECT_MOE_MA_LIST@@ + # @@INJECT_MOE_MSFA_LIST@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_TMA_C_LISTS@@ + tma_c_desc = tma_c_descs[0] + # @@TMA_STORE_ONLY:END@@ mma_warp_id = 4 tma_warp_id = 5 @@ -153,6 +167,10 @@ def _kernel( nvvm.prefetch_tensormap(tma_b_descs[_j].get_ptr()) nvvm.prefetch_tensormap(tma_sfb_descs[_j].get_ptr()) + # @@TMA_STORE_ONLY:BEGIN@@ + nvvm.prefetch_tensormap(tma_c_desc.get_ptr()) + # @@TMA_STORE_ONLY:END@@ + a_pattern = 0 for n_idx in cutlass.range_constexpr(cluster_n): a_pattern = a_pattern | (1 << (n_idx * cluster_m)) @@ -196,6 +214,32 @@ def _kernel( ) for _ in range(num_a_operands) ] + tma_sfa_desc_smem_list = [ + cutlass.Array( + cutlass.Int64, + TENSOR_MAP_QWORDS, + space=cutlass.AddressSpace.smem, + alignment=128, + ) + for _ in range(num_a_operands) + ] + + # @@TMA_STORE_ONLY:BEGIN@@ + # One epilogue subtile = one MMA-M block x 32 cols; the M blocks reuse it. + epi_subtile_elems = epi_tile_mn[0] * epi_tile_mn[1] + smem_d_ptr = cutlass.Array( + cd_dtype, + epi_subtile_elems * EPI_SMEM_STAGES, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + tma_c_desc_smem = cutlass.Array( + cutlass.Int64, + TENSOR_MAP_QWORDS, + space=cutlass.AddressSpace.smem, + alignment=128, + ) + # @@TMA_STORE_ONLY:END@@ sA_elems = sA_packed_elems sB_elems = sB_packed_elems @@ -522,7 +566,7 @@ def _kernel( lane = tidx % 32 block_linear = bidx + bidy * gridx - cta_desc_base_list = [a_tma_workspace.iterator.raw_ptr() + (block_linear * num_a_operands + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands)] + cta_desc_base_list = [a_tma_workspace.iterator.raw_ptr() + (block_linear * moe_desc_slots + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands)] a_desc_tma_ptr_list = [ cute.make_ptr( cutlass.Int64, @@ -531,10 +575,24 @@ def _kernel( ) for _ai in range(num_a_operands) ] + sfa_desc_base_list = [ + a_tma_workspace.iterator.raw_ptr() + (block_linear * moe_desc_slots + num_a_operands + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands) + ] + sfa_desc_tma_ptr_list = [ + cute.make_ptr( + cutlass.Int64, + sfa_desc_base_list[_ai].toint(), + mem_space=cute.AddressSpace.generic, + ) + for _ai in range(num_a_operands) + ] + sfa_block_bytes = 512 * (((k // block_size) + 3) // 4) previous_group_begin = cutlass.Int32(-1) if elect_one: for _ai in cutlass.range_constexpr(num_a_operands): _copy_tensormap_to_workspace(tma_a_descs[_ai].get_ptr(), tma_a_desc_smem_list[_ai]) + for _ai in cutlass.range_constexpr(num_a_operands): + _copy_tensormap_to_workspace(tma_sfa_descs[_ai].get_ptr(), tma_sfa_desc_smem_list[_ai]) nvvm.bar_warp_sync(0xFFFFFFFF) while is_valid != 0: @@ -563,7 +621,7 @@ def _kernel( coord_m_group = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] coord_n_pair = tile_n * cgrp_tile_mnk[1] + n_rank * logical_cta_tile_n - sfa_m_block = start_sf_block_m + coord_m_group // 128 + sfa_m_block = coord_m_group // 128 sfb_n_block = coord_n_pair // 128 if group_begin != previous_group_begin: @@ -580,6 +638,18 @@ def _kernel( (cta_desc_base_list[_ai] + lane).store((tma_a_desc_smem_list[_ai].subview(lane)).load()) nvvm.bar_warp_sync(0xFFFFFFFF) _fence_tensormap_release() + for _ai in cutlass.range_constexpr(num_a_operands): + _fence_tensormap_acquire(sfa_desc_tma_ptr_list[_ai]) + for _ai in cutlass.range_constexpr(num_a_operands): + if elect_one: + sfa_base = mSFA_list[_ai].iterator.raw_ptr().toint() + start_sf_block_m * sfa_block_bytes + _replace_tensormap_global_address(tma_sfa_desc_smem_list[_ai], sfa_base) + _replace_tensormap_global_dim_2(tma_sfa_desc_smem_list[_ai], cute.ceil_div(group_end - group_begin, 128)) + nvvm.bar_warp_sync(0xFFFFFFFF) + if lane < TENSOR_MAP_QWORDS: + (sfa_desc_base_list[_ai] + lane).store((tma_sfa_desc_smem_list[_ai].subview(lane)).load()) + nvvm.bar_warp_sync(0xFFFFFFFF) + _fence_tensormap_release() for k_tile_idx in range(num_k_tiles): stage = ab_iter % ab_stages @@ -620,7 +690,7 @@ def _kernel( 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(), + sfa_desc_tma_ptr_list[_ai], (0, coord_sf_k, sfa_m_block, cutlass.Int32(0)), sf_full_mbar_ptr.subview(stage), [], @@ -1023,6 +1093,24 @@ def _kernel( lane = tidx % 32 # @@EPILOGUE_SETUP:END@@ + # @@TMA_STORE_ONLY:BEGIN@@ + epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) + # The routed output is a single (1, S, N) tensor, so the batch coord is fixed. + tile_l = cutlass.Int32(0) + epi_block_linear = bidx + bidy * gridx + d_desc_base = a_tma_workspace.iterator.raw_ptr() + (epi_block_linear * moe_desc_slots + num_a_operands * 2) * TENSOR_MAP_QWORDS + d_desc_tma_ptr = cute.make_ptr( + cutlass.Int64, + d_desc_base.toint(), + mem_space=cute.AddressSpace.generic, + ) + previous_group_end = cutlass.Int32(-1) + if warp_idx == 0: + if elect_one: + _copy_tensormap_to_workspace(tma_c_desc.get_ptr(), tma_c_desc_smem) + nvvm.bar_warp_sync(0xFFFFFFFF) + # @@TMA_STORE_ONLY:END@@ + while is_valid != 0: while not nvvm.mbarrier_try_wait_parity( sched_full_mbar_ptr.subview(sched_stage), @@ -1046,6 +1134,22 @@ def _kernel( if is_valid != 0: coord_m_tile = group_begin + tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] + # @@TMA_STORE_ONLY:BEGIN@@ + # Re-dimension D to this group's last row so the hardware clips the + # ragged tail; the base stays put, so the store coords are global. + if warp_idx == 0: + if group_end != previous_group_end: + previous_group_end = group_end + nvvm.cp_async_bulk_wait_group(0, read=True) + _fence_tensormap_acquire(d_desc_tma_ptr) + if elect_one: + _replace_tensormap_global_dim_1(tma_c_desc_smem, group_end) + nvvm.bar_warp_sync(0xFFFFFFFF) + if lane < TENSOR_MAP_QWORDS: + (d_desc_base + lane).store((tma_c_desc_smem.subview(lane)).load()) + nvvm.bar_warp_sync(0xFFFFFFFF) + _fence_tensormap_release() + # @@TMA_STORE_ONLY:END@@ # @@EPILOGUE_DRAIN:BEGIN@@ coord_n_c = tile_n * cgrp_tile_mnk[1] + n_rank * pair_n_size @@ -1115,6 +1219,43 @@ def _kernel( 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 * subtile_w) + + vec_f32 = c_rmem_vec + col_j = col + linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 + + # @@INJECT_EPILOGUE@@ + + 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( + barrier_id=EPI_SYNC_BAR_ID, + thread_count=num_epilogue_warps * 32, + ) + + if warp_idx == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_global_shared_cta( + d_desc_tma_ptr, + smem_subtile_ptr, + (col, coord_m, tile_l), + ) + if elect_one: + nvvm.cp_async_bulk_commit_group() + nvvm.cp_async_bulk_wait_group(EPI_SMEM_STAGES - 1, read=True) + + nvvm.barrier_cta_sync( + barrier_id=EPI_SYNC_BAR_ID, + thread_count=num_epilogue_warps * 32, + ) + # @@TMA_STORE_ONLY:END@@ + + # @@STG_ONLY:BEGIN@@ if row_active and row < group_end: for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize @@ -1124,6 +1265,7 @@ 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@@ @@ -1147,6 +1289,9 @@ def _host( # @@INJECT_HOST_AB_PARAMS@@ # @@INJECT_HOST_TAP_PARAMS@@ # @@INJECT_HOST_AUX_PARAMS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_PARAMS@@ + # @@TMA_STORE_ONLY:END@@ stream: _cuda.CUstream, ) -> None: # @@INJECT_HOST_AB_LISTS@@ @@ -1180,6 +1325,23 @@ def _host( # @@INJECT_HOST_REDUCTION_STRIDES@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_LISTS@@ + c = _tma_c_outputs[0] + tma_c_desc = _tma.create_tensor_map_tiled( + global_address=c.iterator.toint(), + dtype=cd_tma_dtype, + global_dims=[n, m, 1], + global_strides=[ + out_stride_m_0 * cd_dtype.width // 128, + out_stride_l_0 * cd_dtype.width // 128, + ], + box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], + swizzle=epi_tma_swizzle, + ) + tma_c_desc_list = [tma_c_desc] + # @@TMA_STORE_ONLY:END@@ + tma_a_desc_list = [] for _a_idx, _a_op in enumerate(_a_operands): a_stride_m, a_stride_k, a_stride_l = _a_stride_sets[_a_idx] @@ -1293,9 +1455,13 @@ def _host( a_tma_workspace, # @@INJECT_HOST_KERNEL_DESC_PASS@@ # @@INJECT_MOE_HOST_MA_PASS@@ + # @@INJECT_MOE_HOST_MSFA_PASS@@ # @@INJECT_HOST_TAP_PASS@@ # @@INJECT_HOST_REDUCTION_STRIDE_PASS@@ # @@INJECT_HOST_AUX_PASS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_PASS@@ + # @@TMA_STORE_ONLY:END@@ ).launch( grid=grid_shape, block=(threads_per_cta, 1, 1), @@ -1366,7 +1532,7 @@ def _make_fake_sfb(): grid_ctas = grid_num_clusters * cluster_m * cluster_n fake_a_tma_workspace = make_fake_compact_tensor( cutlass.Int64, - (grid_ctas * num_a_operands * 16,), + (grid_ctas * moe_desc_slots * 16,), stride_order=(0,), assumed_align=128, ) @@ -1389,6 +1555,18 @@ def _sym_operand_strides(is_mn_major: bool) -> tuple: # @@INJECT_COMPILE_TAP_FAKES@@ + # @@TMA_STORE_ONLY:BEGIN@@ + def _make_fake_c(): + return make_fake_compact_tensor( + cd_dtype, + (sym_m, sym_n // cd_fake_n_div, 1), + stride_order=(1, 0, 2), + assumed_align=16, + ) + + # @@INJECT_COMPILE_TMA_C_FAKES@@ + # @@TMA_STORE_ONLY:END@@ + problem_size = ( sym_m, sym_n, @@ -1411,6 +1589,9 @@ def _sym_operand_strides(is_mn_major: bool) -> tuple: # @@INJECT_COMPILE_AB_PASS@@ # @@INJECT_COMPILE_TAP_PASS@@ # @@INJECT_COMPILE_AUX_PASS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_COMPILE_TMA_C_PASS@@ + # @@TMA_STORE_ONLY:END@@ stream=_fake_stream, options=frost_compile_options, ) 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 7f0cf84ce..af6c0913c 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 @@ -46,6 +46,10 @@ # @@INJECT_TILE_CONSTANTS@@ +# Tensormap workspace slots per CTA: the A operands, plus the output descriptor +# when the TMA-store epilogue re-dimensions it per routed group. +moe_desc_slots = num_a_operands + (1 if use_tma_store_epi else 0) + SCHED_STAGES = 2 SCHED_SLOT_WORDS = 8 @@ -97,9 +101,16 @@ def _kernel( # @@INJECT_KERNEL_TAP_PARAMS@@ # @@INJECT_KERNEL_REDUCTION_STRIDE_PARAMS@@ # @@INJECT_KERNEL_AUX_PARAMS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_KERNEL_TMA_C_PARAMS@@ + # @@TMA_STORE_ONLY:END@@ ) -> None: # @@INJECT_AB_DESC_LISTS@@ # @@INJECT_MOE_MA_LIST@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_TMA_C_LISTS@@ + tma_c_desc = tma_c_descs[0] + # @@TMA_STORE_ONLY:END@@ mma_warp_id = 4 tma_warp_id = 5 @@ -136,6 +147,10 @@ def _kernel( for _j in cutlass.range_constexpr(num_b_operands): nvvm.prefetch_tensormap(tma_b_descs[_j].get_ptr()) + # @@TMA_STORE_ONLY:BEGIN@@ + nvvm.prefetch_tensormap(tma_c_desc.get_ptr()) + # @@TMA_STORE_ONLY:END@@ + a_pattern = 0 for n_idx in cutlass.range_constexpr(cluster_n): a_pattern = a_pattern | (1 << (n_idx * cluster_m)) @@ -205,6 +220,23 @@ def _kernel( for _ in range(num_a_operands) ] + # @@TMA_STORE_ONLY:BEGIN@@ + # One epilogue subtile = one MMA-M block x 32 cols; the M blocks reuse it. + epi_subtile_elems = epi_tile_mn[0] * epi_tile_mn[1] + smem_d_ptr = cutlass.Array( + cd_dtype, + epi_subtile_elems * EPI_SMEM_STAGES, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + tma_c_desc_smem = cutlass.Array( + cutlass.Int64, + TENSOR_MAP_QWORDS, + space=cutlass.AddressSpace.smem, + alignment=128, + ) + # @@TMA_STORE_ONLY:END@@ + if cutlass.const_expr(ab_empty_full_mask): ab_empty_count = cluster_size else: @@ -391,7 +423,7 @@ def _kernel( lane = tidx % 32 block_linear = bidx + bidy * gridx - cta_desc_base_list = [a_tma_workspace.iterator.raw_ptr() + (block_linear * num_a_operands + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands)] + cta_desc_base_list = [a_tma_workspace.iterator.raw_ptr() + (block_linear * moe_desc_slots + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands)] a_desc_tma_ptr_list = [ cute.make_ptr( cutlass.Int64, @@ -728,6 +760,24 @@ def _kernel( lane = tidx % 32 # @@EPILOGUE_SETUP:END@@ + # @@TMA_STORE_ONLY:BEGIN@@ + epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) + # The routed output is a single (1, S, N) tensor, so the batch coord is fixed. + tile_l = cutlass.Int32(0) + epi_block_linear = bidx + bidy * gridx + d_desc_base = a_tma_workspace.iterator.raw_ptr() + (epi_block_linear * moe_desc_slots + num_a_operands) * TENSOR_MAP_QWORDS + d_desc_tma_ptr = cute.make_ptr( + cutlass.Int64, + d_desc_base.toint(), + mem_space=cute.AddressSpace.generic, + ) + previous_group_end = cutlass.Int32(-1) + if warp_idx == 0: + if elect_one: + _copy_tensormap_to_workspace(tma_c_desc.get_ptr(), tma_c_desc_smem) + nvvm.bar_warp_sync(0xFFFFFFFF) + # @@TMA_STORE_ONLY:END@@ + while not nvvm.mbarrier_try_wait_parity(sched_full_mbar_ptr.subview(sched_stage), sched_full_phase, time_limit=10_000_000): pass _slot = sched_storage.subview(sched_stage * SCHED_SLOT_WORDS) @@ -746,6 +796,22 @@ def _kernel( while is_valid != 0: coord_m_tile = group_begin + tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] + # @@TMA_STORE_ONLY:BEGIN@@ + # Re-dimension D to this group's last row so the hardware clips the + # ragged tail; the base stays put, so the store coords are global. + if warp_idx == 0: + if group_end != previous_group_end: + previous_group_end = group_end + nvvm.cp_async_bulk_wait_group(0, read=True) + _fence_tensormap_acquire(d_desc_tma_ptr) + if elect_one: + _replace_tensormap_global_dim_1(tma_c_desc_smem, group_end) + nvvm.bar_warp_sync(0xFFFFFFFF) + if lane < TENSOR_MAP_QWORDS: + (d_desc_base + lane).store((tma_c_desc_smem.subview(lane)).load()) + nvvm.bar_warp_sync(0xFFFFFFFF) + _fence_tensormap_release() + # @@TMA_STORE_ONLY:END@@ # @@EPILOGUE_DRAIN:BEGIN@@ coord_n_c = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] @@ -795,6 +861,43 @@ def _kernel( 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 * subtile_w) + + vec_f32 = c_rmem_vec + col_j = col + linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 + + # @@INJECT_EPILOGUE@@ + + 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( + barrier_id=EPI_SYNC_BAR_ID, + thread_count=num_epilogue_warps * 32, + ) + + if warp_idx == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_global_shared_cta( + d_desc_tma_ptr, + smem_subtile_ptr, + (col, coord_m, tile_l), + ) + if elect_one: + nvvm.cp_async_bulk_commit_group() + nvvm.cp_async_bulk_wait_group(EPI_SMEM_STAGES - 1, read=True) + + nvvm.barrier_cta_sync( + barrier_id=EPI_SYNC_BAR_ID, + thread_count=num_epilogue_warps * 32, + ) + # @@TMA_STORE_ONLY:END@@ + + # @@STG_ONLY:BEGIN@@ if row_active and row < group_end: for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize @@ -804,6 +907,7 @@ 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@@ @@ -841,6 +945,9 @@ def _host( # @@INJECT_HOST_AB_PARAMS@@ # @@INJECT_HOST_TAP_PARAMS@@ # @@INJECT_HOST_AUX_PARAMS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_PARAMS@@ + # @@TMA_STORE_ONLY:END@@ stream: _cuda.CUstream, ) -> None: # @@INJECT_HOST_AB_LISTS@@ @@ -872,6 +979,23 @@ def _host( _stride_idx += 3 # @@INJECT_HOST_REDUCTION_STRIDES@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_LISTS@@ + c = _tma_c_outputs[0] + tma_c_desc = _tma.create_tensor_map_tiled( + global_address=c.iterator.toint(), + dtype=cd_tma_dtype, + global_dims=[n, m, 1], + global_strides=[ + out_stride_m_0 * cd_dtype.width // 128, + out_stride_l_0 * cd_dtype.width // 128, + ], + box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], + swizzle=epi_tma_swizzle, + ) + tma_c_desc_list = [tma_c_desc] + # @@TMA_STORE_ONLY:END@@ + tma_a_desc_list = [] for _a_idx, _a_op in enumerate(_a_operands): a_stride_m, a_stride_k, a_stride_l = _a_stride_sets[_a_idx] @@ -936,6 +1060,9 @@ def _host( # @@INJECT_HOST_TAP_PASS@@ # @@INJECT_HOST_REDUCTION_STRIDE_PASS@@ # @@INJECT_HOST_AUX_PASS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_PASS@@ + # @@TMA_STORE_ONLY:END@@ ).launch( grid=grid_shape, block=(threads_per_cta, 1, 1), @@ -985,7 +1112,7 @@ def _make_fake_b(): grid_ctas = grid_num_clusters * cluster_m * cluster_n fake_a_tma_workspace = make_fake_compact_tensor( cutlass.Int64, - (grid_ctas * num_a_operands * 16,), + (grid_ctas * moe_desc_slots * 16,), stride_order=(0,), assumed_align=128, ) @@ -1004,6 +1131,18 @@ def _sym_operand_strides(is_mn_major: bool) -> tuple: # @@INJECT_COMPILE_REDUCTION_STRIDE_DECLS@@ # @@INJECT_COMPILE_AB_FAKES@@ # @@INJECT_COMPILE_TAP_FAKES@@ + + # @@TMA_STORE_ONLY:BEGIN@@ + def _make_fake_c(): + return make_fake_compact_tensor( + cd_dtype, + (sym_m, sym_n // cd_fake_n_div, 1), + stride_order=(1, 0, 2), + assumed_align=16, + ) + + # @@INJECT_COMPILE_TMA_C_FAKES@@ + # @@TMA_STORE_ONLY:END@@ problem_size = ( sym_m, sym_n, @@ -1024,6 +1163,9 @@ def _sym_operand_strides(is_mn_major: bool) -> tuple: # @@INJECT_COMPILE_AB_PASS@@ # @@INJECT_COMPILE_TAP_PASS@@ # @@INJECT_COMPILE_AUX_PASS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_COMPILE_TMA_C_PASS@@ + # @@TMA_STORE_ONLY:END@@ stream=_fake_stream, options=frost_compile_options, ) 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 175143dc3..87fec235d 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 @@ -48,6 +48,10 @@ # A TMA tensormap is 128 bytes = 16 int64 qwords. # @@INJECT_TILE_CONSTANTS@@ +# Tensormap workspace slots per CTA: the A operands, plus the output descriptor +# when the TMA-store epilogue re-dimensions it per routed group. +moe_desc_slots = num_a_operands + (1 if use_tma_store_epi else 0) + # Scheduler ring depth. SCHED_STAGES = 2 @@ -102,9 +106,16 @@ def _kernel( # @@INJECT_KERNEL_TAP_PARAMS@@ # @@INJECT_KERNEL_REDUCTION_STRIDE_PARAMS@@ # @@INJECT_KERNEL_AUX_PARAMS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_KERNEL_TMA_C_PARAMS@@ + # @@TMA_STORE_ONLY:END@@ ) -> None: # @@INJECT_AB_DESC_LISTS@@ # @@INJECT_MOE_MA_LIST@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_TMA_C_LISTS@@ + tma_c_desc = tma_c_descs[0] + # @@TMA_STORE_ONLY:END@@ mma_warp_id = 4 tma_warp_id = 5 @@ -145,6 +156,10 @@ def _kernel( for _j in cutlass.range_constexpr(num_b_operands): nvvm.prefetch_tensormap(tma_b_descs[_j].get_ptr()) + # @@TMA_STORE_ONLY:BEGIN@@ + nvvm.prefetch_tensormap(tma_c_desc.get_ptr()) + # @@TMA_STORE_ONLY:END@@ + a_pattern = 0 for n_idx in cutlass.range_constexpr(cluster_n): a_pattern = a_pattern | (1 << (n_idx * cluster_m)) @@ -210,6 +225,23 @@ def _kernel( for _ in range(num_a_operands) ] + # @@TMA_STORE_ONLY:BEGIN@@ + # One epilogue subtile = one MMA-M block x 32 cols; the M blocks reuse it. + epi_subtile_elems = epi_tile_mn[0] * epi_tile_mn[1] + smem_d_ptr = cutlass.Array( + cd_dtype, + epi_subtile_elems * EPI_SMEM_STAGES, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + tma_c_desc_smem = cutlass.Array( + cutlass.Int64, + TENSOR_MAP_QWORDS, + space=cutlass.AddressSpace.smem, + alignment=128, + ) + # @@TMA_STORE_ONLY:END@@ + acc_empty_count = num_epilogue_warps * 2 cta_group = 2 if cutlass.const_expr(ab_empty_full_mask): @@ -405,7 +437,7 @@ def _kernel( lane = tidx % 32 block_linear = bidx + bidy * gridx - cta_desc_base_list = [a_tma_workspace.iterator.raw_ptr() + (block_linear * num_a_operands + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands)] + cta_desc_base_list = [a_tma_workspace.iterator.raw_ptr() + (block_linear * moe_desc_slots + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands)] a_desc_tma_ptr_list = [ cute.make_ptr( cutlass.Int64, @@ -786,6 +818,24 @@ def _kernel( lane = tidx % 32 # @@EPILOGUE_SETUP:END@@ + # @@TMA_STORE_ONLY:BEGIN@@ + epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) + # The routed output is a single (1, S, N) tensor, so the batch coord is fixed. + tile_l = cutlass.Int32(0) + epi_block_linear = bidx + bidy * gridx + d_desc_base = a_tma_workspace.iterator.raw_ptr() + (epi_block_linear * moe_desc_slots + num_a_operands) * TENSOR_MAP_QWORDS + d_desc_tma_ptr = cute.make_ptr( + cutlass.Int64, + d_desc_base.toint(), + mem_space=cute.AddressSpace.generic, + ) + previous_group_end = cutlass.Int32(-1) + if warp_idx == 0: + if elect_one: + _copy_tensormap_to_workspace(tma_c_desc.get_ptr(), tma_c_desc_smem) + nvvm.bar_warp_sync(0xFFFFFFFF) + # @@TMA_STORE_ONLY:END@@ + while is_valid != 0: while not nvvm.mbarrier_try_wait_parity( sched_full_mbar_ptr.subview(sched_stage), @@ -809,6 +859,22 @@ def _kernel( if is_valid != 0: coord_m_tile = group_begin + tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] + # @@TMA_STORE_ONLY:BEGIN@@ + # Re-dimension D to this group's last row so the hardware clips the + # ragged tail; the base stays put, so the store coords are global. + if warp_idx == 0: + if group_end != previous_group_end: + previous_group_end = group_end + nvvm.cp_async_bulk_wait_group(0, read=True) + _fence_tensormap_acquire(d_desc_tma_ptr) + if elect_one: + _replace_tensormap_global_dim_1(tma_c_desc_smem, group_end) + nvvm.bar_warp_sync(0xFFFFFFFF) + if lane < TENSOR_MAP_QWORDS: + (d_desc_base + lane).store((tma_c_desc_smem.subview(lane)).load()) + nvvm.bar_warp_sync(0xFFFFFFFF) + _fence_tensormap_release() + # @@TMA_STORE_ONLY:END@@ # @@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): @@ -864,6 +930,43 @@ def _kernel( 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 * subtile_w) + + vec_f32 = c_rmem_vec + col_j = col + linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 + + # @@INJECT_EPILOGUE@@ + + 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( + barrier_id=EPI_SYNC_BAR_ID, + thread_count=num_epilogue_warps * 32, + ) + + if warp_idx == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_global_shared_cta( + d_desc_tma_ptr, + smem_subtile_ptr, + (col, coord_m, tile_l), + ) + if elect_one: + nvvm.cp_async_bulk_commit_group() + nvvm.cp_async_bulk_wait_group(EPI_SMEM_STAGES - 1, read=True) + + nvvm.barrier_cta_sync( + barrier_id=EPI_SYNC_BAR_ID, + thread_count=num_epilogue_warps * 32, + ) + # @@TMA_STORE_ONLY:END@@ + + # @@STG_ONLY:BEGIN@@ if row_active and row < group_end: for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize @@ -873,6 +976,7 @@ 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@@ @@ -890,6 +994,9 @@ def _host( # @@INJECT_HOST_AB_PARAMS@@ # @@INJECT_HOST_TAP_PARAMS@@ # @@INJECT_HOST_AUX_PARAMS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_PARAMS@@ + # @@TMA_STORE_ONLY:END@@ stream: _cuda.CUstream, ) -> None: # @@INJECT_HOST_AB_LISTS@@ @@ -923,6 +1030,23 @@ def _host( # @@INJECT_HOST_REDUCTION_STRIDES@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_LISTS@@ + c = _tma_c_outputs[0] + tma_c_desc = _tma.create_tensor_map_tiled( + global_address=c.iterator.toint(), + dtype=cd_tma_dtype, + global_dims=[n, m, 1], + global_strides=[ + out_stride_m_0 * cd_dtype.width // 128, + out_stride_l_0 * cd_dtype.width // 128, + ], + box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], + swizzle=epi_tma_swizzle, + ) + tma_c_desc_list = [tma_c_desc] + # @@TMA_STORE_ONLY:END@@ + tma_a_desc_list = [] for _a_idx, _a_op in enumerate(_a_operands): a_stride_m, a_stride_k, a_stride_l = _a_stride_sets[_a_idx] @@ -987,6 +1111,9 @@ def _host( # @@INJECT_HOST_TAP_PASS@@ # @@INJECT_HOST_REDUCTION_STRIDE_PASS@@ # @@INJECT_HOST_AUX_PASS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_PASS@@ + # @@TMA_STORE_ONLY:END@@ ).launch( grid=grid_shape, block=(threads_per_cta, 1, 1), @@ -1036,7 +1163,7 @@ def _make_fake_b(): grid_ctas = grid_num_clusters * cluster_m * cluster_n fake_a_tma_workspace = make_fake_compact_tensor( cutlass.Int64, - (grid_ctas * num_a_operands * 16,), + (grid_ctas * moe_desc_slots * 16,), stride_order=(0,), assumed_align=128, ) @@ -1055,6 +1182,18 @@ def _sym_operand_strides(is_mn_major: bool) -> tuple: # @@INJECT_COMPILE_REDUCTION_STRIDE_DECLS@@ # @@INJECT_COMPILE_AB_FAKES@@ # @@INJECT_COMPILE_TAP_FAKES@@ + + # @@TMA_STORE_ONLY:BEGIN@@ + def _make_fake_c(): + return make_fake_compact_tensor( + cd_dtype, + (sym_m, sym_n // cd_fake_n_div, 1), + stride_order=(1, 0, 2), + assumed_align=16, + ) + + # @@INJECT_COMPILE_TMA_C_FAKES@@ + # @@TMA_STORE_ONLY:END@@ problem_size = ( sym_m, sym_n, @@ -1075,6 +1214,9 @@ def _sym_operand_strides(is_mn_major: bool) -> tuple: # @@INJECT_COMPILE_AB_PASS@@ # @@INJECT_COMPILE_TAP_PASS@@ # @@INJECT_COMPILE_AUX_PASS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_COMPILE_TMA_C_PASS@@ + # @@TMA_STORE_ONLY:END@@ stream=_fake_stream, options=frost_compile_options, ) 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 3f8657f06..847de1c4b 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 @@ -46,6 +46,7 @@ moe_swizzle_tile as _moe_swizzle_tile, replace_tensormap_global_address as _replace_tensormap_global_address, replace_tensormap_global_dim_1 as _replace_tensormap_global_dim_1, + replace_tensormap_global_dim_2 as _replace_tensormap_global_dim_2, tcgen05_alloc as _tcgen05_alloc, tcgen05_dealloc as _tcgen05_dealloc, tcgen05_mma_block_scale as _tcgen05_mma_block_scale, @@ -63,6 +64,10 @@ # it to the per-CTA GMEM workspace the TMA reads. # @@INJECT_TILE_CONSTANTS@@ +# Tensormap workspace slots per CTA: the A operands, plus the output descriptor +# when the TMA-store epilogue re-dimensions it per routed group. +moe_desc_slots = num_a_operands * 2 + (1 if use_tma_store_epi else 0) + 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") @@ -121,12 +126,21 @@ def _kernel( a_tma_workspace: cute.Tensor, # @@INJECT_KERNEL_AB_DESC_PARAMS@@ # @@INJECT_MOE_KERNEL_MA_PARAMS@@ + # @@INJECT_MOE_KERNEL_MSFA_PARAMS@@ # @@INJECT_KERNEL_TAP_PARAMS@@ # @@INJECT_KERNEL_REDUCTION_STRIDE_PARAMS@@ # @@INJECT_KERNEL_AUX_PARAMS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_KERNEL_TMA_C_PARAMS@@ + # @@TMA_STORE_ONLY:END@@ ) -> None: # @@INJECT_AB_DESC_LISTS@@ # @@INJECT_MOE_MA_LIST@@ + # @@INJECT_MOE_MSFA_LIST@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_TMA_C_LISTS@@ + tma_c_desc = tma_c_descs[0] + # @@TMA_STORE_ONLY:END@@ mma_warp_id = 4 tma_warp_id = 5 @@ -163,6 +177,10 @@ def _kernel( nvvm.prefetch_tensormap(tma_b_descs[_j].get_ptr()) nvvm.prefetch_tensormap(tma_sfb_descs[_j].get_ptr()) + # @@TMA_STORE_ONLY:BEGIN@@ + nvvm.prefetch_tensormap(tma_c_desc.get_ptr()) + # @@TMA_STORE_ONLY:END@@ + cluster_linear_init = bidx // cluster_m a_pattern = 0 @@ -214,6 +232,32 @@ def _kernel( ) for _ in range(num_a_operands) ] + tma_sfa_desc_smem_list = [ + cutlass.Array( + cutlass.Int64, + TENSOR_MAP_QWORDS, + space=cutlass.AddressSpace.smem, + alignment=128, + ) + for _ in range(num_a_operands) + ] + + # @@TMA_STORE_ONLY:BEGIN@@ + # One epilogue subtile = one MMA-M block x 32 cols; the M blocks reuse it. + epi_subtile_elems = epi_tile_mn[0] * epi_tile_mn[1] + smem_d_ptr = cutlass.Array( + cd_dtype, + epi_subtile_elems * EPI_SMEM_STAGES, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + tma_c_desc_smem = cutlass.Array( + cutlass.Int64, + TENSOR_MAP_QWORDS, + space=cutlass.AddressSpace.smem, + alignment=128, + ) + # @@TMA_STORE_ONLY:END@@ sA_elems = sA_packed_elems sB_elems = sB_packed_elems @@ -534,7 +578,7 @@ def _kernel( lane = tidx % 32 block_linear = bidx + bidy * gridx - cta_desc_base_list = [a_tma_workspace.iterator.raw_ptr() + (block_linear * num_a_operands + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands)] + cta_desc_base_list = [a_tma_workspace.iterator.raw_ptr() + (block_linear * moe_desc_slots + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands)] a_desc_tma_ptr_list = [ cute.make_ptr( cutlass.Int64, @@ -543,10 +587,24 @@ def _kernel( ) for _ai in range(num_a_operands) ] + sfa_desc_base_list = [ + a_tma_workspace.iterator.raw_ptr() + (block_linear * moe_desc_slots + num_a_operands + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands) + ] + sfa_desc_tma_ptr_list = [ + cute.make_ptr( + cutlass.Int64, + sfa_desc_base_list[_ai].toint(), + mem_space=cute.AddressSpace.generic, + ) + for _ai in range(num_a_operands) + ] + sfa_block_bytes = 512 * (((k // block_size) + 3) // 4) previous_group_begin = cutlass.Int32(-1) if elect_one: for _ai in cutlass.range_constexpr(num_a_operands): _copy_tensormap_to_workspace(tma_a_descs[_ai].get_ptr(), tma_a_desc_smem_list[_ai]) + for _ai in cutlass.range_constexpr(num_a_operands): + _copy_tensormap_to_workspace(tma_sfa_descs[_ai].get_ptr(), tma_sfa_desc_smem_list[_ai]) nvvm.bar_warp_sync(0xFFFFFFFF) while is_valid != 0: @@ -574,7 +632,7 @@ def _kernel( if is_valid != 0: coord_m_group = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] - sfa_m_block = start_sf_block_m + coord_m_group // 128 + sfa_m_block = coord_m_group // 128 sfb_n_block = coord_n_per_cta // 128 if group_begin != previous_group_begin: @@ -591,6 +649,18 @@ def _kernel( (cta_desc_base_list[_ai] + lane).store((tma_a_desc_smem_list[_ai].subview(lane)).load()) nvvm.bar_warp_sync(0xFFFFFFFF) _fence_tensormap_release() + for _ai in cutlass.range_constexpr(num_a_operands): + _fence_tensormap_acquire(sfa_desc_tma_ptr_list[_ai]) + for _ai in cutlass.range_constexpr(num_a_operands): + if elect_one: + sfa_base = mSFA_list[_ai].iterator.raw_ptr().toint() + start_sf_block_m * sfa_block_bytes + _replace_tensormap_global_address(tma_sfa_desc_smem_list[_ai], sfa_base) + _replace_tensormap_global_dim_2(tma_sfa_desc_smem_list[_ai], cute.ceil_div(group_end - group_begin, 128)) + nvvm.bar_warp_sync(0xFFFFFFFF) + if lane < TENSOR_MAP_QWORDS: + (sfa_desc_base_list[_ai] + lane).store((tma_sfa_desc_smem_list[_ai].subview(lane)).load()) + nvvm.bar_warp_sync(0xFFFFFFFF) + _fence_tensormap_release() for k_tile_idx in range(num_k_tiles): stage = ab_iter % ab_stages @@ -629,7 +699,7 @@ def _kernel( 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(), + sfa_desc_tma_ptr_list[_ai], (0, coord_sf_k, sfa_m_block, cutlass.Int32(0)), sf_full_mbar_ptr.subview(stage), [], @@ -1022,6 +1092,24 @@ def _kernel( lane = tidx % 32 # @@EPILOGUE_SETUP:END@@ + # @@TMA_STORE_ONLY:BEGIN@@ + epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) + # The routed output is a single (1, S, N) tensor, so the batch coord is fixed. + tile_l = cutlass.Int32(0) + epi_block_linear = bidx + bidy * gridx + d_desc_base = a_tma_workspace.iterator.raw_ptr() + (epi_block_linear * moe_desc_slots + num_a_operands * 2) * TENSOR_MAP_QWORDS + d_desc_tma_ptr = cute.make_ptr( + cutlass.Int64, + d_desc_base.toint(), + mem_space=cute.AddressSpace.generic, + ) + previous_group_end = cutlass.Int32(-1) + if warp_idx == 0: + if elect_one: + _copy_tensormap_to_workspace(tma_c_desc.get_ptr(), tma_c_desc_smem) + nvvm.bar_warp_sync(0xFFFFFFFF) + # @@TMA_STORE_ONLY:END@@ + while is_valid != 0: while not nvvm.mbarrier_try_wait_parity( sched_full_mbar_ptr.subview(sched_stage), @@ -1045,6 +1133,22 @@ def _kernel( if is_valid != 0: coord_m_tile = group_begin + tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] + # @@TMA_STORE_ONLY:BEGIN@@ + # Re-dimension D to this group's last row so the hardware clips the + # ragged tail; the base stays put, so the store coords are global. + if warp_idx == 0: + if group_end != previous_group_end: + previous_group_end = group_end + nvvm.cp_async_bulk_wait_group(0, read=True) + _fence_tensormap_acquire(d_desc_tma_ptr) + if elect_one: + _replace_tensormap_global_dim_1(tma_c_desc_smem, group_end) + nvvm.bar_warp_sync(0xFFFFFFFF) + if lane < TENSOR_MAP_QWORDS: + (d_desc_base + lane).store((tma_c_desc_smem.subview(lane)).load()) + nvvm.bar_warp_sync(0xFFFFFFFF) + _fence_tensormap_release() + # @@TMA_STORE_ONLY:END@@ # @@EPILOGUE_DRAIN:BEGIN@@ coord_n_c = tile_n * cgrp_tile_mnk[1] + n_rank * cta_tile_mnk[1] @@ -1106,6 +1210,43 @@ def _kernel( 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 * subtile_w) + + vec_f32 = c_rmem_vec + col_j = col + linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 + + # @@INJECT_EPILOGUE@@ + + 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( + barrier_id=EPI_SYNC_BAR_ID, + thread_count=num_epilogue_warps * 32, + ) + + if warp_idx == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_global_shared_cta( + d_desc_tma_ptr, + smem_subtile_ptr, + (col, coord_m, tile_l), + ) + if elect_one: + nvvm.cp_async_bulk_commit_group() + nvvm.cp_async_bulk_wait_group(EPI_SMEM_STAGES - 1, read=True) + + nvvm.barrier_cta_sync( + barrier_id=EPI_SYNC_BAR_ID, + thread_count=num_epilogue_warps * 32, + ) + # @@TMA_STORE_ONLY:END@@ + + # @@STG_ONLY:BEGIN@@ if row_active and row < group_end: for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize @@ -1115,6 +1256,7 @@ 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@@ @@ -1138,6 +1280,9 @@ def _host( # @@INJECT_HOST_AB_PARAMS@@ # @@INJECT_HOST_TAP_PARAMS@@ # @@INJECT_HOST_AUX_PARAMS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_PARAMS@@ + # @@TMA_STORE_ONLY:END@@ stream: _cuda.CUstream, ) -> None: # @@INJECT_HOST_AB_LISTS@@ @@ -1169,6 +1314,23 @@ def _host( _stride_idx += 3 # @@INJECT_HOST_REDUCTION_STRIDES@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_LISTS@@ + c = _tma_c_outputs[0] + tma_c_desc = _tma.create_tensor_map_tiled( + global_address=c.iterator.toint(), + dtype=cd_tma_dtype, + global_dims=[n, m, 1], + global_strides=[ + out_stride_m_0 * cd_dtype.width // 128, + out_stride_l_0 * cd_dtype.width // 128, + ], + box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], + swizzle=epi_tma_swizzle, + ) + tma_c_desc_list = [tma_c_desc] + # @@TMA_STORE_ONLY:END@@ + tma_a_desc_list = [] for _a_idx, _a_op in enumerate(_a_operands): a_stride_m, a_stride_k, a_stride_l = _a_stride_sets[_a_idx] @@ -1282,9 +1444,13 @@ def _host( a_tma_workspace, # @@INJECT_HOST_KERNEL_DESC_PASS@@ # @@INJECT_MOE_HOST_MA_PASS@@ + # @@INJECT_MOE_HOST_MSFA_PASS@@ # @@INJECT_HOST_TAP_PASS@@ # @@INJECT_HOST_REDUCTION_STRIDE_PASS@@ # @@INJECT_HOST_AUX_PASS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_PASS@@ + # @@TMA_STORE_ONLY:END@@ ).launch( grid=grid_shape, block=(threads_per_cta, 1, 1), @@ -1355,7 +1521,7 @@ def _make_fake_sfb(): grid_ctas = grid_num_clusters * cluster_m * cluster_n fake_a_tma_workspace = make_fake_compact_tensor( cutlass.Int64, - (grid_ctas * num_a_operands * 16,), + (grid_ctas * moe_desc_slots * 16,), stride_order=(0,), assumed_align=128, ) @@ -1374,6 +1540,18 @@ def _sym_operand_strides(is_mn_major: bool) -> tuple: # @@INJECT_COMPILE_REDUCTION_STRIDE_DECLS@@ # @@INJECT_COMPILE_AB_FAKES@@ # @@INJECT_COMPILE_TAP_FAKES@@ + + # @@TMA_STORE_ONLY:BEGIN@@ + def _make_fake_c(): + return make_fake_compact_tensor( + cd_dtype, + (sym_m, sym_n // cd_fake_n_div, 1), + stride_order=(1, 0, 2), + assumed_align=16, + ) + + # @@INJECT_COMPILE_TMA_C_FAKES@@ + # @@TMA_STORE_ONLY:END@@ problem_size = ( sym_m, sym_n, @@ -1394,6 +1572,9 @@ def _sym_operand_strides(is_mn_major: bool) -> tuple: # @@INJECT_COMPILE_AB_PASS@@ # @@INJECT_COMPILE_TAP_PASS@@ # @@INJECT_COMPILE_AUX_PASS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_COMPILE_TMA_C_PASS@@ + # @@TMA_STORE_ONLY:END@@ stream=_fake_stream, options=frost_compile_options, ) 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 1d4034a73..62273921a 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 @@ -51,6 +51,7 @@ moe_swizzle_tile as _moe_swizzle_tile, replace_tensormap_global_address as _replace_tensormap_global_address, replace_tensormap_global_dim_1 as _replace_tensormap_global_dim_1, + replace_tensormap_global_dim_2 as _replace_tensormap_global_dim_2, tcgen05_alloc as _tcgen05_alloc, tcgen05_dealloc as _tcgen05_dealloc, tcgen05_mma_block_scale as _tcgen05_mma_block_scale, @@ -65,6 +66,10 @@ # @@INJECT_TILE_CONSTANTS@@ +# Tensormap workspace slots per CTA: the A operands, plus the output descriptor +# when the TMA-store epilogue re-dimensions it per routed group. +moe_desc_slots = num_a_operands * 2 + (1 if use_tma_store_epi else 0) + 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") @@ -122,13 +127,22 @@ def _kernel( a_tma_workspace: cute.Tensor, # @@INJECT_KERNEL_AB_DESC_PARAMS@@ # @@INJECT_MOE_KERNEL_MA_PARAMS@@ + # @@INJECT_MOE_KERNEL_MSFA_PARAMS@@ # @@INJECT_KERNEL_TAP_PARAMS@@ # @@INJECT_KERNEL_REDUCTION_STRIDE_PARAMS@@ # @@INJECT_KERNEL_AUX_PARAMS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_KERNEL_TMA_C_PARAMS@@ + # @@TMA_STORE_ONLY:END@@ ) -> None: # @@INJECT_AB_DESC_LISTS@@ # @@INJECT_MOE_MA_LIST@@ + # @@INJECT_MOE_MSFA_LIST@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_TMA_C_LISTS@@ + tma_c_desc = tma_c_descs[0] + # @@TMA_STORE_ONLY:END@@ mma_warp_id = 4 tma_warp_id = 5 @@ -171,6 +185,10 @@ def _kernel( nvvm.prefetch_tensormap(tma_b_descs[_j].get_ptr()) nvvm.prefetch_tensormap(tma_sfb_descs[_j].get_ptr()) + # @@TMA_STORE_ONLY:BEGIN@@ + nvvm.prefetch_tensormap(tma_c_desc.get_ptr()) + # @@TMA_STORE_ONLY:END@@ + a_pattern = 0 for n_idx in cutlass.range_constexpr(cluster_n): a_pattern = a_pattern | (1 << (n_idx * cluster_m)) @@ -214,6 +232,32 @@ def _kernel( ) for _ in range(num_a_operands) ] + tma_sfa_desc_smem_list = [ + cutlass.Array( + cutlass.Int64, + TENSOR_MAP_QWORDS, + space=cutlass.AddressSpace.smem, + alignment=128, + ) + for _ in range(num_a_operands) + ] + + # @@TMA_STORE_ONLY:BEGIN@@ + # One epilogue subtile = one MMA-M block x 32 cols; the M blocks reuse it. + epi_subtile_elems = epi_tile_mn[0] * epi_tile_mn[1] + smem_d_ptr = cutlass.Array( + cd_dtype, + epi_subtile_elems * EPI_SMEM_STAGES, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + tma_c_desc_smem = cutlass.Array( + cutlass.Int64, + TENSOR_MAP_QWORDS, + space=cutlass.AddressSpace.smem, + alignment=128, + ) + # @@TMA_STORE_ONLY:END@@ sA_elems = sA_packed_elems sB_elems = sB_packed_elems @@ -540,7 +584,7 @@ def _kernel( lane = tidx % 32 block_linear = bidx + bidy * gridx - cta_desc_base_list = [a_tma_workspace.iterator.raw_ptr() + (block_linear * num_a_operands + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands)] + cta_desc_base_list = [a_tma_workspace.iterator.raw_ptr() + (block_linear * moe_desc_slots + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands)] a_desc_tma_ptr_list = [ cute.make_ptr( cutlass.Int64, @@ -549,10 +593,24 @@ def _kernel( ) for _ai in range(num_a_operands) ] + sfa_desc_base_list = [ + a_tma_workspace.iterator.raw_ptr() + (block_linear * moe_desc_slots + num_a_operands + _ai) * TENSOR_MAP_QWORDS for _ai in range(num_a_operands) + ] + sfa_desc_tma_ptr_list = [ + cute.make_ptr( + cutlass.Int64, + sfa_desc_base_list[_ai].toint(), + mem_space=cute.AddressSpace.generic, + ) + for _ai in range(num_a_operands) + ] + sfa_block_bytes = 512 * (((k // block_size) + 3) // 4) previous_group_begin = cutlass.Int32(-1) if elect_one: for _ai in cutlass.range_constexpr(num_a_operands): _copy_tensormap_to_workspace(tma_a_descs[_ai].get_ptr(), tma_a_desc_smem_list[_ai]) + for _ai in cutlass.range_constexpr(num_a_operands): + _copy_tensormap_to_workspace(tma_sfa_descs[_ai].get_ptr(), tma_sfa_desc_smem_list[_ai]) nvvm.bar_warp_sync(0xFFFFFFFF) while is_valid != 0: @@ -581,7 +639,7 @@ def _kernel( coord_m_group = tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] coord_n_per_cta = tile_n * cgrp_tile_mnk[1] + n_rank * logical_cta_tile_n + pair_member * cta_tile_mnk[1] coord_n_pair = tile_n * cgrp_tile_mnk[1] + n_rank * logical_cta_tile_n - sfa_m_block = start_sf_block_m + coord_m_group // 128 + sfa_m_block = coord_m_group // 128 sfb_n_block = coord_n_pair // 128 if group_begin != previous_group_begin: @@ -598,6 +656,18 @@ def _kernel( (cta_desc_base_list[_ai] + lane).store((tma_a_desc_smem_list[_ai].subview(lane)).load()) nvvm.bar_warp_sync(0xFFFFFFFF) _fence_tensormap_release() + for _ai in cutlass.range_constexpr(num_a_operands): + _fence_tensormap_acquire(sfa_desc_tma_ptr_list[_ai]) + for _ai in cutlass.range_constexpr(num_a_operands): + if elect_one: + sfa_base = mSFA_list[_ai].iterator.raw_ptr().toint() + start_sf_block_m * sfa_block_bytes + _replace_tensormap_global_address(tma_sfa_desc_smem_list[_ai], sfa_base) + _replace_tensormap_global_dim_2(tma_sfa_desc_smem_list[_ai], cute.ceil_div(group_end - group_begin, 128)) + nvvm.bar_warp_sync(0xFFFFFFFF) + if lane < TENSOR_MAP_QWORDS: + (sfa_desc_base_list[_ai] + lane).store((tma_sfa_desc_smem_list[_ai].subview(lane)).load()) + nvvm.bar_warp_sync(0xFFFFFFFF) + _fence_tensormap_release() for k_tile_idx in range(num_k_tiles): stage = ab_iter % ab_stages @@ -638,7 +708,7 @@ def _kernel( 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(), + sfa_desc_tma_ptr_list[_ai], (0, coord_sf_k, sfa_m_block, cutlass.Int32(0)), sf_full_mbar_ptr.subview(stage), [], @@ -1081,6 +1151,24 @@ def _kernel( lane = tidx % 32 # @@EPILOGUE_SETUP:END@@ + # @@TMA_STORE_ONLY:BEGIN@@ + epi_stage_idx = cutlass.Int32(EPI_SMEM_STAGES - 1) + # The routed output is a single (1, S, N) tensor, so the batch coord is fixed. + tile_l = cutlass.Int32(0) + epi_block_linear = bidx + bidy * gridx + d_desc_base = a_tma_workspace.iterator.raw_ptr() + (epi_block_linear * moe_desc_slots + num_a_operands * 2) * TENSOR_MAP_QWORDS + d_desc_tma_ptr = cute.make_ptr( + cutlass.Int64, + d_desc_base.toint(), + mem_space=cute.AddressSpace.generic, + ) + previous_group_end = cutlass.Int32(-1) + if warp_idx == 0: + if elect_one: + _copy_tensormap_to_workspace(tma_c_desc.get_ptr(), tma_c_desc_smem) + nvvm.bar_warp_sync(0xFFFFFFFF) + # @@TMA_STORE_ONLY:END@@ + while is_valid != 0: while not nvvm.mbarrier_try_wait_parity( sched_full_mbar_ptr.subview(sched_stage), @@ -1104,6 +1192,22 @@ def _kernel( if is_valid != 0: coord_m_tile = group_begin + tile_m * cgrp_tile_mnk[0] + m_rank * cta_tile_mnk[0] + # @@TMA_STORE_ONLY:BEGIN@@ + # Re-dimension D to this group's last row so the hardware clips the + # ragged tail; the base stays put, so the store coords are global. + if warp_idx == 0: + if group_end != previous_group_end: + previous_group_end = group_end + nvvm.cp_async_bulk_wait_group(0, read=True) + _fence_tensormap_acquire(d_desc_tma_ptr) + if elect_one: + _replace_tensormap_global_dim_1(tma_c_desc_smem, group_end) + nvvm.bar_warp_sync(0xFFFFFFFF) + if lane < TENSOR_MAP_QWORDS: + (d_desc_base + lane).store((tma_c_desc_smem.subview(lane)).load()) + nvvm.bar_warp_sync(0xFFFFFFFF) + _fence_tensormap_release() + # @@TMA_STORE_ONLY:END@@ # @@EPILOGUE_DRAIN:BEGIN@@ coord_n_c = tile_n * cgrp_tile_mnk[1] + n_rank * pair_n_size @@ -1173,6 +1277,43 @@ def _kernel( 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 * subtile_w) + + vec_f32 = c_rmem_vec + col_j = col + linear_idx = tile_l * out_stride_l_0 + row * out_stride_m_0 + col_j * out_stride_n_0 + + # @@INJECT_EPILOGUE@@ + + 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( + barrier_id=EPI_SYNC_BAR_ID, + thread_count=num_epilogue_warps * 32, + ) + + if warp_idx == 0: + if elect_one: + nvvm.cp_async_bulk_tensor_global_shared_cta( + d_desc_tma_ptr, + smem_subtile_ptr, + (col, coord_m, tile_l), + ) + if elect_one: + nvvm.cp_async_bulk_commit_group() + nvvm.cp_async_bulk_wait_group(EPI_SMEM_STAGES - 1, read=True) + + nvvm.barrier_cta_sync( + barrier_id=EPI_SYNC_BAR_ID, + thread_count=num_epilogue_warps * 32, + ) + # @@TMA_STORE_ONLY:END@@ + + # @@STG_ONLY:BEGIN@@ if row_active and row < group_end: for j in cutlass.range_constexpr(subtile_w // vsize): col_j = col + j * vsize @@ -1182,6 +1323,7 @@ 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@@ @@ -1205,6 +1347,9 @@ def _host( # @@INJECT_HOST_AB_PARAMS@@ # @@INJECT_HOST_TAP_PARAMS@@ # @@INJECT_HOST_AUX_PARAMS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_PARAMS@@ + # @@TMA_STORE_ONLY:END@@ stream: _cuda.CUstream, ) -> None: # @@INJECT_HOST_AB_LISTS@@ @@ -1238,6 +1383,23 @@ def _host( # @@INJECT_HOST_REDUCTION_STRIDES@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_LISTS@@ + c = _tma_c_outputs[0] + tma_c_desc = _tma.create_tensor_map_tiled( + global_address=c.iterator.toint(), + dtype=cd_tma_dtype, + global_dims=[n, m, 1], + global_strides=[ + out_stride_m_0 * cd_dtype.width // 128, + out_stride_l_0 * cd_dtype.width // 128, + ], + box_dims=[epi_tile_mn[1], epi_tile_mn[0], 1], + swizzle=epi_tma_swizzle, + ) + tma_c_desc_list = [tma_c_desc] + # @@TMA_STORE_ONLY:END@@ + tma_a_desc_list = [] for _a_idx, _a_op in enumerate(_a_operands): a_stride_m, a_stride_k, a_stride_l = _a_stride_sets[_a_idx] @@ -1351,9 +1513,13 @@ def _host( a_tma_workspace, # @@INJECT_HOST_KERNEL_DESC_PASS@@ # @@INJECT_MOE_HOST_MA_PASS@@ + # @@INJECT_MOE_HOST_MSFA_PASS@@ # @@INJECT_HOST_TAP_PASS@@ # @@INJECT_HOST_REDUCTION_STRIDE_PASS@@ # @@INJECT_HOST_AUX_PASS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_HOST_TMA_C_PASS@@ + # @@TMA_STORE_ONLY:END@@ ).launch( grid=grid_shape, block=(threads_per_cta, 1, 1), @@ -1424,7 +1590,7 @@ def _make_fake_sfb(): grid_ctas = grid_num_clusters * cluster_m * cluster_n fake_a_tma_workspace = make_fake_compact_tensor( cutlass.Int64, - (grid_ctas * num_a_operands * 16,), + (grid_ctas * moe_desc_slots * 16,), stride_order=(0,), assumed_align=128, ) @@ -1447,6 +1613,18 @@ def _sym_operand_strides(is_mn_major: bool) -> tuple: # @@INJECT_COMPILE_TAP_FAKES@@ + # @@TMA_STORE_ONLY:BEGIN@@ + def _make_fake_c(): + return make_fake_compact_tensor( + cd_dtype, + (sym_m, sym_n // cd_fake_n_div, 1), + stride_order=(1, 0, 2), + assumed_align=16, + ) + + # @@INJECT_COMPILE_TMA_C_FAKES@@ + # @@TMA_STORE_ONLY:END@@ + problem_size = ( sym_m, sym_n, @@ -1469,6 +1647,9 @@ def _sym_operand_strides(is_mn_major: bool) -> tuple: # @@INJECT_COMPILE_AB_PASS@@ # @@INJECT_COMPILE_TAP_PASS@@ # @@INJECT_COMPILE_AUX_PASS@@ + # @@TMA_STORE_ONLY:BEGIN@@ + # @@INJECT_COMPILE_TMA_C_PASS@@ + # @@TMA_STORE_ONLY:END@@ stream=_fake_stream, options=frost_compile_options, )