Fix GLM-5.2 packed KV-cache layout - #559
Merged
Merged
Conversation
|
|
Declare the DSA decoder's per-layer packed cache widths through the causal LM task so the unpacking Squeeze operates on its intended singleton head axis. Register and type the IndexShare custom op, and add tiny-config and cache-layout regression coverage. Signed-off-by: GitHub <noreply@github.com> Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
The DSA graph requires the pkg.nxrt IndexShare runtime kernel, which stock ORT used by synthetic parity does not register. Keep graph/checker and dedicated DSA tests active while treating this as a capability skip. Signed-off-by: GitHub <noreply@github.com> Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Adopt main's dedicated GlmMoeDsaTask implementation from #560 and retain focused production-task cache layout regression coverage. Signed-off-by: GitHub <noreply@github.com> Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Copilot
AI
changed the title
[WIP] Fix invalid Squeeze in glm_moe_dsa model
Fix GLM-5.2 packed KV-cache layout
Aug 23, 2026
justinchuby
marked this pull request as ready for review
August 23, 2026 04:40
Performance Comparison
|
Contributor
There was a problem hiding this comment.
Pull request overview
This PR updates the co-located GLM-5.2 DSA graph-build tests to exercise the production GlmMoeDsaTask path and adds a regression assertion that the KV-cache I/O is declared using the packed single-head layout expected by DSA.
Changes:
- Switch the test build path from the generic
"text-generation"task to the dedicated"glm-moe-dsa"task. - Add a regression test that asserts past/present KV-cache shapes are symmetric and use
[B, 1, T, packed_width]with per-layer packed widths derived from the config.
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+259
to
+262
| cache_inputs = {value.name: list(value.shape) for value in onnx_model.graph.inputs[3:]} | ||
| cache_outputs = { | ||
| value.name: list(value.shape) for value in onnx_model.graph.outputs[1:] | ||
| } |
justinchuby
added a commit
that referenced
this pull request
Aug 23, 2026
## Root cause `_DeepSeekMoEFFN.__init__` (src/mobius/models/deepseek.py) accepted `linear_class` but only forwarded it to `_SharedExpertMLP`, not to the `MoELayer` that builds the routed-expert dense loop-over-experts fallback: ```python self.moe = MoELayer(config, gate=gate) # linear_class silently dropped ``` Consequence: a quantized config correctly quantized attention, dense FFN, and the shared expert, but every routed MoE expert's gate/up/down projections stayed plain float `MatMul`. This both loses quantization for the routed experts and breaks the `fuse_dense_moe_to_qmoe` post-hoc rewrite, which only pattern-matches a quantized `MatMulNBits` dense-fallback shape — so a quantized DeepSeek-V2/V3 config would silently fail to fuse into QMoE for its routed experts. DeepSeek-V4 is unaffected: its `DeepSeekV4MoE` uses `expert_factory` to build `_DeepSeekV4Expert`, which independently recomputes its quantized class from `config.quantization` per-expert rather than relying on `MoELayer`'s `linear_class` forwarding — verified via review, no matching fix needed there. ## Fix One line: thread `linear_class` through to the routed-expert `MoELayer`: ```python self.moe = MoELayer(config, gate=gate, linear_class=linear_class) ``` ## Tests Added `test_deepseek_moe_ffn_linear_class_reaches_routed_experts` in `deepseek_test.py`, which isolates the assertion on `module.moe.experts` (routed) separately from `module.shared_experts` (already worked before this fix), rather than relying on an aggregate creation count — so it can't pass for the wrong reason. - Verified the new test **fails** on pre-fix code (`assert False` at the routed-experts isinstance check) and **passes** post-fix. - Full suite: `PYTHONPATH=src python3 -m pytest src/` → **3249 passed, 9 skipped**, no failures. - `ruff check` / `ruff format --check` on changed files: clean. - Independently reviewed (separate reviewer pass over the diff); verdict: **Approve**, no findings. ## Scope Deliberately minimal and focused: does not touch GLM-5.2 registration, DeepSeek-V4 QMoE export, YaRN, or any other follow-up — those are tracked/ landed separately (#548, #550, #555, #559, #560). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.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
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.
GLM-5.2 declared standard multi-head KV-cache inputs while DSA expected packed single-head caches, causing an invalid
Squeezeand incomplete output shape inference.Cache wiring
GlmMoeDsaTask.Regression coverage