[Bugfix] Support parameterized SwiGLU in SM90 FP8 MegaMoE

[shiyang814-cpu]

Summary

This PR adds parameterized SwiGLU support to the SM90 FP8 MegaMoE kernel.

The existing implementation assumes standard SwiGLU:

SiLU(gate) * up

Some MoE models use parameterized OAI-SwiGLU instead:

clamp(gate, max=limit)
* sigmoid(alpha * clamp(gate, max=limit))
* (clamp(up, -limit, limit) + up_bias)

For example, MiniMax-M3 uses:

alpha = 1.702
up_bias = 1.0
clamp = 7.0

Without this change, the SM90 MegaMoE kernel silently applies the wrong activation semantics and produces corrupted model outputs.

Changes

• Add activation_alpha and activation_up_bias to the Python FP8 MegaMoE API.
• Propagate both parameters through the TVM FFI and SM90 host API.
• Include the activation parameters in the JIT kernel instantiation and cache key.
• Apply the same parameterized SwiGLU semantics in both:
  • swapAB epilogue
  • non-swapAB epilogue
• Accept swigluoai in the SM90 MegaMoE buffer and execution APIs.
• Preserve backward compatibility with:
  • activation_alpha=1.0
  • activation_up_bias=0.0
• Extend the existing SM90 MegaMoE accuracy suite with an OAI-SwiGLU scenario.

Root Cause

The SM90 FP8 MegaMoE epilogues hardcoded standard SiLU:

gate * sigmoid(gate) * up

Although an activation clamp was already configurable, the sigmoid alpha and up-projection bias were not represented by the public API or JIT template.

As a result, enabling MegaMoE for models using OAI-SwiGLU changed the model's mathematical semantics. The issue affected both epilogue implementations, including the optimized swapAB path.

Validation

The test was first run against an unmodified build from the latest dev branch and failed with:

activation == "swiglu"

After the fix, the SM90 MegaMoE output was compared against the PyTorch FP32/BF16 reference implementation on 8 Hopper GPUs.

Results:

Default swapAB + OAI-SwiGLU:
calc_diff = 0.0000

Forced non-swapAB + OAI-SwiGLU:
calc_diff = 0.0000

Standard SwiGLU regression:
calc_diff = 0.0000

The non-swapAB path was forced with:

DG_SM90_FP8_SWAP_AB=0

The TVM FFI extension and the generated SM90 CUDA JIT kernels were rebuilt from the patched source before running the tests.

Compatibility

Existing callers are unaffected because the new parameters default to standard SwiGLU behavior:

activation_alpha=1.0
activation_up_bias=0.0

No behavior changes are expected for existing DeepSeek MegaMoE workloads.

[shiyang814-cpu]

hi @Fridge003 , could you please review this SM90 FP8 MegaMoE correctness fix when you have a chance? Thanks!