analytical-viz: partial batch (B) scaffold — TopologyConfig field + KV / early stages

WIP toward batch-size support. This first commit is behavior-preserving
at the default B=1 (every new multiplication is by max(1, cfg.topo.b) =
1 today) so all existing tests pass. Follow-up commits will:
  - scale the remaining stages (S6/S7/S8/S9/S10 + C1/C2/C3 + FFN + FFN AR)
  - add a batch selectbox to the sidebar
  - forward b through auto_suggest / auto_explore / auto_hardware

Changes so far:
  - TopologyConfig: new b: int = 1 field (batch size).
  - memory_layout.per_pe_kv_cache_bytes: * B (each concurrent request
    keeps its own KV cache slice).
  - stage_rmsnorm / stage_wq / stage_wkv / stage_kv_append /
    _per_hop_qkT_pv: FLOPs and activation memory scaled by B; weight
    bytes stay fixed (weights shared across the batch).

Verified: 24 pytest tests still pass at default B=1.

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

tests/analytical_visualization/memory_layout.py

@@ -74,17 +74,19 @@ def per_pe_kv_cache_bytes(cfg: FullConfig) -> int:
       TP > H_kv, each KV head is duplicated across TP/H_kv ranks so
       per-PE storage doesn't shrink below 1 head.
     - PP: only layers_per_stage layers stored per PE.
+    - Batch (B): each concurrent request keeps its own KV cache slice.
     """
     m = cfg.model
     pp = cfg.topo.pp
     tp = cfg.topo.tp
+    B = max(1, cfg.topo.b)
     if cfg.topo.kv_shard_mode == "replicate":
         hkv_per_pe_bytes = max(1.0, m.h_kv / tp)
     else:
         hkv_per_pe_bytes = m.h_kv / tp
     layers_per_stage = (m.layers + pp - 1) // pp
     per_layer = 2 * cfg.topo.s_local * hkv_per_pe_bytes * m.d_head * m.bytes_per_elem
-    return int(per_layer * layers_per_stage)
+    return int(per_layer * layers_per_stage * B)
 
 
 def per_pe_transient_bytes(cfg: FullConfig) -> int:

tests/analytical_visualization/model_config.py

@@ -43,6 +43,7 @@ class TopologyConfig:
     pp: int = 1                 # pipeline stages (layer shard)
     dp: int = 1                 # data-parallel replicas
     ep: int = 1                 # expert-parallel degree (MoE only)
+    b: int = 1                  # batch size (concurrent requests per PP stage)
     s_kv: int = 4096
     mode: str = "decode"        # "decode" or "prefill"
     kv_shard_mode: str = "split"  # "split" or "replicate" (used when TP > H_kv)

tests/analytical_visualization/stage_latencies.py

@@ -50,8 +50,10 @@ def stage_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  # load x + load weight
+    B = max(1, cfg.topo.b)
-    flops = 4 * T_q * d
+    # Activation memory (x) 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
     cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
     vis, bnd = _visible(cmp_s, mem_s, 0)
@@ -76,9 +78,11 @@ def stage_wq(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 * d * (hq_per_pe * dh)
+    # FLOPs scale with batch; weight bytes fixed (shared across batch).
+    flops = 2 * B * T_q * d * (hq_per_pe * dh)
     weight_B = d * (hq_per_pe * dh) * b
     cmp_s, mem_s = _gemm_time(flops, weight_B, cfg)
     vis, bnd = _visible(cmp_s, mem_s, 0)
@@ -98,9 +102,10 @@ def stage_wkv(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)
     hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
     dh = cfg.model.d_head
-    flops_one = 2 * T_q * d * (hkv_per_pe * dh)
+    flops_one = 2 * B * T_q * d * (hkv_per_pe * dh)
     weight_B_one = d * (hkv_per_pe * dh) * b
     flops = 2 * flops_one
     weight_B = 2 * weight_B_one
@@ -120,9 +125,10 @@ def stage_wkv(cfg: FullConfig) -> StageCost:
 def stage_kv_append(cfg: FullConfig) -> StageCost:
     T_q = cfg.topo.T_q
     b = cfg.model.bytes_per_elem
+    B = max(1, cfg.topo.b)
     hkv_per_pe = max(1, cfg.model.h_kv // cfg.topo.tp)
     dh = cfg.model.d_head
-    bytes_ = 2 * T_q * hkv_per_pe * dh * b
+    bytes_ = 2 * B * T_q * hkv_per_pe * dh * b
     mem_s = bytes_ / cfg.machine.bw_hbm
     return StageCost(
         name="S4 KV cache append",
@@ -141,7 +147,8 @@ def _per_hop_qkT_pv(cfg: FullConfig) -> tuple[float, str]:
     S_local = cfg.topo.s_local
     dh = cfg.model.d_head
     hq_per_pe = cfg.h_q_per_pe
-    flops = 2 * T_q * S_local * dh * hq_per_pe
+    B = max(1, cfg.topo.b)
+    flops = 2 * B * T_q * S_local * dh * hq_per_pe
     cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
     formula = f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops:.2g} FLOPs/hop"
     return cmp_s, formula