# `llm-mem` User Guide (English)

This guide is written for **end users** of the packaged CLI tool.

For maintainers/reviewers: the internal principles doc (formulas, op-by-op model, assumptions)
lives in the SGLang repository at `docs/internal_principles.md` and is **not** shipped in the
PyPI artifacts.

---

## 1) What is `llm-mem`?

`llm-mem` is a small, dependency-light CLI tool to:

- Estimate inference VRAM usage (**weights / KV cache / index cache / activations**).
- Recommend **tp/dp/sp/ep** configurations that fit a given hardware budget.
- (Optional) Use a simplified **roofline + collectives** model to compare decode/prefill time across configs.
- Emit an Obsidian-friendly Markdown report (+ SVG plots).

This tool is designed for **capacity planning and configuration comparison**.
It is **not** a cycle-accurate simulator, and the activation number is a **heuristic upper bound**.

---

## 2) Installation

### Option A: Install from source (editable)

From the SGLang repo root:

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

Optional: install Hugging Face helper (recommended for private models / caching):

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

### Option B: Install from PyPI (after release)

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

### CLI entrypoints

This package provides:

- `llm-mem` (primary CLI)
- `python -m llm_mem_planner` (module entrypoint)
- `llm-mem-planner` (alias CLI, useful if you have another `llm-mem` installed)

Note: there is an existing PyPI project named `llm-mem` (unrelated). If you have it installed,
prefer using `llm-mem-planner` or `python -m llm_mem_planner` to avoid command name collisions.

---

## 3) Quickstart

### 3.1 Estimate from a Hugging Face model repo

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

If the repo is private, pass a token:

```bash
llm-mem --model <org>/<repo> --hf-token $HF_TOKEN --seq-len 20480 --batch 128 --tp 32 --dp 4 --sp 1 --ep 1
```

### 3.2 Estimate from a local `config.json`

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

### 3.3 Offline quickstart (no network)

If you just want to verify the tool works end-to-end without downloading any model configs,
use the 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
```

### 3.4 Generate an Obsidian report (Markdown + SVG)

```bash
llm-mem \
  --model deepseek-ai/DeepSeek-V3.2 \
  --seq-len 102400 --batch 50 \
  --tp 32 --dp 4 --sp 1 --ep 1 \
  --gpu-mem-gib 96 --mem-fraction 0.96 \
  --obsidian
```

This writes:

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

---

## 4) Key concepts and flags

### 4.1 `--batch` and `--batch-scope`

The tool needs a batch size to estimate KV / index / activations.

- `--batch-scope global` (default):
  - `--batch` means the **global batch** across DP-attn.
  - Local batch per rank is `local_batch = global_batch / dp`.
- `--batch-scope local`:
  - `--batch` means the **local batch per DP rank**.
  - Global batch is `global_batch ≈ local_batch * dp`.

If you are unsure, use the default (`global`) and treat it as the number of sequences you want
to serve simultaneously.

### 4.2 Parallelism dimensions (`tp/dp/sp/ep/moe_dp`)

The CLI uses SGLang-style parallelism naming:

- `tp`: tensor-parallel world size (total GPUs for one replica)
- `dp`: attention data parallel (dp-attn)
- `sp`: context/sequence parallel (prefill-cp)
- `ep`: expert parallel (MoE all-to-all group size)
- `moe-dp`: MoE replication factor (expert data-parallel)

Derived:

- `attn_tp = tp / (dp * sp)`
- `moe_tp  = tp / (ep * moe_dp)`

The tool will fail-fast if the divisibility constraints do not hold.

### 4.3 Dtypes

- `--weight-dtype`: weights dtype override (default: config dtype)
- `--kv-dtype`: KV dtype override (default: config dtype)
- `--act-dtype`: activation dtype override (default: config dtype)

### 4.4 Activation working-set

Activation is estimated as a **working-set upper bound** for one forward step:

- prefill: `tokens_per_step ≈ min(chunked_prefill_size / dp, local_batch * local_seq_len)`
- decode: `tokens_per_step ≈ local_batch`

Tune the estimate with:

- `--chunked-prefill-size` (default: 8192)
- `--misc-gib` (reserve runtime overhead)

### 4.5 mHC (optional)

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

- `--mhc-streams S` (alias: `--hc-num-streams`)
- `--mhc-branches-per-layer K` (typical `K=2` for `attn + mlp`)
- `--mhc-branch-mode single|per_stream`

Assumptions:

- `single` (default): attention/MLP branches stay single-stream; KV/index cache unchanged.
- `per_stream`: each stream runs its own attention/MLP; KV/index + activations scale with streams.

---

### 4.6 Output interpretation (single-config mode)

Single-config mode prints blocks in order. The most important fields:

- `== Weight Memory ==`
  - `weights_unique`: full model weights size (one replica).
  - `weights_per_rank`: per-rank estimate after sharding/replication assumptions.
- `== KV Cache (MLA latent cache) ==`
  - `kv_unique`: KV cache size for the full replica batch/seq.
  - `kv_per_rank`: per-rank KV after localizing batch (`dp`) and sequence (`sp`).
  - When NSA is enabled, the same line also prints `index_unique` / `index_per_rank`.
- `== Activation Working Set (heuristic, per rank) ==`
  - A conservative working-set upper bound for one forward step (prefill and decode).
  - Use this for capacity planning; it is not exact peak memory.

If you provide `--gpu-mem-gib`, the tool also computes a budget:

- `budget_gib = gpu_mem_gib * mem_fraction`
- `misc_gib` is reserved runtime overhead

Recommendation mode uses this budget to decide `fits` for scanned configs.

## 5) Recommendation mode (`--recommend`)

Recommendation mode scans `(dp, sp)` under a fixed `tp` and filters configs that fit your VRAM budget.

Typical usage:

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

Outputs:

- `parallelism_report.md`
- `parallelism_dp_scan.svg`
- `parallelism_dp_scan_perf.svg` (only when perf spec is provided)

---

## 6) Perf model (optional): roofline + collectives

To enable the perf model, provide the hardware spec:

- `--gpu-tflops`
- `--gpu-hbm-gibps`
- `--intra-gibps`, `--inter-gibps`
- (optional) `--intra-us`, `--inter-us`
- (optional) efficiency knobs: `--flops-eff`, `--hbm-eff`, `--comm-eff`

How to interpret:

- The model is a **lower bound** and is most useful to compare configs and find knees.
- Treat absolute numbers as optimistic; tune the efficiency knobs to match reality.

---

## 7) Troubleshooting

### “Config missing MLA attention dims …”

The tool currently expects MLA-style config fields:

`num_attention_heads, q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim`.

If your model uses a different attention architecture, it is not supported by this estimator yet.

### “Invalid parallelism: require tp % (dp*sp) == 0 …”

Make sure `(dp * sp)` divides `tp`. Similarly, `(ep * moe_dp)` must divide `tp`.

### Hugging Face download issues

- If the model is private, provide `--hf-token` or set `HF_TOKEN`.
- If you prefer no dependencies, the tool can download `config.json` via plain HTTPS, but private repos still need a token.

### “Unsupported dtype: …”

The tool is strict on dtype strings to avoid silent under/over-estimation.

Supported examples:

- `float16` / `fp16` / `half`
- `bfloat16` / `bf16`
- `float32` / `fp32`
- `float8` / `fp8`
- `int8` / `uint8`
- `int4` / `fp4`

It also accepts `torch.<dtype>` (e.g. `torch.bfloat16`).

### “Failed to parse JSON in --config file …”

- Ensure the file is valid JSON (no trailing commas).
- Ensure it is a HuggingFace-style `config.json` (or an equivalent minimal subset).
