add comments

mayank31398 · Aug 9, 2024 · ef91329 · ef91329
1 parent d40db0e
commit ef91329
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/kernel_hyperdrive/scattermoe/torch_implementation.py b/kernel_hyperdrive/scattermoe/torch_implementation.py
@@ -131,6 +131,10 @@ def _compute_experts(
             router_weights, selected_experts
         )
 
+        # batch_index -> (num_tokens * top_k)
+        # batch_gates -> (num_tokens * top_k)
+        # num_experts_per_token -> (num_experts)
+
         expert_inputs = hidden_states[batch_index]
 
         hidden_states = self.c_fc(expert_inputs, num_experts_per_token)