[Bugfix] Fix accuracy issue of TRTLLM FP8 MOE and improve logging

[pavanimajety]

Purpose

Fixes #25189: accuracy issues for TRTLLM DSR1 Latency kernels
Bugs introduced in #23991 and #23640

Also fixes incorrect logging prints for which kernels are used. When Flashinfer MOE is enabled, Deep Gemm is automatically disabled for SM100.

(EngineCore_DP0 pid=101213) (Worker_TP3 pid=101225) INFO 09-29 10:14:46 [fp8.py:454] Detected Blackwell GPUs, using FlashInfer TensorRT-LLM kernels for FP8 MOE.
(EngineCore_DP0 pid=101213) (Worker_TP3 pid=101225) INFO 09-29 10:14:46 [fp8.py:468] DeepGemm disabled: FlashInfer MOE is enabled.

Test Plan

LM Eval with VLLM_USE_FLASHINFER_MOE_FP8=1 and VLLM_FLASHINFER_MOE_BACKEND="latency"

root@umbriel-b200-017:/workspace/scratch-pmaj-1/gh-pm-vllm# ./run_lm_eval.sh                                                                        
+ export FORCE_NUM_KV_SPLITS=1                                                                                                                      
+ FORCE_NUM_KV_SPLITS=1                                                                                                                             
+ export VLLM_USE_FLASHINFER_MOE_FP8=1                                                                                                              
+ VLLM_USE_FLASHINFER_MOE_FP8=1                                                                                                                     
+ export VLLM_FLASHINFER_MOE_BACKEND=latency                                                                                                        
+ VLLM_FLASHINFER_MOE_BACKEND=latency                                                                                                               
+ echo 'Running lm_eval with:'                                                                                                                      
Running lm_eval with:                                                                                                                               
                     
  Model: /models/DSR1-FP8/models--deepseek-ai--DeepSeek-R1-0528/snapshots/4236a6af538feda4548eca9ab308586007567f52/                                 
  Tensor Parallel Size: 8                                                                                                                           
  Quantization: fp8                                                                                                                                 
  GPU Memory Utilization: 0.90                                                                                                                      
                                                                      
+ lm_eval --model vllm --model_args pretrained=/models/DSR1-FP8/models--deepseek-ai--DeepSeek-R1-0528/snapshots/4236a6af538feda4548eca9ab30858600756
7f52/,quantization=fp8,tensor_parallel_size=8,gpu_memory_utilization=0.90,add_bos_token=True --gen_kwargs temperature=0.0,max_gen_toks=32768 --trust
_remote_code --tasks gsm8k --num_fewshot 5 --batch_size 200 --limit 1319 --output_path results/4236a6af538feda4548eca9ab308586007567f52-gsm8k-1319sa
mples_20250929_101419 --log_samples

Test Result

vllm (pretrained=/models/DSR1-FP8/models--deepseek-ai--DeepSeek-R1-0528/snapshots/4236a6af538feda4548eca9ab308586007567f52/,quantization=fp8,tensor_parallel_size=8,gpu_memory_utilization=0.90,add_bos_token=True,trust_remote_code=True), gen_kwargs: (temperature=0.0,max_gen_toks=32768), limit: 1319.0, num_fewshot: 5, batch_size: 200
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9666|±  |0.0049|
|     |       |strict-match    |     5|exact_match|↑  |0.9651|±  |0.0051|

---

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
• [N/A] (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 important fixes for FP8 MoE, particularly for TRT-LLM latency kernels. The changes correctly address an accuracy issue by ensuring e_score_correction_bias has the expected dtype, and fix a critical control flow bug by adding an early return after the MoE operation is complete. Additionally, the logic for selecting between FlashInfer, DeepGEMM, and Cutlass kernels is improved to prevent conflicts, disabling other kernels when FlashInfer MoE is active. The logging has also been updated to be more informative about which kernel is being used. The changes are well-structured and significantly improve the correctness and robustness of the FP8 MoE implementation.

[gemini-code-assist]

This is a great fix that addresses two important issues:

1. Casting e_score_correction_bias to x.dtype resolves a potential dtype mismatch that could lead to accuracy problems, which is critical for correctness.
2. Adding the return statement corrects a significant control flow bug. Previously, the code would fall through and execute select_experts and other logic even after the complete MoE operation was performed by flashinfer_fused_moe_blockscale_fp8. This early return ensures the function exits correctly.

This change is crucial for both correctness and logic of the FP8 MoE path.

[alexm-redhat]

LGTM

[mgoin]

Note for future self: we should clean up these logs now that VLLM_USE_DEEP_GEMM=1 by default

[mgoin]

Is this a bug too? It seems like we need to return here rather than write to result

[pavanimajety]

Yeah, it can potentially be a bug, havent' tested the path - so I was hestiant to change it.