icer/qwen36_01
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

revert moe router topk softmax experiment

Browse Source
This commit is contained in:
the magician
2026-07-15 04:21:03 +08:00
parent eea0653ace
commit cb962589e9
1 changed files with 6 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
qwen3_6_scripts/qwen3_5.py
View File

@@ -849,13 +849,12 @@ class Qwen3_5MoeSparseBlock(nn.Module):
Output is partial (pre-all-reduce), same contract as FusedMoE Output is partial (pre-all-reduce), same contract as FusedMoE
with reduce_results=False. with reduce_results=False.
""" """
# Routing: top-k logits -> softmax over selected experts. # Routing: softmax -> topk -> renormalise
# This is mathematically equivalent to softmax(all experts) -> top-k -> routing_weights = torch.softmax(router_logits.float(), dim=-1)
# renormalise, but avoids a full (T, num_experts) softmax in decode. topk_weights, topk_ids = torch.topk(
topk_logits, topk_ids = torch.topk( routing_weights, self.top_k, dim=-1) # (T, top_k)
router_logits.float(), self.top_k, dim=-1) # (T, top_k) topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
topk_weights = torch.softmax( topk_weights = topk_weights.to(hidden_states.dtype)
topk_logits, dim=-1).to(hidden_states.dtype) # (T, top_k)
w13 = self.experts.w13_weight # (E, 2*I, H) w13 = self.experts.w13_weight # (E, 2*I, H)
w2 = self.experts.w2_weight # (E, H, I) w2 = self.experts.w2_weight # (E, H, I)