feat: support MiniMax-Text-01 (hybrid Lightning/full-attention MoE, MXFP4, TP/PP)

[Kritace]

Summary

• Add end-to-end support for MiniMax-Text-01, a 456B-class MoE model that
  interleaves Lightning (linear) attention (70 layers) with full
  attention (10 layers) and uses per-layer MoE with a shared expert pool.
• New model directory csrc/models/minimax_text_01/* (14 files):
  • minimax_text_01_linear_attention — Lightning attention with
    token-by-token recurrent decode and a chunked (block-scan) prefill,
    mathematically equivalent to the HF reference (A/B diff ~1e-4);
  • minimax_text_01_attention (full attention, GQA + partial RoPE),
    minimax_text_01_fused_moe_experts / sparse_moe_block (TP-aware,
    MXFP4-capable MoE), minimax_text_01_decoder_layer,
    minimax_text_01_for_causal_lm, minimax_text_01_allocate_kv_cache_tensors
    (hybrid cache: full layers → paged/static KV, linear layers → recurrent
    state pool, both split per PP stage).
• Python integration via the framework's standard extension points (no changes
  to shared model code):
  • python/infinilm/modeling_utils.py — _remap_minimax_text_01 weight
    remapper (qkv split, linear-attn rename, MoE packing);
  • python/infinilm/infer_engine.py — recognize attn_type_list-based
    models as mamba-cache models (MiniMax is identified and driven through the
    existing hybrid-cache engine path);
  • python/infinilm/processors/minimax_text_01_processor.py — request↔state
    slot mapping for the Lightning state pool.
• Distributed (TP/PP) and MXFP4 quantization are supported; Lightning prefill
  is chunked (default chunk 64, INFINILM_LIGHTNING_CHUNK / config override to
  fall back to the token-by-token path).

Motivation

MiniMax-Text-01 is a widely used open-weight model with a novel hybrid
attention architecture (Lightning linear attention + full attention + MoE) that
InfiniLM did not previously support. This PR adapts it following the project's
model-adapter conventions (MODELS.md / CONTRIBUTING.md), keeping all changes
non-invasive (no edits to framework-shared code, llama_legacy/ or
auto_config.py). A chunked prefill optimization is included because the
token-by-token Lightning recurrence was the prefill bottleneck
(~6.3× speed-up at seq=2048 on a single 4090D).

Type of Change

• feat — new feature / new model
• fix — bug fix
• perf — performance improvement (no behavioral change; chunked prefill
  is numerically equivalent to the recurrent path, A/B diff ~1e-4)
• refactor / test / docs / build / ci / chore
• Breaking change

Test Results of Involved Models on Supported Platforms

This is a vastly new model structure (first hybrid Lightning/full-attention
MoE with a recurrent state pool in this repo), so per the PR template this PR
declares partial support and requests confirmation from an admin. The
model's official weights are bf16-only ≈913 GB (456B params), which cannot be
fetched/run on the test platforms available to this contributor; all numeric
correctness checks therefore use the canonical small config
(hidden=512, 4 layers, attn_type [0,1,1,0] — 2 Lightning + 2 full layers)
with synthetic random weights, which exercises the exact same code paths. The
demo prints sampled tokens as <tokN> placeholders (the synthetic checkpoint
uses a placeholder tokenizer; the sampled ids are valid 0–200063). Model
correctness is established by the numeric logits matrix below — free-text
quality cannot be judged without the real 456B weights.

Test	Result	Notes
Single request — examples/test_infer.py	Passed (synthetic small weights)	End-to-end generation over the full LLM path (Paged KV cache + Mamba state pool, num_blocks=128): prompt rendered, 32 tokens sampled, response printed. Reproduced in three configurations with bit-identical sampled token sequences (local RTX 3050 / server single 4090D / server TP=2), greedy determinism confirming cross-environment/TP consistency — see Run Logs below.
Offline performance — examples/bench.py	Passed (workflow-level, synthetic small weights)	bench.py runs end-to-end and reports prefill/decode metrics on a 4090D: input_len 16 → prefill TTFT 295.9 ms, decode avg ITL 2.61 ms (~383 tok/s); input_len 32 → TTFT 3.2 ms, decode ITL 2.03 ms (~494 tok/s). Real-weight numbers require the official 456B checkpoint.
Sanity — test/bench/test_benchmark.py	N/A (skipped)	Accuracy harness (C-Eval/MMLU) needs external datasets and real model weights; with synthetic random weights accuracy is meaningless, so this is intentionally skipped rather than reported.
Service — inference_server.py + scripts/test_perf.py	Passed (workflow-level, synthetic small weights)	OpenAI-compatible server boots and serves MiniMax: GET /v1/models → minimax_text_01_small; POST /v1/chat/completions → 200 with a 16-token completion (finish_reason=length, usage 14/16). The full test_perf.py load test needs real weights to be meaningful.

Platforms tested:

• Single-request demo (examples/test_infer.py, synthetic small weights):
  • local WSL — RTX 3050 (4 GB), ~1.4 s for 32 tokens;
  • remote 4×4090D cluster — single GPU ~0.36 s, TP=2 ~0.50 s (same 32 tokens).
• Distributed/quantized numeric matrix: 4×4090D cluster (below).

Real-456B end-to-end tests remain blocked by the ~913 GB bf16 (or ~228 GB
MXFP4, needs 6×80GB+) official checkpoint — please tag a reviewer with access
to such resources if this must be completed in-repo.

Numeric correctness matrix (all PASS, synthetic small weights vs. HF
reference and vs. single-GPU baseline):

Check	Detail	Result
Prefill logits vs. HF	max abs diff ≈ 0.0044 (bf16)	PASS
Decode continuation vs. HF	max abs diff ≈ 0.0005	PASS
Distributed matrix (4×4090D)	bf16 + MXFP4 × prefill + decode; TP2×PP1 / TP2×PP2 / TP4×PP1	PASS (≤5e-4 vs. single GPU)
MXFP4 precision	cross-seed, and decode continuation	PASS
Real-config key audit	all 80 layers / 8243 safetensor keys classified (linear/MoE/keep)	PASS
Real-width distributed (4×4090D)	hidden 6144 / 64 heads / 8 KV / vocab 200064, 2 layers × 8 experts; TP2×PP1 / TP2×PP2 / TP4×PP1 vs. single GPU	PASS (≤1e-3)

Run Logs

Raw console output from the runs in the table above. The three single-request
runs (local / server single-GPU / server TP=2) produced bit-identical
sampled token sequences; each full 32-token response is shown below.

Single request — 4×4090D server, single GPU (total_time: 360.27 ms):

Using Paged KV Cache with num_blocks=512
Using Mamba cache with num_blocks=128, zero_state_index=0
LLMEngine initialized with model at test/models/minimax_text_01_small on device cuda, enable_graph=False
=================== start generate ====================
===Query===
<|im_start|>user
How are you<|im_end|>
<|im_start|>assistant
===Response===
<tok20007><tok27662><tok2927><tok27408><tok26317><tok16100><tok30952><tok14209><tok850><tok26317><tok16100><tok30952><tok28530><tok13160><tok27781><tok28746><tok13160><tok27781><tok28746><tok13160><tok27781><tok28746><tok13160><tok27781><tok28746><tok13160><tok27781><tok28746><tok13160><tok27781><tok28746><tok13160>
total_time: 360.27 ms

Single request — 4×4090D server, TP=2 distributed (total_time: 498.21 ms):

Using Paged KV Cache with num_blocks=512
Intra-node TP communicator established: node_rank=0, local_ranks=2
Using Mamba cache with num_blocks=128, zero_state_index=0
LLMEngine initialized with model at test/models/minimax_text_01_small on device cuda, enable_graph=False
=================== start generate ====================
===Query===
<|im_start|>user
How are you<|im_end|>
<|im_start|>assistant
===Response===
<tok20007><tok27662><tok2927><tok27408><tok26317><tok16100><tok30952><tok14209><tok850><tok26317><tok16100><tok30952><tok28530><tok13160><tok27781><tok28746><tok13160><tok27781><tok28746><tok13160><tok27781><tok28746><tok13160><tok27781><tok28746><tok13160><tok27781><tok28746><tok13160><tok27781><tok28746><tok13160>
total_time: 498.21 ms

Single request — local WSL, RTX 3050 (total_time: 1384.23 ms):

Using Paged KV Cache with num_blocks=512
Using Mamba cache with num_blocks=128, zero_state_index=0
LLMEngine initialized with model at test/models/minimax_text_01_small on device cuda, enable_graph=False
=================== start generate ====================
===Query===
<|im_start|>user
How are you<|im_end|>
<|im_start|>assistant
===Response===
<tok20007><tok27662><tok2927><tok27408><tok26317><tok16100><tok30952><tok14209><tok850><tok26317><tok16100><tok30952><tok28530><tok13160><tok27781><tok28746><tok13160><tok27781><tok28746><tok13160><tok27781><tok28746><tok13160><tok27781><tok28746><tok13160><tok27781><tok28746><tok13160><tok27781><tok28746><tok13160>
total_time: 1384.23 ms

Offline performance — examples/bench.py (4×4090D server, single GPU):

Processing : {'idx': 0, 'batch_size': 1, 'input_len': 16, 'output_len': 16}
=================== start generate ====================
Prefill TTFT: 295.91 ms  Throughput: 54.07 tok/s
Decode  Avg ITL: 2.61 ms   Throughput: 382.86 tok/s
total_time: 335.31 ms
Processing : {'idx': 1, 'batch_size': 1, 'input_len': 32, 'output_len': 16}
=================== start generate ====================
Prefill TTFT: 3.22 ms  Throughput: 9936.76 tok/s
Decode  Avg ITL: 2.03 ms   Throughput: 493.62 tok/s
total_time: 33.93 ms

Service — inference_server.py (OpenAI-compatible, curl against port
8011):

$ curl /v1/models
{"object":"list","data":[{"id":"minimax_text_01_small","object":"model","owned_by":"infinilm"}]}

$ curl /v1/chat/completions -d '{"model":"minimax_text_01_small","messages":[{"role":"user","content":"How are you"}],"max_tokens":16}'
{"id":"cmpl-1685a79189fc4ebdb3693d0bccd33e60","object":"chat.completion","created":1788459466,"model":"minimax_text_01_small","system_fingerprint":null,"choices":[{"index":0,"message":{"role":"assistant","content":"<tok20007><tok27662><tok2927><tok27408><tok26317><tok16100><tok30952><tok14209><tok850><tok26317><tok16100><tok30952><tok28530><tok13160><tok27781>"},"logprobs":null,"finish_reason":"length"}],"usage":{"prompt_tokens":14,"completion_tokens":16,"total_tokens":30}}

Reproduction / Deployment Command

Once the official checkpoint (or an MXFP4-converted copy) is available, the
model runs through the standard entry points on NVIDIA, e.g. distributed:

python examples/test_infer.py --model=/path/to/MiniMax-Text-01/ \
    --enable-paged-attn --disable-prefix-caching \
    --tensor-parallel-size=8 --pipeline-parallel-size=4 \
    --node-rank=[rank_id] --master-addr=[addr] --master-port=[port] \
    --max-new-tokens=256

Synthetic-weight reproduction used for this PR (single-request demo above) is
examples/test_infer.py --model <small-synthetic-checkpoint> --enable-paged-attn --disable-prefix-caching, which completes end-to-end on 1 GPU and under TP=2.

Benchmark / Performance Impact

Chunked Lightning prefill vs. the previous token-by-token recurrence
(single 4090D, small model, 3 runs min):

seq	token-by-token	chunked (64)	speed-up
2048	273 ms	42 ms	~6.3×

• Chunk-size curve: 32 → 26.5 ms, 64 → 24.6/49.1/96.3 ms, 128 → 22.1/43.3/87.8
  ms (seq 1024/2048/4096); default 64, upper bound 512, chunk 0/1 keeps
  the recurrent path. Chunked vs. recurrent outputs differ only by bf16
  accumulation order (max ≈1.2e-4).

Notes for Reviewers

• Non-invasive scope: model code under csrc/models/minimax_text_01/; the
  only shared-file touches are standard extension points (weight-remap entry,
  mamba-cache detection in infer_engine.py, new processor) — same pattern as
  qwen3_next/kimi_k3/videonsa.
• Config details: attention_bias read from config (official hard-codes
  false); partial RoPE (rotary_dim); postnorm residual via
  layernorm_*_alpha/beta.
• MXFP4: triggered by quant_method == "quark" (packed per-expert MoE +
  linear), reusing project MXFP4 ops; converter kept out-of-PR (no data files).
• Out of scope / follow-up: minimax_m1 (same architecture) registration;
  real-456B end-to-end + service once weights are available; the
  paged-compiler/graph path is not exercised for the Lightning state pool
  (mamba-cache) — the model runs eager (chunked prefill + recurrent decode).

CI / ChatOps

CI does not run automatically on pull requests. Trigger it manually (Actions →
CI, branch: this PR's head) or ask a maintainer to /retest. Formatting
check (scripts/format.py --check) passes for all changed files
(clang-format-21 / ruff).

---

Checklist

Title, Branch, and Commits

• PR title follows Conventional Commits — feat: support MiniMax-Text-01 (hybrid Lightning/full-attention MoE, MXFP4, TP/PP).
• Branch name follows <type>/xxx-... — feat/minimax-text-01.
• Commit messages follow Conventional Commits — single squashable feat:
  commit (17 files vs. upstream).
• No stray merge commits / fixup! / squash! / wip.

Scope and Design

• Changes are minimal (new model dir + the standard Python extension
  points only; exactly 17 files vs. upstream).

• No dead code / debug prints / unowned TODOs.

• No unrelated formatting churn (only the new dir's own formatting).

• N/A — no public API change (new files only; no signature/behavior change to
  existing callers).

General Code Hygiene

• Comments only where the why is non-obvious.
• Files end with a single trailing newline; no trailing whitespace/BOM.
• Identifiers in comments/errors wrapped in backticks.
• All comments and error messages in English, complete sentences
  (41 Chinese comment lines translated in this branch).

C++ Specific

• Google C++ Style; LLVM error-message wording.
• Initializer-list order matches member declaration order (verified: the
  only init-list initializes a single member; others assign in the body).
• No raw new/delete; RAII/smart pointers only.
• Changed files formatted by scripts/format.py (clang-format-21, --check
  passes).
• No changes/reference to csrc/models/llama_legacy/.

Python Specific

• PEP 8; complete English sentences; backticked references.
• Changed files formatted by scripts/format.py (ruff passes).
• No changes/reference to python/infinilm/auto_config.py.

Testing

• Single request (examples/test_infer.py) passed (synthetic weights;
  local + server + server TP=2, see table/logs).
• Offline performance (examples/bench.py) passed (workflow-level,
  synthetic weights; metrics in table/logs).
• Service (inference_server.py) passed (workflow-level OpenAI smoke).
• Sanity skipped — needs datasets + real weights (reason in table).
• A reviewer with real-456B-weight access tagged for the remaining
  end-to-end tests — author to tag one when opening the PR.

Build, CI, and Tooling

• Builds cleanly on NVIDIA (local final-code build; same code built & run
  on the 4×4090D cluster).
• CI triggered manually (Actions → CI) or /retest requested — to be done
  once the PR is open.

Documentation

• N/A — no README.md/CONTRIBUTING.md behavior or build-flag change outside
  the new model dir (the model list is not maintained in README).

Security and Safety

• No secrets/internal URLs/personal hardware identifiers committed.
• No third-party code added (reuses existing project operators only).
• No unsafe pointer arithmetic / uninitialized reads / missing bounds
  checks (only pointer casts are read-only index views).