feat(mamba2): integrate model loading and recurrent inference

[big-hip]

Summary

Support the prepared state-spaces/mamba2-130m checkpoint through existing inference, benchmark and service entrypoints. Add a Mamba-2 model and processor, weight mapping/preparation, indexed convolution/SSM state, tensor-parallel projections and protected Decode capture. Reuse the existing linear, convolution, normalization, state-pool and graph mechanisms. Prefill remains eager.

The first commit fixes shared packed-linear layouts and loading tied weights before capture; the second adds the model; a focused follow-up rejects unsupported pipeline service configuration before worker setup. The branch contains 18 changed files (+1394/-32), including 636 test lines and 65 README lines.

Branch: feat/mamba2-model-inference. Base: InfiniLM-v0.2.9c.

Motivation

Add a pure SSM model architecture using device computation on NVIDIA and MetaX. For 130M TP1 BF16, recurrent state occupies approximately 18.25 MiB per request independent of processed history length; the existing scheduler still enforces its logical page budget.

Dependencies: InfiniCore #1562 (scan kernels) and #1560 (graph lifetimes/recording); strict MetaX FP32 validation additionally uses #1561. The three Core diffs are disjoint and were validated together; the model needs their combined runtime as described below.

Type of Change

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

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

Follow-up 5f89ae3d: the existing unsupported-service configuration regression now covers PP and fails if it reaches ModelRunner. Before the fix, the PP case reached worker setup (1 failed, 4 passed); after the fix, 5 passed. Repository formatting passed. These configuration checks load no model and are not a fresh GPU/model performance run. The C500 42-test result and performance evidence below refer to the earlier d337a05f revision; model math and device code are unchanged.

Platform	Validation
C500 50% compute / 32000 MiB slice	Current retained model suite plus shared-layout regression: 42 passed, BF16 TP1, eager Prefill + Decode graphs; format checks passed.
Two A6000 48GB	Prior matching-production-source FP32/FP16/BF16 TP1/TP2, graph/state/cancellation checks; not rerun in this session because GPUs were busy.
C500 FP32/FP16	Prior matching-production-source eager/Decode-graph matrix passed; not rerun this session.
C500 TP2, 910B, S4000, PP	Not validated/supported by this adaptation.

Core graph/precision checks passed separately (13 cases). Retained model checks cover recurrence, dynamic request reordering, zero-state isolation, recapture with live requests, state-pool rebuild, capacity exhaustion, cancellation and EOS/length release. Default repeat loops are 16 Decode steps with recapture at step 8 and 12 service cycles; this is bounded regression coverage, not a soak test.

Prior matching-source entrypoint checks completed text inference, offline benchmark, evaluation adapter and HTTP service. The five MMLU samples scored 0/5: this demonstrates adapter execution only, not acceptable task accuracy. Those entrypoints were not all rerun in this cleanup. Saved output is attached below as a rendered image, with immutable logs and provenance links; it is not a fresh CI screenshot.

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.

Benchmark / Performance Impact

Archived C500 BF16 TP1 internal comparison: 128 input tokens, batch 1, 128 greedy output tokens, 16 Decode warmup tokens and three repetitions. Enabling existing Decode graphs reduced the median of per-run mean ITLs from 12.54 to 7.67 ms; finite-request output rate rose from 79.48 to 128.95 token/s. Prefill was eager in both configurations; capture/loading/tokenization were excluded. This is not a new Prefill-graph result 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. It is a bounded numerical-quality check, not universal logit equality or a broad quality evaluation. Official NVIDIA graph inference remained faster in the archived A6000 comparison.

Notes for Reviewers

Current scope: pure Mamba-2, one B/C group, convolution width 4, 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. No CPU scan/update fallback. Low-precision logits can vary across compute shapes.

No research logs, weights, profiles or machine-specific build patches are included in the source diff. Native implementation and model math match the previously tested snapshot; the only subsequent production change is the two-line Python PP preflight check. Fresh full builds of every split branch remain pending.

CI / ChatOps

Current 5f89ae3d fork push checks passed: CI formatting and Ruff. The upstream CI run reports action_required and needs repository-side approval before jobs start.

Local targeted regressions, repository formatting and whitespace checks passed. This remains a Draft PR: fresh independent accelerator builds and remaining checklist audits are pending. Fork formatting is not a full hardware build. No maintainer review is requested yet.

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

• For any platform that could not be tested, an explicit reason is given in the table and a reviewer with access has been tagged.
• 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 has been triggered manually (Actions → CI on this branch), or /retest was requested.

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.
• Third-party code is license-compatible and attributed.
• No unsafe pointer arithmetic, uninitialized reads, or missing bounds checks were introduced.