[Perf] Tune DeepSeek V4 fused MoE for C32/C64 decode

Enable 16-row FP8-activation/FP4-weight FlyDSL kernels for the gfx950 DeepSeek V4 7168/512/E385/top-k7 sparse-decode shape, selecting them only for the C32 and C64 production rows. Round activation-scale resource bounds to 32 rows so the BM16 path preserves the shuffled scale-buffer allocation.

This standalone patch is based on AITER main 97bde6f. C4, C16, and C128+ rows remain unchanged. The separate XCD remainder fix is intentionally excluded because both selected BM16 kernels use XCD0.

Isolated stage timing improved from 142.2405 to 82.0800 us at C32 (1.73x) and from 249.7220 to 129.3620 us at C64 (1.93x). Repeated outputs were bitwise stable, and GSM8K completed 1,319 samples with 95.375% strict match and 95.299% flexible extraction.

AI assistance was used to implement, test, benchmark, and document this change.

Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: Fangzhou Ai <fangzhouai@gmail.com>

Author: Fangzhou-Ai

aiter/configs/model_configs/dsv4_fp8fp4_tuned_fmoe.csv

@@ -44,8 +44,8 @@ gfx950,256,2,7168,512,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3
 gfx950,256,4,7168,512,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,21.9888,flydsl_moe1_afp8_wfp4_bf16_t32x128x256_gui_fp8,0.0%,17.2479,flydsl_moe2_afp8_wfp4_bf16_t32x256x256_atomic_bnt2_xcd4,0.0%,39.2367,0,0,0,15.71,108035.45,
 gfx950,256,8,7168,512,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,37.0281,flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w2_gui_fp8,0.0%,25.8309,flydsl_moe2_afp8_wfp4_bf16_t32x128x256_atomic_bnt2_persist,0.0%,62.859,0,0,0,19.62,67437.29,
 gfx950,256,16,7168,512,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,71.2231,flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w2_gui_fp8,0.0%,41.8664,flydsl_moe2_afp8_wfp4_bf16_t32x128x256_atomic_bnt2,0.0%,113.0895,0,0,0,21.81,37485.47,
-gfx950,256,32,7168,512,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,134.8852,flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w3_gui_xcd4_fp8,0.0%,71.8893,flydsl_moe2_afp8_wfp4_bf16_t32x256x256_atomic_bnt2_xcd4_persist,0.0%,206.7745,0,0,0,23.85,20503.29,
-gfx950,256,64,7168,512,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,232.2826,flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w3_gui_xcd4_fp8,0.0%,119.2028,flydsl_moe2_afp8_wfp4_bf16_t32x256x256_atomic_bnt2_xcd4,0.0%,351.4854,0,0,0,28.07,12063.79,
+gfx950,256,32,7168,512,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,16,0,54.3600,flydsl_moe1_afp8_wfp4_bf16_t16x128x256_w3_gui_fp8,0.0%,27.7200,flydsl_moe2_afp8_wfp4_bf16_t16x256x256_atomic,0.0%,82.0800,0,0,0,60.09,51651.52,
+gfx950,256,64,7168,512,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,16,0,86.8410,flydsl_moe1_afp8_wfp4_bf16_t16x128x256_w3_gui_fp8,0.0%,42.5210,flydsl_moe2_afp8_wfp4_bf16_t16x256x256_atomic,0.0%,129.3620,0,0,0,76.26,32778.13,
 gfx950,256,128,7168,512,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,233.1225,flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w3_gui_fp8,0.0%,121.6478,flydsl_moe2_afp8_wfp4_bf16_t32x256x256_atomic_bnt2_xcd4,0.0%,354.7703,0,0,0,55.61,11955.96,
 gfx950,256,256,7168,512,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,233.9255,flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w4_gui_fp8,0.0%,127.8476,flydsl_moe2_afp8_wfp4_bf16_t32x256x256_atomic_bnt2_xcd4,0.0%,361.7731,0,0,0,109.07,11732.14,
 gfx950,256,512,7168,512,385,7,ActivationType.Silu,torch.bfloat16,torch.float8_e4m3fn,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,236.6982,flydsl_moe1_afp8_wfp4_bf16_t32x128x256_w3_gui_fp8,0.0%,136.2029,flydsl_moe2_afp8_wfp4_bf16_t32x256x256_atomic_bnt2_persist,0.0%,372.9011,0,0,0,211.64,11396.8,

aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py

@@ -112,7 +112,7 @@ def compile_mixed_moe_gemm1(
     is_f4_b = b_dtype == "fp4"
     is_f8_b = b_dtype == "fp8"
 
-    sort_block_m = max(32, tile_m)
+    sort_block_m = tile_m
     num_waves = min(4, tile_n // 32)
     # accumulators are reduced in LDS before the epilogue. k_wave=1 keeps the
     num_n_waves = num_waves
@@ -559,7 +559,12 @@ def ptr_buffer_resource(ptr, num_records_bytes):
             if const_expr(not a_scale_one):
                 c32 = arith.constant(32, index=True)
                 kblk = k_in // c32
-                sx_nbytes_idx = sorted_m * kblk
+                scale_rows = (
+                    (sorted_m + arith.constant(31, index=True))
+                    // arith.constant(32, index=True)
+                    * arith.constant(32, index=True)
+                )
+                sx_nbytes_idx = scale_rows * kblk
                 sx_nbytes_i32 = arith.index_cast(T.i32, sx_nbytes_idx)
                 sx_rsrc = ptr_buffer_resource(arg_scale_x, sx_nbytes_i32)
 
@@ -3296,7 +3301,12 @@ def check_c_k_valid_gate(base_k):
             if const_expr(is_f4_a or is_f8_a):
                 # #3476: use 256-padded K/32 to match host scale padding.
                 kblk = arith.constant(scale_kblk_padded, index=True)
-                sx_nbytes_idx = num_valid_idx * kblk
+                scale_rows = (
+                    (num_valid_idx + arith.constant(31, index=True))
+                    // arith.constant(32, index=True)
+                    * arith.constant(32, index=True)
+                )
+                sx_nbytes_idx = scale_rows * kblk
                 sx_nbytes_i32 = arith.index_cast(T.i32, sx_nbytes_idx)
                 sx_rsrc = ptr_buffer_resource(arg_scale_x, sx_nbytes_i32)
             else:

aiter/ops/flydsl/moe_kernels.py

@@ -82,7 +82,7 @@ def get_flydsl_stage1_kernels(
 
     tile_ns = [32, 64, 128] if is_fp4_b else [128]
     tile_ks = [256]
-    tile_ms = [32, 64, 128]
+    tile_ms = [16, 32, 64, 128] if a_dtype == "fp8" and is_fp4_b else [32, 64, 128]
 
     waves_per_eus = [1, 2, 3, 4]
     k_batches = [1, 2, 4, 7, 14]

op_tests/test_moe_2stage.py

@@ -480,15 +480,18 @@ def weight_per_128x128_quant(weight, quant_dtype):
             "us_stage2": us2_stage,
         }
 
+    test_graph = args.bm16_scale_boundary
     out2_ck, us2 = run_perftest(
         fused_moe,
         input,
         w1_qt_aiter,
         w2_qt_aiter,
         topk_weights,
         topk_ids,
-        num_iters=5,
-        num_warmup=2,
+        num_iters=2 if test_graph else 5,
+        num_warmup=1 if test_graph else 2,
+        num_rotate_args=1 if test_graph else 0,
+        testGraph=test_graph,
         **_fused_moe_kwargs,
     )
     # Regression guard for aiter #3117 (MXFP4 fused-MoE stage2 EP-prefill):
@@ -688,6 +691,11 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor):
     action="store_true",
     help="Skip the original hardcoded shape sweep and skinny tests.",
 )
+parser.add_argument(
+    "--bm16-scale-boundary",
+    action="store_true",
+    help="Run only the deterministic BM16 tiled-scale boundary regression.",
+)
 parser.add_argument(
     "--swiglu-limit",
     "-sl",
@@ -1030,14 +1038,89 @@ def _kw(
                     ), extras
 
 
+def test_bm16_tiled_scale_boundary():
+    """Validate tuned BM16 dispatch and the 33-row scale boundary."""
+    if get_gfx() != "gfx950":
+        aiter.logger.info("skip BM16 tiled-scale boundary test on %s", get_gfx())
+        return
+
+    from aiter.ops.flydsl.moe_kernels import get_flydsl_kernel_params
+
+    expected_kernels = (
+        "flydsl_moe1_afp8_wfp4_bf16_t16x128x256_w3_gui_fp8",
+        "flydsl_moe2_afp8_wfp4_bf16_t16x256x256_atomic",
+    )
+    for token in (32, 64):
+        metadata = get_2stage_cfgs(
+            get_padded_M(token),
+            7168,
+            512,
+            385,
+            7,
+            torch.bfloat16,
+            dtypes.fp8,
+            dtypes.fp4x2,
+            aiter.QuantType.per_1x32,
+            True,
+            aiter.ActivationType.Silu,
+            False,
+            0,
+            0,
+            True,
+            GateMode.INTERLEAVE.value,
+        )
+        assert metadata.block_m == 16
+        stages = (metadata.stage1, metadata.stage2)
+        kernel_names = tuple(stage.keywords["kernelName"] for stage in stages)
+        assert kernel_names == expected_kernels
+        stage1_params = get_flydsl_kernel_params(kernel_names[0])
+        stage2_params = get_flydsl_kernel_params(kernel_names[1])
+        assert stage1_params is not None
+        assert stage2_params is not None
+        assert stage1_params["tile_m"] == 16
+        assert stage1_params["waves_per_eu"] == 3
+        assert stage2_params["tile_m"] == 16
+        assert stage2_params["tile_n"] == 256
+        assert stage2_params.get("xcd_swizzle", 0) == 0
+
+    old_moe_bound = os.environ.get("AITER_BF16_FP8_MOE_BOUND")
+    os.environ["AITER_BF16_FP8_MOE_BOUND"] = "0"
+    torch.manual_seed(0)
+    try:
+        test_fmoe(
+            dtype=torch.bfloat16,
+            token=33,
+            model_dim=7168,
+            inter_dim=512,
+            E=385,
+            topk=7,
+            actType=aiter.ActivationType.Silu,
+            gateMode=GateMode.INTERLEAVE.value,
+            qType=aiter.QuantType.per_1x32,
+            AQDType=dtypes.fp8,
+            WQDType=dtypes.fp4x2,
+            use_g1u1=True,
+            strict_accuracy=True,
+            check_aot_cache=False,
+        )
+    finally:
+        if old_moe_bound is None:
+            os.environ.pop("AITER_BF16_FP8_MOE_BOUND", None)
+        else:
+            os.environ["AITER_BF16_FP8_MOE_BOUND"] = old_moe_bound
+
+
 # ---------------------------------------------------------------------------
 # Run
 # ---------------------------------------------------------------------------
 _case_iters = []
-if not args.no_flydsl_csv:
-    _case_iters.append(_iter_csv_cases())
-if not args.no_legacy:
-    _case_iters.append(_iter_legacy_cases())
+if args.bm16_scale_boundary:
+    test_bm16_tiled_scale_boundary()
+else:
+    if not args.no_flydsl_csv:
+        _case_iters.append(_iter_csv_cases())
+    if not args.no_legacy:
+        _case_iters.append(_iter_legacy_cases())
 case_iter = itertools.chain(*_case_iters)
 
 _csv_out = os.environ.get("AITER_TUNED_OP_BENCH_CSV", "tuned_op_bench.csv")