diff --git a/.github/workflows/claude-review.yml b/.github/workflows/claude-review.yml index e6c8a00..35bb735 100644 --- a/.github/workflows/claude-review.yml +++ b/.github/workflows/claude-review.yml @@ -16,11 +16,17 @@ jobs: with: fetch-depth: 0 + # Code first, prose second. The diff is truncated below when it is large, + # and `git diff` orders by path β€” so a PR that touches `docs/` and `gitm/` + # used to spend its whole budget on documentation and never show the + # reviewer a line of Python. Every review of PR #104 opened by calling it + # "a documentation-only diff" while 2k lines of planner changes sat past + # the cut. Emitting code first means truncation drops prose instead. - name: Get PR diff run: | - git diff origin/${{ github.base_ref }}...HEAD \ - -- '*.py' '*.sh' '*.md' '*.yaml' '*.yml' \ - > pr_diff.txt + BASE="origin/${{ github.base_ref }}" + git diff "$BASE"...HEAD -- '*.py' '*.sh' '*.yaml' '*.yml' > pr_diff.txt + git diff "$BASE"...HEAD -- '*.md' >> pr_diff.txt echo "Diff size: $(wc -c < pr_diff.txt) bytes" - name: Run Claude Review @@ -44,7 +50,7 @@ jobs: fi python3 << 'EOF' - import json, os, subprocess, sys, urllib.request + import json, os, subprocess, sys, time, urllib.error, urllib.request with open("pr_diff.txt") as f: diff = f.read() @@ -64,6 +70,8 @@ jobs: "**⚑ Performance** β€” unnecessary CPU↔GPU transfers, missed parallelism\n" "**πŸ“Š Reproducibility** β€” seed handling, non-determinism risks\n" "**πŸ’‘ Suggestions** β€” missing error handling, untested edge cases\n\n" + "The diff may be truncated; review what is present and do not " + "infer anything from what is missing.\n\n" f"```diff\n{diff}\n```" ) }] @@ -80,13 +88,34 @@ jobs: } ) - try: - with urllib.request.urlopen(req) as resp: - data = json.load(resp) - comment = data["content"][0]["text"] - except urllib.error.HTTPError as e: - print(f"API error {e.code}: {e.read().decode()}", file=sys.stderr) - sys.exit(1) + # Retry the transient classes only. 429 and 5xx are load; 4xx otherwise + # (auth, billing, malformed) will fail identically on every attempt, so + # retrying them just burns runner minutes. + comment = None + reason = None + for attempt in range(3): + try: + with urllib.request.urlopen(req, timeout=180) as resp: + comment = json.load(resp)["content"][0]["text"] + break + except urllib.error.HTTPError as e: + reason = f"HTTP {e.code}: {e.read().decode()[:300]}" + if e.code not in (429, 500, 502, 503, 504) or attempt == 2: + break + time.sleep(5 * (attempt + 1)) + except (urllib.error.URLError, TimeoutError, KeyError, IndexError) as e: + reason = f"{type(e).__name__}: {e}" + if attempt == 2: + break + time.sleep(5 * (attempt + 1)) + + # An advisory reviewer that cannot reach its API has found nothing; it + # has not found a problem. Failing the check here blocks a PR on someone + # else's outage or on an expired key, which is not a signal about the + # code. Say so in the log and pass. + if comment is None: + print(f"::warning::Claude review skipped β€” {reason}", file=sys.stderr) + sys.exit(0) body = f"## πŸ€– Claude Code Review\n\n{comment}" subprocess.run([ diff --git a/.github/workflows/gemini-pr-review.yml b/.github/workflows/gemini-pr-review.yml index 2576518..0c8cd1d 100644 --- a/.github/workflows/gemini-pr-review.yml +++ b/.github/workflows/gemini-pr-review.yml @@ -18,7 +18,12 @@ jobs: with: fetch-depth: 0 # Fetches full history so git diff can calculate cleanly + # Advisory, so an upstream outage must not block a PR. This action exits + # non-zero when the Gemini API answers 503 "experiencing high demand", + # which it did five times on PR #104 alone β€” a red check that says nothing + # about the branch. The review still posts when the API answers. - name: Run Gemini Code Review + continue-on-error: true uses: sshnaidm/gemini-code-review-action@v2 with: gemini-key: ${{ secrets.GEMINI_API_KEY }} diff --git a/docs/glm-5.2/DESIGN-NOTE.md b/docs/glm-5.2/DESIGN-NOTE.md new file mode 100644 index 0000000..b6ed2ac --- /dev/null +++ b/docs/glm-5.2/DESIGN-NOTE.md @@ -0,0 +1,1141 @@ +# GLM-5.2 β€” Working Design Note + +**Predicted execution model for Z.ai GLM-5.2 on 8Γ—H200 SXM, TP8 / EP8, FP8** + +Built from the model repos' own files β€” `config.json` and +`model.safetensors.index.json` for both `zai-org/GLM-5.2` (bf16) and +`zai-org/GLM-5.2-FP8`, plus the vendor's published vLLM recipe for the deployment +shape. **No traces.** Every number is a roofline floor at vendor peak: a lower +bound on time, not a target. + +Reproduce any figure here (against this branch β€” the planner is actively changing, +so check `git rev-parse HEAD` matches if a number disagrees): + +```bash +gitm plan glm-5.2-fp8 --gpu H200 --batch 32 --kv-len 8192 --tp 8 --ep 8 +gitm plan glm-5.2-fp8 --gpu H200 --batch 32 --kv-len 8192 --tp 8 --ep 8 --spec-tokens 5 +gitm plan glm-5.2-fp8 --gpu H200 --prefill-tokens 8192 --batch 0 --kv-len 0 --tp 8 --ep 8 +gitm plan glm-5.2-fp8 --gpu H200 --sweep 1,4,16,32,64,128,256 --kv-len 8192 --tp 8 --ep 8 +gitm plan glm-5.2 --gpu H200 --batch 32 --kv-len 8192 --tp 8 --ep 8 # what bf16 costs +``` + +## Hardware assumption: 8Γ—H200 SXM, NVLink, TP8 / EP8, FP8 weights and KV + +Unlike a hardware assumption inferred from a checkpoint, **this one is the +vendor's own** β€” `recipes.vllm.ai/zai-org/GLM-5.2` publishes it verbatim. If +production differs, *graph topology does not change*; only these constants and +some bound labels do. Regions whose label would flip are marked βš‘ throughout. + +| Constant | Value used | Note | +| -------------------- | --------------------------------- | ------------------------------------------------------------- | +| FP8 e4m3 tensor peak | **1,979 TFLOP/s** | datasheet says 3,958 **"with sparsity"** β€” halved, see A2 | +| BF16 tensor peak | **989.5 TFLOP/s** | same halving | +| FP32 CUDA-core peak | **67 TFLOP/s** | the router runs here β€” see G4, and it is 3.4Γ— the old default | +| HBM3e | **4.8 TB/s** | as published | +| NVLink | **900 GB/s** per GPU | the catalogue's bidirectional convention | +| Memory | 141 GB Γ— 8 = 1,128 GB | one node | +| Kernel launch | ~2 Β΅s graph-replay / ~5 Β΅s eager | βš‘ the crossover hinge β€” see Β§5 rank 3 | + +The recipe, verbatim, because every constant above and every bound label in Β§4 +assumes it β€” and because Β§6.2's C6 ("the engine's launch arguments, as text") is +worth more than most of the traces: + +```bash +vllm serve zai-org/GLM-5.2-FP8 --kv-cache-dtype fp8 --tensor-parallel-size 8 --speculative-config.method mtp --speculative-config.num_speculative_tokens 5 --tool-call-parser glm47 --reasoning-parser glm45 --enable-auto-tool-choice +``` + +`--enable-expert-parallel` is **not** in it. Β§4 prices the EP8 shape anyway, +because it is the shape that makes the `moe_all_to_all` rows exist at all β€” +TP8-only removes them and doubles the per-rank expert bank instead. Which of the +two is running is capture C5, and it re-ranks the largest line in prefill. + +``` +FP8 ridge = 1,979e12 / 4.8e12 = 412 FLOP/byte +BF16 ridge = 989.5e12 / 4.8e12 = 206 FLOP/byte +FP32 ridge = 67e12 / 4.8e12 = 14 FLOP/byte +``` + +**You need all three.** The backbone GEMMs and the experts are FP8; `lm_head`, +`embed_tokens`, the MTP `eh_proj` and β€” the one worth naming β€” the **lightning +indexer** are BF16; the router is FP32. Β§1's precision table shows why, and G5 is +the planner change that made it representable. + +--- + +## 1. Layer-by-layer architecture map + +### Headline structure + +| Property | Value | +| --------------------------- | ------------------------------------------------------------------------ | +| Layers | **78** transformer + **1** MTP draft module | +| Attention | **MLA + DeepSeek Sparse Attention on every layer** β€” no schedule at all | +| KV latent | `kv_lora_rank=512` + `qk_rope_head_dim=64` = **576 elems/token/layer** | +| Indexer | 32 heads Γ— 128 dim, keeps **`index_topk=2048`** positions for the core | +| **IndexShare** | **21 of 78 layers** compute the index; **57 reuse** a neighbour's | +| Dense MLP | **3** β€” layers 0, 1, 2 (`first_k_dense_replace: 3`), `intermediate=12288` | +| MoE layers | **75** β€” layers 3–77, plus the MTP module | +| Experts / top-k | 256 routed, top-8, **1 shared**, `moe_intermediate_size=2048` | +| Routing | sigmoid scoring, `noaux_tc`, `routed_scaling_factor=2.5`, **fp32 router** | +| hidden_size | 6144 | +| Q heads / q_lora / kv_lora | 64 / 2048 / 512 | +| qk_nope / qk_rope / v_head | **192 / 64 / 256** β€” the value width differs from the score width | +| RoPE | ΞΈ=8e6, `rope_interleave`, `indexer_rope_interleave` | +| MTP | 1 module, `index_share_for_mtp_iteration: true`, **carries a full MoE** | +| Vocab | 154,880, untied `lm_head` | +| Max context | 1,048,576 | +| Total / active params | **744 B** published + a **9.9 B** MTP block / ~39 B active | + +### Semantics read from the checkpoint, not guessed + +From `config.json`: `indexer_types[i]` is `"full"` (the layer computes its own +top-2048) or `"shared"` (it **reuses the previous `full` layer's top-k** β€” the +semantics `transformers` documents, verbatim); `mlp_layer_types[i]` is +`"dense"` | `"sparse"`, agreeing with `first_k_dense_replace: 3`; and +`moe_router_dtype: "float32"` makes the router fp32 **on every variant**, being a +field of the base config rather than of any quantisation config. + +**The schedule is proven from the weight map, not inferred.** Indexer tensors +(`*.indexer.wq_b`, `.wk`, `.weights_proj`, `.k_norm`) exist on exactly the 21 +`full` layers, on **none** of the 57 `shared` layers, and on **none** of the MTP +module. Pricing all 78 at full rate β€” the naive reading of `index_topk` β€” +overstates the indexer ~3.7Γ—. (Z.ai publishes **2.9Γ—** for +[IndexShare](https://huggingface.co/papers/2603.12201); that is whole-model +per-token FLOPs at 1M context, where 3.7Γ— is the indexer's own ratio, 78 Γ· 21.) + +Why "read verbatim" is not pedantry: this catalogue entry carried the schedule one +entry short for a while. Layer 77 took the modulo fallback, landed on `shared`, +and the count still came out 21 β€” the right answer from evidence that was not +there, with a byte-identical floor. The loader now refuses a short schedule. + +Attention shapes, per token per layer: + +| Quantity | Shape | Purpose | +| ------------------------ | -------------------- | -------------------------------------------------- | +| Q latent (`q_a`) | 2048 | replicated low-rank query | +| Q per-head (`q_b`) | 64 Γ— 256 = 16,384 | 192 nope + 64 rope | +| **KV cache entry** | **512 + 64 = 576** | **one latent for all 64 heads** β€” the MLA point | +| `kv_b` output | 64 Γ— (192+256) = 28,672 | reconstructed K_nope and V | +| Attention output | 64 Γ— 256 = 16,384 | one 256-d result per Q head | +| After `o_proj` | 6,144 | back to `d_model` | +| Index key (`full` only) | 128 | cached alongside the latent | + +`num_key_value_heads: 64` is a **red herring**: a GQA reading gives 28,672 +elements per token per layer against the real 576 β€” **50Γ— on the single quantity +decode is bound by**. + +### The 79-row table, collapsed to four archetypes + +The 78 layers plus the draft module are exactly four shapes. Everything not listed +is byte-for-byte identical between them. + +| Archetype | Count | Attn | Indexer | MLP | KV elems/tok | Collectives per layer | +| --------- | ----- | ---- | -------- | ------------- | ------------ | --------------------- | +| `Ld,f` | **3** | MLA+DSA | **full** | dense 12288 | 576 + 128 | 2Γ— all-reduce | +| `Ls,f` | **18**| MLA+DSA | **full** | MoE 256Γ—2048 | 576 + 128 | 2Γ— all-reduce (+2Γ— a2a under EP) | +| `Ls,sh` | **57**| MLA+DSA | shared | MoE 256Γ—2048 | 576 | 2Γ— all-reduce (+2Γ— a2a under EP) | +| `Lmtp` | **1**, Γ—D | MLA+DSA | shared | MoE 256Γ—2048 | 576 | 2Γ— all-reduce per stage | + +`Ld,f` and `Ls,f` are both full-indexer archetypes β€” 3 + 18 = **21**, matching Β§1. +The transformer stack is 3 + 18 + 57 = **78**; `Lmtp` is a 79th block that exists +in the checkpoint and runs **only under a speculative config**, once per drafted +token (Β§3.3). + +The two schedules do not line up β€” a 3-layer dense prefix, a 3-layer `full` +prefix, then an IndexShare period of 4 β€” so no single modulo rule reproduces +either, which is why both are read verbatim (Β§7, G3). And note the absences: no +sliding window, no compression schedule, no attention-type alternation, no +encoders. Every layer runs the same attention. + +### Verification β€” three independent checks + +| check | predicted | published | error | +| --- | --- | --- | --- | +| bf16 checkpoint | 1,508.1 GB | 1,506,659,919,872 B (282 shards) | **+0.08 %** | +| fp8 checkpoint | 755.9 GB | 753,329,940,480 B (141 shards) | **+0.34 %** | +| params, MTP block removed | **744.2 B** | **744 B** (Z.ai) | **+0.03 %** | + +The third is the interesting one. The checkpoint is 753.3 B by its own bytes but +Z.ai publishes 744 B β€” the gap is the MTP block, which the published figure +excludes. So the block is **9.9 B**, where a *dense* draft head would be **0.23 B**: +the draft carrying a full 256-expert mixture (Β§2.4) is confirmed twice, once from +the weight map and once from arithmetic on a number published for another reason. + +Two precisions agreeing to under half a percent also rules out sparsity β€” a +2:4-compressed checkpoint would be roughly half the fp8 size (A2). + +### KV cache β€” the number that drives decode + +``` +per layer per token, elements: + every layer kv_lora_rank 512 + qk_rope 64 = 576 ← one latent, all 64 heads + full-indexer layers, additionally + 128 ← the cached index key + +whole model, bytes per token of context (fp8 latent, bf16 rope key + index key): + an fp8 weight costs 1.000244 B, not 1 B β€” the 128Γ—128 block scale (Β§7.0) + 78 Γ— (512Γ—1.000244 + 64Γ—2) + 21 Γ— 128Γ—1.000244 = 52,618 B/token + bf16 throughout: 95,232 B/token + + ⚠ 1.000244 is a *weight*-side constant borrowed for cache bytes. An fp8 KV + cache carries a per-token or per-tensor scale, not 128Γ—128 blocks, so the + true figure is nearer 78Γ—640 + 21Γ—128 = 52,608. The 10-byte gap is 0.02% + and moves nothing; it is named because the constant is the wrong one, not + because the number is. +``` + +| Context | fp8 KV | bf16 KV | +| ------- | ------- | ------- | +| 8,192 | 0.43 GB | 0.78 GB | +| 131,072 | 6.90 GB | 12.48 GB | +| 1,048,576 | **55.2 GB** | **99.9 GB** | + +**Replicated, not sharded**: one shared latent cannot be split, so every rank +reads the whole cache β€” **TP buys no KV bandwidth here**. At 1M that is 55 GB per +rank on top of a 96 GB weight share, which is why the vendor recipe reaches for +B200s and `--max-num-seqs 32` for full context. + +### FP8 β€” what is and is not quantized + +| Component | Precision | Evidence | +| -------------------------------------------------- | ------------------------------- | ------------------------------------------------- | +| `q_a`/`q_b`/`kv_a`/`kv_b`, **`o_proj`**, dense FFN | **FP8 e4m3**, 128Γ—128 block | absent from `modules_to_not_convert` | +| routed experts, shared expert | **FP8 e4m3**, 128Γ—128 block | absent from `modules_to_not_convert` | +| `lm_head`, `embed_tokens` | **BF16** | named in `modules_to_not_convert` | +| **lightning indexer** β€” `indexers_proj`, `wq_b`, `wk`, `weights_proj` (+ `k_norm`) | **BF16** | named in `modules_to_not_convert` | +| MTP `eh_proj`, `enorm`, `hnorm` | **BF16** | named in `modules_to_not_convert` | +| MoE router (`mlp.gate` + `e_score_correction_bias`) | **FP32** | `moe_router_dtype: "float32"`, base config | +| all norms | **BF16** | named in `modules_to_not_convert` | + +**Three precisions inside one attention block**, and the layout **inverts** the +familiar fp8-backbone pattern: here `o_proj` is *inside* the quantised set and the +*indexer* is outside it. Pricing the indexer at fp8 halves the weight traffic of +the one attention node whose cost grows with context β€” and at 1M context that node +is 54 % of the step (Β§4.1). This is why the planner grew +`op_dtype_overrides` (Β§7, G5). + +--- + +## 2. Per-phase execution diagrams + +### 2.1 Prefill β€” a chunk of P tokens against C cached + +```mermaid +%%{init: {'theme':'neutral'}}%% +flowchart TD + T["input_ids"] --> EMB["embed_tokens gather"] + EMB --> L0["layers 0-2 β€” MLA+DSA + DENSE FFN"] + L0 --> LB["layers 3-77 β€” MLA+DSA + MoE 256/top-8"] + LB --> FN["final RMSNorm β€” LAST TOKEN OF EACH PROMPT ONLY"] + FN --> LM["lm_head BF16 β€” 239.5 MB for one row per request"] + + subgraph LB["one MoE layer (node list: Appendix A.1)"] + direction TB + QKV["q_a β†’ q_b Β· kv_a β†’ latent[576] β†’ CACHE WRITE Β· kv_b
fp8, M=P β€” COMPUTE-BOUND"] --> IX + IX{"full indexer layer?"} + IX -- "21 layers" --> IS["index_score over the WHOLE history
O(PΒ·C + PΒ²/2) Γ— 32 heads β€” the quadratic lives HERE"] + IX -- "57 layers" --> RE["reuse the group's selection β€” NO KERNEL"] + IS --> ATT + RE --> ATT{"attention core over ≀2048 selected keys
FLOPs capped · BYTES ARE NOT"} + ATT -.->|"⚠ bytes = whole cache once per REQUEST;
a tiled kernel re-reads per block (A9/Q12) β€” bounded at 1.1x"| ATT + ATT --> AR1{{"o_proj β†’ all_reduce β€” 174 MB, BANDWIDTH-bound"}} + AR1 --> G["router GEMM FP32 β†’ fused gating β†’ top-8 of 256
DATA-DEPENDENT SHAPE"] + G --> A2A{{"EP dispatch all-to-all β€” 1.39 GB/layer, the largest single term"}} + A2A --> EG["grouped GEMM fp8 Γ—3 + SwiGLU β€” ALL 256 experts hit"] + EG --> A2B{{"EP combine β†’ all_reduce #2"}} + end +``` + +**The structural claim of prefill:** with 256 experts and top-8, a chunk of P +tokens issues `8P` tokenβ†’expert assignments. Once `8P ≫ 256` β€” P above a few +hundred β€” every expert receives at least one token, so **every layer reads its +entire expert bank**, 1.26 GB per rank per layer at EP8, **constant in P**. That is +95.7 GB per pass, 23 % of all prefill traffic. But it is not the top line: **the EP +all-to-all is**, at 105.7 GB and 44 % of predicted prefill time. Three quarters of +prefill cost is the MoE path, and under expert parallelism most of that is *wire*, +not DRAM. + +**And DSA inverts the usual prefill story.** In a dense model prefill attention is +the `O(PΒ²)` term. Here the *core* is capped at 2,048 selected keys per query, so it +is linear in context past 2,048 β€” and the quadratic has moved into the **indexer +scan**, which IndexShare then pays on only 21 of 78 layers. `index_topk` bounds the +core's FLOPs in both phases; it does **not** bound its bytes at prefill (Β§2.2). + +### 2.2 Decode β€” steady state, B sequences, one token each + +Identical node set to prefill. Four nodes change **kind**: + +| operator | prefill class | decode class | why the class itself changes | +| ------------------- | ------------------------------------- | ----------------------------------------------------- | -------------------------------------------------------------------- | +| attention core | tiled over query blocks, causal | **paged decode attention** over a top-k block table | one query row against a gathered selection; no tiling over queries | +| indexer scan | `O(PΒ·C + PΒ²/2)`, **compute-bound** | `O(BΒ·C)`, **memory-bound** β€” streams the whole key set | the query count collapses from P to B; the key set does not | +| **attention bytes** | whole cache, once **per request** | **top-2048 window, per sequence** | prefill queries' selections union to everything; one query's does not | +| every GEMM | M = P β‰ˆ 8192, compute-bound | M = B, **weight-streaming** | AI falls ~1,400 β†’ ~60; same kernel name, different regime | +| collectives | **bandwidth** β€” 174 MB, 30.5 ms wire | **latency** β€” 688 kB, a ring floor | payload 250Γ— apart; the EP a2a goes from the top line to 2.8 % | +| `lm_head` | one row per **request** | **every row, every step** | the epilogue is free in prefill and is not in decode | + +Plus one epilogue change: **in prefill only the last position of each prompt runs +`lm_head`. At decode every row is a last position**, so it reads 1.9 GB of bf16 +vocabulary weights *every step* rather than once per request. + +``` + [hidden BF16 BΓ—6144] + β”‚ + RMSNorm ─▢ q_a fp8 (REPLICATED, 2048) ─▢ q_a_layernorm ─▢ q_b fp8 (64Γ—256) + β”‚ + kv_a fp8 ─▢ latent[512] + rope key[64] ─▢ KV-cache APPEND (576 elems/seq) + β”‚ + β”Œβ”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ full-indexer layer (21 of 78)? β”‚ + β”‚ yes β†’ wq_b/wk/gate BF16, then β”‚ ← the ONLY term that grows with S. + β”‚ score the WHOLE history β”‚ 0.9 % of the step at 8K, + β”‚ BΓ—CΓ—128 bytes, 32 heads, β”‚ 13.0 % at 128K, 54.4 % at 1M + β”‚ no β†’ reuse. NO KERNEL. (57 layers) β”‚ + β””β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + attention core over ≀2048 selected entries Β·Β· 42 MB/layer β€” FLAT IN CONTEXT + β”‚ and NOT divided by TP + o_proj fp8 (16384β†’6144) + β”‚ + all_reduce #1 Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β·Β· 688 kB Β· LATENCY-bound [stream: unresolved] + β”‚ + RMSNorm ─▢ router GEMM FP32 ─▢ sigmoid+bias ─▢ top-8 of 256 + β”‚ + β”œβ”€β–Ά per-expert histogram ─▢ ❓ D2H sync ❓ Β· S1 Β· 76/token if real Β· conf LOW + β”‚ ← the ONLY data-dependent shape + β”‚ in the graph. Blocks CUDA-graph + β”‚ capture. Β§5 rank 2. + EP dispatch a2a ─▢ permute ─▢ grouped GEMM fp8 Γ—2 ─▢ SiLU ─▢ grouped GEMM fp8 + β”‚ 163/256 experts woken at B=32 Β· 771 MB/layer Β· 85 % of DRAM + scatter-add Γ—2.5 ─▢ EP combine a2a ─▢ all_reduce #2 + β–Ό + … Γ—78 layers, then: + final RMSNorm (ALL B rows) ─▢ lm_head 1.9 GB BF16 ─▢ sample ─▢ D2H ─▢ scheduler gap +``` + +### 2.3 Encoders β€” there are none, and the absence is worth stating + +GLM-5.2 is a **text-only** decoder: no vision tower, no audio encoder, no +multimodal scatter into `inputs_embeds`. The absence is load-bearing three times +over. **Prefill starts at `embed_tokens`** β€” the encoderβ†’backbone seam that +dominates a multimodal prefill does not exist. **Prompt length is the only +input-side variable**, so Β§4.4 is one row shorter and the rest better constrained. +And **no second model is hiding off-checkpoint**, so every FLOP in a trace should +map to a node in Β§3 β€” which makes an unexplained kernel block a far stronger +signal than it would be elsewhere (Β§6.4 row 6). + +The GLM family does ship vision variants. They are a different checkpoint with a +different `model_type`, and `is_glm_moe_dsa_config` declines them rather than +pricing a tower it never read. + +### 2.4 MTP-on decode β€” draft and verify + +```mermaid +%%{init: {'theme':'neutral'}}%% +flowchart LR + subgraph V["VANILLA DECODE"] + direction TB + v1["1 row per seq"] --> v2["78 layers
1,591 nodes"] --> v3["lm_head"] --> v4["sample"] + v4 --> v5(("1 token")) + end + subgraph D["DRAFT β€” 5 SERIAL stages, 115 nodes"] + direction TB + d0["h at last accepted pos
+ its token"] --> eh["eh_proj [12288β†’6144] BF16"] + eh --> d1["MTP block
MLA+DSA (shared index) + FULL MoE"] + d1 --> dl1["lm_head 1.9 GB"] --> ds1["argmax β†’ D2H"] --> d2["stage 2 …"] + d2 --> dt(("5 draft tokens")) + end + subgraph W["VERIFY = decode at 1+D rows"] + direction TB + w1["6 rows per seq"] --> w2["THE SAME 78 layers
THE SAME 1,591 nodes"] + w2 --> w3["lm_head, 6 rows"] --> w4["compare vs draft"] + w4 --> w5(("1..6 accepted")) + end + v5 -.->|replaced by| dt + dt --> w1 + w5 -->|"h of last accepted"| d0 +``` + +**Verify is not a new graph.** It is the decode graph with the row dimension +multiplied by `1+D`. The only genuinely new subgraph is the draft chain β€” 115 nodes +against the backbone's 1,591. + +**The dependency point:** the loop is strictly serial and now has `D+1` sampling +points instead of one. Nothing in stage `k` can start before stage `k-1`'s token id +exists. **Those five gaps are architectural** (Β§6.1) and no scheduling closes them. + +**And here GLM departs sharply from the MTP designs this template was written +against.** Per the weight map, the MTP module is `enorm` + `hnorm` + `eh_proj` +`[12288, 6144]` BF16 + one MLA+DSA attention block + **a full `mlp.experts.*` bank +of 256 experts**. It carries no indexer (consistent with +`index_share_for_mtp_iteration: true`) and no `lm_head` of its own β€” it shares the +backbone's. So the intuition that "the draft is a small dense copy of the big +model" is **wrong here**: the draft is one *full* MoE layer, and it draws on the +expert bank once per stage. Β§3.3 puts a number on it. + +--- + +## 3. Predicted execution graph + +Detailed enough to put a trace next to. The 79 blocks collapse to **four +archetypes** (Β§1) plus a prologue, an epilogue and, under MTP, a 115-node draft +chain. All figures at **B=32, S=8192, TP8/EP8, FP8 weights and KV, per rank** +unless stated. **Per-node tables are in Appendix A**; what stays here is argued +rather than looked up. + +### 3.1 Prologue and epilogue + +The step does not begin at layer 0 or end at layer 77. + +| id | operator | kernel class | shape | FLOPs | bytes | bound | +| --- | ------------------- | ------------------------- | ------------------------------ | -------- | ------------- | ------ | +| D0 | `embed_tokens` | gather (index_select) | `[32] β†’ [32,6144]` | 0 F | 0.393 MB | launch | +| E0 | `rms_norm` | fused RMSNorm | `[32,6144]`, **logits rows only** | 0.6 MF | 0.786 MB | launch | +| E1 | `lm_head` | GEMM (BF16), tall-skinny | `[32,6144] Γ— [6144,19360]` | 7.61 GF | **239.5 MB** | memory | +| E2 | `logits_all_gather` | collective (all-gather) | `[32,19360] β†’ [32,154880]` fp32 | 0 F | **17.3 MB** | memory | + +**Both vocabulary tensors shard, not just `lm_head`.** vLLM builds the input +embedding as `VocabParallelEmbedding`, split by vocabulary across TP ranks the +same way the output projection is, so `model_weight_bytes` divides both by `tp` β€” +as the dense-MoE and hybrid families already do. A framework that replicated the +input table instead would hold `tp Γ— 1.9 GB / 2` more per node. + +**D0 reads what it selects, not the table** β€” 393 kB at 32 rows; the resident +1.9 GB matters for the fit math, not the step. **E0 runs over `logits_rows`**, so +at prefill it is one row per prompt and not the chunk. **E2 exists because E1 is +vocabulary-sharded**: 17.3 MB gathered across 8 ranks before anything can be +sampled β€” small in bytes, unavoidable in position, and at 19.3 Β΅s the second most +expensive single node in a decode step after the expert bank. + +### 3.2 What prefill changes + +Same node ids, same order; `M` becomes `P`, and the four nodes in Β§2.2's table +change kind. At P = 8,192 in one chunk: + +- **Every projection crosses into compute-bound** β€” `q_a`/`o_proj` at AI 1,403, + `q_b` 963, `kv_a` 488, all past the fp8 ridge of 412. +- **The indexer scan flips from memory to compute by two orders of magnitude**: + the same keys, read once and scored by 8,192 queries instead of 32 β€” 5.77 TF + against 22 MB. +- **The attention core's bytes and FLOPs stop moving together.** FLOPs stay capped + at 2,048 keys per query; bytes do not, because 8,192 queries each select a + *different* 2,048 and their union is the whole cache. `index_topk` bounds + prefill FLOPs, not prefill bytes β€” a path copied from a dense family charges + `P Γ— index_topk` and understates long-context prefill traffic by C/2048. +- **The collectives change character.** `all_reduce` goes 688 kB β†’ 174 MB each; + the EP all-to-all 5.5 MB/layer β†’ **1.39 GB/layer**, 105.7 GB per pass. +- **The epilogue inverts:** `lm_head` reads 239.5 MB to produce one row per + *request*, where at decode every row is a last position. + +**The chunking is an assumption, and the vendor recipe does not pin it.** Nothing +in that `vllm serve` line sets `--max-num-batched-tokens`, so the figures below β€” +which assume the whole 8,192-token prompt arrives as **one chunk** β€” are the +best case. Confirm the engine's actual value (C6); the same prompt at a 2,048 +default costs **707 GB and 319 ms**, not 422 GB and 264 ms. + +| Prefill, P=8192, C=0, TP8/EP8, per rank | value | +| --- | --- | +| predicted floor | **264.0 ms** for the chunk (31.0 k tok/s) | +| FLOPs | **118.4 TF** | +| bytes β€” **HBM** | **289.0 GB** | +| bytes β€” **interconnect** | **133.2 GB** | +| AI against HBM | **410**, against an fp8 ridge of **412** β€” *at* the ridge, not below it | +| where the time goes | **wire 148.0 ms (56 %)** Β· HBM + compute 116.0 ms (44 %) | + +**Two byte pools, and only one of them answers to the HBM ridge.** An earlier +version of this table divided FLOPs by *all* 422 GB and reported AI 281, which +mixed 133 GB of NVLink payload into a denominator the ridge derives from HBM +bandwidth. Against HBM alone the pass sits at AI 410 β€” balanced on that axis to +within half a percent β€” and the thing actually setting the floor is the wire. + +> ⚠ **Dense intuition says prefill is compute-bound. Under EP8 it is +> communication-bound**, and the two claims are not close: 56 % of the predicted +> floor is interconnect time. On the HBM axis the pass is balanced (AI 410 vs +> ridge 412), so neither "compute-bound" nor "memory-bound" is the right label for +> it β€” the label is *comm*. Two structural reasons: 256 experts Γ— top-8 reads the +> whole bank per layer regardless of P, and expert parallelism turns the dispatch +> into an all-to-all that is 105.7 GB per pass. +> +> `confidence: high for the arithmetic, medium for the conclusion` β€” the soft +> links are that essentially all 256 experts are hit (A5) and that EP is on at all +> (Q1). Under TP8-only the wire term largely disappears and the pass reverts to +> HBM-bound; that is the same fork Β§4.2 and capture C5 turn on. + +**Chunk size multiplies the whole MoE term.** The same 8,192 tokens: + +| chunking | 1 Γ— 8,192 | 2 Γ— 4,096 | 4 Γ— 2,048 | 8 Γ— 1,024 | 64 Γ— 128 | +| --- | --- | --- | --- | --- | --- | +| bytes | **422 GB** | 517 GB | 707 GB | 1,088 GB | **6,313 GB** | +| floor | 264 ms | 281 ms | 319 ms | 396 ms | **1,543 ms** | + +**14.9Γ— the bytes for identical FLOPs** β€” the bank is read per *chunk*. The rule is +derivable rather than tuned: the expert bank costs `95.7 GB Γ— ceil(P/C)` per prompt +whatever C is, so prefill bytes scale as `1/C` until C is small enough that the +per-chunk activation terms stop mattering. Pick `--max-num-batched-tokens` as large +as decode latency tolerates; there is no prefill-side reason to make it small. + +### 3.3 MTP β€” the whole-step economics + +At B=32, S=8192, D=5 (the vendor recipe's `num_speculative_tokens`): + +| pass | backbone | + draft | = nodes | bytes | floor | +| --- | ---: | ---: | ---: | --- | --- | +| vanilla decode (D=0) | 1,591 | 0 | **1,591** | 67.34 GB | 16.254 ms | +| MTP step (D=5) | 1,591 | 115 (5 Γ— 23) | **1,706** | 112.77 GB | **28.135 ms** | +| β€” the verify pass alone | 1,591 | β€” | 1,591 | 106.43 GB | 26.652 ms | + +Verify is the backbone at **6Γ— the rows and the same node count** β€” more work per +kernel, not more kernels. That is why its bytes rise (106.43 vs 67.34 GB) while +its node count does not move. +| β€” the draft chain alone | β€” | 115 | 115 | 6.34 GB | 1.483 ms | + +The backbone is the same **1,591 nodes in every row** β€” the 78 transformer layers +plus prologue and epilogue. Verify is that backbone at 1+D rows, not a second +graph. A draft stage is **23 nodes**: the 20 of a shared-indexer MoE layer (A.1) +plus `rms_norm`, `mtp_eh_proj` and `lm_head` (A.4), identical every stage, so +`5 Γ— 23 = 115` is exact. + +**At D=0 the draft head does not run at all.** `num_nextn_predict_layers: 1` says +the block is in the checkpoint; it does not say the engine executes it. Without a +speculative config nothing is drafted, so no stage is emitted β€” the weights stay +resident and none of their kernels launch. + +**Cost ratio 1.73Γ— for up to 6 tokens.** Where the extra 45.43 GB goes: + +- **Verify, +39.1 GB**, almost all one line: the expert union saturates, so 6Γ— the + rows costs 1.57Γ— the expert bytes (163 β†’ 256 distinct). KV read does **not** + move β€” 0.43 GB either way, read per *sequence* β€” and neither does `lm_head`. +- **The draft chain, +6.34 GB**, all of it new work. A stage is **0.297 ms** and + the chain is 5 of them: `moe_routed` 0.161 ms, `lm_head` 0.050, `mtp_eh_proj` + 0.032, the rest of the attention block 0.054. The expert bank is **54 %** of a stage, + not all of it β€” the vocabulary projection (17 %) and the `[12288β†’6144]` fusion + (11 %) are another 28 % together, and neither gets cheaper for being a draft. + Itemised per node in **Appendix A.4**. Each of 5 stages draws on a full 256-expert bank, + linear in D with no saturation to help, so **the draft is 5.3 % of the MTP step + where a dense-draft model's would be 1–2 %**. + +**~86 % of the price of speculation is the MoE expert bank** β€” charged because +more rows and stages touch more experts, not because more work is done per token. + +**Break-even.** Accepted tokens follow a **prefix chain**: the verifier walks the +draft in order and stops at the first rejection, so token *k* counts only if +1…*k*βˆ’1 did. The expectation is `Ξ£ αⁱ = (1βˆ’Ξ±^(D+1))/(1βˆ’Ξ±)`, not `1 + DΒ·Ξ±` β€” that +would be the answer for *independent* draws, and at D=5 it overstates accepted +tokens by 1.8Γ— at Ξ±=0.5. + +| Ξ± | 0.0 | 0.5 | 0.7 | 0.9 | break-even | +| --- | --- | --- | --- | --- | --- | +| accepted tokens/step, `(1βˆ’Ξ±βΆ)/(1βˆ’Ξ±)` | 1.000 | 1.969 | 2.941 | 4.686 | **1.731** | +| **tok/s** = `32 Γ— Ξ£ αⁱ Γ· 28.135 ms` | 1,137 | **2,239** | **3,345** | **5,329** | **Ξ± > 0.426**, where this row crosses 1,969 | +| MTP off, for comparison = `32 Γ· 16.254 ms` | 1,969 | 1,969 | 1,969 | 1,969 | β€” | + +Each row divides its own token count by its own step time, which is what makes +them comparable: break-even is where they meet, at `Ξ£ αⁱ = 28.135/16.254 = 1.731`. + +`BatchConfig.tokens_per_step` computes the chain, so +`gitm plan --spec-tokens 5 --acceptance-rate 0.5` prints the 2,239 above. It used +to compute `1 + DΒ·Ξ±` β€” an earlier revision of this note carried a warning about +its own tool's output, which was the wrong place to fix it. The linear form was +shared with every family and wrong for all of them, since any verifier accepts a +prefix; a tree-attention scheme verifying several candidate continuations at once +would need its own term and does not have one here. + +Z.ai claims the GLM-5.2 MTP layer raises accepted length up to 20 % over its +predecessor, which likely clears even the chained bar β€” but **Ξ± is a serving +observable this graph does not predict** (Β§6.2, C4). + +βš‘ **All of this assumes the step is memory-bound.** At B ≀ 8 it is not (Β§4.1), and +in the launch regime the draft's 115 extra launches are pure cost against a step +that was never moving bytes. The sign of the MTP decision flips with batch. + +### 3.4 Predicted synchronization points + +| # | Where | Kind | conf | Trace signature if real | +| ------ | --------------------------- | ----------------------------------------------------------------- | -------------------------- | ------------------------------------------------------------------------------------ | +| **S1** | after the gating kernel | host readback of the expert histogram to size the grouped GEMM | **low** | **76 D2H per decoded token** β€” fatal for graph capture. One count resolves **both** rank 2 and rank 3 (Β§5.2) | +| S2 | around each all-reduce | stream-to-stream event wait | medium | 158 event pairs/step; at 688 kB **the gap *is* the cost** | +| S3 | around each EP all-to-all | dispatch/combine barrier | medium | 76 more, and **none on the 3 dense layers**; an imbalanced rank stalls every other one | +| S4 | sampling / detokenisation | D2H of sampled ids every step | **high** | one D2H + host round-trip per step; unavoidable, but its *placement* decides overlap | +| S5 | scheduler / block manager | host-side work between steps | medium | a CPU-shaped gap between steps, growing with batch churn | +| S6 | after each draft `argmax` | D2H of drafted ids, Γ—D per step | medium | 5 extra D2H + host round-trips; at B=1 can exceed the draft's own compute | +| S7 | verify β†’ accept/reject | **host-visible, variable-length** result | **high** it exists | a small kernel + a D2H whose *value* decides how far the sequence advanced | +| S8 | KV rollback | discard rejected rows | **low** on mechanism | pointer rewind (free) or real memmove (not free) β€” the trace tells you which | +| S9 | indexer selection handoff | the `full` layer's top-k must be visible to its 3 `shared` layers | **low** | if it round-trips the host, IndexShare costs a sync it should not β€” **21 per step** | + +**S5** is the decode-specific one worth chasing: ~1,591 kernels in ~16 ms, then +control returns to a Python scheduler β€” if the scheduler is slower than the step, +no kernel-level work matters. **S7** is the one that breaks CUDA graphs: MTP adds a +per-step, host-visible, data-dependent sequence length, so a capturing stack needs +two shapes plus a padded accept path, or no capture. **S9** is GLM-specific and +cheap: if the `full` layer's selection round-trips the host, IndexShare costs 21 +syncs a step to save 57 kernels. + +--- + +## 4. Execution-bound / roofline hypotheses + +Five labels: **compute Β· memory-bandwidth Β· communication Β· launch/sync/latency Β· +mixed**. Every row names its precision, the peak it is bounded against, and the +variable that flips it. Labels are **against peak** β€” a realistic achievable +fraction is never used to move a row across a boundary. + +Two notations on purpose: Β§4.1 is a **node** table (decode concentrates cost, so +the question is which nodes own the step), Β§4.2 a **region** table (for prefill and +MTP the question is what flips the label, which needs a flip column, not a +ranking). Β§4.3 reverse-indexes that column. + +### 4.1 Decode as a node table β€” B=32, S=8192, TP8/EP8, FP8, per rank + +Ridge 412 F/B (fp8) Β· 206 (bf16) Β· 14 (fp32); launch floor 2.0 Β΅s (graph-replay). +Shapes and dtypes per node are in Appendix A.1; this is the same nodes sorted by +what they cost. + +| node | bytes | AI | Γ—N | Ξ£ ms | bound | share | +| --- | ---: | ---: | ---: | ---: | --- | ---: | +| `moe_routed` | 771.3 MB | 3.1 | 75 | **12.052** | memory | **74.1 %** | +| `attn_score_value` | 42.0 MB | 12.8 | 78 | 0.682 | memory | 4.2 % | +| `moe_all_to_all` | 5.5 MB | β€” | 75 | 0.459 | comm | 2.8 % | +| `rms_norm` | 1.6 MB | 0.8 | 157 | 0.314 | launch | 1.9 % | +| `act_quant` | 0.6 MB | 0.7 | 156 | 0.312 | launch | 1.9 % | +| `moe_router` (GEMM + gating) | 3.4 MB | 15.0 | 150 | 0.300 | launch | 1.8 % | +| `attn_q_a` Β· `attn_out_proj` | 13.1 MB each | 61.4 | 78 each | 0.213 each | memory | 1.3 % each | +| 9 more nodes at the launch floor | <5 MB | β€” | 78 each | 0.156 each | launch | 1.0 % each | +| `attn_index_score` | 33.6 MB | 64.0 | 21 | 0.147 | memory | 0.9 % ← the only term that grows with S | +| `lm_head` Β· `logits_all_gather` | 239.5 / 17.3 MB | β€” | 1 / 1 | 0.050 / 0.019 | memory | 0.4 % total | +| **1,591 nodes** | **67.34 GB** | | | **16.254** | | **1,969 tok/s** [A8][EP] | + +[EP] this table prices **EP8**, which the vendor recipe does not ask for β€” it sets +no `--enable-expert-parallel`. Under TP8-only the `moe_all_to_all` row disappears +and the per-rank expert bank doubles instead: a different graph, not a corrected +one (Q1, capture C5). Every EP-dependent figure in Β§3 and Β§4 is conditional on +that flag. + +[A8] every MoE byte term assumes `ep_imbalance = 1.0`. Real skew touches *fewer* +distinct experts, so the prediction over-states traffic and therefore over-states +time: **the throughput figures are conservative**, not optimistic. What skew adds +instead is grouped-GEMM tail latency, which this graph does not model at all. + +| facet | nodes | Ξ£ ms | share | | +| --- | ---: | ---: | ---: | --- | +| memory | 434 | 13.940 | 85.8 % | five node types | +| launch | 1,157 | 2.314 | 14.2 % | 73 % of all nodes, a seventh of the time | +| compute | 0 | 0.000 | 0.0 % | the entire roofline claim, one row | + +**A MoE layer is 20 kernels and two of them cost anything** (Appendix A.1). Two +small rows are kept anyway: the second `moe_router` instance is the fused gating +kernel β€” **the only data-dependent shape in the graph** β€” and `attn_index_score` +is the only term that grows with S, which does not stay small: + +| context S | `attn_index_score` | share of step | step floor | +| --------- | ------------------ | ------------- | ---------- | +| 8,192 | 0.147 ms | 0.9 % | 16.254 ms | +| 131,072 | 2.349 ms | 12.7 % | 18.457 ms | +| **1,048,576** | **18.795 ms** | **53.9 %** | 34.902 ms | + +**At 1M the indexer scan is the largest node in the step** β€” and it is the node +IndexShare already cut 3.7Γ—. "Flat in context" is true of the attention *core* and +false of the step. + +**The batch story, and where the labels flip:** + +| B | floor | tok/s | launch nodes | launch time | compute nodes | +| --- | ---------- | ----- | ------------ | ----------- | ------------- | +| 1 | 3.813 ms | 262 | 1,335 | 2.670 ms = **70 %** | 0 | +| 4 | 5.478 ms | 730 | 1,334 | 2.668 ms = 49 % | 0 | +| 16 | 11.061 ms | 1,447 | 1,157 | 2.314 ms = 21 % | 0 | +| 32 | 16.254 ms | 1,969 | 1,157 | 2.314 ms = 14 % | 0 | +| 64 | 22.028 ms | 2,905 | 1,082 | 2.164 ms = 10 % | 75 | +| 128 | 27.245 ms | 4,698 | 926 | 1.852 ms = 7 % | 75 | +| 256 | 33.977 ms | 7,534 | 695 | 1.390 ms = 4 % | 76 | + +**Below Bβ‰ˆ16 the step is launch-bound in memory-bound clothes.** At B=1 that +70 % already assumes CUDA-graph replay; at the eager 5 Β΅s it is **85 %**, and the +whole low-batch analysis changes sign (A4, Β§5 rank 3). + + + +### 4.2 Prefill and MTP β€” regions and what flips them + +Prefill at **P = 8,192, C = 0**; MTP at **D = 5, B = 32, S = 8,192**. All per rank +at TP8/EP8, FP8. + +| Phase | Region | Bound | Why (point at a number) | Precision / peak | Flip variable | +|---|---|---|---|---|---| +| **Pre** | **EP dispatch/combine all-to-all** | **comm β€” BANDWIDTH** | 1.39 GB/layer Γ— 76 = **105.7 GB**, 117.4 ms = **44.5 % of the step** | BF16 payload | βš‘ **fp8 dispatch halves it**; EP degree; TP-only removes the node and doubles the bank | +| **Pre** | **MoE expert grouped GEMMs** | compute (AI 485) | 1.26 GB/layer **constant in P**, 95.7 GB/pass = 23 % of traffic | **FP8 block-scaled** | **chunk size** (14.9Γ— across 1β†’64 chunks); imbalance | +| **Pre** | **router GEMM** | **compute** | 26 GF/layer at **fp32's 67 TF/s** β†’ 29.0 ms = **11.0 %** | ⚠ **FP32** β€” see Q3 | βš‘ whether the engine runs the GEMM in fp32 or only accumulates there | +| **Pre** | `all_reduce` Γ—158 | **comm β€” BANDWIDTH** | 174 MB each, 27.5 GB/pass = 30.5 ms | BF16 payload | P; below Pβ‰ˆ256 flips to latency | +| **Pre** | projections (`q_a`,`o_proj`,`q_b`,`kv_a`,`kv_b`) | **compute** | AI 1,403 / 963 / 488 vs the fp8 ridge 412 | FP8 e4m3 | prompt length β€” below Pβ‰ˆ512 memory-bound | +| **Pre** | indexer scan Γ—21 | **compute** | 5.77 TF against 22 MB of keys β€” `O(PΒ·C + PΒ²/2)` Γ— 32 heads, and **2.2 % of the step**. **The quadratic lives here, not in the core** | **BF16** arithmetic, fp8 keys β€” two dtypes, one node | P **and** C | +| **Pre** | attention core Γ—79 | compute, **linear in C** | FLOPs capped at 2,048 keys/query; bytes are the whole cache **once per request β€” an optimistic floor** (A9/Q12), but a bounded one: this node is 0.10 % of prefill bytes, so the worst tiling is 1.1Γ— on the step | FP8 KV | βš‘ **`index_topk`**; C; ⚠ **per-request vs per-tile** | +| **Pre** | permute / combine | memory | 8.5 GB/layer-pass on the `8P`-row expanded tensor, near-zero FLOPs | BF16 activations | top-k; **expert imbalance** | +| **Pre** | *everything else* β€” embed gather, norms, `lm_head`, gating | memory or launch | each under 1.5 %; `lm_head` reads 239.5 MB for **one row per request** | BF16; FP32 top-k | none of them flips | +| **Pre** | **whole prefill pass** | βš‘ **comm** under EP8 | **56 % of the floor is wire**. On the HBM axis alone AI 410 vs ridge 412 β€” balanced, not memory-bound | mixed FP8/BF16/FP32 | **EP on/off** (Q1); prompt length; chunk size | +| **MTP** | verify expert GEMMs | **memory** | 256 distinct experts at 192 rows vs 163 at 32 β€” **+57 %/layer, ~85 % of MTP's cost** | FP8 block-scaled | `B(1+D)` vs E=256; **D**; imbalance | +| **MTP** | draft expert GEMMs Γ—D | **memory** | the MTP block carries a **full 256-expert bank**; 5 stages Γ— 163 experts, **linear in D, no saturation** | FP8 block-scaled | **D**; batch | +| **MTP** | draft `lm_head` Γ—5 | memory | 239.5 MB Γ— 5 = **1.20 GB = 19 % of the draft's bytes** | **BF16** | sharded sampling; draft vocab | +| **MTP** | draft `eh_proj` Γ—5 | memory | `[12288,6144]` BF16, **replicated per rank** | **BF16** β€” in `modules_to_not_convert` | whether it is TP-sharded | +| **MTP** | verify attention | **memory, unchanged** | 0.43 GB β€” read **per sequence, not per row**; 1+D rows share one block table | FP8 KV | seq length; explicitly *not* D | +| **MTP** | accept/reject + KV rollback | launch + **host sync** | tiny tensors, but a data-dependent host-visible seq length (S7) | n/a | pointer rewind vs memmove | +| **MTP** | **whole MTP step** | βš‘ **memory above Bβ‰ˆ16, launch below β€” and the two regimes disagree about whether MTP helps** | **1.73Γ—** cost for ≀6 tokens at B=32; break-even Ξ± = **0.426** on a prefix chain (Β§3.3) | mixed | **graph capture**; batch; Ξ±; D | + +**Hardware sensitivity:** nothing flips between H200 SXM and H20 on the compute +rows β€” but H20's much lower FP8 peak moves every prefill projection further into +compute-bound, and its bandwidth moves the decode floor directly. + +### 4.3 Flip-variable index + +Β§4.2's flip column, reverse-indexed to the eight variables that move more than one +row, with the magnitude each is worth: + +| Flip variable | Direction and magnitude | +| --- | --- | +| **Decode batch B** | expert bytes are **sub-linear**: 8 distinct experts at B=1, 163 at 32, 252 at 128. 253 β†’ 7,418 tok/s across 1β†’256, and the whole-step label goes launch β†’ memory at Bβ‰ˆ16 | +| **Sequence length S** | moves the indexer scan and **nothing else**: 0.9 % of the step at 8K β†’ 12.7 % at 128K β†’ **53.9 % at 1M** | +| **Prompt length P** | every prefill projection (memoryβ†’compute above Pβ‰ˆ512), the router, the indexer's quadratic, all-reduce (latencyβ†’bandwidth above Pβ‰ˆ256) | +| βš‘ **Chunk size** | **14.9Γ—** on prefill bytes across 1β†’64 chunks, for identical FLOPs | +| βš‘ **CUDA-graph capture** | 69 % of the B=1 floor, 85 % at eager 5 Β΅s. Decides whether MTP is a 3Γ— win or a net loss | +| βš‘ **EP vs TP** | the a2a is 45 % of prefill under EP8, but the per-rank bank is **8Γ— smaller** β€” a trade, not a cost, since the bank is 85 % of decode DRAM | +| βš‘ **Precision, and KV dtype** | 1.79Γ— on the decode floor and **10.7 β†’ 5.4 H200s** for weights; 55 GB vs 100 GB of KV per rank at 1M | +| **D and acceptance Ξ±** | 1.73Γ— cost at D=5; break-even Ξ± = **0.426**, 2,239 β†’ 5,329 tok/s across Ξ± (Β§3.3) | + +Two more move single rows and are named where they appear: **absorbed-vs-unabsorbed +MLA** (Β±2Γ— on `attn_kv_b` + `attn_out_proj`, Q1) and **expert imbalance** (skew +*reduces* bytes while lengthening the grouped-GEMM tail and stalling EP ranks). + +--- + +## 5. Ranked headroom hypotheses + +Ranked by expected recoverable time Γ— confidence. **All are hypotheses from the +graph**, to be confirmed against a capture. + +| Rank | Region | Prediction | Why | Evidence to inspect | What would prove it wrong | +|---|---|---|---|---|---| +| **1** | **EP all-to-all at prefill** | **β‰₯44 % of prefill time is wire, and roughly half of it is recoverable** | 105.7 GB/pass in BF16. An fp8 dispatch halves the payload outright; overlapping dispatch with the shared-expert GEMM hides more | NCCL kernel duration vs payload at prefill (Β§6.4 row 5); whether dispatch is bf16 | duration β‰ͺ payload/900 GB/s β†’ the engine already fuses or compresses it | +| **2** | **Grouped-GEMM group sizing (the fork)** | either **76 D2H per token** (no graph capture possible) **or** fixed-capacity padding (**all 256 experts read every step**) | The only data-dependent shape in the graph is the fused gating kernel. A stack does one or the other β€” see Β§5.2 | `cuda_api_sum` D2H count per decode step | 0 D2H **and** MoE bytes that track `distinct_experts(B)` β†’ a device-side path, nothing to recover | +| **3** | βš‘ **CUDA-graph capture at low batch** | **63 % of the B=1 floor is launch overhead** (81 % at eager 5 Β΅s) | 1,033 nodes Γ— 2 Β΅s = 2.07 ms against a 1.24 ms memory term | `cuda_api_sum` launch count with `--cuda-graph-trace=node`; expect **1** `cudaGraphLaunch` | already captured β†’ this rank is worth nothing, and rank 2's fork is already resolved to "padded" | +| **4** | **The indexer scan at long context** | at 1M it is **54 % of the step**, and IndexShare's 3.7Γ— is already banked | 90.2 GB/step of index keys at S=1M. Levers: fp8 index keys (2Γ—), temporal reuse across steps, a smaller `index_topk_freq` group | `dram__bytes_read.sum` on the scan vs 90.2 GB; whether keys are fp8 | keys already fp8 and no temporal reuse available β†’ architectural | +| **5** | **`moe_routed` β€” the memory-bound heart** | 74.1 % of decode; **EPLB placement and the grouped-GEMM backend are the levers** | Weight traffic scales with *distinct* experts woken, not with FLOPs. Good placement cuts per-rank distinct traffic; DeepGEMM cuts the constant | measured per-rank expert traffic vs `distinct_experts`; **the real EP imbalance** | measured traffic already at the union prediction with balance β‰ˆ1.0 β†’ the bank is the bank | +| **6** | βš‘ **fp32 router at prefill** | **11.0 % of prefill** rests on the reading that the router *GEMM* runs in fp32 | 26 GF/layer at 67 TF/s. If only the softmax/top-k accumulates in fp32 and the GEMM is bf16, this row shrinks ~15Γ— | the engine's router implementation; `cuda_gpu_kern_sum` dtype of the gate GEMM | the GEMM is genuinely fp32 β†’ architectural, and the row stays | +| **7** | **Decode collective placement** | 158 all-reduces + 76 all-to-alls + 1 logits all-gather per step on the compute stream, **latency-bound**, with idle SMs around them | 688 kB payloads: the gap *is* the cost. Nothing prevents other layers' work overlapping | NCCL ranges and stream ids on the timeline (S2/S3) | already on a separate stream with overlap β†’ nothing to recover | +| **8** | **Chunk size** | a decode-latency-protecting `--max-num-batched-tokens` can cost **14.9Γ— the prefill bytes** | the expert bank is read per chunk | total prefill MoE DRAM Γ· 95.7 GB β€” the quotient **is** the chunk count | quotient β‰ˆ 1 β†’ prefill is already unchunked | +| **9** | **`attn_q_a` / `attn_kv_a` replication** | 2.8 % of decode is paid **in full on every rank** and TP does not reduce it | they produce the shared latent, which has nothing to split | per-rank duration of `q_a` vs `q_b` under TP8 | already sharded via DP-attention β†’ the graph is wrong here, not the engine | +| **10** | **IndexShare selection handoff (S9)** | if the top-k round-trips the host, **21 syncs/step** to save 57 kernels | the selection must reach three downstream layers | D2H count attributable to the indexer region | 0 β†’ device-side, and IndexShare is pure win | + +### 5.1 Gate check + +Every row above maps to a category GitM can observe and act on β€” launch gaps (3), +needless syncs (2, 10), serialised work and stream/overlap misses (1, 7), +collective placement (1, 7), dispatch/combine cost (1, 5), routing imbalance (5), +phase transitions (8), precision selection (4, 6). No row sits outside the list. + +### 5.2 Ranks 2 and 3 are the same fork, not two independent bets + +Either group sizes are resolved on the **host** (rank 2) β€” exact shapes, no +padding, but a sync per layer and no graph β€” *or* the kernel pads to **fixed +capacity** (rank 3): no sync, capturable, but all 256 experts read every step. +**You cannot pay both, and you cannot escape both without a device-side grouped +GEMM.** Counting D2H per step is the cheapest measurement in this document. + +### 5.3 Deliberately excluded β€” architectural, not recoverable + +These look alarming on a timeline and are not actionable: **the five MTP draft +gaps** (stage `k` consumes stage `k-1`'s id β€” a producerβ†’consumer edge no +scheduling closes), **the accept/reject readback** (S7), **the sampling D2H** (S4 β€” +one host round-trip per step is the floor for any autoregressive decoder), **the +KV cache replicated across TP ranks** (one shared MLA latent cannot be split), and +**the indexer scan growing with context** (IndexShare already cut it 3.7Γ—; the +remainder is the cost of selecting from an uncompressed history). + +The draft-gap row covers the **draft chain only**. The verify pass is one backbone +forward with no cross-stage dependency, so its collectives are as overlappable as +any other step's (rank 7) β€” a gap around them is not excused by this section. + +--- + +## 6. Validation plan + +Assume the Nsight Systems / CUPTI trace arrives tomorrow. + +### 6.1 The classification rule β€” *unexpected β‰  recoverable* + +**A precondition, not a footnote.** Steps (1) and (2) below require telling a +serial gap from a parallelisable one, and today the graph cannot: `total_pred_s` +is a sum, and `expected_stream_id` is written but read by nothing (Β§7.1 G10). Until +a capture carries stream assignment, ranks 1 and 7 are judged from the timeline by +hand rather than by the rule. That is a gate on Β§6, not just roadmap work. + +Three questions, in order; the first that answers decides it. +**(1) Is there a producerβ†’consumer edge across the gap?** Yes β†’ **architectural**; +the Β§2/Β§3 graphs exist precisely to answer this without guessing. **(2) Does the +gap scale with something the deployment controls** β€” batch, chunk size, graph +capture, KV dtype, D, EP degree, stream assignment? Yes β†’ **recoverable**, and name +the knob *and* the expected delta. **(3) Would the gap survive a perfect +implementation of the same model?** Yes β†’ **architectural**; no β†’ **recoverable**. + +Two worked examples, because the rule is easy to agree with and hard to apply: + +- **The draft chain shows five gaps with no kernel spanning them.** Q1: *yes* β€” + stage `k` consumes stage `k-1`'s token id. **Architectural.** +- **The 235 decode collectives sit on the compute stream with idle gaps around + them.** Q1: *no* β€” the output feeds the next layer, but nothing prevents *other* + layers' work overlapping. Q2: *yes* β€” stream assignment. **Recoverable** (rank 7). + +**The trap runs in both directions.** The accept/reject readback will look alarming +and is architectural. The 1,591 kernel launches are entirely *expected* from Β§3 and +are the largest recoverable item at low batch. **Neither surprise nor familiarity +is evidence.** + +### 6.2 Capture plan β€” request these before anyone opens a timeline + +| # | Capture | Why | What dies without it | +|---|---|---|---| +| **C1** | Decode, **B ∈ {1, 8, 32, 128}**, S fixed at 8k | the coupon-collector curve, the launch/memory crossover, collective latency share | ranks 2, 3, 5, 7 β€” the entire low-batch story | +| **C2** | Decode, **S ∈ {8k, 131k, 1M}**, B fixed at 32 | separates the indexer scan from everything else | the whole Β§4.1 context table; rank 4 | +| **C3** | Prefill, **P ∈ {512, 8192}** Γ— **chunked / unchunked** | the chunking multiplier and the AI curve | ranks 1, 6, 8 | +| **C4** | **MTP on and off** at identical B and S, **with the engine's acceptance metric** | isolates draft cost from verify cost, and Ξ± is the only number here that cannot be predicted | all of Β§2.4/Β§3.3 | +| **C5** | **TP8-only vs TP8/EP8** at the same B, S | the a2a-vs-bank trade | rank 1, and the EP recommendation | +| **C6** | **The engine's launch arguments and version, as text** | chunk size, graph capture, KV dtype, D, TP/EP, whether MLA is absorbed, router dtype | roughly half of every table in Β§4–§5 | + +**C6 is not a trace and is worth more than most of the traces.** + +``` +nsys profile --trace=cuda,nvtx,osrt,cublas --cuda-graph-trace=node ... +``` + +**`--cuda-graph-trace=node` is load-bearing.** Without it a captured graph appears +as **one** timeline blob and the kernel count β€” the thing rank 3 turns on β€” is +unobservable. + +### 6.3 Instrument map β€” three tools, three questions + +**`nsys`** answers *where are the gaps, syncs and serialisations* β€” timeline, API +calls, launch counts, D2H, NCCL ranges, stream assignment, CPU scheduler time. +**`ncu`** answers *how many bytes did that kernel move* β€” per-kernel DRAM, L2, +achieved bandwidth, tensor-pipe activity. **CUPTI activity records** answer *what +happened across the run, cheaply*, without `ncu`'s serialising replay β€” and GitM's +`spec_decode` bucket already separates the MTP scaffolding from ordinary sampling. + +Counters: `dram__bytes_read.sum`, `dram__bytes_write.sum`, `lts__t_bytes.sum` (L2, +the escape hatch for the expert-bank claim), `gpu__time_duration.sum`, tensor-pipe +active %. Exact names vary by architecture and `ncu` version. + +**Two cautions.** `ncu` serialises kernels and destroys the overlap information +ranks 1 and 7 depend on β€” profile bytes with `ncu`, overlap with `nsys`. And DRAM +counters miss L2: a low reading is ambiguous between "did not read it" and "read +it from cache". + +### 6.4 Trace triage β€” what to measure, in order + +**Both branches are written before the data arrives.** That is the entire point. +Key: **R:** recoverable β†’ the rank it feeds Β· **A:** architectural, do not chase Β· +**F:** whole-model falsifier. + +| # | Measure | Scope | Expected | Deviation β†’ meaning | +|---|---|---|---|---| +| **0** | **Launch args, as text** β€” not a measurement | C6 | chunk size, graph capture, KV dtype, D, TP/EP, absorbed MLA | Resolves or reframes **ranks 1, 2, 3, 6, 8 before a timeline is opened** | +| **1** | `cuda_api_sum` β€” **D2H count per decode step** | C1 | **0** in the MoE region | **R:** 76/token β†’ host-resolved group sizes, no graph capture β†’ **rank 2**. **R:** 0 but MoE bytes flat in B β†’ padded capacity β†’ **rank 3** instead (Β§5.2). **A:** none | +| **2** | `cuda_api_sum` β€” **launches per step**, needs `--cuda-graph-trace=node` | C1 | **1** `cudaGraphLaunch` | **R:** ~1,591 individual launches + CPU gaps below Bβ‰ˆ16 β†’ **rank 3**. **F1:** count wildly off 1,591 β†’ the lowering in Β§3 is wrong by an order of magnitude | +| **3** | `dram__bytes_read.sum` β€” **MoE region** | C1, C3 | 1.26 GB/layer/rank = **95.7 GB/pass**; **β‰ˆ85 %** of decode | **R:** prefill total Γ· 95.7 GB > 1 β†’ chunked re-read, and the quotient **is** the chunk count β†’ **rank 8**. **R:** flat in B β†’ padded capacity β†’ **rank 3**. **F2:** much lower β†’ L2 residency, and the central claim of both phases is wrong | +| **4** | `dram__bytes_read.sum` β€” **indexer region**, swept in S | C2 | 33.6 MB/layer at 8K β†’ **90.2 GB/step at 1M**, on **21 layers only** | **R:** keys read on 78 layers β†’ IndexShare is not being honoured, a hidden 3.7Γ— β†’ **rank 4**. **R:** 2Γ— the prediction β†’ keys are bf16 where fp8 would do. **A:** growth on 21 layers is the architecture | +| **5** | **NCCL kernel duration vs payload** | C1, C3, C5 | prefill **∝ payload** (~900 GB/s, 117 ms a2a); decode **flat, ~2 Β΅s Γ— 235** | **R:** prefill a2a at bf16 payload β†’ fp8 dispatch β†’ **rank 1**. **R:** decode collectives on the compute stream with idle SMs β†’ **rank 7**. **A:** the ring latency floor | +| **6** | **Kernel-name coverage** β€” every kernel maps to a Β§3 node | C1 | **complete** | **A/F:** an unmapped kernel block is not headroom, it is a node this graph does not have β€” and unlike a multimodal model (Β§2.3) there is no external encoder to explain it away, so it is a **model-validity failure** | +| **7** | `cuda_gpu_kern_sum` β€” **attention core duration vs S** | C2 | **flat** from 8K to 1M | **R:** grows with S β†’ `index_topk` is not being applied and the core is reading the whole cache. **A:** flat β€” that is DSA working | +| **8** | **Tensor-pipe active %** | C1, C3 | **near-idle at every decode batch**; prefill well below peak with DRAM and NVLink busy | **R:** pipe busy at low batch β†’ something does far more FLOPs than the graph predicts. **A:** near-idle at decode β€” that is what decode *is* | +| **9** | **CPU thread sampling, between steps** | C1 | no gap between step *N* and *N*+1 | **R:** CPU-shaped inter-step gap with the scheduler hot β†’ scheduler-bound (S5). **A:** the single sampling D2H | +| **10** | **`spec_decode` bucket counts + acceptance** | C4 | **D = 5** draft stages, 5 extra `lm_head`-shaped GEMMs, verify KV **flat** as D rises | **R:** KV scales with 1+D β†’ rows treated as sequences; use a multi-query kernel. **R:** one `lm_head` for five stages β†’ the draft samples on a sharded vocab already. **A:** the five serial gaps | +| **11** | **Router GEMM dtype** | C3 | fp32 if the config is honoured | **R:** bf16 GEMM with fp32 accumulate β†’ **rank 6 evaporates and Β§4.2's 11.0 % row shrinks 15Γ—.** **A:** genuinely fp32 β†’ it is the model | + +**Rows 0–3 are the thirty-minute version.** + +--- + +## 7. GitM planner gaps β€” and what this branch changed + +Read against `GitM-Labs/runtime` @ `main`. Β§7.0 is what the planner already got +right; Β§7.1 is what GLM-5.2 broke; Β§7.2 is the code that now exists. + +### 7.0 What the planner already gets right + +Listed first because several findings here turned out to be things GitM already +models, and proposing them as gaps would waste the pilot's time: `positions` vs +`sequences` (a multi-row verify reads the cache **once per sequence** β€” Β§3.3's +central MTP result, and the planner had it first); the coupon-collector +distinct-expert term in `roofline.distinct_experts`; EP-vs-TP as a collective +trade with `ep_imbalance` **calibrated from a trace, not predicted**; the +three-way compute/memory/**launch** bound via `serial_launches`; fp8 block-scale +overhead at 1.000244 bytes/weight; `has_fallback_peaks` / `has_unpriced_collectives` +as self-reported debt; per-checkpoint `provenance`; explicit per-layer schedules +over modulo rules; and prefill as `rows = positions + prefill_tokens` with +`logits_rows` for the epilogue. + +**This is a planner built by someone who has been wrong about these before.** The +gaps below are narrower because of it. + +### 7.1 The gaps GLM-5.2 exposed + +| # | What needs representing | Why the abstraction broke | The extension | Shipped? | +|---|---|---|---|---| +| **G1** | **Three precisions in one block**: fp8 backbone + experts, **bf16 indexer / `lm_head` / `eh_proj`**, **fp32 router** | `GlmMoeDsaModelSpec` carried one `weight_dtype`, and every `add(...)` passed it. No way to say "this op runs at a different width". The indexer was priced at fp8 β€” **half its real weight traffic on the node that owns 54 % of a 1M-context step** | `op_dtype_overrides: tuple[tuple[str, str], ...]`, consulted by `add()` and by `model_weight_bytes` before the family default. Read from `quantization_config.modules_to_not_convert` + `moe_router_dtype`, never assumed | **yes** | +| **G2** | **Prefill, with DSA asymptotics that invert a dense model's** | The family was decode-only. The obvious fix β€” copy `hybrid_graph`'s prefill path β€” produces a **confidently wrong** graph: `BatchConfig.attention_qk_pairs` is the *dense causal* count, which over-charges the DSA core by `C/index_topk` (64Γ— at 128K) and, worse, under-charges its **bytes**, because at prefill the queries' selections union to the whole cache | `core_qk_pairs` / `core_read_entries` / `index_scan_pairs` / `index_scan_entries` β€” four helpers rather than one, because FLOPs and bytes stop moving together on this architecture | **yes** | +| **G3** | **Two per-layer schedules that do not align** β€” dense/sparse (3 + 75) and full/shared indexer (3 + period-4) | Already handled via `mlp_layer_types` / `indexer_types`, read verbatim. Worth recording as a *near*-gap: a modulo rule fitted to either one alone misplaces layers while producing an entirely plausible total | none needed | n/a | +| **G4** | **An fp32 peak for a modern SKU** | `hardware_spec_for` left `peak_flops_fp32_per_s` at the A100 default (19.5 TF/s) with the comment *"nothing currently predicts fp32 kernels"*. GLM's router does. On an H200 that default is **3.4Γ— low**, enough to move the router's bound label | `_FP32_PEAKS` keyed by the same SKU substrings, CUDA-core rates (H200 67 TF/s), wired through `hardware_spec_for` | **yes** | +| **G5** | **`gitm plan` dropping the launch bound and mispricing the ridge** | `_render_table` recomputed `bound` as compute-vs-memory, **discarding `"launch"` entirely**, and divided the ridge by `peak_flops_bf16_per_s` regardless of op dtype. So **854 launch-bound nodes printed as memory-bound**, against ridge 206 where fp8 answers to 412 | use the node's own `bound`; print one ridge per dtype present in the graph; add a launch-bound count and a `*` marker where an op's instances disagree | **yes** | +| **G6** | **Two collectives per layer, not one** | `_emit_layer` folded the post-attention and post-FFN all-reduces into one node with double the payload. Bytes right, **count wrong** β€” and at 688 kB a decode collective is bounded by its ring latency, so the count *is* the cost | `_emit_collective` called at both sub-block boundaries, emitting `tp_all_reduce_attn` and `tp_all_reduce_mlp` separately, with the EP all-to-all on the MoE half only | **yes** | +| **G7** | **An MTP chain D stages deep, each with its own vocabulary projection** | The graph emitted **one** draft block and **one** `lm_head` for what the vendor recipe runs **five** deep. `lm_head` is 19 % of the draft's bytes, so a D-deep chain was understated by ~5Γ— on its largest term | a stage loop in `predict_glm_graph` driven by `BatchConfig.speculative_tokens`, with `mtp_eh_proj` and an `lm_head` per stage; `--spec-tokens` on the CLI | **yes** | +| **G8** | **A graph that is only its GEMMs, priced against a bound it cannot express** | The family emitted 16 nodes per layer where a layer lowers to ~20 kernels, and the seven missing ones were all pointwise: the norms, the dynamic fp8 activation scaling, the fused gating, the prologue gather and the epilogue's all-gather. Every one is a rounding error in bytes and **a full kernel launch in time** β€” so at B=1 the graph reported a step as memory-bound that is 69 % launches. A roofline with a launch bound and a graph with no launches in it cannot both be right | Emit them. `_pointwise`, `add_rms_norm` (with the residual fused in, as vLLM runs it), `add_act_quant` gated on the consuming GEMM actually being fp8, plus `embed_tokens` / `rms_norm` / `logits_all_gather` around the stack. Node names constrained by G9 | **yes** | +| **G9** | **Op names a capture can actually pair against** | G8's new nodes needed names, and `deviation.classify_op` is a *name guess* (`docs/kernel_identity.md`): a name it cannot classify leaves the predicted node permanently unmatched **and** the real kernel filed as unmodeled β€” two errors in opposite directions, in the diff the family exists to support. Three norm sites are one kernel name; `silu_and_mul` was already claimed by `mlp_gate_up`; `moe_align`/`topk_softmax` were already claimed by `moe_router`, a decision the dense-MoE and hybrid families depend on | Follow the canonical names rather than redefine them: one `rms_norm` op for all three sites, SwiGLU folded back into the GEMM that owns its needle, gating emitted as a second `moe_router` instance. Then `_OP_RULES` gains only what is genuinely new and unclaimed β€” `rms_norm`, `act_quant`, `embed_tokens`, `moe_permute`/`moe_combine`, `attn_index_proj`, `attn_kv_b`, `mtp_eh_proj`. A test asserts every op the graph emits resolves | **yes** | +| **G11** | **An intervention vocabulary that can name the expert term** | `kernels/library.yaml` scopes every lever with `applies_to_kernels`, drawn from a canonical op list that is `qkv_proj Β· attn_score_value Β· attn_out_proj Β· mlp_gate_up Β· mlp_down Β· lm_head`. **`moe_routed` and `moe_shared` are not in it**, so the two entries meant to target expert traffic scope to `[mlp_gate_up, mlp_down]` β€” true of a dense FFN, false of either MoE family. Β§5 rank 5 aims levers at **74 % of a decode step** through tooling that cannot match it | Add the two ops to the vocabulary and re-scope those entries. **Pre-existing and not GLM-specific** β€” `moe_graph.py` emits the same names, so DeepSeek-V4 has it identically | **no** β€” the fix is a shared-vocabulary change and should land where both families' coverage can be checked at once | +| **G10** | **Which adjacent nodes may overlap and which may not** | `Graph.total_pred_s` is a sum, not a DAG. GLM puts both kinds of serialisation in one step β€” the draft chain is genuinely serial, the 158 collectives are not β€” and **both appear as the same positive residual today**. Β§6.1 is unanswerable without telling them apart | *Not shipped.* It is a cross-family IR change and it is already sequenced on the roadmap. What this note adds is the requirement. `expected_stream_id=1` is set on collectives, but note that **nothing reads it today** β€” `optimizer/monitor.py` tests overlap using the *observed* kernel's stream, so the predicted field is carried by the IR and consumed by no one. It is a hook for the invariant in `docs/invariants.md` Β§3, not a wiring of it. **And the field defaults to `0`, which is indistinguishable from an explicit "compute stream"** β€” whoever wires the invariant should make it `int | None` first, or every pre-GLM family silently claims stream 0 | **no β€” deliberately** | + +### 7.2 The one that needed more than a table row + +**G2 must not ship as a copy of another family's prefill path.** Aliasing +`BatchConfig.attention_qk_pairs` into the DSA core is a two-line change producing a +complete, plausible graph β€” and wrong in *both* directions at once: the core's +FLOPs over-charged by context Γ· `index_topk`, its bytes under-charged by the same +ratio, both from the same false premise. **The two mistakes partly cancel in the +total, which is what makes them survivable** β€” and a prefill path that is wrong in +a self-cancelling way is worse than none. Hence four helpers rather than one alias. + +### 7.3 What this branch changed, in kind + +Five things, in the order they re-rank the tables. `git log` has the file list. + +1. **Precision became per-op** (`op_dtype_overrides`), read from + `modules_to_not_convert` and `moe_router_dtype`. Everything downstream is + priced against it, and the FP8 catalogue entry became the one to plan against. +2. **Prefill exists**, with DSA's own asymptotics rather than a dense family's β€” + four helpers, because FLOPs and bytes stop moving together (G2, Β§7.2). +3. **A layer lowers to its kernels, not just its GEMMs** (G8): norms, activation + quantisation, the fused gating, the prologue and epilogue. Without them the + launch bound the roofline already supported had nothing to bind. +4. **Node names follow the pairing contract** (G9) rather than redefining it, so + the per-op residual diff in Β§6 can actually pair what the graph predicts. +5. **The MTP chain is D stages deep**, each with its own vocabulary projection + (G7), driven by `--spec-tokens`. + +Three supporting fixes outside the family: an fp32 peak for modern SKUs (G4), +`gitm plan` keeping the launch bound and pricing the ridge per dtype (G5), and +`BatchConfig.tokens_per_step` counting accepted tokens as a prefix chain rather +than `1 + DΒ·Ξ±` β€” shared with every family and wrong for all of them. + +**Deleted:** three committed JSON node dumps, 29k lines that went stale on every +graph change. The commands at the top of this note regenerate any of them. + +--- + +## 8. Open questions and assumptions + +### 8.1 Open questions, ranked by what they change + +| # | Question | What it changes | How to resolve | +|---|---|---|---| +| **Q1** | Is expert parallelism actually on? | **Rank 1 exists or it does not.** `--enable-expert-parallel` is absent from the vendor recipe; without it the `moe_all_to_all` rows disappear (44 % of prefill) and the per-rank expert bank doubles instead. Not a refinement β€” a different graph | engine launch args (C6); capture C5 | +| **Q2** | Does the engine run **absorbed** MLA at decode? | drops `attn_kv_b` and **doubles** `attn_out_proj`'s input width (16384β†’32768). Β±2Γ— on the #4 and #8 lines | serving image / C6 | +| **Q3** | Is the decode step **CUDA-graph captured**? | at B=1 the difference between 1.24 ms and 3.30 ms per step, and it decides whether MTP is a 3Γ— win or a net loss | engine config + `--cuda-graph-trace=node` | +| **Q4** | Is the **router GEMM** fp32, or only its accumulation? | **11.0 % of prefill.** At bf16 the row shrinks ~15Γ— | the engine's MoE gate implementation | +| **Q5** | Is the grouped GEMM **device-sized or host-sized**? | decides **rank 2 *or* rank 3** β€” the fork in Β§5.2 | trace D2H count, or the kernel source | +| **Q6** | Is the **EP dispatch** bf16 or fp8? | **half of 105.7 GB** at prefill β€” the largest single recoverable number here | serving image / C6 | +| **Q7** | Is the **IndexShare selection** passed device-side? | 21 syncs/step if not (S9) | trace D2H attribution | +| **Q8** | Chunked prefill on, at what chunk size? | **up to 14.9Γ— on prefill bytes** | engine launch args (C6) | +| **Q9** | **Ξ±**, the MTP acceptance rate, in production | the entire MTP decision. Break-even is **0.426**; 0.5β†’0.9 is 2,239β†’5,329 tok/s | engine metrics (C4) β€” **not predictable from a config** | +| **Q10** | Are the **index keys** cached in fp8 or bf16? | 2Γ— on 90.2 GB/step at 1M context | C6 / trace | +| **Q11** | Is `eh_proj` **TP-sharded** in the MTP block? | 151 MB β†’ 19 MB per rank per draft stage | serving image | +| **Q12** | Does the prefill attention kernel read the selected KV once per request, or once per query tile? | Up to 64Γ— on that node β€” but **the node is 0.10 % of prefill bytes**, so even a 128-row tiling takes the step from 422 GB to 448 GB. **1.1Γ—, and it does not move the prefill conclusion.** Listed last because it is bounded, not because it is small | `dram__bytes_read.sum` on the prefill core | + + +### 8.2 Assumptions in force + +| # | Assumption | Status | What would falsify it | +|---|---|---|---| +| **A1** | 8Γ—H200 SXM, NVLink, TP8/EP8, fp8 weights and KV | **From the vendor's own published recipe**, not inferred. If production differs, only the constants section redoes | procurement; C6 | +| **A2** | **Dense FP8 peak β€” halving the datasheet's 3,958 to 1,979** | An inference, held with high confidence. Every tensor-core row is footnoted "with sparsity"; GLM-5.2-FP8 declares no sparsity, and 753.33 GB observed against 755.9 GB predicted **dense** confirms full density β€” a 2:4 checkpoint would be ~half that. Rooflining against the sparse peak would make every region look 2Γ— more memory-bound than it is | a sparsity flag in the serving image, or a sparse GEMM path in the trace | +| **A3** | Collectives are priced bandwidth-plus-one-launch, on an unresolved stream | `estimated=True` throughout; the stream assignment is a guess about a stack nobody has opened, and **rank 7 depends entirely on it** | the trace | +| **A4** | ~2 Β΅s kernel launch (CUDA-graph replay) | Eager is nearer 5 Β΅s. **The factor of 2.5 moves rank 3 from 63 % to 81 % of the B=1 floor** and moves the launch/memory crossover batch | calibrate from launch-to-launch gaps | +| **A5** | Routing is not pathologically concentrated β€” `distinct_experts` assumes a uniform router | Real skew touches *fewer* experts, so this over-predicts traffic β€” the conservative direction. `e_score_correction_bias` exists precisely to spread load | expert-GEMM DRAM read well below 1.26 GB/layer | +| **A6** | The serving path uses a grouped GEMM, not a per-expert loop | Architecture rule; the reference implementation is the *semantics*, not the execution | a per-expert kernel launch pattern in the trace | +| **A7** | 158 collectives per step (2 per layer Γ— 79) | TP convention, now modelled explicitly (G6) | NCCL kernel count per step | +| **A8** | `ep_imbalance = 1.0` | **Declared, not fitted** β€” it is trace-calibrated by design and there are no traces | any measured skew | +| **A9** | The prefill attention core streams the selected cache **once per request** | An optimistic floor (Q12), and a *bounded* one: the node is 0.10 % of prefill bytes, so the worst tiling costs 1.1Γ— on the step | `dram__bytes_read.sum` on the prefill core | +| **A10** | The exact kernel names, everywhere | `confidence: none` throughout. The *class* is justified; the implementation is not knowable without the serving image | β€” | + +--- + +## 9. How to run it + +**Predict-only β€” free, no GPU, and it answers the two questions that gate +everything else.** Does the fp8 shape fit (yes: 755.9 GB of weights on 1,128 GB), +and is the step launch-bound at your batch (yes, below Bβ‰ˆ16). The commands are at +the top of this note. + +**Serve and capture.** The footprint decides the hardware: **fp8 β†’ one 8Γ—H200 +node**, leaving ~370 GB for KV and activations; **bf16 β†’ two nodes** (10.7 H200s +for weights alone). Full 1M context wants B200/B300 for the extra HBM β€” 55 GB of +KV per rank on top of a 96 GB weight share. + +1. An **8Γ—H200 SXM** pod with a **network volume β‰₯ 1 TB** for the 141-shard fp8 + checkpoint. +2. Serve with the vendor recipe quoted in the hardware section β€” Β§4's constants + assume it. Add `--enable-expert-parallel` for the EP8 shape Β§4 prices; + **without it the `moe_all_to_all` rows should not appear at all**, and rank 1 + does not exist. That difference is capture C5. +3. `gitm capture serve` (or `gitm capture attach`) for a bounded decode window. +4. Diff observed-vs-predicted per op. **A residual is a lead, not a defect.** + +--- + +## Appendix A β€” Predicted node tables + +Trace-day reference for Β§3. **B=32, S=8192, TP8/EP8, FP8 weights and KV, per +rank.** The 78 transformer blocks are `Ld,f` Γ—3 + `Ls,f` Γ—18 + `Ls,sh` Γ—57, plus +`Lmtp` Γ—D when drafting is on; A.2–A.4 +are **deltas** from A.1, since everything unlisted is byte-for-byte identical. + +Two columns are omitted rather than repeated. Every node runs on the compute +stream except the collectives (`expected_stream_id=1` β€” a declaration, not yet a +check, Β§7.1 G10). Confidence is **high** throughout, these rows being read from +`config.json` and the tensor index, except the collectives (**medium**, a TP/EP +convention) and the S1 histogram readback (**low**, a hypothesis). + +### A.1 β€” Archetype `Ls,sh`, 57 layers (shared indexer + MoE) + +Read straight off the graph, in issue order β€” every row is a `PredictedNode` at +this shape, and `tests/test_glm_graph.py::test_layer_lowers_to_the_documented_node_sequence` +pins this exact sequence so the code cannot drift from the table. + +**Three op names repeat** (`rms_norm` Γ—2, `act_quant` Γ—2, `moe_router` Γ—2): one +kernel name in a trace is one op here, or the node goes unpaired and the kernel +files as unmodeled. Which instance a launch belongs to is an NVTX question, never +a name question β€” `docs/kernel_identity.md`. + +| id | operator | kernel class | FLOPs | bytes | dtype | t (Β΅s) | bound | +|---|---|---|---|---|---|---|---| +| .1 | `rms_norm` | `fused_add_rms_norm` β€” input norm **with the residual carried in** | 1.2 MF | 1.573 MB | BF16 | 2.00 | launch | +| .2 | `act_quant` | dynamic FP8 quant + per-row scale | 0.4 MF | 0.590 MB | BF16β†’FP8 | 2.00 | launch | +| .3 | `attn_q_a` | GEMM, **replicated** (6144β†’2048) | 805.3 MF | 13.110 MB | FP8 | 2.73 | memory | +| .4 | `attn_q_b` | GEMM, head-sharded (2048β†’2048) | 268.4 MF | 4.457 MB | FP8 | 2.00 | launch | +| .5 | `attn_kv_a` | GEMM + cache append (6144β†’576) | 226.5 MF | 3.990 MB | FP8 | 2.00 | launch | +| .6 | `attn_kv_b` | GEMM, head-sharded (**unabsorbed**) | 117.4 MF | 2.098 MB | FP8 | 2.00 | launch | +| .7 | `attn_score_value` | paged decode attn over ≀2048 entries | 536.9 MF | 41.951 MB | FP8 KV | 8.74 | memory | +| .8 | `attn_qnorm_rope_insert` | fused q/kv norm + partial RoPE + insert | 0.4 MF | 0.410 MB | BF16 | 2.00 | launch | +| .9 | `attn_out_proj` | GEMM, tall-skinny (16384β†’6144) | 805.3 MF | 13.110 MB | FP8 | 2.73 | memory | +| .10 | `tp_all_reduce_attn` | NCCL ring | β€” | 0.688 MB | BF16 | 2.00 | launch | +| .11 | `rms_norm` | `fused_add_rms_norm` β€” post-attention | 1.2 MF | 1.573 MB | BF16 | 2.00 | launch | +| .12 | `moe_router` | GEMM, **replicated** (6144β†’256) | 100.7 MF | 6.701 MB | **FP32** | 2.00 | launch | +| .13 | `moe_router` | fused gating: sigmoid + `e_score_correction_bias` + **top-8 of 256** + renorm | 0 F | 0.034 MB | **FP32** | 2.00 | launch ← **the only data-dependent shape** | +| | ⚠ **Blocks CUDA-graph capture. β†― S1: host readback of the expert histogram? 76 D2H per token if real.** | | | | | | | +| .14 | `act_quant` | dynamic FP8 quant, expert input | 0.4 MF | 0.590 MB | BF16β†’FP8 | 2.00 | launch | +| .15 | `moe_shared` | grouped GEMM Γ—3, always on | 302.0 MF | 5.539 MB | FP8 | 2.00 | launch | +| .16 | `moe_permute` | gather into expert-major order | 0 F | 0.442 MB | BF16 | 2.00 | launch | +| .17 | `moe_routed` | grouped GEMM Γ—3 + SwiGLU, **163 distinct of 256** | 2,416.2 MF | **771.324 MB** | FP8 | **160.69** | **memory** | +| .18 | `moe_combine` | scatter-add Γ— `routed_scaling 2.5` | 3.1 MF | 0.442 MB | BF16 | 2.00 | launch | +| .19 | `moe_all_to_all` | EP dispatch + combine | β€” | 5.505 MB | BF16 | 6.12 | comm | +| .20 | `tp_all_reduce_mlp` | NCCL ring | β€” | 0.688 MB | BF16 | 2.00 | launch | + +**Ξ£ per layer: 5.6 GF, 874.9 MB, 0.211 ms.** + +**Twenty kernels, two of which cost anything.** `.17` is **88 % of the layer's +bytes and 76 % of its time**; `.7`, `.3`, `.9` and `.19` are most of the rest; the +other **15 sit at the 2 Β΅s launch floor** β€” 0.030 ms per layer, 14 % of it. That is +where Β§4.1's launch facet comes from, and folding any of them into the GEMM it +precedes would report the layer as more memory-bound than it is. + +Two folds are deliberate, because they are where a reader will expect a row and +not find one. **SwiGLU is inside `.17`** (and inside `mlp_gate_up` on the dense +layers): `silu_and_mul` is already `mlp_gate_up`'s needle, so a separate node +would have no kernel to pair against. **The residual adds are inside `.1` and +`.11`**: vLLM runs `RMSNorm.forward(x, residual)` as one `fused_add_rms_norm`. + +### A.2 β€” Archetype `Ls,f`, 18 layers (full indexer + MoE) + +**Delta from A.1: two nodes inserted after `.6`.** Everything else is identical. + +| id | operator | kernel class | FLOPs | bytes | dtype | t (Β΅s) | bound | +|---|---|---|---|---|---|---|---| +| .6a | `attn_index_proj` | GEMM Γ—3 β€” `wq_b` (2048β†’4096), `wk` (6144β†’128), `weights_proj` (6144β†’32), **replicated** | 599.8 MF | 19.937 MB | **BF16** | 4.15 | memory | +| .6b | `attn_index_score` | 32 heads score the whole history + top-2048 | 2,147.5 MF | 33.563 MB | **BF16 math, FP8 keys** | 6.99 | memory | + +`.6a` is BF16 because the indexer is named in `modules_to_not_convert`; at FP8 it +would price at 10.0 MB. **`.6b` carries two dtypes answering different questions** β€” +its *bytes* follow how the keys are stored (fp8), its *FLOPs* what the indexer +computes in (bf16). Invisible at decode, where the node is memory-bound at every +context; at prefill it is the difference between 1.2 % and 2.2 % of the step. + +**`.6b` is also the only node in the model that grows with S**, and it does not +stay small: + +| context S | `.6b` bytes/layer | `.6b` time/layer | Ξ£ over 21 layers | share of step | +| --------- | ----------------- | ---------------- | ---------------- | ------------- | +| 8,192 | 33.6 MB | 6.99 Β΅s | 0.147 ms | 0.9 % | +| 131,072 | 537.0 MB | 0.112 ms | 2.349 ms | 12.7 % | +| 1,048,576 | **4,296.0 MB** | **0.895 ms** | **18.795 ms** | **53.9 %** | + +**Ξ£ per layer at S=8192: 8.3 GF, 928.4 MB, 0.222 ms** β€” 5 % more than `Ls,sh`, and +that 5 % is the whole price of IndexShare's 21-of-78 schedule at short context. + +### A.3 β€” Archetype `Ld,f`, 3 layers (full indexer + DENSE FFN) + +**Delta from A.2: `.12`–`.19` replaced by three nodes.** No router, no gating, no +expert bank, no all-to-all β€” and therefore **no data-dependent shape and no +expert-parallel traffic**. These three are the only blocks in the model a CUDA +graph could capture unconditionally. + +| id | operator | kernel class | FLOPs | bytes | dtype | t (Β΅s) | bound | +|---|---|---|---|---|---|---|---| +| .12β€² | `act_quant` | dynamic FP8 quant | 0.4 MF | 0.590 MB | BF16β†’FP8 | 2.00 | launch | +| .13β€² | `mlp_gate_up` | GEMM (6144β†’3072/rank) **+ SwiGLU** | 1,208.2 MF | 19.665 MB | FP8 | 4.10 | memory | +| .14β€² | `mlp_down` | GEMM (1536/rankβ†’6144) | 604.0 MF | 9.931 MB | FP8 | 2.07 | memory | + +**Ξ£ per layer: 7.3 GF, 167.6 MB, 0.052 ms** β€” **a quarter the time of a MoE layer +at a fifth the bytes.** Three of 78 layers are 0.9 % of the step. (Layer 0's entry +norm is the one `rms_norm` in the model with no residual to carry β€” nothing +precedes it.) + +### A.4 β€” MTP draft chain, per stage (Γ—5 at D=5) + +Rows are `[B,Β·]` = `[32,Β·]`, **not** the verify pass's `[192,Β·]`: the draft proposes +for the sequence, not for the verify rows. + +| id | operator | kernel class | FLOPs | bytes | dtype | t (Β΅s) | bound | +|---|---|---|---|---|---|---|---| +| M.1 | `rms_norm` | `enorm` + `hnorm`, two launches | 1.2 MF | 1.573 MB | **BF16** | 4.00 | launch | +| M.2 | `mtp_eh_proj` | GEMM `[12288β†’6144]`, **replicated** | 4,831.8 MF | **152.175 MB** | **BF16** | 31.70 | memory | +| M.3–M.22 | the whole `Ls,sh` block (A.1) | as A.1 | 5.6 GF | 874.9 MB | mixed | 211 | memory | +| M.23 | `lm_head` | GEMM, vocab-sharded `[6144β†’19360]` | 7,612.7 MF | **239.528 MB** | **BF16** | 49.90 | memory | +| M.24 | `argmax` + D2H | sampling + host round-trip | β€” | small | BF16 | β€” | **sync (S6)** | + +**23 emitted nodes** β€” M.1, M.2, the 20 of M.3–M.22, and M.23. M.24 is a +synchronization point, not a graph node, which is why the stage counts 23 and not +24 in Β§3.3. + +**Ξ£ per stage: 18.0 GF, 1,268.2 MB, 0.297 ms β€” Γ—5 = 6.34 GB, 1.483 ms.** + +Three things there are the whole Β§3.3 argument. **M.3–M.22 is a full MoE block** β€” +the module carries its own 256-expert bank, so 69 % of a stage's bytes are expert +weights re-read every stage, with no saturation to help (32 rows wakes ~163 +experts, five times over). **M.1, M.2 and M.23 are BF16**, all in +`modules_to_not_convert`, and together **31 %** of the stage. And **what is absent** +β€” no `attn_index_proj`, no `attn_index_score` β€” is +`index_share_for_mtp_iteration: true` made visible: the draft inherits the +selection the backbone already paid for. + +### A.5 β€” Node budget for the whole step + +| region | Γ—N | nodes each | Ξ£ nodes | Ξ£ ms | share | +| ------ | -- | ---------- | ------- | ---- | ----- | +| prologue + epilogue | 1 | 4 | 4 | 0.073 | 0.4 % | +| `Ld,f` dense layers | 3 | 17 | 51 | 0.155 | 1.0 % | +| `Ls,f` full-indexer MoE | 18 | 22 | 396 | 3.999 | 24.6 % | +| `Ls,sh` shared-indexer MoE | 57 | 20 | 1,140 | 12.028 | 74.0 % | + +| **total** | | | **1,591** | **16.254** | | + +At D=5 a draft region of 5 Γ— 23 = 115 nodes and 1.483 ms appears, and the backbone +runs at 192 rows instead of 32 β€” **1,706 nodes, 28.135 ms.** At D=0 there is no +draft region at all. + +The prologue/epilogue rows are 4 nodes and 0.4 % of the step, and two of them are +on the critical path between the last layer and the sample: `lm_head` re-reads +239.5 MB of vocabulary weights **every step**, and `logits_all_gather` cannot start +until it finishes. diff --git a/gitm/optimizer/deviation.py b/gitm/optimizer/deviation.py index 3834b90..8b24e03 100644 --- a/gitm/optimizer/deviation.py +++ b/gitm/optimizer/deviation.py @@ -88,6 +88,12 @@ # completely different cost curve; and the indexer must never fall through to # a bare "index" rule, which is how it gets misfiled as elementwise in the # coarse taxonomy. + # The indexer's own projections, before the score entry below β€” whose + # "indexer" needle would otherwise claim ``indexers_proj`` and file a GEMM as + # a scan. They are separate kernels with opposite bounds: on GLM-5.2 the + # projection is bf16 and flat in context, the scan grows with it. + "attn_index_proj": ("indexer_proj", "indexers_proj", "index_proj", "wq_b", + "weights_proj"), "attn_index_score": ("indexer", "lightning_index", "index_topk", "topk_indices"), "moe_shared": ("shared_expert", "moe_shared"), # `topkGating` is vLLM's fused routing kernel. Without it the generic "moe" @@ -96,6 +102,17 @@ # traffic dominates the step. "moe_router": ("moe_align", "topk_softmax", "topkgating", "gating", "router", "routing", "sinkhorn", "expert_bias"), + # Dispatch/gather into expert-major order and the weighted scatter back. + # Before the generic "moe" needle below, which would claim both as expert + # GEMMs β€” they move real bytes and do no arithmetic, so folding them into the + # dominant weight-traffic row hides a term that chunk size and expert + # imbalance both move. + # Combine first: "unpermute" contains "permute", so the reverse direction has + # to be tested before the forward one or every combine kernel files as a + # dispatch. + "moe_combine": ("moe_sum", "finalize_moe", "unpermute", "scatter_add", + "index_add", "moe_combine"), + "moe_permute": ("permute", "expert_sort", "shuffle_rows", "gather_rows"), "moe_routed": ("moe", "expert", "grouped_gemm", "group_gemm", "groupedgemm"), "dspark": ("dspark",), @@ -120,17 +137,42 @@ # checkpoints use (`_triton_mrope_forward`). "attn_qnorm_rope_insert": ("qnorm", "q_norm", "qk_norm", "mrope", "rope", "rotary"), + # ── pointwise work that is its own kernel ──────────────────────────────── + # These used to classify as ``None`` β€” "a norm/activation/copy" β€” and land as + # unmodeled. That was right while no graph emitted them. The GLM-5.2 graph + # does, because on a sparse model at low batch the pointwise kernels are the + # majority of the launches and a step bounded by its launches cannot be + # explained by a graph that only has GEMMs in it. + # + # ``rms_norm`` covers every norm site in a block. They are one kernel name in + # the trace, so they are one op here; which site a given launch belongs to is + # recoverable only from an NVTX range (``docs/kernel_identity.md``), never + # from the name. + "rms_norm": ("rms_norm", "rmsnorm", "layernorm", "layer_norm", "fused_add_rms"), + # Dynamic FP8 activation scaling ahead of a quantised GEMM. "dequant" is + # excluded on purpose: it is the epilogue of the GEMM, not this kernel. + "act_quant": ("scaled_fp8_quant", "per_token_quant", "act_quant", + "quant_fp8", "dynamic_scaled"), + # Before ``lm_head``, whose "embed" needle would otherwise claim the input + # gather and attribute it to the vocabulary projection. + "embed_tokens": ("embedding", "embed_tokens", "index_select"), + # ── projections ────────────────────────────────────────────────────────── "attn_q_a": ("q_a_proj", "q_lora", "q_down"), "attn_q_b": ("q_b_proj", "q_up"), - # `kv_b_proj` is absent on purpose: in the absorbed decode form it is folded - # into the query and output projections, so there is no node to map it to and - # a guess would attribute real work to the wrong op. "attn_kv_a": ("kv_a_proj", "kv_lora", "kv_down", "compress_kv"), + # Unabsorbed MLA only: the absorbed decode form folds W^UK into the query and + # W^UV into the output and launches no such kernel, so these needles match + # nothing there and the op stays absent rather than mis-attributing. + "attn_kv_b": ("kv_b_proj", "kv_up", "w_uk", "w_uv"), "qkv_proj": ("qkv",), "attn_out_proj": ("o_proj", "out_proj", "attn_out"), "mlp_gate_up": ("gate_up", "gate_proj", "up_proj", "swiglu", "silu_and_mul"), "mlp_down": ("down_proj", "mlp_down"), + # The MTP block's [2h, h] fusion of the carried hidden state with the + # embedding of the token just drafted. Before ``lm_head``, which is the other + # bf16 GEMM in a draft stage. + "mtp_eh_proj": ("eh_proj", "mtp_proj"), "lm_head": ("lm_head", "logits", "vocab_proj", "embed"), } diff --git a/gitm/planner/glm_graph.py b/gitm/planner/glm_graph.py new file mode 100644 index 0000000..8280fd5 --- /dev/null +++ b/gitm/planner/glm_graph.py @@ -0,0 +1,1169 @@ +"""Predicted execution graph for a GLM-5.2-class (``glm_moe_dsa``) engine step. + +A fork of :mod:`gitm.planner.moe_graph`, specialised for Z.ai's +``GlmMoeDsaForCausalLM``. Both families are sparse-MoE with a lightning indexer, +but the attention differs in kind, not degree, so a shared spec would carry a +field for each that is dead on the other. The two graphs live apart and share +only the canonical op names (so residuals stay comparable) and +:func:`~gitm.planner.roofline.distinct_experts` (one owner, no drift). + +Shape and provenance are in ``docs/glm-5.2/DESIGN-NOTE.md`` and in the catalogue +entries' ``provenance`` blocks; this docstring carries only what constrains an +edit to *this file*. + +Four properties do the work, and each has a plausible wrong reading: + +**MLA + DSA on every layer, no compression schedule.** One KV latent +(``kv_lora_rank``) shared across all query heads. Deriving cache traffic from +``num_key_value_heads * head_dim`` is the classic MLA error and overstates it 50x +on GLM-5.2 β€” the config's ``num_key_value_heads: 64`` is a red herring. + +**IndexShare.** ``indexer_types`` is ``full`` | ``shared`` per layer; a ``shared`` +layer reuses the previous ``full`` layer's top-k and physically carries no indexer +tensors. Emitting an indexer on all 78 layers overstates it ~3.7x. + +**Dense-then-sparse MLP.** ``first_k_dense_replace`` leading layers have no router +and no experts. Modelling them as MoE invents traffic the weight map denies. + +**Three precisions in one block**, read from +``quantization_config.modules_to_not_convert`` and ``moe_router_dtype`` rather +than assumed. On GLM-5.2-FP8 the backbone GEMMs *including* ``o_proj`` are fp8 +while the *indexer*, ``lm_head`` and the MTP ``eh_proj`` are not β€” the inversion +of the usual fp8-backbone layout. :attr:`GlmMoeDsaModelSpec.op_dtype_overrides` +carries it; one ``weight_dtype`` cannot. + +Known limits, stated rather than hidden: + +* **Prefill is not decode with a bigger M.** ``index_topk`` bounds the core's + FLOPs in both phases and its *bytes* in neither: at prefill every query selects + a different top-k and their union is the whole cache. See :func:`core_qk_pairs` + and :func:`core_read_entries` β€” reusing ``BatchConfig.attention_qk_pairs`` here + is wrong in both directions at once, and the errors partly cancel. +* **Uniform routing.** ``distinct_experts`` assumes a balanced router; real skew + touches fewer experts, moving *less* traffic β€” the conservative direction. +* **``ep_imbalance`` is calibrated, not predicted.** It stays 1.0 until a trace + measures it. +* **Collectives are bandwidth-plus-one-launch**, flagged ``estimated``, and still + reported unpriced when the SKU carries no interconnect bandwidth. +* **MTP cost is predicted; acceptance is not.** The graph prices D drafts and a + 1+D verify and leaves the payoff to :attr:`BatchConfig.tokens_per_step`. +* **Node names are constrained by the pairing contract.** Every op emitted here + must be classifiable by :func:`gitm.optimizer.deviation.classify_op`, or the + predicted node goes permanently unmatched while the real kernel files as + unmodeled. That is why three norm sites share one ``rms_norm`` op and why + SwiGLU lives inside the GEMM that owns its needle. See + ``docs/kernel_identity.md``. +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import Any + +from gitm.planner.graph import Graph, PredictedNode +from gitm.planner.roofline import ( + BatchConfig, + HardwareSpec, + ShardingConfig, + _canon_dtype, + distinct_experts, + roofline, + weight_bytes, +) + +FULL_INDEXER = "full" +SHARED_INDEXER = "shared" +DENSE_MLP = "dense" +SPARSE_MLP = "sparse" + + +@dataclass(frozen=True) +class GlmMoeDsaModelSpec: + """Model shape for a GLM-5.2-class (``glm_moe_dsa``) decode step. + + **The defaults below are a small reference shape, not any real checkpoint** β€” + deliberately far too small to be mistaken for GLM-5.2. Real checkpoints live in + ``gitm/planner/models/*.yaml`` (``family: glm_moe_dsa``) or are read from a + config by :func:`spec_from_hf_config`. The comments cite GLM-5.2 as the worked + example precisely where its values differ from these defaults in ways that + matter. + """ + + name: str = "glm-moe-dsa-reference" + hidden: int = 512 + n_layers: int = 6 + vocab: int = 2048 + + # ── MLA attention ──────────────────────────────────────────────────────── + n_heads: int = 8 + #: Query down-projection rank. The query goes ``hidden -> q_lora_rank -> + #: n_heads * q_head_dim``; the middle rank is a real, replicated matrix. + q_lora_rank: int = 256 + #: Compressed KV latent width. The cache holds **one** latent per token per + #: layer, shared across every query head β€” deriving KV traffic from + #: ``num_key_value_heads * head_dim`` instead is the classic MLA error and + #: overstates it by ``n_heads`` (64x on GLM-5.2). + kv_lora_rank: int = 128 + #: Per-head query/key width carrying no rotary embedding. + qk_nope_head_dim: int = 96 + #: Per-head query/key width carrying rotary embedding. The decoupled RoPE key + #: is shared (one per token, MQA-style) alongside the latent in the cache. + qk_rope_head_dim: int = 32 + #: Per-head value width. On GLM-5.2 this (256) differs from ``qk_nope`` (192), + #: so the score and the value read use different per-head widths. + v_head_dim: int = 128 + + # ── DeepSeek Sparse Attention indexer ──────────────────────────────────── + index_n_heads: int = 16 + index_head_dim: int = 64 + #: Positions the indexer keeps for the attention core. The core read is + #: bounded by this once history exceeds it, so attention stops growing with + #: context while the *indexer scan* keeps growing β€” a different node, a + #: different bound. + index_topk: int = 512 + #: Period of the IndexShare grouping: one ``full`` layer that computes the + #: selection, then ``index_topk_freq - 1`` ``shared`` layers that reuse it. + index_topk_freq: int = 4 + #: Per-layer ``full`` | ``shared``, straight from the checkpoint. Authoritative + #: when present β€” the ``shared`` layers carry no indexer weights, so this is + #: read, not a rule guessed from the frequency. Empty falls back to + #: :attr:`index_topk_freq` (first layer of each group is ``full``). + indexer_types: tuple[str, ...] = () + + # ── mixture of experts ─────────────────────────────────────────────────── + n_routed_experts: int = 8 + n_shared_experts: int = 1 + num_experts_per_tok: int = 2 + moe_intermediate_size: int = 256 + #: Dense-FFN width for the leading ``first_k_dense_replace`` layers. + intermediate_size: int = 768 + #: Leading layers with a dense FFN: no router, no experts. + first_k_dense_replace: int = 1 + #: Per-layer ``dense`` | ``sparse``; overrides :attr:`first_k_dense_replace`. + mlp_layer_types: tuple[str, ...] = () + #: Numerics only β€” no effect on the FLOP/byte roofline. + routed_scaling_factor: float = 1.0 + + # ── multi-token prediction ─────────────────────────────────────────────── + num_nextn_predict_layers: int = 0 + #: The MTP block reuses the main model's index instead of recomputing it, + #: and the weight map agrees: it carries no indexer tensors. + index_share_for_mtp_iteration: bool = True + + # ── precision ──────────────────────────────────────────────────────────── + weight_dtype: str = "bf16" + expert_dtype: str = "bf16" + kv_dtype: str = "bf16" + act_dtype: str = "bf16" + #: ``(op_name, dtype)`` for the ops that do not run in :attr:`weight_dtype`. + #: A tuple, not a mapping, so the spec stays hashable. Read from the + #: checkpoint, never assumed β€” see the module docstring. + op_dtype_overrides: tuple[tuple[str, str], ...] = () + + # ── derived shapes / schedule ──────────────────────────────────────────── + + @property + def q_head_dim(self) -> int: + """Per-head query/key width: the nope part plus the RoPE part.""" + return self.qk_nope_head_dim + self.qk_rope_head_dim + + @property + def kv_entry_dim(self) -> int: + """Elements one cached KV entry occupies: the latent plus the shared RoPE key. + + One per token per layer, shared across every query head β€” the whole point + of MLA. Not multiplied by ``n_heads`` or ``num_kv_heads``. + """ + return self.kv_lora_rank + self.qk_rope_head_dim + + def dtype_for(self, op: str, default: str) -> str: + """Precision ``op`` runs in β€” an override if the checkpoint declares one. + + ``default`` is the op's family dtype (:attr:`weight_dtype` for + projections, :attr:`expert_dtype` for the mixture), so an entry is needed + only where the checkpoint departs from it. + """ + for name, dtype in self.op_dtype_overrides: + if name == op: + return dtype + return default + + def mlp_kind(self, layer: int) -> str: + """``"dense"`` | ``"sparse"`` for ``layer``'s FFN.""" + if self.mlp_layer_types and layer < len(self.mlp_layer_types): + return self.mlp_layer_types[layer] + return DENSE_MLP if layer < self.first_k_dense_replace else SPARSE_MLP + + def is_sparse_mlp(self, layer: int) -> bool: + return self.mlp_kind(layer) == SPARSE_MLP + + def indexer_kind(self, layer: int) -> str: + """``"full"`` | ``"shared"`` β€” whether ``layer`` computes its own index. + + Prefers the checkpoint's explicit ``indexer_types``. The frequency fallback + makes the first layer of each ``index_topk_freq``-sized group ``full`` and + the rest ``shared``, which is how GLM-5.2 is laid out past its dense prefix. + """ + if self.indexer_types and layer < len(self.indexer_types): + return self.indexer_types[layer] + n = max(1, self.index_topk_freq) + return FULL_INDEXER if layer % n == 0 else SHARED_INDEXER + + def is_full_indexer(self, layer: int) -> bool: + return self.indexer_kind(layer) == FULL_INDEXER + + @property + def n_full_indexer_layers(self) -> int: + return sum(1 for i in range(self.n_layers) if self.is_full_indexer(i)) + + @property + def n_sparse_mlp_layers(self) -> int: + return sum(1 for i in range(self.n_layers) if self.is_sparse_mlp(i)) + + @property + def top_k(self) -> int: + return min(self.num_experts_per_tok, self.n_routed_experts) + + +def kv_entry_bytes(spec: GlmMoeDsaModelSpec) -> float: + """Bytes one cached KV entry occupies. + + The latent plus the decoupled RoPE key. DeepSeek-family checkpoints keep the + RoPE dimensions in bf16 while the latent may be quantised; on a pure-bf16 + checkpoint both are 2 bytes and this reduces to ``kv_entry_dim * 2``. The split + is kept so an fp8-KV *serving* config prices the two halves correctly. + """ + rope = spec.qk_rope_head_dim * weight_bytes("bf16") + latent = spec.kv_lora_rank * weight_bytes(spec.kv_dtype) + return latent + rope + + +def _capped_prefix_sum(context: int, tokens: int, cap: int) -> float: + """``sum(min(context + i, cap) for i in 1..tokens)`` in closed form. + + How many positions a causal chunk of ``tokens`` queries actually attends to + when each query is capped at ``cap`` selected keys. Written out rather than + looped because a prefill chunk is 8,192 queries and this is called per layer. + """ + if tokens <= 0: + return 0.0 + c, n, k = float(context), float(tokens), float(cap) + # Queries whose own history is still under the cap attend to all of it. + uncapped = max(0.0, min(n, k - c)) + total = uncapped * (2.0 * c + uncapped + 1.0) / 2.0 + return total + (n - uncapped) * k + + +def core_qk_pairs(spec: GlmMoeDsaModelSpec, batch: BatchConfig) -> float: + """Query-key products the attention *core* evaluates this step. + + Deliberately **not** :attr:`BatchConfig.attention_qk_pairs`, which is the + dense causal count. DSA hands the core at most ``index_topk`` selected + positions per query, so past 2,048 tokens of context the core is linear in + context where a dense model is quadratic β€” the whole point of the + architecture, and the term a mechanical copy of another family's prefill path + would get wrong by ``kv_len / index_topk`` (64x at 128K). + """ + k = spec.index_topk + pairs = batch.positions_per_step * float(min(batch.kv_cache_len, k)) + if batch.is_prefill: + pairs += _capped_prefix_sum(batch.prefill_context, batch.prefill_tokens, k) + return pairs + + +def core_read_entries(spec: GlmMoeDsaModelSpec, batch: BatchConfig) -> float: + """Cached KV entries the core actually reads β€” the phases disagree here. + + **Decode:** one query row per sequence, reading its own top-``index_topk`` + selection. ``sequences x min(kv_len, index_topk)`` β€” flat in context. + + **Prefill:** every query in the chunk selects a *different* top-k, and their + union over a few thousand queries is the whole history. The kernel therefore + streams the entire cache once per request: top-k bounds prefill *FLOPs*, not + prefill *bytes*. Charging ``P x index_topk`` here instead would understate + long-context prefill traffic by the ratio of context to 2,048, which is the + single most inviting error on this architecture. + """ + entries = batch.batch * float(min(batch.kv_cache_len, spec.index_topk)) + if batch.is_prefill: + entries += batch.prefill_requests * float( + batch.prefill_context + batch.prefill_tokens + ) + return entries + + +def index_scan_pairs(batch: BatchConfig) -> float: + """Query-key products the *indexer* scores β€” the uncompressed history. + + GLM does not compress before selecting, so every past token is scored. This + is the dense causal count, and it is the term that carries the quadratic at + prefill and the growth-in-context at decode. Paid on ``full`` layers only, + which is what IndexShare is worth. + """ + # Called unwindowed: ``window`` exists for sliding-window families and GLM has + # no window β€” the indexer scores the whole history by construction. + return batch.attention_qk_pairs() + + +def index_scan_entries(batch: BatchConfig) -> float: + """Cached index keys read per step: the whole history, once per sequence.""" + entries = float(batch.batch * batch.kv_cache_len) + if batch.is_prefill: + entries += batch.prefill_requests * float( + batch.prefill_context + batch.prefill_tokens + ) + return entries + + +def model_weight_bytes( + spec: GlmMoeDsaModelSpec, sharding: ShardingConfig | None = None +) -> float: + """Resident weight bytes on one rank. + + Decides whether a deployment shape fits at all, which the timing graph cannot. + Experts dominate overwhelmingly: 75 sparse layers x 256 experts x three + ``hidden x moe_intermediate`` matrices is the great majority of the checkpoint, + and it is what makes a ~753B model activate ~25B per token. + + Validated against ground truth: GLM-5.2 predicts within a few percent of the + published 1.507 TB checkpoint (``model.safetensors.index.json`` ``total_size``), + the residual being norms, biases and the MTP head this rolls in coarsely. + """ + sh = sharding or ShardingConfig() + tp = max(1, sh.tp) + es = max(1, sh.expert_shards) + ew = weight_bytes(spec.dtype_for("moe_routed", spec.expert_dtype)) + sw = weight_bytes(spec.dtype_for("moe_shared", spec.expert_dtype)) + ww = weight_bytes(spec.weight_dtype) + rw = weight_bytes(spec.dtype_for("moe_router", spec.weight_dtype)) + iw = weight_bytes(spec.dtype_for("attn_index_proj", spec.weight_dtype)) + lw = weight_bytes(spec.dtype_for("lm_head", spec.weight_dtype)) + tw = weight_bytes(spec.dtype_for("embed_tokens", spec.weight_dtype)) + h = spec.hidden + inter = spec.moe_intermediate_size + + n_sparse = spec.n_sparse_mlp_layers + spec.num_nextn_predict_layers + n_dense = spec.n_layers - spec.n_sparse_mlp_layers + # The MTP block counts here only if it recomputes the index. With + # ``index_share_for_mtp_iteration`` it reuses the main model's selection and + # carries no indexer tensors β€” which is exactly what the weight map shows, and + # what ``_emit_layer`` already honours by emitting no indexer node for it. + # Counting it anyway put one indexer's weights (18.7 MB bf16) in the footprint + # that the checkpoint does not contain, and contradicted the graph beside it. + n_full_idx = spec.n_full_indexer_layers + ( + 0 if spec.index_share_for_mtp_iteration else spec.num_nextn_predict_layers + ) + n_attn = spec.n_layers + spec.num_nextn_predict_layers + + experts = n_sparse * spec.n_routed_experts * 3 * h * inter * ew / es + shared_exp = n_sparse * spec.n_shared_experts * 3 * h * inter * sw / tp + router = n_sparse * h * spec.n_routed_experts * rw # replicated + dense_ffn = n_dense * 3 * h * spec.intermediate_size * ww / tp + + # MLA projections, per attention layer. q_a and kv_a are replicated (they + # produce the shared latent, which has nothing to split); q_b, o_proj and the + # per-head value up-projection shard on heads. + attn_per_layer = ( + h * spec.q_lora_rank # q_a, replicated + + spec.q_lora_rank * spec.n_heads * spec.q_head_dim / tp # q_b + + h * spec.kv_entry_dim # kv_a (latent + rope key), replicated + + spec.kv_lora_rank * spec.n_heads * (spec.qk_nope_head_dim + spec.v_head_dim) / tp # kv_b + + spec.n_heads * spec.v_head_dim * h / tp # o_proj + ) + # Indexer weights live only on ``full`` layers (proven: ``shared`` layers carry + # none). Replicated across ranks, as vLLM builds the indexer ReplicatedLinear. + # ``wq_b`` from the query latent, ``wk`` from hidden, and the per-head gate. + indexer_per_full = ( + spec.q_lora_rank * spec.index_n_heads * spec.index_head_dim + + h * spec.index_head_dim + + h * spec.index_n_heads + ) + + # Untied input embedding and vocabulary projection: two tensors, priced + # separately because the checkpoint names them separately and could quantise + # one without the other. Both stay wide on GLM-5.2-FP8, which is 1.9 GB of + # resident footprint an fp8 read would halve on paper and not on disk. + embed = spec.vocab * h * (tw + lw) / tp + + return ( + experts + + shared_exp + + router + + n_full_idx * indexer_per_full * iw + + embed + + (n_attn * attn_per_layer + dense_ffn) * ww + ) + + +def kv_bytes_per_token(spec: GlmMoeDsaModelSpec) -> float: + """KV bytes each additional token of context costs, across the whole model. + + One MLA latent (plus the indexer key on ``full`` layers) per token per layer. + Flat across layers β€” there is no compression schedule to sum over β€” but the + indexer key is only cached where an indexer runs. + """ + kw = weight_bytes(spec.kv_dtype) + latent = spec.n_layers * kv_entry_bytes(spec) + index_keys = spec.n_full_indexer_layers * spec.index_head_dim * kw + return latent + index_keys + + +def _linear(rows: float, k: int, n: int, act_b: float, w_b: float) -> tuple[float, float]: + """(flops, bytes) for a ``(rows, k) @ (k, n)`` projection. + + Bytes count the activation in, the weights, and the activation out. At decode + ``rows`` is small and the weight term dominates β€” which is why weight dtype, + not activation dtype, sets the floor for every projection here. + """ + return 2.0 * rows * k * n, act_b * rows * k + w_b * k * n + act_b * rows * n + + +def _pointwise(rows: float, elems: float, act_b: float, *, ops: float = 1.0) -> tuple[float, float]: + """(flops, bytes) for an elementwise kernel over ``rows x elems``. + + Read once, written once β€” the traffic that makes a norm or a residual add + cost anything at all. ``ops`` is arithmetic per element (a residual add is 1, + an RMSNorm is ~3 counting the reduction and the rescale). + + These nodes are individually a rounding error and collectively the low-batch + story: at batch 1 there is no useful arithmetic and no useful bandwidth in any + of them, so each one costs exactly one kernel launch and the step is the sum + of its launches. A graph that folds them into the GEMM they precede cannot + show that, and will report a decode step as memory-bound when it is not. + """ + n = rows * elems + return ops * n, 2.0 * n * act_b + + +def _emit_layer( + g: Graph, + spec: GlmMoeDsaModelSpec, + hw: HardwareSpec, + layer: int, + *, + batch: BatchConfig, + sh: ShardingConfig, + prefix: str = "", + force_full_indexer: bool | None = None, +) -> None: + """Append one transformer layer's predicted nodes to ``g``. + + ``batch`` is the phase, already adjusted by the caller: decode, chunked + prefill, or a draft stage (prefill stripped). The node *set* is identical + across all three; only the class of four of them changes. + + Three row counts, and conflating any two is a real error: ``rows`` (positions + computed β€” every projection scales with it), ``batch.batch`` (distinct KV + caches read β€” charged per *sequence*, which is what makes a 1+D verify pay), + and ``batch.prefill_requests`` (the denominator for anything read once per + request). + """ + h = spec.hidden + aw = weight_bytes(spec.act_dtype) + wd, ed = spec.weight_dtype, spec.expert_dtype + tp = max(1, sh.tp) + es = max(1, sh.expert_shards) + + rows = float(batch.positions_per_step + batch.prefill_tokens) + + def add( + op: str, flops: float, byts: float, dtype: str, + *, estimated: bool = False, serial_launches: int = 1, + ) -> None: + """Emit one node, at the precision the checkpoint says the op runs in. + + ``serial_launches`` defaults to 1 because every node here *is* one + dependent kernel launch: it consumes the previous node's output, so its + wall time cannot fall below the launch overhead however few bytes it + moves. On a decode step that floor is what the small pointwise and + routing nodes are actually bounded by, and omitting it does not make the + prediction slightly optimistic β€” it makes a whole bound label absent. + """ + name = f"{prefix}{op}" + g.nodes.append( + PredictedNode( + name, layer, + roofline( + name, flops, byts, hw, spec.dtype_for(op, dtype), + estimated=estimated, serial_launches=serial_launches, + ), + ) + ) + + def w_bytes(op: str, default: str) -> float: + """Bytes per stored weight for ``op``, after any precision override.""" + return weight_bytes(spec.dtype_for(op, default)) + + def add_pointwise(op: str, elems: float, *, ops: float = 1.0) -> None: + f_p, b_p = _pointwise(rows, elems, aw, ops=ops) + add(op, f_p, b_p, spec.act_dtype) + + def add_rms_norm(*, with_residual: bool) -> None: + """One ``rms_norm`` node per norm site β€” and the residual add is inside it. + + vLLM runs ``RMSNorm.forward(x, residual)`` as a single + ``fused_add_rms_norm`` kernel, so a separate residual node would predict a + launch that never happens. ``with_residual`` adds the extra read the fused + form does, and costs nothing else. + + All three norm sites in a layer share this op name deliberately. They are + the *same kernel* in the trace β€” ``classify_op`` matches on the kernel + name, and three distinct op names for one name would leave two of them + permanently unmatched while the third absorbed all three sites' time. The + position that distinguishes them is carried by ``layer`` and by issue + order, which is where it belongs; see ``docs/kernel_identity.md`` on why a + name is a guess and an NVTX range is an identity. + """ + elems = h * (2.0 if with_residual else 1.0) + f_p, b_p = _pointwise(rows, elems, aw, ops=3.0) + add("rms_norm", f_p, b_p, spec.act_dtype) + + def add_act_quant(op: str, elems: float, gemm_op: str) -> None: + """Dynamic FP8 activation scaling ahead of a quantised GEMM. + + GLM-5.2-FP8 declares ``activation_scheme: "dynamic"``, so the activation + is quantised at run time β€” a pointwise pass plus a per-row reduction for + the scale, as its own kernel, once per group of fp8 GEMMs that share an + input. Emitted only where the consuming GEMM is actually fp8: on the bf16 + checkpoint there is nothing to quantise and the kernel does not exist, + which is the sort of difference a single model-wide dtype cannot express. + """ + if _canon_dtype(spec.dtype_for(gemm_op, wd)) != "fp8": + return + # Read the bf16 activation, write the fp8 one plus its scales. + add(op, 2.0 * rows * elems, + rows * elems * (aw + 1.0) + rows * 4.0, spec.act_dtype) + + # ── MLA attention: low-rank query, compressed KV latent ────────────────── + add_rms_norm(with_residual=layer > 0) + add_act_quant("act_quant", h, "attn_q_a") + + # q_a and kv_a are replicated across TP ranks: they produce the shared latent, + # which has nothing to split when there is one KV latent. Every rank pays them + # in full, so TP's speedup on attention is strictly less than ``tp``. + f, b = _linear(rows, h, spec.q_lora_rank, aw, w_bytes("attn_q_a", wd)) + add("attn_q_a", f, b, wd) + + f, b = _linear( + rows, spec.q_lora_rank, spec.n_heads * spec.q_head_dim // tp, aw, + w_bytes("attn_q_b", wd), + ) + add("attn_q_b", f, b, wd) + + # The compressed latent plus the decoupled RoPE key, and the cache write for + # the positions just computed. One projection (no CSA/HCA overlap here). + f, b = _linear(rows, h, spec.kv_entry_dim, aw, w_bytes("attn_kv_a", wd)) + add("attn_kv_a", f, b + rows * kv_entry_bytes(spec), wd) + + # Reconstruct per-head K_nope and V from the cached latent (W^UK, W^UV), + # modelled *unabsorbed*: its own GEMM, and attn_out_proj stays narrow. An + # engine that absorbs MLA drops this node and doubles attn_out_proj's input + # width instead β€” a serving variant, flagged in the catalogue provenance. + f, b = _linear( + rows, spec.kv_lora_rank, + spec.n_heads * (spec.qk_nope_head_dim + spec.v_head_dim) // tp, aw, + w_bytes("attn_kv_b", wd), + ) + add("attn_kv_b", f, b, wd) + + # ── indexer: only ``full`` layers run one ──────────────────────────────── + # ``shared`` layers reuse the group's selection and carry no indexer weights, + # so they emit no indexer node. Emitting one would put a kernel in the graph + # that never ran and inflate the indexer's share ~4x. + runs_indexer = ( + spec.is_full_indexer(layer) if force_full_indexer is None else force_full_indexer + ) + scan_pairs = index_scan_pairs(batch) + if runs_indexer and scan_pairs > 0: + # The indexer query comes off the same latent as the attention query, so + # only the up-projection is charged (``wq_b``), plus the per-token key + # (``wk`` β€” one key per token, MQA-style, not one per index head) and the + # per-head gate (``weights_proj``). Not divided by ``tp``: vLLM builds the + # indexer as ReplicatedLinear, so every rank runs the whole thing. bf16 on + # the FP8 checkpoint β€” the indexer is named in ``modules_to_not_convert``. + idx_w = w_bytes("attn_index_proj", wd) + f_q, b_q = _linear( + rows, spec.q_lora_rank, spec.index_n_heads * spec.index_head_dim, aw, idx_w + ) + f_k, b_k = _linear(rows, h, spec.index_head_dim, aw, idx_w) + f_g, b_g = _linear(rows, h, spec.index_n_heads, aw, idx_w) + add( + "attn_index_proj", + f_q + f_k + f_g, + b_q + b_k + b_g + rows * spec.index_head_dim * weight_bytes(spec.kv_dtype), + wd, + ) + + # Two dtypes, two questions: the *bytes* follow how the keys are stored + # (``kv_dtype``), the *FLOPs* follow what the indexer computes in β€” and the + # indexer is one of the modules the quantiser skipped. Invisible at decode + # (memory-bound at every context); doubles the prefill row. + add( + "attn_index_score", + # Every one of the 32 index heads scores each candidate against the + # single shared 128-d key ``wk`` produces per token β€” MQA-style, which + # is why the key term below is not multiplied by the head count and + # this one is. Dropping ``index_n_heads`` here understates the scan 32x + # and would leave it looking free at every context. + 2.0 * scan_pairs * spec.index_n_heads * spec.index_head_dim, + # Index keys live in the cache: read once per sequence (or per + # prefilling request), not per position, and replicated across ranks + # alongside the KV latent. + index_scan_entries(batch) * spec.index_head_dim + * weight_bytes(spec.kv_dtype), + spec.dtype_for("attn_index_proj", wd), + ) + + # ── attention core over the selected positions ────────────────────────── + # FLOPs follow the *selected* pairs (top-k bounded); bytes follow what the + # kernel must stream, which at prefill is the whole cache and at decode is one + # top-k window per sequence. The two do not move together on this + # architecture, which is why they are separate helpers. + heads = max(1, spec.n_heads // tp) + pairs = core_qk_pairs(spec, batch) + qk = 2.0 * pairs * heads * spec.q_head_dim + pv = 2.0 * pairs * heads * spec.v_head_dim + add( + "attn_score_value", + qk + pv, + # One latent, shared by every query head, so *not* multiplied by n_heads. + # Deliberately not divided by ``tp`` either: a single shared latent cannot + # be split, so the cache is replicated and every rank reads all of it β€” + # tensor parallelism buys no KV bandwidth on this architecture. + core_read_entries(spec, batch) * kv_entry_bytes(spec), + wd, + ) + + # RMSNorm on every query head and the single KV latent, plus partial RoPE on + # the last ``qk_rope_head_dim`` dims and the cache insert β€” one fused kernel. + normed = rows * (heads * spec.q_head_dim + spec.kv_entry_dim) + roped = rows * (heads + 1) * spec.qk_rope_head_dim + add( + "attn_qnorm_rope_insert", + 3.0 * normed + 6.0 * roped, + 2.0 * normed * aw + 2.0 * roped * aw, + spec.act_dtype, + ) + + # Per-head value space back to hidden. No o_lora/o_groups here (unlike V4). + # On the FP8 checkpoint this one *is* quantised β€” the opposite of the + # fp8-backbone checkpoints that keep o_proj wide. + f, b = _linear( + rows, spec.n_heads * spec.v_head_dim // tp, h, aw, w_bytes("attn_out_proj", wd) + ) + add("attn_out_proj", f, b, wd) + + _emit_collective(g, spec, hw, layer, "tp_all_reduce_attn", rows, sh, prefix) + add_rms_norm(with_residual=True) + + # ── FFN: dense on the leading layers, mixture on the rest ──────────────── + if not spec.is_sparse_mlp(layer): + # Dense FFN (first_k_dense_replace). gate+up then down over the wide + # intermediate. Canonical dense-graph names so residuals stay comparable. + inter = spec.intermediate_size + add_act_quant("act_quant", h, "mlp_gate_up") + # SwiGLU stays inside ``mlp_gate_up``: ``silu_and_mul`` is already one of + # that op's needles in ``deviation._OP_RULES``, so a separate node would + # be a prediction the pairing has no way to receive. + f_gu, b_gu = _linear(rows, h, 2 * inter // tp, aw, w_bytes("mlp_gate_up", wd)) + f_act, b_act = _pointwise(rows, inter / tp, aw, ops=4.0) + add("mlp_gate_up", f_gu + f_act, b_gu + b_act, wd) + f_d, b_d = _linear(rows, inter // tp, h, aw, w_bytes("mlp_down", wd)) + add("mlp_down", f_d, b_d, wd) + _emit_collective(g, spec, hw, layer, "tp_all_reduce_mlp", rows, sh, prefix) + return + + # Router is replicated: every rank scores every expert to know what to keep. + # fp32 on every GLM-5.2 variant (``moe_router_dtype``) β€” a model fact, not a + # quantisation choice, and the reason this op carries its own dtype. + f, b = _linear(rows, h, spec.n_routed_experts, aw, w_bytes("moe_router", wd)) + add("moe_router", f, b, wd) + + # vLLM's fused gating kernel (sigmoid + noaux_tc bias + top-8 + renorm): a + # second launch after the router GEMM, sharing its op name because both + # classify to ``moe_router`` β€” a mapping the other families depend on, and a + # private name here would emit a node no capture can pair against. + # + # Its own node regardless: the **only data-dependent shape in the step**, and + # so the thing that decides whether the step is CUDA-graph capturable. + add( + "moe_router", + 3.0 * rows * spec.n_routed_experts, + rows * (2.0 * spec.n_routed_experts + 2.0 * spec.top_k) * aw, + spec.act_dtype, + ) + + add_act_quant("act_quant", h, "moe_routed") + + inter = spec.moe_intermediate_size + per_expert_weights = 3.0 * h * inter + per_position_flops = 6.0 * h * inter # 2 * (gate + up + down) * h * inter + ew = w_bytes("moe_routed", ed) + + if spec.n_shared_experts > 0: + sw = w_bytes("moe_shared", ed) + add( + "moe_shared", + per_position_flops * rows * spec.n_shared_experts / tp, + per_expert_weights * spec.n_shared_experts * sw / tp + + aw * (rows * h * 2 + rows * inter * 2 * spec.n_shared_experts / tp), + ed, + ) + + # Gather into expert-major order. Zero arithmetic, a rounding error at + # decode β€” but the expanded tensor is ``rows x top_k`` wide, so at an + # 8,192-token chunk this and its scatter move hundreds of MB per layer for no + # FLOPs. Folded into the expert GEMM, that share of prefill is invisible. + expanded = rows * spec.top_k + add("moe_permute", 0.0, aw * (rows * h + expanded * h) / es, spec.act_dtype) + + # The saturating set-union: FLOPs scale with rows x top_k, weight traffic + # with how many *distinct* experts the batch woke β€” shared with the dense-MoE + # roofline so there is one owner for the term. The argument is *rows*, not + # sequences: under speculative decoding and under prefill every extra row is + # another draw on the expert bank, which is exactly why a 1+D verify costs + # more than a decode without doing more work per token. + distinct = distinct_experts( + int(rows), spec.n_routed_experts, spec.num_experts_per_tok + ) + skew = sh.ep_imbalance if sh.ep > 1 else 1.0 + # SwiGLU is inside this node, not beside it: ``silu_and_mul`` is one of + # ``mlp_gate_up``'s needles in ``deviation._OP_RULES``, so a separate SwiGLU + # node would be a prediction with no observed kernel to pair against. + swiglu_f, swiglu_b = _pointwise(rows, spec.top_k * inter / es, aw, ops=4.0) + add( + "moe_routed", + per_position_flops * rows * spec.num_experts_per_tok * skew / es + swiglu_f, + per_expert_weights * distinct * ew * skew / es + + aw * (rows * inter * 2 * spec.num_experts_per_tok / es) + + swiglu_b, + ed, + ) + + # Weighted scatter-add back to ``rows x hidden``, including the + # routed_scaling_factor multiply. + add( + "moe_combine", + 2.0 * expanded * h, + aw * (expanded * h + rows * h) / es, + spec.act_dtype, + ) + + _emit_collective( + g, spec, hw, layer, "tp_all_reduce_mlp", rows, sh, prefix, + dispatches_experts=True, + ) + + +def _emit_collective( + g: Graph, + spec: GlmMoeDsaModelSpec, + hw: HardwareSpec, + layer: int, + op: str, + rows: float, + sh: ShardingConfig, + prefix: str, + *, + dispatches_experts: bool = False, +) -> None: + """Emit the cross-rank traffic that closes one sub-block. + + **Two per layer, not one.** A tensor-parallel layer all-reduces after + ``o_proj`` and again after the FFN combine; folding them into a single node + with double the payload gets the bytes right and the *count* wrong β€” and at + decode payloads a collective is bounded by its ring latency rather than by its + bytes, so the count is the cost. Under expert parallelism the MoE half + additionally dispatches and combines across expert ranks. + + ``dispatches_experts`` gates the expert-parallel all-to-all: only a mixture + layer sends tokens to expert ranks, and charging the dense layers for one puts + wire traffic on a block with no experts to send to. + + ``serial_launches`` is withheld when the SKU carries no interconnect + bandwidth, so an unpriced collective still predicts zero time and + :attr:`Graph.has_unpriced_collectives` keeps reporting it. A latency floor + there would convert "cannot price this" into "costs two microseconds". + """ + tp = max(1, sh.tp) + if tp <= 1 and sh.ep <= 1: + return + aw = weight_bytes(spec.act_dtype) + link = replace(hw, peak_mem_bw_bytes_per_s=hw.interconnect_bw_bytes_per_s) + priced = hw.interconnect_bw_bytes_per_s > 0 + + def add_link(name_op: str, byts: float) -> None: + name = f"{prefix}{name_op}" + g.nodes.append( + PredictedNode( + name, layer, + roofline( + name, 0.0, byts, link, spec.act_dtype, + estimated=True, serial_launches=1 if priced else 0, + ), + # The one region this graph expects off the compute stream. + # A declaration, not a check: nothing reads this field today + # (monitor.py tests overlap on the *observed* stream), so it is a + # hook for the invariant in docs/invariants.md Β§3, not a wiring. + expected_stream_id=1, + ) + ) + + if sh.ep > 1 and dispatches_experts: + off_rank = (sh.ep - 1) / sh.ep + add_link( + "moe_all_to_all", + 2.0 * rows * spec.num_experts_per_tok * spec.hidden * aw * off_rank, + ) + if tp > 1: + add_link(op, (2.0 * (tp - 1) / tp) * rows * spec.hidden * aw) + + +def predict_glm_graph( + model: GlmMoeDsaModelSpec | None = None, + hw: HardwareSpec | None = None, + batch: BatchConfig | None = None, + sharding: ShardingConfig | None = None, +) -> Graph: + """Emit a predicted execution graph for one GLM-5.2-class engine step, per rank. + + One step, three passes, and they are not three graphs: + + **The backbone** runs over every position β€” decode positions plus any prefill + chunk riding along. Under speculative decoding that is ``batch x (1 + D)``: + **verify is not a new graph, it is this one at 1+D rows**. + + **The draft chain** is the one genuinely new subgraph: a single MTP module + invoked ``D`` times serially, EAGLE-style, each stage running its own + vocabulary projection over the whole untied matrix. + + **The epilogue** projects only :attr:`BatchConfig.logits_rows` β€” one row per + prefilling *request* plus every decode position, not the whole chunk. + + With ``sharding`` at its default the graph is whole-model; given a real + sharding it predicts what *one rank* does. + """ + spec = model or GlmMoeDsaModelSpec() + hw = hw or HardwareSpec() + batch = batch or BatchConfig() + sh = sharding or ShardingConfig() + + # Refuse a sharding the model cannot take. Head counts floor-divide throughout, + # so ``tp > n_heads`` silently prices the whole attention path at zero work. + if spec.n_heads % max(1, sh.tp) != 0: + raise ValueError( + f"tensor-parallel size {sh.tp} does not divide {spec.n_heads} attention " + "heads β€” every head-sharded op would floor to zero work" + ) + if spec.qk_rope_head_dim > spec.q_head_dim: + raise ValueError( + f"qk_rope_head_dim ({spec.qk_rope_head_dim}) exceeds the query head width " + f"({spec.q_head_dim}) β€” the RoPE slice cannot exceed the head it slices" + ) + if spec.n_layers <= 0: + raise ValueError("n_layers must be positive β€” an empty model predicts nothing") + + g = Graph(model=spec, hw=hw, batch=batch, sharding=sh) # type: ignore[arg-type] + aw = weight_bytes(spec.act_dtype) + rows = float(batch.positions_per_step + batch.prefill_tokens) + + # Prologue: one gather from the untied input embedding. No FLOPs, and the bytes + # are the rows it touches, not the 1.9 GB table β€” an index_select reads what it + # selects. It reads at the *table's* width and writes at the activation width; + # those are the same on GLM-5.2, and would not be on a checkpoint that + # quantised the embedding. + embed_dtype = spec.dtype_for("embed_tokens", spec.weight_dtype) + g.nodes.append( + PredictedNode( + "embed_tokens", None, + roofline( + "embed_tokens", 0.0, + rows * spec.hidden * (weight_bytes(embed_dtype) + aw), hw, + embed_dtype, serial_launches=1, + ), + ) + ) + + for layer in range(spec.n_layers): + _emit_layer(g, spec, hw, layer, batch=batch, sh=sh) + + # Epilogue: the final norm runs over every row that needs logits, not every + # row in the step β€” the same count lm_head uses, and at prefill that is one + # row per prompt rather than the whole chunk. + logit_rows = float(batch.logits_rows) + f_n, b_n = _pointwise(logit_rows, 2.0 * spec.hidden, aw, ops=3.0) + g.nodes.append( + PredictedNode( + "rms_norm", None, + roofline("rms_norm", f_n, b_n, hw, spec.act_dtype, serial_launches=1), + ) + ) + + lm_w = weight_bytes(spec.dtype_for("lm_head", spec.weight_dtype)) + lm_dtype = spec.dtype_for("lm_head", spec.weight_dtype) + + def add_lm_head(rows: float, layer: int | None) -> None: + f, b = _linear(rows, spec.hidden, spec.vocab // max(1, sh.tp), aw, lm_w) + g.nodes.append( + PredictedNode( + "lm_head", layer, + roofline("lm_head", f, b, hw, lm_dtype, serial_launches=1), + ) + ) + + add_lm_head(batch.logits_rows, None) + + # The vocabulary projection is sharded across TP ranks, so the ranks must + # gather each other's slices before sampling. FP32 logits, full vocabulary β€” + # 19.8 MB at 32 rows, which is small in bytes and is one more unavoidable + # collective on the critical path between the last layer and the sample. + if sh.tp > 1: + link = replace(hw, peak_mem_bw_bytes_per_s=hw.interconnect_bw_bytes_per_s) + priced = hw.interconnect_bw_bytes_per_s > 0 + g.nodes.append( + PredictedNode( + "logits_all_gather", None, + roofline( + "logits_all_gather", 0.0, + batch.logits_rows * spec.vocab * 4.0 * (sh.tp - 1) / sh.tp, + link, "fp32", estimated=True, + serial_launches=1 if priced else 0, + ), + expected_stream_id=1, + ) + ) + + # ── the draft chain ────────────────────────────────────────────────────── + # One MTP module invoked once per drafted token, serially, at one row per + # sequence and no prefill β€” a draft proposes continuations, so on a + # pure-prefill step it does not run at all (nodes that never launch would show + # up as a launch facet made of absent kernels). + # + # Emitted as a *shared* block: ``index_share_for_mtp_iteration`` says the + # iteration reuses the main selection, and the weight map agrees β€” no + # ``self_attn.indexer.*`` tensors, exactly like the 57 ``shared`` layers. + # + # It is not a smaller copy of the model. The block carries a full + # ``mlp.experts.*`` bank, so its cost is expert weight traffic paid ``D`` times + # over, not arithmetic. + # ``num_nextn_predict_layers`` says the block exists in the checkpoint; it does + # not say the engine runs it. Drafting happens only under a speculative config, + # so at D=0 there are zero drafted tokens and zero stages β€” the module's weights + # stay resident (``model_weight_bytes`` counts them) and none of its kernels + # launch. Forcing one stage here charged a pure decode step 0.3 ms of drafting + # that a server without ``--speculative-config`` never does. + if ( + spec.num_nextn_predict_layers > 0 + and batch.speculative_tokens > 0 + and batch.positions_per_step > 0 + ): + draft_batch = replace(batch, prefill_tokens=0, speculative_tokens=0) + for stage in range(batch.speculative_tokens): + _emit_layer( + g, spec, hw, spec.n_layers + stage, + batch=draft_batch, sh=sh, + force_full_indexer=not spec.index_share_for_mtp_iteration, + ) + # ``enorm`` and ``hnorm``: the MTP block normalises the embedding and + # the carried hidden state separately before fusing them. Two kernels, + # bf16 on the FP8 checkpoint, and they sit inside the serial chain. + f_n2, b_n2 = _pointwise(draft_batch.batch, 2.0 * spec.hidden, aw, ops=3.0) + g.nodes.append( + PredictedNode( + "rms_norm", spec.n_layers + stage, + roofline("rms_norm", f_n2, b_n2, hw, spec.act_dtype, + serial_launches=2), + ) + ) + # ``eh_proj``: the [2*hidden, hidden] fusion of the previous hidden + # state with the embedding of the token just drafted. bf16 on the FP8 + # checkpoint (named in ``modules_to_not_convert``), and replicated per + # rank unless the engine shards it. + eh_w = weight_bytes(spec.dtype_for("mtp_eh_proj", spec.weight_dtype)) + f, b = _linear(draft_batch.batch, 2 * spec.hidden, spec.hidden, aw, eh_w) + g.nodes.append( + PredictedNode( + "mtp_eh_proj", spec.n_layers + stage, + roofline( + "mtp_eh_proj", f, b, hw, + spec.dtype_for("mtp_eh_proj", spec.weight_dtype), + serial_launches=1, + ), + ) + ) + # There is no ``mtp.*.lm_head`` in the checkpoint β€” the draft shares the + # backbone's, which means it re-reads the same vocabulary weights and + # gets no cheaper for being a draft. + add_lm_head(draft_batch.batch, spec.n_layers + stage) + + return g + + +def is_glm_moe_dsa_config(cfg: dict[str, Any]) -> bool: + """True for the GLM-5.2-class checkpoints this module models. + + The discriminator is unambiguous and cheap: the checkpoint declares + ``model_type == "glm_moe_dsa"`` (equivalently ``GlmMoeDsaForCausalLM`` in + ``architectures``). Both this family and DeepSeek-V4 carry ``index_topk`` and + ``n_routed_experts``, so a structural test would collide β€” the model_type is + the clean separator, and this check must run *before* ``is_sparse_moe_config`` + in :func:`gitm.planner.registry.detect_family`. + """ + if str(cfg.get("model_type", "")).lower() == "glm_moe_dsa": + return True + archs = cfg.get("architectures") or [] + return any("glmmoedsa" in str(a).lower() for a in archs) + + +#: Which graph op each ``modules_to_not_convert`` entry belongs to. Substring +#: match against the tensor name, first hit wins. Norms and biases are omitted +#: deliberately β€” they are not nodes in this graph, so a precision for them would +#: have nothing to price. +_UNQUANTISED_OPS: tuple[tuple[str, str], ...] = ( + ("lm_head", "lm_head"), + ("embed_tokens", "embed_tokens"), + ("eh_proj", "mtp_eh_proj"), + # The indexer's *projections*, named specifically. A bare "indexer" needle also + # matches ``indexer.k_norm``, and a norm carries no information about the + # projection's width β€” every fp8 scheme leaves norms wide, so a checkpoint that + # skipped only the norm and quantised ``indexers_proj`` would be read as + # leaving the whole indexer bf16. + ("indexers_proj", "attn_index_proj"), + ("indexer.wq_b", "attn_index_proj"), + ("indexer.wk", "attn_index_proj"), + ("indexer.weights_proj", "attn_index_proj"), + ("mlp.gate", "moe_router"), +) + + +def _op_dtype_overrides( + cfg: dict[str, Any], q: dict[str, Any], weight_dtype: str, model_dtype: str +) -> tuple[tuple[str, str], ...]: + """Per-op precision, read from the checkpoint rather than assumed. + + Two sources answering different questions. + ``quantization_config.modules_to_not_convert`` is *what the quantiser + skipped* β€” on GLM-5.2-FP8 that is ``lm_head``, ``embed_tokens``, the MTP + ``eh_proj`` and, the one worth naming, the **lightning indexer**. + ``moe_router_dtype`` is *what the model computes in regardless*: fp32 on every + variant, so it is emitted with or without a quantisation config. + + Empty when the checkpoint says nothing β€” inventing entries would put a + precision in the graph that the checkpoint never declared. + """ + found: dict[str, str] = {} + + skipped = q.get("modules_to_not_convert") or q.get("ignored_layers") or [] + if isinstance(skipped, list | tuple) and weight_dtype != model_dtype: + for tensor in skipped: + name = str(tensor).lower() + for needle, op in _UNQUANTISED_OPS: + if needle in name: + found.setdefault(op, model_dtype) + + router = str(cfg.get("moe_router_dtype") or "").lower() + if router.startswith("float32") or router == "fp32": + found["moe_router"] = "fp32" + + return tuple(sorted(found.items())) + + +def spec_from_hf_config( + cfg: dict[str, Any], *, name: str | None = None +) -> GlmMoeDsaModelSpec: + """Build a :class:`GlmMoeDsaModelSpec` from a HuggingFace ``config.json``. + + Reads the checkpoint's declared shape so the graph cannot drift from the model. + The ``indexer_types`` and ``mlp_layer_types`` arrays are read verbatim β€” they + are the IndexShare and dense/sparse schedules, and a rule guessed from the + frequencies would misplace layers while producing a plausible total. + """ + + def _int(key: str, default: int) -> int: + try: + return int(cfg.get(key, default)) + except (TypeError, ValueError): + return default + + def _req(key: str) -> int: + v = cfg.get(key) + if v is None: + raise ValueError( + f"config declares no {key!r}; refusing to substitute a default, " + "which would predict a model this checkpoint is not" + ) + try: + return int(v) + except (TypeError, ValueError): + raise ValueError(f"config field {key!r} is not an integer: {v!r}") from None + + n_layers = _req("num_hidden_layers") + + def _types(key: str) -> tuple[str, ...]: + v = cfg.get(key) + if not isinstance(v, list | tuple): + return () + # Validate the *raw* length before any slicing: GLM-5.2's schedules are + # exactly num_hidden_layers long, and an over-long array is a sign the + # config is not what the caller thinks β€” worth an error, not a silent + # truncation that leaves the trailing layers on the last entry's kind. + if len(v) != n_layers: + raise ValueError( + f"{key} has {len(v)} entries for {n_layers} layers β€” the " + "schedule must cover the model exactly" + ) + return tuple(str(t) for t in v) + + dtype = str(cfg.get("dtype") or cfg.get("torch_dtype") or "bf16").lower() + if dtype.startswith("bfloat"): + dtype = "bf16" + elif dtype.startswith("float16") or dtype == "half": + dtype = "fp16" + elif dtype.startswith("float32"): + dtype = "fp32" + + # Quantisation, if the checkpoint declares any. The base GLM-5.2 release + # carries none β€” an all-bf16 read matches its 1.507 TB on disk β€” while + # GLM-5.2-FP8 declares e4m3 with 128x128 weight blocks. Read rather than + # assumed in either direction: defaulting to fp8 would manufacture headroom on + # the bf16 checkpoint, and defaulting to bf16 would double every weight term on + # the one the vendor actually recommends deploying. + q = cfg.get("quantization_config") or {} + weight_dtype = str(q.get("quant_method", dtype)).lower() if q else dtype + expert_dtype = str(cfg.get("expert_dtype", weight_dtype)).lower() + overrides = _op_dtype_overrides(cfg, q, weight_dtype, dtype) + + return GlmMoeDsaModelSpec( + name=name or str(cfg.get("_name_or_path") or cfg.get("model_type") or "glm-moe-dsa"), + hidden=_req("hidden_size"), + n_layers=n_layers, + vocab=_req("vocab_size"), + n_heads=_req("num_attention_heads"), + q_lora_rank=_req("q_lora_rank"), + kv_lora_rank=_req("kv_lora_rank"), + qk_nope_head_dim=_int("qk_nope_head_dim", _int("head_dim", 128)), + qk_rope_head_dim=_int("qk_rope_head_dim", 64), + v_head_dim=_int("v_head_dim", _int("head_dim", 128)), + index_n_heads=_int("index_n_heads", 32), + index_head_dim=_int("index_head_dim", 128), + index_topk=_int("index_topk", 2048), + index_topk_freq=_int("index_topk_freq", 4), + indexer_types=_types("indexer_types"), + n_routed_experts=_req("n_routed_experts"), + n_shared_experts=_int("n_shared_experts", 1), + num_experts_per_tok=_req("num_experts_per_tok"), + moe_intermediate_size=_req("moe_intermediate_size"), + intermediate_size=_int("intermediate_size", 12288), + first_k_dense_replace=_int("first_k_dense_replace", 0), + mlp_layer_types=_types("mlp_layer_types"), + routed_scaling_factor=float(cfg.get("routed_scaling_factor", 1.0) or 1.0), + num_nextn_predict_layers=_int("num_nextn_predict_layers", 0), + index_share_for_mtp_iteration=bool(cfg.get("index_share_for_mtp_iteration", True)), + weight_dtype=weight_dtype, + expert_dtype=expert_dtype, + op_dtype_overrides=overrides, + # The config declares no cache dtype; bf16 is the model fact. A served + # deployment may pick fp8 β€” that is a serving decision, set at deploy time. + kv_dtype=dtype, + act_dtype=dtype, + ) diff --git a/gitm/planner/model_catalogue.py b/gitm/planner/model_catalogue.py index d140e14..6db2921 100644 --- a/gitm/planner/model_catalogue.py +++ b/gitm/planner/model_catalogue.py @@ -9,7 +9,7 @@ CATALOGUE_DIR = Path(__file__).resolve().parent / "models" #: Families a catalogue entry may declare, and the spec each one builds. -_FAMILIES = ("hybrid", "sparse_moe") +_FAMILIES = ("hybrid", "sparse_moe", "glm_moe_dsa") def available() -> list[str]: @@ -69,13 +69,34 @@ def _expand_layer_types(value: Any, n_layers: int) -> tuple[str, ...]: return expanded -def load_entry(name_or_path: str | Path) -> dict[str, Any]: - """The raw catalogue entry, validated for structure but not yet a spec.""" +def load_entry(name_or_path: str | Path, _seen: frozenset[str] = frozenset()) -> dict[str, Any]: + """The raw catalogue entry, validated for structure but not yet a spec. + + ``extends: `` merges this entry's ``spec`` over that of another. It + exists for the case where two entries describe *the same architecture at a + different precision* β€” a bf16 release and its FP8 sibling, which share a + 78-entry indexer schedule and a 78-entry MLP schedule verbatim. Copying those + into both files is 158 lines of duplicated evidence that can drift apart + silently, and it hides the thing worth seeing: the two entries differ only in + their dtypes. ``provenance`` is deliberately *not* merged β€” each checkpoint + was validated against its own published size and has its own open questions. + """ path = _resolve(name_or_path) data = yaml.safe_load(path.read_text()) or {} if not isinstance(data, dict): raise ValueError(f"{path}: expected a mapping at the top level") + base_name = data.pop("extends", None) + if base_name is not None: + key = str(base_name) + if key in _seen: + raise ValueError(f"{path}: 'extends' cycle through {key!r}") + base = load_entry(key, _seen | {key}) + merged = dict(base.get("spec") or {}) + merged.update(data.get("spec") or {}) + data = {**{k: v for k, v in base.items() if k != "provenance"}, **data} + data["spec"] = merged + family = data.get("family") if family not in _FAMILIES: raise ValueError( @@ -102,6 +123,8 @@ def load_spec(name_or_path: str | Path): if family == "hybrid": from gitm.planner.hybrid_graph import HybridMoEModelSpec as cls + elif family == "glm_moe_dsa": + from gitm.planner.glm_graph import GlmMoeDsaModelSpec as cls # type: ignore[assignment] else: from gitm.planner.roofline import SparseMoEModelSpec as cls # type: ignore[assignment] @@ -110,9 +133,26 @@ def load_spec(name_or_path: str | Path): raw["layer_types"] = _expand_layer_types( raw["layer_types"], int(raw.get("n_layers", 0)) ) - for key in ("compress_ratios", "dspark_layer_ids"): + # Per-layer schedule lists that must reach the frozen dataclass as tuples. A + # list would make the spec unhashable; a dropped tuple-coercion here is how a + # schedule silently arrives as the wrong type. + for key in ("compress_ratios", "dspark_layer_ids", "indexer_types", "mlp_layer_types"): if key in raw and isinstance(raw[key], list): raw[key] = tuple(raw[key]) + # Per-layer schedules must cover the model exactly. ``spec_from_hf_config`` + # already refuses a short one; the catalogue path did not, and a schedule one + # entry short does not fail β€” the missing layers fall through to the modulo + # fallback and can land on the right answer by luck, which is a plausible + # total resting on evidence that is not there. That is the exact failure the + # explicit schedules exist to prevent, so it is an error here too. + n_layers = int(raw.get("n_layers", 0) or 0) + for key in ("indexer_types", "mlp_layer_types"): + sched = raw.get(key) + if sched and n_layers and len(sched) != n_layers: + raise ValueError( + f"{name_or_path}: {key} has {len(sched)} entries for {n_layers} " + "layers β€” the schedule must cover the model exactly" + ) # YAML gives lists; the spec is a frozen dataclass and therefore hashable, so # every collection field has to land as something hashable. A list here does # not fail at load β€” it fails later, at the first ``hash(spec)``, a long way @@ -152,6 +192,11 @@ def predict( return predict_hybrid_graph(spec, hw, batch, sharding, **kwargs), family + if family == "glm_moe_dsa": + from gitm.planner.glm_graph import predict_glm_graph + + return predict_glm_graph(spec, hw, batch, sharding, **kwargs), family + from gitm.planner.moe_graph import predict_moe_graph return predict_moe_graph(spec, hw, batch, sharding, **kwargs), family diff --git a/gitm/planner/models/glm-5.2-fp8.yaml b/gitm/planner/models/glm-5.2-fp8.yaml new file mode 100644 index 0000000..69a4575 --- /dev/null +++ b/gitm/planner/models/glm-5.2-fp8.yaml @@ -0,0 +1,103 @@ +# zai-org/GLM-5.2-FP8 β€” the vendor's recommended deployment shape. +# Source: https://huggingface.co/zai-org/GLM-5.2-FP8/blob/main/config.json +# +# Identical architecture to `glm-5.2`, which this extends; the difference is +# precision, and precision is the difference between a model that fits one +# 8xH200 node and one that does not. 753.33 GB on disk against the bf16 +# release's 1.507 TB. Everything not listed below is inherited, so the fields +# that ARE listed are exactly what the FP8 checkpoint changes. + +name: zai-org/GLM-5.2-FP8 +family: glm_moe_dsa +extends: glm-5.2 +description: > + GLM-5.2 quantised to FP8 e4m3 with 128x128 weight blocks. The backbone GEMMs, + the dense FFN, the shared expert and all routed experts are fp8; lm_head, + embed_tokens, the MTP eh_proj and the lightning indexer stay bf16; the router + is fp32. This is the entry to plan against β€” the vendor recommends it and it + is what fits 8xH200. + +spec: + # FP8 e4m3, 128x128 block-scaled, dynamic activation scaling. + weight_dtype: fp8 + expert_dtype: fp8 + + # A serving choice, not a model fact β€” but it is the one the vendor recipe + # makes (`--kv-cache-dtype fp8`), and it halves the latent half of every cached + # entry. The decoupled RoPE key stays bf16 either way; kv_entry_bytes prices + # the two halves separately for exactly this reason. + kv_dtype: fp8 + + # NOT overridden, and deliberately so: act_dtype stays bf16, inherited from + # `glm-5.2`. Weights are fp8 on disk; activations are bf16 in flight and get + # quantised per-GEMM at run time (activation_scheme "dynamic"), which is the + # act_quant node. Every norm, permute and scatter therefore runs bf16 here. + # Listed as a comment because `extends` makes an inherited value invisible, + # and three fp8 lines above it invite the wrong inference. + + # From quantization_config.modules_to_not_convert, plus the router's own dtype + # from the base config. Without these the indexer β€” the one attention node + # whose cost grows with context β€” is priced at half its real weight traffic, + # and lm_head at half of 154,880 x 6,144. + op_dtype_overrides: + - [attn_index_proj, bf16] + - [embed_tokens, bf16] + - [lm_head, bf16] + - [mtp_eh_proj, bf16] + - [moe_router, fp32] + +provenance: + verified: + - claim: predicted weight bytes match the published FP8 checkpoint + detail: > + model.safetensors totals 753,329,940,480 B (753.33 GB) across 141 shards. + model_weight_bytes predicts 755.9 GB β€” +0.34%, and the same shape at bf16 + predicts 1.5079 TB against the bf16 release's 1,506,659,919,872 B (+0.08%). + Two checkpoints at two precisions agreeing to under half a percent is a + stronger check on the shape arithmetic than either one alone. + - claim: the unquantised op list is read, not assumed + detail: > + modules_to_not_convert names lm_head, embed_tokens, eh_proj, enorm/hnorm, + every layernorm, mlp.gate (+ e_score_correction_bias) and the indexer's + projections and k_norm. o_proj is absent from that list β€” it IS quantised + here, the opposite of the fp8-backbone checkpoints that keep the output + projection wide. + - claim: activations are quantised at run time, so the quant kernel is real + detail: > + activation_scheme "dynamic" β€” the graph emits an act_quant node ahead of + each group of fp8 GEMMs. On the bf16 entry there is nothing to quantise + and the node does not exist. + + estimated: + - field: kv_dtype + value: fp8 + detail: > + A deployment choice, taken from the vendor's vLLM recipe + (`--kv-cache-dtype fp8`), not from the checkpoint. Set it back to bf16 to + price the conservative cache; on this architecture that moves the latent + half of a 576-element entry, not the RoPE half. + - field: MLA absorption (attn_kv_b vs attn_out_proj width) + value: unabsorbed β€” kv_b runs as its own GEMM, o_proj narrow + detail: > + An engine that absorbs MLA folds W^UK into the query and W^UV into the + output, dropping attn_kv_b and doubling attn_out_proj's input width to + n_heads*kv_lora_rank (32768->6144). Same resident weights either way; the + two readings move attn_out_proj by 2x, so it re-ranks the attention side. + - field: expert-parallel imbalance + value: 1.0 (perfect balance) + detail: > + Trace-calibrated by design, declared while the no-traces constraint holds. + Real skew moves less weight traffic than predicted (fewer distinct + experts) but lengthens the grouped-GEMM tail. + + unmodelled: + - Absorbed-MLA decode. See the estimated entry above; the one open question + that re-ranks the attention side by 2x in either direction. + - Acceptance rate under MTP. The graph prices D drafts and a 1+D verify; what + fraction is kept is a serving observable, not a config-derivable one. + - index_topk_freq as a *temporal* reuse across decode steps. Modelled as the + spatial layer-group period, which the weight map proves; extra step-to-step + reuse would only reduce indexer cost. + - Expert-capacity padding. If the grouped-GEMM backend pads to a fixed + capacity, every step reads all 256 experts and the distinct_experts union + term is an underestimate at low batch β€” the fork one D2H count settles. diff --git a/gitm/planner/models/glm-5.2.yaml b/gitm/planner/models/glm-5.2.yaml new file mode 100644 index 0000000..bf97f65 --- /dev/null +++ b/gitm/planner/models/glm-5.2.yaml @@ -0,0 +1,200 @@ +# zai-org/GLM-5.2 β€” GlmMoeDsaForCausalLM (MLA + DeepSeek Sparse Attention + MoE). +# Source: https://huggingface.co/zai-org/GLM-5.2/blob/main/config.json +# Read from config.json and model.safetensors.index.json only β€” no traces. + +name: zai-org/GLM-5.2 +family: glm_moe_dsa +description: > + 78 layers, hidden 6144, ~754B total / ~40B active per token. MLA attention + (kv_lora_rank 512, one latent shared across 64 query heads) with a DeepSeek + Sparse Attention lightning indexer selecting top-2048. IndexShare: only 21 of + 78 layers compute the index; the other 57 reuse a neighbour's selection and + carry no indexer weights. First 3 layers dense, remaining 75 MoE (256 experts, + top-8, one shared). One MTP draft head. bf16 throughout. + +spec: + hidden: 6144 + n_layers: 78 + vocab: 154880 + + # MLA attention. head_dim (192) is the nope width; q_head_dim widens to + # 192+64=256, and v_head_dim (256) differs from it β€” score and value reads use + # different per-head widths. num_key_value_heads is 64 in the config but is a + # red herring: the cache holds one kv_lora_rank latent per token, shared across + # all 64 query heads, so KV traffic derives from kv_lora_rank, never 64*head_dim. + n_heads: 64 + q_lora_rank: 2048 + kv_lora_rank: 512 + qk_nope_head_dim: 192 + qk_rope_head_dim: 64 + v_head_dim: 256 + + # DeepSeek Sparse Attention indexer. index_topk_freq 4 is the IndexShare period: + # one 'full' layer computes the top-2048 selection, the next three reuse it. + index_n_heads: 32 + index_head_dim: 128 + index_topk: 2048 + index_topk_freq: 4 + # Read verbatim from the checkpoint (period-4 past the dense prefix). The 'shared' + # layers physically carry no indexer weights β€” this schedule is proven, not fitted. + indexer_types: + - full + - full + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + - full + - shared + - shared + - shared + + # Mixture of experts. moe_intermediate_size 2048 per expert; the first 3 layers + # are dense FFN (intermediate_size 12288) with no router or experts. + n_routed_experts: 256 + n_shared_experts: 1 + num_experts_per_tok: 8 + moe_intermediate_size: 2048 + intermediate_size: 12288 + first_k_dense_replace: 3 + routed_scaling_factor: 2.5 + + # Multi-token prediction. One MTP module, invoked once per drafted token + # (EAGLE-style); the vendor recipe runs it 5 deep. index_share_for_mtp_iteration: + # the draft reuses the main model's selection, and the weight map agrees β€” the + # MTP block carries no indexer tensors, exactly like the 57 'shared' layers. + num_nextn_predict_layers: 1 + index_share_for_mtp_iteration: true + + # The bf16 release: no quantization_config, and 1.507 TB on disk matches an + # all-bf16 read. This is the model fact. It is NOT the deployment shape β€” the + # vendor ships and recommends zai-org/GLM-5.2-FP8 (753.33 GB), which is the + # `glm-5.2-fp8` entry beside this one and the one to plan against. Keep this + # entry for what it answers: what the unquantised model costs, and therefore + # what quantisation is actually worth. + weight_dtype: bf16 + expert_dtype: bf16 + kv_dtype: bf16 + act_dtype: bf16 + + # fp32 even here. moe_router_dtype is a field of the base config, not of any + # quantisation config, so the router's precision is a property of the model and + # not of the deployment β€” the one op that does not follow weight_dtype on an + # otherwise uniformly-bf16 checkpoint. + op_dtype_overrides: + - [moe_router, fp32] + +provenance: + verified: + - claim: IndexShare is read from the weight map, not inferred + detail: > + model.safetensors.index.json carries indexer tensors (*.indexer.*) on + exactly the 21 layers whose indexer_types is 'full' (0,1,2,6,10,...,74) + and on none of the 57 'shared' layers. A shared layer that recomputed the + index would need those weights; it does not have them. + - claim: dense/sparse MLP split matches the tensor map + detail: > + Layers 0-2 carry no expert tensors (dense FFN); layers 3-77 and the MTP + layer carry experts.* and shared_experts.*. Matches first_k_dense_replace 3. + - claim: predicted weight bytes match the published checkpoint + detail: > + model.safetensors.index.json total_size is 1,506,659,919,872 B (1.507 TB) + across 282 shards. model_weight_bytes at bf16 lands within a few percent; + the residual is norms/biases and the coarse MTP roll-in. + + estimated: + - field: index_share_for_mtp_iteration handling + value: MTP emitted as a shared (no-indexer) block + detail: > + The MTP layer carries indexer weights in the map, but the config sets + index_share_for_mtp_iteration=true, so the iteration reuses the main + model's index. Emitting the draft as a full-indexer block would over- + predict a scan the runtime skips; the alternative reading is a headroom + lever to check once a trace exists. + - field: MLA absorption (attn_kv_b vs attn_out_proj width) + value: unabsorbed β€” kv_b runs as its own GEMM, o_proj narrow + detail: > + The graph models MLA unabsorbed: attn_kv_b reconstructs per-head K/V from + the latent and attn_out_proj is n_heads*v_head_dim -> hidden (16384->6144). + A serving engine that absorbs MLA (the common vLLM decode path) folds W^UK + into the query and W^UV into the output, dropping attn_kv_b and doubling + attn_out_proj's input width to n_heads*kv_lora_rank (32768->6144). Same + resident weights either way (validated footprint holds), but the two + readings move attn_out_proj by 2x β€” which is exactly headroom item #3 to + settle against a capture: does the engine run absorbed MLA? + + unmodelled: + - index_topk_freq as a *temporal* reuse across decode steps. Modelled here as + the spatial layer-group period (which the weight map proves); any additional + step-to-step reuse would only reduce indexer cost further. + - Absorbed-MLA decode. See the estimated entry above. + - Acceptance rate under MTP. The graph prices D drafts and a 1+D verify; the + fraction kept is a serving observable, not a config-derivable one. diff --git a/gitm/planner/registry.py b/gitm/planner/registry.py index 0135ce5..8087a8d 100644 --- a/gitm/planner/registry.py +++ b/gitm/planner/registry.py @@ -10,12 +10,20 @@ def detect_family(cfg: dict[str, Any]) -> str: - """``"hybrid"`` | ``"sparse_moe"`` | ``"dense"`` for a HuggingFace config.""" + """``"hybrid"`` | ``"glm_moe_dsa"`` | ``"sparse_moe"`` | ``"dense"`` for a config.""" + from gitm.planner.glm_graph import is_glm_moe_dsa_config from gitm.planner.hybrid_graph import is_hybrid_moe_config from gitm.planner.moe_graph import is_sparse_moe_config + # The hybrid guard reads ``num_experts``; GLM and V4 both spell it + # ``n_routed_experts``, so they fall through it. GLM must be tested *before* + # sparse_moe: both carry ``index_topk`` + ``n_routed_experts``, so the + # structural sparse-MoE test would claim GLM first β€” the model_type check is + # the clean separator and has to win. if is_hybrid_moe_config(cfg): return "hybrid" + if is_glm_moe_dsa_config(cfg): + return "glm_moe_dsa" if is_sparse_moe_config(cfg): return "sparse_moe" return "dense" @@ -28,6 +36,10 @@ def spec_from_hf_config(cfg: dict[str, Any], *, name: str | None = None): from gitm.planner.hybrid_graph import spec_from_hf_config as _hybrid return _hybrid(cfg, name=name) + if family == "glm_moe_dsa": + from gitm.planner.glm_graph import spec_from_hf_config as _glm + + return _glm(cfg, name=name) if family == "sparse_moe": from gitm.planner.moe_graph import spec_from_hf_config as _sparse @@ -61,15 +73,21 @@ def predict_for_config( from gitm.planner.hybrid_graph import spec_from_hf_config as _hybrid return predict_hybrid_graph(_hybrid(cfg, name=name), hw, batch, sharding), family + if family == "glm_moe_dsa": + from gitm.planner.glm_graph import predict_glm_graph + from gitm.planner.glm_graph import spec_from_hf_config as _glm + + return predict_glm_graph(_glm(cfg, name=name), hw, batch, sharding), family if family == "sparse_moe": from gitm.planner.moe_graph import predict_moe_graph from gitm.planner.moe_graph import spec_from_hf_config as _sparse return predict_moe_graph(_sparse(cfg, name=name), hw, batch, sharding), family raise NotImplementedError( - f"{name or 'this checkpoint'} is neither a hybrid linear-attention MoE nor a " - "DeepSeek-V4-class sparse-MoE checkpoint. The dense graph models it, but has " - "no config reader β€” construct a ModelSpec and call predict_graph directly." + f"{name or 'this checkpoint'} is neither a hybrid linear-attention MoE, a " + "GLM-5.2-class glm_moe_dsa, nor a DeepSeek-V4-class sparse-MoE checkpoint. The " + "dense graph models it, but has no config reader β€” construct a ModelSpec and " + "call predict_graph directly." ) @@ -102,6 +120,12 @@ def add_plan_arguments(ap: argparse.ArgumentParser) -> argparse.ArgumentParser: help="Context already cached before this chunk (0 for a first chunk).") ap.add_argument("--prefill-requests", type=int, default=1, help="How many prompts those tokens belong to β€” sets lm_head rows.") + ap.add_argument("--spec-tokens", type=int, default=0, + help="Speculative (MTP) draft tokens per step. Adds a D-deep " + "draft chain and makes the backbone a 1+D-row verify.") + ap.add_argument("--acceptance-rate", type=float, default=0.0, + help="Fraction of drafted tokens the verifier keeps. Only " + "affects the reported token rate, never the step floor.") ap.add_argument("--launch-overhead", type=float, default=None, help="Seconds per dependent kernel launch. Default 2e-6 " "(CUDA-graph replay); eager is nearer 5e-6, and the " @@ -164,6 +188,10 @@ def _predict(spec, family: str, hw, batch, sharding): from gitm.planner.hybrid_graph import predict_hybrid_graph return predict_hybrid_graph(spec, hw, batch, sharding) + if family == "glm_moe_dsa": + from gitm.planner.glm_graph import predict_glm_graph + + return predict_glm_graph(spec, hw, batch, sharding) from gitm.planner.moe_graph import predict_moe_graph return predict_moe_graph(spec, hw, batch, sharding) @@ -235,16 +263,36 @@ def ridge_for(dtype: str) -> float: ) n_compute = sum(1 for n in g.nodes if n.prediction.bound == "compute") + n_launch = sum(1 for n in g.nodes if n.prediction.bound == "launch") out += [ "", f" floor {total * 1e3:.3f} ms/step " + ( f"{g.batch.prefill_tokens / total:,.0f} tok/s " f"prefilling {g.batch.prefill_tokens:,} tokens" if g.batch.is_prefill - else f"{g.batch.batch / total:,.0f} tok/s at batch {g.batch.batch}" + # ``tokens_per_step`` is the accepted-token count: the batch on a + # plain decode step, and the prefix-chain expectation once drafting is + # on. Reporting ``batch / total`` there would price D drafts and then + # credit none of them. + else f"{g.batch.tokens_per_step / total:,.0f} tok/s at batch " + f"{g.batch.batch}" + + (f", D={g.batch.speculative_tokens} " + f"alpha={g.batch.acceptance_rate:g}" + if g.batch.speculative_tokens > 0 else "") ), - f" {len(g.nodes)} nodes, {n_compute} compute-bound", + f" {len(g.nodes)} nodes, {n_compute} compute-bound, " + f"{n_launch} launch-bound", ] + if any(b for b in bounds.values() if len(b) > 1): + out.append(" * this op's instances do not share a bound β€” the label is " + "the majority one") + if g.batch.speculative_tokens > 0 and g.batch.acceptance_rate <= 0: + # A speculative step with no acceptance rate given prices the work and + # reports one accepted token, which is the floor rather than the outcome. + out.append( + f" ! speculative step (D={g.batch.speculative_tokens}) with no " + "--acceptance-rate: the rate above assumes every draft is rejected" + ) if g.has_unpriced_collectives: out.append(" ! collectives unpriced β€” this SKU has no interconnect bandwidth " "in the catalogue") @@ -314,7 +362,10 @@ def main(argv: list[str] | None = None) -> int: f"kv_len={args.kv_len}, TP={args.tp} EP={args.ep}") print(f" {'batch':>7s} {'ms/step':>10s} {'tok/s':>12s} {'compute-bound':>14s}") for b in sizes: - g = _predict(spec, family, hw, BatchConfig(batch=b, kv_cache_len=args.kv_len), + g = _predict(spec, family, hw, + BatchConfig(batch=b, kv_cache_len=args.kv_len, + speculative_tokens=args.spec_tokens, + acceptance_rate=args.acceptance_rate), sharding) cb = sum(1 for n in g.nodes if n.prediction.bound == "compute") print(f" {b:7d} {g.total_pred_s * 1e3:9.3f} " @@ -323,6 +374,8 @@ def main(argv: list[str] | None = None) -> int: batch = BatchConfig( batch=args.batch, kv_cache_len=args.kv_len, + speculative_tokens=args.spec_tokens, + acceptance_rate=args.acceptance_rate, prefill_tokens=args.prefill_tokens, prefill_context=args.prefill_context, prefill_requests=args.prefill_requests, ) @@ -339,6 +392,8 @@ def main(argv: list[str] | None = None) -> int: "hardware": hw.name, "sharding": {"tp": args.tp, "ep": args.ep, "dp": args.dp}, "batch": {"batch": args.batch, "kv_cache_len": args.kv_len, + "speculative_tokens": args.spec_tokens, + "acceptance_rate": args.acceptance_rate, "prefill_tokens": args.prefill_tokens, "prefill_context": args.prefill_context, "prefill_requests": args.prefill_requests}, diff --git a/gitm/planner/roofline.py b/gitm/planner/roofline.py index 844d0fc..bc2696d 100644 --- a/gitm/planner/roofline.py +++ b/gitm/planner/roofline.py @@ -609,8 +609,23 @@ def tokens_per_step(self) -> float: Always at least ``batch``: the non-speculative token is verified, not drafted, so it is never rejected. + + The speculative term is a **prefix chain**, not a product. A verifier + walks the draft in order and stops at the first rejection, so draft token + *k* is kept only if 1…*k*-1 were also kept: the expectation is + ``sum(alpha**i for i in 0..D)``, not ``1 + D*alpha``. The two are far + apart where it matters β€” at D=5, alpha=0.5 the linear form claims 3.5 + accepted tokens against a real 1.97, overstating throughput 1.8x and + putting break-even at less than a third of its true value. + + This models a single-chain verifier (EAGLE/MTP-style), which is what every + family here drafts with. A tree-attention scheme that verifies several + candidate continuations at once accepts more than one chain and would need + its own term. """ - return self.batch * (1.0 + max(0, self.speculative_tokens) * self.acceptance_rate) + d = max(0, self.speculative_tokens) + a = self.acceptance_rate + return self.batch * sum(a ** i for i in range(d + 1)) @dataclass(frozen=True) diff --git a/tests/fixtures/importers/mixed_dump/nsys_2024_min.sqlite b/tests/fixtures/importers/mixed_dump/nsys_2024_min.sqlite deleted file mode 100644 index 20d880c..0000000 Binary files a/tests/fixtures/importers/mixed_dump/nsys_2024_min.sqlite and /dev/null differ diff --git a/tests/fixtures/importers/mixed_dump/torch_trace_min.json.gz b/tests/fixtures/importers/mixed_dump/torch_trace_min.json.gz deleted file mode 100644 index 647757e..0000000 Binary files a/tests/fixtures/importers/mixed_dump/torch_trace_min.json.gz and /dev/null differ diff --git a/tests/fixtures/importers/nsys_2023_min.sqlite b/tests/fixtures/importers/nsys_2023_min.sqlite deleted file mode 100644 index c06b110..0000000 Binary files a/tests/fixtures/importers/nsys_2023_min.sqlite and /dev/null differ diff --git a/tests/fixtures/importers/nsys_2024_min.sqlite b/tests/fixtures/importers/nsys_2024_min.sqlite index 20d880c..dff4aa3 100644 Binary files a/tests/fixtures/importers/nsys_2024_min.sqlite and b/tests/fixtures/importers/nsys_2024_min.sqlite differ diff --git a/tests/fixtures/importers/nsys_2025_min.sqlite b/tests/fixtures/importers/nsys_2025_min.sqlite deleted file mode 100644 index 7192209..0000000 Binary files a/tests/fixtures/importers/nsys_2025_min.sqlite and /dev/null differ diff --git a/tests/fixtures/importers/parity_nsys.sqlite b/tests/fixtures/importers/parity_nsys.sqlite deleted file mode 100644 index 9b64d72..0000000 Binary files a/tests/fixtures/importers/parity_nsys.sqlite and /dev/null differ diff --git a/tests/fixtures/importers/torch_trace_min.json.gz b/tests/fixtures/importers/torch_trace_min.json.gz deleted file mode 100644 index 647757e..0000000 Binary files a/tests/fixtures/importers/torch_trace_min.json.gz and /dev/null differ diff --git a/tests/test_deviation_alignment.py b/tests/test_deviation_alignment.py index b884cee..6243710 100644 --- a/tests/test_deviation_alignment.py +++ b/tests/test_deviation_alignment.py @@ -38,7 +38,12 @@ def test_classify_op(): assert classify_op("triton_qkv_proj_gemm") == "qkv_proj" assert classify_op("cutlass_down_proj_kernel") == "mlp_down" assert classify_op("lm_head_logits") == "lm_head" - assert classify_op("triton_rms_norm") is None # not a modeled op + # Modelled since the GLM-5.2 graph began emitting the pointwise work: on a + # sparse model at low batch the norms are most of the launches, and a step + # bounded by its launches cannot be explained by a graph of GEMMs alone. + # All three norm sites in a block share this op β€” they are one kernel name, + # and only an NVTX range can say which site a launch belongs to. + assert classify_op("triton_rms_norm") == "rms_norm" def test_classify_op_matches_real_vllm_kernel_names(): diff --git a/tests/test_glm_graph.py b/tests/test_glm_graph.py new file mode 100644 index 0000000..188d492 --- /dev/null +++ b/tests/test_glm_graph.py @@ -0,0 +1,761 @@ +"""The GLM-5.2 (``glm_moe_dsa``) decode graph, pinned against the checkpoint. + +Every assertion guards a term that separates GLM-5.2 from the DeepSeek-V4 sparse +family it was forked from, or a wiring seam a plausible-but-wrong graph would slip +through: + +* IndexShare β€” only ``full`` layers emit indexer nodes; ``shared`` layers reuse + the selection and carry no indexer weights, +* MLA KV traffic scales with the shared latent, never ``n_heads``, +* the dense prefix runs an FFN, not a mixture, +* precision is read per op, not per model β€” bf16 by default, and on the FP8 + checkpoint fp8 everywhere the quantiser went and bf16 where it did not, +* the predicted footprint matches *both* published checkpoints, 1.507 TB bf16 and + 753.33 GB fp8, +* prefill and decode disagree about what ``index_topk`` buys β€” it bounds the core + in both phases but bounds the *bytes* in only one, +* the MTP chain is D stages deep with D vocabulary projections, not one of each, +* ``detect_family`` routes ``glm_moe_dsa`` before the structural sparse-MoE test. +""" + +from __future__ import annotations + +import contextlib +import io +import os +import tempfile +from dataclasses import replace + +import pytest +import yaml + +from gitm.planner.glm_graph import ( + GlmMoeDsaModelSpec, + core_read_entries, + is_glm_moe_dsa_config, + kv_entry_bytes, + model_weight_bytes, + predict_glm_graph, + spec_from_hf_config, +) +from gitm.planner.model_catalogue import available, load_entry, load_spec, predict +from gitm.planner.registry import detect_family +from gitm.planner.roofline import BatchConfig, ShardingConfig + +# GLM-5.2's shape, trimmed to the keys the planner reads. Arrays are the real +# checkpoint's schedules (period-4 IndexShare past a 3-layer dense prefix). +GLM_CONFIG = { + "model_type": "glm_moe_dsa", + "architectures": ["GlmMoeDsaForCausalLM"], + "hidden_size": 6144, + "num_hidden_layers": 78, + "num_attention_heads": 64, + "num_key_value_heads": 64, + "head_dim": 192, + "q_lora_rank": 2048, + "kv_lora_rank": 512, + "qk_nope_head_dim": 192, + "qk_rope_head_dim": 64, + "v_head_dim": 256, + "index_n_heads": 32, + "index_head_dim": 128, + "index_topk": 2048, + "index_topk_freq": 4, + "indexer_types": ( + ["full", "full", "full"] + + ["shared", "shared", "shared", "full"] * 18 + + ["shared", "shared", "shared"] + ), + "n_routed_experts": 256, + "n_shared_experts": 1, + "num_experts_per_tok": 8, + "moe_intermediate_size": 2048, + "intermediate_size": 12288, + "first_k_dense_replace": 3, + "mlp_layer_types": ["dense", "dense", "dense"] + ["sparse"] * 75, + "routed_scaling_factor": 2.5, + "moe_router_dtype": "float32", + "num_nextn_predict_layers": 1, + "index_share_for_mtp_iteration": True, + "dtype": "bfloat16", + "vocab_size": 154880, +} + + +def _spec() -> GlmMoeDsaModelSpec: + return spec_from_hf_config(GLM_CONFIG, name="GLM-5.2") + + +def _ops(g, op: str) -> list: + return [n for n in g.nodes if n.op == op] + + +def test_config_reader_reads_schedules_verbatim(): + spec = _spec() + assert spec.n_layers == 78 + assert spec.n_full_indexer_layers == 21 # 3 dense-prefix + 18 period-4 + assert spec.n_sparse_mlp_layers == 75 + # The dense prefix is dense; a mid-stack layer is sparse. + assert not spec.is_sparse_mlp(0) and spec.is_sparse_mlp(40) + # Read verbatim, not derived: layer 2 is 'full' though 2 % 4 != 0. + assert spec.is_full_indexer(2) + assert not spec.is_full_indexer(3) + + +def test_indexshare_shared_layers_emit_no_indexer(): + """The mechanism the fork exists to price: 57 of 78 layers skip the indexer.""" + spec = _spec() + g = predict_glm_graph(spec, batch=BatchConfig(batch=1, kv_cache_len=4096)) + n_proj = len(_ops(g, "attn_index_proj")) + n_score = len(_ops(g, "attn_index_score")) + # One per full-indexer layer, and the MTP head shares (no indexer node). + assert n_proj == n_score == spec.n_full_indexer_layers == 21 + # Every layer still runs the attention core over the selected positions. + assert len(_ops(g, "attn_score_value")) == spec.n_layers + + +def test_mla_kv_traffic_uses_shared_latent_not_heads(): + """The classic MLA error: charging KV as ``n_heads * head_dim``. + + The cache holds one latent per token, shared across all 64 query heads, so the + per-entry byte count is ``kv_lora_rank + qk_rope`` β€” independent of ``n_heads``. + """ + spec = _spec() + entry = kv_entry_bytes(spec) + assert entry == (512 + 64) * 2 # bf16 latent + bf16 rope key + # A per-head (GQA-style) K+V reading would be tens of times larger. + gqa_wrong = spec.n_heads * (spec.q_head_dim + spec.v_head_dim) * 2 + assert gqa_wrong > 25 * entry + + +def test_dense_prefix_runs_ffn_not_mixture(): + spec = _spec() + g = predict_glm_graph(spec, batch=BatchConfig(batch=1, kv_cache_len=4096)) + # Exactly the 3 dense layers carry an mlp_gate_up/down; the router runs on the + # 75 sparse layers plus the sparse MTP head. + assert len(_ops(g, "mlp_gate_up")) == 3 + # Two per sparse block: the h->256 GEMM, then the fused gating kernel that + # scores and selects. Same op name because the trace cannot tell them apart + # (see MOE_LAYER_NODES), so the count is doubled, not the node list. + assert len(_ops(g, "moe_router")) == 2 * spec.n_sparse_mlp_layers + + +def test_precision_is_bf16_no_fp4_leak(): + """Forked from a fp4-expert default; a leak would deflate the dominant term.""" + spec = _spec() + assert spec.weight_dtype == "bf16" + assert spec.expert_dtype == "bf16" + assert spec.kv_dtype == "bf16" + # fp32 even on the unquantised checkpoint: moe_router_dtype is a base-config + # field, so the router's precision is a model fact, not a deployment one. + assert spec.dtype_for("moe_router", spec.weight_dtype) == "fp32" + + +def test_fp8_checkpoint_reads_what_the_quantiser_skipped(): + """One dtype per model is a fiction here; the checkpoint says which ops differ. + + ``modules_to_not_convert`` is the authority, and the interesting entries are + the ones that invert the usual fp8-backbone layout: ``o_proj`` is quantised + (absent from the list) while the *indexer* is not. + """ + cfg = dict(GLM_CONFIG) + cfg["quantization_config"] = { + "quant_method": "fp8", + "fmt": "e4m3", + "weight_block_size": [128, 128], + "modules_to_not_convert": [ + "model.layers.0.input_layernorm", + "model.layers.47.mlp.gate.e_score_correction_bias", + "model.layers.74.self_attn.indexers_proj", + "model.layers.74.self_attn.indexer.k_norm", + "model.layers.78.eh_proj", + "lm_head", + ], + } + spec = spec_from_hf_config(cfg, name="GLM-5.2-FP8") + assert spec.weight_dtype == "fp8" and spec.expert_dtype == "fp8" + # Quantised: the backbone GEMMs and the experts. + for op in ("attn_q_b", "attn_kv_b", "attn_out_proj", "moe_routed"): + assert spec.dtype_for(op, spec.weight_dtype) == "fp8", op + # Skipped: the vocabulary projection, the MTP fusion, and the indexer. + for op in ("lm_head", "mtp_eh_proj", "attn_index_proj"): + assert spec.dtype_for(op, spec.weight_dtype) == "bf16", op + assert spec.dtype_for("moe_router", spec.weight_dtype) == "fp32" + + +def test_a_skipped_norm_does_not_imply_a_skipped_projection(): + """``indexer.k_norm`` in the not-convert list says nothing about the GEMM. + + Every fp8 scheme leaves norms wide, so a norm appearing there carries no + information. A bare ``indexer`` needle matched it and marked the whole indexer + bf16 β€” right for GLM-5.2, where ``indexers_proj`` is listed too, and wrong for + any checkpoint that skipped only the norm. + """ + cfg = dict(GLM_CONFIG) + cfg["quantization_config"] = { + "quant_method": "fp8", + "modules_to_not_convert": ["model.layers.74.self_attn.indexer.k_norm"], + } + spec = spec_from_hf_config(cfg) + assert spec.dtype_for("attn_index_proj", spec.weight_dtype) == "fp8" + + # Name the projection and it is honoured. + cfg["quantization_config"]["modules_to_not_convert"].append( + "model.layers.74.self_attn.indexers_proj" + ) + spec = spec_from_hf_config(cfg) + assert spec.dtype_for("attn_index_proj", spec.weight_dtype) == "bf16" + + +def test_embed_tokens_carries_its_own_declared_precision(): + """The untied halves are two tensors and the checkpoint names them separately. + + ``embed_tokens`` used to map onto the ``lm_head`` op, which priced the pair + together. That is right for GLM-5.2, where both are in + ``modules_to_not_convert`` β€” and silently wrong for any checkpoint that + quantised one and not the other, with ``dtype_for("embed_tokens")`` answering + fp8 on a model that explicitly does not convert it. + """ + fp8 = load_spec("glm-5.2-fp8") + assert fp8.dtype_for("embed_tokens", fp8.weight_dtype) == "bf16" + assert fp8.dtype_for("lm_head", fp8.weight_dtype) == "bf16" + + # The gather reads the table, so the node runs at the table's width. + node = [n for n in predict_glm_graph(fp8).nodes if n.op == "embed_tokens"][0] + assert node.prediction.dtype == "bf16" + + # And the override is load-bearing: quantising the embedding must move the + # footprint by the size of the table, not by nothing. + quantised = replace( + fp8, + op_dtype_overrides=tuple( + o for o in fp8.op_dtype_overrides if o[0] != "embed_tokens" + ), + ) + table = fp8.vocab * fp8.hidden # one byte per element saved at fp8 + assert model_weight_bytes(fp8) - model_weight_bytes(quantised) == pytest.approx( + table, rel=0.01 + ) + + +@pytest.mark.parametrize("shares", [True, False]) +def test_footprint_counts_the_indexers_the_graph_emits(shares): + """The graph and the footprint must agree on how many indexers exist. + + They are computed independently β€” one walks layers emitting nodes, the other + sums shapes β€” so they can disagree silently. They did: the footprint counted + an indexer for the MTP block while the graph, correctly, emitted none for it, + because ``index_share_for_mtp_iteration`` means it reuses the main selection + and carries no indexer tensors. 18.7 MB of weights the checkpoint does not + have, and a contradiction of the weight-map evidence the note rests on. + + Asserted as an identity rather than a constant, so it holds either way round. + """ + spec = replace(_spec(), index_share_for_mtp_iteration=shares) + # With a draft stage actually running β€” at D=0 the block's weights are + # resident but none of its kernels launch, so node count and footprint + # legitimately diverge there and the identity is about the stage that runs. + g = predict_glm_graph(spec, batch=BatchConfig(batch=1, kv_cache_len=4096, + speculative_tokens=1)) + emitted = len(_ops(g, "attn_index_proj")) + expected = spec.n_full_indexer_layers + (0 if shares else spec.num_nextn_predict_layers) + assert emitted == expected + + # And the footprint moves by exactly one indexer between the two readings. + one = ( + spec.q_lora_rank * spec.index_n_heads * spec.index_head_dim + + spec.hidden * spec.index_head_dim + + spec.hidden * spec.index_n_heads + ) * 2 # bf16 + shared_spec = replace(spec, index_share_for_mtp_iteration=True) + own_spec = replace(spec, index_share_for_mtp_iteration=False) + assert model_weight_bytes(own_spec) - model_weight_bytes(shared_spec) == pytest.approx(one) + + +def test_fp8_footprint_matches_published_checkpoint(): + """The same shape arithmetic, checked against a second published precision. + + Two checkpoints agreeing to under half a percent is a stronger check than + either alone: an error in the shape would have to be precision-proportional + to survive both. + """ + published_fp8 = 753_329_940_480 # 141 shards, zai-org/GLM-5.2-FP8 + spec = replace( + _spec(), weight_dtype="fp8", expert_dtype="fp8", + op_dtype_overrides=( + ("attn_index_proj", "bf16"), ("lm_head", "bf16"), + ("mtp_eh_proj", "bf16"), ("moe_router", "fp32"), + ), + ) + assert abs(model_weight_bytes(spec) / published_fp8 - 1.0) < 0.01 + # And the overrides are load-bearing, not decorative: pricing lm_head and the + # indexer at fp8 loses ~1 GB of real resident weight. + naive = replace(spec, op_dtype_overrides=()) + assert model_weight_bytes(spec) - model_weight_bytes(naive) > 1e9 + + +def test_prefill_core_streams_the_cache_that_decode_only_samples(): + """``index_topk`` bounds the core's FLOPs in both phases β€” its bytes in one. + + At decode a sequence reads its own top-2048 selection. At prefill every query + in the chunk selects a different top-2048 and their union is the whole + history, so the kernel streams the entire cache. A prefill path copied from a + dense family would charge ``P x index_topk`` here and understate long-context + prefill traffic by the ratio of context to 2,048. + """ + spec = _spec() + ctx = 65536 + dec = BatchConfig(batch=1, kv_cache_len=ctx) + pre = BatchConfig(batch=1, kv_cache_len=ctx, prefill_tokens=4096, + prefill_context=ctx, prefill_requests=1) + assert core_read_entries(spec, dec) == spec.index_topk + # Prefill adds the whole history once per request, not another top-k window. + assert core_read_entries(spec, pre) == spec.index_topk + ctx + 4096 + + +def test_prefill_scales_projections_by_chunk_but_not_the_epilogue(): + """Rows and logits rows are different numbers, and lm_head follows the second.""" + spec = _spec() + g = predict_glm_graph( + spec, + batch=BatchConfig(batch=1, kv_cache_len=4096, prefill_tokens=8192, + prefill_requests=2), + ) + dec = predict_glm_graph(spec, batch=BatchConfig(batch=1, kv_cache_len=4096)) + + def flops(graph, op, layer=None): + return sum( + n.prediction.flops for n in graph.nodes + if n.op == op and (layer is None or n.layer == layer) + ) + + # A backbone projection scales with every row in the step: 1 decode position + # plus the 8,192-token chunk riding along with it. Read off one layer β€” the + # draft stage in the same graph runs at one row and no prefill, which is the + # point of keeping the two row counts apart. + assert flops(g, "attn_q_a", layer=3) == pytest.approx( + 8193 * flops(dec, "attn_q_a", layer=3) + ) + # The vocabulary projection scales with rows that need logits: 1 per + # prefilling request plus the decode position, so 3 β€” not 8193. Charging the + # chunk here is the largest single error available on this path. + epi_pre = [n for n in g.nodes if n.op == "lm_head" and n.layer is None] + epi_dec = [n for n in dec.nodes if n.op == "lm_head" and n.layer is None] + assert epi_pre[0].prediction.flops == pytest.approx( + 3 * epi_dec[0].prediction.flops, rel=1e-6 + ) + + +def test_mtp_chain_is_d_deep_with_its_own_vocab_projection(): + """Verify is the backbone at 1+D rows; the draft is D serial stages. + + ``num_nextn_predict_layers`` is 1 β€” one *module*, invoked once per drafted + token. Emitting one draft block and one lm_head for a D-deep chain understates + the draft by D, and the vocabulary projection is the majority of its bytes. + """ + spec = _spec() + d = 5 # the vendor recipe's --speculative-config.num_speculative_tokens + g = predict_glm_graph( + spec, batch=BatchConfig(batch=8, kv_cache_len=8192, speculative_tokens=d) + ) + # One epilogue projection for the verify pass, plus one per draft stage. + assert len(_ops(g, "lm_head")) == 1 + d + assert len(_ops(g, "mtp_eh_proj")) == d + # The backbone still runs once β€” at 1+D rows, not 1+D times. + assert len(_ops(g, "attn_q_a")) == spec.n_layers + d + + # The draft carries no indexer: index_share_for_mtp_iteration is true and the + # MTP block has no indexer tensors in the weight map. + assert len(_ops(g, "attn_index_proj")) == spec.n_full_indexer_layers + + # And the draft is not a small copy of the model: its expert bank is a full + # 256-expert mixture, so the chain's cost is weight traffic paid D times. + assert len(_ops(g, "moe_routed")) == spec.n_sparse_mlp_layers + d + + +def test_two_collectives_per_layer_not_one(): + """A TP layer all-reduces after o_proj and again after the FFN combine. + + Folding them into one node with double the payload gets the bytes right and + the count wrong β€” and at decode payloads a collective is bounded by its ring + latency, so the count is the cost. + """ + spec = _spec() + g = predict_glm_graph( + spec, batch=BatchConfig(batch=1, kv_cache_len=4096), + sharding=ShardingConfig(tp=8), + ) + n_blocks = spec.n_layers # no draft stage without --spec-tokens + assert len(_ops(g, "tp_all_reduce_attn")) == n_blocks + assert len(_ops(g, "tp_all_reduce_mlp")) == n_blocks + + +def test_dense_layers_dispatch_no_experts(): + """Only a mixture layer sends tokens to expert ranks. + + The three dense-FFN layers compute their whole FFN locally. Charging them an + expert-parallel all-to-all puts wire traffic on a block with no experts to + send anything to β€” and on this model that node is half of prefill, so a + spurious three layers of it is not a rounding error. + """ + spec = _spec() + g = predict_glm_graph( + spec, batch=BatchConfig(batch=1, kv_cache_len=4096), + sharding=ShardingConfig(tp=8, ep=8), + ) + assert len(_ops(g, "moe_all_to_all")) == spec.n_sparse_mlp_layers + assert not [n for n in g.nodes if n.op == "moe_all_to_all" and n.layer < 3] + + +def test_unpriced_collectives_stay_visible_under_the_launch_floor(): + """A launch floor must not quietly price a collective the SKU cannot price. + + ``has_unpriced_collectives`` detects a node that moves bytes in zero time. If + every collective carried a 2 us launch cost, an SKU with no interconnect + bandwidth would report a priced graph and credit a sharded deployment with a + nearly-free all-reduce. + """ + from gitm.planner.roofline import HardwareSpec + + g = predict_glm_graph( + _spec(), HardwareSpec(interconnect_bw_bytes_per_s=0.0), + batch=BatchConfig(batch=1, kv_cache_len=4096), + sharding=ShardingConfig(tp=8), + ) + assert g.has_unpriced_collectives + + +def test_footprint_matches_published_checkpoint(): + """Predicted weight bytes within a few percent of the 1.507 TB on disk.""" + published = 1_506_659_919_872 # model.safetensors.index.json total_size + wb = model_weight_bytes(_spec()) + assert abs(wb / published - 1.0) < 0.03 + + +def test_attention_core_flat_indexer_scan_grows_with_context(): + """DSA: the core is bounded by top-k; only the indexer scan grows.""" + spec = _spec() + short = predict_glm_graph(spec, batch=BatchConfig(batch=1, kv_cache_len=4096)) + long = predict_glm_graph(spec, batch=BatchConfig(batch=1, kv_cache_len=131072)) + + def total(g, op): + return sum(n.prediction.t_pred_s for n in g.nodes if n.op == op) + + # Core read is capped at index_topk (2048) β€” unchanged from 4K to 128K. + assert total(long, "attn_score_value") == pytest.approx( + total(short, "attn_score_value"), rel=1e-9 + ) + # The scan scores the whole history, so its *work* grows with context β€” 32x + # from 4K to 128K. Asserted on bytes rather than on time: at 4K the scan moves + # a megabyte across 21 layers and is bounded by its kernel launches, not by + # its bytes, so predicted time there is a launch floor and cannot grow 32x. + # That floor is a real property of the node, not an artefact to assert around. + def total_bytes(g, op): + return sum(n.prediction.bytes for n in g.nodes if n.op == op) + + assert total_bytes(long, "attn_index_score") == pytest.approx( + 32 * total_bytes(short, "attn_index_score") + ) + assert total(long, "attn_index_score") > 5 * total(short, "attn_index_score") + short_scan = [n for n in short.nodes if n.op == "attn_index_score"] + assert all(n.prediction.bound == "launch" for n in short_scan) + + +def test_indexer_scans_with_every_index_head(): + """32 query heads against one shared key per token β€” the head count is a factor. + + ``wk`` produces a single 128-d key per token (MQA-style), which is why the key + *bytes* carry no head factor and the score *FLOPs* do. Dropping it understates + the scan 32x and leaves the one node that grows with context looking free. + """ + spec = _spec() + batch = BatchConfig(batch=32, kv_cache_len=8192) + g = predict_glm_graph(spec, batch=batch) + scan = [n for n in g.nodes if n.op == "attn_index_score"][0] + pairs = 32 * 8192 + assert scan.prediction.flops == pytest.approx( + 2.0 * pairs * spec.index_n_heads * spec.index_head_dim + ) + # The cached keys are read once per sequence and are NOT per-head. + assert scan.prediction.bytes == pytest.approx(pairs * spec.index_head_dim * 2) + + +def test_no_draft_head_without_a_speculative_config(): + """``num_nextn_predict_layers`` says the block exists, not that it runs. + + Drafting happens only under a speculative config. At D=0 nothing is drafted, + so no stage is emitted β€” the weights stay resident (``model_weight_bytes`` + still counts them) and none of their kernels launch. Emitting one anyway + charged a pure decode step 0.3 ms of drafting a server without + ``--speculative-config`` never does. + """ + spec = _spec() + batch = BatchConfig(batch=32, kv_cache_len=8192) + g = predict_glm_graph(spec, batch=batch) + assert not [n for n in g.nodes if n.layer is not None and n.layer >= spec.n_layers] + assert len(_ops(g, "lm_head")) == 1 # the epilogue only + assert len(_ops(g, "attn_q_a")) == spec.n_layers + + # The weights are still resident either way β€” that is the distinction. + assert model_weight_bytes(spec) > model_weight_bytes( + replace(spec, num_nextn_predict_layers=0) + ) + + # One stage per drafted token once a speculative config exists. + g5 = predict_glm_graph(spec, batch=replace(batch, speculative_tokens=5)) + assert len(_ops(g5, "mtp_eh_proj")) == 5 + + +def test_a_pure_prefill_step_runs_no_draft_head(): + """The draft proposes continuations; a prefill chunk has nothing to continue. + + Emitting it anyway puts nodes in the graph that never ran, and since a draft + stage is almost all launch cost at one row, it shows up as a launch facet made + of absent kernels. + """ + spec = _spec() + g = predict_glm_graph( + spec, + batch=BatchConfig(batch=0, kv_cache_len=0, prefill_tokens=8192, + prefill_requests=1), + ) + assert not [n for n in g.nodes if n.layer is not None and n.layer >= spec.n_layers] + assert len(_ops(g, "mtp_eh_proj")) == 0 + # The backbone still runs, and the epilogue still projects one row per prompt. + assert len(_ops(g, "attn_q_a")) == spec.n_layers + assert len(_ops(g, "lm_head")) == 1 + + +#: The node a GLM-5.2 MoE layer lowers to, in issue order. Pinned because the +#: design note's whole low-batch argument is a claim about *how many kernels* a +#: layer is, not just how many bytes it moves β€” and a graph that quietly folds the +#: pointwise work into the GEMM it precedes reports a decode step as memory-bound +#: when it is launch-bound. +MOE_LAYER_NODES = ( + "rms_norm", "act_quant", + "attn_q_a", "attn_q_b", "attn_kv_a", "attn_kv_b", + "attn_score_value", "attn_qnorm_rope_insert", "attn_out_proj", + "tp_all_reduce_attn", "rms_norm", + "moe_router", "moe_router", "act_quant", + "moe_shared", "moe_permute", "moe_routed", "moe_combine", + "moe_all_to_all", "tp_all_reduce_mlp", +) + + +def test_layer_lowers_to_the_documented_node_sequence(): + # The FP8 entry, because two of the 24 nodes are the dynamic activation + # quantisation the bf16 checkpoint does not run. + spec = load_spec("glm-5.2-fp8") + g = predict_glm_graph( + spec, batch=BatchConfig(batch=32, kv_cache_len=8192), + sharding=ShardingConfig(tp=8, ep=8), + ) + shared = tuple(n.op for n in g.nodes if n.layer == 5) # Ls,sh + assert shared == MOE_LAYER_NODES + + # A full-indexer layer is the same sequence with two nodes inserted. + full = tuple(n.op for n in g.nodes if n.layer == 6) # Ls,f + assert len(full) == len(shared) + 2 + assert "attn_index_proj" in full and "attn_index_score" in full + + # A dense layer swaps the whole mixture for three nodes, and so is the only + # block in the model with no data-dependent shape and no expert traffic. + dense = tuple(n.op for n in g.nodes if n.layer == 0) # Ld,f + assert {"act_quant", "mlp_gate_up", "mlp_down"} <= set(dense) + assert "moe_router" not in dense and "moe_all_to_all" not in dense + + +def test_every_emitted_op_name_resolves_from_a_kernel_name(): + """A node the pairing cannot receive is a prediction that never gets checked. + + ``classify_op`` is the fallback identity for a capture with no NVTX ranges + (``docs/kernel_identity.md``), and it matches on the kernel name. An op this + graph emits that no kernel name can classify to would sit in the predicted + graph permanently unmatched while the real kernel landed as unmodeled β€” two + errors in opposite directions, and the per-op residual diff this whole family + exists to support would be quietly decorative. + + Collectives are the documented exception: NCCL kernel names carry no hint of + *which* of a layer's two all-reduces they are, so they are matched by the + coarse taxonomy rather than by op. + """ + from gitm.optimizer.deviation import _OP_RULES + + g = predict_glm_graph( + load_spec("glm-5.2-fp8"), + batch=BatchConfig(batch=32, kv_cache_len=8192, speculative_tokens=2), + sharding=ShardingConfig(tp=8, ep=8), + ) + collectives = {"tp_all_reduce_attn", "tp_all_reduce_mlp", "moe_all_to_all", + "logits_all_gather"} + emitted = {n.op for n in g.nodes} - collectives + assert emitted <= set(_OP_RULES), sorted(emitted - set(_OP_RULES)) + + +def test_prologue_and_epilogue_are_nodes(): + """The step does not begin at layer 0 or end at the last one. + + A gather, a final norm, the vocabulary projection and β€” under TP β€” the logits + all-gather that has to complete before anything can be sampled. + """ + spec = _spec() + g = predict_glm_graph( + spec, batch=BatchConfig(batch=32, kv_cache_len=8192), + sharding=ShardingConfig(tp=8), + ) + ends = [n.op for n in g.nodes if n.layer is None] + assert ends == ["embed_tokens", "rms_norm", "lm_head", "logits_all_gather"] + # Without TP there is nothing to gather. + solo = predict_glm_graph(spec, batch=BatchConfig(batch=32, kv_cache_len=8192)) + assert "logits_all_gather" not in [n.op for n in solo.nodes] + + +def test_act_quant_exists_only_where_a_gemm_is_actually_fp8(): + """Dynamic activation scaling is a kernel the bf16 checkpoint does not run. + + ``activation_scheme: "dynamic"`` means the activation is quantised at run + time, once per group of fp8 GEMMs sharing an input. On the unquantised + checkpoint there is nothing to quantise β€” the sort of difference a single + model-wide dtype cannot express. + """ + bf16 = _spec() + fp8 = replace( + bf16, weight_dtype="fp8", expert_dtype="fp8", + op_dtype_overrides=(("lm_head", "bf16"), ("moe_router", "fp32")), + ) + batch = BatchConfig(batch=32, kv_cache_len=8192) + assert not _ops(predict_glm_graph(bf16, batch=batch), "act_quant") + # Two per block: one ahead of the attention GEMMs, one ahead of the FFN β€” + # the two groups of fp8 GEMMs, with bf16 work in between. + assert len(_ops(predict_glm_graph(fp8, batch=batch), "act_quant")) == 2 * fp8.n_layers + + +def test_accepted_tokens_follow_a_prefix_chain(): + """A verifier stops at the first rejection, so acceptance compounds. + + ``1 + D*alpha`` is the independent-draws answer and overstates throughput by + up to 1.8x at D=5; the CLI reported it, and the design note had to carry a + warning about its own output. Fixed at the source instead. + """ + b = BatchConfig(batch=32, speculative_tokens=5, acceptance_rate=0.5) + chain = sum(0.5 ** i for i in range(6)) + assert b.tokens_per_step == pytest.approx(32 * chain) + assert b.tokens_per_step < 32 * (1 + 5 * 0.5) # strictly under the linear form + + # No speculation, no chain: the term degenerates to the batch. + assert BatchConfig(batch=32).tokens_per_step == 32 + # Perfect acceptance keeps every drafted token. + assert BatchConfig( + batch=1, speculative_tokens=5, acceptance_rate=1.0 + ).tokens_per_step == pytest.approx(6) + + +def test_plan_flags_a_speculative_step_with_no_acceptance_rate(): + """Priced without an acceptance rate, the reported rate is a floor.""" + from gitm.planner.registry import main + + buf = io.StringIO() + with contextlib.redirect_stdout(buf): + main(["glm-5.2-fp8", "--gpu", "H200", "--batch", "32", "--kv-len", "8192", + "--tp", "8", "--ep", "8", "--spec-tokens", "5"]) + assert "no --acceptance-rate" in buf.getvalue() + + buf = io.StringIO() + with contextlib.redirect_stdout(buf): + main(["glm-5.2-fp8", "--gpu", "H200", "--batch", "32", "--kv-len", "8192"]) + assert "speculative step" not in buf.getvalue() + + +def test_detect_family_routes_glm_before_sparse_moe(): + """Both families carry index_topk + n_routed_experts; model_type must win.""" + assert detect_family(GLM_CONFIG) == "glm_moe_dsa" + assert is_glm_moe_dsa_config(GLM_CONFIG) + + +@pytest.mark.parametrize("entry", ["glm-5.2", "glm-5.2-fp8"]) +def test_catalogue_entry_loads_and_predicts(entry): + assert entry in available() + spec = load_spec(entry) + assert spec.n_layers == 78 and spec.n_full_indexer_layers == 21 + # Frozen and hashable: a schedule or an override that arrived as a list would + # look fine until something tried to key on the spec. + assert isinstance(spec.op_dtype_overrides, tuple) + assert hash(spec) + g, family = predict(entry, batch=BatchConfig(batch=1, kv_cache_len=4096)) + assert family == "glm_moe_dsa" + assert g.total_pred_s > 0 + + +def test_catalogue_schedules_cover_the_model_exactly(): + """A schedule one entry short does not fail β€” it falls through and may be right. + + The missing layers take the modulo fallback, which on GLM-5.2 happens to + produce the correct count. A plausible total resting on evidence that is not + there is the exact failure explicit schedules exist to prevent, so the + catalogue path refuses it rather than accepting the luck. + """ + published = ["full"] * 3 + (["shared"] * 3 + ["full"]) * 18 + ["shared"] * 3 + for entry in ("glm-5.2", "glm-5.2-fp8"): + spec = load_spec(entry) + assert len(spec.indexer_types) == spec.n_layers == 78 + assert list(spec.indexer_types) == published + + from gitm.planner.model_catalogue import CATALOGUE_DIR + from gitm.planner.model_catalogue import load_spec as _load + + short = yaml.safe_load((CATALOGUE_DIR / "glm-5.2.yaml").read_text()) + short["spec"]["indexer_types"] = short["spec"]["indexer_types"][:-1] + with tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False) as fh: + yaml.safe_dump(short, fh) + path = fh.name + try: + with pytest.raises(ValueError, match="must cover the model exactly"): + _load(path) + finally: + os.unlink(path) + + +def test_fp8_entry_inherits_the_schedule_it_shares(): + """Two precisions of one architecture, and only the dtypes are written twice. + + The FP8 entry ``extends`` the bf16 one. Copying a 78-entry schedule into both + files is duplicated evidence that can drift apart silently β€” and it did: + the schedule was one entry short for a while, in the file that had it twice. + ``provenance`` is deliberately not inherited; each checkpoint is validated + against its own published size. + """ + bf16, fp8 = load_spec("glm-5.2"), load_spec("glm-5.2-fp8") + assert fp8.indexer_types == bf16.indexer_types + assert fp8.n_layers == bf16.n_layers and fp8.hidden == bf16.hidden + assert (fp8.weight_dtype, fp8.kv_dtype) == ("fp8", "fp8") + assert (bf16.weight_dtype, bf16.kv_dtype) == ("bf16", "bf16") + + entry = load_entry("glm-5.2-fp8") + fields = {e["field"] for e in entry["provenance"]["estimated"]} + assert "kv_dtype" in fields # its own, not the base's + + +def test_fp8_entry_is_the_deployable_one(): + """Precision is what decides whether the model fits a node, so it is checked. + + bf16 needs ~11 H200s for weights alone; fp8 fits 8 with room for KV. The two + entries exist to make that comparison, so a drift in either dtype is a real + regression. + """ + bf16, fp8 = load_spec("glm-5.2"), load_spec("glm-5.2-fp8") + assert bf16.weight_dtype == "bf16" and fp8.weight_dtype == "fp8" + per_gpu_h200 = 141e9 + assert model_weight_bytes(bf16) / per_gpu_h200 > 8 + assert model_weight_bytes(fp8) / per_gpu_h200 < 8 + + +def test_tp_must_divide_heads(): + spec = _spec() + with pytest.raises(ValueError, match="does not divide"): + predict_glm_graph(spec, sharding=ShardingConfig(tp=7)) + + +def test_missing_required_field_refuses_default(): + broken = {k: v for k, v in GLM_CONFIG.items() if k != "kv_lora_rank"} + with pytest.raises(ValueError, match="kv_lora_rank"): + spec_from_hf_config(broken) diff --git a/tests/test_moe_graph.py b/tests/test_moe_graph.py index 88bf30c..20dfd36 100644 --- a/tests/test_moe_graph.py +++ b/tests/test_moe_graph.py @@ -410,7 +410,9 @@ def test_tokens_per_step_accounts_for_acceptance(): """Drafts are paid for always and counted only when kept.""" b = BatchConfig(batch=4, speculative_tokens=3, acceptance_rate=0.5) assert b.positions_per_step == 16 # all drafted work is computed - assert b.tokens_per_step == pytest.approx(4 * (1 + 3 * 0.5)) + # A prefix chain, not 1 + D*alpha: the verifier stops at the first rejection, + # so token k counts only if 1..k-1 did. + assert b.tokens_per_step == pytest.approx(4 * (1 + 0.5 + 0.25 + 0.125)) # ── the observed side lines up with the predicted side ────────────────────── @@ -1158,6 +1160,10 @@ def test_defaults_do_not_match_any_catalogued_checkpoint(): """Stronger than a size bound: no field-by-field match with a real entry.""" d = HybridMoEModelSpec() for entry in available(): + # Only hybrid entries are HybridMoEModelSpecs; other families (glm_moe_dsa, + # sparse_moe) have their own reference-default tests and their own fields. + if load_entry(entry).get("family") != "hybrid": + continue hybrid_spec = load_spec(entry) assert (hybrid_spec.hidden, hybrid_spec.n_layers, hybrid_spec.num_experts) != ( d.hidden, d.n_layers, d.num_experts