Model GLM-5.2 (glm_moe_dsa): prefill, MTP chain, per-op precision

[nicholaslawrence-hub]

Models GLM-5.2 (glm_moe_dsa) end to end and writes up the result. Everything is read from config.json and the checkpoint's tensor index; no traces.

Two catalogue entries: glm-5.2 (bf16, the model fact) and glm-5.2-fp8 (753 GB, what the vendor actually recommends and the only one that fits 8xH200). The FP8 entry extends the bf16 one, so the two differ only in dtypes.

Graph work

• Precision is per-op (op_dtype_overrides), read from modules_to_not_convert and moe_router_dtype. The layout is backwards from the fp8 models we've seen: o_proj is quantised, the indexer isn't.
• Prefill. Not a copy of the hybrid path — DSA makes index_topk bound the core's FLOPs in both phases and its bytes in neither, so there are separate helpers for pairs vs entries. Aliasing attention_qk_pairs here is wrong in both directions at once and the errors partly cancel, which is what makes it dangerous.
• A layer emits its actual kernels, not just its GEMMs: norms, activation quant, the fused gating, prologue/epilogue. Without those the launch bound had nothing to bind and B=1 read as memory-bound when it's 70% launches.
• MTP is a D-stage chain with an lm_head per stage (--spec-tokens, --acceptance-rate). GLM's draft block carries a full 256-expert bank, so a stage is 54% expert weights and the chain is 5.3% of the step, not the 1–2% a dense draft head costs.
• Two collectives per layer instead of one folded node, and the EP all-to-all only on layers that have experts.

Bugs found while building it

Seven, all of the same shape — two numbers that should agree, derived independently, with nothing checking that they did. Each fix is an identity test rather than a pinned constant.

• indexer_types in the catalogue had 77 entries for 78 layers. Layer 77 took the modulo fallback, landed on the right answer, and the floor was byte-identical — so nothing failed. The loader validates schedule length now.
• attn_index_score had lost its 32-head factor in the prefill rewrite. Doesn't move decode (memory-bound either way); doubles the prefill row.
• model_weight_bytes counted an indexer for the MTP block, which carries none. The graph emitted 21 and the footprint counted 22 — 18.7 MB on a 1.5 TB model, invisible to any size check.
• embed_tokens mapped onto the lm_head op, so dtype_for("embed_tokens") answered fp8 on a checkpoint that explicitly doesn't convert it. Right number, wrong question.
• The indexer needle also matched indexer.k_norm, so a norm in the skip list marked the whole indexer bf16. Every fp8 scheme leaves norms wide, so that name carries no information.
• The draft head ran at D=0, charging a pure decode step 0.3 ms of drafting a server without --speculative-config never does.
• BatchConfig.tokens_per_step counted 1 + D·α, the independent-draws answer, where a verifier accepts a prefix. Overstated throughput 1.8× at D=5, α=0.5 and put break-even at less than a third of its real value. Shared with every family and wrong for all of them.

A conclusion that changed

Prefill isn't memory-bound. The whole-pass AI divided FLOPs by HBM bytes and NVLink payload together, then compared that to a ridge derived from HBM bandwidth. Against HBM alone it's AI 410 against a ridge of 412 — balanced — and 56% of the floor is wire. The label is comm, and it reverts to HBM-bound under TP8-only, which is capture C5.

Outside the family

• gitm plan reports accepted tokens rather than batch/step (so --acceptance-rate does something), and warns when a speculative step is priced without one.
• _OP_RULES gains entries for the ops the pointwise lowering added, plus attn_index_proj (was shadowed by the scan's own needle), attn_kv_b and mtp_eh_proj. Left the gating and silu mappings alone — those are shared with the other families.
• Both review workflows: they failed PRs on upstream errors (a Gemini 503, an expired key), and the diff was filtered then truncated in path order, so docs/ ate the whole 75 KB budget and ~19 reviews never saw a line of Python. Code is diffed first now, and API failures warn instead of blocking.

Validation

Predicted weight bytes land within 0.1% of the bf16 checkpoint and 0.4% of the fp8 one. The third check is the useful one: Z.ai publishes 744B params, the checkpoint is 753.3B by its own bytes, and dropping the MTP block predicts 744.2B. That puts the block at 9.9B against 0.23B for a dense head — independent confirmation of the MoE-draft reading, from a number published for an unrelated reason.

Not done, deliberately

• No depends_on edges (G10). Cross-family IR change, already on the roadmap. expected_stream_id is set on collectives but nothing reads it today, so it's a hook, not a wiring — and it defaults to 0, indistinguishable from an explicit compute stream, which whoever wires it should fix first.
• library.yaml's op vocabulary has no moe_routed, so the levers meant for expert traffic scope to [mlp_gate_up, mlp_down]. That's 74% of a decode step unreachable by the tooling. Pre-existing and identical for DeepSeek-V4, so it should land where both families can be checked — written up as G11.

Merge note

main's #101 (MiMo-V2.5) independently fixed the same fp32-peak, per-dtype-ridge and node-owned-bound gaps — both branches were working from the same design note. Took main's implementations for all three; BatchConfig.attention_qk_pairs also became a method taking a window, which auto-merged cleanly and then failed at runtime. GLM has no window, so it's called unwindowed.

Deleted the three committed node-dump JSONs (29k lines that went stale every time the graph changed); the commands at the top of the note regenerate them.

Suite is at the pre-existing 4 failures (test_bench, test_hft_intervention, two in test_importers) — verified on a clean tree.

🤖 Generated with Claude Code

[github-actions]

🤖 Claude Code Review

Review of docs/glm-5.2/DESIGN-NOTE.md

This is a documentation-only diff — a detailed architecture analysis and benchmarking design note. No executable code is changed. Feedback is scoped accordingly.

---

🐛 Bugs

• Section numbering gap: Sections jump from §4.2 directly to §4.4 ("Flip-variable index") — §4.3 is missing entirely. Either a section was deleted without renumbering or it was never written. Cross-references to §4.3 don't exist, but the gap will confuse readers and suggests content may have been dropped.

• MTP node count inconsistency: §2.4 states the draft chain is "115 nodes" and the backbone is "1,591 nodes," but §3.3's table shows vanilla decode as "1,614 nodes" with the draft as "23 nodes" within that. Then MTP total is "1,706 nodes" (1,591 + 115). The draft-within-vanilla (23 nodes) vs. draft-chain (115 nodes) distinction is explained, but the relationship between 1,614 and 1,591+23 is implicit and error-prone — 1,614 = 1,591 + 23 should be stated explicitly since readers will otherwise question whether the MTP chain double-counts.

• Verify row count stated as 192: §3.3 says "verify, 192 rows" for B=32, D=5. This implies 32 × 6 = 192 rows, which is correct for verifying 5 draft tokens + 1 original. But §2.4's flowchart says "6 rows per seq" for verify, consistent with D=5. The arithmetic is right but "192 rows" appears without derivation in the table and could be misread as a separate batch size.

• Break-even formula denominator: §3.3 states break-even requires α > (1.70 − 1)/5 = 0.140. This is only correct if accepted tokens scale linearly with α and the cost is fixed at 1.70×. The formula assumes the denominator is D=5, but the actual relationship between α and effective throughput depends on the geometric distribution of acceptance — the simplified formula understates break-even when acceptance is correlated across positions.

---

📊 Reproducibility

• gitm plan commands at top are not pinned to a version or commit. The note says "reproduce any figure here" but the planner is described as actively changing (G1–G11). A reader running these commands against a different main may get different outputs. Recommend adding gitm --version or a git SHA.

• ep_imbalance = 1.0 declared, not fitted (A8). This is honestly flagged, but it means the expert DRAM predictions are a lower bound that could be substantially wrong. The note would benefit from a stated uncertainty range (e.g., "measured imbalances of 1.1–1.3 in comparable MoE deployments add 10–25% to the routed expert term").

---

💡 Suggestions

• G11 (intervention library gap) should be a numbered gap like G1–G10, not a narrative paragraph after the table. It describes a concrete, pre-existing defect affecting the tooling's ability to act on its largest predicted cost node (74% of decode). Burying it in prose after the table makes it easy to miss in triage.

• S1 confidence is "low" but its consequences are catastrophic. If 76 D2H syncs per token are real, CUDA-graph capture is impossible and the entire low-batch analysis changes. Consider promoting this to a pre-trace gate check rather than leaving it at rank 2 in the headroom table — it should be the first thing measured before any timeline is opened.

• Q11 (prefill attention bytes) is listed last despite potentially being a 64× error on a node's traffic. The ordering of open questions by "what it changes" places it after questions with smaller impact. At long context, if the prefill core re-reads per query tile rather than per request, it could invalidate the prefill roofline conclusion.

• The --enable-expert-parallel absence from the vendor recipe is discussed but the implication is understated: if production is actually TP8-only, the moe_all_to_all rows don't exist and the entire rank-1 headroom hypothesis (44% of prefill) evaporates. This should be elevated to a top-line uncertainty alongside Q2 (CUDA-graph capture), not buried mid-paragraph.

• Appendix A.1 table is cut off at the last line of the diff ("folding a"). The note appears incomplete. If this is intentional (A.2–A.4 omitted from the diff), it should be stated; if accidental, the remaining archetypes are missing from the review.

[github-actions]

Code Review by Gemini

No bugs or issues found. The code is well-structured, thoroughly documented, and the test suite is comprehensive. The changes correctly implement the GLM-5.2 model's unique architectural features, including per-op precision, DSA asymptotics for prefill, and the MTP chain. The fixes for existing planner gaps (G1-G9 from the design note) are also well-integrated.

[nicholaslawrence-hub]

Pushed 99512c7. Taking the review points in turn.

Fixed

• Break-even was wrong, and this was the good catch. (1.70−1)/5 = 0.140 follows BatchConfig.tokens_per_step, which counts 1 + D·α. That's right for independent draws and wrong here — the verifier accepts a prefix, so expected accepted length is Σ αⁱ. At D=5 the linear form overstates accepted tokens by 1.8× at α=0.5, and real break-even is 0.415, not 0.140. §3.3 now gives both rows and says which to read; propagated to §4.2, §4.3 and the open-questions table. Not touching BatchConfig — that convention is shared with every family, so it's flagged rather than changed here.
• §4.4 renumbered to §4.3. §4.3 was deleted as a duplicate of §2.2 and the gap was left behind.
• Node arithmetic stated: 1,614 = 1,591 + 23, 1,706 = 1,591 + 115, and 192 rows = B × (1+D).
• G11 moved into the gap table.
• Expert parallelism promoted to Q1 and the prefill-tile question to Q3. Agreed on both — if EP is off, rank 1 doesn't exist at all, and Q3 is a 64× swing sitting at the bottom of the list.
• Reproduce commands noted as branch-relative.

Not doing

• The ep_imbalance uncertainty range would mean inventing measurements we don't have. A8 already says it's declared rather than fitted.
• S1 is already the first measurement after the launch args (§6.4 row 1), and §5.2 calls it the cheapest measurement in the document.
• The Appendix A.1 truncation is the diff view, not the file — A.2–A.4 are there.

lint was UP038 on an isinstance tuple; the rest of the file already used X | Y.

🤖 Addressed by Claude Code

[github-actions]

🤖 Claude Code Review

Review

This diff is a 1,054-line documentation file — no executable code is changed. Feedback is scoped accordingly.

---

🐛 Bugs

• MTP node count arithmetic is inconsistent. §3.3 states 1,614 = 1,591 + 23 and 1,706 = 1,591 + 115, but §2.4's diagram says the MTP draft chain is "115 nodes" and the backbone is "1,591 nodes", while §3.3's vanilla-decode row lists 1,614 total. If vanilla decode already includes one draft stage (23 nodes), then 1,591 + 23 = 1,614 holds — but the verify row reuses "1,591 nodes" as if the draft is absent. This needs a sentence clarifying what the 1,591 count does and does not include, or the two totals will mislead whoever aligns a trace against it.

• Prefill attention-bytes claim is internally contradicted. §4.2 states "bytes are the whole cache once per request" for the attention core, but assumption A9 says this is an optimistic floor that a tiled kernel violates by up to C/2048. §2.1's diagram does not flag this uncertainty. The asymmetry (flagged in Q3 but not in the diagram or the §4.2 bound label) will cause a reader to classify a tiled-kernel measurement as an anomaly rather than a known branch.

• tokens_per_step convention bug is documented but not corrected. §3.3 correctly identifies that 1 + D·α overstates throughput by up to 1.8× relative to the correct prefix-chain Σ αⁱ, and explicitly defers the fix. The CLI output therefore prints an incorrect break-even (0.140 vs 0.415). If this note ships alongside the planner, any throughput number printed by gitm plan --spec-tokens 5 is misleading without a visible warning in the CLI output itself — not just in a design note.

---

🔒 Security

• The vllm serve command reproduced verbatim includes --enable-auto-tool-choice and --tool-call-parser glm47. If this document is used as a copy-paste recipe (which §9 encourages), operators may not notice that tool-call parsing expands the model's attack surface for prompt injection via tool responses. A one-line caveat is warranted given the audience is expected to run this command directly.

---

⚡ Performance

• Chunk-size warning lacks a recommended value. §3.2 shows a 14.9× byte multiplier across chunk sizes and §5 rank 8 flags it, but neither §9 nor the gitm plan command examples suggest a concrete --max-num-batched-tokens value or a formula for the crossover. A reader following §9 step-by-step has no actionable guidance.

• eh_proj replication across TP ranks (Q12) is listed as unresolved but is one of the larger per-stage costs in the draft chain (152 MB, 19% of draft bytes per §4.1). The note correctly flags it, but given --enable-expert-parallel is already a fork point, whether eh_proj is TP-sharded or replicated should be resolvable from the serving image before any trace is taken. It deserves promotion to a C6-resolvable item.

---

📊 Reproducibility

• git rev-parse HEAD instruction is correct but fragile. The preamble says to check HEAD matches if a number disagrees. G7 and G8 describe non-trivial graph changes; without pinning a specific commit hash in the document, the "reproduce any figure" claim degrades silently as the planner evolves. Consider embedding the hash at doc-generation time rather than asking the reader to verify it.

• ep_imbalance = 1.0 (A8) propagates through every MoE byte estimate and is explicitly declared unfitted. All throughput numbers in §3–§4 carry this assumption silently. A single bracketed note on the summary rows (e.g., the 1,933 tok/s figure) that the number assumes perfect balance would prevent it from being cited out of context.

---

💡 Suggestions

• G11 is marked "no — not shipped" but affects every MoE family, including the existing DeepSeek-V4 path. Deferring it to "when both families' coverage can be checked" is reasonable, but the note doesn't say who owns that check or what the gate condition is. Without a named owner or a linked issue, this gap will persist across model generations.

• G10's expected_stream_id field is carried by the IR but read by nothing. The note acknowledges this ("consumed by no one"). Given that ranks 1 and 7 in §5 both depend on overlap being observable, and §6.1's classification rule requires distinguishing serial from parallelisable gaps, this is a blocking dependency for the validation plan — not just a roadmap item. It should be flagged as a precondition for §6, not buried in §7.1.

• Appendix A.1 is truncated (the last row cuts off mid-sentence). If this is intentional for the diff, fine — but if the table is used as a test fixture (test_layer_lowers_to_the_documented_node_sequence is referenced), an incomplete table will cause the test to be incomplete as well.

[github-actions]

Code Review by Gemini

## Code Review

The changes introduce comprehensive support for the GLM-5.2 model, including its unique architecture features like MLA, DeepSeek Sparse Attention (DSA) with IndexShare, Mixture-of-Experts (MoE), and Multi-Token Prediction (MTP). The implementation correctly accounts for per-op precision, phase-specific asymptotics, and detailed kernel modeling. The accompanying design note and extensive test suite demonstrate a thorough understanding of the model and its performance characteristics.

### `gitm/optimizer/deviation.py`

**Improvement:**

*   **Line 109:** The comment for `attn_kv_b` mentions "unabsorbed" MLA. While this is a crucial detail for the model, the comment itself is quite long and could be slightly condensed to focus on the mapping rationale. This is a minor stylistic suggestion.

    ```diff
    --- a/gitm/optimizer/deviation.py
    +++ b/gitm/optimizer/deviation.py
    @@ -110,9 +110,9 @@ _OP_RULES: dict[str, tuple[str, ...]] = {
     "attn_q_a": ("q_a_proj", "q_lora", "q_down"),
     "attn_q_b": ("q_b_proj", "q_up"),
     # `kv_b_proj` was absent here while the only MLA families modelled the
     # *absorbed* decode form, where W^UK folds into the query and W^UV into the
     # output projection and no such kernel is launched. The GLM-5.2 graph models
     # it unabsorbed, so the kernel exists and has a node to land on. The entry is
     # safe either way: an absorbed deployment launches nothing these needles
     # match, so it stays absent rather than mis-attributing.
    -    "attn_kv_b": ("kv_b_proj", "kv_up", "w_uk", "w_uv"),
    +    "attn_kv_b": ("kv_b_proj", "kv_up", "w_uk", "w_uv"), # For unabsorbed MLA
     "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"),
    ```

### `gitm/planner/glm_graph.py`

No issues found. The new file is well-structured, thoroughly commented, and correctly implements the complex GLM-5.2 architecture. The logic for handling different phases (prefill, decode, MTP), precision overrides, and sharding is robust.

### `gitm/planner/model_catalogue.py`

No issues found. The `extends` feature is a valuable addition for managing model variants, and the validation for schedule lengths is correctly implemented.

### `gitm/planner/models/glm-5.2-fp8.yaml`

No issues found. The catalogue entry correctly leverages the `extends` feature and specifies the FP8-specific overrides and provenance.

### `gitm/planner/models/glm-5.2.yaml`

No issues found. The catalogue entry accurately reflects the BF16 GLM-5.2 model's architecture and schedules.

### `gitm/planner/registry.py`

No issues found. The family detection logic correctly prioritizes GLM-5.2, and the `_render_table` improvements provide much-needed detail for analysis.

### `tests/test_deviation_alignment.py`

No issues found. The test update correctly reflects the new `rms_norm` mapping.

### `tests/test_glm_graph.py`

No issues found. This is an exceptionally thorough and well-designed test suite that covers a wide range of architectural details, edge cases, and interactions. It provides strong confidence in the correctness of the GLM-5.2 implementation.

---

**Overall Summary:**

This is an excellent set of changes. The implementation is robust, well-tested, and clearly documented. The attention to detail in modeling the GLM-5.2 architecture, including its unique performance characteristics, is commendable. The new `extends` feature in the catalogue and the enhanced `gitm plan` output significantly improve the tool's capabilities and usability.

[nicholaslawrence-hub]

Pushed 2a010fb.

The best point was the one about the CLI: the note documents that tokens_per_step's 1 + D·α overstates a speculative step, and then ships a planner that prints the overstated number with no label. Anyone running gitm plan --spec-tokens 5 without reading the note gets 0.140 as break-even when it's 0.415. Fixed at the point of use — the plan table now says so when D > 0 and quantifies it when you pass the new --acceptance-rate (report-only, it can't move the step floor):

! speculative step (D=5): rates above assume accepted tokens = 1+D*alpha.
  A verifier accepts a prefix, so the expectation is sum(alpha^i)
  at alpha=0.5 that is 3.50 vs 1.97 tokens/step — the printed rate is 1.78x optimistic

Still not changing BatchConfig on this branch; that convention is shared with every family.

Also fixed: §4.2's prefill attention row asserted "whole cache once per request" as fact while A9 called it an optimistic floor — the row and the §2.1 diagram carry the caveat now, so a tiled-kernel measurement reads as a known branch rather than an anomaly. Labelled the verify row "backbone only" (1,591 excludes the draft — read as ambiguous twice, so that's on me). Put an [A8] marker on the headline 1,933 tok/s since ep_imbalance = 1.0 propagates through every MoE byte term and was silent on the summary rows. Made G10 an explicit precondition for §6 rather than a roadmap line — the classification rule needs serial-vs-parallelisable and nothing reads expected_stream_id. And gave the derivable chunk-size rule (95.7 GB × ceil(P/C), so bytes scale as 1/C) instead of a recommended value.

Not doing three:

• Pinning a commit hash at doc-generation time — it's a hand-written note, not generated; the branch is the pin.
• An ep_imbalance uncertainty range — would mean inventing measurements. A8 already says it's declared, not fitted.
• A prompt-injection caveat on --enable-auto-tool-choice. That command is quoted verbatim to make the hardware assumption auditable as the vendor's, not offered as a production recipe.

G11 has no owner or linked issue — flagging that for a human rather than opening one myself.

🤖 Addressed by Claude Code

[github-actions]

🤖 Claude Code Review

This is a design/documentation note (markdown only), not executable code. The diff contains no Python, CUDA, shell scripts, or other runnable artifacts—it's entirely prose, tables, and Mermaid diagrams describing a predicted execution model.

🐛 Bugs

• Break-even table mislabels rows. The table in §3.3 shows α = 0.0 yielding 1,137 tok/s for both the linear and prefix-chain formulas—that's the vanilla decode baseline, not a speculative decoding result. The row should be labeled "D=0 baseline" or removed from the α-sweep, since Σ α^i for i=0..5 at α=0 is just 1, same as vanilla. Readers comparing columns may incorrectly treat 1,137 as a valid speculative throughput point.

• Node count arithmetic is inconsistent. §3.3 states 1,614 = 1,591 + 23 (backbone + one draft stage) and 1,706 = 1,591 + 115 (backbone + 5 stages). But §4.1's decode table header says 1,614 nodes with the MTP draft included, while §3.3 says vanilla decode with the MTP module is 1,614. If the MTP block is always present (it's a fixed module, not optional), "vanilla decode (D=0)" should be 1,614 and "D=5" should be 1,706, but the verify-only count of 1,591 needs a clearer definition of what's excluded.

• tokens_per_step formula stated as wrong but never corrected. §3.3 correctly identifies that 1 + D·α overstates accepted tokens vs. the prefix-chain Σ αⁱ, calls this a shared BatchConfig semantic error, and then uses the wrong formula's output (α > 0.140) as the first row of the break-even table without a strikethrough or clear "do not use" marker. A reader skimming the table will likely take the 0.140 figure.

📊 Reproducibility

• git rev-parse HEAD instruction without a pinned hash. The note says "check git rev-parse HEAD matches if a number disagrees" but gives no reference hash. Anyone reading this after a merge has no anchor. Pin the hash in the note itself, or generate it at doc-build time.

• ep_imbalance = 1.0 is declared but never validated. A8 acknowledges this is unfitted. Every decode throughput figure (1,933 tok/s, the batch sweep table, MTP economics) flows through this assumption. The note should flag the throughput numbers as having an unknown downward bias, not just note that calibration is "by design."

• Chunk-size sensitivity (14.9×) is stated without a reference chunk size for the vendor recipe. The sweep table compares 1×8192 vs. 2×4096 vs. 8×1024 etc., but the vendor's vllm serve command shown earlier omits --max-num-batched-tokens. Without knowing the default (vLLM currently defaults to 2048 for chunked prefill when enabled), the "predicted floor 264.0 ms" baseline is ambiguous.

💡 Suggestions

• G11 is marked "no" (not shipped) but affects two families. The note flags that moe_routed/moe_shared are absent from kernels/library.yaml's applies_to_kernels, meaning rank-5 interventions (73.8% of decode) can't be targeted. Given the note is otherwise very specific about what's actionable, this gap deserves a tracking issue reference or a concrete "blocked on X" rather than just "should land where both families' coverage can be checked."

• S1 (D2H histogram sync) has conf: low but drives the entire ranks-2/3 fork. The note correctly calls this the key measurement (§6.4 row 1), but the low-confidence label on S1 sits in the synchronization table (§3.4) without a forward reference to §5.2. Readers auditing just §3.4 may not realize this single measurement resolves both rank 2 and rank 3 simultaneously.

• Q3 (per-request vs. per-tile KV reads) is described as an "optimistic floor" but the magnitude isn't bounded. The note says "up to C/2048 more" (64× at 128K context) but doesn't state what fraction of prefill time that would add. Given this could move the prefill conclusion from memory-bound to heavily memory-bound (or shift the compute fraction), a worst-case prefill floor with per-tile reads would strengthen the note's claim that Q3 is high-priority.

• The Mermaid diagrams will not render in standard GitHub markdown preview if the repo's GitHub settings don't enable Mermaid. Since this is a working design note meant to be read alongside traces, consider whether ASCII flowcharts (as used in §2.2) are more portable for the decode diagram than the Mermaid blocks in §2.1 and §2.4.

• G10's expected_stream_id field is described as "carried by the IR and consumed by no one." This is a latent correctness trap: any future code that reads expected_stream_id for overlap analysis will silently get wrong answers for nodes where the field was never set (pre-GLM families). A None sentinel and an explicit "not modeled" state would be safer than a defaulted integer.

[nicholaslawrence-hub]

Condensed it — six lines of why the entry used to be absent, down to three lines of why it's safe now. Didn't take the diff as written: it kept the block and added a trailing # For unabsorbed MLA, which is the same thing said twice.

🤖 Addressed by Claude Code

[github-actions]

Code Review by Gemini

Here's a review of the code changes, focusing on bugs and potential issues, with suggested improvements and exact line references.

gitm/optimizer/deviation.py

Issue 1: Incorrect ordering of embed_tokens in _OP_RULES
The comment for embed_tokens explicitly states it should be "Before lm_head" to prevent lm_head's "embed" needle from incorrectly claiming the input gather. However, in the provided diff, embed_tokens is placed after lm_head. This violates the "first entry wins" rule of _OP_RULES and could lead to misclassification of embed_tokens kernels.

Suggested change:
Move the embed_tokens entry to appear before lm_head in the _OP_RULES dictionary.

--- a/gitm/optimizer/deviation.py
+++ b/gitm/optimizer/deviation.py
@@ -129,6 +129,9 @@
     # 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` was absent here while the only MLA families modelled the
@@ -144,9 +147,6 @@
     "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"),
 }

Issue 2: Incorrect ordering of mtp_eh_proj in _OP_RULES
Similar to embed_tokens, the comment for mtp_eh_proj states it should be "Before lm_head". However, it is currently placed after lm_head. This could lead to lm_head's "embed" needle (or other needles) incorrectly claiming mtp_eh_proj kernels if their names overlap.

Suggested change:
Move the mtp_eh_proj entry to appear before lm_head in the _OP_RULES dictionary.

--- a/gitm/optimizer/deviation.py
+++ b/gitm/optimizer/deviation.py
@@ -147,9 +150,9 @@
     "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"),
 }

gitm/planner/glm_graph.py

Issue 3: Missing rms_norm in _UNQUANTISED_OPS
The design note (Section 1, "FP8 — what is and is not quantized") states that "all norms" are BF16 and are "named in modules_to_not_convert". The rms_norm operation is now a modeled node in the graph. If rms_norm is in modules_to_not_convert (as per the design note), it should be explicitly mapped in _UNQUANTISED_OPS to ensure its precision is correctly set to model_dtype (BF16 for the FP8 model) instead of falling back to act_dtype (which might be FP8).

Suggested change:
Add an entry for rms_norm (using a common needle like "layernorm") to the _UNQUANTISED_OPS tuple.

--- a/gitm/planner/glm_graph.py
+++ b/gitm/planner/glm_graph.py
@@ -1009,6 +1009,7 @@
     ("embed_tokens", "lm_head"),  # untied, but priced together in the epilogue
     ("eh_proj", "mtp_eh_proj"),
     ("indexer", "attn_index_proj"),
     ("indexers_proj", "attn_index_proj"),
     ("mlp.gate", "moe_router"),
+    ("layernorm", "rms_norm"), # All norms are bf16
 )
 
 

Issue 4: Incorrect mapping of embed_tokens in _UNQUANTISED_OPS
The _UNQUANTISED_OPS tuple maps ("embed_tokens", "lm_head"). However, embed_tokens is a distinct prologue operation, and lm_head is an epilogue operation. If embed_tokens is listed in modules_to_not_convert, it should map to the embed_tokens graph op to ensure its precision is correctly overridden, not to lm_head.

Suggested change:
Change the mapping for embed_tokens in _UNQUANTISED_OPS to map to itself.

--- a/gitm/planner/glm_graph.py
+++ b/gitm/planner/glm_graph.py
@@ -1008,7 +1008,7 @@
 #: them:
 _UNQUANTISED_OPS: tuple[tuple[str, str], ...] = (
     ("lm_head", "lm_head"),
-    ("embed_tokens", "lm_head"),  # untied, but priced together in the epilogue
+    ("embed_tokens", "embed_tokens"),
     ("eh_proj", "mtp_eh_proj"),
     ("indexer", "attn_index_proj"),
     ("indexers_proj", "attn_index_proj"),
     ("mlp.gate", "moe_router"),
 )

[github-actions]

🤖 Claude Code Review

Looking at this diff, it's entirely a documentation file (DESIGN-NOTE.md) — no executable code, no Python/CUDA/bash logic changes. Most standard code-review categories don't apply. Here's what's worth flagging:

---

🐛 Bugs

• §3.3 MTP node-count arithmetic is inconsistent. The text states 1,614 = 1,591 + 23 and 1,706 = 1,591 + 115, implying the backbone is 1,591 nodes. But §4.1 says "1,614 nodes" for vanilla decode including the draft. If the MTP module's single stage is 23 nodes and is part of vanilla decode, the backbone is 1,591 — but then the "vanilla decode (D=0)" row in §3.3 should label it 1,614, not use 1,591 as the backbone in the MTP step. The reconciliation note reads self-consistently only if "vanilla" includes the single draft stage, which should be stated explicitly rather than implied.

• Break-even α formula in §3.3 is applied to the wrong baseline. The table compares 1+Dα vs Σ αⁱ against a baseline of 1,933 tok/s, but that baseline is itself computed under the 1+Dα convention. The "honest" break-even α = 0.415 uses the corrected throughput formula for the numerator but the uncorrected formula baked into the denominator, which slightly understates the true break-even.

• §3.2 prefill "AI = 281 vs ridge 412" conclusion depends on EP being on, but §8.1 Q1 notes EP may not be enabled (it's absent from the vendor recipe). If TP8-only, the a2a disappears and the bank doubles — the overall-memory-bound conclusion and the 264 ms floor could both be wrong by a large factor. The note flags this in §8 but the §3.2 summary presents the EP8 number as the prefill result without a ⚑ marker, inconsistent with the stated convention.

---

⚡ Performance

• §3.2 chunking table is the most actionable number in the document but the prescriptive conclusion ("pick --max-num-batched-tokens as large as decode latency tolerates") doesn't account for HBM pressure from activations at large chunk sizes, which can reduce the effective batch for decode. Worth a qualifier.

---

📊 Reproducibility

• git rev-parse HEAD gate is informal. The preamble says to check HEAD matches before trusting any number, but the gitm plan commands have no --version or lockfile equivalent. If the planner is "actively changing," any generated number is silently stale the moment a commit lands. A gitm plan --dump-inputs or similar provenance flag would make the reproduce-any-figure claim actually reproducible.

• ep_imbalance = 1.0 (A8) propagates through every MoE byte estimate with no sensitivity column. Given that §5 rank 5 identifies imbalance as the primary lever on the 73.8% dominant node, a ±20% imbalance band on the key decode floor would materially strengthen the document's claims.

---

💡 Suggestions

• G11 is marked "no — not shipped" but affects DeepSeek-V4 identically. The note says the fix should land where both families can be checked together, but there's no tracking issue or TODO anchor. A reader acting on this document has no way to know if it ever landed. Add a GitHub issue reference or a # TODO(G11) comment in kernels/library.yaml.

• The document cuts off mid-sentence at the end ("Which instance a "). If this is intentional (work in progress), a <!-- TODO --> marker would prevent a future reader from assuming truncation is a rendering artifact.

• §6.4 row 0 says "not a measurement" but is numbered in the measurement triage table, which will confuse anyone working through the list mechanically. Either move it to a preamble or label it C6 (its own capture category) rather than 0.

• nsys command in §6.2 lacks --output and --duration flags, making the "reproduce any figure" goal harder than it needs to be for someone setting up the capture for the first time.

[github-actions]

Code Review by Gemini

The changes introduce support for the GLM-5.2 model, including its unique architecture features like DeepSeek Sparse Attention (DSA), IndexShare, and Multi-Token Prediction (MTP). The commit body and design note are exceptionally detailed, outlining the architectural nuances, performance predictions, and validation steps. The code implements these features, including per-op precision, specific prefill and MTP logic, and detailed node emission.

I've identified one bug related to precision overrides when loading from a raw config.json and a robustness improvement for weight calculation.

Bug: embed_tokens precision override not applied when loading from config.json

The _op_dtype_overrides function in glm_graph.py incorrectly maps the embed_tokens tensor name to the lm_head op name when determining precision overrides from modules_to_not_convert. As a result, the embed_tokens operation, which is explicitly named "embed_tokens" in the graph, will not receive its intended bf16 override (as specified in the design note for FP8 models) and will default to spec.act_dtype (which would be fp8 for FP8 models). This affects models loaded directly from config.json via spec_from_hf_config. Catalogue entries are not affected as they explicitly list op_dtype_overrides.

File: gitm/planner/glm_graph.py

--- a/gitm/planner/glm_graph.py
+++ b/gitm/planner/glm_graph.py
@@ -1090,7 +1090,7 @@ _UNQUANTISED_OPS: tuple[tuple[str, str], ...] = (
 #: 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", "lm_head"),  # untied, but priced together in the epilogue
+    ("embed_tokens", "embed_tokens"),
     ("eh_proj", "mtp_eh_proj"),
     ("indexer", "attn_index_proj"),
     ("indexers_proj", "attn_index_proj"),
     ("mlp.gate", "moe_router"),

Improvement: Robustness in model_weight_bytes for embed_tokens

The calculation for embed weights currently reuses lw (which is derived from lm_head's dtype). While embed_tokens and lm_head are expected to have the same dtype override (bf16 for FP8 models), it's more robust to explicitly derive the dtype for embed_tokens in its own weight calculation. This prevents potential errors if their dtypes were to diverge in future configurations.

File: gitm/planner/glm_graph.py

--- a/gitm/planner/glm_graph.py
+++ b/gitm/planner/glm_graph.py
@@ -260,7 +260,7 @@ def model_weight_bytes(
     # Untied input embedding and vocabulary projection. Both stay wide on the FP8
     # checkpoint (``embed_tokens`` and ``lm_head`` are in ``modules_to_not_convert``),
     # so they are priced at their own width rather than the backbone's — 1.9 GB of
     # the resident footprint that an fp8 read would halve on paper and not on disk.
-    embed = 2.0 * spec.vocab * h * lw / tp
+    embed = 2.0 * spec.vocab * h * weight_bytes(spec.dtype_for("embed_tokens", spec.weight_dtype)) / tp
 
     return (
         experts

Test Update: Verify embed_tokens precision override

To ensure the bug fix and the intended behavior for embed_tokens precision, the existing test test_fp8_checkpoint_reads_what_the_quantiser_skipped should be updated to include embed_tokens in its assertions.

File: tests/test_glm_graph.py

--- a/tests/test_glm_graph.py
+++ b/tests/test_glm_graph.py
@@ -100,6 +100,8 @@ def test_fp8_checkpoint_reads_what_the_quantiser_skipped():
     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"):
+    # Also the input embedding.
+    for op in ("lm_head", "mtp_eh_proj", "attn_index_proj", "embed_tokens"):
         assert spec.dtype_for(op, spec.weight_dtype) == "bf16", op
     assert spec.dtype_for("moe_router", spec.weight_dtype) == "fp32"
 

[nicholaslawrence-hub]

Pushed 4d4b394.

Gemini's _OP_RULES ordering bug isn't one. embed_tokens is line 158, lm_head is line 176 — it's already before. The review read the diff hunks rather than the file. Confirmed by behaviour (embedding_kernel, model.embed_tokens.gather, index_select_cuda, vocab_parallel_embedding all → embed_tokens; lm_head_gemm → lm_head), and test_every_needle_is_reachable exists to catch exactly that class of shadowing.

Bounding Q3 reversed my own call from last round. I promoted it on the "64× error" framing without computing it. The prefill attention core 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 doesn't move the prefill conclusion. Demoted back to last, with the number attached so it doesn't get promoted again on the multiplier alone. My mistake for acting on the magnitude without checking the share.

The chunking catch was the good one. The vendor recipe doesn't set --max-num-batched-tokens, so the 264.0 ms floor was silently assuming one 8,192-token chunk. Now stated, with a 4 × 2,048 row added (707 GB, 319 ms) and a pointer to C6 to confirm what the engine actually uses.

Also: struck the linear break-even row so it can't be skimmed as usable — it's kept only so the discrepancy is recognisable in CLI output, which now warns as well. Footnoted why α=0 reads the same in both columns (it's the cost of drafting for nothing, and it's below the MTP-off baseline — that's the point of the column). Gave S1 a forward reference to the §5.2 fork it resolves. Corrected the ep_imbalance direction: over-predicted traffic means the throughput figures are conservative, not optimistic. And noted in G10 that expected_stream_id defaults to 0, indistinguishable from an explicit compute-stream assignment — whoever wires the invariant should make it int | None first. Good catch, that one's a real latent trap.

Not doing: the pinned commit hash (third time — it's a hand-written note, the branch is the pin) and the Mermaid-portability suggestion (GitHub renders Mermaid natively).

🤖 Addressed by Claude Code

[github-actions]

🤖 Claude Code Review

Review

This is a design/analysis document (markdown), not executable code. The diff contains no Python, CUDA, CuPy, PyTorch, or bash code that runs — it's a working note describing a roofline model for GLM-5.2. Feedback is scoped accordingly.

---

🐛 Bugs

• Truncated document: The diff ends mid-sentence at "A residual is a le..." (bottom of §9). If this is the actual committed state, the file is broken and any tooling that parses it to generate CLI help or structured output will fail silently or crash.

• tokens_per_step formula mismatch (§3.3): The document correctly identifies that the codebase uses 1 + D·α but the right formula for prefix-chain speculative decoding is Σ αⁱ (i=0..D). This is documented as a known bug but left in place because it requires changing shared BatchConfig semantics. The struck-out row is shown in CLI output today, meaning users reading the CLI will see numbers overstated by up to 1.8× with no correction unless they read this document. The note says "the CLI itself now says so" — but if gitm plan --spec-tokens prints a warning inline, that should be verified; the document doesn't show the warning text.

• G10 field default is misleading: expected_stream_id defaults to 0, which is indistinguishable from an explicit compute-stream assignment. The document correctly flags this, but the field is already shipped in that ambiguous state. Any downstream consumer wiring the invariant will silently treat unset nodes as compute-stream, producing false negatives for overlap detection.

• IndexShare schedule off-by-one (historical, described in §1): The note describes a past bug where layer 77 took a modulo fallback and landed on shared, giving the right count from wrong evidence. The fix ("the loader now refuses a short schedule") should have a corresponding test assertion called out in §7 — it isn't mentioned there.

---

🔒 Security

• vllm serve command in §9 instructs readers to add --enable-expert-parallel based on this document's assumptions. If copy-pasted into a CI/CD pipeline or runbook verbatim, it will diverge from the vendor recipe (which deliberately omits that flag). This is noted in the text but easy to miss when the bash block is extracted.

---

⚡ Performance

• attn_q_a/attn_kv_a replication (rank 9, §5): The document notes these are paid in full on every TP rank. If the planner emits these as replicated but the actual engine has already handled this via DP-attention sharding, the prediction will overstate decode bytes for these nodes. The "if it's already sharded, the graph is wrong" note is correct but this should be a higher-priority validation item than rank 9 suggests — it affects 2.8% of decode at every batch size and the discrepancy direction is uncertain.

• mtp_eh_proj TP-sharding question (Q11): 151 MB vs 19 MB per rank per draft stage is a 8× difference, linear in D. At D=5 this is potentially 660 MB of misattributed traffic in the draft chain. It's listed as an open question but the draft cost table in §3.3 appears to assume the un-sharded case without flagging the uncertainty inline.

---

📊 Reproducibility

• git rev-parse HEAD instruction: The document says to check HEAD matches if a number disagrees, but the commands at the top use gitm plan which depends on the planner state. If the planner is "actively changing" (as stated), the document-to-code binding is fragile. A pinned planner version or a --version flag output should be included in the reproduce block.

• ep_imbalance = 1.0 (A8): Declared, not calibrated, and the document is explicit about this. However, §4.1's throughput figures (1,933 tok/s) flow directly from this assumption and are presented without an uncertainty range. Given the document acknowledges real skew reduces traffic (conservative direction for bytes but adds grouped-GEMM tail latency not modeled at all), the headline throughput number could be optimistic in latency-sensitive regimes even if conservative in bytes.

• Chunk size assumption: The 264 ms / 422 GB prefill figures assume a single 8,192-token chunk, but the document acknowledges this is unconfirmed. The table showing 1,543 ms at 64×128 chunking is correct to include, but the default-presented number should carry an explicit [assumed: unchunked] tag in the summary row.

---

💡 Suggestions

• G11 is unshipped but affects moe_routed (73.8% of decode): The intervention vocabulary gap means levers targeting the dominant cost node are currently scoped to dense FFN ops. This should be ranked higher than its placement in the gap table suggests — it's not a "nice to have" for a benchmarking repo, it's a correctness gap in the optimization recommendation path.

• S1 confidence labeled "low" but consequence is severe: A real D2H sync per expert histogram (76/token) would completely block CUDA-graph capture. The low confidence label is appropriate, but this should be the first thing checked in any trace session, not resolved by a side-note. Consider promoting it to the "30-minute version" in §6.4.

• No negative test for the is_glm_moe_dsa_config guard (§2.3): The document notes the planner "declines" vision variants. There's no mention of a test that passes a GLM-vision config and asserts rejection. Given the family names overlap, this is an easy regression path.

• The logits_all_gather in E2: Described as gathering [32, 19360] → [32, 154880] fp32. The intermediate 19360 = 154880 / 8 (vocab sharded across TP8) — this math checks out. But the node is emitted as fp32 output while lm_head is BF16. If the all-gather promotes to fp32 for numerical stability before sampling, that's correct, but it should be explicit in the precision table in §1 rather than only visible in the prologue/epilogue table.

[github-actions]

ERROR: # Code Review by Gemini
An unexpected error occurred in Client call: 503 UNAVAILABLE. {'error': {'code': 503, 'message': 'This model is currently experiencing high demand. Spikes in demand are usually temporary. Please try again later.', 'status': 'UNAVAILABLE'}}

[nicholaslawrence-hub]

Pushed e11ee7b.

embed_tokens precision — real, though not for the stated reason. The node itself was already bf16 (it runs at act_dtype), so nothing was mispriced. But embed_tokens mapped onto the lm_head op, so dtype_for("embed_tokens") returned fp8 on a checkpoint whose modules_to_not_convert explicitly lists it. Right answer by accident, wrong answer to the question. It's now its own op: the footprint prices the two untied tensors separately, the gather reads at the table's width, and the FP8 entry declares it. Footprints unchanged (755.9 / 1508.1 GB) since both halves are bf16 here — and there's a test pinning that quantising the embedding moves the footprint by the size of the table rather than by nothing.

The break-even baseline claim is wrong. 1,933 tok/s is measured at D=0, where tokens_per_step degenerates to batch — no α, no convention applied. Nothing to correct, but I've said so on the line so it doesn't come up a fourth time.

Node counts: relabelled rather than re-explained. Three reviews have now read 1,591 vs 1,614 as ambiguous, which means the note below the table wasn't doing the job. The row now says "vanilla decode (D=0), draft stage included" on the row itself.

The truncation is the diff view, not the file — fourth time. §9 ends "a residual is a lead, not a defect", Appendix A.1–A.5 follow it, and the file is 1,054 lines. Nothing parses this document to generate CLI help.

On the CLI warning being unverified: it's tested — test_plan_warns_that_a_speculative_rate_is_a_ceiling asserts both the warning text and the 1.78x optimistic figure, and that a non-speculative step stays silent.

gemini-review is red on a 503 from Google's API, not on anything here — the action posted the error itself ("This model is currently experiencing high demand"). lint and all three pytest matrix jobs pass.

🤖 Addressed by Claude Code

[github-actions]

Code Review by Gemini

The changes introduce support for the GLM-5.2 model, including its unique architecture features like DeepSeek Sparse Attention (DSA), IndexShare, per-op precision, and a MoE-based Multi-Token Prediction (MTP) draft. The accompanying DESIGN-NOTE.md is exceptionally thorough and provides excellent context for the implementation.

The code is well-structured, highly detailed, and addresses several complex architectural nuances. The tests are comprehensive and cover critical aspects and potential pitfalls.

Here are some minor suggestions and observations:

gitm/planner/glm_graph.py

1. model_weight_bytes - embed calculation

The embed term in model_weight_bytes is calculated as spec.vocab * h * (tw + lw) / tp. This implies that both embed_tokens (input embedding table) and lm_head (vocabulary projection) are sharded by tp.

While lm_head is explicitly stated as vocabulary-sharded in the design note (A.1), the sharding behavior of embed_tokens is not explicitly mentioned. In many frameworks, the input embedding table is replicated across TP ranks if the vocabulary is not sharded, or sharded if the vocabulary is sharded.

If embed_tokens is replicated (i.e., not sharded by tp), its weight contribution to model_weight_bytes should not be divided by tp.

The test test_embed_tokens_carries_its_own_declared_precision passes, but it effectively runs with tp=1 because model_weight_bytes is called without a sharding argument, making sh.tp default to 1. This means the test doesn't fully validate the tp division for embed_tokens.

If embed_tokens is indeed replicated, the line should be adjusted. If it's sharded (consistent with lm_head due to vocabulary sharding), then the current code is correct. Given the design note's detail, it's likely this is an intentional choice for sharding both. However, it's worth double-checking this assumption.

Suggestion:
If embed_tokens is replicated:

--- a/gitm/planner/glm_graph.py
+++ b/gitm/planner/glm_graph.py
@@ -290,7 +290,7 @@
     # 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
-
+    embed = spec.vocab * h * tw + spec.vocab * h * lw / tp
     return (
         experts
         + shared_exp

If embed_tokens is sharded, no change is needed.

gitm/optimizer/deviation.py

1. _OP_RULES - attn_kv_b comment clarity

The comment for attn_kv_b is very detailed and explains the conditional nature of this operation based on MLA absorption. While accurate, it's quite verbose for an inline code comment.

Suggestion (Minor):
Consider shortening the comment or moving some of the detailed explanation to the DESIGN-NOTE.md if it's not already there, keeping the code comment more concise.

--- a/gitm/optimizer/deviation.py
+++ b/gitm/optimizer/deviation.py
@@ -130,9 +130,8 @@
     # 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"),

The current comment is already quite good, so this is a very minor point.

gitm/planner/registry.py

1. _render_table - ridge_line formatting

The ridge_line string formatting uses f"{d} {r:.0f}". If r is a very large number (e.g., for FP4 peaks), .0f might not be ideal as it rounds to the nearest integer. While the current GLM-5.2 model doesn't use FP4, future models might.

Suggestion (Minor):
Consider using a more flexible format specifier like :.1f or :.0e for very large numbers, or adding logic to switch formats based on the magnitude of r. For the current context, .0f is fine as the FP8 ridge is 412.

--- a/gitm/planner/registry.py
+++ b/gitm/planner/registry.py
@@ -169,7 +169,7 @@
         for d, peak in sorted(dtypes.items())
         if peak > 0 and hw.peak_mem_bw_bytes_per_s
     }
-    ridge_line = " · ".join(f"{d} {r:.0f}" for d, r in ridges.items()) or "unpriced"
+    ridge_line = " · ".join(f"{d} {r:.0f}" for d, r in ridges.items()) or "unpriced" # Consider f"{d} {r:.1f}" or f"{d} {r:.0e}" for very large numbers
 
     out = [
         f"model     {getattr(spec, 'name', '?')}  [{family}]",

Overall, the changes are robust and well-executed, significantly enhancing the planner's capabilities for GLM-5.2 models. The level of detail in the design note and the code comments is commendable.

[nicholaslawrence-hub]

Pushed 2db5f21. The 1.483 ms is fully accounted — your 0.804 ms is exactly one of its terms.

moe_routed × 5 = 0.803 ms, which matches your calculation to a rounding error. The other 0.680 ms is lm_head (0.250), mtp_eh_proj (0.159) and the attention block (0.271) — all itemised per node in Appendix A.4, which has a full per-stage table.

But that's twice now someone has reconstructed the chain from the expert bank alone and found time missing, so A.4 clearly isn't being reached from §3.3. The split is now in §3.3 itself: 54 % expert bank, 17 % vocabulary projection, 11 % eh_proj, 18 % the rest.

That states the section's own point better than the prose did, incidentally — the draft isn't purely an expert-bank cost, and the two BF16 tensors it drags along (lm_head and the [12288→6144] fusion) are more than a quarter of a stage between them, neither of which gets cheaper for being a draft.

The attn_index_score figures both check out, as your own parentheticals note: 0.147/16.254 = 0.9 % and 18.795/34.902 = 53.9 %.

🤖 Addressed by Claude Code

[github-actions]

🤖 Claude Code Review

Review

This is a design/analysis document, not executable code. No Python, CUDA, CuPy, PyTorch, or bash logic is introduced—only prose, tables, Mermaid diagrams, and CLI invocations. The feedback below is scoped accordingly.

---

🐛 Bugs

• Break-even formula inconsistency. §3.3 states accepted tokens follow (1−α^(D+1))/(1−α) (correct prefix-chain formula), and the table header says (1−α⁶)/(1−α) with D=5, which is correct. However the break-even derivation Σ αⁱ = 28.135/16.254 = 1.731 equals ~1.73 tokens, implying α≈0.426. Verify: (1−0.426⁶)/(1−0.426) ≈ 1.731 ✓. But the table row for α=0.0 shows 1.000 accepted tokens/step—this should be exactly 1 by the formula (α=0 gives (1−0)/1 = 1), which is correct. No bug, but worth noting that the "MTP off" row at 1,969 tok/s uses 32 ÷ 16.254 ms—this assumes vanilla decode time, not MTP-off time, which conflates two scenarios if the engine has overhead from MTP scaffolding even when D=0.

• KV cache byte calculation. The formula 78 × (512×1.000244 + 64×2) + 21 × 128×1.000244 mixes fp8 (1.000244 B/weight) for the latent with bf16 (2 B) for rope keys, then applies fp8 overhead to index keys. If index keys are bf16 (as Q10 asks), the 21-layer term should be 21 × 128 × 2 = 5,376 B/token, not 21 × 128 × 1.000244 ≈ 2,563 B/token—a ~2× error on that term. The document acknowledges this as open question Q10 but the stated total 52,618 B/token uses the fp8 assumption without flagging it in the formula line itself.

• MTP node count. The doc states the draft chain is 5 × 23 = 115 nodes, with 23 = 20 (shared-indexer MoE layer) + rms_norm + mtp_eh_proj + lm_head. That's 20+3=23 ✓. But the verify pass is described as "the same 1,591 nodes" at 1+D=6 rows. If verify also runs lm_head at 6 rows, the epilogue node count should be identical to decode—confirm the 1,591 figure doesn't double-count the epilogue between verify and the backbone count.

---

🔒 Security

• The vllm serve command is reproduced verbatim including --enable-auto-tool-choice. If this document is used as a template for actual deployment scripts, --tool-call-parser glm47 is an unvalidated parser name—ensure the serving image actually ships this parser before copy-pasting.

---

⚡ Performance

• G11 is unshipped and affects 74% of decode time. The document explicitly notes moe_routed and moe_shared are absent from kernels/library.yaml's applies_to_kernels, meaning optimization interventions targeting the dominant term cannot be matched by tooling. This should be prioritized over presentational gaps.

• eh_proj replication (Q11). If eh_proj [12288→6144] is not TP-sharded, each of 5 draft stages pays 151 MB/rank. At 8 ranks that's 1.2 GB of redundant reads per step. The document notes this but marks it low-priority; given it's 11% of draft bytes and trivially shardable, it deserves a higher rank in §5.

---

📊 Reproducibility

• git rev-parse HEAD warning is correct but fragile. The note says "check HEAD matches if a number disagrees"—but the planner commands at the top (gitm plan ...) produce output that isn't pinned to a specific version in CI. If these commands are run against a different planner revision, silent numerical disagreement is the failure mode, not an error.

• ep_imbalance = 1.0 assumption (A8). Stated as conservative (over-states traffic), but coupon-collector distinct-expert counts are sensitive to batch size in a nonlinear way. The document correctly notes this but the sweep table in §4.1 (B=1→256) would benefit from showing the distinct-expert count alongside tok/s to make the nonlinearity visible.

---

💡 Suggestions

• G10's expected_stream_id defaulting to 0. The document correctly identifies that int | None should replace int before wiring the overlap invariant. This is a concrete API contract issue—whoever implements §6.1's classification rule will silently treat all pre-GLM nodes as "on the compute stream," making the architectural-vs-recoverable distinction unreliable for existing families, not just GLM.

• Q1 (EP on/off) gates rank 1, the largest headroom item, but C5 is listed as a lower-priority capture than C1/C2. Given that EP on/off determines whether 44% of prefill time exists at all, C5 should be elevated—arguably it should be resolved by C6 (launch args) before any profiling begins.

• The A4 truncation (~2 µs kernel launch (CUDA-graph replay) is cut off mid-row). This is likely an authoring artifact but leaves assumption A4 without its falsification condition, which breaks the parallel structure of the A-table.

• Chunk-size sensitivity (§3.2) is documented as 14.9× bytes but no mitigation path exists in §5. Rank 8 identifies it but only says "measure the quotient." A concrete recommendation (e.g., "set --max-num-batched-tokens ≥ 8192 unless decode latency SLA forbids it") would make the document actionable for operators.

[github-actions]

Code Review by Gemini

The changes introduce comprehensive support for the GLM-5.2 model, including its unique architecture features like MLA, DeepSeek Sparse Attention (DSA), and Multi-Token Prediction (MTP). The implementation is highly detailed, addressing nuances such as per-op precision, distinct prefill/decode behaviors for DSA, and the MTP prefix chain acceptance rate. The accompanying design note is exceptionally thorough, and the new test suite is robust, covering a wide array of architectural specifics and potential edge cases.

Here are the identified issues and suggested improvements:

---

File: docs/glm-5.2/DESIGN-NOTE.md

Issue 1: Ambiguous unit for parameters (B vs G)

The unit "B" is used for parameters, which typically stands for "Billion" in this context. However, given the scale of LLMs, "G" (Giga) is more commonly used for parameters (e.g., 7B, 70B, 700B parameters). Using "B" for 744 parameters implies 744 Billion, which is an extremely large number for parameters (more likely 744 Giga parameters). This can be confusing.

Suggested change:
Clarify the unit for parameters, likely changing "B" to "G" for Giga.

--- a/docs/glm-5.2/DESIGN-NOTE.md
+++ b/docs/glm-5.2/DESIGN-NOTE.md
@@ -60,7 +60,7 @@
 | 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                |
+| Total / active params       | **744 G** published + a **9.9 G** MTP block / ~39 G active                |
 
 ### Semantics read from the checkpoint, not guessed
 
@@ -109,9 +109,9 @@
 | 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**:
+| params, MTP block removed | **744.2 G** | **744 G** (Z.ai) | **+0.03 %** |
+| The third is the interesting one. The checkpoint is 753.3 G by its own bytes but
+| Z.ai publishes 744 G — the gap is the MTP block, which the published figure
+| excludes. So the block is **9.9 G**, where a *dense* draft head would be **0.23 G**:

---

File: gitm/planner/glm_graph.py

No critical issues found. The implementation is robust and aligns well with the detailed design note.

---

File: gitm/planner/model_catalogue.py

No critical issues found. The extends functionality and schedule length validation are valuable additions.

---

File: gitm/planner/models/glm-5.2-fp8.yaml

No issues found.

---

File: gitm/planner/models/glm-5.2.yaml

No issues found.

---

File: gitm/planner/registry.py

No critical issues found. The improvements to family detection and CLI output are well-implemented.

---

File: gitm/planner/roofline.py

No critical issues found. The fix for tokens_per_step is crucial for accurate MTP throughput calculations.

---

File: gitm/optimizer/deviation.py and tests/test_deviation_alignment.py

No issues found. The expansion of _OP_RULES and corresponding test updates are correct.

---

File: tests/test_glm_graph.py

No issues found. This is an exceptionally comprehensive test suite that significantly increases confidence in the correctness of the GLM-5.2 implementation.

---

[nicholaslawrence-hub]

Pushed 1e73a95. Three clarity fixes, no behaviour change.

Break-even cell now says what it crosses: α > 0.426, where this row crosses 1,969 — the MTP-off rate sitting one row below.

Verify node count. Verify is the backbone at 6× the rows and the same node count — more work per kernel, not more kernels. That's also why its bytes move (106.43 vs 67.34 GB) and its node count doesn't. Now on the row rather than inferable from it.

Embedding sharding. Fair — the note stated it for lm_head and never for embed_tokens, so the code looked unexplained. model_weight_bytes divides both by tp because vLLM builds the input table as VocabParallelEmbedding, split by vocabulary the same way the output projection is, and both sibling families (moe_graph.py:189, hybrid_graph.py:299) already price it that way. Now stated in §3.1. This is the third time it's come up, which was the note's fault rather than the reviewer's.

🤖 Addressed by Claude Code

[github-actions]

Code Review by Gemini

The code changes introduce support for the GLM-5.2 model, including its unique architecture features like DeepSeek Sparse Attention (DSA), IndexShare, Multi-Layer Attention (MLA), and a Mixture-of-Experts (MoE) draft head for Multi-Token Prediction (MTP). The changes are comprehensive, covering model specification, graph prediction, hardware context, catalogue entries, and extensive testing.

The design note is exceptionally detailed and clearly outlines the architectural assumptions, predicted execution model, identified planner gaps, and validation plan. This level of documentation significantly aids in understanding the rationale behind the code changes.

Overall, the implementation is robust and directly addresses the complexities of the GLM-5.2 model. The identified planner gaps (G1-G9) and bugs found during development are explicitly handled in the code, and the new test suite provides excellent coverage for these specific features and fixes.

No critical bugs or issues were found that were not already acknowledged and addressed in the commit message or design note.

---

gitm/planner/glm_graph.py

Improvement: Consistent use of spec.act_dtype for activation-related bytes.

In _emit_layer, the add_act_quant function uses aw (which is weight_bytes(spec.act_dtype)) for activation bytes. However, in the _pointwise helper, act_b is passed, which is also weight_bytes(spec.act_dtype). For consistency and clarity, add_act_quant could directly use aw from the outer scope, similar to how add_pointwise is defined. This is a minor stylistic point.

--- a/gitm/planner/glm_graph.py
+++ b/gitm/planner/glm_graph.py
@@ -500,7 +500,7 @@
         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)
+        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)

[nicholaslawrence-hub]

Pushed 3a1e191. One item had substance.

The fp8 constant in the KV formula is borrowed from the wrong side. weight_bytes("fp8") is 1.000244 because of the 128×128 block scales stored with fp8 weights; an fp8 KV cache carries a per-token or per-tensor scale instead. So 52,618 against a truer 52,608. I've named it in the formula block rather than changed it — 0.02 % moves nothing, and fixing it properly means a cache-specific byte constant no other family has and no measurement here could distinguish. Better as a known approximation than an unexamined one.

The MTP-off row is no longer the conflation you describe. That was true two revisions ago; since db81150 the graph emits no MTP block at D=0, so 16.254 ms is the MTP-off time — the two scenarios are identical by construction now, not assumed equal.

Not changing "B" to "G" for parameters. 744 B means 744 billion, which is the universal convention for LLM parameter counts (7B, 70B, 405B). "744 G parameters" isn't a form anyone uses.

The act_dtype suggestion is a no-op. add_act_quant already takes aw from the enclosing scope exactly as add_pointwise does; the suggested diff only adds a trailing space to a continuation line.

🤖 Addressed by Claude Code

[github-actions]

ERROR: # Code Review by Gemini
An unexpected error occurred in Client call: 503 UNAVAILABLE. {'error': {'code': 503, 'message': 'This model is currently experiencing high demand. Spikes in demand are usually temporary. Please try again later.', 'status': 'UNAVAILABLE'}}

[github-actions]

ERROR: # Code Review by Gemini
An unexpected error occurred in Client call: 503 UNAVAILABLE. {'error': {'code': 503, 'message': 'This model is currently experiencing high demand. Spikes in demand are usually temporary. Please try again later.', 'status': 'UNAVAILABLE'}}

[github-actions]

Code Review by Gemini

This pull request introduces comprehensive support for the GLM-5.2 model, including its unique architectural features like DeepSeek Sparse Attention (DSA), IndexShare, and Multi-Token Prediction (MTP). The changes are well-justified by the detailed DESIGN-NOTE.md and are thoroughly tested. The PR also improves the robustness of the CI/CD review workflows and refines the planner's core logic for better accuracy.

🐛 Bugs

• Incorrect MTP Acceptance Rate Calculation: The BatchConfig.tokens_per_step previously used a linear approximation (1 + D*alpha) for accepted tokens in MTP, which overstates throughput. The fix correctly implements the prefix chain expectation (sum(alpha**i for i in 0..D)).

  • File: gitm/planner/roofline.py
  • Lines to change:

    --- a/gitm/planner/roofline.py
    +++ b/gitm/planner/roofline.py
    @@ -602,22 +602,37 @@ class BatchConfig:
             # ``m`` tokens ramp up from ``ctx+1`` to the window; the rest run flat.
             m = min(p, window - ctx)
             return decode + m * ctx + m * (m + 1) / 2.0 + (p - m) * window
     
         @property
         def tokens_per_step(self) -> float:
             """Accepted output tokens per step — the denominator for per-token cost.
     
             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)
     class RooflinePrediction:
         op: str
         flops: float
         bytes: float
         t_compute_s: float
         t_memory_s: float
         t_pred_s: float

• Incorrect Family Detection Order: The detect_family function in registry.py needed to prioritize glm_moe_dsa over sparse_moe because both share structural features like index_topk and n_routed_experts. Without this, GLM-5.2 models could be misclassified.

  • File: gitm/planner/registry.py
  • Lines to change:

    --- a/gitm/planner/registry.py
    +++ b/gitm/planner/registry.py
    @@ -3,18 +3,20 @@ from __future__ import annotations
     import argparse
     import json
     from pathlib import Path
     from typing import Any
     
     from gitm.planner.graph import Graph
     from gitm.planner.roofline import BatchConfig, HardwareSpec, ShardingConfig
     
     
     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_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"

• Indexer Schedule Length Validation: The catalogue loader (model_catalogue.py) did not validate that per-layer schedules (indexer_types, mlp_layer_types) exactly matched n_layers. This could lead to silent misconfigurations where missing entries fell back to a modulo rule, potentially yielding plausible but incorrect results.

  • File: gitm/planner/model_catalogue.py
  • Lines to change:

    --- a/gitm/planner/model_catalogue.py
    +++ b/gitm/planner/model_catalogue.py
    @@ -95,6 +95,16 @@ def load_spec(name_or_path: str | Path):
     
     
         # 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
         # from the file that caused it.
         if isinstance(raw.get("dense_layers"), list):
             raw["dense_layers"] = frozenset(int(i) for i in raw["dense_layers"])
         if isinstance(raw.get("op_dtype_overrides"), list):
             raw["op_dtype_overrides"] = tuple(
                 (str(op), str(dt)) for op, dt in raw["op_dtype_overrides"]
             )

⚡ Performance

• Prefill Attention Core Bytes Calculation: The core_read_entries function in glm_graph.py correctly distinguishes between decode (top-k bounded bytes) and prefill (whole cache streamed bytes) for the attention core. This fixes a potential underestimation of prefill traffic for long contexts if a dense model's prefill logic were naively applied.
  • File: gitm/planner/glm_graph.py
  • Lines to change:

    --- /dev/null
    +++ b/gitm/planner/glm_graph.py
    @@ -149,6 +149,20 @@ def core_qk_pairs(spec: GlmMoeDsaModelSpec, batch: BatchConfig) -> float:
         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

💡 Suggestions

• Unmodeled Stream Overlap (G10): The expected_stream_id field is set on collective nodes, but as noted in the DESIGN-NOTE.md, nothing currently reads it to model stream overlap. This is a known architectural limitation of the planner.

  • File: gitm/planner/glm_graph.py
  • Lines to consider:

    # In _emit_collective function
    # 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,

  • Suggestion: This is acknowledged as a roadmap item. Continue to track this as a critical improvement for accurate collective modeling.

• Incomplete Op Vocabulary for Expert Traffic (G11): The library.yaml's op vocabulary currently lacks moe_routed and moe_shared, preventing tooling from accurately targeting expert traffic. This is a pre-existing issue shared with other MoE families.

  • File: gitm/optimizer/deviation.py
  • Lines to consider:

    # In _OP_RULES dictionary
    # "moe_routed": ("moe", "expert", "grouped_gemm", "group_gemm", "groupedgemm"),
    # "moe_shared": ("shared_expert", "moe_shared"),

  • Suggestion: As noted in the DESIGN-NOTE.md, this fix should be implemented where both MoE families can be checked simultaneously.

• Improved CI/CD Review Workflow Robustness: The changes to claude-review.yml and gemini-pr-review.yml significantly improve the robustness of the automated code review process by adding retry logic for API calls and allowing the workflow to continue on error for external service outages.

  • Files: .github/workflows/claude-review.yml, .github/workflows/gemini-pr-review.yml
  • Suggestion: This is a good improvement. Ensure the ANTHROPIC_API_KEY and GEMINI_API_KEY secrets are properly managed and have appropriate permissions.

• Detailed Design Note: The addition of docs/glm-5.2/DESIGN-NOTE.md is an excellent practice. It provides invaluable context, justification, and validation for the model's implementation in the planner.

  • File: docs/glm-5.2/DESIGN-NOTE.md
  • Suggestion: Continue to maintain and update this document as the model or planner's understanding evolves.

• Per-Op Precision Modeling: The introduction of op_dtype_overrides and the logic to parse modules_to_not_convert and moe_router_dtype from config.json is a crucial improvement for accurately modeling mixed-precision models like GLM-5.2-FP8.

  • File: gitm/planner/glm_graph.py
  • Suggestion: This is a strong feature. Ensure that the _UNQUANTISED_OPS mapping in glm_graph.py remains comprehensive and correctly ordered to avoid misclassifying tensor names.

• Comprehensive Test Suite: The new tests/test_glm_graph.py file provides extensive coverage for the GLM-5.2 model's specific architectural details and potential edge cases.

  • File: tests/test_glm_graph.py
  • Suggestion: Continue to expand this test suite as new insights or model variations emerge.

[github-actions]

Code Review by Gemini

The changes introduce support for the GLM-5.2 model, including its unique architecture (MLA + DeepSeek Sparse Attention, IndexShare, MoE, MTP), mixed-precision handling, and specific prefill/decode characteristics. The pull request is exceptionally well-documented, with a detailed design note and comprehensive test coverage that directly validates the architectural claims and bug fixes.

Here's a breakdown of the review:

🐛 Bugs

1. gitm/planner/model_catalogue.py - Schedule Length Validation:
   The commit body mentions: "indexer_types in the catalogue had 77 entries for 78 layers. Layer 77 took the modulo fallback, landed on the right answer, and the floor was byte-identical — so nothing failed. The loader validates schedule length now." This bug is correctly addressed by adding validation for indexer_types and mlp_layer_types length.

   --- a/gitm/planner/model_catalogue.py
   +++ b/gitm/planner/model_catalogue.py
   @@ -109,6 +130,13 @@ def load_spec(name_or_path: str | Path):
        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"
                )

   Suggestion: This fix is correctly implemented.

2. gitm/planner/roofline.py - MTP Acceptance Rate Calculation:
   The commit body states: "BatchConfig.tokens_per_step counted 1 + D·α, the independent-draws answer, where a verifier accepts a prefix. Overstated throughput 1.8× at D=5, α=0.5 and put break-even at less than a third of its real value. Shared with every family and wrong for all of them." This bug is correctly fixed by changing the calculation to a prefix chain sum.

   --- a/gitm/planner/roofline.py
   +++ b/gitm/planner/roofline.py
   @@ -619,7 +619,9 @@ class BatchConfig:
        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))

   Suggestion: This fix is correctly implemented and critical for accurate MTP throughput prediction.

⚡ Performance

1. gitm/planner/registry.py - _render_table Launch-Bound Nodes:
   The commit body mentions: "gitm plan dropping the launch bound and mispricing the ridge... So 854 launch-bound nodes printed as memory-bound, against ridge 206 where fp8 answers to 412." The _render_table function is updated to correctly count and display launch-bound nodes.

   --- a/gitm/planner/registry.py
   +++ b/gitm/planner/registry.py
   @@ -228,10 +256,12 @@ def _render_table(g, hw: HardwareSpec, spec, family: str, note: str) -> str:
        ),
   -    f"  {len(g.nodes)} nodes, {n_compute} compute-bound",
   +    f"  {len(g.nodes)} nodes, {n_compute} compute-bound, "
   +    f"{n_launch} launch-bound",
    ]

   Suggestion: This is a good improvement for accurately reflecting the performance bottlenecks, especially at low batch sizes.

💡 Suggestions

1. gitm/optimizer/deviation.py - _OP_RULES Comment Clarity:
   The comment for rms_norm in _OP_RULES states: "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..." This is good. However, the glm_graph.py _op_dtype_overrides function has a comment for _UNQUANTISED_OPS that says: "Norms and biases are omitted deliberately — they are not nodes in this graph". While _UNQUANTISED_OPS is for projections/GEMMs and rms_norm uses act_dtype, the conflicting phrasing could be confusing.

   --- a/gitm/planner/glm_graph.py
   +++ b/gitm/planner/glm_graph.py
   @@ -1070,8 +1070,9 @@ _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.
   +    # The indexer's *projections*, named specifically. Norms and biases are
   +    # omitted deliberately — they are not nodes in this graph, and their
   +    # precision is determined by `act_dtype` or `moe_router_dtype`. 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"),

   Suggestion: Update the comment in glm_graph.py to clarify that while norms are nodes in the graph, _UNQUANTISED_OPS specifically lists projections/GEMMs whose precision might be overridden, and norms' precision is handled by act_dtype or moe_router_dtype. This would resolve the perceived inconsistency.

2. gitm/planner/glm_graph.py - _emit_collective Stream ID:
   The comment for expected_stream_id=1 states: "A declaration, not a check: nothing reads this field today...". This is a known limitation (G10 in the design note).

   --- a/gitm/planner/glm_graph.py
   +++ b/gitm/planner/glm_graph.py
   @@ -709,6 +709,9 @@ def _emit_collective(
                ),
                # 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,
            )
        )

   Suggestion: This is correctly noted as a future improvement. No immediate code change is needed, but it's good to keep this in mind for future work on stream scheduling and overlap analysis.

Overall Assessment

This is an exceptionally well-executed and documented pull request. The commit message and the DESIGN-NOTE.md provide an unparalleled level of detail, explaining the architectural nuances of GLM-5.2, the challenges encountered, and the rationale behind every significant code change. The extensive test suite in tests/test_glm_graph.py is particularly commendable, as it rigorously validates the complex logic and ensures the model's behavior aligns with the design. The extends feature in the catalogue is a great addition for maintainability.

The identified bugs are correctly addressed, and the performance improvements are well-justified. The suggestions are minor and mostly related to documentation clarity or known future work. This PR sets a high standard for model integration and planning.

[github-actions]

Code Review by Gemini

The changes in this pull request introduce support for the GLM-5.2 model family, including its unique architecture (MLA + DSA, IndexShare, MoE with specific precision overrides), prefill behavior, and multi-token prediction (MTP) chain. The commit message is exceptionally detailed, outlining the design rationale, bugs found and fixed, and open questions. The code reflects a thorough implementation of these architectural details and addresses several subtle issues identified during development.

Here's a breakdown of the review:

🐛 Bugs

The commit body explicitly lists seven bugs found and fixed during development. The code changes directly address all of them:

1. indexer_types in the catalogue had 77 entries for 78 layers.

   • Fix: The load_spec function in gitm/planner/model_catalogue.py now validates that per-layer schedules (indexer_types, mlp_layer_types) exactly match n_layers. spec_from_hf_config in gitm/planner/glm_graph.py also performs this validation.
   • Relevant lines:

     --- a/gitm/planner/model_catalogue.py
     +++ b/gitm/planner/model_catalogue.py
     @@ -109,6 +130,13 @@ def load_spec(name_or_path: str | Path):
          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"
                  )

2. attn_index_score had lost its 32-head factor in the prefill rewrite.

   • Fix: The attn_index_score node's FLOPs calculation in _emit_layer now correctly multiplies by spec.index_n_heads.
   • Relevant lines:

     --- a/gitm/planner/glm_graph.py
     +++ b/gitm/planner/glm_graph.py
     @@ -400,7 +400,7 @@
          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_head_dim,
     +        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.

3. model_weight_bytes counted an indexer for the MTP block, which carries none.

   • Fix: The n_full_idx calculation in model_weight_bytes now correctly excludes the MTP block's indexer if index_share_for_mtp_iteration is true.
   • Relevant lines:

     --- a/gitm/planner/glm_graph.py
     +++ b/gitm/planner/glm_graph.py
     @@ -179,7 +179,7 @@
      
          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.
     -    n_full_idx = spec.n_full_indexer_layers + spec.num_nextn_predict_layers
     +    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

4. embed_tokens mapped onto the lm_head op, so dtype_for("embed_tokens") answered fp8 on a checkpoint that explicitly doesn't convert it.

   • Fix: embed_tokens is now a distinct op in gitm/optimizer/deviation.py (_OP_RULES) and gitm/planner/glm_graph.py (_UNQUANTISED_OPS), allowing it to have its own precision override.
   • Relevant lines:

     --- a/gitm/optimizer/deviation.py
     +++ b/gitm/optimizer/deviation.py
     @@ -145,6 +145,9 @@
      # 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"),

     --- a/gitm/planner/glm_graph.py
     +++ b/gitm/planner/glm_graph.py
     @@ -1090,6 +1090,7 @@
      # 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

5. The indexer needle also matched indexer.k_norm, so a norm in the skip list marked the whole indexer bf16.

   • Fix: The _OP_RULES in gitm/optimizer/deviation.py are reordered, placing specific indexer projection rules (attn_index_proj) before the more general attn_index_score rule. _UNQUANTISED_OPS in glm_graph.py also explicitly names indexer projections.
   • Relevant lines:

     --- a/gitm/optimizer/deviation.py
     +++ b/gitm/optimizer/deviation.py
     @@ -87,6 +87,9 @@

   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"),
   ```

6. The draft head ran at D=0, charging a pure decode step 0.3 ms of drafting a server without --speculative-config never does.

   • Fix: The MTP draft chain in predict_glm_graph is now conditionally emitted only if spec.num_nextn_predict_layers > 0, batch.speculative_tokens > 0, and batch.positions_per_step > 0.
   • Relevant lines:

     --- a/gitm/planner/glm_graph.py
     +++ b/gitm/planner/glm_graph.py
     @@ -900,7 +900,7 @@
      # 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
      ):

7. BatchConfig.tokens_per_step counted 1 + D·α, the independent-draws answer, where a verifier accepts a prefix.

   • Fix: The tokens_per_step property in gitm/planner/roofline.py now correctly calculates accepted tokens using a prefix chain sum (sum(alpha**i for i in 0..D)).
   • Relevant lines:

     --- a/gitm/planner/roofline.py
     +++ b/gitm/planner/roofline.py
     @@ -618,7 +618,9 @@
      # on. Reporting ``batch / total`` there would price D drafts and then
      # credit none of them.
      """
     -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)

🔒 Security

No security-related issues were identified in the changes. The modifications primarily involve model graph prediction logic and CI workflow robustness.

⚡ Performance

The core of this pull request is about accurate performance modeling. Several changes directly improve the fidelity of performance predictions:

• Per-op precision (op_dtype_overrides): This allows for accurate pricing of operations that run at different precisions (e.g., FP8 GEMMs, BF16 indexer, FP32 router), which is critical for GLM-5.2-FP8.
• DSA-specific prefill asymptotics: Separate helpers (core_qk_pairs, core_read_entries, index_scan_pairs, index_scan_entries) correctly model FLOPs and bytes for DeepSeek Sparse Attention, avoiding underestimation of long-context prefill traffic.
• Detailed layer lowering: Emitting all kernels (norms, activation quant, fused gating, prologue/epilogue) instead of just GEMMs provides a more accurate launch-bound analysis, especially at low batch sizes.
• MTP chain modeling: Correctly modeling the D-stage MTP chain with its own vocabulary projections and expert bank access provides a more accurate cost for speculative decoding.
• CI workflow improvements: Reordering git diff to prioritize code files over documentation for truncation ensures that reviewers see relevant code changes first, improving review efficiency. The retry logic for API calls makes the CI more robust against transient external failures.

📊 Reproducibility

• Comprehensive DESIGN-NOTE.md: The new design note provides an exhaustive explanation of the model's architecture, assumptions, predicted execution, and validation plan, significantly enhancing the reproducibility of the analysis.
• Catalogue entries (glm-5.2.yaml, glm-5.2-fp8.yaml): These entries, derived directly from config.json and model.safetensors.index.json, provide a canonical, trace-free definition of the model's shape and properties. The provenance sections detail verification and estimation points.
• extends mechanism: The extends feature in model_catalogue.py allows glm-5.2-fp8.yaml to inherit most of its definition from glm-5.2.yaml, reducing duplication and ensuring consistency between related model variants.

💡 Suggestions

The code is exceptionally well-documented, and the commit message anticipates most potential questions and concerns. The new test suite (tests/test_glm_graph.py) is comprehensive and directly validates the architectural claims and bug fixes.

One minor observation:

• G10 (Stream assignment for collectives): The commit body explicitly states that expected_stream_id is set on collectives but not yet read by monitor.py. This is a known limitation and is clearly documented in the code and the design note. No action is required for this review, as it's a deliberate, acknowledged, and roadmapped item.

Overall, this is a high-quality and impactful pull request that significantly extends the capabilities of the gitm planner.

[github-actions]

ERROR: # Code Review by Gemini
An unexpected error occurred in Client call: Server disconnected without sending a response.

[github-actions]

Code Review by Gemini

The changes introduce support for the GLM-5.2 model, including its unique architecture (MLA+DSA, IndexShare, MoE with specific precision overrides), prefill logic, and multi-token prediction (MTP) chain. The pull request also includes significant improvements to the gitm plan tool, such as per-op precision, more granular node emission, and corrected speculative decoding throughput calculation. The accompanying design note is exceptionally thorough and provides excellent context for all changes.

Here's a review of the changes:

---

🐛 Bugs

The commit body explicitly lists several bugs found and fixed during development, and the code changes reflect these fixes.

• gitm/planner/glm_graph.py
  • Logic Error: The model_weight_bytes function previously counted an indexer for the MTP block, even though index_share_for_mtp_iteration implies it reuses the main model's selection and carries no indexer tensors. This is now correctly conditional.

    --- a/gitm/planner/glm_graph.py
    +++ b/gitm/planner/glm_graph.py
    @@ -240,9 +240,7 @@
     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_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

  • Wrong Assumption: The embed_tokens op was previously mapped onto the lm_head op, leading to incorrect precision attribution if they differed. It now has its own distinct op name and precision lookup.

    --- a/gitm/planner/glm_graph.py
    +++ b/gitm/planner/glm_graph.py
    @@ -709,7 +709,7 @@
     # 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(
    +        roofline( 
                 "embed_tokens", 0.0,
                 rows * spec.hidden * (weight_bytes(embed_dtype) + aw), hw,
                 embed_dtype, serial_launches=1,

  • Wrong Assumption: The draft head was incorrectly running at D=0 (pure decode step) in the graph, charging for work that doesn't occur without a speculative config. The MTP draft chain is now correctly gated by batch.speculative_tokens > 0.

    --- a/gitm/planner/glm_graph.py
    +++ b/gitm/planner/glm_graph.py
    @@ -770,7 +770,7 @@
     # 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

• gitm/planner/model_catalogue.py
  • Off-by-one: The catalogue loader previously did not validate schedule lengths, allowing indexer_types (and mlp_layer_types) to be one entry short. This is now explicitly checked.

    --- a/gitm/planner/model_catalogue.py
    +++ b/gitm/planner/model_catalogue.py
    @@ -109,6 +109,14 @@
     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

• gitm/planner/roofline.py
  • Wrong Assumption: BatchConfig.tokens_per_step incorrectly calculated accepted tokens for speculative decoding as 1 + D*alpha (independent draws) instead of a prefix chain sum. This overstated throughput.

    --- a/gitm/planner/roofline.py
    +++ b/gitm/planner/roofline.py
    @@ -619,7 +619,9 @@
     # on. Reporting ``batch / total`` there would price D drafts and then
     # credit none of them.
     """
    -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))

• gitm/optimizer/deviation.py
  • Logic Error: The indexer needle in _OP_RULES incorrectly matched indexer.k_norm, causing a norm to mark the whole indexer as BF16. The rules are reordered and attn_index_proj is added to correctly classify indexer projections.

    --- a/gitm/optimizer/deviation.py
    +++ b/gitm/optimizer/deviation.py
    @@ -87,6 +87,8 @@

  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"
  ```
  

⚡ Performance

• gitm/planner/glm_graph.py
  • Missed Parallelism: The _emit_collective function now correctly emits two collectives per layer (post-attention and post-FFN all-reduces) instead of folding them into one. This is crucial for accurate latency-bound analysis, as collective count directly impacts cost when latency-bound.

    --- a/gitm/planner/glm_graph.py
    +++ b/gitm/planner/glm_graph.py
    @@ -590,7 +590,7 @@
     )
     add("attn_out_proj", f, b, wd)
     
    -    _emit_collective(g, spec, hw, layer, "tp_all_reduce_attn", rows, sh, prefix)
    +    _emit_collective(g, spec, hw, layer, "tp_all_reduce_attn", rows, sh, prefix, expected_stream_id=1)
     add_rms_norm(with_residual=True)
     
     # ── FFN: dense on the leading layers, mixture on the rest ────────────────
    @@ -610,7 +610,7 @@
     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)
    +        _emit_collective(g, spec, hw, layer, "tp_all_reduce_mlp", rows, sh, prefix, expected_stream_id=1)
     return

    Note: The expected_stream_id=1 was already present in the _emit_collective function, so adding it to the call site is redundant if the function itself sets it. However, the diff shows it being added to the call site, which is fine as it doesn't change behavior. The _emit_collective function itself was updated to set expected_stream_id=1 for all collectives.
  • Unnecessary CPU↔GPU transfers: The _emit_layer function now emits more granular nodes for pointwise operations (norms, activation quant, fused gating). This allows for a more accurate representation of launch-bound work, especially at low batch sizes, where these small kernels contribute significantly to overall latency.

    --- a/gitm/planner/glm_graph.py
    +++ b/gitm/planner/glm_graph.py
    @@ -340,7 +340,7 @@
     # 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)
    +        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:
    @@ -353,7 +353,7 @@
     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,
    +        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 ──────────────────

    This change is reflected throughout the _emit_layer function where add_rms_norm and add_act_quant are called.

💡 Suggestions

• gitm/planner/glm_graph.py
  • Clarity/Consistency: In _emit_collective, the expected_stream_id=1 is set within the function. However, the diff shows it being explicitly passed in the _emit_layer calls. While this doesn't cause an error, it's redundant. It would be cleaner to rely solely on the _emit_collective function to set this property.

    --- a/gitm/planner/glm_graph.py
    +++ b/gitm/planner/glm_graph.py
    @@ -590,7 +590,7 @@
     )
     add("attn_out_proj", f, b, wd)
     
    -    _emit_collective(g, spec, hw, layer, "tp_all_reduce_attn", rows, sh, prefix)
    +    _emit_collective(g, spec, hw, layer, "tp_all_reduce_attn", rows, sh, prefix, expected_stream_id=1)
     add_rms_norm(with_residual=True)
     
     # ── FFN: dense on the leading layers, mixture on the rest ────────────────
    @@ -610,7 +610,7 @@
     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)
    +        _emit_collective(g, spec, hw, layer, "tp_all_reduce_mlp", rows, sh, prefix, expected_stream_id=1)
     return

    Recommendation: Remove expected_stream_id=1 from the call sites in _emit_layer and rely on _emit_collective to set it.
• gitm/planner/model_catalogue.py
  • Robustness: The _expand_layer_types function validates the length of expanded against n_layers. A similar check could be added for op_dtype_overrides to ensure that if an override is specified for a layer-specific op, the layer index is valid. While op_dtype_overrides is a tuple of (str, str) and doesn't directly contain layer indices, if the str part were to include a layer index (e.g., "layer_0.attn_q_a"), this could be a future point of failure. For now, it's fine as is, but worth noting for future extensions.
• gitm/planner/registry.py
  • Error Message Clarity: The NotImplementedError message in predict_for_config is updated to include glm_moe_dsa. This is a good improvement for user feedback.

    --- a/gitm/planner/registry.py
    +++ b/gitm/planner/registry.py
    @@ -66,8 +66,9 @@
     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, 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."
  + "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."
  )
  ```

---

[github-actions]

ERROR: # Code Review by Gemini
An unexpected error occurred in Client call: 503 UNAVAILABLE. {'error': {'code': 503, 'message': 'This model is currently experiencing high demand. Spikes in demand are usually temporary. Please try again later.', 'status': 'UNAVAILABLE'}}

[github-actions]

Code Review by Gemini

🤖 Code Review

This is a substantial and well-documented change, addressing a new model architecture (glm_moe_dsa) with intricate details like per-op precision, DSA prefill asymptotics, and MoE draft chains. The accompanying DESIGN-NOTE.md is excellent and provides invaluable context for understanding the rationale behind the code changes. The fixes to the CI workflows and the BatchConfig.tokens_per_step calculation are also significant improvements.

🐛 Bugs

• gitm/planner/glm_graph.py - _op_dtype_overrides logic for moe_router
  The _op_dtype_overrides function correctly identifies moe_router_dtype from the config. However, the _UNQUANTISED_OPS tuple also contains ("mlp.gate", "moe_router"). If mlp.gate is in modules_to_not_convert (which it is in GLM-5.2-FP8's modules_to_not_convert list in the design note, though not explicitly in the glm-5.2-fp8.yaml overrides), and moe_router_dtype is float32, there's a potential for the mlp.gate entry to override the float32 setting if model_dtype is bf16 or fp16. The found.setdefault(op, model_dtype) ensures the first assignment wins, which is good, but the order of processing skipped items vs. moe_router_dtype matters. Currently, moe_router_dtype is processed after skipped items. If mlp.gate is in skipped and model_dtype is bf16, moe_router would be set to bf16 by setdefault, and then the fp32 from moe_router_dtype would overwrite it. This is fine, but it relies on the setdefault behavior and the order. It's more robust to ensure moe_router_dtype is always the authoritative source for the router's precision.

  --- a/gitm/planner/glm_graph.py
  +++ b/gitm/planner/glm_graph.py
  @@ -600,10 +600,10 @@ def _op_dtype_overrides(
               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"
  +    found.setdefault("moe_router", "fp32") # Use setdefault to avoid overwriting if already set by modules_to_not_convert
  
   return tuple(sorted(found.items()))

  Self-correction: The current code found["moe_router"] = "fp32" will overwrite any previous setdefault call for moe_router. This means moe_router_dtype is indeed authoritative, which is the desired behavior as per the design note ("fp32 on every variant, being a field of the base config rather than of any quantisation config"). So, the current implementation is correct for the stated intent. No change needed here.

• gitm/planner/glm_graph.py - _emit_collective serial_launches logic
  The _emit_collective function sets serial_launches=1 if priced else 0. While the intent is to avoid pricing collectives when interconnect bandwidth is zero, setting serial_launches=0 means the collective will contribute 0 to the t_pred_s even if t_compute_s or t_memory_s are non-zero (which they won't be for collectives, as FLOPs are 0 and bytes are handled by t_memory_s from link). However, serial_launches=0 also means it won't contribute to the launch-bound count, which might be misleading if the collective would incur a launch overhead if it were priced. The Graph.has_unpriced_collectives flag correctly indicates this, but the serial_launches value itself might be slightly inconsistent with the general definition of a node always having at least one launch.

  This is a minor point, as the has_unpriced_collectives flag is the primary signal. The current approach is consistent with the goal of not misrepresenting the cost.

⚡ Performance

• gitm/planner/glm_graph.py - _emit_layer add_act_quant op parameter
  The add_act_quant function takes op as a parameter, but then uses op directly in the add call. This means the op_dtype_overrides check for act_quant would need to be on the generic act_quant name, not on a specific instance like attn_act_quant. This is consistent with the rms_norm approach, but it's worth noting that if there were a need for different act_quant dtypes (e.g., one bf16, one fp8), the current system wouldn't support it. Given the current model, this is fine, but it's a design constraint.

  def add_act_quant(op: str, elems: float, gemm_op: str) -> None:
      # ...
      add(op, 2.0 * rows * elems,
          rows * elems * (aw + 1.0) + rows * 4.0, spec.act_dtype)

  The op parameter to add_act_quant is actually the name of the operation to be added to the graph (e.g., "act_quant"). The gemm_op is used to check the dtype of the consuming GEMM. So, the op_dtype_overrides would apply to "act_quant". This is consistent. No change needed.

💡 Suggestions

• gitm/planner/glm_graph.py - _emit_layer add_rms_norm with_residual parameter
  The add_rms_norm function takes with_residual: bool to adjust the bytes. For the first layer (layer == 0), with_residual is False. For subsequent layers, it's True. This is correct. However, the ops parameter for _pointwise is hardcoded to 3.0. If the residual add itself contributes FLOPs, this might be slightly underestimated. Typically, a fused add-norm would count the add's FLOPs. This is a minor detail, as norms are usually launch-bound.

  --- a/gitm/planner/glm_graph.py
  +++ b/gitm/planner/glm_graph.py
  @@ -240,7 +240,7 @@
       """
       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)
  +    add("rms_norm", f_p, b_p, spec.act_dtype) # Consider if ops should be 3.0 + (1.0 if with_residual else 0.0)
  
   def add_act_quant(op: str, elems: float, gemm_op: str) -> None:
       """Dynamic FP8 activation scaling ahead of a quantised GEMM.

  Self-correction: The ops=3.0 for RMSNorm is a common approximation for the arithmetic involved (sum of squares, sqrt, divide, multiply). An additional +1.0 for the residual add is plausible, but given norms are launch-bound, it's unlikely to change the overall bound or predicted time significantly. The current value is reasonable.

• gitm/planner/glm_graph.py - _emit_layer attn_kv_b comment
  The comment for attn_kv_b mentions "unabsorbed" MLA. It would be helpful to explicitly state what "absorbed" MLA would look like in terms of node changes (e.g., "An engine that absorbs MLA drops this node and doubles attn_out_proj's input width instead"). The design note already covers this, but a concise summary in the code comment would improve readability for someone only looking at the graph implementation.

  --- a/gitm/planner/glm_graph.py
  +++ b/gitm/planner/glm_graph.py
  @@ -300,7 +300,8 @@
   # 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.
  +# width instead (e.g., `attn_out_proj` input becomes `n_heads * (v_head_dim + kv_lora_rank)`).
  +# This is 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,

• gitm/planner/glm_graph.py - _emit_layer moe_router serial_launches
  The moe_router nodes (both the GEMM and the fused gating) are currently emitted with serial_launches=1. The design note mentions that the fused gating kernel is the "only data-dependent shape in the graph" and "Blocks CUDA-graph capture." If it blocks graph capture, it might imply a host synchronization or a series of launches that are not fully parallelizable. While serial_launches=1 is the default, it might be worth considering if this specific node (the fused gating) should have a higher serial_launches count or a specific estimated=True flag if its behavior is highly variable or involves host interaction. However, the design note also states "76 D2H per decoded token — fatal for graph capture" as a hypothesis (S1), implying it's not yet confirmed. So, serial_launches=1 is a reasonable default until empirical data suggests otherwise.

• gitm/planner/glm_graph.py - predict_glm_graph embed_tokens serial_launches
  The embed_tokens node is emitted with serial_launches=1. This is generally fine for a single kernel launch. However, if the embedding table is very large and not fully resident, or if the gather operation itself is complex, it might be more than a single launch. Given it's a gather, serial_launches=1 is a reasonable default.

• gitm/planner/glm_graph.py - _UNQUANTISED_OPS mlp.gate entry
  The _UNQUANTISED_OPS tuple includes ("mlp.gate", "moe_router"). This is to ensure that if mlp.gate is in modules_to_not_convert, the moe_router op is marked as model_dtype. This is correct. However, the moe_router op itself is explicitly set to fp32 later in _op_dtype_overrides based on moe_router_dtype. This means the mlp.gate entry in _UNQUANTISED_OPS is effectively redundant for moe_router's final dtype, as the moe_router_dtype will always override it. It doesn't cause a bug, but it could be removed for clarity if moe_router_dtype is always the source of truth for the router's precision.

  --- a/gitm/planner/glm_graph.py
  +++ b/gitm/planner/glm_graph.py
  @@ -579,7 +579,7 @@
   # 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"),
  +# ("mlp.gate", "moe_router"), # Redundant as moe_router_dtype is authoritative
   )
  
  

  Self-correction: The design note states "moe_router_dtype: "float32" makes the router fp32 on every variant, being a field of the base config rather than of any quantisation config." This implies moe_router_dtype is the definitive source for the router's precision, overriding any modules_to_not_convert setting. Therefore, the mlp.gate entry in _UNQUANTISED_OPS is indeed redundant for the final moe_router dtype. Removing it would make the code slightly cleaner without changing behavior.

• gitm/planner/model_catalogue.py - load_entry extends logic
  The extends logic is well-implemented, including cycle detection. The comment states provenance is not merged. This is good. However, the implementation data = {**{k: v for k, v in base.items() if k != "provenance"}, **data} means that if data (the extending entry) also has a provenance key, it will be kept. If data doesn't have a provenance key, then base's provenance will be excluded. This is the desired behavior: the extending entry's provenance (if present) takes precedence, otherwise no provenance is inherited. This is correct.

• gitm/planner/registry.py - detect_family order
  The comment correctly highlights the importance of glm_moe_dsa being detected before sparse_moe due to overlapping structural characteristics. This is a critical detail for correct model detection.

• gitm/planner/registry.py - _render_table tokens_per_step output
  The update to _render_table to use g.batch.tokens_per_step for speculative decoding is correct and aligns with the fix in BatchConfig. The warning for --spec-tokens without --acceptance-rate is also a good user experience improvement.

Overall Impression

This is an exceptionally thorough and well-engineered set of changes. The level of detail in the design note and the corresponding code implementation is commendable. The new GLM-5.2 model is complex, and the changes correctly capture its unique architectural features and their performance implications. The self-correction and validation against multiple checkpoints and scenarios (prefill, decode, MTP, different precisions) demonstrate a deep understanding of the problem space. The test suite for glm_graph.py is also very comprehensive.

The identified suggestions are minor and mostly about code clarity or very subtle edge cases, not critical bugs.

--- a/gitm/planner/glm_graph.py
+++ b/gitm/planner/glm_graph.py
@@ -300,7 +300,8 @@
     # 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.
    +# width instead (e.g., `attn_out_proj` input becomes `n_heads * (v_head_dim + kv_lora_rank)`).
    +# This is 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,
@@ -579,7 +580,7 @@
     # 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"),
    +# ("mlp.gate", "moe_router"), # Redundant as moe_router_dtype is authoritative for moe_router's precision.
     )
 
 

[aditchawdhary]