A high-throughput and memory-efficient inference and serving engine for LLMs
[EP] Query NIXL EP top-k index dtype (#45298)
Signed-off-by: Itay Alroy <ialroy@nvidia.com>
Itay Alroy committed
Jun 16, 2026 at 01:50 UTC
16e91176cf77bf0f40ae48da22365a5e21b517af
1 file changed
+3
-1
vllm/model_executor/layers/fused_moe/prepare_finalize/nixl_ep.py
+3
-1
@@ -28,6 +28,8 @@ logger = init_logger(__name__)
28
# NIXL EP kernels quantize dispatch inputs in 128 element chunks.
29
NIXL_EP_QUANT_BLOCK_SIZE = 128
30
NIXL_EP_QUANT_BLOCK_SHAPE = [NIXL_EP_QUANT_BLOCK_SIZE, NIXL_EP_QUANT_BLOCK_SIZE]
31
+NIXL_EP_TOPK_INDICES_DTYPE = getattr(nixl_ep, "topk_idx_t", torch.int64)
32
+assert isinstance(NIXL_EP_TOPK_INDICES_DTYPE, torch.dtype)
33
34
35
def dequant_fp8(
@@ -151,7 +153,7 @@ class NixlEPPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
153
all2all_manager.commit_staged_state()
154
155
def topk_indices_dtype(self) -> torch.dtype | None:
154
- return torch.int64
156
+ return NIXL_EP_TOPK_INDICES_DTYPE
157
158
def _map_global_to_physical_ids(self, topk_ids: torch.Tensor) -> torch.Tensor:
159
if self.global_to_physical is None: