optimize moe tiny batch tokenwise path

qwen3_6_scripts/qwen3_5.py

@@ -47,6 +47,8 @@ logger = init_logger(__name__)
 _ENGINEX_PROFILE_ENABLED = os.getenv("ENGINEX_PROFILE_DECODE", "0") == "1"
 _ENGINEX_PROFILE_EVERY = int(os.getenv("ENGINEX_PROFILE_EVERY", "32"))
 _ENGINEX_PROFILE_SYNC = os.getenv("ENGINEX_PROFILE_SYNC", "1") != "0"
+_ENGINEX_MOE_TINY_IMPL = os.getenv("ENGINEX_MOE_TINY_IMPL", "tokenwise")
+_ENGINEX_MOE_TINY_MAX = int(os.getenv("ENGINEX_MOE_TINY_MAX", "4"))
 _enginex_profile_stats: Dict[str, List[float]] = {}
 _enginex_profile_steps = 0
 _enginex_profile_mode = "unknown"
@@ -853,11 +855,35 @@ class Qwen3_5MoeSparseBlock(nn.Module):
 
                 out = (expert_out * ws.unsqueeze(-1)).sum(0, keepdim=True).to(
                     hidden_states.dtype)                           # (1, H)
-        elif T <= 2:
-            # Fast path: tiny decode batch. With max_num_seqs=2, normal decode
-            # often has T=2 and should not fall back to the Python per-expert
-            # loop used for long prefill chunks.
-            with _enginex_profile("moe.routed_tiny_batch_experts"):
+        elif T <= _ENGINEX_MOE_TINY_MAX and _ENGINEX_MOE_TINY_IMPL == "tokenwise":
+            # Fast path: tiny decode batch. Compute each token with the proven
+            # T==1 large-F.linear path, avoiding the per-expert Python loop.
+            # This usually beats batched bmm for very small T because the first
+            # projection becomes T larger GEMMs instead of T*K tiny GEMMs.
+            with _enginex_profile("moe.routed_tiny_tokenwise_experts"):
+                H = hidden_states.shape[-1]
+                pieces = []
+                for t in range(T):
+                    eids = topk_ids[t]
+                    ws = topk_weights[t].to(hidden_states.dtype)
+                    w13_sel = w13[eids]                         # (K, 2*I, H)
+                    w2_sel = w2[eids]                           # (K, H, I)
+
+                    gate_up = F.linear(
+                        hidden_states[t:t + 1],
+                        w13_sel.reshape(-1, H),
+                    ).view(self.top_k, -1)                      # (K, 2*I)
+                    gate, up = gate_up.chunk(2, dim=-1)
+                    act = F.silu(gate) * up
+                    expert_out = torch.bmm(
+                        w2_sel, act.unsqueeze(-1)).squeeze(-1)  # (K, H)
+                    pieces.append((expert_out * ws.unsqueeze(-1)).sum(0))
+
+                out = torch.stack(pieces, dim=0).to(hidden_states.dtype)
+        elif T <= _ENGINEX_MOE_TINY_MAX:
+            # Alternative tiny decode batch implementation. Kept for A/B
+            # testing via ENGINEX_MOE_TINY_IMPL=bmm.
+            with _enginex_profile("moe.routed_tiny_bmm_experts"):
                 H = hidden_states.shape[-1]
                 flat_eids = topk_ids.reshape(-1)                  # (T*K,)
                 flat_ws = topk_weights.reshape(-1).to(hidden_states.dtype)