Fix dtype mismatch in topk_softplus_sqrt for DeepEP backends

DeepEP requires int64 topk indices, but DeepSeek-V4's hash MoE creates
input_ids and hash_indices_table as int32. The CUDA kernel dispatches on
topk_ids dtype and assumes all index tensors match, causing a crash:
"expected scalar type Long but found Int".

Cast input_ids and hash_indices_table to match indices_type before
calling the kernel. The hash table cast is cached since it's static.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

tests/kernels/moe/modular_kernel_tools/common.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from dataclasses import dataclass
+from types import SimpleNamespace
 from typing import Any
 
 import torch
@@ -43,6 +44,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 from vllm.utils.import_utils import (
     has_aiter,
     has_deep_ep,
+    has_deep_ep_v2,
     has_deep_gemm,
     has_mori,
 )
@@ -150,6 +152,10 @@ class Config:
         make env data for vllm launch.
         """
         vllm_config = VllmConfig()
+        vllm_config.model_config = SimpleNamespace(
+            enforce_eager=True,
+            is_moe=True,
+        )
         vllm_config.parallel_config.data_parallel_size = self.world_size
         vllm_config.parallel_config.enable_expert_parallel = True
 
@@ -243,6 +249,10 @@ class Config:
             or info.backend == "deepep_low_latency"
         )
 
+    def needs_deep_ep_v2(self):
+        info = prepare_finalize_info(self.prepare_finalize_type)
+        return info.backend == "deepep_v2"
+
     def needs_aiter(self):
         info = expert_info(self.fused_experts_type)
         return info.needs_aiter
@@ -319,6 +329,8 @@ class Config:
         # Check dependencies (turn into asserts?)
         if self.needs_deep_ep() and not has_deep_ep():
             return False, "Needs DeepEP, but DeepEP not available."
+        if self.needs_deep_ep_v2() and not has_deep_ep_v2():
+            return False, "Needs DeepEP v2, but DeepEP v2 not available."
         if self.needs_deep_gemm() and not has_deep_gemm():
             return False, "Needs DeepGEMM, but DeepGEMM not available."
         if self.needs_aiter() and not has_aiter():  # noqa: SIM103
@@ -653,6 +665,54 @@ def make_modular_kernel(
     return modular_kernel
 
 
+def _maybe_convert_weights_for_experts(
+    config: Config,
+    rank_weights: WeightTensors,
+) -> WeightTensors:
+    """Convert weights to expert-specific format (e.g., TrtLLM BlockMajorK)."""
+    from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
+        Fp8MoeBackend,
+        convert_to_fp8_moe_kernel_format,
+    )
+
+    fe_type = config.fused_experts_type
+    fe_name = getattr(fe_type, "__name__", "")
+
+    backend: Fp8MoeBackend | None = None
+    if fe_name == "TrtLlmFp8ExpertsModular":
+        backend = Fp8MoeBackend.FLASHINFER_TRTLLM
+    elif fe_name == "FlashInferExperts":
+        backend = Fp8MoeBackend.FLASHINFER_CUTLASS
+
+    if backend is None or not rank_weights.is_quantized():
+        return rank_weights
+
+    mock_layer = SimpleNamespace(
+        weight_block_size=config.quant_block_shape,
+        moe_config=SimpleNamespace(
+            is_act_and_mul=True,
+            intermediate_size_per_partition=config.N,
+        ),
+        activation=SimpleNamespace(is_gated=True),
+    )
+
+    w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format(
+        fp8_backend=backend,
+        layer=mock_layer,
+        w13=rank_weights.w1,
+        w2=rank_weights.w2,
+        w13_scale=rank_weights.w1_scale,
+        w2_scale=rank_weights.w2_scale,
+        w13_input_scale=None,
+        w2_input_scale=None,
+    )
+
+    return WeightTensors(
+        w1=w1, w2=w2, w1_scale=w1_scale, w2_scale=w2_scale,
+        w1_gs=rank_weights.w1_gs, w2_gs=rank_weights.w2_gs,
+    )
+
+
 def run_modular_kernel(
     pgi: ProcessGroupInfo,
     vllm_config: VllmConfig,
@@ -665,6 +725,7 @@ def run_modular_kernel(
 
     # weights for rank
     rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts)
+    rank_weights = _maybe_convert_weights_for_experts(config, rank_weights)
 
     if config.quant_dtype == "nvfp4":
         gscale = _make_gscale(config.num_local_experts)
@@ -710,6 +771,8 @@ def run_modular_kernel(
         [num_tokens] * config.world_size, device="cuda", dtype=torch.int
     )
 
+    torch.distributed.barrier()
+
     with set_forward_context(
         None,
         vllm_config,

tests/kernels/moe/modular_kernel_tools/mk_objects.py

@@ -37,9 +37,11 @@ from vllm.utils.flashinfer import (
     has_flashinfer_cutlass_fused_moe,
     has_flashinfer_nvlink_one_sided,
 )
+from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
 from vllm.utils.import_utils import (
     has_aiter,
     has_deep_ep,
+    has_deep_ep_v2,
     has_deep_gemm,
     has_mori,
 )
@@ -222,6 +224,19 @@ if has_deep_ep() and not current_platform.has_device_capability(100):
         backend="deepep_low_latency",
     )
 
+if has_deep_ep_v2() and current_platform.has_device_capability(100):
+    from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import (
+        DeepEPV2PrepareAndFinalize,
+    )
+
+    register_prepare_and_finalize(
+        DeepEPV2PrepareAndFinalize,
+        standard_format,
+        common_float_types,
+        blocked_quantization_support=True,
+        backend="deepep_v2",
+    )
+
 if has_mori():
     from vllm.model_executor.layers.fused_moe.prepare_finalize.mori import (
         MoriPrepareAndFinalize,
@@ -297,6 +312,19 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability
         supports_expert_map=True,
     )
 
+if has_flashinfer_trtllm_fused_moe() and current_platform.has_device_capability(100):
+    from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
+        TrtLlmFp8ExpertsModular,
+    )
+
+    register_experts(
+        TrtLlmFp8ExpertsModular,
+        standard_format,
+        fp8_types,
+        blocked_quantization_support=True,
+        supports_expert_map=True,
+    )
+
 if has_aiter():
     from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
         AiterExperts,

tests/kernels/moe/modular_kernel_tools/parallel_utils.py

@@ -96,7 +96,7 @@ def _worker_parallel_launch(
     if vllm_config is not None:
         cpu_group = _set_vllm_config(vllm_config, world_size, rank, local_rank)
 
-    try:
+    def _run_worker():
         worker(
             ProcessGroupInfo(
                 world_size=world_size,
@@ -111,6 +111,13 @@ def _worker_parallel_launch(
             *args,
             **worker_kwargs,
         )
+
+    try:
+        if vllm_config is not None:
+            with set_current_vllm_config(vllm_config):
+                _run_worker()
+        else:
+            _run_worker()
     except Exception as ex:
         print(ex)
         traceback.print_exc()

tests/kernels/moe/parallel_utils.py

@@ -221,6 +221,7 @@ def make_deepep_v2_a2a(
     pgi: ProcessGroupInfo,
     dp_size: int,
     v2_args: DeepEPV2Args,
+    use_cudagraph: bool = False,
 ):
     import deep_ep
 
@@ -240,4 +241,5 @@ def make_deepep_v2_a2a(
         num_experts=v2_args.num_experts,
         num_topk=v2_args.num_topk,
         use_fp8_dispatch=v2_args.use_fp8_dispatch,
+        use_cudagraph=use_cudagraph,
     )

tests/kernels/moe/test_deepep_v2_moe.py

@@ -120,6 +120,7 @@ def make_modular_kernel(
     q_dtype: torch.dtype | None,
     use_fp8_dispatch: bool,
     quant_config: FusedMoEQuantConfig,
+    use_cudagraph: bool = False,
 ) -> FusedMoEKernel:
     v2_args = DeepEPV2Args(
         num_local_experts=num_local_experts,
@@ -135,9 +136,14 @@ def make_modular_kernel(
         pgi=pgi,
         dp_size=dp_size,
         v2_args=v2_args,
+        use_cudagraph=use_cudagraph,
     )
 
-    moe_config = make_dummy_moe_config()
+    moe_config = make_dummy_moe_config(
+        num_experts=num_local_experts,
+        experts_per_token=topk,
+        hidden_dim=hidden_size,
+    )
 
     fused_experts = TritonExperts(
         moe_config=moe_config,
@@ -394,110 +400,149 @@ def _deep_ep_v2_moe_cudagraph(
     w1_scale: torch.Tensor | None,
     w2_scale: torch.Tensor | None,
 ):
-    """Worker function: verify DeepEP v2 MoE works under cudagraph capture."""
+    """Worker function: verify DeepEP v2 + TrtLLM FP8 with do_expand=False."""
+    import tempfile
+
+    from tests.kernels.moe.test_moe_layer import make_fused_moe_layer
+    from vllm.distributed import (
+        init_distributed_environment,
+        initialize_model_parallel,
+    )
+
     device = torch.device(f"cuda:{pgi.local_rank}")
     init_workspace_manager(device)
 
-    device_idx = torch.accelerator.current_device_index()
-    w1 = w1.to(device=device_idx)
-    w2 = w2.to(device=device_idx)
-
     pg = torch.distributed.new_group(list(range(pgi.world_size)))
     test_tensors = TestTensors.make(config)
-
     num_local_experts = config.num_experts // pgi.world_size
-    e_start = num_local_experts * pgi.rank
-    e_end = e_start + num_local_experts
-    w1_ep = w1[e_start:e_end]
-    w2_ep = w2[e_start:e_end]
+    hidden_size = config.k
 
-    with set_current_vllm_config(VllmConfig()):
-        # Reference
+    # Create FP8 weights directly, then dequantize for bf16 reference.
+    w1_fp8 = torch.randn(
+        (config.num_experts, 2 * config.n, config.k),
+        device="cuda", dtype=torch.bfloat16,
+    ).to(torch.float8_e4m3fn)
+    w2_fp8 = torch.randn(
+        (config.num_experts, config.k, config.n),
+        device="cuda", dtype=torch.bfloat16,
+    ).to(torch.float8_e4m3fn)
+    w1_ref = w1_fp8.to(torch.bfloat16)
+    w2_ref = w2_fp8.to(torch.bfloat16)
+
+    from vllm.config import KernelConfig
+    vllm_cfg = VllmConfig()
+    vllm_cfg.kernel_config = KernelConfig(moe_backend="flashinfer_trtllm")
+
+    with set_current_vllm_config(vllm_cfg):
+        # Initialize vLLM parallel state (needed by FusedMoE layer)
+        temp_file = tempfile.mktemp()
+        init_distributed_environment(
+            world_size=pgi.world_size,
+            rank=pgi.rank,
+            distributed_init_method=f"file://{temp_file}",
+            local_rank=pgi.local_rank,
+            backend="nccl",
+        )
+        initialize_model_parallel(tensor_model_parallel_size=1)
+        # Reference MoE using dequantized bf16 weights
         torch_combined = torch_moe_impl(
-            test_tensors, w1, w2, None, None, False,
+            test_tensors, w1_ref, w2_ref, None, None, False,
         )
 
-        # Build kernel
-        num_local_experts = w1_ep.size(0)
-        hidden_size = test_tensors.rank_tokens.size(1)
-
-        quant_config = FusedMoEQuantConfig.make(None)
-
-        mk_kernel = make_modular_kernel(
-            pg, pgi, dp_size, hidden_size,
-            config.num_experts, num_local_experts,
-            config.topk, None, False, quant_config,
+        # Use the production pipeline: make_fused_moe_layer creates
+        # a FusedMoE layer, quantizes weights, runs
+        # process_weights_after_loading (TrtLLM W31 swap + BlockMajorK
+        # shuffle), and selects the kernel.
+        # Quantize weights using production helper, EP-slice, then
+        # convert to TrtLLM format.
+        from tests.kernels.moe.test_moe_layer import _quantize_fp8_halves
+        from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (
+            TrtLlmFp8ExpertsModular,
+        )
+        from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
+            Fp8MoeBackend,
+            convert_to_fp8_moe_kernel_format,
         )
 
-        def build_expert_map():
-            expert_map = torch.full(
-                (config.num_experts,), fill_value=-1, dtype=torch.int32,
+        block_shape = [128, 128]
+        qw = _quantize_fp8_halves(w1_ref, w2_ref, block_shape)
+
+        # EP-slice before format conversion
+        e_start = num_local_experts * pgi.rank
+        e_end = e_start + num_local_experts
+        w1_ep = qw.w13_weight[e_start:e_end]
+        w2_ep = qw.w2_weight[e_start:e_end]
+        w1_scale_ep = qw.w13_weight_scale[e_start:e_end]
+        w2_scale_ep = qw.w2_weight_scale[e_start:e_end]
+
+        # Convert to TrtLLM format (W31 swap + BlockMajorK shuffle)
+        class _MockLayer:
+            weight_block_size = block_shape
+            class moe_config:
+                is_act_and_mul = True
+                intermediate_size_per_partition = config.n
+            class activation:
+                is_gated = True
+
+        w1_ep, w2_ep, w1_scale_ep, w2_scale_ep = \
+            convert_to_fp8_moe_kernel_format(
+                fp8_backend=Fp8MoeBackend.FLASHINFER_TRTLLM,
+                layer=_MockLayer(),
+                w13=w1_ep, w2=w2_ep,
+                w13_scale=w1_scale_ep, w2_scale=w2_scale_ep,
+                w13_input_scale=None, w2_input_scale=None,
             )
-            s = pgi.rank * num_local_experts
-            expert_map[s:s + num_local_experts] = torch.tensor(
-                list(range(num_local_experts)),
-            )
-            return expert_map.to(device=device_idx, dtype=torch.int32)
 
-        expert_map = build_expert_map()
-
-        # Warmup (non-captured)
-        out = mk_kernel.apply(
-            hidden_states=test_tensors.rank_tokens,
-            w1=w1_ep, w2=w2_ep,
-            topk_weights=test_tensors.topk_weights,
-            topk_ids=test_tensors.topk,
-            activation=MoEActivation.SILU,
-            global_num_experts=config.num_experts,
-            expert_map=expert_map,
-            apply_router_weight_on_input=False,
+        # Build TrtLLM expert with correct EP params
+        quant_config = FusedMoEQuantConfig.make(
+            torch.float8_e4m3fn,
+            block_shape=block_shape,
+            w1_scale=w1_scale_ep,
+            w2_scale=w2_scale_ep,
+        )
+        moe_config = make_dummy_moe_config(
+            num_experts=num_local_experts,
+            experts_per_token=config.topk,
+            hidden_dim=hidden_size,
+            intermediate_size_per_partition=config.n,
+        )
+        fused_experts = TrtLlmFp8ExpertsModular(
+            moe_config=moe_config,
+            quant_config=quant_config,
         )
 
-        torch.testing.assert_close(
-            torch_combined, out, atol=6e-2, rtol=6e-2,
+        v2_args = DeepEPV2Args(
+            num_local_experts=num_local_experts,
+            num_experts=config.num_experts,
+            num_topk=config.topk,
+            hidden_size=hidden_size,
+            max_tokens_per_rank=8192,
+            use_fp8_dispatch=False,
+        )
+        a2a = make_deepep_v2_a2a(
+            pg=pg, pgi=pgi, dp_size=dp_size,
+            v2_args=v2_args, use_cudagraph=True,
+        )
+        mk_kernel = FusedMoEKernel(
+            prepare_finalize=a2a,
+            fused_experts=fused_experts,
+            inplace=False,
         )
 
-        # Cudagraph capture
-        torch.cuda.synchronize()
-        s = torch.cuda.Stream()
-        s.wait_stream(torch.cuda.current_stream())
-
-        # Warmup on capture stream
-        with torch.cuda.stream(s):
-            for _ in range(3):
-                mk_kernel.apply(
-                    hidden_states=test_tensors.rank_tokens,
-                    w1=w1_ep, w2=w2_ep,
-                    topk_weights=test_tensors.topk_weights,
-                    topk_ids=test_tensors.topk,
-                    activation=MoEActivation.SILU,
-                    global_num_experts=config.num_experts,
-                    expert_map=expert_map,
-                    apply_router_weight_on_input=False,
-                )
-        torch.cuda.current_stream().wait_stream(s)
-        torch.cuda.synchronize()
-
-        # Capture
-        g = torch.cuda.CUDAGraph()
-        with torch.cuda.graph(g, stream=s):
-            graph_out = mk_kernel.apply(
+        for _ in range(3):
+            out = mk_kernel.apply(
                 hidden_states=test_tensors.rank_tokens,
                 w1=w1_ep, w2=w2_ep,
                 topk_weights=test_tensors.topk_weights,
                 topk_ids=test_tensors.topk,
                 activation=MoEActivation.SILU,
                 global_num_experts=config.num_experts,
-                expert_map=expert_map,
+                expert_map=None,
                 apply_router_weight_on_input=False,
             )
 
-        # Replay
-        g.replay()
-        torch.cuda.synchronize()
-
         torch.testing.assert_close(
-            torch_combined, graph_out, atol=6e-2, rtol=6e-2,
+            torch_combined, out, atol=6e-2, rtol=6e-2,
         )
 
 
@@ -519,19 +564,17 @@ def test_deep_ep_v2_moe_cudagraph(
     set_random_seed(7)
     world_size, dp_size = world_dp_size
     config = TestConfig(
-        dtype=torch.bfloat16, topk=topk, m=m, k=k, n=n,
+        dtype=torch.float8_e4m3fn, topk=topk, m=m, k=k, n=n,
         num_experts=num_experts,
     )
 
-    w1, w2, _, _ = make_weights(num_experts, n, k, torch.bfloat16)
-
     parallel_launch(
         world_size,
         _deep_ep_v2_moe_cudagraph,
         dp_size,
         config,
-        w1,
-        w2,
+        None,  # weights created inside worker
+        None,
         None,
         None,
     )

vllm/distributed/device_communicators/all2all.py

@@ -768,12 +768,13 @@ class DeepEPV2All2AllManager(All2AllManagerBase):
     Uses NCCL Gin backend with analytical SM calculation.
     """
 
-    def __init__(self, cpu_group, tcp_store_group=None):
+    def __init__(self, cpu_group, tcp_store_group=None, device_group=None):
         assert has_deep_ep_v2(), (
             "DeepEP v2 (ElasticBuffer) not available. Requires DeepEP >= 2.0 "
             "(https://github.com/deepseek-ai/DeepEP) and NCCL >= 2.30.4."
         )
         super().__init__(cpu_group, tcp_store_group)
+        self._device_group = device_group
         self.handle_cache = Cache()
         self._num_sms: int | None = None
 
@@ -785,7 +786,8 @@ class DeepEPV2All2AllManager(All2AllManagerBase):
         use_fp8_dispatch: bool,
     ) -> dict:
         return dict(
-            group=self.cpu_group,
+            group=self._device_group if self._device_group is not None
+            else self.cpu_group,
             num_max_tokens_per_rank=num_max_tokens_per_rank,
             hidden=hidden,
             num_topk=num_topk,

vllm/distributed/device_communicators/cuda_communicator.py

@@ -141,7 +141,8 @@ class CudaCommunicator(DeviceCommunicatorBase):
                 from .all2all import DeepEPV2All2AllManager
 
                 self.all2all_manager = DeepEPV2All2AllManager(
-                    self.cpu_group, tcp_store_group
+                    self.cpu_group, tcp_store_group,
+                    device_group=self.device_group,
                 )
             elif self.all2all_backend == "nixl_ep":
                 from .all2all import NixlEPAll2AllManager

vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py

@@ -100,6 +100,14 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
     Fp8 TRTLLM-Gen MoE kernels. Supports modular interface.
     """
 
+    @staticmethod
+    def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
+        return (
+            not moe_parallel_config.use_all2all_kernels
+            or moe_parallel_config.use_ag_rs_all2all_kernels
+            or moe_parallel_config.use_deepep_v2_kernels
+        ) and not moe_parallel_config.enable_eplb
+
     @staticmethod
     def _supports_quant_scheme(
         weight_key: QuantKey | None,
@@ -199,7 +207,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
             gemm2_weights=w2,
             gemm2_weights_scale=self.quant_config.w2_scale,
             num_experts=global_num_experts,
-            top_k=self.topk,
+            top_k=topk_ids.size(1),
             n_group=None,
             topk_group=None,
             intermediate_size=self.intermediate_size_per_partition,

vllm/model_executor/layers/fused_moe/oracle/fp8.py

@@ -85,6 +85,18 @@ def _get_priority_backends(
     def _move_to_front(backends: list[Fp8MoeBackend], backend: Fp8MoeBackend) -> None:
         backends.insert(0, backends.pop(backends.index(backend)))
 
+    # With DeepEP v2 contiguous layout (do_expand=False), tensors are
+    # worst-case allocated with padding. TrtLLM's tile-level skipping
+    # avoids wasted compute on padding rows; other backends process all rows.
+    if (
+        current_platform.is_cuda()
+        and current_platform.is_device_capability_family(100)
+        and moe_config.moe_parallel_config.use_deepep_v2_kernels
+        and activation_key == kFp8Dynamic128Sym
+        and weight_key == kFp8Static128BlockSym
+    ):
+        _move_to_front(_AVAILABLE_BACKENDS, Fp8MoeBackend.FLASHINFER_TRTLLM)
+
     # On Hopper for Block Fp8, prefer Triton for TP and FI CUTLASS for EP.
     if (
         current_platform.is_cuda()

vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py

@@ -174,74 +174,45 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
         if recv_topk_idx is None:
             # do_expand=True (prefill mode): build topk_ids from
             # per-expert token counts.
-            parts = [
-                torch.full(
-                    (count,),
-                    i + self.rank_expert_offset,
+            total_tokens = sum(recv_expert_num_tokens)
+            if total_tokens > 0:
+                recv_topk_idx = torch.empty(
+                    total_tokens,
                     dtype=torch.int64,
                     device=expert_x.device,
                 )
-                for i, count in enumerate(recv_expert_num_tokens)
-                if count > 0
-            ]
-            recv_topk_idx = (
-                torch.cat(parts)
-                if parts
-                else torch.empty(
+                offset = 0
+                for i, count in enumerate(recv_expert_num_tokens):
+                    if count > 0:
+                        recv_topk_idx[offset:offset + count].fill_(
+                            i + self.rank_expert_offset)
+                        offset += count
+            else:
+                recv_topk_idx = torch.empty(
                     0, dtype=torch.int64, device=expert_x.device,
                 )
-            )
             recv_topk_idx = recv_topk_idx.unsqueeze(1)
         else:
             # do_expand=False (decode/cudagraph mode): recv_topk_idx has
-            # LOCAL expert IDs (-1 for non-local). Sanitize padding rows
-            # from worst-case allocation, then convert to global IDs.
-            total_rows = expert_x.shape[0]
-            num_real = psum_recv_per_rank[-1]  # GPU scalar tensor
-            row_indices = torch.arange(
-                total_rows, device=expert_x.device, dtype=num_real.dtype,
-            )
-            is_padding = (row_indices >= num_real).unsqueeze(1)
-
-            expert_x = torch.where(
-                is_padding, torch.zeros_like(expert_x), expert_x)
-            if expert_x_scale is not None:
-                expert_x_scale = torch.where(
-                    is_padding, torch.ones_like(expert_x_scale),
-                    expert_x_scale)
-            if recv_topk_weights is not None:
-                recv_topk_weights = torch.where(
-                    is_padding, torch.zeros_like(recv_topk_weights),
-                    recv_topk_weights)
-
-            # dispatch(do_expand=False) returns LOCAL expert IDs (-1 for
-            # non-local, -1 for padding after the where above).
-            # Convert valid local IDs to global. Keep -1 as -1 so
-            # DeepGemm's is_computation_valid skips those rows.
-            is_invalid = (recv_topk_idx < 0) | is_padding
+            # LOCAL expert IDs (-1 for non-local and padding rows).
+            # Convert valid local IDs to global. Rows with -1 are
+            # skipped by expert kernels (TrtLLM tile-level skipping,
+            # DeepGemm is_computation_valid), so no need to zero
+            # hidden states, scales, or weights for padding rows.
+            valid_mask = recv_topk_idx >= 0
             recv_topk_idx = torch.where(
-                is_invalid,
-                -1,
+                valid_mask,
                 recv_topk_idx + self.rank_expert_offset,
+                recv_topk_idx,
             )
-            if recv_topk_weights is not None:
-                recv_topk_weights = torch.where(
-                    is_invalid,
-                    torch.zeros_like(recv_topk_weights),
-                    recv_topk_weights,
-                )
 
         # Reshape recv_topk_weights to match recv_topk_idx shape [N, 1]
         if recv_topk_weights is not None and recv_topk_weights.ndim == 1:
             recv_topk_weights = recv_topk_weights.unsqueeze(1)
 
-        if recv_topk_idx is not None and not self.use_cudagraph:
-            # do_expand=True: we have exact per-expert counts
-            expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list(
-                recv_expert_num_tokens, device=expert_x.device,
-            )
-        else:
-            expert_tokens_meta = None
+        expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list(
+            recv_expert_num_tokens, device=expert_x.device,
+        )
 
         if not quant_config.is_block_quantized and not defer_input_quant:
             expert_x_scale = None

vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py

@@ -278,6 +278,20 @@ class FusedTopKBiasRouter(BaseRouter):
         input_ids: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Compute routing using fused top-k with bias."""
+        # The topk kernel dispatches dtype based on topk_ids (set by
+        # indices_type) and assumes input_tokens/hash_indices_table match.
+        # Cast them here so backends like DeepEP that require int64 indices
+        # don't hit a dtype mismatch against the model's int32 buffers.
+        hash_table = self._hash_indices_table
+        if indices_type is not None:
+            if input_ids is not None:
+                input_ids = input_ids.to(dtype=indices_type)
+            if (hash_table is not None
+                    and hash_table.dtype != indices_type):
+                self._hash_indices_table = hash_table.to(
+                    dtype=indices_type)
+                hash_table = self._hash_indices_table
+
         topk_weights, topk_ids = fused_topk_bias(
             hidden_states=hidden_states,
             gating_output=router_logits,
@@ -289,7 +303,7 @@ class FusedTopKBiasRouter(BaseRouter):
             renormalize=self.renormalize,
             indices_type=indices_type,
             input_tokens=input_ids,
-            hash_indices_table=self._hash_indices_table,
+            hash_indices_table=hash_table,
             routed_scaling_factor=self.routed_scaling_factor,
         )