# llm-mem-planner

A small, dependency-light CLI tool to:

- Estimate inference VRAM usage (**weights / KV cache / index cache / activations**).
- Recommend **tp/dp/sp/ep** configurations under a given hardware budget.
- (Optional) Use a simple **roofline + communication** model to estimate **decode/prefill time**
  and find the **knee points** for multi-node configs.
- Emit an **Obsidian-friendly** Markdown report (+ SVG plot) for note taking.

This tool is **standalone** and does **not** require installing SGLang or its heavy dependencies.

## Docs

- User guide (English): `docs/user_guide.md`

## Install

From PyPI (after release):

```bash
pip install llm-mem-planner
```

From source (editable):

```bash
pip install -e tools/llm-mem-planner
```

Optional: install `huggingface_hub` to support private repos / caching:

```bash
pip install -e "tools/llm-mem-planner[hf]"
```

## Usage

Estimate from a HuggingFace repo id:

```bash
llm-mem --model zai-org/GLM-5 --seq-len 20480 --batch 128 --tp 32 --dp 4 --sp 1 --ep 1
```

Estimate from a local config:

```bash
llm-mem --config /path/to/config.json --seq-len 20480 --batch 128 --tp 32 --dp 4 --sp 1 --ep 1
```

Offline quickstart (no network, uses bundled minimal configs):

```bash
cd tools/llm-mem-planner
python -m llm_mem_planner \
  --config tests/data/minimal_dense_config.json \
  --seq-len 2048 --batch 16 \
  --tp 8 --dp 2 --sp 1 --ep 1
```

MoE enabled:

```bash
python -m llm_mem_planner \
  --config tests/data/minimal_moe_config.json \
  --seq-len 2048 --batch 16 \
  --tp 8 --dp 2 --sp 1 --ep 1 --moe-dp 1
```

MoE + NSA (index cache should show up):

```bash
python -m llm_mem_planner \
  --config tests/data/minimal_moe_nsa_config.json \
  --seq-len 2048 --batch 16 \
  --tp 8 --dp 2 --sp 1 --ep 1 --moe-dp 1
```

If you have another `llm-mem` installed (there is an unrelated PyPI project with the same CLI name),
use one of the alternative entrypoints:

```bash
llm-mem-planner --help
python -m llm_mem_planner --help
```

Write a single-config **Obsidian-friendly** report (+ SVG breakdown):

```bash
llm-mem \
  --model zai-org/GLM-5 \
  --seq-len 102400 --batch 50 \
  --tp 32 --dp 4 --sp 1 --ep 1 \
  --gpu-mem-gib 96 --mem-fraction 0.96 \
  --gpu-tflops 1000 --gpu-hbm-gibps 2000 \
  --intra-gibps 200 --inter-gibps 50 \
  --weight-dtype int8 --act-dtype int8 --kv-dtype float16 \
  --obsidian
```

This writes:

- `mem_report.md`
- `mem_report.svg`

## mHC (optional)

If your model uses **mHC (multi-stream HyperConnections / multi residual streams)**, you can model
the extra working-set by setting:

- `--mhc-streams S` (alias: `--hc-num-streams`): number of residual streams (`S`).
- `--mhc-branches-per-layer K`: number of HyperConnections modules per transformer layer
  (typical `K=2` for `attn + mlp`).
- `--mhc-branch-mode single|per_stream`: whether attention/MLP branches are modeled as single-stream
  (mhc-torch style) or executed per stream.

Assumptions:
- Default (`--mhc-branch-mode single`): attention/MLP **branch stays single-stream** (so **KV/index cache is unchanged**).
- Optional (`--mhc-branch-mode per_stream`): each stream runs its own attention/MLP (so **KV/index + activations scale with `S`**).
- Activations: the hidden/residual working-set scales ~`S×` (tool scales the `hidden` term by `S`);
  in `per_stream` mode, the tool also scales the `attn`/`moe` terms by `S` (heuristic upper bound).
- Weights: mHC adds small replicated logits with params/module `S^2 + 2S` (negligible vs LLM weights).
- Perf model (roofline) currently **does not** include mHC overhead (treat perf estimates as optimistic).

Example:

```bash
llm-mem --model deepseek-ai/DeepSeek-V3.2 --seq-len 20480 --batch 128 --tp 32 --dp 4 --sp 1 --ep 1 --mhc-streams 4
```

Per-stream branch example (KV + act scale with streams):

```bash
llm-mem --model deepseek-ai/DeepSeek-V3.2 --seq-len 20480 --batch 128 --tp 32 --dp 4 --sp 1 --ep 1 --mhc-streams 4 --mhc-branch-mode per_stream
```

Recommend configs + generate Obsidian report:

```bash
llm-mem \
  --model deepseek-ai/DeepSeek-V3.2 \
  --recommend \
  --nodes 4 --gpus-per-node 8 --gpu-mem-gib 141 \
  --gpu-tflops 1000 --gpu-hbm-gibps 3000 \
  --intra-gibps 200 --inter-gibps 50 \
  --tp 32 --seq-len 20480 --batch 128 \
  --sp-candidates 1,2 \
  --preference balanced \
  --obsidian
```

The recommendation mode writes:

- `parallelism_report.md`
- `parallelism_dp_scan.svg`
- `parallelism_dp_scan_perf.svg` (when perf model enabled)

## Output interpretation (high level)

In **single-config** mode, the CLI prints (in order):

- **Model Summary**: dims parsed from config (`hidden_size`, `layers`, MLA dims, MoE/NSA flags).
- **Parallelism**: requested `tp/dp/sp/ep/moe_dp` and derived `attn_tp` / `moe_tp`.
- **Parameter Breakdown**: approximate parameter counts for each block (weights only, not runtime memory).
- **Weight Memory**:
  - `weights_unique`: total weights size for one full replica (no sharding).
  - `weights_per_rank`: per-rank estimate after sharding/replication assumptions.
- **KV Cache** (MLA latent): prints `kv_unique` and `kv_per_rank`, plus `index_*` when NSA is enabled.
- **Activation Working Set**: heuristic per-rank working set for prefill and decode (upper bound for capacity planning).

In **recommend** mode, the tool scans `(dp, sp)` (under a fixed `tp`) and prints the top feasible configs
under the memory budget (`gpu_mem_gib * mem_fraction`, minus `misc_gib`).

## Notes / Assumptions

- KV cache is estimated using the **MLA latent cache** shape used by DeepSeek-V3/V3.2 and GLM-5,
  not the traditional MHA `(K+V)` cache formula.
- The activation number is a **working-set upper bound** intended for capacity planning, not an exact peak.
- The perf model is a **rough lower-bound**:
  - `t_op = max(t_compute, t_mem, t_comm)`
  - `t_comm ≈ bytes / bw + steps * latency`
  - Use it to **compare configs** and locate **critical points**, not as an exact latency predictor.
