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.
This was referenced Sep 17, 2026
Consolidate InfiniTensor#1561 into the Mamba-2 backend contribution. Preserve the default TF32 behavior and the existing process-isolated strict-FP32 regression; no scan algorithm changes.
big-hip
marked this pull request as ready for review
September 20, 2026 11:15
3 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
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.
Problem and result
Closes #1571. Target:
InfiniLM-v0.2.9c.InfiniLM's Mamba-2 model needs indexed SSD scan and single-token recurrent updates that keep Conv/SSM request state on the device. This PR supplies the scan operator and the small MetaX precision prerequisite needed for strict FP32 continuation checks, paired with InfiniTensor/InfiniLM#575.
Implementation and review order
INFINIOP_METAX_ALLOW_TF32=0, preserving default TF32 and existing FP16/BF16 accumulation. The descriptor captures this policy; set it before process startup.These are two commits within one feature PR. The precision commit consolidates #1561 verbatim; scan math and its tests are unchanged. Graph allocation/recording prerequisites are centralized in #1566, shared with Qwen MTP. There is no dependency on a separate #1560 or #1561 PR after consolidation.
The operator documents contiguous shapes, distinct valid destinations, read-only zero state, unsupported cross-request aliasing, state size ≤256 and dtype/shape bounds. Core tests use an independent token recurrence and compare the entire state pool, including untouched rows. Internal 256-token scan chunks do not implement scheduler chunked Prefill.
Validation
Focused checks and byte-equivalence audit; C500/NVIDIA model and operator evidence. No logs, weights or benchmark experiments enter this source diff.
f634435b: fork Ruff passed; Linux/Windows build matrix passed all four Ubuntu/Windows build and CPU-test jobs. The existing workflow labels jobs debug/release but does not passmatrix.typeto the build command; these results do not certify distinct build modes or GPU execution. Original scan and precision branch matrices both completed successfully. Upstream jobs still require maintainer approval.Scope and dependencies
Request code review for NVIDIA/MetaX operator correctness and precision control. Merge requires the final required checks and reviewer approval; review-ready is not a merge claim. Full model Decode-graph integration additionally uses #1566 and InfiniTensor/InfiniLM#575. Other backends, PP, scheduler chunking and recurrent-state prefix reuse are outside this operator's scope. The precision switch is a correctness control, not a performance claim.
@spike-zhu and @wooway777 please review the shared operator contract and MetaX policy. The original head branch is retained to preserve this PR and its history; no replacement PR or force push is used.