From 0d1a15e04894b831eb3e8fc0049faf0d1f730221 Mon Sep 17 00:00:00 2001 From: wangyicong Date: Fri, 22 May 2026 20:56:14 +0800 Subject: [PATCH] [Bugfix] Fix hash topk dtype mismatch Align hash routing metadata to the topk index dtype before invoking the sqrt-softplus topk custom op. DeepEP may request int64 topk indices while DeepSeek V4 hash MoE metadata remains int32, but the CUDA op dispatches hash metadata from the topk index dtype. Co-authored-by: Trae AI Signed-off-by: wangyicong (cherry picked from commit 01e7bebf34b4ce992c07b01ecb0d9c21a5890980) --- vllm/_custom_ops.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 84f944df2bf8..ddfc01a8c510 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2414,6 +2414,13 @@ def topk_hash_softplus_sqrt( input_tokens: torch.Tensor | None = None, hash_indices_table: torch.Tensor | None = None, ) -> None: + if hash_indices_table is not None: + assert input_tokens is not None + if input_tokens.dtype != topk_indices.dtype: + input_tokens = input_tokens.to(dtype=topk_indices.dtype) + if hash_indices_table.dtype != topk_indices.dtype: + hash_indices_table = hash_indices_table.to(dtype=topk_indices.dtype) + torch.ops._moe_C.topk_softplus_sqrt( topk_weights, topk_indices,