feat(mamba2): integrate model loading and recurrent inference

[big-hip]

Summary

Support the prepared state-spaces/mamba2-130m checkpoint through the existing inference, benchmark and service entrypoints, with eager Prefill and indexed recurrent Decode on NVIDIA and MetaX. Reuse existing parallel linears, causal convolution, normalization, request-state management and Decode graphs.

Model-specific code stays in csrc/models/mamba2/. Shared correctness changes isolate packed-linear layouts, finish tied-weight loading before capture, and preserve active recurrent/Attention caches when recapturing graphs. Unsupported service configurations fail before worker setup. The source diff contains 18 files (+1424/-32), including three regression modules (653 lines) and 65 README lines. No experiment logs, weights, profiles or fork-specific workflows are included.

Base: InfiniLM-v0.2.9c. Dependencies: InfiniCore #1562 for scan kernels/MetaX precision policy and #1566 for graph allocation/recording correctness. The separate InfiniLM Qwen MTP PR is not required.

Motivation

Add a pure SSM architecture without CPU model-compute fallback. For 130M TP1 BF16, convolution plus FP32 SSM state is approximately 18.25 MiB per request, independent of processed history length. The existing scheduler still applies its logical page budget; the state pool contains max(2, num_blocks // 4) rows including reserved zero state.

Type of Change

• feat — new model architecture.
• fix — required loading/layout and graph-capture correctness.
• Breaking change.

Test Results of Involved Models on Supported Platforms (Please attach screenshots)

Current source 8a591cb1, combined Core runtime a3ac4df4 + scan/MetaX f634435b:

Platform	Validation and limits
A6000 BF16 TP1	43 passed, including model/service and ordinary Qwen compatibility.
Two A6000 BF16 TP2	43 passed; dynamic graph indices, cancellation/release and request-state checks included.
A6000 FP32 TP1	3 passed, 35 intentionally deselected: continuation, graph recapture and pool rebuild.
C500 BF16 TP1, MACA 3.7.0.38	43 passed on the current head, including ordinary Qwen recapture; standard text inference with Decode graphs also completed.
C500 FP32 and FP16 TP1	3 passed each, 35 intentionally deselected per dtype: continuation, graph recapture and pool rebuild.
C500 Core prerequisites	54 scan cases + 17 regressions passed, covering state correctness, descriptor validation, GEMM precision and dynamic cat graph replay.
C500 TP2, Ascend, Moore, PP	Outside the validated/supported scope of this adaptation.

LM was rebuilt from an empty build directory against the consolidated Core headers, then linked to and tested with the newly built combined runtime. Core used an isolated integration worktree and an existing build cache. Loaded library paths/hashes are recorded. The A6000 workstation is shared; these are correctness results, not new performance or physical peak-memory measurements. No test was skipped in the listed A6000 runs. TP/dtype labels refer to Mamba; the ordinary synthetic Qwen2 fixture is FP16 TP1 in both suites.

The retained checks cover continuation, TP, dynamic graph request indices, batch fallback to device eager, live-request recapture, pool rebuild, admission under state exhaustion, cancellation, EOS/length release and slot reuse. Ordinary Qwen exercises the shared loader/packing/capture changes. This is bounded regression coverage, not a soak test or certification of every Mamba-2 checkpoint.

The recapture regression failed before 8a591cb1: warmup overwrote physical KV page zero of a live ordinary Qwen request, producing 127/128 logits outside the original 1e-3 tolerance. The same check passes after saving/restoring that page; no tolerance was relaxed. Synchronization/cache backup is restricted to capture and does not add per-Decode copying.

Current-source review, raw logs, JUnit results and build provenance.

Earlier matching-model-source entrypoint checks covered text inference, offline benchmark, evaluation adapter and HTTP service. The five-sample MMLU run scored 0/5; it demonstrates adapter execution only. The current C500 run reran examples/test_infer.py; the offline/evaluation/HTTP entrypoint evidence remains archived. Current service lifecycle coverage uses LLMEngine directly. Archived logs and provenance.

Fresh C500 validation used one nominal 64 GiB card, MACA 3.7.0.38 / driver 3.8.30, PyTorch 2.8.0+metax3.7.0.7, Python 3.12.11 and Transformers 4.56.2. Both projects built from empty build directories using cached Boost/pybind11 dependencies. Source checksums matched the recorded revisions before and after testing; no production code or numerical assertion changed. No C500 test was skipped. Environment, current-source logs, commands and benchmark records.

Benchmark / Performance Impact

Current C500/MACA 3.7 check: no measured Decode-graph speedup. BF16 TP1, batch one, 128 input + 128 greedy output tokens, 16-token Decode warmup; two process pairs in eager/graph then graph/eager order, three measured requests per process (six per mode). Prefill was eager in both. Capture, loading and tokenization were excluded.

Median across six runs per mode	Eager	Decode graph
Finite-request output rate	120.45 token/s	117.64 token/s
Per-run mean ITL	8.22 ms	8.42 ms

All 12 measured output hashes matched. Graph output rate was about 2.3% lower overall; the two ordered pairs differed in magnitude, so this bounded measurement does not establish a general regression cause. Keep the existing eager default for this measured configuration. No new process peak-memory or sustained HTTP-throughput claim is made.

Archived C500 BF16 TP1 comparison: 50% compute / 32,000 MiB slice, MACA 3.5.3, PyTorch 2.8.0+metax3.5.3.9; 128 input tokens, batch one, 128 greedy output tokens, 16 Decode warmup tokens and three repetitions. Existing Decode graphs reduced the median of per-run mean ITLs from 12.54 to 7.67 ms; finite-request output rate increased from 79.48 to 128.95 token/s. Prefill was eager in both configurations. Capture, loading and tokenization were excluded. This is not a new throughput measurement or sustained service benchmark.

Fixed WikiText subset: 8 passages / 4099 prediction targets; C500 BF16 mean NLL 3.14331950 versus archived official NVIDIA BF16 3.13971278. This is a bounded quality check; low-precision cross-shape logits are not bitwise identical. Official NVIDIA graph inference remained faster in the archived A6000 comparison.

Source review and current-source NVIDIA/MetaX TP1 acceptance are complete. The C500 verification blocker is resolved. Upstream CI approval and the two prerequisite Core PRs remain merge requirements; fork success does not imply upstream CI has run.

Notes for Reviewers

Scope: pure Mamba-2, one B/C group, convolution width four, head-wise D, gated RMSNorm after gating, and unbounded time steps. No hybrid Attention/SSM, quantization, scheduler chunked Prefill, prefix-state snapshots, remote state transfer, speculative rollback or PP. Ascend and Moore scan backends are not provided by this adaptation.

Shared loader/linear/compiler code also affects ordinary models. The Qwen compatibility regression covers the relevant tied-head, pre-transpose and recapture paths; it does not imply that all existing models/platforms were rerun. See the review evidence for exact source and runtime revisions.

CI / ChatOps

Current-head fork format and Ruff checks passed. The standard LM push CI only checks formatting; its Ruff job checks the formatter script. These are not GPU tests.

Independent fork CPU validation passed: both Core and LM builds succeeded, followed by 33 passed, zero failed, zero skipped configuration, loading and service-preflight checks. The downloaded artifact confirms the combined Core source tree matches the local GPU build. This separate validation branch pins both Core dependencies and the exact LM head; no fork workflow enters this source PR. Earlier workflow setup failures (missing xxhash, then incompatible optional xmake caching) were corrected without changing production code or test assertions. CPU contracts do not replace accelerator tests.

Core prerequisite CI: #1562 and #1566. These workflows build/test CPU on Ubuntu and Windows. GPU validation is recorded separately above. Upstream action_required still needs maintainer approval; fork results do not override upstream workflow policy.

Checklist

Every contributor must verify every item below before requesting
review. Tick each box only after the check has actually been performed —
do not tick speculatively. If an item truly does not apply, replace the
checkbox with N/A and briefly explain why in an inline comment.

Title, Branch, and Commits

• PR title follows Conventional Commits (e.g. feat(nvidia): …, fix(cuda/gemm): …).
• Branch name follows <type>/xxx-yyyy-zzzz where <type> matches the PR title's Conventional Commits type and words are joined with hyphens (see CONTRIBUTING.md §Branches).
• Each commit message follows Conventional Commits.
• Small PR is a single squashable commit; or, for a large PR, every commit is meaningful, well-formed, and independently reviewable (see CONTRIBUTING.md §Pull Requests).
• The branch is based cleanly on the requested InfiniLM-v0.2.9c (270feb3e) without merge commits; the remote base was verified unchanged at publication.
• No fixup! / squash! / wip commits remain.
• N/A: This branch does not use the legacy issue format.

Scope and Design

• Changes are minimal — nothing unrelated to the stated motivation was added (CONTRIBUTING.md §Code/General).
• No dead code, commented-out blocks, debug prints, printf/std::cout/print(...) left behind, or TODO without an owner and issue link.
• No unrelated formatting churn that would obscure the diff.
• Public API changes (if any) are intentional, documented, and reflected in affected callers/tests.

General Code Hygiene (applies to all languages)

• The code is self-explanatory; comments were added only where the why is non-obvious (CONTRIBUTING.md §Code/General).
• Every modified or added file ends with a single trailing newline (CONTRIBUTING.md §Code/General).
• No trailing whitespace, tab/space mixing, or stray BOMs.
• Identifiers in comments and error messages are wrapped in backticks (e.g. the `seqlens_k` tensor) (CONTRIBUTING.md §Code/General).
• All comments and error messages are in English (CONTRIBUTING.md §Code/General).
• Comments and error messages are complete sentences — capitalized first letter, terminal punctuation — unless the language/framework convention says otherwise (CONTRIBUTING.md §Code/General; §Python).

C++ Specific (if C++ files changed)

• Code follows the Google C++ Style Guide strictly.
• Error and warning message wording follows the LLVM Coding Standards (CONTRIBUTING.md §C++).
• Constructor initializer list order matches member declaration order (CONTRIBUTING.md §C++).
• No raw new/delete; RAII / smart pointers / existing allocators are used.
• Changed files are formatted by scripts/format.py.
• No changes/reference to csrc/models/llama_legacy/.

Python Specific (if Python files changed)

• Code is PEP 8 compliant.
• Comments are complete English sentences, starting with a capital letter and ending with punctuation; Markdown backticks are used for code references (CONTRIBUTING.md §Python).
• Docstrings (if any) follow PEP 257 (CONTRIBUTING.md §Python).
• Changed files are formatted by scripts/format.py.
• No changes/reference to python/infinilm/auto_config.py.

Testing

• N/A: No platform within the claimed adaptation scope remains untested: NVIDIA and MetaX TP1 current-source checks passed. C500 TP2, PP and Ascend/Moore scan backends are explicitly outside the validated/supported scope above; no unavailable-platform reviewer substitution is needed.
• Passed single request test (examples/test_infer.py), or specify the reason for skipping.
• Passed offline performance test (examples/bench.py), or specify the reason for skipping.
• Passed sanity test (test/bench/test_benchmark.py), or specify the reason for skipping.
• Passed service test (python/infinilm/server/inference_server.py + scripts/test_perf.py), or specify the reason for skipping.

Build, CI, and Tooling

• The project builds cleanly from a fresh directory on at least one affected platform.
• CI was explicitly triggered in the contributor fork with pinned LM/Core sources, as requested while upstream external-PR CI requires approval. The maintainer-only hardware workflow was not misreported as having run.

Documentation

• README.md, CONTRIBUTING.md, or inline docs updated when behavior, build flags, or developer workflow changed.
• N/A: No user-visible breaking change; the new architecture is opt-in and unsupported variants are rejected.

Security and Safety

• No secrets, access tokens, internal URLs, customer data, or personal hardware identifiers have been committed.
• N/A: No third-party implementation was copied or vendored; existing framework dependencies are reused. This review does not certify repository-wide licensing.
• No unsafe pointer arithmetic, uninitialized reads, or missing bounds checks were introduced.