Fix/dsv4 flash eagle dummy ima (#25892)
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user