minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[NIXL] Query EP top-k index dtype (#35294)

Browse Source
This commit is contained in:
Yoray Zack
2026-08-19 22:31:05 +08:00
committed by GitHub
parent 23f2320c95
commit 1ef7882a5b
1 changed files with 8 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/token_dispatcher/nixl.py
+8 -1
View File
@@ -32,6 +32,13 @@ try:
except ImportError: except ImportError:
use_nixl = False 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__) logger = logging.getLogger(__name__)
NixlEPDispatchOutput = DeepEPLLDispatchOutput NixlEPDispatchOutput = DeepEPLLDispatchOutput
@@ -288,7 +295,7 @@ class _NixlEPDispatcherImpl(_NixlEPDispatcherImplBase):
): ):
buffer = self._get_buffer() buffer = self._get_buffer()
topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids 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() state = NixlEPBuffer._state()
dispatch_ep_size = state.dispatch_ep_size dispatch_ep_size = state.dispatch_ep_size
num_local_experts = state.num_local_experts num_local_experts = state.num_local_experts