fix(sparse-mla): release preparation, cache and unused projection owners - #801
Merged
lukealonso merged 9 commits intoSep 21, 2026
Merged
Conversation
…zation Retain synthetic query and selection tensors through the prepared call closures, not through the published plan. Runtime kernels bind fresh caller-owned inputs; the compiled launcher does not need warmup tensor storage. This prevents per-layer probe retention during profiling and serving preparation without changing attention arithmetic or scratch ownership. Validation: the installed-wheel negative control fails the lifetime assertion; the patched CPU contract and native decode/prefill live-input CUDA-graph tests pass (3 cases). The TP2 allocator trace records 2,785,017,856 query bytes and 348,637,184 selection bytes alive when preparation exhausts memory. Whole-model validation requires the source-built wheel. Ruff and applicable hooks pass; six pre-existing mypy findings match the recorded parent control. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Martin Vit <martin@voipmonitor.org>
Exercise 4096 query rows and 64 heads through the prepared sparse-MLA launcher, then replay changed selections in a CUDA graph. This covers the declaration shape that exhausted memory with retained warmup probes. The lifetime contract plus three native shape cases pass on stock SM120. Serving source is unchanged from 25443f3; its published wheel remains the implementation under qualification. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Martin Vit <martin@voipmonitor.org>
…tion Check CUDA-graph replay with different query values, selected indices and active selection lengths after warmup probe storage is released. Uniform and saturated two-token attention cases have exact BF16 results. All four lifetime/native tests pass on SM120, including the 4096-row 64-head GLM shape; production source is unchanged. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Martin Vit <martin@voipmonitor.org>
|
Warning Review limit reachedNext included review available in 53 minutes. View limit detailsLimit details: You’ve used the included review currently available. You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. Review configuration: ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (9)
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review. 📝 WalkthroughWalkthroughThe B12X sparse MLA preparation path no longer publishes warmup probe tensors as owners. Tests verify probe release and CUDA graph replay with varied live inputs. ChangesSparse MLA probe lifetime
Priority: ⬇️ Low Estimated code review effort: 3 (Moderate) | ~20 minutes Change: Bug fix Suggested reviewers: Merge Risk: ⚪ Minimal · up to Warmup tensors are released only after preparation completes, and runtime launches use fresh caller inputs. No actionable merge-blocking risk remains. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
This was referenced Sep 18, 2026
MQA-only sparse attention retains its independently owned W_UK_T/W_UV matrices but releases the MXFP8 linear holder, packed weights and preparation provider. Reload repacks the linear and preserves absorbed matrix addresses. Profiling reserves expanded dense context only for backends that can execute dense MHA. Adapts the projection lifetime and profiling portions of 7f4aecc to the Karmic preparation API. Five focused CPU/native GPU tests pass, including exact MXFP8 absorption and reload with both sparse-only and dense-prefill configurations. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Martin Vit <martin@voipmonitor.org> (cherry picked from commit 3afee67)
B12X performs physical selection through its prepared indexer or caller-owned scratch. Do not allocate the generic group conversion buffers or events for that backend. Other sparse backends retain group allocation and notification behavior. Validation: four construction/notification contracts pass for implicit and explicit builders; stock TP2 allocation inventory identified twelve unused 33,611,788-byte buffers. Whole-model capacity qualification is separate. Existing sparse-MLA typing findings reproduce before this change; other hooks pass. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Martin Vit <martin@voipmonitor.org> (cherry picked from commit 7e8284e)
Release backend cache-writer bindings and C4 indexer views when the owning MLA layer releases its KV pool. Preserve cache geometry and allow binding a replacement pool without retaining profiling storage. Validation: component ownership and native graph/rebind tests; both ownership regressions fail on the unmodified adapter. The composed wheel passes 52 KDA/cache tests and its allocation census confirms release of 441188352 profiling bytes per TP2 rank. Six pre-existing mypy findings remain; other hooks pass. Whole-model TP2 capacity remains insufficient for independent reasons. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Martin Vit <martin@voipmonitor.org> (cherry picked from commit 48fdea5)
Declare cache-writer formats without token-sized CUDA allocations. Keep execution probes and cache-page backups inside synchronized preparation callbacks, not published plan owners. Generic prefill workspace is allocated only when a generic prefill backend consumes it. Cache precision, kernel selection and live-input graph bindings are unchanged. Validation: eight focused ownership/workspace cases pass, including native FP8/NVFP4 writers and changed-input CUDA graph replay. Sparse suite: 150 pass; six pre-existing query-discovery/indexer failures remain, with query-discovery identity failures confirmed on parent 4aeec44. Six memory negative controls fail on that parent. Six existing mypy findings remain; other applicable hooks pass. Full-model TP2/batch3072 qualification is pending. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Martin Vit <martin@voipmonitor.org> (cherry picked from commit d40c86f)
When an indexer owns the selection buffer, retain only the indexer reference. A redundant explicit buffer otherwise pins draft storage after the indexer switches to target-owned storage. Explicit-buffer-only layers and setter precedence are unchanged. Validation: four lazy-read, owner-release and explicit-override cases pass; all applicable pre-commit hooks pass. The qualified TP2 allocation census identifies a 25202688-byte draft buffer retained only by the unused fallback field. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Martin Vit <martin@voipmonitor.org> (cherry picked from commit ef1aeaf)
Exercise construction both with and without a released prior cache. Keep indexer plan and owner-release assertions independent of shared projection storage so memory PRs compose without conflicting test edits. Production code is unchanged. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Martin Vit <martin@voipmonitor.org>
voipmonitor
changed the title
Author
|
@coderabbitai full review Please review the complete proposed change, including ownership, graph lifetime and the recorded validation. Serving changes remain for maintainer review; integration qualification is not approval. |
|
|
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.
What this fixes
Return sparse-MLA profiling and preparation memory after its last consumer.
The adapter keeps live caller-owned binding without retaining synthetic token
banks, cache-page backups or stale cache-writer/indexer plans.
buffers. Absorbed projection owners are released; dense MHA retains its weights.
Kernel selection, precision, graph addresses and cache arithmetic are unchanged.
This includes and supersedes #804; do not merge #804 separately.
Validation
Status: qualified focused ownership/graph checks and composed Spark TP2 serving.
TP2 prefill geometry. Cache release/rebind tests cover native graph replay.
fail on the parent. Four shadowed-buffer tests pass with a failing ownership
control. The shadowed buffer alone retains 25,202,688 bytes per TP2 rank.
is not an all-green suite. Existing type findings are retained separately.
installed composed wheel on remote GPU5. This is a functional Max-Q test.
Only the reorganized ownership test differs. No whole-model speedup is
attributed to this PR alone.
Tests are in
tests/v1/attention/test_b12x_sparse_mla_api.py,test_sparse_mla_backends.py,test_mla_backends.py, andtests/models/test_glm5next_pooled_indexer.py.All original commit attribution is retained. #803 addresses KDA rather than
sparse MLA; #798 owns shared model workspaces rather than these cache lifetimes.
No additional equivalent KK PR remains after #804 is superseded.
OpenAI Codex assisted with implementation and validation; maintainer review is
required. The target is
dev/karmic-kraken.