[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. Preserve the shuffled activation-scale buffer physical allocation by rounding descriptor bounds to 32 rows. This adds no serving environment knob, dependency, or public API.

The commit is based directly on AITER main 45ef6f3 with FlyDSL 0.2.4. C4, C16, and C128+ rows are unchanged. The unrelated mixed-MoE XCD remap fix is intentionally excluded because both selected BM16 kernels use XCD0.

Kernel efficiency was remeasured after rebasing. Both sides use the same main-based source, compiler, full production routes, live non-unit activation scales, and full-size weights; only the upstream-main BM32 versus proposed BM16 configuration rows differ. Per-stage HIP event timing excludes sorting and quantization and uses 200 warmups, 1,000 iterations, and three repeats.

At C32, stage 1 improves from 93.5605 to 54.3600 us (-41.90%), stage 2 from 48.6800 to 27.7200 us (-43.06%), and combined stages from 142.2405 to 82.0800 us (-42.29%, 1.73x). At C64, stage 1 improves from 161.8010 to 86.8410 us (-46.33%), stage 2 from 87.9210 to 42.5210 us (-51.64%), and combined stages from 249.7220 to 129.3620 us (-48.20%, 1.93x). Every per-stage repeat range is below 0.3%.

Repeated stage outputs are bitwise stable and explicit two-stage output is bitwise identical to fused output. Exact route replay is bitwise identical to the upstream-main BM32 row at C32 and to the safe BM32/XCD0 control at C64. The upstream C64 row uses an unsafe non-divisible XCD4 mapping, while the selected BM16 route uses XCD0.

Full GSM8K validation of the same BM16 implementation completed all 1,319 samples: strict match 1258/1319 (95.375%) and flexible extraction 1257/1319 (95.299%), both above the required 94% threshold. No new end-to-end performance claim is made after the AITER-version rebase because it would not be an apples-to-apples comparison.

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