Metadata-Version: 2.4
Name: flash-mamba
Version: 0.1.0
Summary: A lightweight convenience wrapper for the official mamba-ssm package
Author: Jellyfish042
License-Expression: MIT
Keywords: mamba,mamba-ssm,pytorch,state-space-model
Classifier: Development Status :: 3 - Alpha
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Operating System :: OS Independent
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Provides-Extra: ssm
Requires-Dist: mamba-ssm<3,>=2.3.2.post1; extra == "ssm"
Provides-Extra: dev
Requires-Dist: pytest>=8; extra == "dev"
Requires-Dist: build>=1.2; extra == "dev"
Requires-Dist: twine>=6; extra == "dev"
Requires-Dist: ruff>=0.11; extra == "dev"
Dynamic: license-file

# flash-mamba

`flash-mamba` 是官方 [mamba-ssm](https://github.com/state-spaces/mamba) 的轻量包装库。
统一创建 Mamba / Mamba2，提供懒加载和环境诊断，返回官方模型实例。
训练、推理缓存、参数默认值、`state_dict` 和内核行为均由官方实现决定。

这是独立项目，与官方团队没有隶属关系。项目名称不代表新增加速内核或性能提升。
v0.1.0 的封装目标是 PyPI `mamba-ssm==2.3.2.post1`；Mamba3 不在本次范围内。
当前本机只验证了包装层和发行包，真实 GPU 验证需在 Linux 环境完成。

## 安装

包装层需要 Python 3.10+，自身没有运行时依赖，可在 Windows 导入和运行诊断。
发行包发布前，从本地安装：

```bash
python -m pip install .
```

发布到 PyPI 后才可使用：

```bash
python -m pip install flash-mamba
```

实际模型运行需要先在 **Linux（Windows 用户可使用 WSL2）** 安装适配设备的
PyTorch，再安装官方 GPU 后端。包装层不提供 CPU 替代实现或原生 Windows 内核。
CUDA/ROCm、编译工具和模型的具体要求以
[官方安装文档](https://github.com/state-spaces/mamba#installation) 为准。

```bash
# 先按 https://pytorch.org/get-started/locally/ 安装对应的 GPU 版 PyTorch
python -m pip install packaging ninja wheel setuptools
python -m pip install "mamba-ssm[causal-conv1d]==2.3.2.post1" --no-build-isolation
python -m pip install .
flash-mamba doctor
```

`ssm` extra 声明 `mamba-ssm>=2.3.2.post1,<3`。安装前仍需先配置 GPU 版 PyTorch
和官方构建依赖：

```bash
python -m pip install ".[ssm]" --no-build-isolation
```

显式两步安装更方便排查官方编译或二进制兼容问题。包装层不会自动安装、下载或
替换 PyTorch。官方 `main` 与 PyPI 发行版的安装默认值可能不同；当前 `main`
默认不构建 `selective_scan_cuda`，Mamba-1 使用该扩展时需要显式开启：

```bash
# 仅当选择官方最新源码时使用；并非 v0.1.0 已验证的发行版组合
MAMBA_KEEP_CUDA_BUILD=TRUE python -m pip install \
  "mamba-ssm[causal-conv1d] @ git+https://github.com/state-spaces/mamba.git" \
  --no-build-isolation
```

## 使用

```python
import torch
from flash_mamba import create_mamba

x = torch.randn(2, 64, 128, device="cuda")
model = create_mamba(
    d_model=128,
    variant="mamba2",  # 默认；Mamba-1 使用 "mamba"
    d_state=64,
    d_conv=4,
    expand=2,
    headdim=64,
    device="cuda",
)
y = model(x)
assert y.shape == x.shape
y.square().mean().backward()
```

`create_mamba` 只校验正整数 `d_model` 和 `variant`；其余关键字直接传给官方构造函数。
不额外添加残差、归一化、转置或模型包装。Mamba2 的 `d_ssm`（默认为
`expand * d_model`，分布式时还受 `world_size` 影响）必须满足官方 `headdim`
约束；例如 `d_model=16, expand=2, headdim=64` 不合法。

也可以直接导入官方类，第一次访问时才加载 `mamba_ssm`：

```python
from flash_mamba import Mamba, Mamba2

model = Mamba(d_model=128, device="cuda")
```

返回的实例可直接调用 `.to()`、`.parameters()`、`.state_dict()`、
`.allocate_inference_cache()` 和官方支持的 forward 关键字。
`import flash_mamba` 本身不导入 PyTorch 或 `mamba_ssm`。
后端导入失败会抛出 `BackendUnavailableError`，保留原异常链；构造函数参数错误
和模型执行错误直接沿用官方异常。

## 环境诊断

```bash
flash-mamba doctor
flash-mamba doctor --json
python -m flash_mamba doctor --json
```

检查已安装版本、PyTorch/官方后端能否导入、CUDA/ROCm 是否可访问。
退出码 0 表示可以尝试 GPU smoke test；退出码 1 表示检查未通过。
成功导入不等于内核可运行，真实前向/反向需执行下面的测试。
`flash_mamba.diagnose()` 返回相同的字典结果，不修改环境。

## 开发与验证

```bash
python -m pip install -e ".[dev]"
python -m ruff check .
python -m ruff format --check .
python -m pytest
python -m build
python -m twine check --strict dist/*
python scripts/verify_dist.py
```

仓库约定：没有其他环境要求时，以上 Python 命令在 conda `torch2` 中运行。
默认测试不依赖官方后端，检查懒加载、参数传递、异常处理、诊断和 CLI。
发行包验证在临时虚拟环境中安装 wheel 以及从 sdist 构建的 wheel。

在已配置官方库的 Linux GPU 环境中，必须执行真实测试再确认模型可运行：

```bash
FLASH_MAMBA_RUN_GPU_TESTS=1 python -m pytest tests/test_gpu.py -v
```

显式开启后，环境缺失或模型运行失败会导致测试失败，不会静默跳过。
测试比较工厂实例与官方实例的前向结果、参数/输入梯度，并检查 checkpoint
可直接互载。未开启时，GPU 测试明确标记为 skipped。

## 发布

发布准备、TestPyPI/PyPI 命令和 GitHub Trusted Publishing 配置见
[发布说明](docs/RELEASING.md)。CI 检查包装层并构建发行包，不能替代真实 GPU 测试。

## 许可与来源

本项目的新包装代码使用 MIT 许可证；不复制、不重新分发官方源码或 CUDA 内核。
`mamba-ssm` 由 Albert Gu、Tri Dao 等作者开发，使用 Apache-2.0 许可证。
使用该模型进行研究时，请按[官方说明](https://github.com/state-spaces/mamba#citation)
引用相关论文。
