import torch
import torch.nn as nn
import torch.nn.functional as F
class SparseMoERouter(nn.Module):
def __init__(self, d_model: int, num_experts: int = 8, top_k: int = 2):
super().__init__()
self.num_experts = num_experts
self.top_k = top_k
self.gate = nn.Linear(d_model, num_experts, bias=False)
# Parallel expert networks
self.experts = nn.ModuleList([
nn.Sequential(
nn.Linear(d_model, d_model * 4),
nn.SiLU(),
nn.Linear(d_model * 4, d_model)
) for _ in range(num_experts)
])
def forward(self, x: torch.Tensor):
# x: [Batch, Seq_len, d_model]
batch_size, seq_len, d_model = x.shape
x_flat = x.view(-1, d_model)
# Router logits & Top-K gating
logits = self.gate(x_flat)
weights, indices = torch.topk(F.softmax(logits, dim=-1), self.top_k, dim=-1)
out = torch.zeros_like(x_flat)
for i, expert in enumerate(self.experts):
# Find tokens dispatched to expert i
batch_idx, kth_idx = torch.where(indices == i)
if len(batch_idx) > 0:
expert_in = x_flat[batch_idx]
expert_out = expert(expert_in)
gating_weight = weights[batch_idx, kth_idx].unsqueeze(-1)
out[batch_idx] += expert_out * gating_weight
return out.view(batch_size, seq_len, d_model)