Skip to content

issue/1571 feat(mamba2): add NVIDIA and MetaX inference operators - #1562

Open
big-hip wants to merge 2 commits into
InfiniTensor:InfiniLM-v0.2.9cfrom
big-hip:feat/mamba2-device-scan
Open

big-hip wants to merge 2 commits into
InfiniTensor:InfiniLM-v0.2.9cfrom
big-hip:feat/mamba2-device-scan

Conversation

@big-hip

@big-hip big-hip commented Sep 17, 2026

Copy link
Copy Markdown

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

  1. Add one shared NVIDIA/MetaX scan algorithm with thin platform launch wrappers, FP32 state, packed request offsets, indexed source/destination rows and bounded workspace; expose it through existing C/C++/Python and graph plan/run interfaces.
  2. Permit strict MetaX FP32 GEMM with 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.

  • Current retained NVIDIA scan suite: 54 output/state cases passed and 9 invalid-descriptor cases passed, using matching integrated source/runtime on A6000; no fresh isolated combined-head accelerator build claim.
  • Previously validated C500/NVIDIA operator matrix and C500 model integration use the same scan implementation. The unchanged MetaX precision regression passed two isolated processes covering FP32/FP16/BF16 and rows 1/3/257. MetaX was not rerun in this consolidation.
  • Project formatting with clang-format 21.1.8/Ruff 0.15.20 and whitespace checks pass.
  • Head 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 pass matrix.type to 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.

Saved C500/model and focused-test evidence, not a fresh CI screenshot

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.

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.
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 big-hip changed the title feat(mamba2): add indexed scan kernels for NVIDIA and MetaX issue/1571 feat(mamba2): add NVIDIA and MetaX inference operators Sep 20, 2026
@big-hip
big-hip marked this pull request as ready for review September 20, 2026 11:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant