[NemotronH] Do not force router to run in fp32

[roikoren755]

Purpose

Current code forces the MoE router computation to FP32, even though checkpoints have it in bfloat16. This takes up about 40% of the forward pass, under normal workloads, and does not provide an accuracy boost.

This PR removes this limitations.

Test Plan

No additional tests, all tests pass, accuracy does not degrade

Test Result

All tests pass, accuracy did not degrade

Running GSM8K, got the following results.

PR:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.5671|±  |0.0136|
|     |       |strict-match    |     5|exact_match|↑  |0.8431|±  |0.0100|

Main:

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.5572|±  |0.0137|
|     |       |strict-match    |     5|exact_match|↑  |0.8453|±  |0.0100|

---

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 introduces a valuable performance optimization by removing the forced casting of MoE router logits to float32. The changes in nemotron_h.py correctly implement this, and the special case for DeepSeekV3 is properly handled in flashinfer_trtllm_moe.py. I've found one minor issue: a leftover debug print statement that should be removed.

[gemini-code-assist]

This print statement appears to be for debugging purposes and should be removed before merging to avoid polluting logs.

[tomeras91]

I remember something about it being important that the bias is in FP32.. I understand in this case we first cast logits to FP32 (since we're using DS routing) so the bias is actually in FP32, but doesn't it make more sense to cast logits to bias dtype instead of the other way around?

[mgoin]

LGTM, nice find! We maybe should add an assert on the bias, but this roughly matches other trtllm moe impls. Do you have any perf result? You mentioned this takes 40% of time, which doesn't make sense to me

[roikoren755]

LGTM, nice find! We maybe should add an assert on the bias, but this roughly matches other trtllm moe impls. Do you have any perf result? You mentioned this takes 40% of time, which doesn't make sense to me

40% might have been an exaggeration 😅
But in an example workload, with TP8, we saw the following kernel distribution:

Keep in mind this is with an older commit, before the TRTLLM-Gen kernels were merged. The second kernel, taking ~19% of the profile, is the fp32 router