From 1ef7882a5b76f3eca14b1abd37ffa1da3c353e1c Mon Sep 17 00:00:00 2001 From: Yoray Zack Date: Wed, 19 Aug 2026 17:31:05 +0300 Subject: [PATCH] [NIXL] Query EP top-k index dtype (#35294) --- python/sglang/srt/layers/moe/token_dispatcher/nixl.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/nixl.py b/python/sglang/srt/layers/moe/token_dispatcher/nixl.py index 93a72433e..2f2e3e849 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/nixl.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/nixl.py @@ -32,6 +32,13 @@ try: except ImportError: use_nixl = False +try: + from nixl_ep import topk_idx_t as NIXL_EP_TOPK_INDICES_DTYPE +except ImportError: + NIXL_EP_TOPK_INDICES_DTYPE = torch.int64 + +assert isinstance(NIXL_EP_TOPK_INDICES_DTYPE, torch.dtype) + logger = logging.getLogger(__name__) NixlEPDispatchOutput = DeepEPLLDispatchOutput @@ -288,7 +295,7 @@ class _NixlEPDispatcherImpl(_NixlEPDispatcherImplBase): ): buffer = self._get_buffer() topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids - topk_ids = topk_ids.to(torch.int64) + topk_ids = topk_ids.to(NIXL_EP_TOPK_INDICES_DTYPE) state = NixlEPBuffer._state() dispatch_ep_size = state.dispatch_ep_size num_local_experts = state.num_local_experts