"""CuTeDSL W4A16 NVFP4/BF16 W4A16 MoE kernels.""" from __future__ import annotations import math import os from dataclasses import dataclass, replace from typing import NamedTuple import cuda.bindings.driver as cuda import cutlass import cutlass.cute as cute import torch from cutlass.cutlass_dsl import Int32, Int64, Uint32 from b12x.cute.compiler import KernelCompileSpec, compile as b12x_compile from b12x.cute.fp4 import ( atomic_add_global_i32, bf16_mma_m16n8k16_f32, bf16_mma_rhs_fragments_as_mma_a_m16n8k16_f32, bfloat2_broadcast_lane, bfloat2_mul, bfloat2_to_float2_scaled, broadcast_f32_to_half2, broadcast_f32_to_bfloat2, cp_async4_shared_global, cp_async4_shared_global_pred, fabs_f32, fmax_f32, f16_mma_m16n8k16_f32, f16_mma_rhs_fragments_as_mma_a_m16n8k16_f32, half2_to_float2_scaled, get_ptr_as_int64, half2_mul, ld_global_acquire_i32, ld_global_nc_u32, ld_global_v4_f32, ld_shared_f32, ld_shared_i32_relaxed, ld_shared_u32, ld_shared_v2_u32, ld_shared_v4_f32, ld_shared_v4_u32, ldmatrix_m8n8x4_b16, ldmatrix_m8n8x2_b16, packed_dequant_e2m1x4_to_bfloat2x2, packed_dequant_e2m1x4_to_half2x2, packed_dequant_e4m3x4_to_bfloat2x2, packed_dequant_e4m3x4_to_half2x2, packed_dequant_e8m0x4_to_bfloat2x2, packed_dequant_e8m0x4_to_half2x2, packed_dequant_nf3x8_to_bfloat2x4, nf3_codebook_pools, pack_f32x2_to_bfloat2, pack_f32x2_to_f16x2, red_add_global_bf16x2, red_add_global_release_i32, red_max_global_f32_nonnegative, shared_ptr_to_u32, st_global_v4_u32, st_global_i32, st_global_v4_f32, st_shared_bf16_from_f32, st_shared_f16_from_f32, st_shared_f32, st_shared_i32, st_shared_u32, st_shared_v4_f32, st_shared_v4_u32, threadfence, warp_reduce, ) from b12x.cute.utils import current_cuda_stream, make_ptr from b12x.moe.fused.w4a16.route_pack import ( pack_topk_routes_by_expert as _pack_topk_routes_by_expert, ) from b12x.moe.fused.w4a16.host import ( _W4A16_ALLOWED_ROUTED_SIZES, max_packed_route_slots, packed_gemm_scratch_elements, plan_w4a16_buffers, select_route_block_size_m, validate_activation, ) from b12x.cute.runtime_control import raise_if_kernel_resolution_frozen from b12x.moe.fused.activations import ( SWIGLUOAI_UNINTERLEAVE, is_gated_moe_activation, normalize_moe_activation, normalize_swiglu_alpha_for_activation, normalize_swiglu_beta_for_activation, normalize_swiglu_limit_for_activation, ) from b12x.moe.fused.micro import MoEMicroKernelBackend _ALLOWED_ROUTED_SIZES = _W4A16_ALLOWED_ROUTED_SIZES _PACK_FACTOR = 8 _STAGES = 4 _E8M0_LOGICAL_TAIL_SCALE_N_ALIGNMENT = 64 _DEVICE_MAX_REG_BYTES = 255 * 1024 _DEFAULT_MAX_SHARED_MEM = 101_376 _WEIGHT_LAYOUTS = {"packed", "modelopt", "nf3_2p1"} _MODEL_OPT_W13_LAYOUTS = {"w13", "w31"} # NF3 codebook: 8 bf16 codepoints for the 3-bit ("nf3_2p1") weight layout. The # device dequant reads these as prmt pools (see nf3_codebook_pools); the host # injects the 4 pool words as compile-time constants into the GEMM. _NF3_CODEBOOK = (-1.0, -0.6047, -0.3563, -0.1275, 0.1275, 0.3563, 0.6047, 1.0) _SCALE_FORMATS = { "e4m3_k16": "e4m3_k16", "e8m0_k32": "e8m0_k32", "e4m3_k32": "e4m3_k32", } _E8M0_K32_FP16_GLOBAL_COMPENSATION = float(2.0**7) _E8M0_K32_BF16_GLOBAL_COMPENSATION = float(2.0**119) _MAX_DIRECT_TOPK_ROUTE_M = 6 _W4A16_SMALL_M_DIRECT_MAX_M = 8 # TC-decode: a small-M decode specialization that runs on the PACKED W4A16 # object (the same weights/scales the prefill GEMM uses). It reuses the packed # tensor-core MMA inner loop but folds the top-k sum into the FC2 store # epilogue (dropping the separate top-k-sum launch). It never regresses vs the # packed GEMM within its supported M range, so it is ALWAYS used for the small-M # direct-topk decode sizes — there is no opt-in/opt-out switch. # TC-decode is available for the whole small-M direct-topk range. Its only hard # ceiling is the direct-topk route cap (above it, expert route-packing wins). # _TC_DECODE_M is retained for callers/tests that enumerate the supported sizes. _TC_DECODE_MAX_M = _W4A16_SMALL_M_DIRECT_MAX_M _TC_DECODE_M = tuple(range(1, _TC_DECODE_MAX_M + 1)) def _m_specialization_key(size_m: int) -> int: """M bucket for branches that genuinely specialize on token count.""" return 1 if int(size_m) == 1 else 0 def _fake_m_for_specialization(size_m: int) -> int: return 1 if _m_specialization_key(size_m) == 1 else 2 # The W4A16 launch model chooses blocks/SM from static resource usage # for each specialization. These register counts were measured from the local # SM121 JIT output and keep launch occupancy stable across refactors. _W4A16_REGS_SM121 = { (256, 1, 8, 8, True): 118, (256, 1, 16, 4, True): 118, (256, 1, 16, 8, True): 118, (256, 1, 32, 2, True): 118, (128, 1, 4, 8, True): 118, (128, 1, 8, 4, True): 120, (256, 1, 8, 8, False): 158, (128, 1, 4, 8, False): 154, (128, 1, 8, 4, False): 143, (256, 2, 16, 4, False): 212, (128, 2, 4, 8, False): 215, (128, 2, 8, 4, False): 214, (256, 3, 16, 4, False): 249, (128, 3, 4, 8, False): 249, (128, 3, 8, 4, False): 250, (256, 4, 16, 4, False): 255, (128, 4, 4, 8, False): 255, (128, 4, 8, 4, False): 255, } _SMALL_BATCH_TILE_CONFIGS = ( (128, 128, 256), (64, 128, 128), (128, 64, 128), ) _LARGE_BATCH_TILE_CONFIGS = ( (64, 256, 256), (64, 128, 128), (128, 64, 128), ) def _covering_count(total: int, quantum: int) -> int: return (total + quantum - 1) // quantum def _e8m0_logical_tail_scale_n(size_n: int) -> int: return ( (int(size_n) + _E8M0_LOGICAL_TAIL_SCALE_N_ALIGNMENT - 1) // _E8M0_LOGICAL_TAIL_SCALE_N_ALIGNMENT ) * _E8M0_LOGICAL_TAIL_SCALE_N_ALIGNMENT def _normalize_swiglu_limit(swiglu_limit: float | None) -> float | None: return normalize_swiglu_limit_for_activation("silu", swiglu_limit) def _normalize_activation_swiglu_params( activation: str, swiglu_limit: float | None, swiglu_alpha: float | None, swiglu_beta: float | None, ) -> tuple[float | None, float, float]: activation = normalize_moe_activation(activation) return ( normalize_swiglu_limit_for_activation(activation, swiglu_limit), normalize_swiglu_alpha_for_activation(activation, swiglu_alpha), normalize_swiglu_beta_for_activation(activation, swiglu_beta), ) def _w4a16_num_regs( *, cta_threads: int, cta_m_blocks: int, cta_n_blocks: int, cta_k_blocks: int, uses_m_block_8: bool, ) -> int: key = ( int(cta_threads), int(cta_m_blocks), int(cta_n_blocks), int(cta_k_blocks), bool(uses_m_block_8), ) try: return _W4A16_REGS_SM121[key] except KeyError as exc: raise ValueError( f"missing W4A16 register count for NVFP4 BF16 specialization {key}" ) from exc def _shared_memory_footprint( *, cta_m_blocks: int, tile_n: int, tile_k: int, scale_format: str = "e4m3_k16", ) -> int: cta_m = int(cta_m_blocks) * 16 cta_n = int(tile_n) cta_k = int(tile_k) sh_block_meta_size = cta_m * 16 sh_a_size = _STAGES * (cta_m * cta_k) * 2 sh_b_size = _STAGES * (cta_k * cta_n // _PACK_FACTOR) * 4 sh_red_size = cta_m * (cta_n + 8) * 2 sh_bias_size = cta_n * 2 tmp_size = min(sh_b_size, sh_red_size) + sh_bias_size tmp_size = max(max(sh_b_size, sh_red_size), tmp_size) sh_s_size = ( _covering_count(cta_k, _scale_group_size(scale_format)) * cta_n * 2 * _STAGES ) return tmp_size + sh_a_size + sh_s_size + sh_block_meta_size def _determine_blocks_per_sm( *, problem_m: int, problem_n: int, top_k: int, cta_threads: int, cta_m_blocks: int, tile_n: int, tile_k: int, uses_m_block_8: bool, sms: int, max_shared_mem: int, scale_format: str = "e4m3_k16", ) -> int: num_regs = _w4a16_num_regs( cta_threads=cta_threads, cta_m_blocks=cta_m_blocks, cta_n_blocks=tile_n // 16, cta_k_blocks=tile_k // 16, uses_m_block_8=uses_m_block_8, ) register_bytes = max(num_regs, 1) * int(cta_threads) * 4 smem_bytes = _shared_memory_footprint( cta_m_blocks=cta_m_blocks, tile_n=tile_n, tile_k=tile_k, scale_format=scale_format, ) blocks_per_sm_limit = min( _DEVICE_MAX_REG_BYTES // register_bytes, int(max_shared_mem) // (smem_bytes + 1536), ) if uses_m_block_8: # Small-M (moe_block_size==8) TC-decode is weight-bandwidth/overhead # bound, not parallelism bound (only m*top_k route-blocks of GEMM work). # The fused FC1->activation->FC2 path crosses several grid barriers whose # tid==0 atomic-counter increment serializes across all grid_x CTAs, so an # oversized grid pays barrier-atomic latency proportional to grid_x for no # extra GEMM throughput. Pin one persistent CTA per SM to minimize the # barrier participant count while still covering the machine for the # I_tp=1024 GEMMs. The split-K persistent loop is grid_x-agnostic, so this # is numerically identical. blocks_per_sm_limit = 1 elif cta_m_blocks == 1: blocks_per_sm_limit = max(min(blocks_per_sm_limit, 4), 1) else: blocks_per_sm_limit = max(min(blocks_per_sm_limit, 2), 1) work_cta_count = (int(problem_n) // int(tile_n)) * int(problem_m) * int(top_k) * 4 if work_cta_count < int(sms) * blocks_per_sm_limit: blocks_per_sm_limit = max(work_cta_count // int(sms), 1) return int(blocks_per_sm_limit) def _candidate_tile_fits( *, problem_n: int, problem_k: int, cta_m_blocks: int, tile_n: int, tile_k: int, cta_threads: int, max_shared_mem: int, scale_format: str = "e4m3_k16", allow_logical_tail: bool = False, ) -> bool: if int(tile_k) == -1 or int(tile_n) == -1 or int(cta_threads) == -1: return False scale_group_size = _scale_group_size(scale_format) exact_n = int(problem_n) % int(tile_n) == 0 exact_k = int(problem_k) % int(tile_k) == 0 exact_scale_k = int(problem_k) % scale_group_size == 0 if not allow_logical_tail: if not exact_n or not exact_k: return False if not exact_scale_k or int(tile_k) % scale_group_size != 0: return False elif ( _normalize_scale_format(scale_format) != "e8m0_k32" or int(problem_n) % 16 != 0 or int(problem_k) % 8 != 0 or int(tile_n) % 16 != 0 or int(tile_k) % scale_group_size != 0 ): return False if int(tile_n) < 64 or int(tile_k) < 64 or int(cta_threads) < 128: return False smem_bytes = _shared_memory_footprint( cta_m_blocks=cta_m_blocks, tile_n=tile_n, tile_k=tile_k, scale_format=scale_format, ) return smem_bytes <= int(max_shared_mem) def _select_tile_config( *, problem_m: int, problem_n: int, problem_k: int, top_k: int, moe_block_size: int, sms: int, max_shared_mem: int, required_cta_threads: int | None = None, scale_format: str = "e4m3_k16", allow_logical_tail: bool = False, ) -> tuple[int, int, int, int]: cta_m_blocks = _covering_count(moe_block_size, 16) uses_m_block_8 = moe_block_size == 8 configs = ( _LARGE_BATCH_TILE_CONFIGS if cta_m_blocks > 1 else _SMALL_BATCH_TILE_CONFIGS ) best_blocks_per_sm = 0 best_tile_config: tuple[int, int, int, int] | None = None for tile_k, tile_n, cta_threads in configs: if required_cta_threads is not None and int(cta_threads) != int( required_cta_threads ): continue if not _candidate_tile_fits( problem_n=problem_n, problem_k=problem_k, cta_m_blocks=cta_m_blocks, tile_n=tile_n, tile_k=tile_k, cta_threads=cta_threads, max_shared_mem=int(max_shared_mem) - 512, scale_format=scale_format, allow_logical_tail=allow_logical_tail, ): continue occupancy_problem_n = ( _covering_count(int(problem_n), int(tile_n)) * int(tile_n) if allow_logical_tail else int(problem_n) ) blocks_per_sm_limit = _determine_blocks_per_sm( problem_m=problem_m, problem_n=occupancy_problem_n, top_k=top_k, cta_threads=cta_threads, cta_m_blocks=cta_m_blocks, tile_n=tile_n, tile_k=tile_k, uses_m_block_8=uses_m_block_8, sms=sms, max_shared_mem=max_shared_mem, scale_format=scale_format, ) if blocks_per_sm_limit > best_blocks_per_sm: best_blocks_per_sm = blocks_per_sm_limit best_tile_config = (tile_k, tile_n, cta_threads, blocks_per_sm_limit) if best_tile_config is None: cta_thread_msg = ( "" if required_cta_threads is None else f", required_cta_threads={required_cta_threads}" ) raise ValueError( "no valid W4A16 tile config for " f"M/N/K={problem_m}/{problem_n}/{problem_k}, moe_block_size={moe_block_size}" f"{cta_thread_msg}" ) return best_tile_config @dataclass(frozen=True) class W4A16GemmCompileResult: compiled: object tile_n: int tile_k: int moe_block_size: int max_m_blocks: int blocks_per_sm: int weight_layout: str = "packed" scale_format: str = "e4m3_k16" w13_layout: str = "w13" @dataclass(frozen=True) class W4A16ActivationCompileResult: compiled: object rows: int intermediate_size: int activation: str swiglu_limit: float | None swiglu_alpha: float swiglu_beta: float @dataclass(frozen=True) class W4A16TopKSumCompileResult: compiled: object m: int topk: int hidden_size: int @dataclass(frozen=True) class W4A16FusedMoeCompileResult: compiled: object size_m: int hidden_size: int intermediate_size: int num_experts: int top_k: int activation: str apply_router_weight_on_input: bool zero_fc2_output: bool element_dtype: str fast_math: bool swiglu_limit: float | None swiglu_alpha: float swiglu_beta: float fc1_tile_n: int fc1_tile_k: int fc2_tile_n: int fc2_tile_k: int moe_block_size: int max_m_blocks: int blocks_per_sm: int weight_layout: str = "packed" w13_layout: str = "w13" direct_topk_routes: bool = False scale_format: str = "e4m3_k16" tc_decode_fused_sum: bool = False collect_activation_amax: bool = False @dataclass(frozen=True) class _W4A16GemmLaunch: kernel: W4A16GemmCompileResult c_tmp: torch.Tensor class _W4A16SmallMDirectLaunch(NamedTuple): compiled: object grid_x: int m: int hidden_size: int intermediate_size: int num_experts: int topk: int activation: str fast_math: bool topk_ids_dtype: torch.dtype class _W4A16SmallMDirectKernel(MoEMicroKernelBackend): """Decode-sized W4A16 specialization using the native ModelOpt layout.""" _SUPPORTED_M = tuple(range(1, _W4A16_SMALL_M_DIRECT_MAX_M + 1)) @classmethod def is_supported( cls, *, m: int, hidden_size: int, intermediate_size: int, topk: int, num_experts: int, scale_format: str = "e4m3_k16", ) -> bool: # E8M0 reads the shared packed scale grid. Both block axes must be /32. if scale_format == "e8m0_k32": scale_block_ok = ( int(hidden_size) % 32 == 0 and int(intermediate_size) % 32 == 0 ) else: scale_block_ok = True return ( int(m) in cls._SUPPORTED_M and int(m) <= _W4A16_SMALL_M_DIRECT_MAX_M and int(hidden_size) > 0 and int(hidden_size) % 128 == 0 and int(intermediate_size) > 0 and int(intermediate_size) % 16 == 0 and scale_block_ok and 0 < int(topk) <= 32 and int(num_experts) > 0 and MoEMicroKernelBackend.is_supported( int(m), int(hidden_size), int(intermediate_size), int(topk), int(num_experts), ) ) def __init__( self, *, activation: str, fast_math: bool, share_input_across_experts: bool, share_expert_scales: bool, single_token: bool, scale_format: str = "e4m3_k16", swiglu_limit: float | None = None, swiglu_alpha: float | None = None, swiglu_beta: float | None = None, w13_layout: str = "w13", ): super().__init__( sf_vec_size=16, mma_tiler_mn=(64, 128), output_tile_count_n=1, fast_math=fast_math, activation=activation, share_input_across_experts=share_input_across_experts, share_expert_scales=share_expert_scales, single_token=single_token, dynamic_down_scale=False, w4a16_mode=True, scale_format=scale_format, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, w13_layout=w13_layout, ) class W4A16GemmKernel: def __init__( self, *, size_m: int, size_n: int, size_k: int, num_experts: int, top_k: int, mul_topk_weights: bool, tile_n: int, tile_k: int, moe_block_size: int, max_m_blocks: int, element_dtype: str = "bf16", epilogue_activation: str | None = None, weight_layout: str = "packed", scale_format: str = "e4m3_k16", w13_layout: str = "w13", source_n_rotation: int = 0, single_token_route_fast_path: bool = False, direct_topk_routes: bool = False, fused_topk_sum: bool = False, fused_sum_topk: int = 1, ): if element_dtype not in {"bf16", "fp16"}: raise ValueError(f"unsupported element_dtype {element_dtype!r}") if weight_layout not in _WEIGHT_LAYOUTS: raise ValueError(f"unsupported W4A16 weight_layout {weight_layout!r}") scale_format = _normalize_scale_format(scale_format) if weight_layout == "modelopt": if w13_layout not in _MODEL_OPT_W13_LAYOUTS: raise ValueError(f"unsupported W4A16 w13_layout {w13_layout!r}") else: w13_layout = "packed" source_n_rotation = 0 if weight_layout == "nf3_2p1": if scale_format != "e4m3_k32": raise ValueError( "nf3_2p1 W4A16 weights require scale_format='e4m3_k32'" ) if element_dtype != "bf16": raise ValueError( "nf3_2p1 W4A16 weights are bf16-activation only (v1)" ) if epilogue_activation not in (None, "relu2"): raise ValueError( "W4A16 GEMM epilogue activation currently supports only relu2" ) if tile_n % 16 != 0 or tile_k % 16 != 0: raise ValueError("tile_n/tile_k must be multiples of 16") scale_group_size = _scale_group_size(scale_format) has_n_tile_tail = int(size_n) % int(tile_n) != 0 has_k_tile_tail = int(size_k) % int(tile_k) != 0 has_scale_k_tail = int(size_k) % scale_group_size != 0 has_logical_tail = ( has_n_tile_tail or has_k_tile_tail or has_scale_k_tail ) if has_logical_tail: if weight_layout != "modelopt" or scale_format != "e8m0_k32": raise ValueError( "W4A16 logical tail tiles are only supported for native " "E8M0 K/32 weights" ) if int(size_n) % 16 != 0: raise ValueError( "native E8M0 W4A16 tail tiles require size_n % 16 == 0" ) if int(size_k) % 8 != 0: raise ValueError( "native E8M0 W4A16 tail tiles require size_k % 8 == 0" ) if int(tile_k) % scale_group_size != 0: raise ValueError( "native E8M0 W4A16 tail tiles require tile_k multiples of 32" ) else: if size_n % tile_n != 0: raise ValueError("size_n must be divisible by tile_n") if size_k % tile_k != 0: raise ValueError("size_k must be divisible by tile_k") if scale_format == "e8m0_k32" and ( size_k % 32 != 0 or tile_k % 32 != 0 ): raise ValueError( "E8M0 K/32 W4A16 scales require size_k/tile_k multiples of 32" ) if scale_format == "e4m3_k32" and ( size_k % 32 != 0 or tile_k % 32 != 0 ): raise ValueError( "E4M3 K/32 W4A16 scales require size_k/tile_k multiples of 32" ) if moe_block_size not in _ALLOWED_ROUTED_SIZES: raise ValueError(f"unsupported moe_block_size {moe_block_size}") if moe_block_size != 8 and moe_block_size % 16 != 0: raise ValueError("moe_block_size must be 8 or a multiple of 16") cta_threads = tile_n * tile_k // 64 if cta_threads not in (128, 256): raise ValueError("W4A16 GEMM expects 128 or 256 CTA threads") self.size_m = int(size_m) self.size_n = int(size_n) self.size_k = int(size_k) self.num_experts = int(num_experts) self.top_k = int(top_k) self.mul_topk_weights = bool(mul_topk_weights) self.tile_n = int(tile_n) self.tile_k = int(tile_k) self.cta_n_blocks = int(tile_n // 16) self.cta_k_blocks = int(tile_k // 16) self.cta_threads = int(cta_threads) self.moe_block_size = int(moe_block_size) self.element_dtype = element_dtype self.is_fp16 = element_dtype == "fp16" self.epilogue_relu2 = epilogue_activation == "relu2" self.weight_layout = weight_layout self.weight_layout_nf3 = weight_layout == "nf3_2p1" self.scale_format = scale_format self.scale_format_e8m0_k32 = scale_format == "e8m0_k32" # k32 scale cadence (two K16 rows share one scale group). Shared by the # native E8M0 K/32 path and the NF3 e4m3-style K/32 path; the scale # DECODE stays keyed on scale_format_e8m0_k32 (e4m3_k32 uses the e4m3 # decode arm, its 2**116 compensation lives in the global_scale tensor). self.scale_k32 = scale_format in ("e8m0_k32", "e4m3_k32") # NF3 codebook prmt pools (host-computed compile-time constants injected # into the device dequant). Only consumed by the nf3_2p1 path; cheap to # compute unconditionally. ( self._nf3_cb0, self._nf3_cb1, self._nf3_cb2, self._nf3_cb3, ) = nf3_codebook_pools(_NF3_CODEBOOK) self.scale_group_size = int(scale_group_size) self.scale_k_groups = _covering_count(self.size_k, self.scale_group_size) self.n_tiles = _covering_count(self.size_n, self.tile_n) self.k_tiles = _covering_count(self.size_k, self.tile_k) self.covered_size_n = self.n_tiles * self.tile_n self.covered_size_k = self.k_tiles * self.tile_k self.has_n_tile_tail = bool(has_n_tile_tail) self.has_k_tile_tail = bool(has_k_tile_tail) self.has_scale_k_tail = bool(has_scale_k_tail) self.has_logical_tail = bool(has_logical_tail) self.scale_size_n = ( _e8m0_logical_tail_scale_n(self.size_n) if self.scale_format_e8m0_k32 and self.has_n_tile_tail else self.size_n ) self.scale_n_groups = self.scale_size_n // 16 self.w13_layout = w13_layout self.source_n_rotation = int(source_n_rotation) self.single_token_route_fast_path = bool(single_token_route_fast_path) self.direct_topk_routes = bool(direct_topk_routes) self.fused_topk_sum = bool(fused_topk_sum) self.fused_sum_topk = int(fused_sum_topk) if self.fused_topk_sum and not self.direct_topk_routes: raise ValueError("fused_topk_sum requires direct_topk_routes") if self.fused_topk_sum and self.fused_sum_topk < 1: raise ValueError("fused_sum_topk must be >= 1") self.cta_m_blocks = int(_covering_count(moe_block_size, 16)) self.uses_m_block_8 = moe_block_size == 8 self.max_m_blocks = int(max_m_blocks) if torch.cuda.is_available(): props = torch.cuda.get_device_properties(torch.cuda.current_device()) self.sms = int(props.multi_processor_count) max_shared_mem = int( getattr(props, "shared_memory_per_block_optin", _DEFAULT_MAX_SHARED_MEM) ) else: self.sms = 120 max_shared_mem = _DEFAULT_MAX_SHARED_MEM self.blocks_per_sm = _determine_blocks_per_sm( problem_m=self.size_m, problem_n=self.covered_size_n, top_k=self.top_k, cta_threads=self.cta_threads, cta_m_blocks=self.cta_m_blocks, tile_n=self.tile_n, tile_k=self.tile_k, uses_m_block_8=self.uses_m_block_8, sms=self.sms, max_shared_mem=max_shared_mem, scale_format=scale_format, ) # W4A16 shared-memory geometry, in int4 units unless noted. self.a_sh_stride = 16 * self.cta_k_blocks // 8 self.a_sh_stage = self.a_sh_stride * (16 * self.cta_m_blocks) self.a_gl_rd_delta_o = 16 * self.cta_k_blocks // 8 self.a_sh_wr_delta = self.a_sh_stride * ( self.cta_threads // self.a_gl_rd_delta_o ) self.a_sh_wr_iters = _covering_count(self.a_sh_stage, self.a_sh_wr_delta) self.a_sh_rd_delta_i = self.a_sh_stride * 16 self.b_sh_stride = ((self.cta_n_blocks * 16) * 16 // _PACK_FACTOR) // 4 self.b_thread_vecs = 1 self.b_sh_stride_threads = self.b_sh_stride self.b_sh_stage = self.b_sh_stride * self.cta_k_blocks self.b_sh_wr_iters = self.b_sh_stage // self.cta_threads # NF3 ("nf3_2p1") stages 12-byte triples per 32-code unit instead of the # 16-byte int4 unit; the per-thread unit count is unchanged, only the # bytes-per-unit and the 16-byte cp.async chunk count differ. self.b_unit_bytes = 12 if self.weight_layout_nf3 else 16 self.b_sh_stage_bytes = self.b_sh_stage * self.b_unit_bytes if self.weight_layout_nf3: if self.b_sh_stage_bytes % 16 != 0: raise ValueError( "nf3_2p1 B stage bytes must be a multiple of 16; " f"got {self.b_sh_stage_bytes}" ) self.b_sh_chunks = self.b_sh_stage_bytes // 16 self.b_sh_wr_iters_nf3 = _covering_count( self.b_sh_chunks, self.cta_threads ) else: self.b_sh_chunks = self.b_sh_stage self.b_sh_wr_iters_nf3 = self.b_sh_wr_iters self.s_sh_stride = 16 * self.cta_n_blocks // 16 self.s_tb_groups = ( self.cta_k_blocks // 2 if self.scale_k32 else self.cta_k_blocks ) self.s_sh_stage = self.s_tb_groups * self.s_sh_stride self.tb_n_warps = self.cta_n_blocks // 4 sh_block_route_indices = self.moe_block_size // 4 sh_rd_block_route_indices = self.moe_block_size // 4 sh_block_topk_weights = self.moe_block_size // 2 self.sh_valid_count_off = ( sh_block_route_indices + sh_rd_block_route_indices + sh_block_topk_weights ) self.sh_route_off = 0 self.sh_rd_route_off = sh_block_route_indices self.sh_topk_off = sh_block_route_indices + sh_rd_block_route_indices sh_red_size = (2 * self.cta_n_blocks + 1) * 16 * self.cta_m_blocks # B region size in int4 (16-byte) units. NF3 units are 12 bytes, so the # region is 0.75x; round the total up to a 16-byte multiple so the # following SMEM regions keep their 16-byte alignment (exact for all # supported tiles since b_sh_stage is a multiple of 4). sh_b_size = _covering_count(_STAGES * self.b_sh_stage_bytes, 16) sh_size_min = min(sh_red_size, sh_b_size) sh_size_max = max(sh_red_size, sh_b_size) sh_bias_size = self.cta_n_blocks * 16 // 8 sh_b_red_bias_size = max(sh_size_max, sh_size_min + sh_bias_size) self.sh_b_off = self.sh_valid_count_off self.sh_red_off = self.sh_valid_count_off self.sh_s_off = self.sh_valid_count_off + sh_b_red_bias_size self.sh_a_off = self.sh_s_off + _STAGES * self.s_sh_stage self.shared_int4 = self.sh_a_off + _STAGES * self.a_sh_stage self.shared_words = self.shared_int4 * 4 @property def __cache_key__(self) -> tuple[object, ...]: return ( self.size_n, self.size_k, self.covered_size_n, self.covered_size_k, self.scale_k_groups, self.has_logical_tail, self.num_experts, self.top_k, self.mul_topk_weights, self.tile_n, self.tile_k, self.cta_threads, self.moe_block_size, self.element_dtype, self.epilogue_relu2, self.weight_layout, self.scale_format, self.w13_layout, self.source_n_rotation, self.single_token_route_fast_path, self.direct_topk_routes, self.fused_topk_sum, self.fused_sum_topk, self.cta_m_blocks, self.uses_m_block_8, self.shared_words, # Launch bounds are part of the compiled kernel. Keep binaries # planned for different residency targets out of the same cache # entry even when their arithmetic geometry otherwise matches. self.blocks_per_sm, ) @cute.jit def _activation_smem_permuted_offset(self, i: Int32) -> Int32: row = i // Int32(self.a_gl_rd_delta_o) return Int32(self.a_gl_rd_delta_o) * row + ( (i - row * Int32(self.a_gl_rd_delta_o)) ^ (row & Int32(7)) ) @cute.jit def _int4_addr(self, smem_base: Int32, int4_off: Int32) -> Int32: return smem_base + int4_off * Int32(16) @cute.jit def _dequant_e2m1x4_to_elem2x2(self, packed: Uint32): if cutlass.const_expr(self.is_fp16): return packed_dequant_e2m1x4_to_half2x2(packed) return packed_dequant_e2m1x4_to_bfloat2x2(packed) @cute.jit def _dequant_e4m3x4_to_elem2x2(self, packed: Uint32): if cutlass.const_expr(self.is_fp16): return packed_dequant_e4m3x4_to_half2x2(packed) return packed_dequant_e4m3x4_to_bfloat2x2(packed) @cute.jit def _dequant_scale_x4_to_elem2x2(self, packed: Uint32): if cutlass.const_expr(self.scale_format_e8m0_k32): if cutlass.const_expr(self.is_fp16): return packed_dequant_e8m0x4_to_half2x2(packed) return packed_dequant_e8m0x4_to_bfloat2x2(packed) if cutlass.const_expr(self.is_fp16): return packed_dequant_e4m3x4_to_half2x2(packed) return packed_dequant_e4m3x4_to_bfloat2x2(packed) @cute.jit def _elem2_mul(self, a: Uint32, b: Uint32) -> Uint32: if cutlass.const_expr(self.is_fp16): return half2_mul(a, b) return bfloat2_mul(a, b) @cute.jit def _broadcast_f32_to_elem2(self, x: cutlass.Float32) -> Uint32: if cutlass.const_expr(self.is_fp16): return broadcast_f32_to_half2(x) return broadcast_f32_to_bfloat2(x) @cute.jit def _pack_f32x2_to_elem2(self, x0: cutlass.Float32, x1: cutlass.Float32) -> Uint32: if cutlass.const_expr(self.is_fp16): return pack_f32x2_to_f16x2(x0, x1) return pack_f32x2_to_bfloat2(x0, x1) @cute.jit def _elem2_to_f32x2(self, packed: Uint32): if cutlass.const_expr(self.is_fp16): return half2_to_float2_scaled(packed, cutlass.Float32(1.0)) return bfloat2_to_float2_scaled(packed, cutlass.Float32(1.0)) @cute.jit def _relu2_elem2(self, packed: Uint32) -> Uint32: x0, x1 = self._elem2_to_f32x2(packed) if x0 < cutlass.Float32(0.0): x0 = cutlass.Float32(0.0) if x1 < cutlass.Float32(0.0): x1 = cutlass.Float32(0.0) return self._pack_f32x2_to_elem2(x0 * x0, x1 * x1) @cute.jit def _st_shared_elem_from_f32(self, addr: Int32, val: cutlass.Float32): if cutlass.const_expr(self.is_fp16): st_shared_f16_from_f32(addr, val) else: st_shared_bf16_from_f32(addr, val) @cute.jit def _mma_m16n8k16_f32( self, d0: cutlass.Float32, d1: cutlass.Float32, d2: cutlass.Float32, d3: cutlass.Float32, a0: Uint32, a1: Uint32, a2: Uint32, a3: Uint32, b0: Uint32, b1: Uint32, ): if cutlass.const_expr(self.is_fp16): return f16_mma_m16n8k16_f32(d0, d1, d2, d3, a0, a1, a2, a3, b0, b1) return bf16_mma_m16n8k16_f32(d0, d1, d2, d3, a0, a1, a2, a3, b0, b1) @cute.jit def _mma_rhs_fragments_as_mma_a_m16n8k16_f32( self, d0: cutlass.Float32, d1: cutlass.Float32, d2: cutlass.Float32, d3: cutlass.Float32, b0_0: Uint32, b1_0: Uint32, b0_1: Uint32, b1_1: Uint32, a0: Uint32, a1: Uint32, ): if cutlass.const_expr(self.is_fp16): return f16_mma_rhs_fragments_as_mma_a_m16n8k16_f32( d0, d1, d2, d3, b0_0, b1_0, b0_1, b1_1, a0, a1 ) return bf16_mma_rhs_fragments_as_mma_a_m16n8k16_f32( d0, d1, d2, d3, b0_0, b1_0, b0_1, b1_1, a0, a1 ) @cute.jit def __call__( self, a_bf16_flat: cute.Tensor, b_i32_flat: cute.Tensor, c_bf16_flat: cute.Tensor, scales_i32_flat: cute.Tensor, global_scale: cute.Tensor, packed_route_indices: cute.Tensor, block_expert_ids: cute.Tensor, packed_route_count: cute.Tensor, topk_weights_flat: cute.Tensor, c_tmp_f32_flat: cute.Tensor, locks_i32_flat: cute.Tensor, active_m: cutlass.Int32, grid_x: cutlass.Int32, stream: cuda.CUstream, ): grid = (grid_x, 1, 1) self.kernel( a_bf16_flat, b_i32_flat, c_bf16_flat, scales_i32_flat, global_scale, packed_route_indices, block_expert_ids, packed_route_count, topk_weights_flat, c_tmp_f32_flat, locks_i32_flat, active_m, ).launch( grid=grid, block=[self.cta_threads, 1, 1], min_blocks_per_mp=self.blocks_per_sm, stream=stream, ) @cute.kernel def kernel( self, a_bf16_flat: cute.Tensor, b_i32_flat: cute.Tensor, c_bf16_flat: cute.Tensor, scales_i32_flat: cute.Tensor, global_scale: cute.Tensor, packed_route_indices: cute.Tensor, block_expert_ids: cute.Tensor, packed_route_count: cute.Tensor, topk_weights_flat: cute.Tensor, c_tmp_f32_flat: cute.Tensor, locks_i32_flat: cute.Tensor, active_m: cutlass.Int32, ): tidx, _, _ = cute.arch.thread_idx() bidx, _, _ = cute.arch.block_idx() tid = Int32(tidx) cta = Int32(bidx) smem = cutlass.utils.SmemAllocator() @cute.struct class Storage: words: cute.struct.Align[ cute.struct.MemRange[cutlass.Uint32, self.shared_words], 1024, ] storage = smem.allocate(Storage) smem_base = shared_ptr_to_u32(storage.words.data_ptr()) grid_x, _, _ = cute.arch.grid_dim() self._run_persistent_gemm( a_bf16_flat, b_i32_flat, c_bf16_flat, scales_i32_flat, global_scale, packed_route_indices, block_expert_ids, packed_route_count, topk_weights_flat, c_tmp_f32_flat, locks_i32_flat, smem_base, tid, cta, Int32(grid_x), Int32(active_m), ) @cute.jit def _run_persistent_gemm( self, a_bf16_flat: cute.Tensor, b_i32_flat: cute.Tensor, c_bf16_flat: cute.Tensor, scales_i32_flat: cute.Tensor, global_scale: cute.Tensor, packed_route_indices: cute.Tensor, block_expert_ids: cute.Tensor, packed_route_count: cute.Tensor, topk_weights_flat: cute.Tensor, c_tmp_f32_flat: cute.Tensor, locks_i32_flat: cute.Tensor, smem_base: Int32, tid: Int32, cta: Int32, grid_x: Int32, active_size_m: Int32, ): n_tiles = Int32(self.n_tiles) route_blocks = active_size_m * Int32(self.top_k) if cutlass.const_expr(not self.direct_topk_routes): route_blocks = packed_route_count[Int32(0)].to(Int32) // Int32( self.moe_block_size ) k_tiles = Int32(self.k_tiles) global_mn_tiles = route_blocks * n_tiles tail_mn_tiles = global_mn_tiles full_grid_mn_iters = Int32(0) force_one_tile_per_cta = Int32(0) if cutlass.const_expr(self.uses_m_block_8): # TC-decode small-M: when every mn-tile fits inside the launched grid # (FC1 has only route_blocks*n_tiles tiles, far fewer than grid_x), # the default tail path fans each mn-tile across multiple CTAs along # K and pays a lock-serialized cross-CTA split-K finalize plus the # reduction-turn handshake. Instead give the first global_mn_tiles # CTAs exactly one full mn-tile (all k_tiles, reduce_slice_count==1, # no finalize, no lock traffic) and idle the rest. grid_x and the # grid-barrier participant count are unchanged (so FC2 coverage is # untouched); only FC1's intra-GEMM work partition changes. Numerically # identical: a single CTA computes the whole K-reduction per tile. if global_mn_tiles <= grid_x: force_one_tile_per_cta = Int32(1) if force_one_tile_per_cta != Int32(0): tail_mn_tiles = Int32(0) full_grid_mn_iters = Int32(1) elif global_mn_tiles > grid_x: tail_mn_tiles = global_mn_tiles - (global_mn_tiles // grid_x) * grid_x if tail_mn_tiles * Int32(3) <= grid_x: tail_mn_tiles += grid_x full_grid_mn_iters = (global_mn_tiles - tail_mn_tiles) // grid_x iters = (k_tiles * tail_mn_tiles + grid_x - Int32(1)) // grid_x lock_slot = Int32(0) if tail_mn_tiles >= grid_x: lock_slot = cta else: lock_slot = (iters * cta) // k_tiles - Int32(1) in_tail_region = Int32(0) has_work = Int32(1) work_mn_tile = cta reduce_k_tile = Int32(0) route_block_idx = Int32(0) output_n_tile = Int32(0) if iters == Int32(0) and full_grid_mn_iters == Int32(0): has_work = Int32(0) while has_work != Int32(0): reduce_tile_count = Int32(0) reduce_slice_count = Int32(1) reduce_slice_idx = Int32(0) if in_tail_region == Int32(0) and full_grid_mn_iters > Int32(0): route_block_idx = work_mn_tile // n_tiles output_n_tile = work_mn_tile - route_block_idx * n_tiles reduce_k_tile = Int32(0) reduce_tile_count = k_tiles full_grid_mn_iters -= Int32(1) else: if in_tail_region == Int32(0): in_tail_region = Int32(1) tail_mn_base = global_mn_tiles - tail_mn_tiles cta_iter_start = iters * cta work_mn_tile = cta_iter_start // k_tiles reduce_k_tile = cta_iter_start - work_mn_tile * k_tiles global_mn_tile = work_mn_tile + tail_mn_base route_block_idx = global_mn_tile // n_tiles output_n_tile = global_mn_tile - route_block_idx * n_tiles if work_mn_tile < tail_mn_tiles and iters > Int32(0): reduce_tile_count = iters * (cta + Int32(1)) - ( k_tiles * work_mn_tile + reduce_k_tile ) if reduce_tile_count < Int32(0): reduce_tile_count = Int32(0) if reduce_k_tile + reduce_tile_count > k_tiles: reduce_tile_count = k_tiles - reduce_k_tile if reduce_tile_count > Int32(0): first_reduce_boundary = iters * ( (k_tiles * work_mn_tile + iters - Int32(1)) // iters ) if first_reduce_boundary <= k_tiles * (work_mn_tile + Int32(1)): reduce_boundary_offset = ( first_reduce_boundary - k_tiles * work_mn_tile ) reduce_slice_count = ( k_tiles - reduce_boundary_offset + iters - Int32(1) ) // iters if reduce_boundary_offset > Int32(0): reduce_slice_count += Int32(1) reduce_boundary_delta = iters * cta - first_reduce_boundary if reduce_boundary_delta < Int32(0): reduce_slice_idx = reduce_slice_count - Int32(1) else: if reduce_boundary_offset == Int32( 0 ) and reduce_boundary_delta == Int32(0): reduce_slice_idx = reduce_slice_count - Int32(1) else: reduce_slice_idx = ( reduce_slice_count - Int32(1) - reduce_boundary_delta // iters ) if reduce_boundary_offset > Int32(0): reduce_slice_idx -= Int32(1) if tail_mn_tiles >= grid_x: if reduce_slice_count > Int32( 1 ) and reduce_slice_idx == reduce_slice_count - Int32(1): lock_slot += Int32(1) else: lock_slot += Int32(1) else: has_work = Int32(0) else: has_work = Int32(0) if ( has_work != Int32(0) and reduce_tile_count > Int32(0) and route_block_idx < route_blocks ): if cutlass.const_expr(self.direct_topk_routes): expert_idx = packed_route_indices[route_block_idx].to(Int32) else: expert_idx = block_expert_ids[route_block_idx].to(Int32) if expert_idx >= Int32(0): self._run_tile( a_bf16_flat, b_i32_flat, c_bf16_flat, scales_i32_flat, global_scale, packed_route_indices, topk_weights_flat, c_tmp_f32_flat, locks_i32_flat, smem_base, tid, route_block_idx, expert_idx, output_n_tile, reduce_k_tile, reduce_tile_count, reduce_slice_count, reduce_slice_idx, lock_slot, active_size_m, ) if has_work != Int32(0): if in_tail_region == Int32(0): work_mn_tile += grid_x else: reduce_k_tile = Int32(0) work_mn_tile += Int32(1) output_n_tile += Int32(1) if output_n_tile == n_tiles: output_n_tile = Int32(0) route_block_idx += Int32(1) @cute.jit def _read_moe_block_data( self, packed_route_indices: cute.Tensor, topk_weights_flat: cute.Tensor, smem_base: Int32, tid: Int32, route_block_idx: Int32, global_scale_f32: cutlass.Float32, active_size_m: Int32, ) -> Int32: if cutlass.const_expr(self.direct_topk_routes): if tid == Int32(0): idx = route_block_idx st_shared_i32(smem_base + Int32(self.sh_route_off * 16), idx) rd_row = idx // Int32(self.top_k) st_shared_i32(smem_base + Int32(self.sh_rd_route_off * 16), rd_row) if cutlass.const_expr(self.mul_topk_weights): topk = topk_weights_flat[idx].to(cutlass.Float32) * global_scale_f32 st_shared_u32( smem_base + Int32(self.sh_topk_off * 16), self._broadcast_f32_to_elem2(topk), ) cute.arch.sync_threads() return Int32(1) if cutlass.const_expr(self.single_token_route_fast_path): if tid == Int32(0): idx = packed_route_indices[ route_block_idx * Int32(self.moe_block_size) ].to(Int32) st_shared_i32(smem_base + Int32(self.sh_route_off * 16), idx) rd_row = idx // Int32(self.top_k) st_shared_i32(smem_base + Int32(self.sh_rd_route_off * 16), rd_row) if cutlass.const_expr(self.mul_topk_weights): safe_idx = idx if idx >= active_size_m * Int32(self.top_k): safe_idx = Int32(0) topk = ( topk_weights_flat[safe_idx].to(cutlass.Float32) * global_scale_f32 ) st_shared_u32( smem_base + Int32(self.sh_topk_off * 16), self._broadcast_f32_to_elem2(topk), ) cute.arch.sync_threads() return Int32(1) route_indices_int4_addr = self._int4_addr( smem_base, Int32(self.sh_route_off) + tid ) route_indices_gmem = get_ptr_as_int64( packed_route_indices, route_block_idx * Int32(self.moe_block_size) + tid * Int32(4), ) cp_async4_shared_global_pred( route_indices_int4_addr, route_indices_gmem, (tid < Int32(self.moe_block_size // 4)).to(Int32), ) cute.arch.cp_async_commit_group() cute.arch.cp_async_wait_group(0) cute.arch.sync_threads() if tid >= Int32(self.cta_threads - 32): size_per_thread = _covering_count(self.moe_block_size, 32) lane = tid - Int32(self.cta_threads - 32) local_count = Int32(0) for i in cutlass.range_constexpr(size_per_thread): j = lane * Int32(size_per_thread) + Int32(i) if j < Int32(self.moe_block_size): idx = ld_shared_i32_relaxed( smem_base + Int32(self.sh_route_off * 16) + j * Int32(4) ) if idx < active_size_m * Int32(self.top_k): local_count += Int32(1) valid = cute.arch.warp_redux_sync(local_count, "add") if lane == Int32(0): st_shared_i32(smem_base + Int32(self.sh_valid_count_off * 16), valid) if tid < Int32(self.moe_block_size): idx = ld_shared_i32_relaxed( smem_base + Int32(self.sh_route_off * 16) + tid * Int32(4) ) rd_row = idx // Int32(self.top_k) st_shared_i32( smem_base + Int32(self.sh_rd_route_off * 16) + tid * Int32(4), rd_row, ) if cutlass.const_expr(self.mul_topk_weights): safe_idx = idx if idx >= active_size_m * Int32(self.top_k): safe_idx = Int32(0) topk = ( topk_weights_flat[safe_idx].to(cutlass.Float32) * global_scale_f32 ) packed_topk = self._broadcast_f32_to_elem2(topk) # top-k weights are cached as packed element pairs. topk_word_addr = ( smem_base + Int32(self.sh_topk_off * 16) + tid * Int32(4) ) st_shared_u32(topk_word_addr, packed_topk) cute.arch.sync_threads() valid_count = ld_shared_i32_relaxed( smem_base + Int32(self.sh_valid_count_off * 16) ) cute.arch.sync_threads() return valid_count @cute.jit def _run_tile( self, a_bf16_flat: cute.Tensor, b_i32_flat: cute.Tensor, c_bf16_flat: cute.Tensor, scales_i32_flat: cute.Tensor, global_scale: cute.Tensor, packed_route_indices: cute.Tensor, topk_weights_flat: cute.Tensor, c_tmp_f32_flat: cute.Tensor, locks_i32_flat: cute.Tensor, smem_base: Int32, tid: Int32, route_block_idx: Int32, expert_idx: Int32, output_n_tile: Int32, reduce_k_tile: Int32, reduce_tile_count: Int32, reduce_slice_count: Int32, reduce_slice_idx: Int32, lock_slot: Int32, active_size_m: Int32, ): if cutlass.const_expr(self.uses_m_block_8): self._run_tile_m8( a_bf16_flat, b_i32_flat, c_bf16_flat, scales_i32_flat, global_scale, packed_route_indices, topk_weights_flat, c_tmp_f32_flat, locks_i32_flat, smem_base, tid, route_block_idx, expert_idx, output_n_tile, reduce_k_tile, reduce_tile_count, reduce_slice_count, reduce_slice_idx, lock_slot, active_size_m, ) else: self._run_tile_large_m( a_bf16_flat, b_i32_flat, c_bf16_flat, scales_i32_flat, global_scale, packed_route_indices, topk_weights_flat, c_tmp_f32_flat, locks_i32_flat, smem_base, tid, route_block_idx, expert_idx, output_n_tile, reduce_k_tile, reduce_tile_count, reduce_slice_count, reduce_slice_idx, lock_slot, active_size_m, ) @cute.jit def _tile_common_prologue( self, global_scale: cute.Tensor, packed_route_indices: cute.Tensor, topk_weights_flat: cute.Tensor, smem_base: Int32, tid: Int32, route_block_idx: Int32, expert_idx: Int32, output_n_tile: Int32, active_size_m: Int32, ): global_scale_f32 = global_scale[expert_idx].to(cutlass.Float32) if cutlass.const_expr(self.scale_format_e8m0_k32): if cutlass.const_expr(self.is_fp16): global_scale_f32 *= cutlass.Float32(_E8M0_K32_FP16_GLOBAL_COMPENSATION) else: global_scale_f32 *= cutlass.Float32(_E8M0_K32_BF16_GLOBAL_COMPENSATION) block_valid_rows = self._read_moe_block_data( packed_route_indices, topk_weights_flat, smem_base, tid, route_block_idx, global_scale_f32, active_size_m, ) ( a_gl_stride, b_gl_stride, s_gl_stride, scales_expert_off, b_gl_rd_base, a_gl_rd_row, a_gl_rd_col0, a_sh_wr, a_rows_per_iter, b_sh_rd, s_sh_rd, ) = self._tile_stream_offsets(tid, expert_idx, output_n_tile) return ( global_scale_f32, block_valid_rows, a_gl_stride, b_gl_stride, s_gl_stride, scales_expert_off, b_gl_rd_base, a_gl_rd_row, a_gl_rd_col0, a_sh_wr, a_rows_per_iter, b_sh_rd, s_sh_rd, ) @cute.jit def _tile_stream_offsets(self, tid: Int32, expert_idx: Int32, output_n_tile: Int32): a_gl_stride = Int32(self.size_k // 8) b_gl_stride = Int32(16 * self.size_n // (_PACK_FACTOR * 4)) s_gl_stride = Int32(self.scale_n_groups) if cutlass.const_expr(self.scale_k32): if cutlass.const_expr( self.scale_format_e8m0_k32 and self.has_logical_tail ): scales_expert_stride = Int32(self.scale_k_groups * self.scale_n_groups) else: scales_expert_stride = Int32((self.size_n * self.size_k) // (32 * 16)) else: scales_expert_stride = Int32((self.size_n * self.size_k) // (16 * 16)) b_expert_off = ( Int32((self.size_n * self.size_k) // (_PACK_FACTOR * 4)) * expert_idx ) scales_expert_off = scales_expert_stride * expert_idx a_gl_rd_row = tid // Int32(self.a_gl_rd_delta_o) a_gl_rd_col0 = tid - a_gl_rd_row * Int32(self.a_gl_rd_delta_o) a_sh_wr = Int32(self.a_sh_stride) * (tid // Int32(self.a_gl_rd_delta_o)) + ( tid - (tid // Int32(self.a_gl_rd_delta_o)) * Int32(self.a_gl_rd_delta_o) ) a_rows_per_iter = Int32(self.cta_threads // self.a_gl_rd_delta_o) if cutlass.const_expr(self.cta_threads <= self.b_sh_stride): b_gl_rd_base = tid else: b_gl_rd_base = b_gl_stride * (tid // Int32(self.b_sh_stride)) + ( tid % Int32(self.b_sh_stride) ) b_gl_rd_base += b_expert_off + Int32(self.b_sh_stride) * output_n_tile b_sh_rd = tid b_sh_rd += (b_sh_rd // Int32(self.b_sh_stride)) * Int32( self.b_sh_stride * (self.b_sh_wr_iters - 1) ) s_sh_rd = Int32(8) * ((tid // Int32(32)) % Int32(self.tb_n_warps)) + ( tid & Int32(31) ) // Int32(4) return ( a_gl_stride, b_gl_stride, s_gl_stride, scales_expert_off, b_gl_rd_base, a_gl_rd_row, a_gl_rd_col0, a_sh_wr, a_rows_per_iter, b_sh_rd, s_sh_rd, ) @cute.jit def _a_shared_read_offset(self, tid: Int32, lanes_per_row: cutlass.Constexpr[int]): a_sh_rd = Int32(self.a_sh_stride) * ( (tid & Int32(31)) % Int32(lanes_per_row) ) + (tid & Int32(31)) // Int32(lanes_per_row) a_sh_rd += ( Int32(2) * ((tid // Int32(32)) // Int32(self.tb_n_warps)) * Int32(self.b_sh_wr_iters) ) return a_sh_rd @cute.jit def _run_tile_m8( self, a_bf16_flat: cute.Tensor, b_i32_flat: cute.Tensor, c_bf16_flat: cute.Tensor, scales_i32_flat: cute.Tensor, global_scale: cute.Tensor, packed_route_indices: cute.Tensor, topk_weights_flat: cute.Tensor, c_tmp_f32_flat: cute.Tensor, locks_i32_flat: cute.Tensor, smem_base: Int32, tid: Int32, route_block_idx: Int32, expert_idx: Int32, output_n_tile: Int32, reduce_k_tile: Int32, reduce_tile_count: Int32, reduce_slice_count: Int32, reduce_slice_idx: Int32, lock_slot: Int32, active_size_m: Int32, ): ( global_scale_f32, block_valid_rows, a_gl_stride, b_gl_stride, s_gl_stride, scales_expert_off, b_gl_rd_base, a_gl_rd_row, a_gl_rd_col0, a_sh_wr, a_rows_per_iter, b_sh_rd, s_sh_rd, ) = self._tile_common_prologue( global_scale, packed_route_indices, topk_weights_flat, smem_base, tid, route_block_idx, expert_idx, output_n_tile, active_size_m, ) a_sh_rd = self._a_shared_read_offset(tid, 8) acc = cute.make_rmem_tensor((4, 4), cutlass.Float32) acc.fill(0.0) k_tiles = reduce_tile_count self._prefetch_initial_tiles( a_bf16_flat, b_i32_flat, scales_i32_flat, smem_base, tid, k_tiles, reduce_k_tile, block_valid_rows, a_gl_stride, b_gl_stride, s_gl_stride, scales_expert_off, b_gl_rd_base, a_gl_rd_row, a_gl_rd_col0, a_sh_wr, a_rows_per_iter, output_n_tile, expert_idx, ) b_scale_cur = cute.make_rmem_tensor((2, 4), Uint32) b_scale_next = cute.make_rmem_tensor((2, 4), Uint32) self._load_b_scale_register_bundle( b_scale_cur, smem_base, tid, b_sh_rd, s_sh_rd, Int32(0), Int32(0), ) a_regs_cur = cute.make_rmem_tensor((2,), Uint32) a_regs_next = cute.make_rmem_tensor((2,), Uint32) self._load_a_register_bundle( a_regs_cur, smem_base, a_sh_rd, Int32(0), Int32(0), True, ) self._run_mma_pipeline( a_bf16_flat, b_i32_flat, scales_i32_flat, smem_base, tid, acc, b_scale_cur, b_scale_next, a_regs_cur, a_regs_next, b_sh_rd, s_sh_rd, a_sh_rd, k_tiles, reduce_k_tile, block_valid_rows, a_gl_stride, b_gl_stride, s_gl_stride, scales_expert_off, b_gl_rd_base, a_gl_rd_row, a_gl_rd_col0, a_sh_wr, a_rows_per_iter, output_n_tile, expert_idx, True, ) self._finish_tile( acc, c_bf16_flat, c_tmp_f32_flat, locks_i32_flat, smem_base, tid, output_n_tile, block_valid_rows, global_scale_f32, reduce_slice_count, reduce_slice_idx, lock_slot, True, ) @cute.jit def _run_tile_large_m( self, a_bf16_flat: cute.Tensor, b_i32_flat: cute.Tensor, c_bf16_flat: cute.Tensor, scales_i32_flat: cute.Tensor, global_scale: cute.Tensor, packed_route_indices: cute.Tensor, topk_weights_flat: cute.Tensor, c_tmp_f32_flat: cute.Tensor, locks_i32_flat: cute.Tensor, smem_base: Int32, tid: Int32, route_block_idx: Int32, expert_idx: Int32, output_n_tile: Int32, reduce_k_tile: Int32, reduce_tile_count: Int32, reduce_slice_count: Int32, reduce_slice_idx: Int32, lock_slot: Int32, active_size_m: Int32, ): ( global_scale_f32, block_valid_rows, a_gl_stride, b_gl_stride, s_gl_stride, scales_expert_off, b_gl_rd_base, a_gl_rd_row, a_gl_rd_col0, a_sh_wr, a_rows_per_iter, b_sh_rd, s_sh_rd, ) = self._tile_common_prologue( global_scale, packed_route_indices, topk_weights_flat, smem_base, tid, route_block_idx, expert_idx, output_n_tile, active_size_m, ) a_sh_rd = self._a_shared_read_offset(tid, 16) acc = cute.make_rmem_tensor((self.cta_m_blocks, 4, 2, 4), cutlass.Float32) acc.fill(0.0) k_tiles = reduce_tile_count self._prefetch_initial_tiles( a_bf16_flat, b_i32_flat, scales_i32_flat, smem_base, tid, k_tiles, reduce_k_tile, block_valid_rows, a_gl_stride, b_gl_stride, s_gl_stride, scales_expert_off, b_gl_rd_base, a_gl_rd_row, a_gl_rd_col0, a_sh_wr, a_rows_per_iter, output_n_tile, expert_idx, ) b_scale_cur = cute.make_rmem_tensor((2, 4), Uint32) b_scale_next = cute.make_rmem_tensor((2, 4), Uint32) self._load_b_scale_register_bundle( b_scale_cur, smem_base, tid, b_sh_rd, s_sh_rd, Int32(0), Int32(0), ) a_regs = cute.make_rmem_tensor((self.cta_m_blocks, 4), Uint32) a_regs_next = cute.make_rmem_tensor((self.cta_m_blocks, 4), Uint32) self._load_a_register_bundle( a_regs, smem_base, a_sh_rd, Int32(0), Int32(0), False, ) self._run_mma_pipeline( a_bf16_flat, b_i32_flat, scales_i32_flat, smem_base, tid, acc, b_scale_cur, b_scale_next, a_regs, a_regs_next, b_sh_rd, s_sh_rd, a_sh_rd, k_tiles, reduce_k_tile, block_valid_rows, a_gl_stride, b_gl_stride, s_gl_stride, scales_expert_off, b_gl_rd_base, a_gl_rd_row, a_gl_rd_col0, a_sh_wr, a_rows_per_iter, output_n_tile, expert_idx, False, ) self._finish_tile( acc, c_bf16_flat, c_tmp_f32_flat, locks_i32_flat, smem_base, tid, output_n_tile, block_valid_rows, global_scale_f32, reduce_slice_count, reduce_slice_idx, lock_slot, False, ) @cute.jit def _run_mma_pipeline( self, a_bf16_flat: cute.Tensor, b_i32_flat: cute.Tensor, scales_i32_flat: cute.Tensor, smem_base: Int32, tid: Int32, acc: cute.Tensor, b_scale_cur: cute.Tensor, b_scale_next: cute.Tensor, a_regs_cur: cute.Tensor, a_regs_next: cute.Tensor, b_sh_rd: Int32, s_sh_rd: Int32, a_sh_rd: Int32, k_tiles: Int32, reduce_k_tile: Int32, block_valid_rows: Int32, a_gl_stride: Int32, b_gl_stride: Int32, s_gl_stride: Int32, scales_expert_off: Int32, b_gl_rd_base: Int32, a_gl_rd_row: Int32, a_gl_rd_col0: Int32, a_sh_wr: Int32, a_rows_per_iter: Int32, output_n_tile: Int32, expert_idx: Int32, uses_m_block_8: cutlass.Constexpr[bool], ): b_frag = cute.make_rmem_tensor((2, 2), Uint32) tile_idx = Int32(0) while tile_idx < k_tiles: for pipe in cutlass.range_constexpr(_STAGES): if tile_idx < k_tiles: for kk in cutlass.range_constexpr(self.b_sh_wr_iters): self._load_next_fragment_bundle( b_scale_next, a_regs_next, smem_base, tid, b_sh_rd, s_sh_rd, a_sh_rd, pipe, kk, tile_idx, k_tiles, uses_m_block_8, ) self._prefetch_pipeline_step( a_bf16_flat, b_i32_flat, scales_i32_flat, smem_base, tid, pipe, kk, tile_idx, k_tiles, reduce_k_tile, block_valid_rows, a_gl_stride, b_gl_stride, s_gl_stride, scales_expert_off, b_gl_rd_base, a_gl_rd_row, a_gl_rd_col0, a_sh_wr, a_rows_per_iter, output_n_tile, expert_idx, ) for jj in cutlass.range_constexpr(4): if cutlass.const_expr(self.weight_layout_nf3): # NF3 triple: lo16 = half (jj % 2) of word # (jj // 2); hi8 = byte jj of the hi word; scale # register is the same per-jj lane as packed. lo_w = b_scale_cur[0, jj // 2] lo16 = (lo_w >> Uint32(16 * (jj % 2))) & Uint32(0xFFFF) hi8 = (b_scale_cur[0, 2] >> Uint32(8 * jj)) & Uint32( 0xFF ) s = b_scale_cur[1, jj] self._scaled_dequant_b_fragment_nf3( b_frag, lo16, hi8, s ) else: q, s = self._select_b_scale_register(jj, b_scale_cur) self._scaled_dequant_b_fragment(b_frag, q, s) if cutlass.const_expr(uses_m_block_8): self._mma_accumulate_m8( acc, jj, a_regs_cur, b_frag, ) else: for mb in cutlass.range_constexpr(self.cta_m_blocks): self._mma_accumulate_large_m( acc, a_regs_cur, mb, jj, b_frag, ) if cutlass.const_expr(uses_m_block_8): self._copy_a_register_bundle( a_regs_cur, a_regs_next, uses_m_block_8, ) self._copy_b_scale_register_bundle( b_scale_cur, b_scale_next ) else: self._copy_b_scale_register_bundle( b_scale_cur, b_scale_next ) self._copy_a_register_bundle( a_regs_cur, a_regs_next, uses_m_block_8, ) tile_idx += Int32(1) cute.arch.sync_threads() if tile_idx < k_tiles: self._load_b_scale_register_bundle( b_scale_cur, smem_base, tid, b_sh_rd, s_sh_rd, Int32(0), Int32(0), ) self._load_a_register_bundle( a_regs_cur, smem_base, a_sh_rd, Int32(0), Int32(0), uses_m_block_8, ) @cute.jit def _finish_tile( self, acc: cute.Tensor, c_bf16_flat: cute.Tensor, c_tmp_f32_flat: cute.Tensor, locks_i32_flat: cute.Tensor, smem_base: Int32, tid: Int32, output_n_tile: Int32, block_valid_rows: Int32, global_scale_f32: cutlass.Float32, reduce_slice_count: Int32, reduce_slice_idx: Int32, lock_slot: Int32, uses_m_block_8: cutlass.Constexpr[bool], ): if cutlass.const_expr(uses_m_block_8): self._fold_cta_partials_m8(acc, smem_base, tid) else: self._fold_cta_partials_large_m(acc, smem_base, tid) if reduce_slice_count > Int32(1): self._wait_for_reduction_turn( locks_i32_flat, lock_slot, reduce_slice_idx, tid ) self._combine_splitk_accumulators( acc, c_tmp_f32_flat, block_valid_rows, lock_slot, reduce_slice_idx, reduce_slice_count, tid, uses_m_block_8, ) self._publish_reduction_turn( locks_i32_flat, lock_slot, reduce_slice_idx == reduce_slice_count - Int32(1), tid, ) if reduce_slice_idx == reduce_slice_count - Int32(1): if cutlass.const_expr(uses_m_block_8): self._store_tile_m8( acc, c_bf16_flat, smem_base, tid, output_n_tile, block_valid_rows, global_scale_f32, ) else: self._store_tile_large_m( acc, c_bf16_flat, smem_base, tid, output_n_tile, block_valid_rows, global_scale_f32, ) @cute.jit def _wait_for_reduction_turn( self, locks_i32_flat: cute.Tensor, lock_slot: Int32, count: Int32, tid: Int32, ): lock_addr = get_ptr_as_int64(locks_i32_flat, lock_slot) if tid == Int32(0): state = Int32(-1) while state != count: state = ld_global_acquire_i32(lock_addr) cute.arch.sync_threads() @cute.jit def _publish_reduction_turn( self, locks_i32_flat: cute.Tensor, lock_slot: Int32, reset, tid: Int32, ): lock_addr = get_ptr_as_int64(locks_i32_flat, lock_slot) cute.arch.sync_threads() if tid == Int32(0): if reset: st_global_i32(lock_addr, Int32(0)) else: red_add_global_release_i32(lock_addr, Int32(1)) @cute.jit def _merge_splitk_vec4( self, c_tmp_f32_flat: cute.Tensor, f32_off: Int32, reduce_slice_idx: Int32, reduce_slice_count: Int32, c0: cutlass.Float32, c1: cutlass.Float32, c2: cutlass.Float32, c3: cutlass.Float32, ): if reduce_slice_idx != Int32(0): r0, r1, r2, r3 = ld_global_v4_f32(get_ptr_as_int64(c_tmp_f32_flat, f32_off)) c0 = c0 + r0 c1 = c1 + r1 c2 = c2 + r2 c3 = c3 + r3 if reduce_slice_idx != reduce_slice_count - Int32(1): st_global_v4_f32( get_ptr_as_int64(c_tmp_f32_flat, f32_off), c0, c1, c2, c3, ) return c0, c1, c2, c3 @cute.jit def _combine_splitk_accumulators( self, acc: cute.Tensor, c_tmp_f32_flat: cute.Tensor, block_valid_rows: Int32, lock_slot: Int32, reduce_slice_idx: Int32, reduce_slice_count: Int32, tid: Int32, uses_m_block_8: cutlass.Constexpr[bool], ): active_threads = Int32(32 * self.tb_n_warps) c_size_int4 = Int32((self.cta_m_blocks * 16 * self.cta_n_blocks * 16) // 4) c_cur_offset = lock_slot * c_size_int4 if cutlass.const_expr(uses_m_block_8): if tid < active_threads: for jj in cutlass.range_constexpr(4): k = jj * 2 acc[jj, 0], acc[jj, 1], acc[jj, 2], acc[jj, 3] = ( self._merge_splitk_slot( c_tmp_f32_flat, c_cur_offset, active_threads, Int32(k), tid, reduce_slice_idx, reduce_slice_count, acc[jj, 0], acc[jj, 1], acc[jj, 2], acc[jj, 3], ) ) else: lane_row = (tid & Int32(31)) // Int32(4) if tid < active_threads: for k in cutlass.range_constexpr(self.cta_m_blocks * 8): mb = k // 8 flat_j = k % 8 jj = flat_j // 2 half = flat_j % 2 row_valid = Int32(mb * 16) + lane_row < block_valid_rows if row_valid: ( acc[mb, jj, half, 0], acc[mb, jj, half, 1], acc[mb, jj, half, 2], acc[mb, jj, half, 3], ) = self._merge_splitk_slot( c_tmp_f32_flat, c_cur_offset, active_threads, Int32(k), tid, reduce_slice_idx, reduce_slice_count, acc[mb, jj, half, 0], acc[mb, jj, half, 1], acc[mb, jj, half, 2], acc[mb, jj, half, 3], ) @cute.jit def _merge_splitk_slot( self, c_tmp_f32_flat: cute.Tensor, c_cur_offset: Int32, active_threads: Int32, slot: Int32, tid: Int32, reduce_slice_idx: Int32, reduce_slice_count: Int32, c0: cutlass.Float32, c1: cutlass.Float32, c2: cutlass.Float32, c3: cutlass.Float32, ): int4_off = c_cur_offset + active_threads * slot + tid return self._merge_splitk_vec4( c_tmp_f32_flat, int4_off * Int32(4), reduce_slice_idx, reduce_slice_count, c0, c1, c2, c3, ) @cute.jit def _load_a_registers_large_m( self, smem_base: Int32, a_sh_rd: Int32, pipe: Int32, kk: Int32, m_block: Int32, ): a_addr = self._int4_addr( smem_base, Int32(self.sh_a_off) + pipe * Int32(self.a_sh_stage) + self._activation_smem_permuted_offset( Int32(2) * kk + m_block * Int32(self.a_sh_rd_delta_i) + a_sh_rd ), ) return ldmatrix_m8n8x4_b16(a_addr) @cute.jit def _load_a_registers_m8( self, smem_base: Int32, a_sh_rd: Int32, pipe: Int32, kk: Int32, ): a_addr = self._int4_addr( smem_base, Int32(self.sh_a_off) + pipe * Int32(self.a_sh_stage) + self._activation_smem_permuted_offset(Int32(2) * kk + a_sh_rd), ) return ldmatrix_m8n8x2_b16(a_addr) @cute.jit def _load_a_registers_large_m_bundle( self, regs: cute.Tensor, smem_base: Int32, a_sh_rd: Int32, pipe: Int32, kk: Int32, ): for mb in cutlass.range_constexpr(self.cta_m_blocks): a0, a1, a2, a3 = self._load_a_registers_large_m( smem_base, a_sh_rd, pipe, kk, Int32(mb), ) regs[mb, 0] = a0 regs[mb, 1] = a1 regs[mb, 2] = a2 regs[mb, 3] = a3 @cute.jit def _load_a_registers_m8_bundle( self, regs: cute.Tensor, smem_base: Int32, a_sh_rd: Int32, pipe: Int32, kk: Int32, ): a0, a1 = self._load_a_registers_m8(smem_base, a_sh_rd, pipe, kk) regs[0] = a0 regs[1] = a1 @cute.jit def _clear_a_register_bundle_large_m(self, regs: cute.Tensor): for mb in cutlass.range_constexpr(self.cta_m_blocks): for reg in cutlass.range_constexpr(4): regs[mb, reg] = Uint32(0) @cute.jit def _clear_a_register_bundle_m8(self, regs: cute.Tensor): for reg in cutlass.range_constexpr(2): regs[reg] = Uint32(0) @cute.jit def _copy_a_register_bundle_large_m(self, dst: cute.Tensor, src: cute.Tensor): for mb in cutlass.range_constexpr(self.cta_m_blocks): for reg in cutlass.range_constexpr(4): dst[mb, reg] = src[mb, reg] @cute.jit def _copy_a_register_bundle_m8(self, dst: cute.Tensor, src: cute.Tensor): for reg in cutlass.range_constexpr(2): dst[reg] = src[reg] @cute.jit def _load_a_register_bundle( self, regs: cute.Tensor, smem_base: Int32, a_sh_rd: Int32, pipe: Int32, kk: Int32, uses_m_block_8: cutlass.Constexpr[bool], ): if cutlass.const_expr(uses_m_block_8): self._load_a_registers_m8_bundle(regs, smem_base, a_sh_rd, pipe, kk) else: self._load_a_registers_large_m_bundle(regs, smem_base, a_sh_rd, pipe, kk) @cute.jit def _clear_a_register_bundle( self, regs: cute.Tensor, uses_m_block_8: cutlass.Constexpr[bool], ): if cutlass.const_expr(uses_m_block_8): self._clear_a_register_bundle_m8(regs) else: self._clear_a_register_bundle_large_m(regs) @cute.jit def _copy_a_register_bundle( self, dst: cute.Tensor, src: cute.Tensor, uses_m_block_8: cutlass.Constexpr[bool], ): if cutlass.const_expr(uses_m_block_8): self._copy_a_register_bundle_m8(dst, src) else: self._copy_a_register_bundle_large_m(dst, src) @cute.jit def _load_b_scale_registers( self, smem_base: Int32, tid: Int32, b_sh_rd: Int32, s_sh_rd: Int32, pipe: Int32, kk: Int32, ): if cutlass.const_expr(self.weight_layout == "modelopt"): q0, q1, q2, q3 = self._load_b_registers_modelopt_shared( smem_base, b_sh_rd, pipe, kk, ) elif cutlass.const_expr(self.weight_layout_nf3): # NF3-MAPPING-V1: the thread's 32-code unit index within the pipe's # B region equals the stock staged-unit index (algebraically # identical to b_sh_stride*kk + b_sh_rd): cur_group_id is the K16 row # within the tile, (tid % b_sh_stride) is the N-pair. Each unit is a # 12-byte (lo0, lo1, hi) triple; three scalar u32 loads over stride-3 # words are bank-conflict free (gcd(3, 32) == 1). q3 is unused. nf3_warp_id = tid // Int32(32) nf3_warp_row = nf3_warp_id // Int32(self.tb_n_warps) nf3_group = Int32(self.b_sh_wr_iters) * nf3_warp_row + kk nf3_unit = nf3_group * Int32(self.b_sh_stride) + ( tid % Int32(self.b_sh_stride) ) b_addr = ( smem_base + Int32(self.sh_b_off * 16) + pipe * Int32(self.b_sh_stage_bytes) + nf3_unit * Int32(12) ) q0 = ld_shared_u32(b_addr) q1 = ld_shared_u32(b_addr + Int32(4)) q2 = ld_shared_u32(b_addr + Int32(8)) q3 = Uint32(0) else: b_addr = self._int4_addr( smem_base, Int32(self.sh_b_off) + pipe * Int32(self.b_sh_stage) + Int32(self.b_sh_stride) * kk + b_sh_rd, ) q0, q1, q2, q3 = ld_shared_v4_u32(b_addr) warp_id = tid // Int32(32) warp_row = warp_id // Int32(self.tb_n_warps) cur_group_id = Int32(self.b_sh_wr_iters) * warp_row + kk if cutlass.const_expr(self.scale_k32): scale_group_id = cur_group_id // Int32(2) else: scale_group_id = cur_group_id s_addr = ( smem_base + Int32(self.sh_s_off * 16) + pipe * Int32(self.s_sh_stage * 16) + (s_sh_rd + scale_group_id * Int32(2 * self.s_sh_stride)) * Int32(8) ) s_pack0, s_pack1 = ld_shared_v2_u32(s_addr) s0, s1 = self._dequant_scale_x4_to_elem2x2(s_pack0) s2, s3 = self._dequant_scale_x4_to_elem2x2(s_pack1) return q0, q1, q2, q3, s0, s1, s2, s3 @cute.jit def _load_b_scale_register_bundle( self, regs: cute.Tensor, smem_base: Int32, tid: Int32, b_sh_rd: Int32, s_sh_rd: Int32, pipe: Int32, kk: Int32, ): q0, q1, q2, q3, s0, s1, s2, s3 = self._load_b_scale_registers( smem_base, tid, b_sh_rd, s_sh_rd, pipe, kk, ) regs[0, 0] = q0 regs[0, 1] = q1 regs[0, 2] = q2 regs[0, 3] = q3 regs[1, 0] = s0 regs[1, 1] = s1 regs[1, 2] = s2 regs[1, 3] = s3 @cute.jit def _clear_b_scale_register_bundle(self, regs: cute.Tensor): for row in cutlass.range_constexpr(2): for col in cutlass.range_constexpr(4): regs[row, col] = Uint32(0) @cute.jit def _copy_b_scale_register_bundle(self, dst: cute.Tensor, src: cute.Tensor): for row in cutlass.range_constexpr(2): for col in cutlass.range_constexpr(4): dst[row, col] = src[row, col] @cute.jit def _select_b_scale_register(self, jj: cutlass.Constexpr[int], regs: cute.Tensor): return regs[0, jj], regs[1, jj] @cute.jit def _load_next_fragment_bundle( self, b_scale_next: cute.Tensor, a_regs_next: cute.Tensor, smem_base: Int32, tid: Int32, b_sh_rd: Int32, s_sh_rd: Int32, a_sh_rd: Int32, pipe: cutlass.Constexpr[int], kk: cutlass.Constexpr[int], tile_idx: Int32, k_tiles: Int32, uses_m_block_8: cutlass.Constexpr[bool], ): self._clear_b_scale_register_bundle(b_scale_next) self._clear_a_register_bundle(a_regs_next, uses_m_block_8) if cutlass.const_expr(kk + 1 < self.b_sh_wr_iters): if tile_idx < k_tiles: self._load_b_scale_register_bundle( b_scale_next, smem_base, tid, b_sh_rd, s_sh_rd, Int32(pipe), Int32(kk + 1), ) self._load_a_register_bundle( a_regs_next, smem_base, a_sh_rd, Int32(pipe), Int32(kk + 1), uses_m_block_8, ) else: next_tile = tile_idx + Int32(1) if next_tile < k_tiles: self._load_b_scale_register_bundle( b_scale_next, smem_base, tid, b_sh_rd, s_sh_rd, Int32((pipe + 1) % _STAGES), Int32(0), ) self._load_a_register_bundle( a_regs_next, smem_base, a_sh_rd, Int32((pipe + 1) % _STAGES), Int32(0), uses_m_block_8, ) @cute.jit def _scaled_dequant_b_fragment(self, frag: cute.Tensor, q: Uint32, s: Uint32): bq1 = q bq0 = bq1 << Uint32(8) b0_0, b0_1 = self._dequant_e2m1x4_to_elem2x2(bq0) b1_0, b1_1 = self._dequant_e2m1x4_to_elem2x2(bq1) s_lane0 = bfloat2_broadcast_lane(s, Int32(0)) s_lane1 = bfloat2_broadcast_lane(s, Int32(1)) b0_0 = self._elem2_mul(b0_0, s_lane0) b0_1 = self._elem2_mul(b0_1, s_lane0) b1_0 = self._elem2_mul(b1_0, s_lane1) b1_1 = self._elem2_mul(b1_1, s_lane1) frag[0, 0] = b0_0 frag[0, 1] = b0_1 frag[1, 0] = b1_0 frag[1, 1] = b1_1 @cute.jit def _scaled_dequant_b_fragment_nf3( self, frag: cute.Tensor, lo16: Uint32, hi8: Uint32, s: Uint32 ): # NF3: 8 3-bit codes (2 lo-bits in lo16, 1 hi-bit in hi8) -> 4 bf16x2 # codebook fragments, in the SAME element order the packed path uses: # o0=(cb[c0],cb[c1]) -> frag[0,0], o1 -> frag[0,1], o2 -> frag[1,0], # o3 -> frag[1,1]. frag[0,*] takes the N-col-0 scale lane, frag[1,*] the # N-col-1 lane (identical broadcast to _scaled_dequant_b_fragment). o0, o1, o2, o3 = packed_dequant_nf3x8_to_bfloat2x4( lo16, hi8, Uint32(self._nf3_cb0), Uint32(self._nf3_cb1), Uint32(self._nf3_cb2), Uint32(self._nf3_cb3), ) s_lane0 = bfloat2_broadcast_lane(s, Int32(0)) s_lane1 = bfloat2_broadcast_lane(s, Int32(1)) frag[0, 0] = self._elem2_mul(o0, s_lane0) frag[0, 1] = self._elem2_mul(o1, s_lane0) frag[1, 0] = self._elem2_mul(o2, s_lane1) frag[1, 1] = self._elem2_mul(o3, s_lane1) @cute.jit def _mma_accumulate_m8( self, acc: cute.Tensor, jj: cutlass.Constexpr[int], a_regs: cute.Tensor, b_frag: cute.Tensor, ): d0, d1, d2, d3 = self._mma_rhs_fragments_as_mma_a_m16n8k16_f32( acc[jj, 0], acc[jj, 1], acc[jj, 2], acc[jj, 3], b_frag[0, 0], b_frag[1, 0], b_frag[0, 1], b_frag[1, 1], a_regs[0], a_regs[1], ) acc[jj, 0] = d0 acc[jj, 1] = d1 acc[jj, 2] = d2 acc[jj, 3] = d3 @cute.jit def _mma_accumulate_large_m( self, acc: cute.Tensor, a_regs: cute.Tensor, mb: cutlass.Constexpr[int], jj: cutlass.Constexpr[int], b_frag: cute.Tensor, ): d0, d1, d2, d3 = self._mma_m16n8k16_f32( acc[mb, jj, 0, 0], acc[mb, jj, 0, 1], acc[mb, jj, 0, 2], acc[mb, jj, 0, 3], a_regs[mb, 0], a_regs[mb, 1], a_regs[mb, 2], a_regs[mb, 3], b_frag[0, 0], b_frag[0, 1], ) acc[mb, jj, 0, 0] = d0 acc[mb, jj, 0, 1] = d1 acc[mb, jj, 0, 2] = d2 acc[mb, jj, 0, 3] = d3 d0, d1, d2, d3 = self._mma_m16n8k16_f32( acc[mb, jj, 1, 0], acc[mb, jj, 1, 1], acc[mb, jj, 1, 2], acc[mb, jj, 1, 3], a_regs[mb, 0], a_regs[mb, 1], a_regs[mb, 2], a_regs[mb, 3], b_frag[1, 0], b_frag[1, 1], ) acc[mb, jj, 1, 0] = d0 acc[mb, jj, 1, 1] = d1 acc[mb, jj, 1, 2] = d2 acc[mb, jj, 1, 3] = d3 @cute.jit def _source_n_from_logical(self, logical_n: Int32) -> Int32: source_n = logical_n if cutlass.const_expr(self.source_n_rotation != 0): source_n += Int32(self.source_n_rotation) if source_n >= Int32(self.size_n): source_n -= Int32(self.size_n) return source_n @cute.jit def _stage_b_tile_modelopt_native( self, b_u8_flat: cute.Tensor, smem_addr: Int32, expert_idx: Int32, output_n_tile: Int32, tile_idx: Int32, local_int4: Int32, ): chunks_per_row = Int32(self.tile_k // 32) local_n = local_int4 // chunks_per_row local_k_vec = local_int4 - local_n * chunks_per_row logical_n = output_n_tile * Int32(self.tile_n) + local_n source_n = self._source_n_from_logical(logical_n) packed_cols = Int32(self.size_k // 2) tile_byte = tile_idx * Int32(self.tile_k // 2) + local_k_vec * Int32(16) byte_offset = ( Int64(expert_idx) * Int64(self.size_n * (self.size_k // 2)) + Int64(source_n) * Int64(packed_cols) + Int64(tile_byte) ) if cutlass.const_expr(self.has_logical_tail): valid_n = logical_n < Int32(self.size_n) valid_bytes = packed_cols - tile_byte if valid_bytes > Int32(16): valid_bytes = Int32(16) if valid_n and valid_bytes >= Int32(16) and cutlass.const_expr( not self.has_scale_k_tail ): cp_async4_shared_global( smem_addr, get_ptr_as_int64(b_u8_flat, byte_offset), ) else: v0 = Uint32(0) v1 = Uint32(0) v2 = Uint32(0) v3 = Uint32(0) if valid_n and valid_bytes > Int32(0): src = get_ptr_as_int64(b_u8_flat, byte_offset) if valid_bytes >= Int32(4): v0 = ld_global_nc_u32(src) if valid_bytes >= Int32(8): v1 = ld_global_nc_u32(src + Int64(4)) if valid_bytes >= Int32(12): v2 = ld_global_nc_u32(src + Int64(8)) if valid_bytes >= Int32(16): v3 = ld_global_nc_u32(src + Int64(12)) st_shared_v4_u32(smem_addr, v0, v1, v2, v3) else: cp_async4_shared_global( smem_addr, get_ptr_as_int64(b_u8_flat, byte_offset), ) @cute.jit def _load_modelopt_shared_byte( self, smem_base: Int32, pipe: Int32, n_tile: Int32, k_tile: Int32, warp_id: Int32, tc_col: Int32, tc_row: Int32, n_delta: cutlass.Constexpr[int], k_delta: cutlass.Constexpr[int], ) -> Uint32: local_n = ( n_tile * Int32(64) + warp_id * Int32(16) + tc_col + Int32(n_delta) ) local_k = k_tile * Int32(16) + tc_row + Int32(k_delta) byte_offset = local_n * Int32(self.tile_k // 2) + local_k // Int32(2) word_byte_offset = byte_offset - (byte_offset & Int32(3)) word = ld_shared_u32( smem_base + Int32(self.sh_b_off * 16) + pipe * Int32(self.b_sh_stage * 16) + word_byte_offset ) shift = Uint32((byte_offset - word_byte_offset) * Int32(8)) return (word >> shift) & Uint32(0xFF) @cute.jit def _pack_modelopt_byte_pair( self, word: Uint32, q: Uint32, low_shift: cutlass.Constexpr[int], high_shift: cutlass.Constexpr[int], ) -> Uint32: low = q & Uint32(0xF) high = (q >> Uint32(4)) & Uint32(0xF) return word | (low << Uint32(low_shift)) | (high << Uint32(high_shift)) @cute.jit def _load_modelopt_shared_packed_word_for_lane( self, smem_base: Int32, pipe: Int32, n_tile: Int32, k_tile: Int32, warp_id: Int32, tc_col: Int32, tc_row: Int32, ) -> Uint32: q0 = self._load_modelopt_shared_byte( smem_base, pipe, n_tile, k_tile, warp_id, tc_col, tc_row, 0, 0 ) q1 = self._load_modelopt_shared_byte( smem_base, pipe, n_tile, k_tile, warp_id, tc_col, tc_row, 0, 8 ) q2 = self._load_modelopt_shared_byte( smem_base, pipe, n_tile, k_tile, warp_id, tc_col, tc_row, 8, 0 ) q3 = self._load_modelopt_shared_byte( smem_base, pipe, n_tile, k_tile, warp_id, tc_col, tc_row, 8, 8 ) word = Uint32(0) word = self._pack_modelopt_byte_pair(word, q0, 0, 16) word = self._pack_modelopt_byte_pair(word, q1, 4, 20) word = self._pack_modelopt_byte_pair(word, q2, 8, 24) word = self._pack_modelopt_byte_pair(word, q3, 12, 28) return word @cute.jit def _load_b_registers_modelopt_shared( self, smem_base: Int32, b_sh_rd: Int32, pipe: Int32, kk: Int32, ): packed_word_index = (Int32(self.b_sh_stride) * kk + b_sh_rd) * Int32(4) words_per_k_tile = Int32((self.tile_n // 64) * 128) k_tile = packed_word_index // words_per_k_tile pos_in_k_tile = packed_word_index - k_tile * words_per_k_tile n_tile = pos_in_k_tile // Int32(128) pos = pos_in_k_tile - n_tile * Int32(128) th_id = pos // Int32(4) tc_col = th_id // Int32(4) tc_row = (th_id - tc_col * Int32(4)) * Int32(2) q0 = self._load_modelopt_shared_packed_word_for_lane( smem_base, pipe, n_tile, k_tile, Int32(0), tc_col, tc_row ) q1 = self._load_modelopt_shared_packed_word_for_lane( smem_base, pipe, n_tile, k_tile, Int32(1), tc_col, tc_row ) q2 = self._load_modelopt_shared_packed_word_for_lane( smem_base, pipe, n_tile, k_tile, Int32(2), tc_col, tc_row ) q3 = self._load_modelopt_shared_packed_word_for_lane( smem_base, pipe, n_tile, k_tile, Int32(3), tc_col, tc_row ) return q0, q1, q2, q3 @cute.jit def _stage_k_tile_async( self, a_bf16_flat: cute.Tensor, b_i32_flat: cute.Tensor, scales_i32_flat: cute.Tensor, smem_base: Int32, tid: Int32, pipe: Int32, tile_idx: Int32, block_valid_rows: Int32, a_gl_stride: Int32, b_gl_stride: Int32, s_gl_stride: Int32, scales_expert_off: Int32, b_gl_rd_base: Int32, a_gl_rd_row: Int32, a_gl_rd_col0: Int32, a_sh_wr: Int32, a_rows_per_iter: Int32, output_n_tile: Int32, expert_idx: Int32, ): for i in cutlass.range_constexpr(self.a_sh_wr_iters): row = a_rows_per_iter * Int32(i) + a_gl_rd_row route_index = Int32(0) if row < Int32(self.moe_block_size): route_index = ld_shared_i32_relaxed( smem_base + Int32(self.sh_rd_route_off * 16) + row * Int32(4) ) a_int4 = ( route_index * a_gl_stride + tile_idx * Int32(self.a_gl_rd_delta_o) + a_gl_rd_col0 ) a_dst = self._int4_addr( smem_base, Int32(self.sh_a_off) + pipe * Int32(self.a_sh_stage) + self._activation_smem_permuted_offset( Int32(i * self.a_sh_wr_delta) + a_sh_wr ), ) if cutlass.const_expr(self.has_k_tile_tail): a_k_int4 = tile_idx * Int32(self.a_gl_rd_delta_o) + a_gl_rd_col0 if row < block_valid_rows and a_k_int4 < a_gl_stride: cp_async4_shared_global( a_dst, get_ptr_as_int64(a_bf16_flat, a_int4 * Int32(8)), ) else: st_shared_v4_u32( a_dst, Uint32(0), Uint32(0), Uint32(0), Uint32(0) ) else: cp_async4_shared_global_pred( a_dst, get_ptr_as_int64(a_bf16_flat, a_int4 * Int32(8)), (row < block_valid_rows).to(Int32), ) if cutlass.const_expr(self.weight_layout_nf3): # NF3 flat-span staging: the packer lays the per-(expert, # output_n_tile, k-stage) B block out as ONE contiguous 12-byte-triple # span (n_tile-major, then K16-row, then N-pair), so we copy it # verbatim as 16-byte cp.async chunks. SMEM byte X of the pipe's B # region == global byte X of the span, so the register read (unit*12) # indexes it directly. b_gl_rd_base is unused on this path. nf3_units_per_expert = (self.size_n * self.size_k) // 32 nf3_ntile_stride = (self.size_k // 16) * (self.tile_n // 2) nf3_span_base_unit = ( Int32(nf3_units_per_expert) * expert_idx + Int32(nf3_ntile_stride) * output_n_tile + tile_idx * Int32(self.cta_k_blocks * self.b_sh_stride) ) for i in cutlass.range_constexpr(self.b_sh_wr_iters_nf3): nf3_chunk = Int32(i * self.cta_threads) + tid b_dst = ( smem_base + Int32(self.sh_b_off * 16) + pipe * Int32(self.b_sh_stage_bytes) + nf3_chunk * Int32(16) ) # int32-element index of the chunk's first word: # (span_base_unit*12 + chunk*16) / 4 = span_base_unit*3 + chunk*4. b_src_i32 = nf3_span_base_unit * Int32(3) + nf3_chunk * Int32(4) cp_async4_shared_global_pred( b_dst, get_ptr_as_int64(b_i32_flat, b_src_i32), (nf3_chunk < Int32(self.b_sh_chunks)).to(Int32), ) for i in cutlass.range_constexpr( 0 if self.weight_layout_nf3 else self.b_sh_wr_iters ): b_src_int4 = ( b_gl_rd_base + tile_idx * Int32(self.cta_k_blocks) * b_gl_stride + Int32(i * (self.cta_threads // self.b_sh_stride)) * b_gl_stride ) b_dst = self._int4_addr( smem_base, Int32(self.sh_b_off) + pipe * Int32(self.b_sh_stage) + Int32(i * self.cta_threads) + tid, ) if cutlass.const_expr(self.weight_layout == "packed"): cp_async4_shared_global( b_dst, get_ptr_as_int64(b_i32_flat, b_src_int4 * Int32(4)), ) else: self._stage_b_tile_modelopt_native( b_i32_flat, b_dst, expert_idx, output_n_tile, tile_idx, Int32(i * self.cta_threads) + tid, ) if tid < Int32(self.s_sh_stage): s_k_group = tile_idx * Int32(self.s_tb_groups) + tid // Int32( self.s_sh_stride ) s_n_group = Int32(self.s_sh_stride) * output_n_tile + ( tid % Int32(self.s_sh_stride) ) s_src_int4 = scales_expert_off + s_gl_stride * s_k_group + s_n_group s_dst = self._int4_addr( smem_base, Int32(self.sh_s_off) + pipe * Int32(self.s_sh_stage) + tid, ) if cutlass.const_expr(self.has_logical_tail): if s_k_group < Int32(self.scale_k_groups) and s_n_group < Int32( self.scale_n_groups ): cp_async4_shared_global( s_dst, get_ptr_as_int64(scales_i32_flat, s_src_int4 * Int32(4)), ) else: st_shared_v4_u32( s_dst, Uint32(0), Uint32(0), Uint32(0), Uint32(0) ) else: cp_async4_shared_global( s_dst, get_ptr_as_int64(scales_i32_flat, s_src_int4 * Int32(4)), ) cute.arch.cp_async_commit_group() @cute.jit def _prefetch_pipeline_step( self, a_bf16_flat: cute.Tensor, b_i32_flat: cute.Tensor, scales_i32_flat: cute.Tensor, smem_base: Int32, tid: Int32, pipe: cutlass.Constexpr[int], kk: cutlass.Constexpr[int], tile_idx: Int32, k_tiles: Int32, reduce_k_tile: Int32, block_valid_rows: Int32, a_gl_stride: Int32, b_gl_stride: Int32, s_gl_stride: Int32, scales_expert_off: Int32, b_gl_rd_base: Int32, a_gl_rd_row: Int32, a_gl_rd_col0: Int32, a_sh_wr: Int32, a_rows_per_iter: Int32, output_n_tile: Int32, expert_idx: Int32, ): if cutlass.const_expr(kk == self.b_sh_wr_iters - 2): self._prefetch_lookahead_tile( a_bf16_flat, b_i32_flat, scales_i32_flat, smem_base, tid, pipe, tile_idx, k_tiles, reduce_k_tile, block_valid_rows, a_gl_stride, b_gl_stride, s_gl_stride, scales_expert_off, b_gl_rd_base, a_gl_rd_row, a_gl_rd_col0, a_sh_wr, a_rows_per_iter, output_n_tile, expert_idx, ) @cute.jit def _prefetch_initial_tiles( self, a_bf16_flat: cute.Tensor, b_i32_flat: cute.Tensor, scales_i32_flat: cute.Tensor, smem_base: Int32, tid: Int32, k_tiles: Int32, reduce_k_tile: Int32, block_valid_rows: Int32, a_gl_stride: Int32, b_gl_stride: Int32, s_gl_stride: Int32, scales_expert_off: Int32, b_gl_rd_base: Int32, a_gl_rd_row: Int32, a_gl_rd_col0: Int32, a_sh_wr: Int32, a_rows_per_iter: Int32, output_n_tile: Int32, expert_idx: Int32, ): for pipe in cutlass.range_constexpr(_STAGES - 1): if Int32(pipe) < k_tiles: self._stage_k_tile_async( a_bf16_flat, b_i32_flat, scales_i32_flat, smem_base, tid, Int32(pipe), reduce_k_tile + Int32(pipe), block_valid_rows, a_gl_stride, b_gl_stride, s_gl_stride, scales_expert_off, b_gl_rd_base, a_gl_rd_row, a_gl_rd_col0, a_sh_wr, a_rows_per_iter, output_n_tile, expert_idx, ) else: cute.arch.cp_async_commit_group() cute.arch.cp_async_wait_group(_STAGES - 2) cute.arch.sync_threads() @cute.jit def _prefetch_lookahead_tile( self, a_bf16_flat: cute.Tensor, b_i32_flat: cute.Tensor, scales_i32_flat: cute.Tensor, smem_base: Int32, tid: Int32, pipe: cutlass.Constexpr[int], tile_idx: Int32, k_tiles: Int32, reduce_k_tile: Int32, block_valid_rows: Int32, a_gl_stride: Int32, b_gl_stride: Int32, s_gl_stride: Int32, scales_expert_off: Int32, b_gl_rd_base: Int32, a_gl_rd_row: Int32, a_gl_rd_col0: Int32, a_sh_wr: Int32, a_rows_per_iter: Int32, output_n_tile: Int32, expert_idx: Int32, ): fetch_tile = tile_idx + Int32(_STAGES - 1) if fetch_tile < k_tiles: self._stage_k_tile_async( a_bf16_flat, b_i32_flat, scales_i32_flat, smem_base, tid, Int32((pipe + _STAGES - 1) % _STAGES), reduce_k_tile + fetch_tile, block_valid_rows, a_gl_stride, b_gl_stride, s_gl_stride, scales_expert_off, b_gl_rd_base, a_gl_rd_row, a_gl_rd_col0, a_sh_wr, a_rows_per_iter, output_n_tile, expert_idx, ) else: cute.arch.cp_async_commit_group() cute.arch.cp_async_wait_group(_STAGES - 2) cute.arch.sync_threads() @cute.jit def _reduction_offsets(self, tid: Int32): red_idx = tid // Int32(self.b_sh_stride_threads) red_sh_stride = Int32(self.b_sh_stride_threads * 4 * 2) red_sh_delta = Int32(self.b_sh_stride_threads) red_sh_rd = red_sh_stride * (tid // Int32(self.b_sh_stride_threads)) + ( tid % Int32(self.b_sh_stride_threads) ) return red_idx, red_sh_stride, red_sh_delta, red_sh_rd @cute.jit def _fold_cta_partials_m8(self, acc: cute.Tensor, smem_base: Int32, tid: Int32): red_off = self.cta_threads // self.b_sh_stride_threads // 2 if cutlass.const_expr(red_off >= 1): red_idx, red_sh_stride, red_sh_delta, red_sh_rd = self._reduction_offsets( tid ) if cutlass.const_expr(red_off == 2): if Int32(2) <= red_idx and red_idx < Int32(4): for jj in cutlass.range_constexpr(4): red_sh_wr = red_sh_delta * Int32(jj * 2) + ( red_sh_rd - red_sh_stride * Int32(2) ) st_shared_v4_f32( self._int4_addr( smem_base, Int32(self.sh_red_off) + red_sh_wr ), acc[jj, 0], acc[jj, 1], acc[jj, 2], acc[jj, 3], ) cute.arch.sync_threads() if Int32(1) <= red_idx and red_idx < Int32(2): for jj in cutlass.range_constexpr(4): red_sh_wr = red_sh_delta * Int32(jj * 2) + ( red_sh_rd - red_sh_stride ) if cutlass.const_expr(red_off > 1): rd_addr = self._int4_addr( smem_base, Int32(self.sh_red_off) + red_sh_delta * Int32(jj * 2) + red_sh_rd, ) wr_addr = self._int4_addr( smem_base, Int32(self.sh_red_off) + red_sh_wr ) r0, r1, r2, r3 = ld_shared_v4_f32(rd_addr) w0, w1, w2, w3 = ld_shared_v4_f32(wr_addr) acc[jj, 0] = acc[jj, 0] + r0 + w0 acc[jj, 1] = acc[jj, 1] + r1 + w1 acc[jj, 2] = acc[jj, 2] + r2 + w2 acc[jj, 3] = acc[jj, 3] + r3 + w3 st_shared_v4_f32( self._int4_addr(smem_base, Int32(self.sh_red_off) + red_sh_wr), acc[jj, 0], acc[jj, 1], acc[jj, 2], acc[jj, 3], ) cute.arch.sync_threads() if red_idx == Int32(0): for jj in cutlass.range_constexpr(4): rd_addr = self._int4_addr( smem_base, Int32(self.sh_red_off) + red_sh_delta * Int32(jj * 2) + red_sh_rd, ) r0, r1, r2, r3 = ld_shared_v4_f32(rd_addr) acc[jj, 0] = acc[jj, 0] + r0 acc[jj, 1] = acc[jj, 1] + r1 acc[jj, 2] = acc[jj, 2] + r2 acc[jj, 3] = acc[jj, 3] + r3 cute.arch.sync_threads() @cute.jit def _output_store_cursor(self, tid: Int32, output_n_tile: Int32): c_gl_stride = Int32(self.size_n // 8) c_sh_stride = Int32(2 * self.cta_n_blocks + 1) c_gl_wr_delta = c_gl_stride * Int32(self.cta_threads // (2 * self.cta_n_blocks)) c_sh_rd_delta = c_sh_stride * Int32( self.cta_threads // (2 * self.cta_n_blocks) ) c_gl_wr = ( c_gl_stride * (tid // Int32(2 * self.cta_n_blocks)) + (tid % Int32(2 * self.cta_n_blocks)) + Int32(2 * self.cta_n_blocks) * output_n_tile ) c_sh_rd = c_sh_stride * (tid // Int32(2 * self.cta_n_blocks)) + ( tid % Int32(2 * self.cta_n_blocks) ) return c_gl_stride, c_sh_stride, c_gl_wr_delta, c_sh_rd_delta, c_gl_wr, c_sh_rd @cute.jit def _output_store_cursor_tail(self, tid: Int32, output_n_tile: Int32): c_gl_stride = Int32(self.size_n // 8) c_gl_stride_covered = Int32(self.covered_size_n // 8) c_sh_stride = Int32(2 * self.cta_n_blocks + 1) c_gl_wr_delta = c_gl_stride_covered * Int32( self.cta_threads // (2 * self.cta_n_blocks) ) c_sh_rd_delta = c_sh_stride * Int32(self.cta_threads // (2 * self.cta_n_blocks)) c_gl_wr = ( c_gl_stride_covered * (tid // Int32(2 * self.cta_n_blocks)) + (tid % Int32(2 * self.cta_n_blocks)) + Int32(2 * self.cta_n_blocks) * output_n_tile ) c_sh_rd = c_sh_stride * (tid // Int32(2 * self.cta_n_blocks)) + ( tid % Int32(2 * self.cta_n_blocks) ) return ( c_gl_stride, c_gl_stride_covered, c_sh_stride, c_gl_wr_delta, c_sh_rd_delta, c_gl_wr, c_sh_rd, ) @cute.jit def _drain_output_smem( self, c_bf16_flat: cute.Tensor, smem_base: Int32, c_gl_stride: Int32, c_gl_wr: Int32, c_gl_wr_delta: Int32, c_sh_rd: Int32, c_sh_rd_delta: Int32, block_valid_rows: Int32, store_iters: cutlass.Constexpr[int], ): for _ in cutlass.range_constexpr(store_iters): row = c_gl_wr // c_gl_stride if row < block_valid_rows: route_index = ld_shared_i32_relaxed( smem_base + Int32(self.sh_route_off * 16) + row * Int32(4) ) true_idx = route_index * c_gl_stride + (c_gl_wr % c_gl_stride) q0, q1, q2, q3 = ld_shared_v4_u32( self._int4_addr(smem_base, Int32(self.sh_red_off) + c_sh_rd) ) if cutlass.const_expr(self.mul_topk_weights): scale_bf2 = ld_shared_u32( smem_base + Int32(self.sh_topk_off * 16) + row * Int32(4) ) q0 = self._elem2_mul(q0, scale_bf2) q1 = self._elem2_mul(q1, scale_bf2) q2 = self._elem2_mul(q2, scale_bf2) q3 = self._elem2_mul(q3, scale_bf2) if cutlass.const_expr(self.epilogue_relu2): q0 = self._relu2_elem2(q0) q1 = self._relu2_elem2(q1) q2 = self._relu2_elem2(q2) q3 = self._relu2_elem2(q3) if cutlass.const_expr(self.fused_topk_sum): # Fold the per-route partials into the per-token output in # the epilogue (drops the separate top-k-sum launch). The # output must be zeroed before launch; each route slot maps # to token = route_index // top_k. bf16x2 add lands two # consecutive hidden lanes per word. token_idx = route_index // Int32(self.fused_sum_topk) out_idx = token_idx * c_gl_stride + (c_gl_wr % c_gl_stride) out_addr = get_ptr_as_int64(c_bf16_flat, out_idx * Int32(8)) red_add_global_bf16x2(out_addr, q0) red_add_global_bf16x2(out_addr + Int64(4), q1) red_add_global_bf16x2(out_addr + Int64(8), q2) red_add_global_bf16x2(out_addr + Int64(12), q3) else: st_global_v4_u32( get_ptr_as_int64(c_bf16_flat, true_idx * Int32(8)), q0, q1, q2, q3, ) c_gl_wr += c_gl_wr_delta c_sh_rd += c_sh_rd_delta cute.arch.sync_threads() @cute.jit def _drain_output_smem_tail( self, c_bf16_flat: cute.Tensor, smem_base: Int32, c_gl_stride: Int32, c_gl_stride_covered: Int32, c_gl_wr: Int32, c_gl_wr_delta: Int32, c_sh_rd: Int32, c_sh_rd_delta: Int32, block_valid_rows: Int32, store_iters: cutlass.Constexpr[int], ): for _ in cutlass.range_constexpr(store_iters): row = c_gl_wr // c_gl_stride_covered col_word = c_gl_wr - row * c_gl_stride_covered if row < block_valid_rows and col_word < c_gl_stride: route_index = ld_shared_i32_relaxed( smem_base + Int32(self.sh_route_off * 16) + row * Int32(4) ) true_idx = route_index * c_gl_stride + col_word q0, q1, q2, q3 = ld_shared_v4_u32( self._int4_addr(smem_base, Int32(self.sh_red_off) + c_sh_rd) ) if cutlass.const_expr(self.mul_topk_weights): scale_bf2 = ld_shared_u32( smem_base + Int32(self.sh_topk_off * 16) + row * Int32(4) ) q0 = self._elem2_mul(q0, scale_bf2) q1 = self._elem2_mul(q1, scale_bf2) q2 = self._elem2_mul(q2, scale_bf2) q3 = self._elem2_mul(q3, scale_bf2) if cutlass.const_expr(self.epilogue_relu2): q0 = self._relu2_elem2(q0) q1 = self._relu2_elem2(q1) q2 = self._relu2_elem2(q2) q3 = self._relu2_elem2(q3) if cutlass.const_expr(self.fused_topk_sum): token_idx = route_index // Int32(self.fused_sum_topk) out_idx = token_idx * c_gl_stride + col_word out_addr = get_ptr_as_int64(c_bf16_flat, out_idx * Int32(8)) red_add_global_bf16x2(out_addr, q0) red_add_global_bf16x2(out_addr + Int64(4), q1) red_add_global_bf16x2(out_addr + Int64(8), q2) red_add_global_bf16x2(out_addr + Int64(12), q3) else: st_global_v4_u32( get_ptr_as_int64(c_bf16_flat, true_idx * Int32(8)), q0, q1, q2, q3, ) c_gl_wr += c_gl_wr_delta c_sh_rd += c_sh_rd_delta cute.arch.sync_threads() @cute.jit def _store_tile_m8( self, acc: cute.Tensor, c_bf16_flat: cute.Tensor, smem_base: Int32, tid: Int32, output_n_tile: Int32, block_valid_rows: Int32, global_scale_f32: cutlass.Float32, ): if cutlass.const_expr(self.has_n_tile_tail): ( c_gl_stride, c_gl_stride_covered, c_sh_stride, c_gl_wr_delta, c_sh_rd_delta, c_gl_wr, c_sh_rd, ) = self._output_store_cursor_tail(tid, output_n_tile) else: ( c_gl_stride, c_sh_stride, c_gl_wr_delta, c_sh_rd_delta, c_gl_wr, c_sh_rd, ) = self._output_store_cursor(tid, output_n_tile) c_gl_stride_covered = c_gl_stride c_sh_wr = ( Int32(8) * c_sh_stride * (((tid & Int32(31)) % Int32(4)) * Int32(2)) + (tid & Int32(31)) // Int32(4) + Int32(64) * (tid // Int32(32)) ) if tid // Int32(32) < Int32(self.tb_n_warps): write_scale = cutlass.Float32(1.0) if cutlass.const_expr(not self.mul_topk_weights): write_scale = global_scale_f32 for jj in cutlass.range_constexpr(4): wr = c_sh_wr + Int32(16 * jj) self._st_shared_elem_from_f32( smem_base + Int32(self.sh_red_off * 16) + (wr * Int32(2)), acc[jj, 0] * write_scale, ) self._st_shared_elem_from_f32( smem_base + Int32(self.sh_red_off * 16) + ((wr + Int32(8) * c_sh_stride) * Int32(2)), acc[jj, 1] * write_scale, ) self._st_shared_elem_from_f32( smem_base + Int32(self.sh_red_off * 16) + ((wr + Int32(8)) * Int32(2)), acc[jj, 2] * write_scale, ) self._st_shared_elem_from_f32( smem_base + Int32(self.sh_red_off * 16) + ((wr + Int32(8) + Int32(8) * c_sh_stride) * Int32(2)), acc[jj, 3] * write_scale, ) cute.arch.sync_threads() store_iters = _covering_count( 16, self.cta_threads // (2 * self.cta_n_blocks) ) if cutlass.const_expr(self.has_n_tile_tail): self._drain_output_smem_tail( c_bf16_flat, smem_base, c_gl_stride, c_gl_stride_covered, c_gl_wr, c_gl_wr_delta, c_sh_rd, c_sh_rd_delta, block_valid_rows, store_iters, ) else: self._drain_output_smem( c_bf16_flat, smem_base, c_gl_stride, c_gl_wr, c_gl_wr_delta, c_sh_rd, c_sh_rd_delta, block_valid_rows, store_iters, ) @cute.jit def _fold_cta_partials_large_m( self, acc: cute.Tensor, smem_base: Int32, tid: Int32 ): red_off = self.cta_threads // self.b_sh_stride_threads // 2 if cutlass.const_expr(red_off >= 1): red_idx, red_sh_stride, red_sh_delta, red_sh_rd = self._reduction_offsets( tid ) for mb in cutlass.range_constexpr(self.cta_m_blocks): if cutlass.const_expr(red_off == 2): if Int32(2) <= red_idx and red_idx < Int32(4): for flat_j in cutlass.range_constexpr(8): jj = flat_j // 2 half = flat_j % 2 red_sh_wr = red_sh_delta * Int32(flat_j) + ( red_sh_rd - red_sh_stride * Int32(2) ) st_shared_v4_f32( self._int4_addr( smem_base, Int32(self.sh_red_off) + red_sh_wr ), acc[mb, jj, half, 0], acc[mb, jj, half, 1], acc[mb, jj, half, 2], acc[mb, jj, half, 3], ) cute.arch.sync_threads() if Int32(1) <= red_idx and red_idx < Int32(2): for flat_j in cutlass.range_constexpr(8): jj = flat_j // 2 half = flat_j % 2 red_sh_wr = red_sh_delta * Int32(flat_j) + ( red_sh_rd - red_sh_stride ) if cutlass.const_expr(red_off > 1): rd_addr = self._int4_addr( smem_base, Int32(self.sh_red_off) + red_sh_delta * Int32(flat_j) + red_sh_rd, ) wr_addr = self._int4_addr( smem_base, Int32(self.sh_red_off) + red_sh_wr, ) r0, r1, r2, r3 = ld_shared_v4_f32(rd_addr) w0, w1, w2, w3 = ld_shared_v4_f32(wr_addr) acc[mb, jj, half, 0] = acc[mb, jj, half, 0] + r0 + w0 acc[mb, jj, half, 1] = acc[mb, jj, half, 1] + r1 + w1 acc[mb, jj, half, 2] = acc[mb, jj, half, 2] + r2 + w2 acc[mb, jj, half, 3] = acc[mb, jj, half, 3] + r3 + w3 st_shared_v4_f32( self._int4_addr( smem_base, Int32(self.sh_red_off) + red_sh_wr ), acc[mb, jj, half, 0], acc[mb, jj, half, 1], acc[mb, jj, half, 2], acc[mb, jj, half, 3], ) cute.arch.sync_threads() if red_idx == Int32(0): for flat_j in cutlass.range_constexpr(8): jj = flat_j // 2 half = flat_j % 2 rd_addr = self._int4_addr( smem_base, Int32(self.sh_red_off) + red_sh_delta * Int32(flat_j) + red_sh_rd, ) r0, r1, r2, r3 = ld_shared_v4_f32(rd_addr) acc[mb, jj, half, 0] = acc[mb, jj, half, 0] + r0 acc[mb, jj, half, 1] = acc[mb, jj, half, 1] + r1 acc[mb, jj, half, 2] = acc[mb, jj, half, 2] + r2 acc[mb, jj, half, 3] = acc[mb, jj, half, 3] + r3 cute.arch.sync_threads() @cute.jit def _write_bf16x2_shared( self, smem_base: Int32, half2_idx: Int32, c0: cutlass.Float32, c1: cutlass.Float32, write_scale: cutlass.Float32, ): packed = self._pack_f32x2_to_elem2(c0 * write_scale, c1 * write_scale) st_shared_u32( smem_base + Int32(self.sh_red_off * 16) + half2_idx * Int32(4), packed, ) @cute.jit def _store_tile_large_m( self, acc: cute.Tensor, c_bf16_flat: cute.Tensor, smem_base: Int32, tid: Int32, output_n_tile: Int32, block_valid_rows: Int32, global_scale_f32: cutlass.Float32, ): if cutlass.const_expr(self.has_n_tile_tail): ( c_gl_stride, c_gl_stride_covered, c_sh_stride, c_gl_wr_delta, c_sh_rd_delta, c_gl_wr, c_sh_rd, ) = self._output_store_cursor_tail(tid, output_n_tile) else: ( c_gl_stride, c_sh_stride, c_gl_wr_delta, c_sh_rd_delta, c_gl_wr, c_sh_rd, ) = self._output_store_cursor(tid, output_n_tile) c_gl_stride_covered = c_gl_stride c_sh_wr = ( Int32(4) * c_sh_stride * ((tid & Int32(31)) // Int32(4)) + (tid & Int32(31)) % Int32(4) + Int32(32) * (tid // Int32(32)) ) if tid // Int32(32) < Int32(self.tb_n_warps): write_scale = cutlass.Float32(1.0) if cutlass.const_expr(not self.mul_topk_weights): write_scale = global_scale_f32 for mb in cutlass.range_constexpr(self.cta_m_blocks): for jj in cutlass.range_constexpr(4): wr = c_sh_wr + Int32(8 * jj) self._write_bf16x2_shared( smem_base, wr, acc[mb, jj, 0, 0], acc[mb, jj, 0, 1], write_scale, ) self._write_bf16x2_shared( smem_base, wr + (Int32(4) * c_sh_stride) * Int32(8) + Int32(0), acc[mb, jj, 0, 2], acc[mb, jj, 0, 3], write_scale, ) self._write_bf16x2_shared( smem_base, wr + Int32(4), acc[mb, jj, 1, 0], acc[mb, jj, 1, 1], write_scale, ) self._write_bf16x2_shared( smem_base, wr + (Int32(4) * c_sh_stride) * Int32(8) + Int32(4), acc[mb, jj, 1, 2], acc[mb, jj, 1, 3], write_scale, ) c_sh_wr += Int32(16 * (4 * (2 * self.cta_n_blocks + 1))) cute.arch.sync_threads() store_iters = _covering_count( 16 * self.cta_m_blocks, self.cta_threads // (2 * self.cta_n_blocks), ) if cutlass.const_expr(self.has_n_tile_tail): self._drain_output_smem_tail( c_bf16_flat, smem_base, c_gl_stride, c_gl_stride_covered, c_gl_wr, c_gl_wr_delta, c_sh_rd, c_sh_rd_delta, block_valid_rows, store_iters, ) else: self._drain_output_smem( c_bf16_flat, smem_base, c_gl_stride, c_gl_wr, c_gl_wr_delta, c_sh_rd, c_sh_rd_delta, block_valid_rows, store_iters, ) class W4A16FusedMoeKernel: def __init__( self, *, size_m: int, hidden_size: int, intermediate_size: int, num_experts: int, top_k: int, activation: str, apply_router_weight_on_input: bool, zero_fc2_output: bool, fc1_tile_n: int, fc1_tile_k: int, fc2_tile_n: int, fc2_tile_k: int, moe_block_size: int, max_m_blocks: int, element_dtype: str = "bf16", fast_math: bool = True, swiglu_limit: float | None = None, swiglu_alpha: float | None = None, swiglu_beta: float | None = None, weight_layout: str = "packed", scale_format: str = "e4m3_k16", w13_layout: str = "w13", direct_topk_routes: bool = False, tc_decode_fused_sum: bool = False, tc_zero_output: bool = True, collect_activation_amax: bool = False, ): activation = normalize_moe_activation(activation) is_gated = validate_activation(activation) swiglu_limit, swiglu_alpha, swiglu_beta = _normalize_activation_swiglu_params( activation, swiglu_limit, swiglu_alpha, swiglu_beta, ) if weight_layout not in _WEIGHT_LAYOUTS: raise ValueError(f"unsupported W4A16 weight_layout {weight_layout!r}") scale_format = _normalize_scale_format(scale_format) if weight_layout == "modelopt": if w13_layout not in _MODEL_OPT_W13_LAYOUTS: raise ValueError(f"unsupported W4A16 w13_layout {w13_layout!r}") else: w13_layout = "packed" self.tc_decode_fused_sum = bool(tc_decode_fused_sum) # When two TC-decode launches share one pre-zeroed output (the NF3 hybrid # runs an NVFP4 launch then an NF3 launch into the same tensor), only the # first must zero it. Default True preserves single-launch behavior. self.tc_zero_output = bool(tc_zero_output) self.collect_activation_amax = bool(collect_activation_amax) if self.collect_activation_amax and bool(direct_topk_routes): raise ValueError("activation amax collection requires route-packed W4A16") if self.collect_activation_amax and self.tc_decode_fused_sum: raise ValueError("activation amax collection is incompatible with TC-decode") if self.tc_decode_fused_sum and not bool(direct_topk_routes): raise ValueError("tc_decode_fused_sum requires direct_topk_routes") if self.tc_decode_fused_sum and element_dtype != "bf16": raise ValueError("tc_decode_fused_sum currently requires bf16 activations") fc1_cols = int(intermediate_size) * (2 if is_gated else 1) routed_rows = int(size_m) * int(top_k) self.size_m = int(size_m) self.hidden_size = int(hidden_size) self.intermediate_size = int(intermediate_size) self.fc1_cols = int(fc1_cols) self.num_experts = int(num_experts) self.top_k = int(top_k) self.moe_block_size = int(moe_block_size) self.activation = activation self.activation_is_gated = is_gated self.activation_is_swigluoai = activation == SWIGLUOAI_UNINTERLEAVE self.has_swiglu_limit = swiglu_limit is not None self.swiglu_limit = 0.0 if swiglu_limit is None else swiglu_limit self.swiglu_alpha = float(swiglu_alpha) self.swiglu_beta = float(swiglu_beta) self.weight_layout = weight_layout self.scale_format = scale_format self.w13_layout = w13_layout self.apply_router_weight_on_input = bool(apply_router_weight_on_input) self.zero_fc2_output = bool(zero_fc2_output) self.element_dtype = element_dtype self.is_fp16 = element_dtype == "fp16" self.fast_math = bool(fast_math) self.direct_topk_routes = bool(direct_topk_routes) fc1_source_n_rotation = ( int(intermediate_size) if ( weight_layout == "modelopt" and w13_layout == "w13" and is_gated ) else 0 ) self.fc1 = W4A16GemmKernel( size_m=size_m, size_n=fc1_cols, size_k=hidden_size, num_experts=num_experts, top_k=top_k, mul_topk_weights=bool(apply_router_weight_on_input), tile_n=fc1_tile_n, tile_k=fc1_tile_k, moe_block_size=moe_block_size, max_m_blocks=max_m_blocks, element_dtype=element_dtype, epilogue_activation=None if is_gated else "relu2", weight_layout=weight_layout, scale_format=scale_format, w13_layout=w13_layout, source_n_rotation=fc1_source_n_rotation, single_token_route_fast_path=size_m == 1 and not self.direct_topk_routes, direct_topk_routes=self.direct_topk_routes, ) self.fc2 = W4A16GemmKernel( size_m=routed_rows, size_n=hidden_size, size_k=intermediate_size, num_experts=num_experts, top_k=1, mul_topk_weights=not bool(apply_router_weight_on_input), tile_n=fc2_tile_n, tile_k=fc2_tile_k, moe_block_size=moe_block_size, max_m_blocks=max_m_blocks, element_dtype=element_dtype, weight_layout=weight_layout, scale_format=scale_format, w13_layout=w13_layout, single_token_route_fast_path=size_m == 1 and not self.direct_topk_routes, direct_topk_routes=self.direct_topk_routes, fused_topk_sum=self.tc_decode_fused_sum, fused_sum_topk=int(top_k), ) self.cta_threads = max(self.fc1.cta_threads, self.fc2.cta_threads) if self.fc1.cta_threads != self.fc2.cta_threads: raise ValueError( "fused W4A16 kernel expects matching FC1/FC2 thread counts" ) self.sms = self.fc1.sms self.blocks_per_sm = min(self.fc1.blocks_per_sm, self.fc2.blocks_per_sm) self.shared_words = max(self.fc1.shared_words, self.fc2.shared_words) self.barrier_count_off = self.sms * 4 self.barrier_sense_off = self.sms * 4 + 1 @property def __cache_key__(self) -> tuple[object, ...]: return ( self.hidden_size, self.intermediate_size, self.fc1_cols, self.num_experts, self.top_k, self.activation, self.activation_is_gated, self.activation_is_swigluoai, self.has_swiglu_limit, self.swiglu_limit, self.swiglu_alpha, self.swiglu_beta, self.weight_layout, self.scale_format, self.apply_router_weight_on_input, self.zero_fc2_output, self.element_dtype, self.fast_math, self.direct_topk_routes, self.tc_zero_output, self.collect_activation_amax, self.fc1.__cache_key__, self.fc2.__cache_key__, self.cta_threads, self.sms, self.shared_words, self.blocks_per_sm, ) @cute.jit def _cast_elem(self, x: cutlass.Float32): if cutlass.const_expr(self.is_fp16): return cutlass.Float16(x) return cutlass.BFloat16(x) @cute.jit def _clamp_swiglu_inputs( self, gate: cutlass.Float32, up: cutlass.Float32, ): if cutlass.const_expr(self.has_swiglu_limit): limit = cutlass.Float32(self.swiglu_limit) neg_limit = cutlass.Float32(-self.swiglu_limit) if gate > limit: gate = limit if up > limit: up = limit if up < neg_limit: up = neg_limit return gate, up @cute.jit def __call__( self, a_bf16_ptr: cute.Pointer, w13_i32_flat: cute.Tensor, w2_i32_flat: cute.Tensor, fc1_bf16_flat: cute.Tensor, activated_bf16_flat: cute.Tensor, fc2_bf16_flat: cute.Tensor, w13_scales_i32_flat: cute.Tensor, w2_scales_i32_flat: cute.Tensor, w13_global_scale: cute.Tensor, w2_global_scale: cute.Tensor, packed_route_indices: cute.Tensor, block_expert_ids: cute.Tensor, packed_route_count: cute.Tensor, activation_amax_flat: cute.Tensor, layer_idx: cutlass.Int32, topk_weights_ptr: cute.Pointer, fc1_c_tmp_f32_flat: cute.Tensor, fc2_c_tmp_f32_flat: cute.Tensor, locks_i32_flat: cute.Tensor, active_m: cutlass.Int32, grid_x: cutlass.Int32, stream: cuda.CUstream, ): a_bf16_flat = cute.make_tensor( a_bf16_ptr, layout=cute.make_layout( (active_m * Int32(self.hidden_size),), stride=(1,) ), ) topk_weights_flat = cute.make_tensor( topk_weights_ptr, layout=cute.make_layout((active_m * Int32(self.top_k),), stride=(1,)), ) grid = (grid_x, 1, 1) self.kernel( a_bf16_flat, w13_i32_flat, w2_i32_flat, fc1_bf16_flat, activated_bf16_flat, fc2_bf16_flat, w13_scales_i32_flat, w2_scales_i32_flat, w13_global_scale, w2_global_scale, packed_route_indices, block_expert_ids, packed_route_count, activation_amax_flat, layer_idx, topk_weights_flat, fc1_c_tmp_f32_flat, fc2_c_tmp_f32_flat, locks_i32_flat, active_m, ).launch( grid=grid, block=[self.cta_threads, 1, 1], min_blocks_per_mp=self.blocks_per_sm, stream=stream, ) @cute.kernel def kernel( self, a_bf16_flat: cute.Tensor, w13_i32_flat: cute.Tensor, w2_i32_flat: cute.Tensor, fc1_bf16_flat: cute.Tensor, activated_bf16_flat: cute.Tensor, fc2_bf16_flat: cute.Tensor, w13_scales_i32_flat: cute.Tensor, w2_scales_i32_flat: cute.Tensor, w13_global_scale: cute.Tensor, w2_global_scale: cute.Tensor, packed_route_indices: cute.Tensor, block_expert_ids: cute.Tensor, packed_route_count: cute.Tensor, activation_amax_flat: cute.Tensor, layer_idx: cutlass.Int32, topk_weights_flat: cute.Tensor, fc1_c_tmp_f32_flat: cute.Tensor, fc2_c_tmp_f32_flat: cute.Tensor, locks_i32_flat: cute.Tensor, active_m: cutlass.Int32, ): tidx, _, _ = cute.arch.thread_idx() bidx, _, _ = cute.arch.block_idx() grid_x_raw, _, _ = cute.arch.grid_dim() tid = Int32(tidx) cta = Int32(bidx) grid_x = Int32(grid_x_raw) smem = cutlass.utils.SmemAllocator() @cute.struct class Storage: words: cute.struct.Align[ cute.struct.MemRange[cutlass.Uint32, self.shared_words], 1024, ] storage = smem.allocate(Storage) smem_base = shared_ptr_to_u32(storage.words.data_ptr()) if cutlass.const_expr(self.tc_decode_fused_sum): # The TC-decode FC2 epilogue atomically accumulates per-route # partials directly into the per-token output, so the output must be # pre-zeroed. Previously this was a SEPARATE host-side output.zero_() # kernel launch on the latency-bound decode critical path (an extra # launch + its grid-fill memset before the fused kernel even starts). # Fold it into the fused kernel prologue here: every CTA zeroes a # grid-strided slice of the output BEFORE FC1, and the EXISTING # post-FC1 grid barrier (already required to order FC1 writes before # the activation/FC2 read) makes all zero stores globally visible # before the first FC2 atomic -- so no extra barrier is added. The # tiny m*hidden bf16 memset (decode: <=4*4096 elems) is dwarfed by # FC1's whole-K FP4-weight stream, but we delete one whole kernel # launch from the per-decode chain. The TC-decode output is per-token # (top_k routes atomically summed into the SAME token row), so the # zero span is active_m*hidden_size -- NOT the per-route # active_m*top_k*hidden_size of _zero_fc2_output. # tc_zero_output=False skips the zero (a paired earlier launch has # already zeroed the shared output); the grid barrier below is # unconditional so ordering is preserved either way. if cutlass.const_expr(self.tc_zero_output): zidx = cta * Int32(self.cta_threads) + tid zstride = grid_x * Int32(self.cta_threads) ztotal = active_m * Int32(self.hidden_size) zzero = self._cast_elem(cutlass.Float32(0.0)) while zidx < ztotal: fc2_bf16_flat[zidx] = zzero zidx += zstride if cutlass.const_expr(self.activation_is_gated): self.fc1._run_persistent_gemm( a_bf16_flat, w13_i32_flat, fc1_bf16_flat, w13_scales_i32_flat, w13_global_scale, packed_route_indices, block_expert_ids, packed_route_count, topk_weights_flat, fc1_c_tmp_f32_flat, locks_i32_flat, smem_base, tid, cta, grid_x, active_m, ) self._grid_barrier(locks_i32_flat, tid, grid_x) self._run_activation( fc1_bf16_flat, activated_bf16_flat, tid, cta, grid_x, active_m, ) else: self.fc1._run_persistent_gemm( a_bf16_flat, w13_i32_flat, activated_bf16_flat, w13_scales_i32_flat, w13_global_scale, packed_route_indices, block_expert_ids, packed_route_count, topk_weights_flat, fc1_c_tmp_f32_flat, locks_i32_flat, smem_base, tid, cta, grid_x, active_m, ) self._grid_barrier(locks_i32_flat, tid, grid_x) if cutlass.const_expr(self.collect_activation_amax): self._collect_activation_amax_epilogue( a_bf16_flat, activated_bf16_flat, packed_route_indices, block_expert_ids, packed_route_count, activation_amax_flat, layer_idx, smem_base, tid, cta, grid_x, active_m, ) if cutlass.const_expr(self.zero_fc2_output): self._zero_fc2_output(fc2_bf16_flat, tid, cta, grid_x, active_m) self._grid_barrier(locks_i32_flat, tid, grid_x) self.fc2._run_persistent_gemm( activated_bf16_flat, w2_i32_flat, fc2_bf16_flat, w2_scales_i32_flat, w2_global_scale, packed_route_indices, block_expert_ids, packed_route_count, topk_weights_flat, fc2_c_tmp_f32_flat, locks_i32_flat, smem_base, tid, cta, grid_x, active_m * Int32(self.top_k), ) @cute.jit def _grid_barrier( self, locks_i32_flat: cute.Tensor, tid: Int32, grid_x: Int32, ): cute.arch.sync_threads() if tid == Int32(0): count_addr = get_ptr_as_int64(locks_i32_flat, Int32(self.barrier_count_off)) sense_addr = get_ptr_as_int64(locks_i32_flat, Int32(self.barrier_sense_off)) old_sense = ld_global_acquire_i32(sense_addr) old_count = atomic_add_global_i32(count_addr, Int32(1)) if old_count == grid_x - Int32(1): st_global_i32(count_addr, Int32(0)) threadfence() red_add_global_release_i32(sense_addr, Int32(1)) else: sense = old_sense while sense == old_sense: sense = ld_global_acquire_i32(sense_addr) cute.arch.sync_threads() @cute.jit def _zero_fc2_output( self, fc2_bf16_flat: cute.Tensor, tid: Int32, cta: Int32, grid_x: Int32, active_m: cutlass.Int32, ): idx = cta * Int32(self.cta_threads) + tid stride = grid_x * Int32(self.cta_threads) total = active_m * Int32(self.top_k * self.hidden_size) zero = self._cast_elem(cutlass.Float32(0.0)) while idx < total: fc2_bf16_flat[idx] = zero idx += stride @cute.jit def _reduce_and_red_activation_amax( self, local_max: cutlass.Float32, activation_amax_flat: cute.Tensor, layer_idx: cutlass.Int32, expert_idx: Int32, slot: Int32, smem_base: Int32, tid: Int32, ): lane = tid & Int32(31) warp_idx = tid // Int32(32) warp_amax = warp_reduce(local_max, fmax_f32) if lane == Int32(0): st_shared_f32(smem_base + warp_idx * Int32(4), warp_amax) cute.arch.sync_threads() if warp_idx == Int32(0): block_amax = cutlass.Float32(0.0) if lane < Int32(self.cta_threads // 32): block_amax = ld_shared_f32(smem_base + lane * Int32(4)) block_amax = warp_reduce(block_amax, fmax_f32) if lane == Int32(0) and block_amax > cutlass.Float32(0.0): out_idx = ( (layer_idx * Int32(self.num_experts) + expert_idx) * Int32(2) + slot ) red_max_global_f32_nonnegative( get_ptr_as_int64(activation_amax_flat, out_idx), block_amax, ) cute.arch.sync_threads() @cute.jit def _collect_activation_amax_epilogue( self, a_bf16_flat: cute.Tensor, activated_bf16_flat: cute.Tensor, packed_route_indices: cute.Tensor, block_expert_ids: cute.Tensor, packed_route_count: cute.Tensor, activation_amax_flat: cute.Tensor, layer_idx: cutlass.Int32, smem_base: Int32, tid: Int32, cta: Int32, grid_x: Int32, active_m: cutlass.Int32, ): live_routes = active_m * Int32(self.top_k) route_count = packed_route_count[Int32(0)].to(Int32) route_blocks = ( route_count + Int32(self.moe_block_size) - Int32(1) ) // Int32(self.moe_block_size) expert_idx = cta while expert_idx < Int32(self.num_experts): local_fc1 = cutlass.Float32(0.0) local_fc2 = cutlass.Float32(0.0) has_route_block = Int32(0) block_idx = Int32(0) while block_idx < route_blocks: block_expert = block_expert_ids[block_idx].to(Int32) if block_expert == expert_idx: has_route_block = Int32(1) route_pos = block_idx * Int32(self.moe_block_size) route_stop = route_pos + Int32(self.moe_block_size) if route_stop > route_count: route_stop = route_count while route_pos < route_stop: route_idx = packed_route_indices[route_pos].to(Int32) if route_idx >= Int32(0) and route_idx < live_routes: token_idx = route_idx // Int32(self.top_k) fc1_col = tid while fc1_col < Int32(self.hidden_size): v1 = a_bf16_flat[ token_idx * Int32(self.hidden_size) + fc1_col ].to(cutlass.Float32) local_fc1 = fmax_f32(local_fc1, fabs_f32(v1)) fc1_col += Int32(self.cta_threads) fc2_col = tid while fc2_col < Int32(self.intermediate_size): v2 = activated_bf16_flat[ route_idx * Int32(self.intermediate_size) + fc2_col ].to(cutlass.Float32) local_fc2 = fmax_f32(local_fc2, fabs_f32(v2)) fc2_col += Int32(self.cta_threads) route_pos += Int32(1) block_idx += Int32(1) if has_route_block != Int32(0): self._reduce_and_red_activation_amax( local_fc1, activation_amax_flat, layer_idx, expert_idx, Int32(0), smem_base, tid, ) self._reduce_and_red_activation_amax( local_fc2, activation_amax_flat, layer_idx, expert_idx, Int32(1), smem_base, tid, ) expert_idx += grid_x @cute.jit def _run_activation( self, fc1_bf16_flat: cute.Tensor, activated_bf16_flat: cute.Tensor, tid: Int32, cta: Int32, grid_x: Int32, active_m: cutlass.Int32, ): idx = cta * Int32(self.cta_threads) + tid stride = grid_x * Int32(self.cta_threads) total = active_m * Int32(self.top_k * self.intermediate_size) while idx < total: if cutlass.const_expr(self.activation_is_gated): row = idx // Int32(self.intermediate_size) col = idx - row * Int32(self.intermediate_size) base = row * Int32(self.fc1_cols) gate = fc1_bf16_flat[base + col].to(cutlass.Float32) up = fc1_bf16_flat[base + Int32(self.intermediate_size) + col].to( cutlass.Float32 ) gate, up = self._clamp_swiglu_inputs(gate, up) sigmoid_arg = gate up_term = up if cutlass.const_expr(self.activation_is_swigluoai): sigmoid_arg = cutlass.Float32(self.swiglu_alpha) * gate up_term = up + cutlass.Float32(self.swiglu_beta) if cutlass.const_expr(self.fast_math): exp_neg_gate = cute.math.exp(-sigmoid_arg, fastmath=True) else: exp_neg_gate = cute.math.exp(-sigmoid_arg, fastmath=False) silu = gate / (cutlass.Float32(1.0) + exp_neg_gate) if cutlass.const_expr(self.activation_is_swigluoai): activated_bf16_flat[idx] = self._cast_elem(silu * up_term) else: activated_bf16_flat[idx] = self._cast_elem( self._cast_elem(silu) * self._cast_elem(up_term) ) else: x = fc1_bf16_flat[idx].to(cutlass.Float32) if x < cutlass.Float32(0.0): x = cutlass.Float32(0.0) activated_bf16_flat[idx] = self._cast_elem(x * x) idx += stride class W4A16ActivationKernel: def __init__( self, *, rows: int, intermediate_size: int, activation: str, element_dtype: str = "bf16", fast_math: bool = True, swiglu_limit: float | None = None, swiglu_alpha: float | None = None, swiglu_beta: float | None = None, ): activation = normalize_moe_activation(activation) is_gated = validate_activation(activation) swiglu_limit, swiglu_alpha, swiglu_beta = _normalize_activation_swiglu_params( activation, swiglu_limit, swiglu_alpha, swiglu_beta, ) if element_dtype not in {"bf16", "fp16"}: raise ValueError(f"unsupported element_dtype {element_dtype!r}") if rows <= 0 or intermediate_size <= 0: raise ValueError("rows and intermediate_size must be positive") self.rows = int(rows) self.intermediate_size = int(intermediate_size) self.activation = activation self.is_gated = is_gated self.is_swigluoai = activation == SWIGLUOAI_UNINTERLEAVE self.has_swiglu_limit = swiglu_limit is not None self.swiglu_limit = 0.0 if swiglu_limit is None else swiglu_limit self.swiglu_alpha = float(swiglu_alpha) self.swiglu_beta = float(swiglu_beta) self.element_dtype = element_dtype self.is_fp16 = element_dtype == "fp16" self.fast_math = bool(fast_math) self.cta_threads = 256 @property def __cache_key__(self) -> tuple[object, ...]: return ( self.intermediate_size, self.activation, self.is_gated, self.is_swigluoai, self.has_swiglu_limit, self.swiglu_limit, self.swiglu_alpha, self.swiglu_beta, self.element_dtype, self.fast_math, self.cta_threads, ) @cute.jit def _cast_elem(self, x: cutlass.Float32): if cutlass.const_expr(self.is_fp16): return cutlass.Float16(x) return cutlass.BFloat16(x) @cute.jit def _clamp_swiglu_inputs( self, gate: cutlass.Float32, up: cutlass.Float32, ): if cutlass.const_expr(self.has_swiglu_limit): limit = cutlass.Float32(self.swiglu_limit) neg_limit = cutlass.Float32(-self.swiglu_limit) if gate > limit: gate = limit if up > limit: up = limit if up < neg_limit: up = neg_limit return gate, up @cute.jit def __call__( self, fc1_flat: cute.Tensor, activated_flat: cute.Tensor, active_rows: cutlass.Int32, stream: cuda.CUstream, ): total = active_rows * Int32(self.intermediate_size) grid = (_covering_count(total, self.cta_threads), 1, 1) self.kernel(fc1_flat, activated_flat, active_rows).launch( grid=grid, block=[self.cta_threads, 1, 1], stream=stream, ) @cute.kernel def kernel( self, fc1_flat: cute.Tensor, activated_flat: cute.Tensor, active_rows: cutlass.Int32, ): tidx, _, _ = cute.arch.thread_idx() bidx, _, _ = cute.arch.block_idx() idx = Int32(bidx) * Int32(self.cta_threads) + Int32(tidx) total = Int32(active_rows) * Int32(self.intermediate_size) if idx < total: if cutlass.const_expr(self.is_gated): row = idx // Int32(self.intermediate_size) col = idx - row * Int32(self.intermediate_size) base = row * Int32(2 * self.intermediate_size) gate = fc1_flat[base + col].to(cutlass.Float32) up = fc1_flat[base + Int32(self.intermediate_size) + col].to( cutlass.Float32 ) gate, up = self._clamp_swiglu_inputs(gate, up) sigmoid_arg = gate up_term = up if cutlass.const_expr(self.is_swigluoai): sigmoid_arg = cutlass.Float32(self.swiglu_alpha) * gate up_term = up + cutlass.Float32(self.swiglu_beta) if cutlass.const_expr(self.fast_math): exp_neg_gate = cute.math.exp(-sigmoid_arg, fastmath=True) else: exp_neg_gate = cute.math.exp(-sigmoid_arg, fastmath=False) silu = gate / (cutlass.Float32(1.0) + exp_neg_gate) if cutlass.const_expr(self.is_swigluoai): activated_flat[idx] = self._cast_elem(silu * up_term) else: activated_flat[idx] = self._cast_elem( self._cast_elem(silu) * self._cast_elem(up_term) ) else: x = fc1_flat[idx].to(cutlass.Float32) if x < cutlass.Float32(0.0): x = cutlass.Float32(0.0) activated_flat[idx] = self._cast_elem(x * x) class W4A16TopKSumKernel: def __init__( self, *, topk: int, hidden_size: int, element_dtype: str = "bf16" ): if element_dtype not in {"bf16", "fp16"}: raise ValueError(f"unsupported element_dtype {element_dtype!r}") if topk <= 0 or hidden_size <= 0: raise ValueError("topk and hidden_size must be positive") self.topk = int(topk) self.hidden_size = int(hidden_size) self.element_dtype = element_dtype self.is_fp16 = element_dtype == "fp16" self.cta_threads = 256 @cute.jit def _cast_elem(self, x: cutlass.Float32): if cutlass.const_expr(self.is_fp16): return cutlass.Float16(x) return cutlass.BFloat16(x) @cute.jit def __call__( self, fc2_ptr: cute.Pointer, output_ptr: cute.Pointer, active_m: cutlass.Int32, stream: cuda.CUstream, ): fc2_flat = cute.make_tensor( fc2_ptr, layout=cute.make_layout( (active_m * Int32(self.topk * self.hidden_size),), stride=(1,) ), ) output_flat = cute.make_tensor( output_ptr, layout=cute.make_layout((active_m * Int32(self.hidden_size),), stride=(1,)), ) total = active_m * Int32(self.hidden_size) grid = (_covering_count(total, self.cta_threads), 1, 1) self.kernel(fc2_flat, output_flat, active_m).launch( grid=grid, block=[self.cta_threads, 1, 1], stream=stream, ) @cute.kernel def kernel( self, fc2_flat: cute.Tensor, output_flat: cute.Tensor, active_m: cutlass.Int32, ): tidx, _, _ = cute.arch.thread_idx() bidx, _, _ = cute.arch.block_idx() idx = Int32(bidx) * Int32(self.cta_threads) + Int32(tidx) total = active_m * Int32(self.hidden_size) if idx < total: token = idx // Int32(self.hidden_size) col = idx - token * Int32(self.hidden_size) acc = cutlass.Float32(0.0) for route in cutlass.range_constexpr(self.topk): row = token * Int32(self.topk) + Int32(route) acc += fc2_flat[row * Int32(self.hidden_size) + col].to(cutlass.Float32) output_flat[idx] = self._cast_elem(acc) _CACHE: dict[tuple, W4A16GemmCompileResult] = {} _FUSED_CACHE: dict[tuple, W4A16FusedMoeCompileResult] = {} _ACTIVATION_CACHE: dict[tuple, W4A16ActivationCompileResult] = {} _SUM_CACHE: dict[tuple, W4A16TopKSumCompileResult] = {} _SMALL_M_DIRECT_CACHE: dict[tuple, _W4A16SmallMDirectLaunch] = {} def _normalize_element_dtype(dtype: torch.dtype) -> str: if dtype == torch.bfloat16: return "bf16" if dtype == torch.float16: return "fp16" raise TypeError(f"unsupported W4A16 activation dtype {dtype}") def _normalize_scale_format(scale_format: str) -> str: try: return _SCALE_FORMATS[scale_format.lower()] except KeyError as exc: raise ValueError( "scale_format must be one of 'e4m3_k16', 'e8m0_k32', or 'e4m3_k32', " f"got {scale_format!r}" ) from exc def _scale_group_size(scale_format: str) -> int: return ( 32 if _normalize_scale_format(scale_format) in ("e8m0_k32", "e4m3_k32") else 16 ) def _scale_fake_int32_elements( *, num_experts: int, size_k: int, size_n: int, scale_format: str, allow_k_tail: bool = False, allow_n_tail: bool = False, ) -> int: group_size = _scale_group_size(scale_format) if int(size_n) % 16 != 0: raise ValueError( f"W4A16 {scale_format} scales require size_n divisible by 16" ) if int(size_k) % group_size != 0 and not allow_k_tail: raise ValueError( f"W4A16 {scale_format} scales require size_k divisible by {group_size}, " f"got {size_k}" ) groups_k = ( _covering_count(int(size_k), group_size) if allow_k_tail else int(size_k) // group_size ) scale_size_n = ( _e8m0_logical_tail_scale_n(int(size_n)) if allow_n_tail else int(size_n) ) return int(num_experts) * groups_k * (scale_size_n // 4) def _cutlass_element_dtype(element_dtype: str): if element_dtype == "bf16": return cutlass.BFloat16 if element_dtype == "fp16": return cutlass.Float16 raise ValueError(f"unsupported element_dtype {element_dtype!r}") def _small_m_direct_supported( *, m: int, hidden_size: int, intermediate_size: int, num_experts: int, topk: int, activation: str, apply_router_weight_on_input: bool, swiglu_limit: float | None, swiglu_alpha: float, swiglu_beta: float, element_dtype: str, weight_layout: str, w13_layout: str, scale_format: str = "e4m3_k16", expert_map: torch.Tensor | None = None, ) -> bool: if os.environ.get("B12X_W4A16_SMALL_M_DIRECT", "1") == "0": return False activation = normalize_moe_activation(activation) swiglu_limit, swiglu_alpha, swiglu_beta = _normalize_activation_swiglu_params( activation, swiglu_limit, swiglu_alpha, swiglu_beta, ) if activation == SWIGLUOAI_UNINTERLEAVE and swiglu_limit is None: return False return ( element_dtype == "bf16" # Packed serving weights use the W4A16 TC-decode path for small M. # The direct micro kernel is kept limited to native ModelOpt weights. and weight_layout == "modelopt" and w13_layout in ("w13", "w31") and not bool(apply_router_weight_on_input) and (swiglu_limit is None or activation in ("silu", SWIGLUOAI_UNINTERLEAVE)) and expert_map is None and _W4A16SmallMDirectKernel.is_supported( m=m, hidden_size=hidden_size, intermediate_size=intermediate_size, topk=topk, num_experts=num_experts, scale_format=scale_format, ) ) def _compile_w4a16_small_m_direct( *, m: int, hidden_size: int, intermediate_size: int, num_experts: int, topk: int, activation: str, fast_math: bool, topk_ids_dtype: torch.dtype, device: torch.device | None, scale_format: str = "e4m3_k16", swiglu_limit: float | None = None, swiglu_alpha: float | None = None, swiglu_beta: float | None = None, w13_layout: str = "w13", ) -> _W4A16SmallMDirectLaunch: if topk_ids_dtype not in (torch.int32, torch.int64): raise TypeError("small-M W4A16 direct path requires int32/int64 topk_ids") activation = normalize_moe_activation(activation) swiglu_limit, swiglu_alpha, swiglu_beta = _normalize_activation_swiglu_params( activation, swiglu_limit, swiglu_alpha, swiglu_beta, ) cache_key = ( "w4a16_small_m_direct", None if device is None else int(device.index or 0), int(m), int(hidden_size), int(intermediate_size), int(num_experts), int(topk), activation, bool(fast_math), topk_ids_dtype, scale_format, swiglu_limit, swiglu_alpha, swiglu_beta, w13_layout, ) cached = _SMALL_M_DIRECT_CACHE.get(cache_key) if cached is not None: return cached kernel = _W4A16SmallMDirectKernel( activation=activation, fast_math=bool(fast_math), share_input_across_experts=(int(m) == 1), share_expert_scales=True, single_token=(int(m) == 1), scale_format=scale_format, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, w13_layout=w13_layout, ) kernel.configure( int(m), int(hidden_size), int(intermediate_size), int(topk), int(num_experts), device=device, ) def dummy(dt): return make_ptr(dt, 16, cute.AddressSpace.gmem, assumed_align=16) ids_dtype = cutlass.Int32 if topk_ids_dtype == torch.int32 else cutlass.Int64 barrier_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (1,), assumed_align=4, ) raise_if_kernel_resolution_frozen( "cute.compile", target=kernel, cache_key=cache_key ) compiled = b12x_compile( kernel, dummy(cutlass.BFloat16), dummy(cutlass.Uint8), dummy(cutlass.Uint8), dummy(cutlass.Float32), dummy(cutlass.Float32), dummy(cutlass.Float32), dummy(cutlass.Uint32), dummy(cutlass.Uint8), dummy(cutlass.Uint8), dummy(cutlass.Float32), dummy(ids_dtype), dummy(cutlass.Float32), dummy(cutlass.BFloat16), barrier_fake, barrier_fake, Int32(m), Int32(kernel.grid_x), current_cuda_stream(), compile_spec=KernelCompileSpec.from_facts( "moe.w4a16.small_m_direct", 1, ("device_index", None if device is None else int(device.index or 0)), ("m", int(m)), ("hidden_size", int(hidden_size)), ("intermediate_size", int(intermediate_size)), ("num_experts", int(num_experts)), ("topk", int(topk)), ("activation", activation), ("fast_math", bool(fast_math)), ("topk_ids_dtype", str(topk_ids_dtype)), ("scale_format", scale_format), ("swiglu_limit", swiglu_limit), ("swiglu_alpha", swiglu_alpha), ("swiglu_beta", swiglu_beta), ("w13_layout", w13_layout), ("grid_x", int(kernel.grid_x)), ), ) launch = _W4A16SmallMDirectLaunch( compiled=compiled, grid_x=int(kernel.grid_x), m=int(m), hidden_size=int(hidden_size), intermediate_size=int(intermediate_size), num_experts=int(num_experts), topk=int(topk), activation=activation, fast_math=bool(fast_math), topk_ids_dtype=topk_ids_dtype, ) _SMALL_M_DIRECT_CACHE[cache_key] = launch return launch def compile_w4a16_gemm( *, size_m: int, size_n: int, size_k: int, num_experts: int, top_k: int, mul_topk_weights: bool, tile_n: int, tile_k: int, moe_block_size: int, max_m_blocks: int, element_dtype: str = "bf16", scale_format: str = "e4m3_k16", ) -> W4A16GemmCompileResult: scale_format = _normalize_scale_format(scale_format) cutlass_dtype = _cutlass_element_dtype(element_dtype) if torch.cuda.is_available(): device = int(torch.cuda.current_device()) else: device = None kernel = W4A16GemmKernel( size_m=size_m, size_n=size_n, size_k=size_k, num_experts=num_experts, top_k=top_k, mul_topk_weights=mul_topk_weights, tile_n=tile_n, tile_k=tile_k, moe_block_size=moe_block_size, max_m_blocks=max_m_blocks, element_dtype=element_dtype, scale_format=scale_format, ) cache_key = ( "w4a16_gemm", device, kernel.__cache_key__, ) cached = _CACHE.get(cache_key) if cached is not None: return replace( cached, max_m_blocks=max_m_blocks, blocks_per_sm=kernel.blocks_per_sm, ) compile_size_m = _fake_m_for_specialization(size_m) compile_route_blocks = 1 compile_route_slots = compile_route_blocks * int(moe_block_size) a_fake = cute.runtime.make_fake_compact_tensor( cutlass_dtype, (compile_size_m * size_k,), assumed_align=16, ) b_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (num_experts * (size_k // 16) * (size_n // 16 * 32),), assumed_align=16, ) c_fake = cute.runtime.make_fake_compact_tensor( cutlass_dtype, (compile_size_m * top_k * size_n,), assumed_align=16, ) scales_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, ( _scale_fake_int32_elements( num_experts=num_experts, size_k=size_k, size_n=size_n, scale_format=scale_format, ), ), assumed_align=16, ) global_scale_fake = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (num_experts,), assumed_align=16, ) packed_routes_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (compile_route_slots,), assumed_align=16, ) block_experts_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (compile_route_blocks,), assumed_align=16, ) route_count_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (1,), assumed_align=4, ) topk_fake = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (compile_size_m * top_k,), assumed_align=4, ) c_tmp_fake = cute.runtime.make_fake_compact_tensor( cutlass.Float32, ( max( size_n * compile_route_slots, 4 * 256 * moe_block_size * 256, ), ), assumed_align=16, ) locks_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (4 * 256,), assumed_align=16, ) raise_if_kernel_resolution_frozen( "cute.compile", target=kernel, cache_key=cache_key ) compiled = b12x_compile( kernel, a_fake, b_fake, c_fake, scales_fake, global_scale_fake, packed_routes_fake, block_experts_fake, route_count_fake, topk_fake, c_tmp_fake, locks_fake, Int32(compile_size_m), Int32(1), current_cuda_stream(), compile_spec=KernelCompileSpec.from_key( "moe.w4a16.gemm", 1, cache_key, ), ) result = W4A16GemmCompileResult( compiled=compiled, tile_n=tile_n, tile_k=tile_k, moe_block_size=moe_block_size, max_m_blocks=max_m_blocks, blocks_per_sm=kernel.blocks_per_sm, scale_format=scale_format, ) _CACHE[cache_key] = result return result def compile_w4a16_fused_moe( *, size_m: int, hidden_size: int, intermediate_size: int, num_experts: int, top_k: int, activation: str, apply_router_weight_on_input: bool, zero_fc2_output: bool, moe_block_size: int, max_m_blocks: int, element_dtype: str = "bf16", fast_math: bool = True, sms: int, max_shared_mem: int, swiglu_limit: float | None = None, swiglu_alpha: float | None = None, swiglu_beta: float | None = None, weight_layout: str = "packed", scale_format: str = "e4m3_k16", w13_layout: str = "w13", direct_topk_routes: bool = False, tc_decode_fused_sum: bool = False, collect_activation_amax: bool = False, force_tile_config: tuple[int, int, int, int] | None = None, ) -> W4A16FusedMoeCompileResult: scale_format = _normalize_scale_format(scale_format) cutlass_dtype = _cutlass_element_dtype(element_dtype) device = int(torch.cuda.current_device()) if torch.cuda.is_available() else None activation = normalize_moe_activation(activation) is_gated = validate_activation(activation) swiglu_limit, swiglu_alpha, swiglu_beta = _normalize_activation_swiglu_params( activation, swiglu_limit, swiglu_alpha, swiglu_beta, ) if weight_layout not in _WEIGHT_LAYOUTS: raise ValueError(f"unsupported W4A16 weight_layout {weight_layout!r}") if weight_layout == "modelopt": if w13_layout not in _MODEL_OPT_W13_LAYOUTS: raise ValueError(f"unsupported W4A16 w13_layout {w13_layout!r}") else: w13_layout = "packed" direct_topk_routes = bool(direct_topk_routes) tc_decode_fused_sum = bool(tc_decode_fused_sum) collect_activation_amax = bool(collect_activation_amax) if collect_activation_amax and (direct_topk_routes or tc_decode_fused_sum): raise ValueError( "W4A16 activation amax collection requires the route-packed fused path" ) # The TC-decode path validates M in {1,2,4,8} itself and uses direct-topk # routing for the whole {1,2,4,8} range, so it lifts the default decode cap. direct_topk_m_cap = ( _W4A16_SMALL_M_DIRECT_MAX_M if tc_decode_fused_sum else _MAX_DIRECT_TOPK_ROUTE_M ) if direct_topk_routes and ( int(size_m) > direct_topk_m_cap or weight_layout not in ("packed", "nf3_2p1") or bool(zero_fc2_output) ): raise ValueError( "direct_topk_routes is only valid for small-M packed W4A16 without expert_map" ) fc1_cols = int(intermediate_size) * (2 if is_gated else 1) routed_rows = int(size_m) * int(top_k) allow_native_logical_tail = ( weight_layout == "modelopt" and scale_format == "e8m0_k32" and int(intermediate_size) % 32 != 0 ) fc1_tile_k, fc1_tile_n, fc1_cta_threads, _ = _select_tile_config( problem_m=size_m, problem_n=fc1_cols, problem_k=hidden_size, top_k=top_k, moe_block_size=moe_block_size, sms=sms, max_shared_mem=max_shared_mem, scale_format=scale_format, allow_logical_tail=allow_native_logical_tail, ) fc2_tile_k, fc2_tile_n, fc2_cta_threads, _ = _select_tile_config( problem_m=routed_rows, problem_n=hidden_size, problem_k=intermediate_size, top_k=1, moe_block_size=moe_block_size, sms=sms, max_shared_mem=max_shared_mem, scale_format=scale_format, allow_logical_tail=allow_native_logical_tail, ) if fc1_cta_threads != fc2_cta_threads: common_cta_threads = min(fc1_cta_threads, fc2_cta_threads) fc1_tile_k, fc1_tile_n, fc1_cta_threads, _ = _select_tile_config( problem_m=size_m, problem_n=fc1_cols, problem_k=hidden_size, top_k=top_k, moe_block_size=moe_block_size, sms=sms, max_shared_mem=max_shared_mem, required_cta_threads=common_cta_threads, scale_format=scale_format, allow_logical_tail=allow_native_logical_tail, ) fc2_tile_k, fc2_tile_n, fc2_cta_threads, _ = _select_tile_config( problem_m=routed_rows, problem_n=hidden_size, problem_k=intermediate_size, top_k=1, moe_block_size=moe_block_size, sms=sms, max_shared_mem=max_shared_mem, required_cta_threads=common_cta_threads, scale_format=scale_format, allow_logical_tail=allow_native_logical_tail, ) if fc1_cta_threads != fc2_cta_threads: raise ValueError( "fused W4A16 FC1/FC2 selected different thread counts: " f"{fc1_cta_threads} vs {fc2_cta_threads}" ) # TC-decode FC1 wide-N override (single-wave-collapse sizes only): in the m8 # fused path the host right-sizes grid_x to the FC1 mn-tile count and forces # one whole mn-tile per CTA over the full K. With the default fc1_tile_n=128, # FC1 (N=fc1_cols=2*intermediate_size) produces size_m*top_k*(fc1_cols/128) # mn-tiles. For TP=2 I_tp=1024 (fc1_cols=2048 => 16 n-tiles/route) bs=1 is # 96 tiles (<= 188 SMs, one wave); bs=2 is 192 -- JUST over the single-wave # SM cap -- forcing a 2-wave launch of grid_x=96 (half the machine idle per # wave, a serialized second FC1 wave of pure tail latency on the bandwidth- # bound decode). Widening FC1 to tile_n=256 (256-wide N slab per CTA over # full K; tile_k=64 keeps cta_threads=256) HALVES FC1's mn-tile count, so # bs=2 collapses to 96 tiles = exactly ONE wave, removing that whole second # FC1 wave. The narrower tile_k=64 is slower per-tile, so we apply this ONLY # where it turns a 2-wave launch into a 1-wave launch: the default 128-wide # FC1 spans 2 waves (sms < default_mn_tiles <= 2*sms) AND the wide tile fits # in one (default/2 <= sms). bs=1 (one wave already) and bs>=4 (still multi- # wave after halving) keep the faster default 128x128 tile. # Guarded by fc1_cols%256==0, smem-fit, and 256-thread geometry so the fused # FC1/FC2 single-thread-geometry contract is preserved. default_fc1_mn_tiles = ( int(size_m) * int(top_k) * (int(fc1_cols) // int(fc1_tile_n)) if fc1_tile_n > 0 else 0 ) if ( bool(tc_decode_fused_sum) and int(fc1_cols) % 256 == 0 and fc1_tile_n == 128 and (fc1_tile_n * fc1_tile_k) // 64 == 256 and int(sms) < default_fc1_mn_tiles <= 2 * int(sms) and (default_fc1_mn_tiles // 2) <= int(sms) ): wide_fc1_tile_k = 64 if _candidate_tile_fits( problem_n=fc1_cols, problem_k=hidden_size, cta_m_blocks=_covering_count(moe_block_size, 16), tile_n=256, tile_k=wide_fc1_tile_k, cta_threads=256, max_shared_mem=int(max_shared_mem) - 512, scale_format=scale_format, ): fc1_tile_n = 256 fc1_tile_k = wide_fc1_tile_k fc1_cta_threads = 256 # TC-decode FC2 wide-N override: in the m8 fused path the host right-sizes # grid_x to the FC1 mn-tile count (one persistent wave, no split-K). FC2 # (N=hidden_size, K=intermediate_size) with the default tile_n=128 produces # route_blocks*(hidden_size/128) mn-tiles -- roughly double FC1's count -- # so FC2 would need ~2 persistent waves while FC1 fits in 1; that second # FC2 wave is pure serialized tail latency on the bandwidth-bound decode. # Selecting tile_n=256 for FC2 (a 256-wide N slab per CTA over full K) # halves FC2's mn-tile count so it also fits one wave. Guarded by smem-fit, # hidden_size%256==0, and matching cta_threads so the fused single # thread-geometry contract is preserved. if ( bool(tc_decode_fused_sum) and int(hidden_size) % 256 == 0 and fc2_tile_n == 128 and fc1_cta_threads == 256 ): wide_fc2_tile_k = 64 if _candidate_tile_fits( problem_n=hidden_size, problem_k=intermediate_size, cta_m_blocks=_covering_count(moe_block_size, 16), tile_n=256, tile_k=wide_fc2_tile_k, cta_threads=256, max_shared_mem=int(max_shared_mem) - 512, scale_format=scale_format, ): fc2_tile_n = 256 fc2_tile_k = wide_fc2_tile_k fc2_cta_threads = 256 # TC-decode FC2 ultra-wide override (perfect wave-balance with FC1): the # persistent grid_x is right-sized to FC1's mn-tile count. After the FC1/FC2 # wide-N (tile_n=256) overrides, bs=2 still has FC1=route_blocks*8 tiles vs # FC2=route_blocks*16 tiles -- FC2 is DOUBLE FC1 and thus needs a second # persistent wave at grid_x sized to FC1. That FC2 second wave is pure # serialized tail latency on the bandwidth-bound decode. Widening FC2 to # tile_n=512 (a 512-wide N slab per CTA over full K) halves FC2's mn-tile # count again so FC2 == FC1's tile count and fits the SAME single wave. A # 512-wide N tile needs tile_k=32 to keep cta_threads=256 (512*32/64); that # is below the generic tile_k>=64 fits-floor, so we validate the footprint # directly here. tile_k=32 == the e8m0_k32 scale group, so cta_k_blocks=2 # with one e8m0 scale group per k-tile -- the existing scale layout is a # clean covering. Fire ONLY when it drops FC2 from >1 wave to FC1's wave # count (bs=2). Numerically identical: only the FC2 output-tile width changes. fc1_mn_after = ( int(size_m) * int(top_k) * (int(fc1_cols) // int(fc1_tile_n)) if fc1_tile_n > 0 else 0 ) fc2_mn_after = ( int(size_m) * int(top_k) * (int(hidden_size) // int(fc2_tile_n)) if fc2_tile_n > 0 else 0 ) if ( bool(tc_decode_fused_sum) and int(hidden_size) % 512 == 0 and fc2_tile_n == 256 and fc1_cta_threads == 256 and fc1_mn_after > 0 and fc2_mn_after > fc1_mn_after and (fc2_mn_after // 2) <= fc1_mn_after and fc1_mn_after <= int(sms) ): ultra_fc2_tile_k = 32 ultra_smem = _shared_memory_footprint( cta_m_blocks=_covering_count(moe_block_size, 16), tile_n=512, tile_k=ultra_fc2_tile_k, scale_format=scale_format, ) if ( int(intermediate_size) % ultra_fc2_tile_k == 0 and ultra_smem <= int(max_shared_mem) - 512 ): fc2_tile_n = 512 fc2_tile_k = ultra_fc2_tile_k fc2_cta_threads = 256 if force_tile_config is not None: # Explicit (fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n) pin. The # NF3 ("nf3_2p1") flat-span weight layout is packed for a specific CTA # N-tile, so hybrid deployments pin ONE tile config across every m # regime instead of the m-dependent auto selection (and TC-decode # wide-N overrides) above. Overrides whatever was selected; the tiles # land in the GEMM cache keys, so no cache collision is possible. fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n = ( int(v) for v in force_tile_config ) fc1_cta_threads = (fc1_tile_n * fc1_tile_k) // 64 fc2_cta_threads = (fc2_tile_n * fc2_tile_k) // 64 if fc1_cta_threads != fc2_cta_threads: raise ValueError( "force_tile_config FC1/FC2 thread counts must match, got " f"{fc1_cta_threads} vs {fc2_cta_threads}" ) for name, forced_pn, forced_pk, forced_tn, forced_tk in ( ("fc1", fc1_cols, hidden_size, fc1_tile_n, fc1_tile_k), ("fc2", hidden_size, intermediate_size, fc2_tile_n, fc2_tile_k), ): if not _candidate_tile_fits( problem_n=forced_pn, problem_k=forced_pk, cta_m_blocks=_covering_count(moe_block_size, 16), tile_n=forced_tn, tile_k=forced_tk, cta_threads=fc1_cta_threads, max_shared_mem=int(max_shared_mem) - 512, scale_format=scale_format, allow_logical_tail=allow_native_logical_tail, ): raise ValueError( f"force_tile_config {name} tile " f"(tile_k={forced_tk}, tile_n={forced_tn}) does not fit " f"problem N/K={forced_pn}/{forced_pk} at " f"moe_block_size={moe_block_size}" ) kernel = W4A16FusedMoeKernel( size_m=size_m, hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, top_k=top_k, activation=activation, apply_router_weight_on_input=apply_router_weight_on_input, zero_fc2_output=zero_fc2_output, fc1_tile_n=fc1_tile_n, fc1_tile_k=fc1_tile_k, fc2_tile_n=fc2_tile_n, fc2_tile_k=fc2_tile_k, moe_block_size=moe_block_size, max_m_blocks=max_m_blocks, element_dtype=element_dtype, fast_math=fast_math, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, weight_layout=weight_layout, scale_format=scale_format, w13_layout=w13_layout, direct_topk_routes=direct_topk_routes, tc_decode_fused_sum=tc_decode_fused_sum, collect_activation_amax=collect_activation_amax, ) cache_key = ( "w4a16_fused_moe", device, kernel.__cache_key__, ) cached = _FUSED_CACHE.get(cache_key) if cached is not None: return replace( cached, size_m=size_m, max_m_blocks=max_m_blocks, blocks_per_sm=kernel.blocks_per_sm, ) if (not collect_activation_amax) and _small_m_direct_supported( m=size_m, hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, topk=top_k, activation=activation, apply_router_weight_on_input=bool(apply_router_weight_on_input), swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, element_dtype=element_dtype, weight_layout=weight_layout, w13_layout=w13_layout, scale_format=scale_format, ): for ids_dtype in (torch.int32, torch.int64): _compile_w4a16_small_m_direct( m=size_m, hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, topk=top_k, activation=activation, fast_math=fast_math, topk_ids_dtype=ids_dtype, scale_format=scale_format, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, w13_layout=w13_layout, device=torch.device("cuda", device) if device is not None else None, ) compile_size_m = _fake_m_for_specialization(size_m) compile_routed_rows = int(compile_size_m) * int(top_k) compile_route_blocks = compile_routed_rows if direct_topk_routes else 1 compile_route_slots = compile_route_blocks * int(moe_block_size) packed_route_fake_elements = ( compile_routed_rows if direct_topk_routes else compile_route_slots ) a_fake = make_ptr(cutlass_dtype, 16, cute.AddressSpace.gmem, assumed_align=16) if weight_layout == "modelopt": w13_fake = cute.runtime.make_fake_compact_tensor( cutlass.Uint8, (num_experts * fc1_cols * (hidden_size // 2),), assumed_align=16, ) w2_fake = cute.runtime.make_fake_compact_tensor( cutlass.Uint8, (num_experts * hidden_size * (intermediate_size // 2),), assumed_align=16, ) elif weight_layout == "nf3_2p1": # NF3: int32, 3 words per 32-code unit; (size_n // 2) units per K16 row. w13_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (num_experts * (hidden_size // 16) * (fc1_cols // 2) * 3,), assumed_align=16, ) w2_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (num_experts * (intermediate_size // 16) * (hidden_size // 2) * 3,), assumed_align=16, ) else: w13_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (num_experts * (hidden_size // 16) * (fc1_cols // 16 * 32),), assumed_align=16, ) w2_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (num_experts * (intermediate_size // 16) * (hidden_size // 16 * 32),), assumed_align=16, ) fc1_fake = cute.runtime.make_fake_compact_tensor( cutlass_dtype, (compile_routed_rows * fc1_cols,), assumed_align=16, ) activated_fake = cute.runtime.make_fake_compact_tensor( cutlass_dtype, (compile_routed_rows * intermediate_size,), assumed_align=16, ) fc2_fake = cute.runtime.make_fake_compact_tensor( cutlass_dtype, (compile_routed_rows * hidden_size,), assumed_align=16, ) w13_scales_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, ( _scale_fake_int32_elements( num_experts=num_experts, size_k=hidden_size, size_n=fc1_cols, scale_format=scale_format, allow_n_tail=allow_native_logical_tail, ), ), assumed_align=16, ) w2_scales_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, ( _scale_fake_int32_elements( num_experts=num_experts, size_k=intermediate_size, size_n=hidden_size, scale_format=scale_format, allow_k_tail=allow_native_logical_tail, ), ), assumed_align=16, ) w13_global_fake = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (num_experts,), assumed_align=16, ) w2_global_fake = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (num_experts,), assumed_align=16, ) packed_routes_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (packed_route_fake_elements,), assumed_align=16, ) block_experts_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (compile_route_blocks,), assumed_align=16, ) route_count_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (1,), assumed_align=4, ) activation_amax_fake = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (num_experts * 2,), assumed_align=4, ) topk_fake = make_ptr(cutlass.Float32, 4, cute.AddressSpace.gmem, assumed_align=4) fc1_c_tmp_fake = cute.runtime.make_fake_compact_tensor( cutlass.Float32, ( max( fc1_cols * compile_route_slots, 4 * 256 * moe_block_size * 256, ), ), assumed_align=16, ) fc2_c_tmp_fake = cute.runtime.make_fake_compact_tensor( cutlass.Float32, ( max( hidden_size * compile_route_slots, 4 * 256 * moe_block_size * 256, ), ), assumed_align=16, ) locks_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (4 * 256 + 2,), assumed_align=16, ) raise_if_kernel_resolution_frozen( "cute.compile", target=kernel, cache_key=cache_key ) compiled = b12x_compile( kernel, a_fake, w13_fake, w2_fake, fc1_fake, activated_fake, fc2_fake, w13_scales_fake, w2_scales_fake, w13_global_fake, w2_global_fake, packed_routes_fake, block_experts_fake, route_count_fake, activation_amax_fake, 0, topk_fake, fc1_c_tmp_fake, fc2_c_tmp_fake, locks_fake, 1, 1, current_cuda_stream(), compile_spec=KernelCompileSpec.from_key( "moe.w4a16.fused_moe", 1, cache_key, ), ) result = W4A16FusedMoeCompileResult( compiled=compiled, size_m=size_m, hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, top_k=top_k, activation=activation, apply_router_weight_on_input=bool(apply_router_weight_on_input), zero_fc2_output=bool(zero_fc2_output), element_dtype=element_dtype, fast_math=bool(fast_math), swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, fc1_tile_n=fc1_tile_n, fc1_tile_k=fc1_tile_k, fc2_tile_n=fc2_tile_n, fc2_tile_k=fc2_tile_k, moe_block_size=moe_block_size, max_m_blocks=max_m_blocks, blocks_per_sm=kernel.blocks_per_sm, weight_layout=weight_layout, w13_layout=w13_layout, direct_topk_routes=kernel.direct_topk_routes, scale_format=scale_format, tc_decode_fused_sum=bool(tc_decode_fused_sum), collect_activation_amax=collect_activation_amax, ) _FUSED_CACHE[cache_key] = result return result def clear_w4a16_kernel_cache() -> None: _CACHE.clear() _FUSED_CACHE.clear() _ACTIVATION_CACHE.clear() _SUM_CACHE.clear() _SMALL_M_DIRECT_CACHE.clear() def compile_w4a16_activation( *, rows: int, intermediate_size: int, activation: str, element_dtype: str = "bf16", fast_math: bool = True, swiglu_limit: float | None = None, swiglu_alpha: float | None = None, swiglu_beta: float | None = None, ) -> W4A16ActivationCompileResult: cutlass_dtype = _cutlass_element_dtype(element_dtype) activation = normalize_moe_activation(activation) is_gated = validate_activation(activation) swiglu_limit, swiglu_alpha, swiglu_beta = _normalize_activation_swiglu_params( activation, swiglu_limit, swiglu_alpha, swiglu_beta, ) kernel = W4A16ActivationKernel( rows=rows, intermediate_size=intermediate_size, activation=activation, element_dtype=element_dtype, fast_math=fast_math, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, ) cache_key = ( "w4a16_activation", kernel.__cache_key__, ) cached = _ACTIVATION_CACHE.get(cache_key) if cached is not None: return replace(cached, rows=rows) w13_shards = 2 if is_gated else 1 compile_rows = _fake_m_for_specialization(rows) fc1_fake = cute.runtime.make_fake_compact_tensor( cutlass_dtype, (compile_rows * w13_shards * intermediate_size,), assumed_align=16, ) activated_fake = cute.runtime.make_fake_compact_tensor( cutlass_dtype, (compile_rows * intermediate_size,), assumed_align=16, ) raise_if_kernel_resolution_frozen( "cute.compile", target=kernel, cache_key=cache_key ) compiled = b12x_compile( kernel, fc1_fake, activated_fake, Int32(compile_rows), current_cuda_stream(), compile_spec=KernelCompileSpec.from_key( "moe.w4a16.activation", 1, cache_key, ), ) result = W4A16ActivationCompileResult( compiled=compiled, rows=rows, intermediate_size=intermediate_size, activation=activation, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, ) _ACTIVATION_CACHE[cache_key] = result return result def compile_w4a16_topk_sum( *, m: int, topk: int, hidden_size: int, element_dtype: str = "bf16", ) -> W4A16TopKSumCompileResult: cutlass_dtype = _cutlass_element_dtype(element_dtype) cache_key = ("w4a16_topk_sum", element_dtype, topk, hidden_size) cached = _SUM_CACHE.get(cache_key) if cached is not None: return cached fc2_fake = make_ptr(cutlass_dtype, 16, cute.AddressSpace.gmem, assumed_align=16) output_fake = make_ptr(cutlass_dtype, 16, cute.AddressSpace.gmem, assumed_align=16) kernel = W4A16TopKSumKernel( topk=topk, hidden_size=hidden_size, element_dtype=element_dtype, ) raise_if_kernel_resolution_frozen( "cute.compile", target=kernel, cache_key=cache_key ) compiled = b12x_compile( kernel, fc2_fake, output_fake, 1, current_cuda_stream(), compile_spec=KernelCompileSpec.from_key( "moe.w4a16.topk_sum", 1, cache_key, ), ) result = W4A16TopKSumCompileResult( compiled=compiled, m=0, topk=topk, hidden_size=hidden_size, ) _SUM_CACHE[cache_key] = result return result def _w4a16_small_m_direct_launch_flat( a_input: torch.Tensor, w13_u8: torch.Tensor, w13_scale_u8: torch.Tensor, w13_global_scale: torch.Tensor, w2_global_scale: torch.Tensor, inter_u32: torch.Tensor, w2_u8: torch.Tensor, w2_scale_u8: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, output: torch.Tensor, barrier_count: torch.Tensor, barrier_epoch: torch.Tensor, m: int, hidden_size: int, intermediate_size: int, num_experts: int, topk: int, activation: str, fast_math: bool, scale_format: str, has_swiglu_limit: bool, swiglu_limit_value: float, swiglu_alpha: float, swiglu_beta: float, w13_layout: str, stream_int: int, ) -> None: swiglu_limit = float(swiglu_limit_value) if has_swiglu_limit else None direct_launch = _compile_w4a16_small_m_direct( m=m, hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, topk=topk, activation=activation, fast_math=bool(fast_math), topk_ids_dtype=topk_ids.dtype, scale_format=scale_format, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, w13_layout=w13_layout, device=a_input.device, ) def ptr(dt, tensor: torch.Tensor): return make_ptr(dt, tensor.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) ids_dtype = cutlass.Int64 if topk_ids.dtype == torch.int64 else cutlass.Int32 direct_launch.compiled( ptr(cutlass.BFloat16, a_input), ptr(cutlass.Uint8, w13_u8), ptr(cutlass.Uint8, w13_scale_u8.view(torch.uint8)), ptr(cutlass.Float32, w13_global_scale), ptr(cutlass.Float32, w13_global_scale), ptr(cutlass.Float32, w2_global_scale), ptr(cutlass.Uint32, inter_u32.view(torch.uint32)), ptr(cutlass.Uint8, w2_u8), ptr(cutlass.Uint8, w2_scale_u8.view(torch.uint8)), ptr(cutlass.Float32, w2_global_scale), ptr(ids_dtype, topk_ids), ptr(cutlass.Float32, topk_weights), ptr(cutlass.BFloat16, output), barrier_count, barrier_epoch, Int32(m), Int32(direct_launch.grid_x), cuda.CUstream(stream_int), ) @torch.library.custom_op( "b12x::w4a16_small_m_direct_launch", mutates_args="unknown", ) def _w4a16_small_m_direct_launch_op( a_input: torch.Tensor, w13_u8: torch.Tensor, w13_scale_u8: torch.Tensor, w13_global_scale: torch.Tensor, w2_global_scale: torch.Tensor, inter_u32: torch.Tensor, w2_u8: torch.Tensor, w2_scale_u8: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, output: torch.Tensor, barrier_count: torch.Tensor, barrier_epoch: torch.Tensor, m: int, hidden_size: int, intermediate_size: int, num_experts: int, topk: int, activation: str, fast_math: bool, scale_format: str, has_swiglu_limit: bool, swiglu_limit_value: float, swiglu_alpha: float, swiglu_beta: float, w13_layout: str, stream_int: int, ) -> None: _w4a16_small_m_direct_launch_flat( a_input=a_input, w13_u8=w13_u8, w13_scale_u8=w13_scale_u8, w13_global_scale=w13_global_scale, w2_global_scale=w2_global_scale, inter_u32=inter_u32, w2_u8=w2_u8, w2_scale_u8=w2_scale_u8, topk_ids=topk_ids, topk_weights=topk_weights, output=output, barrier_count=barrier_count, barrier_epoch=barrier_epoch, m=m, hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, topk=topk, activation=activation, fast_math=fast_math, scale_format=scale_format, has_swiglu_limit=has_swiglu_limit, swiglu_limit_value=swiglu_limit_value, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, w13_layout=w13_layout, stream_int=stream_int, ) @_w4a16_small_m_direct_launch_op.register_fake def _w4a16_small_m_direct_launch_fake( a_input: torch.Tensor, w13_u8: torch.Tensor, w13_scale_u8: torch.Tensor, w13_global_scale: torch.Tensor, w2_global_scale: torch.Tensor, inter_u32: torch.Tensor, w2_u8: torch.Tensor, w2_scale_u8: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor, output: torch.Tensor, barrier_count: torch.Tensor, barrier_epoch: torch.Tensor, m: int, hidden_size: int, intermediate_size: int, num_experts: int, topk: int, activation: str, fast_math: bool, scale_format: str, has_swiglu_limit: bool, swiglu_limit_value: float, swiglu_alpha: float, swiglu_beta: float, w13_layout: str, stream_int: int, ) -> None: return None def _w4a16_fused_moe_launch_flat( a_input: torch.Tensor, w13_arg: torch.Tensor, w2_arg: torch.Tensor, fc1_out: torch.Tensor, activated: torch.Tensor, fc2_out: torch.Tensor, w13_scale_i32: torch.Tensor, w2_scale_i32: torch.Tensor, w13_global_scale: torch.Tensor, w2_global_scale: torch.Tensor, packed_route_indices: torch.Tensor, block_expert_ids: torch.Tensor, packed_route_count: torch.Tensor, activation_amax: torch.Tensor | None, layer_idx: int, topk_weights: torch.Tensor, fc1_scratch: torch.Tensor, fc2_scratch: torch.Tensor, workspace: torch.Tensor, m: int, size_m: int, hidden_size: int, intermediate_size: int, num_experts: int, topk: int, activation: str, apply_router_weight_on_input: bool, zero_fc2_output: bool, moe_block_size: int, max_m_blocks: int, element_dtype: str, fast_math: bool, sms: int, max_shared_mem: int, has_swiglu_limit: bool, swiglu_limit_value: float, swiglu_alpha: float, swiglu_beta: float, weight_layout: str, scale_format: str, w13_layout: str, direct_topk_routes: bool, tc_decode_fused_sum: bool, collect_activation_amax: bool, stream_int: int, ) -> None: swiglu_limit = float(swiglu_limit_value) if has_swiglu_limit else None collect_activation_amax = bool(collect_activation_amax) if collect_activation_amax and activation_amax is None: raise ValueError("activation_amax is required for calibrated W4A16 launch") activation_amax_arg = ( activation_amax.view(-1) if activation_amax is not None else w13_global_scale ) fused = compile_w4a16_fused_moe( size_m=size_m, hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, top_k=topk, activation=activation, apply_router_weight_on_input=bool(apply_router_weight_on_input), zero_fc2_output=bool(zero_fc2_output), moe_block_size=moe_block_size, max_m_blocks=max_m_blocks, element_dtype=element_dtype, fast_math=bool(fast_math), sms=sms, max_shared_mem=max_shared_mem, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, weight_layout=weight_layout, scale_format=scale_format, w13_layout=w13_layout, direct_topk_routes=bool(direct_topk_routes), tc_decode_fused_sum=bool(tc_decode_fused_sum), collect_activation_amax=collect_activation_amax, ) fused.compiled( make_ptr( _cutlass_element_dtype(element_dtype), a_input.data_ptr(), cute.AddressSpace.gmem, assumed_align=16, ), w13_arg, w2_arg, fc1_out, activated, fc2_out, w13_scale_i32, w2_scale_i32, w13_global_scale, w2_global_scale, packed_route_indices, block_expert_ids, packed_route_count, activation_amax_arg, int(layer_idx), make_ptr( cutlass.Float32, topk_weights.data_ptr(), cute.AddressSpace.gmem, assumed_align=4, ), fc1_scratch, fc2_scratch, workspace, m, _w4a16_fused_persistent_grid_x( fused=fused, m=m, topk=topk, intermediate_size=intermediate_size, activation=activation, direct_topk_routes=bool(direct_topk_routes), sms=sms, ), cuda.CUstream(stream_int), ) def _w4a16_fused_persistent_grid_x( *, fused: W4A16FusedMoeCompileResult, m: int, topk: int, intermediate_size: int, activation: str, direct_topk_routes: bool, sms: int, ) -> int: """Right-size the persistent grid for the fused FC1+FC2 launch. The default over-subscribes the cooperative grid (sms*blocks_per_sm CTAs) and forces the larger GEMM into the cross-CTA split-K tail (lock-serialized finalize + c_tmp gmem round-trip) whenever its mn-tile count exceeds the grid. For the small-M direct-topk decode the host knows the exact FC1 mn-tile count, so pick the fewest full persistent waves that fit the co-residency cap and set grid_x to that wave's tile count. Then every CTA owns whole FC1 (and FC2, whose tile count is an integer multiple) mn-tiles over the full K -- no split-K reduction, no lock traffic -- with fewer grid-barrier participants, while staying <= the cap so the cooperative barrier never deadlocks. Falls back to the default for the route-pack path where the host cannot know route_blocks ahead of the launch. """ cap = int(sms) * int(fused.blocks_per_sm) if not direct_topk_routes or m <= 0: return max(cap, 1) is_gated = is_gated_moe_activation(activation) fc1_cols = int(intermediate_size) * (2 if is_gated else 1) fc1_tile_n = int(getattr(fused, "fc1_tile_n", 0)) if fc1_tile_n <= 0 or fc1_cols % fc1_tile_n != 0: return max(cap, 1) n_tiles = fc1_cols // fc1_tile_n route_blocks = int(m) * int(topk) fc1_mn_tiles = route_blocks * n_tiles if fc1_mn_tiles <= 0 or cap <= 0: return max(cap, 1) waves = (fc1_mn_tiles + cap - 1) // cap if waves <= 0: return max(cap, 1) # Only commit to the right-sized grid when every wave is a whole tile cover # (no remainder => no split-K tail); otherwise keep the safe default. if fc1_mn_tiles % waves != 0: return max(cap, 1) grid_x = fc1_mn_tiles // waves if grid_x < 1 or grid_x > cap: return max(cap, 1) return grid_x @torch.library.custom_op( "b12x::w4a16_fused_moe_launch", mutates_args="unknown", ) def _w4a16_fused_moe_launch_op( a_input: torch.Tensor, w13_arg: torch.Tensor, w2_arg: torch.Tensor, fc1_out: torch.Tensor, activated: torch.Tensor, fc2_out: torch.Tensor, w13_scale_i32: torch.Tensor, w2_scale_i32: torch.Tensor, w13_global_scale: torch.Tensor, w2_global_scale: torch.Tensor, packed_route_indices: torch.Tensor, block_expert_ids: torch.Tensor, packed_route_count: torch.Tensor, topk_weights: torch.Tensor, fc1_scratch: torch.Tensor, fc2_scratch: torch.Tensor, workspace: torch.Tensor, m: int, size_m: int, hidden_size: int, intermediate_size: int, num_experts: int, topk: int, activation: str, apply_router_weight_on_input: bool, zero_fc2_output: bool, moe_block_size: int, max_m_blocks: int, element_dtype: str, fast_math: bool, sms: int, max_shared_mem: int, has_swiglu_limit: bool, swiglu_limit_value: float, swiglu_alpha: float, swiglu_beta: float, weight_layout: str, scale_format: str, w13_layout: str, direct_topk_routes: bool, tc_decode_fused_sum: bool, stream_int: int, ) -> None: _w4a16_fused_moe_launch_flat( a_input=a_input, w13_arg=w13_arg, w2_arg=w2_arg, fc1_out=fc1_out, activated=activated, fc2_out=fc2_out, w13_scale_i32=w13_scale_i32, w2_scale_i32=w2_scale_i32, w13_global_scale=w13_global_scale, w2_global_scale=w2_global_scale, packed_route_indices=packed_route_indices, block_expert_ids=block_expert_ids, packed_route_count=packed_route_count, activation_amax=None, layer_idx=0, topk_weights=topk_weights, fc1_scratch=fc1_scratch, fc2_scratch=fc2_scratch, workspace=workspace, m=m, size_m=size_m, hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, topk=topk, activation=activation, apply_router_weight_on_input=apply_router_weight_on_input, zero_fc2_output=zero_fc2_output, moe_block_size=moe_block_size, max_m_blocks=max_m_blocks, element_dtype=element_dtype, fast_math=fast_math, sms=sms, max_shared_mem=max_shared_mem, has_swiglu_limit=has_swiglu_limit, swiglu_limit_value=swiglu_limit_value, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, weight_layout=weight_layout, scale_format=scale_format, w13_layout=w13_layout, direct_topk_routes=direct_topk_routes, tc_decode_fused_sum=tc_decode_fused_sum, collect_activation_amax=False, stream_int=stream_int, ) @_w4a16_fused_moe_launch_op.register_fake def _w4a16_fused_moe_launch_fake( a_input: torch.Tensor, w13_arg: torch.Tensor, w2_arg: torch.Tensor, fc1_out: torch.Tensor, activated: torch.Tensor, fc2_out: torch.Tensor, w13_scale_i32: torch.Tensor, w2_scale_i32: torch.Tensor, w13_global_scale: torch.Tensor, w2_global_scale: torch.Tensor, packed_route_indices: torch.Tensor, block_expert_ids: torch.Tensor, packed_route_count: torch.Tensor, topk_weights: torch.Tensor, fc1_scratch: torch.Tensor, fc2_scratch: torch.Tensor, workspace: torch.Tensor, m: int, size_m: int, hidden_size: int, intermediate_size: int, num_experts: int, topk: int, activation: str, apply_router_weight_on_input: bool, zero_fc2_output: bool, moe_block_size: int, max_m_blocks: int, element_dtype: str, fast_math: bool, sms: int, max_shared_mem: int, has_swiglu_limit: bool, swiglu_limit_value: float, swiglu_alpha: float, swiglu_beta: float, weight_layout: str, scale_format: str, w13_layout: str, direct_topk_routes: bool, tc_decode_fused_sum: bool, stream_int: int, ) -> None: return None @torch.library.custom_op( "b12x::w4a16_fused_moe_calibrated_launch", mutates_args="unknown", ) def _w4a16_fused_moe_calibrated_launch_op( a_input: torch.Tensor, w13_arg: torch.Tensor, w2_arg: torch.Tensor, fc1_out: torch.Tensor, activated: torch.Tensor, fc2_out: torch.Tensor, w13_scale_i32: torch.Tensor, w2_scale_i32: torch.Tensor, w13_global_scale: torch.Tensor, w2_global_scale: torch.Tensor, packed_route_indices: torch.Tensor, block_expert_ids: torch.Tensor, packed_route_count: torch.Tensor, activation_amax: torch.Tensor, layer_idx: int, topk_weights: torch.Tensor, fc1_scratch: torch.Tensor, fc2_scratch: torch.Tensor, workspace: torch.Tensor, m: int, size_m: int, hidden_size: int, intermediate_size: int, num_experts: int, topk: int, activation: str, apply_router_weight_on_input: bool, zero_fc2_output: bool, moe_block_size: int, max_m_blocks: int, element_dtype: str, fast_math: bool, sms: int, max_shared_mem: int, has_swiglu_limit: bool, swiglu_limit_value: float, swiglu_alpha: float, swiglu_beta: float, weight_layout: str, scale_format: str, w13_layout: str, stream_int: int, ) -> None: _w4a16_fused_moe_launch_flat( a_input=a_input, w13_arg=w13_arg, w2_arg=w2_arg, fc1_out=fc1_out, activated=activated, fc2_out=fc2_out, w13_scale_i32=w13_scale_i32, w2_scale_i32=w2_scale_i32, w13_global_scale=w13_global_scale, w2_global_scale=w2_global_scale, packed_route_indices=packed_route_indices, block_expert_ids=block_expert_ids, packed_route_count=packed_route_count, activation_amax=activation_amax, layer_idx=layer_idx, topk_weights=topk_weights, fc1_scratch=fc1_scratch, fc2_scratch=fc2_scratch, workspace=workspace, m=m, size_m=size_m, hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=num_experts, topk=topk, activation=activation, apply_router_weight_on_input=apply_router_weight_on_input, zero_fc2_output=zero_fc2_output, moe_block_size=moe_block_size, max_m_blocks=max_m_blocks, element_dtype=element_dtype, fast_math=fast_math, sms=sms, max_shared_mem=max_shared_mem, has_swiglu_limit=has_swiglu_limit, swiglu_limit_value=swiglu_limit_value, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, weight_layout=weight_layout, scale_format=scale_format, w13_layout=w13_layout, direct_topk_routes=False, tc_decode_fused_sum=False, collect_activation_amax=True, stream_int=stream_int, ) @_w4a16_fused_moe_calibrated_launch_op.register_fake def _w4a16_fused_moe_calibrated_launch_fake( a_input: torch.Tensor, w13_arg: torch.Tensor, w2_arg: torch.Tensor, fc1_out: torch.Tensor, activated: torch.Tensor, fc2_out: torch.Tensor, w13_scale_i32: torch.Tensor, w2_scale_i32: torch.Tensor, w13_global_scale: torch.Tensor, w2_global_scale: torch.Tensor, packed_route_indices: torch.Tensor, block_expert_ids: torch.Tensor, packed_route_count: torch.Tensor, activation_amax: torch.Tensor, layer_idx: int, topk_weights: torch.Tensor, fc1_scratch: torch.Tensor, fc2_scratch: torch.Tensor, workspace: torch.Tensor, m: int, size_m: int, hidden_size: int, intermediate_size: int, num_experts: int, topk: int, activation: str, apply_router_weight_on_input: bool, zero_fc2_output: bool, moe_block_size: int, max_m_blocks: int, element_dtype: str, fast_math: bool, sms: int, max_shared_mem: int, has_swiglu_limit: bool, swiglu_limit_value: float, swiglu_alpha: float, swiglu_beta: float, weight_layout: str, scale_format: str, w13_layout: str, stream_int: int, ) -> None: return None def _w4a16_topk_sum_launch_flat( fc2_out: torch.Tensor, output: torch.Tensor, m: int, topk: int, hidden_size: int, element_dtype: str, stream_int: int, ) -> None: sum_kernel = compile_w4a16_topk_sum( m=m, topk=topk, hidden_size=hidden_size, element_dtype=element_dtype, ) sum_kernel.compiled( make_ptr( _cutlass_element_dtype(element_dtype), fc2_out.data_ptr(), cute.AddressSpace.gmem, assumed_align=16, ), make_ptr( _cutlass_element_dtype(element_dtype), output.data_ptr(), cute.AddressSpace.gmem, assumed_align=16, ), m, cuda.CUstream(stream_int), ) @torch.library.custom_op( "b12x::w4a16_topk_sum_launch", mutates_args=("output",), ) def _w4a16_topk_sum_launch_op( fc2_out: torch.Tensor, output: torch.Tensor, m: int, topk: int, hidden_size: int, element_dtype: str, stream_int: int, ) -> None: _w4a16_topk_sum_launch_flat( fc2_out=fc2_out, output=output, m=m, topk=topk, hidden_size=hidden_size, element_dtype=element_dtype, stream_int=stream_int, ) @_w4a16_topk_sum_launch_op.register_fake def _w4a16_topk_sum_launch_fake( fc2_out: torch.Tensor, output: torch.Tensor, m: int, topk: int, hidden_size: int, element_dtype: str, stream_int: int, ) -> None: return None def _get_c_tmp( elements: int, *, device: torch.device, scratch: torch.Tensor | None = None, ) -> torch.Tensor: if scratch is not None: if scratch.dtype != torch.float32: raise TypeError("W4A16 c_tmp scratch buffers must be torch.float32") if scratch.device != device: raise ValueError(f"W4A16 c_tmp scratch buffers must be on {device}") if not scratch.is_contiguous(): raise ValueError("W4A16 c_tmp scratch buffers must be contiguous") if int(scratch.numel()) >= int(elements): return scratch[: int(elements)] if torch.cuda.is_current_stream_capturing(): raise RuntimeError( "W4A16 GEMM scratch is not initialized for CUDA graph capture; " "provide a preallocated fc*_c_tmp workspace with sufficient capacity" ) return torch.empty((elements,), dtype=torch.float32, device=device) def _validate_topk_ids( topk_ids: torch.Tensor, *, require_cuda: bool, require_contiguous: bool = True, ) -> None: if topk_ids.dtype not in (torch.int32, torch.int64): raise TypeError("topk_ids must be torch.int32 or torch.int64") if require_cuda and not topk_ids.is_cuda: raise ValueError("topk_ids must be a CUDA tensor") if require_contiguous and not topk_ids.is_contiguous(): raise ValueError("topk_ids must be contiguous") def _validate_expert_map( expert_map: torch.Tensor | None, *, device: torch.device | None = None, exact_num_experts: int | None = None, ) -> None: if expert_map is None: return if expert_map.dtype != torch.int32: raise TypeError("expert_map must be torch.int32") if device is None: if not expert_map.is_cuda: raise ValueError("expert_map must be a CUDA tensor") elif expert_map.device != device: raise ValueError("expert_map must be on the same device as a_input") if exact_num_experts is None: if expert_map.ndim != 1 or not expert_map.is_contiguous(): raise ValueError("expert_map must be a contiguous rank-1 tensor") return if not expert_map.is_contiguous(): raise ValueError("expert_map must be contiguous") if expert_map.ndim != 1 or int(expert_map.numel()) != int(exact_num_experts): raise ValueError( f"expert_map must have shape {(int(exact_num_experts),)}, got {tuple(expert_map.shape)}" ) def _validate_activation_amax( activation_amax: torch.Tensor | None, *, layer_idx: int | None, num_experts: int, device: torch.device, ) -> int | None: if activation_amax is None: if layer_idx is not None: raise ValueError("layer_idx requires activation_amax") return None if layer_idx is None: raise ValueError("layer_idx is required when activation_amax is provided") if activation_amax.dtype != torch.float32: raise TypeError("activation_amax must be torch.float32") if activation_amax.device != device: raise ValueError("activation_amax must be on the same device as a_input") if not activation_amax.is_cuda: raise ValueError("activation_amax must be a CUDA tensor") if not activation_amax.is_contiguous(): raise ValueError("activation_amax must be contiguous") if activation_amax.ndim != 3 or int(activation_amax.shape[2]) != 2: raise ValueError( "activation_amax must have shape [num_layers, num_experts, 2]" ) if int(activation_amax.shape[1]) < int(num_experts): raise ValueError( "activation_amax expert dimension is smaller than the local expert count" ) layer = int(layer_idx) if layer < 0 or layer >= int(activation_amax.shape[0]): raise ValueError( f"layer_idx {layer} is out of bounds for activation_amax with " f"{int(activation_amax.shape[0])} layers" ) return layer def _compile_w4a16_gemm_launch( *, size_m: int, size_n: int, size_k: int, num_experts: int, top_k: int, mul_topk_weights: bool, moe_block_size: int, max_m_blocks: int, element_dtype: str, packed_route_indices: torch.Tensor, sms: int, max_shared_mem: int, device: torch.device, c_tmp: torch.Tensor | None = None, scale_format: str = "e4m3_k16", ) -> _W4A16GemmLaunch: tile_k, tile_n, _, _ = _select_tile_config( problem_m=size_m, problem_n=size_n, problem_k=size_k, top_k=top_k, moe_block_size=moe_block_size, sms=sms, max_shared_mem=max_shared_mem, scale_format=scale_format, ) kernel = compile_w4a16_gemm( size_m=size_m, size_n=size_n, size_k=size_k, num_experts=num_experts, top_k=top_k, mul_topk_weights=mul_topk_weights, tile_n=tile_n, tile_k=tile_k, moe_block_size=moe_block_size, max_m_blocks=max_m_blocks, element_dtype=element_dtype, scale_format=scale_format, ) c_tmp = _get_c_tmp( packed_gemm_scratch_elements( size_n=size_n, route_slots=int(packed_route_indices.numel()), moe_block_size=moe_block_size, sms=sms, ), device=device, scratch=c_tmp, ) return _W4A16GemmLaunch(kernel=kernel, c_tmp=c_tmp) def pack_topk_routes_by_expert( topk_ids: torch.Tensor, block_size: int, num_experts: int, *, expert_map: torch.Tensor | None = None, packed_route_indices: torch.Tensor | None = None, block_expert_ids: torch.Tensor | None = None, packed_route_count: torch.Tensor | None = None, expert_offsets: torch.Tensor | None = None, stream: cuda.CUstream | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Group top-k routes by expert and pad each group to the GEMM M-block size.""" _validate_topk_ids(topk_ids, require_cuda=True) if block_size <= 0: raise ValueError(f"block_size must be positive, got {block_size}") if num_experts <= 0: raise ValueError(f"num_experts must be positive, got {num_experts}") _validate_expert_map(expert_map, exact_num_experts=int(num_experts)) del stream return _pack_topk_routes_by_expert( topk_ids, int(block_size), int(num_experts), expert_map=expert_map, packed_route_indices=packed_route_indices, block_expert_ids=block_expert_ids, packed_route_count=packed_route_count, expert_offsets=expert_offsets, ) def run_w4a16_moe( a_input: torch.Tensor, prepared, topk_weights: torch.Tensor, topk_ids: torch.Tensor, *, activation: str, intermediate_cache13: torch.Tensor, intermediate_cache2: torch.Tensor, output: torch.Tensor, fc1_c_tmp: torch.Tensor | None = None, fc2_c_tmp: torch.Tensor | None = None, packed_route_indices: torch.Tensor | None = None, block_expert_ids: torch.Tensor | None = None, packed_route_count: torch.Tensor | None = None, expert_offsets: torch.Tensor | None = None, expert_map: torch.Tensor | None = None, activation_amax: torch.Tensor | None = None, layer_idx: int | None = None, apply_router_weight_on_input: bool = False, fast_math: bool = True, swiglu_limit: float | None = None, swiglu_alpha: float | None = None, swiglu_beta: float | None = None, fused_launch: W4A16FusedMoeCompileResult | None = None, topk_sum_launch: W4A16TopKSumCompileResult | None = None, stream: cuda.CUstream | None = None, ) -> torch.Tensor: activation = normalize_moe_activation(activation) is_gated = validate_activation(activation) swiglu_limit, swiglu_alpha, swiglu_beta = _normalize_activation_swiglu_params( activation, swiglu_limit, swiglu_alpha, swiglu_beta, ) element_dtype = _normalize_element_dtype(a_input.dtype) if output.dtype != a_input.dtype: raise TypeError(f"output must have dtype {a_input.dtype}, got {output.dtype}") prepared_dtype = getattr(prepared, "params_dtype", a_input.dtype) if prepared_dtype != a_input.dtype: raise TypeError( f"prepared weights were built for {prepared_dtype}, but a_input has dtype {a_input.dtype}" ) weight_layout = getattr(prepared, "weight_layout", "packed") if weight_layout not in _WEIGHT_LAYOUTS: raise ValueError(f"unsupported W4A16 weight_layout {weight_layout!r}") scale_format = _normalize_scale_format( getattr(prepared, "scale_format", None) or ( "e8m0_k32" if getattr(prepared, "source_format", "") == "fp4_e8m0_k32" else "e4m3_k16" ) ) w13_layout = getattr( prepared, "w13_layout", "w13" if weight_layout == "modelopt" else "packed", ) if weight_layout == "modelopt": if w13_layout not in _MODEL_OPT_W13_LAYOUTS: raise ValueError(f"unsupported W4A16 w13_layout {w13_layout!r}") else: w13_layout = "packed" if topk_weights.dtype != torch.float32: raise TypeError("topk_weights must be torch.float32") _validate_topk_ids(topk_ids, require_cuda=False, require_contiguous=False) if ( not a_input.is_contiguous() or not topk_weights.is_contiguous() or not topk_ids.is_contiguous() ): raise ValueError("a_input, topk_weights, and topk_ids must be contiguous") _validate_expert_map(expert_map, device=a_input.device) m, hidden_size = a_input.shape topk = int(topk_ids.shape[1]) if tuple(topk_weights.shape) != (m, topk): raise ValueError(f"topk_weights must have shape {(m, topk)}") if int(prepared.hidden_size) != hidden_size: raise ValueError("prepared hidden_size does not match a_input") if bool(prepared.is_gated) != is_gated: raise ValueError("prepared weights do not match activation") if tuple(output.shape) != (m, hidden_size): raise ValueError(f"output must have shape {(m, hidden_size)}") if expert_map is not None and int(expert_map.numel()) < int(prepared.num_experts): raise ValueError("expert_map cannot be shorter than the local expert count") layer_idx_int = _validate_activation_amax( activation_amax, layer_idx=layer_idx, num_experts=int(prepared.num_experts), device=a_input.device, ) collect_activation_amax = activation_amax is not None route_num_experts = ( int(expert_map.numel()) if expert_map is not None else int(prepared.num_experts) ) if fused_launch is None: block_size_m = select_route_block_size_m(m, topk, route_num_experts) else: block_size_m = int(fused_launch.moe_block_size) if block_size_m not in _ALLOWED_ROUTED_SIZES: raise ValueError(f"unsupported W4A16 moe_block_size={block_size_m}") stream = current_cuda_stream() if stream is None else stream if (not collect_activation_amax) and _small_m_direct_supported( m=m, hidden_size=hidden_size, intermediate_size=int(prepared.intermediate_size), num_experts=int(prepared.num_experts), topk=topk, activation=activation, apply_router_weight_on_input=bool(apply_router_weight_on_input), swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, element_dtype=element_dtype, weight_layout=weight_layout, w13_layout=w13_layout, scale_format=scale_format, expert_map=expert_map, ): if topk_ids.dtype not in (torch.int32, torch.int64): raise TypeError("W4A16 small-M direct path requires int32/int64 topk_ids") if not topk_ids.is_cuda: raise ValueError("W4A16 small-M direct path requires CUDA topk_ids") if not intermediate_cache2.is_contiguous() or not output.is_contiguous(): raise ValueError( "W4A16 small-M direct path requires contiguous intermediate_cache2 and output" ) if intermediate_cache2.dtype != a_input.dtype: raise TypeError(f"intermediate_cache2 must be {a_input.dtype}") if int(prepared.workspace.numel()) < 2: raise ValueError("prepared W4A16 workspace is too small for small-M direct") intermediate_size = int(prepared.intermediate_size) fc2_n_chunks = ((intermediate_size // 2) + 127) // 128 inter_u32_per_m = fc2_n_chunks * 128 * topk inter_u32 = intermediate_cache2.view(-1).view(torch.uint32) if int(inter_u32.numel()) < m * inter_u32_per_m: raise ValueError( "intermediate_cache2 is smaller than the W4A16 small-M direct scratch " f"requirement: have_u32={int(inter_u32.numel())}, " f"need_u32={m * inter_u32_per_m}" ) micro_w13_scale = getattr(prepared, "micro_w13_scale", None) micro_w2_scale = getattr(prepared, "micro_w2_scale", None) micro_w13_global = getattr(prepared, "micro_w13_global_scale", None) micro_w2_global = getattr(prepared, "micro_w2_global_scale", None) if ( micro_w13_scale is None or micro_w2_scale is None or micro_w13_global is None or micro_w2_global is None ): raise RuntimeError( "W4A16 small-M direct path requires prepared micro scale metadata" ) barrier_count = prepared.workspace[-2:-1] barrier_epoch = prepared.workspace[-1:] barrier_count.zero_() barrier_epoch.zero_() torch.ops.b12x.w4a16_small_m_direct_launch( a_input, prepared.w13.view(torch.uint8), micro_w13_scale, micro_w13_global, micro_w2_global, inter_u32[: m * inter_u32_per_m], prepared.w2.view(torch.uint8), micro_w2_scale, topk_ids, topk_weights, output, barrier_count, barrier_epoch, m, hidden_size, intermediate_size, int(prepared.num_experts), topk, activation, bool(fast_math), scale_format, swiglu_limit is not None, float(swiglu_limit or 0.0), float(swiglu_alpha), float(swiglu_beta), w13_layout, int(stream), ) return output # TC-decode: small-M packed decode that folds the top-k sum into the FC2 # store epilogue. Reuses the packed tensor-core MMA path; only the launch # scheduling/epilogue changes. Requires the packed object, bf16 gated # activation, int32 routes, no expert_map, and a runtime-compiled launch. # A preplanned launch built with the TC-decode fused-sum epilogue carries # ``tc_decode_fused_sum``; accept it through the binding path. A runtime # ``fused_launch is None`` (e.g. the standalone benchmark) compiles its own. preplanned_tc_decode = bool(getattr(fused_launch, "tc_decode_fused_sum", False)) use_tc_decode = bool( (not collect_activation_amax) and (fused_launch is None or preplanned_tc_decode) and weight_layout in ("packed", "nf3_2p1") and expert_map is None and is_gated and element_dtype == "bf16" and topk_ids.dtype in (torch.int32, torch.int64) and topk_ids.is_cuda and int(m) <= _TC_DECODE_MAX_M ) if use_tc_decode and topk_ids.dtype != torch.int32: # The inline direct-topk route path needs int32 route indices. topk_ids = topk_ids.to(torch.int32) direct_topk_eligible = ( (not collect_activation_amax) and (m <= _MAX_DIRECT_TOPK_ROUTE_M or use_tc_decode) and weight_layout in ("packed", "nf3_2p1") and expert_map is None ) use_direct_topk_routes = bool( direct_topk_eligible and topk_ids.dtype == torch.int32 and ( fused_launch is None or bool(getattr(fused_launch, "direct_topk_routes", False)) ) ) if ( bool(getattr(fused_launch, "direct_topk_routes", False)) and not use_direct_topk_routes ): raise RuntimeError( "preplanned W4A16 direct top-k routing requires small-M packed " "int32 topk_ids without expert_map" ) # TC-decode requires the inline direct-topk route path (no route-pack). use_tc_decode = bool(use_tc_decode and use_direct_topk_routes) # A preplanned TC-decode launch atomically accumulates FC2 partials into the # (pre-zeroed) output and emits no separate top-k sum. If it was selected but # the decode preconditions don't hold, running it would corrupt the output, # so fail loudly instead. if preplanned_tc_decode and not use_tc_decode: raise RuntimeError( "preplanned TC-decode W4A16 launch requires small-M packed bf16 " f"decode (m <= {_TC_DECODE_MAX_M}, cuda int32/int64 topk_ids, " "no expert_map)" ) route_slots_for_scratch = int(m) * int(topk) * int(block_size_m) required_m_blocks = int(m) * int(topk) if use_direct_topk_routes else 0 if fused_launch is not None and not use_direct_topk_routes: route_slots_capacity = max_packed_route_slots( int(fused_launch.size_m) * int(topk), int(block_size_m), route_num_experts, ) route_blocks_capacity = (route_slots_capacity + int(block_size_m) - 1) // int( block_size_m ) if packed_route_indices is not None: if int(packed_route_indices.numel()) < route_slots_capacity: raise ValueError( "packed_route_indices is smaller than the selected W4A16 launch capacity" ) packed_route_indices = packed_route_indices[:route_slots_capacity] if block_expert_ids is not None: if int(block_expert_ids.numel()) < route_blocks_capacity: raise ValueError( "block_expert_ids is smaller than the selected W4A16 launch capacity" ) block_expert_ids = block_expert_ids[:route_blocks_capacity] if use_direct_topk_routes: packed_route_indices = topk_ids.view(-1) if block_expert_ids is None: block_expert_ids = packed_route_indices if packed_route_count is None: packed_route_count = packed_route_indices else: packed_route_indices, block_expert_ids, packed_route_count = ( pack_topk_routes_by_expert( topk_ids, block_size_m, route_num_experts, expert_map=expert_map, packed_route_indices=packed_route_indices, block_expert_ids=block_expert_ids, packed_route_count=packed_route_count, expert_offsets=expert_offsets, stream=stream, ) ) route_slots_for_scratch = int(packed_route_indices.numel()) required_m_blocks = int(block_expert_ids.numel()) props = torch.cuda.get_device_properties(a_input.device) sms = int(props.multi_processor_count) max_shared_mem = int( getattr(props, "shared_memory_per_block_optin", _DEFAULT_MAX_SHARED_MEM) ) buffer_plan = plan_w4a16_buffers( prepared, m=m, topk=topk, route_num_experts=route_num_experts, sms=sms, ) intermediate_size = int(prepared.intermediate_size) fc1_cols = buffer_plan.fc1_cols if intermediate_cache13.numel() < buffer_plan.intermediate_cache13_elements: raise ValueError( f"intermediate_cache13 has {intermediate_cache13.numel()} elements; " f"need at least {buffer_plan.intermediate_cache13_elements}" ) if intermediate_cache2.numel() < buffer_plan.intermediate_cache2_elements: raise ValueError( f"intermediate_cache2 has {intermediate_cache2.numel()} elements; " f"need at least {buffer_plan.intermediate_cache2_elements}" ) if ( intermediate_cache13.dtype != a_input.dtype or intermediate_cache2.dtype != a_input.dtype ): raise TypeError(f"intermediate caches must be {a_input.dtype}") if ( not intermediate_cache13.is_contiguous() or not intermediate_cache2.is_contiguous() or not output.is_contiguous() ): raise ValueError("intermediate caches and output must be contiguous") intermediate_cache13_flat = intermediate_cache13.view(-1) intermediate_cache2_flat = intermediate_cache2.view(-1) if int(prepared.workspace.numel()) < sms * 4 + 2: raise ValueError("prepared W4A16 workspace is too small for fused FC1+FC2") if fused_launch is None: fused = compile_w4a16_fused_moe( size_m=m, hidden_size=hidden_size, intermediate_size=intermediate_size, num_experts=int(prepared.num_experts), top_k=topk, activation=activation, apply_router_weight_on_input=bool(apply_router_weight_on_input), zero_fc2_output=expert_map is not None, moe_block_size=block_size_m, max_m_blocks=int(required_m_blocks), element_dtype=element_dtype, fast_math=bool(fast_math), sms=sms, max_shared_mem=max_shared_mem, swiglu_limit=swiglu_limit, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, weight_layout=weight_layout, scale_format=scale_format, w13_layout=w13_layout, direct_topk_routes=use_direct_topk_routes, tc_decode_fused_sum=use_tc_decode, collect_activation_amax=collect_activation_amax, ) else: if int(fused_launch.size_m) < m: raise RuntimeError( "preplanned W4A16 fused MoE launch capacity is smaller than requested rows: " f"requested={m}, planned={int(fused_launch.size_m)}" ) expected_fused = ( hidden_size, intermediate_size, int(prepared.num_experts), topk, activation, bool(apply_router_weight_on_input), expert_map is not None, element_dtype, bool(fast_math), swiglu_limit, float(swiglu_alpha), float(swiglu_beta), weight_layout, scale_format, w13_layout, bool(use_direct_topk_routes), bool(collect_activation_amax), block_size_m, ) actual_fused = ( int(fused_launch.hidden_size), int(fused_launch.intermediate_size), int(fused_launch.num_experts), int(fused_launch.top_k), fused_launch.activation, bool(fused_launch.apply_router_weight_on_input), bool(fused_launch.zero_fc2_output), fused_launch.element_dtype, bool(fused_launch.fast_math), fused_launch.swiglu_limit, float(fused_launch.swiglu_alpha), float(fused_launch.swiglu_beta), getattr(fused_launch, "weight_layout", "packed"), getattr(fused_launch, "scale_format", "e4m3_k16"), getattr( fused_launch, "w13_layout", "w13" if getattr(fused_launch, "weight_layout", "packed") == "modelopt" else "packed", ), bool(getattr(fused_launch, "direct_topk_routes", False)), bool(getattr(fused_launch, "collect_activation_amax", False)), int(fused_launch.moe_block_size), ) if actual_fused != expected_fused or int(fused_launch.max_m_blocks) < int( required_m_blocks ): raise RuntimeError( "preplanned W4A16 fused MoE launch does not match requested contract: " f"requested={expected_fused + (int(required_m_blocks),)}, " f"planned={actual_fused + (int(fused_launch.max_m_blocks),)}" ) fused = fused_launch capacity_m = int(fused.size_m) capacity_routed_rows = capacity_m * topk if intermediate_cache13_flat.numel() < capacity_routed_rows * max(fc1_cols, hidden_size): raise ValueError( "intermediate_cache13 is smaller than the selected W4A16 launch capacity: " f"capacity_rows={capacity_m}, topk={topk}" ) if intermediate_cache2_flat.numel() < capacity_routed_rows * intermediate_size: raise ValueError( "intermediate_cache2 is smaller than the selected W4A16 launch capacity: " f"capacity_rows={capacity_m}, topk={topk}" ) fc1_out = intermediate_cache13_flat[: capacity_routed_rows * fc1_cols] activated = intermediate_cache2_flat[: capacity_routed_rows * intermediate_size] if use_tc_decode: # FC2 atomically accumulates per-route partials directly into the # per-token output, so the output is the FC2 store target and must be # pre-zeroed. The fused tc_decode kernel now zeroes the output in its # own prologue (before FC1, made visible by the existing post-FC1 grid # barrier), so the separate host-side output.zero_() launch is removed # from the decode critical path here. This drops the separate top-k-sum # launch as well. fc2_out = output.view(-1) else: fc2_out = intermediate_cache13_flat[: capacity_routed_rows * hidden_size] fc1_scratch = _get_c_tmp( packed_gemm_scratch_elements( size_n=fc1_cols, route_slots=int(route_slots_for_scratch), moe_block_size=block_size_m, sms=sms, ), device=a_input.device, scratch=fc1_c_tmp, ) fc2_scratch = _get_c_tmp( packed_gemm_scratch_elements( size_n=hidden_size, route_slots=int(route_slots_for_scratch), moe_block_size=block_size_m, sms=sms, ), device=a_input.device, scratch=fc2_c_tmp, ) if weight_layout == "modelopt": w13_arg = prepared.w13.view(torch.uint8).view(-1) w2_arg = prepared.w2.view(torch.uint8).view(-1) else: w13_arg = prepared.w13.view(torch.int32).view(-1) w2_arg = prepared.w2.view(torch.int32).view(-1) launch_common = ( a_input, w13_arg, w2_arg, fc1_out, activated, fc2_out, prepared.w13_scale.view(torch.uint8).view(torch.int32).view(-1), prepared.w2_scale.view(torch.uint8).view(torch.int32).view(-1), prepared.w13_global_scale, prepared.w2_global_scale, packed_route_indices, block_expert_ids, packed_route_count, ) launch_tail = ( topk_weights, fc1_scratch, fc2_scratch, prepared.workspace, m, capacity_m, hidden_size, intermediate_size, int(prepared.num_experts), topk, activation, bool(apply_router_weight_on_input), expert_map is not None, block_size_m, int(fused.max_m_blocks), element_dtype, bool(fast_math), sms, max_shared_mem, swiglu_limit is not None, float(swiglu_limit or 0.0), float(swiglu_alpha), float(swiglu_beta), weight_layout, scale_format, w13_layout, ) if collect_activation_amax: assert activation_amax is not None assert layer_idx_int is not None torch.ops.b12x.w4a16_fused_moe_calibrated_launch( *launch_common, activation_amax, int(layer_idx_int), *launch_tail, int(stream), ) else: torch.ops.b12x.w4a16_fused_moe_launch( *launch_common, *launch_tail, bool(use_direct_topk_routes), bool(use_tc_decode), int(stream), ) if use_tc_decode: # FC2 already wrote the top-k-summed result into `output`. return output if topk_sum_launch is not None: expected_sum = (topk, hidden_size) actual_sum = ( int(topk_sum_launch.topk), int(topk_sum_launch.hidden_size), ) if actual_sum != expected_sum: raise RuntimeError( "preplanned W4A16 top-k sum launch does not match requested contract: " f"requested={expected_sum}, planned={actual_sum}" ) torch.ops.b12x.w4a16_topk_sum_launch( fc2_out, output, m, topk, hidden_size, element_dtype, int(stream), ) return output __all__ = [ "W4A16ActivationCompileResult", "W4A16FusedMoeCompileResult", "W4A16GemmCompileResult", "W4A16TopKSumCompileResult", "W4A16FusedMoeKernel", "W4A16ActivationKernel", "W4A16GemmKernel", "W4A16TopKSumKernel", "clear_w4a16_kernel_cache", "compile_w4a16_activation", "compile_w4a16_fused_moe", "compile_w4a16_gemm", "compile_w4a16_topk_sum", "pack_topk_routes_by_expert", "run_w4a16_moe", ]