MoE Top-K Gating

题面

对形状 [M,E] 的 logits，每行取 Top‑K 索引与对应值，对值做 softmax 得到混合权重，分别写入 topk_indices 与 topk_weights。

Implementation Requirements

• External libraries are not permitted
• The solve function signature must remain unchanged
• 结果写入 topk_weights 与 topk_indices

Examples

见页面示例（K=2 时权重约为 [0.2689,0.7311] 等）。

Constraints

• 1 ≤ M ≤ 10,000；1 ≤ E ≤ 256；1 ≤ k ≤ E
• logits 为 float32；indices 为 int32；Performance: M=1,024, k=2