ops: add selectable SiTU-GLU activation

[Mantissagithub]

summary

This adds Sigmoid Tanh Unit GLU (SiTU-GLU) as a separate MoK activation while preserving SwiGLU as the default.

paper motivation

The Kimi K3 technical report observes that both multiplicative factors in SwiGLU are unbounded. Coincident large coordinates can therefore produce activation outliers and increase overflow risk in low-precision arithmetic. SiTU-GLU retains the sigmoid gate while smoothly capping the linear factor of the gate and the up branch independently.

Kimi K3, Section 2.3.2: SiTU-GLU motivation and Equation 12.

For gate pre-activation $g$ and up pre-activation $u$, the operation is

$$ h = \left[4\tanh\left(\frac{g}{4}\right)\sigma(g)\right] \left[25\tanh\left(\frac{u}{25}\right)\right]. $$

Kimi K3 fixes the gate soft cap to $\beta_1=4$ and the up-branch soft cap to $\beta_2=25$.

Kimi K3, Section 2.3.2: soft-cap values used by the model.

The scaled tanh remains approximately linear near the origin, keeping the local response close to SwiGLU. At large magnitude, both branches saturate and every output coordinate is bounded by $\beta_1\beta_2=100$.

Kimi K3, Figure 4: SiTU-GLU follows SwiGLU near the origin while remaining bounded.

changes

• add MoKConfig.activation with "swiglu" as the backward-compatible default and "situglu" as the new option
• instantiate distinct SiTU-GLU and SwiGLU native operations rather than replacing or aliasing the existing entry points
• implement SiTU-GLU forward and closed-form backward terms in FP32 inside the BF16 and MXFP8 megakernel paths
• cover shared and routed experts, including routed replay and router-gradient accumulation
• register Python custom operations, mutation schemas, fake implementations, and forward-context activation checks
• extend the BF16 reference path and add formula, bound, derivative, gradcheck, selector, schema, and distributed functional coverage
• document the selector and equation in the README

validation

No performance benchmarks were run for this PR. Validation is limited to the reference mathematics and successful SM100 compilation because no compatible SM100/SM103 system was available.

• python3 -m pytest --noconftest tests/test_situglu.py -q → 5 passed
• complete CUDA 13.0 SM100 compile and link with sm_100a code generation → passed
• ptxas reported zero spill stores and zero spill loads for all generated SiTU-GLU and SwiGLU forward/backward variants
• git diff --check → passed

The distributed BF16/MXFP8 end-to-end test is included but has not been executed on physical SM100/SM103 hardware. Throughput and latency therefore remain unmeasured. The available local GPU is SM89, which cannot run this megakernel.

Closes #7

---

Note

Medium Risk
Touches production megakernel forward/backward (including MXFP8 and router grads). Default SwiGLU is preserved, but the new path is numerically sensitive and activation-checkpoint recompute is not implemented for SiTU-GLU.

Overview
Adds SiTU-GLU as an opt-in MoE activation (MoKConfig.activation="situglu") while keeping SwiGLU as the default. The bounded GLU from the Kimi K3 report is compiled as separate BF16/MXFP8 forward and backward ops, not as a replacement of existing kernels.

The megakernel is templated on GluActivation. Device code uses FP32 situglu_forward / situglu_backward in shared and routed paths, including MXFP8 dequant/requant and router-gradient accumulation. Python registers parallel custom ops, fake impls, and mutation schemas, and stores the chosen activation on MoKForwardContext so backward cannot mix activations.

recompute_forward_context still supports SwiGLU only. Tests cover the formula, bounds, derivatives, gradcheck, schemas, and a distributed BF16/MXFP8 e2e path.

^{Reviewed by Cursor Bugbot for commit 45addfd. Bugbot is set up for automated code reviews on this repo. Configure here.}

[Mantissagithub]

@StuartSul, whenever you have a moment, i’d really appreciate it if you could take a look. thank you!

[cursor]

Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.

^{Bugbot Autofix is ON, but it could not run because the branch was deleted or merged before autofix could start.}

^{Reviewed by Cursor Bugbot for commit 2f4d24e. Configure here.}

[Mantissagithub]

rebased this onto the latest main and updated itt. would really appreciate it if you could take a look when you get a chance. cc @StuartSul @HuiyingLi, thank you!