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):

pip install llm-mem-planner

From source (editable):

pip install -e tools/llm-mem-planner

Optional: install huggingface_hub to support private repos / caching:

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

Usage

Estimate from a HuggingFace repo id:

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:

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):

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:

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):

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:

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

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

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:

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):

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:

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.