[Issue]: [gfx950][fused MoE] Kimi-K3 A8W4 SiTUv2 token=8192 intermittently returns NaN

[fsx950223]

Summary

On current main, the Kimi-K3 A8W4 SiTUv2 fused-MoE shape below can return NaN when stage1/stage2 biases are disabled. The issue reproduces with the production-style 5 warmup + 20 dispatch sequence; a token=2048 control remains finite.

This report covers the full two-stage output. It does not establish GEMM1 as the source: the failing token=8192 row uses a reduce_persist stage2 kernel, while the finite token=2048 row uses atomic stage2.

Environment

• AITER commit: 4045b2de7b2dfa6f9e5567ddb9b6b75340c5b119 (current origin/main when tested)
• GPU arch: gfx950:sramecc+:xnack- (MI355X environment)
• Python: 3.12.3
• PyTorch: 2.10.0+rocm7.2.4.git3d3aa833
• HIP runtime reported by PyTorch: 7.2.53211

Reproduction

In op_tests/test_moe_2stage.py, mirror the bias-disabled serving/report setup by assigning:

exp_bias1 = None
exp_bias2 = None

after those tensors are created. Also override run_perftest before the case loop so the final result is taken after the same 5 warmups + 20 measured dispatches:

def run_perftest(func, *args, **kwargs):
    for key in (
        "num_iters",
        "num_warmup",
        "testGraph",
        "num_rotate_args",
        "needTrace",
        "use_cuda_event",
        "return_kernel_times",
    ):
        kwargs.pop(key, None)
    for _ in range(5):
        result = func(*args, **kwargs)
    torch.cuda.synchronize()
    for _ in range(20):
        result = func(*args, **kwargs)
    torch.cuda.synchronize()
    return result, 1.0

Run from the repository root:

AITER_SITUV2_A8W4=1 \
AITER_CONFIG_FMOE=aiter/configs/model_configs/kimik3_a8w4_tuned_fmoe.csv \
AITER_MOE_EXPERT_BALANCE=true \
FLYDSL_RUNTIME_ENABLE_CACHE=1 \
python3 op_tests/test_moe_2stage.py \
  -q 7 -a situv2 -dim 3584,512 -e 896 -k 16 -t 8192 \
  -hip 192,128 --beta 2.0 --linear-beta 1.5 --no-flydsl-csv

The legacy case is non-strict, so the process may exit 0 despite logging the NaN.

Actual result

[fused_moe] using 2stage
  kernelName1='flydsl_moe1_afp8_wfp4_bf16_t64x128x256_w4_bnt0_gui_fp8'
  kernelName2='flydsl_moe2_afp8_wfp4_bf16_t64x128x256_reduce_persist'
output contains NaN! (possible aiter #3117 stage2 K-pad regression)
accuracy check failed (non-strict): output contains NaN
logits_diff: nan

The same bias-disabled harness with token=2048 produced finite output and dispatched:

kernelName1='flydsl_moe1_afp8_wfp4_bf16_t64x128x256_w3_gui'
kernelName2='flydsl_moe2_afp8_wfp4_bf16_t64x128x256_atomic'

Expected result

The fused-MoE output should remain finite for token=8192, including with both biases disabled and model/intermediate padding of 192/128.