[FlyDSL] Support SiTUv2 in the packed-int4 MoE stage1 epilogue

[maeehart]

Motivation

The packed-int4 (a16wi4) FlyDSL stage1 path ignores the requested activation and always computes SiLU.

compile_flydsl_moe_stage1 forwards act, situ_beta and situ_linear_beta to every dispatch branch except b_dtype == "int4". compile_moe_gemm1 has no act parameter, so y = silu(vg) * vu is hardcoded at both live stage1 epilogue sites, the CShuffle one and the direct-store one. Nothing raises. The kernel compiles, runs, and returns fluent text while computing the wrong function.

This affects Kimi-K3 on gfx942, which specifies hidden_act = situ with activation_situ_beta = 4.0 and activation_situ_linear_beta = 25.0. gfx942 has no scaled MXFP4 MFMA, so the MXFP4 experts are requantized to int4 at load and dispatched to this kernel. gfx950 is unaffected, since there bf16 x MXFP4 routes to compile_mixed_moe_gemm1_a16w4, which implements SiTUv2 natively.

Reproducing needs no model. Compile the same shape twice through compile_flydsl_moe_stage1, once with act="silu" and once with act="situv2". The int4 path returns the identical executable for both. The mxfp4 path returns two different ones. Against a SiTUv2 reference the current kernel scores cosine 0.961917, and against a SiLU reference it scores 0.999976.

Technical Details

• Add act, situ_beta and situ_linear_beta to compile_moe_gemm1. Defaults keep every existing caller bit-identical.
• Split the existing silu() into a reusable sigmoid(), add the SiTUv2 form, and route both epilogue sites through a single apply_activation().
• Forward the three values from the int4 branch of compile_flydsl_moe_stage1.
• Include the activation and the exact beta pair in the FlyDSL module name. The lru_cache key already covered the new parameters but the on-disk module name did not, so SiLU and SiTUv2 binaries could collide in the kernel cache.
• Reject activations other than silu and situv2 on this path instead of silently computing SiLU.

tanh is built as 2 * sigmoid(2x) - 1 over the existing exp2 and rcp fast path rather than rocdl.tanh, which lowers to llvm.amdgcn.tanh and only selects on gfx1250. The identity saturates correctly at both ends without clamping, which matters because tanh(up / 25) saturates for any |up| above roughly 100.

situv2 with split-K raises rather than falling back. The split-K path writes raw gate and up partials and defers activation to a reduction kernel that only implements SiLU. This is inert at current Kimi-K3 shapes because get_ksplit returns 0, and it closes the same class of failure against future tuning.

Test Plan

• pytest op_tests/flydsl_tests/test_flydsl_moe_a16wi4.py
• pytest op_tests/flydsl_tests/test_flydsl_moe_module_name.py
• python3 op_tests/test_moe_2stage.py -d bf16 -q 6 -t 1 8 32 -a situv2 --beta 4.0 --linear-beta 25.0
• python3 op_tests/test_moe_2stage.py -d bf16 -q 6 -t 1 8 32 -a silu for the unchanged path

Test Result

Direct SiTUv2 stage1 passes on gfx942. SiLU output is unchanged bit-for-bit for default callers. End to end on 8x MI325X the fix is throughput neutral at -0.07 percent, which is inside run-to-run noise.

When testing SiTUv2 by hand, pass the same betas to the kernel and the reference. torch_moe_stage1 defaults to (2.0, 1.5) and the FlyDSL kernels default to (1.0, 1.0), so a reference built on the defaults compares two different activation functions and reports a large error unrelated to the kernel. #4463 fixes that separately.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4471 --add-label <label>

[Copilot]

Pull request overview

This PR fixes incorrect activation behavior in the FlyDSL packed-int4 (a16wi4) MoE stage1 epilogue by threading activation selection (SiLU vs SiTUv2) and SiTUv2 beta parameters through the int4 stage1 compile path, and by ensuring the compiled module name uniquely encodes activation + betas to prevent kernel-cache collisions.

Changes:

• Add act, situ_beta, and situ_linear_beta to compile_moe_gemm1, implement SiTUv2 epilogue math, and route both epilogue sites through a shared activation helper.
• Forward activation + beta parameters through the int4 branch of compile_flydsl_moe_stage1, and update the FlyDSL kernel module naming to include activation/betas (preventing cache collisions).
• Expand/adjust tests to cover SiTUv2 routing for packed-int4 and to regression-test module-name discrimination.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
op_tests/test_moe_2stage.py	Updates per_1x32 gfx gating and expands SiTUv2 coverage to include packed-int4 cases; adjusts SiTUv2 default sweep beta handling.
op_tests/flydsl_tests/test_flydsl_moe_module_name.py	Adds regression test ensuring stage1 module naming distinguishes activation and beta pairs and is wired into the kernel decorator.
op_tests/flydsl_tests/test_flydsl_moe_a16wi4.py	Adds a focused gfx942 correctness test comparing packed-int4 SiTUv2 stage1 against the torch reference.
csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py	Updates per_1x32 gating to allow packed-int4 weights on non-gfx950.
aiter/ops/flydsl/moe_kernels.py	Forwards act/beta parameters through the a16wi4 (bf16×int4) stage1 compile path.
aiter/ops/flydsl/kernels/moe_gemm_2stage.py	Implements activation selection (SiLU/SiTUv2), adds validation, and encodes activation/betas into the FlyDSL module name.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[maeehart]

Ran the tests on gfx942 (MI325X), building this branch rather than patching an installed AITER.

op_tests/flydsl_tests/test_flydsl_moe_module_name.py   1 passed
op_tests/flydsl_tests/test_flydsl_moe_a16wi4.py        1 passed in 5.45s
same test with Kimi-K3 betas 4.0 / 25.0                1 passed in 5.40s

The a16wi4 test compares the packed-int4 stage1 against torch_moe_stage1 with ActivationType.Situv2 and the same betas on both sides, at atol=0.2, rtol=0.1. It passes at the shipped betas and at K3's production 4.0 and 25.0.

Activation now reaches the kernel and the cache keys separate:

tag silu          : _silu
tag situv2 b4/l25 : _situv2_sb0x1d0000000000000pp2_slb0x1d9000000000000pp4
tag situv2 b2/l25 : _situv2_sb0x1d0000000000000pp1_slb0x1d9000000000000pp4

Three distinct module names, so SiLU and SiTUv2 cannot collide in the on-disk FlyDSL cache and neither can two beta pairs. Requesting an unsupported activation raises ValueError: act must be 'silu' or 'situv2', got 'swiglu' instead of falling through to SiLU.

One note for anyone testing this against an existing install. The patched moe_gemm_2stage.py imports buffer_ops and vector from aiter.ops.flydsl.kernels, which do not exist in v0.1.19, so copying the two changed files over a v0.1.19 install fails at import. Build the branch instead.