Revert "[NemotronH] Do not force router to run in fp32 (#34582)"

[roikoren755]

Purpose

#34582 introduced an accuracy degradation. Working with @robertgshaw2-redhat for a better implementation of this performance improvement in #34302

Test Plan

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request reverts a previous change that introduced an accuracy degradation in the NemotronH model. The changes force the Mixture-of-Experts (MoE) router to operate in float32 precision. This is achieved by explicitly setting the gate layer's parameters to float32 and casting the input hidden states to float32 before the router computation. This is a standard practice to maintain numerical stability in MoE routing and is a correct fix for the reported accuracy issue. The changes are clear, concise, and effectively address the problem. I find no issues with this revert.

[mgoin]

What accuracy regression did it cause and for what model?

[roikoren755]

What accuracy regression did it cause and for what model?

#34704 caught the accuracy degradation for Nemotron 3 Nano FP8, which dropped 4-5%. Running a "standalone" GSM8K the score drops by about 2%. It seems like the BF16 model's accuracy also got a hit, but it stayed inside the tolerance of the test.