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>
This commit is contained in:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user