diff --git a/python/sglang/srt/layers/moe/hash_topk.py b/python/sglang/srt/layers/moe/hash_topk.py index 6b63b286a..cdeca35f7 100644 --- a/python/sglang/srt/layers/moe/hash_topk.py +++ b/python/sglang/srt/layers/moe/hash_topk.py @@ -41,9 +41,27 @@ class HashTopK(nn.Module): torch.empty(vocab_size, topk - num_fused_shared_experts, dtype=torch.int32), requires_grad=False, ) + self._init_default_tid2eid() assert not apply_routed_scaling_factor_on_output, "not implemented" + def _init_default_tid2eid(self) -> None: + topk = self.tid2eid.shape[1] + if topk == 0: + return + + # DummyModelLoader only initializes floating tensors, so keep this int + # lookup table valid until real checkpoints overwrite it. + token_ids = torch.arange( + self.tid2eid.shape[0], dtype=self.tid2eid.dtype, device=self.tid2eid.device + ).unsqueeze(1) + expert_offsets = torch.arange( + topk, dtype=self.tid2eid.dtype, device=self.tid2eid.device + ).unsqueeze(0) + tid2eid = (token_ids + expert_offsets) % self.num_experts + with torch.no_grad(): + self.tid2eid.copy_(tid2eid.to(self.tid2eid.dtype)) + def empty_topk_output(self, device: torch.device): topk = self.topk - self.num_fused_shared_experts topk_weights = torch.empty((0, topk), dtype=torch.float32, device=device)