Fix/dsv4 flash eagle dummy ima (#25892)

This commit is contained in:
Kaixi
2026-05-20 15:41:28 -07:00
committed by GitHub
parent ce7141ef98
commit 9f2bc24b35
+18
View File
@@ -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)