Views
No views yet
d_model) branch running full sliding window attention.B, T, num_blocks). The top-$k$ tokens for each expert are selected independently for each item in the batch based on k = capacity_factor * T / E. This ensures there is absolutely no cross-sequence leakage during routing, fully supporting training with BATCH_SIZE > 1.d_thin to save compute.torch.gather on the batched token sequences, tokens are processed cleanly.d_model and accumulated back into the residual stream via torch.scatter_add_ using the original batch-aware routing gates.1graph TD
2 Input[Input Tokens<br/>B=8, T=1024] --> Embed[Token + Positional Embedding]
3
4 subgraph MoEP-MSIT Layer
5 Embed --> GlobalBlock[Global Branch Block<br/>Full Attention]
6 GlobalBlock --> PreNorm[LayerNorm Pre-MoE]
7 PreNorm --> Router[Batch-Safe Expert Choice Router<br/>top-k selection per sequence]
8
9 PreNorm --> ShrinkProj[Down Projection<br/>d_model to d_thin]
10
11 Router -- Selects Tokens for E1 --> E1[Expert 1: SW=64]
12 Router -- Selects Tokens for E2 --> E2[Expert 2: SW=64]
13 Router -- Selects Tokens for E3 --> E3[Expert 3: SW=16]
14 Router -- Selects Tokens for E4 --> E4[Expert 4: SW=16]
15 Router -- Selects Tokens for E5 --> E5[Expert 5: SW=8]
16 Router -- Selects Tokens for E6 --> E6[Expert 6: SW=8]
17 Router -- Selects Tokens for E7 --> E7[Expert 7: SW=4]
18 Router -- Selects Tokens for E8 --> E8[Expert 8: SW=4]
19
20 ShrinkProj -- torch.gather --> E1 & E2 & E3 & E4 & E5 & E6 & E7 & E8
21
22 E1 & E2 & E3 & E4 & E5 & E6 & E7 & E8 --> GrowProj[Up Projection<br/>d_thin to d_model]
23 GrowProj --> Gate[Multiply by Gate Scores]
24 Gate -- torch.scatter_add_ --> Add[Residual Add]
25 GlobalBlock --> Add
26 end
27
28 Add --> NextLayer[Next Layer or Final Norm]