GLM-5.2 (glm_moe_dsa) inference: IndexShare, indexer RoPE/eps, int8 MLA-KV - #1419
Open
avlp12 wants to merge 7 commits into
Open
GLM-5.2 (glm_moe_dsa) inference: IndexShare, indexer RoPE/eps, int8 MLA-KV#1419avlp12 wants to merge 7 commits into
avlp12 wants to merge 7 commits into
Conversation
…LA-KV Adds inference support for GLM-5.2 (model_type glm_moe_dsa), a DeepSeek-V3.2-style MLA + DSA sparse-attention MoE: - IndexShare: the DSA indexer runs only on "full" layers; "shared" layers reuse the previous full layer's top-k (index_topk_freq). derive_indexer_types() defaults to all-"full" (index_topk_freq=1), preserving stock DeepSeek-V3.2 behavior. - Indexer RoPE/eps: GLM-5.2 uses non-interleaved RoPE + LayerNorm eps 1e-6. deepseek_v32 defaults are unchanged (traditional RoPE, eps 1e-5); glm_moe_dsa overrides them. - int8 MLA latent KV cache via MLACacheList (a CacheList subclass used only by these models): quantizes only the compressed-latent cache, keeps the tiny DSA indexer cache full precision; the latent is dequantized on read. Lets a 1M-token context fit on one machine. - KVCache/ChunkedKVCache.state tolerate an unwritten cache (the indexer KV on shared layers) so mx.eval(cache.state) doesn't deref keys=None. - glm_moe_dsa ModelArgs defaults the fields GLM-5.2's config omits. Adds a glm_moe_dsa entry to test_all_models and a test for the int8 latent-KV path. Resolves ml-explore#879. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
GLM IndexShare leaves the DSA indexer KV empty (keys=None) on "shared" layers. batch_generate (used by mlx_lm.evaluate and batched inference) converts each cache to BatchKVCache and assumes every cache is written, which crashes for glm_moe_dsa / deepseek_v32 with index_topk_freq > 1: - BatchKVCache.state getter/setter: tolerate keys=None (empty state). - BatchKVCache.extract: return an empty per-sequence cache when unwritten. - dynamic_roll: no-op on a None array (left/right padding path). - generate.py BatchStats: guard prompt/generation tps against zero time. Adds test_batch_kv_cache_unwritten covering the four guards. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Contributor
Author
|
Added a follow-up commit extending support to |
Contributor
Author
|
Real-world confirmation that this PR is required to load IndexShare-quantized GLM-5.2 checkpoints with stock
That's exactly Happy to rebase onto current |
Restore GLM-5.2's built-in MTP layer (model.layers.78) instead of stripping it, and use it for lossless self-speculative decoding: - deepseek_v32: sanitize() keeps + remaps layer-num_hidden_layers into an MtpModule when num_nextn_predict_layers > 0 and the checkpoint retains the weights (otherwise strips as before). Model.mtp_forward() (causal-masked for multi-token calls), has_mtp, make_mtp_cache; backbone stashes the pre-final-norm hidden (h_prenorm) that the MTP head consumes. - generate: mtp_speculative_generate_step — chunked prefill of both caches, k chained drafts feeding back the shared_head-NORMED hidden (chaining the raw hidden halves chained acceptance: 0.27 -> 0.75 conditional measured), sampler-equality acceptance (distribution-lossless at any temperature), optional conservative prompt-lookup drafting (--mtp-hybrid), and a single-sync pipelined loop (next drafts+verify enqueued before yielding so detokenizer/wrapper Python overlaps GPU). CLI: --mtp, --mtp-num-draft-tokens, --mtp-hybrid. - server: same flags; MTP requests use fresh caches and skip prompt-cache reuse/store (the spec loop can exit with uncommitted tail entries); batching disabled in MTP mode. - Also syncs the serving stack these changes build on (previously local): int8 MLA-KV moved into cache.py CacheList, DeepseekV32 layer-signature cleanup, optional MLX_LM_EVAL_EVERY_LAYERS command-buffer splitting. Measured (GLM-5.2 744B 3.5bpw, M3 Ultra 512GB, greedy, single request): 22.0 -> 24.6 tok/s (~1.11x), mean accepted length ~2.5, per-position acceptance ~0.88/0.75/0.58; greedy output verified identical to non-speculative decoding, tiny-model e2e (loader/prefill-equivalence/ lossless k=1..3/hybrid/legacy-strip) all pass. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
MTP drafting is validated in the dense-attention regime only. Past index_topk tokens the DSA sparse path engages and MTP acceptance collapses (index_skip_topk_offset handling is not wired yet) — measured 3x slower than plain decoding at a 2190-token prompt. stream_generate now falls back to generate_step with a stderr note instead of degrading. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The earlier 3x slowdown that motivated the guard (49c5505) was misdiagnosed: per-position acceptance in the DSA-sparse regime is fine (measured 0.73-0.82 first-draft at 2750-token context, accept-len 2.1-2.3 end-to-end), and index_skip_topk_offset turned out to be a config-schedule parameter (indexer_types derivation), not an MTP inference knob — it is already honored via the explicit indexer_types. The observed slowdown decomposed into (a) benchmark-host swap pressure and measurement ordering, and (b) the real-but-modest cost of the L>1 sparse-mask verify path (~2.4x a plain step for a 3-token verify), which roughly offsets the accepted length: long-context MTP is approximately performance-neutral, not harmful. A future win is extending the L==1 gather path to small-L verifies so verify cost stops scaling with context. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…text MTP) At context > index_topk, L<=4 forwards (MTP verify, tiny continuations) previously took the L>1 sparse-mask path: per-position argpartition + put_along_axis boolean masks + attention over the FULL key set, on every layer. Amortized over a 2048-token prefill chunk that is fine; paid per speculative-decode iteration it dominated everything — measured 0.26x plain decode at a 2.1K-token context (accept-len 1.8 held; the cost was pure verify overhead, ~7 plain-steps per iteration). L<=4 sparse forwards now gather each position's top-k keys (like the existing L==1 branch) and run batched-matmul attention with k = v = latent + unembed after; gathered scores are re-masked for the threshold-boundary case where a position has fewer than K causally valid keys. Equivalence vs the mask path verified at 1e-06 on a tiny sparse model (L=2/3/4); tiny MTP e2e (lossless k=1..3, hybrid, chunked prefill, legacy strip) all pass. MLXLM_NO_SMALL_L_GATHER=1 restores the old path. Measured (GLM-5.2 4-bit sibling, M3 Ultra 512GB, 2146-token context, prefill-separated timing): mtp k=2 decode 5.35 -> 17.38 tok/s (0.26x -> 0.85x of plain); k=1 0.98x; short-context k=2 unchanged at 1.15x. Remaining long-context gap is acceptance economics (accept-len 1.8 vs ~2.1 break-even), not path cost. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
self.h_prenorm (added for MTP) was a public mx.array attribute, so it
joined Module.parameters() after the first forward — tree_flatten-based
save flows then wrote it into checkpoints, which fail strict loading
("Received 1 parameters not in model: model.h_prenorm"). Renamed to
_h_prenorm (underscored attributes stay private in MLX modules); caught
by a quantize->save->load round-trip test.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Resolves #879 — adds inference support for GLM-5.2 (
model_type: glm_moe_dsa), Zhipu/Z.ai's 744B (~40B active) MoE. It is architecturally DeepSeek-V3.2 (MLA + DSA "lightning indexer" sparse attention), so this extends the existingdeepseek_v32.pyrather than adding a new model.What was missing
The stock
glm_moe_dsa/deepseek_v32port runs the DSA indexer on every layer and threads no indices between layers. GLM-5.2 computes top-k on a "full" layer and reuses it on the following "shared" layers (IndexShare). At ≤2048 ctx this is invisible (top-k selects all tokens → dense attention), but beyond that the port diverges from the reference and forfeits the sparse-attention savings.Changes (inference only)
derive_indexer_types()marks layers full/shared fromindex_topk_freq; "full" layers run the indexer, "shared" layers reuse the previous full layer's top-k (threaded through the decoder).index_topk_freq=1(the default) makes every layer "full", i.e. stock DeepSeek-V3.2 behavior is unchanged.1e-6, distinct from its interleaved main attention (verified against the HF reference to ~1e-7 on post-RoPEq).deepseek_v32defaults are unchanged (traditional=True, eps1e-5);glm_moe_dsaoverrides them.MLACacheList(aCacheListsubclass used only by these models) quantizes only the large compressed-latent cache and keeps the tiny DSA indexer cache full precision; the latent is dequantized on read. This lets a 1M-token context fit on a single machine. The genericCacheListis left untouched, so other hybrid-cache models are unaffected.KVCache/ChunkedKVCache.statetolerate an unwritten sub-cache (the indexer KV on shared layers) somx.eval(cache.state)during generation doesn't dereferencekeys=None.glm_moe_dsaModelArgsdefaults the fields GLM-5.2'sconfig.jsonomits.Tests
test_all_modelsgains aglm_moe_dsaconfig withindex_topk_freq=4(exercises full + shared layers, fp32/fp16 prefill, decode, batch, deepcopy).test_glm_moe_dsa_quantized_kvcovers the int8 latent-KV path: only the latent cache is quantized, the indexer cache stays fp16, and decode after quantization stays finite.All existing
deepseek_v32tests still pass (defaults preserve V3.2 behavior).Scope
Inference + KV-cache only. No training / AWQ / DWQ or distributed/pipeline changes.
🤖 Generated with Claude Code