From 5d7ef48b147a7567cb28bd9e529ac1a7176d7478 Mon Sep 17 00:00:00 2001 From: Mukesh Garg Date: Tue, 28 Jul 2026 22:21:37 -0700 Subject: [PATCH] =?UTF-8?q?analytical-viz:=20partial=20batch=20(B)=20scaff?= =?UTF-8?q?old=20=E2=80=94=20TopologyConfig=20field=20+=20KV=20/=20early?= =?UTF-8?q?=20stages?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../analytical_visualization/memory_layout.py | 4 +++- .../analytical_visualization/model_config.py | 1 + .../stage_latencies.py | 19 +++++++++++++------ 3 files changed, 17 insertions(+), 7 deletions(-) diff --git a/tests/analytical_visualization/memory_layout.py b/tests/analytical_visualization/memory_layout.py index e7c6274..40c48c6 100644 --- a/tests/analytical_visualization/memory_layout.py +++ b/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: diff --git a/tests/analytical_visualization/model_config.py b/tests/analytical_visualization/model_config.py index 3bb7df5..dd4dcda 100644 --- a/tests/analytical_visualization/model_config.py +++ b/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) diff --git a/tests/analytical_visualization/stage_latencies.py b/tests/analytical_visualization/stage_latencies.py index 488e0bd..395be88 100644 --- a/tests/analytical_visualization/stage_latencies.py +++ b/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 - flops = 4 * T_q * d + B = max(1, cfg.topo.b) + # 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