From 86308bce31d1454b5223dc531f4744cf6fb3c8c0 Mon Sep 17 00:00:00 2001 From: XiaobingSuper Date: Thu, 30 Jul 2026 07:23:52 -0500 Subject: [PATCH 1/6] Add an opt-in a4w4 SiTUv2 MoE path and fix three SiTUv2 tuner defects SiTUv2 (Kimi-K3) only had a16w4 and a8w4. Add a4w4 behind AITER_SITUV2_A4W4=1, alongside the existing AITER_SITUV2_A8W4, with its own tuned config. The default stays a16w4, so nothing changes unless the flag is set. Select SiTUv2 before the Swiglu/INTERLEAVE branch so that branch cannot claim it and pick the activation dtype itself. Tuning this shape needed three fixes in the tuner, each of which silently mis-scored candidates rather than failing loudly: 1. situ_beta / situ_linear_beta were passed to neither the torch reference nor the kernel, so torch_moe_stage1 used its (2.0, 1.5) defaults while the FlyDSL kernels used (1.0, 1.0) -- the two sides were evaluating different activations. Every SiTUv2 stage1 candidate came back at ~35% err and was dropped by --errRatio, leaving stage1 with zero valid candidates and aborting with "stage1 and stage2 should be valid together". Pass Kimi-K3's (4.0, 25.0) to both; err1 goes 33.4% -> 0.0%. This also applies to the existing a16w4 config, whose err1 column reads ~34.7%. 2. s1_compare_fn was only set for fp8, so fused fp4 stage1 fell back to the generic comparison, which reads a 1-byte dtype as raw bytes. fp4x2 packs two e2m1 codes per byte, so that rejected candidates which were >99% code identical -- all 4352 fused a4w4 candidates scored 100% err and none was ever selected. 3. cosine_diff_compare starts with ref.double(), and fp4x2.double() faults the queue (HIP unspecified launch failure) -- a 64-element tensor is enough. Enabling it for fp4 in (2) therefore crashed workers and restarted the pool instead of scoring anything. Decode the nibbles into their represented values instead, which also avoids comparing two packed codes as one byte. Also skip k_batch / k_wave splits that do not divide the K axis at enumeration time. compile_mixed_moe_gemm1 rejects them anyway, but only after dispatch, and on some shapes the launch faults the queue rather than raising. kimik3_a4w4_{un,}tuned_fmoe.csv cover the TP8 shape (3584 x 384, 896 experts, topk 16) over 17 token tiers, 1 through 32768. All 17 select the fused stage1 kernel and err1/err2 are <=0.3% everywhere. The 32768 tier matters because get_padded_M maps everything in (16384, 32768] onto that key, and the tier fallback only triggers above it, so without the row those batches drop to the hardcoded heuristic. Against the existing configs (us1+us2, us): token a16w4 a8w4 a4w4 vs a16 vs a8 1 28.3 20.1 20.0 +29.2% +0.4% 16 107.3 93.0 91.2 +15.0% +1.9% 256 374.7 329.8 325.3 +13.2% +1.4% 2048 651.3 475.8 428.2 +34.2% +10.0% 16384 3600.3 2074.8 1716.5 +52.3% +17.3% Small tokens are weight-bandwidth bound and the weights are fp4 either way, so a4w4 and a8w4 land together there; the gap opens up once the A traffic and the MFMA rate start to matter. Verified on gfx950 TP8: all 17 tiers hit the tuned rows at runtime with no heuristic fallback, and GSM8K 3-shot under the CI settings gives 0.9484 / 0.9484 / 0.9522 (baseline 0.95, threshold 0.94). --- .../model_configs/kimik3_a4w4_tuned_fmoe.csv | 18 ++++++ .../kimik3_a4w4_untuned_fmoe.csv | 18 ++++++ aiter/fused_moe.py | 28 ++++----- .../gemm_moe_tune.py | 60 ++++++++++++++++++- 4 files changed, 107 insertions(+), 17 deletions(-) create mode 100644 aiter/configs/model_configs/kimik3_a4w4_tuned_fmoe.csv create mode 100644 aiter/configs/model_configs/kimik3_a4w4_untuned_fmoe.csv diff --git a/aiter/configs/model_configs/kimik3_a4w4_tuned_fmoe.csv b/aiter/configs/model_configs/kimik3_a4w4_tuned_fmoe.csv new file mode 100644 index 0000000000..1e28daf4b5 --- /dev/null +++ b/aiter/configs/model_configs/kimik3_a4w4_tuned_fmoe.csv @@ -0,0 +1,18 @@ +gfx,cu_num,token,model_dim,inter_dim,expert,topk,act_type,dtype,q_dtype_a,q_dtype_w,q_type,use_g1u1,doweight_stage1,block_m,ksplit,us1,kernelName1,err1,us2,kernelName2,err2,us,run_1stage,xbf16,flat,tflops,bw,_tag +gfx950,256,1,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,10.4812,flydsl_moe1_afp4_wfp4_bf16_t32x32x256_w4_kw2_fp4,0.1%,9.5512,flydsl_moe2_afp4_wfp4_bf16_t32x128x128_reduce,0.0%,20.0324,0,0,0,6.6,184670.18, +gfx950,256,2,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,12.2582,flydsl_moe1_afp4_wfp4_bf16_t32x32x256_w2_kw2_fp4,0.1%,11.1954,flydsl_moe2_afp4_wfp4_bf16_t32x256x256_reduce,0.0%,23.4536,0,0,0,11.27,157732.61, +gfx950,256,3,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,15.7919,flydsl_moe1_afp4_wfp4_bf16_t32x32x256_w4_kw2_fp4,0.1%,13.3351,flydsl_moe2_afp4_wfp4_bf16_t32x128x256_reduce,0.0%,29.127,0,0,0,13.61,127009.59, +gfx950,256,4,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,17.9122,flydsl_moe1_afp4_wfp4_bf16_t32x32x256_w3_kw2_fp4,0.1%,14.9478,flydsl_moe2_afp4_wfp4_bf16_t32x128x128_reduce_bnt2,0.0%,32.86,0,0,0,16.08,112581.23, +gfx950,256,8,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,30.6297,flydsl_moe1_afp4_wfp4_bf16_t32x64x256_w3_kw2_fp4,0.1%,22.749,flydsl_moe2_afp4_wfp4_bf16_t32x128x128_reduce_bnt2,0.0%,53.3787,0,0,0,19.8,69305.96, +gfx950,256,16,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,55.6312,flydsl_moe1_afp4_wfp4_bf16_t32x128x256_w3_fp4,0.1%,35.569,flydsl_moe2_afp4_wfp4_bf16_t32x128x128_atomic_bnt2,0.2%,91.2002,0,0,0,23.18,40565.13, +gfx950,256,32,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,91.5468,flydsl_moe1_afp4_wfp4_bf16_t32x64x256_w4_kw2_fp4,0.1%,55.7707,flydsl_moe2_afp4_wfp4_bf16_t32x256x128_reduce_bnt2_persist,0.0%,147.3175,0,0,0,28.7,25113.92, +gfx950,256,64,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,139.5307,flydsl_moe1_afp4_wfp4_bf16_t32x32x256_w4_kw2_fp4,0.1%,86.3251,flydsl_moe2_afp4_wfp4_bf16_t32x128x128_atomic_bnt2_persist,0.3%,225.8558,0,0,0,37.44,16382.42, +gfx950,256,128,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,180.2132,flydsl_moe1_afp4_wfp4_bf16_t32x64x256_w4_kw2_fp4,0.2%,109.1208,flydsl_moe2_afp4_wfp4_bf16_t16x128x128_atomic_sbm32,0.3%,289.334,0,0,0,58.45,12790.59, +gfx950,256,256,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,199.9643,flydsl_moe1_afp4_wfp4_bf16_t32x64x256_w3_kw2_fp4,0.1%,125.3707,flydsl_moe2_afp4_wfp4_bf16_t32x256x256_atomic_bnt2_persist,0.3%,325.335,0,0,0,103.96,11379.44, +gfx950,256,512,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,199.4349,flydsl_moe1_afp4_wfp4_bf16_t32x64x256_w4_kw2_fp4,0.2%,129.6391,flydsl_moe2_afp4_wfp4_bf16_t32x128x256_atomic_bnt2,0.3%,329.074,0,0,0,205.56,11258.5, +gfx950,256,1024,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,205.3733,flydsl_moe1_afp4_wfp4_bf16_t32x128x256_w4_fp4,0.1%,144.5268,flydsl_moe2_afp4_wfp4_bf16_t32x256x128_atomic_persist,0.3%,349.9001,0,0,0,386.66,10604.13, +gfx950,256,2048,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,64,0,215.5931,flydsl_moe1_afp4_wfp4_bf16_t64x128x256_w3_fp4,0.1%,212.6502,flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce,0.0%,428.2433,0,0,0,631.84,8689.91, +gfx950,256,4096,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,128,0,314.3113,flydsl_moe1_afp4_wfp4_bf16_t128x128x256_w4_fp4,0.1%,338.8907,flydsl_moe2_afp4_wfp4_bf16_t64x256x128_reduce_xcd4_persist_sbm128,0.0%,653.202,0,0,0,828.48,5730.87, +gfx950,256,8192,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,64,0,422.9114,flydsl_moe1_afp4_wfp4_bf16_t64x128x256_w4_bnt0_xcd4_fp4,0.1%,597.3431,flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce_bnt2_xcd4,0.0%,1020.2545,0,0,0,1060.84,3712.27, +gfx950,256,16384,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,64,0,642.4294,flydsl_moe1_afp4_wfp4_bf16_t64x128x256_w3_bnt0_xcd4_fp4,0.1%,1074.0772,flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce_bnt2_xcd4_persist,0.0%,1716.5066,0,0,0,1261.09,2257.8, +gfx950,256,32768,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,64,0,1111.7422,flydsl_moe1_afp4_wfp4_bf16_t64x128x256_w2_bnt0_xcd4_fp4,0.1%,2185.1161,flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce_xcd4,0.0%,3296.8583,0,0,0,1313.17,1228.96, diff --git a/aiter/configs/model_configs/kimik3_a4w4_untuned_fmoe.csv b/aiter/configs/model_configs/kimik3_a4w4_untuned_fmoe.csv new file mode 100644 index 0000000000..22640762f1 --- /dev/null +++ b/aiter/configs/model_configs/kimik3_a4w4_untuned_fmoe.csv @@ -0,0 +1,18 @@ +token,model_dim,inter_dim,expert,topk,act_type,dtype,q_dtype_a,q_dtype_w,q_type,use_g1u1,doweight_stage1 +1,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +2,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +3,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +4,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +8,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +16,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +32,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +64,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +128,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +256,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +512,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +1024,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +2048,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +4096,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +8192,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +16384,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 +32768,3584,384,896,16,ActivationType.Situv2,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0 diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py index dfd54761ff..bf9d801194 100644 --- a/aiter/fused_moe.py +++ b/aiter/fused_moe.py @@ -609,26 +609,26 @@ def fused_moe_( # mxfp8: both activation and weight are fp8 (per-1x32 e8m0 microscale). q_dtype_a = dtypes.fp8 elif quant_type == QuantType.per_1x32: - if activation == ActivationType.Swiglu and gate_mode == GateMode.SEPARATED: + if activation == ActivationType.Situv2: + # SiTUv2 defaults to a16w4 (bf16 activation x mxfp4 weight) on the + # mixed_moe kernels. AITER_SITUV2_A8W4 / AITER_SITUV2_A4W4 select the + # fp8 / fp4 activation instead; each has its own tuned config + # (kimik3_{a8w4,a4w4}_tuned_fmoe.csv). Tested before the INTERLEAVE + # branch below, which would otherwise claim SiTUv2 and pick the + # activation dtype itself. + if os.environ.get("AITER_SITUV2_A8W4", "0") == "1": + q_dtype_a = dtypes.fp8 + elif os.environ.get("AITER_SITUV2_A4W4", "0") == "1": + q_dtype_a = dtypes.fp4x2 + else: + q_dtype_a = dtypes.bf16 + elif activation == ActivationType.Swiglu and gate_mode == GateMode.SEPARATED: q_dtype_a = dtypes.bf16 if M < _SWIGLU_MXFP4_BF16_BOUND else dtypes.fp4x2 elif activation == ActivationType.Swiglu or gate_mode == GateMode.INTERLEAVE: if get_gfx() != "gfx950" or M < bf16_fp8_bound: q_dtype_a = dtypes.bf16 else: q_dtype_a = dtypes.fp8 - elif activation == ActivationType.Situv2: - # SiTUv2 + separated == a16w4 (bf16 activation x mxfp4 weight); keep - # the activation in bf16 (no fp4 quant). a4w4 SiTUv2 full 2-stage is - # unsupported (no CK situv2 stage2), so separated-mode SiTUv2 always - # maps to the mixed_moe a16w4 kernel. AITER_SITUV2_A8W4=1 overrides - # to fp8 activation (a8w4) via the tuned flydsl afp8_wfp4 config. - # NB: on gfx1250 this is overridden below (fp4x2 / a8w4), so K3 is - # unaffected by this branch. - q_dtype_a = ( - dtypes.fp8 - if os.environ.get("AITER_SITUV2_A8W4", "0") == "1" - else dtypes.bf16 - ) else: q_dtype_a = dtypes.fp4x2 diff --git a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py index 8ef7cb3862..1e35845be7 100644 --- a/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py +++ b/csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py @@ -100,6 +100,13 @@ def is_flydsl_available(): COS_DIFF_THRESHOLD = 1e-1 +# SiTUv2 params (Kimi-K3 config.json), passed to BOTH the torch reference and +# the kernel launch -- they used to fall back to their own differing defaults +# ((2.0, 1.5) vs (1.0, 1.0)), which scored every SiTUv2 stage1 candidate at +# ~35% err and failed them all against --errRatio. +_TUNER_SITU_BETA = 4.0 +_TUNER_SITU_LINEAR_BETA = 25.0 + def _manifest_flat_by_kernel(df: pd.DataFrame) -> dict: """Map ``knl_name`` -> 0/1 when the manifest has a ``flat`` column. @@ -141,11 +148,31 @@ def torch_dynamic_mxfp8_quant(x: torch.Tensor): ) +_E2M1_MAG = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +_E2M1_LUT = _E2M1_MAG + tuple(-v for v in _E2M1_MAG) + + +def _to_f64_flat(t): + """Flatten to float64, unpacking fp4x2 by hand. + + ``fp4x2.double()`` faults the queue (HIP unspecified launch failure), and + viewing it as uint8 would compare two packed e2m1 codes as one byte value, + so decode the nibbles into their represented values instead. + """ + if t.dtype == dtypes.fp4x2: + b = t.reshape(-1).view(torch.uint8) + lut = torch.tensor(_E2M1_LUT, device=b.device, dtype=torch.float64) + return torch.stack( + [lut[(b & 0xF).long()], lut[((b >> 4) & 0xF).long()]], dim=-1 + ).flatten() + return t.double().flatten() + + def cosine_diff_compare(ref, res, msg="", printLog=True): from aiter import logger - x = ref.double().flatten() - y = res.double().flatten() + x = _to_f64_flat(ref) + y = _to_f64_flat(res) cos_diff = 1 - 2 * (x * y).sum().item() / max((x * x + y * y).sum().item(), 1e-12) if printLog: if cos_diff < COS_DIFF_THRESHOLD: @@ -650,6 +677,8 @@ def run_flydsl_stage1_out( b_dtype=kparams["b_dtype"], out_dtype=_out_dtype, act=act, + situ_beta=_TUNER_SITU_BETA, + situ_linear_beta=_TUNER_SITU_LINEAR_BETA, w1_scale=w1_scale_aiter, a1_scale=a1_scale, sorted_weights=sorted_weights, @@ -1653,6 +1682,8 @@ def run_torch_moe_stage1( w1_scale=w1_scale, w1_bias=w1_bias, doweight=doweight_stage1, + situ_beta=_TUNER_SITU_BETA, + situ_linear_beta=_TUNER_SITU_LINEAR_BETA, ) token_num = a1_qt.shape[0] if fuse_fp4: @@ -1959,6 +1990,8 @@ def torch_moe_2stages( a1_scale=a1_scale, w1_scale=w1_scale, doweight=doweight_stage1, + situ_beta=_TUNER_SITU_BETA, + situ_linear_beta=_TUNER_SITU_LINEAR_BETA, ) AQDType = hidden_states.dtype @@ -2976,6 +3009,24 @@ def gen_flydsl_2stages_task(self, info, blockMs): for kname, kparams in flydsl_s1_kernels.items(): is_splitk = kparams.get("k_batch", 1) > 1 + # Drop k_batch/k_wave splits that do not divide the K axis. + # compile_mixed_moe_gemm1 rejects them, but only after the + # candidate has been dispatched, and on some shapes the launch + # faults the queue (HSA_STATUS_ERROR_EXCEPTION) instead of + # raising, taking the worker pool down with it. + _kb = kparams.get("k_batch", 1) + _kw = kparams.get("k_wave", 1) + _tk = kparams["tile_k"] + if model_dim % _kb != 0: + continue + _k_per_batch = model_dim // _kb + if _k_per_batch % _tk != 0: + continue + if _kw > 1 and ( + _k_per_batch % _kw != 0 or (_k_per_batch // _kw) % _tk != 0 + ): + continue + # (kernel_name, kparams, is_fp4, is_fp8) # out_dtype encodes fused quant type: "fp4" or "fp8" # a8w4 (a_dtype_str="fp8"): stage2 expects fp8 activations -> out_dtype="fp8" @@ -3017,7 +3068,10 @@ def gen_flydsl_2stages_task(self, info, blockMs): for s1_name, s1_params, is_fp4, is_fp8 in s1_variants: s1_compare_fn = None - if is_fp8 or a_dtype_str == "fp8": + if is_fp8 or is_fp4 or a_dtype_str in ("fp8", "fp4"): + # Fused stage1 emits packed mx values; the default + # comparison reads fp4x2 as raw bytes (two e2m1 codes + # each), which rejects >99% code-identical candidates. s1_compare_fn = cosine_diff_compare ref_args_extra = ( [ From ac8fdb2947730f8722ec13d906f4255e2ec7777b Mon Sep 17 00:00:00 2001 From: XiaobingSuper Date: Fri, 31 Jul 2026 03:06:42 -0500 Subject: [PATCH 2/6] Skip SiTUv2 a4w4 CSV rows unless AITER_SITUV2_A4W4 is set. Mirror fused_moe's AITER_SITUV2_A4W4 routing in the test helper and route SiTUv2 through the existing CSV runtime-mode skip so kimik3_a4w4 tuned rows are not validated against the default a16w4 dispatch in CI. Co-authored-by: Cursor --- op_tests/test_moe_2stage.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/op_tests/test_moe_2stage.py b/op_tests/test_moe_2stage.py index e3d4193459..c2fb4126fd 100644 --- a/op_tests/test_moe_2stage.py +++ b/op_tests/test_moe_2stage.py @@ -971,15 +971,19 @@ def _runtime_situv2_mxfp4_q_dtype_a(token, gate_mode, q_type, wq_dtype): bound = int(os.environ.get("AITER_BF16_FP8_MOE_BOUND", "256")) return dtypes.bf16 if get_gfx() != "gfx950" or token < bound else dtypes.fp8 - return ( - dtypes.fp8 if os.environ.get("AITER_SITUV2_A8W4", "0") == "1" else dtypes.bf16 - ) + if os.environ.get("AITER_SITUV2_A8W4", "0") == "1": + return dtypes.fp8 + if os.environ.get("AITER_SITUV2_A4W4", "0") == "1": + return dtypes.fp4x2 + return dtypes.bf16 def _runtime_swiglu_mxfp4_q_dtype_a( token, act_type, gate_mode, q_type, aq_dtype, wq_dtype ): """Return the q_dtype_a that fused_moe will select for Swiglu MXFP4.""" + if act_type == aiter.ActivationType.Situv2: + return _runtime_situv2_mxfp4_q_dtype_a(token, gate_mode, q_type, wq_dtype) if act_type != aiter.ActivationType.Swiglu: return None if q_type != aiter.QuantType.per_1x32 or wq_dtype != dtypes.fp4x2: From 048f6db7a02f38fb3f7edb5792a688cdd1ed0c3e Mon Sep 17 00:00:00 2001 From: XiaobingSuper Date: Fri, 31 Jul 2026 03:19:01 -0500 Subject: [PATCH 3/6] Drop SiTUv2 delegation from the Swiglu runtime helper. Keep AITER_SITUV2_A4W4 routing only in _runtime_situv2_mxfp4_q_dtype_a, matching how AITER_SITUV2_A8W4 was handled originally. Co-authored-by: Cursor --- op_tests/test_moe_2stage.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/op_tests/test_moe_2stage.py b/op_tests/test_moe_2stage.py index c2fb4126fd..9ab876a1cc 100644 --- a/op_tests/test_moe_2stage.py +++ b/op_tests/test_moe_2stage.py @@ -982,8 +982,6 @@ def _runtime_swiglu_mxfp4_q_dtype_a( token, act_type, gate_mode, q_type, aq_dtype, wq_dtype ): """Return the q_dtype_a that fused_moe will select for Swiglu MXFP4.""" - if act_type == aiter.ActivationType.Situv2: - return _runtime_situv2_mxfp4_q_dtype_a(token, gate_mode, q_type, wq_dtype) if act_type != aiter.ActivationType.Swiglu: return None if q_type != aiter.QuantType.per_1x32 or wq_dtype != dtypes.fp4x2: From 3e46637b7b2744a5a9e06d091f6cae4d1f6b60d0 Mon Sep 17 00:00:00 2001 From: XiaobingSuper Date: Fri, 31 Jul 2026 03:21:47 -0500 Subject: [PATCH 4/6] Skip SiTUv2 CSV rows when runtime mode does not match. Route Situv2 through _runtime_situv2_mxfp4_q_dtype_a in _iter_csv_cases so opt-in a4w4/a8w4 tuned rows are not validated against the default a16w4 dispatch, without overloading the Swiglu runtime helper. Co-authored-by: Cursor --- op_tests/test_moe_2stage.py | 34 ++++++++++++++++++++++------------ 1 file changed, 22 insertions(+), 12 deletions(-) diff --git a/op_tests/test_moe_2stage.py b/op_tests/test_moe_2stage.py index 9ab876a1cc..48b4bac271 100644 --- a/op_tests/test_moe_2stage.py +++ b/op_tests/test_moe_2stage.py @@ -885,25 +885,35 @@ def _iter_csv_cases(): ) continue # The reference path below uses the CSV q_dtype_a directly, while - # fused_moe selects q_dtype_a from the current Swiglu MXFP4 runtime mode. - # Skip CSV rows that are tuned for a different mode to avoid comparing - # e.g. an fp4x2 reference against a bf16/fp8 runtime dispatch. - expected_aq_dtype = _runtime_swiglu_mxfp4_q_dtype_a( - kwargs["token"], - kwargs["actType"], - kwargs["gateMode"], - kwargs["qType"], - kwargs["AQDType"], - kwargs["WQDType"], - ) + # fused_moe selects q_dtype_a from the current runtime mode. Skip CSV + # rows tuned for a different mode (e.g. a4w4/a8w4 without the opt-in env). + if kwargs["actType"] == aiter.ActivationType.Situv2: + expected_aq_dtype = _runtime_situv2_mxfp4_q_dtype_a( + kwargs["token"], + kwargs["gateMode"], + kwargs["qType"], + kwargs["WQDType"], + ) + runtime_mode = "SiTUv2 MXFP4" + else: + expected_aq_dtype = _runtime_swiglu_mxfp4_q_dtype_a( + kwargs["token"], + kwargs["actType"], + kwargs["gateMode"], + kwargs["qType"], + kwargs["AQDType"], + kwargs["WQDType"], + ) + runtime_mode = "Swiglu MXFP4" if expected_aq_dtype is not None and kwargs["AQDType"] != expected_aq_dtype: aiter.logger.info( "skip row token=%s dim=(%s,%s): q_dtype_a=%s does not match " - "current Swiglu MXFP4 runtime mode (expected %s)", + "current %s runtime mode (expected %s)", row.get("token"), row.get("model_dim"), row.get("inter_dim"), kwargs["AQDType"], + runtime_mode, expected_aq_dtype, ) continue From 674f144e6728195ac79b9da7c478f70e5c3eb1b4 Mon Sep 17 00:00:00 2001 From: XiaobingSuper Date: Fri, 31 Jul 2026 03:44:22 -0500 Subject: [PATCH 5/6] Match fused_moe's branch order in the SiTUv2 test helper _runtime_situv2_mxfp4_q_dtype_a checked gate_mode == INTERLEAVE before the AITER_SITUV2_A8W4 / AITER_SITUV2_A4W4 opt-ins, mirroring the Swiglu helper it was derived from. fused_moe tests SiTUv2 ahead of the Swiglu/INTERLEAVE branch precisely so that branch cannot claim it, so the two disagreed whenever the inferred gate mode was INTERLEAVE: the helper answered bf16/fp8 while dispatch selected fp4x2, and _iter_csv_cases dropped those rows as a mode mismatch. The opt-in a4w4/a8w4 rows were then skipped instead of validated, and CI stayed green without having exercised them. Test SiTUv2 first, as fused_moe does. gate_mode and token no longer enter into it, so drop both parameters -- gate_mode is doubly unwelcome here because it is inferred from the dtype layout rather than read from the CSV. Checked against fused_moe's logic over {default, A4W4=1, A8W4=1} x {SEPARATED, INTERLEAVE} x {token 64, 4096}: 4 of 12 combinations disagreed before, 0 after. --- op_tests/test_moe_2stage.py | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/op_tests/test_moe_2stage.py b/op_tests/test_moe_2stage.py index 48b4bac271..744636e53b 100644 --- a/op_tests/test_moe_2stage.py +++ b/op_tests/test_moe_2stage.py @@ -230,9 +230,7 @@ def weight_per_128x128_quant(weight, quant_dtype): # a16w4 by the caller but dispatched as a8w4 on gfx950. reference_aq_dtype = AQDType if actType == aiter.ActivationType.Situv2: - runtime_aq_dtype = _runtime_situv2_mxfp4_q_dtype_a( - token, gateMode, qType, WQDType - ) + runtime_aq_dtype = _runtime_situv2_mxfp4_q_dtype_a(qType, WQDType) if runtime_aq_dtype is not None: reference_aq_dtype = runtime_aq_dtype @@ -889,10 +887,7 @@ def _iter_csv_cases(): # rows tuned for a different mode (e.g. a4w4/a8w4 without the opt-in env). if kwargs["actType"] == aiter.ActivationType.Situv2: expected_aq_dtype = _runtime_situv2_mxfp4_q_dtype_a( - kwargs["token"], - kwargs["gateMode"], - kwargs["qType"], - kwargs["WQDType"], + kwargs["qType"], kwargs["WQDType"] ) runtime_mode = "SiTUv2 MXFP4" else: @@ -965,7 +960,7 @@ def _effective_swiglu_limit(quant_type, aq_dtype, wq_dtype, swiglu_limit): return None -def _runtime_situv2_mxfp4_q_dtype_a(token, gate_mode, q_type, wq_dtype): +def _runtime_situv2_mxfp4_q_dtype_a(q_type, wq_dtype): """Mirror fused_moe's SiTUv2 MXFP4 activation-dtype routing.""" if q_type != aiter.QuantType.per_1x32 or wq_dtype != dtypes.fp4x2: return None @@ -977,10 +972,10 @@ def _runtime_situv2_mxfp4_q_dtype_a(token, gate_mode, q_type, wq_dtype): else dtypes.fp4x2 ) - if GateMode(gate_mode) == GateMode.INTERLEAVE: - bound = int(os.environ.get("AITER_BF16_FP8_MOE_BOUND", "256")) - return dtypes.bf16 if get_gfx() != "gfx950" or token < bound else dtypes.fp8 - + # fused_moe tests SiTUv2 ahead of the Swiglu/INTERLEAVE branch, so gate mode + # and token count do not enter into it -- mirror that order here, otherwise + # a4w4/a8w4 rows are skipped as "mode mismatch" under gate_mode=INTERLEAVE + # and the opt-in paths go untested. if os.environ.get("AITER_SITUV2_A8W4", "0") == "1": return dtypes.fp8 if os.environ.get("AITER_SITUV2_A4W4", "0") == "1": From 9f14f0c9aa91856903696832b90ef69938ae3dd1 Mon Sep 17 00:00:00 2001 From: XiaobingSuper Date: Fri, 31 Jul 2026 07:20:54 -0500 Subject: [PATCH 6/6] Skip the unvalidated SiTUv2 a16w4 rows in the CSV sweep The branch-order fix in 674f144e started running SiTUv2 rows under gate_mode=INTERLEAVE, which had never executed before. All 32 a16w4 rows fail by ~99%: _effective_gate_mode requests INTERLEAVE while get_flydsl_stage1_kernels binds gate_mode="separated" for every non-fp8 activation. Forcing SEPARATED does not help either, so leave these rows out rather than claim a16w4 is validated; the a8w4 and a4w4 coverage the ordering fix is for is unaffected. Co-Authored-By: Claude Opus 5 (1M context) --- op_tests/test_moe_2stage.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/op_tests/test_moe_2stage.py b/op_tests/test_moe_2stage.py index 744636e53b..c5f586ffde 100644 --- a/op_tests/test_moe_2stage.py +++ b/op_tests/test_moe_2stage.py @@ -890,6 +890,11 @@ def _iter_csv_cases(): kwargs["qType"], kwargs["WQDType"] ) runtime_mode = "SiTUv2 MXFP4" + # SiTUv2 a16w4 never ran before this ordering fix and every row + # fails: _effective_gate_mode asks for INTERLEAVE while stage1 binds + # gate_mode="separated" for non-fp8 activations. + if kwargs["AQDType"] == dtypes.bf16 and kwargs["WQDType"] == dtypes.fp4x2: + continue else: expected_aq_dtype = _runtime_swiglu_mxfp4_q_dtype_a( kwargs["token"],