Skip to content

feat(nvidia): add Mamba2 and RWKV5 adapters - #591

Open
lgeln10 wants to merge 1 commit into
InfiniTensor:mainfrom
lgeln10:feat/mamba2-rwkv5-nvidia-adapter
Open

lgeln10 wants to merge 1 commit into
InfiniTensor:mainfrom
lgeln10:feat/mamba2-rwkv5-nvidia-adapter

Conversation

@lgeln10

@lgeln10 lgeln10 commented Sep 20, 2026

Copy link
Copy Markdown

Summary

  • Add NVIDIA single-GPU Mamba2 and RWKV5 model adapters with recurrent state rows.
  • Add checkpoint conversion/configuration helpers and request-state isolation.
  • Keep pure recurrent scheduling separate from the existing attention KV BlockManager path.
  • Add unit, GPU correctness, real-model, and raw-token benchmark coverage.

Motivation

This PR adds two recurrent model architectures to InfiniLM for NVIDIA GPUs. Mamba2 and RWKV5 do not consume Transformer attention KV blocks; their per-request convolution/SSM or RWKV state is managed independently. The existing generic Transformer KV admission and reservation policy is unchanged. This work is intentionally separate from the KV admission performance work in PR 590.

Type of Change

  • feat - new model
  • test - model and scheduler coverage
  • docs - usage and benchmark notes

Test Results of Involved Models on Supported Platforms

  • NVIDIA GeForce RTX 4090 D, CUDA, BF16.
  • pytest -q test/models/rwkv5/test_scheduler.py test/models/mamba2/test_scheduler.py: 4 passed.
  • pytest -q test/models/rwkv5/test_adaptation.py test/models/mamba2/test_adaptation.py: 14 passed.
  • INFINILM_RUN_GPU_TESTS=1 pytest -q test/models/rwkv5/test_correctness.py test/models/mamba2/test_correctness.py: 2 passed.
  • Mamba2-130M real-model load/prefill/decode and raw-token benchmark were previously verified on the same GPU.
  • Ascend and other accelerator backends are not included in this PR.

Benchmark / Performance Impact

The included raw-token benchmark reports Mamba2 throughput on RTX 4090 D. The scheduler comparison is an architectural ablation for pure recurrent state versus a forced attention-KV path; it does not modify generic KV over-reservation.

Notes for Reviewers

  • Review the recurrent state indexing and request isolation in the Mamba2 and RWKV5 processors/model runners.
  • Mamba2 currently supports the tested NVIDIA configuration with one scan group; tensor parallelism and graph compilation remain disabled.
  • Official Mamba2 checkpoint preparation is documented because the released config omits several dimensions used by the runtime.

Checklist

  • Branch and commit use Conventional Commits naming.
  • No model weights or generated benchmark result files are included.
  • Local Python syntax, whitespace, and trailing-newline checks passed.
  • NVIDIA scheduler, adaptation, and GPU correctness tests passed.
  • CI has not yet been triggered.

@lgeln10
lgeln10 requested a review from a team September 20, 2026 15:50
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