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

[Bugfix] Accept int64 top-k IDs in FlashInfer routed MoE packer (#32882)

Browse Source
This commit is contained in:
Po-Han Huang (NVIDIA)
2026-09-01 14:17:46 -07:00
committed by GitHub
parent 83a9b5dd88
commit e57e934bcc
1 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/moe/pack_topk_ids.py
+6 -3
View File
@@ -3,6 +3,8 @@
Migrated from ``sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe`` Migrated from ``sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe``
(RFC #29630, Phase 2.5). Used by the FlashInfer TRT-LLM routed-MoE path, which (RFC #29630, Phase 2.5). Used by the FlashInfer TRT-LLM routed-MoE path, which
consumes routing ids and bf16 weights packed as ``(id << 16) | weight_bits``. consumes routing ids and bf16 weights packed as ``(id << 16) | weight_bits``.
Routing ids may be int32 or the int64 dtype produced by ``torch.topk``; they are
converted to int32 by the Triton kernel before packing.
""" """
import torch import torch
@@ -36,9 +38,10 @@ class PackTopkIds:
), f"shape mismatch: {topk_ids.shape=} vs {topk_weights.shape=}" ), f"shape mismatch: {topk_ids.shape=} vs {topk_weights.shape=}"
assert topk_ids.ndim >= 1, f"expected >=1D, got {topk_ids.shape=}" assert topk_ids.ndim >= 1, f"expected >=1D, got {topk_ids.shape=}"
assert ( assert topk_ids.dtype in (
topk_ids.dtype == torch.int32 torch.int32,
), f"topk_ids must be int32, got {topk_ids.dtype}" torch.int64,
), f"topk_ids must be int32 or int64, got {topk_ids.dtype}"
assert ( assert (
topk_weights.dtype == torch.float32 topk_weights.dtype == torch.float32
), f"topk_weights must be float32, got {topk_weights.dtype}" ), f"topk_weights must be float32, got {topk_weights.dtype}"