ywkang/kernbench2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

analytical-viz: batch B in every formula string + rename 'latency-optimal' → 'memory-optimal'

Browse Source 
Two related fixes.

1) Formula strings show B alongside every other symbol. Previously the
   numeric FLOPs/mem_bytes/comm_bytes were scaled by B, but the printed
   formula strings still read like B=1 amounts. That made the per-stage
   table confusing at high B: numbers moved but formulas didn't. Updated
   every attention + FFN stage's `formula` / `flops_formula` /
   `mem_formula` / `comm_formula` to include B explicitly. Weight-bytes
   lines now say "(shared across batch)" / "(weight, B-invariant)" so it's
   obvious what does and doesn't scale.

   Stages touched: S1..S10 (attention), C1/C2/C3 (attn comm), F1..F5
   (FFN), CF1 (FFN AR).

2) Renamed the button/spinner/help copy from "latency-optimal" to
   "memory-optimal" to match the smallest-fit pick semantics I put in
   place earlier. The buttons already picked smallest-fit; the labels
   just still said "latency-optimal" from the previous iteration.

   Left the Auto Hardware sensitivity chart's "latency-optimal baseline"
   label alone — that panel is a HW co-design view where "how fast can
   each HW go?" is the intended question.

Verified: 24 pytest tests still pass.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Mukesh Garg
2026-07-28 23:20:32 -07:00
parent 50aab8d591
commit 47c4f40f4d
2 changed files with 109 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/analytical_visualization/app.py
+13 -12
View File
@@ -218,8 +218,8 @@ with st.sidebar.expander("Parallelism", expanded=True):
st.session_state["cp_placement"] = _default_suggestion.cp_placement st.session_state["cp_placement"] = _default_suggestion.cp_placement
st.session_state["_pl_active_scope"] = "full" st.session_state["_pl_active_scope"] = "full"
st.rerun() st.rerun()
# Three latency-optimal apply buttons — pick a scope, load the # Three memory-optimal apply buttons — pick a scope, load the
# latency-min Pareto config into the sidebar sliders. # smallest-fit Pareto config into the sidebar sliders.
from tests.analytical_visualization.auto_explore import ( from tests.analytical_visualization.auto_explore import (
run_auto_explore as _run_auto_explore_sb, run_auto_explore as _run_auto_explore_sb,
) )
@@ -235,14 +235,14 @@ with st.sidebar.expander("Parallelism", expanded=True):
_sb_clicks[_sb_scope] = st.button( _sb_clicks[_sb_scope] = st.button(
_sb_scope_meta[_sb_scope][2], width='stretch', _sb_scope_meta[_sb_scope][2], width='stretch',
key=f"_sb_apply_{_sb_scope}", key=f"_sb_apply_{_sb_scope}",
help=f"Apply latency-optimal parallelism for " help=f"Apply memory-optimal (smallest-fit) parallelism for "
f"{_sb_scope_meta[_sb_scope][2]} scope.", f"{_sb_scope_meta[_sb_scope][2]} scope.",
) )
for _sb_scope, _sb_click in _sb_clicks.items(): for _sb_scope, _sb_click in _sb_clicks.items():
if not _sb_click: if not _sb_click:
continue continue
_sb_ia, _sb_ff, _sb_lbl = _sb_scope_meta[_sb_scope] _sb_ia, _sb_ff, _sb_lbl = _sb_scope_meta[_sb_scope]
with st.spinner(f"Finding latency-optimal parallelism ({_sb_lbl})..."): with st.spinner(f"Finding memory-optimal parallelism ({_sb_lbl})..."):
_sb_res = _run_auto_explore_sb( _sb_res = _run_auto_explore_sb(
model, _default_machine, s_kv=s_kv, mode=mode, model, _default_machine, s_kv=s_kv, mode=mode,
include_attention=_sb_ia, include_ffn=_sb_ff, b=b_batch, include_attention=_sb_ia, include_ffn=_sb_ff, b=b_batch,
@@ -553,17 +553,18 @@ tab_layout, tab_auto, tab_memory, tab_stages, tab_compare, tab_hw = st.tabs([
# ── TAB 1: Physical layout ────────────────────────────────────── # ── TAB 1: Physical layout ──────────────────────────────────────
with tab_layout: with tab_layout:
# Quick shortcut: three buttons that pick the latency-optimal Pareto # Quick shortcut: three buttons that pick the memory-optimal
# config for a chosen scope and load it into the sidebar. Layout # (smallest-fit) Pareto config for a chosen scope and load it into
# re-renders immediately on the next Streamlit run with the new config. # the sidebar. Layout re-renders immediately on the next Streamlit
# run with the new config.
from tests.analytical_visualization.auto_explore import ( from tests.analytical_visualization.auto_explore import (
run_auto_explore as _run_auto_explore_layout, run_auto_explore as _run_auto_explore_layout,
) )
st.markdown( st.markdown(
"**Quick auto-suggest layout** — pick a scope; the latency-optimal " "**Quick auto-suggest layout** — pick a scope; the memory-optimal "
"Pareto config for that scope is loaded into the sidebar and this " "(smallest-fit) Pareto config for that scope is loaded into the "
"layout redraws around it." "sidebar and this layout redraws around it."
) )
_pl1, _pl2, _pl3, _pl4 = st.columns([1, 1, 1, 1]) _pl1, _pl2, _pl3, _pl4 = st.columns([1, 1, 1, 1])
with _pl1: with _pl1:
@@ -591,7 +592,7 @@ with tab_layout:
if not _pl_clicked: if not _pl_clicked:
continue continue
_pl_ia, _pl_ff, _pl_lbl = _pl_scope_meta[_pl_scope] _pl_ia, _pl_ff, _pl_lbl = _pl_scope_meta[_pl_scope]
with st.spinner(f"Finding latency-optimal parallelism ({_pl_lbl})..."): with st.spinner(f"Finding memory-optimal parallelism ({_pl_lbl})..."):
_pl_res = _run_auto_explore_layout( _pl_res = _run_auto_explore_layout(
model, machine, s_kv=s_kv, mode=mode, model, machine, s_kv=s_kv, mode=mode,
include_attention=_pl_ia, include_ffn=_pl_ff, b=b_batch, include_attention=_pl_ia, include_ffn=_pl_ff, b=b_batch,
@@ -629,7 +630,7 @@ with tab_layout:
# stage tables (attn/ffn/full). # stage tables (attn/ffn/full).
st.session_state["_pl_active_scope"] = _pl_scope st.session_state["_pl_active_scope"] = _pl_scope
st.success( st.success(
f"Loaded latency-optimal {_pl_lbl} config: CP={_pl_best.cp} " f"Loaded memory-optimal {_pl_lbl} config: CP={_pl_best.cp} "
f"TP={_pl_best.tp} PP={_pl_best.pp} DP={_pl_best.dp} → " f"TP={_pl_best.tp} PP={_pl_best.pp} DP={_pl_best.dp} → "
f"{_pl_best.latency_ms:.2f} ms, {_pl_best.pes_used} PEs. " f"{_pl_best.latency_ms:.2f} ms, {_pl_best.pes_used} PEs. "
"Redrawing layout..." "Redrawing layout..."
tests/analytical_visualization/stage_latencies.py
+96 -55
View File
@@ -59,12 +59,13 @@ def stage_rmsnorm(cfg: FullConfig) -> StageCost:
vis, bnd = _visible(cmp_s, mem_s, 0) vis, bnd = _visible(cmp_s, mem_s, 0)
return StageCost( return StageCost(
name="S1 RMSNorm", name="S1 RMSNorm",
formula=f"bytes = {T_q}*{d}*{b}*2 = {bytes_} B / BW_HBM", formula=f"bytes = B*T_q*d*b + T_q*d*b = {bytes_} B / BW_HBM",
compute_s=cmp_s, memory_s=mem_s, comm_s=0.0, compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
bound=bnd, visible_s=vis, bound=bnd, visible_s=vis,
flops=flops, mem_bytes=bytes_, flops=flops, mem_bytes=bytes_,
flops_formula=f"4*T_q*d = 4*{T_q}*{d} = {flops}", flops_formula=f"4*B*T_q*d = 4*{B}*{T_q}*{d} = {flops}",
mem_formula=f"2*T_q*d*b = 2*{T_q}*{d}*{b} = {bytes_} B", mem_formula=f"B*T_q*d*b + T_q*d*b (weight) "
f"= {B}*{T_q}*{d}*{b} + {T_q*d*b} = {bytes_} B",
) )
@@ -88,13 +89,16 @@ def stage_wq(cfg: FullConfig) -> StageCost:
vis, bnd = _visible(cmp_s, mem_s, 0) vis, bnd = _visible(cmp_s, mem_s, 0)
return StageCost( return StageCost(
name="S2 W_Q GEMM", name="S2 W_Q GEMM",
formula=f"FLOPs = 2*{T_q}*{d}*{hq_per_pe*dh} = {flops:.2g}; " formula=f"FLOPs = 2*B*T_q*d*(H_q/TP*d_h) "
f"weight = {weight_B/1e6:.1f} MB", f"= 2*{B}*{T_q}*{d}*{hq_per_pe*dh} = {flops:.2g}; "
f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
compute_s=cmp_s, memory_s=mem_s, comm_s=0.0, compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
bound=bnd, visible_s=vis, bound=bnd, visible_s=vis,
flops=int(flops), mem_bytes=int(weight_B), flops=int(flops), mem_bytes=int(weight_B),
flops_formula=f"2*T_q*d*(H_q/TP*d_h) = 2*{T_q}*{d}*{hq_per_pe*dh} = {flops:.3g}", flops_formula=f"2*B*T_q*d*(H_q/TP*d_h) "
mem_formula=f"d*(H_q/TP*d_h)*b = {d}*{hq_per_pe*dh}*{b} = {weight_B/1e6:.2f} MB", f"= 2*{B}*{T_q}*{d}*{hq_per_pe*dh} = {flops:.3g}",
mem_formula=f"d*(H_q/TP*d_h)*b (weight, B-invariant) "
f"= {d}*{hq_per_pe*dh}*{b} = {weight_B/1e6:.2f} MB",
) )
@@ -113,12 +117,15 @@ def stage_wkv(cfg: FullConfig) -> StageCost:
vis, bnd = _visible(cmp_s, mem_s, 0) vis, bnd = _visible(cmp_s, mem_s, 0)
return StageCost( return StageCost(
name="S3 W_K + W_V GEMM", name="S3 W_K + W_V GEMM",
formula=f"FLOPs per proj = 2*{T_q}*{d}*{hkv_per_pe*dh} = {flops_one:.2g}; x2", formula=f"FLOPs per proj = 2*B*T_q*d*(H_kv/TP*d_h) "
f"= 2*{B}*{T_q}*{d}*{hkv_per_pe*dh} = {flops_one:.2g}; x2",
compute_s=cmp_s, memory_s=mem_s, comm_s=0.0, compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
bound=bnd, visible_s=vis, bound=bnd, visible_s=vis,
flops=int(flops), mem_bytes=int(weight_B), flops=int(flops), mem_bytes=int(weight_B),
flops_formula=f"2*(2*T_q*d*(H_kv/TP*d_h)) = 2*(2*{T_q}*{d}*{hkv_per_pe*dh}) = {flops:.3g}", flops_formula=f"2*(2*B*T_q*d*(H_kv/TP*d_h)) "
mem_formula=f"2*(d*(H_kv/TP*d_h)*b) = 2*({d}*{hkv_per_pe*dh}*{b}) = {weight_B/1e6:.2f} MB", f"= 2*(2*{B}*{T_q}*{d}*{hkv_per_pe*dh}) = {flops:.3g}",
mem_formula=f"2*(d*(H_kv/TP*d_h)*b) (weights, B-invariant) "
f"= 2*({d}*{hkv_per_pe*dh}*{b}) = {weight_B/1e6:.2f} MB",
) )
@@ -132,12 +139,14 @@ def stage_kv_append(cfg: FullConfig) -> StageCost:
mem_s = bytes_ / cfg.machine.bw_hbm mem_s = bytes_ / cfg.machine.bw_hbm
return StageCost( return StageCost(
name="S4 KV cache append", name="S4 KV cache append",
formula=f"bytes = 2*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B", formula=f"bytes = 2*B*T_q*(H_kv/TP)*d_h*b "
f"= 2*{B}*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
compute_s=0, memory_s=mem_s, comm_s=0, compute_s=0, memory_s=mem_s, comm_s=0,
bound="memory", visible_s=mem_s, bound="memory", visible_s=mem_s,
flops=0, mem_bytes=bytes_, flops=0, mem_bytes=bytes_,
flops_formula="0 (no matmul, just cache write)", flops_formula="0 (no matmul, just cache write)",
mem_formula=f"2*T_q*(H_kv/TP)*d_h*b = 2*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B", mem_formula=f"2*B*T_q*(H_kv/TP)*d_h*b "
f"= 2*{B}*{T_q}*{hkv_per_pe}*{dh}*{b} = {bytes_} B",
) )
@@ -172,19 +181,23 @@ def stage_qkT(cfg: FullConfig) -> StageCost:
S_local = cfg.topo.s_local S_local = cfg.topo.s_local
dh = cfg.model.d_head dh = cfg.model.d_head
hq_per_pe = cfg.h_q_per_pe hq_per_pe = cfg.h_q_per_pe
flops_per_hop = 2 * T_q * S_local * dh * hq_per_pe B = max(1, cfg.topo.b)
flops_per_hop = 2 * B * T_q * S_local * dh * hq_per_pe
total_flops = flops_per_hop * passes total_flops = flops_per_hop * passes
_hop_word = "pass" if passes == 1 else "hops" _hop_word = "pass" if passes == 1 else "hops"
return StageCost( return StageCost(
name=f"S5 Q.K^T (x{passes} {_hop_word})", name=f"S5 Q.K^T (x{passes} {_hop_word})",
formula=f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops_per_hop:.2g} FLOPs/hop", formula=f"2*B*T_q*S_local*d_h*(H_q/TP) = "
f"2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe} = "
f"{flops_per_hop:.2g} FLOPs/hop",
compute_s=total, memory_s=0, comm_s=0, compute_s=total, memory_s=0, comm_s=0,
bound="compute", visible_s=total, bound="compute", visible_s=total,
hop_multiplier=passes, hop_multiplier=passes,
flops=int(total_flops), mem_bytes=0, flops=int(total_flops), mem_bytes=0,
flops_formula=( flops_formula=(
f"{passes}*(2*T_q*S_local*d_h*(H_q/TP)) = " f"{passes}*(2*B*T_q*S_local*d_h*(H_q/TP)) = "
f"{passes}*(2*{T_q}*{S_local}*{dh}*{hq_per_pe}) = {total_flops:.3g}" f"{passes}*(2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe}) = "
f"{total_flops:.3g}"
), ),
mem_formula="0 (scores accumulated on-chip)", mem_formula="0 (scores accumulated on-chip)",
) )
@@ -205,15 +218,16 @@ def stage_softmax(cfg: FullConfig) -> StageCost:
_hop_word = "pass" if passes == 1 else "hops" _hop_word = "pass" if passes == 1 else "hops"
return StageCost( return StageCost(
name=f"S6 softmax (x{passes} {_hop_word})", name=f"S6 softmax (x{passes} {_hop_word})",
formula=f"elts/hop = {hq_per_pe}*{T_q}*{S_local} = {elems:.2g}", formula=f"elts/hop = B*(H_q/TP)*T_q*S_local "
f"= {B}*{hq_per_pe}*{T_q}*{S_local} = {elems:.2g}",
compute_s=0, memory_s=total, comm_s=0, compute_s=0, memory_s=total, comm_s=0,
bound="memory", visible_s=total, bound="memory", visible_s=total,
hop_multiplier=passes, hop_multiplier=passes,
flops=0, mem_bytes=int(total_bytes), flops=0, mem_bytes=int(total_bytes),
flops_formula="~O(elts) (negligible)", flops_formula="~O(elts) (negligible)",
mem_formula=( mem_formula=(
f"{passes}*2*b*(H_q/TP)*T_q*S_local = " f"{passes}*2*b*B*(H_q/TP)*T_q*S_local = "
f"{passes}*2*{b}*{hq_per_pe}*{T_q}*{S_local} = " f"{passes}*2*{b}*{B}*{hq_per_pe}*{T_q}*{S_local} = "
f"{total_bytes/1e6:.2f} MB" f"{total_bytes/1e6:.2f} MB"
), ),
) )
@@ -227,19 +241,23 @@ def stage_pv(cfg: FullConfig) -> StageCost:
S_local = cfg.topo.s_local S_local = cfg.topo.s_local
dh = cfg.model.d_head dh = cfg.model.d_head
hq_per_pe = cfg.h_q_per_pe hq_per_pe = cfg.h_q_per_pe
flops_per_hop = 2 * T_q * S_local * dh * hq_per_pe B = max(1, cfg.topo.b)
flops_per_hop = 2 * B * T_q * S_local * dh * hq_per_pe
total_flops = flops_per_hop * passes total_flops = flops_per_hop * passes
_hop_word = "pass" if passes == 1 else "hops" _hop_word = "pass" if passes == 1 else "hops"
return StageCost( return StageCost(
name=f"S7 P.V (x{passes} {_hop_word})", name=f"S7 P.V (x{passes} {_hop_word})",
formula=f"2*{T_q}*{S_local}*{dh}*{hq_per_pe} = {flops_per_hop:.2g} FLOPs/hop", formula=f"2*B*T_q*S_local*d_h*(H_q/TP) = "
f"2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe} = "
f"{flops_per_hop:.2g} FLOPs/hop",
compute_s=total, memory_s=0, comm_s=0, compute_s=total, memory_s=0, comm_s=0,
bound="compute", visible_s=total, bound="compute", visible_s=total,
hop_multiplier=passes, hop_multiplier=passes,
flops=int(total_flops), mem_bytes=0, flops=int(total_flops), mem_bytes=0,
flops_formula=( flops_formula=(
f"{passes}*(2*T_q*S_local*d_h*(H_q/TP)) = " f"{passes}*(2*B*T_q*S_local*d_h*(H_q/TP)) = "
f"{passes}*(2*{T_q}*{S_local}*{dh}*{hq_per_pe}) = {total_flops:.3g}" f"{passes}*(2*{B}*{T_q}*{S_local}*{dh}*{hq_per_pe}) = "
f"{total_flops:.3g}"
), ),
mem_formula="0 (accumulated on-chip)", mem_formula="0 (accumulated on-chip)",
) )
@@ -302,13 +320,15 @@ def stage_merge(cfg: FullConfig) -> StageCost:
return StageCost( return StageCost(
name=f"S8 online-softmax merge (x{max(0, cfg.topo.cp-1)}){name_suffix}", name=f"S8 online-softmax merge (x{max(0, cfg.topo.cp-1)}){name_suffix}",
formula=f"~6*{T_q}*{hq_per_pe}*{dh}*(C-1) = {flops:.2g} FLOPs", formula=f"~6*B*T_q*(H_q/TP)*d_h*(C-1) "
f"= 6*{B}*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} "
f"= {flops:.2g} FLOPs",
compute_s=cmp_s, memory_s=0, comm_s=comm_s, compute_s=cmp_s, memory_s=0, comm_s=comm_s,
bound=bound, visible_s=visible_s, bound=bound, visible_s=visible_s,
flops=int(flops), mem_bytes=0, comm_bytes=comm_bytes, flops=int(flops), mem_bytes=0, comm_bytes=comm_bytes,
flops_formula=( flops_formula=(
f"6*T_q*(H_q/TP)*d_h*(CP-1) = " f"6*B*T_q*(H_q/TP)*d_h*(CP-1) = "
f"6*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} = {flops}" f"6*{B}*{T_q}*{hq_per_pe}*{dh}*{max(0, cfg.topo.cp-1)} = {flops}"
), ),
mem_formula="0 (in-register)", mem_formula="0 (in-register)",
comm_formula=comm_formula or "0", comm_formula=comm_formula or "0",
@@ -324,11 +344,13 @@ def stage_normalize(cfg: FullConfig) -> StageCost:
cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util) cmp_s = flops / (cfg.machine.peak_flops * cfg.machine.compute_util)
return StageCost( return StageCost(
name="S9 normalize O/l", name="S9 normalize O/l",
formula=f"{T_q}*{hq_per_pe}*{dh} = {flops} divisions", formula=f"B*T_q*(H_q/TP)*d_h = {B}*{T_q}*{hq_per_pe}*{dh} "
f"= {flops} divisions",
compute_s=cmp_s, memory_s=0, comm_s=0, compute_s=cmp_s, memory_s=0, comm_s=0,
bound="trivial", visible_s=cmp_s, bound="trivial", visible_s=cmp_s,
flops=int(flops), mem_bytes=0, flops=int(flops), mem_bytes=0,
flops_formula=f"T_q*(H_q/TP)*d_h = {T_q}*{hq_per_pe}*{dh} = {flops}", flops_formula=f"B*T_q*(H_q/TP)*d_h "
f"= {B}*{T_q}*{hq_per_pe}*{dh} = {flops}",
mem_formula="0", mem_formula="0",
) )
@@ -346,13 +368,16 @@ def stage_wo(cfg: FullConfig) -> StageCost:
vis, bnd = _visible(cmp_s, mem_s, 0) vis, bnd = _visible(cmp_s, mem_s, 0)
return StageCost( return StageCost(
name="S10 W_O GEMM", name="S10 W_O GEMM",
formula=f"FLOPs = 2*{T_q}*{hq_per_pe*dh}*{d} = {flops:.2g}; " formula=f"FLOPs = 2*B*T_q*(H_q/TP*d_h)*d "
f"weight = {weight_B/1e6:.1f} MB", f"= 2*{B}*{T_q}*{hq_per_pe*dh}*{d} = {flops:.2g}; "
f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
compute_s=cmp_s, memory_s=mem_s, comm_s=0.0, compute_s=cmp_s, memory_s=mem_s, comm_s=0.0,
bound=bnd, visible_s=vis, bound=bnd, visible_s=vis,
flops=int(flops), mem_bytes=int(weight_B), flops=int(flops), mem_bytes=int(weight_B),
flops_formula=f"2*T_q*(H_q/TP*d_h)*d = 2*{T_q}*{hq_per_pe*dh}*{d} = {flops:.3g}", flops_formula=f"2*B*T_q*(H_q/TP*d_h)*d "
mem_formula=f"(H_q/TP*d_h)*d*b = {hq_per_pe*dh}*{d}*{b} = {weight_B/1e6:.2f} MB", f"= 2*{B}*{T_q}*{hq_per_pe*dh}*{d} = {flops:.3g}",
mem_formula=f"(H_q/TP*d_h)*d*b (weight, B-invariant) "
f"= {hq_per_pe*dh}*{d}*{b} = {weight_B/1e6:.2f} MB",
) )
@@ -404,7 +429,7 @@ def comm_cp_ring(cfg: FullConfig) -> StageCost:
return StageCost( return StageCost(
name=f"C1 O/m/l all-reduce ({tier}, x{cp} ranks)", name=f"C1 O/m/l all-reduce ({tier}, x{cp} ranks)",
formula=(f"decode: gather partial (O,m,l) once at end; " formula=(f"decode: gather partial (O,m,l) once at end; "
f"M = T_q*(H_q/TP)*(d_h*b + 2*b) = {M} B per rank; " f"M = B*T_q*(H_q/TP)*(d_h*b + 2*b) = {M} B per rank; "
f"AR: 2*(CP-1)/CP*M / BW + 2*(CP-1)*alpha"), f"AR: 2*(CP-1)/CP*M / BW + 2*(CP-1)*alpha"),
compute_s=0, memory_s=0, comm_s=ar_time, compute_s=0, memory_s=0, comm_s=ar_time,
bound="comm", visible_s=ar_time, bound="comm", visible_s=ar_time,
@@ -432,16 +457,16 @@ def comm_cp_ring(cfg: FullConfig) -> StageCost:
M_KV = B * ((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" variant_desc = "Q+O/m/l ring"
formula_bytes = ( formula_bytes = (
f"2*T_q*(H_q/TP)*d_h*b + 2*T_q*(H_q/TP)*b " f"B*(2*T_q*(H_q/TP)*d_h*b + 2*T_q*(H_q/TP)*b) "
f"= 2*{T_q}*{hq_per_pe}*{dh}*{b} + 2*{T_q}*{hq_per_pe}*{b} " f"= {B}*(2*{T_q}*{hq_per_pe}*{dh}*{b} + 2*{T_q}*{hq_per_pe}*{b}) "
f"= {M_KV} B/hop" f"= {M_KV} B/hop"
) )
else: else:
M_KV = 2 * B * S_local * hkv_per_pe * dh * b M_KV = 2 * B * S_local * hkv_per_pe * dh * b
variant_desc = "K/V ring" variant_desc = "K/V ring"
formula_bytes = ( formula_bytes = (
f"2*S_local*(H_kv/TP)*d_h*b " f"2*B*S_local*(H_kv/TP)*d_h*b "
f"= 2*{S_local}*{hkv_per_pe}*{dh}*{b} " f"= 2*{B}*{S_local}*{hkv_per_pe}*{dh}*{b} "
f"= {M_KV/1e6:.3f} MB/hop" f"= {M_KV/1e6:.3f} MB/hop"
) )
@@ -527,7 +552,8 @@ def comm_tp_allreduce(cfg: FullConfig) -> StageCost:
total_comm_bytes = int(2 * (tp - 1) / tp * bytes_) total_comm_bytes = int(2 * (tp - 1) / tp * bytes_)
return StageCost( return StageCost(
name=f"C2 TP AllReduce W_O (TP={tp} ranks, {scope})", name=f"C2 TP AllReduce W_O (TP={tp} ranks, {scope})",
formula=f"2*(TP-1)/TP * {T_q}*{d}*{b} B / BW + 2(TP-1)*alpha [{scope}]", formula=f"2*(TP-1)/TP * B*T_q*d*b B / BW + 2(TP-1)*alpha "
f"[B={B}, {scope}]",
compute_s=0, memory_s=0, comm_s=comm_time, compute_s=0, memory_s=0, comm_s=comm_time,
bound="comm", visible_s=comm_time, bound="comm", visible_s=comm_time,
flops=0, mem_bytes=0, comm_bytes=total_comm_bytes, flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
@@ -541,7 +567,8 @@ def comm_tp_allreduce(cfg: FullConfig) -> StageCost:
f"every rank ends up with the full hidden vector for the next\n" f"every rank ends up with the full hidden vector for the next\n"
f"stage (RMSNorm -> FFN). Fires ONCE per layer over {scope} links.\n" f"stage (RMSNorm -> FFN). Fires ONCE per layer over {scope} links.\n"
f"---\n" f"---\n"
f"2*(TP-1)/TP * T_q*d*b = 2*({tp}-1)/{tp} * {T_q}*{d}*{b} " f"2*(TP-1)/TP * B*T_q*d*b "
f"= 2*({tp}-1)/{tp} * {B}*{T_q}*{d}*{b} "
f"= {total_comm_bytes/1e6:.2f} MB over {scope} at " f"= {total_comm_bytes/1e6:.2f} MB over {scope} at "
f"{bw/1e9:.0f} GB/s" f"{bw/1e9:.0f} GB/s"
), ),
@@ -576,7 +603,9 @@ def comm_kv_split_allreduce(cfg: FullConfig) -> StageCost:
total_comm_bytes = int(2 * (split - 1) / split * bytes_per_hop * cfg.topo.cp) total_comm_bytes = int(2 * (split - 1) / split * bytes_per_hop * cfg.topo.cp)
return StageCost( return StageCost(
name=f"C3 Score AllReduce ({split}-way, xCP hops)", name=f"C3 Score AllReduce ({split}-way, xCP hops)",
formula=f"2*({split}-1)/{split} * {hq_per_pe}*{T_q}*{S_local}*{b} / BW + latency", formula=f"2*({split}-1)/{split} * B*(H_q/TP)*T_q*S_local*b / BW "
f"= 2*({split}-1)/{split} * {B}*{hq_per_pe}*{T_q}*{S_local}*{b} "
f"/ BW + latency",
compute_s=0, memory_s=0, comm_s=total, compute_s=0, memory_s=0, comm_s=total,
bound="comm", visible_s=total, bound="comm", visible_s=total,
flops=0, mem_bytes=0, comm_bytes=total_comm_bytes, flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
@@ -647,12 +676,13 @@ def stage_ffn_rmsnorm(cfg: FullConfig) -> StageCost:
mem_s = bytes_ / cfg.machine.bw_hbm mem_s = bytes_ / cfg.machine.bw_hbm
return StageCost( return StageCost(
name="F1 RMSNorm (pre-FFN)", name="F1 RMSNorm (pre-FFN)",
formula=f"{T_q}*{d}*{b}*2 = {bytes_} B / BW_HBM", formula=f"bytes = B*T_q*d*b + T_q*d*b = {bytes_} B / BW_HBM",
compute_s=0, memory_s=mem_s, comm_s=0, compute_s=0, memory_s=mem_s, comm_s=0,
bound="memory", visible_s=mem_s, bound="memory", visible_s=mem_s,
flops=flops, mem_bytes=bytes_, flops=flops, mem_bytes=bytes_,
flops_formula=f"4*T_q*d = 4*{T_q}*{d} = {flops}", flops_formula=f"4*B*T_q*d = 4*{B}*{T_q}*{d} = {flops}",
mem_formula=f"2*T_q*d*b = 2*{T_q}*{d}*{b} = {bytes_} B", mem_formula=f"B*T_q*d*b + T_q*d*b (weight) "
f"= {B}*{T_q}*{d}*{b} + {T_q*d*b} = {bytes_} B",
) )
@@ -668,12 +698,16 @@ def _ffn_gemm(cfg: FullConfig, name: str, ffn_per_pe: int) -> StageCost:
vis, bnd = _visible(cmp_s, mem_s, 0) vis, bnd = _visible(cmp_s, mem_s, 0)
return StageCost( return StageCost(
name=name, name=name,
formula=f"FLOPs = 2*{T_q}*{d}*{ffn_per_pe} = {flops:.2g}; weight = {weight_B/1e6:.1f} MB", formula=f"FLOPs = 2*B*T_q*d*(ffn/div) "
f"= 2*{B}*{T_q}*{d}*{ffn_per_pe} = {flops:.2g}; "
f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
compute_s=cmp_s, memory_s=mem_s, comm_s=0, compute_s=cmp_s, memory_s=mem_s, comm_s=0,
bound=bnd, visible_s=vis, bound=bnd, visible_s=vis,
flops=int(flops), mem_bytes=int(weight_B), flops=int(flops), mem_bytes=int(weight_B),
flops_formula=f"2*T_q*d*(ffn/div) = 2*{T_q}*{d}*{ffn_per_pe} = {flops:.3g}", flops_formula=f"2*B*T_q*d*(ffn/div) "
mem_formula=f"d*(ffn/div)*b = {d}*{ffn_per_pe}*{b} = {weight_B/1e6:.2f} MB", f"= 2*{B}*{T_q}*{d}*{ffn_per_pe} = {flops:.3g}",
mem_formula=f"d*(ffn/div)*b (weight, B-invariant) "
f"= {d}*{ffn_per_pe}*{b} = {weight_B/1e6:.2f} MB",
) )
@@ -698,12 +732,14 @@ def stage_ffn_swiglu(cfg: FullConfig) -> StageCost:
mem_s = bytes_ / cfg.machine.bw_hbm mem_s = bytes_ / cfg.machine.bw_hbm
return StageCost( return StageCost(
name="F4 SwiGLU act", name="F4 SwiGLU act",
formula=f"3*{T_q}*{ffn_per_pe}*{b} = {bytes_} B / BW_HBM", formula=f"3*B*T_q*(ffn/div)*b "
f"= 3*{B}*{T_q}*{ffn_per_pe}*{b} = {bytes_} B / BW_HBM",
compute_s=0, memory_s=mem_s, comm_s=0, compute_s=0, memory_s=mem_s, comm_s=0,
bound="memory", visible_s=mem_s, bound="memory", visible_s=mem_s,
flops=flops, mem_bytes=bytes_, flops=flops, mem_bytes=bytes_,
flops_formula=f"~3*T_q*(ffn/div) = 3*{T_q}*{ffn_per_pe} = {flops}", flops_formula=f"~3*B*T_q*(ffn/div) = 3*{B}*{T_q}*{ffn_per_pe} = {flops}",
mem_formula=f"3*T_q*(ffn/div)*b = 3*{T_q}*{ffn_per_pe}*{b} = {bytes_} B", mem_formula=f"3*B*T_q*(ffn/div)*b "
f"= 3*{B}*{T_q}*{ffn_per_pe}*{b} = {bytes_} B",
) )
@@ -720,12 +756,16 @@ def stage_ffn_down(cfg: FullConfig) -> StageCost:
vis, bnd = _visible(cmp_s, mem_s, 0) vis, bnd = _visible(cmp_s, mem_s, 0)
return StageCost( return StageCost(
name="F5 W_down GEMM", name="F5 W_down GEMM",
formula=f"FLOPs = 2*{T_q}*{ffn_per_pe}*{d} = {flops:.2g}; weight = {weight_B/1e6:.1f} MB", formula=f"FLOPs = 2*B*T_q*(ffn/div)*d "
f"= 2*{B}*{T_q}*{ffn_per_pe}*{d} = {flops:.2g}; "
f"weight = {weight_B/1e6:.1f} MB (shared across batch)",
compute_s=cmp_s, memory_s=mem_s, comm_s=0, compute_s=cmp_s, memory_s=mem_s, comm_s=0,
bound=bnd, visible_s=vis, bound=bnd, visible_s=vis,
flops=int(flops), mem_bytes=int(weight_B), flops=int(flops), mem_bytes=int(weight_B),
flops_formula=f"2*T_q*(ffn/div)*d = 2*{T_q}*{ffn_per_pe}*{d} = {flops:.3g}", flops_formula=f"2*B*T_q*(ffn/div)*d "
mem_formula=f"(ffn/div)*d*b = {ffn_per_pe}*{d}*{b} = {weight_B/1e6:.2f} MB", f"= 2*{B}*{T_q}*{ffn_per_pe}*{d} = {flops:.3g}",
mem_formula=f"(ffn/div)*d*b (weight, B-invariant) "
f"= {ffn_per_pe}*{d}*{b} = {weight_B/1e6:.2f} MB",
) )
@@ -758,7 +798,8 @@ def comm_ffn_allreduce(cfg: FullConfig) -> StageCost:
total_comm_bytes = int(2 * (divisor - 1) / divisor * bytes_) total_comm_bytes = int(2 * (divisor - 1) / divisor * bytes_)
return StageCost( return StageCost(
name=f"CF1 FFN AllReduce (scope={scope}, x{divisor})", name=f"CF1 FFN AllReduce (scope={scope}, x{divisor})",
formula=f"2*({divisor}-1)/{divisor} * {T_q}*{d}*{b} B / BW + latency", formula=f"2*({divisor}-1)/{divisor} * B*T_q*d*b B / BW + latency "
f"[B={B}]",
compute_s=0, memory_s=0, comm_s=comm_time, compute_s=0, memory_s=0, comm_s=comm_time,
bound="comm", visible_s=comm_time, bound="comm", visible_s=comm_time,
flops=0, mem_bytes=0, comm_bytes=total_comm_bytes, flops=0, mem_bytes=0, comm_bytes=total_comm_bytes,
@@ -771,8 +812,8 @@ def comm_ffn_allreduce(cfg: FullConfig) -> StageCost:
f"rank has the full FFN output for the next layer. Larger scope\n" f"rank has the full FFN output for the next layer. Larger scope\n"
f"= less FFN weight memory per PE but bigger AR bill.\n" f"= less FFN weight memory per PE but bigger AR bill.\n"
f"---\n" f"---\n"
f"2*(div-1)/div * T_q*d*b = 2*({divisor}-1)/{divisor} * " f"2*(div-1)/div * B*T_q*d*b = 2*({divisor}-1)/{divisor} * "
f"{T_q}*{d}*{b} = {total_comm_bytes/1e6:.2f} MB at " f"{B}*{T_q}*{d}*{b} = {total_comm_bytes/1e6:.2f} MB at "
f"{bw/1e9:.0f} GB/s (scope={scope})" f"{bw/1e9:.0f} GB/s (scope={scope})"
), ),
) )