[NVIDIA] Fix trtllm fp4 moe backend when used in MTP (#9384)

python/sglang/srt/layers/moe/topk.py

@@ -198,6 +198,7 @@ class TopK(CustomOp):
         correction_bias: Optional[torch.Tensor] = None,
         routed_scaling_factor: Optional[float] = None,
         apply_routed_scaling_factor_on_output: Optional[bool] = False,
+        force_topk: bool = False,
     ):
         # NOTE: scoring_func is not used for now, but we keep it for future use
         # see https://github.com/sgl-project/sglang/pull/4505 for more details
@@ -220,6 +221,7 @@ class TopK(CustomOp):
         )
 
         self.use_triton_kernels = get_moe_runner_backend().is_triton_kernel()
+        self.force_topk = force_topk
 
     def forward_native(
         self,
@@ -254,7 +256,7 @@ class TopK(CustomOp):
                 sm_first=not self.topk_config.renormalize,
             )
             return TritonKernelTopKOutput(routing_data, gather_idx, scatter_idx)
-        elif (
+        elif not self.force_topk and (
             should_use_flashinfer_trtllm_moe()
             or get_moe_runner_backend().is_flashinfer_mxfp4()
         ):