Add opt-in FlashInfer-fast activations for rollout parity

[zianglih]

What does this PR do ?

Adds an opt-in Megatron SwiGLU implementation that matches the fast activation arithmetic used by FlashInfer's SM100 CuTe DSL MoE path.

The implementation is disabled by default and enabled at process startup with:

MILES_USE_FAST_ACTIVATIONS=1

Why

Miles uses Megatron for training and FlashInfer for rollout inference. Megatron's standard F.silu path and FlashInfer's fast exp2/reciprocal approximation can produce different SwiGLU values, which is undesirable for RL rollout/training consistency.

The target arithmetic is the FlashInfer SM100 CuTe sequence:

sigmoid = rcp.approx(1 + exp2.approx(-gate * log2(e)))
output = (sigmoid * gate) * up

Source review found that FlashInfer's CuTe, CUTLASS, CUDA, and TRT-LLM paths use the same broad fast-math family, although this PR claims exact arithmetic alignment only with the explicit SM100 CuTe sequence.

Implementation

• Adds one consolidated fast_activations.py module that owns the environment policy, fast SwiGLU forward, and analytic training backward.
• Selects stable forward/backward aliases at import time for existing fused SwiGLU wrappers.
• Wires the legacy GroupedMLP SwiGLU path through the existing custom autograd wrappers when enabled.
• Adds one consolidated unit-test file.
• Adds a standalone CuTe DSL experiment with plain contiguous layouts and no FlashInfer dependency.
• Documents the source analysis, environment setup, scope, and B200 results.

The experiment uses the DeepSeek-V3 activation shape [7168, 4096] -> [7168, 2048] and excludes GEMM, quantization, routing, permutation, unpermutation, and MoE combine.

Nemotron-3's ReLU2 path was also audited. It uses only max(x, 0) followed by self-multiplication, with no approximate transcendental operation. The B200 experiment found exact agreement, so this PR does not add an alternate ReLU2 implementation.

Validation

Run on NVIDIA B200 with PyTorch 2.11.0+cu130, CUDA 13.0, and CUTLASS DSL 4.5.2:

• Stock FP32 Megatron SwiGLU differed from the CuTe oracle on 9.64% to 37.04% of outputs across four deterministic distributions.
• The opt-in implementation exactly matched all 58,720,256 FP32 and 58,720,256 BF16 candidate outputs.
• Megatron ReLU2 exactly matched all 106,430,464 tested FP32/BF16 outputs.
• Focused unit tests: 7 passed with the environment disabled and 7 passed with it enabled.
• python3 -m py_compile passed for the new implementation, test, and experiment files.
• git diff --check passed.
• A GLM-5.2 5-layer NVFP4 E2E check used SGLang sgl-kernel Top-K, Miles torch Top-K, and zero BF16 boundary layers. Independent temperature-1 runs appeared lower with the fast toggle, but they sampled different responses.
• A fixed-rollout replay then isolated the activation toggle: the strict arm reproduced its original metrics exactly, while the fast arm showed small mixed changes. Across two rollouts, mean absolute log-probability difference was 0.0951% higher, rollout KL was 0.0961% higher, and KL loss was 0.0366% higher. This does not demonstrate an end-to-end improvement.

The full reproduction procedure and result table are in docs/discussions/flashinfer-fast-activations-alignment.md.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Loading

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

The repository autoformatter and pre-commit were not run for this experiment branch.

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

1. Attach the Expert Review label when your PR is ready for review.
2. GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

1. Add Final Review label
2. GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.