diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index 22e2c7c901a6..ba8dca7893a3 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import ast import functools import itertools import math @@ -17,7 +18,12 @@ from ..autotuner import (AutoTuner, ConstraintSpec, DistributedTuningStrategy, DynamicTensorSpec, OptimizationProfile, TunableRunner, TuningConfig) -from ..cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from ..cute_dsl_utils import (IS_CUTLASS_DSL_AVAILABLE, + IS_CUTLASS_DSL_RUBIN_AVAILABLE) +from ..locality_domain.autotune import tune_locality_domain_concurrent +from ..locality_domain.runtime import LocalityDomainRuntime +from ..locality_domain_utils import (get_current_locality_domain, + node_local_max_active_clusters) from ..utils import (ActivationType, deep_gemm_gen_tuning_buckets, fp4_scale_infer_shape, fp8_scale_infer_shape, get_last_power_of_2_num_tokens_buckets, @@ -36,10 +42,84 @@ SWIGLU_LIMIT_SCALAR_DISABLED = -1.0 +def _with_input_cuda_device(function): + """Run a custom-op implementation under its input tensor's CUDA device.""" + + @functools.wraps(function) + def wrapped(input, *args, **kwargs): + with torch.cuda.device(input.device): + return function(input, *args, **kwargs) + + return wrapped + + +def _validate_16_byte_aligned_dense_tensor(tensor: torch.Tensor, + tensor_name: str) -> None: + """Validate the pointer/strides required by dense CuTe TMA operands.""" + if tensor.data_ptr() % 16 != 0: + raise ValueError( + f"{tensor_name} data pointer must be 16-byte aligned, got " + f"data_ptr={tensor.data_ptr()}.") + if tensor.shape[-1] > 1 and tensor.stride(-1) != 1: + raise ValueError( + f"{tensor_name} must have a contiguous innermost dimension, got " + f"shape={tuple(tensor.shape)} and stride={tuple(tensor.stride())}.") + for dim in range(tensor.dim() - 1): + if (tensor.shape[dim] > 1 + and tensor.stride(dim) * tensor.element_size() % 16 != 0): + raise ValueError( + f"{tensor_name} stride in dimension {dim} must preserve " + "16-byte alignment, got " + f"shape={tuple(tensor.shape)} and stride={tuple(tensor.stride())}." + ) + + def _canonicalize_swiglu_limit_scalar(swiglu_limit_scalar: float) -> float: return float("inf") if swiglu_limit_scalar < 0 else swiglu_limit_scalar +def _get_cute_dsl_swap_ab_candidates( + m: int, + output_aligned: bool, + include_alternative: bool = False, +) -> List[bool]: + """Return swap candidates in autotuning preference order. + + Both orientations write the same physical row-major [M, N] output. When + swapping A and B, the kernel sees an [N, M] column-major view, so its + contiguous C dimension is still the original N dimension. Therefore the + 16-byte output alignment requirement does not depend on logical M. + + Base kernels retain the existing M-based performance preference to bound + autotuning cost. Mixed-cluster callers request the alternative orientation + because cluster-grid feasibility depends on which logical axis becomes the + kernel M dimension. + """ + if not output_aligned: + return [] + if m <= 128: + swap_ab_candidates = [True] + elif m >= 256: + swap_ab_candidates = [False] + else: + swap_ab_candidates = [False, True] + if include_alternative and len(swap_ab_candidates) == 1: + swap_ab_candidates.append(not swap_ab_candidates[0]) + return swap_ab_candidates + + +def _get_sm107_nvfp4_default_mma_config( + tile_size: int +) -> Tuple[Tuple[int, int, int], Tuple[int, int, int], Tuple[int, int]]: + """Return the valid fallback MMA and cluster shapes for one routing tile.""" + mma_inst_m = min(tile_size, 256) + return ( + (tile_size, 128, 256), + (mma_inst_m, 128, 128), + (mma_inst_m // 128, 1), + ) + + class GroupedGemmInputsHelper: """Base helper class for grouped GEMM input preparation and tuning. @@ -257,11 +337,27 @@ class GatherGroupedGemmInputsHelper(GroupedGemmInputsHelper): 7: permuted_idx_to_expanded_idx - Token permutation mapping 8: num_non_exiting_tiles - Number of valid tiles 9: global_sf - Global scale factor + 10+: optional output tensors for inplace variants """ # Override: use permuted_idx_to_expanded_idx for shape inference IDX_PERMUTED_IDX_TO_EXPANDED_IDX = 7 IDX_SHAPE_INFER = IDX_PERMUTED_IDX_TO_EXPANDED_IDX + @staticmethod + def _resize_locality_domain_outputs( + m: int, + output_tensor: torch.Tensor, + output_sf_tensor: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + output_m = output_tensor.size(0) + assert output_m > 0 + assert output_sf_tensor.numel() % output_m == 0 + sf_size_per_m = output_sf_tensor.numel() // output_m + if output_m == m: + return output_tensor, output_sf_tensor + return (output_tensor.new_empty((m, output_tensor.size(1))), + output_sf_tensor.new_empty((m * sf_size_per_m, ))) + def inputs_pre_hook(self, inputs: List[torch.Tensor]) -> List[torch.Tensor]: """Pre-hook for gather-based activation fusion kernel. @@ -270,9 +366,22 @@ def inputs_pre_hook(self, inputs: List[torch.Tensor]) -> List[torch.Tensor]: - tile_idx_to_mn_limit - permuted_idx_to_expanded_idx (for gather operation) - num_non_exiting_tiles + + Input layout: + 0: a - Original input activation (not permuted) + 1: b - Weight tensor + 2: a_sf - Scale factor for a + 3: b_sf - Scale factor for b + 4: alpha - Per-expert scaling factor + 5: tile_idx_to_group_idx - Tile to expert mapping + 6: tile_idx_to_mn_limit - Tile M/N limits + 7: permuted_idx_to_expanded_idx - Token permutation mapping + 8: num_non_exiting_tiles - Number of valid tiles + 9: global_sf - Global scale factor + 10+: optional output tensors for inplace variants """ a, b, a_sf, b_sf, alpha, tile_idx_to_group_idx, tile_idx_to_mn_limit, \ - permuted_idx_to_expanded_idx, num_non_exiting_tiles, global_sf = inputs + permuted_idx_to_expanded_idx, num_non_exiting_tiles, global_sf, *others = inputs # Verify permuted_idx_to_expanded_idx index matches the class constant assert inputs[ self. @@ -304,9 +413,13 @@ def inputs_pre_hook(self, inputs: List[torch.Tensor]) -> List[torch.Tensor]: local_num_experts=self.num_local_experts, tile_tokens_dim=self.tile_size, ) + if len(others) >= 2 and others[0] is not None and others[1] is not None: + others = list(others) + others[0], others[1] = self._resize_locality_domain_outputs( + permuted_idx_to_expanded_idx.size(0), others[0], others[1]) return (a, b, a_sf, b_sf, alpha, tile_idx_to_group_idx, tile_idx_to_mn_limit, permuted_idx_to_expanded_idx, - num_non_exiting_tiles, global_sf) + num_non_exiting_tiles, global_sf, *others) def get_dense_gemm_approximate_cta_nums( @@ -356,6 +469,22 @@ def get_dense_gemm_approximate_cta_nums( SinglePassMultiCTARadixTopKClusterKernel, _query_max_cluster_size) from ..cute_dsl_kernels.blackwell.utils import make_ptr + @functools.cache + def _get_full_device_max_active_clusters(device_id: int, + cluster_size: int) -> int: + """Return the cached full-device occupancy for a cluster shape.""" + hardware_info = cutlass.utils.HardwareInfo(device_id=device_id) + return hardware_info.get_max_active_clusters(cluster_size) + + def get_max_activate_clusters(cluster_size): + max_active = _get_full_device_max_active_clusters( + torch.cuda.current_device(), cluster_size) + if get_current_locality_domain() is not None: + node_local = node_local_max_active_clusters(max_active) + max_active = node_local if node_local is not None else max( + 1, max_active // 2) + return max_active + class CuteDSLNVFP4BlackwellRunner(TunableRunner): kernel_class = Sm100BlockScaledPersistentDenseGemmKernel kernel_cache = dict() @@ -392,6 +521,14 @@ def unique_id(self): self.use_tvm_ffi, ) + def __hash__(self): + return hash(self.unique_id()) + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return False + return self.unique_id() == other.unique_id() + def get_valid_tactics( self, inputs: List[torch.Tensor], @@ -430,29 +567,25 @@ def get_valid_tactics( f"(K%32={real_k%32}, expected 0). Skipping all tactics.") return [] - # Optimize swap_ab candidates based on M and N alignment - # swap_ab=False → C is N-major → requires N%8==0 (BF16: 128 bits / 16 bits = 8) - # swap_ab=True → C is M-major → requires M%8==0 - m_aligned = (m % 8 == 0) - n_aligned = (n % 8 == 0) - - if not m_aligned and not n_aligned: + # Both swap orientations use the original N as C's physical + # contiguous dimension and require 16-byte BF16 alignment. + output_aligned = (n % 8 == 0) + if not output_aligned: logger.debug( - f"CuteDSL: Neither M={m} nor N={n} meets 16-byte alignment " - f"(M%8={m%8}, N%8={n%8}). No valid C layout. Skipping all tactics." - ) + f"CuteDSL: Output N={n} does not meet the 16-byte " + f"alignment requirement (N%8={n%8}). Skipping all tactics.") return [] - # Only test swap_ab values that satisfy alignment - swap_ab_candidates = [] - if n_aligned: - swap_ab_candidates.append(False) # N-major layout - if m_aligned: - swap_ab_candidates.append(True) # M-major layout + swap_ab_candidates = _get_cute_dsl_swap_ab_candidates( + m, output_aligned) + if not swap_ab_candidates: + logger.debug(f"CuteDSL: No valid C layout for M={m}, N={n}. " + "Skipping all tactics.") + return [] logger.debug( - f"CuteDSL: M={m}(aligned={m_aligned}), N={n}(aligned={n_aligned}), K={real_k}(aligned=True). " - f"Testing swap_ab={swap_ab_candidates}") + f"CuteDSL: M={m}, N={n}(aligned={output_aligned}), K={real_k}(aligned=True). " + f"Using swap_ab={swap_ab_candidates}") # full shamoo mma_tiler_mn_candidates = [ @@ -476,7 +609,6 @@ def get_valid_tactics( (4, 2), (4, 4), ] - swap_ab_candidates = [True, False] # prune: prefetch is beneficial only when K is large enough use_prefetch_candidates = [True, False] @@ -764,6 +896,18 @@ def _rank_prune_tactics(self, tactics, m, n, real_k): ) return tactics + def should_profile_tactic_in_subprocess( + self, + custom_op: str, + inputs: List[torch.Tensor], + tactic, + tuning_config: TuningConfig, + **kwargs, + ) -> bool: + # get_valid_tactics emits 4 fields: + # (mma_tiler_mn, cluster_shape_mn, swap_ab, use_prefetch). + return isinstance(tactic, tuple) and len(tactic) == 4 + def make_cute_dsl_global_pointer(self, tensor: torch.Tensor, dtype, assumed_align: int): return make_ptr( @@ -4278,11 +4422,17 @@ class CuteDSLFp8BlackwellBmmRunner(TunableRunner): kernel_class = Sm100BlockwiseGemmKernel kernel_cache = dict() + # Keep the output M dimension aligned with input0's bucketed M so + # profiling uses consistent BMM shapes and runtime cache keys can be + # shared by inputs that map to the same bucket. tuning_config = TuningConfig( dynamic_tensor_specs=(DynamicTensorSpec( 0, 1, get_last_power_of_2_num_tokens_buckets, last_positive_power_of_2), ), - constraint_specs=(ConstraintSpec(2, 2, fp8_scale_infer_shape), ), + constraint_specs=(ConstraintSpec(2, 2, fp8_scale_infer_shape), + ConstraintSpec( + 4, 1, + lambda input_shapes: input_shapes[0][1])), ) def __init__(self, @@ -8517,6 +8667,60 @@ def _( # BF16 Dense Persistent BMM (CuTe DSL) for Blackwell # ====================================================================== + def _bf16_preferred_cluster_has_launchable_grid( + m: int, + n: int, + batch_size: int, + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + preferred_cluster_shape_mn: Tuple[int, int], + fallback_cluster_shape_mn: Tuple[int, int], + ) -> bool: + """Return whether mixed preferred/fallback launch has any preferred cluster. + + The preferred-cluster kernel derives preferred_grid.z from the fallback + grid CTA count divided by the preferred cluster size. If the autotuner + profiles a small M bucket and that quotient is zero, CUDA rejects the + launch with cudaErrorInvalidValue. + + The per-CTA M-tile is mma_tiler_mn[0] for 1-CTA MMA but mma_tiler_mn[0]//2 + for 2-CTA MMA (the grid is built from the per-CTA tile), so the CTA-tile + count uses the halved tile when use_2cta_instrs -- otherwise the fallback + CTA count is undercounted and valid preferred-cluster tactics are pruned. + """ + cta_tile_m = mma_tiler_mn[0] // (2 if use_2cta_instrs else 1) + ctas_m = ceil_div(m, cta_tile_m) + ctas_n = ceil_div(n, mma_tiler_mn[1]) + fallback_ctas_m = pad_up(ctas_m, fallback_cluster_shape_mn[0]) + fallback_ctas_n = pad_up(ctas_n, fallback_cluster_shape_mn[1]) + fallback_ctas = fallback_ctas_m * fallback_ctas_n * batch_size + preferred_cluster_ctas = (preferred_cluster_shape_mn[0] * + preferred_cluster_shape_mn[1]) + return fallback_ctas >= preferred_cluster_ctas + + def _bf16_cluster_m_fits( + m: int, + use_2cta_instrs: bool, + mma_tiler_mn: Tuple[int, int], + cluster_shape_mn: Tuple[int, int], + ) -> bool: + """Whether the M dimension provides enough CTA-tiles for the M-cluster. + + An M-cluster wider than the available M CTA-tiles leaves phantom CTAs + whose cluster-multicast peers are never launched, which probabilistically + triggers an illegal memory access / hang during autotuner profiling + (observed on SM107 with the M=1 decode MLA absorb BMM, cluster_m=4). The + per-CTA M-tile is mma_tiler_mn[0] for 1-CTA MMA but mma_tiler_mn[0]//2 for + 2-CTA MMA (the kernel builds the grid from the per-CTA tile), so the + CTA-tile count must use the halved tile when use_2cta_instrs -- otherwise + valid 2-CTA / preferred-cluster tactics (e.g. m=128, tile_m=128, 2cta, + cluster_m=2 -> 2 real M-CTAs) get pruned. N over-padding is handled by the + kernel, so only the M axis is gated; cluster_m=4 / preferred (4,2) stay + available for large-M shapes. + """ + cta_tile_m = mma_tiler_mn[0] // (2 if use_2cta_instrs else 1) + return ceil_div(m, cta_tile_m) >= cluster_shape_mn[0] + class CuteDSLBf16BlackwellBmmRunner(TunableRunner): kernel_class = PersistentDenseGemmKernel kernel_cache = dict() @@ -10562,3 +10766,2699 @@ def _( softmax_stats: Optional[torch.Tensor], ) -> None: return None + + # ============================================================================ + # Rubin (SM107) Support + # ============================================================================ + # The following code provides CuTe DSL GEMM support for Rubin GPUs. + # Requires Rubin support in the public nvidia-cutlass-dsl package. + + if IS_CUTLASS_DSL_RUBIN_AVAILABLE: + # Rubin (SM107) MOE Grouped GEMM Support + # ==================================================================== + # The following code provides CuteDSL NVFP4 grouped GEMM kernels for + # Mixture-of-Experts (MoE) on Rubin GPUs (SM107). + # Two fused kernels are provided: + # 1. Gather + Grouped GEMM + activation fusion (FC1 layer) + # 2. Grouped GEMM + Finalize (scatter-add) fusion (FC2 layer) + + from ..cute_dsl_kernels.rubin.moe.rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion import \ + Sm107BlockScaledContiguousGatherGroupedGemmActFusionKernel + + class Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner( + TunableRunner): + """Rubin runner for gather + grouped GEMM + activation fusion. + + SM107 counterpart to Blackwell's + ``Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner``. + Supports SwiGLU and Relu2. + Key differences from Blackwell: + - Uses LDGSTS (cp.async) for A/SFA loading instead of TMA + - Supports B-reuse pattern (mma_tiler_m = 2 * mma_inst_shape_m) + - Takes mma_inst_shape and mma_tiler as 3-tuples (not 2-tuples) + """ + kernel_class = Sm107BlockScaledContiguousGatherGroupedGemmActFusionKernel + kernel_cache = dict() + tuning_config_cache = dict() + + def __init__( + self, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + scaling_vector_size: int = 16, + activation_type: ActivationType = ActivationType.Swiglu): + super().__init__() + self.num_experts = num_experts + self.top_k = top_k + self.num_local_experts = num_local_experts + self.local_expert_offset = local_expert_offset + self.tile_size = tile_size + self.scaling_vector_size = scaling_vector_size + self.activation_type = ActivationType(int(activation_type)) + if self.activation_type not in (ActivationType.Swiglu, + ActivationType.Relu2): + raise ValueError( + f"Rubin NVFP4 CuteDSL FC1 does not support " + f"{self.activation_type.name}") + self.is_gated = is_gated_activation(self.activation_type) + + if (sm_version := get_sm_version()) != 107: + raise ValueError( + f"{self.__class__.kernel_class.__name__} supports SM 107 (Rubin) only, but got SM {sm_version}" + ) + + if self.tile_size not in (128, 256, 512): + raise ValueError( + f"{self.__class__.kernel_class.__name__} supports tile_size 128, 256 and 512 only, but got {self.tile_size}" + ) + + def unique_id(self): + return ( + self.num_experts, + self.top_k, + self.num_local_experts, + self.local_expert_offset, + self.tile_size, + self.scaling_vector_size, + int(self.activation_type), + ) + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + **kwargs, + ) -> List[Tuple[int, int]]: + a, b, a_sf, b_sf, alpha, tile_idx_to_group_idx, tile_idx_to_mn_limit, permuted_idx_to_expanded_idx, *_ = inputs + # m is the permuted size from permuted_idx_to_expanded_idx, not from a + m = permuted_idx_to_expanded_idx.size(0) + k = a.size(1) * 2 + l, n = b.size(0), b.size(1) # noqa: E741 + + # Rubin 4xFP4 tile sizes: + # Without B-reuse: mma_tiler_m == mma_inst_shape_m + # - (128, 128): 1CTA + # - (256, 256): 2CTA + # With B-reuse: mma_tiler_m == 2 * mma_inst_shape_m + # - (256, 128): 1CTA, B-reuse + # - (512, 256): 2CTA, B-reuse + # Fixed K dimensions for FP4: mma_tiler_k=256, mma_inst_k=128 + mma_tiler_k = 256 + mma_inst_k = 128 + + # (mma_tiler_m, mma_inst_m) candidates + mma_m_candidates = [ + (128, 128), # no B-reuse, 1CTA + (256, 256), # no B-reuse, 2CTA + (256, 128), # B-reuse, 1CTA + (512, 256), # B-reuse, 2CTA + ] + # N dimension candidates (must match between tiler and inst) + mma_n_candidates = [128, 256] + # cluster_M is pinned to the MMA CTA-group size by the + # cta_group_m filter below (1 for 1-CTA, 2 for 2-CTA). + # cluster_N>1 multicasts A across N-CTAs (one GMEM read + # broadcast to the cluster); validated for cluster_M==atom_m. + cluster_shape_mn_candidates = [(1, 1), (1, 2), (1, 4), (2, 1), + (2, 2), (2, 4)] + raster_along_m_candidates = [False] + # A-load path: "cpasync" (per-thread LDGSTS gather) vs "tma" + # (TMA tile::gather4). Autotuned per shape. + a_path_candidates = ["cpasync", "tma"] + + valid_tactics = [] + for (mma_tiler_m, mma_inst_m), mma_n, cluster_shape_mn, \ + raster_along_m, a_path in ( + itertools.product(mma_m_candidates, mma_n_candidates, + cluster_shape_mn_candidates, + raster_along_m_candidates, + a_path_candidates)): + + # The kernel requires each cluster to cover exactly + # one routing tile along M: + # - 1-CTA (mma_inst_m=128): cluster_m must be 1 + # - 2-CTA (mma_inst_m=256): cluster_m must be 2 + # and the CTA-pair's M extent (mma_tiler_m) must equal + # the routing tile_size. Any other combination + # (mma_tiler_m != tile_size, or cluster_m spanning + # multiple independent M-tiles) produces ~40% element + # mismatches empirically. + cta_group_m = 2 if mma_inst_m == 256 else 1 + if cluster_shape_mn[0] != cta_group_m: + continue + if mma_tiler_m != self.tile_size: + continue + # The Rubin epilogue stores full output and SFC subtiles + # only, so partial GEMM N tiles are not valid. + if n % mma_n != 0: + continue + + # cluster_N>1 multicasts A across cluster-N CTAs, each + # computing a DISTINCT N-tile. This is only correct when the + # N-tiles (n // mma_n) divide evenly across the cluster; + # otherwise the multicast misaligns and the output is wrong. + # (The autotuner picks by latency with no correctness check, + # so an un-pruned bad multicast tactic could be selected in + # production — must gate it here.) + if cluster_shape_mn[1] > 1 and ( + n % (mma_n * cluster_shape_mn[1]) != 0): + continue + + # TMA gather4 is stable for per-CTA/2CTA A loads, but the + # multicast variant is not reliable on Rubin: cluster_N>1 + # requires A multicast across N-CTAs. Keep cluster_N + # tactics available through the cpasync A path only. + if a_path == "tma" and cluster_shape_mn[1] > 1: + continue + + mma_tiler = (mma_tiler_m, mma_n, mma_tiler_k) + mma_inst_shape = (mma_inst_m, mma_n, mma_inst_k) + + if self.__class__.kernel_class.can_implement( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + sf_dtype=cutlass.Float8E4M3FN, + sf_vec_size=self.scaling_vector_size, + c_dtype=cutlass.Float4E2M1FN, + mma_inst_shape=mma_inst_shape, + mma_tiler=mma_tiler, + cluster_shape_mn=cluster_shape_mn, + m=m, + n=n, + k=k, + l=l, + a_major="k", + b_major="k", + c_major="n", + ): + valid_tactics.append( + (mma_tiler, mma_inst_shape, cluster_shape_mn, + raster_along_m, a_path)) + + logger.debug( + f"CuteDSL Rubin GatherGroupedGemmSwiglu: Found {len(valid_tactics)} valid tactics " + f"for M={m}, N={n}, K={k}, L={l}") + return valid_tactics + + def get_tuning_config(self) -> TuningConfig: + key = self.unique_id() + if key not in self.__class__.tuning_config_cache: + helper = GatherGroupedGemmInputsHelper( + self.num_experts, self.top_k, self.num_local_experts, + self.local_expert_offset, self.tile_size) + self.__class__.tuning_config_cache[key] = TuningConfig( + # Use permuted_idx_to_expanded_idx (IDX_SHAPE_INFER) for tuning + dynamic_tensor_specs=(DynamicTensorSpec( + GatherGroupedGemmInputsHelper.IDX_SHAPE_INFER, 0, + helper.gen_tuning_buckets, + helper.map_to_tuning_buckets), ), + constraint_specs=( + ConstraintSpec(0, 0, helper.infer_shape_num_tokens), + ConstraintSpec(2, 0, helper.infer_shape_num_tokens), + ConstraintSpec(5, 0, + helper.infer_shape_max_num_tiles), + ConstraintSpec(6, 0, + helper.infer_shape_max_num_tiles), + ), + inputs_pre_hook=helper.inputs_pre_hook, + ) + return self.__class__.tuning_config_cache[key] + + def forward(self, inputs: List[torch.Tensor], + tactic: Optional[tuple], **kwargs) -> torch.Tensor: + a, b, a_sf, b_sf, alpha, tile_idx_to_group_idx, tile_idx_to_mn_limit, permuted_idx_to_expanded_idx, num_non_exiting_tiles, global_sf, output_tensor, output_sf_tensor = inputs + # Verify permuted_idx_to_expanded_idx index matches the class constant + assert inputs[ + GatherGroupedGemmInputsHelper. + IDX_PERMUTED_IDX_TO_EXPANDED_IDX] is permuted_idx_to_expanded_idx + assert a.dtype == torch.float4_e2m1fn_x2 + assert a.dim() == 2 + assert b.dtype == torch.float4_e2m1fn_x2 + assert b.dim() == 3 + assert a_sf.dtype == torch.uint8 + assert a_sf.dim() == 2 + assert b_sf.dtype == torch.uint8 + assert b_sf.dim() == 3 + assert alpha.dtype == torch.float32 + assert alpha.dim() == 1 + + # a.size(0) is orig_m (original input size before gather) + # permuted_idx_to_expanded_idx.size(0) is m (permuted size after gather) + orig_m, k = a.size(0), a.size(1) * 2 + m = permuted_idx_to_expanded_idx.size(0) + l, n = b.size(0), b.size(1) # noqa: E741 + scale_k = k // self.scaling_vector_size + interm_size = n // 2 if self.is_gated else n + assert m % self.tile_size == 0 + assert k % (self.scaling_vector_size * 4) == 0 + n_alignment = self.scaling_vector_size * 4 * (2 if self.is_gated + else 1) + assert n % n_alignment == 0 + assert b.size(2) * 2 == k + assert a_sf.size(0) == orig_m + assert a_sf.size(1) == scale_k + assert b_sf.size(0) == l + assert b_sf.size(1) == n + assert b_sf.size(2) == scale_k + assert alpha.size(0) == l + + num_tiles = m // self.tile_size + assert tile_idx_to_group_idx.dtype == torch.int32 + assert tile_idx_to_group_idx.size() == (num_tiles, ) + assert tile_idx_to_mn_limit.dtype == torch.int32 + assert tile_idx_to_mn_limit.size() == (num_tiles, ) + assert permuted_idx_to_expanded_idx.dtype == torch.int32 + assert permuted_idx_to_expanded_idx.size() == (m, ) + assert num_non_exiting_tiles.dtype == torch.int32 + assert num_non_exiting_tiles.numel() == 1 + assert global_sf.dtype == torch.float32 + assert global_sf.numel() == 1 + + partition_id = kwargs.get("partition_id", -1) + locality_domain_half_gemm = output_tensor is not None or output_sf_tensor is not None + if locality_domain_half_gemm: + if not self.is_gated: + raise ValueError( + "Rubin locality domain half-GEMM currently supports SwiGLU only" + ) + if partition_id < 0 or partition_id >= 2: + raise ValueError( + "partition_id must be 0 or 1 when output tensors are provided." + ) + assert output_tensor is not None and output_sf_tensor is not None + assert output_tensor.dim() == 2 + assert output_tensor.dtype == a.dtype + assert output_tensor.shape[0] == m and output_tensor.shape[ + 1] == interm_size // 2 * 2, f"[locality domain] output_tensor.shape={output_tensor.shape}, m={m}, n={n}" + assert output_sf_tensor.dim() == 1 + sf_locality_domain_total_size = m * interm_size // self.scaling_vector_size + assert output_sf_tensor.shape[ + 0] == sf_locality_domain_total_size * 2 + # c: point into shared buffer at column offset + # (kernel uses full stride via locality_domain_half_gemm + full_c_shape) + c_byte_offset = partition_id * interm_size // 2 # fp4x2 cols + c = output_tensor.view(torch.uint8)[:, c_byte_offset:].view( + torch.float4_e2m1fn_x2) + # Keep the SFC pointer at the start of the full shared + # buffer. The kernel applies the partition's N-tile offset + # in full-layout coordinates. + assert interm_size % 64 == 0 + c_sf = output_sf_tensor + c_sf_n_tile_offset_val = cutlass.Int64(partition_id * + interm_size // 64) + else: + c = torch.empty(m, + interm_size // 2, + dtype=a.dtype, + device=a.device) + c_sf = torch.empty(m * interm_size // + self.scaling_vector_size, + dtype=a_sf.dtype, + device=a_sf.device) + c_sf_n_tile_offset_val = cutlass.Int64(0) + + a_ptr = make_ptr(cutlass.Float4E2M1FN, + a.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32) + b_ptr = make_ptr(cutlass.Float4E2M1FN, + b.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32) + a_sf_ptr = make_ptr(cutlass.Float8E4M3FN, + a_sf.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + b_sf_ptr = make_ptr(cutlass.Float8E4M3FN, + b_sf.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + alpha_ptr = make_ptr(cutlass.Float32, alpha.data_ptr(), + cute.AddressSpace.gmem) + tile_idx_to_group_idx_ptr = make_ptr( + cutlass.Int32, tile_idx_to_group_idx.data_ptr(), + cute.AddressSpace.gmem) + tile_idx_to_mn_limit_ptr = make_ptr( + cutlass.Int32, tile_idx_to_mn_limit.data_ptr(), + cute.AddressSpace.gmem) + permuted_idx_to_expanded_idx_ptr = make_ptr( + cutlass.Int32, permuted_idx_to_expanded_idx.data_ptr(), + cute.AddressSpace.gmem) + num_non_exiting_tiles_ptr = make_ptr( + cutlass.Int32, num_non_exiting_tiles.data_ptr(), + cute.AddressSpace.gmem) + global_sf_ptr = make_ptr(cutlass.Float32, global_sf.data_ptr(), + cute.AddressSpace.gmem) + c_ptr = make_ptr(cutlass.Float4E2M1FN, + c.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32) + c_sf_ptr = make_ptr(cutlass.Float8E4M3FN, + c_sf.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + + torch_stream = torch.cuda.current_stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + if isinstance(tactic, tuple): + mma_tiler, mma_inst_shape, cluster_shape_mn, raster_along_m, a_path = tactic + else: + # Default tactic for Rubin + mma_tiler, mma_inst_shape, cluster_shape_mn = \ + _get_sm107_nvfp4_default_mma_config(self.tile_size) + raster_along_m = False + a_path = "cpasync" + assert mma_tiler[ + 0] >= self.tile_size, f"Tactic ({tactic}) is incompatible with tile size ({self.tile_size})" + + # c_stride_m for locality domain strided output (0 = default contiguous) + c_stride_m_val = cutlass.Int64( + interm_size * + 2) if locality_domain_half_gemm else cutlass.Int64(0) + + max_active_clusters = get_max_activate_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1]) + cache_key = (self.scaling_vector_size, self.tile_size, + self.top_k, mma_tiler, mma_inst_shape, + cluster_shape_mn, raster_along_m, + locality_domain_half_gemm, a_path, + int(self.activation_type), max_active_clusters) + if cache_key not in self.__class__.kernel_cache: + gemm = self.__class__.kernel_class( + sf_vec_size=self.scaling_vector_size, + mma_inst_shape=mma_inst_shape, + mma_tiler=mma_tiler, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=True, + topk=self.top_k, + raster_along_m=raster_along_m, + locality_domain_half_gemm=locality_domain_half_gemm, + a_path=a_path, + activation_type=self.activation_type, + ) + compiled_gemm = cute.compile( + gemm.wrapper, + a_ptr, + b_ptr, + a_sf_ptr, + b_sf_ptr, + c_ptr, + c_sf_ptr, + alpha_ptr, + tile_idx_to_group_idx_ptr, + tile_idx_to_mn_limit_ptr, + permuted_idx_to_expanded_idx_ptr, + num_non_exiting_tiles_ptr, + global_sf_ptr, + orig_m, + m, + n, + k, + l, + tile_size=self.tile_size, + scaling_vector_size=self.scaling_vector_size, + max_active_clusters=max_active_clusters, + stream=stream, + c_stride_m=c_stride_m_val, + c_sf_n_tile_offset=c_sf_n_tile_offset_val, + ) + self.__class__.kernel_cache[cache_key] = compiled_gemm + else: + compiled_gemm = self.__class__.kernel_cache[cache_key] + + compiled_gemm( + a_ptr, + b_ptr, + a_sf_ptr, + b_sf_ptr, + c_ptr, + c_sf_ptr, + alpha_ptr, + tile_idx_to_group_idx_ptr, + tile_idx_to_mn_limit_ptr, + permuted_idx_to_expanded_idx_ptr, + num_non_exiting_tiles_ptr, + global_sf_ptr, + orig_m, + m, + n, + k, + l, + stream=stream, + c_stride_m=c_stride_m_val, + c_sf_n_tile_offset=c_sf_n_tile_offset_val, + ) + return c, c_sf + + def _run_nvfp4_gather_grouped_gemm_act_fusion_rubin( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_tensor: Optional[torch.Tensor], + output_sf_tensor: Optional[torch.Tensor], + scaling_vector_size: int, + partition_id: int, + activation_type: ActivationType, + precomputed_tactic: Optional[str], + tuner_key: str, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + tuner = AutoTuner.get() + if output_tensor is not None or output_sf_tensor is not None: + if output_tensor is None or output_sf_tensor is None: + raise ValueError( + "output_tensor and output_sf_tensor must be provided together." + ) + if partition_id < 0 or partition_id >= 2: + raise ValueError( + "partition_id must be 0 or 1 when output tensors are provided." + ) + elif partition_id != -1: + raise ValueError( + "partition_id must be -1 when output tensors are not provided." + ) + + runner = Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner( + num_experts, + top_k, + num_local_experts, + local_expert_offset, + tile_size, + scaling_vector_size, + activation_type=activation_type, + ) + inputs = [ + input, weight, input_scale, weight_scale, alpha, + tile_idx_to_group_idx, tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, num_non_exiting_tiles, global_sf, + output_tensor, output_sf_tensor + ] + choose_one_kwargs = {} + if output_tensor is not None: + choose_one_kwargs["partition_id"] = partition_id + + if precomputed_tactic is None: + _, best_tactic = tuner.choose_one( + tuner_key, + [runner], + runner.get_tuning_config(), + inputs, + **choose_one_kwargs, + ) + else: + best_tactic = ast.literal_eval(precomputed_tactic) + output, output_sf = runner(inputs, + tactic=best_tactic, + partition_id=partition_id) + if output_tensor is not None: + return None, None + return output, output_sf + + @torch.library.custom_op( + "trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin", + mutates_args=("output_tensor", "output_sf_tensor"), + schema= + "(Tensor input, Tensor weight, Tensor input_scale, Tensor weight_scale, Tensor alpha, " + "Tensor tile_idx_to_group_idx, Tensor tile_idx_to_mn_limit, " + "Tensor permuted_idx_to_expanded_idx, Tensor num_non_exiting_tiles, Tensor global_sf, " + "SymInt num_experts, SymInt top_k, SymInt num_local_experts, " + "SymInt local_expert_offset, SymInt tile_size, " + "Tensor(a16!)? output_tensor, Tensor(a17!)? output_sf_tensor, " + "SymInt scaling_vector_size=16, SymInt partition_id=-1, " + f"SymInt activation_type={int(ActivationType.Swiglu)}, " + "str? precomputed_tactic=None) -> (Tensor?, Tensor?)", + device_types="cuda") + def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_tensor: Optional[torch.Tensor], + output_sf_tensor: Optional[torch.Tensor], + scaling_vector_size: int = 16, + partition_id: int = -1, + activation_type: int = int(ActivationType.Swiglu), + precomputed_tactic: Optional[str] = None, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + return _run_nvfp4_gather_grouped_gemm_act_fusion_rubin( + input, weight, input_scale, weight_scale, alpha, + tile_idx_to_group_idx, tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, num_non_exiting_tiles, global_sf, + num_experts, top_k, num_local_experts, local_expert_offset, + tile_size, output_tensor, output_sf_tensor, + scaling_vector_size, partition_id, + ActivationType(activation_type), precomputed_tactic, + "trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin") + + @torch.library.register_fake( + "trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin") + def _( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_tensor: Optional[torch.Tensor], + output_sf_tensor: Optional[torch.Tensor], + scaling_vector_size: int = 16, + partition_id: int = -1, + activation_type: int = int(ActivationType.Swiglu), + precomputed_tactic: Optional[str] = None, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: + m = permuted_idx_to_expanded_idx.size(0) + n = weight.size(1) + is_gated = is_gated_activation(ActivationType(activation_type)) + interm_size = n // 2 if is_gated else n + if output_tensor is not None or output_sf_tensor is not None: + assert output_tensor is not None + assert output_sf_tensor is not None + return None, None + output = torch.empty(m, + interm_size // 2, + dtype=input.dtype, + device=input.device) + output_scale = torch.empty(m * interm_size // scaling_vector_size, + dtype=input_scale.dtype, + device=input_scale.device) + return output, output_scale + + @torch.library.custom_op( + "trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_locality_domain_inplace_rubin", + mutates_args=("output_tensor", "output_sf_tensor"), + schema= + "(Tensor input, Tensor weight_0, Tensor weight_1, Tensor input_scale, " + "Tensor weight_scale_0, Tensor weight_scale_1, Tensor alpha, " + "Tensor tile_idx_to_group_idx, Tensor tile_idx_to_mn_limit, " + "Tensor permuted_idx_to_expanded_idx, Tensor num_non_exiting_tiles, " + "Tensor global_sf, SymInt num_experts, SymInt top_k, " + "SymInt num_local_experts, SymInt local_expert_offset, " + "SymInt tile_size, Tensor(a!) output_tensor, " + "Tensor(b!) output_sf_tensor, SymInt scaling_vector_size=16, " + f"SymInt activation_type={int(ActivationType.Swiglu)}) -> ()", + device_types="cuda") + def cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_locality_domain_inplace_rubin( + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + input_scale: torch.Tensor, + weight_scale_0: torch.Tensor, + weight_scale_1: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_tensor: torch.Tensor, + output_sf_tensor: torch.Tensor, + scaling_vector_size: int = 16, + activation_type: int = int(ActivationType.Swiglu), + ) -> None: + """Tune and launch both Rubin locality domain NVFP4 MoE FC1 partitions. + + The MoE outer runner invokes this op during preparation so locality domain + resources, tactics, and kernels are ready before CUDA graph capture. + Direct callers must likewise invoke it once before capture. + """ + if weight_0.shape != weight_1.shape: + raise ValueError( + "locality domain NVFP4 MoE FC1 weight shards must have identical " + f"shapes, got {tuple(weight_0.shape)} and " + f"{tuple(weight_1.shape)}.") + if weight_0.dtype != weight_1.dtype: + raise ValueError( + "locality domain NVFP4 MoE FC1 weight shards must have identical " + f"dtypes, got {weight_0.dtype} and {weight_1.dtype}.") + if weight_scale_0.shape != weight_scale_1.shape: + raise ValueError( + "locality domain NVFP4 MoE FC1 weight-scale shards must have " + f"identical shapes, got {tuple(weight_scale_0.shape)} and " + f"{tuple(weight_scale_1.shape)}.") + if weight_scale_0.dtype != weight_scale_1.dtype: + raise ValueError( + "locality domain NVFP4 MoE FC1 weight-scale shards must have " + f"identical dtypes, got {weight_scale_0.dtype} and " + f"{weight_scale_1.dtype}.") + + runtime = LocalityDomainRuntime(num_partitions=2) + # Preserve the pre-refactor leaf namespace so persisted tactic + # caches remain reusable after moving ownership into this op. + tuner_key = ( + "trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin") + op_runner = ( + Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner( + num_experts, + top_k, + num_local_experts, + local_expert_offset, + tile_size, + scaling_vector_size, + activation_type=ActivationType(activation_type), + )) + inputs = [ + input, + weight_0, + input_scale, + weight_scale_0, + alpha, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + global_sf, + output_tensor, + output_sf_tensor, + ] + + def launch_partition( + partition_id: int, + partition_inputs: List[torch.Tensor], + tactic, + ) -> None: + weight = weight_0 if partition_id == 0 else weight_1 + weight_scale = (weight_scale_0 + if partition_id == 0 else weight_scale_1) + torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( + input=partition_inputs[0], + weight=weight, + input_scale=partition_inputs[2], + weight_scale=weight_scale, + alpha=partition_inputs[4], + tile_idx_to_group_idx=partition_inputs[5], + tile_idx_to_mn_limit=partition_inputs[6], + permuted_idx_to_expanded_idx=partition_inputs[7], + num_non_exiting_tiles=partition_inputs[8], + global_sf=partition_inputs[9], + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=local_expert_offset, + tile_size=tile_size, + output_tensor=partition_inputs[10], + output_sf_tensor=partition_inputs[11], + scaling_vector_size=scaling_vector_size, + partition_id=partition_id, + activation_type=activation_type, + precomputed_tactic=repr(tactic), + ) + + runner, best_tactic = tune_locality_domain_concurrent( + tuner_key, + op_runner, + runtime, + 2, + launch_partition, + inputs, + op_runner.get_tuning_config(), + ) + runner(inputs, tactic=best_tactic) + + @torch.library.register_fake( + "trtllm::cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_locality_domain_inplace_rubin" + ) + def _( + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + input_scale: torch.Tensor, + weight_scale_0: torch.Tensor, + weight_scale_1: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + global_sf: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_tensor: torch.Tensor, + output_sf_tensor: torch.Tensor, + scaling_vector_size: int = 16, + activation_type: int = int(ActivationType.Swiglu), + ) -> None: + return None + + # ---------------------------------------------------------------- + # Rubin BF16/FP16 Gather + SwiGLU Fusion (FC1 layer) + # ---------------------------------------------------------------- + from ..cute_dsl_kernels.rubin.moe.rubin_contiguous_gather_grouped_gemm_swiglu_fusion import \ + Sm107ContiguousGatherGroupedGemmSwigluFusionKernel + + class Sm107ContiguousGatherGroupedGemmSwigluFusionRunner(TunableRunner): + """Rubin (SM107) runner for BF16/FP16 gather + grouped GEMM + SwiGLU fusion (FC1). + + Similar to the blockscaled runner but without scale factors. + Uses MmaF16BF16Op (K=16) for BFloat16/Float16 inputs. + """ + kernel_class = Sm107ContiguousGatherGroupedGemmSwigluFusionKernel + kernel_cache = dict() + tuning_config_cache = dict() + + def __init__(self, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + input_dtype: Optional[torch.dtype] = None): + super().__init__() + self.num_experts = num_experts + self.top_k = top_k + self.num_local_experts = num_local_experts + self.local_expert_offset = local_expert_offset + self.tile_size = tile_size + self.input_dtype = input_dtype + + if input_dtype is not None and input_dtype not in ( + torch.bfloat16, torch.float16): + raise ValueError( + f"{self.__class__.kernel_class.__name__} requires BF16 or FP16 input, " + f"but got {input_dtype}") + + if (sm_version := get_sm_version()) != 107: + raise ValueError( + f"{self.__class__.kernel_class.__name__} supports SM 107 (Rubin) only, but got SM {sm_version}" + ) + + if self.tile_size not in (64, 128, 256): + raise ValueError( + f"{self.__class__.kernel_class.__name__} supports tile_size 64, 128 and 256 only, but got {self.tile_size}" + ) + + def unique_id(self): + return ( + self.num_experts, + self.top_k, + self.num_local_experts, + self.local_expert_offset, + self.tile_size, + self.input_dtype, + ) + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + **kwargs, + ) -> List[Tuple[int, int]]: + a, b, alpha, tile_idx_to_group_idx, tile_idx_to_mn_limit, permuted_idx_to_expanded_idx, *_ = inputs + m = permuted_idx_to_expanded_idx.size(0) + k = a.size(1) + l, n = b.size(0), b.size(1) # noqa: E741 + + ab_dtype = cutlass.BFloat16 if a.dtype == torch.bfloat16 else cutlass.Float16 + c_dtype = ab_dtype + + # BF16/FP16 tile sizes (no B-reuse): + # - (64, N): 1CTA + # - (128, N): 1CTA + # - (256, N): 2CTA + # Fixed K dimensions: mma_tiler_k=64, mma_inst_k=16 + mma_tiler_k = 64 + mma_inst_k = 16 + + mma_n_candidates = [128, 256] + raster_along_m_candidates = [False] + + # BF16 (no B-reuse): CTA tile M must equal tile_size. + # 2CTA (CtaGroup.TWO) is only valid for tile_size=256 + # (mma_m=256 triggers CtaGroup.TWO internally). + # cluster_shape_mn is always (max(1, tile_size//128), 1). + mma_m = self.tile_size + cluster_shape_mn = (max(1, self.tile_size // 128), 1) + + valid_tactics = [] + for mma_n, raster_along_m in (itertools.product( + mma_n_candidates, raster_along_m_candidates)): + + mma_tiler = (mma_m, mma_n, mma_tiler_k) + mma_inst_shape = (mma_m, mma_n, mma_inst_k) + + if self.__class__.kernel_class.can_implement( + a_dtype=ab_dtype, + b_dtype=ab_dtype, + c_dtype=c_dtype, + mma_inst_shape=mma_inst_shape, + mma_tiler=mma_tiler, + cluster_shape_mn=cluster_shape_mn, + m=m, + n=n, + k=k, + l=l, + a_major="k", + b_major="k", + c_major="n", + ): + valid_tactics.append((mma_tiler, mma_inst_shape, + cluster_shape_mn, raster_along_m)) + + logger.debug( + f"CuteDSL Rubin BF16 GatherGroupedGemmSwiglu: Found {len(valid_tactics)} valid tactics " + f"for M={m}, N={n}, K={k}, L={l}") + return valid_tactics + + # BF16 input layout (no scale factors): + # 0: a, 1: b, 2: alpha, + # 3: tile_idx_to_group_idx, 4: tile_idx_to_mn_limit, + # 5: permuted_idx_to_expanded_idx, 6: num_non_exiting_tiles + _BF16_IDX_PERMUTED = 5 + + def get_tuning_config(self, + has_output_tensor: bool = False + ) -> TuningConfig: + key = (*self.unique_id(), has_output_tensor) + if key not in self.__class__.tuning_config_cache: + helper = GatherGroupedGemmInputsHelper( + self.num_experts, self.top_k, self.num_local_experts, + self.local_expert_offset, self.tile_size) + # BF16 has permuted_idx at index 5, not 7 (no scale + # factor inputs). Override IDX_SHAPE_INFER so that + # infer_shape_num_tokens / infer_shape_max_num_tiles + # read the correct tensor. + helper.IDX_SHAPE_INFER = self._BF16_IDX_PERMUTED + constraint_specs = [ + ConstraintSpec(0, 0, helper.infer_shape_num_tokens), + ConstraintSpec(3, 0, helper.infer_shape_max_num_tiles), + ConstraintSpec(4, 0, helper.infer_shape_max_num_tiles), + ] + if has_output_tensor: + constraint_specs.append( + ConstraintSpec( + 7, 0, + helper.infer_shape_max_num_permuted_tokens)) + self.__class__.tuning_config_cache[key] = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + self._BF16_IDX_PERMUTED, 0, + helper.gen_tuning_buckets, + helper.map_to_tuning_buckets), ), + constraint_specs=tuple(constraint_specs), + inputs_pre_hook=self._bf16_inputs_pre_hook, + ) + return self.__class__.tuning_config_cache[key] + + def _bf16_inputs_pre_hook( + self, inputs: List[torch.Tensor]) -> List[torch.Tensor]: + """Pre-hook adapted for BF16 input layout (no scale factors).""" + a, b, alpha, tile_idx_to_group_idx, tile_idx_to_mn_limit, \ + permuted_idx_to_expanded_idx, num_non_exiting_tiles, *maybe_output = inputs + + helper = GatherGroupedGemmInputsHelper(self.num_experts, + self.top_k, + self.num_local_experts, + self.local_expert_offset, + self.tile_size) + + max_num_permuted_tokens = permuted_idx_to_expanded_idx.size(0) + num_tokens = helper.infer_num_tokens(max_num_permuted_tokens) + num_tokens_per_expert = helper.generate_num_tokens_per_expert( + num_tokens, approx_max_load=True) + token_selected_experts = helper.generate_token_selected_experts( + num_tokens, num_tokens_per_expert) + + token_selected_experts = token_selected_experts.cuda() + token_final_scales = torch.ones_like(token_selected_experts, + dtype=torch.float32) + + ( + new_tile_idx_to_group_idx, + new_tile_idx_to_mn_limit, + _, + new_permuted_idx_to_expanded_idx, + _, + new_num_non_exiting_tiles, + ) = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=self.num_experts, + top_k=self.top_k, + local_expert_offset=self.local_expert_offset, + local_num_experts=self.num_local_experts, + tile_tokens_dim=self.tile_size, + ) + + updated_inputs = [ + a, + b, + alpha, + new_tile_idx_to_group_idx, + new_tile_idx_to_mn_limit, + new_permuted_idx_to_expanded_idx, + new_num_non_exiting_tiles, + ] + if maybe_output: + updated_inputs.append(maybe_output[0]) + return updated_inputs + + def forward(self, inputs: List[torch.Tensor], + tactic: Optional[tuple], **kwargs) -> torch.Tensor: + (a, b, alpha, tile_idx_to_group_idx, tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, num_non_exiting_tiles, + *maybe_output) = inputs + assert inputs[ + self._BF16_IDX_PERMUTED] is permuted_idx_to_expanded_idx + assert a.dtype in (torch.bfloat16, torch.float16) + assert a.dim() == 2 + assert b.dtype == a.dtype + assert b.dim() == 3 + assert alpha.dtype == torch.float32 + assert alpha.dim() == 1 + + ab_dtype = cutlass.BFloat16 if a.dtype == torch.bfloat16 else cutlass.Float16 + + orig_m, k = a.size(0), a.size(1) + m = permuted_idx_to_expanded_idx.size(0) + l, n = b.size(0), b.size(1) # noqa: E741 + interm_size = n // 2 + assert m % self.tile_size == 0 + assert b.size(2) == k + assert alpha.size(0) == l + + num_tiles = m // self.tile_size + assert tile_idx_to_group_idx.dtype == torch.int32 + assert tile_idx_to_group_idx.size() == (num_tiles, ) + assert tile_idx_to_mn_limit.dtype == torch.int32 + assert tile_idx_to_mn_limit.size() == (num_tiles, ) + assert permuted_idx_to_expanded_idx.dtype == torch.int32 + assert permuted_idx_to_expanded_idx.size() == (m, ) + assert num_non_exiting_tiles.dtype == torch.int32 + assert num_non_exiting_tiles.numel() == 1 + + if maybe_output: + partition_id = kwargs.get("partition_id", -1) + if partition_id < 0 or partition_id >= 2: + raise ValueError( + "partition_id must be 0 or 1 when output_tensor is provided." + ) + output_tensor = maybe_output[0] + assert output_tensor.dim() == 2 + assert output_tensor.dtype == a.dtype + assert output_tensor.shape[0] == m + assert output_tensor.shape[1] == interm_size * 2 + c = output_tensor[:, partition_id * + interm_size:(partition_id + 1) * + interm_size] + c_stride_m_val = cutlass.Int64(output_tensor.shape[1]) + locality_domain_half_gemm = True + else: + c = torch.empty(m, + interm_size, + dtype=a.dtype, + device=a.device) + c_stride_m_val = cutlass.Int64(0) + locality_domain_half_gemm = False + + a_ptr = make_ptr(ab_dtype, + a.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + b_ptr = make_ptr(ab_dtype, + b.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + c_ptr = make_ptr(ab_dtype, + c.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + alpha_ptr = make_ptr(cutlass.Float32, alpha.data_ptr(), + cute.AddressSpace.gmem) + tile_idx_to_group_idx_ptr = make_ptr( + cutlass.Int32, tile_idx_to_group_idx.data_ptr(), + cute.AddressSpace.gmem) + tile_idx_to_mn_limit_ptr = make_ptr( + cutlass.Int32, tile_idx_to_mn_limit.data_ptr(), + cute.AddressSpace.gmem) + permuted_idx_to_expanded_idx_ptr = make_ptr( + cutlass.Int32, permuted_idx_to_expanded_idx.data_ptr(), + cute.AddressSpace.gmem) + num_non_exiting_tiles_ptr = make_ptr( + cutlass.Int32, num_non_exiting_tiles.data_ptr(), + cute.AddressSpace.gmem) + + torch_stream = torch.cuda.current_stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + if isinstance(tactic, tuple): + mma_tiler, mma_inst_shape, cluster_shape_mn, raster_along_m = tactic + else: + mma_tiler = (self.tile_size, 128, 64) + mma_inst_shape = (self.tile_size, 128, 16) + cluster_shape_mn = (max(1, self.tile_size // 128), 1) + raster_along_m = False + + max_active_clusters = get_max_activate_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1]) + cache_key = (a.dtype, self.tile_size, self.top_k, mma_tiler, + mma_inst_shape, cluster_shape_mn, raster_along_m, + locality_domain_half_gemm, max_active_clusters) + if cache_key not in self.__class__.kernel_cache: + gemm = self.__class__.kernel_class( + mma_inst_shape=mma_inst_shape, + mma_tiler=mma_tiler, + cluster_shape_mn=cluster_shape_mn, + vectorized_f32=True, + topk=self.top_k, + raster_along_m=raster_along_m, + ) + compiled_gemm = cute.compile( + gemm.wrapper, + a_ptr, + b_ptr, + c_ptr, + alpha_ptr, + tile_idx_to_group_idx_ptr, + tile_idx_to_mn_limit_ptr, + permuted_idx_to_expanded_idx_ptr, + num_non_exiting_tiles_ptr, + orig_m, + m, + n, + k, + l, + tile_size=self.tile_size, + max_active_clusters=max_active_clusters, + stream=stream, + c_stride_m=c_stride_m_val, + ) + self.__class__.kernel_cache[cache_key] = compiled_gemm + else: + compiled_gemm = self.__class__.kernel_cache[cache_key] + + compiled_gemm( + a_ptr, + b_ptr, + c_ptr, + alpha_ptr, + tile_idx_to_group_idx_ptr, + tile_idx_to_mn_limit_ptr, + permuted_idx_to_expanded_idx_ptr, + num_non_exiting_tiles_ptr, + orig_m, + m, + n, + k, + l, + stream=stream, + c_stride_m=c_stride_m_val, + ) + return c + + @torch.library.custom_op( + "trtllm::cute_dsl_bf16_gather_grouped_gemm_swiglu_rubin", + mutates_args=("output_tensor", ), + schema="(Tensor input, Tensor weight, Tensor alpha, " + "Tensor tile_idx_to_group_idx, Tensor tile_idx_to_mn_limit, " + "Tensor permuted_idx_to_expanded_idx, Tensor num_non_exiting_tiles, " + "SymInt num_experts, SymInt top_k, SymInt num_local_experts, " + "SymInt local_expert_offset, SymInt tile_size, " + "Tensor(a!)? output_tensor, SymInt partition_id, " + "str? precomputed_tactic=None) -> Tensor?", + device_types="cuda") + def cute_dsl_bf16_gather_grouped_gemm_swiglu_rubin( + input: torch.Tensor, + weight: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_tensor: Optional[torch.Tensor], + partition_id: int, + precomputed_tactic: Optional[str] = None, + ) -> Optional[torch.Tensor]: + tuner = AutoTuner.get() + if output_tensor is not None: + if partition_id < 0 or partition_id >= 2: + raise ValueError( + "partition_id must be 0 or 1 when output_tensor is provided." + ) + elif partition_id != -1: + raise ValueError( + "partition_id must be -1 when output_tensor is not provided." + ) + + runner = Sm107ContiguousGatherGroupedGemmSwigluFusionRunner( + num_experts, + top_k, + num_local_experts, + local_expert_offset, + tile_size, + input_dtype=input.dtype) + inputs = [ + input, + weight, + alpha, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + ] + if output_tensor is not None: + inputs.append(output_tensor) + + choose_one_kwargs = {} + if output_tensor is not None: + choose_one_kwargs["partition_id"] = partition_id + + if precomputed_tactic is None: + _, best_tactic = tuner.choose_one( + "trtllm::cute_dsl_bf16_gather_grouped_gemm_swiglu_rubin", + [runner], + runner.get_tuning_config(output_tensor is not None), + inputs, + **choose_one_kwargs, + ) + else: + best_tactic = ast.literal_eval(precomputed_tactic) + + output = runner(inputs, + tactic=best_tactic, + partition_id=partition_id) + if output_tensor is not None: + return None + return output + + @torch.library.register_fake( + "trtllm::cute_dsl_bf16_gather_grouped_gemm_swiglu_rubin") + def _( + input: torch.Tensor, + weight: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_tensor: Optional[torch.Tensor], + partition_id: int, + precomputed_tactic: Optional[str] = None, + ) -> Optional[torch.Tensor]: + m = permuted_idx_to_expanded_idx.size(0) + n = weight.size(1) + interm_size = n // 2 + if output_tensor is not None: + return None + return torch.empty(m, + interm_size, + dtype=input.dtype, + device=input.device) + + @torch.library.custom_op( + "trtllm::cute_dsl_bf16_gather_grouped_gemm_swiglu_locality_domain_inplace_rubin", + mutates_args=("output_tensor", ), + schema="(Tensor input, Tensor weight_0, Tensor weight_1, " + "Tensor alpha, Tensor tile_idx_to_group_idx, " + "Tensor tile_idx_to_mn_limit, " + "Tensor permuted_idx_to_expanded_idx, " + "Tensor num_non_exiting_tiles, SymInt num_experts, SymInt top_k, " + "SymInt num_local_experts, SymInt local_expert_offset, " + "SymInt tile_size, Tensor(a!) output_tensor) -> ()", + device_types="cuda") + def cute_dsl_bf16_gather_grouped_gemm_swiglu_locality_domain_inplace_rubin( + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_tensor: torch.Tensor, + ) -> None: + """Tune and launch both Rubin locality domain BF16 MoE FC1 partitions. + + The MoE outer runner primes this op before CUDA graph capture. + Direct callers must likewise invoke it once before capture. + """ + if weight_0.shape != weight_1.shape: + raise ValueError( + "locality domain BF16 MoE FC1 weight shards must have identical " + f"shapes, got {tuple(weight_0.shape)} and " + f"{tuple(weight_1.shape)}.") + if weight_0.dtype != weight_1.dtype: + raise ValueError( + "locality domain BF16 MoE FC1 weight shards must have identical " + f"dtypes, got {weight_0.dtype} and {weight_1.dtype}.") + + runtime = LocalityDomainRuntime(num_partitions=2) + # Preserve the pre-refactor leaf namespace for cache compatibility. + tuner_key = ( + "trtllm::cute_dsl_bf16_gather_grouped_gemm_swiglu_rubin") + op_runner = Sm107ContiguousGatherGroupedGemmSwigluFusionRunner( + num_experts, + top_k, + num_local_experts, + local_expert_offset, + tile_size, + input_dtype=input.dtype, + ) + inputs = [ + input, + weight_0, + alpha, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + output_tensor, + ] + + def launch_partition( + partition_id: int, + partition_inputs: List[torch.Tensor], + tactic, + ) -> None: + weight = weight_0 if partition_id == 0 else weight_1 + torch.ops.trtllm.cute_dsl_bf16_gather_grouped_gemm_swiglu_rubin( + input=partition_inputs[0], + weight=weight, + alpha=partition_inputs[2], + tile_idx_to_group_idx=partition_inputs[3], + tile_idx_to_mn_limit=partition_inputs[4], + permuted_idx_to_expanded_idx=partition_inputs[5], + num_non_exiting_tiles=partition_inputs[6], + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=local_expert_offset, + tile_size=tile_size, + output_tensor=partition_inputs[7], + partition_id=partition_id, + precomputed_tactic=repr(tactic), + ) + + runner, best_tactic = tune_locality_domain_concurrent( + tuner_key, + op_runner, + runtime, + 2, + launch_partition, + inputs, + op_runner.get_tuning_config(has_output_tensor=True), + ) + runner(inputs, tactic=best_tactic) + + @torch.library.register_fake( + "trtllm::cute_dsl_bf16_gather_grouped_gemm_swiglu_locality_domain_inplace_rubin" + ) + def _( + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_tensor: torch.Tensor, + ) -> None: + return None + + # ---------------------------------------------------------------- + # Rubin Finalize Fusion (FC2 layer: grouped GEMM + scatter-add) + # ---------------------------------------------------------------- + from ..cute_dsl_kernels.rubin.moe.rubin_contiguous_grouped_blockscaled_gemm_finalize_fusion import \ + Sm107BlockScaledContiguousGroupedGemmFinalizeFusionKernel + + class Sm107BlockScaledContiguousGroupedGemmFinalizeFusionRunner( + TunableRunner): + """Rubin (SM107) runner for grouped GEMM + finalize fusion (FC2). + + This is the Rubin counterpart to + Sm100BlockScaledContiguousGroupedGemmFinalizeFusionRunner. + Key differences from Blackwell: + - Takes mma_inst_shape and mma_tiler as 3-tuples (not 2-tuples) + - Kernel __init__ takes topK parameter + - Supports B-reuse pattern + """ + kernel_class = Sm107BlockScaledContiguousGroupedGemmFinalizeFusionKernel + kernel_cache = dict() + tuning_config_cache = dict() + + def __init__(self, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + scaling_vector_size: int = 16): + super().__init__() + self.num_experts = num_experts + self.top_k = top_k + self.num_local_experts = num_local_experts + self.local_expert_offset = local_expert_offset + self.tile_size = tile_size + + assert output_dtype == torch.bfloat16 + self.output_dtype = output_dtype + self.scaling_vector_size = scaling_vector_size + + if (sm_version := get_sm_version()) != 107: + raise ValueError( + f"{self.__class__.kernel_class.__name__} supports SM 107 (Rubin) only, but got SM {sm_version}" + ) + + if self.tile_size not in (128, 256, 512): + raise ValueError( + f"{self.__class__.kernel_class.__name__} supports tile_size 128, 256 and 512 only, but got {self.tile_size}" + ) + + def unique_id(self): + return ( + self.num_experts, + self.top_k, + self.num_local_experts, + self.local_expert_offset, + self.tile_size, + self.output_dtype, + self.scaling_vector_size, + ) + + @staticmethod + def _is_n_tiling_compatible(n: int, mma_n: int, + cluster_n: int) -> bool: + """Return whether the kernel can cover N without a tail tile.""" + return n % (mma_n * cluster_n) == 0 + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + **kwargs, + ) -> List[Tuple[int, int]]: + a, b, *_ = inputs + m, k = a.size(0), a.size(1) * 2 + l, n = b.size(0), b.size(1) # noqa: E741 + + # Rubin FP4 K-dimension: mma_tiler_k=256, mma_inst_k=128 + mma_tiler_k = 256 + mma_inst_k = 128 + + # (mma_tiler_m, mma_inst_m) candidates + mma_m_candidates = [ + (128, 128), # no B-reuse, 1CTA + (256, 256), # no B-reuse, 2CTA + (256, 128), # B-reuse, 1CTA + (512, 256), # B-reuse, 2CTA + ] + # Restrict N candidates to 128 and 256 (matching Blackwell). + # mma_n=192 and mma_n=64 trigger the kernel's special SFB + # slicing paths (cta_tile_shape_n=192 / cta_tile_shape_n=64) + # which cause CUDA_ERROR_ILLEGAL_ADDRESS for certain N + # dimensions (e.g., Qwen3-30B-A3B with N=2048, mma_n=192). + mma_n_candidates = [128, 256] + cluster_shape_mn_candidates = [(1, 1), (2, 1), (1, 2), (2, 2)] + raster_along_m_candidates = [False] + + valid_tactics = [] + for (mma_tiler_m, + mma_inst_m), mma_n, cluster_shape_mn, raster_along_m in ( + itertools.product(mma_m_candidates, mma_n_candidates, + cluster_shape_mn_candidates, + raster_along_m_candidates)): + + # The kernel requires each cluster to cover exactly + # one routing tile along M: + # - 1-CTA (mma_inst_m=128): cluster_m must be 1 + # - 2-CTA (mma_inst_m=256): cluster_m must be 2 + # and the CTA-pair's M extent (mma_tiler_m) must equal + # the routing tile_size. Any other combination + # (mma_tiler_m != tile_size, or cluster_m spanning + # multiple independent M-tiles) produces ~40% element + # mismatches empirically. + cta_group_m = 2 if mma_inst_m == 256 else 1 + if cluster_shape_mn[0] != cta_group_m: + continue + if mma_tiler_m != self.tile_size: + continue + if not self._is_n_tiling_compatible(n, mma_n, + cluster_shape_mn[1]): + continue + + mma_tiler = (mma_tiler_m, mma_n, mma_tiler_k) + mma_inst_shape = (mma_inst_m, mma_n, mma_inst_k) + + if self.__class__.kernel_class.can_implement( + a_dtype=cutlass.Float4E2M1FN, + b_dtype=cutlass.Float4E2M1FN, + sf_dtype=cutlass.Float8E4M3FN, + sf_vec_size=self.scaling_vector_size, + c_dtype=cutlass.BFloat16, + mma_inst_shape=mma_inst_shape, + mma_tiler=mma_tiler, + cluster_shape_mn=cluster_shape_mn, + m=m, + n=n, + k=k, + l=l, + a_major="k", + b_major="k", + c_major="n", + ): + valid_tactics.append((mma_tiler, mma_inst_shape, + cluster_shape_mn, raster_along_m)) + + logger.debug( + f"CuteDSL Rubin GroupedGemmFinalize: Found {len(valid_tactics)} valid tactics " + f"for M={m}, N={n}, K={k}, L={l}") + return valid_tactics + + def get_tuning_config(self) -> TuningConfig: + key = self.unique_id() + if key not in self.__class__.tuning_config_cache: + helper = GroupedGemmInputsHelper(self.num_experts, + self.top_k, + self.num_local_experts, + self.local_expert_offset, + self.tile_size) + self.__class__.tuning_config_cache[key] = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, helper.gen_tuning_buckets, + helper.map_to_tuning_buckets), ), + constraint_specs=( + ConstraintSpec(2, 0, fp4_scale_infer_shape), + ConstraintSpec(5, 0, helper.infer_shape_num_tokens), + ConstraintSpec(6, 0, + helper.infer_shape_max_num_tiles), + ConstraintSpec(7, 0, + helper.infer_shape_max_num_tiles), + ConstraintSpec( + 8, 0, + helper.infer_shape_max_num_permuted_tokens), + ConstraintSpec(10, 0, + helper.infer_shape_num_tokens), + ), + inputs_pre_hook=helper.inputs_pre_hook_finalize_fusion, + ) + return self.__class__.tuning_config_cache[key] + + def forward(self, inputs: List[torch.Tensor], + tactic: Optional[tuple]) -> torch.Tensor: + a, b, a_sf, b_sf, alpha, c, tile_idx_to_group_idx, tile_idx_to_mn_limit, permuted_idx_to_expanded_idx, num_non_exiting_tiles, token_final_scales = inputs + assert a.dtype == torch.float4_e2m1fn_x2 + assert a.dim() == 2 + assert b.dtype == torch.float4_e2m1fn_x2 + assert b.dim() == 3 + assert a_sf.dtype == torch.uint8 + assert a_sf.dim() == 1 + assert b_sf.dtype == torch.uint8 + assert b_sf.dim() == 3 + assert alpha.dtype == torch.float32 + assert alpha.dim() == 1 + + m, k = a.size(0), a.size(1) * 2 + l, n = b.size(0), b.size(1) # noqa: E741 + scale_k = k // self.scaling_vector_size + assert m % self.tile_size == 0 + assert k % (self.scaling_vector_size * 4) == 0 + assert b.size(2) * 2 == k + assert a_sf.size(0) == m * scale_k + assert b_sf.size(0) == l + assert b_sf.size(1) == n + assert b_sf.size(2) == scale_k + assert alpha.size(0) == l + + assert c.dtype == self.output_dtype + assert c.dim() == 2 + num_tokens = c.size(0) + assert c.size(1) == n or c.size(1) == n * 2 + + num_tiles = m // self.tile_size + assert tile_idx_to_group_idx.dtype == torch.int32 + assert tile_idx_to_group_idx.size() == (num_tiles, ) + assert tile_idx_to_mn_limit.dtype == torch.int32 + assert tile_idx_to_mn_limit.size() == (num_tiles, ) + assert permuted_idx_to_expanded_idx.dtype == torch.int32 + assert permuted_idx_to_expanded_idx.size() == (m, ) + assert num_non_exiting_tiles.dtype == torch.int32 + assert num_non_exiting_tiles.numel() == 1 + assert token_final_scales.dtype == torch.float32 + assert token_final_scales.dim() == 2 + assert token_final_scales.size() == (num_tokens, self.top_k) + + locality_domain_id = get_current_locality_domain() + locality_domain_half_gemm = locality_domain_id is not None + if locality_domain_half_gemm: + assert locality_domain_id in ( + 0, + 1), f"Invalid locality domain id: {locality_domain_id}" + assert c.size(1) == n * 2, \ + f"[locality domain] FC2 output must be 2x width: c.size(1)={c.size(1)}, n={n}" + # c: stride directly into shared buffer (kernel uses c_stride_row) + c = c[:, + locality_domain_id * n:(locality_domain_id + 1) * n] + + if isinstance(tactic, tuple): + mma_tiler, mma_inst_shape, cluster_shape_mn, raster_along_m = tactic + else: + # Default tactic for Rubin + mma_tiler, mma_inst_shape, cluster_shape_mn = \ + _get_sm107_nvfp4_default_mma_config(self.tile_size) + raster_along_m = False + assert mma_tiler[ + 0] >= self.tile_size, f"Tactic ({tactic}) is incompatible with tile size ({self.tile_size})" + if not self._is_n_tiling_compatible(n, mma_tiler[1], + cluster_shape_mn[1]): + raise ValueError( + f"Tactic ({tactic}) is incompatible with N={n}: " + f"mma_n={mma_tiler[1]} and cluster_n={cluster_shape_mn[1]} " + "require N to be divisible by their product.") + + a_ptr = make_ptr(cutlass.Float4E2M1FN, + a.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32) + b_ptr = make_ptr(cutlass.Float4E2M1FN, + b.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=32) + a_sf_ptr = make_ptr(cutlass.Float8E4M3FN, + a_sf.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + b_sf_ptr = make_ptr(cutlass.Float8E4M3FN, + b_sf.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + alpha_ptr = make_ptr(cutlass.Float32, alpha.data_ptr(), + cute.AddressSpace.gmem) + tile_idx_to_group_idx_ptr = make_ptr( + cutlass.Int32, tile_idx_to_group_idx.data_ptr(), + cute.AddressSpace.gmem) + tile_idx_to_mn_limit_ptr = make_ptr( + cutlass.Int32, tile_idx_to_mn_limit.data_ptr(), + cute.AddressSpace.gmem) + permuted_idx_to_expanded_idx_ptr = make_ptr( + cutlass.Int32, permuted_idx_to_expanded_idx.data_ptr(), + cute.AddressSpace.gmem) + num_non_exiting_tiles_ptr = make_ptr( + cutlass.Int32, num_non_exiting_tiles.data_ptr(), + cute.AddressSpace.gmem) + token_final_scales_ptr = make_ptr(cutlass.Float32, + token_final_scales.data_ptr(), + cute.AddressSpace.gmem) + c_ptr = make_ptr(cutlass.BFloat16, + c.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + + torch_stream = torch.cuda.current_stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + # c_stride_row for locality domain strided output (0 = default contiguous) + c_stride_row_val = cutlass.Int64( + n * 2) if locality_domain_half_gemm else cutlass.Int64(0) + + max_active_clusters = get_max_activate_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1]) + cache_key = (self.scaling_vector_size, self.tile_size, + self.top_k, mma_tiler, mma_inst_shape, + cluster_shape_mn, raster_along_m, + locality_domain_half_gemm, max_active_clusters) + if cache_key not in self.__class__.kernel_cache: + gemm = self.__class__.kernel_class( + sf_vec_size=self.scaling_vector_size, + mma_inst_shape=mma_inst_shape, + mma_tiler=mma_tiler, + cluster_shape_mn=cluster_shape_mn, + raster_along_m=raster_along_m, + topK=self.top_k, + ) + compiled_gemm = cute.compile( + gemm.wrapper, + a_ptr, + b_ptr, + a_sf_ptr, + b_sf_ptr, + c_ptr, + alpha_ptr, + tile_idx_to_group_idx_ptr, + tile_idx_to_mn_limit_ptr, + permuted_idx_to_expanded_idx_ptr, + num_non_exiting_tiles_ptr, + token_final_scales_ptr, + m, + n, + k, + l, + num_tokens, + self.top_k, + tile_size=self.tile_size, + scaling_vector_size=self.scaling_vector_size, + max_active_clusters=max_active_clusters, + stream=stream, + c_stride_row=c_stride_row_val, + ) + self.__class__.kernel_cache[cache_key] = compiled_gemm + else: + compiled_gemm = self.__class__.kernel_cache[cache_key] + + compiled_gemm( + a_ptr, + b_ptr, + a_sf_ptr, + b_sf_ptr, + c_ptr, + alpha_ptr, + tile_idx_to_group_idx_ptr, + tile_idx_to_mn_limit_ptr, + permuted_idx_to_expanded_idx_ptr, + num_non_exiting_tiles_ptr, + token_final_scales_ptr, + m, + n, + k, + l, + num_tokens, + self.top_k, + stream=stream, + c_stride_row=c_stride_row_val, + ) + # c written via stride — no copy-back needed + return c + + @torch.library.custom_op( + "trtllm::cute_dsl_nvfp4_grouped_gemm_finalize_inplace_rubin", + mutates_args=("output", ), + device_types="cuda") + def cute_dsl_nvfp4_grouped_gemm_finalize_inplace_rubin( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + scaling_vector_size: int = 16, + precomputed_tactic: Optional[str] = None, + ) -> None: + tuner = AutoTuner.get() + + runner = Sm107BlockScaledContiguousGroupedGemmFinalizeFusionRunner( + num_experts, top_k, num_local_experts, local_expert_offset, + tile_size, output_dtype, scaling_vector_size) + + inputs = [ + input, weight, input_scale, weight_scale, alpha, output, + tile_idx_to_group_idx, tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, num_non_exiting_tiles, + token_final_scales + ] + + if precomputed_tactic is None: + _, best_tactic = tuner.choose_one( + "trtllm::cute_dsl_nvfp4_grouped_gemm_finalize_inplace_rubin", + [runner], + runner.get_tuning_config(), + inputs, + ) + else: + best_tactic = ast.literal_eval(precomputed_tactic) + + runner(inputs, tactic=best_tactic) + + @torch.library.custom_op( + "trtllm::cute_dsl_nvfp4_grouped_gemm_finalize_rubin", + mutates_args=(), + device_types="cuda") + def cute_dsl_nvfp4_grouped_gemm_finalize_rubin( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + scaling_vector_size: int = 16, + ) -> torch.Tensor: + num_tokens = token_final_scales.size(0) + n = weight.size(1) + output = torch.zeros(num_tokens, + n, + dtype=output_dtype, + device=input.device) + torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_finalize_inplace_rubin( + input=input, + weight=weight, + input_scale=input_scale, + weight_scale=weight_scale, + alpha=alpha, + output=output, + tile_idx_to_group_idx=tile_idx_to_group_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=local_expert_offset, + tile_size=tile_size, + output_dtype=output_dtype, + scaling_vector_size=scaling_vector_size, + ) + return output + + @torch.library.register_fake( + "trtllm::cute_dsl_nvfp4_grouped_gemm_finalize_inplace_rubin") + def _( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + output: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + scaling_vector_size: int = 16, + precomputed_tactic: Optional[str] = None, + ) -> None: + return + + @torch.library.custom_op( + "trtllm::cute_dsl_nvfp4_grouped_gemm_finalize_locality_domain_inplace_rubin", + mutates_args=("output", ), + schema= + "(Tensor input, Tensor weight_0, Tensor weight_1, Tensor input_scale, " + "Tensor weight_scale_0, Tensor weight_scale_1, Tensor alpha, " + "Tensor(a!) output, Tensor tile_idx_to_group_idx, " + "Tensor tile_idx_to_mn_limit, " + "Tensor expanded_idx_to_permuted_idx, " + "Tensor permuted_idx_to_expanded_idx, " + "Tensor num_non_exiting_tiles, Tensor token_final_scales, " + "SymInt num_experts, SymInt top_k, SymInt num_local_experts, " + "SymInt local_expert_offset, SymInt tile_size, " + "ScalarType output_dtype, SymInt ep_size, " + "bool enable_alltoall=False, SymInt scaling_vector_size=16) -> ()", + device_types="cuda") + def cute_dsl_nvfp4_grouped_gemm_finalize_locality_domain_inplace_rubin( + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + input_scale: torch.Tensor, + weight_scale_0: torch.Tensor, + weight_scale_1: torch.Tensor, + alpha: torch.Tensor, + output: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + ep_size: int, + enable_alltoall: bool = False, + scaling_vector_size: int = 16, + ) -> None: + """Tune and launch both Rubin locality domain NVFP4 MoE FC2 partitions. + + The MoE outer runner primes this op before CUDA graph capture. + Direct callers must likewise invoke it once before capture. + """ + if weight_0.shape != weight_1.shape: + raise ValueError( + "locality domain NVFP4 MoE FC2 weight shards must have identical " + f"shapes, got {tuple(weight_0.shape)} and " + f"{tuple(weight_1.shape)}.") + if weight_0.dtype != weight_1.dtype: + raise ValueError( + "locality domain NVFP4 MoE FC2 weight shards must have identical " + f"dtypes, got {weight_0.dtype} and {weight_1.dtype}.") + if weight_scale_0.shape != weight_scale_1.shape: + raise ValueError( + "locality domain NVFP4 MoE FC2 weight-scale shards must have " + f"identical shapes, got {tuple(weight_scale_0.shape)} and " + f"{tuple(weight_scale_1.shape)}.") + if weight_scale_0.dtype != weight_scale_1.dtype: + raise ValueError( + "locality domain NVFP4 MoE FC2 weight-scale shards must have " + f"identical dtypes, got {weight_scale_0.dtype} and " + f"{weight_scale_1.dtype}.") + if output.dtype != output_dtype: + raise ValueError( + "locality domain NVFP4 MoE FC2 output tensor dtype must match " + f"output_dtype, got {output.dtype} and {output_dtype}.") + + runtime = LocalityDomainRuntime(num_partitions=2) + # Preserve the pre-refactor leaf namespace for cache compatibility. + tuner_key = ( + "trtllm::cute_dsl_nvfp4_grouped_gemm_finalize_inplace_rubin") + op_runner = ( + Sm107BlockScaledContiguousGroupedGemmFinalizeFusionRunner( + num_experts, + top_k, + num_local_experts, + local_expert_offset, + tile_size, + output_dtype, + scaling_vector_size, + )) + inputs = [ + input, + weight_0, + input_scale, + weight_scale_0, + alpha, + output, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + token_final_scales, + ] + + def launch_partition( + partition_id: int, + partition_inputs: List[torch.Tensor], + tactic, + ) -> None: + weight = weight_0 if partition_id == 0 else weight_1 + weight_scale = (weight_scale_0 + if partition_id == 0 else weight_scale_1) + torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_finalize_inplace_rubin( + input=partition_inputs[0], + weight=weight, + input_scale=partition_inputs[2], + weight_scale=weight_scale, + alpha=partition_inputs[4], + output=partition_inputs[5], + tile_idx_to_group_idx=partition_inputs[6], + tile_idx_to_mn_limit=partition_inputs[7], + permuted_idx_to_expanded_idx=partition_inputs[8], + num_non_exiting_tiles=partition_inputs[9], + token_final_scales=partition_inputs[10], + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=local_expert_offset, + tile_size=tile_size, + output_dtype=output_dtype, + scaling_vector_size=scaling_vector_size, + precomputed_tactic=repr(tactic), + ) + + runner, best_tactic = tune_locality_domain_concurrent( + tuner_key, + op_runner, + runtime, + 2, + launch_partition, + inputs, + op_runner.get_tuning_config(), + ) + # Profiling finalize accumulates into the shared output. Preserve + # the selective all-to-all memset semantics when restoring the + # baseline before the actual dual-partition launch. + if AutoTuner.get().is_tuning_mode: + torch.ops.trtllm.moe_output_memset_inplace( + input=output, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + tile_tokens_dim=tile_size, + top_k=top_k, + ep_size=ep_size, + enable_alltoall=enable_alltoall, + ) + runner(inputs, tactic=best_tactic) + + @torch.library.register_fake( + "trtllm::cute_dsl_nvfp4_grouped_gemm_finalize_locality_domain_inplace_rubin" + ) + def _( + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + input_scale: torch.Tensor, + weight_scale_0: torch.Tensor, + weight_scale_1: torch.Tensor, + alpha: torch.Tensor, + output: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + ep_size: int, + enable_alltoall: bool = False, + scaling_vector_size: int = 16, + ) -> None: + return None + + @torch.library.register_fake( + "trtllm::cute_dsl_nvfp4_grouped_gemm_finalize_rubin") + def _( + input: torch.Tensor, + weight: torch.Tensor, + input_scale: torch.Tensor, + weight_scale: torch.Tensor, + alpha: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + scaling_vector_size: int = 16, + ) -> torch.Tensor: + num_tokens = token_final_scales.size(0) + n = weight.size(1) + return torch.empty(num_tokens, + n, + dtype=output_dtype, + device=input.device) + + # ---------------------------------------------------------------- + # Rubin BF16/FP16 Finalize Fusion (FC2 layer: grouped GEMM + scatter-add) + # ---------------------------------------------------------------- + from ..cute_dsl_kernels.rubin.moe.rubin_contiguous_grouped_gemm_finalize_fusion import \ + Sm107ContiguousGroupedGemmFinalizeFusionKernel + + class Sm107ContiguousGroupedGemmFinalizeFusionRunner(TunableRunner): + """Rubin (SM107) runner for BF16/FP16 grouped GEMM + finalize fusion (FC2). + + Similar to the blockscaled finalize runner but without scale factors. + Uses MmaF16BF16Op (K=16) for BFloat16/Float16 inputs. + """ + kernel_class = Sm107ContiguousGroupedGemmFinalizeFusionKernel + kernel_cache = dict() + tuning_config_cache = dict() + + def __init__(self, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + input_dtype: Optional[torch.dtype] = None): + super().__init__() + self.num_experts = num_experts + self.top_k = top_k + self.num_local_experts = num_local_experts + self.local_expert_offset = local_expert_offset + self.tile_size = tile_size + + assert output_dtype in (torch.bfloat16, torch.float16) + self.output_dtype = output_dtype + self.input_dtype = input_dtype + if input_dtype is not None and input_dtype not in ( + torch.bfloat16, torch.float16): + raise ValueError( + f"{self.__class__.kernel_class.__name__} requires BF16 or FP16 input, " + f"but got {input_dtype}") + + if (sm_version := get_sm_version()) != 107: + raise ValueError( + f"{self.__class__.kernel_class.__name__} supports SM 107 (Rubin) only, but got SM {sm_version}" + ) + + if self.tile_size not in (64, 128, 256): + raise ValueError( + f"{self.__class__.kernel_class.__name__} supports tile_size 64, 128 and 256 only, but got {self.tile_size}" + ) + + def unique_id(self): + return ( + self.num_experts, + self.top_k, + self.num_local_experts, + self.local_expert_offset, + self.tile_size, + self.output_dtype, + self.input_dtype, + ) + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + **kwargs, + ) -> List[Tuple[int, int]]: + a, b, *_ = inputs + m, k = a.size(0), a.size(1) + l, n = b.size(0), b.size(1) # noqa: E741 + + ab_dtype = cutlass.BFloat16 if a.dtype == torch.bfloat16 else cutlass.Float16 + c_dtype = cutlass.BFloat16 if self.output_dtype == torch.bfloat16 else cutlass.Float16 + + # BF16/FP16 K-dimension: mma_tiler_k=64, mma_inst_k=16 + mma_tiler_k = 64 + mma_inst_k = 16 + + # BF16 (no B-reuse): CTA tile M must equal tile_size. + # cluster_shape_mn is always (max(1, tile_size//128), 1). + mma_n_candidates = [128, 256] + raster_along_m_candidates = [False] + + mma_m = self.tile_size + cluster_shape_mn = (max(1, self.tile_size // 128), 1) + + valid_tactics = [] + for mma_n, raster_along_m in (itertools.product( + mma_n_candidates, raster_along_m_candidates)): + + mma_tiler = (mma_m, mma_n, mma_tiler_k) + mma_inst_shape = (mma_m, mma_n, mma_inst_k) + + if self.__class__.kernel_class.can_implement( + a_dtype=ab_dtype, + b_dtype=ab_dtype, + c_dtype=c_dtype, + mma_inst_shape=mma_inst_shape, + mma_tiler=mma_tiler, + cluster_shape_mn=cluster_shape_mn, + m=m, + n=n, + k=k, + l=l, + a_major="k", + b_major="k", + c_major="n", + ): + valid_tactics.append((mma_tiler, mma_inst_shape, + cluster_shape_mn, raster_along_m)) + + logger.debug( + f"CuteDSL Rubin BF16 GroupedGemmFinalize: Found {len(valid_tactics)} valid tactics " + f"for M={m}, N={n}, K={k}, L={l}") + return valid_tactics + + def _bf16_inputs_pre_hook_finalize( + self, inputs: List[torch.Tensor]) -> List[torch.Tensor]: + """Pre-hook for BF16 finalize (no scale factors).""" + a, b, output, tile_idx_to_group_idx, \ + tile_idx_to_mn_limit, permuted_idx_to_expanded_idx, \ + num_non_exiting_tiles, token_final_scales = inputs + + helper = GroupedGemmInputsHelper(self.num_experts, self.top_k, + self.num_local_experts, + self.local_expert_offset, + self.tile_size) + num_tokens = helper.infer_num_tokens(a.size(0)) + num_tokens_per_expert = helper.generate_num_tokens_per_expert( + num_tokens, approx_max_load=True) + token_selected_experts = \ + helper.generate_token_selected_experts( + num_tokens, num_tokens_per_expert) + + token_selected_experts = token_selected_experts.cuda() + token_final_scales = torch.ones_like(token_selected_experts, + dtype=torch.float32) + ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles, + ) = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=self.num_experts, + top_k=self.top_k, + local_expert_offset=self.local_expert_offset, + local_num_experts=self.num_local_experts, + tile_tokens_dim=self.tile_size, + ) + return (a, b, output, tile_idx_to_group_idx, + tile_idx_to_mn_limit, permuted_idx_to_expanded_idx, + num_non_exiting_tiles, token_final_scales) + + def get_tuning_config(self) -> TuningConfig: + key = self.unique_id() + if key not in self.__class__.tuning_config_cache: + helper = GroupedGemmInputsHelper(self.num_experts, + self.top_k, + self.num_local_experts, + self.local_expert_offset, + self.tile_size) + # BF16 finalize input layout (8 tensors): + # 0: a, 1: b, 2: c (output), + # 3: tile_idx_to_group_idx, 4: tile_idx_to_mn_limit, + # 5: permuted_idx_to_expanded_idx, 6: num_non_exiting_tiles, + # 7: token_final_scales + self.__class__.tuning_config_cache[key] = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, helper.gen_tuning_buckets, + helper.map_to_tuning_buckets), ), + constraint_specs=( + ConstraintSpec(2, 0, helper.infer_shape_num_tokens), + ConstraintSpec(3, 0, + helper.infer_shape_max_num_tiles), + ConstraintSpec(4, 0, + helper.infer_shape_max_num_tiles), + ConstraintSpec( + 5, 0, + helper.infer_shape_max_num_permuted_tokens), + ConstraintSpec(7, 0, helper.infer_shape_num_tokens), + ), + inputs_pre_hook=self._bf16_inputs_pre_hook_finalize, + ) + return self.__class__.tuning_config_cache[key] + + def forward(self, inputs: List[torch.Tensor], + tactic: Optional[tuple]) -> torch.Tensor: + a, b, c, tile_idx_to_group_idx, tile_idx_to_mn_limit, permuted_idx_to_expanded_idx, num_non_exiting_tiles, token_final_scales = inputs + assert a.dtype in (torch.bfloat16, torch.float16) + assert a.dim() == 2 + assert b.dtype == a.dtype + assert b.dim() == 3 + + ab_dtype = cutlass.BFloat16 if a.dtype == torch.bfloat16 else cutlass.Float16 + c_cutlass_dtype = cutlass.BFloat16 if c.dtype == torch.bfloat16 else cutlass.Float16 + + m, k = a.size(0), a.size(1) + l, n = b.size(0), b.size(1) # noqa: E741 + assert m % self.tile_size == 0 + assert b.size(2) == k + + assert c.dtype == self.output_dtype + assert c.dim() == 2 + num_tokens = c.size(0) + assert c.size(1) == n or c.size(1) == n * 2 + + locality_domain_id = get_current_locality_domain() + locality_domain_half_gemm = locality_domain_id is not None + if locality_domain_half_gemm: + assert locality_domain_id in ( + 0, + 1), f"Invalid locality domain id: {locality_domain_id}" + assert c.size(1) == n * 2, \ + f"[locality domain] BF16 FC2 output must be 2x width: c.size(1)={c.size(1)}, n={n}" + c_stride_row_val = cutlass.Int64(c.size(1)) + c = c[:, + locality_domain_id * n:(locality_domain_id + 1) * n] + else: + c_stride_row_val = cutlass.Int64(0) + + num_tiles = m // self.tile_size + assert tile_idx_to_group_idx.dtype == torch.int32 + assert tile_idx_to_group_idx.size() == (num_tiles, ) + assert tile_idx_to_mn_limit.dtype == torch.int32 + assert tile_idx_to_mn_limit.size() == (num_tiles, ) + assert permuted_idx_to_expanded_idx.dtype == torch.int32 + assert permuted_idx_to_expanded_idx.size() == (m, ) + assert num_non_exiting_tiles.dtype == torch.int32 + assert num_non_exiting_tiles.numel() == 1 + assert token_final_scales.dtype == torch.float32 + assert token_final_scales.dim() == 2 + assert token_final_scales.size() == (num_tokens, self.top_k) + + a_ptr = make_ptr(ab_dtype, + a.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + b_ptr = make_ptr(ab_dtype, + b.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + tile_idx_to_group_idx_ptr = make_ptr( + cutlass.Int32, tile_idx_to_group_idx.data_ptr(), + cute.AddressSpace.gmem) + tile_idx_to_mn_limit_ptr = make_ptr( + cutlass.Int32, tile_idx_to_mn_limit.data_ptr(), + cute.AddressSpace.gmem) + permuted_idx_to_expanded_idx_ptr = make_ptr( + cutlass.Int32, permuted_idx_to_expanded_idx.data_ptr(), + cute.AddressSpace.gmem) + num_non_exiting_tiles_ptr = make_ptr( + cutlass.Int32, num_non_exiting_tiles.data_ptr(), + cute.AddressSpace.gmem) + token_final_scales_ptr = make_ptr(cutlass.Float32, + token_final_scales.data_ptr(), + cute.AddressSpace.gmem) + c_ptr = make_ptr(c_cutlass_dtype, + c.data_ptr(), + cute.AddressSpace.gmem, + assumed_align=16) + + torch_stream = torch.cuda.current_stream() + stream = cuda.CUstream(torch_stream.cuda_stream) + + if isinstance(tactic, tuple): + mma_tiler, mma_inst_shape, cluster_shape_mn, raster_along_m = tactic + else: + # Default tactic for Rubin BF16/FP16 + mma_tiler = (self.tile_size, 128, 64) + mma_inst_shape = (self.tile_size, 128, 16) + cluster_shape_mn = (max(1, self.tile_size // 128), 1) + raster_along_m = False + + max_active_clusters = get_max_activate_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1]) + cache_key = (a.dtype, c.dtype, self.tile_size, self.top_k, + mma_tiler, mma_inst_shape, cluster_shape_mn, + raster_along_m, locality_domain_half_gemm, + max_active_clusters) + if cache_key not in self.__class__.kernel_cache: + gemm = self.__class__.kernel_class( + mma_inst_shape=mma_inst_shape, + mma_tiler=mma_tiler, + cluster_shape_mn=cluster_shape_mn, + raster_along_m=raster_along_m, + topK=self.top_k, + ) + compiled_gemm = cute.compile( + gemm.wrapper, + a_ptr, + b_ptr, + c_ptr, + tile_idx_to_group_idx_ptr, + tile_idx_to_mn_limit_ptr, + permuted_idx_to_expanded_idx_ptr, + num_non_exiting_tiles_ptr, + token_final_scales_ptr, + m, + n, + k, + l, + num_tokens, + self.top_k, + tile_size=self.tile_size, + max_active_clusters=max_active_clusters, + stream=stream, + c_stride_row=c_stride_row_val, + ) + self.__class__.kernel_cache[cache_key] = compiled_gemm + else: + compiled_gemm = self.__class__.kernel_cache[cache_key] + + compiled_gemm( + a_ptr, + b_ptr, + c_ptr, + tile_idx_to_group_idx_ptr, + tile_idx_to_mn_limit_ptr, + permuted_idx_to_expanded_idx_ptr, + num_non_exiting_tiles_ptr, + token_final_scales_ptr, + m, + n, + k, + l, + num_tokens, + self.top_k, + stream=stream, + c_stride_row=c_stride_row_val, + ) + return c + + @torch.library.custom_op( + "trtllm::cute_dsl_bf16_grouped_gemm_finalize_inplace_rubin", + mutates_args=("output", ), + device_types="cuda") + def cute_dsl_bf16_grouped_gemm_finalize_inplace_rubin( + input: torch.Tensor, + weight: torch.Tensor, + output: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + precomputed_tactic: Optional[str] = None, + ) -> None: + tuner = AutoTuner.get() + + runner = Sm107ContiguousGroupedGemmFinalizeFusionRunner( + num_experts, + top_k, + num_local_experts, + local_expert_offset, + tile_size, + output_dtype, + input_dtype=input.dtype) + + inputs = [ + input, weight, output, tile_idx_to_group_idx, + tile_idx_to_mn_limit, permuted_idx_to_expanded_idx, + num_non_exiting_tiles, token_final_scales + ] + + if precomputed_tactic is None: + _, best_tactic = tuner.choose_one( + "trtllm::cute_dsl_bf16_grouped_gemm_finalize_inplace_rubin", + [runner], + runner.get_tuning_config(), + inputs, + ) + else: + best_tactic = ast.literal_eval(precomputed_tactic) + + runner(inputs, tactic=best_tactic) + + @torch.library.custom_op( + "trtllm::cute_dsl_bf16_grouped_gemm_finalize_rubin", + mutates_args=(), + device_types="cuda") + def cute_dsl_bf16_grouped_gemm_finalize_rubin( + input: torch.Tensor, + weight: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + ) -> torch.Tensor: + num_tokens = token_final_scales.size(0) + n = weight.size(1) + output = torch.zeros(num_tokens, + n, + dtype=output_dtype, + device=input.device) + torch.ops.trtllm.cute_dsl_bf16_grouped_gemm_finalize_inplace_rubin( + input=input, + weight=weight, + output=output, + tile_idx_to_group_idx=tile_idx_to_group_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=local_expert_offset, + tile_size=tile_size, + output_dtype=output_dtype, + ) + return output + + @torch.library.register_fake( + "trtllm::cute_dsl_bf16_grouped_gemm_finalize_inplace_rubin") + def _( + input: torch.Tensor, + weight: torch.Tensor, + output: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + precomputed_tactic: Optional[str] = None, + ) -> None: + return + + @torch.library.custom_op( + "trtllm::cute_dsl_bf16_grouped_gemm_finalize_locality_domain_inplace_rubin", + mutates_args=("output", ), + schema="(Tensor input, Tensor weight_0, Tensor weight_1, " + "Tensor(a!) output, Tensor tile_idx_to_group_idx, " + "Tensor tile_idx_to_mn_limit, " + "Tensor expanded_idx_to_permuted_idx, " + "Tensor permuted_idx_to_expanded_idx, " + "Tensor num_non_exiting_tiles, Tensor token_final_scales, " + "SymInt num_experts, SymInt top_k, SymInt num_local_experts, " + "SymInt local_expert_offset, SymInt tile_size, " + "ScalarType output_dtype, SymInt ep_size, " + "bool enable_alltoall=False) -> ()", + device_types="cuda") + def cute_dsl_bf16_grouped_gemm_finalize_locality_domain_inplace_rubin( + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + output: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + ep_size: int, + enable_alltoall: bool = False, + ) -> None: + """Tune and launch both Rubin locality domain BF16 MoE FC2 partitions. + + The MoE outer runner primes this op before CUDA graph capture. + Direct callers must likewise invoke it once before capture. + """ + if weight_0.shape != weight_1.shape: + raise ValueError( + "locality domain BF16 MoE FC2 weight shards must have identical " + f"shapes, got {tuple(weight_0.shape)} and " + f"{tuple(weight_1.shape)}.") + if weight_0.dtype != weight_1.dtype: + raise ValueError( + "locality domain BF16 MoE FC2 weight shards must have identical " + f"dtypes, got {weight_0.dtype} and {weight_1.dtype}.") + if output.dtype != output_dtype: + raise ValueError( + "locality domain BF16 MoE FC2 output tensor dtype must match " + f"output_dtype, got {output.dtype} and {output_dtype}.") + + runtime = LocalityDomainRuntime(num_partitions=2) + # Preserve the pre-refactor leaf namespace for cache compatibility. + tuner_key = ( + "trtllm::cute_dsl_bf16_grouped_gemm_finalize_inplace_rubin") + op_runner = Sm107ContiguousGroupedGemmFinalizeFusionRunner( + num_experts, + top_k, + num_local_experts, + local_expert_offset, + tile_size, + output_dtype, + input_dtype=input.dtype, + ) + inputs = [ + input, + weight_0, + output, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + token_final_scales, + ] + + def launch_partition( + partition_id: int, + partition_inputs: List[torch.Tensor], + tactic, + ) -> None: + weight = weight_0 if partition_id == 0 else weight_1 + torch.ops.trtllm.cute_dsl_bf16_grouped_gemm_finalize_inplace_rubin( + input=partition_inputs[0], + weight=weight, + output=partition_inputs[2], + tile_idx_to_group_idx=partition_inputs[3], + tile_idx_to_mn_limit=partition_inputs[4], + permuted_idx_to_expanded_idx=partition_inputs[5], + num_non_exiting_tiles=partition_inputs[6], + token_final_scales=partition_inputs[7], + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=local_expert_offset, + tile_size=tile_size, + output_dtype=output_dtype, + precomputed_tactic=repr(tactic), + ) + + runner, best_tactic = tune_locality_domain_concurrent( + tuner_key, + op_runner, + runtime, + 2, + launch_partition, + inputs, + op_runner.get_tuning_config(), + ) + # Restore the same selective zero baseline used by the backend + # before the real dual-partition finalize launch. + if AutoTuner.get().is_tuning_mode: + torch.ops.trtllm.moe_output_memset_inplace( + input=output, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + tile_tokens_dim=tile_size, + top_k=top_k, + ep_size=ep_size, + enable_alltoall=enable_alltoall, + ) + runner(inputs, tactic=best_tactic) + + @torch.library.register_fake( + "trtllm::cute_dsl_bf16_grouped_gemm_finalize_locality_domain_inplace_rubin" + ) + def _( + input: torch.Tensor, + weight_0: torch.Tensor, + weight_1: torch.Tensor, + output: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + expanded_idx_to_permuted_idx: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + ep_size: int, + enable_alltoall: bool = False, + ) -> None: + return None + + @torch.library.register_fake( + "trtllm::cute_dsl_bf16_grouped_gemm_finalize_rubin") + def _( + input: torch.Tensor, + weight: torch.Tensor, + tile_idx_to_group_idx: torch.Tensor, + tile_idx_to_mn_limit: torch.Tensor, + permuted_idx_to_expanded_idx: torch.Tensor, + num_non_exiting_tiles: torch.Tensor, + token_final_scales: torch.Tensor, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + tile_size: int, + output_dtype: torch.dtype, + ) -> torch.Tensor: + num_tokens = token_final_scales.size(0) + n = weight.size(1) + return torch.empty(num_tokens, + n, + dtype=output_dtype, + device=input.device) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py index aef0c2e3553e..3941b06d296f 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/blockscaled_contiguous_grouped_gemm_finalize_fusion.py @@ -2465,6 +2465,7 @@ def wrapper( max_active_clusters: cutlass.Constexpr, stream: cuda.CUstream, epilogue_op: cutlass.Constexpr = lambda x: x, + c_stride_row: cutlass.Int64 = cutlass.Int64(0), ): """Single-B wrapper. @@ -2480,8 +2481,15 @@ def wrapper( (32, 4, m // 128, 4, scale_k // 4, 1), order=(2, 1, 4, 0, 3, 5) ), ) + # c supports strided output for locality domain shared buffers. + # c_stride_row: row stride in output elements (0 = default = n) + actual_c_stride_row = n if c_stride_row == 0 else c_stride_row c = cute.make_tensor( - c_ptr, layout=cute.make_ordered_layout((num_tokens, n, 1), order=(1, 0, 2)) + c_ptr, + layout=cute.make_layout( + (num_tokens, n, 1), + stride=(actual_c_stride_row, 1, num_tokens * actual_c_stride_row), + ), ) alpha = cute.make_tensor(alpha_ptr, layout=cute.make_layout((l,))) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/inline_ptx.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/inline_ptx.py new file mode 100644 index 000000000000..321fb2607120 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/inline_ptx.py @@ -0,0 +1,200 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import cutlass +from cutlass._mlir.dialects import cute as _cute_ir +from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir +from cutlass._mlir.dialects import llvm +from cutlass.cute.typing import AddressSpace +from cutlass.cutlass_dsl import dsl_user_op + +# PTX `mbarrier::peer_bit` mask used by TMA gather4 in 2CTA mode: keeps +# all address bits except bit 24 (the peer-CTA bit), so both CTAs' bytes +# flow to the leader CTA's mbar. +_PEER_BIT_MASK = 0xFEFFFFFF + + +@dsl_user_op +def sm100_tma_gather4_load( + tma_atom, + smem_dst_ptr, + mbar_ptr, + col, + r0, + r1, + r2, + r3, + *, + use_cta_group_2: bool = False, + mcast_mask=None, + loc=None, + ip=None, +): + """Issue one TMA TILE_GATHER4 load. + + Emits inline PTX because gather4 has no DSL op. There are four PTX variants: + {1CTA, 2CTA} x {no-mcast, mcast::cluster}. + + - 2CTA: `.cta_group::2`; mbar peer-bit-masked so both CTAs' bytes flow to + the leader's mbar. + - mcast::cluster: adds `.multicast::cluster` + u16 mcast_mask operand. + All CTAs in the mcast group issue with identical params; HW coalesces + into one GMEM load + broadcast. Each CTA's mbar receives full tx bytes. + """ + exec_atom = _cute_nvgpu_ir.atom_make_exec_tma(tma_atom._trait.value, loc=loc, ip=ip) + desc_ptr_ty = _cute_ir.PtrType.get( + _cute_nvgpu_ir.TmaDescriptorTiledType.get(), + AddressSpace.generic, + 64, + ) + desc_cute_ptr = _cute_nvgpu_ir.get_tma_desc_addr(desc_ptr_ty, exec_atom, loc=loc, ip=ip) + desc_i64 = desc_cute_ptr.toint().ir_value(loc=loc, ip=ip) + + smem_dst_int = cutlass.Int32(smem_dst_ptr.toint()) + mbar_int = cutlass.Int32(mbar_ptr.toint()) + if use_cta_group_2: + mbar_int = mbar_int & cutlass.Int32(_PEER_BIT_MASK) + smem_dst_i32 = smem_dst_int.ir_value(loc=loc, ip=ip) + mbar_i32 = mbar_int.ir_value(loc=loc, ip=ip) + col_i32 = cutlass.Int32(col).ir_value(loc=loc, ip=ip) + r0_i32 = cutlass.Int32(r0).ir_value(loc=loc, ip=ip) + r1_i32 = cutlass.Int32(r1).ir_value(loc=loc, ip=ip) + r2_i32 = cutlass.Int32(r2).ir_value(loc=loc, ip=ip) + r3_i32 = cutlass.Int32(r3).ir_value(loc=loc, ip=ip) + cache_hint_i64 = cutlass.Int64(0).ir_value(loc=loc, ip=ip) + + use_mcast = mcast_mask is not None + if use_mcast: + mcast_mask_i16 = cutlass.Int16(mcast_mask).ir_value(loc=loc, ip=ip) + + if use_cta_group_2 and use_mcast: + asm = ( + "cp.async.bulk.tensor.2d.shared::cluster.global" + ".tile::gather4.mbarrier::complete_tx::bytes.multicast::cluster" + ".cta_group::2" + " [$0], [$1, {$3, $4, $5, $6, $7}], [$2], $8;" + ) + operands = [ + smem_dst_i32, + desc_i64, + mbar_i32, + col_i32, + r0_i32, + r1_i32, + r2_i32, + r3_i32, + mcast_mask_i16, + ] + constraints = "r, l, r, r, r, r, r, r, h" + elif use_cta_group_2: + asm = ( + "cp.async.bulk.tensor.2d.shared::cluster.global" + ".tile::gather4.mbarrier::complete_tx::bytes.L2::cache_hint.cta_group::2" + " [$0], [$1, {$3, $4, $5, $6, $7}], [$2], $8;" + ) + operands = [ + smem_dst_i32, + desc_i64, + mbar_i32, + col_i32, + r0_i32, + r1_i32, + r2_i32, + r3_i32, + cache_hint_i64, + ] + constraints = "r, l, r, r, r, r, r, r, l" + elif use_mcast: + asm = ( + "cp.async.bulk.tensor.2d.shared::cluster.global" + ".tile::gather4.mbarrier::complete_tx::bytes.multicast::cluster" + " [$0], [$1, {$3, $4, $5, $6, $7}], [$2], $8;" + ) + operands = [ + smem_dst_i32, + desc_i64, + mbar_i32, + col_i32, + r0_i32, + r1_i32, + r2_i32, + r3_i32, + mcast_mask_i16, + ] + constraints = "r, l, r, r, r, r, r, r, h" + else: + asm = ( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".tile::gather4.mbarrier::complete_tx::bytes.L2::cache_hint" + " [$0], [$1, {$3, $4, $5, $6, $7}], [$2], $8;" + ) + operands = [ + smem_dst_i32, + desc_i64, + mbar_i32, + col_i32, + r0_i32, + r1_i32, + r2_i32, + r3_i32, + cache_hint_i64, + ] + constraints = "r, l, r, r, r, r, r, r, l" + + llvm.inline_asm( + None, + operands, + asm, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@dsl_user_op +def sm100_tcgen05_st_32x32b_x4( + tmem_addr, + r0, + r1, + r2, + r3, + *, + loc=None, + ip=None, +): + """Issue one tcgen05.st.sync.aligned.32x32b.x4.b32. + + Writes 4 32-bit cells per lane to TMEM[lane_offset, col_base..col_base+3]. + Used by SFA transform warps to write LDS+repacked SF data into TMEM, + bypassing cute.copy auto-partition. + """ + addr_i32 = cutlass.Uint32(tmem_addr).ir_value(loc=loc, ip=ip) + r0_i32 = cutlass.Uint32(r0).ir_value(loc=loc, ip=ip) + r1_i32 = cutlass.Uint32(r1).ir_value(loc=loc, ip=ip) + r2_i32 = cutlass.Uint32(r2).ir_value(loc=loc, ip=ip) + r3_i32 = cutlass.Uint32(r3).ir_value(loc=loc, ip=ip) + asm = "tcgen05.st.sync.aligned.32x32b.x4.b32 [$0], {$1, $2, $3, $4};" + llvm.inline_asm( + None, + [addr_i32, r0_i32, r1_i32, r2_i32, r3_i32], + asm, + "r, r, r, r, r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py new file mode 100644 index 000000000000..c26f544877b8 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_gather_grouped_blockscaled_gemm_act_fusion.py @@ -0,0 +1,5058 @@ +# Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import os +import re +from typing import NamedTuple, Optional, Tuple, Type, Union + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.torch as cutlass_torch +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +import cutlass.utils.rubin_helpers as sm107_utils +import torch +from cutlass._mlir.dialects import math, nvvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu.tcgen05.mma import CollectorOp +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.utils.gemm.sm100 import ( + epilogue_smem_copy_and_partition, + transform_partitioned_tensor_layout, +) + +from ....utils import ActivationType, is_gated_activation + +try: + from .custom_pipeline import PipelineCpAsyncUmma + from .inline_ptx import sm100_tcgen05_st_32x32b_x4, sm100_tma_gather4_load + from .utils import TRTLLM_ENABLE_PDL +except ImportError: + from custom_pipeline import PipelineCpAsyncUmma + from inline_ptx import sm100_tcgen05_st_32x32b_x4, sm100_tma_gather4_load + from utils import TRTLLM_ENABLE_PDL + + +# ============================================================================ +# Inline utility functions +# ============================================================================ + + +@dsl_user_op +def fmin( + a: Union[float, cutlass.Float32], + b: Union[float, cutlass.Float32], + *, + nan=False, + loc=None, + ip=None, +) -> cutlass.Float32: + return cutlass.Float32( + nvvm.fmin( + cutlass.Float32(a).ir_value(loc=loc, ip=ip), + cutlass.Float32(b).ir_value(loc=loc, ip=ip), + nan=nan, + loc=loc, + ip=ip, + ) + ) + + +def sigmoid_f32( + a: Union[float, cutlass.Float32], fastmath: bool = False +) -> Union[float, cutlass.Float32]: + """Compute the sigmoid of the input tensor.""" + return cute.arch.rcp_approx(1.0 + cute.math.exp(-a, fastmath=fastmath)) + + +def silu_f32( + a: Union[float, cutlass.Float32], fastmath: bool = False +) -> Union[float, cutlass.Float32]: + """Compute the silu of the input tensor.""" + return a * sigmoid_f32(a, fastmath=fastmath) + + +SUPPORTED_ACTIVATION_TYPES = (ActivationType.Swiglu, ActivationType.Relu2) + + +def validate_activation_type(activation_type) -> ActivationType: + """Normalize and validate the fused FC1 activation.""" + activation_type = ActivationType(int(activation_type)) + assert activation_type in SUPPORTED_ACTIVATION_TYPES, ( + f"Unsupported activation type {activation_type}; " + f"expected one of {SUPPORTED_ACTIVATION_TYPES}" + ) + return activation_type + + +class S2TCopyBundle(NamedTuple): + """Bundle of tiled copy and partitioned tensors for smem-to-tmem copies.""" + + tiled_copy: cute.TiledCopy + sSF_compact: cute.Tensor # Partitioned source (smem) + tSF_compact: cute.Tensor # Partitioned destination (tmem) + + +""" +Rubin (SM107) persistent blockscaled contiguous grouped GEMM with token gather +and fused SwiGLU or Relu2 activation (FC1 of MoE). + +Compute: + acc = alpha * (SFA * A[token_ids]) * (SFB * B) # GEMM + C = up * silu(gate) or relu(acc)^2 # selected activation + + optional NVFP4 quantization (generates SFC) when c_dtype == Float4E2M1FN. + +Shapes: A is M×K×1; B is N×K×L (L = num experts). SwiGLU uses interleaved +[up, gate] weights at granularity=64 and produces M×(N/2)×1; Relu2 uses +plain weights and produces M×N×1. SFA/SFB layouts follow BlockScaledBasicChunk. +token_id_mapping drives the row gather for A/SFA; token_id == -1 marks padding. + +Within a tile, valid_m varies per group; padding rows are handled at load: +TMA gather4 passes -1 to zero-fill; CpAsync predicates on `abs_row < mn_limit`. + +Constraints: A/B share dtype (mxf8 | mxf4 | nvf4); mma_tiler M in {128, 256}; +mma_tiler N in {64, 128, 192, 256}; cluster M/N pow-2, total ≤ 16; +contiguous dim ≥ 16B aligned (16/32 elems for f8/f4). + +For CUDA graph, A/C/SFA/token_id_mapping/tile_idx_to_expert_idx can be padded +to permuted_m; padded tiles are filtered by the scheduler. +""" + + +class Sm107BlockScaledContiguousGatherGroupedGemmActFusionKernel: + """Rubin (SM107) FC1: contiguous grouped blockscaled GEMM with token + gather on A/SFA and activation fusion in the epilogue. + + Supports both SwiGLU (gated) and Relu2 (non-gated) activations. + + Builds on Sm107BlockScaledContiguousGroupedGemmKernel (persistent tile + scheduling, warp specialization, B-reuse, tcgen05.mma block-scale, TMA + B/SFB with M-multicast, per-group alpha). Refer to backbone for those. + + Additions on top of backbone: + - Token gather: A/SFA rows are gathered by token_id_mapping + (token_id == -1 marks padding rows). + - A load path (knob `a_path`): + * cpasync — CpAsync128.CG per-thread (default); separate + a_pipeline; in 2CTA, warp 11 relays per-CTA a_pipeline to a + cluster-wide a_sync_transform_pipeline so MMA cta_group::2 sees + both CTAs' A. + * tma — TMA gather4 with HW multicast; A and B share a single + merged ab_pipeline (no relay warp needed). + SFA is always loaded via CpAsync128.CG, then reorganized into SFA + TMEM by transform warps via LDS + STTM (sfa_transform_pipeline). + - SwiGLU epilogue: C = up * silu(gate), where up/gate come from + interleaved accumulator at granularity=64 → output N is halved. + - Optional NVFP4 quant: when c_dtype == Float4E2M1FN, the epilogue + also generates SFC and quantizes the output. + + Extra warp roles (20 warps total; 4-19 are FC1-only): + - 0-3 epilogue (LDTM → SwiGLU → optional quant → TMA store) + - 4-7 gather A (CpAsync128.CG or TMA gather4) + - 8 MMA + - 9 TMA B / SFB + - 10 scheduler + - 11 cpasync 2CTA A sync-transform relay (idle on 1CTA / tma) + - 12-15 gather SFA (CpAsync128.CG) + - 16-19 SFA transform (LDS + STTM into SFA TMEM) + + :param sf_vec_size: Scale factor vector size (16 or 32). + :param mma_inst_shape: MMA instruction shape (M, N, K). + :param mma_tiler: MMA tiler shape (M, N, K). + :param cluster_shape_mn: Cluster dimensions (M, N). + :param vectorized_f32: Use vectorized f32x2 ops in epilogue. + :param topk: Experts selected per token. + :param raster_along_m: If True, raster tiles along M first. + :param a_path: "cpasync" or "tma" — A load implementation. + """ + + def __init__( + self, + sf_vec_size: int, + mma_inst_shape: Tuple[int, int, int], + mma_tiler: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + vectorized_f32: bool, + topk: cutlass.Int64, + raster_along_m: bool = False, + a_path: str = "cpasync", + use_pdl: bool = True, + locality_domain_half_gemm: bool = False, + activation_type: ActivationType = ActivationType.Swiglu, + ): + self.a_path = a_path + # locality domain half-GEMM: two partitions write their N-half into a shared + # full-width C/SFC buffer at a column offset (see wrapper/__call__). + self.locality_domain_half_gemm = locality_domain_half_gemm + self.sf_vec_size = sf_vec_size + self.topk = topk + self.activation_type = validate_activation_type(activation_type) + self.is_gated = is_gated_activation(self.activation_type) + if locality_domain_half_gemm and not self.is_gated: + raise ValueError("Rubin locality domain half-GEMM currently supports SwiGLU only") + self.acc_dtype = cutlass.Float32 + self.mma_inst_shape = mma_inst_shape + self.mma_tiler = mma_tiler + self.cluster_shape_mn = cluster_shape_mn + self.raster_along_m = raster_along_m + # Honor the TRTLLM_ENABLE_PDL env flag (PDL on by default); a caller + # passing use_pdl=False still disables PDL. + self.use_pdl = use_pdl and TRTLLM_ENABLE_PDL + + self.use_2cta_instrs = mma_inst_shape[0] == 256 + self.cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + self.arch = "sm_107" + self.smem_capacity = utils.get_smem_capacity_in_bytes(self.arch) + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols(self.arch) + + self.occupancy = 1 + self.epilog_warp_id = (0, 1, 2, 3) + self.gather_a_warp_id = ( + 4, + 5, + 6, + 7, + ) + self.mma_warp_id = 8 + self.tma_b_warp_id = 9 + self.sched_warp_id = 10 + # Warp 11: cpasync 2CTA A peer-sync relay (sync_transform_warp_id) / + # idle on tma-A (dummy_warp_id). Slot reserved for SM occupancy. + # sync_transform_warp_id is always defined so downstream traces + # resolve in tma mode (body is const-gated out). + self.sync_transform_warp_id = 11 + if self.a_path == "tma": + self.dummy_warp_id = 11 + self.gather_sfa_warp_id = ( + 12, + 13, + 14, + 15, + ) + # 4 SFA transform warps (LDS source + STTM destination into SFA TMEM). + self.sfa_transform_warp_id = ( + 16, + 17, + 18, + 19, + ) + # Register reconfig (setmaxnreg) per warpgroup. Default 128/thread. + self.num_regs_epilogue_warps = 168 + self.num_regs_gather_a_warps = 80 + self.num_regs_gather_sfa_warps = 80 + self.num_regs_sfa_transform_warps = 48 + self.num_regs_mma_group_warps = 128 + self.threads_per_warp = 32 + # warp 11 slot is always counted in threads_per_cta (SM occupancy) + # regardless of cpasync-A peer-sync role vs TMA-A idle role. + _warp11_id = self.sync_transform_warp_id if self.a_path == "cpasync" else self.dummy_warp_id + self.threads_per_cta = self.threads_per_warp * len( + ( + self.mma_warp_id, + *self.gather_a_warp_id, + self.tma_b_warp_id, + *self.epilog_warp_id, + self.sched_warp_id, + _warp11_id, + *self.gather_sfa_warp_id, + *self.sfa_transform_warp_id, + ) + ) + # warps_wo_sched = tile_info_pipeline consumers (all warps except + # scheduler). Warp 11 counts only on cpasync 2CTA (relay needs tiles); + # excluded on 1CTA or tma-A (warp 11 idle). + if self.use_2cta_instrs and self.a_path == "cpasync": + _wo_sched_warps = ( + *self.epilog_warp_id, + self.mma_warp_id, + self.tma_b_warp_id, + self.sync_transform_warp_id, + *self.gather_a_warp_id, + *self.gather_sfa_warp_id, + *self.sfa_transform_warp_id, + ) + else: + _wo_sched_warps = ( + *self.epilog_warp_id, + self.mma_warp_id, + self.tma_b_warp_id, + *self.gather_a_warp_id, + *self.gather_sfa_warp_id, + *self.sfa_transform_warp_id, + ) + self.warps_wo_sched = len(_wo_sched_warps) + self.threads_wo_sched = self.threads_per_warp * self.warps_wo_sched + + # Set barrier for cta sync, epilogue sync and tmem ptr sync + self.cta_sync_barrier = pipeline.NamedBarrier( + barrier_id=1, + num_threads=self.threads_per_cta, + ) + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=32 * len(self.epilog_warp_id), + ) + # tmem_alloc_barrier participants: epi (allocator) + mma (consumer) + # + transform warps (STTM producers, need TMEM ptr to write SFA). + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=3, + num_threads=32 + * len( + ( + self.mma_warp_id, + *self.epilog_warp_id, + *self.sfa_transform_warp_id, + ) + ), + ) + self.sched_sync_barrier = pipeline.NamedBarrier( + barrier_id=4, + num_threads=self.threads_per_warp, + ) + + self.num_smem_capacity = self.smem_capacity + # num_tmem_alloc_cols already set in __init__ + + self.vectorized_f32 = vectorized_f32 + + # For epilogue compatibility + self.epilogue_warp_id = self.epilog_warp_id + + # B-reuse pattern control + self.enable_breuse = True if mma_tiler[0] // mma_inst_shape[0] == 2 else False + + # Overlapping ACC TMEM: acc[0]/acc[1] share 64 cols. Epilogue + # iterates the overlap region first (reverse for acc[0]) and + # early-releases so MMA can write the next stage. Frees enough TMEM + # for 4 SFA stages. Auto-on for non-breuse cta_tile_N=256. + self.use_overlap_accum = (not self.enable_breuse) and (mma_tiler[1] == 256) + + def _setup_attributes(self): + """Set up configurations that are dependent on GEMM inputs + + This method configures various attributes based on the input tensor properties + (data types, leading dimensions) and kernel settings: + - Configuring tiled MMA + - Computing MMA/cluster/tile shapes + - Computing cluster layout + - Computing multicast CTAs for A/B + - Computing epilogue subtile + - Setting up A/B/C stage counts in shared memory + - Computing A/B/C shared memory layout + - Computing tensor memory allocation columns + """ + + self.mma_inst_shape_sfb = ( + self.mma_inst_shape[0] // (2 if self.use_2cta_instrs else 1), + cute.round_up(self.mma_inst_shape[1], 128), + self.mma_inst_shape[2], + ) + + # Configure tiled mma (Rubin SM107) + tiled_mma = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + atom_layout_mnk=(1, 1, 1), + permutation_mnk=self._get_mma_permutation_mnk(), + ) + + tiled_mma_sfb = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_sfb, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + ) + + # Compute mma/cluster/tile shapes + self.mma_tiler_sfb = ( + self.mma_inst_shape_sfb[0], + self.mma_inst_shape_sfb[1], + self.mma_tiler[2], + ) + + self.mma_tiler_c = ( + self.mma_tiler[0], + self.mma_tiler[1] // 2 if self.is_gated else self.mma_tiler[1], + self.mma_tiler[2], + ) + + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + + # Number of CpAsync128.CG loads per thread for A matrix (each loads 16 M-rows) + self.a_num_loads = self.cta_tile_shape_mnk[0] // 16 + + self.cta_tile_shape_mnk_sfb = ( + self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_sfb[1], + self.mma_tiler_sfb[2], + ) + + self.cta_tile_shape_mnk_c = ( + self.mma_tiler_c[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_c[1], + self.mma_tiler_c[2], + ) + + # Compute SFA tiler for CpAsync gather (use mma_inst_shape for M/N, scaled K for SF) + mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2]) + mma_inst_tile_k = self.mma_tiler[2] // mma_inst_shape_k + self.mma_tiler_sfa = ( + self.mma_inst_shape[0], + self.mma_inst_shape[1], + mma_inst_shape_k * mma_inst_tile_k // 16, + ) + self.cta_tile_shape_mnk_sfa = ( + self.mma_tiler_sfa[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_sfa[1], + self.mma_tiler_sfa[2], + ) + + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma_sfb.thr_id.shape,), + ) + + # Compute number of multicast CTAs for A/B + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.is_b_mcast = self.num_mcast_ctas_b > 1 + # A multicast: cluster_N CTAs share A along N dim. Only meaningful when + # cluster_N > 1. SFA multicast intentionally NOT enabled — was buggy + # on cta_tile_N >= 256 and not worth the complexity. + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + + # Fixed epilogue tile (128, 64). SwiGLU halves N, so the default + # SM107_TILES lookup (keyed on full cta_n) can pick epi_tile_n too + # small (wrong TMA store strides + insufficient SFC for cvt_fptrunc + # 32-bit alignment). (128, 64) works for all configs. + self.epi_tile = (128, 64) + self.epi_tile_n = cute.size(self.epi_tile[1]) + self.epi_tile_cnt = ( + self.cta_tile_shape_mnk_c[0] // cute.size(self.epi_tile[0]), + self.cta_tile_shape_mnk_c[1] // cute.size(self.epi_tile[1]), + ) + + # Setup A/B/C/Scale stage count in shared memory and ACC stage count in tensor memory + ( + self.num_acc_stage, + self.num_ab_stage, + self.num_c_stage, + self.num_tile_stage, + ) = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.cta_tile_shape_mnk, + self.a_dtype, + self.b_dtype, + self.epi_tile, + self.c_dtype, + self.c_layout, + self.sf_dtype, + self.sf_vec_size, + self.smem_capacity, + self.occupancy, + self.enable_breuse, + ) + + # Compute A/B/C/Scale shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.num_ab_stage, + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, + self.mma_tiler, + self.b_dtype, + self.num_ab_stage, + ) + # Canonical SFA SMEM layout (from blockscaled_utils), used only to + # derive tCtSFA_layout below. The actual SFA SMEM uses the linear + # layout built next; this canonical one isn't allocated. + sfa_canon_smem_layout = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + self.num_ab_stage, + ) + + # SFA SMEM is plain linear (M_per_cta, tile_K_sf, stage), no pad. + # Each thread does one CpAsync128.CG (16B = tile_K_sf=16 × FP8) per row. + # Layout exposes (row, k_sf_byte, stage) with byte strides. + sfa_tile_k_sf = self.cta_tile_shape_mnk[2] // self.sf_vec_size + sf_bytes_per_row = sfa_tile_k_sf * self.sf_dtype.width // 8 + sfa_bytes_per_stage = self.cta_tile_shape_mnk[0] * sf_bytes_per_row + self.sfa_smem_layout_staged = cute.make_layout( + (self.cta_tile_shape_mnk[0], sfa_tile_k_sf, self.num_ab_stage), + stride=(sf_bytes_per_row, 1, sfa_bytes_per_stage), + ) + self.sfa_smem_alloc_bytes = self.num_ab_stage * sfa_bytes_per_stage + + self.c_smem_layout_staged = sm100_utils.make_smem_layout_epi( + self.c_dtype, + self.c_layout, + self.epi_tile, + self.num_c_stage, + ) + + # Compute TMEM layouts for SFA/SFB (Rubin precomputed) + self.tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(sfa_canon_smem_layout, (None, None, None, 0)), + ) + self.tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)), + ) + + # Compute TMEM column counts. + # SFA TMEM holds num_sfa_tmem_stage stages, each tCtSFA_layout wide. + # TMEM layout: [acc | sfa (N stages) | sfb]. + self.num_sfa_tmem_cols_per_stage = ( + cute.cosize(cute.recast_layout(32, self.sf_dtype.width, self.tCtSFA_layout)) + & 0x0000FFFF + ) + # SFA TMEM stages: 4 for non-breuse, 1 for breuse (576-col TMEM + # already saturated with 1-stage SFA at 32 cols). At non-breuse + # N=256, use_overlap_accum is auto-on to free TMEM for 4 stages. + self.num_sfa_tmem_stage = 1 if self.enable_breuse else 4 + self.num_sfa_tmem_cols = self.num_sfa_tmem_cols_per_stage * self.num_sfa_tmem_stage + self.num_sfb_tmem_cols = ( + cute.cosize(cute.recast_layout(32, self.sf_dtype.width, self.tCtSFB_layout)) + & 0x0000FFFF + ) + self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols + # use_overlap_accum: pipeline tracks 1 acc stage; physically 2 stages + # share TMEM with one epi_tile_n of overlap. Otherwise: + # tile_N × num_acc_stage × (2 if breuse). + if self.use_overlap_accum: + self.num_acc_stage = 1 # logical 2 via overlap + self.num_accumulator_tmem_cols = self.cta_tile_shape_mnk[1] * 2 - self.epi_tile_n + else: + self.num_accumulator_tmem_cols = ( + self.cta_tile_shape_mnk[1] * self.num_acc_stage * (2 if self.enable_breuse else 1) + ) + # SFA TMEM offset (cols, 32-bit each): right after acc. + self.sfa_tmem_offset = self.num_accumulator_tmem_cols + # Validation: 512 + 32 + 32 = 576 (exact fit for main target on sm_107) + _total_used = ( + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols + self.num_sfb_tmem_cols + ) + if _total_used > self.num_tmem_alloc_cols: + raise ValueError( + f"TMEM overflow: acc({self.num_accumulator_tmem_cols}) + " + f"sfa({self.num_sfa_tmem_cols}) + " + f"sfb({self.num_sfb_tmem_cols}) = {_total_used} > " + f"max {self.num_tmem_alloc_cols}" + ) + + def _get_mma_permutation_mnk(self): + if cutlass.const_expr(self.use_2cta_instrs and self.enable_breuse): + m_layout = cute.make_layout( + shape=(self.mma_inst_shape[0] // 2, 2, 2), + stride=(1, self.mma_inst_shape[0], self.mma_inst_shape[0] // 2), + ) + return (m_layout, self.mma_inst_shape[1], self.mma_inst_shape[2]) + else: + return (1, 1, 1) + + def _is_interleaved_utccp(self) -> bool: + """Enable interleaving UTCCP for Bkeep-Breuse case for 4xFP4 kernel.""" + return self.a_dtype.width == 4 and self.b_dtype.width == 4 and self.enable_breuse + + def _mainloop_s2t_copy_and_partition( + self, + sSF: cute.Tensor, + tSF: cute.Tensor, + ) -> S2TCopyBundle: + """Make tiledCopy for smem to tmem load for scale factor tensor.""" + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + + def appendMNBroadcastMode(smem_layout: cute.Layout): + mn_dim = cute.get(smem_layout, mode=[0, 0]) + mn_dim = cute.append(mn_dim, cute.make_layout((4), stride=(0))) + layout = cute.append(cute.group_modes(mn_dim, 0), cute.get(smem_layout, mode=[0, 1])) + layout = cute.append(cute.group_modes(layout, 0), cute.get(smem_layout, mode=[1])) + layout = cute.append(layout, cute.get(smem_layout, mode=[2])) + layout = cute.append(layout, cute.get(smem_layout, mode=[3])) + return layout + + tCsSF_compact_bcast = cute.make_tensor( + tCsSF_compact.iterator, appendMNBroadcastMode(tCsSF_compact.layout) + ) + + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact_bcast) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + + return S2TCopyBundle(tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) + + def _mainloop_s2t_copies( + self, + stage_idx: int, + sfb_s2t_bundle: S2TCopyBundle, + ): + """Copy SFB from smem to tmem (UTCCP). SFA path now uses LDS+STTM + from transform warps, no UTCCP needed here.""" + s2t_stage_coord = (None, None, None, None, stage_idx) + + cute.copy( + sfb_s2t_bundle.tiled_copy, + sfb_s2t_bundle.sSF_compact[s2t_stage_coord], + sfb_s2t_bundle.tSF_compact, + ) + + def _mainloop_s2t_interleaved_copies( + self, + k_block: int, + stage_idx: int, + sfa_s2t_bundle: S2TCopyBundle, + sfb_s2t_bundle: S2TCopyBundle, + ): + """Interleaved UTCCP for Bkeep-Breuse pattern.""" + s_sfa_crd_keep = (None, 0, None, k_block, stage_idx) + s_sfa_crd_reuse = (None, 1, None, k_block, stage_idx) + s_sfb_crd = (None, None, None, k_block, stage_idx) + + t_sfa_crd_keep = (None, 0, None, k_block) + t_sfa_crd_reuse = (None, 1, None, k_block) + t_sfb_crd = (None, None, None, k_block) + + cute.copy( + sfa_s2t_bundle.tiled_copy, + sfa_s2t_bundle.sSF_compact[s_sfa_crd_keep], + sfa_s2t_bundle.tSF_compact[t_sfa_crd_keep], + ) + cute.copy( + sfb_s2t_bundle.tiled_copy, + sfb_s2t_bundle.sSF_compact[s_sfb_crd], + sfb_s2t_bundle.tSF_compact[t_sfb_crd], + ) + cute.copy( + sfa_s2t_bundle.tiled_copy, + sfa_s2t_bundle.sSF_compact[s_sfa_crd_reuse], + sfa_s2t_bundle.tSF_compact[t_sfa_crd_reuse], + ) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + sfa: cute.Tensor, + sfb: cute.Tensor, + sfc_tensor: Optional[cute.Tensor], + full_c_shape: Optional[cute.Shape], + norm_const_tensor: Optional[cute.Tensor], + tile_idx_to_expert_idx: cute.Tensor, + tile_idx_to_mn_limit: cute.Tensor, + token_id_mapping_tensor: cute.Tensor, + num_non_exiting_tiles: cute.Tensor, + alpha: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + c_sf_n_tile_offset: cutlass.Int64 = cutlass.Int64(0), + ): + """Execute the contiguous grouped GEMM with gather operation and SwiGLU fusion. + + This method performs FC1 layer computation: + 1. GEMM: acc = alpha * (SFA * A[token_ids]) * (SFB * B) + 2. SwiGLU: C = up * silu(gate), where up/gate are extracted from interleaved acc (granularity=64) + 3. Optional Quant: When c_dtype is Float4E2M1FN, generates SFC and quantizes output + + Data loading: + - A and SFA are loaded using CpAsync instructions with token-based gather + - B and SFB are loaded using TMA instructions with multicast + - B weights are interleaved: [up_0:64, gate_64:128, up_128:192, gate_192:256, ...] + + Execution steps: + 1. Setup static attributes before smem/grid computation + 2. Setup TMA load/store atoms for B, SFB, and C (no TMA for A/SFA) + 3. Compute grid size with regard to hardware constraints + 4. Define shared storage for kernel + 5. Launch the kernel synchronously with warp specialization: + - Scheduler warp: Dispatches tile information + - CpAsync warps: Load A and SFA with gather + - A Sync Transform warps: Transform the sync signal of A and SFA from global to + shared memory when use_2cta_instrs is True + - TMA warp: Load B and SFB with multicast + - MMA warp: Perform matrix multiply-accumulate + - Epilogue warps: Apply SwiGLU activation, optional quantization, and store results + + :param a: Input tensor A (MxKx1), will be gathered using token_id_mapping + :type a: cute.Tensor + :param b: Input tensor B (NxKxL), L is the number of experts/groups, weights are interleaved for SwiGLU + :type b: cute.Tensor + :param c: Output tensor C (Mx(N/2)x1), N is halved due to SwiGLU fusion + :type c: cute.Tensor + :param sfa: Scale factor tensor A, will be gathered using token_id_mapping + :type sfa: cute.Tensor + :param sfb: Scale factor tensor B + :type sfb: cute.Tensor + :param sfc_tensor: Scale factor tensor C for quantized output (None if not quantizing) + :type sfc_tensor: Optional[cute.Tensor] + :param norm_const_tensor: Normalization constant for scale factor generation + (None if not quantizing) + :type norm_const_tensor: Optional[cute.Tensor] + :param tile_idx_to_expert_idx: Mapping from tile index to expert ID, + shape (permuted_m/cta_tile_m,) where cta_tile_m is the CTA tile M size + :type tile_idx_to_expert_idx: cute.Tensor + :param tile_idx_to_mn_limit: Mapping from tile index to M-N dimension limit + for boundary checking, shape (permuted_m/cta_tile_m,) + :type tile_idx_to_mn_limit: cute.Tensor + :param token_id_mapping_tensor: Token ID mapping for gather operation, shape (permuted_m,) + :type token_id_mapping_tensor: cute.Tensor + :param num_non_exiting_tiles: Number of valid tiles to process (valid_m/cta_tile_m), shape (1,) + :type num_non_exiting_tiles: cute.Tensor + :param alpha: Alpha tensor for each group + :type alpha: cute.Tensor + :param max_active_clusters: Maximum number of active clusters + :type max_active_clusters: cutlass.Constexpr + :param stream: CUDA stream for asynchronous execution + :type stream: cuda.CUstream + :param epilogue_op: Optional elementwise lambda function to apply to the output tensor + :type epilogue_op: cutlass.Constexpr + :raises TypeError: If input data types are incompatible with the MMA instruction. + """ + # Setup static attributes before smem/grid/tma computation + self.a_dtype: Type[cutlass.Numeric] = a.element_type + self.b_dtype: Type[cutlass.Numeric] = b.element_type + self.c_dtype: Type[cutlass.Numeric] = c.element_type + self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.from_tensor(c) + + # Note: Rubin supports mixed A/B dtypes (e.g., Float8E4M3FN x Float8E5M2) + + # Setup attributes that dependent on gemm inputs + self._setup_attributes() + + # Setup sfb tensor by filling B tensor to scale factor atom layout + # ((Atom_N, Rest_N),(Atom_K, Rest_K),RestL) + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b.shape, self.sf_vec_size) + sfb = cute.make_tensor(sfb.iterator, sfb_layout) + + # Setup sfc tensor by filling C tensor to scale factor atom layout. + # For locality domain, full_c_shape carries the full N dimension so sfc gets the + # correct M-tile stride (two locality domains write their N-half into the shared + # SF buffer without copy-back); None → use c.shape (non-locality domain). + self.generate_sfc = sfc_tensor is not None and norm_const_tensor is not None + if cutlass.const_expr(self.generate_sfc): + sfc_shape = c.shape if full_c_shape is None else full_c_shape + sfc_layout = blockscaled_utils.tile_atom_to_shape_SF(sfc_shape, self.sf_vec_size) + sfc_tensor = cute.make_tensor(sfc_tensor.iterator, sfc_layout) + + atom_layout_mnk = (1, 1, 1) + permutation_mnk = self._get_mma_permutation_mnk() + + tiled_mma = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + atom_layout_mnk=atom_layout_mnk, + permutation_mnk=permutation_mnk, + ) + tiled_mma.set(tcgen05.Field.NEGATE_A, False) + tiled_mma.set(tcgen05.Field.NEGATE_B, False) + + # For 2CTA blockscaled kernels, SFB needs to be replicated across peer CTAs. + tiled_mma_sfb = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_sfb, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + ) + tiled_mma_sfb.set(tcgen05.Field.NEGATE_A, False) + tiled_mma_sfb.set(tcgen05.Field.NEGATE_B, False) + + tiled_mma_bkeep = None + tiled_mma_breuse = None + if cutlass.const_expr(self.enable_breuse): + tiled_mma_bkeep = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.FILL, + atom_layout_mnk=atom_layout_mnk, + permutation_mnk=permutation_mnk, + ) + tiled_mma_bkeep.set(tcgen05.Field.NEGATE_A, False) + tiled_mma_bkeep.set(tcgen05.Field.NEGATE_B, False) + + tiled_mma_breuse = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.LASTUSE, + atom_layout_mnk=atom_layout_mnk, + permutation_mnk=permutation_mnk, + ) + tiled_mma_breuse.set(tcgen05.Field.NEGATE_A, False) + tiled_mma_breuse.set(tcgen05.Field.NEGATE_B, False) + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # cpasync-A: CpAsync128.CG gmem → sA SMEM, no TMA. 4 gather_a warps × 32 + # threads issue cp.async.cg.16B per (token row, k chunk). a_num_loads + # (in _setup_attributes) controls CpAsync iterations per thread per k_tile. + tma_atom_a = None + tma_tensor_a = None + # tma-A mode: build a 2D gather4 TMA atom for A (box_rows = 1, + # SW128-permuted base) so gather4 writes sA at the same SW128 offsets + # UMMA's K_SW128 reads expect. + if cutlass.const_expr(self.a_path == "tma"): + a_2d = a[(None, None, 0)] + a_gather_base = cute.make_layout( + (1, self.mma_tiler[2]), + stride=(self.mma_tiler[2], 1), + ) + sw128 = cute.make_swizzle(3, 4, 3) # K_SW128 + a_gather_smem_layout = cute.make_composed_layout(sw128, 0, a_gather_base) + tma_atom_a, tma_tensor_a = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), + a_2d, + a_gather_smem_layout, + (1, self.mma_tiler[2]), # cta_tiler, box_rows = 1 + ) + # tx_count split evenly across the 4 gather warps. Total = full A + # tile bytes (× atom_thr_size for 2CTA's leader-mbar collapse). + self.tma_gather_num_warps = len(self.gather_a_warp_id) + self.a_num_tma_load_bytes_total = ( + self.cta_tile_shape_mnk[0] + * self.cta_tile_shape_mnk[2] + * self.a_dtype.width + // 8 + * atom_thr_size + ) + self.a_num_tma_load_bytes = self.a_num_tma_load_bytes_total // self.tma_gather_num_warps + + # cpasync-A SFA path is also CpAsync128.CG (no TMA). gather_sfa warps + # issue cp.async per row; one row per thread (128 rows / 128 threads + # in 4 warps). No tma_atom_sfa needed in either A mode. + + # Setup TMA load for B + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # Setup TMA load for SFB + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma.thr_id) + sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + sfb, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # This modifies the layout to handle overlapping 256x(# of scale factors for a single column of B (nNSF)) + # logical blocks for SFB when cta_tile_shape_n=192. + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): + x = tma_tensor_sfb.stride[0][1] + y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4) + + new_shape = ( + (tma_tensor_sfb.shape[0][0], ((2, 2), y)), + tma_tensor_sfb.shape[1], + tma_tensor_sfb.shape[2], + ) + # Use right multiplication for ScaledBasis (3 * x instead of x * 3) + x_times_3 = 3 * x + new_stride = ( + (tma_tensor_sfb.stride[0][0], ((x, x), x_times_3)), + tma_tensor_sfb.stride[1], + tma_tensor_sfb.stride[2], + ) + tma_tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride) + tma_tensor_sfb = cute.make_tensor(tma_tensor_sfb.iterator, tma_tensor_sfb_new_layout) + + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = (b_copy_size + sfb_copy_size) * atom_thr_size + + # Setup TMA store for C + tma_atom_c = None + tma_tensor_c = None + epi_smem_layout = cute.slice_(self.c_smem_layout_staged, (None, None, 0)) + tma_atom_c, tma_tensor_c = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + c, + epi_smem_layout, + self.epi_tile, + ) + + # Compute grid size + self.tile_sched_params, grid = self._compute_grid( + c, + self.cta_tile_shape_mnk_c, + self.cluster_shape_mn, + max_active_clusters, + self.raster_along_m, + ) + + self.buffer_align_bytes = 1024 + + # Define shared storage for kernel. + @cute.struct + class SharedStorageCpasync1cta: + sInfo: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int32, 5 * self.num_tile_stage], + 1, + ] + # cpasync mode: A and B use separate pipelines (CpAsync A is + # CpAsync type, B is TmaUmma type — they can't share one mbar). + # Each mbar set holds num_ab_stage * 2 (full + empty per stage). + a_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + b_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + sfa_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + sfa_transform_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_sfa_tmem_stage * 2] + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + tile_info_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_tile_stage * 2] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + cute.cosize(self.c_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + # sSFA placed BEFORE sA so SFA gets a low SMEM offset: CpAsync128.CG + # destination addr stays under the 248KB threshold (above which the + # .CG cache-mode hint would degrade to a plain CpAsync). + sSFA: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, self.sfa_smem_alloc_bytes], + self.buffer_align_bytes, + ] + sA: cute.struct.Align[ + cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sB: cute.struct.Align[ + cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sSFB: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)], + self.buffer_align_bytes, + ] + + # 2CTA variant: adds a_sync_transform_mbar_ptr for the warp-11 relay + # pipeline (PipelineAsyncUmma) that bridges per-CTA `a_pipeline` to + # the cluster-wide MMA consumer (tcgen05.mma.cta_group::2). + @cute.struct + class SharedStorageCpasync2cta: + sInfo: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int32, 5 * self.num_tile_stage], + 1, + ] + a_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + b_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + sfa_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + sfa_transform_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_sfa_tmem_stage * 2] + a_sync_transform_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + tile_info_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_tile_stage * 2] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + cute.cosize(self.c_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + sSFA: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, self.sfa_smem_alloc_bytes], + self.buffer_align_bytes, + ] + sA: cute.struct.Align[ + cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sB: cute.struct.Align[ + cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sSFB: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)], + self.buffer_align_bytes, + ] + + # TMA-A variant: A and B share one ab_pipeline mbar (4 gather_a + 1 + # tma_b producers arrive on the same mbar). No a_sync_transform mbar + # needed (TMA gather4 has HW `.multicast::cluster`, no peer-sync relay). + # Storage layout is independent of 1CTA/2CTA mode in TMA-A mode. + @cute.struct + class SharedStorageTma: + sInfo: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int32, 5 * self.num_tile_stage], + 1, + ] + ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + sfa_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + sfa_transform_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_sfa_tmem_stage * 2] + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + tile_info_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_tile_stage * 2] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + cute.cosize(self.c_smem_layout_staged.outer), + ], + self.buffer_align_bytes, + ] + # sSFA placed BEFORE sA so SFA gets a low SMEM offset (cpasync.128 + # destination addr stays under the 248KB threshold). + sSFA: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, self.sfa_smem_alloc_bytes], + self.buffer_align_bytes, + ] + sA: cute.struct.Align[ + cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sB: cute.struct.Align[ + cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sSFB: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)], + self.buffer_align_bytes, + ] + + if cutlass.const_expr(self.a_path == "cpasync"): + self.shared_storage = ( + SharedStorageCpasync2cta if self.use_2cta_instrs else SharedStorageCpasync1cta + ) + else: # "tma" + self.shared_storage = SharedStorageTma + + # Launch the kernel synchronously + self.kernel( + tiled_mma, + tiled_mma_bkeep, + tiled_mma_breuse, + tiled_mma_sfb, + a, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + sfa, + tma_atom_sfb, + tma_tensor_sfb, + tma_atom_c, + tma_tensor_c, + sfc_tensor, + norm_const_tensor, + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + token_id_mapping_tensor, + num_non_exiting_tiles, + alpha, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + self.tCtSFA_layout, + self.tCtSFB_layout, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + epilogue_op, + c_sf_n_tile_offset, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + use_pdl=self.use_pdl, + ) + return + + def mainloop_s2t_copy_and_partition( + self, + sSF: cute.Tensor, + tSF: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for smem to tmem load for scale factor tensor, then use it to + partition smem memory (source) and tensor memory (destination). + + :param sSF: The scale factor tensor in smem + :type sSF: cute.Tensor + :param tSF: The scale factor tensor in tmem + :type tSF: cute.Tensor + + :return: A tuple containing (tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) where: + - tiled_copy_s2t: The tiled copy operation for smem to tmem load for scale factor tensor(s2t) + - tCsSF_compact_s2t: The partitioned scale factor tensor in smem + - tSF_compact_s2t: The partitioned scale factor tensor in tmem + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + # (MMA, MMA_MN, MMA_K, STAGE) + tCsSF_compact = cute.filter_zeros(sSF) + # (MMA, MMA_MN, MMA_K) + tCtSF_compact = cute.filter_zeros(tSF) + + # Make S2T CopyAtom and tiledCopy + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K, STAGE) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_) + # ((ATOM_V, REST_V), Rest_Tiler, MMA_MN, MMA_K) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + + return tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t + + # GPU device kernel + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_bkeep: Optional[cute.TiledMma], + tiled_mma_breuse: Optional[cute.TiledMma], + tiled_mma_sfb: cute.TiledMma, + mA_mkl: cute.Tensor, + tma_atom_a: Optional[cute.CopyAtom], + tma_tensor_a: Optional[cute.Tensor], + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + tma_atom_c: cute.CopyAtom, + mC_mnl: cute.Tensor, + mSFC_mnl: Optional[cute.Tensor], + norm_const_tensor: Optional[cute.Tensor], + tile_idx_to_expert_idx: cute.Tensor, + tile_idx_to_mn_limit: cute.Tensor, + token_id_mapping_tensor: cute.Tensor, + num_non_exiting_tiles: cute.Tensor, + alpha: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + tCtSFA_layout: cute.Layout, + tCtSFB_layout: cute.Layout, + c_smem_layout_staged: Union[cute.Layout, cute.ComposedLayout, None], + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + epilogue_op: cutlass.Constexpr, + c_sf_n_tile_offset: cutlass.Int64 = cutlass.Int64(0), + ): + """ + GPU device kernel performing the Persistent batched GEMM computation. + """ + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # + # Prefetch tma desc + # + if warp_idx == self.tma_b_warp_id: + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_sfb) + cpasync.prefetch_descriptor(tma_atom_c) + if cutlass.const_expr(self.a_path == "tma"): + cpasync.prefetch_descriptor(tma_atom_a) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # + # Setup cta/thread coordinates + # + # Coords inside cluster + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + + # Coord inside cta + tidx, _, _ = cute.arch.thread_idx() + + # + # Alloc and init: a+b full/empty, accumulator full/empty, tensor memory dealloc barrier + # + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # (a_pipeline created below alongside b_pipeline.) + + # SFA pipeline (PipelineCpAsync): gather_sfa warps → transform warps. + # Producer: 4 gather_sfa warps × 32 threads (one CpAsync128.CG per row). + # Consumer: 4 transform warps × 32 threads. + # MMA waits on sfa_transform_pipeline downstream (after LDS+STTM). + sfa_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.gather_sfa_warp_id) * self.threads_per_warp, + ) + sfa_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.sfa_transform_warp_id) * self.threads_per_warp, + ) + sfa_pipeline = pipeline.PipelineCpAsync.create( + barrier_storage=storage.sfa_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=sfa_pipeline_producer_group, + consumer_group=sfa_pipeline_consumer_group, + defer_sync=True, + ) + + # SFA transform pipeline: transform warps (STTM) → MMA. PipelineAsyncUmma + # with cta_layout_vmnk so 2CTA peer-CTA arrives route to leader's mbar. + # Producer: 4 transform warps × 32 threads × cta_v_size. Consumer: MMA. + # num_stages = num_sfa_tmem_stage SFA TMEM slots, rotated. + cta_v_size = cute.size(cluster_layout_vmnk, mode=[0]) + sfa_transform_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.sfa_transform_warp_id) * self.threads_per_warp * cta_v_size, + ) + sfa_transform_pipeline_consumer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + sfa_transform_pipeline = pipeline.PipelineAsyncUmma.create( + barrier_storage=storage.sfa_transform_mbar_ptr.data_ptr(), + num_stages=self.num_sfa_tmem_stage, + producer_group=sfa_transform_pipeline_producer_group, + consumer_group=sfa_transform_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # A/B pipeline topology — branch on a_path: + # cpasync-A: separate a_pipeline (CpAsyncUmma) + b_pipeline (TmaUmma) + # + a_sync_transform_pipeline (AsyncUmma, 2CTA peer-sync relay). + # tma-A: single ab_pipeline (TmaUmma) shared by A and B producers. + if cutlass.const_expr(self.a_path == "cpasync"): + # cpasync-A: A producer = 128 threads (4 gather_a warps) issuing + # cp.async.cg.16B; each thread arrives once per stage. Consumer = + # MMA (UMMA), 1 thread. + a_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * len(self.gather_a_warp_id), + ) + a_pipeline = PipelineCpAsyncUmma.create( + barrier_storage=storage.a_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=a_pipeline_producer_group, + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # cpasync 2CTA A sync-transform relay. Producer thread count + # = 1 warp × cta_v_size so the cluster-wide arrive_count matches + # warp 11 across CTAs. + if cutlass.const_expr(self.use_2cta_instrs): + a_sync_transform_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * cta_v_size, + ) + a_sync_transform_pipeline = pipeline.PipelineAsyncUmma.create( + barrier_storage=storage.a_sync_transform_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=a_sync_transform_pipeline_producer_group, + consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread), + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # B/SFB pipeline (TMA → UMMA), 1 producer thread (tma_b warp). + # mcast_mode_mn=(0, 1): A is per-CTA cpasync (no N-multicast); B + # is TMA per-CTA or M-multicast. Default (1,1) would release + # across N peers — wrong since N-peers hold different B tiles. + b_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_b + b_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + b_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.b_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=b_pipeline_producer_group, + consumer_group=b_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + mcast_mode_mn=(0, 1), + defer_sync=True, + ) + else: # tma — merged ab_pipeline (TmaUmma): 4 gather_a + 1 tma_b + # on one mbar. Per-call expected_tx accumulates 4 × a_num_tma_load + # + num_tma_load per stage. Consumer release routes A's N peers + # + B's M peers → consumer_group = num_mcast_ctas_a + num_mcast_ctas_b - 1. + ab_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + len(self.gather_a_warp_id) + 1, # 4 gather_a + 1 tma_b + ) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=0, # per-producer expected_tx overrides drive accumulation + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Pipeline Init: Initialize acc_pipeline (barrier) and states + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilog_warp_id) * (2 if use_2cta_instrs else 1) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Pipeline Init:Initialize tile info pipeline (barrier) and states + tile_info_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * 1, + ) + # All 4 gather A warps consume tile_info in both CpAsync and TMA paths. + tile_info_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_wo_sched, + ) + tile_info_pipeline = pipeline.PipelineAsync.create( + barrier_storage=storage.tile_info_mbar_ptr.data_ptr(), + num_stages=self.num_tile_stage, + producer_group=tile_info_pipeline_producer_group, + consumer_group=tile_info_pipeline_consumer_group, + ) + + # Tensor memory dealloc barrier init + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.epilog_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, + arch=self.arch, + ) + + # Cluster arrive after barrier init (Rubin uses pipeline_init_arrive) + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # + # Setup smem tensor A/B/C/Scale + # + # (EPI_TILE_M, EPI_TILE_N, STAGE) + sC = storage.sC.get_tensor(c_smem_layout_staged.outer, swizzle=c_smem_layout_staged.inner) + # (MMA, MMA_M, MMA_K, STAGE) + sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner) + # (MMA, MMA_N, MMA_K, STAGE) + sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner) + # SFA SMEM (linear+pad layout for TMA gather4). + sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) + # (granularity_n, repeat_n), (granularity_k, repeat_k), num_scale_stage) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) + # (bidx, bidy, bidz, valid, mn_limit) + info_layout = cute.make_layout((5, self.num_tile_stage), stride=(1, 5)) + sInfo = storage.sInfo.get_tensor(info_layout) + + # + # Compute multicast mask for A/B buffer full + # + b_full_mcast_mask = None + sfb_full_mcast_mask = None + if cutlass.const_expr(self.is_b_mcast or use_2cta_instrs): + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1 + ) + # A multicast mask (tma-A only). cpasync-A path has no TMA multicast + # (each thread issues its own cp.async.cg.16B). tma-A broadcasts A + # along cluster N (mcast_mode=2); sm100_tma_gather4_load picks the + # `.multicast::cluster` PTX variant when this mask is non-None. + a_full_mcast_mask = None + if cutlass.const_expr(self.a_path == "tma" and self.is_a_mcast): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + + # + # Local_tile partition global tensors + # + # (bM, bK, loopM, loopK, loopL) + gA_mkl = cute.local_tile( + mA_mkl, + cute.slice_(self.cta_tile_shape_mnk, (None, 0, None)), + (None, None, None), + ) + # (bN, bK, loopN, loopK, loopL) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + + # (bM, bK, RestM, RestK, RestL) + gSFA_mkl = cute.local_tile( + mSFA_mkl, + cute.slice_(self.cta_tile_shape_mnk_sfa, (None, 0, None)), + (None, None, None), + ) + + # (bN, bK, RestN, RestK, RestL) + gSFB_nkl = cute.local_tile( + mSFB_nkl, + cute.slice_(self.mma_tiler_sfb, (0, None, None)), + (None, None, None), + ) + + gToken_ml = cute.local_tile( + token_id_mapping_tensor, + cute.slice_(self.cta_tile_shape_mnk, (None, 0, 0)), + (None,), + ) + + # (bM, bN, loopM, loopN, loopL) + gC_mnl = cute.local_tile( + mC_mnl, cute.slice_(self.mma_tiler_c, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cutlass.Int32(cute.size(gA_mkl, mode=[3])) + + # + # Partition global tensor for TiledMMA_A/B/C + # + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v) + # (MMA, MMA_N, MMA_K, loopN, loopK, loopL) + tCgB = thr_mma.partition_B(gB_nkl) + # (MMA, MMA_N, MMA_K, RestN, RestK, RestL) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + # (MMA, MMA_M, MMA_N, loopM, loopN, loopL) + tCgC = thr_mma.partition_C(gC_mnl) + + # + # Partition global/shared tensor for TMA load B + # + # TMA load B partition_S/D + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), loopM, loopK, loopL) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + # TMA load SFB partition_S/D + sfb_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape + ) + # ((atom_v, rest_v), STAGE) + # ((atom_v, rest_v), RestN, RestK, RestL) + tBsSFB, tBgSFB = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfb, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # + # Partition shared/tensor memory tensor for TiledMMA_A/B/C + # + # (MMA, MMA_M, MMA_K, STAGE) + tCrA = tiled_mma.make_fragment_A(sA) + # (MMA, MMA_N, MMA_K, STAGE) + tCrB = tiled_mma.make_fragment_B(sB) + # (MMA, MMA_M, MMA_N) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + # (MMA, MMA_M, MMA_N, STAGE) + if cutlass.const_expr(self.use_overlap_accum): + # Pipeline tracks 1 stage but TMEM has 2 physical regions overlap- + # ping by 64 cols. Build fragment with 2 stages and stride-hack + # the stage dim to (256 - 64) = 192 cols (= cta_tile_N - overlap). + num_acc_stage_overlapped = 2 + tCtAcc_fake = tiled_mma.make_fragment_C( + cute.append(acc_shape, num_acc_stage_overlapped) + ) + tCtAcc_fake = cute.make_tensor( + tCtAcc_fake.iterator, + cute.make_layout( + tCtAcc_fake.shape, + stride=( + tCtAcc_fake.stride[0], + tCtAcc_fake.stride[1], + tCtAcc_fake.stride[2], + (256 - 64) * tCtAcc_fake.stride[0][1], + ), + ), + ) + else: + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + # + # Cluster wait before tensor memory alloc + # + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + cute.arch.griddepcontrol_wait() + + # + # Specialized Schedule Warp + # + if warp_idx == self.sched_warp_id: + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + # First tile + work_tile = tile_sched.initial_work_tile_info() + + tile_info_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_tile_stage + ) + + num_non_exiting_tiles_value = num_non_exiting_tiles[0] + + if cutlass.const_expr(self.raster_along_m): + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_m = cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape) + if mma_tile_coord_m < num_non_exiting_tiles_value: + tile_info_pipeline.producer_acquire(tile_info_producer_state) + cur_tile_coord = work_tile.tile_idx + expert_idx = tile_idx_to_expert_idx[mma_tile_coord_m] + mn_limit = tile_idx_to_mn_limit[mma_tile_coord_m] + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = cur_tile_coord[0] + sInfo[(1, tile_info_producer_state.index)] = cur_tile_coord[1] + sInfo[(2, tile_info_producer_state.index)] = expert_idx + sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32( + work_tile.is_valid_tile + ) + sInfo[(4, tile_info_producer_state.index)] = mn_limit + # fence view async shared + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + else: + is_continue = cutlass.Boolean(1) + while work_tile.is_valid_tile and is_continue: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_m = cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape) + if mma_tile_coord_m < num_non_exiting_tiles_value: + tile_info_pipeline.producer_acquire(tile_info_producer_state) + cur_tile_coord = work_tile.tile_idx + expert_idx = tile_idx_to_expert_idx[mma_tile_coord_m] + mn_limit = tile_idx_to_mn_limit[mma_tile_coord_m] + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = cur_tile_coord[0] + sInfo[(1, tile_info_producer_state.index)] = cur_tile_coord[1] + sInfo[(2, tile_info_producer_state.index)] = expert_idx + sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32( + work_tile.is_valid_tile + ) + sInfo[(4, tile_info_producer_state.index)] = mn_limit + # fence view async shared + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + else: + is_continue = cutlass.Boolean(0) + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + tile_info_pipeline.producer_acquire(tile_info_producer_state) + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = work_tile.tile_idx[0] + sInfo[(1, tile_info_producer_state.index)] = work_tile.tile_idx[1] + sInfo[(2, tile_info_producer_state.index)] = -1 + sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32(0) + sInfo[(4, tile_info_producer_state.index)] = -1 + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + tile_info_pipeline.producer_tail(tile_info_producer_state) + + # Gather A warps (warps 4-7). cpasync / tma bodies static-gated by + # cutlass.const_expr — only one is traced. cpasync: 4 warps × 32 + # threads issue a_num_loads CpAsync128.CG per k_tile (thread layout + # (16, 8) covers 16 M-rows × 8 K-chunks); padded rows predicate off. + if warp_idx <= self.gather_a_warp_id[-1] and warp_idx >= self.gather_a_warp_id[0]: + cute.arch.setmaxregister_decrease(self.num_regs_gather_a_warps) + if cutlass.const_expr(self.a_path == "cpasync"): + a_atom_copy = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + mA_mkl.element_type, + num_bits_per_copy=128, + ) + a_thread_layout = cute.make_layout((16, 8), stride=(8, 1)) + a_value_layout = cute.make_layout((1, 32), stride=(32, 1)) + a_tiled_copy = cute.make_tiled_copy_tv( + a_atom_copy, + a_thread_layout, + a_value_layout, + ) + tidx_in_warpgroup = tidx % 128 + + sA_tiled = cute.make_tensor( + sA.iterator, + layout=cute.make_layout( + ( + self.cta_tile_shape_mnk[0], + self.cta_tile_shape_mnk[2], + self.num_ab_stage, + ), + stride=( + self.cta_tile_shape_mnk[2], + 1, + self.cta_tile_shape_mnk[0] * self.cta_tile_shape_mnk[2], + ), + ), + ) + a_thr_copy = a_tiled_copy.get_slice(tidx_in_warpgroup) + tAsA_tiled = a_thr_copy.partition_D(sA_tiled) + + a_token_offset_tensor = cute.make_rmem_tensor( + cute.make_layout((self.a_num_loads,)), + cutlass.Int32, + ) + a_predicate_tensor = cute.make_rmem_tensor( + cute.make_layout((self.a_num_loads,)), + cutlass.Boolean, + ) + + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + a_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + tile_info = cute.make_rmem_tensor((5,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(5, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + gToken_ml_tile = gToken_ml[(None, tile_info[0])] + for i in range(self.a_num_loads): + token_ml_tile_offset = (tidx_in_warpgroup // 8) + i * 16 + a_token_offset_tensor[i] = gToken_ml_tile[token_ml_tile_offset] + a_predicate_tensor[i] = ( + cutlass.Boolean(1) + if tile_info[0] * self.cta_tile_shape_mnk[0] + token_ml_tile_offset + < tile_info[4] + else cutlass.Boolean(0) + ) + a_token_offset_tensor[i] = ( + a_token_offset_tensor[i] // self.topk + if tile_info[0] * self.cta_tile_shape_mnk[0] + token_ml_tile_offset + < tile_info[4] + else 0 + ) + + tAgA = gA_mkl[(None, None, 0, None, 0)] + A_gmem_thread_offset = cute.assume((tidx_in_warpgroup % 8) * 32, divby=32) + + a_producer_state.reset_count() + peek_a_empty_status = cutlass.Boolean(1) + if a_producer_state.count < k_tile_cnt: + peek_a_empty_status = a_pipeline.producer_try_acquire(a_producer_state) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + a_pipeline.producer_acquire(a_producer_state, peek_a_empty_status) + + tAgA_ktile = tAgA[(None, None, a_producer_state.count)] + tAsA_ktile = tAsA_tiled[(None, None, None, a_producer_state.index)] + + for i in range(self.a_num_loads): + A_gmem_slice_offset = A_gmem_thread_offset + cute.assume( + a_token_offset_tensor[i] * tAgA_ktile.layout[0].stride, + divby=32, + ) + A_gmem_slice_offset = cute.assume(A_gmem_slice_offset, divby=32) + tAgA_slice_ptr = tAgA_ktile.iterator + A_gmem_slice_offset + tAgA_slice = cute.make_tensor( + tAgA_slice_ptr, layout=cute.make_layout((32,)) + ) + tAsA_slice = cute.make_tensor( + tAsA_ktile[(None, i, None)].iterator, + layout=cute.make_layout((32,)), + ) + a_predicate_slice = cute.make_rmem_tensor( + cute.make_layout((1,)), cutlass.Boolean + ) + a_predicate_slice[0] = a_predicate_tensor[i] + cute.copy_atom_call( + a_atom_copy, + tAgA_slice, + tAsA_slice, + pred=a_predicate_slice, + ) + + a_pipeline.producer_commit(a_producer_state) + + a_producer_state.advance() + peek_a_empty_status = cutlass.Boolean(1) + if a_producer_state.count < k_tile_cnt: + peek_a_empty_status = a_pipeline.producer_try_acquire(a_producer_state) + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(5, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + a_pipeline.producer_tail(a_producer_state) + + # tma-A: 4 warps × elect-one issuing TMA gather4. Each warp owns + # 1/4 of the M rows and issues n_gather_per_warp gather4 calls + # (each pulling 4 rows × cta_tile_K). token_id == -1 → TMA + # zero-fills the row. Signals merged ab_pipeline mbar (shared w/ B). + elif cutlass.const_expr(self.a_path == "tma"): + warp_rel = warp_idx - self.gather_a_warp_id[0] + rows_per_warp = self.cta_tile_shape_mnk[0] // self.tma_gather_num_warps + n_gather_per_warp = rows_per_warp // 4 + + a_row_ids = cute.make_rmem_tensor( + cute.make_layout((rows_per_warp,)), + cutlass.Int32, + ) + + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + ) + work_tile = tile_sched.initial_work_tile_info() + + ab_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + tile_info = cute.make_rmem_tensor((5,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(5, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + gToken_ml_tile = gToken_ml[(None, tile_info[0])] + + # Each warp computes its own row range (1/4 of the tile). + # The routing helper initializes padded mapping entries to + # zero, so use mn_limit as the authoritative valid-row + # predicate before passing -1 to gather4 for zero-fill. + for i in range(rows_per_warp): + row_global = warp_rel * rows_per_warp + i + token_id = gToken_ml_tile[row_global] + valid_row = ( + tile_info[0] * self.cta_tile_shape_mnk[0] + row_global < tile_info[4] + ) + row_id = token_id // self.topk if valid_row else cutlass.Int32(-1) + a_row_ids[i] = cutlass.Int32(-1) if token_id == -1 else row_id + + # A multicast leader gate: when A is N-multicast, only the N=0 + # CTA issues the mcast PTX; HW broadcasts data + mbar tx_count + # to the N-peer. The peer's producer_acquire still sets its + # mbar's expect_tx; that mbar fires via HW routing. + is_a_mcast_leader = (not self.is_a_mcast) or ( + block_in_cluster_coord_vmnk[2] == 0 + ) + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # gather_a contributes a_num_tma_load_bytes per warp to + # the merged ab_pipeline mbar (4 such arrivals per stage). + ab_pipeline.producer_acquire( + ab_producer_state, + expected_tx=self.a_num_tma_load_bytes, + ) + + if is_a_mcast_leader: + with cute.arch.elect_one(): + col_k = k_tile * self.cta_tile_shape_mnk[2] + mbar_ptr = ab_pipeline.producer_get_barrier(ab_producer_state) + stage_base_elements = ( + ab_producer_state.index + * self.cta_tile_shape_mnk[0] + * self.cta_tile_shape_mnk[2] + ) + warp_base_elements = ( + warp_rel * rows_per_warp * self.cta_tile_shape_mnk[2] + ) + + for g in range(n_gather_per_warp): + row_start = g * 4 + dst_offset = ( + stage_base_elements + + warp_base_elements + + row_start * self.cta_tile_shape_mnk[2] + ) + dst_ptr = sA.iterator + dst_offset + sm100_tma_gather4_load( + tma_atom_a, + dst_ptr, + mbar_ptr, + col_k, + a_row_ids[row_start], + a_row_ids[row_start + 1], + a_row_ids[row_start + 2], + a_row_ids[row_start + 3], + use_cta_group_2=self.use_2cta_instrs, + mcast_mask=a_full_mcast_mask, + ) + + ab_pipeline.producer_commit(ab_producer_state) + ab_producer_state.advance() + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(5, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + ab_pipeline.producer_tail(ab_producer_state) + + # Gather SFA warps (warps 12-15): CpAsync128.CG per row. 4 warps × 32 + # threads, each loads one row (16B = 16 FP8 SFs). sSFA is plain linear + # (M, tile_K_sf, stage). tiled_copy_tv for SMEM dest + manual per-thread + # GMEM source tensor. + if warp_idx <= self.gather_sfa_warp_id[-1] and warp_idx >= self.gather_sfa_warp_id[0]: + cute.arch.setmaxregister_decrease(self.num_regs_gather_sfa_warps) + + sfa_tile_k_sf = self.cta_tile_shape_mnk[2] // self.sf_vec_size # 16 + sfa_gather_threads = len(self.gather_sfa_warp_id) * self.threads_per_warp + # Rows per thread = cta_tile_M / gather_threads. + # - non-breuse cta_tile_M=128: 1 row per thread + # - breuse cta_tile_M=256: 2 rows per thread (outer r-loop in inner k loop) + sfa_rows_per_thread = self.cta_tile_shape_mnk[0] // sfa_gather_threads + + # One CpAsync128.CG per thread per row per k_tile. + sfa_atom_copy = cute.make_copy_atom( + cute.nvgpu.cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL), + mSFA_mkl.element_type, + num_bits_per_copy=128, + ) + + sSFA_tiled = cute.make_tensor( + sSFA.iterator, + layout=cute.make_layout( + ( + self.cta_tile_shape_mnk[0], + sfa_tile_k_sf, + self.num_ab_stage, + ), + stride=( + sfa_tile_k_sf, + 1, + self.cta_tile_shape_mnk[0] * sfa_tile_k_sf, + ), + ), + ) + + tidx_in_warpgroup = tidx - self.gather_sfa_warp_id[0] * self.threads_per_warp + + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + sfa_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + tile_info = cute.make_rmem_tensor((5,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(5, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + gToken_ml_tile = gToken_ml[(None, tile_info[0])] + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + sfa_pipeline.producer_acquire(sfa_producer_state) + + # For each row this thread is responsible for + # (1 row at non-breuse cta_tile_M=128; 2 rows at breuse + # cta_tile_M=256). Inner r-loop fully unrolls. + for r in cutlass.range_constexpr(sfa_rows_per_thread): + row_in_cta = tidx_in_warpgroup + r * sfa_gather_threads + + # Per-row token id → SFA gmem row index. + tok = gToken_ml_tile[row_in_cta] + sfa_row_id = cutlass.Int32(-1) if tok == -1 else tok // self.topk + + # OOB predicate: rows past mn_limit do not store. + sfa_pred = cute.make_rmem_tensor(cute.make_layout((1,)), cutlass.Boolean) + sfa_pred[0] = ( + cutlass.Boolean(1) + if ( + tile_info[0] * self.cta_tile_shape_mnk[0] + row_in_cta + < tile_info[4] + ) + else cutlass.Boolean(0) + ) + + # GMEM src for this row + k_tile. + tAgSFA_row = gSFA_mkl[(sfa_row_id, None, 0, None, 0)] + tAgSFA_ktile = tAgSFA_row[(None, k_tile)] + tAgSFA_slice = cute.make_tensor( + tAgSFA_ktile.iterator, + layout=cute.make_layout((sfa_tile_k_sf,)), + ) + + # SMEM dst: direct row-indexed slice (bypass partition_D + # so we naturally handle multiple rows per thread). + tAsSFA_slice = cute.make_tensor( + sSFA_tiled[(row_in_cta, None, sfa_producer_state.index)].iterator, + cute.make_layout((sfa_tile_k_sf,)), + ) + + cute.copy_atom_call( + sfa_atom_copy, + tAgSFA_slice, + tAsSFA_slice, + pred=sfa_pred, + ) + + sfa_pipeline.producer_commit(sfa_producer_state) + sfa_producer_state.advance() + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(5, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + sfa_pipeline.producer_tail(sfa_producer_state) + + # + # SFA Transform warps (warps 16-19) — LDS + STTM consumer + # + if warp_idx >= self.sfa_transform_warp_id[0] and warp_idx <= self.sfa_transform_warp_id[-1]: + cute.arch.setmaxregister_decrease(self.num_regs_sfa_transform_warps) + + if cutlass.const_expr(True): + # Transform warps wait for TMEM alloc and compute + # sfa_tmem_ptr = acc_tmem_ptr + offset for STTM destination. + tmem.wait_for_alloc() + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.sfa_tmem_offset, + dtype=self.sf_dtype, + ) + + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + _ = tile_sched.initial_work_tile_info() + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + sfa_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage + ) + # Producer state on sfa_transform_pipeline: tracks which TMEM + # stage (of num_sfa_tmem_stage) is being filled this iter. + sfa_transform_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, + self.num_sfa_tmem_stage, + ) + tile_info = cute.make_rmem_tensor((5,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(5, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # All 4 transform warps run identical address logic: each does + # the full cta_tile_M LDS+STTM into its own per-warp TMEM lane; + # UMMA reads all 4 to cover M. M-blocks = cta_tile_M / 32 (4 + # non-breuse, 8 breuse). Per warp: num_m_blocks LDS.128 + 4 STTM. + num_m_blocks = self.cta_tile_shape_mnk[0] // 32 + lane_in_warp = tidx % 32 + + while is_valid_tile: + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + sfa_pipeline.consumer_wait(sfa_consumer_state) + + # LDS.128 → RMEM Uint32. Per ci read 16 SFs = 4 u32. + # Issue LDS before producer_acquire so SMEM→RMEM moves + # overlap with prior MMA still holding the next TMEM + # stage; STTM still gates on producer_acquire below. + sfa_rmem_u32 = cute.make_rmem_tensor( + cute.make_layout((num_m_blocks, 4)), # (ci, K-group) + cutlass.Uint32, + ) + # Linear sSFA (M, K_sf, stage). Each LDS reads 16 + # contiguous FP8 SFs (4 u32) at row_global of the + # current stage. + for ci in range(num_m_blocks): + row_global = 32 * ci + lane_in_warp + smem_slice = sSFA[ + ( + row_global, + None, + sfa_consumer_state.index, + ) + ] + smem_slice_u32 = cute.make_tensor( + cute.recast_ptr(smem_slice.iterator, dtype=cutlass.Uint32), + cute.make_layout((4,)), + ) + cute.autovec_copy(smem_slice_u32, sfa_rmem_u32[ci, None]) + + # Acquire SFA TMEM slot (MMA's UMMA consumer release + # frees it after consuming). Done after LDS so the + # SMEM→RMEM stage isn't blocked on TMEM availability. + sfa_transform_pipeline.producer_acquire(sfa_transform_producer_state) + + # STTM via inline PTX. tCtSFA_layout cols: + # non-breuse (M=128): 1 half, cols 0..15 = 4 K-groups + # × 4 M-blocks (gi*4 stride). 4 STTM x4. + # breuse (M=256): 2 halves at +16 col offset (keep + # ci 0..3 / reuse ci 4..7), each 4 STTM x4. 8 STTM x4 total. + stage_idx_in_tmem = sfa_transform_producer_state.index + sfa_tmem_addr_base = ( + acc_tmem_ptr + + self.sfa_tmem_offset + + stage_idx_in_tmem * self.num_sfa_tmem_cols_per_stage + ).toint() + # Number of "halves" (keep/reuse splits) and the stride + # between them. non-breuse: 1 half, no second offset. + # breuse: 2 halves, second at +16 cols. + num_halves = num_m_blocks // 4 # 1 or 2 + half_col_stride = 16 # cols between keep and reuse + for half in range(num_halves): + half_addr_base = sfa_tmem_addr_base + half * half_col_stride + ci_base = half * 4 + for gi in range(4): + sm100_tcgen05_st_32x32b_x4( + half_addr_base + gi * 4, + sfa_rmem_u32[ci_base + 0, gi], + sfa_rmem_u32[ci_base + 1, gi], + sfa_rmem_u32[ci_base + 2, gi], + sfa_rmem_u32[ci_base + 3, gi], + ) + # Make TMEM stores visible to UMMA, then commit the + # transform pipeline producer slot. + cute.arch.fence_view_async_tmem_store() + sfa_transform_pipeline.producer_commit(sfa_transform_producer_state) + sfa_transform_producer_state.advance() + + sfa_pipeline.consumer_release(sfa_consumer_state) + sfa_consumer_state.advance() + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(5, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # Drain transform pipeline before exit. + sfa_transform_pipeline.producer_tail(sfa_transform_producer_state) + + # A Sync-Transform Warp (warp 11). Active only on cpasync 2CTA. + # Consumes per-CTA `a_pipeline` and re-produces cluster-wide + # `a_sync_transform_pipeline` so MMA's cta_group::2 sees both CTAs' A. + # SFA needs no relay — its transform warps + sfa_transform_pipeline + # handle cluster arrives. Idle on 1CTA / tma-A. + if warp_idx == self.sync_transform_warp_id: + if cutlass.const_expr(self.a_path == "cpasync" and self.use_2cta_instrs): + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + a_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage + ) + a_sync_transform_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + # First tile info (only need the validity flag). + valid_tile_info = cute.make_rmem_tensor((1,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + valid_tile_info[0] = sInfo[(3, tile_info_consumer_state.index)] + is_valid_tile = valid_tile_info[0] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + a_consumer_state.reset_count() + peek_a_full_status = cutlass.Boolean(1) + if a_consumer_state.count < k_tile_cnt: + peek_a_full_status = a_pipeline.consumer_try_wait(a_consumer_state) + a_sync_transform_producer_state.reset_count() + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Wait per-CTA A full → commit cluster-wide A sync-transform + # full. We do NOT release a_pipeline here; MMA owns its + # consumer_release so each CTA's producer sees the empty arrive. + a_pipeline.consumer_wait(a_consumer_state, peek_a_full_status) + a_sync_transform_pipeline.producer_commit(a_sync_transform_producer_state) + a_sync_transform_producer_state.advance() + a_consumer_state.advance() + peek_a_full_status = cutlass.Boolean(1) + if a_consumer_state.count < k_tile_cnt: + peek_a_full_status = a_pipeline.consumer_try_wait(a_consumer_state) + + # Advance to next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + valid_tile_info[0] = sInfo[(3, tile_info_consumer_state.index)] + is_valid_tile = valid_tile_info[0] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # Drain A sync-transform pipeline before exit. + a_sync_transform_pipeline.producer_tail(a_sync_transform_producer_state) + + # TMA B/SFB load warp (warp 9). Loads B/SFB GMEM → SMEM with multicast. + if warp_idx == self.tma_b_warp_id: + # B producer signals the pipeline owning A+B (tma → ab_pipeline) + # or B alone (cpasync → b_pipeline). Body is the same; only the + # pipeline alias and producer_acquire's expected_tx differ. + if cutlass.const_expr(self.a_path == "cpasync"): + _b_pipe = b_pipeline + else: + _b_pipe = ab_pipeline + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + # First tile + work_tile = tile_sched.initial_work_tile_info() + + b_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + # Get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + mma_tile_coord_mnl = ( + tile_info[0] // cute.size(tiled_mma.thr_id.shape), + tile_info[1], + tile_info[2], + ) + # + # Slice to per mma tile index + # + # ((atom_v, rest_v), loopK) + tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] + + # Apply SFB slicing hack when cta_tile_shape_n=64 + slice_n = mma_tile_coord_mnl[1] + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): + slice_n = mma_tile_coord_mnl[1] // 2 + + # ((atom_v, rest_v), RestK) + tBgSFB_slice = tBgSFB[(None, slice_n, None, mma_tile_coord_mnl[2])] + + # Peek (try_wait) B buffer empty for k_tile = prefetch_k_tile_cnt + b_producer_state.reset_count() + peek_b_empty_status = cutlass.Boolean(1) + if b_producer_state.count < k_tile_cnt: + peek_b_empty_status = _b_pipe.producer_try_acquire(b_producer_state) + # + # Tma load loop + # + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + # Conditionally wait for B buffer empty. + # tma-A mode passes expected_tx (ab_pipeline has tx_count=0 + # at create and accumulates per producer call); cpasync-A + # mode's b_pipeline has tx_count fixed at create. + if cutlass.const_expr(self.a_path == "cpasync"): + _b_pipe.producer_acquire(b_producer_state, peek_b_empty_status) + else: + _b_pipe.producer_acquire( + b_producer_state, + peek_b_empty_status, + expected_tx=self.num_tma_load_bytes, + ) + + tBgB_k = tBgB_slice[(None, b_producer_state.count)] + tBgSFB_k = tBgSFB_slice[(None, b_producer_state.count)] + tBsB_pipe = tBsB[(None, b_producer_state.index)] + tBsSFB_pipe = tBsSFB[(None, b_producer_state.index)] + + tma_bar = _b_pipe.producer_get_barrier(b_producer_state) + + # TMA load B + cute.copy( + tma_atom_b, + tBgB_k, + tBsB_pipe, + tma_bar_ptr=tma_bar, + mcast_mask=b_full_mcast_mask, + ) + + # TMA load SFB + cute.copy( + tma_atom_sfb, + tBgSFB_k, + tBsSFB_pipe, + tma_bar_ptr=tma_bar, + mcast_mask=sfb_full_mcast_mask, + ) + + # Peek (try_wait) B buffer empty for k_tile + 1 + b_producer_state.advance() + peek_b_empty_status = cutlass.Boolean(1) + if b_producer_state.count < k_tile_cnt: + peek_b_empty_status = _b_pipe.producer_try_acquire(b_producer_state) + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # + # Wait B buffer empty + # + _b_pipe.producer_tail(b_producer_state) + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + # + # Bar sync for retrieve tensor memory ptr from shared mem + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # SFA TMEM base ptr (sf_dtype). Stage-indexed tCtSFA is rebuilt + # per k_tile inside the loop using sfa_transform_consumer_state.index. + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols, + dtype=self.sf_dtype, + ) + + # Make SFB tmem tensor (using precomputed layout) + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + + # SFA TMEM is filled by transform warps via LDS+STTM (no UTCCP). + # Only SFB uses UTCCP. + sfb_s2t_bundle = self._mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + # MMA consumer states. cpasync: separate a / b states + # (+ a_sync_transform in 2CTA). tma: single ab_consumer_state. + if cutlass.const_expr(self.a_path == "cpasync"): + a_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage + ) + b_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage + ) + if cutlass.const_expr(self.use_2cta_instrs): + a_sync_transform_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage + ) + else: # "tma" + ab_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage + ) + # a_consumer_state and b_consumer_state both aliased to + # ab_consumer_state so existing `.index` reads in shared + # SMEM-slice code (sA[stage] / sB[stage] in the MMA mainloop) + # still resolve (A and B share stages in tma-A mode). + a_consumer_state = ab_consumer_state + b_consumer_state = ab_consumer_state + # sfa_transform_pipeline has num_sfa_tmem_stage slots, not num_ab_stage. + sfa_transform_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_sfa_tmem_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + # Get the first tile info from pipeline (scheduler has filtered out tiles >= num_non_exiting_tiles) + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + # Peek (try_wait) A / B / SFA buffer full for k_tile = 0. + # cpasync-A 1CTA: A peek on a_pipeline. + # cpasync-A 2CTA: A peek on a_sync_transform_pipeline (cluster-wide). + # tma-A: A+B peek on ab_pipeline (single shared mbar). + if cutlass.const_expr(self.a_path == "cpasync"): + if cutlass.const_expr(self.use_2cta_instrs): + a_sync_transform_consumer_state.reset_count() + peek_a_sync_transform_full_status = cutlass.Boolean(1) + if a_sync_transform_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_a_sync_transform_full_status = ( + a_sync_transform_pipeline.consumer_try_wait( + a_sync_transform_consumer_state + ) + ) + a_consumer_state.reset_count() + else: + a_consumer_state.reset_count() + peek_a_full_status = cutlass.Boolean(1) + if a_consumer_state.count < k_tile_cnt: + peek_a_full_status = a_pipeline.consumer_try_wait(a_consumer_state) + b_consumer_state.reset_count() + peek_b_full_status = cutlass.Boolean(1) + if b_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_b_full_status = b_pipeline.consumer_try_wait(b_consumer_state) + else: # "tma" + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + sfa_transform_consumer_state.reset_count() + peek_sfa_full_status = cutlass.Boolean(1) + if sfa_transform_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_sfa_full_status = sfa_transform_pipeline.consumer_try_wait( + sfa_transform_consumer_state + ) + + mma_tile_coord_mnl = ( + tile_info[0] // cute.size(tiled_mma.thr_id.shape), + tile_info[1], + tile_info[2], + ) + + # Accumulator stage. Overlap mode: pipeline tracks 1 stage + # but TMEM has 2 regions (acc[0] cols 0..255, acc[1] 192..447); + # use phase to pick. producer_state.phase starts at 1, so XOR + # with 1 to align with TMEM slot 0 on tile 0. + if cutlass.const_expr(self.use_overlap_accum): + acc_stage_index = acc_producer_state.phase ^ 1 + else: + acc_stage_index = acc_producer_state.index + + tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)] + + # Apply TMEM pointer offset hack when cta_tile_shape_n=192 or + # cta_tile_shape_n=64 + tCtSFB_mma = tCtSFB + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): + # If this is an ODD tile, shift the TMEM start address for + # cta_tile_shape_n=192 case by two words + # (ignores first 64 columns of SFB) + offset = ( + cutlass.Int32(2) if mma_tile_coord_mnl[1] % 2 == 1 else cutlass.Int32(0) + ) + shifted_ptr = cute.recast_ptr( + acc_tmem_ptr + + self.num_accumulator_tmem_cols + + self.num_sfa_tmem_cols + + offset, + dtype=self.sf_dtype, + ) + tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout) + elif cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): + # Move in increments of 64 columns of SFB + offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2) + shifted_ptr = cute.recast_ptr( + acc_tmem_ptr + + self.num_accumulator_tmem_cols + + self.num_sfa_tmem_cols + + offset, + dtype=self.sf_dtype, + ) + tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout) + # + # Wait for accumulator buffer empty + # + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + # + # Mma mainloop + # + + for k_tile in cutlass.range(k_tile_cnt): + # Set tensor memory buffer for current tile + # (MMA, MMA_M, MMA_N) + + if is_leader_cta: + # Wait for A / B / SFA buffer full. + # cpasync-A: 2 separate waits (A side + B side). + # 2CTA A wait uses a_sync_transform_pipeline (relay). + # tma-A: single ab_pipeline wait covers A and B. + if cutlass.const_expr(self.a_path == "cpasync"): + if cutlass.const_expr(self.use_2cta_instrs): + a_sync_transform_pipeline.consumer_wait( + a_sync_transform_consumer_state, + peek_a_sync_transform_full_status, + ) + else: + a_pipeline.consumer_wait(a_consumer_state, peek_a_full_status) + b_pipeline.consumer_wait(b_consumer_state, peek_b_full_status) + a_stage_idx = a_consumer_state.index + b_stage_idx = b_consumer_state.index + else: # "tma" + ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) + # Read ab_consumer_state.index directly (NOT via + # the aliased a_/b_consumer_state names): CuteDSL + # caches .index SSA values per variable *name*, + # so aliased reads go stale after ab.advance(). + a_stage_idx = ab_consumer_state.index + b_stage_idx = ab_consumer_state.index + sfa_transform_pipeline.consumer_wait( + sfa_transform_consumer_state, peek_sfa_full_status + ) + + # Rebuild tCtSFA pointing at the current SFA TMEM + # stage (transform warps rotate through stages; + # sfa_transform_consumer_state.index tracks which to read). + tCtSFA = cute.make_tensor( + sfa_tmem_ptr + + sfa_transform_consumer_state.index + * self.num_sfa_tmem_cols_per_stage + * 4, + tCtSFA_layout, + ) + + # SFB UTCCP (SFA is already in TMEM from transform warps). + self._mainloop_s2t_copies(b_stage_idx, sfb_s2t_bundle) + + num_kblocks = cute.size(tCrA, mode=[2]) + + for kblock_idx in cutlass.range(num_kblocks, unroll_full=True): + if cutlass.const_expr( + self.enable_breuse + and cute.size(tCtAcc.layout, mode=[1]) == 2 + and cute.size(tCtAcc.layout, mode=[2]) == 1 + ): + tCtAcc_bkeep = tCtAcc[(None, 0, 0)] + tCtAcc_breuse = tCtAcc[(None, 1, 0)] + + a_kblk_crd_keep = (None, 0, kblock_idx, a_stage_idx) + a_kblk_crd_reuse = (None, 1, kblock_idx, a_stage_idx) + b_kblk_crd = (None, 0, kblock_idx, b_stage_idx) + + sfa_kblk_crd_keep = (None, 0, kblock_idx) + sfa_kblk_crd_reuse = (None, 1, kblock_idx) + sfb_kblk_crd = (None, 0, kblock_idx) + + # Bkeep + tiled_mma_bkeep.set( + tcgen05.Field.ACCUMULATE, + k_tile != 0 or kblock_idx != 0, + ) + cute.gemm( + tiled_mma_bkeep, + tCtAcc_bkeep, + [tCrA[a_kblk_crd_keep], tCtSFA[sfa_kblk_crd_keep]], + [tCrB[b_kblk_crd], tCtSFB_mma[sfb_kblk_crd]], + tCtAcc_bkeep, + ) + # Breuse + tiled_mma_breuse.set( + tcgen05.Field.ACCUMULATE, + k_tile != 0 or kblock_idx != 0, + ) + cute.gemm( + tiled_mma_breuse, + tCtAcc_breuse, + [ + tCrA[a_kblk_crd_reuse], + tCtSFA[sfa_kblk_crd_reuse], + ], + [tCrB[b_kblk_crd], tCtSFB_mma[sfb_kblk_crd]], + tCtAcc_breuse, + ) + else: + a_kblock_coord = (None, None, kblock_idx, a_stage_idx) + b_kblock_coord = (None, None, kblock_idx, b_stage_idx) + sf_kblock_coord = (None, None, kblock_idx) + + tiled_mma.set( + tcgen05.Field.ACCUMULATE, + k_tile != 0 or kblock_idx != 0, + ) + cute.gemm( + tiled_mma, + tCtAcc, + [tCrA[a_kblock_coord], tCtSFA[sf_kblock_coord]], + [tCrB[b_kblock_coord], tCtSFB_mma[sf_kblock_coord]], + tCtAcc, + ) + + # Release A/B/SFA buffer-empty (async arrive). + # cpasync: a_pipeline + a_sync_transform_pipeline + # (2CTA only) + b_pipeline. tma: single ab_pipeline. + if cutlass.const_expr(self.a_path == "cpasync"): + a_pipeline.consumer_release(a_consumer_state) + if cutlass.const_expr(self.use_2cta_instrs): + a_sync_transform_pipeline.consumer_release( + a_sync_transform_consumer_state + ) + b_pipeline.consumer_release(b_consumer_state) + else: # "tma" + ab_pipeline.consumer_release(ab_consumer_state) + sfa_transform_pipeline.consumer_release(sfa_transform_consumer_state) + + # Peek (try_wait) A / B / SFA buffer full for k_tile + 1. + if cutlass.const_expr(self.a_path == "cpasync"): + if cutlass.const_expr(self.use_2cta_instrs): + a_sync_transform_consumer_state.advance() + peek_a_sync_transform_full_status = cutlass.Boolean(1) + if a_sync_transform_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_a_sync_transform_full_status = ( + a_sync_transform_pipeline.consumer_try_wait( + a_sync_transform_consumer_state + ) + ) + a_consumer_state.advance() + else: + a_consumer_state.advance() + peek_a_full_status = cutlass.Boolean(1) + if a_consumer_state.count < k_tile_cnt: + peek_a_full_status = a_pipeline.consumer_try_wait(a_consumer_state) + b_consumer_state.advance() + peek_b_full_status = cutlass.Boolean(1) + if b_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_b_full_status = b_pipeline.consumer_try_wait(b_consumer_state) + else: # "tma" + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + sfa_transform_consumer_state.advance() + peek_sfa_full_status = cutlass.Boolean(1) + if sfa_transform_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_sfa_full_status = sfa_transform_pipeline.consumer_try_wait( + sfa_transform_consumer_state + ) + + # + # Async arrive accumulator buffer full(each kblock) + # + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + + # Peek (try_wait) Acc buffer empty for k_tile = k_tile + 1 + acc_producer_state.advance() + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # + # Wait for accumulator buffer empty + # + acc_pipeline.producer_tail(acc_producer_state) + + # + # Specialized epilogue warps + # + if warp_idx <= self.epilog_warp_id[-1]: + # Register reconfig: epilogue needs many regs for SwiGLU + cute.arch.setmaxregister_increase(self.num_regs_epilogue_warps) + # + # Alloc tensor memory buffer + # + tmem.allocate(self.num_tmem_alloc_cols) + + # + # Bar sync for retrieve tensor memory ptr from shared memory + # + tmem.wait_for_alloc() + + # + # Retrieving tensor memory ptr and make accumulator tensor + # + tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + # (MMA, MMA_M, MMA_N, STAGE) + tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout) + + # Epilogue partition: transform both accumulator and C layout. + # transform_partitioned_tensor_layout merges (MMA_ATOM, MMA_M) → flat M. + tCtAcc_transformed = transform_partitioned_tensor_layout(tCtAcc_base) + tCgC_for_epi = transform_partitioned_tensor_layout(tCgC) + + epi_tidx = tidx % 128 + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc_up, + tTR_rAcc_gate, + ) = self.epilog_tmem_copy_and_partition( + epi_tidx, tCtAcc_transformed, tCgC_for_epi, epi_tile, use_2cta_instrs + ) + + tTR_rC = None + tiled_copy_r2s = None + tRS_rC = None + tRS_sC = None + bSG_sC = None + bSG_gC_partitioned = None + tTR_rC = cute.make_rmem_tensor(tTR_rAcc_up.shape, self.c_dtype) + tiled_copy_r2s, tRS_rC, tRS_sC = epilogue_smem_copy_and_partition( + self, tiled_copy_t2r, tTR_rC, epi_tidx, sC + ) + ( + tma_atom_c, + bSG_sC, + bSG_gC_partitioned, + ) = self.epilog_gmem_copy_and_partition( + epi_tidx, tma_atom_c, tCgC_for_epi, epi_tile, sC + ) + + if cutlass.const_expr(self.generate_sfc): + norm_const = norm_const_tensor[0] + # (EPI_TILE_M, EPI_TILE_N, RestM, RestN, RestL) + gSFC_mnl = cute.local_tile(mSFC_mnl, epi_tile, (None, None, None)) + + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, RestM, RestN, RestL) + tCgSFC_mnl = thr_copy_t2r.partition_D(gSFC_mnl) + tCgSFC_mnl = cute.filter_zeros(tCgSFC_mnl) + # (T2R, T2R_M, T2R_N) + tCrSFC = cute.make_rmem_tensor( + tCgSFC_mnl[(None, None, None, 0, 0, 0)].layout, self.sf_dtype + ) + tCrSFC_pvscale = cute.make_rmem_tensor_like(tCrSFC, cutlass.Float32) + + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + c_pipeline = None + # Threads/warps participating in tma store pipeline + c_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + 32 * len(self.epilog_warp_id), + ) + c_pipeline = pipeline.PipelineTmaStore.create( + num_stages=self.num_c_stage, + producer_group=c_producer_group, + ) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + # Get the first tile info + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + num_prev_subtiles = cutlass.Int32(0) + while is_valid_tile: + mma_tile_coord_mnl = ( + tile_info[0] // cute.size(tiled_mma.thr_id.shape), + tile_info[1], + tile_info[2], + ) + # + # Get alpha for current group + # + + expert_idx = mma_tile_coord_mnl[2] + alpha_val = alpha[expert_idx] + + # + # Slice to per mma tile index + # + bSG_gC = None + # ((ATOM_V, REST_V), EPI_M, EPI_N) + bSG_gC = bSG_gC_partitioned[ + ( + None, + None, + None, + mma_tile_coord_mnl[0], + mma_tile_coord_mnl[1], + 0, + ) + ] + + # Get accumulator stage index. Overlap mode uses phase + # (alternates 0/1) and reverse-iterates acc[0]'s subtiles so + # the overlap region (high cols of acc[0]) is consumed first + # → early-release lets MMA write acc[1] in those cols. + if cutlass.const_expr(self.use_overlap_accum): + # Consumer state starts at phase=0 (per make_pipeline_state) + # so phase directly maps to TMEM slot index. + acc_stage_index = acc_consumer_state.phase + reverse_subtile = ( + cutlass.Boolean(True) if acc_stage_index == 0 else cutlass.Boolean(False) + ) + else: + acc_stage_index = acc_consumer_state.index + + # Set tensor memory buffer for current tile + # (T2R, T2R_M, T2R_N, EPI_M, EPI_M) + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_stage_index)] + + if cutlass.const_expr(self.generate_sfc): + # (T2R, T2R_M, T2R_N, RestM, RestN) + tCgSFC_mn = tCgSFC_mnl[ + ( + None, + None, + None, + None, + None, + 0, + ) + ] + + # + acc_pipeline.consumer_wait(acc_consumer_state) + + # Activation epilogue. SwiGLU consumes interleaved [up, gate] + # accumulator subtiles and halves N; Relu2 consumes each + # accumulator subtile directly and preserves N. + # tTR_tAcc: (T2R, T2R_M, T2R_N, EPI_M, EPI_N, STAGE), sliced on STAGE. + # bSG_gC: ((ATOM_V, REST_V), EPI_M, EPI_N, loopM, loopN, loopL). + interleave_granularity = 64 + gate_offset = interleave_granularity // self.epi_tile_n + epi_m_cnt = cute.size(tTR_tAcc.shape, mode=[3]) + acc_n_subtile_cnt = cute.size(tTR_tAcc.shape, mode=[4]) + out_n_subtile_cnt = acc_n_subtile_cnt // 2 if self.is_gated else acc_n_subtile_cnt + + for epi_m_idx in cutlass.range(epi_m_cnt): + for out_n_idx in cutlass.range(out_n_subtile_cnt): + # Map output N subtile → acc N subtile. Each + # interleave block of 2*gate_offset subtiles is + # [up*gate_offset, gate*gate_offset]. acc[0] in + # overlap mode iterates in reverse (consume high-col + # overlap region first). + if cutlass.const_expr(self.use_overlap_accum): + real_out_n_idx = ( + (out_n_subtile_cnt - 1 - out_n_idx) + if reverse_subtile + else out_n_idx + ) + else: + real_out_n_idx = out_n_idx + if cutlass.const_expr(self.is_gated): + block_idx = real_out_n_idx // gate_offset + within_block = real_out_n_idx % gate_offset + up_n_subtile = block_idx * 2 * gate_offset + within_block + gate_n_subtile = ( + block_idx * 2 * gate_offset + gate_offset + within_block + ) + else: + up_n_subtile = real_out_n_idx + # + # Load accumulator from tensor memory buffer to register + # + tTR_tAcc_mn_up = tTR_tAcc[(None, None, None, epi_m_idx, up_n_subtile)] + + cute.copy(tiled_copy_t2r, tTR_tAcc_mn_up, tTR_rAcc_up) + if cutlass.const_expr(self.is_gated): + tTR_tAcc_mn_gate = tTR_tAcc[ + (None, None, None, epi_m_idx, gate_n_subtile) + ] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn_gate, tTR_rAcc_gate) + + # Overlap mode: after iter 0 the up/gate LDTM has + # covered cols 192..255 (reverse for acc[0], forward + # for acc[1]). Fence + early-release so MMA can write + # the next stage into the overlap region without racing. + if cutlass.const_expr(self.use_overlap_accum): + if out_n_idx == 0: + cute.arch.fence_view_async_tmem_load() + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + acc_vec_up = tTR_rAcc_up.load() + tCompute = cute.make_rmem_tensor(acc_vec_up.shape, self.acc_dtype) + if cutlass.const_expr(self.activation_type == ActivationType.Swiglu): + acc_vec_gate = tTR_rAcc_gate.load() + self._apply_swiglu_epilogue( + acc_vec_up, acc_vec_gate, alpha_val, tCompute + ) + elif cutlass.const_expr(self.activation_type == ActivationType.Relu2): + self._apply_relu2_epilogue(acc_vec_up, alpha_val, tCompute) + + if cutlass.const_expr(self.generate_sfc): + # Float4E2M1FN quantization: per-vector absmax → + # SFC → store SFC to gmem → quantize output by + # reciprocal of SFC. (Subtile is partitioned on N.) + # locality domain: shift the SFC N-subtile by c_sf_n_tile_offset + # so this partition writes into the shared full-width + # SF buffer at the correct N tile (0 in non-locality domain). + sfc_subtile_idx_mn = ( + tile_info[0] * self.epi_tile_cnt[0] + epi_m_idx, + c_sf_n_tile_offset + + tile_info[1] * self.epi_tile_cnt[1] + + real_out_n_idx, + ) + tCgSFC = tCgSFC_mn[ + ( + None, + None, + None, + *sfc_subtile_idx_mn, + ) + ] + + # + # Get absolute max across a vector and Compute SFC + # + tTR_rAcc_frg = cute.logical_divide( + tCompute, cute.make_layout(self.sf_vec_size) + ) + acc_frg = tTR_rAcc_frg.load() + acc_frg = epilogue_op(acc_frg) + + # Apply element-wise absolute value using math.absf (supports vectors) + abs_acc_frg_ir = math.absf(acc_frg.ir_value()) + abs_acc_frg = type(acc_frg)( + abs_acc_frg_ir, acc_frg.shape, acc_frg.dtype + ) + + if cutlass.const_expr(self.vectorized_f32): + for vi in cutlass.range_constexpr(abs_acc_frg.shape[1]): + tCrSFC_pvscale[vi] = abs_acc_frg[None, vi].reduce( + cute.ReductionOp.MAX, + cutlass.Float32(0.0), + 0, # Use 0.0 as init for abs values + ) + for vi in cutlass.range_constexpr(0, abs_acc_frg.shape[1], 2): + tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1] = ( + cute.arch.mul_packed_f32x2( + ( + tCrSFC_pvscale[vi], + tCrSFC_pvscale[vi + 1], + ), + ( + self.get_dtype_rcp_limits(self.c_dtype), + self.get_dtype_rcp_limits(self.c_dtype), + ), + ) + ) + tCrSFC_pvscale[vi], tCrSFC_pvscale[vi + 1] = ( + cute.arch.mul_packed_f32x2( + ( + tCrSFC_pvscale[vi], + tCrSFC_pvscale[vi + 1], + ), + (norm_const, norm_const), + ) + ) + else: + for vi in cutlass.range_constexpr(abs_acc_frg.shape[1]): + tCrSFC_pvscale[vi] = ( + abs_acc_frg[None, vi].reduce( + cute.ReductionOp.MAX, + cutlass.Float32(0.0), + 0, # Use 0.0 as init for abs values + ) + * self.get_dtype_rcp_limits(self.c_dtype) + * norm_const + ) + + # TODO: f32x2 -> f8x2 conversion + tCrSFC.store(tCrSFC_pvscale.load().to(self.sf_dtype)) + + # Store SFC to gmem. + # TODO: predicate (cute.elem_less) + cute.autovec_copy(tCrSFC, tCgSFC) + + # Quantize output and convert to c_dtype. + # TODO: need to add f8x2 -> f32x2 conversion + tCrSFC_qpvscale_up = tCrSFC.load().to(cutlass.Float32) + fp32_max = cutlass.Float32(3.40282346638528859812e38) + if cutlass.const_expr(self.vectorized_f32): + for vi in cutlass.range_constexpr(0, cute.size(tCrSFC), 2): + acc_scale = cute.arch.mul_packed_f32x2( + ( + cute.arch.rcp_approx(tCrSFC_qpvscale_up[vi]), + cute.arch.rcp_approx(tCrSFC_qpvscale_up[vi + 1]), + ), + (norm_const, norm_const), + ) + acc_scale_min0 = fmin(acc_scale[0], fp32_max, nan=True) + acc_scale_min1 = fmin(acc_scale[1], fp32_max, nan=True) + + vec0 = tTR_rAcc_frg[None, vi] + vec1 = tTR_rAcc_frg[None, vi + 1] + for ei in cutlass.range_constexpr(self.sf_vec_size): + vec0[ei], vec1[ei] = cute.arch.mul_packed_f32x2( + (vec0[ei], vec1[ei]), + (acc_scale_min0, acc_scale_min1), + ) + else: + for vi in cutlass.range_constexpr(cute.size(tCrSFC)): + # TODO:Need to add E8M0 rcp approximation + acc_scale = norm_const * cute.arch.rcp_approx( + tCrSFC_qpvscale_up[vi] + ) + acc_scale = fmin(acc_scale, fp32_max, nan=True) + + vec = tTR_rAcc_frg[None, vi] + for ei in cutlass.range_constexpr(self.sf_vec_size): + vec[ei] = vec[ei] * acc_scale + + acc_vec = tiled_copy_r2s.retile(tCompute).load() + tRS_rC.store(acc_vec.to(self.c_dtype)) + else: + # + # Convert to C type + # + acc_vec = tiled_copy_r2s.retile(tCompute).load() + acc_vec = epilogue_op(acc_vec.to(self.c_dtype)) + tRS_rC.store(acc_vec) + + # + # Store C to shared memory + # + num_prev_subtiles = num_prev_subtiles + 1 + c_buffer = num_prev_subtiles % self.num_c_stage + + cute.copy( + tiled_copy_r2s, + tRS_rC, + tRS_sC[(None, None, None, c_buffer)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + self.epilog_sync_barrier.arrive_and_wait() + # + # TMA store C to global memory + # + if warp_idx == self.epilog_warp_id[0]: + cute.copy( + tma_atom_c, + bSG_sC[(None, c_buffer)], + bSG_gC[(None, epi_m_idx, real_out_n_idx)], + ) + # Fence and barrier to make sure shared memory store is visible to TMA store + c_pipeline.producer_commit() + c_pipeline.producer_acquire() + self.epilog_sync_barrier.arrive_and_wait() + + # + # Async arrive accumulator buffer empty. Overlap mode already + # released early inside the subtile loop; skip the final one. + # + if cutlass.const_expr(not self.use_overlap_accum): + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # + # Advance to next tile + # + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy( + "async.shared", + space="cta", + ) + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + # + # Dealloc the tensor memory buffer + # + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(tmem_ptr) + # + # Wait for C store complete + # + c_pipeline.producer_tail() + + cute.arch.griddepcontrol_launch_dependents() + + @cute.jit + def _apply_swiglu_epilogue( + self, + acc_vec_up: cute.Tensor, + acc_vec_gate: cute.Tensor, + alpha_val, + tCompute: cute.Tensor, + ): + """SwiGLU: ``tCompute[i] = (alpha * up[i]) * silu(alpha * gate[i])``.""" + if cutlass.const_expr(self.vectorized_f32): + LOG2_E = cutlass.Float32(1.4426950408889634) + for i in cutlass.range_constexpr(0, cute.size(acc_vec_up.shape), 2): + acc_vec_up_alpha = cute.arch.mul_packed_f32x2( + (acc_vec_up[i], acc_vec_up[i + 1]), + (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)), + ) + acc_vec_gate_alpha = cute.arch.mul_packed_f32x2( + (acc_vec_gate[i], acc_vec_gate[i + 1]), + (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)), + ) + tCompute_log2e = cute.arch.mul_packed_f32x2( + (acc_vec_gate_alpha[0], acc_vec_gate_alpha[1]), + (-LOG2_E, -LOG2_E), + ) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.add_packed_f32x2( + ( + cute.math.exp2(tCompute_log2e[0], fastmath=True), + cute.math.exp2(tCompute_log2e[1], fastmath=True), + ), + (1.0, 1.0), + ) + tCompute[i] = cute.arch.rcp_approx(tCompute[i]) + tCompute[i + 1] = cute.arch.rcp_approx(tCompute[i + 1]) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (acc_vec_gate_alpha[0], acc_vec_gate_alpha[1]), + ) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (tCompute[i], tCompute[i + 1]), + (acc_vec_up_alpha[0], acc_vec_up_alpha[1]), + ) + else: + for i in cutlass.range_constexpr(cute.size(acc_vec_up.shape)): + acc_vec_up_alpha = acc_vec_up[i] * cutlass.Float32(alpha_val) + acc_vec_gate_alpha = acc_vec_gate[i] * cutlass.Float32(alpha_val) + tCompute[i] = acc_vec_up_alpha * silu_f32(acc_vec_gate_alpha, fastmath=True) + + @cute.jit + def _apply_relu2_epilogue( + self, + acc_vec_up: cute.Tensor, + alpha_val, + tCompute: cute.Tensor, + ): + """Relu2: ``tCompute[i] = relu(alpha * up[i]) ** 2``.""" + if cutlass.const_expr(self.vectorized_f32): + for i in cutlass.range_constexpr(0, cute.size(acc_vec_up.shape), 2): + scaled = cute.arch.mul_packed_f32x2( + (acc_vec_up[i], acc_vec_up[i + 1]), + (cutlass.Float32(alpha_val), cutlass.Float32(alpha_val)), + ) + relu0 = cute.arch.fmax(scaled[0], 0.0) + relu1 = cute.arch.fmax(scaled[1], 0.0) + ( + tCompute[i], + tCompute[i + 1], + ) = cute.arch.mul_packed_f32x2( + (relu0, relu1), + (relu0, relu1), + ) + else: + for i in cutlass.range_constexpr(cute.size(acc_vec_up.shape)): + scaled = acc_vec_up[i] * cutlass.Float32(alpha_val) + relu_val = cute.arch.fmax(scaled, 0.0) + tCompute[i] = relu_val * relu_val + + def epilog_tmem_copy_and_partition( + self, + tidx: cutlass.Int32, + tAcc: cute.Tensor, + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + use_2cta_instrs: Union[cutlass.Boolean, bool], + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for tensor memory load, then use it to partition tensor memory + (source) and register array (destination). + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param tAcc: The accumulator tensor to be copied and partitioned + :type tAcc: cute.Tensor + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param use_2cta_instrs: Whether use_2cta_instrs is enabled + :type use_2cta_instrs: bool + + :return: A tuple containing (tiled_copy_t2r, tTR_tAcc, tTR_rAcc_up, tTR_rAcc_gate) where: + - tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + - tTR_tAcc: The partitioned accumulator tensor + - tTR_rAcc_up: The partitioned accumulator tensor for acc up + - tTR_rAcc_gate: The partitioned accumulator tensor for acc gate + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor, cute.Tensor] + """ + # Make tiledCopy for tensor memory load (Rubin uses transformed layout) + copy_atom_t2r = sm100_utils.get_tmem_load_op( + self.cta_tile_shape_mnk, + self.c_layout, + self.c_dtype, + self.acc_dtype, + epi_tile, + use_2cta_instrs, + ) + + # tAcc is already transformed: (M, N, STAGE) layout + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, STAGE) + tAcc_epi = cute.flat_divide( + tAcc, + epi_tile, + ) + # (EPI_TILE_M, EPI_TILE_N) + tiled_copy_t2r = tcgen05.make_tmem_copy(copy_atom_t2r, tAcc_epi[(None, None, 0, 0, 0)]) + + thr_copy_t2r = tiled_copy_t2r.get_slice(tidx) + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, STAGE) + tTR_tAcc = thr_copy_t2r.partition_S(tAcc_epi) + + # gC_mnl is already transformed: (M, N_half, loopM, loopN, loopL) + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL) + gC_mnl_epi = cute.flat_divide(gC_mnl, epi_tile) + + # (T2R, T2R_M, T2R_N, EPI_M, EPI_N, loopM, loopN, loopL) + tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi) + + # (T2R, T2R_M, T2R_N) + tTR_rAcc_up = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype + ) + # (T2R, T2R_M, T2R_N) + tTR_rAcc_gate = cute.make_rmem_tensor( + tTR_gC[(None, None, None, 0, 0, 0, 0, 0)].shape, self.acc_dtype + ) + return tiled_copy_t2r, tTR_tAcc, tTR_rAcc_up, tTR_rAcc_gate + + def epilog_smem_copy_and_partition( + self, + tiled_copy_t2r: cute.TiledCopy, + tTR_rC: cute.Tensor, + tidx: cutlass.Int32, + sC: cute.Tensor, + ) -> Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor]: + """ + Make tiledCopy for shared memory store, then use it to partition register + array (source) and shared memory (destination). + + :param tiled_copy_t2r: The tiled copy operation for tmem to register copy(t2r) + :type tiled_copy_t2r: cute.TiledCopy + :param tTR_rC: The partitioned accumulator tensor + :type tTR_rC: cute.Tensor + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + :type sepi: cute.Tensor + + :return: A tuple containing (tiled_copy_r2s, tRS_rC, tRS_sC) where: + - tiled_copy_r2s: The tiled copy operation for register to smem copy(r2s) + - tRS_rC: The partitioned tensor C (register source) + - tRS_sC: The partitioned tensor C (smem destination) + :rtype: Tuple[cute.TiledCopy, cute.Tensor, cute.Tensor] + """ + copy_atom_r2s = sm100_utils.get_smem_store_op( + self.c_layout, self.c_dtype, self.acc_dtype, tiled_copy_t2r + ) + tiled_copy_r2s = cute.make_tiled_copy_D(copy_atom_r2s, tiled_copy_t2r) + # (R2S, R2S_M, R2S_N, PIPE_D) + thr_copy_r2s = tiled_copy_r2s.get_slice(tidx) + tRS_sC = thr_copy_r2s.partition_D(sC) + # (R2S, R2S_M, R2S_N) + tRS_rC = tiled_copy_r2s.retile(tTR_rC) + return tiled_copy_r2s, tRS_rC, tRS_sC + + def epilog_gmem_copy_and_partition( + self, + tidx: cutlass.Int32, + atom: Union[cute.CopyAtom, cute.TiledCopy], + gC_mnl: cute.Tensor, + epi_tile: cute.Tile, + sC: cute.Tensor, + ) -> Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor]: + """Make tiledCopy for global memory store, then use it to: + - partition register array (source) and global memory (destination) for none TMA store version; + - partition shared memory (source) and global memory (destination) for TMA store version. + + :param tidx: The thread index in epilogue warp groups + :type tidx: cutlass.Int32 + :param atom: The copy_atom_c to be used for TMA store version, or tiled_copy_t2r for none TMA store version + :type atom: cute.CopyAtom or cute.TiledCopy + :param gC_mnl: The global tensor C + :type gC_mnl: cute.Tensor + :param epi_tile: The epilogue tiler + :type epi_tile: cute.Tile + :param sC: The shared memory tensor to be copied and partitioned + :type sC: cute.Tensor + + :return: A tuple containing : + - For TMA store: (tma_atom_c, bSG_sC, bSG_gC) where: + - tma_atom_c: The TMA copy atom + - bSG_sC: The partitioned shared memory tensor C + - bSG_gC: The partitioned global tensor C + :rtype: Tuple[cute.CopyAtom, cute.Tensor, cute.Tensor] + """ + # gC_mnl is already transformed: (M, N_half, loopM, loopN, loopL) + # (EPI_TILE_M, EPI_TILE_N, EPI_M, EPI_N, loopM, loopN, loopL) + gC_epi = cute.flat_divide(gC_mnl, epi_tile) + tma_atom_c = atom + sC_for_tma_partition = cute.group_modes(sC, 0, 2) + gC_for_tma_partition = cute.group_modes(gC_epi, 0, 2) + # ((ATOM_V, REST_V), EPI_M, EPI_N) + # ((ATOM_V, REST_V), EPI_M, EPI_N, loopM, loopN, loopL) + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_atom_c, + 0, + cute.make_layout(1), + sC_for_tma_partition, + gC_for_tma_partition, + ) + return tma_atom_c, bSG_sC, bSG_gC + + @staticmethod + def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + cta_tile_shape_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + epi_tile: cute.Tile, + c_dtype: Type[cutlass.Numeric], + c_layout: utils.LayoutEnum, + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + num_smem_capacity: int, + occupancy: int, + with_breuse: bool = False, + ) -> Tuple[int, int, int]: + """Computes the number of stages for A/B/C operands based on heuristics. + + :param tiled_mma: The tiled MMA object defining the core computation. + :type tiled_mma: cute.TiledMma + :param mma_tiler_mnk: The shape (M, N, K) of the MMA tiler. + :type mma_tiler_mnk: tuple[int, int, int] + :param a_dtype: Data type of operand A. + :type a_dtype: type[cutlass.Numeric] + :param b_dtype: Data type of operand B. + :type b_dtype: type[cutlass.Numeric] + :param epi_tile: The epilogue tile shape. + :type epi_tile: cute.Tile + :param c_dtype: Data type of operand C (output). + :type c_dtype: type[cutlass.Numeric] + :param c_layout: Layout of operand C. + :type c_layout: utils.LayoutEnum + :param sf_dtype: Data type of scale factor. + :type sf_dtype: type[cutlass.Numeric] + :param sf_vec_size: Vector size of scale factor. + :type sf_vec_size: int + :param num_smem_capacity: Total available shared memory capacity in bytes. + :type num_smem_capacity: int + :param occupancy: Target number of CTAs per SM (occupancy). + :type occupancy: int + + :return: A tuple containing the computed number of stages for: + (ACC stages, A/B operand stages, C stages) + :rtype: tuple[int, int, int] + """ + # Default ACC stages + num_acc_stage = 1 if (with_breuse and mma_tiler_mnk[1] in {192, 256}) else 2 + + # Default C stages + num_c_stage = 2 + + # Default Tile info stages + num_tile_stage = 2 + + # Calculate smem layout and size for one stage of A, B, and C + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + tiled_mma, + mma_tiler_mnk, + a_dtype, + 1, # a tmp 1 stage is provided + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, + mma_tiler_mnk, + b_dtype, + 1, # a tmp 1 stage is provided + ) + + sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, + mma_tiler_mnk, + sf_vec_size, + 1, # a tmp 1 stage is provided + ) + + c_smem_layout_staged_one = sm100_utils.make_smem_layout_epi( + c_dtype, + c_layout, + epi_tile, + 1, + ) + + # SFA SMEM is plain linear (M, tile_K_sf), no pad. + # Per stage = cta_tile_M × tile_K_sf bytes (FP8 = 1 byte/element). + sfa_tile_k_sf = cta_tile_shape_mnk[2] // sf_vec_size + sf_bytes_per_row = sfa_tile_k_sf * sf_dtype.width // 8 + sfa_bytes_per_stage_one = cta_tile_shape_mnk[0] * sf_bytes_per_row + + ab_bytes_per_stage = ( + cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + + sfa_bytes_per_stage_one + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + # 1024B alignment + mbar_helpers_bytes = 1024 + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * num_c_stage + + # Calculate A/B stages: + # Start with total smem per CTA (capacity / occupancy) + # Subtract reserved bytes and initial C stages bytes + # Divide remaining by bytes needed per A/B stage + num_ab_stage = ( + num_smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes) + ) // ab_bytes_per_stage + + # Refine epilogue stages: + # Calculate remaining smem after allocating for A/B stages and reserved bytes + # Add remaining unused smem to epilogue + num_c_stage += ( + num_smem_capacity + - occupancy * ab_bytes_per_stage * num_ab_stage + - occupancy * (mbar_helpers_bytes + c_bytes) + ) // (occupancy * c_bytes_per_stage) + return num_acc_stage, num_ab_stage, num_c_stage, num_tile_stage + + @staticmethod + def _compute_grid( + c: cute.Tensor, + cta_tile_shape_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + max_active_clusters: cutlass.Constexpr, + raster_along_m: bool = False, + ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]: + """Use persistent tile scheduler to compute the grid size for the output tensor C. + + :param c: The output tensor C + :type c: cute.Tensor + :param cta_tile_shape_mnk: The shape (M, N, K) of the CTA tile. + :type cta_tile_shape_mnk: tuple[int, int, int] + :param cluster_shape_mn: Shape of each cluster in M, N dimensions. + :type cluster_shape_mn: tuple[int, int] + :param max_active_clusters: Maximum number of active clusters. + :type max_active_clusters: cutlass.Constexpr + + :return: A tuple containing: + - tile_sched_params: Parameters for the persistent tile scheduler. + - grid: Grid shape for kernel launch. + :rtype: Tuple[utils.PersistentTileSchedulerParams, tuple[int, int, int]] + """ + c_shape = cute.slice_(cta_tile_shape_mnk, (None, None, 0)) + gc = cute.zipped_divide(c, tiler=c_shape) + num_ctas_mnl = gc[(0, (None, None, None))].shape + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl, raster_along_m=raster_along_m + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def _get_tma_atom_kind( + atom_sm_cnt: cutlass.Int32, mcast: cutlass.Boolean + ) -> Union[cpasync.CopyBulkTensorTileG2SMulticastOp, cpasync.CopyBulkTensorTileG2SOp]: + """ + Select the appropriate TMA copy atom based on the number of SMs and the multicast flag. + + :param atom_sm_cnt: The number of SMs + :type atom_sm_cnt: cutlass.Int32 + :param mcast: The multicast flag + :type mcast: cutlass.Boolean + + :return: The appropriate TMA copy atom kind + :rtype: cpasync.CopyBulkTensorTileG2SMulticastOp or cpasync.CopyBulkTensorTileG2SOp + + :raise ValueError: If the atom_sm_cnt is invalid + """ + if atom_sm_cnt == 2 and mcast: + return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.TWO) + elif atom_sm_cnt == 2 and not mcast: + return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.TWO) + elif atom_sm_cnt == 1 and mcast: + return cpasync.CopyBulkTensorTileG2SMulticastOp(tcgen05.CtaGroup.ONE) + elif atom_sm_cnt == 1 and not mcast: + return cpasync.CopyBulkTensorTileG2SOp(tcgen05.CtaGroup.ONE) + + raise ValueError(f"Invalid atom_sm_cnt: {atom_sm_cnt} and {mcast}") + + @staticmethod + def get_dtype_rcp_limits(dtype: Type[cutlass.Numeric]) -> float: + """ + Calculates the reciprocal of the maximum absolute value for a given data type. + + :param dtype: Data type + :type dtype: Type[cutlass.Numeric] + + :return: An float representing the reciprocal of the maximum absolute value + :rtype: float + """ + if dtype == cutlass.Float4E2M1FN: + return 1 / 6.0 + if dtype == cutlass.Float8E4M3FN: + return 1 / 448.0 + if dtype == cutlass.Float8E5M2: + return 1 / 128.0 + return 1.0 + + @staticmethod + def is_valid_dtypes_and_scale_factor_vec_size( + ab_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + ) -> bool: + """ + Check if the dtypes are valid + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param sf_dtype: The data type of the scale factor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: The vector size of the scale factor + :type sf_vec_size: int + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + + :return: True if the dtypes are valid, False otherwise + :rtype: bool + """ + is_valid = True + if ab_dtype not in { + cutlass.Float4E2M1FN, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + }: + is_valid = False + + # Check valid sf_vec_size + if sf_vec_size not in {16, 32}: + is_valid = False + + # Check valid sf_dtype + if sf_dtype not in {cutlass.Float8E8M0FNU, cutlass.Float8E4M3FN}: + is_valid = False + + # Check valid sf_dtype and sf_vec_size combinations + if sf_dtype == cutlass.Float8E4M3FN and sf_vec_size == 32: + is_valid = False + if ab_dtype in {cutlass.Float8E5M2, cutlass.Float8E4M3FN} and sf_vec_size == 16: + is_valid = False + + # Check valid c_dtype + if c_dtype not in { + cutlass.Float32, + cutlass.Float16, + cutlass.BFloat16, + cutlass.Float8E5M2, + cutlass.Float8E4M3FN, + cutlass.Float4E2M1FN, + }: + is_valid = False + + return is_valid + + @staticmethod + def is_valid_layouts( + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if layouts and dtypes are valid combinations + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major dimension of the A tensor + :type a_major: str + :param b_major: The major dimension of the B tensor + :type b_major: str + :param c_major: The major dimension of the C tensor + :type c_major: str + + :return: True if the layouts are valid, False otherwise + :rtype: bool + """ + is_valid = True + + if ab_dtype is cutlass.Float4E2M1FN and not (a_major == "k" and b_major == "k"): + is_valid = False + if c_dtype is cutlass.Float4E2M1FN and c_major == "m": + is_valid = False + return is_valid + + @staticmethod + def is_valid_mma_tiler_and_cluster_shape( + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + mma_inst_shape: Tuple[int, int, int], + mma_tiler: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + ) -> bool: + """Check if the mma tiler and cluster shape are valid.""" + # Check valid mma_inst_shape + if mma_inst_shape[0] not in [128, 256]: + return False + # SwiGLU Fusion requires even epi_tile counts + if mma_inst_shape[1] not in [128, 256]: + return False + + # Check valid mma_tiler + if mma_tiler[0] not in [128, 256, 512]: + return False + if mma_tiler[1] not in [128, 256]: + return False + + # Check MMA tiler vs MMA instruction relationship + # mma_tiler[0] == mma_inst_shape[0] (no B-reuse) or 2 * mma_inst_shape[0] (B-reuse) + if mma_tiler[0] not in (mma_inst_shape[0], 2 * mma_inst_shape[0]): + return False + if mma_tiler[1] != mma_inst_shape[1]: + return False + + # Check K-dimension constraints based on data type + if a_dtype in {cutlass.Float8E4M3FN, cutlass.Float8E5M2} and b_dtype in { + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }: + if mma_tiler[2] != 128 or mma_inst_shape[2] != 64: + return False + else: + if mma_tiler[2] != 256 or mma_inst_shape[2] != 128: + return False + + # Cluster-M constraint: cluster_M must EQUAL the MMA CTA-group size along M + # (atom_cta_m): 1 for 1-CTA MMA, 2 for 2-CTA MMA (mma_inst_shape[0] == 256). + # Splitting the gathered-token (M) dimension across MORE cluster CTAs than the + # MMA group (cluster_M > atom_cta_m) is not correctly handled by the gather / + # tile-scheduler row mapping and produces wrong output rows (verified: ~1-2% + # of rows mismatch on the nvf4 accuracy sweep), so reject it. + # Cluster-N multicast of A is unaffected and remains supported. + atom_cta_m = 2 if mma_inst_shape[0] == 256 else 1 + if cluster_shape_mn[0] != atom_cta_m: + return False + + # Check cluster shape validity + def _is_power_of_2(x): + return x > 0 and (x & (x - 1)) == 0 + + if ( + cluster_shape_mn[0] * cluster_shape_mn[1] > 16 + or cluster_shape_mn[0] <= 0 + or cluster_shape_mn[1] <= 0 + or cluster_shape_mn[0] > 4 + or cluster_shape_mn[1] > 4 + or not _is_power_of_2(cluster_shape_mn[0]) + or not _is_power_of_2(cluster_shape_mn[1]) + ): + return False + + return True + + @staticmethod + def is_valid_tensor_alignment( + m: cutlass.Int64, + n: cutlass.Int64, + k: cutlass.Int64, + l: cutlass.Int64, # noqa: E741 + ab_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if the tensor alignment is valid + + :param m: The number of rows in the A tensor + :type m: cutlass.Int64 + :param n: The number of columns in the B tensor + :type n: cutlass.Int64 + :param k: The number of columns in the A tensor + :type k: cutlass.Int64 + :param l: The number of columns in the C tensor + :type l: cutlass.Int64 + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param a_major: The major axis of the A tensor + :type a_major: str + :param b_major: The major axis of the B tensor + :type b_major: str + :param c_major: The major axis of the C tensor + :type c_major: str + + :return: True if the problem shape is valid, False otherwise + :rtype: bool + """ + is_valid = True + + def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contigous_16B_alignment(ab_dtype, a_major == "m", (m, k, l)) + or not check_contigous_16B_alignment(ab_dtype, b_major == "n", (n, k, l)) + or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + is_valid = False + return is_valid + + @classmethod + def can_implement( + cls, + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + mma_inst_shape: Tuple[int, int, int], + mma_tiler: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + m: cutlass.Int64, + n: cutlass.Int64, + k: cutlass.Int64, + l: cutlass.Int64, # noqa: E741 + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """ + Check if the gemm can be implemented + + :param ab_dtype: The data type of the A and B operands + :type ab_dtype: Type[cutlass.Numeric] + :param sf_dtype: The data type of the scale factor + :type sf_dtype: Type[cutlass.Numeric] + :param sf_vec_size: The vector size of the scale factor + :type sf_vec_size: int + :param c_dtype: The data type of the output tensor + :type c_dtype: Type[cutlass.Numeric] + :param mma_tiler_mn: The (M, N) shape of the MMA instruction tiler + :type mma_tiler_mn: Tuple[int, int] + :param cluster_shape_mn: The (ClusterM, ClusterN) shape of the CTA cluster + :type cluster_shape_mn: Tuple[int, int] + :param m: The number of rows in the A tensor + :type m: cutlass.Int64 + :param n: The number of columns in the B tensor + :type n: cutlass.Int64 + :param k: The number of columns in the A tensor + :type k: cutlass.Int64 + :param l: The number of columns in the C tensor + :type l: cutlass.Int64 + :param a_major: The major axis of the A tensor + :type a_major: str + :param b_major: The major axis of the B tensor + :type b_major: str + :param c_major: The major axis of the C tensor + :type c_major: str + + :return: True if the gemm can be implemented, False otherwise + :rtype: bool + """ + # Check data types + if not cls.is_valid_dtypes_and_scale_factor_vec_size( + a_dtype, sf_dtype, sf_vec_size, c_dtype + ): + return False + + # Check layouts + if not cls.is_valid_layouts(a_dtype, c_dtype, a_major, b_major, c_major): + return False + + # Check MMA tiler and cluster shape + if not cls.is_valid_mma_tiler_and_cluster_shape( + a_dtype, b_dtype, mma_inst_shape, mma_tiler, cluster_shape_mn + ): + return False + + # Check tensor alignment + if not cls.is_valid_tensor_alignment( + m, n, k, l, a_dtype, c_dtype, a_major, b_major, c_major + ): + return False + + # Check A/B layout + if not (a_major == "k" and b_major == "k"): + return False + return True + + @cute.jit + def wrapper( + self, + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + a_sf_ptr: cute.Pointer, + b_sf_ptr: cute.Pointer, + c_ptr: cute.Pointer, + c_sf_ptr: cute.Pointer, + alpha_ptr: cute.Pointer, + tile_idx_to_group_idx_ptr: cute.Pointer, + tile_idx_to_mn_limit_ptr: cute.Pointer, + token_id_mapping_ptr: cute.Pointer, + num_non_exiting_tiles_ptr: cute.Pointer, + global_sf_ptr: cute.Pointer, + orig_m: cutlass.Int64, + m: cutlass.Int64, + n: cutlass.Int64, + k: cutlass.Int64, + l: cutlass.Int64, # noqa: E741 + tile_size: cutlass.Constexpr, + scaling_vector_size: cutlass.Constexpr, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + c_stride_m: cutlass.Int64 = cutlass.Int64(0), + c_sf_n_tile_offset: cutlass.Int64 = cutlass.Int64(0), + ): + scale_k = k // scaling_vector_size + interm_size = n // 2 if self.is_gated else n + num_tiles = m // tile_size + a = cute.make_tensor( + a_ptr, layout=cute.make_ordered_layout((orig_m, k, 1), order=(1, 0, 2)) + ) + b = cute.make_tensor(b_ptr, layout=cute.make_ordered_layout((n, k, l), order=(1, 0, 2))) + a_sf = cute.make_tensor( + a_sf_ptr, + layout=cute.make_ordered_layout((orig_m, scale_k, 1), order=(1, 0, 2)), + ) + b_sf = cute.make_tensor( + b_sf_ptr, + layout=cute.make_ordered_layout( + (32, 4, n // 128, 4, scale_k // 4, l), order=(2, 1, 4, 0, 3, 5) + ), + ) + # c: runtime Int64 row stride. For locality domain half-GEMM, two partitions + # interleave their N-halves into one shared full-width buffer + # (c_stride_m = full intermediate size). A runtime stride also avoids a + # cutlass-dsl MLIR alignment bug seen with + # make_layout(..., stride=ordered_layout.stride). c_stride_m == 0 -> + # natural interm_size stride (non-locality domain, == make_ordered_layout). + actual_c_stride_m = interm_size if c_stride_m == 0 else c_stride_m + c = cute.make_tensor( + c_ptr, + layout=cute.make_layout( + (m, interm_size, 1), + stride=(actual_c_stride_m, 1, m * actual_c_stride_m), + ), + ) + # full_c_shape gives SFC the full-N M-tile stride in locality domain mode so the + # shared SF buffer is written without copy-back; None → use c.shape. + if cutlass.const_expr(not self.locality_domain_half_gemm): + full_c_shape = None + else: + full_interm_size = 2 * interm_size + full_c_shape = cute.make_ordered_layout((m, full_interm_size, 1), order=(0, 1, 2)).shape + c_sf = cute.make_tensor( + c_sf_ptr, + layout=cute.make_ordered_layout( + (32, 4, m // 128, 4, interm_size // (scaling_vector_size * 4), l), + order=(2, 1, 4, 0, 3, 5), + ), + ) + alpha = cute.make_tensor(alpha_ptr, layout=cute.make_layout((l,))) + + tile_idx_to_group_idx = cute.make_tensor( + tile_idx_to_group_idx_ptr, layout=cute.make_layout((num_tiles,)) + ) + tile_idx_to_mn_limit = cute.make_tensor( + tile_idx_to_mn_limit_ptr, layout=cute.make_layout((num_tiles,)) + ) + token_id_mapping = cute.make_tensor(token_id_mapping_ptr, layout=cute.make_layout((m,))) + num_non_exiting_tiles = cute.make_tensor( + num_non_exiting_tiles_ptr, layout=cute.make_layout((1,)) + ) + global_sf = cute.make_tensor(global_sf_ptr, layout=cute.make_layout((1,))) + + return self( + a, + b, + c, + a_sf, + b_sf, + c_sf, + full_c_shape, + global_sf, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + token_id_mapping, + num_non_exiting_tiles, + alpha, + max_active_clusters=max_active_clusters, + stream=stream, + epilogue_op=epilogue_op, + c_sf_n_tile_offset=c_sf_n_tile_offset, + ) + + +@cute.jit +def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + sf_ref_tensor: cute.Tensor, + sf_mma_tensor: cute.Tensor, +): + """Convert scale factors from MKL layout to the MMA scale-factor layout.""" + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3) + for i in cutlass.range(cute.size(sf_ref_tensor)): + mkl_coord = sf_ref_tensor.layout.get_hier_coord(i) + sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord] + + +@cute.jit +def cvt_sf_M32x4xrm_K4xrk_L_to_MKL( + sf_swizzled_tensor: cute.Tensor, + sf_unswizzled_tensor: cute.Tensor, +): + """Convert scale factors from the MMA scale-factor layout to MKL layout.""" + sf_swizzled_tensor = cute.group_modes(sf_swizzled_tensor, 0, 3) + sf_swizzled_tensor = cute.group_modes(sf_swizzled_tensor, 1, 3) + for i in cutlass.range(cute.size(sf_unswizzled_tensor)): + mkl_coord = sf_unswizzled_tensor.layout.get_hier_coord(i) + sf_unswizzled_tensor[mkl_coord] = sf_swizzled_tensor[mkl_coord] + + +def create_mask(group_m_list, mma_tiler_m, permuted_m=None): + """Create group metadata for contiguous grouped GEMM with gather.""" + valid_m = 0 + aligned_group_m_list = [] + tile_idx_to_expert_idx = [] + tile_idx_to_mn_limit = [] + + for i, group_m in enumerate(group_m_list): + aligned_group_m = ((group_m + mma_tiler_m - 1) // mma_tiler_m) * mma_tiler_m + aligned_group_m_list.append(aligned_group_m) + + num_tiles_in_group = aligned_group_m // mma_tiler_m + tile_idx_to_expert_idx.extend([i] * num_tiles_in_group) + for tile_idx_in_group in range(num_tiles_in_group): + tile_idx_to_mn_limit.append( + valid_m + min(tile_idx_in_group * mma_tiler_m + mma_tiler_m, group_m) + ) + valid_m += aligned_group_m + + num_non_exiting_tiles = len(tile_idx_to_expert_idx) + + if permuted_m is not None: + if permuted_m < valid_m: + raise ValueError(f"permuted_m ({permuted_m}) must be >= valid_m ({valid_m}).") + if (permuted_m - valid_m) % mma_tiler_m != 0: + raise ValueError( + f"permuted_m ({permuted_m}) must be aligned to tile M " + f"({mma_tiler_m}) after valid_m ({valid_m})." + ) + if permuted_m > valid_m: + num_padding_tiles = (permuted_m - valid_m) // mma_tiler_m + tile_idx_to_expert_idx.extend([int(-2e9)] * num_padding_tiles) + tile_idx_to_mn_limit.extend([int(-2e9)] * num_padding_tiles) + + tile_idx_to_expert_idx = torch.tensor(tile_idx_to_expert_idx, device="cuda", dtype=torch.int32) + num_non_exiting_tiles_tensor = torch.tensor( + [num_non_exiting_tiles], device="cuda", dtype=torch.int32 + ) + tile_idx_to_mn_limit_tensor = torch.tensor( + tile_idx_to_mn_limit, device="cuda", dtype=torch.int32 + ) + + return ( + valid_m, + aligned_group_m_list, + tile_idx_to_expert_idx, + num_non_exiting_tiles_tensor, + tile_idx_to_mn_limit_tensor, + ) + + +def create_scale_factor_tensor(num_groups, mn, k, sf_vec_size, dtype): + def ceil_div(a, b): + return (a + b - 1) // b + + sf_k = ceil_div(k, sf_vec_size) + ref_shape = (num_groups, mn, sf_k) + + atom_m = (32, 4) + atom_k = 4 + mma_shape = ( + num_groups, + ceil_div(mn, atom_m[0] * atom_m[1]), + ceil_div(sf_k, atom_k), + atom_m[0], + atom_m[1], + atom_k, + ) + + ref_permute_order = (1, 2, 0) + mma_permute_order = (3, 4, 1, 5, 2, 0) + + ref_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + ref_shape, + torch.float32, + permute_order=ref_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=1, max_val=3), + ) + + cute_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + mma_shape, + torch.float32, + permute_order=mma_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=0, max_val=1), + ) + + cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + from_dlpack(ref_f32_torch_tensor_cpu), + from_dlpack(cute_f32_torch_tensor_cpu), + ) + + cute_f32_torch_tensor = cute_f32_torch_tensor_cpu.cuda() + + ref_f32_torch_tensor_cpu = ( + ref_f32_torch_tensor_cpu.permute(2, 0, 1) + .unsqueeze(-1) + .expand(num_groups, mn, sf_k, sf_vec_size) + .reshape(num_groups, mn, sf_k * sf_vec_size) + .permute(*ref_permute_order) + ) + ref_f32_torch_tensor_cpu = ref_f32_torch_tensor_cpu[:, :k, :] + + cute_tensor, cute_torch_tensor = cutlass_torch.cute_tensor_like( + cute_f32_torch_tensor_cpu, + dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + + cute_tensor = cutlass_torch.convert_cute_tensor( + cute_f32_torch_tensor, + cute_tensor, + dtype, + is_dynamic_layout=True, + ) + return ref_f32_torch_tensor_cpu, cute_tensor, cute_torch_tensor + + +def create_scale_factor_tensor_unswizzled(num_groups, mn, k, sf_vec_size, dtype): + def ceil_div(a, b): + return (a + b - 1) // b + + sf_k = ceil_div(k, sf_vec_size) + sf_ref = cutlass_torch.matrix( + num_groups, + mn, + sf_k, + False, + cutlass.Float32, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=1, max_val=3), + ) + sf_tensor, sf_torch = cutlass_torch.cute_tensor_like( + sf_ref, dtype, is_dynamic_layout=True, assumed_align=16 + ) + + sf_ref = ( + sf_ref.permute(2, 0, 1) + .unsqueeze(-1) + .expand(num_groups, mn, sf_k, sf_vec_size) + .reshape(num_groups, mn, sf_k * sf_vec_size) + .permute(1, 2, 0) + ) + sf_ref = sf_ref[:, :k, :] + return sf_ref, sf_tensor, sf_torch + + +def create_sf_layout_tensor(num_groups, mn, nk, sf_vec_size): + def ceil_div(a, b): + return (a + b - 1) // b + + sf_k = ceil_div(nk, sf_vec_size) + + atom_m = (32, 4) + atom_k = 4 + mma_shape = ( + num_groups, + ceil_div(mn, atom_m[0] * atom_m[1]), + ceil_div(sf_k, atom_k), + atom_m[0], + atom_m[1], + atom_k, + ) + + mma_permute_order = (3, 4, 1, 5, 2, 0) + + cute_f32_torch_tensor = cutlass_torch.create_and_permute_torch_tensor( + mma_shape, + torch.float32, + permute_order=mma_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=0, max_val=1), + ) + return cute_f32_torch_tensor, sf_k + + +def create_token_id_mapping_tensor(group_m_list, mma_tiler_m, max_token_id, permuted_m=None): + """Create token_id_mapping tensor for gather with random token IDs.""" + valid_m = 0 + for group_m in group_m_list: + valid_m += ((group_m + mma_tiler_m - 1) // mma_tiler_m) * mma_tiler_m + + tensor_m = permuted_m if permuted_m is not None else valid_m + base_data = torch.full((tensor_m,), -1, dtype=torch.int32) + + accumulated_m = 0 + for group_m in group_m_list: + start_idx = accumulated_m + rounded_group_m = ((group_m + mma_tiler_m - 1) // mma_tiler_m) * mma_tiler_m + random_token_ids = torch.randint(0, max_token_id, (group_m,), dtype=torch.int32) + base_data[start_idx : start_idx + group_m] = random_token_ids + accumulated_m += rounded_group_m + + token_id_mapping_ref = base_data.clone() + token_id_mapping_tensor, token_id_mapping_torch = cutlass_torch.cute_tensor_like( + token_id_mapping_ref, cutlass.Int32, is_dynamic_layout=True, assumed_align=4 + ) + return token_id_mapping_ref, token_id_mapping_tensor, token_id_mapping_torch + + +def create_tensors( + num_groups, + group_m_list, + n, + k, + a_major, + b_major, + cd_major, + a_dtype, + b_dtype, + c_dtype, + sf_dtype, + sf_vec_size, + mma_tiler_m, + permuted_m=None, +): + """Create tensors for grouped blockscaled GEMM with gather and SwiGLU fusion.""" + torch.manual_seed(1111) + + alpha_torch_cpu = torch.randn((num_groups,), dtype=torch.float32) + + ( + valid_m, + aligned_group_m_list, + _tile_idx_to_expert_idx, + _num_non_exiting_tiles, + _tile_idx_to_mn_limit, + ) = create_mask(group_m_list, mma_tiler_m, permuted_m) + + max_m = max(group_m_list) + tensor_m = permuted_m if permuted_m is not None else valid_m + + a_torch_cpu = cutlass_torch.matrix(1, max_m, k, a_major == "m", cutlass.Float32) + b_torch_cpu = cutlass_torch.matrix(num_groups, n, k, b_major == "n", cutlass.Float32) + c_torch_cpu = cutlass_torch.matrix(1, tensor_m, n // 2, cd_major == "m", cutlass.Float32) + + a_tensor, a_torch_gpu = cutlass_torch.cute_tensor_like( + a_torch_cpu, a_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, b_torch_gpu = cutlass_torch.cute_tensor_like( + b_torch_cpu, b_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like( + c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + a_tensor.mark_compact_shape_dynamic( + mode=1 if a_major == "k" else 0, + stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), + divisibility=32 if a_dtype == cutlass.Float4E2M1FN else 16, + ) + b_tensor.mark_compact_shape_dynamic( + mode=1 if b_major == "k" else 0, + stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), + divisibility=32 if b_dtype == cutlass.Float4E2M1FN else 16, + ) + c_tensor.mark_compact_shape_dynamic( + mode=1 if cd_major == "n" else 0, + stride_order=(2, 0, 1) if cd_major == "n" else (2, 1, 0), + divisibility=32 if c_dtype == cutlass.Float4E2M1FN else 16, + ) + + sfa_torch_cpu, sfa_tensor, sfa_torch_gpu = create_scale_factor_tensor_unswizzled( + 1, max_m, k, sf_vec_size, sf_dtype + ) + sfb_torch_cpu, sfb_tensor, sfb_torch_gpu = create_scale_factor_tensor( + num_groups, n, k, sf_vec_size, sf_dtype + ) + + token_id_mapping_cpu, token_id_mapping, token_id_mapping_torch = create_token_id_mapping_tensor( + group_m_list, mma_tiler_m, max_token_id=max_m, permuted_m=permuted_m + ) + + tile_idx_to_expert_idx = from_dlpack(_tile_idx_to_expert_idx).mark_layout_dynamic() + tile_idx_to_mn_limit = from_dlpack(_tile_idx_to_mn_limit).mark_layout_dynamic() + num_non_exiting_tiles = from_dlpack(_num_non_exiting_tiles).mark_layout_dynamic() + alpha = from_dlpack(alpha_torch_cpu.cuda()).mark_layout_dynamic() + + sfc_torch_cpu = None + sfc_tensor = None + sfc_torch_gpu = None + norm_const_torch_cpu = None + norm_const_tensor = None + norm_const_torch_gpu = None + n_out = n // 2 + if c_dtype == cutlass.Float4E2M1FN: + sfc_torch_cpu, sfc_tensor, sfc_torch_gpu = create_scale_factor_tensor( + 1, tensor_m, n_out, sf_vec_size, sf_dtype + ) + norm_const_torch_gpu = torch.tensor([1.0], dtype=torch.float32, device="cuda") + norm_const_tensor = from_dlpack(norm_const_torch_gpu).mark_layout_dynamic() + norm_const_torch_cpu = norm_const_torch_gpu.cpu() + + return ( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + sfc_tensor, + norm_const_tensor, + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + token_id_mapping, + num_non_exiting_tiles, + alpha, + a_torch_cpu, + b_torch_cpu, + c_torch_cpu, + sfa_torch_cpu, + sfb_torch_cpu, + sfc_torch_cpu, + norm_const_torch_cpu, + alpha_torch_cpu, + a_torch_gpu, + b_torch_gpu, + c_torch_gpu, + sfa_torch_gpu, + sfb_torch_gpu, + sfc_torch_gpu, + norm_const_torch_gpu, + aligned_group_m_list, + valid_m, + token_id_mapping_cpu, + ) + + +def run( + nkl: Tuple[int, int, int], + group_m_list: Tuple[int, ...], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + a_major: str, + b_major: str, + c_major: str, + mma_inst_shape: Tuple[int, int, int], + mma_tiler: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + tolerance: float, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + permuted_m: Optional[int] = None, + raster_along_m: bool = False, + use_cupti: bool = False, + a_path: str = "cpasync", + use_pdl: bool = True, +): + """Run the Rubin blockscaled contiguous gather grouped GEMM SwiGLU kernel.""" + mma_tiler_m = mma_tiler[0] + + print( + "Running Rubin Persistent Contiguous Grouped BlockScaled GEMM with " + "Gather and SwiGLU Fusion:" + ) + print(f"nkl: {nkl}") + print(f"group_m_list: {group_m_list}") + print( + f"A dtype: {a_dtype}, B dtype: {b_dtype}, C dtype: {c_dtype}, " + f"SF dtype: {sf_dtype}, SF Vec size: {sf_vec_size}" + ) + if permuted_m is not None: + print(f"Padded M (CUDA graph support): {permuted_m}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, C: {c_major}") + print(f"MMA Inst Shape: {mma_inst_shape}, MMA Tiler: {mma_tiler}") + print(f"Cluster Shape: {cluster_shape_mn}") + print(f"Raster along M: {raster_along_m}") + print(f"A path: {a_path}") + print(f"Use PDL: {use_pdl}") + print(f"Use CUPTI: {use_cupti}") + + n, k, num_groups = nkl + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + if not Sm107BlockScaledContiguousGatherGroupedGemmActFusionKernel.can_implement( + a_dtype=a_dtype, + b_dtype=b_dtype, + sf_dtype=sf_dtype, + sf_vec_size=sf_vec_size, + c_dtype=c_dtype, + mma_inst_shape=mma_inst_shape, + mma_tiler=mma_tiler, + cluster_shape_mn=cluster_shape_mn, + m=mma_tiler_m, + n=n, + k=k, + l=num_groups, + a_major=a_major, + b_major=b_major, + c_major=c_major, + ): + raise TypeError( + f"Unsupported testcase a_dtype={a_dtype}, b_dtype={b_dtype}, " + f"sf_dtype={sf_dtype}, sf_vec_size={sf_vec_size}, c_dtype={c_dtype}, " + f"mma_inst_shape={mma_inst_shape}, mma_tiler={mma_tiler}, " + f"cluster_shape_mn={cluster_shape_mn}" + ) + + ( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + sfc_tensor, + norm_const_tensor, + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + token_id_mapping, + num_non_exiting_tiles, + alpha, + a_torch_cpu, + b_torch_cpu, + c_torch_cpu, + sfa_torch_cpu, + sfb_torch_cpu, + sfc_torch_cpu, + norm_const_torch_cpu, + alpha_torch_cpu, + a_torch_gpu, + b_torch_gpu, + c_torch_gpu, + sfa_torch_gpu, + sfb_torch_gpu, + sfc_torch_gpu, + norm_const_torch_gpu, + aligned_group_m_list, + valid_m, + token_id_mapping_cpu, + ) = create_tensors( + num_groups, + group_m_list, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + sf_dtype, + sf_vec_size, + mma_tiler_m, + permuted_m, + ) + + gemm = Sm107BlockScaledContiguousGatherGroupedGemmActFusionKernel( + sf_vec_size, + mma_inst_shape, + mma_tiler, + cluster_shape_mn, + True, + topk=1, + raster_along_m=raster_along_m, + a_path=a_path, + use_pdl=use_pdl, + ) + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + + torch_stream = torch.cuda.current_stream() + current_stream = cuda.CUstream(torch_stream.cuda_stream) + full_c_shape = None + + compiled_gemm = cute.compile( + gemm, + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + sfc_tensor, + full_c_shape, + norm_const_tensor, + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + token_id_mapping, + num_non_exiting_tiles, + alpha, + max_active_clusters, + current_stream, + ) + + compiled_gemm( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + sfc_tensor, + full_c_shape, + norm_const_tensor, + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + token_id_mapping, + num_non_exiting_tiles, + alpha, + current_stream, + ) + + torch.cuda.synchronize() + + if not skip_ref_check: + print("Verifying results...") + interleave_granularity = 64 + n_out = n // 2 + + gemm_result = torch.empty((1, valid_m, n), dtype=torch.float32) + start = 0 + a_torch_cpu_f32 = torch.einsum("mk,mk->mk", a_torch_cpu[:, :, 0], sfa_torch_cpu[:, :, 0]) + for i, group_m in enumerate(aligned_group_m_list): + end = start + group_m + res_a = a_torch_cpu_f32[token_id_mapping_cpu[start:end]] + res_b = torch.einsum("nk,nk->nk", b_torch_cpu[:, :, i], sfb_torch_cpu[:, :, i]) + gemm_result[0, start:end, :] = ( + torch.einsum("mk,nk->mn", res_a, res_b) * alpha_torch_cpu[i] + ) + start = end + + assert n % (2 * interleave_granularity) == 0 + ref = torch.empty((1, valid_m, n_out), dtype=torch.float32) + for n_block in range(0, n, 2 * interleave_granularity): + up_result = gemm_result[0, :, n_block : n_block + interleave_granularity] + gate_result = gemm_result[ + 0, + :, + n_block + interleave_granularity : n_block + 2 * interleave_granularity, + ] + silu_gate = gate_result * torch.sigmoid(gate_result) + output_block = up_result * silu_gate + out_start = n_block // 2 + out_end = out_start + interleave_granularity + ref[0, :, out_start:out_end] = output_block + + ref = ref.permute((1, 2, 0)) + + res = c_torch_cpu.cuda() + cute.testing.convert( + c_tensor, + from_dlpack(res, assumed_align=16).mark_layout_dynamic( + leading_dim=(1 if c_major == "n" else 0) + ), + ) + + res = res[:valid_m] + mask = token_id_mapping_cpu[:valid_m] >= 0 + res = res.cpu()[mask] + ref = ref[mask] + + print(f"valid_m: {valid_m}, ref.shape: {ref.shape}, res.shape: {res.shape}") + + if c_dtype in (cutlass.Float32, cutlass.Float16, cutlass.BFloat16): + torch.testing.assert_close(res.cpu(), ref.cpu(), atol=tolerance, rtol=1e-02) + elif c_dtype in (cutlass.Float8E5M2, cutlass.Float8E4M3FN): + ref_f8_ = torch.empty(*(1, valid_m, n_out), dtype=torch.uint8, device="cuda").permute( + 1, 2, 0 + ) + ref_f8 = from_dlpack(ref_f8_, assumed_align=16).mark_layout_dynamic(leading_dim=1) + ref_f8.element_type = c_dtype + ref_device = ref.cuda() + ref_tensor = from_dlpack(ref_device, assumed_align=16).mark_layout_dynamic( + leading_dim=1 + ) + cute.testing.convert(ref_tensor, ref_f8) + cute.testing.convert(ref_f8, ref_tensor) + torch.testing.assert_close(res.cpu(), ref_device.cpu(), atol=tolerance, rtol=1e-02) + elif c_dtype is cutlass.Float4E2M1FN: + + def ceil_div(a, b): + return (a + b - 1) // b + + def simulate_f8_quantization(tensor_f32, f8_dtype): + shape = tensor_f32.shape + f8_torch = torch.empty(*shape, dtype=torch.uint8, device="cuda") + f8_tensor = from_dlpack(f8_torch, assumed_align=16).mark_layout_dynamic( + leading_dim=1 + ) + f8_tensor.element_type = f8_dtype + f32_device = tensor_f32.cuda() + f32_tensor = from_dlpack(f32_device, assumed_align=16).mark_layout_dynamic( + leading_dim=1 + ) + cute.testing.convert(f32_tensor, f8_tensor) + cute.testing.convert(f8_tensor, f32_tensor) + return f32_device.cpu() + + def simulate_nvfp4_quantization(tensor_f32): + m_dim, n_dim, ng = tensor_f32.shape + ref_f32_torch = cutlass_torch.matrix(ng, m_dim, n_dim, False, cutlass.Float32) + f4_tensor, _ = cutlass_torch.cute_tensor_like( + ref_f32_torch, + cutlass.Float4E2M1FN, + is_dynamic_layout=True, + assumed_align=16, + ) + f32_device = tensor_f32.cuda() + f32_tensor = from_dlpack(f32_device, assumed_align=16).mark_layout_dynamic( + leading_dim=1 + ) + cute.testing.convert(f32_tensor, f4_tensor) + cute.testing.convert(f4_tensor, f32_tensor) + return f32_device.cpu() + + def compute_scale_factor(tensor_f32, sf_vec_size_local, norm_const, rcp_limits): + m_dim, n_dim, ng = tensor_f32.shape + sfn = ceil_div(n_dim, sf_vec_size_local) + padded_n = sfn * sf_vec_size_local + if padded_n > n_dim: + tensor_padded = torch.zeros(m_dim, padded_n, ng, dtype=tensor_f32.dtype) + tensor_padded[:, :n_dim, :] = tensor_f32 + else: + tensor_padded = tensor_f32 + tensor_reshaped = tensor_padded.view(m_dim, sfn, sf_vec_size_local, ng) + abs_max, _ = torch.abs(tensor_reshaped).max(dim=2) + return abs_max * norm_const * rcp_limits + + def apply_quantization_scale(tensor_f32, scale_factor, sf_vec_size_local, norm_const): + m_dim, n_dim, ng = tensor_f32.shape + sfn = scale_factor.shape[1] + fp32_max = torch.tensor(3.40282346638528859812e38, dtype=torch.float32) + scale_rcp = norm_const * scale_factor.reciprocal() + scale_rcp = torch.where(torch.isinf(scale_rcp), fp32_max, scale_rcp) + scale_rcp_expanded = scale_rcp.unsqueeze(2).expand( + m_dim, sfn, sf_vec_size_local, ng + ) + scale_rcp_expanded = scale_rcp_expanded.reshape(m_dim, sfn * sf_vec_size_local, ng) + scale_rcp_expanded = scale_rcp_expanded[:, :n_dim, :] + return tensor_f32 * scale_rcp_expanded + + def unswizzle_kernel_sfc( + sfc_tensor_local, permuted_m_local, n_out_local, sf_vec_size_local + ): + sfn = ceil_div(n_out_local, sf_vec_size_local) + unswizzled_sfc = torch.empty(permuted_m_local, sfn, 1, dtype=torch.float32) + swizzled_sfc_cpu, _ = create_sf_layout_tensor( + 1, permuted_m_local, n_out_local, sf_vec_size_local + ) + swizzled_sfc_tensor, swizzled_sfc_torch = cutlass_torch.cute_tensor_like( + swizzled_sfc_cpu, + cutlass.Float32, + is_dynamic_layout=True, + assumed_align=16, + ) + cute.testing.convert(sfc_tensor_local, swizzled_sfc_tensor) + swizzled_sfc_cpu = swizzled_sfc_torch.cpu() + cvt_sf_M32x4xrm_K4xrk_L_to_MKL( + from_dlpack(swizzled_sfc_cpu), + from_dlpack(unswizzled_sfc), + ) + return unswizzled_sfc + + norm_const = norm_const_torch_cpu.item() + rcp_limits = gemm.get_dtype_rcp_limits(c_dtype) + + ref_sfc_f32 = compute_scale_factor(ref, sf_vec_size, norm_const, rcp_limits) + ref_sfc_f32 = simulate_f8_quantization(ref_sfc_f32, sf_dtype) + + permuted_m_val = token_id_mapping_cpu.shape[0] + kernel_sfc = unswizzle_kernel_sfc(sfc_tensor, permuted_m_val, n_out, sf_vec_size) + torch.testing.assert_close( + ref_sfc_f32, kernel_sfc[:valid_m][mask], atol=tolerance, rtol=1e-02 + ) + print("SFC Tensor comparison passed!") + + ref_scaled = apply_quantization_scale(ref, ref_sfc_f32, sf_vec_size, norm_const) + ref_quantized = simulate_nvfp4_quantization(ref_scaled) + + print("Verifying C Tensor...") + res_cpu = res.cpu() + diff = torch.abs(res_cpu - ref_quantized) + within_tolerance = (diff <= tolerance) | (diff <= torch.abs(ref_quantized) * 1e-02) + pass_rate = within_tolerance.float().mean().item() + print(f"C Tensor pass rate: {pass_rate * 100:.2f}% (threshold: 95%)") + assert pass_rate >= 0.95, ( + f"Only {pass_rate * 100:.2f}% elements within tolerance, expected >= 95%" + ) + + def generate_tensors(): + ( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + sfc_tensor, + norm_const_tensor, + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + token_id_mapping, + num_non_exiting_tiles, + alpha, + *_, + ) = create_tensors( + num_groups, + group_m_list, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + sf_dtype, + sf_vec_size, + mma_tiler_m, + permuted_m, + ) + return cute.testing.JitArguments( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + sfc_tensor, + full_c_shape, + norm_const_tensor, + tile_idx_to_expert_idx, + tile_idx_to_mn_limit, + token_id_mapping, + num_non_exiting_tiles, + alpha, + current_stream, + ) + + workspace_count = 1 + if use_cold_l2: + tensor_m = permuted_m if permuted_m is not None else valid_m + one_workspace_bytes = ( + a_torch_gpu.numel() * a_torch_gpu.element_size() + + b_torch_gpu.numel() * b_torch_gpu.element_size() + + c_torch_gpu.numel() * c_torch_gpu.element_size() + + sfa_torch_gpu.numel() * sfa_torch_gpu.element_size() + + sfb_torch_gpu.numel() * sfb_torch_gpu.element_size() + + ( + sfc_torch_gpu.numel() * sfc_torch_gpu.element_size() + if sfc_torch_gpu is not None + else 0 + ) + + ( + norm_const_torch_gpu.numel() * norm_const_torch_gpu.element_size() + if norm_const_torch_gpu is not None + else 0 + ) + + (tensor_m // mma_tiler_m) * 4 + + (tensor_m // mma_tiler_m) * 4 + + tensor_m * 4 + + 1 * 4 + + alpha_torch_cpu.numel() * alpha_torch_cpu.element_size() + ) + workspace_count = cute.testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = cute.testing.benchmark( + compiled_gemm, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + use_cupti=use_cupti, + ) + + return exec_time + + +def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError as exc: + raise argparse.ArgumentTypeError( + "Invalid format. Expected comma-separated integers." + ) from exc + + +def read_benchmark_file( + filepath: str, +) -> Tuple[Tuple[int, int, int], Tuple[int, ...]]: + """Read benchmark file and return nkl plus per-group M values.""" + problems = [] + try: + with open(filepath, "r") as f: + for line in f: + line = line.strip() + if not line or line.startswith("#"): + continue + parts = line.split() + if len(parts) < 2: + continue + dims = parts[1].split("x") + if len(dims) == 3: + m, n, k = int(dims[0]), int(dims[1]), int(dims[2]) + problems.append((m, n, k)) + + if not problems: + raise ValueError(f"No valid problems found in benchmark file: {filepath}") + + _, n, k = problems[0] + num_groups = len(problems) + m_values = tuple(m for m, _, _ in problems) + + print(f"Loaded {num_groups} problems from benchmark file") + print(f"Using N={n}, K={k}, L={num_groups}") + print(f"M values per group: {m_values}") + + return ((n, k, num_groups), m_values) + + except FileNotFoundError as exc: + raise argparse.ArgumentTypeError(f"Benchmark file not found: {filepath}") from exc + except (OSError, ValueError) as exc: + raise argparse.ArgumentTypeError(f"Error reading benchmark file: {exc}") from exc + + +def parse_benchmark_arg( + arg: str, +) -> Tuple[Tuple[int, int, int], Tuple[int, ...]]: + """Parse benchmark argument string.""" + match_list = re.match(r"\[([\d,\s]+)\]\s*x\s*(\d+)\s*x\s*(\d+)", arg) + if match_list: + m_str = match_list.group(1) + n = int(match_list.group(2)) + k = int(match_list.group(3)) + try: + m_values = tuple(int(x.strip()) for x in m_str.split(",")) + num_groups = len(m_values) + return ((n, k, num_groups), m_values) + except ValueError as exc: + raise argparse.ArgumentTypeError( + f"Invalid integer list in benchmark argument: {arg}" + ) from exc + + parts = arg.split("x") + if len(parts) == 4: + try: + m, n, k, num_groups = [int(x.strip()) for x in parts] + m_values = tuple([m] * num_groups) + return ((n, k, num_groups), m_values) + except ValueError: + pass + + raise argparse.ArgumentTypeError(f"Invalid benchmark argument format. Got: {arg}") + + +def main(): + """Main entry point for running the Rubin blockscaled SwiGLU fusion kernel.""" + parser = argparse.ArgumentParser( + description=("Rubin BlockScaled Contiguous Gather Grouped GEMM with SwiGLU Fusion.") + ) + + parser.add_argument("--nkl", type=parse_comma_separated_ints, default=(256, 512, 1)) + parser.add_argument("--fixed_m", type=int, default=None) + parser.add_argument("--custom_mask", type=parse_comma_separated_ints, default=None) + parser.add_argument("--benchmark", type=str, default=None) + parser.add_argument("--permuted_m", type=int, default=None) + parser.add_argument( + "--mma_inst_shape", type=parse_comma_separated_ints, default=(128, 128, 128) + ) + parser.add_argument("--mma_tiler", type=parse_comma_separated_ints, default=(128, 128, 256)) + parser.add_argument("--cluster_shape_mn", type=parse_comma_separated_ints, default=(1, 1)) + parser.add_argument("--a_dtype", type=cutlass.dtype, default=cutlass.Float4E2M1FN) + parser.add_argument("--b_dtype", type=cutlass.dtype, default=cutlass.Float4E2M1FN) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.BFloat16) + parser.add_argument("--sf_dtype", type=cutlass.dtype, default=cutlass.Float8E4M3FN) + parser.add_argument("--sf_vec_size", type=int, default=16) + parser.add_argument("--a_major", choices=["k"], type=str, default="k") + parser.add_argument("--b_major", choices=["k"], type=str, default="k") + parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n") + parser.add_argument("--tolerance", type=float, default=1e-01) + parser.add_argument("--warmup_iterations", type=int, default=0) + parser.add_argument("--iterations", type=int, default=1) + parser.add_argument("--skip_ref_check", action="store_true") + parser.add_argument("--use_cold_l2", action="store_true", default=False) + parser.add_argument("--raster_along_m", action="store_true", default=False) + parser.add_argument("--use_cupti", action="store_true", default=False) + parser.add_argument( + "--a_path", + choices=["cpasync", "tma"], + default="cpasync", + help=( + "A load path: 'cpasync' = per-thread cp.async.cg.16B; " + "'tma' = TMA gather4. SFA path is always cpasync.128." + ), + ) + parser.add_argument( + "--use_pdl", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "Enable/disable PDL (Programmatic Dependent Launch). " + "Default: on. Use --no-use_pdl to disable." + ), + ) + args = parser.parse_args() + + if args.benchmark: + if os.path.isfile(args.benchmark): + nkl, group_m_list = read_benchmark_file(args.benchmark) + else: + nkl, group_m_list = parse_benchmark_arg(args.benchmark) + else: + if len(args.nkl) != 3: + parser.error("--nkl must contain exactly 3 values") + n, k, num_groups = args.nkl + nkl = (n, k, num_groups) + + if args.custom_mask is not None: + group_m_list = args.custom_mask + if len(group_m_list) != num_groups: + parser.error(f"--custom_mask must have exactly {num_groups} values") + elif args.fixed_m is not None: + group_m_list = tuple([args.fixed_m] * num_groups) + else: + group_m_list = tuple([128] * num_groups) + + if len(args.mma_inst_shape) != 3: + parser.error("--mma_inst_shape must contain exactly 3 values") + if len(args.mma_tiler) != 3: + parser.error("--mma_tiler must contain exactly 3 values") + if len(args.cluster_shape_mn) != 2: + parser.error("--cluster_shape_mn must contain exactly 2 values") + + exec_time = run( + nkl=nkl, + group_m_list=group_m_list, + a_dtype=args.a_dtype, + b_dtype=args.b_dtype, + c_dtype=args.c_dtype, + sf_dtype=args.sf_dtype, + sf_vec_size=args.sf_vec_size, + a_major=args.a_major, + b_major=args.b_major, + c_major=args.c_major, + mma_inst_shape=args.mma_inst_shape, + mma_tiler=args.mma_tiler, + cluster_shape_mn=args.cluster_shape_mn, + tolerance=args.tolerance, + warmup_iterations=args.warmup_iterations, + iterations=args.iterations, + skip_ref_check=args.skip_ref_check, + use_cold_l2=args.use_cold_l2, + permuted_m=args.permuted_m, + raster_along_m=args.raster_along_m, + use_cupti=args.use_cupti, + a_path=args.a_path, + use_pdl=args.use_pdl, + ) + print(f"Execution time: {exec_time:.2f} us") + print("PASS") + + +if __name__ == "__main__": + main() diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_finalize_fusion.py b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_finalize_fusion.py new file mode 100644 index 000000000000..7cf15e17f995 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/rubin/moe/rubin_contiguous_grouped_blockscaled_gemm_finalize_fusion.py @@ -0,0 +1,2826 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# 1. Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# 2. Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# 3. Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +Rubin (SM107) Contiguous Grouped Blockscaled GEMM Kernel with Finalize Fusion + +This module implements a contiguous grouped GEMM kernel for Rubin architecture +with fused MoE finalize (scatter-add) operation. + +Key features: +- Rubin-specific MMA features (B-reuse pattern, CollectorOp, etc.) +- Tile scheduling logic for contiguous grouped GEMM +- Fused finalize operation with atomic add for MoE scatter-add + +The finalize fusion performs: +1. GEMM: C_permuted = alpha * A * SFA * B * SFB +2. Scatter-add: c[token_idx] += token_scale * C_permuted[permuted_row] + +Example usage: + python rubin_contiguous_grouped_blockscaled_gemm_finalize_fusion.py \\ + --ab_dtype Float4E2M1FN --c_dtype BFloat16 \\ + --sf_dtype Float8E4M3FN --sf_vec_size 16 \\ + --mma_inst_shape 256,256,128 --mma_tiler 256,256,256 \\ + --cluster_shape_mn 2,1 --seq_len 4096 \\ + --benchmark 128x7168x2048x8 --iterations 1 +""" + +import argparse +import os +import re +from typing import List, NamedTuple, Optional, Tuple, Type + +import cuda.bindings.driver as cuda +import cutlass +import cutlass.cute as cute +import cutlass.pipeline as pipeline +import cutlass.torch as cutlass_torch +import cutlass.utils as utils +import cutlass.utils.blackwell_helpers as sm100_utils +import cutlass.utils.blockscaled_layout as blockscaled_utils +import cutlass.utils.rubin_helpers as sm107_utils +import torch +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.nvgpu.tcgen05.mma import CollectorOp +from cutlass.cute.runtime import from_dlpack +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait +from cutlass.utils.gemm.sm100 import ( + epilogue_tmem_copy_and_partition, + transform_partitioned_tensor_layout, +) + +# ============================================================================ +# Inline utility functions +# ============================================================================ + + +@dsl_user_op +def blk_reduce_bf16(dst_gemm, src_smem, size, loc=None, ip=None): + """Block reduce for BF16 using cp.reduce.async.bulk.""" + llvm.inline_asm( + None, + [ + dst_gemm.iterator.llvm_ptr, + src_smem.iterator.llvm_ptr, + size.ir_value(), + ], + "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.noftz.bf16 [$0], [$1], $2;", + "l,l,r", + has_side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def blk_reduce_fp32(dst_gemm, src_smem, size, loc=None, ip=None): + """Block reduce for FP32 using cp.reduce.async.bulk.""" + llvm.inline_asm( + None, + [ + dst_gemm.iterator.llvm_ptr, + src_smem.iterator.llvm_ptr, + size.ir_value(), + ], + "cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32 [$0], [$1], $2;", + "l,l,r", + has_side_effects=True, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def blk_reduce_fp16(dst_gemm, src_smem, size, loc=None, ip=None): + """Block reduce for FP16 using cp.reduce.async.bulk.""" + llvm.inline_asm( + None, + [ + dst_gemm.iterator.llvm_ptr, + src_smem.iterator.llvm_ptr, + size.ir_value(), + ], + "cp.reduce.async.bulk.global.shared::cta.bulk_group.noftz.f16 [$0], [$1], $2;", + "l,l,r", + has_side_effects=True, + loc=loc, + ip=ip, + ) + + +class S2TCopyBundle(NamedTuple): + """Bundle of tiled copy and partitioned tensors for smem-to-tmem copies.""" + + tiled_copy: cute.TiledCopy + sSF_compact: cute.Tensor # Partitioned source (smem) + tSF_compact: cute.Tensor # Partitioned destination (tmem) + + +class Sm107BlockScaledContiguousGroupedGemmFinalizeFusionKernel: + """Rubin (SM107) Contiguous Grouped Blockscaled GEMM Kernel with Finalize Fusion. + + This kernel implements batched matrix multiplication (c = scatter_add(alpha * A x SFA x B x SFB * token_scale)) + with contiguous grouped GEMM support and fused MoE finalize for Rubin GPUs. + + Key features: + - Persistent tile scheduling with dedicated scheduler warp + - Warp specialization (scheduler, TMA, MMA, epilogue warps) + - Support for B-reuse pattern (Bkeep-Breuse) + - Per-group alpha scaling + - Fused finalize with atomic add scatter + + :param sf_vec_size: Scale factor vector size (16 or 32) + :param mma_inst_shape: Shape of MMA instruction (M, N, K) + :param mma_tiler: Shape of MMA tiler (M, N, K) + :param cluster_shape_mn: Cluster dimensions (M, N) + :param raster_along_m: If True, raster tiles along M dimension first + """ + + def __init__( + self, + sf_vec_size: int, + mma_inst_shape: Tuple[int, int, int], + mma_tiler: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + raster_along_m: bool = False, + topK: int = 1, + ): + self.sf_vec_size = sf_vec_size + self.acc_dtype = cutlass.Float32 + self.mma_inst_shape = mma_inst_shape + self.mma_tiler = mma_tiler + self.cluster_shape_mn = cluster_shape_mn + self.raster_along_m = raster_along_m + self.topK = topK + + self.use_2cta_instrs = mma_inst_shape[0] == 256 + self.cta_group = tcgen05.CtaGroup.TWO if self.use_2cta_instrs else tcgen05.CtaGroup.ONE + self.arch = "sm_107" + self.smem_capacity = utils.get_smem_capacity_in_bytes(self.arch) + self.num_tmem_alloc_cols = cute.arch.get_max_tmem_alloc_cols(self.arch) + + self.occupancy = 1 + + # Warp IDs for warp specialization + self.epilog_warp_id = (0, 1, 2, 3) + self.mma_warp_id = 4 + self.tma_warp_id = 5 + self.sched_warp_id = 6 + + self.threads_per_warp = 32 + self.threads_per_cta = self.threads_per_warp * len( + ( + *self.epilog_warp_id, + self.mma_warp_id, + self.tma_warp_id, + self.sched_warp_id, + ) + ) + self.threads_wo_sched = self.threads_per_warp * len( + ( + *self.epilog_warp_id, + self.mma_warp_id, + self.tma_warp_id, + ) + ) + + # Set barriers for synchronization + self.epilog_sync_barrier = pipeline.NamedBarrier( + barrier_id=2, + num_threads=self.threads_per_warp * len(self.epilog_warp_id), + ) + self.tmem_alloc_barrier = pipeline.NamedBarrier( + barrier_id=3, + num_threads=self.threads_per_warp * len((self.mma_warp_id, *self.epilog_warp_id)), + ) + self.sched_sync_barrier = pipeline.NamedBarrier( + barrier_id=4, + num_threads=self.threads_per_warp, + ) + + # For epilogue compatibility + self.epilogue_warp_id = self.epilog_warp_id + self.epilog_sync_bar_id = self.epilog_sync_barrier.barrier_id + + # B-reuse pattern control + self.enable_breuse = True if mma_tiler[0] // mma_inst_shape[0] == 2 else False + + def _get_mma_permutation_mnk(self): + if cutlass.const_expr(self.use_2cta_instrs and self.enable_breuse): + m_layout = cute.make_layout( + shape=(self.mma_inst_shape[0] // 2, 2, 2), + stride=(1, self.mma_inst_shape[0], self.mma_inst_shape[0] // 2), + ) + return (m_layout, self.mma_inst_shape[1], self.mma_inst_shape[2]) + else: + return (1, 1, 1) + + @staticmethod + def _compute_stages( + tiled_mma: cute.TiledMma, + mma_tiler_mnk: Tuple[int, int, int], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + cta_tile: Tuple[int, int, int], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + smem_capacity: int, + occupancy: int, + with_breuse: bool, + ) -> Tuple[int, int, int, int]: + """Compute the number of stages for A/B/C/tile_info operands.""" + # ACC stages + # Note that here we have assumed the kernel have access to all TMEM capacity + # associated with sm_107 architecture. + num_acc_stage = 1 if (with_breuse and mma_tiler_mnk[1] in {192, 256}) else 2 + + # Default C stages and tile info stages (5 elements now: bidx, bidy, bidz, valid, mn_limit) + num_c_stage = 1 # Thinking about it + num_tile_stage = 2 + + # Calculate smem layout and size for one stage + a_smem_layout_stage_one = sm100_utils.make_smem_layout_a( + tiled_mma, mma_tiler_mnk, a_dtype, 1 + ) + b_smem_layout_staged_one = sm100_utils.make_smem_layout_b( + tiled_mma, mma_tiler_mnk, b_dtype, 1 + ) + sfa_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, mma_tiler_mnk, sf_vec_size, 1 + ) + sfb_smem_layout_staged_one = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, mma_tiler_mnk, sf_vec_size, 1 + ) + + # Shared memory for epilogue block reduce (if enabled) + swizzled_pad = 16 // (c_dtype.width // 8) + c_smem_layout_staged_one = cute.make_layout( + (cta_tile[0], cta_tile[1]), stride=(cta_tile[1] + swizzled_pad, 1) + ) + + ab_bytes_per_stage = ( + cute.size_in_bytes(a_dtype, a_smem_layout_stage_one) + + cute.size_in_bytes(b_dtype, b_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfa_smem_layout_staged_one) + + cute.size_in_bytes(sf_dtype, sfb_smem_layout_staged_one) + ) + mbar_helpers_bytes = 1024 + c_bytes_per_stage = cute.size_in_bytes(c_dtype, c_smem_layout_staged_one) + c_bytes = c_bytes_per_stage * num_c_stage + + # Calculate A/B/SFA/SFB stages + num_ab_stage = ( + smem_capacity // occupancy - (mbar_helpers_bytes + c_bytes) + ) // ab_bytes_per_stage + + return num_acc_stage, num_ab_stage, num_c_stage, num_tile_stage + + def _setup_attributes(self): + """Set up configurations dependent on GEMM inputs.""" + # Compute mma instruction shapes + self.mma_inst_shape_sfb = ( + self.mma_inst_shape[0] // (2 if self.use_2cta_instrs else 1), + cute.round_up(self.mma_inst_shape[1], 128), + self.mma_inst_shape[2], + ) + + tiled_mma = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + atom_layout_mnk=(1, 1, 1), + permutation_mnk=self._get_mma_permutation_mnk(), + ) + + tiled_mma_sfb = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_sfb, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + ) + + # Compute mma/cluster/tile shapes + self.mma_tiler_sfb = ( + self.mma_inst_shape_sfb[0], + self.mma_inst_shape_sfb[1], + self.mma_tiler[2], + ) + + self.cta_tile_shape_mnk = ( + self.mma_tiler[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler[1], + self.mma_tiler[2], + ) + self.cta_tile_shape_mnk_sfb = ( + self.mma_tiler_sfb[0] // cute.size(tiled_mma.thr_id.shape), + self.mma_tiler_sfb[1], + self.mma_tiler_sfb[2], + ) + + # Compute cluster layout + self.cluster_layout_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma.thr_id.shape,), + ) + self.cluster_layout_sfb_vmnk = cute.tiled_divide( + cute.make_layout((*self.cluster_shape_mn, 1)), + (tiled_mma_sfb.thr_id.shape,), + ) + + # Compute number of multicast CTAs + self.num_mcast_ctas_a = cute.size(self.cluster_layout_vmnk.shape[2]) + self.num_mcast_ctas_b = cute.size(self.cluster_layout_vmnk.shape[1]) + self.num_mcast_ctas_sfb = cute.size(self.cluster_layout_sfb_vmnk.shape[1]) + self.is_a_mcast = self.num_mcast_ctas_a > 1 + self.is_b_mcast = self.num_mcast_ctas_b > 1 + self.is_sfb_mcast = self.num_mcast_ctas_sfb > 1 + + # Compute epilogue subtile + self.epi_tile = sm107_utils.compute_epilogue_tile_shape( + tiled_mma.op, + self.cta_tile_shape_mnk, + self.use_2cta_instrs, + self.c_layout, + self.c_dtype, + ) + self.epi_tile_n = cute.size(self.epi_tile[1]) + + # Setup stage counts + ( + self.num_acc_stage, + self.num_ab_stage, + self.num_c_stage, + self.num_tile_stage, + ) = self._compute_stages( + tiled_mma, + self.mma_tiler, + self.a_dtype, + self.b_dtype, + self.c_dtype, + self.cta_tile_shape_mnk, + self.sf_dtype, + self.sf_vec_size, + self.smem_capacity, + self.occupancy, + self.enable_breuse, + ) + + # Compute A/B/SFA/SFB/C shared memory layout + self.a_smem_layout_staged = sm100_utils.make_smem_layout_a( + tiled_mma, self.mma_tiler, self.a_dtype, self.num_ab_stage + ) + self.b_smem_layout_staged = sm100_utils.make_smem_layout_b( + tiled_mma, self.mma_tiler, self.b_dtype, self.num_ab_stage + ) + self.sfa_smem_layout_staged = blockscaled_utils.make_smem_layout_sfa( + tiled_mma, self.mma_tiler, self.sf_vec_size, self.num_ab_stage + ) + self.sfb_smem_layout_staged = blockscaled_utils.make_smem_layout_sfb( + tiled_mma, self.mma_tiler, self.sf_vec_size, self.num_ab_stage + ) + + # C smem layout for block reduce (if enabled) + swizzled_pad = 16 // (self.c_dtype.width // 8) + self.c_smem_layout_staged = cute.make_layout( + (self.cta_tile_shape_mnk[0], self.cta_tile_shape_mnk[1], self.num_c_stage), + stride=( + self.cta_tile_shape_mnk[1] + swizzled_pad, + 1, + self.cta_tile_shape_mnk[0] * (self.cta_tile_shape_mnk[1] + 8), + ), + ) + + # Compute TMEM layouts for SFA/SFB + self.tCtSFA_layout = blockscaled_utils.make_tmem_layout_sfa( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)), + ) + self.tCtSFB_layout = blockscaled_utils.make_tmem_layout_sfb( + tiled_mma, + self.mma_tiler, + self.sf_vec_size, + cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)), + ) + + # Compute TMEM column counts + # Each column entry in TMEM is 32-bit wide, and so we recast the TMEM layout + # from its original data type to a 32-bit wide data type. Moreover, TMEM + # addresses are expressed as (row << 16) | col, which in CUTE are expressed + # as an affine transformation row * (1<<16) + col, which can be seen as a CUTE + # layout of (row, col):(1<<16, 1). As a result, by masking out the upper 16 bits + # (keeping only the lower 16 bits), we extract the cosize corresponding + # to only the columns. + self.num_sfa_tmem_cols = ( + cute.cosize(cute.recast_layout(32, self.sf_dtype.width, self.tCtSFA_layout)) + & 0x0000FFFF + ) + self.num_sfb_tmem_cols = ( + cute.cosize(cute.recast_layout(32, self.sf_dtype.width, self.tCtSFB_layout)) + & 0x0000FFFF + ) + self.num_sf_tmem_cols = self.num_sfa_tmem_cols + self.num_sfb_tmem_cols + self.num_accumulator_tmem_cols = ( + self.cta_tile_shape_mnk[1] * self.num_acc_stage * (2 if self.enable_breuse else 1) + ) + + # Epilogue vectorization config + if cutlass.const_expr(self.c_dtype == cutlass.BFloat16): + self.element_offset = 8 + self.epi_loop_size = self.epi_tile_n // 8 + elif cutlass.const_expr(self.c_dtype == cutlass.Float32): + self.element_offset = 2 + self.epi_loop_size = self.epi_tile_n // 2 + else: + self.element_offset = 1 + self.epi_loop_size = self.epi_tile_n + + # copy_size is in bytes for cp.reduce.async.bulk instruction + self.copy_size = self.cta_tile_shape_mnk[1] * (self.c_dtype.width // 8) + + def _is_interleaved_utccp(self) -> bool: + """Enable interleaving UTCCP for Bkeep-Breuse case for 4xFP4 kernel.""" + return self.a_dtype.width == 4 and self.b_dtype.width == 4 and self.enable_breuse + + def _mainloop_s2t_copy_and_partition( + self, + sSF: cute.Tensor, + tSF: cute.Tensor, + ) -> S2TCopyBundle: + """Make tiledCopy for smem to tmem load for scale factor tensor.""" + tCsSF_compact = cute.filter_zeros(sSF) + tCtSF_compact = cute.filter_zeros(tSF) + + copy_atom_s2t = cute.make_copy_atom( + tcgen05.Cp4x32x128bOp(self.cta_group), + self.sf_dtype, + ) + tiled_copy_s2t = tcgen05.make_s2t_copy(copy_atom_s2t, tCtSF_compact) + thr_copy_s2t = tiled_copy_s2t.get_slice(0) + + def appendMNBroadcastMode(smem_layout: cute.Layout): + mn_dim = cute.get(smem_layout, mode=[0, 0]) + mn_dim = cute.append(mn_dim, cute.make_layout((4), stride=(0))) + layout = cute.append(cute.group_modes(mn_dim, 0), cute.get(smem_layout, mode=[0, 1])) + layout = cute.append(cute.group_modes(layout, 0), cute.get(smem_layout, mode=[1])) + layout = cute.append(layout, cute.get(smem_layout, mode=[2])) + layout = cute.append(layout, cute.get(smem_layout, mode=[3])) + return layout + + tCsSF_compact_bcast = cute.make_tensor( + tCsSF_compact.iterator, appendMNBroadcastMode(tCsSF_compact.layout) + ) + + tCsSF_compact_s2t_ = thr_copy_s2t.partition_S(tCsSF_compact_bcast) + tCsSF_compact_s2t = tcgen05.get_s2t_smem_desc_tensor(tiled_copy_s2t, tCsSF_compact_s2t_) + tCtSF_compact_s2t = thr_copy_s2t.partition_D(tCtSF_compact) + + return S2TCopyBundle(tiled_copy_s2t, tCsSF_compact_s2t, tCtSF_compact_s2t) + + def _mainloop_s2t_copies( + self, + stage_idx: int, + sfa_s2t_bundle: S2TCopyBundle, + sfb_s2t_bundle: S2TCopyBundle, + ): + """Copy SFA/SFB from smem to tmem.""" + s2t_stage_coord = (None, None, None, None, stage_idx) + + cute.copy( + sfa_s2t_bundle.tiled_copy, + sfa_s2t_bundle.sSF_compact[s2t_stage_coord], + sfa_s2t_bundle.tSF_compact, + ) + cute.copy( + sfb_s2t_bundle.tiled_copy, + sfb_s2t_bundle.sSF_compact[s2t_stage_coord], + sfb_s2t_bundle.tSF_compact, + ) + + def _mainloop_s2t_interleaved_copies( + self, + k_block: int, + stage_idx: int, + sfa_s2t_bundle: S2TCopyBundle, + sfb_s2t_bundle: S2TCopyBundle, + ): + """Interleaved UTCCP for Bkeep-Breuse pattern.""" + s_sfa_crd_keep = (None, 0, None, k_block, stage_idx) + s_sfa_crd_reuse = (None, 1, None, k_block, stage_idx) + s_sfb_crd = (None, None, None, k_block, stage_idx) + + t_sfa_crd_keep = (None, 0, None, k_block) + t_sfa_crd_reuse = (None, 1, None, k_block) + t_sfb_crd = (None, None, None, k_block) + + cute.copy( + sfa_s2t_bundle.tiled_copy, + sfa_s2t_bundle.sSF_compact[s_sfa_crd_keep], + sfa_s2t_bundle.tSF_compact[t_sfa_crd_keep], + ) + cute.copy( + sfb_s2t_bundle.tiled_copy, + sfb_s2t_bundle.sSF_compact[s_sfb_crd], + sfb_s2t_bundle.tSF_compact[t_sfb_crd], + ) + cute.copy( + sfa_s2t_bundle.tiled_copy, + sfa_s2t_bundle.sSF_compact[s_sfa_crd_reuse], + sfa_s2t_bundle.tSF_compact[t_sfa_crd_reuse], + ) + + @cute.jit + def __call__( + self, + a: cute.Tensor, + b: cute.Tensor, + c: cute.Tensor, + sfa: cute.Tensor, + sfb: cute.Tensor, + tile_idx_to_group_idx: cute.Tensor, + num_non_exiting_tiles: cute.Tensor, + tile_idx_to_mn_limit: cute.Tensor, + alpha: cute.Tensor, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + permuted_idx_to_expanded_idx: cute.Tensor, + token_final_scales: cute.Tensor, + epilogue_op: cutlass.Constexpr = lambda x: x, + ): + """Execute the contiguous grouped GEMM with finalize fusion. + + :param a: Input tensor A (permuted_m, k, 1) + :param b: Input tensor B (n, k, l) + :param c: Output tensor (seq_len, n, 1) + :param sfa: Scale factor tensor A + :param sfb: Scale factor tensor B + :param tile_idx_to_group_idx: Mapping from tile index to group ID + :param num_non_exiting_tiles: Number of valid tiles + :param tile_idx_to_mn_limit: M limit for each tile + :param alpha: Alpha tensor for each group + :param max_active_clusters: Maximum number of active clusters + :param stream: CUDA stream + :param permuted_idx_to_expanded_idx: Mapping from permuted row to expanded index + :param token_final_scales: Final scales for each token (seq_len, topK) + :param epilogue_op: Optional epilogue operation + """ + # Setup static attributes + self.a_dtype: Type[cutlass.Numeric] = a.element_type + self.b_dtype: Type[cutlass.Numeric] = b.element_type + self.c_dtype: Type[cutlass.Numeric] = c.element_type + self.sf_dtype: Type[cutlass.Numeric] = sfa.element_type + self.final_scale_dtype = cutlass.Float32 + self.a_major_mode = utils.LayoutEnum.from_tensor(a).mma_major_mode() + self.b_major_mode = utils.LayoutEnum.from_tensor(b).mma_major_mode() + self.c_layout = utils.LayoutEnum.ROW_MAJOR # Always N-major for GEMM output + + # Check data types + if cutlass.const_expr(self.a_dtype != self.b_dtype): + raise TypeError(f"Type must match: {self.a_dtype} != {self.b_dtype}") + + # Setup attributes + self._setup_attributes() + + # Setup sfa/sfb tensor + sfa_layout = blockscaled_utils.tile_atom_to_shape_SF(a.shape, self.sf_vec_size) + sfa = cute.make_tensor(sfa.iterator, sfa_layout) + + sfb_layout = blockscaled_utils.tile_atom_to_shape_SF(b.shape, self.sf_vec_size) + sfb = cute.make_tensor(sfb.iterator, sfb_layout) + + atom_layout_mnk = (1, 1, 1) + permutation_mnk = self._get_mma_permutation_mnk() + + tiled_mma = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + atom_layout_mnk=atom_layout_mnk, + permutation_mnk=permutation_mnk, + ) + + tiled_mma.set(tcgen05.Field.NEGATE_A, False) + tiled_mma.set(tcgen05.Field.NEGATE_B, False) + + tiled_mma_sfb = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + cute.nvgpu.tcgen05.CtaGroup.ONE, + self.mma_inst_shape_sfb, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.DISCARD, + ) + tiled_mma_sfb.set(tcgen05.Field.NEGATE_A, False) + tiled_mma_sfb.set(tcgen05.Field.NEGATE_B, False) + + tiled_mma_bkeep = None + tiled_mma_breuse = None + if cutlass.const_expr(self.enable_breuse): + tiled_mma_bkeep = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.FILL, + atom_layout_mnk=atom_layout_mnk, + permutation_mnk=permutation_mnk, + ) + tiled_mma_bkeep.set(tcgen05.Field.NEGATE_A, False) + tiled_mma_bkeep.set(tcgen05.Field.NEGATE_B, False) + + tiled_mma_breuse = sm107_utils.make_blockscaled_trivial_tiled_mma( + self.a_dtype, + self.b_dtype, + self.a_major_mode, + self.b_major_mode, + self.sf_dtype, + self.sf_vec_size, + self.cta_group, + self.mma_inst_shape, + a_collector_op=CollectorOp.DISCARD, + b_collector_op=CollectorOp.LASTUSE, + atom_layout_mnk=atom_layout_mnk, + permutation_mnk=permutation_mnk, + ) + tiled_mma_breuse.set(tcgen05.Field.NEGATE_A, False) + tiled_mma_breuse.set(tcgen05.Field.NEGATE_B, False) + + atom_thr_size = cute.size(tiled_mma.thr_id.shape) + + # Setup TMA load for A + a_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + a_smem_layout = cute.slice_(self.a_smem_layout_staged, (None, None, None, 0)) + tma_atom_a, tma_tensor_a = cute.nvgpu.make_tiled_tma_atom_A( + a_op, + a, + a_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # Setup TMA load for B + b_op = sm100_utils.cluster_shape_to_tma_atom_B(self.cluster_shape_mn, tiled_mma.thr_id) + b_smem_layout = cute.slice_(self.b_smem_layout_staged, (None, None, None, 0)) + tma_atom_b, tma_tensor_b = cute.nvgpu.make_tiled_tma_atom_B( + b_op, + b, + b_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + ) + + # Setup TMA load for SFA + sfa_op = sm100_utils.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, tiled_mma.thr_id) + sfa_smem_layout = cute.slice_(self.sfa_smem_layout_staged, (None, None, None, 0)) + tma_atom_sfa, tma_tensor_sfa = cute.nvgpu.make_tiled_tma_atom_A( + sfa_op, + sfa, + sfa_smem_layout, + self.mma_tiler, + tiled_mma, + self.cluster_layout_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # Setup TMA load for SFB + sfb_op = sm100_utils.cluster_shape_to_tma_atom_SFB(self.cluster_shape_mn, tiled_mma.thr_id) + sfb_smem_layout = cute.slice_(self.sfb_smem_layout_staged, (None, None, None, 0)) + tma_atom_sfb, tma_tensor_sfb = cute.nvgpu.make_tiled_tma_atom_B( + sfb_op, + sfb, + sfb_smem_layout, + self.mma_tiler_sfb, + tiled_mma_sfb, + self.cluster_layout_sfb_vmnk.shape, + internal_type=cutlass.Int16, + ) + + # Handle cta_tile_shape_n=192 case + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 192): + x = tma_tensor_sfb.stride[0][1] + y = cute.ceil_div(tma_tensor_sfb.shape[0][1], 4) + + new_shape = ( + (tma_tensor_sfb.shape[0][0], ((2, 2), y)), + tma_tensor_sfb.shape[1], + tma_tensor_sfb.shape[2], + ) + x_times_3 = 3 * x + new_stride = ( + (tma_tensor_sfb.stride[0][0], ((x, x), x_times_3)), + tma_tensor_sfb.stride[1], + tma_tensor_sfb.stride[2], + ) + tma_tensor_sfb_new_layout = cute.make_layout(new_shape, stride=new_stride) + tma_tensor_sfb = cute.make_tensor(tma_tensor_sfb.iterator, tma_tensor_sfb_new_layout) + + a_copy_size = cute.size_in_bytes(self.a_dtype, a_smem_layout) + b_copy_size = cute.size_in_bytes(self.b_dtype, b_smem_layout) + sfa_copy_size = cute.size_in_bytes(self.sf_dtype, sfa_smem_layout) + sfb_copy_size = cute.size_in_bytes(self.sf_dtype, sfb_smem_layout) + self.num_tma_load_bytes = ( + a_copy_size + b_copy_size + sfa_copy_size + sfb_copy_size + ) * atom_thr_size + + # Compute grid size based on GEMM shape + gemm_m = a.shape[0] + gemm_n = b.shape[0] + gemm_l = a.shape[2] + gemm_shape = (gemm_m, gemm_n, gemm_l) + + self.tile_sched_params, grid = self._compute_grid( + gemm_shape, + self.cta_tile_shape_mnk, + self.cluster_shape_mn, + max_active_clusters, + self.raster_along_m, + ) + + self.buffer_align_bytes = 1024 + + @cute.struct + class SharedStorage: + # (bidx, bidy, expert_idx, valid, mn_limit) + sInfo: cute.struct.Align[ + cute.struct.MemRange[cutlass.Int32, 5 * self.num_tile_stage], + 1, + ] + ab_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_ab_stage * 2] + acc_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_acc_stage * 2] + tile_info_mbar_ptr: cute.struct.MemRange[cutlass.Int64, self.num_tile_stage * 2] + tmem_dealloc_mbar_ptr: cutlass.Int64 + tmem_holding_buf: cutlass.Int32 + sC: cute.struct.Align[ + cute.struct.MemRange[ + self.c_dtype, + cute.cosize(self.c_smem_layout_staged), + ], + self.buffer_align_bytes, + ] + sA: cute.struct.Align[ + cute.struct.MemRange[self.a_dtype, cute.cosize(self.a_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sB: cute.struct.Align[ + cute.struct.MemRange[self.b_dtype, cute.cosize(self.b_smem_layout_staged.outer)], + self.buffer_align_bytes, + ] + sSFA: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfa_smem_layout_staged)], + self.buffer_align_bytes, + ] + sSFB: cute.struct.Align[ + cute.struct.MemRange[self.sf_dtype, cute.cosize(self.sfb_smem_layout_staged)], + self.buffer_align_bytes, + ] + + self.shared_storage = SharedStorage + + # Launch the kernel + self.kernel( + tiled_mma, + tiled_mma_bkeep, + tiled_mma_breuse, + tiled_mma_sfb, + tma_atom_a, + tma_tensor_a, + tma_atom_b, + tma_tensor_b, + tma_atom_sfa, + tma_tensor_sfa, + tma_atom_sfb, + tma_tensor_sfb, + c, + tile_idx_to_group_idx, + num_non_exiting_tiles, + tile_idx_to_mn_limit, + alpha, + permuted_idx_to_expanded_idx, + token_final_scales, + self.cluster_layout_vmnk, + self.cluster_layout_sfb_vmnk, + self.a_smem_layout_staged, + self.b_smem_layout_staged, + self.sfa_smem_layout_staged, + self.sfb_smem_layout_staged, + self.tCtSFA_layout, + self.tCtSFB_layout, + self.c_smem_layout_staged, + self.epi_tile, + self.tile_sched_params, + epilogue_op, + ).launch( + grid=grid, + block=[self.threads_per_cta, 1, 1], + cluster=(*self.cluster_shape_mn, 1), + smem=self.shared_storage.size_in_bytes(), + stream=stream, + min_blocks_per_mp=1, + ) + return + + @cute.kernel + def kernel( + self, + tiled_mma: cute.TiledMma, + tiled_mma_bkeep: Optional[cute.TiledMma], + tiled_mma_breuse: Optional[cute.TiledMma], + tiled_mma_sfb: cute.TiledMma, + tma_atom_a: cute.CopyAtom, + mA_mkl: cute.Tensor, + tma_atom_b: cute.CopyAtom, + mB_nkl: cute.Tensor, + tma_atom_sfa: cute.CopyAtom, + mSFA_mkl: cute.Tensor, + tma_atom_sfb: cute.CopyAtom, + mSFB_nkl: cute.Tensor, + c: cute.Tensor, + tile_idx_to_group_idx: cute.Tensor, + num_non_exiting_tiles: cute.Tensor, + tile_idx_to_mn_limit: cute.Tensor, + alpha: cute.Tensor, + permuted_idx_to_expanded_idx: cute.Tensor, + token_final_scales: cute.Tensor, + cluster_layout_vmnk: cute.Layout, + cluster_layout_sfb_vmnk: cute.Layout, + a_smem_layout_staged: cute.ComposedLayout, + b_smem_layout_staged: cute.ComposedLayout, + sfa_smem_layout_staged: cute.Layout, + sfb_smem_layout_staged: cute.Layout, + tCtSFA_layout: cute.Layout, + tCtSFB_layout: cute.Layout, + c_smem_layout_staged: cute.Layout, + epi_tile: cute.Tile, + tile_sched_params: utils.PersistentTileSchedulerParams, + epilogue_op: cutlass.Constexpr, + ): + """GPU device kernel for contiguous grouped GEMM with finalize fusion.""" + warp_idx = cute.arch.warp_idx() + warp_idx = cute.arch.make_warp_uniform(warp_idx) + + # Prefetch TMA descriptors + if warp_idx == self.tma_warp_id: + cpasync.prefetch_descriptor(tma_atom_a) + cpasync.prefetch_descriptor(tma_atom_b) + cpasync.prefetch_descriptor(tma_atom_sfa) + cpasync.prefetch_descriptor(tma_atom_sfb) + + use_2cta_instrs = cute.size(tiled_mma.thr_id.shape) == 2 + + # Setup coordinates + bidx, bidy, bidz = cute.arch.block_idx() + mma_tile_coord_v = bidx % cute.size(tiled_mma.thr_id.shape) + is_leader_cta = mma_tile_coord_v == 0 + cta_rank_in_cluster = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + block_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(cta_rank_in_cluster) + block_in_cluster_coord_sfb_vmnk = cluster_layout_sfb_vmnk.get_flat_coord( + cta_rank_in_cluster + ) + tidx, _, _ = cute.arch.thread_idx() + + # Allocate shared memory + smem = utils.SmemAllocator() + storage = smem.allocate(self.shared_storage) + + # Initialize pipelines + ab_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_tma_producer = self.num_mcast_ctas_a + self.num_mcast_ctas_b - 1 + ab_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_tma_producer + ) + ab_pipeline = pipeline.PipelineTmaUmma.create( + barrier_storage=storage.ab_mbar_ptr.data_ptr(), + num_stages=self.num_ab_stage, + producer_group=ab_pipeline_producer_group, + consumer_group=ab_pipeline_consumer_group, + tx_count=self.num_tma_load_bytes, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + acc_pipeline_producer_group = pipeline.CooperativeGroup(pipeline.Agent.Thread) + num_acc_consumer_threads = len(self.epilog_warp_id) * (2 if use_2cta_instrs else 1) + acc_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, num_acc_consumer_threads + ) + acc_pipeline = pipeline.PipelineUmmaAsync.create( + barrier_storage=storage.acc_mbar_ptr.data_ptr(), + num_stages=self.num_acc_stage, + producer_group=acc_pipeline_producer_group, + consumer_group=acc_pipeline_consumer_group, + cta_layout_vmnk=cluster_layout_vmnk, + defer_sync=True, + ) + + # Initialize tile info pipeline + tile_info_pipeline_producer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_per_warp * 1, + ) + tile_info_pipeline_consumer_group = pipeline.CooperativeGroup( + pipeline.Agent.Thread, + self.threads_wo_sched, + ) + tile_info_pipeline = pipeline.PipelineAsync.create( + barrier_storage=storage.tile_info_mbar_ptr.data_ptr(), + num_stages=self.num_tile_stage, + producer_group=tile_info_pipeline_producer_group, + consumer_group=tile_info_pipeline_consumer_group, + ) + + # Initialize tensor memory allocator + tmem = utils.TmemAllocator( + storage.tmem_holding_buf.ptr, + barrier_for_retrieve=self.tmem_alloc_barrier, + allocator_warp_id=self.epilog_warp_id[0], + is_two_cta=use_2cta_instrs, + two_cta_tmem_dealloc_mbar_ptr=storage.tmem_dealloc_mbar_ptr.ptr, + arch=self.arch, + ) + + # Cluster arrive after barrier init + pipeline_init_arrive(cluster_shape_mn=self.cluster_shape_mn, is_relaxed=True) + + # Setup smem tensors + sC = storage.sC.get_tensor(c_smem_layout_staged) + sA = storage.sA.get_tensor(a_smem_layout_staged.outer, swizzle=a_smem_layout_staged.inner) + sB = storage.sB.get_tensor(b_smem_layout_staged.outer, swizzle=b_smem_layout_staged.inner) + sSFA = storage.sSFA.get_tensor(sfa_smem_layout_staged) + sSFB = storage.sSFB.get_tensor(sfb_smem_layout_staged) + info_layout = cute.make_layout((5, self.num_tile_stage), stride=(1, 5)) + sInfo = storage.sInfo.get_tensor(info_layout) + + # Compute multicast masks + a_full_mcast_mask = None + b_full_mcast_mask = None + sfa_full_mcast_mask = None + sfb_full_mcast_mask = None + if cutlass.const_expr(self.is_a_mcast or self.is_b_mcast or use_2cta_instrs): + a_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + b_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=1 + ) + sfa_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_vmnk, block_in_cluster_coord_vmnk, mcast_mode=2 + ) + sfb_full_mcast_mask = cpasync.create_tma_multicast_mask( + cluster_layout_sfb_vmnk, block_in_cluster_coord_sfb_vmnk, mcast_mode=1 + ) + + # Local_tile partition global tensors + gA_mkl = cute.local_tile( + mA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + gB_nkl = cute.local_tile( + mB_nkl, cute.slice_(self.mma_tiler, (0, None, None)), (None, None, None) + ) + gSFA_mkl = cute.local_tile( + mSFA_mkl, cute.slice_(self.mma_tiler, (None, 0, None)), (None, None, None) + ) + gSFB_nkl = cute.local_tile( + mSFB_nkl, + cute.slice_(self.mma_tiler_sfb, (0, None, None)), + (None, None, None), + ) + gC_mnl = cute.local_tile( + c, cute.slice_(self.mma_tiler, (None, None, 0)), (None, None, None) + ) + k_tile_cnt = cutlass.Int32(cute.size(gA_mkl, mode=[3])) + + # Partition global tensors for TiledMMA + thr_mma = tiled_mma.get_slice(mma_tile_coord_v) + thr_mma_sfb = tiled_mma_sfb.get_slice(mma_tile_coord_v) + tCgA = thr_mma.partition_A(gA_mkl) + tCgB = thr_mma.partition_B(gB_nkl) + tCgSFA = thr_mma.partition_A(gSFA_mkl) + tCgSFB = thr_mma_sfb.partition_B(gSFB_nkl) + tCgC = thr_mma.partition_C(gC_mnl) + + # Partition for TMA load + a_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, 0, None, 0)).shape) + tAsA, tAgA = cpasync.tma_partition( + tma_atom_a, + block_in_cluster_coord_vmnk[2], + a_cta_layout, + cute.group_modes(sA, 0, 3), + cute.group_modes(tCgA, 0, 3), + ) + b_cta_layout = cute.make_layout(cute.slice_(cluster_layout_vmnk, (0, None, 0, 0)).shape) + tBsB, tBgB = cpasync.tma_partition( + tma_atom_b, + block_in_cluster_coord_vmnk[1], + b_cta_layout, + cute.group_modes(sB, 0, 3), + cute.group_modes(tCgB, 0, 3), + ) + + sfa_cta_layout = a_cta_layout + tAsSFA, tAgSFA = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfa, + block_in_cluster_coord_vmnk[2], + sfa_cta_layout, + cute.group_modes(sSFA, 0, 3), + cute.group_modes(tCgSFA, 0, 3), + ) + tAsSFA = cute.filter_zeros(tAsSFA) + tAgSFA = cute.filter_zeros(tAgSFA) + + sfb_cta_layout = cute.make_layout( + cute.slice_(cluster_layout_sfb_vmnk, (0, None, 0, 0)).shape + ) + tBsSFB, tBgSFB = cute.nvgpu.cpasync.tma_partition( + tma_atom_sfb, + block_in_cluster_coord_sfb_vmnk[1], + sfb_cta_layout, + cute.group_modes(sSFB, 0, 3), + cute.group_modes(tCgSFB, 0, 3), + ) + tBsSFB = cute.filter_zeros(tBsSFB) + tBgSFB = cute.filter_zeros(tBgSFB) + + # Partition for MMA + tCrA = tiled_mma.make_fragment_A(sA) + tCrB = tiled_mma.make_fragment_B(sB) + acc_shape = tiled_mma.partition_shape_C(self.mma_tiler[:2]) + tCtAcc_fake = tiled_mma.make_fragment_C(cute.append(acc_shape, self.num_acc_stage)) + + # Cluster wait before tensor memory alloc + pipeline_init_wait(cluster_shape_mn=self.cluster_shape_mn) + + # + # Specialized Scheduler warp + # + if warp_idx == self.sched_warp_id: + # + # Persistent tile scheduling loop + # + tile_sched = utils.StaticPersistentTileScheduler.create( + tile_sched_params, cute.arch.block_idx(), cute.arch.grid_dim() + ) + work_tile = tile_sched.initial_work_tile_info() + + tile_info_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_tile_stage + ) + + num_valid_tiles = num_non_exiting_tiles[0] + + if cutlass.const_expr(self.raster_along_m): + while work_tile.is_valid_tile: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_m = cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape) + + expert_idx = tile_idx_to_group_idx[mma_tile_coord_m] + tile_idx = mma_tile_coord_m + + if tile_idx < num_valid_tiles: + mn_limit = tile_idx_to_mn_limit[mma_tile_coord_m] + + tile_info_pipeline.producer_acquire(tile_info_producer_state) + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = cur_tile_coord[0] + sInfo[(1, tile_info_producer_state.index)] = cur_tile_coord[1] + sInfo[(2, tile_info_producer_state.index)] = expert_idx + sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32( + work_tile.is_valid_tile + ) + sInfo[(4, tile_info_producer_state.index)] = mn_limit + cute.arch.fence_proxy("async.shared", space="cta") + + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + else: + is_continue = cutlass.Boolean(1) + while work_tile.is_valid_tile and is_continue: + cur_tile_coord = work_tile.tile_idx + mma_tile_coord_m = cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape) + + expert_idx = tile_idx_to_group_idx[mma_tile_coord_m] + tile_idx = mma_tile_coord_m + + if tile_idx < num_valid_tiles: + mn_limit = tile_idx_to_mn_limit[mma_tile_coord_m] + tile_info_pipeline.producer_acquire(tile_info_producer_state) + with cute.arch.elect_one(): + sInfo[(0, tile_info_producer_state.index)] = cur_tile_coord[0] + sInfo[(1, tile_info_producer_state.index)] = cur_tile_coord[1] + sInfo[(2, tile_info_producer_state.index)] = expert_idx + sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32( + work_tile.is_valid_tile + ) + sInfo[(4, tile_info_producer_state.index)] = mn_limit + cute.arch.fence_proxy("async.shared", space="cta") + + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + else: + is_continue = cutlass.Boolean(0) + + tile_sched.advance_to_next_work() + work_tile = tile_sched.get_current_work() + + # Signal end of work + tile_info_pipeline.producer_acquire(tile_info_producer_state) + with cute.arch.elect_one(): + sInfo[(3, tile_info_producer_state.index)] = cutlass.Int32(0) + cute.arch.fence_proxy("async.shared", space="cta") + self.sched_sync_barrier.arrive_and_wait() + tile_info_pipeline.producer_commit(tile_info_producer_state) + tile_info_producer_state.advance() + tile_info_pipeline.producer_tail(tile_info_producer_state) + + # + # Specialized TMA load warp + # + if warp_idx == self.tma_warp_id: + ab_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_ab_stage + ) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + # Get first tile info (5 elements) + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + mma_tile_coord_mnl = ( + tile_info[0] // cute.size(tiled_mma.thr_id.shape), + tile_info[1], + tile_info[2], + ) + + tAgA_slice = tAgA[(None, mma_tile_coord_mnl[0], None, 0)] + tBgB_slice = tBgB[(None, mma_tile_coord_mnl[1], None, mma_tile_coord_mnl[2])] + tAgSFA_slice = tAgSFA[(None, mma_tile_coord_mnl[0], None, 0)] + + slice_n = mma_tile_coord_mnl[1] + if cutlass.const_expr(self.cta_tile_shape_mnk[1] == 64): + slice_n = mma_tile_coord_mnl[1] // 2 + + tBgSFB_slice = tBgSFB[(None, slice_n, None, mma_tile_coord_mnl[2])] + + ab_producer_state.reset_count() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + + for k_tile in cutlass.range(0, k_tile_cnt, 1, unroll=1): + ab_pipeline.producer_acquire(ab_producer_state, peek_ab_empty_status) + + cute.copy( + tma_atom_a, + tAgA_slice[(None, ab_producer_state.count)], + tAsA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=a_full_mcast_mask, + ) + cute.copy( + tma_atom_b, + tBgB_slice[(None, ab_producer_state.count)], + tBsB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=b_full_mcast_mask, + ) + cute.copy( + tma_atom_sfa, + tAgSFA_slice[(None, ab_producer_state.count)], + tAsSFA[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=sfa_full_mcast_mask, + ) + cute.copy( + tma_atom_sfb, + tBgSFB_slice[(None, ab_producer_state.count)], + tBsSFB[(None, ab_producer_state.index)], + tma_bar_ptr=ab_pipeline.producer_get_barrier(ab_producer_state), + mcast_mask=sfb_full_mcast_mask, + ) + + ab_producer_state.advance() + peek_ab_empty_status = cutlass.Boolean(1) + if ab_producer_state.count < k_tile_cnt: + peek_ab_empty_status = ab_pipeline.producer_try_acquire(ab_producer_state) + + # Get next tile + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + ab_pipeline.producer_tail(ab_producer_state) + + # + # Specialized MMA warp + # + if warp_idx == self.mma_warp_id: + tmem.wait_for_alloc() + + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + sfa_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFA = cute.make_tensor(sfa_tmem_ptr, tCtSFA_layout) + + sfb_tmem_ptr = cute.recast_ptr( + acc_tmem_ptr + self.num_accumulator_tmem_cols + self.num_sfa_tmem_cols, + dtype=self.sf_dtype, + ) + tCtSFB = cute.make_tensor(sfb_tmem_ptr, tCtSFB_layout) + + sfa_s2t_bundle = self._mainloop_s2t_copy_and_partition(sSFA, tCtSFA) + sfb_s2t_bundle = self._mainloop_s2t_copy_and_partition(sSFB, tCtSFB) + + ab_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_ab_stage + ) + acc_producer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Producer, self.num_acc_stage + ) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + # Get first tile info (5 elements) + tile_info = cute.make_rmem_tensor((4,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + mma_tile_coord_mnl = ( + tile_info[0] // cute.size(tiled_mma.thr_id.shape), + tile_info[1], + tile_info[2], + ) + + acc_stage_index = acc_producer_state.index + tCtAcc = tCtAcc_base[(None, None, None, acc_stage_index)] + + ab_consumer_state.reset_count() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt and is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + + if is_leader_cta: + acc_pipeline.producer_acquire(acc_producer_state) + + # TMEM pointer offset for cta_tile_shape_n=192 or 64 + tCtSFB_mma = tCtSFB + if cutlass.const_expr(self.cta_tile_shape_mnk[1] in {64, 192}): + offset = cutlass.Int32((mma_tile_coord_mnl[1] % 2) * 2) + shifted_ptr = cute.recast_ptr( + acc_tmem_ptr + + self.num_accumulator_tmem_cols + + self.num_sfa_tmem_cols + + offset, + dtype=self.sf_dtype, + ) + tCtSFB_mma = cute.make_tensor(shifted_ptr, tCtSFB_layout) + + # MMA mainloop + for k_tile in range(k_tile_cnt): + if is_leader_cta: + ab_pipeline.consumer_wait(ab_consumer_state, peek_ab_full_status) + + if cutlass.const_expr(not self._is_interleaved_utccp()): + self._mainloop_s2t_copies( + ab_consumer_state.index, sfa_s2t_bundle, sfb_s2t_bundle + ) + + num_kblocks = cute.size(tCrA, mode=[2]) + for k_block in cutlass.range(num_kblocks, unroll_full=True): + if cutlass.const_expr( + self.enable_breuse + and cute.size(tCtAcc.layout, mode=[1]) == 2 + and cute.size(tCtAcc.layout, mode=[2]) == 1 + ): + tCtAcc_bkeep = tCtAcc[(None, 0, 0)] + tCtAcc_breuse = tCtAcc[(None, 1, 0)] + + a_kblk_crd_keep = ( + None, + 0, + k_block, + ab_consumer_state.index, + ) + a_kblk_crd_reuse = ( + None, + 1, + k_block, + ab_consumer_state.index, + ) + b_kblk_crd = (None, 0, k_block, ab_consumer_state.index) + + sfa_kblk_crd_keep = (None, 0, k_block) + sfa_kblk_crd_reuse = (None, 1, k_block) + sfb_kblk_crd = (None, 0, k_block) + + if cutlass.const_expr(self._is_interleaved_utccp()): + self._mainloop_s2t_interleaved_copies( + k_block, + ab_consumer_state.index, + sfa_s2t_bundle, + sfb_s2t_bundle, + ) + + # Bkeep + tiled_mma_bkeep.set( + tcgen05.Field.ACCUMULATE, + k_tile != 0 or k_block != 0, + ) + cute.gemm( + tiled_mma_bkeep, + tCtAcc_bkeep, + [tCrA[a_kblk_crd_keep], tCtSFA[sfa_kblk_crd_keep]], + [tCrB[b_kblk_crd], tCtSFB_mma[sfb_kblk_crd]], + tCtAcc_bkeep, + ) + # Breuse + tiled_mma_breuse.set( + tcgen05.Field.ACCUMULATE, + k_tile != 0 or k_block != 0, + ) + cute.gemm( + tiled_mma_breuse, + tCtAcc_breuse, + [ + tCrA[a_kblk_crd_reuse], + tCtSFA[sfa_kblk_crd_reuse], + ], + [tCrB[b_kblk_crd], tCtSFB_mma[sfb_kblk_crd]], + tCtAcc_breuse, + ) + else: + kblk_crd = ( + None, + None, + k_block, + ab_consumer_state.index, + ) + sf_kblk_crd = (None, None, k_block) + + tiled_mma.set( + tcgen05.Field.ACCUMULATE, + k_tile != 0 or k_block != 0, + ) + cute.gemm( + tiled_mma, + tCtAcc, + [tCrA[kblk_crd], tCtSFA[sf_kblk_crd]], + [tCrB[kblk_crd], tCtSFB_mma[sf_kblk_crd]], + tCtAcc, + ) + + ab_pipeline.consumer_release(ab_consumer_state) + + ab_consumer_state.advance() + peek_ab_full_status = cutlass.Boolean(1) + if ab_consumer_state.count < k_tile_cnt: + if is_leader_cta: + peek_ab_full_status = ab_pipeline.consumer_try_wait(ab_consumer_state) + + if is_leader_cta: + acc_pipeline.producer_commit(acc_producer_state) + acc_producer_state.advance() + + # Get next tile (4 elements) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(4, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + acc_pipeline.producer_tail(acc_producer_state) + + # + # Specialized epilogue warps with finalize fusion + # + if warp_idx < self.mma_warp_id: + tmem.allocate(self.num_tmem_alloc_cols) + tmem.wait_for_alloc() + + acc_tmem_ptr = tmem.retrieve_ptr(self.acc_dtype) + tCtAcc_base = cute.make_tensor(acc_tmem_ptr, tCtAcc_fake.layout) + + # Layout transformation for tCtAcc_base and tCgC + # ((MMA_ATOM_M, MMA_ATOM_N), MMA_M, MMA_N, ...rest) + # -> ((MMA_ATOM_M, MMA_M), (MMA_ATOM_N, MMA_N), ...rest) + tCtAcc = transform_partitioned_tensor_layout(tCtAcc_base) + tCgC_transformed = transform_partitioned_tensor_layout(tCgC) + + # Partition for epilogue + epi_tidx = tidx % 128 + ( + tiled_copy_t2r, + tTR_tAcc_base, + tTR_rAcc, + ) = epilogue_tmem_copy_and_partition( + self, epi_tidx, tCtAcc, tCgC_transformed, epi_tile, use_2cta_instrs + ) + + tTR_rC = cute.make_rmem_tensor(tTR_rAcc.shape, self.c_dtype) + + # Setup smem copy for block reduce + atom = cute.make_copy_atom( + cute.nvgpu.CopyUniversalOp(), + self.c_dtype, + ) + tiled_copy_r2s = cute.make_tiled_copy_D(atom, tiled_copy_t2r) + thr_copy_r2s = tiled_copy_r2s.get_slice(epi_tidx) + tRS_sC = thr_copy_r2s.partition_D(sC) + tRS_rC = tiled_copy_r2s.retile(tTR_rC) + + acc_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage + ) + + tile_info_consumer_state = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_tile_stage + ) + + token_idx = cutlass.Int32(0) + token_scale = self.final_scale_dtype(0.0) + + # Get first tile info (5 elements) + tile_info = cute.make_rmem_tensor((5,), cutlass.Int32) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(5, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + while is_valid_tile: + mma_tile_coord_mnl = ( + tile_info[0] // cute.size(tiled_mma.thr_id.shape), + tile_info[1], + tile_info[2], + ) + + expert_idx = mma_tile_coord_mnl[2] + alpha_val = alpha[expert_idx] + + # Compute base row index for this tile + tile_m_start = tile_info[0] * self.cta_tile_shape_mnk[0] + + # Get accumulator stage index + acc_stage_index = acc_consumer_state.index + + # Set tensor memory buffer for current tile + tTR_tAcc = tTR_tAcc_base[(None, None, None, None, None, acc_stage_index)] + + # Wait for accumulator buffer full + acc_pipeline.consumer_wait(acc_consumer_state) + + tTR_tAcc = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) + + # Group tRS_sC modes 1 and 2 (m_iter, n_subtile) into a single 2D mode + # Before: ((1,32), 2, 4, 1) -> After: ((1,32), (2,4), 1) + # This allows using tuple indexing like tTR_tAcc + tRS_sC_grouped = cute.group_modes(tRS_sC, 1, 3) + + # Get m-iteration count and n-subtile count from tTR_tAcc shape + # In B-reuse (cta_tile_m=256): mode-3 is (2,4), m_iter_cnt=2, n_subtile_cnt=4 + # In non-B-reuse (cta_tile_m=128): mode-3 is (1,4), m_iter_cnt=1, n_subtile_cnt=4 + m_iter_cnt = cute.size(tTR_tAcc.shape[3], mode=[0]) + n_subtile_cnt = cute.size(tTR_tAcc.shape[3], mode=[1]) + + # TODO:How could reduction be done with a better perf? + # Process all m-iterations and n-subtiles + for m_iter_idx in cutlass.range(m_iter_cnt): + # Compute row indices for this m-iteration + # Each thread handles row: tile_m_start + m_iter_idx * 128 + epi_tidx + permuted_row = tile_m_start + m_iter_idx * 128 + epi_tidx + expanded_idx = permuted_idx_to_expanded_idx[permuted_row] + is_valid_row = permuted_row < tile_info[4] + + # Compute token info and scaled alpha for this row + token_idx = cutlass.Int32(0) + alpha_val_iter = alpha_val + if is_valid_row: + token_idx = expanded_idx // self.topK + topk_idx = expanded_idx % self.topK + token_scale = token_final_scales[(token_idx, topk_idx)] + alpha_val_iter = alpha_val * token_scale + + for n_iter_idx in cutlass.range(n_subtile_cnt): + # Load accumulator from tensor memory buffer to register + # Index with (m_iter_idx, subtile_idx) for the composite mode-3 + tTR_tAcc_mn = tTR_tAcc[(None, None, None, (m_iter_idx, n_iter_idx))] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + + # For block reduce: retile then load, compute, store + acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() + acc_vec_final = (alpha_val_iter * acc_vec).to(self.c_dtype) + tRS_rC.store(acc_vec_final) + if is_valid_row: + # Use grouped tRS_sC with tuple indexing to preserve rank + # tRS_sC_grouped: ((1,32), (2,4), 1) + # Index: (None, (m_iter_idx, subtile_idx), None) -> ((1,32), 1, 1) rank 3 + cute.copy( + tiled_copy_r2s, + # TODO: check if there is a better way to index and make same rank + tRS_rC[None, 0, 0], + tRS_sC_grouped[(None, (m_iter_idx, n_iter_idx), 0)], + ) + + cute.arch.fence_proxy("async.shared", space="cta") + + # Async arrive accumulator buffer empty + with cute.arch.elect_one(): + acc_pipeline.consumer_release(acc_consumer_state) + acc_consumer_state.advance() + + # TODO:Currently finish all sts and do all reduce, do we have a better way for epilogue overlapping? + # Block reduce for all m-iterations + for m_iter_idx in cutlass.range(m_iter_cnt): + permuted_row = tile_m_start + m_iter_idx * 128 + epi_tidx + is_valid_row = permuted_row < tile_info[4] + + if is_valid_row: + expanded_idx = permuted_idx_to_expanded_idx[permuted_row] + token_idx = expanded_idx // self.topK + coord_n = mma_tile_coord_mnl[1] * self.cta_tile_shape_mnk[1] + scatter_out_offset = cute.domain_offset((token_idx, coord_n, 0), c) + + # sC row index: m_iter_idx * 128 + epi_tidx + sC_row = m_iter_idx * 128 + epi_tidx + + if cutlass.const_expr(self.c_dtype == cutlass.BFloat16): + blk_reduce_bf16( + scatter_out_offset, + sC[sC_row, None, 0], + cutlass.Int32(self.copy_size), + ) + elif cutlass.const_expr(self.c_dtype == cutlass.Float32): + blk_reduce_fp32( + scatter_out_offset, + sC[sC_row, None, 0], + cutlass.Int32(self.copy_size), + ) + elif cutlass.const_expr(self.c_dtype == cutlass.Float16): + blk_reduce_fp16( + scatter_out_offset, + sC[sC_row, None, 0], + cutlass.Int32(self.copy_size), + ) + + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) + self.epilog_sync_barrier.arrive_and_wait() + # ============================================================ + # END OF NEW CODE + # ============================================================ + + # Get next tile (5 elements) + tile_info_pipeline.consumer_wait(tile_info_consumer_state) + for idx in cutlass.range(5, unroll_full=True): + tile_info[idx] = sInfo[(idx, tile_info_consumer_state.index)] + is_valid_tile = tile_info[3] == 1 + cute.arch.fence_proxy("async.shared", space="cta") + tile_info_pipeline.consumer_release(tile_info_consumer_state) + tile_info_consumer_state.advance() + + # Dealloc tensor memory + tmem.relinquish_alloc_permit() + self.epilog_sync_barrier.arrive_and_wait() + tmem.free(acc_tmem_ptr) + + @staticmethod + def _compute_grid( + gemm_shape: Tuple[int, int, int], + cta_tile_shape_mnk: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + max_active_clusters: cutlass.Constexpr, + raster_along_m: bool, + ) -> Tuple[utils.PersistentTileSchedulerParams, Tuple[int, int, int]]: + """Compute grid size based on GEMM shape.""" + (m, n, l) = gemm_shape # noqa: E741 + + num_ctas_m = cute.ceil_div(m, cta_tile_shape_mnk[0]) + num_ctas_n = cute.ceil_div(n, cta_tile_shape_mnk[1]) + num_ctas_l = l + + num_ctas_mnl = (num_ctas_m, num_ctas_n, num_ctas_l) + cluster_shape_mnl = (*cluster_shape_mn, 1) + + tile_sched_params = utils.PersistentTileSchedulerParams( + num_ctas_mnl, cluster_shape_mnl, raster_along_m=raster_along_m + ) + grid = utils.StaticPersistentTileScheduler.get_grid_shape( + tile_sched_params, max_active_clusters + ) + + return tile_sched_params, grid + + @staticmethod + def is_valid_dtypes_and_scale_factor_vec_size( + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + ) -> bool: + """Check if the dtypes and sf_vec_size are valid combinations.""" + valid_combinations = { + # 4xFP4 + (cutlass.Float4E2M1FN, cutlass.Float4E2M1FN, cutlass.Float8E8M0FNU, 16), + (cutlass.Float4E2M1FN, cutlass.Float4E2M1FN, cutlass.Float8E8M0FNU, 32), + (cutlass.Float4E2M1FN, cutlass.Float4E2M1FN, cutlass.Float8E4M3FN, 16), + (cutlass.Float4E2M1FN, cutlass.Float4E2M1FN, cutlass.Float8E4M3FN, 32), + # 2xFP8 + (cutlass.Float8E5M2, cutlass.Float8E5M2, cutlass.Float8E8M0FNU, 32), + (cutlass.Float8E5M2, cutlass.Float8E4M3FN, cutlass.Float8E8M0FNU, 32), + (cutlass.Float8E4M3FN, cutlass.Float8E4M3FN, cutlass.Float8E8M0FNU, 32), + (cutlass.Float8E4M3FN, cutlass.Float8E5M2, cutlass.Float8E8M0FNU, 32), + } + + current_combination = (a_dtype, b_dtype, sf_dtype, sf_vec_size) + if current_combination not in valid_combinations: + return False + + if c_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.BFloat16}: + return False + + return True + + @staticmethod + def is_valid_layouts( + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """Check if layouts and dtypes are valid combinations.""" + if ( + a_dtype is cutlass.Float4E2M1FN + and b_dtype is cutlass.Float4E2M1FN + and not (a_major == "k" and b_major == "k") + ): + return False + + if c_dtype is cutlass.Float4E2M1FN and c_major == "m": + return False + + return True + + @staticmethod + def is_valid_mma_tiler_and_cluster_shape( + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + mma_inst_shape: Tuple[int, int, int], + mma_tiler: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + ) -> bool: + """Check if the mma tiler and cluster shape are valid.""" + # Check valid mma_inst_shape + if mma_inst_shape[0] not in [128, 256]: + return False + if mma_inst_shape[1] not in [64, 128, 192, 256]: + return False + + # Check valid mma_tiler + if mma_tiler[0] not in [128, 256, 512]: + return False + if mma_tiler[1] not in [64, 128, 192, 256]: + return False + + # Check MMA tiler vs MMA instruction relationship + b_reuse = mma_tiler[0] // mma_inst_shape[0] == 2 + if mma_tiler[0] != mma_inst_shape[0] and not b_reuse: + return False + if mma_tiler[1] != mma_inst_shape[1]: + return False + + # Check K-dimension constraints + if a_dtype in {cutlass.Float8E4M3FN, cutlass.Float8E5M2} and b_dtype in { + cutlass.Float8E4M3FN, + cutlass.Float8E5M2, + }: + if mma_tiler[2] != 128 or mma_inst_shape[2] != 64: + return False + else: + if mma_tiler[2] != 256 or mma_inst_shape[2] != 128: + return False + + # Check cluster shape + if cluster_shape_mn[0] % (2 if mma_inst_shape[0] == 256 else 1) != 0: + return False + + # Check cluster shape validity + def is_power_of_2(x): + return x > 0 and (x & (x - 1)) == 0 + + if ( + cluster_shape_mn[0] * cluster_shape_mn[1] > 16 + or cluster_shape_mn[0] <= 0 + or cluster_shape_mn[1] <= 0 + or cluster_shape_mn[0] > 4 + or cluster_shape_mn[1] > 4 + or not is_power_of_2(cluster_shape_mn[0]) + or not is_power_of_2(cluster_shape_mn[1]) + ): + return False + + return True + + @staticmethod + def is_valid_tensor_alignment( + m: int, + n: int, + k: int, + l: int, # noqa: E741 + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """Check if the tensor alignment is valid (16B alignment).""" + + def check_contigous_16B_alignment(dtype, is_mode0_major, tensor_shape): + major_mode_idx = 0 if is_mode0_major else 1 + num_major_elements = tensor_shape[major_mode_idx] + num_contiguous_elements = 16 * 8 // dtype.width + return num_major_elements % num_contiguous_elements == 0 + + if ( + not check_contigous_16B_alignment(a_dtype, a_major == "m", (m, k, l)) + or not check_contigous_16B_alignment(b_dtype, b_major == "n", (n, k, l)) + or not check_contigous_16B_alignment(c_dtype, c_major == "m", (m, n, l)) + ): + return False + + return True + + @cute.jit + def wrapper( + self, + a_ptr: cute.Pointer, + b_ptr: cute.Pointer, + a_sf_ptr: cute.Pointer, + b_sf_ptr: cute.Pointer, + c_ptr: cute.Pointer, + alpha_ptr: cute.Pointer, + tile_idx_to_group_idx_ptr: cute.Pointer, + tile_idx_to_mn_limit_ptr: cute.Pointer, + permuted_idx_to_expanded_idx_ptr: cute.Pointer, + num_non_exiting_tiles_ptr: cute.Pointer, + token_final_scales_ptr: cute.Pointer, + m: cutlass.Int64, + n: cutlass.Int64, + k: cutlass.Int64, + l: cutlass.Int64, # noqa: E741 + num_tokens: cutlass.Int64, + top_k: cutlass.Int64, + tile_size: cutlass.Constexpr, + scaling_vector_size: cutlass.Constexpr, + max_active_clusters: cutlass.Constexpr, + stream: cuda.CUstream, + epilogue_op: cutlass.Constexpr = lambda x: x, + c_stride_row: cutlass.Int64 = cutlass.Int64(0), + ): + scale_k = k // scaling_vector_size + num_tiles = m // tile_size + a = cute.make_tensor(a_ptr, layout=cute.make_ordered_layout((m, k, 1), order=(1, 0, 2))) + b = cute.make_tensor(b_ptr, layout=cute.make_ordered_layout((n, k, l), order=(1, 0, 2))) + a_sf = cute.make_tensor( + a_sf_ptr, + layout=cute.make_ordered_layout( + (32, 4, m // 128, 4, scale_k // 4, 1), order=(2, 1, 4, 0, 3, 5) + ), + ) + b_sf = cute.make_tensor( + b_sf_ptr, + layout=cute.make_ordered_layout( + (32, 4, n // 128, 4, scale_k // 4, l), order=(2, 1, 4, 0, 3, 5) + ), + ) + actual_c_stride_row = n if c_stride_row == 0 else c_stride_row + c = cute.make_tensor( + c_ptr, + layout=cute.make_layout( + (num_tokens, n, 1), + stride=(actual_c_stride_row, 1, num_tokens * actual_c_stride_row), + ), + ) + alpha = cute.make_tensor(alpha_ptr, layout=cute.make_layout((l,))) + + tile_idx_to_group_idx = cute.make_tensor( + tile_idx_to_group_idx_ptr, layout=cute.make_layout((num_tiles,)) + ) + tile_idx_to_mn_limit = cute.make_tensor( + tile_idx_to_mn_limit_ptr, layout=cute.make_layout((num_tiles,)) + ) + permuted_idx_to_expanded_idx = cute.make_tensor( + permuted_idx_to_expanded_idx_ptr, layout=cute.make_layout((m,)) + ) + num_non_exiting_tiles = cute.make_tensor( + num_non_exiting_tiles_ptr, layout=cute.make_layout((1,)) + ) + token_final_scales = cute.make_tensor( + token_final_scales_ptr, + layout=cute.make_ordered_layout((num_tokens, top_k), order=(1, 0)), + ) + + return self( + a, + b, + c, + a_sf, + b_sf, + tile_idx_to_group_idx, + num_non_exiting_tiles, + tile_idx_to_mn_limit, + alpha, + max_active_clusters=max_active_clusters, + stream=stream, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + token_final_scales=token_final_scales, + epilogue_op=epilogue_op, + ) + + @staticmethod + def can_implement( + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + c_dtype: Type[cutlass.Numeric], + mma_inst_shape: Tuple[int, int, int], + mma_tiler: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + m: int, + n: int, + k: int, + l: int, # noqa: E741 + a_major: str, + b_major: str, + c_major: str, + ) -> bool: + """Check if the GEMM can be implemented.""" + # Check data types + if not Sm107BlockScaledContiguousGroupedGemmFinalizeFusionKernel.is_valid_dtypes_and_scale_factor_vec_size( + a_dtype, b_dtype, sf_dtype, sf_vec_size, c_dtype + ): + return False + + # Check layouts + if not Sm107BlockScaledContiguousGroupedGemmFinalizeFusionKernel.is_valid_layouts( + a_dtype, b_dtype, c_dtype, a_major, b_major, c_major + ): + return False + + # Check MMA tiler and cluster shape + if not Sm107BlockScaledContiguousGroupedGemmFinalizeFusionKernel.is_valid_mma_tiler_and_cluster_shape( + a_dtype, b_dtype, mma_inst_shape, mma_tiler, cluster_shape_mn + ): + return False + + # Check tensor alignment + if not Sm107BlockScaledContiguousGroupedGemmFinalizeFusionKernel.is_valid_tensor_alignment( + m, n, k, l, a_dtype, b_dtype, c_dtype, a_major, b_major, c_major + ): + return False + + return True + + +# ============================================================================ +# Run utilities +# ============================================================================ + + +@cute.jit +def cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + sf_ref_tensor: cute.Tensor, + sf_mma_tensor: cute.Tensor, +): + """Convert scale factor tensor from MKL layout to mma specification layout.""" + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 0, 3) + sf_mma_tensor = cute.group_modes(sf_mma_tensor, 1, 3) + for i in cutlass.range(cute.size(sf_ref_tensor)): + mkl_coord = sf_ref_tensor.layout.get_hier_coord(i) + sf_mma_tensor[mkl_coord] = sf_ref_tensor[mkl_coord] + + +def create_mask(group_m_list, cta_tile_mn, permuted_m=None): + """Create mask and group mapping for contiguous grouped GEMM. + + :param group_m_list: List of M values for each group (will be aligned to cta_tile_mn[0] dimension) + :param cta_tile_mn: CTA tile size tuple (M, N) - M dimension used for alignment + :param permuted_m: Optional padded M dimension for cuda_graph support. + + :return: Tuple of (valid_m, aligned_group_m_list, tile_idx_to_expert_idx, + num_non_exiting_tiles, tile_idx_to_mn_limit) + """ + m_aligned = cta_tile_mn[0] + valid_m = 0 + aligned_group_m_list = [] + tile_idx_to_expert_idx = [] + tile_idx_to_mn_limit = [] + + for i, group_m in enumerate(group_m_list): + aligned_group_m = ((group_m + m_aligned - 1) // m_aligned) * m_aligned + aligned_group_m_list.append(aligned_group_m) + + num_tiles_in_group = aligned_group_m // cta_tile_mn[0] + tile_idx_to_expert_idx.extend([i] * num_tiles_in_group) + + tile_idx_to_mn_limit.extend([group_m + valid_m] * num_tiles_in_group) + valid_m += aligned_group_m + + num_non_exiting_tiles = len(tile_idx_to_expert_idx) + + if permuted_m is not None: + if permuted_m < valid_m: + raise ValueError(f"permuted_m ({permuted_m}) must be >= valid_m ({valid_m}).") + if permuted_m > valid_m: + num_padding_tiles = (permuted_m - valid_m) // cta_tile_mn[0] + tile_idx_to_expert_idx.extend([0] * num_padding_tiles) + + tile_idx_to_expert_idx = torch.tensor(tile_idx_to_expert_idx, device="cuda", dtype=torch.int32) + num_non_exiting_tiles_tensor = torch.tensor( + [num_non_exiting_tiles], device="cuda", dtype=torch.int32 + ) + tile_idx_to_mn_limit_tensor = torch.tensor( + tile_idx_to_mn_limit, device="cuda", dtype=torch.int32 + ) + + return ( + valid_m, + aligned_group_m_list, + tile_idx_to_expert_idx, + num_non_exiting_tiles_tensor, + tile_idx_to_mn_limit_tensor, + ) + + +def create_fused_finalize_tensors(seq_len, topK, permuted_m, group_m_list, mma_tiler_mn): + """Create tensors for fused finalize operation.""" + m_aligned = mma_tiler_mn[0] + permuted_idx_to_expanded_idx_tensor = torch.empty( + (permuted_m,), dtype=torch.int32, device="cuda" + ).fill_(-1) + token_final_scales = torch.rand(seq_len, topK).to(dtype=torch.float32).cuda() + token_final_scales = token_final_scales / token_final_scales.sum(dim=1, keepdim=True) + + start_idx = 0 + for group_idx in range(len(group_m_list)): + m_per_group = group_m_list[group_idx] + + if m_per_group > 0: + expert_set_idx = group_idx // topK + k_in_set = group_idx % topK + start_token = expert_set_idx * m_per_group + + token_indices = torch.arange( + start_token, start_token + m_per_group, dtype=torch.int32, device="cuda" + ) + token_indices = token_indices % seq_len + expanded_idx = token_indices * topK + k_in_set + + permuted_idx_to_expanded_idx_tensor[start_idx : (start_idx + m_per_group)] = ( + expanded_idx + ) + m_aligned_per_group = ((m_per_group + m_aligned - 1) // m_aligned) * m_aligned + start_idx += m_aligned_per_group + + return ( + permuted_idx_to_expanded_idx_tensor, + token_final_scales, + from_dlpack(permuted_idx_to_expanded_idx_tensor).mark_layout_dynamic(), + from_dlpack(token_final_scales).mark_layout_dynamic(), + ) + + +def create_scale_factor_tensor(l, mn, k, sf_vec_size, dtype): # noqa: E741 + """Create scale factor tensor with proper layout conversion.""" + + def ceil_div(a, b): + return (a + b - 1) // b + + sf_k = ceil_div(k, sf_vec_size) + ref_shape = (l, mn, sf_k) + + atom_m = (32, 4) + atom_k = 4 + mma_shape = ( + l, + ceil_div(mn, atom_m[0] * atom_m[1]), + ceil_div(sf_k, atom_k), + atom_m[0], + atom_m[1], + atom_k, + ) + + ref_permute_order = (1, 2, 0) + mma_permute_order = (3, 4, 1, 5, 2, 0) + + ref_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + ref_shape, + torch.float32, + permute_order=ref_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=1, max_val=3), + ) + + cute_f32_torch_tensor_cpu = cutlass_torch.create_and_permute_torch_tensor( + mma_shape, + torch.float32, + permute_order=mma_permute_order, + init_type=cutlass_torch.TensorInitType.RANDOM, + init_config=cutlass_torch.RandomInitConfig(min_val=0, max_val=1), + ) + + cvt_sf_MKL_to_M32x4xrm_K4xrk_L( + from_dlpack(ref_f32_torch_tensor_cpu), + from_dlpack(cute_f32_torch_tensor_cpu), + ) + + cute_f32_torch_tensor = cute_f32_torch_tensor_cpu.cuda() + + ref_f32_torch_tensor_cpu = ( + ref_f32_torch_tensor_cpu.permute(2, 0, 1) + .unsqueeze(-1) + .expand(l, mn, sf_k, sf_vec_size) + .reshape(l, mn, sf_k * sf_vec_size) + .permute(*ref_permute_order) + ) + ref_f32_torch_tensor_cpu = ref_f32_torch_tensor_cpu[:, :k, :] + + cute_tensor, cute_torch_tensor = cutlass_torch.cute_tensor_like( + cute_f32_torch_tensor_cpu, + dtype, + is_dynamic_layout=True, + assumed_align=16, + ) + + cute_tensor = cutlass_torch.convert_cute_tensor( + cute_f32_torch_tensor, + cute_tensor, + dtype, + is_dynamic_layout=True, + ) + + return ref_f32_torch_tensor_cpu, cute_tensor, cute_torch_tensor + + +def create_tensors( + l, # noqa: E741 + group_m_list, + n, + k, + a_major, + b_major, + cd_major, + a_dtype, + b_dtype, + c_dtype, + sf_dtype, + sf_vec_size, + mma_tiler_mn, + permuted_m=None, + seq_len=None, +): + """Create tensors for contiguous grouped GEMM with finalize fusion.""" + torch.manual_seed(1111) + + alpha_torch_cpu = torch.ones((l,), dtype=torch.float32) * 0.1 + + ( + valid_m, + aligned_group_m_list, + _tile_idx_to_expert_idx, + _num_non_exiting_tiles, + _tile_idx_to_mn_limit, + ) = create_mask(group_m_list, mma_tiler_mn, permuted_m) + + tensor_m = permuted_m if permuted_m is not None else valid_m + + a_torch_cpu = cutlass_torch.matrix(1, tensor_m, k, a_major == "m", cutlass.Float32) + b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", cutlass.Float32) + c_torch_cpu = cutlass_torch.matrix(1, seq_len, n, cd_major == "m", cutlass.Float32) + + # # Fill A and B with 1 for debugging + # a_torch_cpu.fill_(1.0) + # b_torch_cpu.fill_(1.0) + c_torch_cpu.fill_(0) + + a_tensor, a_torch_gpu = cutlass_torch.cute_tensor_like( + a_torch_cpu, a_dtype, is_dynamic_layout=True, assumed_align=16 + ) + b_tensor, b_torch_gpu = cutlass_torch.cute_tensor_like( + b_torch_cpu, b_dtype, is_dynamic_layout=True, assumed_align=16 + ) + c_tensor, c_torch_gpu = cutlass_torch.cute_tensor_like( + c_torch_cpu, c_dtype, is_dynamic_layout=True, assumed_align=16 + ) + + a_tensor.mark_compact_shape_dynamic( + mode=1 if a_major == "k" else 0, + stride_order=(2, 0, 1) if a_major == "k" else (2, 1, 0), + divisibility=32 if a_dtype == cutlass.Float4E2M1FN else 16, + ) + b_tensor.mark_compact_shape_dynamic( + mode=1 if b_major == "k" else 0, + stride_order=(2, 0, 1) if b_major == "k" else (2, 1, 0), + divisibility=32 if b_dtype == cutlass.Float4E2M1FN else 16, + ) + c_tensor.mark_compact_shape_dynamic( + mode=1 if cd_major == "n" else 0, + stride_order=(2, 0, 1) if cd_major == "n" else (2, 1, 0), + divisibility=32 if c_dtype == cutlass.Float4E2M1FN else 16, + ) + + sfa_torch_cpu, sfa_tensor, sfa_torch_gpu = create_scale_factor_tensor( + 1, tensor_m, k, sf_vec_size, sf_dtype + ) + sfb_torch_cpu, sfb_tensor, sfb_torch_gpu = create_scale_factor_tensor( + l, n, k, sf_vec_size, sf_dtype + ) + + tile_idx_to_expert_idx = from_dlpack(_tile_idx_to_expert_idx).mark_layout_dynamic() + num_non_exiting_tiles = from_dlpack(_num_non_exiting_tiles).mark_layout_dynamic() + tile_idx_to_mn_limit = from_dlpack(_tile_idx_to_mn_limit).mark_layout_dynamic() + + alpha = from_dlpack(alpha_torch_cpu.cuda()).mark_layout_dynamic() + + c_torch_gpu.fill_(0) + + return ( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + tile_idx_to_expert_idx, + num_non_exiting_tiles, + tile_idx_to_mn_limit, + alpha, + a_torch_cpu, + b_torch_cpu, + c_torch_cpu, + sfa_torch_cpu, + sfb_torch_cpu, + alpha_torch_cpu, + a_torch_gpu, + b_torch_gpu, + sfa_torch_gpu, + sfb_torch_gpu, + c_torch_gpu, + aligned_group_m_list, + valid_m, + ) + + +def verify_reference_result( + a_torch_cpu: torch.Tensor, + b_torch_cpu: torch.Tensor, + sfa_torch_cpu: torch.Tensor, + sfb_torch_cpu: torch.Tensor, + alpha_torch_cpu: torch.Tensor, + permuted_idx_to_expanded_idx_torch: torch.Tensor, + token_final_scales_torch: torch.Tensor, + group_m_list: List[int], + aligned_group_m_list: List[int], + c_dtype: torch.dtype, + valid_m: int, + n: int, + topK: int, + seq_len: int, +) -> torch.Tensor: + """Compute reference result for validation.""" + gemm_output = torch.empty((1, valid_m, n), dtype=torch.float32) + valid_mask = torch.zeros((valid_m,), dtype=torch.bool, device="cuda") + + start = 0 + for i, group_m in enumerate(aligned_group_m_list): + end = start + group_m + res_a = torch.einsum( + "mk,mk->mk", + a_torch_cpu[start:end, :, 0], + sfa_torch_cpu[start:end, :, 0], + ) + res_b = torch.einsum("nk,nk->nk", b_torch_cpu[:, :, i], sfb_torch_cpu[:, :, i]) + gemm_output[0, start:end, :] = torch.einsum("mk,nk->mn", res_a, res_b) * alpha_torch_cpu[i] + valid_mask[start : start + group_m_list[i]] = 1 + start = end + + gemm_output = gemm_output.permute((1, 2, 0)).cuda() + + final_output = torch.zeros((seq_len, n), dtype=c_dtype).cuda() + + gemm_output = gemm_output[:valid_m, :, 0].clone() + expanded_idx_all = permuted_idx_to_expanded_idx_torch[:valid_m] + + expanded_idx_valid = expanded_idx_all[valid_mask] + gemm_output_valid = gemm_output[valid_mask] + + token_idx = expanded_idx_valid // topK + topk_idx = expanded_idx_valid % topK + scales = token_final_scales_torch[token_idx, topk_idx] + scaled_output = gemm_output_valid * scales.unsqueeze(1) + scaled_output = scaled_output.to(c_dtype) + + for i in range(len(token_idx)): + final_output[token_idx[i]] += scaled_output[i] + + return final_output + + +def run( + nkl: Tuple[int, int, int], + group_m_list: Tuple[int, ...], + a_dtype: Type[cutlass.Numeric], + b_dtype: Type[cutlass.Numeric], + c_dtype: Type[cutlass.Numeric], + sf_dtype: Type[cutlass.Numeric], + sf_vec_size: int, + a_major: str, + b_major: str, + c_major: str, + mma_inst_shape: Tuple[int, int, int], + mma_tiler: Tuple[int, int, int], + cluster_shape_mn: Tuple[int, int], + tolerance: float, + warmup_iterations: int = 0, + iterations: int = 1, + skip_ref_check: bool = False, + use_cold_l2: bool = False, + permuted_m: int = None, + topK: int = 8, + seq_len: int = 4096, + raster_along_m: bool = False, + use_cupti: bool = False, + **kwargs, +): + """Run the Rubin contiguous grouped GEMM kernel with finalize fusion.""" + mma_tiler_mn = (mma_tiler[0], mma_tiler[1]) + m_aligned = mma_tiler[0] + + print("Running Rubin Persistent Dense Contiguous Grouped GEMM Finalize Fusion test with:") + print(f"nkl: {nkl}") + print(f"group_m_list: {group_m_list}") + print( + f"A dtype: {a_dtype}, B dtype: {b_dtype}, C dtype: {c_dtype}, SF dtype: {sf_dtype}, SF Vec size: {sf_vec_size}" + ) + print(f"Group M alignment: {m_aligned}") + if permuted_m is not None: + print(f"Padded M (CUDA graph support): {permuted_m}") + print(f"Fused finalize enabled with topK={topK}") + print(f"Sequence length: {seq_len}") + print(f"Matrix majors - A: {a_major}, B: {b_major}, Out: {c_major}") + print(f"Mma Inst Shape (M, N, K): {mma_inst_shape}") + print(f"Mma Tiler (M, N, K): {mma_tiler}") + print(f"Cluster Shape (M, N): {cluster_shape_mn}") + print(f"Raster along M: {raster_along_m}") + print(f"Use CUPTI: {'True' if use_cupti else 'False'}") + n, k, l = nkl # noqa: E741 + + if not torch.cuda.is_available(): + raise RuntimeError("GPU is required to run this example!") + + if not Sm107BlockScaledContiguousGroupedGemmFinalizeFusionKernel.can_implement( + a_dtype, + b_dtype, + sf_dtype, + sf_vec_size, + c_dtype, + mma_inst_shape, + mma_tiler, + cluster_shape_mn, + m_aligned, + n, + k, + l, + a_major, + b_major, + c_major, + ): + raise TypeError( + f"Unsupported testcase a_dtype={a_dtype}, b_dtype={b_dtype}, sf_dtype={sf_dtype}, " + f"sf_vec_size={sf_vec_size}, c_dtype={c_dtype}, mma_inst_shape={mma_inst_shape}, " + f"mma_tiler={mma_tiler}, cluster_shape_mn={cluster_shape_mn}, n={n}, k={k}, l={l}, " + f"a_major={a_major}, b_major={b_major}, c_major={c_major}, m_aligned={m_aligned}" + ) + + ( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + tile_idx_to_expert_idx, + num_non_exiting_tiles, + tile_idx_to_mn_limit, + alpha, + a_torch_cpu, + b_torch_cpu, + c_torch_cpu, + sfa_torch_cpu, + sfb_torch_cpu, + alpha_torch_cpu, + a_torch_gpu, + b_torch_gpu, + sfa_torch_gpu, + sfb_torch_gpu, + c_torch_gpu, + aligned_group_m_list, + valid_m, + ) = create_tensors( + l, + group_m_list, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + sf_dtype, + sf_vec_size, + mma_tiler_mn, + permuted_m, + seq_len, + ) + + tensor_m = permuted_m if permuted_m is not None else valid_m + + ( + permuted_idx_to_expanded_idx_torch, + token_final_scales_torch, + permuted_idx_to_expanded_idx, + token_final_scales, + ) = create_fused_finalize_tensors( + seq_len, + topK, + tensor_m, + group_m_list, + mma_tiler_mn, + ) + + gemm = Sm107BlockScaledContiguousGroupedGemmFinalizeFusionKernel( + sf_vec_size, + mma_inst_shape, + mma_tiler, + cluster_shape_mn, + raster_along_m, + topK, + ) + + hardware_info = cutlass.utils.HardwareInfo() + max_active_clusters = hardware_info.get_max_active_clusters( + cluster_shape_mn[0] * cluster_shape_mn[1] + ) + + print(f"max_active_clusters: {max_active_clusters}") + + current_stream = cutlass_torch.default_stream() + + compiled_gemm = cute.compile( + gemm, + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + tile_idx_to_expert_idx, + num_non_exiting_tiles, + tile_idx_to_mn_limit, + alpha, + max_active_clusters, + current_stream, + permuted_idx_to_expanded_idx, + token_final_scales, + ) + + if not skip_ref_check: + compiled_gemm( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + tile_idx_to_expert_idx, + num_non_exiting_tiles, + tile_idx_to_mn_limit, + alpha, + current_stream, + permuted_idx_to_expanded_idx, + token_final_scales, + ) + + torch.cuda.synchronize() + print("Verifying results...") + ref_result = verify_reference_result( + a_torch_cpu, + b_torch_cpu, + sfa_torch_cpu, + sfb_torch_cpu, + alpha_torch_cpu, + permuted_idx_to_expanded_idx_torch, + token_final_scales_torch, + group_m_list, + aligned_group_m_list, + c_torch_gpu.dtype, + valid_m, + n, + topK, + seq_len, + ) + + actual_result = c_torch_gpu[:, :, 0] + if c_dtype in (cutlass.Float32, cutlass.Float16, cutlass.BFloat16): + torch.testing.assert_close(actual_result, ref_result, atol=tolerance, rtol=1e-02) + + def generate_tensors(): + ( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + tile_idx_to_expert_idx, + num_non_exiting_tiles, + tile_idx_to_mn_limit, + alpha, + *_, + ) = create_tensors( + l, + group_m_list, + n, + k, + a_major, + b_major, + c_major, + a_dtype, + b_dtype, + c_dtype, + sf_dtype, + sf_vec_size, + mma_tiler_mn, + permuted_m, + seq_len, + ) + + ( + _, + _, + permuted_idx_to_expanded_idx, + token_final_scales, + ) = create_fused_finalize_tensors( + seq_len, + topK, + tensor_m, + group_m_list, + mma_tiler_mn, + ) + + return cute.testing.JitArguments( + a_tensor, + b_tensor, + c_tensor, + sfa_tensor, + sfb_tensor, + tile_idx_to_expert_idx, + num_non_exiting_tiles, + tile_idx_to_mn_limit, + alpha, + current_stream, + permuted_idx_to_expanded_idx, + token_final_scales, + ) + + workspace_count = 1 + if use_cold_l2: + one_workspace_bytes = ( + a_torch_gpu.numel() * a_torch_gpu.element_size() + + b_torch_gpu.numel() * b_torch_gpu.element_size() + + c_torch_gpu.numel() * c_torch_gpu.element_size() + + sfa_torch_gpu.numel() * sfa_torch_gpu.element_size() + + sfb_torch_gpu.numel() * sfb_torch_gpu.element_size() + + (tensor_m // mma_tiler_mn[0]) * 4 + + 1 * 4 + + alpha_torch_cpu.numel() * alpha_torch_cpu.element_size() + ) + workspace_count = cute.testing.get_workspace_count( + one_workspace_bytes, warmup_iterations, iterations + ) + + exec_time = cute.testing.benchmark( + compiled_gemm, + workspace_generator=generate_tensors, + workspace_count=workspace_count, + stream=current_stream, + warmup_iterations=warmup_iterations, + iterations=iterations, + use_cupti=use_cupti, + ) + + return exec_time + + +def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: + """Parse comma-separated integers from string.""" + try: + return tuple(int(x.strip()) for x in s.split(",")) + except ValueError: + raise argparse.ArgumentTypeError("Invalid format. Expected comma-separated integers.") + + +def read_benchmark_file( + filepath: str, +) -> Tuple[Tuple[int, int, int], Tuple[int, ...]]: + """Read benchmark file and return nkl and group_m_list.""" + problems = [] + try: + with open(filepath, "r") as f: + for line in f: + line = line.strip() + if not line or line.startswith("#"): + continue + + parts = line.split() + if len(parts) < 2: + continue + + dims = parts[1].split("x") + if len(dims) == 3: + m, n, k = int(dims[0]), int(dims[1]), int(dims[2]) + problems.append((m, n, k)) + + if not problems: + raise ValueError(f"No valid problems found in benchmark file: {filepath}") + + m_first, n, k = problems[0] + l = len(problems) # noqa: E741 + + m_values = tuple(m for m, _, _ in problems) + + print(f"Loaded {l} problems from benchmark file") + print(f"Using N={n}, K={k}, L={l}") + print(f"M values per group: {m_values}") + + return ((n, k, l), m_values) + + except FileNotFoundError: + raise argparse.ArgumentTypeError(f"Benchmark file not found: {filepath}") + except Exception as e: + raise argparse.ArgumentTypeError(f"Error reading benchmark file: {e}") + + +def parse_benchmark_arg( + arg: str, +) -> Tuple[Tuple[int, int, int], Tuple[int, ...]]: + """Parse benchmark argument string.""" + match_list = re.match(r"\[([\d,\s]+)\]\s*x\s*(\d+)\s*x\s*(\d+)", arg) + if match_list: + m_str = match_list.group(1) + n = int(match_list.group(2)) + k = int(match_list.group(3)) + try: + m_values = tuple(int(x.strip()) for x in m_str.split(",")) + l = len(m_values) # noqa: E741 + print(f"Parsed benchmark arg: N={n}, K={k}, L={l}") + print(f"M values per group: {m_values}") + return ((n, k, l), m_values) + except ValueError: + raise argparse.ArgumentTypeError(f"Invalid integer list in benchmark argument: {arg}") + + parts = arg.split("x") + if len(parts) == 4: + try: + m, n, k, l = [int(x.strip()) for x in parts] # noqa: E741 + m_values = tuple([m] * l) + print(f"Parsed benchmark arg: M={m}, N={n}, K={k}, L={l}") + return ((n, k, l), m_values) + except ValueError: + pass + + raise argparse.ArgumentTypeError(f"Invalid benchmark argument format. Got: {arg}") + + +def main(): + """Main entry point for running the kernel.""" + parser = argparse.ArgumentParser( + description="Rubin (SM107) BlockScaled Contiguous Grouped GEMM Finalize Fusion kernel." + ) + + parser.add_argument( + "--nkl", + type=parse_comma_separated_ints, + default=(256, 512, 1), + help="nkl dimensions: N, K, L (comma-separated)", + ) + + parser.add_argument( + "--benchmark", + type=str, + default=None, + help="Path to benchmark file or 'MxNxKxL' or '[m0,m1,...]xNxK'.", + ) + + parser.add_argument( + "--permuted_m", + type=int, + default=None, + help="Optional padded M dimension for CUDA graph support.", + ) + + parser.add_argument( + "--seq_len", + type=int, + default=4096, + help="Sequence length for MoE, used by fused finalize.", + ) + + parser.add_argument( + "--topk", + type=int, + default=8, + help="Top-K experts per token (used for fused finalize)", + ) + + parser.add_argument( + "--mma_inst_shape", + type=parse_comma_separated_ints, + default=(256, 256, 128), + help="MMA instruction shape M, N, K (comma-separated).", + ) + + parser.add_argument( + "--mma_tiler", + type=parse_comma_separated_ints, + default=(256, 256, 256), + help="MMA tile shape M, N, K (comma-separated).", + ) + + parser.add_argument( + "--cluster_shape_mn", + type=parse_comma_separated_ints, + default=(2, 1), + help="Cluster shape (comma-separated)", + ) + + parser.add_argument("--a_dtype", type=cutlass.dtype, default=cutlass.Float4E2M1FN) + parser.add_argument("--b_dtype", type=cutlass.dtype, default=cutlass.Float4E2M1FN) + parser.add_argument("--c_dtype", type=cutlass.dtype, default=cutlass.BFloat16) + parser.add_argument("--sf_dtype", type=cutlass.dtype, default=cutlass.Float8E4M3FN) + parser.add_argument("--sf_vec_size", type=int, default=16) + parser.add_argument("--a_major", choices=["k"], type=str, default="k") + parser.add_argument("--b_major", choices=["k"], type=str, default="k") + parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n") + parser.add_argument("--tolerance", type=float, default=1e-01) + parser.add_argument("--warmup_iterations", type=int, default=0) + parser.add_argument("--iterations", type=int, default=1) + parser.add_argument("--skip_ref_check", action="store_true") + parser.add_argument("--use_cold_l2", action="store_true", default=False) + parser.add_argument("--raster_along_m", action="store_true", default=False) + parser.add_argument( + "--use_cupti", action="store_true", default=False, help="Use Cupti profiler" + ) + args = parser.parse_args() + + if args.benchmark: + if os.path.isfile(args.benchmark): + nkl, group_m_list = read_benchmark_file(args.benchmark) + else: + nkl, group_m_list = parse_benchmark_arg(args.benchmark) + else: + parser.error("No benchmark file or benchmark argument provided") + + if len(args.mma_inst_shape) != 3: + parser.error("--mma_inst_shape must contain exactly 3 values") + + if len(args.mma_tiler) != 3: + parser.error("--mma_tiler must contain exactly 3 values") + + if len(args.cluster_shape_mn) != 2: + parser.error("--cluster_shape_mn must contain exactly 2 values") + + exec_time = run( + nkl, + group_m_list, + args.a_dtype, + args.b_dtype, + args.c_dtype, + args.sf_dtype, + args.sf_vec_size, + args.a_major, + args.b_major, + args.c_major, + args.mma_inst_shape, + args.mma_tiler, + args.cluster_shape_mn, + args.tolerance, + args.warmup_iterations, + args.iterations, + args.skip_ref_check, + args.use_cold_l2, + args.permuted_m, + args.topk, + args.seq_len, + args.raster_along_m, + args.use_cupti, + ) + print("exec_time: ", exec_time) + print("PASS") + + +if __name__ == "__main__": + main() diff --git a/tensorrt_llm/_torch/locality_domain/policy.py b/tensorrt_llm/_torch/locality_domain/policy.py index 87cfda0d028d..5b612ce8ee25 100644 --- a/tensorrt_llm/_torch/locality_domain/policy.py +++ b/tensorrt_llm/_torch/locality_domain/policy.py @@ -254,6 +254,7 @@ def plan_moe( moe_backend: str = "CUTEDSL", use_fused_finalize: bool = True, dtype_activation: torch.dtype = torch.bfloat16, + activation: str = "Swiglu", ) -> PartitionPlan: """Decide whether to partition a MoE GroupGemm for locality domain execution. @@ -305,6 +306,14 @@ def plan_moe( enabled=False, reason_if_disabled="locality domain MoE only supports NVFP4 or BF16" ) + # Both locality-domain MoE kernels fuse SwiGLU. Staying unpartitioned is + # the right answer for any other activation, not an error. + if activation != "Swiglu": + return PartitionPlan( + enabled=False, + reason_if_disabled=f"locality domain MoE fuses SwiGLU only, got {activation}", + ) + if op_name not in self.policy.allowed_ops: return PartitionPlan(enabled=False, reason_if_disabled=f"{op_name} not in allowed_ops") diff --git a/tensorrt_llm/_torch/locality_domain_utils.py b/tensorrt_llm/_torch/locality_domain_utils.py index b16e4d7f0987..ce7127ca1141 100644 --- a/tensorrt_llm/_torch/locality_domain_utils.py +++ b/tensorrt_llm/_torch/locality_domain_utils.py @@ -610,3 +610,20 @@ def get_locality_domain_mempool(locality_domain_id: int) -> torch.cuda.MemPool: ) return manager.mempools[pool_key] + + +def _copy_to_new_cuda_allocation(tensor: torch.Tensor) -> torch.Tensor: + """Copy ``tensor`` into a fresh contiguous allocation on the current CUDA device. + + The allocation follows the allocator selected by the surrounding context, + including ``torch.cuda.use_mem_pool``. Unlike ``contiguous().cuda()``, this + always allocates new storage when ``tensor`` is already contiguous and on + the current CUDA device. + """ + copied = torch.empty_like( + tensor, + device="cuda", + memory_format=torch.contiguous_format, + ) + copied.copy_(tensor) + return copied diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 74b6757628bd..d415e4dd6393 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -29,6 +29,7 @@ import transformers from transformers.utils import HF_MODULES_CACHE +from tensorrt_llm._torch.locality_domain.policy import LocalityDomainPolicy from tensorrt_llm._torch.pyexecutor.config_utils import ( get_kimi_linear_num_attention_layers, get_qwen3_hybrid_num_attention_layers, is_kimi_linear, is_nemotron_hybrid, is_qwen3_hybrid, load_pretrained_config) @@ -277,6 +278,10 @@ class ModelConfig(Generic[TConfig]): use_cute_dsl_bf16_bmm: bool = False use_cute_dsl_bf16_gemm: bool = False + # locality domain execution policy (controls partitioned linear/MoE execution) + locality_domain_policy: LocalityDomainPolicy = field( + default_factory=LocalityDomainPolicy) + _frozen: bool = field(default=False, init=False, repr=False) # If true, ONLY the vision encoder part of the full model is loaded/executed. @@ -351,6 +356,8 @@ def get_all_reduce_strategy(strategy: str = "AUTO"): if self.moe_max_num_tokens is None: self.moe_max_num_tokens = self.max_num_tokens * self.mapping.dp_size + self.extra_attrs["locality_domain_policy"] = self.locality_domain_policy + def is_moe_max_num_tokens_default(self) -> bool: """Whether ``moe_max_num_tokens`` was derived rather than configured. diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index cdcf46f892ee..4d137996dd48 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -325,10 +325,10 @@ Each backend's `can_implement(p, d)` classmethod declares what it supports. Sour | Quantization | Cutlass | TRTLLMGen | DeepGemm | DenseGEMM | CuteDSL | MegaMoE-DG | MegaMoE-CuteDSL | Triton | Marlin | Vanilla | |---|---|---|---|---|---|---|---|---|---|---| -| Unquantized (BF16/FP16) | Y (SM80+) | Y (SM100/103, BF16, needs FlashInfer `trtllm_bf16_moe`)§ | N | N | N | N | N | Y (SM90, BF16) | N | Y | +| Unquantized (BF16/FP16) | Y (SM80+) | Y (SM100/103, BF16, needs FlashInfer `trtllm_bf16_moe`)§ | N | N | Y (SM107, BF16, SwiGLU only)¶ | N | N | Y (SM90, BF16) | N | Y | | FP8 QDQ | Y (SM89+) | N | N | N | N | N | N | Y (SM90) | N | Y | | FP8 Block Scales | Y (SM90, SM120) | Y (SM100/103) | Y (SM100/103) | N | N‡ | N | N | N | N | Y | -| NVFP4 | Y (SM100/103/120/121) | Y (SM100/103) | N | Y (SM100/103) | Y (SM100/103/120/121) | N | Y (SM100/103, cu13 cutlass-dsl + NVSHMEM provider; per-expert alpha/norm_const + SwiGLU clamp) | N | Y (SM89-SM99) | Y | +| NVFP4 | Y (SM100/103/120/121) | Y (SM100/103) | N | Y (SM100/103) | Y (SM100/103/107/120/121)¶ | N | Y (SM100/103, cu13 cutlass-dsl + NVSHMEM provider; per-expert alpha/norm_const + SwiGLU clamp) | N | Y (SM89-SM99) | Y | | W4A16 NVFP4 | Y (SM80+, dequant-on-the-fly) | N | N | N | Y (SM120/121 via `CuteDslB12xFusedMoE`, needs flashinfer) | N | N | N | Y (SM89-SM99, BF16) | Y | | W4A8 NVFP4 FP8 | N | Y (SM100/103) | N | N | N | N | N | N | N | N | | W4A16 MXFP4 | Y (SM90) | Y (SM100/103) | N | N | N | N | N | Y (SM90) | N | N | @@ -349,6 +349,23 @@ pre-resolver code raised `RuntimeError` instead. The same path also requires `intermediate_size_per_partition % 128 == 0` (`Bf16MoeLauncher::check_moe`); a non-aligned shard is `SHAPE_UNALIGNED` and falls back to Cutlass. +¶ `CuteDslFusedMoE` on SM107 needs `MoEDep.CUTEDSL_RUBIN` (the installed CuTe DSL +exposes the Rubin helpers) and carries three constraints the other SMs do not: + +- **Fused finalize is mandatory.** There is no unfused FC2 — NVFP4 has no plain + grouped GEMM there, and the BF16 op fuses finalize unconditionally. Disabling + finalize fusion, explicitly or by configuring LoRA, is + `FINALIZE_FUSION_REQUIRED` rather than a late `NotImplementedError`. +- **Unquantized is SwiGLU-only.** `cute_dsl_bf16_gather_grouped_gemm_swiglu_rubin` + takes no activation argument, so any other activation is + `ACTIVATION_UNSUPPORTED`. NVFP4 forwards `activation_type` and serves Relu2 too. +- **Locality domain excludes EPLB and DWDP.** Localized weight shards are built + once from the loaded weights, so they cannot follow expert migration + (`EPLB_UNSUPPORTED`) or parameter rebinding. DWDP is not an error: with + `uses_locality_domain` true, `_should_enable_dwdp` simply returns False. + Both locality-domain kernels also fuse SwiGLU, so `plan_moe` declines any other + activation: an NVFP4 Relu2 layer runs unpartitioned rather than being rejected. + Cutlass covers `W4A16 NVFP4` on a wider SM range than plain `NVFP4` because the two run different kernels: `W4A16NVFP4CutlassFusedMoEMethod` dequantizes the FP4 weights into the activation dtype each forward and then calls the unquantized diff --git a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py index 47c3a07f221b..d02bbf22a9c3 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py @@ -442,6 +442,11 @@ def _should_enable_dwdp(self) -> bool: if not self.backend.capabilities.supports_dwdp: return False + # DWDP rebinds the backend parameters, which would strand the localized + # weight shards. Not enabling it is the correct outcome, not an error. + if self.backend.uses_locality_domain: + return False + quant_config = getattr(self.backend, "quant_config", None) if quant_config is None: quant_config = getattr(self.model_config, "quant_config", None) diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py index 8a890047fde8..0a956593a8be 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py @@ -20,13 +20,21 @@ import torch import torch.nn.functional as F -from tensorrt_llm._utils import is_sm_100f +from tensorrt_llm._utils import get_sm_version, is_sm_100f from tensorrt_llm.models.modeling_utils import QuantAlgo from ...autotuner import (AutoTuner, ConstraintSpec, DynamicTensorSpec, OptimizationProfile, TunableRunner, TuningConfig) from ...custom_ops.cute_dsl_custom_ops import GroupedGemmInputsHelper -from ...cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE +from ...cute_dsl_utils import (IS_CUTLASS_DSL_AVAILABLE, + IS_CUTLASS_DSL_RUBIN_AVAILABLE) +from ...locality_domain.autotune import \ + LocalityDomainConcurrentTunableRunner as \ + _LocalityDomainConcurrentTunableRunner +from ...locality_domain.policy import LocalityDomainExecutionPlanner +from ...locality_domain.runtime import LocalityDomainRuntime +from ...locality_domain_utils import (_copy_to_new_cuda_allocation, + get_reserved_remainder_stream) from ...model_config import ModelConfig from ...utils import (ActivationType, AuxStreamType, EventType, Fp4QuantizedTensor, @@ -40,8 +48,10 @@ MoEStaticCapability, nvfp4_fc1_row_alignment_rejection, require_comm_plan) +from .impl_environment import MoEDep from .interface import _reject -from .quantization import MoEWeightLoadingMode, NVFP4CuteDslFusedMoEMethod +from .quantization import (BF16CuteDslFusedMoEMethod, MoEWeightLoadingMode, + NVFP4CuteDslFusedMoEMethod) from .routing import BaseMoeRoutingMethod # These runners are defined inside cute_dsl_custom_ops' ``if @@ -65,6 +75,36 @@ Sm100BlockScaledContiguousGatherGroupedGemmActFusionRunner) +def _unwrap_locality_domain_runner(runner: TunableRunner) -> TunableRunner: + """Return the kernel runner wrapped by the shared locality domain tuning adapter.""" + if isinstance(runner, _LocalityDomainConcurrentTunableRunner): + return runner.op_runner + return runner + + +def _runner_tactics_match_tile_size( + comb: List[Tuple[TunableRunner, Any]], + outer_runner_type: type, + checked_runner_types: Tuple[type, ...], +) -> bool: + """Check that nested kernel tactics use the outer MoE tile size.""" + tile_size = None + for runner, tactic in comb: + runner = _unwrap_locality_domain_runner(runner) + if isinstance(runner, outer_runner_type): + tile_size = tactic + if tile_size is None: + return True + + for runner, tactic in comb: + runner = _unwrap_locality_domain_runner(runner) + if isinstance(runner, checked_runner_types): + mma_tiler_mn, *_ = tactic + if mma_tiler_mn[0] != tile_size: + return False + return True + + @dataclass class NvFp4WeightView: """Bundles all NVFP4 weight tensors for MoE computation. @@ -267,7 +307,8 @@ def __init__(self, enable_finalize_fusion: bool = True, enable_alltoall: bool = False, output_dtype: torch.dtype = torch.bfloat16, - scaling_vector_size: int = 16): + scaling_vector_size: int = 16, + workload_identity: Optional[Tuple] = None): super().__init__() self.forward_impl = forward_impl self.num_experts = num_experts @@ -280,9 +321,10 @@ def __init__(self, assert output_dtype == torch.bfloat16 self.output_dtype = output_dtype self.scaling_vector_size = scaling_vector_size + self.workload_identity = workload_identity def unique_id(self): - return ( + identity = ( self.num_experts, self.top_k, self.num_local_experts, @@ -292,6 +334,9 @@ def unique_id(self): self.output_dtype, self.scaling_vector_size, ) + if self.workload_identity is not None: + identity += (self.workload_identity, ) + return identity def get_valid_tactics( self, @@ -299,6 +344,13 @@ def get_valid_tactics( profile: OptimizationProfile, **kwargs, ) -> List[int]: + return self._tile_sizes() + + @staticmethod + def _tile_sizes() -> List[int]: + # tile_size=512 is only supported on Rubin (SM107). + if get_sm_version() == 107: + return [128, 256, 512] return [128, 256] def get_tuning_config(self) -> TuningConfig: @@ -325,8 +377,24 @@ def get_tuning_config(self) -> TuningConfig: ) return self.__class__.tuning_config_cache[key] - def forward(self, inputs: List[torch.Tensor], - tactic: Optional[int]) -> torch.Tensor: + def forward(self, + inputs: List[torch.Tensor], + tactic: Optional[int], + do_preparation: bool = False) -> torch.Tensor: + if do_preparation: + if self.workload_identity is not None: + # Inner FC tuning cannot run from inside the CUDA graph used to + # profile an outer tile. Prime every tile's FC1/FC2 cache for + # this optimization profile before outer profiling starts. + for tile_size in self._tile_sizes(): + self.forward_impl( + *inputs, + enable_alltoall=self.enable_alltoall, + tile_size=tile_size, + overlap_moe_output_memset=False, + ) + return inputs[4] + if isinstance(tactic, int) and tactic > 0: tile_size = tactic else: @@ -339,27 +407,176 @@ def forward(self, inputs: List[torch.Tensor], @staticmethod def runner_tactic_comb_checker( comb: List[Tuple[TunableRunner, Any]]) -> bool: - tile_size = None - for runner, tactic in comb: - if isinstance(runner, CuteDslFusedMoENvfp4Runner): - tile_size = tactic - if tile_size is None: - return True - - # Imported here rather than at module scope: these runners are defined - # inside cute_dsl_custom_ops' ``if IS_CUTLASS_DSL_AVAILABLE:`` block, - # which has no else-branch, so at module scope a missing cutlass DSL - # would break every importer of this file -- and create_moe imports it - # eagerly under _torch.models, so that reaches all model startup rather - # than just this backend. Reaching this line means a CuteDSL runner is - # already being tuned, so the DSL is installed. - - for runner, tactic in comb: - if isinstance(runner, _TILE_SIZE_CHECKED_RUNNERS): - mma_tiler_mn, *_ = tactic - if mma_tiler_mn[0] != tile_size: - return False - return True + checked_runner_types = list(_TILE_SIZE_CHECKED_RUNNERS) + if IS_CUTLASS_DSL_RUBIN_AVAILABLE: + from ...custom_ops.cute_dsl_custom_ops import ( + Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner, + Sm107BlockScaledContiguousGroupedGemmFinalizeFusionRunner) + checked_runner_types.extend([ + Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner, + Sm107BlockScaledContiguousGroupedGemmFinalizeFusionRunner, + ]) + + return _runner_tactics_match_tile_size( + comb, + CuteDslFusedMoENvfp4Runner, + tuple(checked_runner_types), + ) + + +class CuteDslFusedMoEBF16InputsHelper(GroupedGemmInputsHelper): + """Helper for CuteDSL BF16 MoE input preprocessing and autotuning.""" + + def __init__(self, num_experts: int, top_k: int, num_local_experts: int, + local_expert_offset: int): + self.num_experts = num_experts + self.top_k = top_k + self.num_local_experts = num_local_experts + self.local_expert_offset = local_expert_offset + + def infer_shape_num_tokens(self, input_shapes: List[torch.Size]) -> int: + return input_shapes[0][0] + + def inputs_pre_hook(self, inputs: List[torch.Tensor]) -> List[torch.Tensor]: + x, token_selected_experts, *others = inputs + num_tokens = token_selected_experts.size(0) + num_tokens_per_expert = self.generate_num_tokens_per_expert( + num_tokens, approx_max_load=True) + + new_token_selected_experts = [] + for i, curr_num_tokens in enumerate(num_tokens_per_expert, + start=self.local_expert_offset): + new_token_selected_experts.extend([i] * curr_num_tokens) + new_token_selected_experts = new_token_selected_experts + [-1] * ( + num_tokens * self.top_k - len(new_token_selected_experts)) + new_token_selected_experts = torch.tensor( + new_token_selected_experts, + dtype=token_selected_experts.dtype, + device=token_selected_experts.device) + new_token_selected_experts = new_token_selected_experts.view( + self.top_k, num_tokens).transpose(0, 1).contiguous() + return x, new_token_selected_experts, *others + + +class CuteDslFusedMoEBF16Runner(TunableRunner): + """Autotuner runner for BF16/FP16 MoE on Rubin (SM107). + + Selects tile_size from {64, 128, 256} and delegates to run_moe_bf16_impl. + """ + tuning_config_cache = dict() + + def __init__(self, + forward_impl: Callable, + num_experts: int, + top_k: int, + num_local_experts: int, + local_expert_offset: int, + enable_alltoall: bool = False, + output_dtype: torch.dtype = torch.bfloat16, + workload_identity: Optional[Tuple] = None): + super().__init__() + self.forward_impl = forward_impl + self.num_experts = num_experts + self.top_k = top_k + self.num_local_experts = num_local_experts + self.local_expert_offset = local_expert_offset + self.enable_alltoall = enable_alltoall + self.output_dtype = output_dtype + self.workload_identity = workload_identity + + def unique_id(self): + identity = ( + self.num_experts, + self.top_k, + self.num_local_experts, + self.local_expert_offset, + self.enable_alltoall, + self.output_dtype, + ) + if self.workload_identity is not None: + identity += (self.workload_identity, ) + return identity + + def get_valid_tactics( + self, + inputs: List[torch.Tensor], + profile: OptimizationProfile, + **kwargs, + ) -> List[int]: + return self._tile_sizes() + + @staticmethod + def _tile_sizes() -> List[int]: + return [64, 128, 256] + + def get_tuning_config(self) -> TuningConfig: + key = self.unique_id() + if key not in self.__class__.tuning_config_cache: + helper = CuteDslFusedMoEBF16InputsHelper(self.num_experts, + self.top_k, + self.num_local_experts, + self.local_expert_offset) + # BF16 inputs: [x, token_selected_experts, token_final_scales, + # moe_output] + self.__class__.tuning_config_cache[key] = TuningConfig( + dynamic_tensor_specs=(DynamicTensorSpec( + 0, 0, get_last_power_of_2_num_tokens_buckets, + last_positive_power_of_2), ), + constraint_specs=( + ConstraintSpec(1, 0, helper.infer_shape_num_tokens), + ConstraintSpec(2, 0, helper.infer_shape_num_tokens), + ConstraintSpec(3, 0, helper.infer_shape_num_tokens), + ), + inputs_pre_hook=helper.inputs_pre_hook, + use_cold_l2_cache=True, + ) + return self.__class__.tuning_config_cache[key] + + def forward(self, + inputs: List[torch.Tensor], + tactic: Optional[int], + do_preparation: bool = False) -> torch.Tensor: + if do_preparation: + if self.workload_identity is not None: + # See the NVFP4 runner: nested FC tuning must complete before + # the outer tile is profiled under CUDA graph capture. + for tile_size in self._tile_sizes(): + self.forward_impl( + *inputs, + enable_alltoall=self.enable_alltoall, + tile_size=tile_size, + overlap_moe_output_memset=False, + ) + return inputs[3] + + if isinstance(tactic, int) and tactic > 0: + tile_size = tactic + else: + tile_size = 128 + return self.forward_impl(*inputs, + enable_alltoall=self.enable_alltoall, + tile_size=tile_size) + + @AutoTuner.TacticsCapture.register_runner_tactic_comb_checker + @staticmethod + def runner_tactic_comb_checker( + comb: List[Tuple[TunableRunner, Any]]) -> bool: + # BF16 GEMM runners that need CTA_M == tile_size. + checked_runner_types = [] + if IS_CUTLASS_DSL_RUBIN_AVAILABLE: + from ...custom_ops.cute_dsl_custom_ops import ( + Sm107ContiguousGatherGroupedGemmSwigluFusionRunner, + Sm107ContiguousGroupedGemmFinalizeFusionRunner) + checked_runner_types.extend([ + Sm107ContiguousGatherGroupedGemmSwigluFusionRunner, + Sm107ContiguousGroupedGemmFinalizeFusionRunner, + ]) + + return _runner_tactics_match_tile_size( + comb, + CuteDslFusedMoEBF16Runner, + tuple(checked_runner_types), + ) class CuteDslFusedMoE(MoEImplBase): @@ -397,6 +614,52 @@ class CuteDslFusedMoE(MoEImplBase): limit_when_absent=float("inf"), ) + def _has_moe_output_memset_aux_stream(self) -> bool: + event_dict = getattr(self, 'event_dict', None) + aux_stream_dict = getattr(self, 'aux_stream_dict', None) + return (event_dict is not None and aux_stream_dict is not None + and EventType.Main in event_dict + and EventType.MoeOutputMemset in event_dict + and AuxStreamType.MoeOutputMemset in aux_stream_dict) + + def _get_reserved_moe_output_memset_stream( + self) -> Optional[torch.cuda.Stream]: + """Resolve and cache the strict split's remainder stream before capture.""" + if hasattr(self, "_cached_reserved_moe_output_memset_stream"): + return self._cached_reserved_moe_output_memset_stream + + stream = None + if self._locality_domain_runtime is not None: + stream = get_reserved_remainder_stream() + self._cached_reserved_moe_output_memset_stream = stream + return stream + + def _moe_output_memset_run_stream(self) -> torch.cuda.Stream: + """Select the remainder stream when present, otherwise the aux stream.""" + remainder_stream = self._get_reserved_moe_output_memset_stream() + if remainder_stream is not None: + return remainder_stream + return self.aux_stream_dict[AuxStreamType.MoeOutputMemset] + + def _locality_domain_workload_identity( + self, input_dtype: torch.dtype) -> Tuple[Any, ...]: + """Describe the sharded MoE workload and its compute topology.""" + if self._locality_domain_runtime is None or self._locality_domain_weight_shards is None: + raise RuntimeError( + "locality domain workload identity requires initialized shards") + shard_identity = tuple(( + tuple(shard['w3_w1_weight'].shape), + str(shard['w3_w1_weight'].dtype), + tuple(shard['w2_weight'].shape), + str(shard['w2_weight'].dtype), + ) for shard in self._locality_domain_weight_shards) + return ( + self._locality_domain_plan.num_partitions, + self._locality_domain_runtime.topology_identity(), + str(input_dtype), + shard_identity, + ) + @classmethod def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: """CuteDSL grouped GEMM: NVFP4 on SM100/SM103, bfloat16 activations.""" @@ -417,11 +680,6 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: f"CuteDslFusedMoE only supports bfloat16 activation (output is hardcoded to bfloat16), " f"got {p.dtype_act}") - # CuteDslFusedMoE does NOT support unquantized mode - if quant_algo is None: - return _reject(MoERejectReason.QUANT_UNSUPPORTED, - "CuteDslFusedMoE does not support unquantized mode") - # CuteDslFusedMoE does NOT support swiglu_gptoss_style if p.swiglu_gptoss_style: return _reject( @@ -429,12 +687,53 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: "CuteDslFusedMoE does not support swiglu_gptoss_style (bias/swiglu with custom alpha/beta/limit)" ) - # NVFP4 - SM in {100, 103} + # Localized weight shards are built once from the loaded weights, so + # they cannot follow EPLB expert migration. + if (d.locality_domain_requested + and d.env.has_dep(MoEDep.LOCALITY_DOMAIN) and d.eplb_enabled): + return _reject( + MoERejectReason.EPLB_UNSUPPORTED, + "locality domain MoE cannot follow EPLB expert migration") + + # SM107 has no unfused FC2: NVFP4 has no plain grouped GEMM there, and + # the BF16 op always fuses finalize. + if sm_version == 107 and not d.fused_finalize_enabled: + return _reject( + MoERejectReason.FINALIZE_FUSION_REQUIRED, + "CuteDslFusedMoE on SM107 only has a fused-finalize FC2") + + if quant_algo is None: + if sm_version != 107: + return _reject( + MoERejectReason.SM_UNSUPPORTED, + f"Unquantized CuteDSL MoE requires SM107, got SM{sm_version}" + ) + if not d.env.has_dep(MoEDep.CUTEDSL_RUBIN): + return _reject( + MoERejectReason.DEP_MISSING, + "Unquantized CuteDSL MoE on SM107 requires Rubin support in CuTe DSL" + ) + # The BF16 FC1 op fuses SwiGLU by name and takes no activation + # argument, unlike its NVFP4 counterpart. + if p.activation != "Swiglu": + return _reject( + MoERejectReason.ACTIVATION_UNSUPPORTED, + f"Unquantized CuteDSL MoE fuses SwiGLU only, got {p.activation}" + ) + return MoEEligibility.ok() + + # NVFP4 - SM in {100, 103, 107} if quant_algo == QuantAlgo.NVFP4: - if sm_version not in {100, 103}: + if sm_version not in {100, 103, 107}: return _reject( MoERejectReason.SM_UNSUPPORTED, - f"NVFP4 requires SM100 or SM103, got SM{sm_version}") + f"NVFP4 requires SM100, SM103, or SM107, got SM{sm_version}" + ) + if sm_version == 107 and not d.env.has_dep(MoEDep.CUTEDSL_RUBIN): + return _reject( + MoERejectReason.DEP_MISSING, + "NVFP4 CuteDSL MoE on SM107 requires Rubin support in CuTe DSL" + ) # process_weights_after_loading() unswizzles the FC1 block scales, # which asserts 128-row tiles; without this gate an unaligned shard # dies mid weight load with a bare swizzle error. @@ -495,8 +794,8 @@ def __init__( self.use_fused_finalize = (not model_config.moe_disable_finalize_fusion and model_config.lora_config is None) - # ``run_moe_nvfp4*`` overlaps the output memset on its own stream. This - # backend never chunks, so it needs no chunking stream. + # Output-memset overlap is independent of MoE chunking, so ensure its + # stream and events exist even if the parent creates no chunking event. if self.aux_stream_dict is None: self.aux_stream_dict = {} if AuxStreamType.MoeOutputMemset not in self.aux_stream_dict: @@ -508,9 +807,32 @@ def __init__( } self._weights_created = False + + self.scaling_vector_size = 16 + # locality domain: fork/join with _locality_domain kernel variants + shared output buffers. + # Weight splitting happens in post_load_weights after normal loading. + self._locality_domain_runtime = None + self._locality_domain_weight_shards = None # set in post_load_weights + planner = LocalityDomainExecutionPlanner( + model_config.locality_domain_policy) + self._locality_domain_plan = planner.plan_moe( + self.quant_config, + moe_backend=model_config.moe_backend, + use_fused_finalize=self.use_fused_finalize, + dtype_activation=self.dtype, + activation=ActivationType(self.activation_type).name, + ) + if self._locality_domain_plan.enabled: + self._locality_domain_runtime = LocalityDomainRuntime( + self._locality_domain_plan.num_partitions) if not model_config.skip_create_weights_in_init: self.create_weights() + def create_weights(self): + if self._weights_created: + return + super().create_weights() + def _build_local_weight_view(self) -> NvFp4WeightView: """Build the weight view from this backend's per-layer weights.""" return NvFp4WeightView( @@ -524,14 +846,27 @@ def _build_local_weight_view(self) -> NvFp4WeightView: slot_start=self.slot_start, ) + @property + def uses_locality_domain(self) -> bool: + return self._locality_domain_plan.enabled + def _get_quant_method(self): if self.quant_config is not None and self.quant_config.layer_quant_mode.has_any_quant( exclude_kv_cache=True): if self.quant_config.layer_quant_mode.has_nvfp4(): return NVFP4CuteDslFusedMoEMethod() - # ``can_implement`` admits NVFP4 only, so selection never lands here. - # Raise rather than fall back: any other method owns a weight layout - # these kernels cannot read. + elif get_sm_version() == 107 and IS_CUTLASS_DSL_RUBIN_AVAILABLE: + # Unquantized on SM107: the BF16 method interleaves FC1 weights for + # the fused gather + grouped GEMM + SwiGLU kernel, which serves no + # other activation. + if self.activation_type != ActivationType.Swiglu: + raise ValueError( + "Unquantized CuteDslFusedMoE fuses SwiGLU only, got " + f"{ActivationType(self.activation_type).name}") + return BF16CuteDslFusedMoEMethod() + # ``can_implement`` admits NVFP4, plus unquantized BF16 on SM107, so + # selection never lands here. Raise rather than fall back: any other + # method owns a weight layout these kernels cannot read. raise ValueError( f"CuteDslFusedMoE only supports NVFP4, got {self.quant_config}") @@ -544,7 +879,8 @@ def _check_configs(self): assert self.routing_method.top_k == 1, "Current walkaround only supports top-1 routing" def supports_moe_output_in_alltoall_workspace(self): - return self.has_nvfp4 + return self.has_nvfp4 or (not self.has_any_quant + and get_sm_version() == 107) def quantize_input(self, x: Union[torch.Tensor, Fp4QuantizedTensor], @@ -579,13 +915,21 @@ def quantize_input(self, # FP8 block scales doesn't support permutation of quantized inputs. # WAR: The quantization is in run_moe_fp8_block_scales. pass + elif not self.has_any_quant: + # Unquantized BF16/FP16: no quantization needed + pass else: raise ValueError( f"{self.__class__.__name__} doesn't support quantization mode {self.quant_config.quant_mode}." ) if x_sf is not None: - x_sf = x_sf.view(x_row, -1) + # ``view(0, -1)`` is ambiguous for an empty micro-batch. The + # scale width is fixed by the logical hidden size, so spell it + # out for both empty and non-empty inputs. + scale_cols = (self.hidden_size + self.scaling_vector_size - + 1) // self.scaling_vector_size + x_sf = x_sf.view(x_row, scale_cols) return x, x_sf def run_moe_nvfp4( @@ -610,8 +954,20 @@ def run_moe_nvfp4( """ assert self.has_nvfp4 assert weight_view is not None + if self.activation_type not in (ActivationType.Swiglu, + ActivationType.Relu2): + raise NotImplementedError( + "CuteDSL NVFP4 FC1 supports only SwiGLU and Relu2; " + f"got {self.activation_type.name}") output_dtype = torch.bfloat16 + use_locality_domain = self._locality_domain_runtime is not None + if use_locality_domain: + if self.activation_type != ActivationType.Swiglu: + raise NotImplementedError( + "Rubin locality domain NVFP4 MoE currently supports SwiGLU only" + ) + if moe_output is None: moe_output = torch.empty( (token_final_scales.size(0), self.hidden_size), @@ -622,9 +978,38 @@ def run_moe_nvfp4( self.hidden_size) assert moe_output.dtype == output_dtype + # Empty micro-batches are valid at the backend boundary. Avoid + # entering autotuning because its synthetic grouped-GEMM inputs + # require at least one output row. + if token_selected_experts.size(0) == 0: + return moe_output + effective_top_k = token_selected_experts.size(-1) - forward_impl = self.run_moe_nvfp4_impl + if use_locality_domain: + forward_impl = self._run_moe_nvfp4_locality_domain + workload_identity = self._locality_domain_workload_identity(x.dtype) + tuner_key = ( + "CuteDslFusedMoE::run_moe_nvfp4::locality_domain_end_to_end") + inputs = [ + x, + token_selected_experts, + token_final_scales, + x_sf, + moe_output, + ] + else: + forward_impl = self.run_moe_nvfp4_impl + workload_identity = None + tuner_key = "CuteDslFusedMoE::run_moe_nvfp4" + inputs = [ + x, + token_selected_experts, + token_final_scales, + x_sf, + moe_output, + weight_view, + ] tuner = AutoTuner.get() runner = CuteDslFusedMoENvfp4Runner( @@ -635,18 +1020,11 @@ def run_moe_nvfp4( local_expert_offset=weight_view.slot_start, enable_finalize_fusion=self.use_fused_finalize, enable_alltoall=enable_alltoall, + workload_identity=workload_identity, ) - inputs = [ - x, - token_selected_experts, - token_final_scales, - x_sf, - moe_output, - weight_view, - ] _, best_tactic = tuner.choose_one( - "CuteDslFusedMoE::run_moe_nvfp4", + tuner_key, [runner], runner.get_tuning_config(), inputs, @@ -666,6 +1044,9 @@ def run_moe_nvfp4_impl( ) -> torch.Tensor: """Non-DWDP NVFP4 MoE implementation using single-tensor ops.""" output_dtype = torch.bfloat16 + sm_version = get_sm_version() + use_rubin = (sm_version == 107 and IS_CUTLASS_DSL_RUBIN_AVAILABLE) + effective_top_k = token_selected_experts.size(1) esp = weight_view.expert_size_per_partition slot_start = weight_view.slot_start @@ -680,15 +1061,35 @@ def run_moe_nvfp4_impl( tile_tokens_dim=tile_size, ) - if self.use_fused_finalize: + has_aux_streams = self._has_moe_output_memset_aux_stream() + if self.use_fused_finalize and has_aux_streams: + memset_stream = self._moe_output_memset_run_stream() self.event_dict[EventType.Main].record() - moe_output.record_stream( - self.aux_stream_dict[AuxStreamType.MoeOutputMemset]) + moe_output.record_stream(memset_stream) + with torch.cuda.stream(memset_stream): + self.event_dict[EventType.Main].wait() + torch.ops.trtllm.moe_output_memset_inplace( + input=moe_output, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + tile_tokens_dim=tile_size, + top_k=effective_top_k, + ep_size=self.mapping.moe_ep_size, + enable_alltoall=enable_alltoall, + ) + self.event_dict[EventType.MoeOutputMemset].record() # Fused gather + GEMM + activation + quantize for FC1. # For gated (SwiGLU): weights are interleaved [up, gate], output is N/2. # For non-gated (Relu2): weights are plain, output is N. - x, x_sf = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell( + gather_act_op = ( + torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin + if use_rubin else torch.ops.trtllm. + cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell) + + gather_act_kwargs = dict( input=x.view(torch.float4_e2m1fn_x2), weight=weight_view.w3_w1_weight.view(torch.float4_e2m1fn_x2), input_scale=x_sf.view(torch.uint8), @@ -704,14 +1105,21 @@ def run_moe_nvfp4_impl( num_local_experts=esp, local_expert_offset=slot_start, tile_size=tile_size, - activation_type=self.activation_type, - swiglu_limit_scalar=self.act_clamp, ) + if use_rubin: + gather_act_kwargs["output_tensor"] = None + gather_act_kwargs["output_sf_tensor"] = None + else: + gather_act_kwargs["activation_type"] = self.activation_type + gather_act_kwargs["swiglu_limit_scalar"] = self.act_clamp + gather_act_kwargs["activation_type"] = self.activation_type + + x, x_sf = gather_act_op(**gather_act_kwargs) if self.use_fused_finalize: - with torch.cuda.stream( - self.aux_stream_dict[AuxStreamType.MoeOutputMemset]): - self.event_dict[EventType.Main].wait() + if has_aux_streams: + self.event_dict[EventType.MoeOutputMemset].wait() + else: torch.ops.trtllm.moe_output_memset_inplace( input=moe_output, tile_idx_to_mn_limit=tile_idx_to_mn_limit, @@ -723,10 +1131,15 @@ def run_moe_nvfp4_impl( ep_size=self.mapping.moe_ep_size, enable_alltoall=enable_alltoall, ) - self.event_dict[EventType.MoeOutputMemset].record() - self.event_dict[EventType.MoeOutputMemset].wait() - torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell( + # FC2: Grouped GEMM + Finalize (scatter-add) fusion + finalize_inplace_op = ( + torch.ops.trtllm. + cute_dsl_nvfp4_grouped_gemm_finalize_inplace_rubin + if use_rubin else torch.ops.trtllm. + cute_dsl_nvfp4_grouped_gemm_finalize_inplace_blackwell) + + finalize_inplace_op( input=x.view(torch.float4_e2m1fn_x2), weight=weight_view.w2_weight.view(torch.float4_e2m1fn_x2), input_scale=x_sf.view(torch.uint8), @@ -746,6 +1159,13 @@ def run_moe_nvfp4_impl( output_dtype=output_dtype, ) else: + if use_rubin: + # Rubin does not have a basic grouped GEMM kernel (without + # fused finalize) yet. Force use_fused_finalize=True for Rubin. + raise NotImplementedError( + "Rubin (SM107) MOE requires use_fused_finalize=True. " + "Basic grouped GEMM without finalize fusion is not yet " + "supported on Rubin.") x = torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_blackwell( input=x.view(torch.float4_e2m1fn_x2), weight=weight_view.w2_weight.view(torch.float4_e2m1fn_x2), @@ -769,6 +1189,474 @@ def run_moe_nvfp4_impl( ) return moe_output + def run_moe_bf16( + self, + x: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: Optional[torch.Tensor], + moe_output: Optional[torch.Tensor] = None, + enable_alltoall: bool = False, + ) -> torch.Tensor: + """Autotuner wrapper for BF16/FP16 MoE on Rubin (SM107).""" + assert not self.has_any_quant + # The FC2 op below always fuses finalize; ``can_implement`` declines + # SM107 when the caller disabled it, so honor that rather than ignore it. + assert self.use_fused_finalize, ( + "BF16 CuteDSL MoE has no unfused FC2 path") + output_dtype = x.dtype + effective_top_k = token_selected_experts.size(-1) + + if moe_output is None: + moe_output = torch.empty( + (token_selected_experts.size(0), self.hidden_size), + dtype=output_dtype, + device=x.device) + else: + assert moe_output.size() == (token_selected_experts.size(0), + self.hidden_size) + assert moe_output.dtype == output_dtype + + if token_selected_experts.size(0) == 0: + return moe_output + + self._ensure_bf16_alpha(x.device) + + use_locality_domain = self._locality_domain_runtime is not None + if use_locality_domain: + forward_impl = self._run_moe_bf16_locality_domain + workload_identity = self._locality_domain_workload_identity(x.dtype) + tuner_key = ( + "CuteDslFusedMoE::run_moe_bf16::locality_domain_end_to_end") + else: + forward_impl = self.run_moe_bf16_impl + workload_identity = None + tuner_key = "CuteDslFusedMoE::run_moe_bf16" + + tuner = AutoTuner.get() + runner = CuteDslFusedMoEBF16Runner( + forward_impl=forward_impl, + num_experts=self.num_slots, + top_k=effective_top_k, + num_local_experts=self.expert_size_per_partition, + local_expert_offset=self.slot_start, + enable_alltoall=enable_alltoall, + output_dtype=output_dtype, + workload_identity=workload_identity, + ) + + inputs = [x, token_selected_experts, token_final_scales, moe_output] + _, best_tactic = tuner.choose_one( + tuner_key, + [runner], + runner.get_tuning_config(), + inputs, + ) + return runner(inputs, tactic=best_tactic) + + def _ensure_bf16_alpha(self, device: torch.device) -> torch.Tensor: + if not hasattr(self, '_bf16_alpha') or self._bf16_alpha is None \ + or self._bf16_alpha.device != device \ + or self._bf16_alpha.size(0) != self.expert_size_per_partition: + self._bf16_alpha = torch.ones(self.expert_size_per_partition, + dtype=torch.float32, + device=device) + return self._bf16_alpha + + def run_moe_bf16_impl( + self, + x: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: Optional[torch.Tensor], + moe_output: torch.Tensor, + enable_alltoall: bool = False, + tile_size: int = 128, + ) -> torch.Tensor: + """BF16/FP16 MoE implementation using CuTE DSL Rubin kernels. + + FC1: gather + grouped GEMM + SwiGLU fusion + FC2: grouped GEMM + finalize (scatter-add) fusion + """ + output_dtype = x.dtype + effective_top_k = token_selected_experts.size(-1) + + # Step 1: moe_sort — identical to NVFP4 path + tile_idx_to_expert_idx, tile_idx_to_mn_limit, expanded_idx_to_permuted_idx, permuted_idx_to_expanded_idx, total_num_padded_tokens, num_non_exiting_tiles = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=self.num_slots, + top_k=effective_top_k, + local_expert_offset=self.slot_start, + local_num_experts=self.expert_size_per_partition, + tile_tokens_dim=tile_size, + ) + + # Step 2: Memset overlap for fused finalize + has_aux = self._has_moe_output_memset_aux_stream() + if has_aux: + self.event_dict[EventType.Main].record() + moe_output.record_stream(self._moe_output_memset_run_stream()) + + # Step 3: Alpha = 1.0 for all local experts (no quantization scaling) + # The wrapper initializes this before autotuner profiling so allocation + # does not happen inside CUDA graph capture. + self._ensure_bf16_alpha(x.device) + + # Step 4: FC1 — BF16 gather + grouped GEMM + SwiGLU + fc1_out = torch.ops.trtllm.cute_dsl_bf16_gather_grouped_gemm_swiglu_rubin( + input=x, + weight=self.w3_w1_weight, + alpha=self._bf16_alpha, + tile_idx_to_group_idx=tile_idx_to_expert_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + num_experts=self.num_slots, + top_k=effective_top_k, + num_local_experts=self.expert_size_per_partition, + local_expert_offset=self.slot_start, + tile_size=tile_size, + output_tensor=None, + partition_id=-1, + ) + + # Step 5: Memset overlap — zero out moe_output rows not touched by + # the finalize kernel (same pattern as NVFP4 path). + if has_aux: + with torch.cuda.stream(self._moe_output_memset_run_stream()): + self.event_dict[EventType.Main].wait() + torch.ops.trtllm.moe_output_memset_inplace( + input=moe_output, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + tile_tokens_dim=tile_size, + top_k=effective_top_k, + ep_size=self.mapping.moe_ep_size, + enable_alltoall=enable_alltoall, + ) + self.event_dict[EventType.MoeOutputMemset].record() + self.event_dict[EventType.MoeOutputMemset].wait() + else: + torch.ops.trtllm.moe_output_memset_inplace( + input=moe_output, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + tile_tokens_dim=tile_size, + top_k=effective_top_k, + ep_size=self.mapping.moe_ep_size, + enable_alltoall=enable_alltoall, + ) + + # Step 6: FC2 — BF16 grouped GEMM + finalize (scatter-add) inplace + torch.ops.trtllm.cute_dsl_bf16_grouped_gemm_finalize_inplace_rubin( + input=fc1_out, + weight=self.w2_weight, + output=moe_output, + tile_idx_to_group_idx=tile_idx_to_expert_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + token_final_scales=token_final_scales, + num_experts=self.num_slots, + top_k=effective_top_k, + num_local_experts=self.expert_size_per_partition, + local_expert_offset=self.slot_start, + tile_size=tile_size, + output_dtype=output_dtype, + ) + return moe_output + + def _run_moe_nvfp4_locality_domain( + self, + x: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: Optional[torch.Tensor], + x_sf: Optional[torch.Tensor] = None, + moe_output: Optional[torch.Tensor] = None, + enable_alltoall: bool = False, + tile_size: int = 128, + overlap_moe_output_memset: bool = True, + ) -> torch.Tensor: + """locality domain path: half-weight children, shared output buffers, fork/join. + + Each child holds localized half-N weights. Both partitions use the + same tuned tactic and write directly into their strided regions of the + shared FC1/FC2 output buffers. + """ + output_dtype = torch.bfloat16 + num_partitions = self._locality_domain_plan.num_partitions + shards = self._locality_domain_weight_shards + effective_top_k = token_selected_experts.size(-1) + + # --- moe_sort (shared, on main stream) --- + (tile_idx_to_expert_idx, tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles) = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=self.num_slots, + top_k=effective_top_k, + local_expert_offset=self.slot_start, + local_num_experts=self.expert_size_per_partition, + tile_tokens_dim=tile_size, + ) + + # --- Allocate shared output --- + if moe_output is None: + moe_output = torch.empty( + (token_selected_experts.size(0), self.hidden_size), + dtype=output_dtype, + device=x.device) + + # --- FC1: gather + grouped GEMM + SwiGLU, fork/join --- + # Each child has half-N weight [num_exp, inner_size, hidden_size]. + # Shared output buffers: + # c: [permute_m, inner_size] — each locality domain writes half via strided layout + # c_sf: [full_sf_size] — each locality domain writes at K-tile offset via full_c_shape + # Kernel uses full_c_shape to compute sfc layout with full M-tile stride, + # so no copy-back or interleave is needed. + m = permuted_idx_to_expanded_idx.size(0) + shard_weight_n = shards[0]['w3_w1_weight'].size(1) # half interleaved N + shard_interm = shard_weight_n // 2 # post-SwiGLU per partition + full_interm = shard_interm * num_partitions + fc1_out = torch.empty(m, + shard_interm // 2 * 2, + dtype=torch.float4_e2m1fn_x2, + device=x.device) + full_sf_size = m * full_interm // self.scaling_vector_size + fc1_out_sf = torch.empty(full_sf_size, + dtype=torch.uint8, + device=x.device) + + assert self.use_fused_finalize, ( + "locality domain MoE requires use_fused_finalize=True on Rubin") + + # Inner FC tactics are prepared before outer CUDA-graph profiling. Keep + # that first-use tuning on the main stream; normal execution still + # overlaps this memset with the already-prepared FC1 composite op. + memset_overlapped = (overlap_moe_output_memset + and self._has_moe_output_memset_aux_stream()) + if memset_overlapped: + memset_stream = self._moe_output_memset_run_stream() + self.event_dict[EventType.Main].record() + moe_output.record_stream(memset_stream) + with torch.cuda.stream(memset_stream): + self.event_dict[EventType.Main].wait() + torch.ops.trtllm.moe_output_memset_inplace( + input=moe_output, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + tile_tokens_dim=tile_size, + top_k=effective_top_k, + ep_size=self.mapping.moe_ep_size, + enable_alltoall=enable_alltoall, + ) + self.event_dict[EventType.MoeOutputMemset].record() + + torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_locality_domain_inplace_rubin( + input=x.view(torch.float4_e2m1fn_x2), + weight_0=shards[0]['w3_w1_weight'].view(torch.float4_e2m1fn_x2), + weight_1=shards[1]['w3_w1_weight'].view(torch.float4_e2m1fn_x2), + input_scale=x_sf.view(torch.uint8), + weight_scale_0=shards[0]['fc1_weight_block'].view(torch.uint8), + weight_scale_1=shards[1]['fc1_weight_block'].view(torch.uint8), + alpha=shards[0]['fc1_global'], + tile_idx_to_group_idx=tile_idx_to_expert_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + global_sf=shards[0]['fc2_input_scale'], + num_experts=self.num_slots, + top_k=effective_top_k, + num_local_experts=self.expert_size_per_partition, + local_expert_offset=self.slot_start, + tile_size=tile_size, + output_tensor=fc1_out, + output_sf_tensor=fc1_out_sf, + scaling_vector_size=self.scaling_vector_size, + activation_type=self.activation_type, + ) + + fc1_out_sf_merged = fc1_out_sf + + if memset_overlapped: + self.event_dict[EventType.MoeOutputMemset].wait() + else: + torch.ops.trtllm.moe_output_memset_inplace( + input=moe_output, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + tile_tokens_dim=tile_size, + top_k=effective_top_k, + ep_size=self.mapping.moe_ep_size, + enable_alltoall=enable_alltoall, + ) + + torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_finalize_locality_domain_inplace_rubin( + input=fc1_out.view(torch.float4_e2m1fn_x2), + weight_0=shards[0]['w2_weight'].view(torch.float4_e2m1fn_x2), + weight_1=shards[1]['w2_weight'].view(torch.float4_e2m1fn_x2), + input_scale=fc1_out_sf_merged.view(torch.uint8), + weight_scale_0=shards[0]['fc2_weight_block'].view(torch.uint8), + weight_scale_1=shards[1]['fc2_weight_block'].view(torch.uint8), + alpha=shards[0]['fc2_global'], + output=moe_output, + tile_idx_to_group_idx=tile_idx_to_expert_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + token_final_scales=token_final_scales, + num_experts=self.num_slots, + top_k=effective_top_k, + num_local_experts=self.expert_size_per_partition, + local_expert_offset=self.slot_start, + tile_size=tile_size, + output_dtype=output_dtype, + ep_size=self.mapping.moe_ep_size, + enable_alltoall=enable_alltoall, + scaling_vector_size=self.scaling_vector_size, + ) + + return moe_output + + def _run_moe_bf16_locality_domain( + self, + x: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: Optional[torch.Tensor], + moe_output: Optional[torch.Tensor] = None, + enable_alltoall: bool = False, + tile_size: int = 128, + overlap_moe_output_memset: bool = True, + ) -> torch.Tensor: + """locality domain path for unquantized BF16 MoE on Rubin.""" + output_dtype = x.dtype + num_partitions = self._locality_domain_plan.num_partitions + shards = self._locality_domain_weight_shards + effective_top_k = token_selected_experts.size(-1) + + (tile_idx_to_expert_idx, tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles) = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=self.num_slots, + top_k=effective_top_k, + local_expert_offset=self.slot_start, + local_num_experts=self.expert_size_per_partition, + tile_tokens_dim=tile_size, + ) + + if moe_output is None: + moe_output = torch.empty( + (token_selected_experts.size(0), self.hidden_size), + dtype=output_dtype, + device=x.device) + else: + assert moe_output.size() == (token_selected_experts.size(0), + self.hidden_size) + assert moe_output.dtype == output_dtype + + self._ensure_bf16_alpha(x.device) + + m = permuted_idx_to_expanded_idx.size(0) + shard_weight_n = shards[0]['w3_w1_weight'].size(1) + shard_interm = shard_weight_n // 2 + full_interm = shard_interm * num_partitions + fc1_out = torch.empty(m, + full_interm, + dtype=output_dtype, + device=x.device) + + assert self.use_fused_finalize, ( + "locality domain MoE requires use_fused_finalize=True on Rubin") + + memset_overlapped = (overlap_moe_output_memset + and self._has_moe_output_memset_aux_stream()) + if memset_overlapped: + memset_stream = self._moe_output_memset_run_stream() + self.event_dict[EventType.Main].record() + moe_output.record_stream(memset_stream) + with torch.cuda.stream(memset_stream): + self.event_dict[EventType.Main].wait() + torch.ops.trtllm.moe_output_memset_inplace( + input=moe_output, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + tile_tokens_dim=tile_size, + top_k=effective_top_k, + ep_size=self.mapping.moe_ep_size, + enable_alltoall=enable_alltoall, + ) + self.event_dict[EventType.MoeOutputMemset].record() + + torch.ops.trtllm.cute_dsl_bf16_gather_grouped_gemm_swiglu_locality_domain_inplace_rubin( + input=x, + weight_0=shards[0]['w3_w1_weight'], + weight_1=shards[1]['w3_w1_weight'], + alpha=self._bf16_alpha, + tile_idx_to_group_idx=tile_idx_to_expert_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + num_experts=self.num_slots, + top_k=effective_top_k, + num_local_experts=self.expert_size_per_partition, + local_expert_offset=self.slot_start, + tile_size=tile_size, + output_tensor=fc1_out, + ) + + if memset_overlapped: + self.event_dict[EventType.MoeOutputMemset].wait() + else: + torch.ops.trtllm.moe_output_memset_inplace( + input=moe_output, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + tile_tokens_dim=tile_size, + top_k=effective_top_k, + ep_size=self.mapping.moe_ep_size, + enable_alltoall=enable_alltoall, + ) + + torch.ops.trtllm.cute_dsl_bf16_grouped_gemm_finalize_locality_domain_inplace_rubin( + input=fc1_out, + weight_0=shards[0]['w2_weight'], + weight_1=shards[1]['w2_weight'], + output=moe_output, + tile_idx_to_group_idx=tile_idx_to_expert_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + token_final_scales=token_final_scales, + num_experts=self.num_slots, + top_k=effective_top_k, + num_local_experts=self.expert_size_per_partition, + local_expert_offset=self.slot_start, + tile_size=tile_size, + output_dtype=output_dtype, + ep_size=self.mapping.moe_ep_size, + enable_alltoall=enable_alltoall, + ) + + return moe_output + def run_moe_fp8_block_scales( self, x: torch.Tensor, @@ -862,7 +1750,7 @@ def run_moe( Run MoE computation with CuteDSL backend. This method encapsulates the core MoE computation logic, handling different - quantization schemes (fp8_block_scales and nvfp4). + quantization schemes (fp8_block_scales, nvfp4, and unquantized BF16). Returns: final_hidden_states tensor. @@ -895,6 +1783,13 @@ def run_moe( token_final_scales=token_final_scales, x_sf=x_sf, enable_alltoall=enable_alltoall) + elif not self.has_any_quant: + return self.run_moe_bf16( + x=x, + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + moe_output=moe_output, + enable_alltoall=enable_alltoall) else: raise ValueError( f"{self.__class__.__name__} doesn't support quantization mode {self.quant_config.quant_mode}." @@ -906,3 +1801,88 @@ def load_weights(self, allow_partial_loading: bool = False): super().load_weights(weights, allow_partial_loading=allow_partial_loading) + # Keep DWDP registration after base weight loading. This preserves + # loaded tensors for collector setup and remains compatible with the + # later locality domain post_load_weights splitting flow. + dwdp_handle_collector = getattr(self, "dwdp_handle_collector", None) + if dwdp_handle_collector is not None: + dwdp_handle_collector.register_weights(self) + + def post_load_weights(self): + super().post_load_weights() + # Split full weights into per-partition halves on localized memory + if self._locality_domain_runtime is not None: + self._locality_domain_weight_shards = self._split_weights_for_locality_domain( + ) + # Weight splitting initializes the process-lifetime locality domain resource. + # Resolve the borrowed remainder stream now, never during capture. + self._get_reserved_moe_output_memset_stream() + self._release_full_weights_after_locality_domain_split() + + def _release_full_weights_after_locality_domain_split(self): + """Release full tensors that are replaced by localized locality domain shards.""" + for param_name in ( + "w3_w1_weight", + "w2_weight", + "w3_w1_weight_scale", + "w2_weight_scale", + ): + param = getattr(self, param_name, None) + if param is None: + continue + setattr( + self, + param_name, + torch.nn.Parameter(param.new_empty(0), requires_grad=False), + ) + self.quant_method.setup_quant_scales(self) + + def _split_weights_for_locality_domain(self): + """Split full N-dimension weights into per-partition halves. + + After normal load_weights + post_load_weights, the full weights + are on self. Split them along dim=1 (N) and allocate halves on + each locality domain partition's localized memory. + """ + num_p = self._locality_domain_plan.num_partitions + shards = [] + for pid in range(num_p): + with self._locality_domain_runtime.partition_weight_context(pid): + n1 = self.w3_w1_weight.size(1) + n2 = self.w2_weight.size(1) + half_n1 = n1 // num_p + half_n2 = n2 // num_p + w3_w1_slice = self.w3_w1_weight[:, pid * half_n1:(pid + 1) * + half_n1] + w2_slice = self.w2_weight[:, pid * half_n2:(pid + 1) * half_n2] + shard = { + 'w3_w1_weight': _copy_to_new_cuda_allocation(w3_w1_slice), + 'w2_weight': _copy_to_new_cuda_allocation(w2_slice), + } + if self.has_nvfp4: + fc1_scale_slice = self.quant_scales.fc1_weight_block[:, + pid * + half_n1: + (pid + + 1) * + half_n1] + fc2_scale_slice = self.quant_scales.fc2_weight_block[:, + pid * + half_n2: + (pid + + 1) * + half_n2] + shard.update({ + 'fc1_weight_block': + _copy_to_new_cuda_allocation(fc1_scale_slice), + 'fc2_weight_block': + _copy_to_new_cuda_allocation(fc2_scale_slice), + 'fc1_global': + self.quant_scales.fc1_global, + 'fc2_global': + self.quant_scales.fc2_global, + 'fc2_input_scale': + self.fc2_input_scale, + }) + shards.append(shard) + return shards diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_blocks.py b/tensorrt_llm/_torch/moe/fused_moe/impl_blocks.py index 52de4d6031e7..67a91433e54e 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_blocks.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_blocks.py @@ -72,6 +72,11 @@ def supports_moe_output_in_alltoall_workspace(self) -> bool: """ return False + @property + def uses_locality_domain(self) -> bool: + """True when this impl runs partitioned across locality domains.""" + return False + def validate_configurable_moe(self, moe: "torch.nn.Module") -> None: """Backend-specific validation hook called by ``ConfigurableMoE``. diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py index cae34e45943c..98990b593df2 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py @@ -230,6 +230,12 @@ class MoEDeployment: eplb_enabled: bool = False # True only for routed-expert LoRA targets. moe_lora_enabled: bool = False + # False when ``moe_disable_finalize_fusion`` is set or any LoRA is + # configured: both need an unfused FC2 so a seam is left for the LoRA GEMM. + fused_finalize_enabled: bool = True + # ``model_config.locality_domain_policy.enabled``. Whether the machine can + # actually serve it is ``env.has_dep(MoEDep.LOCALITY_DOMAIN)``. + locality_domain_requested: bool = False @property def smart_router(self) -> bool: @@ -268,6 +274,9 @@ class MoERejectReason(str, Enum): # EPLB is registered for this layer and the impl cannot lay out slots for # it. Distinct from TOPOLOGY_UNSUPPORTED: the parallel sizes are fine. EPLB_UNSUPPORTED = "eplb_unsupported" + # The impl only has a fused-finalize FC2 epilogue, and the caller disabled + # finalize fusion (explicitly, or implicitly by configuring LoRA). + FINALIZE_FUSION_REQUIRED = "finalize_fusion_required" # Not a capability verdict: the impl could run, but the resolver refuses to # route production traffic there. Kept separate so that "we chose not to" # never reads as "it cannot". diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_environment.py b/tensorrt_llm/_torch/moe/fused_moe/impl_environment.py index c9811c266caa..26687af055bf 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_environment.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_environment.py @@ -40,6 +40,12 @@ class MoEDep(str, Enum): MEGAMOE_CUTEDSL_RUNTIME = "megamoe_cutedsl_runtime" #: The ``trtllm::cute_dsl_megamoe_nvfp4_*`` custom ops are registered. MEGAMOE_CUTEDSL_OP = "megamoe_cutedsl_op" + #: The installed CuTe DSL exposes the Rubin helpers the SM107 CuteDSL + #: kernels are written against. + CUTEDSL_RUBIN = "cutedsl_rubin" + #: Locality-domain execution is usable here: SM107, driver support, and + #: ``DISABLE_LOCALITY_DOMAINS`` unset. + LOCALITY_DOMAIN = "locality_domain" class MoEEnvFlag(str, Enum): @@ -94,6 +100,22 @@ def _probe_megamoe_cutedsl_runtime() -> Tuple[bool, str]: return bool(available), "" if available else str(reason) +def _probe_cutedsl_rubin() -> Tuple[bool, str]: + from ...cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE + + if IS_CUTLASS_DSL_RUBIN_AVAILABLE: + return True, "" + return False, "installed CuTe DSL lacks Rubin helpers" + + +def _probe_locality_domain() -> Tuple[bool, str]: + from ...locality_domain_utils import is_locality_domain_enabled + + if is_locality_domain_enabled(): + return True, "" + return False, "locality domain unsupported or disabled on this machine" + + def _probe_megamoe_cutedsl_op() -> Tuple[bool, str]: # Read the module because registration updates this flag after import. from ..custom_ops import cute_dsl_megamoe_custom_op as megamoe_op @@ -109,6 +131,8 @@ def _probe_megamoe_cutedsl_op() -> Tuple[bool, str]: MoEDep.DEEPGEMM_MEGAMOE: _probe_deepgemm_megamoe, MoEDep.MEGAMOE_CUTEDSL_RUNTIME: _probe_megamoe_cutedsl_runtime, MoEDep.MEGAMOE_CUTEDSL_OP: _probe_megamoe_cutedsl_op, + MoEDep.CUTEDSL_RUBIN: _probe_cutedsl_rubin, + MoEDep.LOCALITY_DOMAIN: _probe_locality_domain, } # Preserve prior defaults when environment variables are unset. diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py index 99972cbabf5d..bfe5ca692710 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py @@ -284,6 +284,7 @@ def build_moe_deployment( if num_slots is None: num_slots = num_experts if num_experts is not None else 0 lora_config = getattr(model_config, "lora_config", None) + locality_domain_policy = getattr(model_config, "locality_domain_policy", None) return MoEDeployment( ep_size=mapping.moe_ep_size, tp_size=mapping.moe_tp_size, @@ -297,6 +298,14 @@ def build_moe_deployment( eplb_enabled=eplb_enabled, # Routed-expert LoRA only; attention-only LoRA stays False. moe_lora_enabled=has_moe_lora_targets(lora_config), + # Same expression the impls use to set ``use_fused_finalize``; any + # LoRA counts, not just routed-expert LoRA. + fused_finalize_enabled=( + not getattr(model_config, "moe_disable_finalize_fusion", False) and lora_config is None + ), + locality_domain_requested=bool( + locality_domain_policy is not None and locality_domain_policy.enabled + ), ) diff --git a/tests/unittest/_torch/moe/moe_test_utils.py b/tests/unittest/_torch/moe/moe_test_utils.py index 7fbd836a2f85..b965700ffa5f 100644 --- a/tests/unittest/_torch/moe/moe_test_utils.py +++ b/tests/unittest/_torch/moe/moe_test_utils.py @@ -1411,6 +1411,7 @@ def should_skip_to_accelerate_ci( swiglu_gptoss_style: bool = False, parallel_mode: Optional[str] = None, activation_type: Optional[ActivationType] = ActivationType.Swiglu, + enable_locality_domains: bool = False, ) -> Optional[str]: """ Skip low-information-density test combinations to accelerate CI. @@ -1419,8 +1420,8 @@ def should_skip_to_accelerate_ci( all combinations run (local exhaustive testing). Rules applied (in order): - 0. Skip unquantized (quant=None) for most paths, but keep TRTLLM BF16 - unquantized coverage enabled. + 0. Skip unquantized (quant=None) for most paths, but keep TRTLLM BF16 and + locality domain CuteDSL BF16 coverage enabled. 0a. MARLIN backend: only NVFP4 on Ada/Hopper (SM89-SM99); skip all other quant_algo / architecture combinations. 1. e256 model: only DeepSeekV3 routing, bfloat16, seq=1, non-gptoss @@ -1438,6 +1439,7 @@ def should_skip_to_accelerate_ci( seq_len: Sequence length swiglu_gptoss_style: Whether using SwiGLU gptoss style parallel_mode: Multi-GPU parallel mode (None for single-GPU tests) + enable_locality_domains: Whether the test enables locality domain execution Returns: Skip reason string if test should be skipped for CI, None otherwise @@ -1449,11 +1451,16 @@ def should_skip_to_accelerate_ci( return None # --- Rule 0: Skip gated and unquantized (quant=None) for most backends --- - # Keep TRTLLM BF16 unquantized enabled to cover FlashInfer BF16 TRTLLM MoE. + # Keep TRTLLM BF16 for FlashInfer and CuteDSL BF16 with locality domain for the + # dedicated locality domain backend matrix. + keeps_unquantized_bf16_coverage = dtype == torch.bfloat16 and ( + backend_type == MoeBackendType.TRTLLM + or (backend_type == MoeBackendType.CUTEDSL and enable_locality_domains) + ) if ( quant_algo is None and is_gated_activation(activation_type) - and not (backend_type == MoeBackendType.TRTLLM and dtype == torch.bfloat16) + and not keeps_unquantized_bf16_coverage ): return "[CI accel] Skip unquantized (quant=None) in CI" diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index eb6fa8c55b79..d3f2caf4d56e 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -19,6 +19,7 @@ import logging import os from contextlib import contextmanager +from pathlib import Path from types import SimpleNamespace from typing import List, Optional, Tuple from unittest.mock import MagicMock @@ -41,8 +42,11 @@ from _torch.moe.quantize_utils import get_test_quant_params from transformers.configuration_utils import PretrainedConfig -from tensorrt_llm._torch.autotuner import AutoTuner, autotune +from tensorrt_llm._torch.autotuner import AutoTuner, OptimizationProfile, autotune from tensorrt_llm._torch.custom_ops.trtllm_gen_custom_ops import _select_explicit_fallback_tactic +from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_RUBIN_AVAILABLE +from tensorrt_llm._torch.locality_domain.policy import LocalityDomainPolicy +from tensorrt_llm._torch.locality_domain_utils import is_locality_domain_enabled from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.moe.fused_moe import ( DeepSeekV3MoeRoutingMethod, @@ -79,7 +83,11 @@ ) from tensorrt_llm._torch.moe.fused_moe.interface import MoE, MoESchedulerKind, MoEWeightLoadingMode from tensorrt_llm._torch.moe.fused_moe.mega_moe import MegaMoECuteDsl, MegaMoEDeepGemm -from tensorrt_llm._torch.moe.fused_moe.moe_resolution import impl_class_for, resolve_moe_impl +from tensorrt_llm._torch.moe.fused_moe.moe_resolution import ( + build_moe_deployment, + impl_class_for, + resolve_moe_impl, +) from tensorrt_llm._torch.moe.fused_moe.quantization import ( FusedMoEMethodBase, NVFP4FusedMoEMethod, @@ -275,10 +283,13 @@ def create_test_backend( swiglu_limit: Optional[torch.Tensor] = None, weight_loading_mode: MoEWeightLoadingMode = MoEWeightLoadingMode.VANILLA, activation_type: ActivationType = ActivationType.Swiglu, + locality_domain_policy: Optional[LocalityDomainPolicy] = None, n_shared_experts: int = 0, ) -> MoE: """Create a MoE backend for testing.""" backend_cls = get_backend_class(backend_type) + if locality_domain_policy is None: + locality_domain_policy = LocalityDomainPolicy(enabled=False) pretrained_config = PretrainedConfig() pretrained_config.num_experts = num_experts @@ -302,6 +313,7 @@ def create_test_backend( quant_config=quant_config, mapping=mapping, moe_backend=moe_backend_value, + locality_domain_policy=locality_domain_policy, ) if n_shared_experts > 0: # The shared-expert-fusion gate runs after the eager create_weights() @@ -328,9 +340,9 @@ def create_test_backend( return backend -# ============================================================================ +# ===================================================================== # Staged post-load hook lifecycle tests -# ============================================================================ +# ===================================================================== # These tests cover staged hook contracts rather than the common backend matrix # below. Keep them grouped so they can move to a dedicated file if the MoE test # layout is split later. @@ -518,9 +530,9 @@ def test_marlin_override_quant_config_degrades_per_layer(): assert report.degraded_from.reason is MoERejectReason.QUANT_UNSUPPORTED -# ============================================================================ +# ===================================================================== # TRTLLM-Gen SiTu backend contract -# ============================================================================ +# ===================================================================== # SiTu rides the generic SwiGLU geometry, so the host-side wiring is easy to # get wrong in ways no shape check catches: it reaches the cubin through the # same ``gemm1_alpha`` / ``gemm1_beta`` slots SwiGLU's constants use, and only @@ -1134,9 +1146,9 @@ def run_backend_moe( return backend.run_moe(MoERunContext(**args), workspace=workspace) -# ============================================================================ +# ===================================================================== # Test Parameters -# ============================================================================ +# ===================================================================== # Quantization algorithms to test QUANT_ALGOS_TO_TEST = [ @@ -1223,6 +1235,79 @@ def run_backend_moe( SWIGLU_COMBOS = CI_SWIGLU_COMBOS if IS_CI_MODE else LOCAL_SWIGLU_COMBOS +def should_skip_locality_domain_param( + backend_type: MoeBackendType, + quant_algo: Optional[QuantAlgo], + activation_type: ActivationType, + swiglu_gptoss_style: bool, + dtype: torch.dtype, +) -> Optional[str]: + """Return a static skip reason for locality domain MoE backend params.""" + if backend_type != MoeBackendType.CUTEDSL: + return "locality domain MoE backend test only supports CuteDSL" + if quant_algo not in (QuantAlgo.NVFP4, None): + return "locality domain MoE backend test only supports NVFP4 or BF16" + # plan_moe only enables the unquantized path for bfloat16 activations. + if quant_algo is None and dtype != torch.bfloat16: + return "unquantized locality domain MoE requires bfloat16" + if activation_type != ActivationType.Swiglu: + return "locality domain MoE backend test only supports SwiGLU" + if swiglu_gptoss_style: + return "locality domain MoE backend test does not cover GPT-OSS SwiGLU style" + return None + + +def should_skip_locality_domain_runtime(enable_locality_domains: bool) -> Optional[str]: + """Return a runtime skip reason for locality domain MoE backend params.""" + if not enable_locality_domains: + return None + if not torch.cuda.is_available(): + return "CUDA is not available" + sm_version = get_sm_version() + if sm_version != 107: + return f"Rubin (SM 107) required, got SM {sm_version}" + if not IS_CUTLASS_DSL_RUBIN_AVAILABLE: + return "public CuteDSL Rubin kernels are not available" + is_locality_domain_enabled.cache_clear() + if not is_locality_domain_enabled(): + return "locality domain is not enabled/supported on this system" + return None + + +def test_ci_acceleration_keeps_only_locality_domain_cutedsl_bf16( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from _torch.moe import moe_test_utils + + monkeypatch.setattr(moe_test_utils, "IS_CI_MODE", True) + model_config = MoeModelConfig(60, 4, 2048, 1408) + common_kwargs = { + "backend_type": MoeBackendType.CUTEDSL, + "quant_algo": None, + "model_config": model_config, + "routing_method_cls": RenormalizeMoeRoutingMethod, + "activation_type": ActivationType.Swiglu, + } + + assert should_skip_to_accelerate_ci(dtype=torch.bfloat16, **common_kwargs) is not None + assert ( + should_skip_to_accelerate_ci( + dtype=torch.bfloat16, + enable_locality_domains=True, + **common_kwargs, + ) + is None + ) + assert ( + should_skip_to_accelerate_ci( + dtype=torch.float16, + enable_locality_domains=True, + **common_kwargs, + ) + is not None + ) + + def generate_test_params() -> List: """ Generate test parameter combinations, filtering out unsupported configurations. @@ -1267,9 +1352,42 @@ def generate_test_params() -> List: swiglu_alpha, swiglu_beta, swiglu_limit, + False, ) params.append(create_test_param(param_values, test_id)) + if quant_algo in (QuantAlgo.NVFP4, None): + swiglu_gptoss_style = ( + swiglu_alpha != 1 or swiglu_beta != 0 or swiglu_limit != float("inf") + ) + locality_domain_skip_reason = should_skip_locality_domain_param( + backend_type, + quant_algo, + ActivationType.Swiglu, + swiglu_gptoss_style, + dtype, + ) + locality_domain_param_values = ( + dtype, + backend_type, + quant_algo, + seq_len, + model_config, + routing_method_cls, + ActivationType.Swiglu, + swiglu_alpha, + swiglu_beta, + swiglu_limit, + True, + ) + params.append( + create_test_param( + locality_domain_param_values, + f"locality_domain=enabled-{test_id}", + locality_domain_skip_reason, + ) + ) + return params @@ -1321,6 +1439,7 @@ def generate_element_wise_test_params() -> List: None, None, None, + False, ) params.append(create_test_param(param_values, test_id)) return params @@ -1329,23 +1448,23 @@ def generate_element_wise_test_params() -> List: TEST_PARAMS += generate_element_wise_test_params() -# ============================================================================ +# ===================================================================== # Test Implementation -# ============================================================================ +# ===================================================================== # # This file provides a UNIFIED TEST FRAMEWORK for testing all MoE backend # implementations through their backend-level interfaces. # -# ============================================================================= +# ====================================================================== # Purpose & Scope -# ============================================================================= +# ====================================================================== # - Test MoE backends via: routing_method.apply -> quantize_input -> run_moe # - Single GPU execution (no multi-GPU/distributed testing) # - Accuracy validation against reference implementations # -# ============================================================================= +# ====================================================================== # Test Coverage Matrix -# ============================================================================= +# ====================================================================== # 1. BACKENDS: CUTLASS, TRTLLM, CUTEDSL, DEEPGEMM # - When using element wise activations (Relu2, Silu), only CUTLASS and TRTLLM # are supported @@ -1378,19 +1497,20 @@ def generate_element_wise_test_params() -> List: # - Real models: Mixtral, DeepSeek, Grok, GPT-OSS # - Boundary cases: prime num_experts, small sizes, top_k=1, top_k=num_experts # -# ============================================================================= +# ====================================================================== # Skip Logic -# ============================================================================= +# ====================================================================== # Tests are automatically skipped for unsupported configurations using: # - backend.can_implement(p, d): declared quant / dtype / SM / dependency support # - should_skip_trtllm(): TRTLLM-specific constraints (num_experts % 4, etc.) # - should_skip_cutedsl(): CuteDSL-specific accuracy issues # - 128-alignment requirements for quantization # -# ============================================================================= +# ====================================================================== @pytest.mark.parametrize( "dtype_activation,backend_type,quant_algo,seq_len,model_config," - "routing_method_cls,activation_type,swiglu_alpha,swiglu_beta,swiglu_limit", + "routing_method_cls,activation_type,swiglu_alpha,swiglu_beta,swiglu_limit," + "enable_locality_domains", TEST_PARAMS, ) def test_moe_backend( @@ -1405,6 +1525,8 @@ def test_moe_backend( swiglu_beta: Optional[float], swiglu_limit: Optional[float], monkeypatch: pytest.MonkeyPatch, + enable_locality_domains: bool, + tmp_path: Path, ): """ Test MoE backend with autotune to capture all tactics. @@ -1431,6 +1553,11 @@ def test_moe_backend( # Default values: alpha=1, beta=0, limit=inf swiglu_gptoss_style = swiglu_alpha != 1 or swiglu_beta != 0 or swiglu_limit != float("inf") + locality_domain_runtime_skip = should_skip_locality_domain_runtime(enable_locality_domains) + if locality_domain_runtime_skip: + pytest.skip(locality_domain_runtime_skip) + locality_domain_policy = LocalityDomainPolicy(enabled=enable_locality_domains) + ci_skip = should_skip_to_accelerate_ci( backend_type=backend_type, quant_algo=quant_algo, @@ -1440,6 +1567,7 @@ def test_moe_backend( seq_len=seq_len, swiglu_gptoss_style=swiglu_gptoss_style, activation_type=activation_type, + enable_locality_domains=enable_locality_domains, ) if ci_skip: pytest.skip(ci_skip) @@ -1522,6 +1650,7 @@ def test_moe_backend( swiglu_limit=swiglu_tensors["swiglu_limit"] if swiglu_tensors else None, weight_loading_mode=weight_loading_mode, activation_type=activation_type, + locality_domain_policy=locality_domain_policy, ) # W4A8_MXFP4_MXFP8 / W4A8_MXFP4_FP8 require backend-layout-aware @@ -1542,6 +1671,9 @@ def test_moe_backend( backend.load_weights([weights]) backend.post_load_weights() backend.cuda() + if enable_locality_domains: + assert backend._locality_domain_runtime is not None + assert backend._locality_domain_weight_shards is not None # Create reference if ref_cls is not None: @@ -1555,6 +1687,8 @@ def test_moe_backend( # Clear autotuner cache before autotune phase AutoTuner.get().clear_cache() + if enable_locality_domains: + AutoTuner.get().reset_statistics() # Get reference output first with torch.inference_mode(): @@ -1583,8 +1717,26 @@ def run_moe(): # Autotune phase: tune kernels to find best tactics # Use cache_path to speed up subsequent runs by reusing tuning results - with torch.inference_mode(), autotune(cache_path="/tmp/moe_autotuner_cache.json"): + cache_path = ( + str(tmp_path / "moe_autotuner_cache.json") + if enable_locality_domains + else "/tmp/moe_autotuner_cache.json" + ) + with torch.inference_mode(), autotune(cache_path=cache_path): _ = run_moe() + if enable_locality_domains: + quant_name = "nvfp4" if quant_algo == QuantAlgo.NVFP4 else "bf16" + expected_tuning_ops = ( + f"CuteDslFusedMoE::run_moe_{quant_name}::locality_domain_end_to_end", + f"trtllm::cute_dsl_{quant_name}_gather_grouped_gemm_" + f"{'act_fusion' if quant_algo == QuantAlgo.NVFP4 else 'swiglu'}" + "_rubin::locality_domain_concurrent", + f"trtllm::cute_dsl_{quant_name}_grouped_gemm_finalize_" + "inplace_rubin::locality_domain_concurrent", + ) + for op_name in expected_tuning_ops: + assert autotuner.stats.tuned_op_profiled_configs.get(op_name, 0) > 0 + assert not autotuner.stats.failed_profiling_count.get(op_name, set()) # flashinfer has no capture and replay mechanisms, so we skip test_all_kernels use_flashinfer = getattr(backend, "use_flashinfer", False) @@ -1595,6 +1747,40 @@ def run_moe(): with AutoTuner.get().capture() as all_tactics, torch.inference_mode(): _ = run_moe() + # Replaying every outer tile is deliberately exhaustive and would + # multiply the inner FC tactic replay for every matrix member. One + # representative production shape per locality domain path covers that + # outer-tile contract; all matrix members still validate their + # tuned/failed statistics and replay the tactics selected for their + # own shape below. + representative_outer_tile_replay = ( + enable_locality_domains + and quant_algo in (QuantAlgo.NVFP4, None) + and seq_len == 1 + and (num_experts, top_k, hidden_size, intermediate_size) == (60, 4, 2048, 1408) + ) + if representative_outer_tile_replay: + # The regular Cartesian replay contains inner FC tactics only + # for the selected outer tile. Exercise every outer tile + # directly after tuning, when all corresponding FC caches have + # been prepared, so a non-winning tile cannot silently regress. + outer_context = all_tactics._captured_contexts[0] + outer_runner = outer_context["runners"][0] + outer_tactics = outer_runner.get_valid_tactics( + outer_context["inputs"], OptimizationProfile() + ) + expected_outer_tactics = ( + {128, 256, 512} if quant_algo == QuantAlgo.NVFP4 else {64, 128, 256} + ) + assert set(outer_tactics) == expected_outer_tactics + for outer_tactic in outer_tactics: + # Direct runner replay reuses the captured inplace output; + # reset it to the fresh-output baseline used by run_moe(). + with torch.inference_mode(): + outer_context["inputs"][-1].zero_() + output = outer_runner(outer_context["inputs"], tactic=outer_tactic) + ref_fused_moe.check_accuracy(output, ref_output) + # Replay phase: test each tactic for correctness # Set fail_fast=True to stop on first failure, False to run all and report summary replay_tactics_and_check( @@ -1614,9 +1800,9 @@ def run_moe(): ref_fused_moe.check_accuracy(output, ref_output) -# ============================================================================ +# ===================================================================== # BF16 (unquantized) TRTLLM-Gen MoE: DeepSeekV3 / Renormalize routing -# ============================================================================ +# ===================================================================== # The main test_moe_backend skips TRTLLM + quant_algo=None, so cover the BF16 # FlashInfer path here (Nemotron-H enablement): DeepSeekV3/Renormalize routing # x Relu2/Swiglu, via both fused and separated routing. @@ -1787,10 +1973,10 @@ def test_trtllm_bf16_dsv3_routing_kimi_k3_shape(seq_len): ) -# ============================================================================ +# ===================================================================== # TRTLLM-Gen shared-expert fusion (migrated from deprecated # tests/unittest/_torch/thop/serial/test_moe.py::TestMoeFP8 fusion coverage) -# ============================================================================ +# ===================================================================== # TRTLLMGenFusedMoE can fold n_shared_experts into the routed grouped GEMM as # always-selected experts (opt-in via TLLM_MOE_ENABLE_SHARED_EXPERT_FUSION=1; # requires FP8_BLOCK_SCALES + dp_size==1 + DeepSeekV3 routing). The fused @@ -2237,6 +2423,24 @@ def test_moe_backend_trtllm_nvfp4_fine_grained(num_tokens: int): ref_fused_moe.check_accuracy(output, ref_output) +def test_build_moe_deployment_carries_finalize_and_locality_domain(): + """The two fields SM107 eligibility reads must survive ``ModelConfig``.""" + from tensorrt_llm._torch.locality_domain.policy import LocalityDomainPolicy + + default = build_moe_deployment(ModelConfig(), num_experts=8) + assert default.fused_finalize_enabled is True + assert default.locality_domain_requested is False + + disabled = build_moe_deployment(ModelConfig(moe_disable_finalize_fusion=True), num_experts=8) + assert disabled.fused_finalize_enabled is False + + requested = build_moe_deployment( + ModelConfig(locality_domain_policy=LocalityDomainPolicy(enabled=True)), + num_experts=8, + ) + assert requested.locality_domain_requested is True + + def _nvfp4_problem(intermediate_size: Optional[int], activation: str) -> MoEProblem: return MoEProblem( quant=QuantAlgo.NVFP4.value, diff --git a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py index 118ac8963a48..73ad8102f004 100644 --- a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py +++ b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py @@ -1,13 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from contextlib import contextmanager +from types import SimpleNamespace + import pytest import torch from utils.util import check_accuracy -from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import GroupedGemmInputsHelper -from tensorrt_llm._torch.cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE -from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl import cute_dsl_nvfp4_grouped_gemm_ref +from tensorrt_llm._torch.autotuner import AutoTuner, OptimizationProfile, TunableRunner +from tensorrt_llm._torch.custom_ops import cute_dsl_custom_ops +from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + GroupedGemmInputsHelper, + _get_sm107_nvfp4_default_mma_config, +) +from tensorrt_llm._torch.cute_dsl_utils import ( + IS_CUTLASS_DSL_AVAILABLE, + IS_CUTLASS_DSL_RUBIN_AVAILABLE, +) +from tensorrt_llm._torch.locality_domain_utils import ( + end_for_all_locality_domain, + get_locality_domain_stream, + is_locality_domain_enabled, + locality_domain_device, + start_for_all_locality_domain, +) +from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl import ( + CuteDslFusedMoE, + _LocalityDomainConcurrentTunableRunner, + _runner_tactics_match_tile_size, + cute_dsl_nvfp4_grouped_gemm_ref, +) from tensorrt_llm._torch.moe.fused_moe.quantization import interleave_linear_and_gate from tensorrt_llm._torch.utils import ( ActivationType, + Fp4QuantizedTensor, is_gated_activation, relu2, swizzle_sf, @@ -941,3 +980,2927 @@ def test_nvfp4_gather_grouped_gemm_act_fusion_blackwell( c_sf_valid = torch.cat(c_sf_valid) c_sf_ref_valid = torch.cat(c_sf_ref_valid) check_accuracy(c_sf_valid, c_sf_ref_valid, atol=1e-4, rtol=1e-4, percent=0.95) + + +# ============================================================================ +# Rubin (SM107) Tests +# ============================================================================ + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on SM 107 (Rubin) GPUs", +) +@pytest.mark.parametrize( + "activation_type", + [ActivationType.Swiglu, ActivationType.Relu2], + ids=["swiglu", "relu2"], +) +@pytest.mark.parametrize("tile_size", [128, 256]) +@pytest.mark.parametrize("ep_size", [1, 8, 32]) +@pytest.mark.parametrize("top_k", [1, 2, 8]) +@pytest.mark.parametrize("num_tokens", [128, 515, 1024, 8192]) +def test_nvfp4_gather_grouped_gemm_act_fusion_rubin( + num_tokens: int, + top_k: int, + ep_size: int, + tile_size: int, + activation_type: ActivationType, +): + """Test gather-based grouped GEMM with fused activation on Rubin (SM107). + + This test validates the gather kernel which: + 1. Uses gather for A/SFA loading with permuted_idx_to_expanded_idx + 2. Performs GEMM with (interleaved for gated) weights + 3. Applies the fused activation (SwiGLU for gated, Relu2 for non-gated) + 4. Quantizes output to FP4 with scale factor generation + """ + is_gated = is_gated_activation(activation_type) + weight_n_multiplier = 2 if is_gated else 1 + sf_vec_size = 16 + hidden_size = 4096 + interm_size = 8192 + num_experts = 256 + num_local_experts = num_experts // ep_size + + # Generate routing information + routing_logits = torch.randn(num_tokens, num_experts, device="cuda") + token_final_scales, token_selected_experts = routing_logits.topk(top_k, dim=-1) + token_selected_experts = token_selected_experts.to(torch.int32) + token_final_scales = token_final_scales.softmax(dim=-1).to(torch.float32) + # Ensure at least one valid token + token_selected_experts[0] = 0 + + ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles, + ) = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=0, + local_num_experts=num_local_experts, + tile_tokens_dim=tile_size, + ) + + max_num_permuted_tokens = permuted_idx_to_expanded_idx.size(0) + num_valid_permuted_tokens = total_num_padded_tokens.item() + + # Create input tensors (original size, not permuted) + a = torch.randint(-5, 5, (num_tokens, hidden_size), dtype=torch.int32, device="cuda").to( + torch.bfloat16 + ) + b = torch.randint( + -5, + 5, + (num_local_experts, interm_size * weight_n_multiplier, hidden_size), + dtype=torch.int32, + device="cuda", + ).to(torch.bfloat16) + + # Quantize inputs to FP4 + a_global_sf = a.abs().max().float() / (448 * 6) + b_global_sf = b.abs().amax(dim=(1, 2)).float() / (448 * 6) + a, a_sf = torch.ops.trtllm.fp4_quantize(a, 1 / a_global_sf, sf_vec_size, False) + a = a.view(torch.float4_e2m1fn_x2) + a_sf_unswizzled = unswizzle_sf(a_sf, (num_tokens + 127) // 128 * 128, hidden_size)[:num_tokens] + b, b_sf = torch.ops.trtllm.fp4_quantize(b, 1 / b_global_sf, sf_vec_size, False) + b = b.view(torch.float4_e2m1fn_x2) + weight_n = interm_size * weight_n_multiplier + b_sf = b_sf.view(num_local_experts, weight_n, hidden_size // sf_vec_size) + alpha = a_global_sf * b_global_sf + + b_kernel = b + b_sf_kernel = b_sf + if is_gated: + b_kernel = interleave_linear_and_gate(b.view(torch.uint8), group_size=64, dim=1).view( + torch.float4_e2m1fn_x2 + ) + b_sf_unswizzled = unswizzle_sf(b_sf, weight_n, hidden_size).view( + num_local_experts, weight_n, hidden_size // sf_vec_size + ) + b_sf_unswizzled = interleave_linear_and_gate(b_sf_unswizzled, group_size=64, dim=1) + b_sf_kernel = swizzle_sf(b_sf_unswizzled, weight_n, hidden_size).view( + num_local_experts, weight_n, hidden_size // sf_vec_size + ) + + # Compute reference: manually gather, compute GEMM, apply activation, then quantize + permuted_idx_to_expanded_idx_list = permuted_idx_to_expanded_idx.cpu().tolist() + tile_idx_to_mn_limit_list = tile_idx_to_mn_limit.cpu().tolist() + + a_gathered = torch.empty(max_num_permuted_tokens, hidden_size // 2, dtype=a.dtype) + a_sf_gathered = torch.empty( + max_num_permuted_tokens, hidden_size // sf_vec_size, dtype=a_sf.dtype + ) + for i in range(num_valid_permuted_tokens): + if i >= tile_idx_to_mn_limit_list[i // tile_size]: + continue + expanded_idx = permuted_idx_to_expanded_idx_list[i] + token_id = expanded_idx // top_k + a_gathered[i] = a[token_id] + a_sf_gathered[i] = a_sf_unswizzled[token_id] + a_gathered = a_gathered.to(a.device) + a_sf_gathered = a_sf_gathered.to(a.device) + + a_sf_gathered_swizzled = swizzle_sf( + a_sf_gathered.view(max_num_permuted_tokens, hidden_size // sf_vec_size), + max_num_permuted_tokens, + hidden_size, + ) + + c_ref = cute_dsl_nvfp4_grouped_gemm_ref( + a_gathered, + b, + a_sf_gathered_swizzled, + b_sf, + alpha, + tile_idx_to_group_idx, + num_non_exiting_tiles, + tile_size=tile_size, + output_dtype=torch.bfloat16, + scaling_vector_size=sf_vec_size, + ) + c_ref = apply_activation_ref(c_ref, activation_type) + global_sf = c_ref[:num_valid_permuted_tokens].abs().max().float() / (448 * 6) + c_ref, c_sf_ref = torch.ops.trtllm.fp4_quantize(c_ref, 1 / global_sf, sf_vec_size, False) + + # Call Rubin gather kernel + c, c_sf = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( + a, + b_kernel, + a_sf_unswizzled, + b_sf_kernel, + alpha, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + torch.tensor([1 / global_sf], dtype=torch.float32, device="cuda"), + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=0, + tile_size=tile_size, + output_tensor=None, + output_sf_tensor=None, + scaling_vector_size=sf_vec_size, + activation_type=activation_type, + ) + + # Verify output (only compare valid tokens, skip padding) + valid_token_mask = torch.zeros(num_valid_permuted_tokens, dtype=torch.bool, device="cuda") + for i in range(num_valid_permuted_tokens): + if i >= tile_idx_to_mn_limit_list[i // tile_size]: + continue + valid_token_mask[i] = True + + num_valid_tokens = valid_token_mask.sum().item() + if num_valid_tokens > 0: + c_valid = c[:num_valid_permuted_tokens].view(torch.uint8)[valid_token_mask] + c_ref_valid = c_ref[:num_valid_permuted_tokens][valid_token_mask] + check_accuracy(c_valid, c_ref_valid, atol=1e-4, rtol=1e-4, percent=0.95) + + c_sf_unswizzled = unswizzle_sf(c_sf, max_num_permuted_tokens, interm_size, sf_vec_size) + c_sf_ref_unswizzled = unswizzle_sf( + c_sf_ref, max_num_permuted_tokens, interm_size, sf_vec_size + ) + + c_sf_valid = [] + c_sf_ref_valid = [] + for i in range(num_valid_permuted_tokens): + if i >= tile_idx_to_mn_limit_list[i // tile_size]: + continue + c_sf_valid.append(c_sf_unswizzled[i]) + c_sf_ref_valid.append(c_sf_ref_unswizzled[i]) + + c_sf_valid = torch.cat(c_sf_valid) + c_sf_ref_valid = torch.cat(c_sf_ref_valid) + check_accuracy(c_sf_valid, c_sf_ref_valid, atol=1e-4, rtol=1e-4, percent=0.95) + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on SM 107 (Rubin) GPUs", +) +@pytest.mark.parametrize( + "num_tokens, num_experts, top_k, hidden_size, interm_size, tile_size", + [ + # DeepSeek V3 Lite-like: small tokens with tile_size=128 and 256 + (16, 72, 6, 2560, 1536, 128), + (16, 72, 6, 2560, 1536, 256), + (8, 72, 6, 2560, 1536, 128), + (8, 72, 6, 2560, 1536, 256), + # Qwen3-30B-A3B-like: hidden_size=2048 (triggered Fix 7) + (8, 128, 8, 2048, 1536, 128), + (8, 128, 8, 2048, 1536, 256), + # Very small: 1 token, tile_size=256 (triggered Fix 9) + (1, 72, 6, 2560, 1536, 128), + (1, 72, 6, 2560, 1536, 256), + # Small tokens, high padding ratio + (4, 72, 6, 2560, 1536, 256), + (2, 128, 8, 2048, 1536, 256), + ], +) +def test_nvfp4_gather_grouped_gemm_swiglu_rubin_small_tokens( + num_tokens: int, + num_experts: int, + top_k: int, + hidden_size: int, + interm_size: int, + tile_size: int, +): + """Test FC1 gather+SwiGLU kernel on Rubin with small num_tokens (high padding ratio). + + Covers configurations that triggered Fix 9 (pad_val crash with tile_size=256) + and Qwen3-30B-A3B shapes (hidden_size=2048, Fix 7). Uses real moe_sort routing + metadata (not synthetic). + """ + sf_vec_size = 16 + num_local_experts = num_experts # ep_size=1 + + routing_logits = torch.randn(num_tokens, num_experts, device="cuda") + token_final_scales, token_selected_experts = routing_logits.topk(top_k, dim=-1) + token_selected_experts = token_selected_experts.to(torch.int32) + token_final_scales = token_final_scales.softmax(dim=-1).to(torch.float32) + + ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles, + ) = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=0, + local_num_experts=num_local_experts, + tile_tokens_dim=tile_size, + ) + + max_num_permuted_tokens = permuted_idx_to_expanded_idx.size(0) + n_tiles = num_non_exiting_tiles.item() + + a = torch.randint(-5, 5, (num_tokens, hidden_size), dtype=torch.int32, device="cuda").to( + torch.bfloat16 + ) + b = torch.randint( + -5, + 5, + (num_local_experts, interm_size * 2, hidden_size), + dtype=torch.int32, + device="cuda", + ).to(torch.bfloat16) + + a_global_sf = a.abs().max().float() / (448 * 6) + b_global_sf = b.abs().amax(dim=(1, 2)).float() / (448 * 6) + a, a_sf = torch.ops.trtllm.fp4_quantize(a, 1 / a_global_sf, sf_vec_size, False) + a = a.view(torch.float4_e2m1fn_x2) + a_sf_unswizzled = unswizzle_sf(a_sf, (num_tokens + 127) // 128 * 128, hidden_size)[:num_tokens] + b, b_sf = torch.ops.trtllm.fp4_quantize(b, 1 / b_global_sf, sf_vec_size, False) + b = b.view(torch.float4_e2m1fn_x2) + b_sf = b_sf.view(num_local_experts, interm_size * 2, hidden_size // sf_vec_size) + alpha = a_global_sf * b_global_sf + + b_interleaved = interleave_linear_and_gate(b.view(torch.uint8), group_size=64, dim=1).view( + torch.float4_e2m1fn_x2 + ) + b_sf_unswizzled = unswizzle_sf(b_sf, interm_size * 2, hidden_size).view( + num_local_experts, interm_size * 2, hidden_size // sf_vec_size + ) + b_sf_unswizzled_interleaved = interleave_linear_and_gate(b_sf_unswizzled, group_size=64, dim=1) + b_sf_interleaved = swizzle_sf(b_sf_unswizzled_interleaved, interm_size * 2, hidden_size).view( + num_local_experts, interm_size * 2, hidden_size // sf_vec_size + ) + + # Compute reference: gather A using permuted_idx, then grouped GEMM + SwiGLU + # Use uint8 + view because torch.zeros doesn't support Float4_e2m1fn_x2 + a_gathered = torch.zeros( + max_num_permuted_tokens, hidden_size // 2, dtype=torch.uint8, device=a.device + ).view(torch.float4_e2m1fn_x2) + a_sf_gathered = torch.zeros( + max_num_permuted_tokens, hidden_size // sf_vec_size, dtype=a_sf.dtype, device=a_sf.device + ) + num_valid_permuted_tokens = n_tiles * tile_size + for i in range(min(num_valid_permuted_tokens, max_num_permuted_tokens)): + expanded_idx = permuted_idx_to_expanded_idx[i].item() + if expanded_idx > 0 or i == 0: + token_id = expanded_idx // top_k + if token_id < num_tokens: + a_gathered[i] = a[token_id] + a_sf_gathered[i] = a_sf_unswizzled[token_id] + + a_sf_gathered_swizzled = swizzle_sf( + a_sf_gathered.view(max_num_permuted_tokens, hidden_size // sf_vec_size), + max_num_permuted_tokens, + hidden_size, + ) + + c_ref = cute_dsl_nvfp4_grouped_gemm_ref( + a_gathered, + b, + a_sf_gathered_swizzled, + b_sf, + alpha, + tile_idx_to_group_idx, + num_non_exiting_tiles, + tile_size=tile_size, + output_dtype=torch.bfloat16, + scaling_vector_size=sf_vec_size, + ) + c_ref = swiglu_ref(c_ref) + global_sf = c_ref[:num_valid_permuted_tokens].abs().max().float() / (448 * 6) + if global_sf == 0: + global_sf = torch.tensor(1.0, dtype=torch.float32, device="cuda") + c_ref, c_sf_ref = torch.ops.trtllm.fp4_quantize(c_ref, 1 / global_sf, sf_vec_size, False) + + c, c_sf = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( + a, + b_interleaved, + a_sf_unswizzled, + b_sf_interleaved, + alpha, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + torch.tensor([1 / global_sf], dtype=torch.float32, device="cuda"), + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=0, + tile_size=tile_size, + output_tensor=None, + output_sf_tensor=None, + scaling_vector_size=sf_vec_size, + activation_type=ActivationType.Swiglu, + ) + + # Verify output for valid (non-padding) tokens + valid_token_mask = permuted_idx_to_expanded_idx[:num_valid_permuted_tokens] != 0 + valid_token_mask[0] = True # index 0 is always valid + num_valid_tokens = valid_token_mask.sum().item() + if num_valid_tokens > 0: + c_valid = c[:num_valid_permuted_tokens].view(torch.uint8)[valid_token_mask] + c_ref_valid = c_ref[:num_valid_permuted_tokens][valid_token_mask] + check_accuracy(c_valid, c_ref_valid, atol=1e-4, rtol=1e-4, percent=0.95) + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on SM 107 (Rubin) GPUs", +) +@pytest.mark.parametrize("tile_size", [128, 256]) +@pytest.mark.parametrize("ep_size", [1, 8, 32]) +@pytest.mark.parametrize("top_k", [1, 2, 8]) +@pytest.mark.parametrize("num_tokens", [128, 515, 1024, 8192]) +def test_nvfp4_grouped_gemm_finalize_rubin( + num_tokens: int, top_k: int, ep_size: int, tile_size: int +): + """Test grouped GEMM with finalize fusion on Rubin (SM107). + + Same test logic as test_nvfp4_grouped_gemm_finalize_blackwell + but calls the Rubin-specific custom op. + """ + sf_vec_size = 16 + hidden_size = 4096 + interm_size = 8192 + num_experts = 256 + num_local_experts = num_experts // ep_size + + routing_logits = torch.randn(num_tokens, num_experts, device="cuda") + token_final_scales, token_selected_experts = routing_logits.topk(top_k, dim=-1) + token_selected_experts = token_selected_experts.to(torch.int32) + token_final_scales = token_final_scales.softmax(dim=-1).to(torch.float32) + + ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles, + ) = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=0, + local_num_experts=num_local_experts, + tile_tokens_dim=tile_size, + ) + + max_num_permuted_tokens = permuted_idx_to_expanded_idx.size(0) + a = torch.randint( + -5, 5, (max_num_permuted_tokens, hidden_size), dtype=torch.int32, device="cuda" + ).to(torch.bfloat16) + b = torch.randint( + -5, + 5, + (num_local_experts, interm_size, hidden_size), + dtype=torch.int32, + device="cuda", + ).to(torch.bfloat16) + + a_global_sf = a.abs().max().float() / (448 * 6) + b_global_sf = b.abs().amax(dim=(1, 2)).float() / (448 * 6) + a, a_sf = torch.ops.trtllm.fp4_quantize(a, 1 / a_global_sf, sf_vec_size, False) + a = a.view(torch.float4_e2m1fn_x2) + b, b_sf = torch.ops.trtllm.fp4_quantize(b, 1 / b_global_sf, sf_vec_size, False) + b = b.view(torch.float4_e2m1fn_x2) + b_sf = b_sf.view(num_local_experts, interm_size, hidden_size // sf_vec_size) + alpha = a_global_sf * b_global_sf + + # Call Rubin finalize kernel + c = torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_finalize_rubin( + a, + b, + a_sf, + b_sf, + alpha, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + token_final_scales, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=0, + tile_size=tile_size, + output_dtype=torch.bfloat16, + scaling_vector_size=sf_vec_size, + ) + + # Compute reference for test_nvfp4_grouped_gemm_finalize_rubin + c_ref = cute_dsl_nvfp4_grouped_gemm_ref( + a, + b, + a_sf, + b_sf, + alpha, + tile_idx_to_group_idx, + num_non_exiting_tiles, + tile_size=tile_size, + output_dtype=torch.bfloat16, + scaling_vector_size=sf_vec_size, + ) + c_ref = torch.ops.trtllm.moe_unpermute( + permuted_input=c_ref, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + topk_scales=token_final_scales, + ) + match_ratio = torch.isclose(c, c_ref, rtol=1.6e-2, atol=1e-5).sum().item() / c.numel() + assert match_ratio > 0.99 + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on SM 107 (Rubin) GPUs", +) +@pytest.mark.parametrize( + "num_tokens, num_experts, top_k, hidden_size, interm_size, tile_size", + [ + # DeepSeek V3 Lite-like: 16 tokens, 72 experts, top_k=6, tile_size=128 + (16, 72, 6, 2560, 1536, 128), + # DeepSeek V3 Lite-like: 16 tokens, 72 experts, top_k=6, tile_size=256 (CRASHES in e2e) + (16, 72, 6, 2560, 1536, 256), + # DeepSeek V3 Lite-like: 8 tokens + (8, 72, 6, 2560, 1536, 128), + (8, 72, 6, 2560, 1536, 256), + # Qwen3-30B-A3B-like: 8 tokens, 128 experts, top_k=8 + (8, 128, 8, 2048, 1536, 128), + (8, 128, 8, 2048, 1536, 256), + # Very small: 1 token + (1, 72, 6, 2560, 1536, 128), + (1, 72, 6, 2560, 1536, 256), + # Small tokens, high padding ratio + (4, 72, 6, 2560, 1536, 256), + (2, 128, 8, 2048, 1536, 256), + ], +) +def test_nvfp4_grouped_gemm_finalize_rubin_small_tokens( + num_tokens: int, + num_experts: int, + top_k: int, + hidden_size: int, + interm_size: int, + tile_size: int, +): + """Test FC2 finalize kernel on Rubin with small num_tokens (high padding ratio). + + This reproduces the crashing e2e configuration where num_tokens is small + relative to num_experts, causing most tile rows to be padding. Uses real + moe_sort routing metadata (not synthetic). + """ + sf_vec_size = 16 + num_local_experts = num_experts # ep_size=1 + + routing_logits = torch.randn(num_tokens, num_experts, device="cuda") + token_final_scales, token_selected_experts = routing_logits.topk(top_k, dim=-1) + token_selected_experts = token_selected_experts.to(torch.int32) + token_final_scales = token_final_scales.softmax(dim=-1).to(torch.float32) + + ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles, + ) = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=0, + local_num_experts=num_local_experts, + tile_tokens_dim=tile_size, + ) + + max_num_permuted_tokens = permuted_idx_to_expanded_idx.size(0) + n_tiles = num_non_exiting_tiles.item() + print( + f"\n[test_small_tokens] num_tokens={num_tokens}, num_experts={num_experts}, " + f"top_k={top_k}, tile_size={tile_size}, " + f"total_padded={max_num_permuted_tokens}, num_tiles={n_tiles}" + ) + + a = torch.randint( + -5, 5, (max_num_permuted_tokens, hidden_size), dtype=torch.int32, device="cuda" + ).to(torch.bfloat16) + b = torch.randint( + -5, + 5, + (num_local_experts, interm_size, hidden_size), + dtype=torch.int32, + device="cuda", + ).to(torch.bfloat16) + + a_global_sf = a.abs().max().float() / (448 * 6) + b_global_sf = b.abs().amax(dim=(1, 2)).float() / (448 * 6) + a, a_sf = torch.ops.trtllm.fp4_quantize(a, 1 / a_global_sf, sf_vec_size, False) + a = a.view(torch.float4_e2m1fn_x2) + b, b_sf = torch.ops.trtllm.fp4_quantize(b, 1 / b_global_sf, sf_vec_size, False) + b = b.view(torch.float4_e2m1fn_x2) + b_sf = b_sf.view(num_local_experts, interm_size, hidden_size // sf_vec_size) + alpha = a_global_sf * b_global_sf + + # Call Rubin finalize kernel with real moe_sort routing data + c = torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_finalize_rubin( + a, + b, + a_sf, + b_sf, + alpha, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + token_final_scales, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=0, + tile_size=tile_size, + output_dtype=torch.bfloat16, + scaling_vector_size=sf_vec_size, + ) + + # Compute reference for test_nvfp4_grouped_gemm_finalize_rubin_small_tokens + c_ref = cute_dsl_nvfp4_grouped_gemm_ref( + a, + b, + a_sf, + b_sf, + alpha, + tile_idx_to_group_idx, + num_non_exiting_tiles, + tile_size=tile_size, + output_dtype=torch.bfloat16, + scaling_vector_size=sf_vec_size, + ) + c_ref = torch.ops.trtllm.moe_unpermute( + permuted_input=c_ref, + expanded_idx_to_permuted_idx=expanded_idx_to_permuted_idx, + topk_scales=token_final_scales, + ) + match_ratio = torch.isclose(c, c_ref, rtol=1.6e-2, atol=1e-5).sum().item() / c.numel() + assert match_ratio > 0.99 + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on SM 107 (Rubin) GPUs", +) +@pytest.mark.parametrize("tile_size", [64, 128, 256]) +@pytest.mark.parametrize("ep_size", [1, 8, 32]) +@pytest.mark.parametrize("top_k", [1, 2, 8]) +@pytest.mark.parametrize("num_tokens", [128, 515, 1024, 8192]) +def test_bf16_gather_grouped_gemm_swiglu_rubin( + num_tokens: int, + top_k: int, + ep_size: int, + tile_size: int, +): + """Test BF16 gather-based grouped GEMM with SwiGLU fusion on Rubin (SM107). + + Uses torch.ops.trtllm.cute_dsl_bf16_gather_grouped_gemm_swiglu_rubin. + No scale factors or quantization — direct BF16 inputs/outputs. + """ + hidden_size = 4096 + interm_size = 8192 + num_experts = 256 + num_local_experts = num_experts // ep_size + interleave_granularity = 32 + + # Generate routing information + routing_logits = torch.randn(num_tokens, num_experts, device="cuda") + token_final_scales, token_selected_experts = routing_logits.topk(top_k, dim=-1) + token_selected_experts = token_selected_experts.to(torch.int32) + token_final_scales = token_final_scales.softmax(dim=-1).to(torch.float32) + token_selected_experts[0] = 0 + + ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles, + ) = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=0, + local_num_experts=num_local_experts, + tile_tokens_dim=tile_size, + ) + + max_num_permuted_tokens = permuted_idx_to_expanded_idx.size(0) + num_valid_permuted_tokens = total_num_padded_tokens.item() + + # Create BF16 input tensors + a = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda") + b = torch.randn( + num_local_experts, interm_size * 2, hidden_size, dtype=torch.bfloat16, device="cuda" + ) + alpha = torch.ones(num_local_experts, dtype=torch.float32, device="cuda") + + # Interleave weights for SwiGLU: [up_0:64, gate_64:128, ...] + b_interleaved = interleave_linear_and_gate( + b.view(torch.uint8), group_size=interleave_granularity, dim=1 + ).view(torch.bfloat16) + + # Compute reference: gather A, GEMM per group, SwiGLU + permuted_idx_list = permuted_idx_to_expanded_idx.cpu().tolist() + tile_group_list = tile_idx_to_group_idx.cpu().tolist() + tile_mn_limit_list = tile_idx_to_mn_limit.cpu().tolist() + + c_ref = torch.zeros(max_num_permuted_tokens, interm_size, dtype=torch.float32, device="cuda") + for tile_idx in range(num_non_exiting_tiles.item()): + group_idx = tile_group_list[tile_idx] + mn_limit = tile_mn_limit_list[tile_idx] + start = tile_idx * tile_size + end = min(start + tile_size, mn_limit) + + for i in range(start, end): + token_id = permuted_idx_list[i] // top_k + a_row = a[token_id].float() + gemm_row = a_row @ b_interleaved[group_idx].float().T * alpha[group_idx].item() + # SwiGLU on interleaved result + out_row = torch.zeros(interm_size, dtype=torch.float32, device="cuda") + for n_block in range(0, interm_size * 2, 2 * interleave_granularity): + up = gemm_row[n_block : n_block + interleave_granularity] + gate = gemm_row[ + n_block + interleave_granularity : n_block + 2 * interleave_granularity + ] + out_start = n_block // 2 + out_row[out_start : out_start + interleave_granularity] = up * ( + gate * torch.sigmoid(gate) + ) + c_ref[i] = out_row + + # Build valid mask for accuracy checking + valid_mask = torch.zeros(num_valid_permuted_tokens, dtype=torch.bool, device="cuda") + for i in range(num_valid_permuted_tokens): + if i < tile_mn_limit_list[i // tile_size]: + valid_mask[i] = True + c_ref_valid = c_ref[:num_valid_permuted_tokens][valid_mask] + + # Even-tile padding for Rubin cluster sync + kernel_nnet = ((num_non_exiting_tiles + 1) // 2) * 2 + + # Test all valid autotuner candidate tactics via direct runner call + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + Sm107ContiguousGatherGroupedGemmSwigluFusionRunner, + ) + + runner = Sm107ContiguousGatherGroupedGemmSwigluFusionRunner( + num_experts, top_k, num_local_experts, 0, tile_size + ) + inputs = [ + a, + b_interleaved, + alpha, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + kernel_nnet, + ] + tactics = runner.get_valid_tactics(inputs, None) + assert len(tactics) > 0, f"No valid tactics for tile_size={tile_size}" + + failed = [] + for tactic in tactics: + mma_tiler, _, cluster, _ = tactic + label = f"mma={mma_tiler[:2]} cluster={cluster}" + with torch.inference_mode(): + c = runner.forward(inputs, tactic=tactic) + c_valid = c[:num_valid_permuted_tokens].float()[valid_mask] + if c_ref_valid.numel() > 0: + match = ( + torch.isclose(c_valid, c_ref_valid, rtol=1e-2, atol=0.5).sum().item() + / c_ref_valid.numel() + ) + if match < 0.95: + failed.append(f"{label}: match={match:.4f}") + assert not failed, ( + f"tile_size={tile_size}: {len(failed)}/{len(tactics)} tactics failed:\n " + + "\n ".join(failed) + ) + + +def _skip_if_no_locality_domain(): + is_locality_domain_enabled.cache_clear() + if not is_locality_domain_enabled(): + pytest.skip("locality domain localization is not enabled/supported on this system") + + +def _setup_locality_domain_routing(num_tokens, num_experts, num_local_experts, top_k, tile_size): + torch.manual_seed(42) + routing_logits = torch.randn(num_tokens, num_experts, device="cuda") + token_final_scales, token_selected_experts = routing_logits.topk(top_k, dim=-1) + token_selected_experts = token_selected_experts.to(torch.int32) + token_final_scales = token_final_scales.softmax(dim=-1).to(torch.float32) + token_selected_experts[0] = 0 + + ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles, + ) = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=0, + local_num_experts=num_local_experts, + tile_tokens_dim=tile_size, + ) + + return ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles, + token_final_scales, + ) + + +def _valid_permuted_token_mask(tile_idx_to_mn_limit, num_valid, tile_size): + tile_idx_to_mn_limit_list = tile_idx_to_mn_limit.cpu().tolist() + valid_mask = torch.zeros(num_valid, dtype=torch.bool, device=tile_idx_to_mn_limit.device) + for row_idx in range(num_valid): + if row_idx < tile_idx_to_mn_limit_list[row_idx // tile_size]: + valid_mask[row_idx] = True + return valid_mask + + +def _create_quantized_locality_domain_inputs(num_tokens, hidden_size, sf_vec_size=16, seed=42): + torch.manual_seed(seed) + a = torch.randint(-5, 5, (num_tokens, hidden_size), dtype=torch.int32, device="cuda").to( + torch.bfloat16 + ) + a_global_sf = a.abs().max().float() / (448 * 6) + a_fp4, a_sf = torch.ops.trtllm.fp4_quantize(a, 1 / a_global_sf, sf_vec_size, False) + a_fp4 = a_fp4.view(torch.float4_e2m1fn_x2) + a_sf_unswizzled = unswizzle_sf(a_sf, (num_tokens + 127) // 128 * 128, hidden_size)[:num_tokens] + return a_fp4, a_sf_unswizzled, a_global_sf + + +def _create_quantized_locality_domain_weights( + num_local_experts, interm_size, hidden_size, sf_vec_size=16, seed=123 +): + torch.manual_seed(seed) + weight = torch.randint( + -5, 5, (num_local_experts, interm_size * 2, hidden_size), dtype=torch.int32, device="cuda" + ).to(torch.bfloat16) + weight_global_sf = weight.abs().amax(dim=(1, 2)).float() / (448 * 6) + weight_fp4, weight_sf = torch.ops.trtllm.fp4_quantize( + weight, 1 / weight_global_sf, sf_vec_size, False + ) + weight_fp4 = weight_fp4.view(torch.float4_e2m1fn_x2) + weight_sf = weight_sf.view(num_local_experts, interm_size * 2, hidden_size // sf_vec_size) + + weight_interleaved = interleave_linear_and_gate( + weight_fp4.view(torch.uint8), group_size=64, dim=1 + ).view(torch.float4_e2m1fn_x2) + weight_sf_unswizzled = unswizzle_sf(weight_sf, interm_size * 2, hidden_size).view( + num_local_experts, interm_size * 2, hidden_size // sf_vec_size + ) + weight_sf_unswizzled_interleaved = interleave_linear_and_gate( + weight_sf_unswizzled, group_size=64, dim=1 + ) + weight_sf_interleaved = swizzle_sf( + weight_sf_unswizzled_interleaved, interm_size * 2, hidden_size + ).view(num_local_experts, interm_size * 2, hidden_size // sf_vec_size) + + return weight_interleaved, weight_sf_interleaved, weight_global_sf + + +def test_moe_output_memset_aux_stream_guard_requires_full_aux_state(): + from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl import CuteDslFusedMoE + from tensorrt_llm._torch.utils import AuxStreamType, EventType + + backend = object.__new__(CuteDslFusedMoE) + backend.event_dict = {EventType.Main: object()} + backend.aux_stream_dict = {} + + assert not backend._has_moe_output_memset_aux_stream() + + backend.event_dict[EventType.MoeOutputMemset] = object() + backend.aux_stream_dict[AuxStreamType.MoeOutputMemset] = object() + + assert backend._has_moe_output_memset_aux_stream() + + +def test_runner_tactics_match_tile_size_unwraps_locality_domain_runner(): + class OuterRunner: + pass + + class FakeOpRunner(TunableRunner): + def get_valid_tactics(self, inputs, profile, **kwargs): + return [] + + def forward(self, inputs, tactic): + return None + + op_runner = FakeOpRunner() + runner = _LocalityDomainConcurrentTunableRunner( + op_runner, + SimpleNamespace(num_partitions=2), + 2, + lambda *_: None, + ) + outer_runner = OuterRunner() + matching_tactic = ((128, 128, 64), (128, 128, 16), (1, 1), False) + mismatched_tactic = ((64, 128, 64), (64, 128, 16), (1, 1), False) + assert _runner_tactics_match_tile_size( + [(outer_runner, 128), (runner, matching_tactic)], + OuterRunner, + (FakeOpRunner,), + ) + assert not _runner_tactics_match_tile_size( + [(outer_runner, 128), (runner, mismatched_tactic)], + OuterRunner, + (FakeOpRunner,), + ) + + +def test_locality_domain_concurrent_tunable_runner_delegates_and_launches_all_partitions(): + class FakeOpRunner(TunableRunner): + def __init__(self): + self.tactics_call = None + + def unique_id(self): + return ("fake-op", 7) + + def get_valid_tactics(self, inputs, profile, **kwargs): + self.tactics_call = (id(inputs), profile, kwargs) + return [-1, ("fast", 128)] + + def forward(self, /, inputs, *, tactic=-1, **kwargs): + raise AssertionError("the wrapper must call launch_fn, not the op runner") + + class FakeRuntime: + def __init__(self, topology): + self.topology = topology + self.events = [] + + def topology_identity(self): + return self.topology + + def fork(self): + self.events.append(("fork",)) + + @contextmanager + def partition_context(self, partition_id): + self.events.append(("enter", partition_id)) + yield + self.events.append(("exit", partition_id)) + + def join(self): + self.events.append(("join",)) + + op_runner = FakeOpRunner() + runtime = FakeRuntime(((100, 212), (100, 212), (12, 212))) + inputs = [torch.empty(1)] + + def launch(partition_id, launch_inputs, tactic): + runtime.events.append(("launch", partition_id, id(launch_inputs), tactic)) + + runner = _LocalityDomainConcurrentTunableRunner(op_runner, runtime, 3, launch) + profile = OptimizationProfile() + + assert runner.get_valid_tactics(inputs, profile, marker="x") == [-1, ("fast", 128)] + assert op_runner.tactics_call == (id(inputs), profile, {"marker": "x"}) + assert runner.unique_id()[0] == op_runner.unique_id() + assert runner.op_runner is op_runner + + other_runtime = FakeRuntime(((106, 212), (106, 212), (0, 212))) + other_runner = _LocalityDomainConcurrentTunableRunner(op_runner, other_runtime, 3, launch) + assert runner.unique_id() != other_runner.unique_id() + + tactic = ("chosen", 256) + assert runner(inputs, tactic=tactic) is None + assert runtime.events == [ + ("fork",), + ("enter", 0), + ("launch", 0, id(inputs), tactic), + ("exit", 0), + ("enter", 1), + ("launch", 1, id(inputs), tactic), + ("exit", 1), + ("enter", 2), + ("launch", 2, id(inputs), tactic), + ("exit", 2), + ("join",), + ] + + +@pytest.mark.parametrize( + "runner_name,num_inputs,output_idx,tile_sizes", + [ + ("CuteDslFusedMoENvfp4Runner", 5, 4, [128, 256, 512]), + ("CuteDslFusedMoEBF16Runner", 4, 3, [64, 128, 256]), + ], +) +def test_locality_domain_outer_preparation_disables_memset_overlap( + monkeypatch, + runner_name: str, + num_inputs: int, + output_idx: int, + tile_sizes: list[int], +): + from tensorrt_llm._torch.moe.fused_moe import fused_moe_cute_dsl + + monkeypatch.setattr(fused_moe_cute_dsl, "get_sm_version", lambda: 107) + preparation_calls = [] + + def forward_impl(*args, **kwargs): + preparation_calls.append((args, kwargs)) + return args[output_idx] + + runner_cls = getattr(fused_moe_cute_dsl, runner_name) + runner = runner_cls( + forward_impl, + num_experts=1, + top_k=1, + num_local_experts=1, + local_expert_offset=0, + workload_identity=("locality_domain",), + ) + inputs = [torch.empty(0) for _ in range(num_inputs)] + + result = runner(inputs, tactic=-1, do_preparation=True) + + assert result is inputs[output_idx] + assert [kwargs["tile_size"] for _, kwargs in preparation_calls] == tile_sizes + for args, kwargs in preparation_calls: + assert args == tuple(inputs) + assert kwargs["enable_alltoall"] is False + assert kwargs["overlap_moe_output_memset"] is False + + +def _cute_dsl_eligibility(**deployment_kwargs): + """``can_implement`` verdict for NVFP4 on an SM107 machine with Rubin DSL.""" + from tensorrt_llm._torch.moe.fused_moe.impl_contract import ( + MoEDeployment, + MoEEnvironment, + MoEProblem, + ) + from tensorrt_llm._torch.moe.fused_moe.impl_environment import MoEDep + + problem = MoEProblem(quant="NVFP4", dtype_act=torch.bfloat16) + deployment = MoEDeployment( + ep_size=1, + tp_size=1, + parallel_size=1, + use_dp=False, + num_slots=8, + env=MoEEnvironment( + sm=107, + available_deps=(MoEDep.CUTEDSL_RUBIN.value, MoEDep.LOCALITY_DOMAIN.value), + ), + **deployment_kwargs, + ) + return CuteDslFusedMoE.can_implement(problem, deployment) + + +def test_cute_dsl_locality_domain_rejects_eplb(): + """Localized shards cannot follow EPLB migration, so selection declines.""" + from tensorrt_llm._torch.moe.fused_moe.impl_contract import MoERejectReason + + verdict = _cute_dsl_eligibility(locality_domain_requested=True, eplb_enabled=True) + assert not verdict.eligible + assert verdict.reject_reason is MoERejectReason.EPLB_UNSUPPORTED + assert _cute_dsl_eligibility(locality_domain_requested=True).eligible + assert _cute_dsl_eligibility(eplb_enabled=True).eligible + + +def test_cute_dsl_sm107_requires_fused_finalize(): + """SM107 has no unfused FC2, so disabling finalize fusion declines.""" + from tensorrt_llm._torch.moe.fused_moe.impl_contract import MoERejectReason + + verdict = _cute_dsl_eligibility(fused_finalize_enabled=False) + assert not verdict.eligible + assert verdict.reject_reason is MoERejectReason.FINALIZE_FUSION_REQUIRED + + +def test_cute_dsl_unquantized_rejects_non_swiglu(): + """The BF16 FC1 op fuses SwiGLU by name; NVFP4 serves both activations.""" + from tensorrt_llm._torch.moe.fused_moe.impl_contract import ( + MoEDeployment, + MoEEnvironment, + MoEProblem, + MoERejectReason, + ) + from tensorrt_llm._torch.moe.fused_moe.impl_environment import MoEDep + + env = MoEEnvironment(sm=107, available_deps=(MoEDep.CUTEDSL_RUBIN.value,)) + deployment = MoEDeployment( + ep_size=1, tp_size=1, parallel_size=1, use_dp=False, num_slots=8, env=env + ) + + def verdict(quant): + return CuteDslFusedMoE.can_implement( + MoEProblem(quant=quant, dtype_act=torch.bfloat16, activation="Relu2"), deployment + ) + + assert not verdict(None).eligible + assert verdict(None).reject_reason is MoERejectReason.ACTIVATION_UNSUPPORTED + assert verdict("NVFP4").eligible + + +def test_cute_dsl_locality_domain_disables_dwdp(): + """DWDP rebinds parameters, so it stays off rather than raising.""" + from tensorrt_llm._torch.moe.fused_moe.configurable_moe import ConfigurableMoE + + moe = SimpleNamespace( + backend=SimpleNamespace( + capabilities=SimpleNamespace(supports_dwdp=True), + uses_locality_domain=True, + ) + ) + assert ConfigurableMoE._should_enable_dwdp(moe) is False + + +@pytest.mark.parametrize("quantized", [True, False], ids=["nvfp4", "bf16"]) +def test_cute_dsl_moe_zero_tokens_short_circuits_autotune(quantized: bool): + hidden_size = 8 + top_k = 2 + token_selected_experts = torch.empty((0, top_k), dtype=torch.int32) + token_final_scales = torch.empty((0, top_k), dtype=torch.float32) + moe_output = torch.empty((0, hidden_size), dtype=torch.bfloat16) + + if quantized: + backend = SimpleNamespace( + has_nvfp4=True, + has_deepseek_fp8_block_scales=False, + has_any_quant=True, + activation_type=ActivationType.Swiglu, + _locality_domain_runtime=object(), + hidden_size=hidden_size, + scaling_vector_size=4, + ) + x, x_sf = CuteDslFusedMoE.quantize_input( + backend, + Fp4QuantizedTensor( + fp4_tensor=torch.empty((0, hidden_size // 2), dtype=torch.uint8), + scaling_factor=torch.empty(0, dtype=torch.uint8), + is_sf_swizzled=False, + ), + ) + assert x_sf.shape == (0, hidden_size // backend.scaling_vector_size) + output = CuteDslFusedMoE.run_moe_nvfp4( + backend, + x, + token_selected_experts, + token_final_scales, + x_sf=x_sf, + moe_output=moe_output, + weight_view=object(), + ) + else: + backend = SimpleNamespace( + has_any_quant=False, + hidden_size=hidden_size, + use_fused_finalize=True, + ) + output = CuteDslFusedMoE.run_moe_bf16( + backend, + torch.empty((0, hidden_size), dtype=torch.bfloat16), + token_selected_experts, + token_final_scales, + moe_output=moe_output, + ) + + assert output is moe_output + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on Rubin (SM 107) GPUs", +) +def test_cute_dsl_nvfp4_quantize_empty_input_rubin(): + hidden_size = 32 + scaling_vector_size = 16 + backend = SimpleNamespace( + has_nvfp4=True, + has_deepseek_fp8_block_scales=False, + has_any_quant=True, + fc31_input_scale=torch.ones(1, dtype=torch.float32, device="cuda"), + hidden_size=hidden_size, + scaling_vector_size=scaling_vector_size, + ) + + x, x_sf = CuteDslFusedMoE.quantize_input( + backend, + torch.empty((0, hidden_size), dtype=torch.bfloat16, device="cuda"), + ) + + assert x.shape == (0, hidden_size // 2) + assert x_sf.shape == (0, hidden_size // scaling_vector_size) + + +def test_sm107_nvfp4_tile512_fallback_uses_two_cta_cluster(): + mma_tiler, mma_inst_shape, cluster_shape_mn = _get_sm107_nvfp4_default_mma_config(512) + + assert mma_tiler == (512, 128, 256) + assert mma_inst_shape == (256, 128, 128) + assert cluster_shape_mn == (2, 1) + + +def _make_sm107_nvfp4_finalize_inputs(n: int = 384) -> list[torch.Tensor]: + m = 128 + packed_k = 32 + scale_k = packed_k * 2 // 16 + num_tokens = 2 + return [ + torch.empty((m, packed_k), dtype=torch.float4_e2m1fn_x2), + torch.empty((1, n, packed_k), dtype=torch.float4_e2m1fn_x2), + torch.empty(m * scale_k, dtype=torch.uint8), + torch.empty((1, n, scale_k), dtype=torch.uint8), + torch.empty(1, dtype=torch.float32), + torch.empty((num_tokens, n), dtype=torch.bfloat16), + torch.empty(1, dtype=torch.int32), + torch.empty(1, dtype=torch.int32), + torch.empty(m, dtype=torch.int32), + torch.empty(1, dtype=torch.int32), + torch.empty((num_tokens, 1), dtype=torch.float32), + ] + + +def _make_sm107_nvfp4_finalize_runner(monkeypatch): + monkeypatch.setattr(cute_dsl_custom_ops, "get_sm_version", lambda: 107) + runner_cls = cute_dsl_custom_ops.Sm107BlockScaledContiguousGroupedGemmFinalizeFusionRunner + return runner_cls( + num_experts=1, + top_k=1, + num_local_experts=1, + local_expert_offset=0, + tile_size=128, + output_dtype=torch.bfloat16, + ) + + +@pytest.mark.skipif( + not IS_CUTLASS_DSL_RUBIN_AVAILABLE, + reason="This test requires the public Rubin CuTe DSL package", +) +def test_sm107_nvfp4_finalize_get_valid_tactics_filters_n_tiling(monkeypatch): + runner = _make_sm107_nvfp4_finalize_runner(monkeypatch) + monkeypatch.setattr(runner.kernel_class, "can_implement", lambda **_: True) + + tactics = runner.get_valid_tactics(_make_sm107_nvfp4_finalize_inputs(), OptimizationProfile()) + + assert tactics + assert all(384 % (tactic[0][1] * tactic[2][1]) == 0 for tactic in tactics) + + +@pytest.mark.skipif( + not IS_CUTLASS_DSL_RUBIN_AVAILABLE, + reason="This test requires the public Rubin CuTe DSL package", +) +def test_sm107_nvfp4_finalize_forward_rejects_incompatible_n_tiling(monkeypatch): + runner = _make_sm107_nvfp4_finalize_runner(monkeypatch) + incompatible_tactic = ((128, 128, 256), (128, 128, 128), (1, 2), False) + + with pytest.raises(ValueError, match="incompatible with N=384"): + runner(_make_sm107_nvfp4_finalize_inputs(), tactic=incompatible_tactic) + + +@pytest.mark.skipif( + not IS_CUTLASS_DSL_RUBIN_AVAILABLE, + reason="This test requires the public Rubin CuTe DSL package", +) +def test_sm107_nvfp4_finalize_accepts_two_cta_n_tiling(): + runner_cls = cute_dsl_custom_ops.Sm107BlockScaledContiguousGroupedGemmFinalizeFusionRunner + + assert runner_cls._is_n_tiling_compatible(n=256, mma_n=128, cluster_n=2) + + +def _get_registered_rubin_moe_op(name: str): + try: + return getattr(torch.ops.trtllm, name) + except AttributeError: + pytest.skip("public Rubin CuTe DSL MoE ops are not registered") + + +def _assert_rubin_moe_op_schema( + op_name: str, + argument_names: tuple[str, ...], + default_values: dict[str, bool | int | None], + mutated_arguments: set[str], + return_types: tuple[str, ...], +) -> None: + schema = _get_registered_rubin_moe_op(op_name).default._schema + + assert tuple(argument.name for argument in schema.arguments) == argument_names + assert { + argument.name: argument.default_value + for argument in schema.arguments + if argument.has_default_value() + } == default_values + assert { + argument.name + for argument in schema.arguments + if argument.alias_info is not None and argument.alias_info.is_write + } == mutated_arguments + assert tuple(str(result.type) for result in schema.returns) == return_types + + +@pytest.mark.parametrize( + "op_name,argument_names,default_values,mutated_arguments,return_types", + [ + pytest.param( + "cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin", + ( + "input", + "weight", + "input_scale", + "weight_scale", + "alpha", + "tile_idx_to_group_idx", + "tile_idx_to_mn_limit", + "permuted_idx_to_expanded_idx", + "num_non_exiting_tiles", + "global_sf", + "num_experts", + "top_k", + "num_local_experts", + "local_expert_offset", + "tile_size", + "output_tensor", + "output_sf_tensor", + "scaling_vector_size", + "partition_id", + "activation_type", + "precomputed_tactic", + ), + { + "scaling_vector_size": 16, + "partition_id": -1, + "activation_type": int(ActivationType.Swiglu), + "precomputed_tactic": None, + }, + {"output_tensor", "output_sf_tensor"}, + ("Optional[Tensor]", "Optional[Tensor]"), + id="nvfp4_fc1", + ), + pytest.param( + "cute_dsl_nvfp4_grouped_gemm_finalize_inplace_rubin", + ( + "input", + "weight", + "input_scale", + "weight_scale", + "alpha", + "output", + "tile_idx_to_group_idx", + "tile_idx_to_mn_limit", + "permuted_idx_to_expanded_idx", + "num_non_exiting_tiles", + "token_final_scales", + "num_experts", + "top_k", + "num_local_experts", + "local_expert_offset", + "tile_size", + "output_dtype", + "scaling_vector_size", + "precomputed_tactic", + ), + {"scaling_vector_size": 16, "precomputed_tactic": None}, + {"output"}, + (), + id="nvfp4_fc2", + ), + pytest.param( + "cute_dsl_bf16_gather_grouped_gemm_swiglu_rubin", + ( + "input", + "weight", + "alpha", + "tile_idx_to_group_idx", + "tile_idx_to_mn_limit", + "permuted_idx_to_expanded_idx", + "num_non_exiting_tiles", + "num_experts", + "top_k", + "num_local_experts", + "local_expert_offset", + "tile_size", + "output_tensor", + "partition_id", + "precomputed_tactic", + ), + {"precomputed_tactic": None}, + {"output_tensor"}, + ("Optional[Tensor]",), + id="bf16_fc1", + ), + pytest.param( + "cute_dsl_bf16_grouped_gemm_finalize_inplace_rubin", + ( + "input", + "weight", + "output", + "tile_idx_to_group_idx", + "tile_idx_to_mn_limit", + "permuted_idx_to_expanded_idx", + "num_non_exiting_tiles", + "token_final_scales", + "num_experts", + "top_k", + "num_local_experts", + "local_expert_offset", + "tile_size", + "output_dtype", + "precomputed_tactic", + ), + {"precomputed_tactic": None}, + {"output"}, + (), + id="bf16_fc2", + ), + ], +) +def test_rubin_moe_leaf_schema( + op_name: str, + argument_names: tuple[str, ...], + default_values: dict[str, bool | int | None], + mutated_arguments: set[str], + return_types: tuple[str, ...], +) -> None: + _assert_rubin_moe_op_schema( + op_name, + argument_names, + default_values, + mutated_arguments, + return_types, + ) + + +def _make_fake_rubin_moe_tensors(quantized: bool) -> SimpleNamespace: + dtype = torch.uint8 if quantized else torch.bfloat16 + fc1_n = 32 if quantized else 8 + fc2_n = 32 if quantized else 4 + output_n = 32 if quantized else 8 + return SimpleNamespace( + input=torch.empty((2, 4), dtype=dtype, device="cuda"), + fc1_weight=torch.empty((1, fc1_n, 4), dtype=dtype, device="cuda"), + fc2_weight=torch.empty((1, fc2_n, 4), dtype=dtype, device="cuda"), + input_scale=torch.empty(2, dtype=torch.uint8, device="cuda"), + fc1_weight_scale=torch.empty((1, fc1_n, 1), dtype=torch.uint8, device="cuda"), + fc2_weight_scale=torch.empty((1, fc2_n, 1), dtype=torch.uint8, device="cuda"), + alpha=torch.empty(1, dtype=torch.float32, device="cuda"), + tile_idx_to_group_idx=torch.empty(1, dtype=torch.int32, device="cuda"), + tile_idx_to_mn_limit=torch.empty(1, dtype=torch.int32, device="cuda"), + expanded_idx_to_permuted_idx=torch.empty((2, 1), dtype=torch.int32, device="cuda"), + permuted_idx_to_expanded_idx=torch.empty(2, dtype=torch.int32, device="cuda"), + num_non_exiting_tiles=torch.empty(1, dtype=torch.int32, device="cuda"), + global_sf=torch.empty(1, dtype=torch.float32, device="cuda"), + token_final_scales=torch.empty((2, 1), dtype=torch.float32, device="cuda"), + fc1_output=torch.empty((2, 8), dtype=dtype, device="cuda"), + fc1_output_sf=torch.empty(2, dtype=torch.uint8, device="cuda"), + output=torch.empty((2, output_n), dtype=torch.bfloat16, device="cuda"), + ) + + +def test_rubin_moe_precomputed_tactic_fake_signatures(): + try: + from torch._subclasses.fake_tensor import FakeTensorMode + except ImportError: + pytest.skip("FakeTensorMode is not available") + + fc1_op = _get_registered_rubin_moe_op("cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin") + fc2_op = _get_registered_rubin_moe_op("cute_dsl_nvfp4_grouped_gemm_finalize_inplace_rubin") + + with FakeTensorMode(): + tensors = _make_fake_rubin_moe_tensors(quantized=True) + + output, output_sf = fc1_op( + input=tensors.input, + weight=tensors.fc1_weight, + input_scale=tensors.input_scale, + weight_scale=tensors.fc1_weight_scale, + alpha=tensors.alpha, + tile_idx_to_group_idx=tensors.tile_idx_to_group_idx, + tile_idx_to_mn_limit=tensors.tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=tensors.permuted_idx_to_expanded_idx, + num_non_exiting_tiles=tensors.num_non_exiting_tiles, + global_sf=tensors.global_sf, + num_experts=1, + top_k=1, + num_local_experts=1, + local_expert_offset=0, + tile_size=128, + output_tensor=None, + output_sf_tensor=None, + scaling_vector_size=16, + partition_id=-1, + activation_type=int(ActivationType.Swiglu), + precomputed_tactic=repr(-1), + ) + assert output.shape == (2, 8) + assert output_sf.shape == (2,) + + result = fc2_op( + input=tensors.input, + weight=tensors.fc2_weight, + input_scale=tensors.input_scale, + weight_scale=tensors.fc2_weight_scale, + alpha=tensors.alpha, + output=tensors.output, + tile_idx_to_group_idx=tensors.tile_idx_to_group_idx, + tile_idx_to_mn_limit=tensors.tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=tensors.permuted_idx_to_expanded_idx, + num_non_exiting_tiles=tensors.num_non_exiting_tiles, + token_final_scales=tensors.token_final_scales, + num_experts=1, + top_k=1, + num_local_experts=1, + local_expert_offset=0, + tile_size=128, + output_dtype=torch.bfloat16, + scaling_vector_size=16, + precomputed_tactic=repr(-1), + ) + assert result is None + + +def test_rubin_bf16_moe_precomputed_tactic_fake_signatures(): + try: + from torch._subclasses.fake_tensor import FakeTensorMode + except ImportError: + pytest.skip("FakeTensorMode is not available") + + fc1_op = _get_registered_rubin_moe_op("cute_dsl_bf16_gather_grouped_gemm_swiglu_rubin") + fc2_op = _get_registered_rubin_moe_op("cute_dsl_bf16_grouped_gemm_finalize_inplace_rubin") + + with FakeTensorMode(): + tensors = _make_fake_rubin_moe_tensors(quantized=False) + + result = fc1_op( + input=tensors.input, + weight=tensors.fc1_weight, + alpha=tensors.alpha, + tile_idx_to_group_idx=tensors.tile_idx_to_group_idx, + tile_idx_to_mn_limit=tensors.tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=tensors.permuted_idx_to_expanded_idx, + num_non_exiting_tiles=tensors.num_non_exiting_tiles, + num_experts=1, + top_k=1, + num_local_experts=1, + local_expert_offset=0, + tile_size=128, + output_tensor=tensors.fc1_output, + partition_id=0, + precomputed_tactic=repr(-1), + ) + assert result is None + + result = fc2_op( + input=tensors.input, + weight=tensors.fc2_weight, + output=tensors.output, + tile_idx_to_group_idx=tensors.tile_idx_to_group_idx, + tile_idx_to_mn_limit=tensors.tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=tensors.permuted_idx_to_expanded_idx, + num_non_exiting_tiles=tensors.num_non_exiting_tiles, + token_final_scales=tensors.token_final_scales, + num_experts=1, + top_k=1, + num_local_experts=1, + local_expert_offset=0, + tile_size=128, + output_dtype=torch.bfloat16, + precomputed_tactic=repr(-1), + ) + assert result is None + + +@pytest.mark.parametrize( + "op_name,argument_names,default_values,mutated_arguments", + [ + pytest.param( + "cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_locality_domain_inplace_rubin", + ( + "input", + "weight_0", + "weight_1", + "input_scale", + "weight_scale_0", + "weight_scale_1", + "alpha", + "tile_idx_to_group_idx", + "tile_idx_to_mn_limit", + "permuted_idx_to_expanded_idx", + "num_non_exiting_tiles", + "global_sf", + "num_experts", + "top_k", + "num_local_experts", + "local_expert_offset", + "tile_size", + "output_tensor", + "output_sf_tensor", + "scaling_vector_size", + "activation_type", + ), + { + "scaling_vector_size": 16, + "activation_type": int(ActivationType.Swiglu), + }, + {"output_tensor", "output_sf_tensor"}, + id="nvfp4_fc1", + ), + pytest.param( + "cute_dsl_nvfp4_grouped_gemm_finalize_locality_domain_inplace_rubin", + ( + "input", + "weight_0", + "weight_1", + "input_scale", + "weight_scale_0", + "weight_scale_1", + "alpha", + "output", + "tile_idx_to_group_idx", + "tile_idx_to_mn_limit", + "expanded_idx_to_permuted_idx", + "permuted_idx_to_expanded_idx", + "num_non_exiting_tiles", + "token_final_scales", + "num_experts", + "top_k", + "num_local_experts", + "local_expert_offset", + "tile_size", + "output_dtype", + "ep_size", + "enable_alltoall", + "scaling_vector_size", + ), + {"enable_alltoall": False, "scaling_vector_size": 16}, + {"output"}, + id="nvfp4_fc2", + ), + pytest.param( + "cute_dsl_bf16_gather_grouped_gemm_swiglu_locality_domain_inplace_rubin", + ( + "input", + "weight_0", + "weight_1", + "alpha", + "tile_idx_to_group_idx", + "tile_idx_to_mn_limit", + "permuted_idx_to_expanded_idx", + "num_non_exiting_tiles", + "num_experts", + "top_k", + "num_local_experts", + "local_expert_offset", + "tile_size", + "output_tensor", + ), + {}, + {"output_tensor"}, + id="bf16_fc1", + ), + pytest.param( + "cute_dsl_bf16_grouped_gemm_finalize_locality_domain_inplace_rubin", + ( + "input", + "weight_0", + "weight_1", + "output", + "tile_idx_to_group_idx", + "tile_idx_to_mn_limit", + "expanded_idx_to_permuted_idx", + "permuted_idx_to_expanded_idx", + "num_non_exiting_tiles", + "token_final_scales", + "num_experts", + "top_k", + "num_local_experts", + "local_expert_offset", + "tile_size", + "output_dtype", + "ep_size", + "enable_alltoall", + ), + {"enable_alltoall": False}, + {"output"}, + id="bf16_fc2", + ), + ], +) +def test_rubin_moe_locality_domain_composite_schema( + op_name: str, + argument_names: tuple[str, ...], + default_values: dict[str, bool | int | None], + mutated_arguments: set[str], +) -> None: + _assert_rubin_moe_op_schema( + op_name, + argument_names, + default_values, + mutated_arguments, + (), + ) + + +def test_rubin_nvfp4_moe_locality_domain_composite_fake_signatures(): + try: + from torch._subclasses.fake_tensor import FakeTensorMode + except ImportError: + pytest.skip("FakeTensorMode is not available") + + fc1_op = _get_registered_rubin_moe_op( + "cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_locality_domain_inplace_rubin" + ) + fc2_op = _get_registered_rubin_moe_op( + "cute_dsl_nvfp4_grouped_gemm_finalize_locality_domain_inplace_rubin" + ) + + with FakeTensorMode(): + tensors = _make_fake_rubin_moe_tensors(quantized=True) + + result = fc1_op( + input=tensors.input, + weight_0=tensors.fc1_weight, + weight_1=tensors.fc1_weight, + input_scale=tensors.input_scale, + weight_scale_0=tensors.fc1_weight_scale, + weight_scale_1=tensors.fc1_weight_scale, + alpha=tensors.alpha, + tile_idx_to_group_idx=tensors.tile_idx_to_group_idx, + tile_idx_to_mn_limit=tensors.tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=tensors.permuted_idx_to_expanded_idx, + num_non_exiting_tiles=tensors.num_non_exiting_tiles, + global_sf=tensors.global_sf, + num_experts=1, + top_k=1, + num_local_experts=1, + local_expert_offset=0, + tile_size=128, + output_tensor=tensors.fc1_output, + output_sf_tensor=tensors.fc1_output_sf, + scaling_vector_size=16, + activation_type=int(ActivationType.Swiglu), + ) + assert result is None + + result = fc2_op( + input=tensors.input, + weight_0=tensors.fc2_weight, + weight_1=tensors.fc2_weight, + input_scale=tensors.input_scale, + weight_scale_0=tensors.fc2_weight_scale, + weight_scale_1=tensors.fc2_weight_scale, + alpha=tensors.alpha, + output=tensors.output, + tile_idx_to_group_idx=tensors.tile_idx_to_group_idx, + tile_idx_to_mn_limit=tensors.tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx=tensors.expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx=tensors.permuted_idx_to_expanded_idx, + num_non_exiting_tiles=tensors.num_non_exiting_tiles, + token_final_scales=tensors.token_final_scales, + num_experts=1, + top_k=1, + num_local_experts=1, + local_expert_offset=0, + tile_size=128, + output_dtype=torch.bfloat16, + ep_size=1, + enable_alltoall=False, + scaling_vector_size=16, + ) + assert result is None + + +def test_rubin_bf16_moe_locality_domain_composite_fake_signatures(): + try: + from torch._subclasses.fake_tensor import FakeTensorMode + except ImportError: + pytest.skip("FakeTensorMode is not available") + + fc1_op = _get_registered_rubin_moe_op( + "cute_dsl_bf16_gather_grouped_gemm_swiglu_locality_domain_inplace_rubin" + ) + fc2_op = _get_registered_rubin_moe_op( + "cute_dsl_bf16_grouped_gemm_finalize_locality_domain_inplace_rubin" + ) + + with FakeTensorMode(): + tensors = _make_fake_rubin_moe_tensors(quantized=False) + + result = fc1_op( + input=tensors.input, + weight_0=tensors.fc1_weight, + weight_1=tensors.fc1_weight, + alpha=tensors.alpha, + tile_idx_to_group_idx=tensors.tile_idx_to_group_idx, + tile_idx_to_mn_limit=tensors.tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=tensors.permuted_idx_to_expanded_idx, + num_non_exiting_tiles=tensors.num_non_exiting_tiles, + num_experts=1, + top_k=1, + num_local_experts=1, + local_expert_offset=0, + tile_size=128, + output_tensor=tensors.fc1_output, + ) + assert result is None + + result = fc2_op( + input=tensors.input, + weight_0=tensors.fc2_weight, + weight_1=tensors.fc2_weight, + output=tensors.output, + tile_idx_to_group_idx=tensors.tile_idx_to_group_idx, + tile_idx_to_mn_limit=tensors.tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx=tensors.expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx=tensors.permuted_idx_to_expanded_idx, + num_non_exiting_tiles=tensors.num_non_exiting_tiles, + token_final_scales=tensors.token_final_scales, + num_experts=1, + top_k=1, + num_local_experts=1, + local_expert_offset=0, + tile_size=128, + output_dtype=torch.bfloat16, + ep_size=1, + enable_alltoall=False, + ) + assert result is None + + +@pytest.mark.parametrize( + "composite_name,leaf_name,runner_name,quantized,is_fc1", + [ + ( + "cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_locality_domain_inplace_rubin", + "cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin", + "Sm107BlockScaledContiguousGatherGroupedGemmActFusionRunner", + True, + True, + ), + ( + "cute_dsl_nvfp4_grouped_gemm_finalize_locality_domain_inplace_rubin", + "cute_dsl_nvfp4_grouped_gemm_finalize_inplace_rubin", + "Sm107BlockScaledContiguousGroupedGemmFinalizeFusionRunner", + True, + False, + ), + ( + "cute_dsl_bf16_gather_grouped_gemm_swiglu_locality_domain_inplace_rubin", + "cute_dsl_bf16_gather_grouped_gemm_swiglu_rubin", + "Sm107ContiguousGatherGroupedGemmSwigluFusionRunner", + False, + True, + ), + ( + "cute_dsl_bf16_grouped_gemm_finalize_locality_domain_inplace_rubin", + "cute_dsl_bf16_grouped_gemm_finalize_inplace_rubin", + "Sm107ContiguousGroupedGemmFinalizeFusionRunner", + False, + False, + ), + ], +) +def test_rubin_moe_locality_domain_composite_owns_concurrent_tuning( + monkeypatch, + composite_name: str, + leaf_name: str, + runner_name: str, + quantized: bool, + is_fc1: bool, +): + composite_op = getattr(cute_dsl_custom_ops, composite_name, None) + if composite_op is None: + pytest.skip("public Rubin CuTe DSL MoE composite op is not registered") + + tuning_config = object() + chosen_tactic = ("chosen", 128) + runner_instances = [] + runtime_instances = [] + tune_calls = [] + concurrent_calls = [] + leaf_calls = [] + memset_calls = [] + execution_order = [] + tuner_state = SimpleNamespace(is_tuning_mode=not is_fc1) + + class FakeOpRunner: + def __init__(self, *args, **kwargs): + self.init_call = (args, kwargs) + self.tuning_config_calls = [] + runner_instances.append(self) + + def get_tuning_config(self, *args, **kwargs): + self.tuning_config_calls.append((args, kwargs)) + return tuning_config + + class FakeRuntime: + def __init__(self, num_partitions: int): + self.num_partitions = num_partitions + runtime_instances.append(self) + + class FakeConcurrentRunner: + def __init__(self, launch_partition): + self.launch_partition = launch_partition + + def __call__(self, inputs, *, tactic): + execution_order.append("launch") + concurrent_calls.append((inputs, tactic)) + for partition_id in range(2): + self.launch_partition(partition_id, inputs, tactic) + + def fake_tune_locality_domain_concurrent( + op_name, + op_runner, + runtime, + num_partitions, + launch_partition, + inputs, + actual_tuning_config, + ): + execution_order.append("tune") + tune_calls.append( + ( + op_name, + op_runner, + runtime, + num_partitions, + launch_partition, + inputs, + actual_tuning_config, + ) + ) + return FakeConcurrentRunner(launch_partition), chosen_tactic + + def fake_leaf_op(*args, **kwargs): + leaf_calls.append((args, kwargs)) + + def fake_moe_output_memset(*args, **kwargs): + execution_order.append("reset") + memset_calls.append((args, kwargs)) + + monkeypatch.setattr(cute_dsl_custom_ops, runner_name, FakeOpRunner) + monkeypatch.setattr(cute_dsl_custom_ops, "LocalityDomainRuntime", FakeRuntime) + monkeypatch.setattr( + cute_dsl_custom_ops, + "tune_locality_domain_concurrent", + fake_tune_locality_domain_concurrent, + ) + monkeypatch.setattr(cute_dsl_custom_ops, "get_sm_version", lambda: 107) + monkeypatch.setattr( + AutoTuner, + "get", + staticmethod(lambda: tuner_state), + ) + monkeypatch.setattr(torch.ops.trtllm, leaf_name, fake_leaf_op) + monkeypatch.setattr( + torch.ops.trtllm, + "moe_output_memset_inplace", + fake_moe_output_memset, + ) + + dtype = torch.uint8 if quantized else torch.bfloat16 + input_tensor = torch.empty((2, 4), dtype=dtype) + weight_0 = torch.empty((1, 8, 4), dtype=dtype) + weight_1 = torch.empty_like(weight_0) + alpha = torch.empty((1,), dtype=torch.float32) + tile_idx_to_group_idx = torch.empty((1,), dtype=torch.int32) + tile_idx_to_mn_limit = torch.empty((1,), dtype=torch.int32) + expanded_idx_to_permuted_idx = torch.empty((2, 1), dtype=torch.int32) + permuted_idx_to_expanded_idx = torch.empty((2,), dtype=torch.int32) + num_non_exiting_tiles = torch.empty((1,), dtype=torch.int32) + token_final_scales = torch.empty((2, 1), dtype=torch.float32) + output = torch.empty((2, 16), dtype=torch.bfloat16) + + common_kwargs = { + "input": input_tensor, + "weight_0": weight_0, + "weight_1": weight_1, + "tile_idx_to_group_idx": tile_idx_to_group_idx, + "tile_idx_to_mn_limit": tile_idx_to_mn_limit, + "permuted_idx_to_expanded_idx": permuted_idx_to_expanded_idx, + "num_non_exiting_tiles": num_non_exiting_tiles, + "num_experts": 1, + "top_k": 1, + "num_local_experts": 1, + "local_expert_offset": 0, + "tile_size": 128, + } + if is_fc1: + call_kwargs = { + **common_kwargs, + "alpha": alpha, + "output_tensor": output if not quantized else torch.empty((2, 8), dtype=torch.uint8), + } + if quantized: + weight_scale_0 = torch.empty((1, 8, 1), dtype=torch.uint8) + weight_scale_1 = torch.empty_like(weight_scale_0) + call_kwargs.update( + { + "input_scale": torch.empty((2,), dtype=torch.uint8), + "weight_scale_0": weight_scale_0, + "weight_scale_1": weight_scale_1, + "global_sf": torch.empty((1,), dtype=torch.float32), + "output_sf_tensor": torch.empty((2,), dtype=torch.uint8), + "scaling_vector_size": 16, + "activation_type": int(ActivationType.Swiglu), + } + ) + else: + call_kwargs = { + **common_kwargs, + "output": output, + "expanded_idx_to_permuted_idx": expanded_idx_to_permuted_idx, + "token_final_scales": token_final_scales, + "output_dtype": torch.bfloat16, + "ep_size": 4, + "enable_alltoall": True, + } + if quantized: + weight_scale_0 = torch.empty((1, 8, 1), dtype=torch.uint8) + weight_scale_1 = torch.empty_like(weight_scale_0) + call_kwargs.update( + { + "input_scale": torch.empty((2,), dtype=torch.uint8), + "weight_scale_0": weight_scale_0, + "weight_scale_1": weight_scale_1, + "alpha": alpha, + "scaling_vector_size": 16, + } + ) + + # Exercise the registered Python implementation directly; dispatcher and + # fake contracts are covered separately by the schema/fake tests above. + assert composite_op._init_fn(**call_kwargs) is None + + assert len(runner_instances) == 1 + assert len(runtime_instances) == 1 + assert runtime_instances[0].num_partitions == 2 + assert len(tune_calls) == 1 + ( + tune_key, + tuned_op_runner, + tuned_runtime, + num_partitions, + _, + tuned_inputs, + actual_tuning_config, + ) = tune_calls[0] + assert tune_key == f"trtllm::{leaf_name}" + assert tuned_op_runner is runner_instances[0] + assert tuned_runtime is runtime_instances[0] + assert num_partitions == 2 + assert actual_tuning_config is tuning_config + assert len(runner_instances[0].tuning_config_calls) == 1 + runner_args, runner_kwargs = runner_instances[0].init_call + if quantized and is_fc1: + assert runner_args == (1, 1, 1, 0, 128, 16) + assert runner_kwargs == {"activation_type": ActivationType.Swiglu} + elif quantized: + assert runner_args == (1, 1, 1, 0, 128, torch.bfloat16, 16) + assert not runner_kwargs + elif is_fc1: + assert runner_args == (1, 1, 1, 0, 128) + assert runner_kwargs == {"input_dtype": torch.bfloat16} + assert runner_instances[0].tuning_config_calls == [((), {"has_output_tensor": True})] + else: + assert runner_args == (1, 1, 1, 0, 128, torch.bfloat16) + assert runner_kwargs == {"input_dtype": torch.bfloat16} + if quantized or not is_fc1: + assert runner_instances[0].tuning_config_calls == [((), {})] + assert len(concurrent_calls) == 1 + assert concurrent_calls[0][0] is tuned_inputs + assert concurrent_calls[0][1] == chosen_tactic + + assert len(leaf_calls) == 2 + assert leaf_calls[0][1]["weight"] is weight_0 + assert leaf_calls[1][1]["weight"] is weight_1 + for _, leaf_kwargs in leaf_calls: + assert leaf_kwargs["input"] is input_tensor + assert leaf_kwargs["tile_idx_to_group_idx"] is tile_idx_to_group_idx + assert leaf_kwargs["tile_idx_to_mn_limit"] is tile_idx_to_mn_limit + assert leaf_kwargs["permuted_idx_to_expanded_idx"] is permuted_idx_to_expanded_idx + assert leaf_kwargs["num_non_exiting_tiles"] is num_non_exiting_tiles + assert leaf_kwargs["num_experts"] == 1 + assert leaf_kwargs["top_k"] == 1 + assert leaf_kwargs["num_local_experts"] == 1 + assert leaf_kwargs["local_expert_offset"] == 0 + assert leaf_kwargs["tile_size"] == 128 + assert leaf_kwargs["precomputed_tactic"] == repr(chosen_tactic) + if is_fc1: + assert [kwargs["partition_id"] for _, kwargs in leaf_calls] == [0, 1] + assert all(kwargs["alpha"] is alpha for _, kwargs in leaf_calls) + assert all( + kwargs["output_tensor"] is call_kwargs["output_tensor"] for _, kwargs in leaf_calls + ) + assert not memset_calls + assert execution_order == ["tune", "launch"] + else: + assert all("partition_id" not in kwargs for _, kwargs in leaf_calls) + assert all(kwargs["output"] is output for _, kwargs in leaf_calls) + assert all(kwargs["token_final_scales"] is token_final_scales for _, kwargs in leaf_calls) + assert all(kwargs["output_dtype"] == torch.bfloat16 for _, kwargs in leaf_calls) + assert execution_order == ["tune", "reset", "launch"] + assert len(memset_calls) == 1 + reset_args, reset_kwargs = memset_calls[0] + assert not reset_args + assert reset_kwargs == { + "input": output, + "tile_idx_to_mn_limit": tile_idx_to_mn_limit, + "expanded_idx_to_permuted_idx": expanded_idx_to_permuted_idx, + "permuted_idx_to_expanded_idx": permuted_idx_to_expanded_idx, + "num_non_exiting_tiles": num_non_exiting_tiles, + "tile_tokens_dim": 128, + "top_k": 1, + "ep_size": 4, + "enable_alltoall": True, + } + execution_order.clear() + memset_calls.clear() + tuner_state.is_tuning_mode = False + assert composite_op._init_fn(**call_kwargs) is None + assert execution_order == ["tune", "launch"] + assert not memset_calls + if quantized: + assert leaf_calls[0][1]["weight_scale"] is weight_scale_0 + assert leaf_calls[1][1]["weight_scale"] is weight_scale_1 + assert all(kwargs["input_scale"] is call_kwargs["input_scale"] for _, kwargs in leaf_calls) + assert all(kwargs["alpha"] is alpha for _, kwargs in leaf_calls) + assert all(kwargs["scaling_vector_size"] == 16 for _, kwargs in leaf_calls) + if is_fc1: + assert all(kwargs["global_sf"] is call_kwargs["global_sf"] for _, kwargs in leaf_calls) + assert all( + kwargs["output_sf_tensor"] is call_kwargs["output_sf_tensor"] + for _, kwargs in leaf_calls + ) + assert all( + kwargs["activation_type"] == int(ActivationType.Swiglu) for _, kwargs in leaf_calls + ) + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on Rubin (SM 107) GPUs", +) +@pytest.mark.parametrize("tile_size", [128]) +@pytest.mark.parametrize("ep_size", [1, 8]) +@pytest.mark.parametrize("top_k", [1, 2]) +@pytest.mark.parametrize("num_tokens", [128, 515]) +def test_nvfp4_gather_grouped_gemm_swiglu_locality_domain_rubin( + num_tokens: int, top_k: int, ep_size: int, tile_size: int +): + _skip_if_no_locality_domain() + + sf_vec_size = 16 + hidden_size = 2048 + interm_size = 1536 + num_experts = 256 + num_local_experts = num_experts // ep_size + + ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles, + _, + ) = _setup_locality_domain_routing(num_tokens, num_experts, num_local_experts, top_k, tile_size) + max_num_permuted_tokens = permuted_idx_to_expanded_idx.size(0) + + a_fp4, a_sf_unswizzled, a_global_sf = _create_quantized_locality_domain_inputs( + num_tokens, hidden_size, sf_vec_size + ) + weight, weight_sf, weight_global_sf = _create_quantized_locality_domain_weights( + num_local_experts, interm_size, hidden_size, sf_vec_size + ) + alpha = a_global_sf * weight_global_sf + global_sf = torch.tensor([1.0], dtype=torch.float32, device="cuda") + + c_ref, c_sf_ref = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( + a_fp4, + weight, + a_sf_unswizzled, + weight_sf, + alpha, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + global_sf, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=0, + tile_size=tile_size, + scaling_vector_size=sf_vec_size, + output_tensor=None, + output_sf_tensor=None, + partition_id=-1, + activation_type=ActivationType.Swiglu, + ) + + half_weight_n = weight.size(1) // 2 + c_locality_domain = torch.empty( + max_num_permuted_tokens, interm_size // 2, dtype=a_fp4.dtype, device=a_fp4.device + ) + c_sf_locality_domain = torch.empty( + max_num_permuted_tokens * interm_size // sf_vec_size, dtype=torch.uint8, device=a_fp4.device + ) + + start_for_all_locality_domain() + try: + for locality_domain_id in range(2): + with locality_domain_device(locality_domain_id): + with torch.cuda.stream(get_locality_domain_stream(locality_domain_id)): + weight_shard = weight.view(torch.uint8)[ + :, + locality_domain_id * half_weight_n : (locality_domain_id + 1) + * half_weight_n, + ] + weight_shard = weight_shard.contiguous().view(torch.float4_e2m1fn_x2) + weight_sf_shard = weight_sf[ + :, + locality_domain_id * half_weight_n : (locality_domain_id + 1) + * half_weight_n, + ].contiguous() + torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin( + a_fp4, + weight_shard, + a_sf_unswizzled, + weight_sf_shard, + alpha, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + global_sf, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=0, + tile_size=tile_size, + scaling_vector_size=sf_vec_size, + output_tensor=c_locality_domain, + output_sf_tensor=c_sf_locality_domain, + partition_id=locality_domain_id, + activation_type=ActivationType.Swiglu, + ) + finally: + end_for_all_locality_domain() + torch.cuda.synchronize() + + num_valid = total_num_padded_tokens.item() + valid_mask = _valid_permuted_token_mask(tile_idx_to_mn_limit, num_valid, tile_size) + c_valid = c_locality_domain.view(torch.uint8)[:num_valid][valid_mask] + c_ref_valid = c_ref.view(torch.uint8)[:num_valid][valid_mask] + assert c_valid.any() + torch.testing.assert_close(c_valid, c_ref_valid) + + c_sf_locality_domain = unswizzle_sf( + c_sf_locality_domain, max_num_permuted_tokens, interm_size, sf_vec_size + ) + c_sf_ref = unswizzle_sf(c_sf_ref, max_num_permuted_tokens, interm_size, sf_vec_size) + torch.testing.assert_close( + c_sf_locality_domain[:num_valid][valid_mask], c_sf_ref[:num_valid][valid_mask] + ) + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on Rubin (SM 107) GPUs", +) +@pytest.mark.parametrize("tile_size", [128]) +@pytest.mark.parametrize("ep_size", [1, 8]) +@pytest.mark.parametrize("top_k", [1, 2]) +@pytest.mark.parametrize("num_tokens", [128, 515]) +def test_nvfp4_grouped_gemm_finalize_locality_domain_rubin( + num_tokens: int, top_k: int, ep_size: int, tile_size: int +): + _skip_if_no_locality_domain() + + sf_vec_size = 16 + hidden_size = 2048 + interm_size = 1536 + num_experts = 256 + num_local_experts = num_experts // ep_size + + ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + _, + num_non_exiting_tiles, + token_final_scales, + ) = _setup_locality_domain_routing(num_tokens, num_experts, num_local_experts, top_k, tile_size) + max_num_permuted_tokens = permuted_idx_to_expanded_idx.size(0) + + torch.manual_seed(99) + fc2_input_bf16 = ( + torch.randn(max_num_permuted_tokens, interm_size, dtype=torch.bfloat16, device="cuda") + * 0.05 + ) + fc2_global_sf = fc2_input_bf16.abs().max().float() / (448 * 6) + fc2_input, fc2_input_sf = torch.ops.trtllm.fp4_quantize( + fc2_input_bf16, 1 / fc2_global_sf, sf_vec_size, False + ) + fc2_input = fc2_input.view(torch.float4_e2m1fn_x2) + + fc2_weight_bf16 = ( + torch.randn( + num_local_experts, hidden_size, interm_size, dtype=torch.bfloat16, device="cuda" + ) + * 0.05 + ) + fc2_weight_global_sf = fc2_weight_bf16.abs().amax(dim=(1, 2)).float() / (448 * 6) + fc2_weight, fc2_weight_sf = torch.ops.trtllm.fp4_quantize( + fc2_weight_bf16, 1 / fc2_weight_global_sf, sf_vec_size, False + ) + fc2_weight = fc2_weight.view(torch.float4_e2m1fn_x2) + fc2_weight_sf = fc2_weight_sf.view(num_local_experts, hidden_size, interm_size // sf_vec_size) + fc2_alpha = fc2_global_sf * fc2_weight_global_sf + + output_ref = torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_finalize_rubin( + input=fc2_input, + weight=fc2_weight, + input_scale=fc2_input_sf.view(torch.uint8), + weight_scale=fc2_weight_sf.view(torch.uint8), + alpha=fc2_alpha, + tile_idx_to_group_idx=tile_idx_to_group_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=0, + tile_size=tile_size, + output_dtype=torch.bfloat16, + ) + + output = torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda") + half_hidden = hidden_size // 2 + start_for_all_locality_domain() + try: + for locality_domain_id in range(2): + with locality_domain_device(locality_domain_id): + with torch.cuda.stream(get_locality_domain_stream(locality_domain_id)): + weight_shard = fc2_weight.view(torch.uint8)[ + :, locality_domain_id * half_hidden : (locality_domain_id + 1) * half_hidden + ] + weight_shard = weight_shard.contiguous().view(torch.float4_e2m1fn_x2) + weight_sf_shard = fc2_weight_sf[ + :, locality_domain_id * half_hidden : (locality_domain_id + 1) * half_hidden + ] + torch.ops.trtllm.cute_dsl_nvfp4_grouped_gemm_finalize_inplace_rubin( + input=fc2_input, + weight=weight_shard, + input_scale=fc2_input_sf.view(torch.uint8), + weight_scale=weight_sf_shard.contiguous().view(torch.uint8), + alpha=fc2_alpha, + output=output, + tile_idx_to_group_idx=tile_idx_to_group_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=0, + tile_size=tile_size, + output_dtype=torch.bfloat16, + ) + finally: + end_for_all_locality_domain() + torch.cuda.synchronize() + + assert output[:, :half_hidden].any() + assert output[:, half_hidden:].any() + torch.testing.assert_close(output, output_ref, rtol=1e-2, atol=0.15) + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on Rubin (SM 107) GPUs", +) +@pytest.mark.parametrize("tile_size", [128]) +@pytest.mark.parametrize("ep_size", [1, 8]) +@pytest.mark.parametrize("top_k", [1, 2]) +def test_bf16_gather_grouped_gemm_swiglu_locality_domain_rubin( + top_k: int, ep_size: int, tile_size: int +): + _skip_if_no_locality_domain() + + num_tokens = 128 + hidden_size = 2048 + interm_size = 1536 + num_experts = 256 + num_local_experts = num_experts // ep_size + + ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles, + _, + ) = _setup_locality_domain_routing(num_tokens, num_experts, num_local_experts, top_k, tile_size) + max_num_permuted_tokens = permuted_idx_to_expanded_idx.size(0) + + torch.manual_seed(7) + input_tensor = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda") * 0.05 + weight = ( + torch.randn( + num_local_experts, interm_size * 2, hidden_size, dtype=torch.bfloat16, device="cuda" + ) + * 0.05 + ) + weight = interleave_linear_and_gate(weight, group_size=32, dim=1) + alpha = torch.ones(num_local_experts, dtype=torch.float32, device="cuda") + + output_ref = torch.ops.trtllm.cute_dsl_bf16_gather_grouped_gemm_swiglu_rubin( + input=input_tensor, + weight=weight, + alpha=alpha, + tile_idx_to_group_idx=tile_idx_to_group_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=0, + tile_size=tile_size, + output_tensor=None, + partition_id=-1, + ) + + output = torch.empty(max_num_permuted_tokens, interm_size, dtype=torch.bfloat16, device="cuda") + half_weight_n = weight.size(1) // 2 + start_for_all_locality_domain() + try: + for locality_domain_id in range(2): + with locality_domain_device(locality_domain_id): + with torch.cuda.stream(get_locality_domain_stream(locality_domain_id)): + torch.ops.trtllm.cute_dsl_bf16_gather_grouped_gemm_swiglu_rubin( + input=input_tensor, + weight=weight[ + :, + locality_domain_id * half_weight_n : (locality_domain_id + 1) + * half_weight_n, + ].contiguous(), + alpha=alpha, + tile_idx_to_group_idx=tile_idx_to_group_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=0, + tile_size=tile_size, + output_tensor=output, + partition_id=locality_domain_id, + ) + finally: + end_for_all_locality_domain() + torch.cuda.synchronize() + + num_valid = total_num_padded_tokens.item() + valid_mask = _valid_permuted_token_mask(tile_idx_to_mn_limit, num_valid, tile_size) + torch.testing.assert_close( + output[:num_valid][valid_mask], output_ref[:num_valid][valid_mask], rtol=1e-2, atol=0.15 + ) + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on Rubin (SM 107) GPUs", +) +@pytest.mark.parametrize("tile_size", [128]) +@pytest.mark.parametrize("ep_size", [1, 8]) +@pytest.mark.parametrize("top_k", [1, 2]) +def test_bf16_grouped_gemm_finalize_locality_domain_rubin(top_k: int, ep_size: int, tile_size: int): + _skip_if_no_locality_domain() + + num_tokens = 128 + hidden_size = 2048 + interm_size = 1536 + num_experts = 256 + num_local_experts = num_experts // ep_size + + ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + _, + num_non_exiting_tiles, + token_final_scales, + ) = _setup_locality_domain_routing(num_tokens, num_experts, num_local_experts, top_k, tile_size) + max_num_permuted_tokens = permuted_idx_to_expanded_idx.size(0) + + torch.manual_seed(11) + input_tensor = ( + torch.randn(max_num_permuted_tokens, interm_size, dtype=torch.bfloat16, device="cuda") + * 0.05 + ) + weight = ( + torch.randn( + num_local_experts, hidden_size, interm_size, dtype=torch.bfloat16, device="cuda" + ) + * 0.05 + ) + + output_ref = torch.ops.trtllm.cute_dsl_bf16_grouped_gemm_finalize_rubin( + input=input_tensor, + weight=weight, + tile_idx_to_group_idx=tile_idx_to_group_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=0, + tile_size=tile_size, + output_dtype=torch.bfloat16, + ) + + output = torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda") + half_hidden = hidden_size // 2 + start_for_all_locality_domain() + try: + for locality_domain_id in range(2): + with locality_domain_device(locality_domain_id): + with torch.cuda.stream(get_locality_domain_stream(locality_domain_id)): + torch.ops.trtllm.cute_dsl_bf16_grouped_gemm_finalize_inplace_rubin( + input=input_tensor, + weight=weight[ + :, + locality_domain_id * half_hidden : (locality_domain_id + 1) + * half_hidden, + ].contiguous(), + output=output, + tile_idx_to_group_idx=tile_idx_to_group_idx, + tile_idx_to_mn_limit=tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx=permuted_idx_to_expanded_idx, + num_non_exiting_tiles=num_non_exiting_tiles, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + num_local_experts=num_local_experts, + local_expert_offset=0, + tile_size=tile_size, + output_dtype=torch.bfloat16, + ) + finally: + end_for_all_locality_domain() + torch.cuda.synchronize() + + assert output[:, :half_hidden].any() + assert output[:, half_hidden:].any() + torch.testing.assert_close(output, output_ref, rtol=1e-2, atol=0.15) + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on Rubin (SM 107) GPUs", +) +@pytest.mark.parametrize("num_tokens", [128]) +@pytest.mark.parametrize("top_k", [1, 2]) +def test_moe_module_locality_domain_correctness_rubin(num_tokens: int, top_k: int): + _skip_if_no_locality_domain() + + from _torch.moe.quantize_utils import get_test_quant_params + from transformers.configuration_utils import PretrainedConfig + + from tensorrt_llm._torch.locality_domain.policy import LocalityDomainPolicy + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.moe.fused_moe import RenormalizeMoeRoutingMethod + from tensorrt_llm._torch.moe.fused_moe.create_moe import create_moe_backend + from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl import CuteDslFusedMoE + from tensorrt_llm._utils import mpi_rank + from tensorrt_llm.mapping import Mapping + from tensorrt_llm.models.modeling_utils import QuantAlgo + + hidden_size = 2048 + intermediate_size = 1536 + num_experts = 256 + dtype = torch.bfloat16 + + mapping = Mapping() + mapping.rank = mpi_rank() + + with torch.device(f"cuda:{mapping.rank}"): + torch.manual_seed(0) + torch.cuda.manual_seed(0) + + routing_method = RenormalizeMoeRoutingMethod(top_k=top_k) + input_tensor = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda") + router_logits = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda") + + quantize_util_cls, quant_config, quant_kwargs = get_test_quant_params( + QuantAlgo.NVFP4, input_tensor, "CUTEDSL" + ) + quantize_util = quantize_util_cls( + num_experts=num_experts, + dtype=dtype, + intermediate_size=intermediate_size, + hidden_size=hidden_size, + quant_config=quant_config, + bias=False, + swiglu_gptoss_style=False, + ) + weights = quantize_util.create_weights(**quant_kwargs) + + pretrained_config = PretrainedConfig() + pretrained_config.num_experts = num_experts + pretrained_config.hidden_size = hidden_size + pretrained_config.intermediate_size = intermediate_size + pretrained_config.torch_dtype = dtype + + def create_backend(enable_locality_domains: bool): + model_config = ModelConfig( + pretrained_config=pretrained_config, + quant_config=quant_config, + mapping=mapping, + moe_backend="CUTEDSL", + locality_domain_policy=LocalityDomainPolicy(enabled=enable_locality_domains), + ) + backend = create_moe_backend( + moe_cls=CuteDslFusedMoE, + routing_method=routing_method, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + reduce_results=True, + model_config=model_config, + init_load_balancer=False, + ) + backend.load_weights([weights]) + source_storage_ptrs = {} + if enable_locality_domains and top_k == 2: + source_storage_ptrs = { + "w3_w1_weight": backend.w3_w1_weight.untyped_storage().data_ptr(), + "w2_weight": backend.w2_weight.untyped_storage().data_ptr(), + "fc1_weight_block": backend.quant_scales.fc1_weight_block.untyped_storage().data_ptr(), + "fc2_weight_block": backend.quant_scales.fc2_weight_block.untyped_storage().data_ptr(), + } + backend.post_load_weights() + backend.cuda() + return backend, source_storage_ptrs + + base_backend, _ = create_backend(False) + locality_domain_backend, source_storage_ptrs = create_backend(True) + + if top_k == 2: + assert locality_domain_backend._locality_domain_runtime is not None + assert locality_domain_backend._locality_domain_weight_shards is not None + assert hasattr(locality_domain_backend, "_cached_reserved_moe_output_memset_stream") + for shard in locality_domain_backend._locality_domain_weight_shards: + for name, source_storage_ptr in source_storage_ptrs.items(): + assert shard[name].untyped_storage().data_ptr() != source_storage_ptr + for param_name in ( + "w3_w1_weight", + "w2_weight", + "w3_w1_weight_scale", + "w2_weight_scale", + ): + assert getattr(locality_domain_backend, param_name).numel() == 0 + assert locality_domain_backend.quant_scales.fc1_weight_block.numel() == 0 + assert locality_domain_backend.quant_scales.fc2_weight_block.numel() == 0 + + with torch.inference_mode(): + base_output = base_backend.forward_chunk(input_tensor, router_logits) + locality_domain_output = locality_domain_backend.forward_chunk( + input_tensor, router_logits + ) + + torch.cuda.synchronize() + torch.testing.assert_close(base_output, locality_domain_output, rtol=1e-2, atol=0.15) + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on Rubin (SM 107) GPUs", +) +def test_moe_module_bf16_locality_domain_lifecycle_and_forward_chunk_rubin(): + _skip_if_no_locality_domain() + + from _torch.moe.quantize_utils import get_test_quant_params + from transformers.configuration_utils import PretrainedConfig + + from tensorrt_llm._torch.locality_domain.policy import LocalityDomainPolicy + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.moe.fused_moe import RenormalizeMoeRoutingMethod + from tensorrt_llm._torch.moe.fused_moe.create_moe import create_moe_backend + from tensorrt_llm._torch.moe.fused_moe.fused_moe_cute_dsl import CuteDslFusedMoE + from tensorrt_llm._utils import mpi_rank + from tensorrt_llm.mapping import Mapping + + hidden_size = 2048 + intermediate_size = 1536 + num_experts = 256 + num_tokens = 128 + top_k = 2 + dtype = torch.bfloat16 + + mapping = Mapping() + mapping.rank = mpi_rank() + + with torch.device(f"cuda:{mapping.rank}"): + torch.manual_seed(0) + torch.cuda.manual_seed(0) + + routing_method = RenormalizeMoeRoutingMethod(top_k=top_k) + input_tensor = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda") + router_logits = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda") + + quantize_util_cls, quant_config, quant_kwargs = get_test_quant_params( + None, input_tensor, "CUTEDSL" + ) + quantize_util = quantize_util_cls( + num_experts=num_experts, + dtype=dtype, + intermediate_size=intermediate_size, + hidden_size=hidden_size, + quant_config=quant_config, + bias=False, + swiglu_gptoss_style=False, + ) + weights = quantize_util.create_weights(**quant_kwargs) + + pretrained_config = PretrainedConfig() + pretrained_config.num_experts = num_experts + pretrained_config.hidden_size = hidden_size + pretrained_config.intermediate_size = intermediate_size + pretrained_config.torch_dtype = dtype + + def create_backend(enable_locality_domains: bool): + model_config = ModelConfig( + pretrained_config=pretrained_config, + quant_config=quant_config, + mapping=mapping, + moe_backend="CUTEDSL", + locality_domain_policy=LocalityDomainPolicy(enabled=enable_locality_domains), + ) + backend = create_moe_backend( + moe_cls=CuteDslFusedMoE, + routing_method=routing_method, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + reduce_results=True, + model_config=model_config, + init_load_balancer=False, + ) + backend.load_weights([weights]) + full_w3_w1 = None + full_w2 = None + source_storage_ptrs = () + if enable_locality_domains: + full_w3_w1 = backend.w3_w1_weight.data.clone() + full_w2 = backend.w2_weight.data.clone() + source_storage_ptrs = ( + backend.w3_w1_weight.untyped_storage().data_ptr(), + backend.w2_weight.untyped_storage().data_ptr(), + ) + backend.post_load_weights() + backend.cuda() + return backend, full_w3_w1, full_w2, source_storage_ptrs + + base_backend, _, _, _ = create_backend(False) + locality_domain_backend, full_w3_w1, full_w2, source_storage_ptrs = create_backend(True) + + assert locality_domain_backend._locality_domain_runtime is not None + assert locality_domain_backend._locality_domain_weight_shards is not None + assert hasattr(locality_domain_backend, "_cached_reserved_moe_output_memset_stream") + + shards = locality_domain_backend._locality_domain_weight_shards + for shard in shards: + assert shard["w3_w1_weight"].untyped_storage().data_ptr() != source_storage_ptrs[0] + assert shard["w2_weight"].untyped_storage().data_ptr() != source_storage_ptrs[1] + assert torch.equal(torch.cat([s["w3_w1_weight"] for s in shards], dim=1), full_w3_w1) + assert torch.equal(torch.cat([s["w2_weight"] for s in shards], dim=1), full_w2) + assert locality_domain_backend.w3_w1_weight.numel() == 0 + assert locality_domain_backend.w2_weight.numel() == 0 + + # Keep the production-shape lifecycle and public forward_chunk + # integration here. Broad accuracy, autotune, capture, and outer-tile + # replay are covered by the unified backend matrix. + with torch.inference_mode(): + base_output = base_backend.forward_chunk(input_tensor, router_logits) + locality_domain_output = locality_domain_backend.forward_chunk( + input_tensor, router_logits + ) + + torch.cuda.synchronize() + torch.testing.assert_close(base_output, locality_domain_output, rtol=1e-2, atol=0.15) + + +@pytest.mark.skipif( + get_sm_version() != 107, + reason="This test is only supported on SM 107 (Rubin) GPUs", +) +@pytest.mark.parametrize("tile_size", [64, 128, 256]) +@pytest.mark.parametrize("ep_size", [1, 8]) +@pytest.mark.parametrize("top_k", [1, 2, 8]) +@pytest.mark.parametrize("num_tokens", [128, 515, 1024]) +def test_bf16_grouped_gemm_finalize_rubin( + num_tokens: int, top_k: int, ep_size: int, tile_size: int +): + """Test BF16 grouped GEMM with finalize fusion on Rubin (SM107). + + Uses torch.ops.trtllm.cute_dsl_bf16_grouped_gemm_finalize_rubin. + No scale factors or quantization — direct BF16 inputs/outputs. + """ + hidden_size = 4096 + interm_size = 8192 + num_experts = 256 + num_local_experts = num_experts // ep_size + + # Generate routing information + routing_logits = torch.randn(num_tokens, num_experts, device="cuda") + token_final_scales, token_selected_experts = routing_logits.topk(top_k, dim=-1) + token_selected_experts = token_selected_experts.to(torch.int32) + token_final_scales = token_final_scales.softmax(dim=-1).to(torch.float32) + + ( + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + expanded_idx_to_permuted_idx, + permuted_idx_to_expanded_idx, + total_num_padded_tokens, + num_non_exiting_tiles, + ) = torch.ops.trtllm.moe_sort( + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + num_experts=num_experts, + top_k=top_k, + local_expert_offset=0, + local_num_experts=num_local_experts, + tile_tokens_dim=tile_size, + ) + + max_num_permuted_tokens = permuted_idx_to_expanded_idx.size(0) + + # Create BF16 input tensors (FC2: interm_size -> hidden_size) + a = torch.randn(max_num_permuted_tokens, interm_size, dtype=torch.bfloat16, device="cuda") + b = torch.randn( + num_local_experts, hidden_size, interm_size, dtype=torch.bfloat16, device="cuda" + ) + # Compute reference: per-group GEMM + scatter-add finalize + tile_group_list = tile_idx_to_group_idx.cpu().tolist() + tile_mn_limit_list = tile_idx_to_mn_limit.cpu().tolist() + permuted_idx_list = permuted_idx_to_expanded_idx.cpu().tolist() + + c_permuted = torch.zeros( + max_num_permuted_tokens, hidden_size, dtype=torch.float32, device="cuda" + ) + for tile_idx in range(num_non_exiting_tiles.item()): + group_idx = tile_group_list[tile_idx] + mn_limit = tile_mn_limit_list[tile_idx] + start = tile_idx * tile_size + end = min(start + tile_size, mn_limit) + + for i in range(start, end): + a_row = a[i].float() + gemm_row = a_row @ b[group_idx].float().T + c_permuted[i] = gemm_row + + # Scatter-add with token_final_scales + c_ref = torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda") + for tile_idx in range(num_non_exiting_tiles.item()): + mn_limit = tile_mn_limit_list[tile_idx] + start = tile_idx * tile_size + end = min(start + tile_size, mn_limit) + for i in range(start, end): + expanded_idx = permuted_idx_list[i] + token_idx = expanded_idx // top_k + topk_idx = expanded_idx % top_k + scale = token_final_scales[token_idx, topk_idx].item() + c_ref[token_idx] += (c_permuted[i] * scale).to(torch.bfloat16) + + # Even-tile padding for Rubin cluster sync + kernel_nnet = ((num_non_exiting_tiles + 1) // 2) * 2 + + # Test all valid autotuner candidate tactics via direct runner call + from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import ( + Sm107ContiguousGroupedGemmFinalizeFusionRunner, + ) + + runner = Sm107ContiguousGroupedGemmFinalizeFusionRunner( + num_experts, top_k, num_local_experts, 0, tile_size, torch.bfloat16 + ) + + tactics = runner.get_valid_tactics( + [ + a, + b, + torch.zeros_like(c_ref), + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + num_non_exiting_tiles, + token_final_scales, + ], + None, + ) + assert len(tactics) > 0, f"No valid tactics for tile_size={tile_size}" + + failed = [] + for tactic in tactics: + mma_tiler, _, cluster, _ = tactic + label = f"mma={mma_tiler[:2]} cluster={cluster}" + output = torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device="cuda") + inputs = [ + a, + b, + output, + tile_idx_to_group_idx, + tile_idx_to_mn_limit, + permuted_idx_to_expanded_idx, + kernel_nnet, + token_final_scales, + ] + with torch.inference_mode(): + c = runner.forward(inputs, tactic=tactic) + match = torch.isclose(c, c_ref, rtol=1.6e-2, atol=1e-1).sum().item() / c_ref.numel() + if match < 0.95: + failed.append(f"{label}: match={match:.4f}") + assert not failed, ( + f"tile_size={tile_size}: {len(failed)}/{len(tactics)} tactics failed:\n " + + "\n ".join(failed) + ) diff --git a/tests/unittest/_torch/thop/parallel/test_locality_domain_planner.py b/tests/unittest/_torch/thop/parallel/test_locality_domain_planner.py index f479f241ab23..a749ffc429a9 100644 --- a/tests/unittest/_torch/thop/parallel/test_locality_domain_planner.py +++ b/tests/unittest/_torch/thop/parallel/test_locality_domain_planner.py @@ -699,6 +699,16 @@ def test_moe_disabled_without_fused_finalize(self, mock_locality_domain): assert not plan.enabled assert "fused finalize" in plan.reason_if_disabled + @patch( + "tensorrt_llm._torch.locality_domain_utils.is_locality_domain_enabled", return_value=True + ) + def test_moe_disabled_for_non_swiglu_activation(self, mock_locality_domain): + """Both locality-domain MoE kernels fuse SwiGLU; others stay unpartitioned.""" + planner = LocalityDomainExecutionPlanner(LocalityDomainPolicy(enabled=True)) + plan = planner.plan_moe(_FakeMoeQuantConfig(nvfp4=True), activation="Relu2") + assert not plan.enabled + assert "SwiGLU only" in plan.reason_if_disabled + @patch( "tensorrt_llm._torch.locality_domain_utils.is_locality_domain_enabled", return_value=True ) diff --git a/tests/unittest/_torch/thop/parallel/test_locality_domain_utils.py b/tests/unittest/_torch/thop/parallel/test_locality_domain_utils.py index 12357c757507..a5b7b5480104 100644 --- a/tests/unittest/_torch/thop/parallel/test_locality_domain_utils.py +++ b/tests/unittest/_torch/thop/parallel/test_locality_domain_utils.py @@ -42,6 +42,7 @@ ) from tensorrt_llm._torch.locality_domain.runtime import LocalityDomainRuntime from tensorrt_llm._torch.locality_domain_utils import ( + _copy_to_new_cuda_allocation, get_locality_domain_compute_sm_counts, get_locality_domain_mempool, get_locality_domain_stream, @@ -843,9 +844,18 @@ def test_concurrent_stream_operations(self, check_locality_domain_support): class TestLocalityDomainMempoolAllocation: """Tests for LOCALITY_DOMAIN memory pool allocation and deallocation.""" - # TODO: restore test_copy_to_new_cuda_allocation_does_not_alias_contiguous_input - # alongside the Linear/MoE wire-up. It covers _copy_to_new_cuda_allocation, which - # lives in modules/linear.py and lands with those call sites rather than here. + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + def test_copy_to_new_cuda_allocation_does_not_alias_contiguous_input(self): + source_storage = torch.arange(32, device="cuda") + source = source_storage[:16] + assert source.is_contiguous() + assert source.untyped_storage().data_ptr() == source_storage.untyped_storage().data_ptr() + + copied = _copy_to_new_cuda_allocation(source) + + assert copied.is_contiguous() + assert copied.untyped_storage().data_ptr() != source.untyped_storage().data_ptr() + torch.testing.assert_close(copied, source) def test_allocate_tensor_with_mempool(self, check_locality_domain_support): """Test allocating a tensor using LOCALITY_DOMAIN mempool."""