[gfx1250][flydsl]moe group gemm swiglu limit for dsv4

.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",

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,
                     ),
                 )

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
+32768,7168,1536,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0

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

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,

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(