[Bugfix] Preserve fp32 MoE router precision across model gates

[seongyun1104]

Motivation

Addresses the independent model fixes called out in #38695.

Several MoE models still discard router precision at model-specific boundaries:

• MiMo-V2, Ernie4, and Kimi Linear can construct the checkpoint correction bias in bf16, even though expert selection consumes it in fp32-sensitive score arithmetic.
• Bailing MoE and Bailing MoE Linear compute with fp32 router weights when configured, then cast the logits back to the hidden-state dtype before top-k.

These paths either collapse checkpoint bias values or throw away the fp32 router output immediately before expert selection.

LLaDA2 originally had the same router-output downcast, but #31768 fixed that path independently. The current PR therefore leaves LLaDA2 untouched.

Modifications

• Keep e_score_correction_bias parameters in fp32 for MiMo-V2, Ernie4, and Kimi Linear.
• Return the gate GEMM output at the router-weight dtype for Bailing MoE and Bailing MoE Linear.

The patch is limited to five model files. It does not introduce a shared RouterGate, add a kernel, or change expert computation.

Accuracy Tests

No new unit test is included, following the maintainer guidance in #38695 that these model-file fixes do not require one. The changes only preserve the dtype already selected for the correction bias or router GEMM output.

Speed Tests and Profiling

Not applicable. No kernel is added. The Bailing paths remove a dtype-conversion operation when the router is configured for fp32.

Test Plan

• Rebased onto current main (961bcf481d).
• Python bytecode compilation for all five changed files.
• pre-commit run --files <all changed files>.
• git diff --check.

Checklist

• Format code according to the contribution guide.
• Follow the requested model-file-only scope.

---

CI States

Latest PR Test (Base): ❌ Run #38344434492
Latest PR Test (Extra): ❌ Run #38344434204
Latest PR Test (AMD ROCm 10): ❌ Run #38344434321

[seongyun1104]

@b8zhong this implements the two independent model-level precision fixes you called out in #38695: keeping correction bias fp32 in MiMo-V2/Ernie4/Kimi Linear, and removing the fp32-router-logit downcast in Bailing/Bailing Linear/LLaDA2.

@Fridge003 @ishandhanani @Qiaolin-Yu could one of the NV/model-specific optimization oncalls review the scope and contract? The patch is limited to the six model gates plus one registered CPU regression test. All changed files pass the full local pre-commit hook set, bytecode compilation, and git diff --check.

The current red matrix entries are PR-gate authorization failures and the registered CPU test was skipped. If the direction looks right, could a maintainer authorize base-a-test-cpu / add the run-ci label? Thanks.

[seongyun1104]

@b8zhong I rebased this PR onto current main (0b635266d4) and narrowed it according to your guidance in #38695.

• Removed the added unit-test file; the diff is now limited to the six model files.
• Kept the two independent fixes identified in the issue: FP32 correction biases for MiMo-V2/Ernie4/Kimi Linear, and removal of the router-logit downcast for Bailing MoE/Bailing MoE Linear/LLaDA2.
• Resolved the MiMo-V2 rebase conflict without restoring the upstream-removed reduce_moe_output path.
• All pre-commit hooks for the changed files and git diff --check pass.

The resulting patch is 9 insertions and 26 deletions and remains separate from the shared RouterGate refactor. Could you please take another look at the updated minimal scope when convenient? I can still split the two fix groups if that is specifically preferred.

[seongyun1104]

Rebased onto current main at 961bcf4.

The only conflict was LLaDA2: #31768 has independently landed the same router-output precision contract, so I dropped that hunk rather than duplicate it. The PR is now limited to the five still-unfixed model paths (Bailing MoE, Bailing MoE Linear, Ernie4, Kimi Linear, and MiMo-V2).

Validation after rebase: Python bytecode compilation, repository-pinned pre-commit on all changed files, and git diff --check all pass.

[seongyun1104]

Rebased onto current main at 4d5692b6cd after the shared model-stage refactor in #43460. The refactor touched four of the same model files but not the router-output or correction-bias contracts, so the five-path precision patch remains unchanged in scope.

Post-rebase validation: repository pre-commit on all five files, Python bytecode compilation, and git diff --check all pass. The current head is f4b4bee257.