Bộ Định Tuyến Chuyên Gia Top-K (MoE Router) Mixtral 8x7B

Mục tiêu kiến thức

• Trong mô hình Thưa (Sparse Mixture-of-Experts như Mixtral 8x7B hay DeepSeek-V2/V3), thay vì dùng một tầng FFN chung, mô hình có E mạng chuyên gia (experts).
• Với mỗi token x, một bộ định tuyến (Router / Gate) tính toán điểm tương thích:

H(x) = x W_g (W_g ∈ R^{d_model × E})

• Bộ định tuyến chọn ra K chuyên gia có điểm số cao nhất (ví dụ Mixtral chọn K = 2 trong tổng số E = 8 chuyên gia).
• Trọng số định tuyến sau khi chọn top-k được chuẩn hóa lại bằng Softmax:

p_i = Softmax(TopK(H(x)))_i

• Token chỉ được gửi đến K chuyên gia này, giảm 75% chi phí tính toán FLOPs tại mỗi bước suy luận!

Yêu cầu

Xây dựng lớp MoETopKRouter(torch.nn.Module):

class MoETopKRouter(nn.Module):
    def __init__(self, d_model: int, num_experts: int, top_k: int = 2):
        pass

    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        # Returns: (topk_weights, topk_indices, router_logits)
        pass

• Đầu vào: x shape (B, S, d_model).
• Đầu ra:
• topk_weights: Tensor shape (B, S, top_k) chứa trọng số đã chuẩn hóa Softmax (∑_{top_k} p = 1.0).
• topk_indices: Tensor shape (B, S, top_k) chứa chỉ số của K chuyên gia được chọn.
• router_logits: Tensor shape (B, S, num_experts) chứa logits thô x W_g.

Ràng buộc

• Thời gian chạy tối đa: 3000ms.
• Giới hạn bộ nhớ: 512MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.