Fix uniform_random routing simulation to sample without replacement (#43680)
Signed-off-by: Elvir Crncevic <elvircrn@gmail.com> Signed-off-by: Elvir Crnčević <elvircrn@gmail.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
This commit is contained in:
co-authored by
Claude Opus 4.6
gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Tyler Michael Smith
parent
fac808b36f
commit
8fea1d306f
@@ -131,14 +131,10 @@ class DistributionBasedRouting(RoutingStrategy):
|
||||
"""Sample expert IDs based on the specified distribution."""
|
||||
|
||||
if self.distribution == "uniform":
|
||||
# Uniform random sampling
|
||||
return torch.randint(
|
||||
low=0,
|
||||
high=num_experts,
|
||||
size=(num_tokens, top_k),
|
||||
dtype=indices_type,
|
||||
device=device,
|
||||
)
|
||||
# Generate random scores, and take the top-k to avoid duplicate topk_ids
|
||||
scores = torch.rand(num_tokens, num_experts, device=device)
|
||||
_, topk_ids = torch.topk(scores, top_k, dim=-1)
|
||||
return topk_ids.to(indices_type)
|
||||
|
||||
elif self.distribution == "normal":
|
||||
# For normal distribution, sample continuous values and map to
|
||||
|
||||
Reference in New Issue
Block a user