A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 12:30 UTC
[ModelBash][DSR1 NVFp4] Removed Bf16 Bias Cast (#34298)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Robert Shaw committed
Feb 11, 2026 at 00:00 UTC
9b17c57460bb5f6595f27b43e43caba144a8ec3c
1 file changed
+4
-7
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+4
-7
@@ -25,6 +25,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
25
from vllm.platforms import current_platform
26
27
if TYPE_CHECKING:
28
+ from vllm.model_executor.layers.fused_moe.layer import FusedMoE
29
from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
30
NvFp4MoeBackend,
31
)
@@ -316,11 +317,7 @@ def flashinfer_trtllm_fp4_moe(
317
if use_llama4_routing:
318
routing_method_type = flashinfer.RoutingMethodType.Llama4
319
319
- # Prepare routing bias
320
- routing_bias = e_score_correction_bias
321
- if routing_bias is not None:
322
- routing_bias = routing_bias.to(torch.bfloat16)
323
-
320
+ # Cast to Fp32 (required by kernel).
321
router_logits = (
322
router_logits.to(torch.float32)
323
if routing_method_type == RoutingMethodType.DeepSeekV3
@@ -330,7 +327,7 @@ def flashinfer_trtllm_fp4_moe(
327
# Call TRT-LLM FP4 block-scale MoE kernel
328
out = flashinfer.fused_moe.trtllm_fp4_block_scale_moe(
329
routing_logits=router_logits,
333
- routing_bias=routing_bias,
330
+ routing_bias=e_score_correction_bias,
331
hidden_states=hidden_states_fp4,
332
hidden_states_scale=hidden_states_scale_linear_fp4.view(
333
torch.float8_e4m3fn
@@ -447,7 +444,7 @@ def flashinfer_trtllm_fp4_routed_moe(
444
445
def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
446
backend: "NvFp4MoeBackend",
450
- layer: torch.nn.Module,
447
+ layer: "FusedMoE",
448
w13: torch.Tensor,
449
w13_scale: torch.Tensor,
450
w13_scale_2: torch.Tensor,