mukesh
9fdde44922
New tests/analytical_visualization/ module - an interactive dashboard for exploring memory / latency tradeoffs of transformer inference on the SIP architecture. Highlights: - 30+ model presets (Qwen, Llama 2/3/3.1, Mistral, Gemma 2, Phi 3, DeepSeek MLA, Mixtral/Qwen 3 MoE, Grok-1, ...) - Placement toggles: TP and CP each on PE-level vs cube-level - CP ring variant: K/V ring vs Q+O/m/l ring (prefill); in decode the O/m/l all-reduce is folded into S8 (no separate C1 row) - SIP interconnect: ring / mesh2d / torus2d with matching link drawing - Per-stage latency table with compute + memory + comm formulas, auto-scaled ns/us/ms, colored by dominant bound - Ring attention loop indicator on the pipeline diagram (purple arc over S5-S8 with 'xN hops' badge) - Tensor sharding view with optional physical PE/cube annotations - Replication-waste + optimization-hints panel - Save & compare configurations (config1, config2, ...): summary table plus side-by-side per-stage attention and FFN latency, best-in-row highlighting - Symbol glossary with current values for every symbol used in formulas Not tied to production sim_engine or runtime API; purely analytical tooling for design-space exploration. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
107 lines
3.6 KiB
Python
107 lines
3.6 KiB
Python
"""Auto-suggest (CP, TP, PP) that fits the model in the per-PE budget.
|
|
|
|
Strategy: iterate over candidate (CP, TP, PP) triples in order of
|
|
increasing total PE count and return the first one that satisfies
|
|
memory + kernel-support constraints with >10% slack. If none fits, return
|
|
the best-effort configuration and flag over-budget.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, replace
|
|
from typing import Iterable
|
|
|
|
from .model_config import FullConfig, ModelConfig, TopologyConfig, MachineParams
|
|
from .memory_layout import compute_memory
|
|
|
|
|
|
# Candidate parallelism dimensions (powers of 2 mostly, up to sensible limits).
|
|
_TP_OPTIONS = (1, 2, 4, 8, 16, 32)
|
|
_CP_OPTIONS = (1, 2, 4, 8, 16, 32, 64, 96)
|
|
_PP_OPTIONS = (1, 2, 4, 8, 16, 32)
|
|
|
|
|
|
@dataclass
|
|
class Suggestion:
|
|
cp: int
|
|
tp: int
|
|
pp: int
|
|
weights_gb: float
|
|
kv_gb: float
|
|
transient_gb: float
|
|
slack_gb: float
|
|
fits: bool
|
|
pes_used: int
|
|
sips_used: int
|
|
reason: str = ""
|
|
|
|
|
|
def _iter_candidates(model: ModelConfig) -> Iterable[tuple[int, int, int]]:
|
|
"""Yield (CP, TP, PP) in order of increasing PE count."""
|
|
triples: list[tuple[int, int, int]] = []
|
|
for tp in _TP_OPTIONS:
|
|
for cp in _CP_OPTIONS:
|
|
for pp in _PP_OPTIONS:
|
|
# PP must not exceed layer count.
|
|
if pp > model.layers:
|
|
continue
|
|
# Skip TP > H_q * some factor (unrealistic).
|
|
if tp > model.h_q * 4:
|
|
continue
|
|
triples.append((cp, tp, pp))
|
|
# Sort by pe_count then by (pp, tp, cp) — prefer smaller PP first
|
|
# (avoids pipeline bubbles), then smaller TP (avoids head-dim split).
|
|
triples.sort(key=lambda t: (t[0] * t[1] * t[2], t[2], t[1], t[0]))
|
|
for t in triples:
|
|
yield t
|
|
|
|
|
|
def _score_candidate(cp: int, tp: int, pp: int,
|
|
model: ModelConfig, machine: MachineParams,
|
|
s_kv: int, mode: str,
|
|
slack_frac: float = 0.10) -> Suggestion:
|
|
topo = TopologyConfig(cp=cp, tp=tp, pp=pp, s_kv=s_kv, mode=mode)
|
|
cfg = FullConfig(model=model, topo=topo, machine=machine)
|
|
mem = compute_memory(cfg)
|
|
|
|
fits = (not mem.over_budget
|
|
and mem.slack_bytes >= slack_frac * mem.budget_bytes)
|
|
reason = ""
|
|
if mem.over_budget:
|
|
reason = (f"weights+KV+transient ({mem.used_bytes/1e9:.2f} GB) "
|
|
f"exceeds budget ({mem.budget_bytes/1e9:.2f} GB)")
|
|
elif not fits:
|
|
reason = f"slack ({mem.slack_bytes/1e9:.2f} GB) below 10% of budget"
|
|
return Suggestion(
|
|
cp=cp, tp=tp, pp=pp,
|
|
weights_gb=mem.weights_bytes / 1e9,
|
|
kv_gb=mem.kv_cache_bytes / 1e9,
|
|
transient_gb=mem.transient_bytes / 1e9,
|
|
slack_gb=mem.slack_bytes / 1e9,
|
|
fits=fits,
|
|
pes_used=topo.total_pes,
|
|
sips_used=topo.sips_used,
|
|
reason=reason,
|
|
)
|
|
|
|
|
|
def auto_suggest(model: ModelConfig, machine: MachineParams,
|
|
s_kv: int, mode: str = "decode",
|
|
slack_frac: float = 0.10) -> Suggestion:
|
|
"""Return the smallest-PE-count (CP, TP, PP) that fits.
|
|
|
|
If no candidate fits, returns the best-effort one (highest slack,
|
|
even if negative) with fits=False.
|
|
"""
|
|
best_fit: Suggestion | None = None
|
|
best_effort: Suggestion | None = None
|
|
|
|
for cp, tp, pp in _iter_candidates(model):
|
|
s = _score_candidate(cp, tp, pp, model, machine, s_kv, mode, slack_frac)
|
|
if s.fits and best_fit is None:
|
|
best_fit = s
|
|
break # candidates are pre-sorted by PE count
|
|
if best_effort is None or s.slack_gb > best_effort.slack_gb:
|
|
best_effort = s
|
|
|
|
return best_fit or best_effort # type: ignore
|