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),
|
torch.empty(vocab_size, topk - num_fused_shared_experts, dtype=torch.int32),
|
||||||
requires_grad=False,
|
requires_grad=False,
|
||||||
)
|
)
|
||||||
|
self._init_default_tid2eid()
|
||||||
|
|
||||||
assert not apply_routed_scaling_factor_on_output, "not implemented"
|
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):
|
def empty_topk_output(self, device: torch.device):
|
||||||
topk = self.topk - self.num_fused_shared_experts
|
topk = self.topk - self.num_fused_shared_experts
|
||||||
topk_weights = torch.empty((0, topk), dtype=torch.float32, device=device)
|
topk_weights = torch.empty((0, topk), dtype=torch.float32, device=device)
|
||||||
|
|||||||
Reference in New Issue
Block a user