add moe bucket profiling and record bottlenecks

qwen3_6_scripts/qwen3_5.py

@@ -51,12 +51,79 @@ import os as _os_prof
 import sys as _sys_prof
 import time as _time_prof
 _PROF_TRACE = _os_prof.environ.get("PROF_TRACE", "0") == "1"
+_PROF_MOE_SUMMARY = _os_prof.environ.get("PROF_MOE_SUMMARY", "0") == "1"
+_PROF_MOE_SYNC = _os_prof.environ.get("PROF_MOE_SYNC", "1") != "0"
+_PROF_MOE_INTERVAL = int(_os_prof.environ.get("PROF_MOE_INTERVAL", "400"))
+_PROF_MOE_STATS = {}
 _t0 = _time_prof.time()
 def _prof(msg):
     if _PROF_TRACE:
         print(f"[PROF] {_time_prof.time()-_t0:7.2f}s {msg}", file=_sys_prof.stderr, flush=True)
 
 
+def _prof_moe_time():
+    if not _PROF_MOE_SUMMARY:
+        return None
+    if _PROF_MOE_SYNC:
+        try:
+            torch.cuda.synchronize()
+        except Exception:
+            pass
+    return _time_prof.perf_counter()
+
+
+def _prof_moe_bucket(tokens):
+    return "decode" if int(tokens) == 1 else "prefill"
+
+
+def _prof_moe_stats(bucket):
+    stats = _PROF_MOE_STATS.get(bucket)
+    if stats is None:
+        stats = {
+            "calls": 0,
+            "tokens": 0,
+            "router_gate": 0.0,
+            "routed": 0.0,
+            "shared": 0.0,
+            "all_reduce": 0.0,
+            "total": 0.0,
+        }
+        _PROF_MOE_STATS[bucket] = stats
+    return stats
+
+
+def _prof_moe_add(name, start, bucket):
+    if not _PROF_MOE_SUMMARY or start is None:
+        return None
+    end = _prof_moe_time()
+    _prof_moe_stats(bucket)[name] += end - start
+    return end
+
+
+def _prof_moe_finish(total_start, tokens, bucket):
+    if not _PROF_MOE_SUMMARY or total_start is None:
+        return
+    end = _prof_moe_time()
+    stats = _prof_moe_stats(bucket)
+    stats["total"] += end - total_start
+    stats["calls"] += 1
+    stats["tokens"] += int(tokens)
+    if stats["calls"] % max(1, _PROF_MOE_INTERVAL) != 0:
+        return
+    calls = stats["calls"]
+    parts = [
+        f"pid={_os_prof.getpid()}",
+        f"bucket={bucket}",
+        f"calls={calls}",
+        f"tokens={stats['tokens']}",
+    ]
+    for name in ("router_gate", "routed", "shared", "all_reduce", "total"):
+        avg_ms = stats[name] * 1000.0 / calls
+        parts.append(f"{name}_avg_ms={avg_ms:.3f}")
+    print("[PROF_MOE_SUMMARY] " + " ".join(parts),
+          file=_sys_prof.stderr, flush=True)
+
+
 # ---------------------------------------------------------------------------
 # Pure-PyTorch DeltaNet kernels (fallbacks from transformers 5.2.0)
 # ---------------------------------------------------------------------------
@@ -960,8 +1027,14 @@ class Qwen3_5MoeSparseBlock(nn.Module):
         return out
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        _moe_total_t = _prof_moe_time()
+        _tokens = hidden_states.shape[0]
+        _prof_bucket = _prof_moe_bucket(_tokens)
+        _section_t = _moe_total_t
         router_logits, _ = self.gate(hidden_states)
+        _section_t = _prof_moe_add("router_gate", _section_t, _prof_bucket)
         routed_out = self._pure_pytorch_experts(hidden_states, router_logits)
+        _section_t = _prof_moe_add("routed", _section_t, _prof_bucket)
         _prof(f"  MoE routed-experts done tokens={hidden_states.shape[0]}")
 
         gate_up, _ = self.shared_expert_gate_up(hidden_states)
@@ -970,12 +1043,15 @@ class Qwen3_5MoeSparseBlock(nn.Module):
         # Scalar sigmoid gate (Qwen2-MoE / Qwen3.5-MoE style)
         gate_score, _ = self.shared_expert_gate(hidden_states)  # (T, 1)
         shared_out = shared_out * torch.sigmoid(gate_score)
+        _section_t = _prof_moe_add("shared", _section_t, _prof_bucket)
 
         out = routed_out + shared_out
         if self.experts.tp_size > 1:
             _prof(f"  MoE all_reduce start tp_size={self.experts.tp_size}")
             out = tensor_model_parallel_all_reduce(out)
             _prof(f"  MoE all_reduce done")
+        _section_t = _prof_moe_add("all_reduce", _section_t, _prof_bucket)
+        _prof_moe_finish(_moe_total_t, _tokens, _prof_bucket)
         return out