Conversation
Add the descriptor, workspace and graph-aware operator interfaces for packed prefill and single-step updates. Share device scan kernels and verify outputs plus the entire state pool against an independent recurrence.
41 tasks
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
1 participant
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.
Add packed Mamba-2 SSD scanning and single-step recurrent updates with FP32 state, indexed request rows and device-side computation. NVIDIA and MetaX share the scan kernels and workspace contract through platform launch wrappers. Integrate with the existing C/C++/Python operator and graph plan/run interfaces.
Branch:
feat/mamba2-device-scan. Base:InfiniLM-v0.2.9c.Twenty files (+851/-0): 647 implementation/interface lines and 204 test lines. The test uses an independent token recurrence and compares the whole state pool, including untouched rows. Public interfaces document contiguous tensor shapes, valid distinct destinations, read-only zero state and unsupported cross-request aliasing.
Validation: the unchanged operator sources previously passed 54 output/state cases on C500 and NVIDIA plus descriptor negative cases. The current C500 real-model test run passed with the same operator binary. This session did not rerun the full standalone operator matrix or rebuild this isolated branch. Formatting and whitespace checks passed.
Full Mamba Decode-graph use also needs the separate graph-lifetime/recording fixes in #1560. This operator PR does not include them or the MetaX GEMM precision switch. Scheduler chunked Prefill, state prefix caching, PP and additional device backends are outside its scope. The scan kernel's internal 256-token chunks do not implement scheduler chunking.
Published as a draft while independent branch builds and upstream CI are pending.
Saved output, source provenance, device conditions and archived measurements: immutable evidence. Attachments live on a separate fork branch and are not part of this source diff.
CI at publication: fork Ruff passed; the independent fork CPU build/test matrix is in progress. The upstream run reports
action_requiredand needs repository-side approval. No successful upstream build is claimed.