analytical-viz: complete batch (B) scaling in remaining stages + comm

Finishes the batch scaffold. All stages that produce per-request work
now scale flops / activation-memory / comm-bytes by B; weight bytes stay
fixed (shared across the batch).

Attention block, previously partial:
  - S6 stage_softmax: elems + bytes × B
  - S8 stage_merge: flops × B; O/m/l AR bytes M × B
  - S9 stage_normalize: flops × B
  - S10 stage_wo: flops × B
  - C1 comm_cp_ring: decode O/m/l AR M × B; prefill K/V ring + Q+O/m/l
    variants both × B
  - C2 comm_tp_allreduce (W_O output): bytes × B
  - C3 comm_kv_split_allreduce (head-split scores): bytes_per_hop × B

FFN block, previously untouched:
  - F1 stage_ffn_rmsnorm: activation × B (weight fixed)
  - F2 stage_ffn_gate (via _ffn_gemm): flops × B
  - F3 stage_ffn_up: same
  - F4 stage_ffn_swiglu: elems + flops × B
  - F5 stage_ffn_down: flops × B
  - CF1 comm_ffn_allreduce (batched FFN output): bytes × B

Verified with a smoke check on Qwen 3 8B / 128K decode:
  B=1  per-layer visible latency:  363 us
  B=8                                716 us  (sub-linear — many stages
                                              stay weight-bound)
  B=64                              4023 us  (approaching linear scaling
                                              as batch dominates)

All 24 pytest tests still pass at default B=1 (backward compat).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

tests/analytical_visualization/stage_latencies.py

@@ -195,7 +195,8 @@ def stage_softmax(cfg: FullConfig) -> StageCost:
     S_local = cfg.topo.s_local
     hq_per_pe = cfg.h_q_per_pe
     b = cfg.model.bytes_per_elem
-    elems = hq_per_pe * T_q * S_local
+    B = max(1, cfg.topo.b)
+    elems = B * hq_per_pe * T_q * S_local
     bytes_ = elems * b * 2
     mem_s_per_hop = bytes_ / cfg.machine.bw_hbm
     passes = _cp_compute_passes(cfg)
@@ -254,7 +255,8 @@ def stage_merge(cfg: FullConfig) -> StageCost:
     hq_per_pe = cfg.h_q_per_pe
     dh = cfg.model.d_head
     b = cfg.model.bytes_per_elem
-    flops = 6 * T_q * hq_per_pe * dh * max(0, cfg.topo.cp - 1)
+    B = max(1, cfg.topo.b)
+    flops = 6 * B * T_q * hq_per_pe * dh * max(0, cfg.topo.cp - 1)
     cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
 
     comm_s = 0.0
@@ -264,8 +266,8 @@ def stage_merge(cfg: FullConfig) -> StageCost:
 
     if cfg.topo.mode == "decode" and cfg.topo.cp > 1:
         cp = cfg.topo.cp
-        # (O + m + l) bytes per rank
+        # (O + m + l) bytes per rank — scales with B (per-request partials).
-        M = (T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)
+        M = B * ((T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b))
         if cfg.topo.cp_placement == "pe":
             bw, alpha, tier = cfg.machine.bw_intra, cfg.machine.alpha_intra, "intra-cube"
         elif cfg.topo.sips_used > 1:
@@ -317,7 +319,8 @@ def stage_normalize(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     hq_per_pe = cfg.h_q_per_pe
     dh = cfg.model.d_head
-    flops = T_q * hq_per_pe * dh
+    B = max(1, cfg.topo.b)
+    flops = B * T_q * hq_per_pe * dh
     cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
     return StageCost(
         name="S9 normalize O/l",
@@ -334,9 +337,10 @@ def stage_wo(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
     hq_per_pe = cfg.h_q_per_pe
     dh = cfg.model.d_head
-    flops = 2 * T_q * (hq_per_pe * dh) * d
+    flops = 2 * B * T_q * (hq_per_pe * dh) * d
     weight_B = (hq_per_pe * dh) * d * b
     cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
     vis, bnd = _visible(cmp_s, mem_s, 0)
@@ -376,12 +380,13 @@ def comm_cp_ring(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     dh = cfg.model.d_head
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
 
     # ── Decode: single O/m/l all-reduce after local compute ──
     if cfg.topo.mode == "decode":
         cp = cfg.topo.cp
-        # per-rank O + m + l bytes
+        # per-rank O + m + l bytes — scales with B (per-request partials).
-        M = (T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)
+        M = B * ((T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b))
         if cfg.topo.cp_placement == "pe":
             bw = cfg.machine.bw_intra
             alpha = cfg.machine.alpha_intra
@@ -424,7 +429,7 @@ def comm_cp_ring(cfg: FullConfig) -> StageCost:
 
     # ── Prefill: per-hop ring (K/V or Q+O/m/l) concurrent with S5-S7 ──
     if cfg.topo.cp_ring_variant == "qoml":
-        M_KV = (2 * T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b)
+        M_KV = B * ((2 * T_q * hq_per_pe * dh * b) + (2 * T_q * hq_per_pe * b))
         variant_desc = "Q+O/m/l ring"
         formula_bytes = (
             f"2*T_q*(H_q/TP)*d_h*b + 2*T_q*(H_q/TP)*b "
@@ -432,7 +437,7 @@ def comm_cp_ring(cfg: FullConfig) -> StageCost:
             f"= {M_KV} B/hop"
         )
     else:
-        M_KV = 2 * S_local * hkv_per_pe * dh * b
+        M_KV = 2 * B * S_local * hkv_per_pe * dh * b
         variant_desc = "K/V ring"
         formula_bytes = (
             f"2*S_local*(H_kv/TP)*d_h*b "
@@ -508,7 +513,8 @@ def comm_tp_allreduce(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
-    bytes_ = T_q * d * b
+    B = max(1, cfg.topo.b)
+    bytes_ = B * T_q * d * b
     tp = cfg.topo.tp
     tier = cfg.topo.tp_link_tier()  # "intra" | "inter" | "intersip"
     if tier == "intra":
@@ -559,8 +565,9 @@ def comm_kv_split_allreduce(cfg: FullConfig) -> StageCost:
     S_local = cfg.topo.s_local
     hq_per_pe = cfg.h_q_per_pe
     b = cfg.model.bytes_per_elem
-    # score bytes per rank per hop
+    B = max(1, cfg.topo.b)
-    bytes_per_hop = hq_per_pe * T_q * S_local * b
+    # score bytes per rank per hop — scales with B (per-request scores).
+    bytes_per_hop = B * hq_per_pe * T_q * S_local * b
     # split-group AllReduce (intra-cube assumed if group fits)
     bw = cfg.machine.bw_intra
     alpha = cfg.machine.alpha_intra
@@ -633,8 +640,10 @@ def stage_ffn_rmsnorm(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
-    bytes_ = T_q * d * b * 2
+    B = max(1, cfg.topo.b)
-    flops = 4 * T_q * d
+    # Activation memory scales with B; weight (once) is fixed.
+    bytes_ = B * T_q * d * b + T_q * d * b
+    flops = 4 * B * T_q * d
     mem_s = bytes_ / cfg.machine.bw_hbm
     return StageCost(
         name="F1 RMSNorm (pre-FFN)",
@@ -651,7 +660,9 @@ def _ffn_gemm(cfg: FullConfig, name: str, ffn_per_pe: int) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
-    flops = 2 * T_q * d * ffn_per_pe
+    B = max(1, cfg.topo.b)
+    # FLOPs scale with batch; weight (shared) is fixed.
+    flops = 2 * B * T_q * d * ffn_per_pe
     weight_B = d * ffn_per_pe * b
     cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
     vis, bnd = _visible(cmp_s, mem_s, 0)
@@ -681,8 +692,9 @@ def stage_ffn_swiglu(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
     b = cfg.model.bytes_per_elem
-    bytes_ = 3 * T_q * ffn_per_pe * b
+    B = max(1, cfg.topo.b)
-    flops = 3 * T_q * ffn_per_pe  # gate * silu(up) approximated as 3 flops/elt
+    bytes_ = 3 * B * T_q * ffn_per_pe * b
+    flops = 3 * B * T_q * ffn_per_pe  # gate * silu(up) approximated as 3 flops/elt
     mem_s = bytes_ / cfg.machine.bw_hbm
     return StageCost(
         name="F4 SwiGLU act",
@@ -700,8 +712,9 @@ def stage_ffn_down(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
     ffn_per_pe = max(1, cfg.model.ffn_dim // (cfg.ffn_shard_divisor * max(1, cfg.topo.ep)))
-    flops = 2 * T_q * ffn_per_pe * d
+    flops = 2 * B * T_q * ffn_per_pe * d
     weight_B = ffn_per_pe * d * b
     cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
     vis, bnd = _visible(cmp_s, mem_s, 0)
@@ -730,7 +743,9 @@ def comm_ffn_allreduce(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     d = cfg.model.hidden
     b = cfg.model.bytes_per_elem
-    bytes_ = T_q * d * b
+    B = max(1, cfg.topo.b)
+    # AR reduces the batched FFN output — bytes scale with B.
+    bytes_ = B * T_q * d * b
     # Choose BW/alpha tier based on scope (rough): TP=intra-cube (or inter-cube),
     # +CP=inter-SIP possible, +DP=inter-SIP always.
     if "DP" in scope or "CP" in scope: