[Hamerlate]

Stacked on #121: this PR adds the last three commits. It will be retargeted to main, and rebased, once #121 merges.

Summary

Adds analytic_models/distributed: TTFT and decode throughput of a decoder served by several PLENA chips under tensor, expert, data, pipeline and context parallelism. Each chip is PerfModel compute priced against a stacked_dram memory system. The network side is DeepStack's (arXiv:2604.04750): hierarchical NoC topologies and routing, collective algorithms, pipeline transfers, MoE routing statistics and NoC energy. It is vendored as the DeepStack git submodule (tile-ai/DeepStack at 17e21bf) and imported, not copied.

commit	change
a35565ac refactor(stacked_dram)	expose the per-stage roofline as price_stage and PhaseTotals; no behaviour change
c7eddbef build	add the DeepStack submodule and SciPy, its only dependency PLENA lacked; ruff skips the submodule
ce1adbcd feat(analytic)	the model, CLI, illustrative NoC profiles, tests and README

How a layer is costed

On each device a layer is one roofline stage (compute against its DRAM traffic) followed by its collectives (--comm-overlap none, the default) or overlapped with them (full).

part	per-device work	collective (DeepStack)
attention	heads and KV heads split over TP (KV heads replicated beyond their count), KV sequence split over CP	TP all-reduce; with CP, all-reduce of the partial outputs (decode) or all-gather of K and V (prefill)
dense FFN	intermediate size split over TP	TP all-reduce
routed experts	the busiest EP rank's tokens and token-expert pairs	dispatch and combine all-to-all in each EP group; TP all-reduce when the experts stay tensor parallel
RMSNorm, residual	replicated across TP	none
embedding, LM head	first stage (prefill lookup) and last stage, vocabulary split over TP	none; sampling reduces the logits locally

• Layer compositions are PLENA's: llama follows llama_model.py, and moe follows gpt_oss_model.py (full or sliding-window attention per layer, routed experts, the LM head in prefill).
• Plan semantics follow DeepStack's decode driver: EP ranks act as data parallel outside MoE layers; MoE layers fold TP ranks into EP ranks by default (--moe-tp-mode replace, DeepStack's replace_only) or keep the experts tensor parallel (keep); pipeline stages hold ceil(layers / pp) layers.
• MoE routing is balanced (PerfModel.mlp_moe's assumption, exact on one device), random, or DeepStack's qwen3_235b and deepseek_v3 traces. As in DeepStack, the busiest rank's load comes from the routed tokens below 4096 micro-batch tokens and from DeepStack's closed-form imbalance estimate above; the all-to-all switches at 8192.
• Pipelining uses micro-batches of ceil(batch / pp) sequences. The period is the slowest stage plus one stage-to-stage transfer. TTFT is one micro-batch's prefill through every stage plus its first decode step, TPS is one micro-batch of tokens per period, and per-sequence TPS is one token every pp periods (DeepStack's utps).

git submodule update --init DeepStack
just distributed-describe-noc analytic_models/distributed/examples/illustrative_gpu_cluster_32.json
just distributed-estimate llama-3.1-70b analytic_models/distributed/examples/illustrative_gpu_cluster_32.json \
    --profile analytic_models/stacked_dram/examples/fictional_stacked_dram.json --tp 8 --pp 4
just distributed-estimate gpt-oss-20b analytic_models/distributed/examples/illustrative_gpu_node_8.json \
    --fixed-bandwidth-gbs 1000 --tp 2 --ep 4 --routing random
just test-distributed

Validation

• 64 unit tests. A fresh clone with the submodule initialised passes all 139 analytic tests.
• One device reproduces stacked_dram.estimate_decoder_latency exactly (three memories, LM head on and off) and, with unbounded bandwidth, llama_model.py and gpt_oss_model.py (relative error ≤ 2.2e-16).
• TP, CP, PP and EP times and NoC energy equal DeepStack's wrappers called directly, including the ring fallback for non-power-of-two TP groups.
• ruff format --check . and ruff check . pass with the locked ruff.

Modeling choices worth a look

1. RMSNorm is replicated across TP and needs no collective. DeepStack shards it over the hidden dimension and all-reduces its statistics, so with TP > 1 this charges less communication than DeepStack.
2. TTFT sums the stages plus pp − 1 transfers. DeepStack multiplies one stage time by pp, which counts one more transfer.
3. Where TP ranks fold into EP ranks, the TP all-reduce after attention stands in for the reduce-scatter and all-gather around the experts, as in DeepStack's decode driver.
4. Balanced routing spreads a token's choices over EP ranks. DeepStack's all-to-all counts one transfer per (token, destination rank), so consecutive expert ids would undercount the dispatch traffic.
5. NoC energy follows the all-reduce algorithm DeepStack picks for time. On switches that is often the one-shot all-to-all, which moves the most bytes.
6. Pipeline transfers send the full activation from every TP rank, as in DeepStack.
7. Scope: CP covers dense full-attention models only, shared experts are not modelled, and the LM head follows the single-chip scripts unless --include-lm-head is given.

Notes

• import analytic_models does not need the submodule; analytic_models.distributed does (git submodule update --init DeepStack, or PLENA_DEEPSTACK_ROOT for another checkout). The import shim loads DeepStack source modules only and fails if any of DeepStack's bundled reference binaries gets loaded.
• The example NoC profiles are illustrative. Two follow DeepStack's public h200x32 GPU preset and two its examples/custom_noc.py, all with DeepStack's example energy coefficients; none describes PLENA hardware.
• The DeepStack changes this relies on are on tile-ai/DeepStack main: b7adc20 keeps switch instances apart in all-switch hierarchies, and 17e21bf implements all_gather_wrapper.
• CI triggers on pull requests to main, so it runs here once this PR is retargeted.

[GeorgeWu1204]

lgtm, approve. Thanks Zhiwen!