diff --git a/.github/scripts/compare_benchmark.py b/.github/scripts/compare_benchmark.py index 8cf4229549..56c0f40e3d 100644 --- a/.github/scripts/compare_benchmark.py +++ b/.github/scripts/compare_benchmark.py @@ -48,7 +48,6 @@ "preshuffle", "strict_accuracy", "check_aot_cache", - "swiglu_limit", # Source cols folded into derived `hip` column below "hidden_pad", "intermediate_pad", diff --git a/aiter/aot/flydsl/moe.py b/aiter/aot/flydsl/moe.py index dd3379a25f..bac48f1c4c 100644 --- a/aiter/aot/flydsl/moe.py +++ b/aiter/aot/flydsl/moe.py @@ -49,6 +49,7 @@ compile_flydsl_moe_stage1, compile_flydsl_moe_stage2, get_flydsl_kernel_params, + runtime_swiglu_limit, ) # Keep the default AOT coverage aligned with runtime config resolution. @@ -58,22 +59,6 @@ MOE_AOT_ARCH_DEFAULT = "gfx950" -def _parse_optional_float(value, source: str) -> float | None: - if value is None: - return None - value = str(value).strip() - if value == "": - return None - try: - return float(value) - except ValueError as e: - raise ValueError(f"{source} must be a float, got {value!r}") from e - - -def _row_swiglu_limit(row: dict[str, str]) -> float: - return _parse_optional_float(row.get("swiglu_limit"), "swiglu_limit") or 0.0 - - def parse_csv(csv_path: str): """Parse the CSV and return a list of unique compile jobs. @@ -107,7 +92,6 @@ def parse_csv(csv_path: str): q_type = row.get("q_type", "") dtype = row.get("dtype", "") q_dtype_w = row.get("q_dtype_w", "") - swiglu_limit = _row_swiglu_limit(row) # Cover both runtime bias choices for fp4-weight MoE. Model configs # share kernel families, and runtime bias selection can vary by # activation dtype/model semantics. @@ -151,7 +135,6 @@ def parse_csv(csv_path: str): "enable_bias": enable_bias, "token_num": token, "block_m": block_m, - "swiglu_limit": swiglu_limit, } # Stage2 needs to know whether stage1 fuses fp4/fp8 quant — # this changes the shape of a2_scale (sorted scale buffer @@ -209,7 +192,6 @@ def _precompile_to_cache( xcd_swizzle: int = 0, enable_bias: bool = False, stage1_fuse_quant=None, - swiglu_limit: float = 0.0, k_wave: int = 1, # Stage2-only kernel tuning knobs (registered by the production-variant # entries in `get_flydsl_stage2_kernels`). Forwarded into @@ -356,7 +338,7 @@ def _make_a2_scale_for_stage2(): _padded_rows * _padded_cols, dtype=torch.uint8, device=dev ) if a_dtype == "fp8": - if act == "silu" and swiglu_limit == 0.0: + if act == "silu": # fused_moe_2stages uses fused_quant_fp8_sort for this path. rows = (max_num_tokens_padded + 31) // 32 * 32 cols = (inter_dim + 31) // 32 @@ -529,6 +511,7 @@ def _make_a_user(a_dtype_user_shape): else torch.empty(0, device=dev) ), stream=0, + swiglu_limit=runtime_swiglu_limit(None, act), ) else: args = _s1_args_std( @@ -569,7 +552,6 @@ def _make_a_user(a_dtype_user_shape): enable_bias=(kernel_bias is not None), a_scale_one=a_scale_one, xcd_swizzle=xcd_swizzle, - swiglu_limit=swiglu_limit, k_wave=k_wave, ) _run_compiled(exe, args) @@ -591,7 +573,6 @@ def _make_a_user(a_dtype_user_shape): gui_layout=gui_layout, act=act, enable_bias=False, - swiglu_limit=swiglu_limit, ) _run_compiled( silu_fused, @@ -605,6 +586,7 @@ def _make_a_user(a_dtype_user_shape): _ptr_view_safe(torch.empty(0, device=dev, dtype=torch.float32)), tokens, sorted_token_ids.shape[0], + runtime_swiglu_limit(None, act), 0, ), ) diff --git a/aiter/configs/model_configs/dsv4_fp8fp4_untuned_fmoe.csv b/aiter/configs/model_configs/dsv4_fp8fp4_untuned_fmoe.csv index bb279a43b6..60eaf47a81 100644 --- a/aiter/configs/model_configs/dsv4_fp8fp4_untuned_fmoe.csv +++ b/aiter/configs/model_configs/dsv4_fp8fp4_untuned_fmoe.csv @@ -46,4 +46,4 @@ token,model_dim,inter_dim,expert,topk,act_type,dtype,q_dtype_a,q_dtype_w,q_type, 4096,7168,1536,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 8192,7168,1536,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 16384,7168,1536,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 -32768,7168,1536,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 \ No newline at end of file +32768,7168,1536,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py index f02926953f..2c87e81fd9 100644 --- a/aiter/fused_moe.py +++ b/aiter/fused_moe.py @@ -283,7 +283,7 @@ def fused_moe( bias1=None, bias2=None, splitk=0, - swiglu_limit=0.0, + swiglu_limit=None, gate_mode: Optional[str] = GateMode.SEPARATED.value, ): if not block_size_M: @@ -339,7 +339,7 @@ def fused_moe_fake( intermediate_pad: int = 0, bias1: Optional[torch.Tensor] = None, bias2: Optional[torch.Tensor] = None, - swiglu_limit: float = 0.0, + swiglu_limit: Optional[float] = None, gate_mode: str = GateMode.SEPARATED.value, ) -> torch.Tensor: device = topk_ids.device @@ -375,7 +375,7 @@ def fused_moe_( intermediate_pad: int = 0, bias1: Optional[torch.Tensor] = None, bias2: Optional[torch.Tensor] = None, - swiglu_limit: float = 0.0, + swiglu_limit: Optional[float] = None, gate_mode: str = GateMode.SEPARATED.value, ) -> torch.Tensor: # We do such convert since custom_op schema restriction on block_size_M, and Enum type @@ -472,7 +472,9 @@ def fused_moe_( bias1=bias1, bias2=bias2, gate_mode=gate_mode, + swiglu_limit=swiglu_limit, ) + if grouped_a8w4_out is not None: return grouped_a8w4_out @@ -919,7 +921,7 @@ def _flydsl_stage1_wrapper( out_scale_sorted=None, bias1=None, topk_ids=None, - swiglu_limit: float = 0.0, + swiglu_limit: Optional[float] = None, inter_dim_pad: int = 0, model_dim_pad: int = 0, **_kwargs, @@ -1760,7 +1762,7 @@ def fused_moe_2stages( bias2=None, topk_ids=None, topk_weights=None, - swiglu_limit=0.0, + swiglu_limit=None, gate_mode=GateMode.SEPARATED.value, expert_mask=None, ): @@ -2168,8 +2170,8 @@ def torch_moe( # temp workaround for swiglu -def swiglu(x_glu, x_linear, alpha: float = 1.702, limit: float = 7.0): - if limit == 0.0: +def swiglu(x_glu, x_linear, alpha: float = 1.702, limit: Optional[float] = 7.0): + if limit is None: limit = 7.0 # Clamp the input values x_glu = x_glu.clamp(min=None, max=limit) @@ -2193,7 +2195,7 @@ def torch_moe_stage1( w1_scale=None, # [expert, inter_dim, 1] w1_bias=None, # [expert, inter_dim, 1] doweight=False, - swiglu_limit=0.0, + swiglu_limit=None, ): quant_type = quant_remap.get(quant_type, quant_type) ctype = dtypes.fp32 # compute type @@ -2300,7 +2302,7 @@ def torch_moe_stage1( if use_swiglu: out = swiglu(gate, up, limit=swiglu_limit) else: - if swiglu_limit != 0: + if swiglu_limit: gate = gate.clamp(min=None, max=swiglu_limit) up = up.clamp(min=-swiglu_limit, max=swiglu_limit) out = torch_act(gate) * up diff --git a/aiter/ops/flydsl/grouped_moe_gfx1250.py b/aiter/ops/flydsl/grouped_moe_gfx1250.py index bd82a671c1..45217f4349 100644 --- a/aiter/ops/flydsl/grouped_moe_gfx1250.py +++ b/aiter/ops/flydsl/grouped_moe_gfx1250.py @@ -327,6 +327,7 @@ def _maybe_grouped_gfx1250_a8w4_moe( bias1: Optional[torch.Tensor], bias2: Optional[torch.Tensor], gate_mode: GateMode = GateMode.SEPARATED, + swiglu_limit: Optional[float] = None, ): def _grouped_dbg(msg: str, stacklevel: int = 1): if os.environ.get("AITER_GROUPED_DEBUG", "0") not in ( @@ -788,6 +789,7 @@ def _quantize_mxfp8_payload(x: torch.Tensor, last_dim: int): model_dim, E, stream=torch.cuda.current_stream(), + swiglu_limit=swiglu_limit, _m_tile_prefix=m_tile_prefix, _m_tile_map=m_tile_map, bias=_bias1_arg, diff --git a/aiter/ops/flydsl/kernels/gemm_mxscale_gfx1250.py b/aiter/ops/flydsl/kernels/gemm_mxscale_gfx1250.py index dae35a5e3f..e383bbf599 100644 --- a/aiter/ops/flydsl/kernels/gemm_mxscale_gfx1250.py +++ b/aiter/ops/flydsl/kernels/gemm_mxscale_gfx1250.py @@ -579,6 +579,7 @@ def kernel_mxscale_gemm( i32_m_tile_bound: fx.Int32, i32_m: fx.Int32, i32_n: fx.Int32, + f32_swiglu_limit: fx.Float32, ): # Enable back-to-back WMMA issue (SCHED_MODE bit[4] = DISABLE_VALU_STALL) rocdl.disable_xdl_arb_stall() @@ -1626,12 +1627,14 @@ def _stage1_silu_elem(g): def _stage1_act_mul_scalar(g, u): one = arith.constant(1.0, type=T.f32) alpha = arith.constant(1.702, type=T.f32) - limit = arith.constant(7.0, type=T.f32) - neg_limit = arith.constant(-7.0, type=T.f32) neg_log2e = arith.constant(-1.4426950408889634, type=T.f32) + # Runtime clamp bound: host passes the limit (7.0 default for + # swiglu) or +inf to disable clamping (silu without a limit). + # min(x, lim) == -max(-x, -lim), expressed via wrapped maximumf. + neg_lim = -f32_swiglu_limit + g = -((-g).maximumf(neg_lim)) + u = (-((-u).maximumf(neg_lim))).maximumf(neg_lim) if const_expr(stage1_act_mode == "swiglu"): - g = arith.minimumf(g, limit) - u = arith.maximumf(arith.minimumf(u, limit), neg_limit) emu = llvm.call_intrinsic( T.f32, "llvm.amdgcn.exp2.f32", [g * alpha * neg_log2e], [], [] ) @@ -3130,6 +3133,7 @@ def launch_mxscale_gemm( i32_m: fx.Int32, i32_n: fx.Int32, stream: fx.Stream, + swiglu_limit_f: fx.Float32 = fx.Float32(float("inf")), ): _ = cache_tag ctx = CompilationContext.get_current() @@ -3158,6 +3162,7 @@ def launch_mxscale_gemm( i32_m, i32_m, i32_n, + swiglu_limit_f, ) for op in ctx.gpu_module_body.operations: if const_expr( @@ -3195,6 +3200,7 @@ def launch_mxscale_gemm_masked( i32_m: fx.Int32, i32_n: fx.Int32, stream: fx.Stream, + swiglu_limit_f: fx.Float32 = fx.Float32(float("inf")), ): _ = cache_tag ctx = CompilationContext.get_current() @@ -3228,6 +3234,7 @@ def launch_mxscale_gemm_masked( i32_m_tile_bound, i32_m, i32_n, + swiglu_limit_f, ) for op in ctx.gpu_module_body.operations: if const_expr( @@ -3264,6 +3271,7 @@ def launch_mxscale_gemm_masked_persistent( i32_m: fx.Int32, i32_n: fx.Int32, stream: fx.Stream, + swiglu_limit_f: fx.Float32 = fx.Float32(float("inf")), ): _ = cache_tag ctx = CompilationContext.get_current() @@ -3289,6 +3297,7 @@ def launch_mxscale_gemm_masked_persistent( i32_m, i32_m, i32_n, + swiglu_limit_f, ) for op in ctx.gpu_module_body.operations: if const_expr( @@ -3317,6 +3326,7 @@ def launch_mxscale_gemm_bias( i32_m: fx.Int32, i32_n: fx.Int32, stream: fx.Stream, + swiglu_limit_f: fx.Float32 = fx.Float32(float("inf")), ): _ = cache_tag ctx = CompilationContext.get_current() @@ -3345,6 +3355,7 @@ def launch_mxscale_gemm_bias( i32_m, i32_m, i32_n, + swiglu_limit_f, ) for op in ctx.gpu_module_body.operations: if const_expr( @@ -3383,6 +3394,7 @@ def launch_mxscale_gemm_masked_bias( i32_m: fx.Int32, i32_n: fx.Int32, stream: fx.Stream, + swiglu_limit_f: fx.Float32 = fx.Float32(float("inf")), ): _ = cache_tag ctx = CompilationContext.get_current() @@ -3416,6 +3428,7 @@ def launch_mxscale_gemm_masked_bias( i32_m_tile_bound, i32_m, i32_n, + swiglu_limit_f, ) for op in ctx.gpu_module_body.operations: if const_expr( @@ -3453,6 +3466,7 @@ def launch_mxscale_gemm_masked_persistent_bias( i32_m: fx.Int32, i32_n: fx.Int32, stream: fx.Stream, + swiglu_limit_f: fx.Float32 = fx.Float32(float("inf")), ): _ = cache_tag ctx = CompilationContext.get_current() @@ -3478,6 +3492,7 @@ def launch_mxscale_gemm_masked_persistent_bias( i32_m, i32_m, i32_n, + swiglu_limit_f, ) for op in ctx.gpu_module_body.operations: if const_expr( diff --git a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py index c1be89dbbc..f5ed73d354 100644 --- a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py +++ b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py @@ -95,7 +95,6 @@ def compile_mixed_moe_gemm1( gate_mode: GateMode = GateMode.SEPARATED, a_scale_one: bool = False, xcd_swizzle: int = 0, - swiglu_limit: float = 0.0, k_wave: int = 1, ): """Compile stage1 kernel: act(X @ W_gate.T, X @ W_up.T) -> [tokens*topk, inter_dim].""" @@ -389,6 +388,7 @@ def moe_gemm1( i32_n_in: fx.Int32, i32_k_in: fx.Int32, i32_size_expert_ids_in: fx.Int32, + f32_swiglu_limit: fx.Float32, ): tokens_in = arith.index_cast(ir.IndexType.get(), i32_tokens_in.ir_value()) @@ -397,6 +397,11 @@ def moe_gemm1( size_expert_ids_in = arith.index_cast( ir.IndexType.get(), i32_size_expert_ids_in.ir_value() ) + # Runtime clamp bound for the activation. Host passes the configured + # swiglu_limit (7.0 default for swiglu) or +inf to disable clamping. + # ``-lim`` is precomputed once; ``min(x, lim) == -max(-x, -lim)`` so + # the kernel uses only the wrapped maximumf/negation ops. + swiglu_neg_limit = -f32_swiglu_limit x_elem = T.f8 f32 = T.f32 @@ -1899,15 +1904,22 @@ def silu_elem(g): sig = llvm.call_intrinsic(f32, "llvm.amdgcn.rcp.f32", [den], [], []) return g * sig + def _clamp_gate(x): + # min(x, lim) == -max(-x, -lim); upper bound only. + return -((-x).maximumf(swiglu_neg_limit)) + + def _clamp_lin(x): + # clamp to [-lim, lim]. + return (-((-x).maximumf(swiglu_neg_limit))).maximumf( + swiglu_neg_limit + ) + def silu_mul_vec4(gate_v4, up_v4): """Element-wise silu(gate) * up on vec4_f32. - When swiglu_limit != 0, clamp gate <= limit and - -limit <= up <= limit before applying silu(gate) * up. + Clamp gate <= limit and -limit <= up <= limit (runtime limit; + +inf disables the clamp) before applying silu(gate) * up. """ result_elems = [] - if const_expr(swiglu_limit != 0): - limit = arith.constant(float(swiglu_limit), type=f32) - neg_limit = arith.constant(-float(swiglu_limit), type=f32) for ei in range_constexpr(4): g = vector.extract( gate_v4, static_position=[ei], dynamic_position=[] @@ -1915,29 +1927,21 @@ def silu_mul_vec4(gate_v4, up_v4): u = vector.extract( up_v4, static_position=[ei], dynamic_position=[] ) - if const_expr(swiglu_limit != 0): - g = arith.minimumf(g, limit) - u = arith.minimumf(u, limit) - u = arith.maximumf(u, neg_limit) + g = _clamp_gate(g) + u = _clamp_lin(u) result_elems.append(silu_elem(g) * u) return vector.from_elements(vec4_f32, result_elems) def swiglu_mul_vec4(gate_v4, up_v4): """Element-wise swiglu(gate, up) on vec4_f32. swiglu(g, u) = g * sigmoid(alpha * g) * (u + 1) - When swiglu_limit != 0, clamp gate <= limit and - -limit <= up <= limit before the activation. + Clamp gate <= limit and -limit <= up <= limit (runtime limit, + 7.0 default) before the activation. """ result_elems = [] alpha = arith.constant(1.702, type=f32) one = arith.constant(1.0, type=f32) neg_log2e = arith.constant(-1.4426950408889634, type=f32) - if const_expr(swiglu_limit != 0): - limit = arith.constant(float(swiglu_limit), type=f32) - neg_limit = arith.constant(-float(swiglu_limit), type=f32) - else: - limit = arith.constant(float(7.0), type=f32) - neg_limit = arith.constant(-float(7.0), type=f32) for ei in range_constexpr(4): g = vector.extract( @@ -1946,9 +1950,8 @@ def swiglu_mul_vec4(gate_v4, up_v4): u = vector.extract( up_v4, static_position=[ei], dynamic_position=[] ) - g = arith.minimumf(g, limit) - u = arith.minimumf(u, limit) - u = arith.maximumf(u, neg_limit) + g = _clamp_gate(g) + u = _clamp_lin(u) t = g * alpha * neg_log2e emu = llvm.call_intrinsic( f32, "llvm.amdgcn.exp2.f32", [t], [], [] @@ -1975,16 +1978,8 @@ def act_elem(g, u): alpha = arith.constant(1.702, type=f32) one = arith.constant(1.0, type=f32) neg_log2e = arith.constant(-1.4426950408889634, type=f32) - lim = arith.constant( - float(swiglu_limit) if swiglu_limit != 0 else 7.0, type=f32 - ) - nlim = arith.constant( - -float(swiglu_limit) if swiglu_limit != 0 else -7.0, - type=f32, - ) - g = arith.minimumf(g, lim) - u = arith.minimumf(u, lim) - u = arith.maximumf(u, nlim) + g = _clamp_gate(g) + u = _clamp_lin(u) t = g * alpha * neg_log2e emu = llvm.call_intrinsic( f32, "llvm.amdgcn.exp2.f32", [t], [], [] @@ -1995,12 +1990,8 @@ def act_elem(g, u): ) return g * sig * (u + one) else: - if const_expr(swiglu_limit != 0): - lim = arith.constant(float(swiglu_limit), type=f32) - nlim = arith.constant(-float(swiglu_limit), type=f32) - g = arith.minimumf(g, lim) - u = arith.minimumf(u, lim) - u = arith.maximumf(u, nlim) + g = _clamp_gate(g) + u = _clamp_lin(u) return silu_elem(g) * u kwave_fused = const_expr( @@ -2854,6 +2845,7 @@ def launch_mixed_moe_gemm1( i32_inter_in: fx.Int32, i32_k_in: fx.Int32, i32_size_expert_ids_in: fx.Int32, + f32_swiglu_limit: fx.Float32, stream: fx.Stream, ): _ = cache_tag @@ -2908,6 +2900,7 @@ def launch_mixed_moe_gemm1( i32_inter_in, i32_k_in, i32_size_expert_ids_in, + f32_swiglu_limit, ).launch(grid=(gx, gy, k_batch), block=(total_threads, 1, 1), stream=stream) return launch_mixed_moe_gemm1 @@ -3137,7 +3130,6 @@ def moe_gemm2( i32 = T.i32 i64 = T.i64 vec4_f32 = T.vec(4, f32) - vec4_i32 = T.vec(4, i32) vec16_elems = 16 if a_elem_bytes == 1 else 8 vec8_elems = 8 if a_elem_bytes == 1 else 4 vec4_elems = 4 if a_elem_bytes == 1 else 2 diff --git a/aiter/ops/flydsl/kernels/moe_grouped_gemm_mxscale_gfx1250.py b/aiter/ops/flydsl/kernels/moe_grouped_gemm_mxscale_gfx1250.py index 8fc42e3f5d..2d61f3a908 100644 --- a/aiter/ops/flydsl/kernels/moe_grouped_gemm_mxscale_gfx1250.py +++ b/aiter/ops/flydsl/kernels/moe_grouped_gemm_mxscale_gfx1250.py @@ -65,6 +65,7 @@ class _GroupedA8W4Config: persistent_workers: Optional[int] = None data_format: str = "a8w4" act: str = "silu" + swiglu_limit: float | None = None stage1_weight_layout: str = "gguu" @@ -336,11 +337,20 @@ def _check_stage1_args( ) -def _apply_gate_up(gate: torch.Tensor, up: torch.Tensor, act: str) -> torch.Tensor: +def _apply_gate_up( + gate: torch.Tensor, + up: torch.Tensor, + act: str, + swiglu_limit: float | None = None, +) -> torch.Tensor: + _lim = 7.0 if swiglu_limit is None else float(swiglu_limit) if act == "swiglu": - gate = gate.clamp(max=7.0) - up = up.clamp(min=-7.0, max=7.0) + gate = gate.clamp(max=_lim) + up = up.clamp(min=-_lim, max=_lim) return gate * torch.sigmoid(1.702 * gate) * (up + 1.0) + if swiglu_limit is not None: + gate = gate.clamp(max=_lim) + up = up.clamp(min=-_lim, max=_lim) return torch.nn.functional.silu(gate) * up @@ -377,6 +387,7 @@ def stage1_finalize_act_kernel( arg_y: fx.Tensor, arg_tmp: fx.Tensor, arg_masked_m: fx.Tensor, + swiglu_limit_f: fx.Float32, ): elem_ty = T.bf16 if out_dtype == "bf16" else T.f16 tx = arith.index_cast(T.index, _raw(gpu.thread_id("x"))) @@ -435,12 +446,14 @@ def stage1_finalize_act_kernel( u = up_h.extf(T.f32) one = arith.constant(1.0, type=T.f32) neg_log2e = arith.constant(-1.4426950408889634, type=T.f32) + # Runtime clamp bound: host passes the limit (7.0 default for + # swiglu) or +inf to disable clamping (silu without a limit). + # min(x, lim) == -max(-x, -lim), expressed via wrapped maximumf. + neg_lim = -swiglu_limit_f + g = -((-g).maximumf(neg_lim)) + u = (-((-u).maximumf(neg_lim))).maximumf(neg_lim) if const_expr(act == "swiglu"): - limit = arith.constant(7.0, type=T.f32) - neg_limit = arith.constant(-7.0, type=T.f32) alpha = arith.constant(1.702, type=T.f32) - g = arith.minimumf(g, limit) - u = arith.maximumf(arith.minimumf(u, limit), neg_limit) t = g * alpha * neg_log2e emu = llvm.call_intrinsic( T.f32, "llvm.amdgcn.exp2.f32", [t], [], [] @@ -468,6 +481,7 @@ def launch_stage1_finalize_act( arg_y: fx.Tensor, arg_tmp: fx.Tensor, arg_masked_m: fx.Tensor, + swiglu_limit_f: fx.Float32, stream: fx.Stream, ): ctx = CompilationContext.get_current() @@ -476,7 +490,9 @@ def launch_stage1_finalize_act( gx = (arith.index(total_elems) + arith.index(block_threads - 1)) // arith.index( block_threads ) - launcher = stage1_finalize_act_kernel(arg_y, arg_tmp, arg_masked_m) + launcher = stage1_finalize_act_kernel( + arg_y, arg_tmp, arg_masked_m, swiglu_limit_f + ) launcher.launch( grid=(_raw(gx), 1, 1), block=(block_threads, 1, 1), @@ -521,6 +537,7 @@ def stage1_finalize_act_bias_kernel( arg_tmp: fx.Tensor, arg_bias: fx.Tensor, arg_masked_m: fx.Tensor, + swiglu_limit_f: fx.Float32, ): elem_ty = T.bf16 if out_dtype == "bf16" else T.f16 tx = arith.index_cast(T.index, _raw(gpu.thread_id("x"))) @@ -597,12 +614,14 @@ def stage1_finalize_act_bias_kernel( u = up_h.extf(T.f32) + up_bias_h.extf(T.f32) one = arith.constant(1.0, type=T.f32) neg_log2e = arith.constant(-1.4426950408889634, type=T.f32) + # Runtime clamp bound: host passes the limit (7.0 default for + # swiglu) or +inf to disable clamping (silu without a limit). + # min(x, lim) == -max(-x, -lim), expressed via wrapped maximumf. + neg_lim = -swiglu_limit_f + g = -((-g).maximumf(neg_lim)) + u = (-((-u).maximumf(neg_lim))).maximumf(neg_lim) if const_expr(act == "swiglu"): - limit = arith.constant(7.0, type=T.f32) - neg_limit = arith.constant(-7.0, type=T.f32) alpha = arith.constant(1.702, type=T.f32) - g = arith.minimumf(g, limit) - u = arith.maximumf(arith.minimumf(u, limit), neg_limit) t = g * alpha * neg_log2e emu = llvm.call_intrinsic( T.f32, "llvm.amdgcn.exp2.f32", [t], [], [] @@ -631,6 +650,7 @@ def launch_stage1_finalize_act_bias( arg_tmp: fx.Tensor, arg_bias: fx.Tensor, arg_masked_m: fx.Tensor, + swiglu_limit_f: fx.Float32, stream: fx.Stream, ): ctx = CompilationContext.get_current() @@ -640,7 +660,7 @@ def launch_stage1_finalize_act_bias( block_threads ) launcher = stage1_finalize_act_bias_kernel( - arg_y, arg_tmp, arg_bias, arg_masked_m + arg_y, arg_tmp, arg_bias, arg_masked_m, swiglu_limit_f ) launcher.launch( grid=(_raw(gx), 1, 1), @@ -915,6 +935,7 @@ def launch( experts_arg, *, stream=None, + swiglu_limit=None, _gemm_events=None, _m_tile_prefix=None, _m_tile_map=None, @@ -951,6 +972,12 @@ def launch( _check_bias_args("bias", bias, (cfg.experts, 2 * cfg.inter_dim), y) if stream is None: stream = torch.cuda.current_stream() + # Runtime clamp bound passed to the act epilogue / finalize kernels. + # swiglu defaults to 7.0; silu without a limit uses +inf (no clamp). + if cfg.act == "swiglu": + _swiglu_lim_rt = float(swiglu_limit) if swiglu_limit else 7.0 + else: + _swiglu_lim_rt = float(swiglu_limit) if swiglu_limit else float("inf") fused_gemm = _get_fused_base_bias() if bias is not None else _get_fused_base() use_fused_gemm = ( fused_gemm is not None @@ -996,6 +1023,7 @@ def launch( cfg.max_m, fused_n, stream, + _swiglu_lim_rt, ) else: _run_compiled( @@ -1011,6 +1039,7 @@ def launch( cfg.max_m, fused_n, stream, + _swiglu_lim_rt, ) else: _run_compiled( @@ -1026,6 +1055,7 @@ def launch( cfg.max_m, 2 * cfg.inter_dim, stream, + _swiglu_lim_rt, ) if _gemm_events is not None: _gemm_events[1].record(stream) @@ -1056,6 +1086,7 @@ def launch( contiguous_m, fused_n, stream, + _swiglu_lim_rt, ) else: _run_compiled( @@ -1072,6 +1103,7 @@ def launch( contiguous_m, fused_n, stream, + _swiglu_lim_rt, ) else: _run_compiled( @@ -1088,6 +1120,7 @@ def launch( contiguous_m, 2 * cfg.inter_dim, stream, + _swiglu_lim_rt, ) if _gemm_events is not None: _gemm_events[1].record(stream) @@ -1114,6 +1147,7 @@ def launch( cfg.max_m, fused_n, stream, + _swiglu_lim_rt, ) else: _run_compiled( @@ -1130,6 +1164,7 @@ def launch( cfg.max_m, fused_n, stream, + _swiglu_lim_rt, ) else: _run_compiled( @@ -1146,6 +1181,7 @@ def launch( cfg.max_m, 2 * cfg.inter_dim, stream, + _swiglu_lim_rt, ) if _gemm_events is not None: _gemm_events[1].record(stream) @@ -1158,9 +1194,17 @@ def launch( if _skip_epilogue: return tmp if bias is not None: - _run_compiled(_get_finalize_act_bias(), y, tmp, bias, masked_m, stream) + _run_compiled( + _get_finalize_act_bias(), + y, + tmp, + bias, + masked_m, + _swiglu_lim_rt, + stream, + ) else: - _run_compiled(_get_finalize_act(), y, tmp, masked_m, stream) + _run_compiled(_get_finalize_act(), y, tmp, masked_m, _swiglu_lim_rt, stream) return y return launch @@ -1281,6 +1325,7 @@ def launch( if cfg.split_k > 1: y.zero_() gemm = _get_base_bias() if bias is not None else _get_base() + _no_act_swiglu_lim = float("inf") if cfg.grouped_persistent_m: m_tile_prefix = _m_tile_prefix if m_tile_prefix is None: @@ -1305,6 +1350,7 @@ def launch( cfg.max_m, cfg.model_dim, stream, + _no_act_swiglu_lim, ) else: _run_compiled( @@ -1320,6 +1366,7 @@ def launch( cfg.max_m, cfg.model_dim, stream, + _no_act_swiglu_lim, ) if _gemm_events is not None: _gemm_events[1].record(stream) @@ -1348,6 +1395,7 @@ def launch( contiguous_m, cfg.model_dim, stream, + _no_act_swiglu_lim, ) else: _run_compiled( @@ -1364,6 +1412,7 @@ def launch( contiguous_m, cfg.model_dim, stream, + _no_act_swiglu_lim, ) if _gemm_events is not None: _gemm_events[1].record(stream) @@ -1389,6 +1438,7 @@ def launch( cfg.max_m, cfg.model_dim, stream, + _no_act_swiglu_lim, ) else: _run_compiled( @@ -1405,6 +1455,7 @@ def launch( cfg.max_m, cfg.model_dim, stream, + _no_act_swiglu_lim, ) if _gemm_events is not None: _gemm_events[1].record(stream) diff --git a/aiter/ops/flydsl/kernels/silu_and_mul_fq.py b/aiter/ops/flydsl/kernels/silu_and_mul_fq.py index da9f792bc4..0619ae9481 100644 --- a/aiter/ops/flydsl/kernels/silu_and_mul_fq.py +++ b/aiter/ops/flydsl/kernels/silu_and_mul_fq.py @@ -49,7 +49,6 @@ def build_silu_and_mul_fq_module( gui_layout: bool = False, act: str = "silu", enable_bias: bool = False, - swiglu_limit: float = 0.0, ): """Return a JIT launcher for fused gate activation + optional quant + scale sort. @@ -124,6 +123,7 @@ def silu_and_mul_fq_kernel( topk_ids: fx.Pointer, bias: fx.Pointer, token_num: Int32, + swiglu_limit_f: fx.Float32, ): bid = fx.block_idx.x tid = fx.thread_idx.x @@ -277,12 +277,16 @@ def _load_bias_scalar(offset): swiglu_neg_alpha_log2e = arith.constant( -1.4426950408889634 * 1.702, type=f32 ) - if const_expr(swiglu_limit != 0): - _limit = arith.constant(float(swiglu_limit), type=f32) - _neg_limit = arith.constant(-float(swiglu_limit), type=f32) - else: - _limit = arith.constant(7.0, type=f32) - _neg_limit = arith.constant(-7.0, type=f32) + # ``swiglu_limit`` is a runtime f32 scalar. The host passes the + # clamp bound (7.0 default for swiglu) or +inf to disable the + # clamp (silu without a configured limit). ``min(x, lim)`` is + # expressed via the wrapped ``maximumf`` + negation so the kernel + # never bakes the limit as a compile-time constant. + _neg_limit = -swiglu_limit_f + + def _fmin(x): + # min(x, lim) == -max(-x, -lim) + return -((-x).maximumf(_neg_limit)) act_vals = [] for vi in range_constexpr(VEC): @@ -299,19 +303,13 @@ def _load_bias_scalar(offset): u = u + _load_bias_scalar( bias_row + inter_dim_i32 + bias_col ) - gate = g - linear = u - t = gate * neg_log2e + # gate: upper-clamped only; linear: clamped to [-lim, lim]. + gate = _fmin(g) + linear = _fmin(u).maximumf(_neg_limit) if const_expr(act == "swiglu"): - gate = arith.minimumf(gate, _limit) - linear = arith.minimumf(linear, _limit) - linear = arith.maximumf(linear, _neg_limit) - t = gate * swiglu_neg_alpha_log2e - elif const_expr(swiglu_limit != 0 and act != "swiglu"): - gate = arith.minimumf(gate, _limit) - linear = arith.minimumf(linear, _limit) - linear = arith.maximumf(linear, _neg_limit) t = gate * swiglu_neg_alpha_log2e + else: + t = gate * neg_log2e emu = llvm.call_intrinsic( f32, "llvm.amdgcn.exp2.f32", [t], [], [] @@ -556,6 +554,7 @@ def launch_silu_and_mul_fq( bias: fx.Pointer, token_num: fx.Int32, num_sorted_rows: fx.Int32, + swiglu_limit_f: fx.Float32, stream: fx.Stream = fx.Stream(None), ): ctx = CompilationContext.get_current() @@ -572,6 +571,7 @@ def launch_silu_and_mul_fq( topk_ids, bias, token_num, + swiglu_limit_f, ) launcher.launch( grid=(idx_rows, 1, 1), diff --git a/aiter/ops/flydsl/moe_kernels.py b/aiter/ops/flydsl/moe_kernels.py index 40a1c56ad4..6693a86ecc 100644 --- a/aiter/ops/flydsl/moe_kernels.py +++ b/aiter/ops/flydsl/moe_kernels.py @@ -367,7 +367,6 @@ def compile_flydsl_moe_stage1( enable_bias: bool = False, a_scale_one: bool = False, xcd_swizzle: int = 0, - swiglu_limit: float = 0.0, k_wave: int = 1, ): """Compile stage1 kernel (cached via underlying lru_cache).""" @@ -399,7 +398,6 @@ def compile_flydsl_moe_stage1( enable_bias=enable_bias, a_scale_one=a_scale_one, xcd_swizzle=xcd_swizzle, - swiglu_limit=swiglu_limit, k_wave=k_wave, ) elif a_dtype == "bf16" and b_dtype == "int4": @@ -541,6 +539,19 @@ def _ptr_view_safe(t: torch.Tensor): return flyc.from_c_void_p(fx.Uint8, view.data_ptr()) +def runtime_swiglu_limit(swiglu_limit: Optional[float], act: str) -> float: + """Normalize swiglu_limit into the runtime f32 clamp bound passed to kernels. + + The kernels always clamp using this value, so "no clamp" is encoded as +inf: + - swiglu: defaults to 7.0 when unset (matches the reference ``swiglu()``). + - silu: clamps only when a positive limit is configured, else +inf + (matches the reference's ``if swiglu_limit:`` truthiness). + """ + if act == "swiglu": + return float(swiglu_limit) if swiglu_limit else 7.0 + return float(swiglu_limit) if swiglu_limit else float("inf") + + def _s1_args_fp4( out, a, @@ -559,6 +570,7 @@ def _s1_args_fp4( dev, bias=None, stream=None, + swiglu_limit=float("inf"), ): empty_f32 = torch.empty(0, device=dev, dtype=torch.float32) _bias = bias if bias is not None else empty_f32 @@ -580,6 +592,7 @@ def _s1_args_fp4( n_in, k_in, size_expert_ids_in, + float(swiglu_limit), stream, ) @@ -1012,7 +1025,6 @@ def _get_compiled_silu_fused( gui_layout: bool = False, act: str = "silu", enable_bias: bool = False, - swiglu_limit: float = 0.0, ): """Compile and cache the fused gate activation + quant + scale-sort kernel.""" from aiter.ops.flydsl.kernels.silu_and_mul_fq import build_silu_and_mul_fq_module @@ -1024,7 +1036,6 @@ def _get_compiled_silu_fused( gui_layout, act=act, enable_bias=enable_bias, - swiglu_limit=swiglu_limit, ) @@ -1098,6 +1109,7 @@ def flydsl_silu_and_mul_interleaved( _ptr_view_safe(empty_f32), token_num, num_sorted_rows, + float("inf"), torch.cuda.current_stream(), ), ) @@ -1137,7 +1149,7 @@ def flydsl_moe_stage1( topk_ids: Optional[torch.Tensor] = None, a_scale_one: bool = False, xcd_swizzle: int = 0, - swiglu_limit: float = 0.0, + swiglu_limit: Optional[float] = None, k_wave: int = 1, ): """Fused gate+up GEMM (MOE stage1). @@ -1260,6 +1272,7 @@ def flydsl_moe_stage1( use_mx_gemm = b_dtype in ("fp4", "fp8") _n_in = inter_dim * 2 if use_mx_gemm else inter_dim _k_in = model_dim + _swiglu_limit_val = runtime_swiglu_limit(swiglu_limit, act) if use_mx_gemm: args = _s1_args_fp4( @@ -1283,6 +1296,7 @@ def flydsl_moe_stage1( if kernel_bias is not None else torch.empty(0, device=dev) ), + swiglu_limit=_swiglu_limit_val, ) else: args = _s1_args_std( @@ -1325,7 +1339,6 @@ def flydsl_moe_stage1( enable_bias=(kernel_bias is not None), a_scale_one=a_scale_one, xcd_swizzle=xcd_swizzle, - swiglu_limit=swiglu_limit, k_wave=k_wave, ) _run_compiled(exe, args) @@ -1359,7 +1372,6 @@ def flydsl_moe_stage1( gui_layout=True, act=act, enable_bias=use_splitk_bias, - swiglu_limit=swiglu_limit, ) _run_compiled( _silu_fused_k, @@ -1373,6 +1385,7 @@ def flydsl_moe_stage1( _ptr_view_safe(bias_arg), token_num, num_sorted_rows, + _swiglu_limit_val, torch.cuda.current_stream(), ), ) @@ -1384,7 +1397,6 @@ def flydsl_moe_stage1( gui_layout=True, act=act, enable_bias=use_splitk_bias, - swiglu_limit=swiglu_limit, ) _run_compiled( _silu_fused_k, @@ -1398,6 +1410,7 @@ def flydsl_moe_stage1( _ptr_view_safe(bias_arg), token_num, num_sorted_rows, + _swiglu_limit_val, torch.cuda.current_stream(), ), ) @@ -1407,7 +1420,6 @@ def flydsl_moe_stage1( topk, act=act, enable_bias=use_splitk_bias, - swiglu_limit=swiglu_limit, ) _run_compiled( _silu_fused_k, @@ -1421,6 +1433,7 @@ def flydsl_moe_stage1( _ptr_view_safe(bias_arg), token_num, num_sorted_rows, + _swiglu_limit_val, torch.cuda.current_stream(), ), ) diff --git a/op_tests/flydsl_tests/test_silu_and_mul_fq.py b/op_tests/flydsl_tests/test_silu_and_mul_fq.py index 5acdf6a5ef..423426e1a4 100644 --- a/op_tests/flydsl_tests/test_silu_and_mul_fq.py +++ b/op_tests/flydsl_tests/test_silu_and_mul_fq.py @@ -72,6 +72,7 @@ def test_flydsl_swiglu_fused_fp4_quant_matches_reference( torch.empty(0, dtype=torch.float32, device=device), token_num, sorted_ids.numel(), + 7.0, torch.cuda.current_stream(), ), ) diff --git a/op_tests/test_flydsl_grouped_gemm_gfx1250.py b/op_tests/test_flydsl_grouped_gemm_gfx1250.py index 47210c3689..465abdb07e 100644 --- a/op_tests/test_flydsl_grouped_gemm_gfx1250.py +++ b/op_tests/test_flydsl_grouped_gemm_gfx1250.py @@ -170,10 +170,10 @@ def _per_1x32_fp8_dequant(x: torch.Tensor) -> torch.Tensor: a1_scale=stage1_hidden_scale, w1_scale=w1_scale, w1_bias=w1_bias, - # torch_moe_stage1 also applies swiglu_limit as a generic gate/up - # clamp in the non-SwiGLU branch. The grouped FlyDSL SiLU epilogue - # does *not* clamp, so only pass the limit for true SwiGLU. - swiglu_limit=swiglu_limit if activation == ActivationType.Swiglu else 0.0, + # swiglu_limit clamps gate/up for both SwiGLU and SiLU: the grouped + # FlyDSL epilogue now applies it in either branch, so the reference + # passes it through unconditionally to stay in sync. + swiglu_limit=swiglu_limit, ) if data_format == "a4w4": # Match the grouped a4w4 path again: stage2 input is MXFP4. @@ -504,6 +504,12 @@ def test_grouped_a4w4_swiglu_matches_torch_ref(layout): run_moe("a4w4", layout=layout, activation=ActivationType.Swiglu) +@pytest.mark.parametrize("layout", ["gguu", "gugu"]) +@pytest.mark.parametrize("activation", [ActivationType.Silu, ActivationType.Swiglu]) +def test_grouped_a4w4_swiglu_limit_clamps(layout, activation): + run_moe("a4w4", layout=layout, activation=activation, swiglu_limit=1.0) + + # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- diff --git a/op_tests/test_moe_2stage.py b/op_tests/test_moe_2stage.py index dbfb496231..05bb659343 100644 --- a/op_tests/test_moe_2stage.py +++ b/op_tests/test_moe_2stage.py @@ -74,7 +74,7 @@ def test_fmoe( preshuffle=True, strict_accuracy=True, check_aot_cache=True, - swiglu_limit=0.0, + swiglu_limit=None, ): if get_gfx() not in ["gfx950"] and qType in [aiter.QuantType.per_1x32]: return @@ -586,8 +586,8 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor): "--swiglu-limit", "-sl", type=float, - default=0.0, - help="Limit the number of experts for swiglu activation type. Default is 0.0.", + default=None, + help="swiglu/silu clamp limit. Default None means the kernel default (7.0).", ) args = parser.parse_args() @@ -652,6 +652,9 @@ def _row_to_kwargs(row): hidden_pad=0, intermediate_pad=0, preshuffle=True, + swiglu_limit=_effective_swiglu_limit( + q_type, aq_dtype, wq_dtype, args.swiglu_limit + ), ) @@ -726,7 +729,7 @@ def _effective_gate_mode(aq_dtype, wq_dtype): def _effective_swiglu_limit(quant_type, aq_dtype, wq_dtype, swiglu_limit): if (quant_type, aq_dtype, wq_dtype) in (_PER1X32_BF16_FP4, _PER1X32_FP8_FP4): return swiglu_limit - return 0.0 + return None def _runtime_swiglu_mxfp4_q_dtype_a( @@ -781,6 +784,9 @@ def _kw( doweight_stage1=doweight_stage1, strict_accuracy=False, check_aot_cache=False, + swiglu_limit=_effective_swiglu_limit( + quant_type, aq_dtype, wq_dtype, args.swiglu_limit + ), **over, ) @@ -908,12 +914,6 @@ def _write_bench_csv(rows): seen = 0 for kwargs, extras in case_iter: seen += 1 - swiglu_limit = _effective_swiglu_limit( - kwargs["qType"], - kwargs["AQDType"], - kwargs["WQDType"], - args.swiglu_limit, - ) _old_moe_bound = os.environ.get("AITER_BF16_FP8_MOE_BOUND") _force_moe_bound_zero = ( kwargs["qType"], @@ -928,7 +928,7 @@ def _write_bench_csv(rows): if kwargs.get("check_aot_cache", False) else test_fmoe ) - ret = run_test_fmoe(**kwargs, swiglu_limit=swiglu_limit) + ret = run_test_fmoe(**kwargs) finally: if _force_moe_bound_zero: if _old_moe_bound is None: