[Bug] Fix TRTLLM Block FP8 MoE Monolithic

[wzhao18]

Purpose

Fix #36295

Test Plan

vllm serve deepseek-ai/DeepSeek-R1 -tp 8

lm_eval --model local-completions --model_args "base_url=http://0.0.0.0:8000/v1/completions,max_length=8192,tokenized_requests=False,tokenizer_backend=None,num_concurrent=32" --tasks gsm8k --num_fewshot 5

Test Result

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.953|±  |0.0058|
|     |       |strict-match    |     5|exact_match|↑  |0.953|±  |0.0058|

---

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 PR fixes a dtype mismatch bug for e_score_correction_bias in the TRT-LLM FP8 MoE monolithic path. The change in _apply_per_block is correct. I've identified a similar issue in _apply_per_tensor that is not covered by this PR and have left a comment with a suggested fix. Addressing this will ensure consistency and prevent potential runtime errors.

[ProExpertProg]

Can you try vllm serve deepseek-ai/DeepSeek-V3 -tp=8 as well?

[wzhao18]

@ProExpertProg I am experiencing OOM issue on top of tree. I can roll back to an older commit and test but it would be best if we could test on top-of-tree.

[ProExpertProg]

Tested your commit on top of my branch in #36551, that's good enough for me:

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.97	±	0.0171
		strict-match	5	exact_match	↑	0.97	±	0.0171

[robertgshaw2-redhat]

we need to backport this to 0.17.1

[pavanimajety]

LGTM, can we add a note in the PR or comment in the code saying what determines the e_score_correction_bias datatype? It's typically based on the checkpoint and kernel checkpoint, so hardcoding with hidden_states or router_logits, could work for a few models while not work for the others.
For example deepseek uses -

model.layers.3.mlp.gate.e_score_correction_bias | [256] | F32

[wzhao18]

Talked to Pavani offline. I think we could merge this patch for now. but it may be worth checking whether we should instead make change to the kernel side to support the built-in dtype for e_score_correction_bias from the model checkpoint.

[ProExpertProg]

This PR on B200:

local-completions (pretrained=deepseek-ai/DeepSeek-V3,base_url=http://0.0.0.0:8000/v1/completions,num_concurrent=50,max_retries=3), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.95	±	0.006
		strict-match	5	exact_match	↑	0.95	±	0.006

[robertgshaw2-redhat]

do we know which PR introduced this issue?

the monolithic refactor did not change the behavior here AFAICT

• wzhao18@97995f6#diff-0bca428f0755cb2a74957adf7250557ef5069290db40e919744308bd36e50001L220-L221

[robertgshaw2-redhat]

do we know which PR introduced this issue?

the monolithic refactor did not change the behavior here AFAICT

• wzhao18@97995f6#diff-0bca428f0755cb2a74957adf7250557ef5069290db40e919744308bd36e50001L220-L221

ah I see, its because now the hidden states are quantized. I dont understand why this is not caught by the tests, we run this for Qwen3 FP8 Block on B200

[robertgshaw2-redhat]

oh, its because it does not have e_score_correction_bias

[robertgshaw2-redhat]

this works, but is a hack. Do we know what dtype is required by the kernel here? Is it just bf16?

[wzhao18]

The reason for this regression per @robertgshaw2-redhat 's investigation is because the MoE refactor changes where the input quantization happens:

# prior
bias = bias.to(x.dtype)
x_q, x_scale = quantize(x)
kernel()

# after:
x << input to function is quantized
bias = bias.to(x.dtype)  # oops, cast to fp8
kernel()

Based on discussion, we remove the casting for the routing bias.

[robertgshaw2-redhat]

we discovered this cast is not needed due to trtllm kernel supporting fp32 inputs now. We originally added this cast here - ef28354

but this is not needed anymore

[ProExpertProg]

local-completions (pretrained=deepseek-ai/DeepSeek-V3,base_url=http://0.0.0.0:8000/v1/completions,num_concurrent=50,max_retries=3), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9462	±	0.0062
		strict-match	5	exact_match	↑	0.9462	±	0.0062

[robertgshaw2-redhat]

LGTM, thanks for the fix

[robertgshaw2-redhat]

verified same accuracy with FP32 and Bf16