[Hamerlate]

Summary

Adds analytic_models/stacked_dram: a 3D-stacked DRAM memory model and a memory-aware TTFT/TPS estimate for one PLENA chip.

PerfModel charges H_PREFETCH_* a single cycle, so the TTFT and TPS of llama_model.py assume the memory system always keeps up. This adds the missing memory term. A memory is either a fixed-bandwidth baseline (for example a datasheet HBM number) or a DRAM stack whose usable bandwidth follows from its height, vertical connectivity, buffering and thermal limits, so the same decoder can be compared across stack designs.

The memory mechanisms are ported from DeepStack (Mo et al., MICRO 2026, arXiv:2604.04750), and every module names the DeepStack source it follows (tile-ai/DeepStack@8509061). The package has no runtime or test dependency on DeepStack. No hardware values ship with it: a design is a JSON profile, and the two example profiles are deliberately fictional.

Mechanisms

mechanism	model	DeepStack source
peak bandwidth	connected layers × channels × bytes per transfer × transfers per clock × memory clock, or a direct per-layer peak	src/deepstack/mosaic/arch/custom_profile.py
connectivity efficiency	1 while at most half the layers are connected, then linear down to the fully connected efficiency	same
fully connected efficiency	row read cycles / (row read + recharge cycles) from bank timing, or given directly	same
buffering cap (Little's law)	bytes in flight per requester = bandwidth / requesters × round-trip latency × buffering factor; a smaller buffer scales the usable bandwidth down	src/deepstack/mosaic/dse_space/case_study_dram_layer/dram_layer_config.py
thermal derating	junction resistance grows with each layer; the dynamic power left sets the compute clock scale	same
wave quantisation	each transfer rounds up to transaction bytes × channels	src/tilesight/tilesight/fused_op_dtype_wave/matmul_fused_op_new_api_wave.py
capacity, energy	layers × capacity per layer; read and write bits × pJ/bit	src/deepstack/mosaic/arch/custom_profile.py, src/deepstack/mosaic/cost/energy.py

GPU-only fields (L2 multiplier, uncached utilisation, the bank-oracle clock) are dropped, the memory clock moves into the DRAM configuration, and DeepStack's per-SM quantities become per-requester (DMA stream) quantities.

Estimate

Each stage, meaning the embedding lookup and every transformer block, takes max(cycles / f, bytes / usable bandwidth) (stage-roofline, the default, which is how DeepStack prices an operator) or their sum (serial). The compute term is llama_model.py's PerfModel calls, including its ×2 decode issue factor. Traffic per block:

• weights are read once per forward pass;
• decode reads the layer's whole KV cache and writes one new K/V row;
• prefill re-reads K/V once per MLEN-row query tile (the H_PREFETCH_M loop of flash_attention);
• activations reach DRAM only where PerfModel spills them past Vector SRAM (both RMSNorms, the projection, the residual);
• storage precisions come from [ANALYTIC.PRECISION] in plena_settings.toml (MXFP8 by default).

The result reports TTFT, TPS, per-phase compute and memory time with the memory-bound share, DRAM bytes and energy, weights + KV against capacity, and warnings (buffering cap, thermal derating, head_dim × heads ≠ hidden_size). MoE decoders are rejected here; the multi-chip follow-up handles them.

just stacked-dram-describe analytic_models/stacked_dram/examples/fictional_stacked_dram.json --total-layers 4,8,12
just stacked-dram-estimate llama-3.1-8b analytic_models/stacked_dram/examples/fictional_stacked_dram.json
just test-stacked-dram

Validation

• 75 unit tests: configuration rules, each mechanism, profiles (strict keys, JSON booleans rejected as numbers), the estimate and the CLI.
• With unbounded bandwidth the estimate reproduces llama_model.py's TTFT and TPS (relative error ≤ 1e-12).
• Outside the repository, the mechanisms were checked against DeepStack@8509061 on 153 values (largest relative error 1.7e-16).
• ruff format --check . and ruff check . pass with the locked ruff.

Also in this PR

093de4ea makes analytic_models importable and installable: main's analytic_models/__init__.py imports memory and utilisation, which do not exist, and pyproject.toml does not include analytic_models*. #103 and #118 carry the same fix. Whichever lands first, the others resolve the conflict by keeping the union of imports.

Follow-ups

• feat(analytic): add a multi-chip model on DeepStack's network model #122, stacked on this PR: a multi-chip model on DeepStack's network model.
• A MemoryProvider adapter, so that the compiler-trace latency model of Add compiler-derived compute and HBM latency models #103 can price HBM traffic with a stack profile.