diff --git a/benchmarks/routines/moe.py b/benchmarks/routines/moe.py index 091f3480d38..e0dae2a5f0c 100644 --- a/benchmarks/routines/moe.py +++ b/benchmarks/routines/moe.py @@ -1711,16 +1711,21 @@ def testB12xFusedMoe(args): f"intermediate={intermediate_size}, experts={num_experts}, top_k={top_k}" ) - # b12x supports SwiGLU (gated) and ReLU2 (non-gated) + # b12x supports SwiGLU / GeGLU (gated) and ReLU2 (non-gated) activation_type = args.activation_type - _ACT_STR = {ActivationType.Swiglu: "silu", ActivationType.Relu2: "relu2"} + _ACT_STR = { + ActivationType.Swiglu: "silu", + ActivationType.Geglu: "gelu_tanh", + ActivationType.GegluTanh: "gelu_tanh", + ActivationType.Relu2: "relu2", + } if activation_type not in _ACT_STR: raise ValueError( - f"b12x_fused_moe only supports Swiglu and Relu2 activations, " - f"got {activation_type.name}" + f"b12x_fused_moe only supports Swiglu, Geglu, GegluTanh, and Relu2 " + f"activations, got {activation_type.name}" ) activation_str = _ACT_STR[activation_type] - is_gated = activation_type == ActivationType.Swiglu + is_gated = activation_type.is_gated # Create b12x-specific NVFP4 test data (weights quantized, input stays bf16) tensors = _create_nvfp4_moe_test_data( diff --git a/flashinfer/cute_dsl/fp4_common.py b/flashinfer/cute_dsl/fp4_common.py index 76697963c42..e7264b67c72 100644 --- a/flashinfer/cute_dsl/fp4_common.py +++ b/flashinfer/cute_dsl/fp4_common.py @@ -222,6 +222,22 @@ def ld_global_nc_v4_u32( ) +@dsl_user_op +def prefetch_global_l2(base_ptr: Int64, *, loc=None, ip=None) -> None: + """Prefetch a global memory line into L2.""" + llvm.inline_asm( + None, + [Int64(base_ptr).ir_value(loc=loc, ip=ip)], + "prefetch.global.L2 [$0];", + "l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + @dsl_user_op def st_global_u64(base_ptr: Int64, value: Uint64, *, loc=None, ip=None): """Store 64 bits to global memory.""" @@ -257,15 +273,18 @@ def st_global_u32(base_ptr: Int64, value: Uint32, *, loc=None, ip=None): @dsl_user_op -def get_ptr_as_int64(tensor: cute.Tensor, offset: Int32, *, loc=None, ip=None) -> Int64: +def get_ptr_as_int64(tensor: cute.Tensor, offset, *, loc=None, ip=None) -> Int64: """Get the memory address of tensor[offset] as Int64. + ``offset`` may be Int32 or Int64 and is added to the iterator unchanged, + so 64-bit offsets are not truncated. + WARNING: This uses ptrtoint which strips address space information. For SMEM tensors, the resulting Int64 is a raw SMEM offset that does NOT work with generic-addressing loads (ld.v4.u32). Use only with explicit address-space loads (ld.global.*) or for global memory tensors. """ - elem_ptr = tensor.iterator + Int32(offset) + elem_ptr = tensor.iterator + offset ptr_int = llvm.ptrtoint(T.i64(), elem_ptr.llvm_ptr, loc=loc, ip=ip) return Int64(ptr_int) @@ -1133,6 +1152,150 @@ def ue8m0_to_output_scale(ue8m0_val: Uint32, *, loc=None, ip=None) -> Float32: ) +@dsl_user_op +def pow2_ceil_ue8m0( + scale: Float32, + *, + loc=None, + ip=None, +) -> Tuple[Float32, Uint32]: + """Round a positive FP32 ``scale`` up to a power of two, bit-exactly. + + Returns ``(rounded_fp32, ue8m0_byte)`` with ``ue8m0_byte`` the biased + exponent of the rounded value. Unlike ``cvt_f32_to_ue8m0``, this matches + the fp8_quant.cuh and scale_convert.cuh fp32_to_ue8m0 references + bit-for-bit. + """ + result = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.i32()]), + [Float32(scale).ir_value(loc=loc, ip=ip)], + """ + { + .reg .pred p_mant; + .reg .b32 bits, mant; + + // bits = __float_as_uint(scale) + mov.b32 bits, $2; + // mant = bits & 0x007FFFFF (mantissa field) + and.b32 mant, bits, 8388607; + setp.ne.u32 p_mant, mant, 0; + // if (mant) bits = (bits + 0x00800000) & 0x7F800000 + @p_mant add.u32 bits, bits, 8388608; + @p_mant and.b32 bits, bits, 2139095040; + // rounded fp32 scale = __uint_as_float(bits) + mov.b32 $0, bits; + // ue8m0 = (bits >> 23) & 0xFF + bfe.u32 $1, bits, 23, 8; + } + """, + "=f,=r,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + rounded = llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip) + ue8m0 = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip) + return Float32(rounded), Uint32(ue8m0) + + +@dsl_user_op +def cvt_e8m0_to_f32(e8m0_val: Uint32, *, loc=None, ip=None) -> Float32: + """Convert a single E8M0 scale byte to its true f32 value 2**(byte-127). + + Byte 0 maps to 0.0. + """ + return Float32( + llvm.inline_asm( + T.f32(), + [Uint32(e8m0_val).ir_value(loc=loc, ip=ip)], + """ + { + .reg .pred p0; + .reg .u32 b0; + .reg .s32 e0; + .reg .f32 ef0; + and.b32 b0, $1, 0x000000ff; + setp.eq.u32 p0, b0, 0; + cvt.s32.u32 e0, b0; + sub.s32 e0, e0, 127; + cvt.rn.f32.s32 ef0, e0; + ex2.approx.f32 $0, ef0; + selp.f32 $0, 0f00000000, $0, p0; + } + """, + "=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def cvt_e8m0x4_to_f32x4( + packed: Uint32, *, loc=None, ip=None +) -> Tuple[Float32, Float32, Float32, Float32]: + """Decode 4 E8M0 scale bytes (packed u32) to 4 x true f32 (2**(byte-127)).""" + result = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.f32(), T.f32(), T.f32()]), + [Uint32(packed).ir_value(loc=loc, ip=ip)], + """ + { + .reg .pred p0, p1, p2, p3; + .reg .u32 b0, b1, b2, b3; + .reg .s32 e0, e1, e2, e3; + .reg .f32 ef0, ef1, ef2, ef3; + and.b32 b0, $4, 0x000000ff; + shr.u32 b1, $4, 8; + and.b32 b1, b1, 0x000000ff; + shr.u32 b2, $4, 16; + and.b32 b2, b2, 0x000000ff; + shr.u32 b3, $4, 24; + setp.eq.u32 p0, b0, 0; + setp.eq.u32 p1, b1, 0; + setp.eq.u32 p2, b2, 0; + setp.eq.u32 p3, b3, 0; + cvt.s32.u32 e0, b0; + cvt.s32.u32 e1, b1; + cvt.s32.u32 e2, b2; + cvt.s32.u32 e3, b3; + sub.s32 e0, e0, 127; + sub.s32 e1, e1, 127; + sub.s32 e2, e2, 127; + sub.s32 e3, e3, 127; + cvt.rn.f32.s32 ef0, e0; + cvt.rn.f32.s32 ef1, e1; + cvt.rn.f32.s32 ef2, e2; + cvt.rn.f32.s32 ef3, e3; + ex2.approx.f32 $0, ef0; + ex2.approx.f32 $1, ef1; + ex2.approx.f32 $2, ef2; + ex2.approx.f32 $3, ef3; + selp.f32 $0, 0f00000000, $0, p0; + selp.f32 $1, 0f00000000, $1, p1; + selp.f32 $2, 0f00000000, $2, p2; + selp.f32 $3, 0f00000000, $3, p3; + } + """, + "=f,=f,=f,=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [2], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [3], loc=loc, ip=ip)), + ) + + # ============================================================================= # E2M1 Conversion # ============================================================================= @@ -2043,6 +2206,38 @@ def cvt_e4m3_to_f32_via_f16(fp8_val: Uint32, *, loc=None, ip=None) -> Float32: ) +@dsl_user_op +def cvt_w4a16_packed_e4m3_scale_to_f32( + packed_byte: Uint32, *, loc=None, ip=None +) -> Float32: + """Scalar mirror of packed_dequant_e4m3x4_to_bfloat2x2 for W4A16 scales.""" + return Float32( + llvm.inline_asm( + T.f32(), + [Uint32(packed_byte).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b32 bits, tmp; + .reg .b16 bf; + and.b32 bits, $1, 0x80; + shl.b32 bits, bits, 7; + and.b32 tmp, $1, 0x7f; + shl.b32 tmp, tmp, 4; + or.b32 bits, bits, tmp; + cvt.u16.u32 bf, bits; + cvt.f32.bf16 $0, bf; + } + """, + "=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + @dsl_user_op def cvt_s0e5m3_to_f16x2_broadcast(fp8_val: Uint32, *, loc=None, ip=None) -> Uint32: """Convert one S0E5M3 scale byte to an f16x2 with the scale in both lanes. @@ -2168,6 +2363,35 @@ def quant_dequant_2( return f0 * sf_f32, f1 * sf_f32 +@cute.jit +def mx_scale_from_amax32(amax: Float32) -> Tuple[Float32, Float32]: + """Per-32 MXFP8 block scale from a precomputed amax. + + Returns ``(scale, inv_scale)`` with ``scale = pow2_ceil(amax/448)`` + (zero amax yields (0, 0), silencing the block). + """ + rounded, byte = pow2_ceil_ue8m0(amax * Float32(1.0 / FLOAT8_E4M3_MAX)) + inv_scale = ue8m0_to_output_scale(byte) + return rounded, inv_scale + + +@cute.jit +def quant_dequant_e4m3_2( + v0: Float32, + v1: Float32, + inv_scale: Float32, + scale: Float32, +) -> Tuple[Float32, Float32]: + """Quantize-dequantize a pair through E4M3 with a power-of-two block scale. + + Keeps a8_mx decode numerics matched to the w4a8 prefill activation + quantizer. + """ + q0 = fp8_e4m3_to_f32(cvt_f32_to_e4m3(v0 * inv_scale)) * scale + q1 = fp8_e4m3_to_f32(cvt_f32_to_e4m3(v1 * inv_scale)) * scale + return q0, q1 + + @dsl_user_op def fp4_dot4_sum( u_packed: Uint32, diff --git a/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/__init__.py b/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/__init__.py index 5c4dcd9d601..fe7b4c47f0b 100644 --- a/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/__init__.py +++ b/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/__init__.py @@ -2,6 +2,7 @@ from .moe_static_kernel import MoEStaticKernel from .moe_micro_kernel import MoEMicroKernel +from .moe_direct_micro_kernel import MoEDirectMicroKernel from .moe_dynamic_kernel import MoEDynamicKernel from .moe_dispatch import ( Sm120StaticMoEWorkspace, @@ -9,6 +10,7 @@ allocate_sm120_moe_workspace, allocate_sm120_static_workspace, allocate_sm120_dynamic_workspace, + clear_sm120_moe_caches, launch_sm120_static_moe, launch_sm120_dynamic_moe, launch_sm120_moe, @@ -18,12 +20,14 @@ __all__ = [ "MoEStaticKernel", "MoEMicroKernel", + "MoEDirectMicroKernel", "MoEDynamicKernel", "Sm120StaticMoEWorkspace", "Sm120DynamicMoEWorkspace", "allocate_sm120_moe_workspace", "allocate_sm120_static_workspace", "allocate_sm120_dynamic_workspace", + "clear_sm120_moe_caches", "launch_sm120_static_moe", "launch_sm120_dynamic_moe", "launch_sm120_moe", diff --git a/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_activation.py b/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_activation.py index 66629c79f1f..25ac981a3a9 100644 --- a/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_activation.py +++ b/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_activation.py @@ -1,7 +1,9 @@ -"""Activation helper functions for the b12x fused-MoE kernels.""" +"""Activation helper functions and metadata for the b12x fused-MoE kernels.""" from __future__ import annotations +import math + import cutlass import cutlass.cute as cute from cutlass import Float32 @@ -9,8 +11,87 @@ from flashinfer.cute_dsl.fp4_common import fmax_f32, fmin_f32 +SWIGLUOAI_UNINTERLEAVE = "swigluoai_uninterleave" +SWIGLUOAI_DEFAULT_LIMIT = 7.0 +SWIGLUOAI_DEFAULT_ALPHA = 1.702 +SWIGLUOAI_DEFAULT_BETA = 1.0 + +SUPPORTED_MOE_ACTIVATIONS = frozenset( + {"silu", "gelu_tanh", "relu2", SWIGLUOAI_UNINTERLEAVE} +) +GATED_MOE_ACTIVATIONS = frozenset({"silu", "gelu_tanh", SWIGLUOAI_UNINTERLEAVE}) + + +def normalize_moe_activation(activation: str) -> str: + activation = str(activation) + if activation not in SUPPORTED_MOE_ACTIVATIONS: + raise ValueError(f"unsupported activation {activation!r}") + return activation + + +def is_gated_moe_activation(activation: str) -> bool: + return normalize_moe_activation(activation) in GATED_MOE_ACTIVATIONS + + +def normalize_swiglu_limit_for_activation( + activation: str, + swiglu_limit: float | None, +) -> float | None: + activation = normalize_moe_activation(activation) + if swiglu_limit is None: + return SWIGLUOAI_DEFAULT_LIMIT if activation == SWIGLUOAI_UNINTERLEAVE else None + if activation not in GATED_MOE_ACTIVATIONS: + raise ValueError("swiglu_limit requires a gated MoE activation") + limit = float(swiglu_limit) + if not math.isfinite(limit) or limit <= 0.0: + raise ValueError(f"swiglu_limit must be positive and finite, got {limit}") + return limit + + +def normalize_swiglu_alpha_for_activation( + activation: str, + swiglu_alpha: float | None, +) -> float: + activation = normalize_moe_activation(activation) + if activation != SWIGLUOAI_UNINTERLEAVE: + # Accept the module default so callers can pass it unconditionally. + if swiglu_alpha is not None and float(swiglu_alpha) not in ( + 1.0, + SWIGLUOAI_DEFAULT_ALPHA, + ): + raise ValueError( + "swiglu_alpha is only configurable for swigluoai_uninterleave" + ) + return 1.0 + alpha = SWIGLUOAI_DEFAULT_ALPHA if swiglu_alpha is None else float(swiglu_alpha) + if not math.isfinite(alpha): + raise ValueError(f"swiglu_alpha must be finite, got {alpha}") + return alpha + + +def normalize_swiglu_beta_for_activation( + activation: str, + swiglu_beta: float | None, +) -> float: + activation = normalize_moe_activation(activation) + if activation != SWIGLUOAI_UNINTERLEAVE: + # Accept the module default so callers can pass it unconditionally. + if swiglu_beta is not None and float(swiglu_beta) not in ( + 0.0, + SWIGLUOAI_DEFAULT_BETA, + ): + raise ValueError( + "swiglu_beta is only configurable for swigluoai_uninterleave" + ) + return 0.0 + beta = SWIGLUOAI_DEFAULT_BETA if swiglu_beta is None else float(swiglu_beta) + if not math.isfinite(beta): + raise ValueError(f"swiglu_beta must be finite, got {beta}") + return beta + + def is_gated_activation(activation: str) -> bool: - return activation in ("silu", "gelu_tanh", "swigluoai_uninterleave") + return activation in GATED_MOE_ACTIVATIONS def gated_activation_f32( diff --git a/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py b/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py index 691a50825b3..c8d1c9cffd6 100644 --- a/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py +++ b/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_direct_micro_kernel.py @@ -1,19 +1,15 @@ -""" -MoEDirectMicroKernel - direct routed NVFP4/W4A16 micro MoE kernel for SM120/SM121. - -Ported from the b12x kernel library to FlashInfer. +"""MoEDirectMicroKernel: direct-routed NVFP4 MoE decode kernel for SM12x. -This direct micro backend is the low-latency path for very small routed decode -batches. Unlike the compact static/dynamic kernels, it performs the FC1 and FC2 -work with warp-level direct dot products instead of staging full MMA tiles. -The W4A16 micro kernel subclasses this implementation with BF16 intermediate -storage enabled. +Unlike MoEMicroKernel (Triton route pre-pass, expert-major packed A), this +kernel consumes raw per-token topk_ids/topk_weights with no routing +pre-pass and computes both GEMMs as software fp4 dot products on CUDA +cores, targeting tiny decode batches. """ from __future__ import annotations from dataclasses import dataclass -from typing import Tuple, cast +from typing import Tuple import cutlass import cutlass.cute as cute @@ -24,25 +20,30 @@ from flashinfer.cute_dsl.utils import ( current_cuda_stream, - get_max_active_clusters, get_num_sm, + get_max_active_clusters, + make_ptr, ) from flashinfer.cute_dsl.fp4_common import ( atomic_add_global_i32, cvt_e4m3_to_f32_via_f16, cvt_e4m3x4_to_f32x4, + cvt_e8m0_to_f32, + cvt_e8m0x4_to_f32x4, cvt_f32_to_e4m3, + cvt_w4a16_packed_e4m3_scale_to_f32, fmax_f32, - fp4_dot4_sum, fp4_dot4_sum_f32acc, - fp4_dot8_sum, fp4_dot8_sum_f32acc, get_ptr_as_int64, ld_global_acquire_i32, ld_global_nc_u32, ld_global_nc_v4_u32, + mx_scale_from_amax32, nvfp4_scale_from_amax, + quant_dequant_e4m3_2, pack_f32x2_to_f16x2, + prefetch_global_l2, quant_dequant_2, spin_wait_global_eq_i32, st_global_i32, @@ -50,6 +51,14 @@ threadfence, warp_reduce, ) +from .moe_activation 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, +) _BLOCK_SIZE = 16 @@ -59,6 +68,7 @@ _BLOCK_DIM = _NUM_WARPS * 32 _K_PER_CTA = 16 _MAX_DIRECT_K_SEGMENTS = 12 +_W4A16_PACKED_E4M3_SCALE_CACHE_VERSION = 1 def _direct_k_segments_supported(k_segments: int) -> bool: @@ -69,6 +79,10 @@ def _align_up(value: int, align: int) -> int: return ((int(value) + int(align) - 1) // int(align)) * int(align) +def _direct_k_segments_for_k(k: int) -> int: + return _align_up(k, 32 * _BLOCK_SIZE) // (32 * _BLOCK_SIZE) + + def _fc1_chunks_for_m(m: int, n: int) -> int: rows_per_warp = max(1, int(m)) rows_per_chunk = max(_BLOCK_SIZE, _NUM_WARPS * rows_per_warp) @@ -96,7 +110,9 @@ class _ShapeConfig: w1_sf_cols: int w2_sf_rows: int w2_sf_cols: int + k_blocks: int k_segments: int + k_segments_aligned: bool fc1_chunks: int fc1_chunks_per_block: int i_chunk: int @@ -124,7 +140,9 @@ def _make_shape_config( w1_sf_cols = _align_up(k // _BLOCK_SIZE, 4) w2_sf_rows = _align_up(k, 128) w2_sf_cols = _align_up(n // _BLOCK_SIZE, 4) - k_segments = k // (32 * _BLOCK_SIZE) + k_blocks = k // _BLOCK_SIZE + k_segments = _direct_k_segments_for_k(k) + k_segments_aligned = k_blocks == k_segments * 32 fc1_chunks = _fc1_chunks_for_m(m, n) fc1_chunks_per_block = 16 # chunks per 128-wide swizzle block i_chunk = n // fc1_chunks @@ -146,7 +164,9 @@ def _make_shape_config( w1_sf_cols=w1_sf_cols, w2_sf_rows=w2_sf_rows, w2_sf_cols=w2_sf_cols, + k_blocks=k_blocks, k_segments=k_segments, + k_segments_aligned=k_segments_aligned, fc1_chunks=fc1_chunks, fc1_chunks_per_block=fc1_chunks_per_block, i_chunk=i_chunk, @@ -179,7 +199,9 @@ def _remake_shape_config_fc1(cfg: _ShapeConfig, fc1_chunks: int) -> _ShapeConfig w1_sf_cols=cfg.w1_sf_cols, w2_sf_rows=cfg.w2_sf_rows, w2_sf_cols=cfg.w2_sf_cols, + k_blocks=cfg.k_blocks, k_segments=cfg.k_segments, + k_segments_aligned=cfg.k_segments_aligned, fc1_chunks=fc1_chunks, fc1_chunks_per_block=cfg.fc1_chunks_per_block, i_chunk=i_chunk, @@ -192,46 +214,6 @@ def _remake_shape_config_fc1(cfg: _ShapeConfig, fc1_chunks: int) -> _ShapeConfig ) -@cute.jit -def _block_dot_hfma2( - u_a: Uint32, - u_b: Uint32, - smem_xh: cute.Tensor, - xh_base: Int32, -) -> Float32: - xh0 = Uint32(smem_xh[xh_base + Int32(0)]) - xh1 = Uint32(smem_xh[xh_base + Int32(1)]) - xh2 = Uint32(smem_xh[xh_base + Int32(2)]) - xh3 = Uint32(smem_xh[xh_base + Int32(3)]) - xh4 = Uint32(smem_xh[xh_base + Int32(4)]) - xh5 = Uint32(smem_xh[xh_base + Int32(5)]) - xh6 = Uint32(smem_xh[xh_base + Int32(6)]) - xh7 = Uint32(smem_xh[xh_base + Int32(7)]) - return fp4_dot8_sum(u_a, u_b, xh0, xh1, xh2, xh3, xh4, xh5, xh6, xh7) - - -@cute.jit -def _block_dot_hfma2_pair( - up_a: Uint32, - up_b: Uint32, - gate_a: Uint32, - gate_b: Uint32, - smem_xh: cute.Tensor, - xh_base: Int32, -) -> Tuple[Float32, Float32]: - xh0 = Uint32(smem_xh[xh_base + Int32(0)]) - xh1 = Uint32(smem_xh[xh_base + Int32(1)]) - xh2 = Uint32(smem_xh[xh_base + Int32(2)]) - xh3 = Uint32(smem_xh[xh_base + Int32(3)]) - xh4 = Uint32(smem_xh[xh_base + Int32(4)]) - xh5 = Uint32(smem_xh[xh_base + Int32(5)]) - xh6 = Uint32(smem_xh[xh_base + Int32(6)]) - xh7 = Uint32(smem_xh[xh_base + Int32(7)]) - up = fp4_dot8_sum(up_a, up_b, xh0, xh1, xh2, xh3, xh4, xh5, xh6, xh7) - gate = fp4_dot8_sum(gate_a, gate_b, xh0, xh1, xh2, xh3, xh4, xh5, xh6, xh7) - return up, gate - - @cute.jit def _block_dot_hfma2_f32acc( u_a: Uint32, @@ -272,35 +254,6 @@ def _block_dot_hfma2_pair_f32acc( return up, gate -@cute.jit -def _block_dot4( - u_val: Uint32, - smem_xh: cute.Tensor, - xh_base: Int32, -) -> Float32: - xh0 = Uint32(smem_xh[xh_base + Int32(0)]) - xh1 = Uint32(smem_xh[xh_base + Int32(1)]) - xh2 = Uint32(smem_xh[xh_base + Int32(2)]) - xh3 = Uint32(smem_xh[xh_base + Int32(3)]) - return fp4_dot4_sum(u_val, xh0, xh1, xh2, xh3) - - -@cute.jit -def _block_dot4_pair( - up_val: Uint32, - gate_val: Uint32, - smem_xh: cute.Tensor, - xh_base: Int32, -) -> Tuple[Float32, Float32]: - xh0 = Uint32(smem_xh[xh_base + Int32(0)]) - xh1 = Uint32(smem_xh[xh_base + Int32(1)]) - xh2 = Uint32(smem_xh[xh_base + Int32(2)]) - xh3 = Uint32(smem_xh[xh_base + Int32(3)]) - up = fp4_dot4_sum(up_val, xh0, xh1, xh2, xh3) - gate = fp4_dot4_sum(gate_val, xh0, xh1, xh2, xh3) - return up, gate - - @cute.jit def _block_dot4_f32acc( u_val: Uint32, @@ -363,7 +316,12 @@ def _token_wait_fc1_ready( class MoEDirectMicroKernel: - """Decode-focused compact MoE kernel for SM120.""" + """Decode-focused direct-routed MoE kernel for SM12x. + + Scale contract: w1_alphas/input_gs/down_input_scale are per-expert + [weight_E] f32 tensors, with input_gs and down_input_scale in multiplier + form; reciprocal-form scales must be inverted host-side before launch. + """ def __init__( self, @@ -377,24 +335,98 @@ def __init__( share_expert_scales: bool = False, single_token: bool = False, dynamic_down_scale: bool = False, + compile_time_phase: int = 0, w4a16_mode: bool = False, + a8_mx_mode: bool = False, + scale_format: str = "e4m3_k16", + e8m0_scale_layout: str = "packed", + swiglu_limit: float | None = None, + swiglu_alpha: float | None = None, + swiglu_beta: float | None = None, + w13_layout: str = "w13", ): - if activation not in {"silu", "relu2"}: - raise ValueError(f"unsupported activation {activation!r}") + activation = normalize_moe_activation(activation) + if int(compile_time_phase) not in {0, 1, 2}: + raise ValueError(f"unsupported direct micro phase {compile_time_phase!r}") + if scale_format not in {"e4m3_k16", "e8m0_k32"}: + raise ValueError(f"unsupported micro scale_format {scale_format!r}") + if w4a16_mode and a8_mx_mode: + raise ValueError("w4a16_mode and a8_mx_mode are mutually exclusive") + if scale_format == "e8m0_k32" and not (w4a16_mode or a8_mx_mode): + raise ValueError("e8m0_k32 scales require the W4A16 or a8_mx micro mode") + if e8m0_scale_layout not in {"packed", "logical"}: + raise ValueError( + f"unsupported micro e8m0_scale_layout {e8m0_scale_layout!r}" + ) + swiglu_limit = normalize_swiglu_limit_for_activation(activation, swiglu_limit) + swiglu_alpha = normalize_swiglu_alpha_for_activation(activation, swiglu_alpha) + swiglu_beta = normalize_swiglu_beta_for_activation(activation, swiglu_beta) + if w13_layout not in {"w13", "w31"}: + raise ValueError(f"unsupported micro w13_layout {w13_layout!r}") + self.scale_format = scale_format + self.scale_format_e8m0_k32 = scale_format == "e8m0_k32" + self.e8m0_scale_layout = e8m0_scale_layout + self.e8m0_scale_layout_logical = e8m0_scale_layout == "logical" + self.w13_layout = w13_layout + # "w31" (gate_up) keeps the gate half first; "w13" (up_gate) keeps up first. + self.w13_gate_first = w13_layout == "w31" + self.has_swiglu_limit = swiglu_limit is not None + self.swiglu_limit = 0.0 if swiglu_limit is None else float(swiglu_limit) + self.swiglu_alpha = float(swiglu_alpha) + self.swiglu_beta = float(swiglu_beta) self.sf_vec_size = sf_vec_size - self.fast_math = fast_math + # Accepted for call compatibility with the MMA micro kernels; this + # CUDA-core body has no fast-math variant, so the flag is a no-op. + del fast_math self.activation = activation - self.is_gated = activation == "silu" + self.is_gated = is_gated_moe_activation(activation) + self.is_swigluoai = activation == SWIGLUOAI_UNINTERLEAVE + self.is_gelu_tanh = activation == "gelu_tanh" self.share_input_across_experts = share_input_across_experts self.share_expert_scales = share_expert_scales self.single_token = single_token self.dynamic_down_scale = dynamic_down_scale + self.compile_time_phase = int(compile_time_phase) self.w4a16_mode = w4a16_mode - self._cfg = cast(_ShapeConfig, None) + # a8_mx: quantize-dequantize activations through E4M3 with per-32 + # UE8M0 block scales (no global scale) so decode numerics track the + # w4a8 prefill recipe. Same f16 dot-product math and weights. + self.a8_mx_mode = a8_mx_mode + self._cfg: _ShapeConfig | None = None self.m_const = 0 self.m1_fc2_onepass = False + self.m1_fc2_rows_per_cta = _K_PER_CTA * 2 + self.launch_block_dim = _BLOCK_DIM self.grid_x = 0 + @property + def __cache_key__(self): + return ( + self.sf_vec_size, + self.activation, + self.is_gelu_tanh, + self.share_input_across_experts, + self.share_expert_scales, + self.single_token, + self.dynamic_down_scale, + self.compile_time_phase, + self.w4a16_mode, + self.a8_mx_mode, + self.scale_format, + self.e8m0_scale_layout, + self.w13_layout, + self.has_swiglu_limit, + self.swiglu_limit, + self.swiglu_alpha, + self.swiglu_beta, + _W4A16_PACKED_E4M3_SCALE_CACHE_VERSION, + self._cfg, + self.m_const, + self.m1_fc2_onepass, + self.m1_fc2_rows_per_cta, + self.launch_block_dim, + ) + @cute.jit def _fp4_dot4_for_math( self, @@ -404,10 +436,98 @@ def _fp4_dot4_for_math( x2: Uint32, x3: Uint32, ) -> Float32: - if cutlass.const_expr(self.fast_math): - return fp4_dot4_sum(u_packed, x0, x1, x2, x3) return fp4_dot4_sum_f32acc(u_packed, x0, x1, x2, x3) + @cute.jit + def _scale_byte_to_f32(self, byte: Uint32) -> Float32: + """Decode one block-scale byte to f32 (E8M0 in MXFP4 mode, else E4M3).""" + if cutlass.const_expr(self.scale_format_e8m0_k32): + return cvt_e8m0_to_f32(byte) + if cutlass.const_expr(self.w4a16_mode): + return cvt_w4a16_packed_e4m3_scale_to_f32(byte) + return cvt_e4m3_to_f32_via_f16(byte) + + @cute.jit + def _scale_word_to_f32x4( + self, word: Uint32 + ) -> Tuple[Float32, Float32, Float32, Float32]: + """Decode 4 packed block-scale bytes to f32x4 (E8M0 in MXFP4 mode).""" + if cutlass.const_expr(self.scale_format_e8m0_k32): + return cvt_e8m0x4_to_f32x4(word) + return cvt_e4m3x4_to_f32x4(word) + + @cute.jit + def _packed_scale_col(self, n: Int32) -> Int32: + """Column of output row n in the packed E8M0 [K/32, N] scale grid. + For N a multiple of 64 the Marlin permute reduces to a per-row column + permutation: (n & ~63) | ((n&7)<<3) | swap_bits01((n>>3)&7).""" + hi = (n >> Int32(3)) & Int32(7) + hi_sw = ( + (hi & Int32(4)) + | ((hi & Int32(1)) << Int32(1)) + | ((hi >> Int32(1)) & Int32(1)) + ) + return (n & ~Int32(63)) | ((n & Int32(7)) << Int32(3)) | hi_sw + + @cute.jit + def _ld_e8m0_scale( + self, + base_addr: Int64, + ebase: Int64, + kb32: Int32, + out_row: Int32, + n_cols: Int32, + k32_cols: Int32, + ) -> Float32: + """E8M0 scale from either packed [K/32, N] or logical [N, K/32].""" + if cutlass.const_expr(self.e8m0_scale_layout_logical): + addr = base_addr + ebase + Int64(out_row) * Int64(k32_cols) + Int64(kb32) + else: + col = self._packed_scale_col(out_row) + addr = base_addr + ebase + Int64(kb32) * Int64(n_cols) + Int64(col) + word = ld_global_nc_u32(addr & ~Int64(3)) + byte = (word >> Uint32((addr & Int64(3)) * Int64(8))) & Uint32(0xFF) + return cvt_e8m0_to_f32(byte) + + @cute.jit + def _packed_e4m3_scale_col(self, n: Int32) -> Int32: + """Column of row n in _permute_nvfp4_scales' packed E4M3 grid.""" + x = n & Int32(63) + perm_pos = ((x & Int32(7)) << Int32(3)) + (x >> Int32(3)) + rem = perm_pos & Int32(3) + rem_swapped = rem + if rem == Int32(1): + rem_swapped = Int32(2) + elif rem == Int32(2): + rem_swapped = Int32(1) + return (n & ~Int32(63)) + (perm_pos & ~Int32(3)) + rem_swapped + + @cute.jit + def _ld_e4m3_packed_scale( + self, + base_addr: Int64, + ebase: Int64, + kb16: Int32, + out_row: Int32, + n_cols: Int32, + ) -> Float32: + col = self._packed_e4m3_scale_col(out_row) + return self._ld_e4m3_packed_scale_col(base_addr, ebase, kb16, col, n_cols) + + @cute.jit + def _ld_e4m3_packed_scale_col( + self, + base_addr: Int64, + ebase: Int64, + kb16: Int32, + col: Int32, + n_cols: Int32, + ) -> Float32: + addr = base_addr + ebase + Int64(kb16) * Int64(n_cols) + Int64(col) + word = ld_global_nc_u32(addr & ~Int64(3)) + byte = (word >> Uint32((addr & Int64(3)) * Int64(8))) & Uint32(0xFF) + return cvt_w4a16_packed_e4m3_scale_to_f32(byte) + @cute.jit def _block_dot_hfma2_for_math( self, @@ -416,8 +536,6 @@ def _block_dot_hfma2_for_math( smem_xh: cute.Tensor, xh_base: Int32, ) -> Float32: - if cutlass.const_expr(self.fast_math): - return _block_dot_hfma2(u_a, u_b, smem_xh, xh_base) return _block_dot_hfma2_f32acc(u_a, u_b, smem_xh, xh_base) @cute.jit @@ -430,12 +548,33 @@ def _block_dot_hfma2_pair_for_math( smem_xh: cute.Tensor, xh_base: Int32, ) -> Tuple[Float32, Float32]: - if cutlass.const_expr(self.fast_math): - return _block_dot_hfma2_pair(up_a, up_b, gate_a, gate_b, smem_xh, xh_base) return _block_dot_hfma2_pair_f32acc( up_a, up_b, gate_a, gate_b, smem_xh, xh_base ) + @cute.jit + def _block_dot_hfma2_pair_regs_for_math( + self, + up_a: Uint32, + up_b: Uint32, + gate_a: Uint32, + gate_b: Uint32, + xh0: Uint32, + xh1: Uint32, + xh2: Uint32, + xh3: Uint32, + xh4: Uint32, + xh5: Uint32, + xh6: Uint32, + xh7: Uint32, + ) -> Tuple[Float32, Float32]: + """Paired up/gate fp4_dot8 over pre-loaded activation registers.""" + up = fp4_dot8_sum_f32acc(up_a, up_b, xh0, xh1, xh2, xh3, xh4, xh5, xh6, xh7) + gate = fp4_dot8_sum_f32acc( + gate_a, gate_b, xh0, xh1, xh2, xh3, xh4, xh5, xh6, xh7 + ) + return up, gate + @cute.jit def _block_dot4_for_math( self, @@ -443,8 +582,6 @@ def _block_dot4_for_math( smem_xh: cute.Tensor, xh_base: Int32, ) -> Float32: - if cutlass.const_expr(self.fast_math): - return _block_dot4(u_val, smem_xh, xh_base) return _block_dot4_f32acc(u_val, smem_xh, xh_base) @cute.jit @@ -455,8 +592,6 @@ def _block_dot4_pair_for_math( smem_xh: cute.Tensor, xh_base: Int32, ) -> Tuple[Float32, Float32]: - if cutlass.const_expr(self.fast_math): - return _block_dot4_pair(up_val, gate_val, smem_xh, xh_base) return _block_dot4_pair_f32acc(up_val, gate_val, smem_xh, xh_base) @classmethod @@ -468,11 +603,13 @@ def is_supported( num_topk: int, weight_E: int, ) -> bool: - if m not in (1, 2, 4, 8): + # The m tokens' activations stay resident; 8 is the register budget + # ceiling. FC1/FC2 are generic in m, so any 1 <= m <= 8 is correct. + if not (1 <= m <= 8): return False - if k <= 0 or k % (32 * _BLOCK_SIZE) != 0 or k % 128 != 0: + if k <= 0 or k % _BLOCK_SIZE != 0 or k % 128 != 0: return False - if k // _BLOCK_SIZE > 32 * _MAX_DIRECT_K_SEGMENTS: + if _direct_k_segments_for_k(k) > _MAX_DIRECT_K_SEGMENTS: return False if n <= 0 or n % _BLOCK_SIZE != 0: return False @@ -482,7 +619,7 @@ def is_supported( i_chunk = n // fc1_chunks if i_chunk % _BLOCK_SIZE != 0: return False - k_segments = k // (32 * _BLOCK_SIZE) + k_segments = _direct_k_segments_for_k(k) return ( _direct_k_segments_supported(k_segments) and 0 < num_topk <= 32 @@ -504,23 +641,53 @@ def configure( m=m, k=k, n=n, num_topk=num_topk, weight_E=weight_E, is_gated=self.is_gated ) num_fc1_chunks = _fc1_chunks_for_m(m, n) + if self.w4a16_mode and m == 1 and n <= 2048: + # 4 rows/warp only helps the k_segments==8 aligned gated path (its + # reg-hoist + dual-dot assume 4 rows). The k_segments==12 path is + # scale-load limited; 1 row/warp keeps those loads out of the row loop. + rows_per_warp_div = 2 + if cfg.k_segments_aligned and cfg.k_segments == 8 and self.is_gated: + rows_per_warp_div = 4 + elif cfg.k_segments_aligned and cfg.k_segments == 12 and self.is_gated: + rows_per_warp_div = 1 + num_fc1_chunks = min(num_fc1_chunks, n // (rows_per_warp_div * _BLOCK_SIZE)) if self.w4a16_mode and m > 1: - # Keep one FC1 row per warp for W4A16 multi-token decode so the - # direct kernel stays within the CUTLASS 4.5 launch resource limit. - num_fc1_chunks = max(num_fc1_chunks, n // _BLOCK_SIZE) + # Keep W4A16 multi-token FC1 chunks narrow enough to stay within + # the 512-thread launch register limit. + num_fc1_chunks = max(num_fc1_chunks, n // (_BLOCK_SIZE * 2)) + if self.a8_mx_mode: + # Per-32 self-ranging blocks: chunks must hold whole 32-blocks. + # The standalone FC1 phase never forms the FC2 per-32 activation + # scale, so it can keep the native 16-row tile (twice the grid). + a8_chunk_rows = 16 if self.compile_time_phase == 1 else 32 + a8_chunks = max(1, min(num_fc1_chunks, n // a8_chunk_rows)) + while a8_chunks > 1 and ( + n % a8_chunks != 0 or (n // a8_chunks) % a8_chunk_rows != 0 + ): + a8_chunks -= 1 + num_fc1_chunks = a8_chunks cfg = _remake_shape_config_fc1(cfg, num_fc1_chunks) fc1_tasks = m * cfg.num_topk * cfg.fc1_chunks w4a16_rowpair_fc2 = bool(self.w4a16_mode and m > 1 and cfg.fc2_n_chunks == 1) + m1_half_cta_fc2 = bool(self.compile_time_phase == 2 and m == 1) + m1_fc2_rows = _K_PER_CTA if m1_half_cta_fc2 else _K_PER_CTA * 2 if m == 1: - fc2_tasks = cfg.k_dim // (_K_PER_CTA * 2) + fc2_tasks = cfg.k_dim // m1_fc2_rows elif w4a16_rowpair_fc2: fc2_tasks = (m * cfg.k_dim) // (_K_PER_CTA * 2) else: fc2_tasks = (m * cfg.k_dim) // (_K_PER_CTA * 4) if max_active_ctas is None: max_active_ctas = min(get_num_sm(device), get_max_active_clusters(1)) - if m == 1 or m == 2: + if self.compile_time_phase == 1: + # A standalone FC1 phase has no cooperative-grid requirement. + grid_x = max(1, fc1_tasks) + elif self.compile_time_phase == 2: + # Likewise FC2 can expose every output-row task directly once + # FC1 has completed in a prior launch. + grid_x = max(1, fc2_tasks) + elif m in (1, 2): grid_x = max(1, min(int(max_active_ctas), max(fc1_tasks, fc2_tasks))) elif num_fc1_chunks < 16: grid_x = max(1, min(int(max_active_ctas), fc2_tasks)) @@ -528,9 +695,15 @@ def configure( grid_x = max(1, min(int(max_active_ctas), fc1_tasks, fc2_tasks)) m1_fc2_onepass = bool(m == 1 and grid_x >= fc2_tasks) + if self.a8_mx_mode and self.compile_time_phase != 1 and cfg.i_chunk % 32 != 0: + # The per-32 block scale reads the partner 16-block; the FC2 + # intermediate chunk must hold whole 32-blocks. + raise ValueError("a8_mx micro mode requires i_chunk % 32 == 0") self._cfg = cfg - self.m_const = m + self.m_const = m if m in (1, 9) else 0 self.m1_fc2_onepass = m1_fc2_onepass + self.m1_fc2_rows_per_cta = m1_fc2_rows + self.launch_block_dim = _K_PER_CTA * 16 if m1_half_cta_fc2 else _BLOCK_DIM self.grid_x = grid_x @cute.jit @@ -570,27 +743,43 @@ def _m1_fc2_rowpair_narrow( scatter_output: cute.Tensor, ): cfg = self._cfg - k_chunk_off = fc2_task * Int32(_K_PER_CTA * 2) + k_chunk_off = fc2_task * Int32(self.m1_fc2_rows_per_cta) k_row0 = k_chunk_off + warp_id * Int32(2) k_row1 = k_row0 + Int32(1) lane_byte_off = Int64(lane) * Int64(4) sf_cols = Int32(cfg.w2_sf_cols) + num_cb = sf_cols >> Int32(2) lane_cb = lane >> Int32(3) + w_valid = Int32(1) if lane_cb < num_cb else Int32(0) lane_mode_c = (lane >> Int32(1)) & Int32(3) bsf_byte_shift = lane_mode_c * Int32(8) out_acc0 = Float32(0.0) out_acc1 = Float32(0.0) + k_col0 = Int32(0) + k_col1 = Int32(0) + if cutlass.const_expr(self.w4a16_mode): + k_col0 = self._packed_e4m3_scale_col(k_row0) + k_col1 = self._packed_e4m3_scale_col(k_row1) for kk in cutlass.range_constexpr(cfg.num_topk): eid_addr = Int32(kk) eid = Int32(topk_ids[eid_addr]) router_w = topk_weights[eid_addr] - alpha_fc2 = w2_alphas[eid] - scale_lane = alpha_fc2 * router_w + if cutlass.const_expr( + self.w4a16_mode + and (not self.is_gated) + and cfg.k_dim == 2688 + and cfg.n == 1856 + ): + scale_lane = router_w + else: + alpha_fc2 = w2_alphas[eid] + scale_lane = alpha_fc2 * router_w ebase_w = Int64(eid) * Int64(cfg.k_dim * cfg.n_half) ebase_sf = Int64(eid) * Int64(cfg.w2_sf_rows * cfg.w2_sf_cols) + ebase_sf_packed_e4m3 = Int64(eid) * Int64((cfg.n // 16) * cfg.k_dim) row_rb0 = k_row0 >> Int32(7) row_mode_a0 = (k_row0 >> Int32(5)) & Int32(3) @@ -605,21 +794,62 @@ def _m1_fc2_rowpair_narrow( xh2 = Uint32(intermediate[kk_off + Int32(2 * 32) + lane]) xh3 = Uint32(intermediate[kk_off + Int32(3 * 32) + lane]) - u_packed0 = ld_global_nc_u32( - w2_base_addr - + ebase_w - + Int64(k_row0) * Int64(cfg.n_half) - + lane_byte_off - ) - bsf_off0 = ( - Int64(row_rb0) * Int64(sf_cols * 128) - + Int64(lane_cb) * Int64(512) - + Int64(row_mode_32_0) * Int64(16) - + Int64(row_mode_a0) * Int64(4) + u_packed0 = ( + ld_global_nc_u32( + w2_base_addr + + ebase_w + + Int64(k_row0) * Int64(cfg.n_half) + + lane_byte_off + ) + if w_valid > Int32(0) + else Uint32(0) ) - sf_word0 = ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off0) - bsf_byte0 = (sf_word0 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) - bsf_f0 = cvt_e4m3_to_f32_via_f16(bsf_byte0) + if cutlass.const_expr(self.scale_format_e8m0_k32): + ebase_w2p = Int64(eid) * Int64((cfg.n // 32) * cfg.k_dim) + kb32_i = lane >> Int32(2) + bsf_f0 = ( + self._ld_e8m0_scale( + w2s_base_addr, + ebase_w2p, + kb32_i, + k_row0, + Int32(cfg.k_dim), + Int32(cfg.n // 32), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + elif cutlass.const_expr(self.w4a16_mode): + kb16_i = lane >> Int32(1) + bsf_f0 = ( + self._ld_e4m3_packed_scale_col( + w2s_base_addr, + ebase_sf_packed_e4m3, + kb16_i, + k_col0, + Int32(cfg.k_dim), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + else: + bsf_off0 = ( + Int64(row_rb0) * Int64(sf_cols * 128) + + Int64(lane_cb) * Int64(512) + + Int64(row_mode_32_0) * Int64(16) + + Int64(row_mode_a0) * Int64(4) + ) + sf_word0 = ( + ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off0) + if w_valid > Int32(0) + else Uint32(0) + ) + bsf_byte0 = (sf_word0 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) + bsf_f0 = ( + self._scale_byte_to_f32(bsf_byte0) + if w_valid > Int32(0) + else Float32(0.0) + ) out_acc0 = ( out_acc0 + bsf_f0 @@ -627,21 +857,60 @@ def _m1_fc2_rowpair_narrow( * scale_lane ) - u_packed1 = ld_global_nc_u32( - w2_base_addr - + ebase_w - + Int64(k_row1) * Int64(cfg.n_half) - + lane_byte_off - ) - bsf_off1 = ( - Int64(row_rb1) * Int64(sf_cols * 128) - + Int64(lane_cb) * Int64(512) - + Int64(row_mode_32_1) * Int64(16) - + Int64(row_mode_a1) * Int64(4) + u_packed1 = ( + ld_global_nc_u32( + w2_base_addr + + ebase_w + + Int64(k_row1) * Int64(cfg.n_half) + + lane_byte_off + ) + if w_valid > Int32(0) + else Uint32(0) ) - sf_word1 = ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off1) - bsf_byte1 = (sf_word1 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) - bsf_f1 = cvt_e4m3_to_f32_via_f16(bsf_byte1) + if cutlass.const_expr(self.scale_format_e8m0_k32): + bsf_f1 = ( + self._ld_e8m0_scale( + w2s_base_addr, + ebase_w2p, + kb32_i, + k_row1, + Int32(cfg.k_dim), + Int32(cfg.n // 32), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + elif cutlass.const_expr(self.w4a16_mode): + kb16_i = lane >> Int32(1) + bsf_f1 = ( + self._ld_e4m3_packed_scale_col( + w2s_base_addr, + ebase_sf_packed_e4m3, + kb16_i, + k_col1, + Int32(cfg.k_dim), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + else: + bsf_off1 = ( + Int64(row_rb1) * Int64(sf_cols * 128) + + Int64(lane_cb) * Int64(512) + + Int64(row_mode_32_1) * Int64(16) + + Int64(row_mode_a1) * Int64(4) + ) + sf_word1 = ( + ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off1) + if w_valid > Int32(0) + else Uint32(0) + ) + bsf_byte1 = (sf_word1 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) + bsf_f1 = ( + self._scale_byte_to_f32(bsf_byte1) + if w_valid > Int32(0) + else Float32(0.0) + ) out_acc1 = ( out_acc1 + bsf_f1 @@ -670,7 +939,7 @@ def _m1_fc2_rowpair_wide( scatter_output: cute.Tensor, ): cfg = self._cfg - k_chunk_off = fc2_task * Int32(_K_PER_CTA * 2) + k_chunk_off = fc2_task * Int32(self.m1_fc2_rows_per_cta) k_row0 = k_chunk_off + warp_id * Int32(2) k_row1 = k_row0 + Int32(1) @@ -679,20 +948,35 @@ def _m1_fc2_rowpair_wide( sf_cols = Int32(cfg.w2_sf_cols) num_cb = sf_cols >> Int32(2) lane_cb = lane >> Int32(3) + w_valid = Int32(1) if lane_cb < num_cb else Int32(0) lane_mode_c = (lane >> Int32(1)) & Int32(3) bsf_byte_shift = lane_mode_c * Int32(8) out_acc0 = Float32(0.0) out_acc1 = Float32(0.0) + k_col0 = Int32(0) + k_col1 = Int32(0) + if cutlass.const_expr(self.w4a16_mode): + k_col0 = self._packed_e4m3_scale_col(k_row0) + k_col1 = self._packed_e4m3_scale_col(k_row1) for kk in cutlass.range_constexpr(cfg.num_topk): eid_addr = Int32(kk) eid = Int32(topk_ids[eid_addr]) router_w = topk_weights[eid_addr] - alpha_fc2 = w2_alphas[eid] - scale_lane = alpha_fc2 * router_w + if cutlass.const_expr( + self.w4a16_mode + and (not self.is_gated) + and cfg.k_dim == 2688 + and cfg.n == 1856 + ): + scale_lane = router_w + else: + alpha_fc2 = w2_alphas[eid] + scale_lane = alpha_fc2 * router_w ebase_w = Int64(eid) * Int64(cfg.k_dim * cfg.n_half) ebase_sf = Int64(eid) * Int64(cfg.w2_sf_rows * cfg.w2_sf_cols) + ebase_sf_packed_e4m3 = Int64(eid) * Int64((cfg.n // 16) * cfg.k_dim) row_rb0 = k_row0 >> Int32(7) row_mode_a0 = (k_row0 >> Int32(5)) & Int32(3) @@ -704,13 +988,37 @@ def _m1_fc2_rowpair_wide( for nc in cutlass.range_constexpr(cfg.fc2_n_chunks): chunk_base = Int32(nc) * Int32(128) kk_off = Int32(kk) * n_u32_per_expert + chunk_base - xh0 = Uint32(intermediate[kk_off + Int32(0 * 32) + lane]) - xh1 = Uint32(intermediate[kk_off + Int32(1 * 32) + lane]) - xh2 = Uint32(intermediate[kk_off + Int32(2 * 32) + lane]) - xh3 = Uint32(intermediate[kk_off + Int32(3 * 32) + lane]) - cb_idx = Int32(nc) * Int32(4) + lane_cb w_valid = Int32(1) if cb_idx < num_cb else Int32(0) + # A last 256-wide chunk overhanging a non-256-aligned n reads + # the uninitialized intermediate tail; the weight is masked to + # 0 but 0 * NaN = NaN, so mask the activation read too. + if cutlass.const_expr((cfg.w2_sf_cols >> 2) < cfg.fc2_n_chunks * 4): + xh0 = ( + Uint32(intermediate[kk_off + Int32(0 * 32) + lane]) + if w_valid > Int32(0) + else Uint32(0) + ) + xh1 = ( + Uint32(intermediate[kk_off + Int32(1 * 32) + lane]) + if w_valid > Int32(0) + else Uint32(0) + ) + xh2 = ( + Uint32(intermediate[kk_off + Int32(2 * 32) + lane]) + if w_valid > Int32(0) + else Uint32(0) + ) + xh3 = ( + Uint32(intermediate[kk_off + Int32(3 * 32) + lane]) + if w_valid > Int32(0) + else Uint32(0) + ) + else: + xh0 = Uint32(intermediate[kk_off + Int32(0 * 32) + lane]) + xh1 = Uint32(intermediate[kk_off + Int32(1 * 32) + lane]) + xh2 = Uint32(intermediate[kk_off + Int32(2 * 32) + lane]) + xh3 = Uint32(intermediate[kk_off + Int32(3 * 32) + lane]) u_packed0 = ( ld_global_nc_u32( w2_base_addr @@ -722,23 +1030,52 @@ def _m1_fc2_rowpair_wide( if w_valid > Int32(0) else Uint32(0) ) - bsf_off0 = ( - Int64(row_rb0) * Int64(sf_cols * 128) - + Int64(cb_idx) * Int64(512) - + Int64(row_mode_32_0) * Int64(16) - + Int64(row_mode_a0) * Int64(4) - ) - sf_word0 = ( - ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off0) - if w_valid > Int32(0) - else Uint32(0) - ) - bsf_byte0 = (sf_word0 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) - bsf_f0 = ( - cvt_e4m3_to_f32_via_f16(bsf_byte0) - if w_valid > Int32(0) - else Float32(0.0) - ) + if cutlass.const_expr(self.scale_format_e8m0_k32): + ebase_w2p = Int64(eid) * Int64((cfg.n // 32) * cfg.k_dim) + kb32_i = (chunk_base + lane * Int32(4)) >> Int32(4) + bsf_f0 = ( + self._ld_e8m0_scale( + w2s_base_addr, + ebase_w2p, + kb32_i, + k_row0, + Int32(cfg.k_dim), + Int32(cfg.n // 32), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + elif cutlass.const_expr(self.w4a16_mode): + kb16_i = (chunk_base + lane * Int32(4)) >> Int32(3) + bsf_f0 = ( + self._ld_e4m3_packed_scale_col( + w2s_base_addr, + ebase_sf_packed_e4m3, + kb16_i, + k_col0, + Int32(cfg.k_dim), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + else: + bsf_off0 = ( + Int64(row_rb0) * Int64(sf_cols * 128) + + Int64(cb_idx) * Int64(512) + + Int64(row_mode_32_0) * Int64(16) + + Int64(row_mode_a0) * Int64(4) + ) + sf_word0 = ( + ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off0) + if w_valid > Int32(0) + else Uint32(0) + ) + bsf_byte0 = (sf_word0 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) + bsf_f0 = ( + self._scale_byte_to_f32(bsf_byte0) + if w_valid > Int32(0) + else Float32(0.0) + ) out_acc0 = ( out_acc0 + bsf_f0 @@ -757,23 +1094,50 @@ def _m1_fc2_rowpair_wide( if w_valid > Int32(0) else Uint32(0) ) - bsf_off1 = ( - Int64(row_rb1) * Int64(sf_cols * 128) - + Int64(cb_idx) * Int64(512) - + Int64(row_mode_32_1) * Int64(16) - + Int64(row_mode_a1) * Int64(4) - ) - sf_word1 = ( - ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off1) - if w_valid > Int32(0) - else Uint32(0) - ) - bsf_byte1 = (sf_word1 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) - bsf_f1 = ( - cvt_e4m3_to_f32_via_f16(bsf_byte1) - if w_valid > Int32(0) - else Float32(0.0) - ) + if cutlass.const_expr(self.scale_format_e8m0_k32): + bsf_f1 = ( + self._ld_e8m0_scale( + w2s_base_addr, + ebase_w2p, + kb32_i, + k_row1, + Int32(cfg.k_dim), + Int32(cfg.n // 32), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + elif cutlass.const_expr(self.w4a16_mode): + kb16_i = (chunk_base + lane * Int32(4)) >> Int32(3) + bsf_f1 = ( + self._ld_e4m3_packed_scale_col( + w2s_base_addr, + ebase_sf_packed_e4m3, + kb16_i, + k_col1, + Int32(cfg.k_dim), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + else: + bsf_off1 = ( + Int64(row_rb1) * Int64(sf_cols * 128) + + Int64(cb_idx) * Int64(512) + + Int64(row_mode_32_1) * Int64(16) + + Int64(row_mode_a1) * Int64(4) + ) + sf_word1 = ( + ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off1) + if w_valid > Int32(0) + else Uint32(0) + ) + bsf_byte1 = (sf_word1 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) + bsf_f1 = ( + self._scale_byte_to_f32(bsf_byte1) + if w_valid > Int32(0) + else Float32(0.0) + ) out_acc1 = ( out_acc1 + bsf_f1 @@ -788,7 +1152,7 @@ def _m1_fc2_rowpair_wide( scatter_output[k_row1] = BFloat16(sum_warp1) @cute.jit - def _m2_fc2_rowpair_narrow( + def _m2_fc2_rowquad_narrow( self, fc2_task: Int32, warp_id: Int32, @@ -802,21 +1166,27 @@ def _m2_fc2_rowpair_narrow( scatter_output: cute.Tensor, ): cfg = self._cfg - rows_per_cta = Int32(_K_PER_CTA * 2) + rows_per_cta = Int32(_K_PER_CTA * 4) linear_row_base = fc2_task * rows_per_cta t = linear_row_base // Int32(cfg.k_dim) k_chunk_off = linear_row_base - t * Int32(cfg.k_dim) - k_row0 = k_chunk_off + warp_id * Int32(2) + k_row0 = k_chunk_off + warp_id * Int32(4) k_row1 = k_row0 + Int32(1) + k_row2 = k_row0 + Int32(2) + k_row3 = k_row0 + Int32(3) lane_byte_off = Int64(lane) * Int64(4) token_inter_base = t * Int32(cfg.inter_u32) sf_cols = Int32(cfg.w2_sf_cols) + num_cb = sf_cols >> Int32(2) lane_cb = lane >> Int32(3) + w_valid = Int32(1) if lane_cb < num_cb else Int32(0) lane_mode_c = (lane >> Int32(1)) & Int32(3) bsf_byte_shift = lane_mode_c * Int32(8) out_acc0 = Float32(0.0) out_acc1 = Float32(0.0) + out_acc2 = Float32(0.0) + out_acc3 = Float32(0.0) row_rb0 = k_row0 >> Int32(7) row_mode_a0 = (k_row0 >> Int32(5)) & Int32(3) @@ -824,13 +1194,27 @@ def _m2_fc2_rowpair_narrow( row_rb1 = k_row1 >> Int32(7) row_mode_a1 = (k_row1 >> Int32(5)) & Int32(3) row_mode_32_1 = k_row1 & Int32(31) + row_rb2 = k_row2 >> Int32(7) + row_mode_a2 = (k_row2 >> Int32(5)) & Int32(3) + row_mode_32_2 = k_row2 & Int32(31) + row_rb3 = k_row3 >> Int32(7) + row_mode_a3 = (k_row3 >> Int32(5)) & Int32(3) + row_mode_32_3 = k_row3 & Int32(31) for kk in cutlass.range_constexpr(cfg.num_topk): eid_addr = t * Int32(cfg.num_topk) + Int32(kk) eid = Int32(topk_ids[eid_addr]) router_w = topk_weights[eid_addr] - alpha_fc2 = w2_alphas[eid] - scale_lane = alpha_fc2 * router_w + if cutlass.const_expr( + self.w4a16_mode + and (not self.is_gated) + and cfg.k_dim == 2688 + and cfg.n == 1856 + ): + scale_lane = router_w + else: + alpha_fc2 = w2_alphas[eid] + scale_lane = alpha_fc2 * router_w ebase_w = Int64(eid) * Int64(cfg.k_dim * cfg.n_half) ebase_sf = Int64(eid) * Int64(cfg.w2_sf_rows * cfg.w2_sf_cols) @@ -841,11 +1225,15 @@ def _m2_fc2_rowpair_narrow( xh2 = Uint32(intermediate[kk_off + Int32(2 * 32) + lane]) xh3 = Uint32(intermediate[kk_off + Int32(3 * 32) + lane]) - u_packed0 = ld_global_nc_u32( - w2_base_addr - + ebase_w - + Int64(k_row0) * Int64(cfg.n_half) - + lane_byte_off + u_packed0 = ( + ld_global_nc_u32( + w2_base_addr + + ebase_w + + Int64(k_row0) * Int64(cfg.n_half) + + lane_byte_off + ) + if w_valid > Int32(0) + else Uint32(0) ) bsf_off0 = ( Int64(row_rb0) * Int64(sf_cols * 128) @@ -853,9 +1241,17 @@ def _m2_fc2_rowpair_narrow( + Int64(row_mode_32_0) * Int64(16) + Int64(row_mode_a0) * Int64(4) ) - sf_word0 = ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off0) + sf_word0 = ( + ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off0) + if w_valid > Int32(0) + else Uint32(0) + ) bsf_byte0 = (sf_word0 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) - bsf_f0 = cvt_e4m3_to_f32_via_f16(bsf_byte0) + bsf_f0 = ( + self._scale_byte_to_f32(bsf_byte0) + if w_valid > Int32(0) + else Float32(0.0) + ) out_acc0 = ( out_acc0 + bsf_f0 @@ -863,11 +1259,15 @@ def _m2_fc2_rowpair_narrow( * scale_lane ) - u_packed1 = ld_global_nc_u32( - w2_base_addr - + ebase_w - + Int64(k_row1) * Int64(cfg.n_half) - + lane_byte_off + u_packed1 = ( + ld_global_nc_u32( + w2_base_addr + + ebase_w + + Int64(k_row1) * Int64(cfg.n_half) + + lane_byte_off + ) + if w_valid > Int32(0) + else Uint32(0) ) bsf_off1 = ( Int64(row_rb1) * Int64(sf_cols * 128) @@ -875,9 +1275,17 @@ def _m2_fc2_rowpair_narrow( + Int64(row_mode_32_1) * Int64(16) + Int64(row_mode_a1) * Int64(4) ) - sf_word1 = ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off1) + sf_word1 = ( + ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off1) + if w_valid > Int32(0) + else Uint32(0) + ) bsf_byte1 = (sf_word1 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) - bsf_f1 = cvt_e4m3_to_f32_via_f16(bsf_byte1) + bsf_f1 = ( + self._scale_byte_to_f32(bsf_byte1) + if w_valid > Int32(0) + else Float32(0.0) + ) out_acc1 = ( out_acc1 + bsf_f1 @@ -885,15 +1293,87 @@ def _m2_fc2_rowpair_narrow( * scale_lane ) + u_packed2 = ( + ld_global_nc_u32( + w2_base_addr + + ebase_w + + Int64(k_row2) * Int64(cfg.n_half) + + lane_byte_off + ) + if w_valid > Int32(0) + else Uint32(0) + ) + bsf_off2 = ( + Int64(row_rb2) * Int64(sf_cols * 128) + + Int64(lane_cb) * Int64(512) + + Int64(row_mode_32_2) * Int64(16) + + Int64(row_mode_a2) * Int64(4) + ) + sf_word2 = ( + ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off2) + if w_valid > Int32(0) + else Uint32(0) + ) + bsf_byte2 = (sf_word2 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) + bsf_f2 = ( + self._scale_byte_to_f32(bsf_byte2) + if w_valid > Int32(0) + else Float32(0.0) + ) + out_acc2 = ( + out_acc2 + + bsf_f2 + * self._fp4_dot4_for_math(u_packed2, xh0, xh1, xh2, xh3) + * scale_lane + ) + + u_packed3 = ( + ld_global_nc_u32( + w2_base_addr + + ebase_w + + Int64(k_row3) * Int64(cfg.n_half) + + lane_byte_off + ) + if w_valid > Int32(0) + else Uint32(0) + ) + bsf_off3 = ( + Int64(row_rb3) * Int64(sf_cols * 128) + + Int64(lane_cb) * Int64(512) + + Int64(row_mode_32_3) * Int64(16) + + Int64(row_mode_a3) * Int64(4) + ) + sf_word3 = ( + ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off3) + if w_valid > Int32(0) + else Uint32(0) + ) + bsf_byte3 = (sf_word3 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) + bsf_f3 = ( + self._scale_byte_to_f32(bsf_byte3) + if w_valid > Int32(0) + else Float32(0.0) + ) + out_acc3 = ( + out_acc3 + + bsf_f3 + * self._fp4_dot4_for_math(u_packed3, xh0, xh1, xh2, xh3) + * scale_lane + ) + sum_warp0 = cute.arch.warp_reduction_sum(out_acc0) sum_warp1 = cute.arch.warp_reduction_sum(out_acc1) + sum_warp2 = cute.arch.warp_reduction_sum(out_acc2) + sum_warp3 = cute.arch.warp_reduction_sum(out_acc3) if lane == Int32(0): out_base = t * Int32(cfg.k_dim) scatter_output[out_base + k_row0] = BFloat16(sum_warp0) scatter_output[out_base + k_row1] = BFloat16(sum_warp1) + scatter_output[out_base + k_row2] = BFloat16(sum_warp2) + scatter_output[out_base + k_row3] = BFloat16(sum_warp3) @cute.jit - def _m2_fc2_rowquad_narrow( + def _m2_fc2_rowpair_narrow( self, fc2_task: Int32, warp_id: Int32, @@ -907,25 +1387,28 @@ def _m2_fc2_rowquad_narrow( scatter_output: cute.Tensor, ): cfg = self._cfg - rows_per_cta = Int32(_K_PER_CTA * 4) + rows_per_cta = Int32(_K_PER_CTA * 2) linear_row_base = fc2_task * rows_per_cta t = linear_row_base // Int32(cfg.k_dim) k_chunk_off = linear_row_base - t * Int32(cfg.k_dim) - k_row0 = k_chunk_off + warp_id * Int32(4) + k_row0 = k_chunk_off + warp_id * Int32(2) k_row1 = k_row0 + Int32(1) - k_row2 = k_row0 + Int32(2) - k_row3 = k_row0 + Int32(3) lane_byte_off = Int64(lane) * Int64(4) token_inter_base = t * Int32(cfg.inter_u32) sf_cols = Int32(cfg.w2_sf_cols) + num_cb = sf_cols >> Int32(2) lane_cb = lane >> Int32(3) + w_valid = Int32(1) if lane_cb < num_cb else Int32(0) lane_mode_c = (lane >> Int32(1)) & Int32(3) bsf_byte_shift = lane_mode_c * Int32(8) out_acc0 = Float32(0.0) out_acc1 = Float32(0.0) - out_acc2 = Float32(0.0) - out_acc3 = Float32(0.0) + k_col0 = Int32(0) + k_col1 = Int32(0) + if cutlass.const_expr(self.w4a16_mode): + k_col0 = self._packed_e4m3_scale_col(k_row0) + k_col1 = self._packed_e4m3_scale_col(k_row1) row_rb0 = k_row0 >> Int32(7) row_mode_a0 = (k_row0 >> Int32(5)) & Int32(3) @@ -933,22 +1416,25 @@ def _m2_fc2_rowquad_narrow( row_rb1 = k_row1 >> Int32(7) row_mode_a1 = (k_row1 >> Int32(5)) & Int32(3) row_mode_32_1 = k_row1 & Int32(31) - row_rb2 = k_row2 >> Int32(7) - row_mode_a2 = (k_row2 >> Int32(5)) & Int32(3) - row_mode_32_2 = k_row2 & Int32(31) - row_rb3 = k_row3 >> Int32(7) - row_mode_a3 = (k_row3 >> Int32(5)) & Int32(3) - row_mode_32_3 = k_row3 & Int32(31) for kk in cutlass.range_constexpr(cfg.num_topk): eid_addr = t * Int32(cfg.num_topk) + Int32(kk) eid = Int32(topk_ids[eid_addr]) router_w = topk_weights[eid_addr] - alpha_fc2 = w2_alphas[eid] - scale_lane = alpha_fc2 * router_w + if cutlass.const_expr( + self.w4a16_mode + and (not self.is_gated) + and cfg.k_dim == 2688 + and cfg.n == 1856 + ): + scale_lane = router_w + else: + alpha_fc2 = w2_alphas[eid] + scale_lane = alpha_fc2 * router_w ebase_w = Int64(eid) * Int64(cfg.k_dim * cfg.n_half) ebase_sf = Int64(eid) * Int64(cfg.w2_sf_rows * cfg.w2_sf_cols) + ebase_sf_packed_e4m3 = Int64(eid) * Int64((cfg.n // 16) * cfg.k_dim) kk_off = token_inter_base + Int32(kk) * Int32(128) xh0 = Uint32(intermediate[kk_off + Int32(0 * 32) + lane]) @@ -956,21 +1442,62 @@ def _m2_fc2_rowquad_narrow( xh2 = Uint32(intermediate[kk_off + Int32(2 * 32) + lane]) xh3 = Uint32(intermediate[kk_off + Int32(3 * 32) + lane]) - u_packed0 = ld_global_nc_u32( - w2_base_addr - + ebase_w - + Int64(k_row0) * Int64(cfg.n_half) - + lane_byte_off - ) - bsf_off0 = ( - Int64(row_rb0) * Int64(sf_cols * 128) - + Int64(lane_cb) * Int64(512) - + Int64(row_mode_32_0) * Int64(16) - + Int64(row_mode_a0) * Int64(4) + u_packed0 = ( + ld_global_nc_u32( + w2_base_addr + + ebase_w + + Int64(k_row0) * Int64(cfg.n_half) + + lane_byte_off + ) + if w_valid > Int32(0) + else Uint32(0) ) - sf_word0 = ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off0) - bsf_byte0 = (sf_word0 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) - bsf_f0 = cvt_e4m3_to_f32_via_f16(bsf_byte0) + if cutlass.const_expr(self.scale_format_e8m0_k32): + ebase_w2p = Int64(eid) * Int64((cfg.n // 32) * cfg.k_dim) + kb32_i = lane >> Int32(2) + bsf_f0 = ( + self._ld_e8m0_scale( + w2s_base_addr, + ebase_w2p, + kb32_i, + k_row0, + Int32(cfg.k_dim), + Int32(cfg.n // 32), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + elif cutlass.const_expr(self.w4a16_mode): + kb16_i = lane >> Int32(1) + bsf_f0 = ( + self._ld_e4m3_packed_scale_col( + w2s_base_addr, + ebase_sf_packed_e4m3, + kb16_i, + k_col0, + Int32(cfg.k_dim), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + else: + bsf_off0 = ( + Int64(row_rb0) * Int64(sf_cols * 128) + + Int64(lane_cb) * Int64(512) + + Int64(row_mode_32_0) * Int64(16) + + Int64(row_mode_a0) * Int64(4) + ) + sf_word0 = ( + ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off0) + if w_valid > Int32(0) + else Uint32(0) + ) + bsf_byte0 = (sf_word0 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) + bsf_f0 = ( + self._scale_byte_to_f32(bsf_byte0) + if w_valid > Int32(0) + else Float32(0.0) + ) out_acc0 = ( out_acc0 + bsf_f0 @@ -978,21 +1505,60 @@ def _m2_fc2_rowquad_narrow( * scale_lane ) - u_packed1 = ld_global_nc_u32( - w2_base_addr - + ebase_w - + Int64(k_row1) * Int64(cfg.n_half) - + lane_byte_off - ) - bsf_off1 = ( - Int64(row_rb1) * Int64(sf_cols * 128) - + Int64(lane_cb) * Int64(512) - + Int64(row_mode_32_1) * Int64(16) - + Int64(row_mode_a1) * Int64(4) + u_packed1 = ( + ld_global_nc_u32( + w2_base_addr + + ebase_w + + Int64(k_row1) * Int64(cfg.n_half) + + lane_byte_off + ) + if w_valid > Int32(0) + else Uint32(0) ) - sf_word1 = ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off1) - bsf_byte1 = (sf_word1 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) - bsf_f1 = cvt_e4m3_to_f32_via_f16(bsf_byte1) + if cutlass.const_expr(self.scale_format_e8m0_k32): + bsf_f1 = ( + self._ld_e8m0_scale( + w2s_base_addr, + ebase_w2p, + kb32_i, + k_row1, + Int32(cfg.k_dim), + Int32(cfg.n // 32), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + elif cutlass.const_expr(self.w4a16_mode): + kb16_i = lane >> Int32(1) + bsf_f1 = ( + self._ld_e4m3_packed_scale_col( + w2s_base_addr, + ebase_sf_packed_e4m3, + kb16_i, + k_col1, + Int32(cfg.k_dim), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + else: + bsf_off1 = ( + Int64(row_rb1) * Int64(sf_cols * 128) + + Int64(lane_cb) * Int64(512) + + Int64(row_mode_32_1) * Int64(16) + + Int64(row_mode_a1) * Int64(4) + ) + sf_word1 = ( + ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off1) + if w_valid > Int32(0) + else Uint32(0) + ) + bsf_byte1 = (sf_word1 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) + bsf_f1 = ( + self._scale_byte_to_f32(bsf_byte1) + if w_valid > Int32(0) + else Float32(0.0) + ) out_acc1 = ( out_acc1 + bsf_f1 @@ -1000,60 +1566,12 @@ def _m2_fc2_rowquad_narrow( * scale_lane ) - u_packed2 = ld_global_nc_u32( - w2_base_addr - + ebase_w - + Int64(k_row2) * Int64(cfg.n_half) - + lane_byte_off - ) - bsf_off2 = ( - Int64(row_rb2) * Int64(sf_cols * 128) - + Int64(lane_cb) * Int64(512) - + Int64(row_mode_32_2) * Int64(16) - + Int64(row_mode_a2) * Int64(4) - ) - sf_word2 = ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off2) - bsf_byte2 = (sf_word2 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) - bsf_f2 = cvt_e4m3_to_f32_via_f16(bsf_byte2) - out_acc2 = ( - out_acc2 - + bsf_f2 - * self._fp4_dot4_for_math(u_packed2, xh0, xh1, xh2, xh3) - * scale_lane - ) - - u_packed3 = ld_global_nc_u32( - w2_base_addr - + ebase_w - + Int64(k_row3) * Int64(cfg.n_half) - + lane_byte_off - ) - bsf_off3 = ( - Int64(row_rb3) * Int64(sf_cols * 128) - + Int64(lane_cb) * Int64(512) - + Int64(row_mode_32_3) * Int64(16) - + Int64(row_mode_a3) * Int64(4) - ) - sf_word3 = ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off3) - bsf_byte3 = (sf_word3 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) - bsf_f3 = cvt_e4m3_to_f32_via_f16(bsf_byte3) - out_acc3 = ( - out_acc3 - + bsf_f3 - * self._fp4_dot4_for_math(u_packed3, xh0, xh1, xh2, xh3) - * scale_lane - ) - sum_warp0 = cute.arch.warp_reduction_sum(out_acc0) sum_warp1 = cute.arch.warp_reduction_sum(out_acc1) - sum_warp2 = cute.arch.warp_reduction_sum(out_acc2) - sum_warp3 = cute.arch.warp_reduction_sum(out_acc3) if lane == Int32(0): out_base = t * Int32(cfg.k_dim) scatter_output[out_base + k_row0] = BFloat16(sum_warp0) scatter_output[out_base + k_row1] = BFloat16(sum_warp1) - scatter_output[out_base + k_row2] = BFloat16(sum_warp2) - scatter_output[out_base + k_row3] = BFloat16(sum_warp3) @cute.jit def _m2_fc2_rowquad_wide( @@ -1091,6 +1609,15 @@ def _m2_fc2_rowquad_wide( out_acc1 = Float32(0.0) out_acc2 = Float32(0.0) out_acc3 = Float32(0.0) + k_col0 = Int32(0) + k_col1 = Int32(0) + k_col2 = Int32(0) + k_col3 = Int32(0) + if cutlass.const_expr(self.w4a16_mode): + k_col0 = self._packed_e4m3_scale_col(k_row0) + k_col1 = self._packed_e4m3_scale_col(k_row1) + k_col2 = self._packed_e4m3_scale_col(k_row2) + k_col3 = self._packed_e4m3_scale_col(k_row3) row_rb0 = k_row0 >> Int32(7) row_mode_a0 = (k_row0 >> Int32(5)) & Int32(3) @@ -1109,22 +1636,158 @@ def _m2_fc2_rowquad_wide( eid_addr = t * Int32(cfg.num_topk) + Int32(kk) eid = Int32(topk_ids[eid_addr]) router_w = topk_weights[eid_addr] - alpha_fc2 = w2_alphas[eid] - scale_lane = alpha_fc2 * router_w + if cutlass.const_expr( + self.w4a16_mode + and (not self.is_gated) + and cfg.k_dim == 2688 + and cfg.n == 1856 + ): + scale_lane = router_w + else: + alpha_fc2 = w2_alphas[eid] + scale_lane = alpha_fc2 * router_w ebase_w = Int64(eid) * Int64(cfg.k_dim * cfg.n_half) ebase_sf = Int64(eid) * Int64(cfg.w2_sf_rows * cfg.w2_sf_cols) + ebase_sf_packed_e4m3 = Int64(eid) * Int64((cfg.n // 16) * cfg.k_dim) for nc in cutlass.range_constexpr(cfg.fc2_n_chunks): chunk_base = Int32(nc) * Int32(128) - kk_off = token_inter_base + Int32(kk) * n_u32_per_expert + chunk_base - xh0 = Uint32(intermediate[kk_off + Int32(0 * 32) + lane]) - xh1 = Uint32(intermediate[kk_off + Int32(1 * 32) + lane]) - xh2 = Uint32(intermediate[kk_off + Int32(2 * 32) + lane]) - xh3 = Uint32(intermediate[kk_off + Int32(3 * 32) + lane]) - cb_idx = Int32(nc) * Int32(4) + lane_cb w_valid = Int32(1) if cb_idx < num_cb else Int32(0) + if cutlass.const_expr(nc + 1 < cfg.fc2_n_chunks): + next_cb_idx = Int32(nc + 1) * Int32(4) + lane_cb + if next_cb_idx < num_cb: + next_chunk_base = Int32(nc + 1) * Int32(128) + prefetch_global_l2( + w2_base_addr + + ebase_w + + Int64(k_row0) * Int64(cfg.n_half) + + Int64(next_chunk_base) + + lane_byte_off, + ) + prefetch_global_l2( + w2_base_addr + + ebase_w + + Int64(k_row1) * Int64(cfg.n_half) + + Int64(next_chunk_base) + + lane_byte_off, + ) + prefetch_global_l2( + w2_base_addr + + ebase_w + + Int64(k_row2) * Int64(cfg.n_half) + + Int64(next_chunk_base) + + lane_byte_off, + ) + prefetch_global_l2( + w2_base_addr + + ebase_w + + Int64(k_row3) * Int64(cfg.n_half) + + Int64(next_chunk_base) + + lane_byte_off, + ) + elif cutlass.const_expr(kk + 1 < cfg.num_topk): + next_eid_addr = t * Int32(cfg.num_topk) + Int32(kk + 1) + next_eid = Int32(topk_ids[next_eid_addr]) + next_ebase_w = Int64(next_eid) * Int64(cfg.k_dim * cfg.n_half) + next_ebase_sf = Int64(next_eid) * Int64( + cfg.w2_sf_rows * cfg.w2_sf_cols + ) + next_cb_idx = lane_cb + if next_cb_idx < num_cb: + prefetch_global_l2( + w2_base_addr + + next_ebase_w + + Int64(k_row0) * Int64(cfg.n_half) + + lane_byte_off, + ) + prefetch_global_l2( + w2_base_addr + + next_ebase_w + + Int64(k_row1) * Int64(cfg.n_half) + + lane_byte_off, + ) + prefetch_global_l2( + w2_base_addr + + next_ebase_w + + Int64(k_row2) * Int64(cfg.n_half) + + lane_byte_off, + ) + prefetch_global_l2( + w2_base_addr + + next_ebase_w + + Int64(k_row3) * Int64(cfg.n_half) + + lane_byte_off, + ) + if cutlass.const_expr( + (not self.w4a16_mode) and (not self.scale_format_e8m0_k32) + ): + next_bsf_off0 = ( + Int64(row_rb0) * Int64(sf_cols * 128) + + Int64(next_cb_idx) * Int64(512) + + Int64(row_mode_32_0) * Int64(16) + + Int64(row_mode_a0) * Int64(4) + ) + next_bsf_off1 = ( + Int64(row_rb1) * Int64(sf_cols * 128) + + Int64(next_cb_idx) * Int64(512) + + Int64(row_mode_32_1) * Int64(16) + + Int64(row_mode_a1) * Int64(4) + ) + next_bsf_off2 = ( + Int64(row_rb2) * Int64(sf_cols * 128) + + Int64(next_cb_idx) * Int64(512) + + Int64(row_mode_32_2) * Int64(16) + + Int64(row_mode_a2) * Int64(4) + ) + next_bsf_off3 = ( + Int64(row_rb3) * Int64(sf_cols * 128) + + Int64(next_cb_idx) * Int64(512) + + Int64(row_mode_32_3) * Int64(16) + + Int64(row_mode_a3) * Int64(4) + ) + prefetch_global_l2( + w2s_base_addr + next_ebase_sf + next_bsf_off0 + ) + prefetch_global_l2( + w2s_base_addr + next_ebase_sf + next_bsf_off1 + ) + prefetch_global_l2( + w2s_base_addr + next_ebase_sf + next_bsf_off2 + ) + prefetch_global_l2( + w2s_base_addr + next_ebase_sf + next_bsf_off3 + ) + kk_off = token_inter_base + Int32(kk) * n_u32_per_expert + chunk_base + # See _m1_fc2_rowpair_wide: mask the intermediate tail read for + # non-256-aligned n (0 weight * NaN tail = NaN). constexpr-gated. + if cutlass.const_expr((cfg.w2_sf_cols >> 2) < cfg.fc2_n_chunks * 4): + xh0 = ( + Uint32(intermediate[kk_off + Int32(0 * 32) + lane]) + if w_valid > Int32(0) + else Uint32(0) + ) + xh1 = ( + Uint32(intermediate[kk_off + Int32(1 * 32) + lane]) + if w_valid > Int32(0) + else Uint32(0) + ) + xh2 = ( + Uint32(intermediate[kk_off + Int32(2 * 32) + lane]) + if w_valid > Int32(0) + else Uint32(0) + ) + xh3 = ( + Uint32(intermediate[kk_off + Int32(3 * 32) + lane]) + if w_valid > Int32(0) + else Uint32(0) + ) + else: + xh0 = Uint32(intermediate[kk_off + Int32(0 * 32) + lane]) + xh1 = Uint32(intermediate[kk_off + Int32(1 * 32) + lane]) + xh2 = Uint32(intermediate[kk_off + Int32(2 * 32) + lane]) + xh3 = Uint32(intermediate[kk_off + Int32(3 * 32) + lane]) u_packed0 = ( ld_global_nc_u32( w2_base_addr @@ -1136,23 +1799,52 @@ def _m2_fc2_rowquad_wide( if w_valid > Int32(0) else Uint32(0) ) - bsf_off0 = ( - Int64(row_rb0) * Int64(sf_cols * 128) - + Int64(cb_idx) * Int64(512) - + Int64(row_mode_32_0) * Int64(16) - + Int64(row_mode_a0) * Int64(4) - ) - sf_word0 = ( - ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off0) - if w_valid > Int32(0) - else Uint32(0) - ) - bsf_byte0 = (sf_word0 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) - bsf_f0 = ( - cvt_e4m3_to_f32_via_f16(bsf_byte0) - if w_valid > Int32(0) - else Float32(0.0) - ) + if cutlass.const_expr(self.scale_format_e8m0_k32): + ebase_w2p = Int64(eid) * Int64((cfg.n // 32) * cfg.k_dim) + kb32_i = (chunk_base + lane * Int32(4)) >> Int32(4) + bsf_f0 = ( + self._ld_e8m0_scale( + w2s_base_addr, + ebase_w2p, + kb32_i, + k_row0, + Int32(cfg.k_dim), + Int32(cfg.n // 32), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + elif cutlass.const_expr(self.w4a16_mode): + kb16_i = (chunk_base + lane * Int32(4)) >> Int32(3) + bsf_f0 = ( + self._ld_e4m3_packed_scale_col( + w2s_base_addr, + ebase_sf_packed_e4m3, + kb16_i, + k_col0, + Int32(cfg.k_dim), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + else: + bsf_off0 = ( + Int64(row_rb0) * Int64(sf_cols * 128) + + Int64(cb_idx) * Int64(512) + + Int64(row_mode_32_0) * Int64(16) + + Int64(row_mode_a0) * Int64(4) + ) + sf_word0 = ( + ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off0) + if w_valid > Int32(0) + else Uint32(0) + ) + bsf_byte0 = (sf_word0 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) + bsf_f0 = ( + self._scale_byte_to_f32(bsf_byte0) + if w_valid > Int32(0) + else Float32(0.0) + ) out_acc0 = ( out_acc0 + bsf_f0 @@ -1171,23 +1863,50 @@ def _m2_fc2_rowquad_wide( if w_valid > Int32(0) else Uint32(0) ) - bsf_off1 = ( - Int64(row_rb1) * Int64(sf_cols * 128) - + Int64(cb_idx) * Int64(512) - + Int64(row_mode_32_1) * Int64(16) - + Int64(row_mode_a1) * Int64(4) - ) - sf_word1 = ( - ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off1) - if w_valid > Int32(0) - else Uint32(0) - ) - bsf_byte1 = (sf_word1 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) - bsf_f1 = ( - cvt_e4m3_to_f32_via_f16(bsf_byte1) - if w_valid > Int32(0) - else Float32(0.0) - ) + if cutlass.const_expr(self.scale_format_e8m0_k32): + bsf_f1 = ( + self._ld_e8m0_scale( + w2s_base_addr, + ebase_w2p, + kb32_i, + k_row1, + Int32(cfg.k_dim), + Int32(cfg.n // 32), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + elif cutlass.const_expr(self.w4a16_mode): + kb16_i = (chunk_base + lane * Int32(4)) >> Int32(3) + bsf_f1 = ( + self._ld_e4m3_packed_scale_col( + w2s_base_addr, + ebase_sf_packed_e4m3, + kb16_i, + k_col1, + Int32(cfg.k_dim), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + else: + bsf_off1 = ( + Int64(row_rb1) * Int64(sf_cols * 128) + + Int64(cb_idx) * Int64(512) + + Int64(row_mode_32_1) * Int64(16) + + Int64(row_mode_a1) * Int64(4) + ) + sf_word1 = ( + ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off1) + if w_valid > Int32(0) + else Uint32(0) + ) + bsf_byte1 = (sf_word1 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) + bsf_f1 = ( + self._scale_byte_to_f32(bsf_byte1) + if w_valid > Int32(0) + else Float32(0.0) + ) out_acc1 = ( out_acc1 + bsf_f1 @@ -1206,23 +1925,50 @@ def _m2_fc2_rowquad_wide( if w_valid > Int32(0) else Uint32(0) ) - bsf_off2 = ( - Int64(row_rb2) * Int64(sf_cols * 128) - + Int64(cb_idx) * Int64(512) - + Int64(row_mode_32_2) * Int64(16) - + Int64(row_mode_a2) * Int64(4) - ) - sf_word2 = ( - ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off2) - if w_valid > Int32(0) - else Uint32(0) - ) - bsf_byte2 = (sf_word2 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) - bsf_f2 = ( - cvt_e4m3_to_f32_via_f16(bsf_byte2) - if w_valid > Int32(0) - else Float32(0.0) - ) + if cutlass.const_expr(self.scale_format_e8m0_k32): + bsf_f2 = ( + self._ld_e8m0_scale( + w2s_base_addr, + ebase_w2p, + kb32_i, + k_row2, + Int32(cfg.k_dim), + Int32(cfg.n // 32), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + elif cutlass.const_expr(self.w4a16_mode): + kb16_i = (chunk_base + lane * Int32(4)) >> Int32(3) + bsf_f2 = ( + self._ld_e4m3_packed_scale_col( + w2s_base_addr, + ebase_sf_packed_e4m3, + kb16_i, + k_col2, + Int32(cfg.k_dim), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + else: + bsf_off2 = ( + Int64(row_rb2) * Int64(sf_cols * 128) + + Int64(cb_idx) * Int64(512) + + Int64(row_mode_32_2) * Int64(16) + + Int64(row_mode_a2) * Int64(4) + ) + sf_word2 = ( + ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off2) + if w_valid > Int32(0) + else Uint32(0) + ) + bsf_byte2 = (sf_word2 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) + bsf_f2 = ( + self._scale_byte_to_f32(bsf_byte2) + if w_valid > Int32(0) + else Float32(0.0) + ) out_acc2 = ( out_acc2 + bsf_f2 @@ -1241,23 +1987,50 @@ def _m2_fc2_rowquad_wide( if w_valid > Int32(0) else Uint32(0) ) - bsf_off3 = ( - Int64(row_rb3) * Int64(sf_cols * 128) - + Int64(cb_idx) * Int64(512) - + Int64(row_mode_32_3) * Int64(16) - + Int64(row_mode_a3) * Int64(4) - ) - sf_word3 = ( - ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off3) - if w_valid > Int32(0) - else Uint32(0) - ) - bsf_byte3 = (sf_word3 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) - bsf_f3 = ( - cvt_e4m3_to_f32_via_f16(bsf_byte3) - if w_valid > Int32(0) - else Float32(0.0) - ) + if cutlass.const_expr(self.scale_format_e8m0_k32): + bsf_f3 = ( + self._ld_e8m0_scale( + w2s_base_addr, + ebase_w2p, + kb32_i, + k_row3, + Int32(cfg.k_dim), + Int32(cfg.n // 32), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + elif cutlass.const_expr(self.w4a16_mode): + kb16_i = (chunk_base + lane * Int32(4)) >> Int32(3) + bsf_f3 = ( + self._ld_e4m3_packed_scale_col( + w2s_base_addr, + ebase_sf_packed_e4m3, + kb16_i, + k_col3, + Int32(cfg.k_dim), + ) + if w_valid > Int32(0) + else Float32(0.0) + ) + else: + bsf_off3 = ( + Int64(row_rb3) * Int64(sf_cols * 128) + + Int64(cb_idx) * Int64(512) + + Int64(row_mode_32_3) * Int64(16) + + Int64(row_mode_a3) * Int64(4) + ) + sf_word3 = ( + ld_global_nc_u32(w2s_base_addr + ebase_sf + bsf_off3) + if w_valid > Int32(0) + else Uint32(0) + ) + bsf_byte3 = (sf_word3 >> Uint32(bsf_byte_shift)) & Uint32(0xFF) + bsf_f3 = ( + self._scale_byte_to_f32(bsf_byte3) + if w_valid > Int32(0) + else Float32(0.0) + ) out_acc3 = ( out_acc3 + bsf_f3 @@ -1300,6 +2073,22 @@ def kernel( bidx_x, _, _ = cute.arch.block_idx() tidx, _, _ = cute.arch.thread_idx() gdim_x, _, _ = cute.arch.grid_dim() + if cutlass.const_expr(self.compile_time_phase == 2): + self._run_fc2( + Int32(bidx_x), + Int32(gdim_x), + tidx // Int32(32), + tidx % Int32(32), + m_val, + w2_weights, + w2_scales, + w2_alphas, + intermediate, + topk_ids, + topk_weights, + scatter_output, + ) + return is_cta_leader = Int32(1) if Int32(tidx) == Int32(0) else Int32(0) m1_epoch0 = Int32(0) if cutlass.const_expr(self.m_const == 1): @@ -1328,7 +2117,7 @@ def kernel( # =================================================================== # PHASE 1: FC1 over route-order tasks # =================================================================== - fc1_task_count = Int32(self.m_const * cfg.num_topk * cfg.fc1_chunks) + fc1_task_count = m_val * Int32(cfg.num_topk * cfg.fc1_chunks) fc1_task = Int32(bidx_x) if cutlass.const_expr(cfg.k_segments == 2): buf_idx = Int32(0) @@ -1350,7 +2139,26 @@ def kernel( v0 = Float32(a_input[x_base + Int32(i * 2)]) v1 = Float32(a_input[x_base + Int32(i * 2 + 1)]) smem_xh[phys_base + Int32(i)] = pack_f32x2_to_f16x2(v0, v1) - else: + if cutlass.const_expr(self.a8_mx_mode): + # Per-32 UE8M0 + E4M3 quantize-dequant (w4a8 prefill numerics). + blk_peak = Float32(0.0) + pair_delta = ( + Int32(1) - Int32(2) * (in_blk & Int32(1)) + ) * Int32(_BLOCK_SIZE) + for i in cutlass.range_constexpr(_BLOCK_SIZE): + v = Float32(a_input[x_base + Int32(i)]) + w = Float32(a_input[x_base + pair_delta + Int32(i)]) + blk_peak = fmax_f32(blk_peak, fmax_f32(v, -v)) + blk_peak = fmax_f32(blk_peak, fmax_f32(w, -w)) + scale32, inv32 = mx_scale_from_amax32(blk_peak) + for i in cutlass.range_constexpr(_BLOCK_SIZE // 2): + v0 = Float32(a_input[x_base + Int32(i * 2)]) + v1 = Float32(a_input[x_base + Int32(i * 2 + 1)]) + f0, f1 = quant_dequant_e4m3_2(v0, v1, inv32, scale32) + smem_xh[phys_base + Int32(i)] = pack_f32x2_to_f16x2(f0, f1) + if cutlass.const_expr( + (not self.w4a16_mode) and (not self.a8_mx_mode) + ): blk_peak = Float32(0.0) for i in cutlass.range_constexpr(_BLOCK_SIZE): v = Float32(a_input[x_base + Int32(i)]) @@ -1362,7 +2170,7 @@ def kernel( q_scale = nvfp4_scale_from_amax(blk_peak, gs_fc1_0) if q_scale > Float32(_FP8_E4M3_MAX): q_scale = Float32(_FP8_E4M3_MAX) - sf_val = cvt_e4m3_to_f32_via_f16(cvt_f32_to_e4m3(q_scale)) + sf_val = self._scale_byte_to_f32(cvt_f32_to_e4m3(q_scale)) eff_scale = Float32(0.0) if gs_fc1_0 != Float32(0.0): eff_scale = sf_val / gs_fc1_0 @@ -1378,17 +2186,42 @@ def kernel( else: prev_t = Int32(-1) while fc1_task < fc1_task_count: - route_idx = fc1_task // Int32(cfg.fc1_chunks) - chunk_idx = fc1_task - route_idx * Int32(cfg.fc1_chunks) + if cutlass.const_expr( + self.w4a16_mode + and self.m_const == 9 + and (not cfg.k_segments_aligned) + and cfg.k_segments == 6 + and cfg.k_blocks == 168 + ): + route_count = m_val * Int32(cfg.num_topk) + chunk_idx = fc1_task // route_count + route_idx = fc1_task - chunk_idx * route_count + else: + route_idx = fc1_task // Int32(cfg.fc1_chunks) + chunk_idx = fc1_task - route_idx * Int32(cfg.fc1_chunks) t = route_idx // Int32(cfg.num_topk) k_idx = route_idx - t * Int32(cfg.num_topk) i_chunk_off = chunk_idx * Int32(cfg.i_chunk) eid_addr = t * Int32(cfg.num_topk) + k_idx eid = Int32(topk_ids[eid_addr]) - alpha_fc1 = w1_alphas[eid] - gs_fc1 = input_gs[eid] - gs_fc2 = down_input_scale[eid] + if cutlass.const_expr( + self.w4a16_mode + and (not self.is_gated) + and cfg.k_dim == 2688 + and cfg.n == 1856 + ): + alpha_fc1 = Float32(1.0) + gs_fc1 = Float32(1.0) + gs_fc2 = Float32(1.0) + else: + alpha_fc1 = w1_alphas[eid] + gs_fc1 = input_gs[eid] + gs_fc2 = down_input_scale[eid] + if cutlass.const_expr(self.a8_mx_mode): + # a8_mx activations are self-ranging: fold the calibrated + # input global scale out of the combined nvfp4 alpha. + alpha_fc1 = alpha_fc1 * gs_fc1 # ---- Input quantization ---- if cutlass.const_expr(cfg.k_segments != 2): @@ -1408,7 +2241,28 @@ def kernel( smem_xh[phys_base + Int32(i)] = pack_f32x2_to_f16x2( v0, v1 ) - else: + if cutlass.const_expr(self.a8_mx_mode): + # Per-32 UE8M0 + E4M3 quantize-dequant (w4a8 prefill numerics). + blk_peak = Float32(0.0) + pair_delta = ( + Int32(1) - Int32(2) * (in_blk & Int32(1)) + ) * Int32(_BLOCK_SIZE) + for i in cutlass.range_constexpr(_BLOCK_SIZE): + v = Float32(a_input[x_base + Int32(i)]) + w = Float32(a_input[x_base + pair_delta + Int32(i)]) + blk_peak = fmax_f32(blk_peak, fmax_f32(v, -v)) + blk_peak = fmax_f32(blk_peak, fmax_f32(w, -w)) + scale32, inv32 = mx_scale_from_amax32(blk_peak) + for i in cutlass.range_constexpr(_BLOCK_SIZE // 2): + v0 = Float32(a_input[x_base + Int32(i * 2)]) + v1 = Float32(a_input[x_base + Int32(i * 2 + 1)]) + f0, f1 = quant_dequant_e4m3_2(v0, v1, inv32, scale32) + smem_xh[phys_base + Int32(i)] = pack_f32x2_to_f16x2( + f0, f1 + ) + if cutlass.const_expr( + (not self.w4a16_mode) and (not self.a8_mx_mode) + ): blk_peak = Float32(0.0) for i in cutlass.range_constexpr(_BLOCK_SIZE): v = Float32(a_input[x_base + Int32(i)]) @@ -1420,7 +2274,7 @@ def kernel( q_scale = nvfp4_scale_from_amax(blk_peak, gs_fc1) if q_scale > Float32(_FP8_E4M3_MAX): q_scale = Float32(_FP8_E4M3_MAX) - sf_val = cvt_e4m3_to_f32_via_f16(cvt_f32_to_e4m3(q_scale)) + sf_val = self._scale_byte_to_f32(cvt_f32_to_e4m3(q_scale)) eff_scale = Float32(0.0) if gs_fc1 != Float32(0.0): eff_scale = sf_val / gs_fc1 @@ -1441,6 +2295,8 @@ def kernel( # ---- FC1 weight load + dot product ---- ebase_w = Int64(eid) * Int64(cfg.two_n) * Int64(cfg.k_half) ebase_sf = Int64(eid) * Int64(cfg.w1_sf_rows * cfg.w1_sf_cols) + ebase_sf_packed = Int64(eid) * Int64((cfg.k_dim // 32) * cfg.two_n) + ebase_sf_packed_e4m3 = Int64(eid) * Int64((cfg.k_dim // 16) * cfg.two_n) thread_byte_off = Int64(lane) * Int64(cfg.k_half // 32) xh_buf_base = Int32(0) if cutlass.const_expr(cfg.k_segments == 2): @@ -1451,21 +2307,100 @@ def kernel( xh_buf_base + lane_seg_base * Int32(_BLOCK_SIZE // 2) + lane_pad_base ) + # The smem activation layout depends only on the lane and token, + # so the k_segments==8 gated path's 4 rows/warp re-read the same + # 64 words; hoist them into registers once per warp-task. + hoist_xh = cutlass.const_expr( + cfg.k_segments_aligned and cfg.k_segments == 8 and self.is_gated + ) + if cutlass.const_expr(hoist_xh): + xa0 = Uint32(smem_xh[xh_base_t + Int32(0)]) + xa1 = Uint32(smem_xh[xh_base_t + Int32(1)]) + xa2 = Uint32(smem_xh[xh_base_t + Int32(2)]) + xa3 = Uint32(smem_xh[xh_base_t + Int32(3)]) + xa4 = Uint32(smem_xh[xh_base_t + Int32(4)]) + xa5 = Uint32(smem_xh[xh_base_t + Int32(5)]) + xa6 = Uint32(smem_xh[xh_base_t + Int32(6)]) + xa7 = Uint32(smem_xh[xh_base_t + Int32(7)]) + xa8 = Uint32(smem_xh[xh_base_t + Int32(8)]) + xa9 = Uint32(smem_xh[xh_base_t + Int32(9)]) + xa10 = Uint32(smem_xh[xh_base_t + Int32(10)]) + xa11 = Uint32(smem_xh[xh_base_t + Int32(11)]) + xa12 = Uint32(smem_xh[xh_base_t + Int32(12)]) + xa13 = Uint32(smem_xh[xh_base_t + Int32(13)]) + xa14 = Uint32(smem_xh[xh_base_t + Int32(14)]) + xa15 = Uint32(smem_xh[xh_base_t + Int32(15)]) + xa16 = Uint32(smem_xh[xh_base_t + Int32(16)]) + xa17 = Uint32(smem_xh[xh_base_t + Int32(17)]) + xa18 = Uint32(smem_xh[xh_base_t + Int32(18)]) + xa19 = Uint32(smem_xh[xh_base_t + Int32(19)]) + xa20 = Uint32(smem_xh[xh_base_t + Int32(20)]) + xa21 = Uint32(smem_xh[xh_base_t + Int32(21)]) + xa22 = Uint32(smem_xh[xh_base_t + Int32(22)]) + xa23 = Uint32(smem_xh[xh_base_t + Int32(23)]) + xa24 = Uint32(smem_xh[xh_base_t + Int32(24)]) + xa25 = Uint32(smem_xh[xh_base_t + Int32(25)]) + xa26 = Uint32(smem_xh[xh_base_t + Int32(26)]) + xa27 = Uint32(smem_xh[xh_base_t + Int32(27)]) + xa28 = Uint32(smem_xh[xh_base_t + Int32(28)]) + xa29 = Uint32(smem_xh[xh_base_t + Int32(29)]) + xa30 = Uint32(smem_xh[xh_base_t + Int32(30)]) + xa31 = Uint32(smem_xh[xh_base_t + Int32(31)]) + xa32 = Uint32(smem_xh[xh_base_t + Int32(32)]) + xa33 = Uint32(smem_xh[xh_base_t + Int32(33)]) + xa34 = Uint32(smem_xh[xh_base_t + Int32(34)]) + xa35 = Uint32(smem_xh[xh_base_t + Int32(35)]) + xa36 = Uint32(smem_xh[xh_base_t + Int32(36)]) + xa37 = Uint32(smem_xh[xh_base_t + Int32(37)]) + xa38 = Uint32(smem_xh[xh_base_t + Int32(38)]) + xa39 = Uint32(smem_xh[xh_base_t + Int32(39)]) + xa40 = Uint32(smem_xh[xh_base_t + Int32(40)]) + xa41 = Uint32(smem_xh[xh_base_t + Int32(41)]) + xa42 = Uint32(smem_xh[xh_base_t + Int32(42)]) + xa43 = Uint32(smem_xh[xh_base_t + Int32(43)]) + xa44 = Uint32(smem_xh[xh_base_t + Int32(44)]) + xa45 = Uint32(smem_xh[xh_base_t + Int32(45)]) + xa46 = Uint32(smem_xh[xh_base_t + Int32(46)]) + xa47 = Uint32(smem_xh[xh_base_t + Int32(47)]) + xa48 = Uint32(smem_xh[xh_base_t + Int32(48)]) + xa49 = Uint32(smem_xh[xh_base_t + Int32(49)]) + xa50 = Uint32(smem_xh[xh_base_t + Int32(50)]) + xa51 = Uint32(smem_xh[xh_base_t + Int32(51)]) + xa52 = Uint32(smem_xh[xh_base_t + Int32(52)]) + xa53 = Uint32(smem_xh[xh_base_t + Int32(53)]) + xa54 = Uint32(smem_xh[xh_base_t + Int32(54)]) + xa55 = Uint32(smem_xh[xh_base_t + Int32(55)]) + xa56 = Uint32(smem_xh[xh_base_t + Int32(56)]) + xa57 = Uint32(smem_xh[xh_base_t + Int32(57)]) + xa58 = Uint32(smem_xh[xh_base_t + Int32(58)]) + xa59 = Uint32(smem_xh[xh_base_t + Int32(59)]) + xa60 = Uint32(smem_xh[xh_base_t + Int32(60)]) + xa61 = Uint32(smem_xh[xh_base_t + Int32(61)]) + xa62 = Uint32(smem_xh[xh_base_t + Int32(62)]) + xa63 = Uint32(smem_xh[xh_base_t + Int32(63)]) + for r_iter in cutlass.range_constexpr(cfg.rows_per_warp_fc1): i_local = warp_id * Int32(cfg.rows_per_warp_fc1) + Int32(r_iter) i = i_chunk_off + i_local if cutlass.const_expr(self.is_gated): + # Physical FC1-half rows for output channel i. "w13" (up_gate) + # keeps up in the first half; "w31" (gate_up) swaps them. + if cutlass.const_expr(self.w13_gate_first): + row_u = Int32(cfg.n) + i + row_g = i + else: + row_u = i + row_g = Int32(cfg.n) + i up_byte_addr = ( w1_base_addr + ebase_w - + Int64(i) * Int64(cfg.k_half) + + Int64(row_u) * Int64(cfg.k_half) + thread_byte_off ) - row_g = Int32(cfg.n) + i - rb_u = i >> Int32(7) - mode_a_u = (i >> Int32(5)) & Int32(3) - mode_32_u = i & Int32(31) + rb_u = row_u >> Int32(7) + mode_a_u = (row_u >> Int32(5)) & Int32(3) + mode_32_u = row_u & Int32(31) bsf_base_u = Int64(rb_u) * Int64(cfg.w1_sf_cols * 128) + Int64( mode_32_u * Int32(16) + mode_a_u * Int32(4) ) @@ -1475,6 +2410,16 @@ def kernel( bsf_base_g = Int64(rb_g) * Int64(cfg.w1_sf_cols * 128) + Int64( mode_32_g * Int32(16) + mode_a_g * Int32(4) ) + scale_row_u = row_u + scale_row_g = row_g + if cutlass.const_expr( + self.w4a16_mode and (not self.w13_gate_first) + ): + # Main W4A16 scales are packed in kernel-native gate/up + # order. ModelOpt "w13" weights remain up/gate, so only + # the scale rows need the half swap in the micro path. + scale_row_u = row_g + scale_row_g = row_u else: row_g = i rb_g = row_g >> Int32(7) @@ -1483,6 +2428,14 @@ def kernel( bsf_base_g = Int64(rb_g) * Int64(cfg.w1_sf_cols * 128) + Int64( mode_32_g * Int32(16) + mode_a_g * Int32(4) ) + scale_row_g = row_g + scale_col_g = Int32(0) + if cutlass.const_expr(self.w4a16_mode): + scale_col_g = self._packed_e4m3_scale_col(scale_row_g) + if cutlass.const_expr(self.is_gated): + scale_col_u = Int32(0) + if cutlass.const_expr(self.w4a16_mode): + scale_col_u = self._packed_e4m3_scale_col(scale_row_u) gate_byte_addr = ( w1_base_addr + ebase_w @@ -1491,7 +2444,7 @@ def kernel( ) col_blk_off = Int64(lane) * Int64((cfg.k_segments // 4) * 512) - if cutlass.const_expr(cfg.k_segments == 8): + if cutlass.const_expr(cfg.k_segments_aligned and cfg.k_segments == 8): if cutlass.const_expr(self.is_gated): uw_a0, uw_a1, uw_a2, uw_a3 = ld_global_nc_v4_u32(up_byte_addr) uw_b0, uw_b1, uw_b2, uw_b3 = ld_global_nc_v4_u32( @@ -1504,16 +2457,118 @@ def kernel( up_byte_addr + Int64(48) ) - bsf_addr_u_a = ( - w1s_base_addr + ebase_sf + bsf_base_u + col_blk_off - ) - bsf_addr_u_b = bsf_addr_u_a + Int64(512) - sf_u0, sf_u1, sf_u2, sf_u3 = cvt_e4m3x4_to_f32x4( - ld_global_nc_u32(bsf_addr_u_a) - ) - sf_u4, sf_u5, sf_u6, sf_u7 = cvt_e4m3x4_to_f32x4( - ld_global_nc_u32(bsf_addr_u_b) - ) + if cutlass.const_expr(self.scale_format_e8m0_k32): + kbb = (lane * Int32(cfg.k_segments)) >> Int32(1) + nc1 = Int32(cfg.two_n) + u_k0 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb + Int32(0), + row_u, + nc1, + Int32(cfg.k_dim // 32), + ) + u_k1 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb + Int32(1), + row_u, + nc1, + Int32(cfg.k_dim // 32), + ) + u_k2 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb + Int32(2), + row_u, + nc1, + Int32(cfg.k_dim // 32), + ) + u_k3 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb + Int32(3), + row_u, + nc1, + Int32(cfg.k_dim // 32), + ) + sf_u0 = u_k0 + sf_u1 = u_k0 + sf_u2 = u_k1 + sf_u3 = u_k1 + sf_u4 = u_k2 + sf_u5 = u_k2 + sf_u6 = u_k3 + sf_u7 = u_k3 + elif cutlass.const_expr(self.w4a16_mode): + k16_u = lane * Int32(cfg.k_segments) + sf_u0 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(0), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u1 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(1), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u2 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(2), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u3 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(3), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u4 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(4), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u5 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(5), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u6 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(6), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u7 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(7), + scale_col_u, + Int32(cfg.two_n), + ) + else: + bsf_addr_u_a = ( + w1s_base_addr + ebase_sf + bsf_base_u + col_blk_off + ) + bsf_addr_u_b = bsf_addr_u_a + Int64(512) + sf_u0, sf_u1, sf_u2, sf_u3 = self._scale_word_to_f32x4( + ld_global_nc_u32(bsf_addr_u_a) + ) + sf_u4, sf_u5, sf_u6, sf_u7 = self._scale_word_to_f32x4( + ld_global_nc_u32(bsf_addr_u_b) + ) gw_a0, gw_a1, gw_a2, gw_a3 = ld_global_nc_v4_u32(gate_byte_addr) gw_b0, gw_b1, gw_b2, gw_b3 = ld_global_nc_v4_u32( @@ -1526,14 +2581,118 @@ def kernel( gate_byte_addr + Int64(48) ) - bsf_addr_g_a = w1s_base_addr + ebase_sf + bsf_base_g + col_blk_off - bsf_addr_g_b = bsf_addr_g_a + Int64(512) - sf_g0, sf_g1, sf_g2, sf_g3 = cvt_e4m3x4_to_f32x4( - ld_global_nc_u32(bsf_addr_g_a) - ) - sf_g4, sf_g5, sf_g6, sf_g7 = cvt_e4m3x4_to_f32x4( - ld_global_nc_u32(bsf_addr_g_b) - ) + if cutlass.const_expr(self.scale_format_e8m0_k32): + kbbg = (lane * Int32(cfg.k_segments)) >> Int32(1) + nc1g = Int32(cfg.two_n) + g_k0 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbbg + Int32(0), + row_g, + nc1g, + Int32(cfg.k_dim // 32), + ) + g_k1 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbbg + Int32(1), + row_g, + nc1g, + Int32(cfg.k_dim // 32), + ) + g_k2 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbbg + Int32(2), + row_g, + nc1g, + Int32(cfg.k_dim // 32), + ) + g_k3 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbbg + Int32(3), + row_g, + nc1g, + Int32(cfg.k_dim // 32), + ) + sf_g0 = g_k0 + sf_g1 = g_k0 + sf_g2 = g_k1 + sf_g3 = g_k1 + sf_g4 = g_k2 + sf_g5 = g_k2 + sf_g6 = g_k3 + sf_g7 = g_k3 + elif cutlass.const_expr(self.w4a16_mode): + k16_g = lane * Int32(cfg.k_segments) + sf_g0 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(0), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g1 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(1), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g2 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(2), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g3 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(3), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g4 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(4), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g5 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(5), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g6 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(6), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g7 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(7), + scale_col_g, + Int32(cfg.two_n), + ) + else: + bsf_addr_g_a = ( + w1s_base_addr + ebase_sf + bsf_base_g + col_blk_off + ) + bsf_addr_g_b = bsf_addr_g_a + Int64(512) + sf_g0, sf_g1, sf_g2, sf_g3 = self._scale_word_to_f32x4( + ld_global_nc_u32(bsf_addr_g_a) + ) + sf_g4, sf_g5, sf_g6, sf_g7 = self._scale_word_to_f32x4( + ld_global_nc_u32(bsf_addr_g_b) + ) if cutlass.const_expr(not self.is_gated): partial_gate = ( @@ -1571,29 +2730,117 @@ def kernel( ) ) elif cutlass.const_expr(self.is_gated): - dot_u0, dot_g0 = self._block_dot_hfma2_pair_for_math( - uw_a0, uw_a1, gw_a0, gw_a1, smem_xh, xh_base_t + Int32(0) + dot_u0, dot_g0 = self._block_dot_hfma2_pair_regs_for_math( + uw_a0, + uw_a1, + gw_a0, + gw_a1, + xa0, + xa1, + xa2, + xa3, + xa4, + xa5, + xa6, + xa7, ) - dot_u1, dot_g1 = self._block_dot_hfma2_pair_for_math( - uw_a2, uw_a3, gw_a2, gw_a3, smem_xh, xh_base_t + Int32(8) + dot_u1, dot_g1 = self._block_dot_hfma2_pair_regs_for_math( + uw_a2, + uw_a3, + gw_a2, + gw_a3, + xa8, + xa9, + xa10, + xa11, + xa12, + xa13, + xa14, + xa15, ) - dot_u2, dot_g2 = self._block_dot_hfma2_pair_for_math( - uw_b0, uw_b1, gw_b0, gw_b1, smem_xh, xh_base_t + Int32(16) + dot_u2, dot_g2 = self._block_dot_hfma2_pair_regs_for_math( + uw_b0, + uw_b1, + gw_b0, + gw_b1, + xa16, + xa17, + xa18, + xa19, + xa20, + xa21, + xa22, + xa23, ) - dot_u3, dot_g3 = self._block_dot_hfma2_pair_for_math( - uw_b2, uw_b3, gw_b2, gw_b3, smem_xh, xh_base_t + Int32(24) + dot_u3, dot_g3 = self._block_dot_hfma2_pair_regs_for_math( + uw_b2, + uw_b3, + gw_b2, + gw_b3, + xa24, + xa25, + xa26, + xa27, + xa28, + xa29, + xa30, + xa31, ) - dot_u4, dot_g4 = self._block_dot_hfma2_pair_for_math( - uw_c0, uw_c1, gw_c0, gw_c1, smem_xh, xh_base_t + Int32(32) + dot_u4, dot_g4 = self._block_dot_hfma2_pair_regs_for_math( + uw_c0, + uw_c1, + gw_c0, + gw_c1, + xa32, + xa33, + xa34, + xa35, + xa36, + xa37, + xa38, + xa39, ) - dot_u5, dot_g5 = self._block_dot_hfma2_pair_for_math( - uw_c2, uw_c3, gw_c2, gw_c3, smem_xh, xh_base_t + Int32(40) + dot_u5, dot_g5 = self._block_dot_hfma2_pair_regs_for_math( + uw_c2, + uw_c3, + gw_c2, + gw_c3, + xa40, + xa41, + xa42, + xa43, + xa44, + xa45, + xa46, + xa47, ) - dot_u6, dot_g6 = self._block_dot_hfma2_pair_for_math( - uw_d0, uw_d1, gw_d0, gw_d1, smem_xh, xh_base_t + Int32(48) + dot_u6, dot_g6 = self._block_dot_hfma2_pair_regs_for_math( + uw_d0, + uw_d1, + gw_d0, + gw_d1, + xa48, + xa49, + xa50, + xa51, + xa52, + xa53, + xa54, + xa55, ) - dot_u7, dot_g7 = self._block_dot_hfma2_pair_for_math( - uw_d2, uw_d3, gw_d2, gw_d3, smem_xh, xh_base_t + Int32(56) + dot_u7, dot_g7 = self._block_dot_hfma2_pair_regs_for_math( + uw_d2, + uw_d3, + gw_d2, + gw_d3, + xa56, + xa57, + xa58, + xa59, + xa60, + xa61, + xa62, + xa63, ) partial_up = ( sf_u0 * dot_u0 @@ -1615,7 +2862,7 @@ def kernel( + sf_g6 * dot_g6 + sf_g7 * dot_g7 ) - elif cutlass.const_expr(cfg.k_segments == 6): + elif cutlass.const_expr(cfg.k_segments_aligned and cfg.k_segments == 6): xh_off0 = Int32(0) xh_off1 = Int32(8) + ( (lane_seg_base + Int32(1)) // Int32(8) - lane_pad_base @@ -1645,54 +2892,137 @@ def kernel( up_byte_addr + Int64(32) ) - sf_word_u_a = ld_global_nc_u32( - w1s_base_addr + ebase_sf + bsf_base_u + scale_pair_off - ) - sf_word_u_b = ld_global_nc_u32( - w1s_base_addr - + ebase_sf - + bsf_base_u - + scale_pair_off - + Int64(512) - ) sf_u0 = Float32(0.0) sf_u1 = Float32(0.0) sf_u2 = Float32(0.0) sf_u3 = Float32(0.0) sf_u4 = Float32(0.0) sf_u5 = Float32(0.0) - if scale_lane_mod == Int32(0): - sf_u0 = cvt_e4m3_to_f32_via_f16(sf_word_u_a & Uint32(0xFF)) - sf_u1 = cvt_e4m3_to_f32_via_f16( - (sf_word_u_a >> Uint32(8)) & Uint32(0xFF) + if cutlass.const_expr(self.scale_format_e8m0_k32): + kbb_u = lane_seg_base >> Int32(1) + nc_u = Int32(cfg.two_n) + u_k0 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_u + Int32(0), + row_u, + nc_u, + Int32(cfg.k_dim // 32), ) - sf_u2 = cvt_e4m3_to_f32_via_f16( - (sf_word_u_a >> Uint32(16)) & Uint32(0xFF) + u_k1 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_u + Int32(1), + row_u, + nc_u, + Int32(cfg.k_dim // 32), ) - sf_u3 = cvt_e4m3_to_f32_via_f16( - (sf_word_u_a >> Uint32(24)) & Uint32(0xFF) + u_k2 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_u + Int32(2), + row_u, + nc_u, + Int32(cfg.k_dim // 32), ) - sf_u4 = cvt_e4m3_to_f32_via_f16(sf_word_u_b & Uint32(0xFF)) - sf_u5 = cvt_e4m3_to_f32_via_f16( - (sf_word_u_b >> Uint32(8)) & Uint32(0xFF) + sf_u0 = u_k0 + sf_u1 = u_k0 + sf_u2 = u_k1 + sf_u3 = u_k1 + sf_u4 = u_k2 + sf_u5 = u_k2 + elif cutlass.const_expr(self.w4a16_mode): + sf_u0 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base + Int32(0), + scale_col_u, + Int32(cfg.two_n), ) - else: - sf_u0 = cvt_e4m3_to_f32_via_f16( - (sf_word_u_a >> Uint32(16)) & Uint32(0xFF) + sf_u1 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base + Int32(1), + scale_col_u, + Int32(cfg.two_n), ) - sf_u1 = cvt_e4m3_to_f32_via_f16( - (sf_word_u_a >> Uint32(24)) & Uint32(0xFF) + sf_u2 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base + Int32(2), + scale_col_u, + Int32(cfg.two_n), ) - sf_u2 = cvt_e4m3_to_f32_via_f16(sf_word_u_b & Uint32(0xFF)) - sf_u3 = cvt_e4m3_to_f32_via_f16( - (sf_word_u_b >> Uint32(8)) & Uint32(0xFF) + sf_u3 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base + Int32(3), + scale_col_u, + Int32(cfg.two_n), ) - sf_u4 = cvt_e4m3_to_f32_via_f16( - (sf_word_u_b >> Uint32(16)) & Uint32(0xFF) + sf_u4 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base + Int32(4), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u5 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base + Int32(5), + scale_col_u, + Int32(cfg.two_n), + ) + else: + sf_word_u_a = ld_global_nc_u32( + w1s_base_addr + ebase_sf + bsf_base_u + scale_pair_off ) - sf_u5 = cvt_e4m3_to_f32_via_f16( - (sf_word_u_b >> Uint32(24)) & Uint32(0xFF) + sf_word_u_b = ld_global_nc_u32( + w1s_base_addr + + ebase_sf + + bsf_base_u + + scale_pair_off + + Int64(512) ) + if scale_lane_mod == Int32(0): + sf_u0 = self._scale_byte_to_f32( + sf_word_u_a & Uint32(0xFF) + ) + sf_u1 = self._scale_byte_to_f32( + (sf_word_u_a >> Uint32(8)) & Uint32(0xFF) + ) + sf_u2 = self._scale_byte_to_f32( + (sf_word_u_a >> Uint32(16)) & Uint32(0xFF) + ) + sf_u3 = self._scale_byte_to_f32( + (sf_word_u_a >> Uint32(24)) & Uint32(0xFF) + ) + sf_u4 = self._scale_byte_to_f32( + sf_word_u_b & Uint32(0xFF) + ) + sf_u5 = self._scale_byte_to_f32( + (sf_word_u_b >> Uint32(8)) & Uint32(0xFF) + ) + else: + sf_u0 = self._scale_byte_to_f32( + (sf_word_u_a >> Uint32(16)) & Uint32(0xFF) + ) + sf_u1 = self._scale_byte_to_f32( + (sf_word_u_a >> Uint32(24)) & Uint32(0xFF) + ) + sf_u2 = self._scale_byte_to_f32( + sf_word_u_b & Uint32(0xFF) + ) + sf_u3 = self._scale_byte_to_f32( + (sf_word_u_b >> Uint32(8)) & Uint32(0xFF) + ) + sf_u4 = self._scale_byte_to_f32( + (sf_word_u_b >> Uint32(16)) & Uint32(0xFF) + ) + sf_u5 = self._scale_byte_to_f32( + (sf_word_u_b >> Uint32(24)) & Uint32(0xFF) + ) gw_a0, gw_a1, gw_a2, gw_a3 = ld_global_nc_v4_u32(gate_byte_addr) gw_b0, gw_b1, gw_b2, gw_b3 = ld_global_nc_v4_u32( @@ -1702,54 +3032,131 @@ def kernel( gate_byte_addr + Int64(32) ) - sf_word_g_a = ld_global_nc_u32( - w1s_base_addr + ebase_sf + bsf_base_g + scale_pair_off - ) - sf_word_g_b = ld_global_nc_u32( - w1s_base_addr - + ebase_sf - + bsf_base_g - + scale_pair_off - + Int64(512) - ) sf_g0 = Float32(0.0) sf_g1 = Float32(0.0) sf_g2 = Float32(0.0) sf_g3 = Float32(0.0) sf_g4 = Float32(0.0) sf_g5 = Float32(0.0) - if scale_lane_mod == Int32(0): - sf_g0 = cvt_e4m3_to_f32_via_f16(sf_word_g_a & Uint32(0xFF)) - sf_g1 = cvt_e4m3_to_f32_via_f16( - (sf_word_g_a >> Uint32(8)) & Uint32(0xFF) + if cutlass.const_expr(self.scale_format_e8m0_k32): + kbb_g = lane_seg_base >> Int32(1) + nc_g = Int32(cfg.two_n) + g_k0 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_g + Int32(0), + row_g, + nc_g, + Int32(cfg.k_dim // 32), ) - sf_g2 = cvt_e4m3_to_f32_via_f16( - (sf_word_g_a >> Uint32(16)) & Uint32(0xFF) + g_k1 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_g + Int32(1), + row_g, + nc_g, + Int32(cfg.k_dim // 32), ) - sf_g3 = cvt_e4m3_to_f32_via_f16( - (sf_word_g_a >> Uint32(24)) & Uint32(0xFF) + g_k2 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_g + Int32(2), + row_g, + nc_g, + Int32(cfg.k_dim // 32), ) - sf_g4 = cvt_e4m3_to_f32_via_f16(sf_word_g_b & Uint32(0xFF)) - sf_g5 = cvt_e4m3_to_f32_via_f16( - (sf_word_g_b >> Uint32(8)) & Uint32(0xFF) + sf_g0 = g_k0 + sf_g1 = g_k0 + sf_g2 = g_k1 + sf_g3 = g_k1 + sf_g4 = g_k2 + sf_g5 = g_k2 + elif cutlass.const_expr(self.w4a16_mode): + sf_g0 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base + Int32(0), + scale_col_g, + Int32(cfg.two_n), ) - else: - sf_g0 = cvt_e4m3_to_f32_via_f16( - (sf_word_g_a >> Uint32(16)) & Uint32(0xFF) + sf_g1 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base + Int32(1), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g2 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base + Int32(2), + scale_col_g, + Int32(cfg.two_n), ) - sf_g1 = cvt_e4m3_to_f32_via_f16( - (sf_word_g_a >> Uint32(24)) & Uint32(0xFF) + sf_g3 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base + Int32(3), + scale_col_g, + Int32(cfg.two_n), ) - sf_g2 = cvt_e4m3_to_f32_via_f16(sf_word_g_b & Uint32(0xFF)) - sf_g3 = cvt_e4m3_to_f32_via_f16( - (sf_word_g_b >> Uint32(8)) & Uint32(0xFF) + sf_g4 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base + Int32(4), + scale_col_g, + Int32(cfg.two_n), ) - sf_g4 = cvt_e4m3_to_f32_via_f16( - (sf_word_g_b >> Uint32(16)) & Uint32(0xFF) + sf_g5 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base + Int32(5), + scale_col_g, + Int32(cfg.two_n), + ) + else: + sf_word_g_a = ld_global_nc_u32( + w1s_base_addr + ebase_sf + bsf_base_g + scale_pair_off ) - sf_g5 = cvt_e4m3_to_f32_via_f16( - (sf_word_g_b >> Uint32(24)) & Uint32(0xFF) + sf_word_g_b = ld_global_nc_u32( + w1s_base_addr + + ebase_sf + + bsf_base_g + + scale_pair_off + + Int64(512) ) + if scale_lane_mod == Int32(0): + sf_g0 = self._scale_byte_to_f32(sf_word_g_a & Uint32(0xFF)) + sf_g1 = self._scale_byte_to_f32( + (sf_word_g_a >> Uint32(8)) & Uint32(0xFF) + ) + sf_g2 = self._scale_byte_to_f32( + (sf_word_g_a >> Uint32(16)) & Uint32(0xFF) + ) + sf_g3 = self._scale_byte_to_f32( + (sf_word_g_a >> Uint32(24)) & Uint32(0xFF) + ) + sf_g4 = self._scale_byte_to_f32(sf_word_g_b & Uint32(0xFF)) + sf_g5 = self._scale_byte_to_f32( + (sf_word_g_b >> Uint32(8)) & Uint32(0xFF) + ) + else: + sf_g0 = self._scale_byte_to_f32( + (sf_word_g_a >> Uint32(16)) & Uint32(0xFF) + ) + sf_g1 = self._scale_byte_to_f32( + (sf_word_g_a >> Uint32(24)) & Uint32(0xFF) + ) + sf_g2 = self._scale_byte_to_f32(sf_word_g_b & Uint32(0xFF)) + sf_g3 = self._scale_byte_to_f32( + (sf_word_g_b >> Uint32(8)) & Uint32(0xFF) + ) + sf_g4 = self._scale_byte_to_f32( + (sf_word_g_b >> Uint32(16)) & Uint32(0xFF) + ) + sf_g5 = self._scale_byte_to_f32( + (sf_word_g_b >> Uint32(24)) & Uint32(0xFF) + ) if cutlass.const_expr(not self.is_gated): partial_gate = ( @@ -1813,7 +3220,9 @@ def kernel( + sf_g4 * dot_g4 + sf_g5 * dot_g5 ) - elif cutlass.const_expr(cfg.k_segments == 12): + elif cutlass.const_expr( + cfg.k_segments_aligned and cfg.k_segments == 12 + ): xh_off0 = Int32(0) xh_off1 = Int32(8) + ( (lane_seg_base + Int32(1)) // Int32(8) - lane_pad_base @@ -1867,20 +3276,170 @@ def kernel( up_byte_addr + Int64(80) ) - bsf_addr_u_a = ( - w1s_base_addr + ebase_sf + bsf_base_u + col_blk_off - ) - bsf_addr_u_b = bsf_addr_u_a + Int64(512) - bsf_addr_u_c = bsf_addr_u_a + Int64(1024) - sf_u0, sf_u1, sf_u2, sf_u3 = cvt_e4m3x4_to_f32x4( - ld_global_nc_u32(bsf_addr_u_a) - ) - sf_u4, sf_u5, sf_u6, sf_u7 = cvt_e4m3x4_to_f32x4( - ld_global_nc_u32(bsf_addr_u_b) - ) - sf_u8, sf_u9, sf_u10, sf_u11 = cvt_e4m3x4_to_f32x4( - ld_global_nc_u32(bsf_addr_u_c) - ) + if cutlass.const_expr(self.scale_format_e8m0_k32): + kbb_u = (lane * Int32(cfg.k_segments)) >> Int32(1) + nc_u = Int32(cfg.two_n) + u_k0 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_u + Int32(0), + row_u, + nc_u, + Int32(cfg.k_dim // 32), + ) + u_k1 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_u + Int32(1), + row_u, + nc_u, + Int32(cfg.k_dim // 32), + ) + u_k2 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_u + Int32(2), + row_u, + nc_u, + Int32(cfg.k_dim // 32), + ) + u_k3 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_u + Int32(3), + row_u, + nc_u, + Int32(cfg.k_dim // 32), + ) + u_k4 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_u + Int32(4), + row_u, + nc_u, + Int32(cfg.k_dim // 32), + ) + u_k5 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_u + Int32(5), + row_u, + nc_u, + Int32(cfg.k_dim // 32), + ) + sf_u0 = u_k0 + sf_u1 = u_k0 + sf_u2 = u_k1 + sf_u3 = u_k1 + sf_u4 = u_k2 + sf_u5 = u_k2 + sf_u6 = u_k3 + sf_u7 = u_k3 + sf_u8 = u_k4 + sf_u9 = u_k4 + sf_u10 = u_k5 + sf_u11 = u_k5 + elif cutlass.const_expr(self.w4a16_mode): + k16_u = lane * Int32(cfg.k_segments) + sf_u0 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(0), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u1 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(1), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u2 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(2), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u3 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(3), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u4 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(4), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u5 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(5), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u6 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(6), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u7 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(7), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u8 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(8), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u9 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(9), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u10 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(10), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u11 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(11), + scale_col_u, + Int32(cfg.two_n), + ) + else: + bsf_addr_u_a = ( + w1s_base_addr + ebase_sf + bsf_base_u + col_blk_off + ) + bsf_addr_u_b = bsf_addr_u_a + Int64(512) + bsf_addr_u_c = bsf_addr_u_a + Int64(1024) + sf_u0, sf_u1, sf_u2, sf_u3 = self._scale_word_to_f32x4( + ld_global_nc_u32(bsf_addr_u_a) + ) + sf_u4, sf_u5, sf_u6, sf_u7 = self._scale_word_to_f32x4( + ld_global_nc_u32(bsf_addr_u_b) + ) + sf_u8, sf_u9, sf_u10, sf_u11 = self._scale_word_to_f32x4( + ld_global_nc_u32(bsf_addr_u_c) + ) gw_a0, gw_a1, gw_a2, gw_a3 = ld_global_nc_v4_u32(gate_byte_addr) gw_b0, gw_b1, gw_b2, gw_b3 = ld_global_nc_v4_u32( @@ -1899,18 +3458,170 @@ def kernel( gate_byte_addr + Int64(80) ) - bsf_addr_g_a = w1s_base_addr + ebase_sf + bsf_base_g + col_blk_off - bsf_addr_g_b = bsf_addr_g_a + Int64(512) - bsf_addr_g_c = bsf_addr_g_a + Int64(1024) - sf_g0, sf_g1, sf_g2, sf_g3 = cvt_e4m3x4_to_f32x4( - ld_global_nc_u32(bsf_addr_g_a) - ) - sf_g4, sf_g5, sf_g6, sf_g7 = cvt_e4m3x4_to_f32x4( - ld_global_nc_u32(bsf_addr_g_b) - ) - sf_g8, sf_g9, sf_g10, sf_g11 = cvt_e4m3x4_to_f32x4( - ld_global_nc_u32(bsf_addr_g_c) - ) + if cutlass.const_expr(self.scale_format_e8m0_k32): + kbb_g = (lane * Int32(cfg.k_segments)) >> Int32(1) + nc_g = Int32(cfg.two_n) + g_k0 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_g + Int32(0), + row_g, + nc_g, + Int32(cfg.k_dim // 32), + ) + g_k1 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_g + Int32(1), + row_g, + nc_g, + Int32(cfg.k_dim // 32), + ) + g_k2 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_g + Int32(2), + row_g, + nc_g, + Int32(cfg.k_dim // 32), + ) + g_k3 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_g + Int32(3), + row_g, + nc_g, + Int32(cfg.k_dim // 32), + ) + g_k4 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_g + Int32(4), + row_g, + nc_g, + Int32(cfg.k_dim // 32), + ) + g_k5 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + kbb_g + Int32(5), + row_g, + nc_g, + Int32(cfg.k_dim // 32), + ) + sf_g0 = g_k0 + sf_g1 = g_k0 + sf_g2 = g_k1 + sf_g3 = g_k1 + sf_g4 = g_k2 + sf_g5 = g_k2 + sf_g6 = g_k3 + sf_g7 = g_k3 + sf_g8 = g_k4 + sf_g9 = g_k4 + sf_g10 = g_k5 + sf_g11 = g_k5 + elif cutlass.const_expr(self.w4a16_mode): + k16_g = lane * Int32(cfg.k_segments) + sf_g0 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(0), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g1 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(1), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g2 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(2), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g3 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(3), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g4 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(4), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g5 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(5), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g6 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(6), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g7 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(7), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g8 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(8), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g9 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(9), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g10 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(10), + scale_col_g, + Int32(cfg.two_n), + ) + sf_g11 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(11), + scale_col_g, + Int32(cfg.two_n), + ) + else: + bsf_addr_g_a = ( + w1s_base_addr + ebase_sf + bsf_base_g + col_blk_off + ) + bsf_addr_g_b = bsf_addr_g_a + Int64(512) + bsf_addr_g_c = bsf_addr_g_a + Int64(1024) + sf_g0, sf_g1, sf_g2, sf_g3 = self._scale_word_to_f32x4( + ld_global_nc_u32(bsf_addr_g_a) + ) + sf_g4, sf_g5, sf_g6, sf_g7 = self._scale_word_to_f32x4( + ld_global_nc_u32(bsf_addr_g_b) + ) + sf_g8, sf_g9, sf_g10, sf_g11 = self._scale_word_to_f32x4( + ld_global_nc_u32(bsf_addr_g_c) + ) if cutlass.const_expr(not self.is_gated): partial_gate = ( @@ -2028,7 +3739,7 @@ def kernel( + sf_g10 * dot_g10 + sf_g11 * dot_g11 ) - elif cutlass.const_expr(cfg.k_segments == 2): + elif cutlass.const_expr(cfg.k_segments_aligned and cfg.k_segments == 2): if cutlass.const_expr(self.is_gated): uw_a0, uw_a1, uw_a2, uw_a3 = ld_global_nc_v4_u32(up_byte_addr) gw_a0, gw_a1, gw_a2, gw_a3 = ld_global_nc_v4_u32(gate_byte_addr) @@ -2038,24 +3749,78 @@ def kernel( sf_shift1 = sf_shift0 + Uint32(8) if cutlass.const_expr(self.is_gated): - sf_word_u = ld_global_nc_u32( - w1s_base_addr + ebase_sf + bsf_base_u + col_blk_off_2 + if cutlass.const_expr(self.scale_format_e8m0_k32): + sf_u0 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + lane, + row_u, + Int32(cfg.two_n), + Int32(cfg.k_dim // 32), + ) + sf_u1 = sf_u0 + elif cutlass.const_expr(self.w4a16_mode): + k16_u = lane * Int32(2) + sf_u0 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(0), + scale_col_u, + Int32(cfg.two_n), + ) + sf_u1 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_u + Int32(1), + scale_col_u, + Int32(cfg.two_n), + ) + else: + sf_word_u = ld_global_nc_u32( + w1s_base_addr + ebase_sf + bsf_base_u + col_blk_off_2 + ) + sf_u0 = self._scale_byte_to_f32( + (sf_word_u >> sf_shift0) & Uint32(0xFF) + ) + sf_u1 = self._scale_byte_to_f32( + (sf_word_u >> sf_shift1) & Uint32(0xFF) + ) + if cutlass.const_expr(self.scale_format_e8m0_k32): + sf_g0 = self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + lane, + row_g, + Int32(cfg.two_n), + Int32(cfg.k_dim // 32), ) - sf_u0 = cvt_e4m3_to_f32_via_f16( - (sf_word_u >> sf_shift0) & Uint32(0xFF) + sf_g1 = sf_g0 + elif cutlass.const_expr(self.w4a16_mode): + k16_g = lane * Int32(2) + sf_g0 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(0), + scale_col_g, + Int32(cfg.two_n), ) - sf_u1 = cvt_e4m3_to_f32_via_f16( - (sf_word_u >> sf_shift1) & Uint32(0xFF) + sf_g1 = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + k16_g + Int32(1), + scale_col_g, + Int32(cfg.two_n), + ) + else: + sf_word_g = ld_global_nc_u32( + w1s_base_addr + ebase_sf + bsf_base_g + col_blk_off_2 + ) + sf_g0 = self._scale_byte_to_f32( + (sf_word_g >> sf_shift0) & Uint32(0xFF) + ) + sf_g1 = self._scale_byte_to_f32( + (sf_word_g >> sf_shift1) & Uint32(0xFF) ) - sf_word_g = ld_global_nc_u32( - w1s_base_addr + ebase_sf + bsf_base_g + col_blk_off_2 - ) - sf_g0 = cvt_e4m3_to_f32_via_f16( - (sf_word_g >> sf_shift0) & Uint32(0xFF) - ) - sf_g1 = cvt_e4m3_to_f32_via_f16( - (sf_word_g >> sf_shift1) & Uint32(0xFF) - ) seg_blk0 = lane * Int32(2) seg_blk1 = seg_blk0 + Int32(1) @@ -2089,66 +3854,577 @@ def kernel( partial_gate = sf_g0 * (dot_g0a + dot_g0b) + sf_g1 * ( dot_g1a + dot_g1b ) - else: - partial_up = Float32(0.0) + elif cutlass.const_expr( + (not self.is_gated) + and (not cfg.k_segments_aligned) + and cfg.k_segments == 6 + and cfg.k_blocks == 168 + ): + tail_seg_count = Int32(6) + lane_seg_base_tail = lane * Int32(6) + if lane >= Int32(16): + if lane < Int32(24): + tail_seg_count = Int32(5) + lane_seg_base_tail = Int32(96) + (lane - Int32(16)) * Int32( + 5 + ) + else: + tail_seg_count = Int32(4) + lane_seg_base_tail = Int32(136) + ( + lane - Int32(24) + ) * Int32(4) + lane_pad_base_tail = lane_seg_base_tail // Int32(8) + xh_base_tail = ( + xh_buf_base + + lane_seg_base_tail * Int32(_BLOCK_SIZE // 2) + + lane_pad_base_tail + ) + xh_off0 = Int32(0) + xh_off1 = Int32(8) + ( + (lane_seg_base_tail + Int32(1)) // Int32(8) - lane_pad_base_tail + ) + xh_off2 = Int32(16) + ( + (lane_seg_base_tail + Int32(2)) // Int32(8) - lane_pad_base_tail + ) + xh_off3 = Int32(24) + ( + (lane_seg_base_tail + Int32(3)) // Int32(8) - lane_pad_base_tail + ) + xh_off4 = Int32(32) + ( + (lane_seg_base_tail + Int32(4)) // Int32(8) - lane_pad_base_tail + ) + xh_off5 = Int32(40) + ( + (lane_seg_base_tail + Int32(5)) // Int32(8) - lane_pad_base_tail + ) + + gate_row_addr_u6 = ( + w1_base_addr + ebase_w + Int64(row_g) * Int64(cfg.k_half) + ) + lane_byte_base = Int64(lane_seg_base_tail) * Int64(_BLOCK_SIZE // 2) + gw_a0 = Uint32(0) + gw_a1 = Uint32(0) + gw_a2 = Uint32(0) + gw_a3 = Uint32(0) + gw_b0 = Uint32(0) + gw_b1 = Uint32(0) + gw_b2 = Uint32(0) + gw_b3 = Uint32(0) + gw_c0 = Uint32(0) + gw_c1 = Uint32(0) + gw_c2 = Uint32(0) + gw_c3 = Uint32(0) + if tail_seg_count == Int32(5): + gw_a0 = ld_global_nc_u32(gate_row_addr_u6 + lane_byte_base) + gw_a1 = ld_global_nc_u32( + gate_row_addr_u6 + lane_byte_base + Int64(4) + ) + gw_a2 = ld_global_nc_u32( + gate_row_addr_u6 + lane_byte_base + Int64(8) + ) + gw_a3 = ld_global_nc_u32( + gate_row_addr_u6 + lane_byte_base + Int64(12) + ) + gw_b0 = ld_global_nc_u32( + gate_row_addr_u6 + lane_byte_base + Int64(16) + ) + gw_b1 = ld_global_nc_u32( + gate_row_addr_u6 + lane_byte_base + Int64(20) + ) + gw_b2 = ld_global_nc_u32( + gate_row_addr_u6 + lane_byte_base + Int64(24) + ) + gw_b3 = ld_global_nc_u32( + gate_row_addr_u6 + lane_byte_base + Int64(28) + ) + gw_c0 = ld_global_nc_u32( + gate_row_addr_u6 + lane_byte_base + Int64(32) + ) + gw_c1 = ld_global_nc_u32( + gate_row_addr_u6 + lane_byte_base + Int64(36) + ) + else: + gw_a0, gw_a1, gw_a2, gw_a3 = ld_global_nc_v4_u32( + gate_row_addr_u6 + lane_byte_base + ) + gw_b0, gw_b1, gw_b2, gw_b3 = ld_global_nc_v4_u32( + gate_row_addr_u6 + lane_byte_base + Int64(16) + ) + if tail_seg_count == Int32(6): + gw_c0, gw_c1, gw_c2, gw_c3 = ld_global_nc_v4_u32( + gate_row_addr_u6 + lane_byte_base + Int64(32) + ) + + scale_pair_off = Int64(lane_seg_base_tail // Int32(4)) * Int64(512) + scale_lane_mod = lane_seg_base_tail % Int32(4) + sf_g0 = Float32(0.0) + sf_g1 = Float32(0.0) + sf_g2 = Float32(0.0) + sf_g3 = Float32(0.0) + sf_g4 = Float32(0.0) + sf_g5 = Float32(0.0) + if cutlass.const_expr(self.w4a16_mode): + sf_g0 = ( + self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base_tail + Int32(0), + scale_col_g, + Int32(cfg.two_n), + ) + if tail_seg_count > Int32(0) + else Float32(0.0) + ) + sf_g1 = ( + self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base_tail + Int32(1), + scale_col_g, + Int32(cfg.two_n), + ) + if tail_seg_count > Int32(1) + else Float32(0.0) + ) + sf_g2 = ( + self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base_tail + Int32(2), + scale_col_g, + Int32(cfg.two_n), + ) + if tail_seg_count > Int32(2) + else Float32(0.0) + ) + sf_g3 = ( + self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base_tail + Int32(3), + scale_col_g, + Int32(cfg.two_n), + ) + if tail_seg_count > Int32(3) + else Float32(0.0) + ) + sf_g4 = ( + self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base_tail + Int32(4), + scale_col_g, + Int32(cfg.two_n), + ) + if tail_seg_count > Int32(4) + else Float32(0.0) + ) + sf_g5 = ( + self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + lane_seg_base_tail + Int32(5), + scale_col_g, + Int32(cfg.two_n), + ) + if tail_seg_count > Int32(5) + else Float32(0.0) + ) + else: + sf_word_g_a = ld_global_nc_u32( + w1s_base_addr + ebase_sf + bsf_base_g + scale_pair_off + ) + sf_word_g_b = ( + ld_global_nc_u32( + w1s_base_addr + + ebase_sf + + bsf_base_g + + scale_pair_off + + Int64(512) + ) + if tail_seg_count > Int32(4) + else Uint32(0) + ) + sf_a0, sf_a1, sf_a2, sf_a3 = self._scale_word_to_f32x4( + sf_word_g_a + ) + sf_b0, sf_b1, sf_b2, sf_b3 = self._scale_word_to_f32x4( + sf_word_g_b + ) + if tail_seg_count == Int32(6): + if scale_lane_mod == Int32(0): + sf_g0 = sf_a0 + sf_g1 = sf_a1 + sf_g2 = sf_a2 + sf_g3 = sf_a3 + sf_g4 = sf_b0 + sf_g5 = sf_b1 + else: + sf_g0 = sf_a2 + sf_g1 = sf_a3 + sf_g2 = sf_b0 + sf_g3 = sf_b1 + sf_g4 = sf_b2 + sf_g5 = sf_b3 + elif tail_seg_count == Int32(5): + if scale_lane_mod == Int32(0): + sf_g0 = sf_a0 + sf_g1 = sf_a1 + sf_g2 = sf_a2 + sf_g3 = sf_a3 + sf_g4 = sf_b0 + elif scale_lane_mod == Int32(1): + sf_g0 = sf_a1 + sf_g1 = sf_a2 + sf_g2 = sf_a3 + sf_g3 = sf_b0 + sf_g4 = sf_b1 + elif scale_lane_mod == Int32(2): + sf_g0 = sf_a2 + sf_g1 = sf_a3 + sf_g2 = sf_b0 + sf_g3 = sf_b1 + sf_g4 = sf_b2 + else: + sf_g0 = sf_a3 + sf_g1 = sf_b0 + sf_g2 = sf_b1 + sf_g3 = sf_b2 + sf_g4 = sf_b3 + else: + sf_g0 = sf_a0 + sf_g1 = sf_a1 + sf_g2 = sf_a2 + sf_g3 = sf_a3 + partial_gate = Float32(0.0) - for seg in cutlass.range_constexpr(cfg.k_segments): - seg_byte_off = Int64(seg * (_BLOCK_SIZE // 2)) - scale_col = lane * Int32(cfg.k_segments) + Int32(seg) - sf_group_off = Int64(scale_col // Int32(4)) * Int64(512) - sf_shift = Uint32((scale_col % Int32(4)) * Int32(8)) - xh_base = ( - xh_buf_base - + scale_col * Int32(_BLOCK_SIZE // 2) - + scale_col // Int32(8) - ) - - gw0 = ld_global_nc_u32(gate_byte_addr + seg_byte_off) - gw1 = ld_global_nc_u32(gate_byte_addr + seg_byte_off + Int64(4)) - sf_word_g = ld_global_nc_u32( - w1s_base_addr + ebase_sf + bsf_base_g + sf_group_off + if tail_seg_count == Int32(6): + partial_gate = ( + sf_g0 + * self._block_dot_hfma2_for_math( + gw_a0, gw_a1, smem_xh, xh_base_tail + xh_off0 + ) + + sf_g1 + * self._block_dot_hfma2_for_math( + gw_a2, gw_a3, smem_xh, xh_base_tail + xh_off1 + ) + + sf_g2 + * self._block_dot_hfma2_for_math( + gw_b0, gw_b1, smem_xh, xh_base_tail + xh_off2 + ) + + sf_g3 + * self._block_dot_hfma2_for_math( + gw_b2, gw_b3, smem_xh, xh_base_tail + xh_off3 + ) + + sf_g4 + * self._block_dot_hfma2_for_math( + gw_c0, gw_c1, smem_xh, xh_base_tail + xh_off4 + ) + + sf_g5 + * self._block_dot_hfma2_for_math( + gw_c2, gw_c3, smem_xh, xh_base_tail + xh_off5 + ) + ) + elif tail_seg_count == Int32(5): + partial_gate = ( + sf_g0 + * self._block_dot_hfma2_for_math( + gw_a0, gw_a1, smem_xh, xh_base_tail + xh_off0 + ) + + sf_g1 + * self._block_dot_hfma2_for_math( + gw_a2, gw_a3, smem_xh, xh_base_tail + xh_off1 + ) + + sf_g2 + * self._block_dot_hfma2_for_math( + gw_b0, gw_b1, smem_xh, xh_base_tail + xh_off2 + ) + + sf_g3 + * self._block_dot_hfma2_for_math( + gw_b2, gw_b3, smem_xh, xh_base_tail + xh_off3 + ) + + sf_g4 + * self._block_dot_hfma2_for_math( + gw_c0, gw_c1, smem_xh, xh_base_tail + xh_off4 + ) ) - sf_g = cvt_e4m3_to_f32_via_f16( - (sf_word_g >> sf_shift) & Uint32(0xFF) + else: + partial_gate = ( + sf_g0 + * self._block_dot_hfma2_for_math( + gw_a0, gw_a1, smem_xh, xh_base_tail + xh_off0 + ) + + sf_g1 + * self._block_dot_hfma2_for_math( + gw_a2, gw_a3, smem_xh, xh_base_tail + xh_off1 + ) + + sf_g2 + * self._block_dot_hfma2_for_math( + gw_b0, gw_b1, smem_xh, xh_base_tail + xh_off2 + ) + + sf_g3 + * self._block_dot_hfma2_for_math( + gw_b2, gw_b3, smem_xh, xh_base_tail + xh_off3 + ) ) + else: + partial_up = Float32(0.0) + partial_gate = Float32(0.0) + if cutlass.const_expr(cfg.k_segments_aligned): + for seg in cutlass.range_constexpr(cfg.k_segments): + seg_byte_off = Int64(seg * (_BLOCK_SIZE // 2)) + scale_col = lane * Int32(cfg.k_segments) + Int32(seg) + sf_group_off = Int64(scale_col // Int32(4)) * Int64(512) + sf_shift = Uint32((scale_col % Int32(4)) * Int32(8)) + xh_base = ( + xh_buf_base + + scale_col * Int32(_BLOCK_SIZE // 2) + + scale_col // Int32(8) + ) + + gw0 = ld_global_nc_u32(gate_byte_addr + seg_byte_off) + gw1 = ld_global_nc_u32( + gate_byte_addr + seg_byte_off + Int64(4) + ) + if cutlass.const_expr(self.w4a16_mode): + sf_g = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + scale_col, + scale_col_g, + Int32(cfg.two_n), + ) + else: + sf_word_g = ld_global_nc_u32( + w1s_base_addr + ebase_sf + bsf_base_g + sf_group_off + ) + sf_g = self._scale_byte_to_f32( + (sf_word_g >> sf_shift) & Uint32(0xFF) + ) + if cutlass.const_expr(self.is_gated): + uw0 = ld_global_nc_u32(up_byte_addr + seg_byte_off) + uw1 = ld_global_nc_u32( + up_byte_addr + seg_byte_off + Int64(4) + ) + if cutlass.const_expr(self.w4a16_mode): + sf_u = self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + scale_col, + scale_col_u, + Int32(cfg.two_n), + ) + else: + sf_word_u = ld_global_nc_u32( + w1s_base_addr + + ebase_sf + + bsf_base_u + + sf_group_off + ) + sf_u = self._scale_byte_to_f32( + (sf_word_u >> sf_shift) & Uint32(0xFF) + ) + dot_u, dot_g = self._block_dot_hfma2_pair_for_math( + uw0, uw1, gw0, gw1, smem_xh, xh_base + ) + partial_up = partial_up + sf_u * dot_u + partial_gate = partial_gate + sf_g * dot_g + else: + partial_gate = ( + partial_gate + + sf_g + * self._block_dot_hfma2_for_math( + gw0, + gw1, + smem_xh, + xh_base, + ) + ) + else: + gate_row_addr = ( + w1_base_addr + ebase_w + Int64(row_g) * Int64(cfg.k_half) + ) if cutlass.const_expr(self.is_gated): - uw0 = ld_global_nc_u32(up_byte_addr + seg_byte_off) - uw1 = ld_global_nc_u32( - up_byte_addr + seg_byte_off + Int64(4) + up_row_addr = ( + w1_base_addr + + ebase_w + + Int64(row_u) * Int64(cfg.k_half) ) - sf_word_u = ld_global_nc_u32( - w1s_base_addr + ebase_sf + bsf_base_u + sf_group_off + for seg in cutlass.range_constexpr(cfg.k_segments): + scale_col = lane * Int32(cfg.k_segments) + Int32(seg) + valid_seg = ( + Int32(1) + if scale_col < Int32(cfg.k_blocks) + else Int32(0) ) - sf_u = cvt_e4m3_to_f32_via_f16( - (sf_word_u >> sf_shift) & Uint32(0xFF) + seg_byte_off = Int64(scale_col) * Int64(_BLOCK_SIZE // 2) + sf_group_off = Int64(scale_col // Int32(4)) * Int64(512) + sf_shift = Uint32((scale_col % Int32(4)) * Int32(8)) + xh_base = ( + xh_buf_base + + scale_col * Int32(_BLOCK_SIZE // 2) + + scale_col // Int32(8) ) - dot_u, dot_g = self._block_dot_hfma2_pair_for_math( - uw0, uw1, gw0, gw1, smem_xh, xh_base + xh_base = xh_base if valid_seg > Int32(0) else xh_buf_base + + gw0 = ( + ld_global_nc_u32(gate_row_addr + seg_byte_off) + if valid_seg > Int32(0) + else Uint32(0) ) - partial_up = partial_up + sf_u * dot_u - partial_gate = partial_gate + sf_g * dot_g - else: - partial_gate = ( - partial_gate - + sf_g - * self._block_dot_hfma2_for_math( - gw0, - gw1, - smem_xh, - xh_base, + gw1 = ( + ld_global_nc_u32( + gate_row_addr + seg_byte_off + Int64(4) ) + if valid_seg > Int32(0) + else Uint32(0) ) + if cutlass.const_expr(self.scale_format_e8m0_k32): + sf_g = ( + self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + scale_col >> Int32(1), + row_g, + Int32(cfg.two_n), + Int32(cfg.k_dim // 32), + ) + if valid_seg > Int32(0) + else Float32(0.0) + ) + elif cutlass.const_expr(self.w4a16_mode): + sf_g = ( + self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + scale_col, + scale_col_g, + Int32(cfg.two_n), + ) + if valid_seg > Int32(0) + else Float32(0.0) + ) + else: + sf_word_g = ( + ld_global_nc_u32( + w1s_base_addr + + ebase_sf + + bsf_base_g + + sf_group_off + ) + if valid_seg > Int32(0) + else Uint32(0) + ) + sf_g = ( + self._scale_byte_to_f32( + (sf_word_g >> sf_shift) & Uint32(0xFF) + ) + if valid_seg > Int32(0) + else Float32(0.0) + ) + + if cutlass.const_expr(self.is_gated): + uw0 = ( + ld_global_nc_u32(up_row_addr + seg_byte_off) + if valid_seg > Int32(0) + else Uint32(0) + ) + uw1 = ( + ld_global_nc_u32( + up_row_addr + seg_byte_off + Int64(4) + ) + if valid_seg > Int32(0) + else Uint32(0) + ) + if cutlass.const_expr(self.scale_format_e8m0_k32): + sf_u = ( + self._ld_e8m0_scale( + w1s_base_addr, + ebase_sf_packed, + scale_col >> Int32(1), + row_u, + Int32(cfg.two_n), + Int32(cfg.k_dim // 32), + ) + if valid_seg > Int32(0) + else Float32(0.0) + ) + elif cutlass.const_expr(self.w4a16_mode): + sf_u = ( + self._ld_e4m3_packed_scale_col( + w1s_base_addr, + ebase_sf_packed_e4m3, + scale_col, + scale_col_u, + Int32(cfg.two_n), + ) + if valid_seg > Int32(0) + else Float32(0.0) + ) + else: + sf_word_u = ( + ld_global_nc_u32( + w1s_base_addr + + ebase_sf + + bsf_base_u + + sf_group_off + ) + if valid_seg > Int32(0) + else Uint32(0) + ) + sf_u = ( + self._scale_byte_to_f32( + (sf_word_u >> sf_shift) & Uint32(0xFF) + ) + if valid_seg > Int32(0) + else Float32(0.0) + ) + dot_u, dot_g = self._block_dot_hfma2_pair_for_math( + uw0, uw1, gw0, gw1, smem_xh, xh_base + ) + partial_up = partial_up + sf_u * dot_u + partial_gate = partial_gate + sf_g * dot_g + else: + partial_gate = ( + partial_gate + + sf_g + * self._block_dot_hfma2_for_math( + gw0, + gw1, + smem_xh, + xh_base, + ) + ) # ---- Activation + intermediate quant ---- gate_red = cute.arch.warp_reduction_sum(partial_gate) * alpha_fc1 if cutlass.const_expr(self.is_gated): up_red = cute.arch.warp_reduction_sum(partial_up) * alpha_fc1 if lane == Int32(0): if cutlass.const_expr(self.is_gated): + if cutlass.const_expr(self.has_swiglu_limit): + limit = Float32(self.swiglu_limit) + neg_limit = Float32(-self.swiglu_limit) + if gate_red > limit: + gate_red = limit + if up_red > limit: + up_red = limit + if up_red < neg_limit: + up_red = neg_limit + sigmoid_arg = gate_red + up_term = up_red + if cutlass.const_expr(self.is_swigluoai): + sigmoid_arg = Float32(self.swiglu_alpha) * gate_red + up_term = up_red + Float32(self.swiglu_beta) + elif cutlass.const_expr(self.is_gelu_tanh): + # sigmoid(2z) = 0.5*(1+tanh(z)) turns the shared + # sigmoid*gate*up form into tanh-approx GELU. + sigmoid_arg = Float32(2.0 * 0.7978845608028654) * ( + gate_red + + Float32(0.044715) * gate_red * gate_red * gate_red + ) sigmoid = Float32(1.0) / ( - Float32(1.0) + cute.math.exp(-gate_red, fastmath=False) + Float32(1.0) + cute.math.exp(-sigmoid_arg, fastmath=False) ) - activated = sigmoid * gate_red * up_red + activated = sigmoid * gate_red * up_term else: relu_val = fmax_f32(gate_red, Float32(0.0)) activated = relu_val * relu_val @@ -2184,7 +4460,28 @@ def kernel( smem_xh[next_buf_base + phys_base + Int32(i)] = ( pack_f32x2_to_f16x2(v0, v1) ) - else: + if cutlass.const_expr(self.a8_mx_mode): + # Per-32 UE8M0 + E4M3 quantize-dequant (w4a8 prefill numerics). + blk_peak = Float32(0.0) + pair_delta = ( + Int32(1) - Int32(2) * (in_blk & Int32(1)) + ) * Int32(_BLOCK_SIZE) + for i in cutlass.range_constexpr(_BLOCK_SIZE): + v = Float32(a_input[x_base + Int32(i)]) + w = Float32(a_input[x_base + pair_delta + Int32(i)]) + blk_peak = fmax_f32(blk_peak, fmax_f32(v, -v)) + blk_peak = fmax_f32(blk_peak, fmax_f32(w, -w)) + scale32, inv32 = mx_scale_from_amax32(blk_peak) + for i in cutlass.range_constexpr(_BLOCK_SIZE // 2): + v0 = Float32(a_input[x_base + Int32(i * 2)]) + v1 = Float32(a_input[x_base + Int32(i * 2 + 1)]) + f0, f1 = quant_dequant_e4m3_2(v0, v1, inv32, scale32) + smem_xh[next_buf_base + phys_base + Int32(i)] = ( + pack_f32x2_to_f16x2(f0, f1) + ) + if cutlass.const_expr( + (not self.w4a16_mode) and (not self.a8_mx_mode) + ): blk_peak = Float32(0.0) for i in cutlass.range_constexpr(_BLOCK_SIZE): v = Float32(a_input[x_base + Int32(i)]) @@ -2196,7 +4493,7 @@ def kernel( q_scale = nvfp4_scale_from_amax(blk_peak, gs_fc1_next) if q_scale > Float32(_FP8_E4M3_MAX): q_scale = Float32(_FP8_E4M3_MAX) - sf_val = cvt_e4m3_to_f32_via_f16(cvt_f32_to_e4m3(q_scale)) + sf_val = self._scale_byte_to_f32(cvt_f32_to_e4m3(q_scale)) eff_scale = Float32(0.0) if gs_fc1_next != Float32(0.0): eff_scale = sf_val / gs_fc1_next @@ -2248,6 +4545,44 @@ def kernel( fc2_rescale = gs_fc2 / gs_fc2_eff if tidx < Int32(cfg.inter_blocks): mid_blk = tidx + if cutlass.const_expr(self.a8_mx_mode): + # Per-32 UE8M0 + E4M3 quantize-dequant of the FC2 input + # (self-ranging: no global scale, no dynamic rescale). + blk_peak = Float32(0.0) + pair_delta = (Int32(1) - Int32(2) * (mid_blk & Int32(1))) * Int32( + _BLOCK_SIZE + ) + for i in cutlass.range_constexpr(_BLOCK_SIZE): + v = smem_int[mid_blk * Int32(_BLOCK_SIZE) + Int32(i)] + w = smem_int[ + mid_blk * Int32(_BLOCK_SIZE) + pair_delta + Int32(i) + ] + blk_peak = fmax_f32(blk_peak, fmax_f32(v, -v)) + blk_peak = fmax_f32(blk_peak, fmax_f32(w, -w)) + scale32, inv32 = mx_scale_from_amax32(blk_peak) + for i in cutlass.range_constexpr(_BLOCK_SIZE // 2): + v0 = smem_int[mid_blk * Int32(_BLOCK_SIZE) + Int32(i * 2)] + v1 = smem_int[mid_blk * Int32(_BLOCK_SIZE) + Int32(i * 2 + 1)] + f0, f1 = quant_dequant_e4m3_2(v0, v1, inv32, scale32) + # The combined nvfp4 FC2 alpha is 1/(gs_fc2 * gs_w2); + # a8_mx quantizes without the global scale, so fold it + # into the dequantized values here (alpha-equivalent). + f0 = f0 * gs_fc2 + f1 = f1 * gs_fc2 + half_base = chunk_idx * Int32( + cfg.i_chunk // 2 + ) + mid_blk * Int32(_BLOCK_SIZE // 2) + n_blk = half_base // Int32(128) + h_local = half_base - n_blk * Int32(128) + h_i = h_local + Int32(i) + packed_idx = ( + t * Int32(cfg.inter_u32) + + k_idx * Int32(cfg.fc2_n_chunks * 128) + + n_blk * Int32(128) + + (h_i % Int32(4)) * Int32(32) + + (h_i // Int32(4)) + ) + intermediate[packed_idx] = pack_f32x2_to_f16x2(f0, f1) if cutlass.const_expr(self.w4a16_mode): for i in cutlass.range_constexpr(_BLOCK_SIZE // 2): v0 = smem_int[mid_blk * Int32(_BLOCK_SIZE) + Int32(i * 2)] @@ -2266,7 +4601,7 @@ def kernel( + (h_i // Int32(4)) ) intermediate[packed_idx] = pack_f32x2_to_f16x2(v0, v1) - else: + if cutlass.const_expr((not self.w4a16_mode) and (not self.a8_mx_mode)): blk_peak = Float32(0.0) for i in cutlass.range_constexpr(_BLOCK_SIZE): v = smem_int[mid_blk * Int32(_BLOCK_SIZE) + Int32(i)] @@ -2278,7 +4613,7 @@ def kernel( q_scale = nvfp4_scale_from_amax(blk_peak, gs_fc2_eff) if q_scale > Float32(_FP8_E4M3_MAX): q_scale = Float32(_FP8_E4M3_MAX) - sf_val = cvt_e4m3_to_f32_via_f16(cvt_f32_to_e4m3(q_scale)) + sf_val = self._scale_byte_to_f32(cvt_f32_to_e4m3(q_scale)) eff_scale = Float32(0.0) if gs_fc2_eff != Float32(0.0): eff_scale = sf_val / gs_fc2_eff @@ -2311,6 +4646,9 @@ def kernel( buf_idx = Int32(1) - buf_idx fc1_task += Int32(gdim_x) + if cutlass.const_expr(self.compile_time_phase == 1): + return + if cutlass.const_expr(self.m_const == 1): _token_publish_fc1_ready( barrier_count, @@ -2326,14 +4664,45 @@ def kernel( barrier_count, barrier_epoch, Int32(gdim_x), is_cta_leader ) - # =================================================================== - # PHASE 2: FC2 output - # =================================================================== + self._run_fc2( + bidx_x, + gdim_x, + warp_id, + lane, + m_val, + w2_weights, + w2_scales, + w2_alphas, + intermediate, + topk_ids, + topk_weights, + scatter_output, + ) + + @cute.jit + def _run_fc2( + self, + bidx_x: Int32, + gdim_x: Int32, + warp_id: Int32, + lane: Int32, + m_val: Int32, + w2_weights: cute.Tensor, + w2_scales: cute.Tensor, + w2_alphas: cute.Tensor, + intermediate: cute.Tensor, + topk_ids: cute.Tensor, + topk_weights: cute.Tensor, + scatter_output: cute.Tensor, + ): + # FC2 is factored out so it can also run as a second, non-cooperative + # launch after a standalone FC1 phase. + cfg = self._cfg w2_base_addr = w2_weights.iterator.toint() w2s_base_addr = w2_scales.iterator.toint() # ---- m==1 FC2 rowpair ---- if cutlass.const_expr(self.m_const == 1): - fc2_chunks_m1 = Int32(cfg.k_dim // (_K_PER_CTA * 2)) + fc2_chunks_m1 = Int32(cfg.k_dim // self.m1_fc2_rows_per_cta) if cutlass.const_expr(self.m1_fc2_onepass): fc2_task = Int32(bidx_x) if fc2_task < fc2_chunks_m1: @@ -2396,9 +4765,10 @@ def kernel( # ---- m>=2 FC2 rowquad ---- else: - fc2_task_count = Int32((self.m_const * cfg.k_dim) // (_K_PER_CTA * 4)) if cutlass.const_expr(self.w4a16_mode and cfg.fc2_n_chunks == 1): - fc2_task_count = Int32((self.m_const * cfg.k_dim) // (_K_PER_CTA * 2)) + fc2_task_count = (m_val * Int32(cfg.k_dim)) // Int32(_K_PER_CTA * 2) + else: + fc2_task_count = (m_val * Int32(cfg.k_dim)) // Int32(_K_PER_CTA * 4) fc2_task = Int32(bidx_x) while fc2_task < fc2_task_count: if cutlass.const_expr(self.w4a16_mode and cfg.fc2_n_chunks == 1): @@ -2445,7 +4815,7 @@ def kernel( @cute.jit def __call__( self, - x: cute.Tensor, + x_ptr: cute.Pointer, w1_ptr: cute.Pointer, w1s_ptr: cute.Pointer, w1a_ptr: cute.Pointer, @@ -2458,23 +4828,32 @@ def __call__( tid_ptr: cute.Pointer, tw_ptr: cute.Pointer, out_ptr: cute.Pointer, - barrier_count_ptr: cute.Pointer, - barrier_epoch_ptr: cute.Pointer, + barrier_count: cute.Tensor, + barrier_epoch: cute.Tensor, m_val: Int32, grid_x: Int32, stream, ): cfg = self._cfg - a_input = cute.make_tensor( - x.iterator, cute.make_layout(Int32(m_val * cfg.k_dim)) - ) + a_input = cute.make_tensor(x_ptr, cute.make_layout(Int32(m_val * cfg.k_dim))) w1_weights = cute.make_tensor( w1_ptr, cute.make_layout(Int64(cfg.weight_E * cfg.two_n * cfg.k_half)) ) - w1_scales = cute.make_tensor( - w1s_ptr, - cute.make_layout(Int64(cfg.weight_E * cfg.w1_sf_rows * cfg.w1_sf_cols)), - ) + if cutlass.const_expr(self.scale_format_e8m0_k32): + w1_scales = cute.make_tensor( + w1s_ptr, + cute.make_layout(Int64(cfg.weight_E * (cfg.k_dim // 32) * cfg.two_n)), + ) + elif cutlass.const_expr(self.w4a16_mode): + w1_scales = cute.make_tensor( + w1s_ptr, + cute.make_layout(Int64(cfg.weight_E * (cfg.k_dim // 16) * cfg.two_n)), + ) + else: + w1_scales = cute.make_tensor( + w1s_ptr, + cute.make_layout(Int64(cfg.weight_E * cfg.w1_sf_rows * cfg.w1_sf_cols)), + ) w1_alphas = cute.make_tensor(w1a_ptr, cute.make_layout(Int32(cfg.weight_E))) input_gs = cute.make_tensor(a1_ptr, cute.make_layout(Int32(cfg.weight_E))) down_input_scale = cute.make_tensor( @@ -2486,10 +4865,21 @@ def __call__( w2_weights = cute.make_tensor( w2_ptr, cute.make_layout(Int64(cfg.weight_E * cfg.k_dim * cfg.n_half)) ) - w2_scales = cute.make_tensor( - w2s_ptr, - cute.make_layout(Int64(cfg.weight_E * cfg.w2_sf_rows * cfg.w2_sf_cols)), - ) + if cutlass.const_expr(self.scale_format_e8m0_k32): + w2_scales = cute.make_tensor( + w2s_ptr, + cute.make_layout(Int64(cfg.weight_E * (cfg.n // 32) * cfg.k_dim)), + ) + elif cutlass.const_expr(self.w4a16_mode): + w2_scales = cute.make_tensor( + w2s_ptr, + cute.make_layout(Int64(cfg.weight_E * (cfg.n // 16) * cfg.k_dim)), + ) + else: + w2_scales = cute.make_tensor( + w2s_ptr, + cute.make_layout(Int64(cfg.weight_E * cfg.w2_sf_rows * cfg.w2_sf_cols)), + ) w2_alphas = cute.make_tensor(w2a_ptr, cute.make_layout(Int32(cfg.weight_E))) topk_ids_tensor = cute.make_tensor( tid_ptr, cute.make_layout(Int32(m_val * cfg.num_topk)) @@ -2500,13 +4890,6 @@ def __call__( scatter_output_tensor = cute.make_tensor( out_ptr, cute.make_layout(Int32(m_val * cfg.k_dim)) ) - barrier_slots = m_val * Int32(cfg.num_topk + 16) - barrier_count = cute.make_tensor( - barrier_count_ptr, cute.make_layout(barrier_slots) - ) - barrier_epoch = cute.make_tensor( - barrier_epoch_ptr, cute.make_layout(barrier_slots) - ) self.kernel( a_input, @@ -2527,8 +4910,15 @@ def __call__( m_val, ).launch( grid=(grid_x, Int32(1), Int32(1)), - block=(_BLOCK_DIM, 1, 1), - smem=0, + block=(self.launch_block_dim, 1, 1), + # The fused and FC1-only bodies use 512-thread CTAs; + # the FC2-only m=1 specialization is a 256-thread independent CTA. + # One block per SM preserves each variant's register budget. + min_blocks_per_mp=1, + # The fused phase crosses a software all-CTA barrier between FC1 + # and FC2, so require whole-grid admission; resident CTAs must not + # spin while peers remain queued. Split phases have no barrier. + cooperative=self.compile_time_phase == 0, stream=stream, ) @@ -2554,28 +4944,212 @@ def launch( m: int, grid_x: int, ): + def ptr(dt, t): + return make_ptr(dt, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + + ids_dtype = cutlass.Int64 if topk_ids.dtype == torch.int64 else cutlass.Int32 stream = current_cuda_stream() compiled_fn( - x, - w1_fp4.data_ptr(), - w1_blockscale.view(torch.uint8).data_ptr(), - w1_alphas.data_ptr(), - a1_gscale.data_ptr(), - a2_gscale.data_ptr(), - inter_fp32.view(torch.uint32).data_ptr(), - w2_fp4.data_ptr(), - w2_blockscale.view(torch.uint8).data_ptr(), - w2_alphas.data_ptr(), - topk_ids.data_ptr(), - topk_weights.data_ptr(), - out.data_ptr(), - barrier_count.data_ptr(), - barrier_epoch.data_ptr(), + ptr(cutlass.BFloat16, x), + ptr(cutlass.Uint8, w1_fp4), + ptr(cutlass.Uint8, w1_blockscale.view(torch.uint8)), + ptr(cutlass.Float32, w1_alphas), + ptr(cutlass.Float32, a1_gscale), + ptr(cutlass.Float32, a2_gscale), + ptr(cutlass.Uint32, inter_fp32.view(torch.uint32)), + ptr(cutlass.Uint8, w2_fp4), + ptr(cutlass.Uint8, w2_blockscale.view(torch.uint8)), + ptr(cutlass.Float32, w2_alphas), + ptr(ids_dtype, topk_ids), + ptr(cutlass.Float32, topk_weights), + ptr(cutlass.BFloat16, out), + barrier_count, + barrier_epoch, Int32(m), Int32(grid_x), stream, ) -__all__ = ["MoEDirectMicroKernel"] +# --------------------------------------------------------------------------- +# Host-side helpers for the dispatch layer +# --------------------------------------------------------------------------- + + +def build_direct_micro_kernel( + weight_E: int, + m: int, + k: int, + n: int, + num_topk: int, + *, + activation: str = "silu", + fast_math: bool = False, + share_input_across_experts: bool = False, + share_expert_scales: bool = False, + single_token: bool = False, + dynamic_down_scale: bool = False, + compile_time_phase: int = 0, + w4a16_mode: bool = False, + a8_mx_mode: bool = False, + scale_format: str = "e4m3_k16", + e8m0_scale_layout: str = "packed", + w13_layout: str = "w13", + swiglu_limit: float | None = None, + swiglu_alpha: float | None = None, + swiglu_beta: float | None = None, + max_active_ctas: int | None = None, + device: torch.device | None = None, +) -> MoEDirectMicroKernel: + """Construct and configure a direct micro kernel for one problem shape. + + Returns the configured (uncompiled) kernel; the caller keys its compile + cache on ``kernel.__cache_key__`` plus the topk_ids dtype and launches + with ``kernel.grid_x``. + """ + kernel = MoEDirectMicroKernel( + 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=dynamic_down_scale, + compile_time_phase=compile_time_phase, + w4a16_mode=w4a16_mode, + a8_mx_mode=a8_mx_mode, + scale_format=scale_format, + e8m0_scale_layout=e8m0_scale_layout, + w13_layout=w13_layout, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + ) + kernel.configure( + m, k, n, num_topk, weight_E, max_active_ctas=max_active_ctas, device=device + ) + return kernel + + +def compile_direct_micro_kernel( + kernel: MoEDirectMicroKernel, + *, + topk_ids_dtype: torch.dtype = torch.int32, + options: str | None = None, +): + """cute.compile a configured direct micro kernel against fake pointers. + + Returns the compiled callable; launch via ``MoEDirectMicroKernel.launch``. + """ + + 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, + ) + compile_kwargs = {} + if options: + compile_kwargs["options"] = options + compile_m = int(kernel.m_const) if int(kernel.m_const) != 0 else 8 + return cute.compile( + kernel, + dummy(cutlass.BFloat16), # x_ptr + dummy(cutlass.Uint8), # w1_ptr + dummy(cutlass.Uint8), # w1s_ptr + dummy(cutlass.Float32), # w1a_ptr + dummy(cutlass.Float32), # a1_ptr + dummy(cutlass.Float32), # a2_ptr + dummy(cutlass.Uint32), # inter_ptr + dummy(cutlass.Uint8), # w2_ptr + dummy(cutlass.Uint8), # w2s_ptr + dummy(cutlass.Float32), # w2a_ptr + dummy(ids_dtype), # tid_ptr + dummy(cutlass.Float32), # tw_ptr + dummy(cutlass.BFloat16), # out_ptr + barrier_fake, # barrier_count + barrier_fake, # barrier_epoch + Int32(compile_m), # m_val + Int32(1), # grid_x + current_cuda_stream(), # stream + **compile_kwargs, + ) + + +_PROBE_FAILURE_WARNED = False + + +def compiled_direct_micro_accepts_block_dim(compiled, block_dim: int) -> bool: + """Return whether the compiled direct micro kernel can launch ``block_dim`` + threads (register pressure can cap CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK + below the 512-thread CTA the fused body wants). Callers should cache the + result per compiled kernel.""" + global _PROBE_FAILURE_WARNED + try: + from cuda.bindings import driver, runtime + + executor = compiled.to(None) + kernel_info = getattr(compiled, "kernel_info", None) or {} + kernel_name = next(iter(kernel_info.keys()), None) + if kernel_name is None and hasattr(compiled, "_get_name"): + kernel_name = compiled._get_name() + if isinstance(kernel_name, str): + kernel_name = kernel_name.encode() + if kernel_name is None: + raise RuntimeError("compiled micro kernel did not expose a kernel name") + + jit_module = getattr(executor, "jit_module", None) + cuda_library = getattr(jit_module, "cuda_library", None) + if isinstance(cuda_library, (list, tuple)): + cuda_library = cuda_library[0] if cuda_library else None + if cuda_library is None: + cuda_library = getattr(executor, "kernel", None) + if cuda_library is None: + raise RuntimeError("compiled micro kernel did not expose a CUDA library") + + err, kernel = runtime.cudaLibraryGetKernel(cuda_library, kernel_name) + if err != runtime.cudaError_t.cudaSuccess: + raise RuntimeError(f"cudaLibraryGetKernel failed with {err}") + cu_kernel = driver.CUkernel(int(kernel)) + err, max_threads = driver.cuKernelGetAttribute( + driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK, + cu_kernel, + 0, + ) + if err != driver.CUresult.CUDA_SUCCESS: + raise RuntimeError(f"cuKernelGetAttribute failed with {err}") + return int(max_threads) >= int(block_dim) + except (AttributeError, KeyError, TypeError): + # Expected introspection misses on this DSL version; fall back to the + # MMA micro kernel silently. + return False + except Exception as exc: + # Anything else means the probe itself broke (e.g. a DSL internals + # change). Warn once so the direct micro backend is not silently + # disabled, but keep the safe fallback. + if not _PROBE_FAILURE_WARNED: + _PROBE_FAILURE_WARNED = True + import warnings + + warnings.warn( + "compiled_direct_micro_accepts_block_dim probe failed " + f"({type(exc).__name__}: {exc}); disabling the direct micro " + "MoE backend for this process.", + RuntimeWarning, + stacklevel=2, + ) + return False + + +__all__ = [ + "MoEDirectMicroKernel", + "build_direct_micro_kernel", + "compile_direct_micro_kernel", + "compiled_direct_micro_accepts_block_dim", +] diff --git a/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py b/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py index 32603b82a92..c912609c42e 100644 --- a/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py +++ b/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dispatch.py @@ -24,8 +24,14 @@ get_num_sm, make_ptr, ) -from .moe_activation import is_gated_activation -from .moe_dynamic_kernel import MoEDynamicKernel +from .moe_activation import SWIGLUOAI_UNINTERLEAVE, is_gated_activation +from .moe_direct_micro_kernel import ( + MoEDirectMicroKernel, + build_direct_micro_kernel, + compile_direct_micro_kernel, + compiled_direct_micro_accepts_block_dim, +) +from .moe_dynamic_kernel import _TASK_SLICE_CHUNK, MoEDynamicKernel from .moe_micro_kernel import MoEMicroKernel from .moe_static_kernel import MoEStaticKernel from .moe_w4a16_fp4_helpers import swizzle_block_scale @@ -49,7 +55,9 @@ _NVFP4_BLOCK_SIZE = 16 _LEVEL_TILE_M = 128 _LEVEL_TILE_N = 128 -_DYNAMIC_SLICE_CHUNK = 1 +# Must equal the kernel's task materialization granularity or the task +# queue is mis-sized. +_DYNAMIC_SLICE_CHUNK = _TASK_SLICE_CHUNK SF_VEC_SIZE = 16 _FORCE_MOE_W4A16_ENV = "FLASHINFER_B12X_FORCE_MOE_W4A16" _MICRO_SHARE_INPUT_ACROSS_EXPERTS = ( @@ -59,6 +67,16 @@ # Micro kernel cutover thresholds (routed pairs) _MICRO_COMPACT_CUTOVER_PAIRS = 20 _MICRO_COMPACT_CUTOVER_PAIRS_MULTI_TOPK = 40 +# The micro kernel's per-token staging assumes decode-sized batches. +_MICRO_MAX_TOKENS = 8 +# Direct micro takes the smallest decode batches ahead of the MMA micro +# kernel, and only at small intermediate sizes where its CUDA-core dots +# keep up with per-token work. Measured on GB10, pending other GPUs. +_DIRECT_MICRO_CUTOVER_PAIRS = 32 +_DIRECT_MICRO_MAX_N = 512 +# Test/bench hook: force one backend ("direct_micro", "micro", "static", +# "dynamic"). Deliberately module-level (a monkeypatch target), not an env var. +_FORCED_BACKEND: str | None = None _STATIC_COMPACT_CUTOVER_PAIRS_DEFAULT = 640 _STATIC_COMPACT_CUTOVER_PAIRS = _STATIC_COMPACT_CUTOVER_PAIRS_DEFAULT _STATIC_COMPACT_CUTOVER_PAIRS_CACHE: Dict[str, int] = {} @@ -90,6 +108,12 @@ (512, 175), (640, 188), ) +# Workloads at or below the static cutover (640 routed pairs by default) +# take the static kernel, so only the 1024 entry is normally reachable. +_DYNAMIC_MAC_LADDER: Tuple[Tuple[int, int], ...] = ( + (640, 188), + (1024, 147), +) def _lookup_mac_ladder( @@ -106,6 +130,20 @@ def _align_up(value: int, alignment: int) -> int: return ((value + alignment - 1) // alignment) * alignment +# The kernels index the packed activation and scale planes with 32-bit +# offsets, so reject any workspace plane large enough to overflow them. +_RUNTIME_MEMREF_LIMIT = (1 << 31) - 1 + + +def _check_memref_limit(name: str, elements: int) -> None: + if elements > _RUNTIME_MEMREF_LIMIT: + raise ValueError( + f"{name} needs {elements} elements, which exceeds the 2^31-1 " + "runtime memref limit. Reduce the token chunk or expert count " + "for this launch." + ) + + def _first_env(*names: str) -> str | None: for name in names: value = os.environ.get(name) @@ -182,14 +220,6 @@ def _is_w4a16(activation_precision: str) -> bool: return _normalize_activation_precision(activation_precision) == "bf16" -def _level_tile_m(activation_precision: str = "fp4") -> int: - if _is_w4a16(activation_precision): - raise ValueError( - "internal routing error: quant_mode='w4a16' reached the NVFP4 tile selector" - ) - return _LEVEL_TILE_M - - def _level_tile_n(activation_precision: str = "fp4") -> int: if _is_w4a16(activation_precision): raise ValueError( @@ -198,6 +228,31 @@ def _level_tile_n(activation_precision: str = "fp4") -> int: return _LEVEL_TILE_N +def _select_dynamic_tile_m( + routed_rows: int, + num_experts: int, + activation: str = "silu", +) -> int: + """Pick the dynamic kernel's M-tile from routed rows per expert. + + Small tiles cut per-expert tail padding for sparse routing; 128 amortizes + best for dense prefill (crossovers measured on gated NVFP4). Workspace + sizing and the kernel build must both derive the tile from this function, + or the scratch is mis-sized for what the kernel indexes. + """ + if not is_gated_activation(activation): + return _LEVEL_TILE_M + routed_rows = max(1, int(routed_rows)) + num_experts = max(1, int(num_experts)) + if routed_rows < 15 * num_experts: + return 16 + if routed_rows < 48 * num_experts: + return 32 + if routed_rows < 96 * num_experts: + return 64 + return _LEVEL_TILE_M + + def _get_static_compact_cutover_pairs(activation_precision: str = "fp4") -> int: activation_precision = _normalize_activation_precision(activation_precision) cached = _STATIC_COMPACT_CUTOVER_PAIRS_CACHE.get(activation_precision) @@ -219,7 +274,12 @@ def _get_static_compact_cutover_pairs(activation_precision: str = "fp4") -> int: return cached -def _select_moe_mma_tiler_mn(routed_rows: int, n: int) -> Tuple[int, int]: +def _select_moe_mma_tiler_mn( + routed_rows: int, + n: int, + *, + resident_clusters: int | None = None, +) -> Tuple[int, int]: """Select optimal MoE tile shape based on routed rows and N dimension. Uses narrower 64x128 tiles when routed_rows <= 128 and default 128x128 @@ -227,13 +287,19 @@ def _select_moe_mma_tiler_mn(routed_rows: int, n: int) -> Tuple[int, int]: """ sm_count = get_num_sm(torch.device("cuda")) coarse_tile = (128, 128) + if routed_rows <= 32 and n <= 256: + return (64, 128) + if resident_clusters is not None and resident_clusters < sm_count: + return coarse_tile coarse_tiles = ((routed_rows + coarse_tile[0] - 1) // coarse_tile[0]) * ( (n + coarse_tile[1] - 1) // coarse_tile[1] ) # Single-token decode often lands exactly on the "half the machine" # boundary. Keeping the coarse 128x128 tile there leaves the M dimension # badly underfilled, so take the narrow 64x128 tile inclusive of equality. - if routed_rows <= 128 and coarse_tiles <= max(1, sm_count // 2): + if routed_rows <= 64 or ( + routed_rows <= 128 and coarse_tiles <= max(1, sm_count // 2) + ): return (64, 128) return (128, 128) @@ -287,6 +353,21 @@ class Sm120StaticMoEWorkspace: packed_a_flat: torch.Tensor | None = None scale_flat: torch.Tensor | None = None + # Direct micro planes (allocated only when the shape can take that path). + dm_barrier_count: torch.Tensor | None = None + dm_barrier_epoch: torch.Tensor | None = None + dm_intermediate: torch.Tensor | None = None + dm_input_gs: torch.Tensor | None = None + dm_down_input_scale: torch.Tensor | None = None + + +def _direct_micro_candidate(k: int, n: int, num_topk: int, weight_E: int) -> bool: + """Whether any m in the tiny-decode band can run the direct micro kernel.""" + return any( + MoEDirectMicroKernel.is_supported(m, k, n, num_topk, weight_E) + for m in range(1, _MICRO_MAX_TOKENS + 1) + ) + def allocate_sm120_static_workspace( *, @@ -309,6 +390,8 @@ def allocate_sm120_static_workspace( rows_pad_k = _align_up(max_rows, 128) cols_pad_k = _align_up(k // _NVFP4_BLOCK_SIZE, 4) + _check_memref_limit("static packed_input", state_E * max_rows * (k // 2)) + _check_memref_limit("static packed_input_scale", state_E * rows_pad_k * cols_pad_k) packed_input = torch.empty( state_E, max_rows, k // 2, dtype=torch.uint8, device=device ) @@ -354,6 +437,33 @@ def allocate_sm120_static_workspace( cute.AddressSpace.gmem, assumed_align=16, ) + + # Direct micro reads weights by global expert id, so its planes are only + # useful without EP remapping. + if state_E == weight_E and _direct_micro_candidate(k, n, num_topk, weight_E): + dm_rows = min(max_rows, _MICRO_MAX_TOKENS * num_topk) + # The epoch-based barriers restore their slots after each launch, so + # the zeroed allocation is the only reset needed (graph-replay safe). + dm_slots = dm_rows + _MICRO_MAX_TOKENS * 16 + fc2_n_chunks = (n // 2 + 127) // 128 + # The fused kernel binds the intermediate as m * num_topk * + # fc2_n_chunks * 128 u32 words; size for the largest supported m. + dm_inter = _MICRO_MAX_TOKENS * num_topk * fc2_n_chunks * 128 + workspace.dm_barrier_count = torch.zeros( + dm_slots, dtype=torch.int32, device=device + ) + workspace.dm_barrier_epoch = torch.zeros( + dm_slots, dtype=torch.int32, device=device + ) + workspace.dm_intermediate = torch.empty( + dm_inter, dtype=torch.float32, device=device + ) + workspace.dm_input_gs = torch.empty( + weight_E, dtype=torch.float32, device=device + ) + workspace.dm_down_input_scale = torch.empty( + weight_E, dtype=torch.float32, device=device + ) return workspace @@ -523,7 +633,7 @@ def _get_static_kernel( routed_rows = m * num_topk mma_tiler_mn = (128, 128) if activation_precision == "fp4" and num_topk > 1: - mma_tiler_mn = _select_moe_mma_tiler_mn(routed_rows, n) + mma_tiler_mn = _select_moe_mma_tiler_mn(routed_rows, n, resident_clusters=mac) cache_key = ( "static", @@ -975,6 +1085,93 @@ def _get_micro_kernel( return result +# The launch cache skips the per-launch build/configure; the kernel cache +# dedupes compiles across keys that configure to the same artifact +# (m=2..8 differ only in grid_x). +_DIRECT_MICRO_LAUNCH_CACHE: Dict[Tuple, Tuple] = {} +_DIRECT_MICRO_KERNEL_CACHE: Dict[Tuple, Tuple] = {} + + +def _get_direct_micro_kernel( + weight_E: int, + m: int, + k: int, + n: int, + num_topk: int, + *, + topk_ids_dtype: torch.dtype = torch.int32, + fast_math: bool = True, + share_input_across_experts: bool = False, + share_expert_scales: bool = False, + activation: str = "silu", + swiglu_alpha: float = 1.702, + swiglu_beta: float = 1.0, + swiglu_limit: float | None = None, + device: torch.device | None = None, +): + """Compile (or retrieve cached) the SM120 direct micro MoE kernel. + + Returns (compiled, grid_x, accepts_block_dim). + """ + if activation != SWIGLUOAI_UNINTERLEAVE: + # The kernel constructor only accepts configurable swiglu parameters + # for swigluoai; other activations use its normalized defaults + # (accept-and-ignore, matching the MMA kernels). + swiglu_alpha = None + swiglu_beta = None + swiglu_limit = None + launch_key = ( + weight_E, + m, + k, + n, + num_topk, + topk_ids_dtype, + fast_math, + share_input_across_experts, + share_expert_scales, + activation, + swiglu_alpha, + swiglu_beta, + swiglu_limit, + str(_canonical_cuda_device(device)) if device is not None else None, + ) + cached = _DIRECT_MICRO_LAUNCH_CACHE.get(launch_key) + if cached is not None: + return cached + kernel = build_direct_micro_kernel( + weight_E, + m, + k, + n, + num_topk, + activation=activation, + fast_math=fast_math, + share_input_across_experts=share_input_across_experts, + share_expert_scales=share_expert_scales, + single_token=m == 1, + swiglu_limit=swiglu_limit, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + device=device, + ) + compile_key = ("direct_micro", kernel.__cache_key__, topk_ids_dtype) + entry = _DIRECT_MICRO_KERNEL_CACHE.get(compile_key) + if entry is None: + compiled = compile_direct_micro_kernel(kernel, topk_ids_dtype=topk_ids_dtype) + # Register pressure can cap the launchable CTA below the fused body's + # 512 threads; probe once per compiled kernel. + accepts = compiled_direct_micro_accepts_block_dim( + compiled, kernel.launch_block_dim + ) + entry = (compiled, accepts) + _DIRECT_MICRO_KERNEL_CACHE[compile_key] = entry + compiled, accepts = entry + cached = (compiled, kernel.grid_x, accepts) + _DIRECT_MICRO_LAUNCH_CACHE[launch_key] = cached + return cached + + # --------------------------------------------------------------------------- # Launch # --------------------------------------------------------------------------- @@ -1012,12 +1209,15 @@ def launch_sm120_static_moe( swiglu_limit: float | None = None, activation_precision: str = "fp4", ) -> torch.Tensor: - """Launch the SM120 static or micro MoE kernel. + """Launch the SM120 static, micro, or direct micro MoE kernel. - Selects the micro kernel for tiny decode batches (routed_rows <= 20-40) - and the static kernel otherwise. The micro path runs a Triton pre-pass - to compact routing IDs before launching. + The direct micro kernel takes tiny decode batches (m <= 8, routed_rows + < 64) when it supports the shape, the MMA micro kernel takes the rest of + its band (routed_rows <= 20-40), and the static kernel takes the rest. + The MMA micro path runs a Triton pre-pass to compact routing IDs before + launching; direct micro routes on global expert ids directly. """ + _check_memref_limit("scatter_output", scatter_output.numel()) activation_precision = _normalize_activation_precision(activation_precision) if activation_precision == "bf16": raise ValueError( @@ -1039,19 +1239,6 @@ def launch_sm120_static_moe( input_gs = _expand_to_experts(input_gs, num_experts) down_input_scale = _expand_to_experts(down_input_scale, num_experts) - # Decide micro vs static - micro_cutover = _MICRO_COMPACT_CUTOVER_PAIRS - if top_k > 1: - micro_cutover = _MICRO_COMPACT_CUTOVER_PAIRS_MULTI_TOPK - use_micro = activation_precision == "fp4" and routed_rows <= micro_cutover - - sm_count = get_num_sm(torch.device("cuda")) - base_mac = min(get_max_active_clusters(1), sm_count) - tuned_static_mac = _lookup_mac_ladder(_STATIC_MAC_LADDER, routed_rows) - static_mac = min(tuned_static_mac or base_mac, base_mac) - if activation_precision == "fp4" and not use_micro and routed_rows < 40: - static_mac = min(static_mac, 64) - # Shared-scale flags let compact W4A4 micro match the ReLU2 single-token # specialization. share_input_across_experts = ( @@ -1064,6 +1251,131 @@ def launch_sm120_static_moe( activation == "relu2" and input_gs_is_shared and down_input_scale_is_shared ) + # Direct micro takes its band before the MMA micro decision. It reads + # weights by global expert id, so EP shapes keep the compact path. + use_direct_micro = ( + activation_precision == "fp4" + and workspace.state_E == num_experts + and workspace.dm_barrier_count is not None + and workspace.dm_barrier_count.numel() >= routed_rows + num_tokens * 16 + and num_tokens <= _MICRO_MAX_TOKENS + and routed_rows < _DIRECT_MICRO_CUTOVER_PAIRS + and n <= _DIRECT_MICRO_MAX_N + and MoEDirectMicroKernel.is_supported(num_tokens, k, n, top_k, num_experts) + ) + if _FORCED_BACKEND is not None: + if _FORCED_BACKEND == "direct_micro": + if workspace.dm_barrier_count is None or not ( + MoEDirectMicroKernel.is_supported(num_tokens, k, n, top_k, num_experts) + ): + raise ValueError( + "forced direct_micro backend cannot run this shape " + f"(m={num_tokens}, k={k}, n={n}, top_k={top_k})" + ) + if workspace.dm_barrier_count.numel() < routed_rows + num_tokens * 16: + raise ValueError( + "forced direct_micro backend exceeds the workspace barrier " + f"capacity ({workspace.dm_barrier_count.numel()} slots < " + f"{routed_rows} routed rows + {num_tokens * 16})" + ) + use_direct_micro = True + else: + use_direct_micro = False + if use_direct_micro: + compiled, grid_x, block_ok = _get_direct_micro_kernel( + num_experts, + num_tokens, + k, + n, + top_k, + topk_ids_dtype=flat_ids.dtype, + fast_math=fast_math, + share_input_across_experts=share_input_across_experts, + share_expert_scales=share_expert_scales, + activation=activation, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + swiglu_limit=swiglu_limit, + device=a.device, + ) + if not block_ok: + if _FORCED_BACKEND == "direct_micro": + raise RuntimeError("compiled direct micro MoE kernel cannot launch") + use_direct_micro = False + if use_direct_micro: + # The kernel takes multiplier-form scales only; invert reciprocal + # inputs into the persistent workspace planes (zeros stay zero, + # matching the MMA kernels). + if input_scales_are_reciprocal: + workspace.dm_input_gs.copy_( + torch.where(input_gs != 0, 1.0 / input_gs, input_gs) + ) + workspace.dm_down_input_scale.copy_( + torch.where( + down_input_scale != 0, 1.0 / down_input_scale, down_input_scale + ) + ) + launch_gs = workspace.dm_input_gs + launch_down = workspace.dm_down_input_scale + else: + launch_gs = input_gs + launch_down = down_input_scale + MoEDirectMicroKernel.launch( + compiled, + x=a, + w1_fp4=weights.w1_storage, + w1_blockscale=weights.w1_scale_storage, + w1_alphas=weights.w1_alpha, + a1_gscale=launch_gs, + a2_gscale=launch_down, + inter_fp32=workspace.dm_intermediate, + w2_fp4=weights.w2_storage, + w2_blockscale=weights.w2_scale_storage, + w2_alphas=weights.w2_alpha, + topk_ids=flat_ids, + topk_weights=flat_weights, + out=scatter_output, + barrier_count=workspace.dm_barrier_count, + barrier_epoch=workspace.dm_barrier_epoch, + m=num_tokens, + grid_x=grid_x, + ) + return scatter_output + + # Decide micro vs static + micro_cutover = _MICRO_COMPACT_CUTOVER_PAIRS + if top_k > 1: + micro_cutover = _MICRO_COMPACT_CUTOVER_PAIRS_MULTI_TOPK + use_micro = ( + activation_precision == "fp4" + and num_tokens <= _MICRO_MAX_TOKENS + and routed_rows <= micro_cutover + ) + if _FORCED_BACKEND is not None: + if _FORCED_BACKEND == "micro": + # Forced mode raises on correctness violations, never falls back. + if num_tokens > _MICRO_MAX_TOKENS: + raise ValueError( + f"forced micro backend supports at most {_MICRO_MAX_TOKENS} " + f"tokens (got {num_tokens})" + ) + if flat_ids.numel() > workspace.compact_topk_ids.numel(): + raise ValueError( + "forced micro backend exceeds the workspace compact-id " + f"capacity ({workspace.compact_topk_ids.numel()} < " + f"{flat_ids.numel()})" + ) + use_micro = True + else: + use_micro = False + + sm_count = get_num_sm(torch.device("cuda")) + base_mac = min(get_max_active_clusters(1), sm_count) + tuned_static_mac = _lookup_mac_ladder(_STATIC_MAC_LADDER, routed_rows) + static_mac = min(tuned_static_mac or base_mac, base_mac) + if activation_precision == "fp4" and not use_micro and routed_rows < 40: + static_mac = min(static_mac, 64) + if use_micro: assert flat_ids.numel() <= workspace.compact_topk_ids.numel(), ( f"compact_topk_ids buffer too small: " @@ -1193,6 +1505,11 @@ def select_sm120_moe_backend( mode = _normalize_quant_mode(quant_mode, activation_precision) if mode == "w4a16": return "w4a16" + if _FORCED_BACKEND == "dynamic": + return "dynamic" + if _FORCED_BACKEND in ("static", "micro", "direct_micro"): + # Both micro variants launch through the static workspace path. + return "static" routed_rows = num_tokens * num_topk if routed_rows <= _get_static_compact_cutover_pairs("fp4"): return "static" @@ -1228,20 +1545,16 @@ class Sm120DynamicMoEWorkspace: routed_rows_capacity: int physical_tiles_capacity: int task_capacity: int + # The M-tile the geometry above was sized for; launches must build the + # kernel with the same tile. + tile_m: int = _LEVEL_TILE_M expert_write_rows: torch.Tensor expert_tile_base: torch.Tensor pair_head: torch.Tensor - producers_done_count: torch.Tensor - all_work_published: torch.Tensor task_head: torch.Tensor task_tail: torch.Tensor - task_ready: torch.Tensor task_expert: torch.Tensor - task_m_tile: torch.Tensor - task_slice_begin: torch.Tensor - task_slice_count: torch.Tensor task_valid_rows: torch.Tensor - tile_write_count: torch.Tensor # Views packed_a_view: torch.Tensor | None = None @@ -1286,6 +1599,7 @@ def allocate_sm120_dynamic_workspace( num_topk: int, device: torch.device, activation_precision: str = "fp4", + activation: str = "silu", ) -> Sm120DynamicMoEWorkspace: """Allocate workspace buffers for the SM120 dynamic MoE kernel.""" activation_precision = _normalize_activation_precision(activation_precision) @@ -1294,7 +1608,7 @@ def allocate_sm120_dynamic_workspace( "allocate_sm120_dynamic_workspace only supports quant_mode='nvfp4'; " "use allocate_sm120_moe_workspace(..., quant_mode='w4a16') for W4A16." ) - tile_m = _level_tile_m(activation_precision) + tile_m = _select_dynamic_tile_m(routed_rows, state_E, activation) physical_tiles, _, max_tasks = _dynamic_task_geometry( state_E, n, @@ -1303,7 +1617,12 @@ def allocate_sm120_dynamic_workspace( tile_n=_level_tile_n(activation_precision), ) rows_padded = physical_tiles * tile_m + # The kernel addresses activation scales in 128-row SF atoms regardless of + # tile_m, so the scale plane must cover the last partial atom. + scale_rows = _align_up(rows_padded, 128) cols_pad_k = _align_up(k // _NVFP4_BLOCK_SIZE, 4) + _check_memref_limit("dynamic packed_input", rows_padded * (k // 2)) + _check_memref_limit("dynamic packed_input_scale", scale_rows * cols_pad_k) packed_input = torch.empty(1, rows_padded, k // 2, dtype=torch.uint8, device=device) workspace = Sm120DynamicMoEWorkspace( @@ -1318,29 +1637,23 @@ def allocate_sm120_dynamic_workspace( routed_rows_capacity=routed_rows, physical_tiles_capacity=physical_tiles, task_capacity=max_tasks, + tile_m=tile_m, row_counts=torch.zeros(state_E, dtype=torch.int32, device=device), token_map=torch.zeros(rows_padded, dtype=torch.int32, device=device), token_weights=torch.zeros(rows_padded, dtype=torch.float32, device=device), packed_input=packed_input, packed_input_scale=torch.empty( - rows_padded, cols_pad_k, dtype=torch.uint8, device=device + scale_rows, cols_pad_k, dtype=torch.uint8, device=device ), barrier_count=torch.zeros(1, dtype=torch.int32, device=device), barrier_epoch=torch.zeros(1, dtype=torch.int32, device=device), expert_write_rows=torch.zeros(state_E, dtype=torch.int32, device=device), expert_tile_base=torch.zeros(state_E + 1, dtype=torch.int32, device=device), pair_head=torch.zeros(1, dtype=torch.int32, device=device), - producers_done_count=torch.zeros(1, dtype=torch.int32, device=device), - all_work_published=torch.zeros(1, dtype=torch.int32, device=device), task_head=torch.zeros(1, dtype=torch.int32, device=device), task_tail=torch.zeros(1, dtype=torch.int32, device=device), - task_ready=torch.zeros(max_tasks, dtype=torch.int32, device=device), task_expert=torch.zeros(max_tasks, dtype=torch.int32, device=device), - task_m_tile=torch.zeros(max_tasks, dtype=torch.int32, device=device), - task_slice_begin=torch.zeros(max_tasks, dtype=torch.int32, device=device), - task_slice_count=torch.zeros(max_tasks, dtype=torch.int32, device=device), task_valid_rows=torch.zeros(max_tasks, dtype=torch.int32, device=device), - tile_write_count=torch.zeros(physical_tiles, dtype=torch.int32, device=device), ) # Finalize views @@ -1392,17 +1705,10 @@ def __call__( barrier_count: cute.Tensor, barrier_epoch: cute.Tensor, pair_head: cute.Tensor, - producers_done_count: cute.Tensor, - all_work_published: cute.Tensor, task_head: cute.Tensor, task_tail: cute.Tensor, - task_ready_ptr: cute.Pointer, task_expert_ptr: cute.Pointer, - task_m_tile_ptr: cute.Pointer, - task_slice_begin_ptr: cute.Pointer, - task_slice_count_ptr: cute.Pointer, task_valid_rows_ptr: cute.Pointer, - tile_write_count_ptr: cute.Pointer, b_w13: cute.Tensor, sfb_w13_ptr: cute.Pointer, b_down: cute.Tensor, @@ -1421,7 +1727,6 @@ def __call__( max_rows: cutlass.Int32, rows_padded: cutlass.Int32, max_tasks: cutlass.Int32, - max_phys_tiles: cutlass.Int32, max_active_clusters: cutlass.Constexpr, stream, ): @@ -1452,9 +1757,12 @@ def __call__( (rows_padded * self._packed_storage_cols,), stride=(1,) ), ) + # Activation scales live in 128-row SF atoms; the plane is allocated + # through the last partial atom even when rows_padded is not aligned. + scale_rows = ((rows_padded + 127) // 128) * 128 scale_storage = cute.make_tensor( scale_storage_ptr, - layout=cute.make_layout((rows_padded * self._cols_pad_k,), stride=(1,)), + layout=cute.make_layout((scale_rows * self._cols_pad_k,), stride=(1,)), ) token_map = cute.make_tensor( token_map_ptr, layout=cute.make_layout((rows_padded,), stride=(1,)) @@ -1462,28 +1770,12 @@ def __call__( token_weights_t = cute.make_tensor( token_weights_ptr, layout=cute.make_layout((rows_padded,), stride=(1,)) ) - task_ready = cute.make_tensor( - task_ready_ptr, layout=cute.make_layout((max_tasks,), stride=(1,)) - ) task_expert = cute.make_tensor( task_expert_ptr, layout=cute.make_layout((max_tasks,), stride=(1,)) ) - task_m_tile = cute.make_tensor( - task_m_tile_ptr, layout=cute.make_layout((max_tasks,), stride=(1,)) - ) - task_slice_begin = cute.make_tensor( - task_slice_begin_ptr, layout=cute.make_layout((max_tasks,), stride=(1,)) - ) - task_slice_count = cute.make_tensor( - task_slice_count_ptr, layout=cute.make_layout((max_tasks,), stride=(1,)) - ) task_valid_rows = cute.make_tensor( task_valid_rows_ptr, layout=cute.make_layout((max_tasks,), stride=(1,)) ) - tile_write_count = cute.make_tensor( - tile_write_count_ptr, - layout=cute.make_layout((max_phys_tiles,), stride=(1,)), - ) self._kernel( a_input, topk_ids, @@ -1495,17 +1787,10 @@ def __call__( barrier_count, barrier_epoch, pair_head, - producers_done_count, - all_work_published, task_head, task_tail, - task_ready, task_expert, - task_m_tile, - task_slice_begin, - task_slice_count, task_valid_rows, - tile_write_count, b_w13, sfb_w13_ptr, b_down, @@ -1545,6 +1830,7 @@ def _get_dynamic_kernel( swiglu_limit: float | None = None, activation_precision: str = "fp4", share_input_across_experts: bool = False, + tile_m: int = _LEVEL_TILE_M, ): """Compile (or retrieve cached) the SM120 dynamic MoE kernel.""" activation_precision = _normalize_activation_precision(activation_precision) @@ -1557,11 +1843,12 @@ def _get_dynamic_kernel( ) sf_vec_size = 16 sm_count = get_num_sm(torch.device("cuda")) - mac = min(get_max_active_clusters(1), sm_count) - mma_tiler_mn = ( - _level_tile_m(activation_precision), - _level_tile_n(activation_precision), - ) + base_mac = min(get_max_active_clusters(1), sm_count) + tuned_mac = _lookup_mac_ladder(_DYNAMIC_MAC_LADDER, m * num_topk) + mac = min(tuned_mac or base_mac, base_mac) + # tile_m comes from the workspace's shared selection so the kernel's task + # and scale indexing matches the allocated scratch geometry. + mma_tiler_mn = (tile_m, _level_tile_n(activation_precision)) cache_key = ( "dynamic", @@ -1648,12 +1935,6 @@ def _get_dynamic_kernel( pair_head_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (1,), assumed_align=4 ) - producers_done_count_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Int32, (1,), assumed_align=4 - ) - all_work_published_fake = cute.runtime.make_fake_compact_tensor( - cutlass.Int32, (1,), assumed_align=4 - ) task_head_fake = cute.runtime.make_fake_compact_tensor( cutlass.Int32, (1,), assumed_align=4 ) @@ -1661,27 +1942,12 @@ def _get_dynamic_kernel( cutlass.Int32, (1,), assumed_align=4 ) - task_ready_fake = make_ptr( - cutlass.Int32, 4, cute.AddressSpace.gmem, assumed_align=4 - ) task_expert_fake = make_ptr( cutlass.Int32, 4, cute.AddressSpace.gmem, assumed_align=4 ) - task_m_tile_fake = make_ptr( - cutlass.Int32, 4, cute.AddressSpace.gmem, assumed_align=4 - ) - task_slice_begin_fake = make_ptr( - cutlass.Int32, 4, cute.AddressSpace.gmem, assumed_align=4 - ) - task_slice_count_fake = make_ptr( - cutlass.Int32, 4, cute.AddressSpace.gmem, assumed_align=4 - ) task_valid_rows_fake = make_ptr( cutlass.Int32, 4, cute.AddressSpace.gmem, assumed_align=4 ) - tile_write_count_fake = make_ptr( - cutlass.Int32, 4, cute.AddressSpace.gmem, assumed_align=4 - ) b_w13_fake = cute.runtime.make_fake_compact_tensor( weight_dtype, @@ -1736,17 +2002,10 @@ def _get_dynamic_kernel( barrier_count_fake, barrier_epoch_fake, pair_head_fake, - producers_done_count_fake, - all_work_published_fake, task_head_fake, task_tail_fake, - task_ready_fake, task_expert_fake, - task_m_tile_fake, - task_slice_begin_fake, - task_slice_count_fake, task_valid_rows_fake, - tile_write_count_fake, b_w13_fake, sfb_w13_fake, b_down_fake, @@ -1764,7 +2023,6 @@ def _get_dynamic_kernel( 1, 1, 1, - 1, 1, # runtime Int32 placeholders mac, current_cuda_stream(), @@ -1808,6 +2066,7 @@ def launch_sm120_dynamic_moe( raise ValueError( "internal routing error: quant_mode='w4a16' reached the NVFP4 dynamic launcher" ) + _check_memref_limit("scatter_output", scatter_output.numel()) flat_ids = topk_ids.view(-1).to(torch.int32) flat_weights = topk_weights.view(-1).to(torch.float32) input_gs_is_shared = input_gs.numel() == 1 @@ -1832,6 +2091,7 @@ def launch_sm120_dynamic_moe( swiglu_limit=swiglu_limit, activation_precision=activation_precision, share_input_across_experts=input_gs_is_shared, + tile_m=workspace.tile_m, ) # Dynamic kernel: runtime-shaped args are DataPointer (pass data_ptr()), @@ -1847,17 +2107,10 @@ def launch_sm120_dynamic_moe( workspace.barrier_count, workspace.barrier_epoch, workspace.pair_head, - workspace.producers_done_count, - workspace.all_work_published, workspace.task_head, workspace.task_tail, - workspace.task_ready.data_ptr(), workspace.task_expert.data_ptr(), - workspace.task_m_tile.data_ptr(), - workspace.task_slice_begin.data_ptr(), - workspace.task_slice_count.data_ptr(), workspace.task_valid_rows.data_ptr(), - workspace.tile_write_count.data_ptr(), weights.w13_fp4, weights._w13_sf_storage.data_ptr(), weights.down_fp4, @@ -1874,9 +2127,8 @@ def launch_sm120_dynamic_moe( workspace.token_weights.data_ptr(), num_tokens, workspace.max_rows, - workspace.physical_tiles_capacity * _level_tile_m(activation_precision), + workspace.physical_tiles_capacity * workspace.tile_m, workspace.task_capacity, - workspace.physical_tiles_capacity, ) compiled(*runtime_args, current_cuda_stream()) @@ -2281,12 +2533,27 @@ def _launch_sm120_w4a16_moe( Sm120W4A16MoEWorkspace, ] -# Keyed by (state_E, weight_E, k, n, top_k, device, backend). -# Stores the workspace with the largest max_rows seen for each key. -# Grows monotonically — never shrinks within a process. +# Stores the workspace with the largest capacity seen per key and never +# shrinks within a process. clear_sm120_moe_caches() releases everything. _WORKSPACE_CACHE: Dict[Tuple, _Sm120Workspace] = {} +def clear_sm120_moe_caches() -> None: + """Release every module-level SM12x MoE cache. + + References held by callers are unaffected. + """ + _WORKSPACE_CACHE.clear() + _WEIGHT_CACHE.clear() + _W4A16_WEIGHT_CACHE.clear() + _PADDED_WEIGHT_CACHE.clear() + _STATIC_KERNEL_CACHE.clear() + _MICRO_KERNEL_CACHE.clear() + _DIRECT_MICRO_LAUNCH_CACHE.clear() + _DIRECT_MICRO_KERNEL_CACHE.clear() + _DYNAMIC_KERNEL_CACHE.clear() + + def allocate_sm120_moe_workspace( *, state_E: int, @@ -2343,6 +2610,7 @@ def allocate_sm120_moe_workspace( num_topk=num_topk, device=device, activation_precision=activation_precision, + activation=activation, ) if backend == "static": return allocate_sm120_static_workspace( @@ -2382,6 +2650,13 @@ def _get_cached_workspace( """ quant_mode = _normalize_quant_mode(quant_mode, activation_precision) activation_precision = _activation_precision_from_quant_mode(quant_mode) + # Key dynamic workspaces on the tile band of this call's routed_rows; a + # larger cached workspace must not pin small calls to its 128 tile. + tile_m = ( + _select_dynamic_tile_m(max(1, routed_rows), state_E, activation) + if backend == "dynamic" and quant_mode != "w4a16" + else None + ) cache_key = ( state_E, weight_E, @@ -2392,12 +2667,17 @@ def _get_cached_workspace( backend, quant_mode, activation, + tile_m, ) cached = _WORKSPACE_CACHE.get(cache_key) if cached is not None: - if isinstance(cached, (Sm120DynamicMoEWorkspace, Sm120W4A16MoEWorkspace)): - if cached.routed_rows_capacity >= max(1, routed_rows): # type: ignore[union-attr] + if isinstance(cached, Sm120DynamicMoEWorkspace): + if cached.routed_rows_capacity >= max(1, routed_rows): + assert tile_m is None or cached.tile_m == tile_m + return cached + elif isinstance(cached, Sm120W4A16MoEWorkspace): + if cached.routed_rows_capacity >= max(1, routed_rows): return cached else: if cached.max_rows >= max(1, routed_rows): @@ -2680,6 +2960,8 @@ def launch_sm120_moe( "num_local_experts == num_experts because dynamic expert " "buffers are indexed by global topk ids." ) + # A pre-allocated dynamic workspace keeps its stored tile_m even + # for smaller calls; its geometry was sized for that tile. backend = "dynamic" else: backend = "static" diff --git a/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py b/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py index ed8ea5a1b1a..42dcac2d1ba 100644 --- a/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py +++ b/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_dynamic_kernel.py @@ -103,7 +103,17 @@ _SF_VEC_SIZE = 16 _TASK_SLICE_CHUNK = 1 _PRODUCER_PAIRS_PER_WARP = 2 -_FC2_TILE_RECIP_GS_NUM = 6.0 * 448.0 + +_WORK_ITEM_FIELDS = 5 +_WORK_EXPERT = 0 +_WORK_M_TILE = 1 +_WORK_SLICE_BEGIN = 2 +_WORK_SLICE_COUNT = 3 +_WORK_VALID_ROWS = 4 + +_CTRL_HAS_WORK = 0 +_CTRL_DONE = 1 +_CTRL_WORK_BEGIN = 2 class DynamicLaunchParams: @@ -247,25 +257,6 @@ def _threadfence(*, loc=None, ip=None): ) -@dsl_user_op -def _atomic_cas_global_i32(addr, compare, value, *, loc=None, ip=None): - return Int32( - llvm.inline_asm( - T.i32(), - [ - Int64(addr).ir_value(loc=loc, ip=ip), - Int32(compare).ir_value(loc=loc, ip=ip), - Int32(value).ir_value(loc=loc, ip=ip), - ], - "atom.global.cas.b32 $0, [$1], $2, $3;", - "=r,l,r,r", - has_side_effects=True, - is_align_stack=False, - asm_dialect=llvm.AsmDialect.AD_ATT, - ) - ) - - class MoEDynamicKernel: """Queue-driven first-pass dynamic MoE kernel.""" @@ -288,20 +279,45 @@ def __init__( self.acc_dtype = cutlass.Float32 self.sf_vec_size = sf_vec_size self.input_scales_are_reciprocal = input_scales_are_reciprocal - self.fast_math = fast_math self.activation = activation self.is_gated = is_gated_activation(activation) + # relu2's squared outputs need the exact quantizer and scale math. + self.fast_math = bool(fast_math) and self.is_gated self.swiglu_alpha = float(swiglu_alpha) self.swiglu_beta = float(swiglu_beta) self.swiglu_limit = float(swiglu_limit) if swiglu_limit is not None else None self.share_input_across_experts = share_input_across_experts tile_k = sf_vec_size * 8 self.tile_shape_mnk = (mma_tiler_mn[0], mma_tiler_mn[1], tile_k) + # Scale factors come in 128-row atoms, so for sub-128 MMA tiles the + # TMA atoms and smem are built at max(128, tile) and the kernel + # offsets into the shared block by `*_tiles_per_block`. + self.sa_tile_shape_mk = (max(128, mma_tiler_mn[0]), tile_k) + self.sa_tiles_per_block = self.sa_tile_shape_mk[0] // mma_tiler_mn[0] + self.sfa_tile_shape_mk = (max(128, mma_tiler_mn[0]), tile_k) + self.sfa_tiles_per_block = self.sfa_tile_shape_mk[0] // mma_tiler_mn[0] + self.sfb_tile_shape_nk = (max(128, mma_tiler_mn[1]), tile_k) + self.sfb_tiles_per_block = self.sfb_tile_shape_nk[0] // mma_tiler_mn[1] self.cluster_shape_mnk = (1, 1, 1) self.cluster_shape_mn = (1, 1) self.epi_tile = (mma_tiler_mn[0], mma_tiler_mn[1]) self.occupancy = 1 - self.num_mma_warps = 4 + # Smaller tiles need a matching MMA atom shape so the SF smem layout + # and V-map stay consistent. + if mma_tiler_mn[0] == 128: + self.atom_shape = (2, 2, 1) + self.num_mma_warps = 4 + elif mma_tiler_mn[0] == 64: + self.atom_shape = (4, 2, 1) + self.num_mma_warps = 8 + elif mma_tiler_mn[0] == 32: + self.atom_shape = (2, 2, 1) + self.num_mma_warps = 4 + elif mma_tiler_mn[0] == 16: + self.atom_shape = (1, 2, 1) + self.num_mma_warps = 2 + else: + raise ValueError(f"unsupported dynamic MMA tile_m {mma_tiler_mn[0]}") self.tma_load_warp_id = self.num_mma_warps self.num_threads_per_warp = 32 self.threads_per_cta = (self.num_mma_warps + 1) * self.num_threads_per_warp @@ -345,7 +361,7 @@ def _setup_attributes(self, hidden_size: int): self.acc_dtype, self.sf_dtype, ) - atom_layout = cute.make_layout((2, 2, 1)) + atom_layout = cute.make_layout(self.atom_shape) permutation_mnk = sm120_utils.get_permutation_mnk( self.tile_shape_mnk, self.sf_vec_size, @@ -358,19 +374,27 @@ def _setup_attributes(self, hidden_size: int): ) self.mma_atom = cute.make_mma_atom(mma_op) self.cta_layout_mnk = cute.make_layout(self.cluster_shape_mnk) - self.num_m_tiles = self.tile_shape_mnk[0] // (16 * 4) - self.num_n_tiles = self.tile_shape_mnk[1] // (8 * 2) + self.num_m_tiles = self.tile_shape_mnk[0] // (16 * self.atom_shape[0]) + self.num_n_tiles = self.tile_shape_mnk[1] // (8 * self.atom_shape[1]) self.num_k_blocks = self.tile_shape_mnk[2] // 64 + # A/SFA smem hold the whole 128-row block for sub-128 MMA tiles (the + # SF smem helper also rejects tile M % 64 != 0), so build all smem + # layouts at the block shape; sub-128 tasks slice into it. + smem_tile_shape_mnk = ( + self.sa_tile_shape_mk[0], + self.tile_shape_mnk[1], + self.tile_shape_mnk[2], + ) sfa_smem = sm120_make_smem_layout_sfa( self.tiled_mma, - self.tile_shape_mnk, + smem_tile_shape_mnk, self.sf_vec_size, 1, ) sfb_smem = sm120_make_smem_layout_sfb( self.tiled_mma, - self.tile_shape_mnk, + smem_tile_shape_mnk, self.sf_vec_size, 1, ) @@ -387,11 +411,40 @@ def _setup_attributes(self, hidden_size: int): self.smem_capacity, self.occupancy, ) - # The gated path stages a second weight pipeline (sB_up / sSFB_up) that - # _compute_stages doesn't model. Extra smem usage requires capping at 2 - # to stay within the SM12x smem budget. - if self.is_gated: - self.ab_stage = max(1, min(self.ab_stage, 2)) + # dense._compute_stages models a single B/SFB buffer, but the gated + # path double-buffers both. Recompute the stage cap from the real + # per-stage footprint so shapes with room keep their extra stages. + nb = 2 if self.is_gated else 1 + n_pipe = 3 if self.is_gated else 2 + # sA smem holds the whole 128-row block, so size the stage from it. + a_bytes = ( + self.sa_tile_shape_mk[0] + * self.sa_tile_shape_mk[1] + * self.a_dtype.width + // 8 + ) + b_bytes = ( + cute.size(cute.slice_(self.tile_shape_mnk, (0, None, None))) + * self.b_dtype.width + // 8 + ) + sfa_bytes = ( + cute.size(cute.filter_zeros(sfa_smem).shape) * self.sf_dtype.width // 8 + ) + sfb_bytes = ( + cute.size(cute.filter_zeros(sfb_smem).shape) * self.sf_dtype.width // 8 + ) + per_stage = a_bytes + nb * b_bytes + sfa_bytes + nb * sfb_bytes + n_pipe * 2 * 8 + fixed = ( + self.tile_shape_mnk[0] * self.tile_shape_mnk[1] * 2 # sC (bf16 epi) + + 2 * (self.num_mma_warps + 1) * 32 * 4 # route caches + + self.tile_shape_mnk[0] * 8 # scatter caches + + 8 * 8 + + 1024 + + 8 * 1024 # ctrl/mbar/align slack + ) + max_fit = max(1, (self.smem_capacity - fixed) // per_stage) + self.ab_stage = min(self.ab_stage, max_fit) # ab_stage must divide k_tile_cnt evenly to avoid pipeline phase mismatch. # _compute_stages returns the max that fits in smem, but it may not # divide k_tile_cnt. Round down to the nearest divisor. @@ -406,7 +459,7 @@ def _setup_attributes(self, hidden_size: int): self.sfb_smem_layout_staged, self.epi_smem_layout_staged, ) = self._dense_cls._make_smem_layouts( - self.tile_shape_mnk, + smem_tile_shape_mnk, self.epi_tile, self.a_dtype, self.a_layout, @@ -443,14 +496,9 @@ def _resident_grid_barrier( cute.arch.sync_threads() @cute.jit - def _publish_ready_tasks( + def _publish_deferred_tasks( self, - task_tail: cute.Tensor, - task_ready: cute.Tensor, task_expert: cute.Tensor, - task_m_tile: cute.Tensor, - task_slice_begin: cute.Tensor, - task_slice_count: cute.Tensor, task_valid_rows: cute.Tensor, gate_tile_cnt: Int32, slice_chunk: Int32, @@ -459,60 +507,66 @@ def _publish_ready_tasks( valid_rows: Int32, ): num_groups = (gate_tile_cnt + slice_chunk - Int32(1)) // slice_chunk - start = atomic_add_global_i32(get_ptr_as_int64(task_tail, Int32(0)), num_groups) + start = m_tile_idx * num_groups g = Int32(0) while g < num_groups: slot = start + g - slice_begin = g * slice_chunk - slice_count = gate_tile_cnt - slice_begin - if slice_count > slice_chunk: - slice_count = slice_chunk task_expert[slot] = expert_idx - task_m_tile[slot] = m_tile_idx - task_slice_begin[slot] = slice_begin - task_slice_count[slot] = slice_count task_valid_rows[slot] = valid_rows g += Int32(1) - _threadfence() - - g = Int32(0) - while g < num_groups: - slot = start + g - _st_global_release_i32(get_ptr_as_int64(task_ready, slot), Int32(1)) - g += Int32(1) - @cute.jit - def _publish_deferred_tasks( + def _decode_materialized_work_item( self, + work_item: cute.Tensor, task_expert: cute.Tensor, - task_m_tile: cute.Tensor, - task_slice_begin: cute.Tensor, - task_slice_count: cute.Tensor, task_valid_rows: cute.Tensor, - gate_tile_cnt: Int32, + slot: Int32, + num_groups: Int32, slice_chunk: Int32, - expert_idx: Int32, - m_tile_idx: Int32, - valid_rows: Int32, + gate_tile_cnt: Int32, ): - num_groups = (gate_tile_cnt + slice_chunk - Int32(1)) // slice_chunk - start = m_tile_idx * num_groups + """Decode a deferred task from its slot index. + + Slots are ordered by m-tile, then group, so only expert and + valid-row metadata need global storage. + """ + + m_tile = slot // num_groups + group = slot - m_tile * num_groups + slice_begin = group * slice_chunk + slice_count = gate_tile_cnt - slice_begin + if slice_count > slice_chunk: + slice_count = slice_chunk + work_item[_WORK_EXPERT] = task_expert[slot].to(Int32) + work_item[_WORK_M_TILE] = m_tile + work_item[_WORK_SLICE_BEGIN] = slice_begin + work_item[_WORK_SLICE_COUNT] = slice_count + work_item[_WORK_VALID_ROWS] = task_valid_rows[slot].to(Int32) - g = Int32(0) - while g < num_groups: - slot = start + g - slice_begin = g * slice_chunk - slice_count = gate_tile_cnt - slice_begin - if slice_count > slice_chunk: - slice_count = slice_chunk - task_expert[slot] = expert_idx - task_m_tile[slot] = m_tile_idx - task_slice_begin[slot] = slice_begin - task_slice_count[slot] = slice_count - task_valid_rows[slot] = valid_rows - g += Int32(1) + @cute.jit + def _store_shared_work_item( + self, + ctrl_base_addr: Int32, + work_item: cute.Tensor, + ): + for field in cutlass.range_constexpr(_WORK_ITEM_FIELDS): + _st_shared_i32( + ctrl_base_addr + Int32((_CTRL_WORK_BEGIN + field) * 4), + work_item[field], + ) + + @cute.jit + def _load_shared_work_item( + self, + work_item: cute.Tensor, + ctrl_base_addr: Int32, + ): + for field in cutlass.range_constexpr(_WORK_ITEM_FIELDS): + work_item[field] = _ld_shared_i32( + ctrl_base_addr + Int32((_CTRL_WORK_BEGIN + field) * 4) + ) @cute.jit def __call__( @@ -527,17 +581,10 @@ def __call__( barrier_count: cute.Tensor, # [1] int32 (host-zeroed) barrier_epoch: cute.Tensor, # [1] int32 (host-zeroed) pair_head: cute.Tensor, # [1] int32 - producers_done_count: cute.Tensor, # [1] int32 - all_work_published: cute.Tensor, # [1] int32 task_head: cute.Tensor, # [1] int32 task_tail: cute.Tensor, # [1] int32 - task_ready: cute.Tensor, # [max_tasks] int32 task_expert: cute.Tensor, # [max_tasks] int32 - task_m_tile: cute.Tensor, # [max_tasks] int32 - task_slice_begin: cute.Tensor, # [max_tasks] int32 - task_slice_count: cute.Tensor, # [max_tasks] int32 task_valid_rows: cute.Tensor, # [max_tasks] int32 - tile_write_count: cute.Tensor, # [E * max_m_tiles] int32 b_w13: cute.Tensor, # [2*I_tp, K, E] (gated) or [I_tp, K, E] (relu2) sfb_w13_ptr: cute.Pointer, # scale factors for w13 b_down: cute.Tensor, # [K, I_tp, E] @@ -582,13 +629,13 @@ def __call__( tma_a, gA = self._dense_cls._make_tma_atoms_and_tensors( packed_a, self.a_smem_layout_staged, - (self.tile_shape_mnk[0], self.tile_shape_mnk[2]), + self.sa_tile_shape_mk, 1, ) tma_sfa, gSFA = self._dense_cls._make_tma_atoms_and_tensors( sfa_tensor, self.sfa_smem_layout_staged, - (self.tile_shape_mnk[0], self.tile_shape_mnk[2]), + self.sfa_tile_shape_mk, 1, internal_type=cutlass.Int16, ) @@ -603,7 +650,7 @@ def __call__( tma_sfb_w13, gSFB_w13 = self._dense_cls._make_tma_atoms_and_tensors( sfb_w13_tensor, self.sfb_smem_layout_staged, - (self.tile_shape_mnk[1], self.tile_shape_mnk[2]), + self.sfb_tile_shape_nk, 1, internal_type=cutlass.Int16, ) @@ -621,7 +668,7 @@ def __call__( tma_sfb_down, gSFB_down = self._dense_cls._make_tma_atoms_and_tensors( sfb_down_tensor, self.sfb_smem_layout_staged, - (self.tile_shape_mnk[1], self.tile_shape_mnk[2]), + self.sfb_tile_shape_nk, 1, internal_type=cutlass.Int16, ) @@ -640,17 +687,10 @@ def __call__( barrier_count, barrier_epoch, pair_head, - producers_done_count, - all_work_published, task_head, task_tail, - task_ready, task_expert, - task_m_tile, - task_slice_begin, - task_slice_count, task_valid_rows, - tile_write_count, tma_a, gA, tma_sfa, @@ -685,6 +725,9 @@ def __call__( grid=grid, block=[self.threads_per_cta, 1, 1], cluster=[1, 1, 1], + # A regular launch beside other stream work can admit only part + # of the grid, deadlocking the software grid barriers below. + cooperative=True, stream=stream, ) @@ -699,17 +742,10 @@ def kernel( barrier_count: cute.Tensor, barrier_epoch: cute.Tensor, pair_head: cute.Tensor, - producers_done_count: cute.Tensor, - all_work_published: cute.Tensor, task_head: cute.Tensor, task_tail: cute.Tensor, - task_ready: cute.Tensor, task_expert: cute.Tensor, - task_m_tile: cute.Tensor, - task_slice_begin: cute.Tensor, - task_slice_count: cute.Tensor, task_valid_rows: cute.Tensor, - tile_write_count: cute.Tensor, tma_a: cute.CopyAtom, mA: cute.Tensor, tma_sfa: cute.CopyAtom, @@ -749,7 +785,7 @@ def kernel( _, _, gdim_z = cute.arch.grid_dim() warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) - is_cta_leader = Int32(1) if Int32(tidx) == Int32(0) else Int32(0) + is_cta_leader = Int32(Int32(tidx) == Int32(0)) if warp_idx == 0: cpasync.prefetch_descriptor(tma_a) @@ -958,11 +994,13 @@ class StorageRelu2: num_k_tiles = (cols + Int32(63)) // Int32(64) route_gate_tile_cnt = launch_params.gate_tile_cnt task_slice_chunk = Int32(_TASK_SLICE_CHUNK) - full_tile_publish_enabled = Int32(0) + materialized_num_groups = ( + route_gate_tile_cnt + task_slice_chunk - Int32(1) + ) // task_slice_chunk - # Phase 0: cooperative init — zero routing state, queue state, and output. - task_capacity = Int32(task_ready.shape[0]) - tile_write_slots = Int32(tile_write_count.shape[0]) + # Phase 0: cooperative init. Zero routing state, queue state, and + # output. Task metadata slots are overwritten before the second grid + # barrier and need no clear. i = flat_tid while i < num_experts: row_counts[i] = Int32(0) @@ -971,7 +1009,7 @@ class StorageRelu2: if flat_tid < num_experts + Int32(1): expert_tile_base[flat_tid] = Int32(0) - scatter_total_u32 = num_tokens * cols_u32 + scatter_total_u32 = Int32(scatter_output.shape[0]) * cols_u32 scatter_vecs = scatter_total_u32 // Int32(4) zero_u32 = Uint32(0) zv = flat_tid @@ -990,26 +1028,8 @@ class StorageRelu2: scatter_output_u32[j // cols_u32, j % cols_u32] = Uint32(0) j += flat_stride - k = flat_tid - while k < task_capacity: - task_ready[k] = Int32(0) - task_expert[k] = Int32(0) - task_m_tile[k] = Int32(0) - task_slice_begin[k] = Int32(0) - task_slice_count[k] = Int32(0) - task_valid_rows[k] = Int32(0) - k += flat_stride - - if full_tile_publish_enabled > Int32(0): - tw = flat_tid - while tw < tile_write_slots: - tile_write_count[tw] = Int32(0) - tw += flat_stride - if flat_tid == Int32(0): pair_head[Int32(0)] = Int32(0) - producers_done_count[Int32(0)] = Int32(0) - all_work_published[Int32(0)] = Int32(0) task_head[Int32(0)] = Int32(0) task_tail[Int32(0)] = Int32(0) @@ -1141,18 +1161,17 @@ class StorageRelu2: phys_row = _ld_shared_i32( route_phys_rows_addr + slot * Int32(4) ) - phys_tile = phys_row // Int32(self.tile_shape_mnk[0]) - tile_row = phys_row - phys_tile * Int32( - self.tile_shape_mnk[0] - ) route_output_base[cache_slot] = ( phys_row * output_bytes_per_row ) + # Scale storage is tiled in 128-row SF atoms, + # independently of the MMA tile. + sf_atom = phys_row >> Int32(7) + sf_row = phys_row & Int32(127) route_scale_base[cache_slot] = ( - phys_tile * num_k_tiles * Int32(32 * 4 * 4) - + (tile_row % Int32(32)) * Int32(4 * 4) - + ((tile_row % Int32(32 * 4)) // Int32(32)) - * Int32(4) + sf_atom * num_k_tiles * Int32(32 * 4 * 4) + + (sf_row % Int32(32)) * Int32(4 * 4) + + (sf_row // Int32(32)) * Int32(4) ) sf_idx = lane_id @@ -1228,12 +1247,6 @@ class StorageRelu2: phys_row = _ld_shared_i32( route_phys_rows_addr + slot * Int32(4) ) - phys_tile = phys_row // Int32( - self.tile_shape_mnk[0] - ) - tile_row = phys_row - phys_tile * Int32( - self.tile_shape_mnk[0] - ) output_offset = ( phys_row * output_bytes_per_row + sf_idx * Int32(8) @@ -1244,14 +1257,16 @@ class StorageRelu2: ), packed64, ) + # Scale storage is tiled in 128-row SF + # atoms, not MMA tiles. k_tile_idx = sf_idx // Int32(4) - outer_m_idx = tile_row % Int32(32) - inner_m_idx = (tile_row % Int32(32 * 4)) // Int32( - 32 - ) + sf_atom = phys_row >> Int32(7) + sf_row = phys_row & Int32(127) + outer_m_idx = sf_row % Int32(32) + inner_m_idx = sf_row // Int32(32) inner_k_idx = sf_idx % Int32(4) scale_offset = ( - phys_tile * num_k_tiles * Int32(32 * 4 * 4) + sf_atom * num_k_tiles * Int32(32 * 4 * 4) + k_tile_idx * Int32(32 * 4 * 4) + outer_m_idx * Int32(4 * 4) + inner_m_idx * Int32(4) @@ -1261,44 +1276,6 @@ class StorageRelu2: topk_slot += Int32(1) sf_idx += Int32(32) - if full_tile_publish_enabled > Int32(0): - cute.arch.sync_warp() - _threadfence() - cute.arch.sync_warp() - - if lane_id == Int32(0): - topk_slot = Int32(0) - while topk_slot < num_topk: - slot = route_slot_base + topk_slot - phys_row = _ld_shared_i32( - route_phys_rows_addr + slot * Int32(4) - ) - expert_id = _ld_shared_i32( - route_expert_ids_addr + slot * Int32(4) - ) - phys_tile = phys_row // Int32( - self.tile_shape_mnk[0] - ) - completed = atomic_add_global_i32( - get_ptr_as_int64(tile_write_count, phys_tile), - Int32(1), - ) + Int32(1) - if completed == Int32(self.tile_shape_mnk[0]): - self._publish_ready_tasks( - task_tail, - task_ready, - task_expert, - task_m_tile, - task_slice_begin, - task_slice_count, - task_valid_rows, - route_gate_tile_cnt, - task_slice_chunk, - expert_id, - phys_tile, - Int32(self.tile_shape_mnk[0]), - ) - topk_slot += Int32(1) else: warp_item = Int32(0) while warp_item < Int32(_PRODUCER_PAIRS_PER_WARP): @@ -1378,48 +1355,25 @@ class StorageRelu2: packed64, ) + # Scale storage is tiled in 128-row SF atoms, + # not MMA tiles. k_tile_idx = sf_idx // Int32(4) - outer_m_idx = row % Int32(32) - inner_m_idx = (row % Int32(32 * 4)) // Int32(32) inner_k_idx = sf_idx % Int32(4) + phys_row = phys_tile * Int32( + self.tile_shape_mnk[0] + ) + row % Int32(self.tile_shape_mnk[0]) + sf_atom = phys_row >> Int32(7) + sf_row = phys_row & Int32(127) scale_offset = ( - phys_tile * num_k_tiles * Int32(32 * 4 * 4) + sf_atom * num_k_tiles * Int32(32 * 4 * 4) + k_tile_idx * Int32(32 * 4 * 4) - + outer_m_idx * Int32(4 * 4) - + inner_m_idx * Int32(4) + + (sf_row % Int32(32)) * Int32(4 * 4) + + (sf_row // Int32(32)) * Int32(4) + inner_k_idx ) scale_storage[scale_offset] = scale_byte sf_idx += Int32(32) - if full_tile_publish_enabled > Int32(0): - cute.arch.sync_warp() - # When the whole launch has fewer than one M-tile of routed - # rows, only the final partial-tile flush can publish work. - # Skip the per-row fence/counter path in that common micro case. - _threadfence() - cute.arch.sync_warp() - - if lane_id == Int32(0): - completed = atomic_add_global_i32( - get_ptr_as_int64(tile_write_count, phys_tile), - Int32(1), - ) + Int32(1) - if completed == Int32(self.tile_shape_mnk[0]): - self._publish_ready_tasks( - task_tail, - task_ready, - task_expert, - task_m_tile, - task_slice_begin, - task_slice_count, - task_valid_rows, - route_gate_tile_cnt, - task_slice_chunk, - expert_id, - phys_tile, - Int32(self.tile_shape_mnk[0]), - ) warp_item += Int32(1) cute.arch.sync_threads() @@ -1429,99 +1383,51 @@ class StorageRelu2: _threadfence() cute.arch.sync_threads() - if full_tile_publish_enabled == Int32(0): - # Micro batches cannot fill a full M tile, so overlap is impossible. - # Rendezvous once, publish the final partial tiles, then consume. - self._resident_grid_barrier( - barrier_count, - barrier_epoch, - Int32(gdim_z), - is_cta_leader, - ) + # Rendezvous once, publish every physical tile, then consume the + # finished queue. + self._resident_grid_barrier( + barrier_count, + barrier_epoch, + Int32(gdim_z), + is_cta_leader, + ) - if is_cta_leader > Int32(0): - expert_flush = Int32(bidz) - while expert_flush < num_experts: - rows_remaining = row_counts[expert_flush] - m_tile_offset = Int32(0) - while rows_remaining > Int32(0): - valid_rows = rows_remaining - if valid_rows > Int32(self.tile_shape_mnk[0]): - valid_rows = Int32(self.tile_shape_mnk[0]) - self._publish_deferred_tasks( - task_expert, - task_m_tile, - task_slice_begin, - task_slice_count, - task_valid_rows, - route_gate_tile_cnt, - task_slice_chunk, - expert_flush, - expert_tile_base[expert_flush] + m_tile_offset, - valid_rows, - ) - rows_remaining -= Int32(self.tile_shape_mnk[0]) - m_tile_offset += Int32(1) - expert_flush += Int32(gdim_z) - - if flat_tid == Int32(0): - num_groups = ( - route_gate_tile_cnt + task_slice_chunk - Int32(1) - ) // task_slice_chunk - st_global_i32( - get_ptr_as_int64(task_tail, Int32(0)), - expert_tile_base[num_experts] * num_groups, - ) + if is_cta_leader > Int32(0): + expert_flush = Int32(bidz) + while expert_flush < num_experts: + rows_remaining = row_counts[expert_flush] + m_tile_offset = Int32(0) + while rows_remaining > Int32(0): + valid_rows = rows_remaining + if valid_rows > Int32(self.tile_shape_mnk[0]): + valid_rows = Int32(self.tile_shape_mnk[0]) + self._publish_deferred_tasks( + task_expert, + task_valid_rows, + route_gate_tile_cnt, + task_slice_chunk, + expert_flush, + expert_tile_base[expert_flush] + m_tile_offset, + valid_rows, + ) + rows_remaining -= Int32(self.tile_shape_mnk[0]) + m_tile_offset += Int32(1) + expert_flush += Int32(gdim_z) - self._resident_grid_barrier( - barrier_count, - barrier_epoch, - Int32(gdim_z), - is_cta_leader, - ) - if flat_tid == Int32(0): - _st_global_release_i32( - get_ptr_as_int64(all_work_published, Int32(0)), - Int32(1), - ) - elif is_cta_leader > Int32(0): - prev_done = atomic_add_global_i32( - get_ptr_as_int64(producers_done_count, Int32(0)), - Int32(1), + if flat_tid == Int32(0): + st_global_i32( + get_ptr_as_int64(task_tail, Int32(0)), + expert_tile_base[num_experts] * materialized_num_groups, ) - if prev_done == Int32(gdim_z) - Int32(1): - expert_flush = Int32(0) - while expert_flush < num_experts: - rows = row_counts[expert_flush] - rem = rows % Int32(self.tile_shape_mnk[0]) - if rem != Int32(0): - partial_m_tile = expert_tile_base[expert_flush] + rows // Int32( - self.tile_shape_mnk[0] - ) - self._publish_ready_tasks( - task_tail, - task_ready, - task_expert, - task_m_tile, - task_slice_begin, - task_slice_count, - task_valid_rows, - route_gate_tile_cnt, - task_slice_chunk, - expert_flush, - partial_m_tile, - rem, - ) - expert_flush += Int32(1) - _threadfence() - _st_global_release_i32( - get_ptr_as_int64(all_work_published, Int32(0)), - Int32(1), - ) - gA = cute.local_tile( - mA, cute.slice_(self.tile_shape_mnk, (None, 0, None)), (None, None, None) + self._resident_grid_barrier( + barrier_count, + barrier_epoch, + Int32(gdim_z), + is_cta_leader, ) + + gA = cute.local_tile(mA, self.sa_tile_shape_mk, (None, None, None)) # Tiled view over w13. # Gated: [2*I_tp, K, E] packed as [up, gate] across N. # Up tiles: N-indices 0..gate_tile_cnt-1 @@ -1532,13 +1438,11 @@ class StorageRelu2: cute.slice_(self.tile_shape_mnk, (0, None, None)), (None, None, None), ) - gSFA = cute.local_tile( - mSFA, cute.slice_(self.tile_shape_mnk, (None, 0, None)), (None, None, None) - ) + # A/SF tiles use the 128-row block shape; for sub-128 MMA tiles one + # block backs `sa_tiles_per_block` MMA tiles (offset applied below). + gSFA = cute.local_tile(mSFA, self.sfa_tile_shape_mk, (None, None, None)) gSFB_w13_tiled = cute.local_tile( - mSFB_w13, - cute.slice_(self.tile_shape_mnk, (0, None, None)), - (None, None, None), + mSFB_w13, self.sfb_tile_shape_nk, (None, None, None) ) thr_mma = tiled_mma.get_slice(tidx) @@ -1627,11 +1531,29 @@ class StorageRelu2: tBgSFB_down = cute.filter_zeros(tBgSFB_down) # MMA fragment partitions - tCsA = thr_mma.partition_A(sA) + # sA/sSFA hold the whole 128-row block; slice to the tile_m sub-tile + # the V-map expects (the per-task offset is applied at consumption). + if cutlass.const_expr(self.sa_tiles_per_block > 1): + sA_part = cute.local_tile( + sA, + cute.slice_(self.tile_shape_mnk, (None, 0, None)), + (Int32(0), 0, None), + ) + else: + sA_part = sA + tCsA = thr_mma.partition_A(sA_part) tCrA = tiled_mma.make_fragment_A(tCsA[None, None, None, 0]) + if cutlass.const_expr(self.sfa_tiles_per_block > 1): + sSFA_part = cute.local_tile( + sSFA, + cute.slice_(self.tile_shape_mnk, (None, 0, None)), + (Int32(0), 0, None), + ) + else: + sSFA_part = sSFA tCrSFA = self._dense_cls._partition_fragment_SFA( self, # type: ignore[arg-type] - sSFA[None, None, 0], + sSFA_part[None, None, 0], thr_mma, tidx, ) @@ -1726,7 +1648,7 @@ class StorageRelu2: thr_ld_A = smem_copy_A.get_slice(tidx) thr_ld_B = smem_copy_B.get_slice(tidx) - csA = thr_ld_A.partition_S(sA) + csA = thr_ld_A.partition_S(sA_part) crA = thr_ld_A.retile(tCrA) csB = thr_ld_B.partition_S(sB) csB_up = thr_ld_B.partition_S(sB_up) @@ -1734,7 +1656,7 @@ class StorageRelu2: thr_ld_SFA = smem_copy_SFA.get_slice(tidx) thr_ld_SFB = smem_copy_SFB.get_slice(tidx) - csSFA = thr_ld_SFA.partition_S(sSFA) + csSFA = thr_ld_SFA.partition_S(sSFA_part) crSFA = thr_ld_SFA.retile(tCrSFA) csSFB = thr_ld_SFB.partition_S(sSFB) csSFB_up = thr_ld_SFB.partition_S(sSFB_up) @@ -1752,77 +1674,40 @@ class StorageRelu2: # Consumer steady state: pop one ready task per CTA, then let # the MMA warps and DMA warp cooperate on that task. # =================================================================== + work_item = cute.make_rmem_tensor((_WORK_ITEM_FIELDS,), Int32) consumer_live = Int32(1) while consumer_live > Int32(0): if is_cta_leader > Int32(0): - _st_shared_i32(ctrl_base_addr + Int32(0), Int32(0)) # has_task - _st_shared_i32(ctrl_base_addr + Int32(4), Int32(0)) # done - _st_shared_i32(ctrl_base_addr + Int32(28), Int32(0)) # claimed slot - if full_tile_publish_enabled == Int32(0): - tail = _ld_global_acquire_i32(get_ptr_as_int64(task_tail, Int32(0))) - slot = atomic_add_global_i32( - get_ptr_as_int64(task_head, Int32(0)), - Int32(1), + _st_shared_i32(ctrl_base_addr + Int32(_CTRL_HAS_WORK * 4), Int32(0)) + _st_shared_i32(ctrl_base_addr + Int32(_CTRL_DONE * 4), Int32(0)) + slot = atomic_add_global_i32( + get_ptr_as_int64(task_head, Int32(0)), Int32(1) + ) + # task_tail is final after the pre-consume grid barrier, so an + # out-of-range slot is a definitive done signal. + tail = _ld_global_acquire_i32(get_ptr_as_int64(task_tail, Int32(0))) + if slot < tail: + self._decode_materialized_work_item( + work_item, + task_expert, + task_valid_rows, + slot, + materialized_num_groups, + task_slice_chunk, + route_gate_tile_cnt, ) - if slot < tail: - _st_shared_i32(ctrl_base_addr + Int32(0), Int32(1)) - _st_shared_i32(ctrl_base_addr + Int32(28), slot) - _st_shared_i32(ctrl_base_addr + Int32(8), task_expert[slot]) - _st_shared_i32(ctrl_base_addr + Int32(12), task_m_tile[slot]) - _st_shared_i32( - ctrl_base_addr + Int32(16), task_slice_begin[slot] - ) - _st_shared_i32( - ctrl_base_addr + Int32(20), task_slice_count[slot] - ) - _st_shared_i32( - ctrl_base_addr + Int32(24), task_valid_rows[slot] - ) - else: - _st_shared_i32(ctrl_base_addr + Int32(4), Int32(1)) + self._store_shared_work_item(ctrl_base_addr, work_item) + _st_shared_i32(ctrl_base_addr + Int32(_CTRL_HAS_WORK * 4), Int32(1)) else: - head = _ld_global_acquire_i32(get_ptr_as_int64(task_head, Int32(0))) - tail = _ld_global_acquire_i32(get_ptr_as_int64(task_tail, Int32(0))) - if head < tail: - prev_head = _atomic_cas_global_i32( - get_ptr_as_int64(task_head, Int32(0)), - head, - head + Int32(1), - ) - if prev_head == head: - _spin_wait_global_eq_i32( - get_ptr_as_int64(task_ready, head), Int32(0) - ) - _st_shared_i32(ctrl_base_addr + Int32(0), Int32(1)) - _st_shared_i32(ctrl_base_addr + Int32(28), head) - _st_shared_i32(ctrl_base_addr + Int32(8), task_expert[head]) - _st_shared_i32( - ctrl_base_addr + Int32(12), task_m_tile[head] - ) - _st_shared_i32( - ctrl_base_addr + Int32(16), task_slice_begin[head] - ) - _st_shared_i32( - ctrl_base_addr + Int32(20), task_slice_count[head] - ) - _st_shared_i32( - ctrl_base_addr + Int32(24), task_valid_rows[head] - ) - else: - if _ld_global_acquire_i32( - get_ptr_as_int64(all_work_published, Int32(0)) - ) > Int32(0): - _st_shared_i32(ctrl_base_addr + Int32(4), Int32(1)) + _st_shared_i32(ctrl_base_addr + Int32(_CTRL_DONE * 4), Int32(1)) cute.arch.sync_threads() - has_task = _ld_shared_i32(ctrl_base_addr + Int32(0)) - is_done = _ld_shared_i32(ctrl_base_addr + Int32(4)) - if has_task > Int32(0) and full_tile_publish_enabled > Int32(0): - claimed_slot = _ld_shared_i32(ctrl_base_addr + Int32(28)) - _ld_global_acquire_i32(get_ptr_as_int64(task_ready, claimed_slot)) + has_task = _ld_shared_i32(ctrl_base_addr + Int32(_CTRL_HAS_WORK * 4)) + is_done = _ld_shared_i32(ctrl_base_addr + Int32(_CTRL_DONE * 4)) if has_task > Int32(0): - task_m_tile_idx_cache = _ld_shared_i32(ctrl_base_addr + Int32(12)) - task_valid_rows_cache = _ld_shared_i32(ctrl_base_addr + Int32(24)) + self._load_shared_work_item(work_item, ctrl_base_addr) + task_m_tile_idx_cache = work_item[_WORK_M_TILE] + task_valid_rows_cache = work_item[_WORK_VALID_ROWS] tile_m_base_cache = task_m_tile_idx_cache * Int32( self.tile_shape_mnk[0] ) @@ -1842,15 +1727,44 @@ class StorageRelu2: if is_done > Int32(0): consumer_live = Int32(0) elif warp_idx < self.num_mma_warps: - task_expert_idx = _ld_shared_i32(ctrl_base_addr + Int32(8)) - task_m_tile_idx = _ld_shared_i32(ctrl_base_addr + Int32(12)) - task_slice_begin_idx = _ld_shared_i32(ctrl_base_addr + Int32(16)) - task_slice_count_val = _ld_shared_i32(ctrl_base_addr + Int32(20)) - task_valid_rows_val = _ld_shared_i32(ctrl_base_addr + Int32(24)) + task_expert_idx = work_item[_WORK_EXPERT] + task_m_tile_idx = work_item[_WORK_M_TILE] + task_slice_begin_idx = work_item[_WORK_SLICE_BEGIN] + task_slice_count_val = work_item[_WORK_SLICE_COUNT] + task_valid_rows_val = work_item[_WORK_VALID_ROWS] alpha_value = alpha[task_expert_idx].to(cutlass.Float32) valid_rows = task_valid_rows_val + # FC1's activation rows sit at offset (task_m_tile_idx % + # sfa_tiles_per_block) within the shared 128-row block; FC2 + # re-slices at offset 0 since its intermediate is written there. + if cutlass.const_expr(self.sfa_tiles_per_block > 1): + _fc1_off = task_m_tile_idx % Int32(self.sfa_tiles_per_block) + _sA_il = cute.local_tile( + sA, + cute.slice_(self.tile_shape_mnk, (None, 0, None)), + (_fc1_off, 0, None), + ) + tCrA = tiled_mma.make_fragment_A( + thr_mma.partition_A(_sA_il)[None, None, None, 0] + ) + csA = thr_ld_A.partition_S(_sA_il) + crA = thr_ld_A.retile(tCrA) + _sSFA_il = cute.local_tile( + sSFA, + cute.slice_(self.tile_shape_mnk, (None, 0, None)), + (_fc1_off, 0, None), + ) + tCrSFA = self._dense_cls._partition_fragment_SFA( + self, # type: ignore[arg-type] + _sSFA_il[None, None, 0], + thr_mma, + tidx, + ) + csSFA = thr_ld_SFA.partition_S(_sSFA_il) + crSFA = thr_ld_SFA.retile(tCrSFA) + _is_m_major = self.c_layout.is_m_major_c() copy_atom_r2s = cute.make_copy_atom( cute.nvgpu.CopyUniversalOp(), @@ -2315,6 +2229,34 @@ class StorageRelu2: warp_m_base = (warp_in_tile >> Int32(1)) * Int32(64) warp_n_base = (warp_in_tile & Int32(1)) * Int32(64) + # FC2's intermediate was quant-written to the head of the + # shared 128-row block, so phase B re-slices at offset 0 + # rather than FC1's per-task offset. + if cutlass.const_expr(self.sfa_tiles_per_block > 1): + _sA_p2 = cute.local_tile( + sA, + cute.slice_(self.tile_shape_mnk, (None, 0, None)), + (Int32(0), 0, None), + ) + tCrA = tiled_mma.make_fragment_A( + thr_mma.partition_A(_sA_p2)[None, None, None, 0] + ) + csA = thr_ld_A.partition_S(_sA_p2) + crA = thr_ld_A.retile(tCrA) + _sSFA_p2 = cute.local_tile( + sSFA, + cute.slice_(self.tile_shape_mnk, (None, 0, None)), + (Int32(0), 0, None), + ) + tCrSFA = self._dense_cls._partition_fragment_SFA( + self, # type: ignore[arg-type] + _sSFA_p2[None, None, 0], + thr_mma, + tidx, + ) + csSFA = thr_ld_SFA.partition_S(_sSFA_p2) + crSFA = thr_ld_SFA.retile(tCrSFA) + csA_phase2 = csA[None, None, None, 0] csSFA_phase2 = csSFA[None, None, None, 0] @@ -2549,13 +2491,17 @@ class StorageRelu2: slice_idx += Int32(1) elif warp_idx == self.tma_load_warp_id: - task_expert_idx = _ld_shared_i32(ctrl_base_addr + Int32(8)) - task_m_tile_idx = _ld_shared_i32(ctrl_base_addr + Int32(12)) - task_slice_begin_idx = _ld_shared_i32(ctrl_base_addr + Int32(16)) - task_slice_count_val = _ld_shared_i32(ctrl_base_addr + Int32(20)) - - tAgA_mk = tAgA[(None, task_m_tile_idx, None, Int32(0))] - tAgSFA_mk = tAgSFA[(None, task_m_tile_idx, None, Int32(0))] + task_expert_idx = work_item[_WORK_EXPERT] + task_m_tile_idx = work_item[_WORK_M_TILE] + task_slice_begin_idx = work_item[_WORK_SLICE_BEGIN] + task_slice_count_val = work_item[_WORK_SLICE_COUNT] + + # gA/gSFA are tiled in 128-row blocks; the fragment partition + # selects the within-block sub-tile. + sa_block_idx = task_m_tile_idx // Int32(self.sa_tiles_per_block) + tAgA_mk = tAgA[(None, sa_block_idx, None, Int32(0))] + sfa_block_idx = task_m_tile_idx // Int32(self.sfa_tiles_per_block) + tAgSFA_mk = tAgSFA[(None, sfa_block_idx, None, Int32(0))] slice_idx = Int32(0) while slice_idx < task_slice_count_val: intermediate_slice = task_slice_begin_idx + slice_idx diff --git a/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_micro_kernel.py b/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_micro_kernel.py index 6b1183b4c1f..af4775ce906 100644 --- a/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_micro_kernel.py +++ b/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_micro_kernel.py @@ -380,9 +380,10 @@ def __init__( self.acc_dtype = cutlass.Float32 self.sf_vec_size = sf_vec_size self.input_scales_are_reciprocal = input_scales_are_reciprocal - self.fast_math = fast_math self.activation = activation self.is_gated = is_gated_activation(activation) + # relu2's squared outputs need the exact quantizer and scale math. + self.fast_math = bool(fast_math) and self.is_gated self.swiglu_alpha = float(swiglu_alpha) self.swiglu_beta = float(swiglu_beta) self.swiglu_limit = float(swiglu_limit) if swiglu_limit is not None else None @@ -771,6 +772,9 @@ def __call__( grid=grid, block=[self.threads_per_cta, 1, 1], cluster=[1, 1, 1], + # A regular launch beside other stream work can admit only part + # of the grid, deadlocking the software grid barriers below. + cooperative=True, stream=stream, ) @@ -824,7 +828,7 @@ def kernel( _, _, gdim_z = cute.arch.grid_dim() warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) - is_cta_leader = Int32(1) if Int32(tidx) == Int32(0) else Int32(0) + is_cta_leader = Int32(Int32(tidx) == Int32(0)) if warp_idx == 0: cpasync.prefetch_descriptor(tma_a) diff --git a/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_static_kernel.py b/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_static_kernel.py index df8fde7f4d6..6cb17a48105 100644 --- a/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_static_kernel.py +++ b/flashinfer/fused_moe/cute_dsl/blackwell_sm12x/moe_static_kernel.py @@ -113,12 +113,13 @@ quantize_block_fp4, quantize_block_fp4_fast, get_ptr_as_int64, + ld_shared_i32_relaxed, st_global_f32, st_global_i32, shared_ptr_to_u32, st_shared_u8, st_global_u64, - scatter_add_bf16x2, + scatter_add_v4_bf16x2, ) from flashinfer.gemm.kernels.dense_blockscaled_gemm_sm120_b12x import ( Sm120B12xBlockScaledDenseGemmKernel as DenseGemmKernel, @@ -357,9 +358,10 @@ def __init__( self.acc_dtype = cutlass.Float32 self.sf_vec_size = sf_vec_size self.input_scales_are_reciprocal = input_scales_are_reciprocal - self.fast_math = fast_math self.activation = activation self.is_gated = is_gated_activation(activation) + # relu2's squared outputs need the exact quantizer and scale math. + self.fast_math = bool(fast_math) and self.is_gated self.swiglu_alpha = float(swiglu_alpha) self.swiglu_beta = float(swiglu_beta) self.swiglu_limit = float(swiglu_limit) if swiglu_limit is not None else None @@ -741,6 +743,9 @@ def __call__( grid=grid, block=[self.threads_per_cta, 1, 1], cluster=[1, 1, 1], + # A regular launch beside other stream work can admit only part + # of the grid, deadlocking the software grid barriers below. + cooperative=True, stream=stream, ) @@ -794,7 +799,7 @@ def kernel( _, _, gdim_z = cute.arch.grid_dim() warp_idx = cute.arch.warp_idx() warp_idx = cute.arch.make_warp_uniform(warp_idx) - is_cta_leader = Int32(1) if Int32(tidx) == Int32(0) else Int32(0) + is_cta_leader = Int32(Int32(tidx) == Int32(0)) if warp_idx == 0: cpasync.prefetch_descriptor(tma_a) @@ -2032,15 +2037,14 @@ class StorageRelu2: tRS_sD[(None, None, None, epi_buffer)], ) cute.arch.fence_proxy("async.shared", space="cta") - # No cross-warp barrier needed before scatter: - # StMatrix is warp-local, and each warp only reads - # its own 64×64 quadrant of sC below. + # The 8-wide reads from sC can cross another warp's + # stores, so wait for all MMA warps. + self.epilog_sync_barrier.arrive_and_wait() rows_offset = Int32(epi_m) * Int32(self.epi_tile[0]) # Per-warp scatter: each warp scatters its own quadrant - # of sC (64 M-rows × 64 N-cols). No cross-warp read - # dependencies, so no pre-scatter barrier is needed. + # of sC (64 M-rows × 64 N-cols). warp_epi_rows = ( valid_rows - tile_m_base - rows_offset - warp_m_base ) @@ -2049,50 +2053,90 @@ class StorageRelu2: if warp_epi_rows < Int32(0): warp_epi_rows = Int32(0) - pair_idx = lane_id - while pair_idx < warp_epi_rows * Int32(32): - local_row = pair_idx >> Int32(5) # / 32 - local_pair_col = pair_idx & Int32(31) # % 32 - global_col = ( - tile_n_base_cur - + warp_n_base - + local_pair_col * Int32(2) - ) + tile_vec_cols = Int32(64) // Int32(8) + vec_idx = lane_id + while vec_idx < warp_epi_rows * tile_vec_cols: + local_row = vec_idx // tile_vec_cols + local_vec_col = vec_idx - local_row * tile_vec_cols + local_col = warp_n_base + local_vec_col * Int32(8) + global_col = tile_n_base_cur + local_col cached_row = rows_offset + warp_m_base + local_row - # Only lane 0 loads tok/wv from gmem; broadcast via shuffle. - tok = Int32(0) - wv = cutlass.Float32(0.0) - if lane_id == Int32(0): - tok = _ld_shared_i32( - scatter_tok_base_addr + cached_row * Int32(4) - ) - wv = _ld_shared_f32( - scatter_weight_base_addr + cached_row * Int32(4) - ) - tok = cute.arch.shuffle_sync(tok, Int32(0)) - wv = cute.arch.shuffle_sync(wv, Int32(0)) + tok = ld_shared_i32_relaxed( + scatter_tok_base_addr + cached_row * Int32(4) + ) + wv = _ld_shared_f32( + scatter_weight_base_addr + cached_row * Int32(4) + ) sc_v0 = cutlass.Float32( sC[ warp_m_base + local_row, - warp_n_base + local_pair_col * Int32(2), + local_col, epi_buffer, ] ) sc_v1 = cutlass.Float32( sC[ warp_m_base + local_row, - warp_n_base + local_pair_col * Int32(2) + Int32(1), + local_col + Int32(1), + epi_buffer, + ] + ) + sc_v2 = cutlass.Float32( + sC[ + warp_m_base + local_row, + local_col + Int32(2), + epi_buffer, + ] + ) + sc_v3 = cutlass.Float32( + sC[ + warp_m_base + local_row, + local_col + Int32(3), + epi_buffer, + ] + ) + sc_v4 = cutlass.Float32( + sC[ + warp_m_base + local_row, + local_col + Int32(4), + epi_buffer, + ] + ) + sc_v5 = cutlass.Float32( + sC[ + warp_m_base + local_row, + local_col + Int32(5), + epi_buffer, + ] + ) + sc_v6 = cutlass.Float32( + sC[ + warp_m_base + local_row, + local_col + Int32(6), + epi_buffer, + ] + ) + sc_v7 = cutlass.Float32( + sC[ + warp_m_base + local_row, + local_col + Int32(7), epi_buffer, ] ) - scatter_add_bf16x2( + scatter_add_v4_bf16x2( get_ptr_as_int64( scatter_output, tok * scatter_N + global_col ), wv * sc_v0, wv * sc_v1, + wv * sc_v2, + wv * sc_v3, + wv * sc_v4, + wv * sc_v5, + wv * sc_v6, + wv * sc_v7, ) - pair_idx += Int32(self.num_threads_per_warp) + vec_idx += Int32(self.num_threads_per_warp) # Post-scatter barrier: needed to ensure all warps # finish scatter before next output tile's pipeline ops diff --git a/tests/moe/test_b12x_fused_moe.py b/tests/moe/test_b12x_fused_moe.py index d5ea517e01e..da88e688411 100644 --- a/tests/moe/test_b12x_fused_moe.py +++ b/tests/moe/test_b12x_fused_moe.py @@ -1273,6 +1273,84 @@ def test_w4a16_is_more_accurate_than_w4a4(self, num_tokens: int): f"a16={a16_mse.item():.6f}, a4={a4_mse.item():.6f}" ) + def test_dynamic_tile_ladder_reengages_after_large_call(self, monkeypatch): + """A cached large dynamic workspace must not pin later small calls to + the 128 M-tile: each tile band keeps its own cached workspace, and the + kernel build reads the tile from the workspace it runs with.""" + from flashinfer import b12x_fused_moe + from flashinfer.fused_moe.cute_dsl.blackwell_sm12x import moe_dispatch + + monkeypatch.setattr(moe_dispatch, "_FORCED_BACKEND", "dynamic") + + hidden_size, intermediate_size = 256, 512 + num_experts, top_k = 8, 2 + # routed_rows = 1024 >= 96 * 8 selects the 128 tile; 16 < 15 * 8 + # selects the 16 tile. + large_tokens, small_tokens = 512, 8 + + def run(num_tokens: int) -> None: + tensors = create_moe_tensors( + num_tokens=num_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + num_local_experts=num_experts, + top_k=top_k, + ) + result = b12x_fused_moe( + x=tensors["x_bf16"], + w1_weight=tensors["w1_weight"], + w1_weight_sf=tensors["w1_weight_sf"], + w1_alpha=tensors["w1_alpha"], + fc2_input_scale=tensors["fc2_input_scale"], + w2_weight=tensors["w2_weight"], + w2_weight_sf=tensors["w2_weight_sf"], + w2_alpha=tensors["w2_alpha"], + token_selected_experts=tensors["token_selected_experts"], + token_final_scales=tensors["token_final_scales"], + num_experts=num_experts, + top_k=top_k, + ) + ref_output = compute_reference_moe_fp4( + hidden_states=tensors["x_bf16"].float().cuda(), + gemm1_weights=tensors["w1_weight_bf16"].float().cuda(), + gemm2_weights=tensors["w2_weight_bf16"].float().cuda(), + token_selected_experts=tensors["token_selected_experts"], + token_final_scales=tensors["token_final_scales"], + num_tokens=num_tokens, + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + fc2_input_scale=tensors["fc2_input_scale"], + ) + passed, percent_within, atol = check_accuracy(result, ref_output) + assert passed, ( + f"Dynamic tile ladder: {percent_within * 100:.2f}% within " + f"tolerance (atol={atol:.4f}, tokens={num_tokens})" + ) + + def cached_workspace(num_tokens: int): + return moe_dispatch._get_cached_workspace( + backend="dynamic", + state_E=num_experts, + weight_E=num_experts, + routed_rows=num_tokens * top_k, + k=hidden_size, + n=intermediate_size, + num_topk=top_k, + device=torch.device("cuda", torch.cuda.current_device()), + ) + + run(large_tokens) + ws_large = cached_workspace(large_tokens) + assert ws_large.tile_m == 128 + + run(small_tokens) + ws_small = cached_workspace(small_tokens) + assert ws_small is not ws_large + assert ws_small.tile_m == 16 + # ============================================================================= # Test Class: Wrapper API (B12xMoEWrapper) @@ -1784,6 +1862,91 @@ def test_micro_wrapper_accuracy(self): f"(atol={atol:.4f})" ) + @pytest.mark.parametrize( + "activation,num_tokens,top_k", + [ + ("silu", 1, 2), + ("silu", 2, 2), + ("silu", 8, 2), + ("silu", 1, 8), + ("silu", 2, 8), + ("silu", 8, 8), + ("gelu_tanh", 1, 2), + ("gelu_tanh", 8, 2), + ("swigluoai_uninterleave", 1, 2), + ("swigluoai_uninterleave", 8, 2), + ], + ) + def test_direct_micro_forced_accuracy( + self, monkeypatch, activation: str, num_tokens: int, top_k: int + ): + """Accuracy of the forced direct micro backend on tiny decode shapes. + + The forced hook raises instead of falling back, so a pass means the + direct micro kernel produced the result. + """ + from flashinfer import b12x_fused_moe + from flashinfer.fused_moe.cute_dsl.blackwell_sm12x import moe_dispatch + + monkeypatch.setattr(moe_dispatch, "_FORCED_BACKEND", "direct_micro") + + hidden_size, intermediate_size = 256, 512 + num_experts = 256 + swiglu_limit = 7.0 if activation == "swigluoai_uninterleave" else None + + tensors = create_moe_tensors( + num_tokens=num_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + num_local_experts=num_experts, + top_k=top_k, + ) + + result = b12x_fused_moe( + x=tensors["x_bf16"], + w1_weight=tensors["w1_weight"], + w1_weight_sf=tensors["w1_weight_sf"], + w1_alpha=tensors["w1_alpha"], + fc2_input_scale=tensors["fc2_input_scale"], + w2_weight=tensors["w2_weight"], + w2_weight_sf=tensors["w2_weight_sf"], + w2_alpha=tensors["w2_alpha"], + token_selected_experts=tensors["token_selected_experts"], + token_final_scales=tensors["token_final_scales"], + num_experts=num_experts, + top_k=top_k, + activation=activation, + swiglu_limit=swiglu_limit, + ) + + assert result.shape == (num_tokens, hidden_size) + assert not torch.isnan(result).any() + assert not torch.isinf(result).any() + + ref_output = compute_reference_moe_fp4( + hidden_states=tensors["x_bf16"].float().cuda(), + gemm1_weights=tensors["w1_weight_bf16"].float().cuda(), + gemm2_weights=tensors["w2_weight_bf16"].float().cuda(), + token_selected_experts=tensors["token_selected_experts"], + token_final_scales=tensors["token_final_scales"], + num_tokens=num_tokens, + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + fc2_input_scale=tensors["fc2_input_scale"], + activation=activation, + swiglu_limit=swiglu_limit, + ) + + passed, percent_within, atol = check_accuracy(result, ref_output) + assert passed, ( + f"Direct micro: {percent_within * 100:.2f}% within tolerance " + f"(atol={atol:.4f}, act={activation}, tokens={num_tokens}, " + f"top_k={top_k})" + ) + @pytest.mark.parametrize("num_tokens", [1, 2, 4]) def test_w4a16_direct_micro_functional_accuracy(self, num_tokens: int): """Accuracy test for the W4A16 small-batch route-packing path.""" @@ -2379,6 +2542,125 @@ def test_relu2_micro_accuracy(self): f"(atol={atol:.4f})" ) + @pytest.mark.parametrize("num_tokens", [1, 4]) + def test_relu2_direct_micro_forced_accuracy(self, monkeypatch, num_tokens: int): + """Forced direct micro backend with the non-gated ReLU2 activation.""" + from flashinfer import b12x_fused_moe + from flashinfer.fused_moe.cute_dsl.blackwell_sm12x import moe_dispatch + + monkeypatch.setattr(moe_dispatch, "_FORCED_BACKEND", "direct_micro") + + hidden_size, intermediate_size = 256, 512 + num_experts, top_k = 256, 2 + + tensors = create_relu2_moe_tensors( + num_tokens=num_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + num_local_experts=num_experts, + top_k=top_k, + ) + + result = b12x_fused_moe( + x=tensors["x_bf16"], + w1_weight=tensors["w1_weight"], + w1_weight_sf=tensors["w1_weight_sf"], + w1_alpha=tensors["w1_alpha"], + fc2_input_scale=tensors["fc2_input_scale"], + w2_weight=tensors["w2_weight"], + w2_weight_sf=tensors["w2_weight_sf"], + w2_alpha=tensors["w2_alpha"], + token_selected_experts=tensors["token_selected_experts"], + token_final_scales=tensors["token_final_scales"], + num_experts=num_experts, + top_k=top_k, + activation="relu2", + ) + + assert result.shape == (num_tokens, hidden_size) + assert not torch.isnan(result).any() + + ref_output = compute_reference_moe_relu2( + hidden_states=tensors["x_bf16"].float().cuda(), + fc1_weights=tensors["w1_weight_bf16"].float().cuda(), + fc2_weights=tensors["w2_weight_bf16"].float().cuda(), + token_selected_experts=tensors["token_selected_experts"], + token_final_scales=tensors["token_final_scales"], + num_tokens=num_tokens, + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + fc2_input_scale=tensors["fc2_input_scale"], + ) + + passed, percent_within, atol = check_accuracy(result, ref_output) + assert passed, ( + f"ReLU2 direct micro: {percent_within * 100:.2f}% within tolerance " + f"(atol={atol:.4f}, tokens={num_tokens})" + ) + + def test_relu2_direct_micro_ignores_swiglu_limit(self, monkeypatch): + """A caller-supplied swiglu_limit must be accepted and ignored for + relu2 on the direct micro path, like on the MMA kernels.""" + from flashinfer import b12x_fused_moe + from flashinfer.fused_moe.cute_dsl.blackwell_sm12x import moe_dispatch + + monkeypatch.setattr(moe_dispatch, "_FORCED_BACKEND", "direct_micro") + + num_tokens, hidden_size, intermediate_size = 2, 256, 512 + num_experts, top_k = 256, 2 + + tensors = create_relu2_moe_tensors( + num_tokens=num_tokens, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + num_local_experts=num_experts, + top_k=top_k, + ) + + result = b12x_fused_moe( + x=tensors["x_bf16"], + w1_weight=tensors["w1_weight"], + w1_weight_sf=tensors["w1_weight_sf"], + w1_alpha=tensors["w1_alpha"], + fc2_input_scale=tensors["fc2_input_scale"], + w2_weight=tensors["w2_weight"], + w2_weight_sf=tensors["w2_weight_sf"], + w2_alpha=tensors["w2_alpha"], + token_selected_experts=tensors["token_selected_experts"], + token_final_scales=tensors["token_final_scales"], + num_experts=num_experts, + top_k=top_k, + activation="relu2", + swiglu_limit=7.0, + ) + + assert result.shape == (num_tokens, hidden_size) + assert not torch.isnan(result).any() + + ref_output = compute_reference_moe_relu2( + hidden_states=tensors["x_bf16"].float().cuda(), + fc1_weights=tensors["w1_weight_bf16"].float().cuda(), + fc2_weights=tensors["w2_weight_bf16"].float().cuda(), + token_selected_experts=tensors["token_selected_experts"], + token_final_scales=tensors["token_final_scales"], + num_tokens=num_tokens, + num_experts=num_experts, + top_k=top_k, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + fc2_input_scale=tensors["fc2_input_scale"], + ) + + passed, percent_within, atol = check_accuracy(result, ref_output) + assert passed, ( + f"ReLU2 with ignored swiglu_limit: {percent_within * 100:.2f}% " + f"within tolerance (atol={atol:.4f})" + ) + def test_relu2_w4a16_direct_micro_accuracy(self): """Accuracy test for ReLU2 with the W4A16 route-packing small-batch path.""" from flashinfer import b12x_fused_moe