perf(moe): speed up grouped MoE routing (NVIDIA#20)

## Summary

This PR tightens the routing/unrouting path in
`Qwen3VLMoeTextExpertsGroupedMm`:

- keep routed token indices as a 1D vector instead of expanding them
across the hidden dimension
- use `index_select` / `index_add_` for route and combine instead of
hidden-size-expanded `gather` / `scatter_add_`
- update the standalone MoE comparison harness so it calls the current
expert API and checks grouped output against the PyTorch reference

The old path was doing a lot of unnecessary index traffic. That is
mostly invisible for smaller routing configs, but it becomes very
noticeable when `top_k` and the expert count go up.

## Benchmarks

Environment: RTX 6000 Ada, `torch==2.10.0+cu128`, BF16 unless noted.

Focused MoE harness:

```text
python -m cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe_test

naive:      13.81 ms
optimized:   3.92 ms
speedup:     3.53x
relative diff: 0.00531
```

Synthetic decoder stack throughput, including attention + norms + MoE,
old grouped routing vs this PR:

```text
E128_K8_I768, 8 layers, seq=1024
old: 134.58 ms,  7,608.7 tok/s
new:  42.16 ms, 24,289.3 tok/s
speedup: 3.19x, +219.2%

E128_K8_I768, 12 layers, seq=1024
old: 136.89 ms,  7,480.3 tok/s
new:  63.33 ms, 16,170.2 tok/s
speedup: 2.16x, +116.2%

E128_K8_I1408, 8 layers, seq=1024
old: 103.66 ms,  9,878.7 tok/s
new:  52.81 ms, 19,391.2 tok/s
speedup: 1.96x, +96.3%
```

For the default-ish `E=60, top_k=4` shape, I did not see a meaningful
total-throughput improvement. The gain scales with routing pressure;
high-expert / high-`top_k` configurations are where the old
expanded-index route/combine path becomes expensive.

## Tests

```text
LD_LIBRARY_PATH= uv run --no-sync python -m ruff check \
  cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe.py \
  cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe_test.py

LD_LIBRARY_PATH= uv run --no-sync python -m cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe_test

LD_LIBRARY_PATH= uv run --no-sync python -m cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe_bench \
  --compare --num-tokens 256 --num-iters 10 --num-warmup 3 --no-compile --dtype fp32

LD_LIBRARY_PATH= uv run --no-sync python -m cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe_bench \
  --compare --num-tokens 2048 --num-iters 20 --num-warmup 5 --no-compile --dtype bf16

LD_LIBRARY_PATH= uv run --no-sync python -m cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe_bench \
  --num-tokens 1024 --num-iters 5 --num-warmup 2 --dtype bf16 --backward
```

Signed-off-by: Shivanjan Chakravorty <shivanjanc@nvidia.com>
Co-authored-by: lfengad <liangf@nvidia.com>

Author: Glitchfix

cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe.py

@@ -104,15 +104,14 @@ def forward(
         sentinel = torch.tensor([num_tokens], device=hidden_states.device)  # for padding slots
         token_indices_ext = torch.cat([token_indices_sorted, sentinel])
         combined_indices = token_indices_ext[permuted_indices.long()]
-        combined_indices = combined_indices.unsqueeze(-1).expand(-1, dim)
 
         # Pad scores with a zero sentinel so padding slots contribute nothing
         scores_ext = torch.cat([topk_scores_sorted, topk_scores_sorted.new_zeros(1)])
         combined_scores = scores_ext[permuted_indices.long()]  # [padded_len]
 
         # Single gather (with a zero-padded sentinel row)
         input_padded = torch.cat([hidden_states, hidden_states.new_zeros(1, dim)])
-        routed_input = input_padded.gather(dim=0, index=combined_indices)
+        routed_input = input_padded.index_select(dim=0, index=combined_indices)
 
         # Run experts
         routed_output = _run_experts_grouped_mm(
@@ -125,7 +124,7 @@ def forward(
         )
 
         output_padded = torch.zeros_like(input_padded)
-        output_padded.scatter_add_(dim=0, index=combined_indices, src=routed_output)
+        output_padded.index_add_(dim=0, index=combined_indices, source=routed_output)
         return output_padded[:-1]
 
     def _reorder_tokens(
@@ -220,9 +219,8 @@ def forward(
                 assert weighted_output.dtype == hidden_states.dtype
                 next_states.index_add_(0, token_idx, weighted_output)
         else:
-            hidden_states = hidden_states.repeat(self.num_experts, 1)  # [num_experts*num_tokens,hidden_size]
-            hidden_states = hidden_states.view(
-                self.num_experts, -1, self.hidden_size
+            hidden_states = hidden_states.unsqueeze(0).expand(
+                self.num_experts, -1, -1
             )  # [num_experts,num_tokens,hidden_size]
             gate_up = torch.bmm(hidden_states, self.gate_up_proj)  # [num_experts,num_tokens,2*moe_intermediate_size]
             gate, up = gate_up.chunk(

cosmos_framework/model/vfm/vlm/qwen3_vl_moe/moe_test.py

@@ -12,18 +12,26 @@
 from cosmos_framework.model.vfm.vlm.qwen3_vl_moe.moe import create_text_experts
 
 
-def run_moe(mod: nn.Module, hidden_states: torch.Tensor, topk_scores: torch.Tensor, expert_indices: torch.Tensor):
+def run_moe(
+    mod: nn.Module,
+    hidden_states: torch.Tensor,
+    topk_scores: torch.Tensor,
+    expert_indices: torch.Tensor,
+    num_tokens_per_expert: torch.Tensor,
+):
     num_warmup_iterations = 10
     num_timing_iterations = 100
 
     for _ in range(num_warmup_iterations):
         with torch.no_grad():
-            output = mod(hidden_states, topk_scores, expert_indices)
+            output = mod(hidden_states, topk_scores, expert_indices, num_tokens_per_expert)
+    torch.cuda.synchronize()
 
     start_time = time.time()
     for _ in range(num_timing_iterations):
         with torch.no_grad():
-            output = mod(hidden_states, topk_scores, expert_indices)
+            output = mod(hidden_states, topk_scores, expert_indices, num_tokens_per_expert)
+    torch.cuda.synchronize()
     end_time = time.time()
 
     time_taken = (end_time - start_time) / num_timing_iterations
@@ -46,7 +54,7 @@ def main():
     control = create_text_experts(config, implementation_type="naive")
     exp = create_text_experts(config, implementation_type="grouped_mm")
 
-    control.init_weights()
+    control.init_weights(torch.device("cpu"))
     exp.load_state_dict(control.state_dict())
 
     control = control.to(device="cuda", dtype=torch.bfloat16)
@@ -58,31 +66,36 @@ def main():
         dtype=torch.bfloat16,
         device="cuda",
     )
-    topk_scores = torch.randn(
+    topk_scores = torch.rand(
         num_tokens,
         config.num_experts_per_tok,
         dtype=torch.bfloat16,
         device="cuda",
     )
     topk_scores = topk_scores / topk_scores.sum(dim=-1, keepdim=True)
-    expert_indices = torch.randint(
-        0,
-        config.num_experts,
-        (num_tokens, config.num_experts_per_tok),
-        dtype=torch.int64,
-        device="cuda",
+    expert_indices = torch.stack(
+        [torch.randperm(config.num_experts, device="cuda")[: config.num_experts_per_tok] for _ in range(num_tokens)]
+    ).to(torch.int64)
+    num_tokens_per_expert = torch.histc(
+        expert_indices.to(dtype=torch.int32).view(-1),
+        bins=config.num_experts,
+        min=0,
+        max=config.num_experts - 1,
     )
 
     print(
         f"hidden_states: {hidden_states.norm().detach().cpu().item()} {hidden_states.shape} {hidden_states.dtype} {hidden_states.device}"
     )
 
-    control_output, control_time_taken = run_moe(control, hidden_states, topk_scores, expert_indices)
-    exp_output, exp_time_taken = run_moe(exp, hidden_states, topk_scores, expert_indices)
+    control_output, control_time_taken = run_moe(
+        control, hidden_states, topk_scores, expert_indices, num_tokens_per_expert
+    )
+    exp_output, exp_time_taken = run_moe(exp, hidden_states, topk_scores, expert_indices, num_tokens_per_expert)
 
     diff = (control_output.detach().cpu() - exp_output.detach().cpu()).norm() / control_output.detach().cpu().norm()
     print(f"Diff: {diff}")
     print(f"Speedup: {control_time_taken / exp_time_taken}")
+    torch.testing.assert_close(control_output, exp_output, rtol=1e-2, atol=1e-2)
 
 
 if __name__ == "__main__":