flash-mamba 0.1.0__tar.gz

flash_mamba-0.1.0/CHANGELOG.md

@@ -0,0 +1,9 @@
+# Changelog
+
+## 0.1.0
+
+- Add `create_mamba()` for official Mamba and Mamba2 instances.
+- Expose lazy `Mamba` / `Mamba2` imports and actionable backend import errors.
+- Add read-only environment diagnostics through Python and `flash-mamba doctor`.
+- Package a pure Python wheel and sdist; keep GPU backend installation separate.
+- Add wrapper tests, opt-in upstream GPU comparisons, CI, and release instructions.

flash_mamba-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Flash Mamba contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

flash_mamba-0.1.0/MANIFEST.in

@@ -0,0 +1,5 @@
+include CHANGELOG.md
+include VALIDATION.md
+recursive-include docs *.md
+recursive-include scripts *.py
+recursive-include tests *.py

flash_mamba-0.1.0/PKG-INFO

@@ -0,0 +1,170 @@
+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)
+引用相关论文。

flash_mamba-0.1.0/README.md

@@ -0,0 +1,146 @@
+# 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)
+引用相关论文。

flash_mamba-0.1.0/VALIDATION.md

@@ -0,0 +1,45 @@
+# v0.1.0 验证记录
+
+日期：2026-10-12（Asia/Shanghai）。以下是本机执行结果，不代表 GitHub CI 已运行。
+
+## 环境
+
+- Windows，conda `torch2`，Python 3.12.0。
+- PyTorch `2.13.0+cu130`，CUDA 13.0，NVIDIA GeForce RTX 2060。
+- 未安装 `mamba-ssm`、`causal-conv1d` 或 Triton。
+- 最初一次 PyTorch 导入遇到 WinError 4551；后续导入和 CUDA 设备查询成功。
+- 已安装 WSL2 Ubuntu-24.04，能识别 GPU；其 PATH 中没有 conda，未检测到常规
+  位置的 `torch2` 环境。没有为测试额外安装 Linux GPU 软件栈。
+
+## 已通过
+
+- `python -m ruff check .`。
+- `python -m ruff format --check .`。
+- `python -m pytest`：32 passed，2 skipped（显式 opt-in 的 GPU 测试）。
+- `python -m build`：生成 `flash_mamba-0.1.0-py3-none-any.whl` 和
+  `flash_mamba-0.1.0.tar.gz`；隔离构建环境不导入 PyTorch。
+- `python -m twine check --strict dist/*`：wheel / sdist 均通过。
+- `python scripts/verify_dist.py`：分别在全新虚拟环境安装 wheel，以及从 sdist
+  重建的 wheel；导入、版本一致性、console script、模块入口、缺少后端时的
+  JSON 诊断和退出码均通过。
+- `python -m flash_mamba doctor --json`：能够检查本机环境，按预期返回退出码 1，
+  报告官方库缺失和原生 Windows 平台限制。
+
+## 尚未验证
+
+真实 `mamba-ssm` 的前向/反向、官方 checkpoint 兼容性以及具体 GPU 内核兼容性。
+包装层单元测试使用替身验证接口边界，不能代替真实模型测试。
+
+在已配置官方后端的 Linux GPU 环境执行：
+
+```bash
+cd /mnt/e/codebase/flash_mamba
+python -m pip install -e ".[dev]"
+FLASH_MAMBA_RUN_GPU_TESTS=1 python -m pytest tests/test_gpu.py -v
+```
+
+运行前按 README 安装对应的 GPU 版 PyTorch 和官方后端，激活该 Linux 环境。
+若环境有 conda `torch2`，优先使用该环境。GPU 测试明确开启后不自动跳过依赖失败。
+
+尚未执行 TestPyPI/PyPI 上传、GitHub Actions、远程 Git push。发布说明在
+[docs/RELEASING.md](docs/RELEASING.md)。

flash_mamba-0.1.0/docs/RELEASING.md

@@ -0,0 +1,83 @@
+# 发布 flash-mamba 0.1.0
+
+源码和本地发行包已准备；上传 PyPI、创建远程仓库和 GPU 验证是独立步骤。
+没有配置远程仓库地址，所以元数据没有虚构的 Repository URL。
+创建仓库后可在 `pyproject.toml` 的 `[project.urls]` 中添加实际地址，并重新构建。
+
+## 本地检查
+
+在项目默认 conda 环境 `torch2` 中运行；Linux GPU 集成测试应使用安装了官方
+GPU 后端的对应环境。PowerShell 示例：
+
+```powershell
+conda run -n torch2 python -m pip install -e ".[dev]"
+conda run -n torch2 python -m ruff check .
+conda run -n torch2 python -m ruff format --check .
+conda run -n torch2 python -m pytest
+conda run -n torch2 python -m build
+conda run -n torch2 python -m twine check --strict dist/*
+conda run -n torch2 python scripts/verify_dist.py
+git diff --cached --check
+```
+
+真实模型兼容性尚需在 Linux GPU 上运行（Bash）：
+
+```bash
+FLASH_MAMBA_RUN_GPU_TESTS=1 python -m pytest tests/test_gpu.py -v
+```
+
+本地 Windows 的发行包验证通过不能代替此步骤。发布前应记录 GPU、驱动、
+PyTorch/CUDA、`mamba-ssm` 和 `causal-conv1d` 版本及真实测试结果。
+
+## Git 提交
+
+本地仓库准备好后可执行：
+
+```bash
+git add .
+git diff --cached --stat
+git commit -m "feat: add lightweight mamba-ssm wrapper and release packaging"
+```
+
+实际远程地址由维护者指定，之后再设置 `origin` 并 push。不要把 `dist/`、
+虚拟环境、`.pypirc` 或 PyPI token 提交到 Git。
+
+## TestPyPI / PyPI
+
+查询不到项目不代表名称一定可注册，发布前确认账户有 `flash-mamba` 名称权限。
+TestPyPI 和 PyPI 是独立账户/项目。先上传到 TestPyPI：
+
+```bash
+python -m twine upload --repository testpypi dist/flash_mamba-0.1.0*
+```
+
+在新的临时虚拟环境验证 TestPyPI 上的纯包装包，避免同时从 TestPyPI 解析依赖：
+
+```bash
+python -m pip install --index-url https://test.pypi.org/simple/ --no-deps flash-mamba==0.1.0
+python -c "import flash_mamba; print(flash_mamba.__version__)"
+flash-mamba --version
+```
+
+确认发行物和 GPU 验证结果后上传同一批文件到 PyPI：
+
+```bash
+python -m twine upload dist/flash_mamba-0.1.0*
+```
+
+使用 Trusted Publishing 或在本机认证提示中输入 token；不要写入仓库或命令示例。
+PyPI 版本不能覆盖。后续发布同时修改 `pyproject.toml` 与
+`src/flash_mamba/__init__.py` 的版本，更新 CHANGELOG，清理旧发行物再构建。
+
+## GitHub Trusted Publishing
+
+仓库提供 `.github/workflows/publish.yml`，只能手动触发。
+在目标 TestPyPI/PyPI 项目中配置 Trusted Publisher：填写真实 GitHub 用户/组织、
+仓库名、workflow 文件名 `publish.yml`，environment 分别为 `testpypi` / `pypi`。
+在 GitHub 同名 environment 设置保护规则；生产发布应要求人工审核。
+
+从已审查且完成 GPU 验证的提交运行 workflow：首次选 `testpypi` 验证，之后选
+`pypi`。workflow 自己执行 lint、包装层测试、构建、元数据检查及临时环境安装；
+它不会安装和测试 GPU 后端，不应把其成功当作模型兼容性的证据。
+
+GitHub 发布和本地上传二选一，以避免重复上传同一版本。

flash_mamba-0.1.0/pyproject.toml

@@ -0,0 +1,44 @@
+[build-system]
+requires = ["setuptools>=77", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "flash-mamba"
+version = "0.1.0"
+description = "A lightweight convenience wrapper for the official mamba-ssm package"
+readme = "README.md"
+requires-python = ">=3.10"
+license = "MIT"
+license-files = ["LICENSE"]
+authors = [{ name = "Jellyfish042" }]
+keywords = ["mamba", "mamba-ssm", "pytorch", "state-space-model"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3 :: Only",
+    "Operating System :: OS Independent",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = []
+
+[project.optional-dependencies]
+ssm = ["mamba-ssm>=2.3.2.post1,<3"]
+dev = ["pytest>=8", "build>=1.2", "twine>=6", "ruff>=0.11"]
+
+[project.scripts]
+flash-mamba = "flash_mamba.cli:main"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-ra --strict-markers"
+markers = ["gpu: real upstream CUDA forward/backward tests (explicit opt-in)"]
+
+[tool.ruff]
+target-version = "py310"
+line-length = 100
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP", "B"]

flash_mamba-0.1.0/scripts/verify_dist.py

@@ -0,0 +1,94 @@
+"""Install both distribution forms into fresh environments, without the GPU backend."""
+
+import json
+import os
+import subprocess
+import sys
+import tempfile
+import venv
+from pathlib import Path
+
+CHECK_IMPORT = """
+import importlib.abc
+import importlib.metadata
+import pathlib
+import sys
+class RejectBackend(importlib.abc.MetaPathFinder):
+    def find_spec(self, fullname, path=None, target=None):
+        if fullname.split('.')[0] in {'torch', 'mamba_ssm'}:
+            raise AssertionError('Unexpected backend import: ' + fullname)
+sys.meta_path.insert(0, RejectBackend())
+import flash_mamba
+assert flash_mamba.__version__ == importlib.metadata.version('flash-mamba')
+assert pathlib.Path(flash_mamba.__file__).is_relative_to(pathlib.Path(sys.prefix))
+assert callable(flash_mamba.create_mamba)
+print('Installed package import OK:', flash_mamba.__version__)
+"""
+
+
+def run(command: list[str], cwd: Path) -> None:
+    subprocess.run(command, cwd=cwd, check=True)
+
+
+def verify_wheel(wheel: Path, temporary: Path, label: str) -> None:
+    environment = temporary / label
+    venv.EnvBuilder(with_pip=True).create(environment)
+    executable = environment / ("Scripts/python.exe" if os.name == "nt" else "bin/python")
+    python = [str(executable), "-I"]
+    run([*python, "-m", "pip", "install", "--no-deps", "--no-index", str(wheel)], temporary)
+    run([*python, "-c", CHECK_IMPORT], temporary)
+    run([*python, "-m", "flash_mamba", "--version"], temporary)
+    console = environment / ("Scripts/flash-mamba.exe" if os.name == "nt" else "bin/flash-mamba")
+    run([str(console), "--version"], temporary)
+    doctor = subprocess.run(
+        [*python, "-m", "flash_mamba", "doctor", "--json"],
+        cwd=temporary,
+        text=True,
+        capture_output=True,
+        check=False,
+    )
+    if doctor.returncode != 1:
+        raise RuntimeError(
+            f"Expected missing-backend exit status 1: {doctor.stdout}\n{doctor.stderr}"
+        )
+    report = json.loads(doctor.stdout)
+    if report["torch_importable"] or report["backend_importable"]:
+        raise RuntimeError("Fresh environment unexpectedly has backend dependencies")
+    print(
+        f"{label}: wheel install, lazy import, CLI and missing-backend diagnostics OK", flush=True
+    )
+
+
+def main() -> None:
+    root = Path(__file__).resolve().parents[1]
+    wheels = list((root / "dist").glob("*.whl"))
+    sources = list((root / "dist").glob("*.tar.gz"))
+    if len(wheels) != 1 or len(sources) != 1:
+        raise RuntimeError("Build exactly one wheel and one sdist in dist/ before verification")
+    with tempfile.TemporaryDirectory(prefix="flash-mamba-dist-") as directory:
+        temporary = Path(directory)
+        verify_wheel(wheels[0], temporary, "wheel")
+        rebuilt = temporary / "rebuilt"
+        # Use the caller's build tools; no runtime backend dependency is resolved.
+        run(
+            [
+                sys.executable,
+                "-m",
+                "pip",
+                "wheel",
+                "--no-deps",
+                "--no-build-isolation",
+                "--wheel-dir",
+                str(rebuilt),
+                str(sources[0]),
+            ],
+            temporary,
+        )
+        built_wheels = list(rebuilt.glob("*.whl"))
+        if len(built_wheels) != 1:
+            raise RuntimeError("Expected exactly one wheel rebuilt from sdist")
+        verify_wheel(built_wheels[0], temporary, "sdist")
+
+
+if __name__ == "__main__":
+    main()

flash_mamba-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

flash_mamba-0.1.0/src/flash_mamba/__init__.py

@@ -0,0 +1,20 @@
+"""Convenience APIs for official Mamba blocks; importing this package is lightweight."""
+
+from typing import Any
+
+from ._backend import BackendUnavailableError, load_class
+from .diagnostics import diagnose
+from .factory import create_mamba
+
+__version__ = "0.1.0"
+__all__ = ["BackendUnavailableError", "Mamba", "Mamba2", "create_mamba", "diagnose"]
+
+
+def __getattr__(name: str) -> Any:
+    if name in {"Mamba", "Mamba2"}:
+        return load_class(name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(__all__))

flash_mamba-0.1.0/src/flash_mamba/__main__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+raise SystemExit(main())

flash_mamba-0.1.0/src/flash_mamba/_backend.py

@@ -0,0 +1,32 @@
+"""Load upstream only when a model is requested."""
+
+import importlib
+from typing import Any
+
+INSTALL_HINT = (
+    "On Linux, install a CUDA-enabled PyTorch first, then run "
+    '`python -m pip install "mamba-ssm[causal-conv1d]==2.3.2.post1" '
+    "--no-build-isolation`. For current upstream source builds, Mamba-1 also needs "
+    "MAMBA_KEEP_CUDA_BUILD=TRUE to enable selective_scan_cuda. "
+    "Use `flash-mamba doctor` to inspect the environment. "
+    "Native Windows/CPU execution is not provided by this wrapper."
+)
+
+
+class BackendUnavailableError(RuntimeError):
+    """The official backend or one of its native dependencies could not be loaded."""
+
+
+def load_class(name: str) -> Any:
+    try:
+        backend = importlib.import_module("mamba_ssm")
+    except (ImportError, OSError, RuntimeError) as exc:
+        raise BackendUnavailableError(
+            f"Cannot import official mamba_ssm ({type(exc).__name__}: {exc}). {INSTALL_HINT}"
+        ) from exc
+    try:
+        return getattr(backend, name)
+    except AttributeError as exc:
+        raise BackendUnavailableError(
+            f"Installed mamba_ssm does not expose {name}. {INSTALL_HINT}"
+        ) from exc

flash_mamba-0.1.0/src/flash_mamba/cli.py

@@ -0,0 +1,36 @@
+"""Command-line environment checks."""
+
+import argparse
+import json
+
+from . import __version__
+from .diagnostics import diagnose
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="flash-mamba")
+    parser.add_argument("--version", action="version", version=__version__)
+    commands = parser.add_subparsers(dest="command", required=True)
+    doctor = commands.add_parser(
+        "doctor", help="Inspect upstream dependencies and GPU availability"
+    )
+    doctor.add_argument("--json", action="store_true", help="Print machine-readable JSON")
+    args = parser.parse_args(argv)
+    report = diagnose()
+    if args.json:
+        print(json.dumps(report, ensure_ascii=True, indent=2))
+    else:
+        print(f"Python {report['python']} on {report['platform']}")
+        for package, version in report["packages"].items():
+            print(f"  {package}: {version or 'not installed'}")
+        print(f"CUDA/ROCm device available: {report['cuda_available']}")
+        print(f"Official backend importable: {report['backend_importable']}")
+        print(f"Ready for GPU smoke test: {report['ready_for_smoke_test']}")
+        for message in report["errors"]:
+            print(f"ERROR: {message}")
+        for message in report["warnings"]:
+            print(f"WARNING: {message}")
+        print(report["note"])
+        if not report["ready_for_smoke_test"]:
+            print(report["install_hint"])
+    return 0 if report["ready_for_smoke_test"] else 1

flash_mamba-0.1.0/src/flash_mamba/diagnostics.py

@@ -0,0 +1,65 @@
+"""Read-only diagnostics; these checks do not run kernels or install packages."""
+
+import importlib
+import platform
+from importlib import metadata
+from typing import Any
+
+from ._backend import INSTALL_HINT
+
+
+def diagnose() -> dict[str, Any]:
+    """Report imports and CUDA availability, without claiming kernel compatibility."""
+    report: dict[str, Any] = {
+        "python": platform.python_version(),
+        "platform": platform.system(),
+        "packages": {},
+        "torch_importable": False,
+        "backend_importable": False,
+        "cuda_available": False,
+        "cuda_version": None,
+        "hip_version": None,
+        "models": [],
+        "errors": [],
+        "warnings": [],
+        "ready_for_smoke_test": False,
+    }
+    for package in ("flash-mamba", "torch", "mamba-ssm", "causal-conv1d", "triton"):
+        try:
+            report["packages"][package] = metadata.version(package)
+        except metadata.PackageNotFoundError:
+            report["packages"][package] = None
+    if report["platform"] != "Linux":
+        report["warnings"].append("Use Linux/WSL2 for the upstream GPU backend.")
+    try:
+        torch = importlib.import_module("torch")
+        report["torch_importable"] = True
+        report["cuda_version"] = torch.version.cuda
+        report["hip_version"] = getattr(torch.version, "hip", None)
+        report["cuda_available"] = bool(torch.cuda.is_available())
+    except Exception as exc:
+        report["errors"].append(f"torch: {type(exc).__name__}: {exc}")
+    # Avoid importing a dependent library again if PyTorch already failed.
+    if report["torch_importable"]:
+        try:
+            backend = importlib.import_module("mamba_ssm")
+            report["models"] = [name for name in ("Mamba", "Mamba2") if hasattr(backend, name)]
+            missing = sorted({"Mamba", "Mamba2"} - set(report["models"]))
+            if missing:
+                report["errors"].append(f"mamba_ssm: missing exports: {', '.join(missing)}")
+            else:
+                report["backend_importable"] = True
+        except Exception as exc:
+            report["errors"].append(f"mamba_ssm: {type(exc).__name__}: {exc}")
+    if not report["cuda_available"]:
+        report["warnings"].append("PyTorch cannot access a CUDA/ROCm device.")
+    report["ready_for_smoke_test"] = (
+        report["platform"] == "Linux"
+        and report["torch_importable"]
+        and report["backend_importable"]
+        and report["cuda_available"]
+        and not report["errors"]
+    )
+    report["note"] = "Imports and device availability do not verify forward/backward kernels."
+    report["install_hint"] = INSTALL_HINT
+    return report

flash_mamba-0.1.0/src/flash_mamba/factory.py

@@ -0,0 +1,25 @@
+"""A small factory that preserves upstream model behavior and checkpoint keys."""
+
+from typing import Any, Literal
+
+from ._backend import load_class
+
+
+def create_mamba(
+    d_model: int,
+    *,
+    variant: Literal["mamba", "mamba2"] = "mamba2",
+    **kwargs: Any,
+) -> Any:
+    """Return an official Mamba/Mamba2 instance with kwargs passed through unchanged.
+
+    The default is Mamba2. Model execution follows upstream platform requirements.
+    Upstream controls defaults, device/dtype, dimension constraints, and validation
+    of all arguments other than d_model and variant.
+    """
+    if isinstance(d_model, bool) or not isinstance(d_model, int) or d_model <= 0:
+        raise ValueError("d_model must be a positive integer")
+    if variant not in ("mamba", "mamba2"):
+        raise ValueError("variant must be 'mamba' or 'mamba2'")
+    model_class = load_class("Mamba" if variant == "mamba" else "Mamba2")
+    return model_class(d_model=d_model, **kwargs)

flash_mamba-0.1.0/src/flash_mamba.egg-info/PKG-INFO

@@ -0,0 +1,170 @@
+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)
+引用相关论文。

flash_mamba-0.1.0/src/flash_mamba.egg-info/SOURCES.txt

@@ -0,0 +1,23 @@
+CHANGELOG.md
+LICENSE
+MANIFEST.in
+README.md
+VALIDATION.md
+pyproject.toml
+docs/RELEASING.md
+scripts/verify_dist.py
+src/flash_mamba/__init__.py
+src/flash_mamba/__main__.py
+src/flash_mamba/_backend.py
+src/flash_mamba/cli.py
+src/flash_mamba/diagnostics.py
+src/flash_mamba/factory.py
+src/flash_mamba.egg-info/PKG-INFO
+src/flash_mamba.egg-info/SOURCES.txt
+src/flash_mamba.egg-info/dependency_links.txt
+src/flash_mamba.egg-info/entry_points.txt
+src/flash_mamba.egg-info/requires.txt
+src/flash_mamba.egg-info/top_level.txt
+tests/test_diagnostics.py
+tests/test_factory.py
+tests/test_gpu.py

flash_mamba-0.1.0/src/flash_mamba.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

flash_mamba-0.1.0/src/flash_mamba.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+flash-mamba = flash_mamba.cli:main

flash_mamba-0.1.0/src/flash_mamba.egg-info/requires.txt

@@ -0,0 +1,9 @@
+
+[dev]
+pytest>=8
+build>=1.2
+twine>=6
+ruff>=0.11
+
+[ssm]
+mamba-ssm<3,>=2.3.2.post1

flash_mamba-0.1.0/src/flash_mamba.egg-info/top_level.txt

@@ -0,0 +1 @@
+flash_mamba

flash_mamba-0.1.0/tests/test_diagnostics.py

@@ -0,0 +1,110 @@
+import importlib
+import json
+from importlib import metadata
+from types import SimpleNamespace
+
+import pytest
+
+from flash_mamba import cli, diagnose
+from flash_mamba import diagnostics as module
+
+
+@pytest.fixture
+def healthy_environment(monkeypatch):
+    torch = SimpleNamespace(
+        version=SimpleNamespace(cuda="12.8", hip=None),
+        cuda=SimpleNamespace(is_available=lambda: True),
+    )
+    backend = SimpleNamespace(Mamba=object(), Mamba2=object())
+    imports = {"torch": torch, "mamba_ssm": backend}
+    monkeypatch.setattr(importlib, "import_module", lambda name: imports[name])
+    monkeypatch.setattr(module.platform, "system", lambda: "Linux")
+    monkeypatch.setattr(metadata, "version", lambda name: "1.0")
+    return imports
+
+
+def test_healthy_imports_and_cuda_only_qualify_for_smoke_test(healthy_environment):
+    report = diagnose()
+    assert report["ready_for_smoke_test"] is True
+    assert report["errors"] == []
+    assert report["models"] == ["Mamba", "Mamba2"]
+    assert report["cuda_version"] == "12.8"
+    assert "do not verify" in report["note"]
+    json.dumps(report)
+
+
+def test_missing_packages_and_blocked_torch_are_reported(monkeypatch):
+    imports = []
+
+    def missing_version(name):
+        raise metadata.PackageNotFoundError(name)
+
+    def blocked_import(name):
+        imports.append(name)
+        raise OSError("WinError 4551: c10.dll blocked")
+
+    monkeypatch.setattr(metadata, "version", missing_version)
+    monkeypatch.setattr(importlib, "import_module", blocked_import)
+    monkeypatch.setattr(module.platform, "system", lambda: "Windows")
+    report = diagnose()
+    assert imports == ["torch"]
+    assert report["ready_for_smoke_test"] is False
+    assert report["torch_importable"] is False
+    assert all(version is None for version in report["packages"].values())
+    assert "c10.dll blocked" in report["errors"][0]
+    assert any("WSL2" in warning for warning in report["warnings"])
+
+
+def test_backend_import_failure_is_reported(healthy_environment):
+    del healthy_environment["mamba_ssm"]
+    report = diagnose()
+    assert report["torch_importable"] is True
+    assert report["backend_importable"] is False
+    assert report["ready_for_smoke_test"] is False
+    assert "mamba_ssm" in report["errors"][0]
+
+
+def test_missing_class_is_not_considered_ready(healthy_environment):
+    del healthy_environment["mamba_ssm"].Mamba2
+    report = diagnose()
+    assert report["backend_importable"] is False
+    assert "missing exports: Mamba2" in report["errors"][0]
+
+
+def test_no_gpu_is_not_considered_ready(healthy_environment):
+    healthy_environment["torch"].cuda.is_available = lambda: False
+    report = diagnose()
+    assert report["backend_importable"] is True
+    assert report["ready_for_smoke_test"] is False
+
+
+def test_rocm_is_reported(healthy_environment):
+    healthy_environment["torch"].version.cuda = None
+    healthy_environment["torch"].version.hip = "6.4"
+    assert diagnose()["hip_version"] == "6.4"
+
+
+def test_cli_json_and_exit_status(healthy_environment, capsys):
+    assert cli.main(["doctor", "--json"]) == 0
+    report = json.loads(capsys.readouterr().out)
+    assert report["ready_for_smoke_test"] is True
+    healthy_environment["torch"].cuda.is_available = lambda: False
+    assert cli.main(["doctor", "--json"]) == 1
+    assert json.loads(capsys.readouterr().out)["ready_for_smoke_test"] is False
+
+
+def test_cli_readable_output(healthy_environment, capsys):
+    healthy_environment["torch"].cuda.is_available = lambda: False
+    assert cli.main(["doctor"]) == 1
+    output = capsys.readouterr().out
+    assert "Ready for GPU smoke test: False" in output
+    assert "--no-build-isolation" in output
+
+
+def test_cli_version(capsys):
+    with pytest.raises(SystemExit) as caught:
+        cli.main(["--version"])
+    assert caught.value.code == 0
+    from flash_mamba import __version__
+
+    assert capsys.readouterr().out.strip() == __version__

flash_mamba-0.1.0/tests/test_factory.py

@@ -0,0 +1,122 @@
+import importlib
+import subprocess
+import sys
+from types import SimpleNamespace
+
+import pytest
+
+import flash_mamba
+from flash_mamba import BackendUnavailableError, create_mamba
+
+
+def test_import_does_not_load_backend():
+    code = """
+import importlib.abc
+import sys
+class RejectBackend(importlib.abc.MetaPathFinder):
+    def find_spec(self, fullname, path=None, target=None):
+        if fullname.split('.')[0] in {'torch', 'mamba_ssm'}:
+            raise AssertionError('Unexpected backend import: ' + fullname)
+sys.meta_path.insert(0, RejectBackend())
+import flash_mamba
+assert callable(flash_mamba.create_mamba)
+assert 'Mamba2' in dir(flash_mamba)
+"""
+    subprocess.run([sys.executable, "-c", code], check=True)
+
+
+@pytest.fixture
+def backend(monkeypatch):
+    class Mamba:
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+
+    class Mamba2(Mamba):
+        pass
+
+    module = SimpleNamespace(Mamba=Mamba, Mamba2=Mamba2)
+    monkeypatch.setitem(sys.modules, "mamba_ssm", module)
+    return module
+
+
+@pytest.mark.parametrize("variant, export", [("mamba", "Mamba"), ("mamba2", "Mamba2")])
+def test_returns_upstream_instance_and_preserves_kwargs(backend, variant, export):
+    dtype = object()
+    process_group = object()
+    kwargs = {"device": "cuda:1", "dtype": dtype, "process_group": process_group, "layer_idx": 3}
+    model = create_mamba(128, variant=variant, **kwargs)
+    assert type(model) is getattr(backend, export)
+    assert model.kwargs == {"d_model": 128, **kwargs}
+
+
+def test_default_variant_keeps_upstream_defaults(backend):
+    model = create_mamba(128)
+    assert type(model) is backend.Mamba2
+    assert model.kwargs == {"d_model": 128}
+
+
+def test_lazy_exports_are_original_classes(backend):
+    assert flash_mamba.Mamba is backend.Mamba
+    assert flash_mamba.Mamba2 is backend.Mamba2
+
+
+@pytest.mark.parametrize("d_model", [0, -1, 3.5, True, None, "128"])
+def test_invalid_dimension_fails_before_backend_import(monkeypatch, d_model):
+    def unexpected_import(name):
+        raise AssertionError(f"Imported {name} for invalid arguments")
+
+    monkeypatch.setattr(importlib, "import_module", unexpected_import)
+    with pytest.raises(ValueError, match="positive integer"):
+        create_mamba(d_model)
+
+
+@pytest.mark.parametrize("variant", ["Mamba2", "mamba3", None, [], {}])
+def test_invalid_variant_fails_before_backend_import(monkeypatch, variant):
+    def unexpected_import(name):
+        raise AssertionError(f"Imported {name} for invalid arguments")
+
+    monkeypatch.setattr(importlib, "import_module", unexpected_import)
+    with pytest.raises(ValueError, match="variant"):
+        create_mamba(128, variant=variant)
+
+
+@pytest.mark.parametrize(
+    "error",
+    [
+        ModuleNotFoundError("mamba_ssm"),
+        ImportError("selective_scan_cuda"),
+        OSError("DLL blocked"),
+        RuntimeError("ABI mismatch"),
+    ],
+)
+def test_backend_errors_keep_original_cause(monkeypatch, error):
+    def broken_import(name):
+        raise error
+
+    monkeypatch.setattr(importlib, "import_module", broken_import)
+    with pytest.raises(BackendUnavailableError, match="flash-mamba doctor") as caught:
+        create_mamba(128)
+    assert caught.value.__cause__ is error
+
+
+def test_missing_export_has_actionable_error(monkeypatch):
+    monkeypatch.setitem(sys.modules, "mamba_ssm", SimpleNamespace())
+    with pytest.raises(BackendUnavailableError, match="does not expose Mamba2"):
+        create_mamba(128)
+
+
+def test_constructor_failure_is_not_misreported_as_install_problem(monkeypatch):
+    original = TypeError("unexpected keyword argument")
+
+    def broken_constructor(**kwargs):
+        raise original
+
+    monkeypatch.setitem(sys.modules, "mamba_ssm", SimpleNamespace(Mamba2=broken_constructor))
+    with pytest.raises(TypeError) as caught:
+        create_mamba(128, unknown=True)
+    assert caught.value is original
+
+
+def test_unknown_attribute_uses_normal_python_error():
+    with pytest.raises(AttributeError, match="no attribute"):
+        _ = flash_mamba.unavailable_attribute

flash_mamba-0.1.0/tests/test_gpu.py

@@ -0,0 +1,50 @@
+"""Real upstream comparisons: opt-in must fail if the GPU environment is broken."""
+
+import os
+
+import pytest
+
+pytestmark = [
+    pytest.mark.gpu,
+    pytest.mark.skipif(
+        os.environ.get("FLASH_MAMBA_RUN_GPU_TESTS") != "1",
+        reason="Set FLASH_MAMBA_RUN_GPU_TESTS=1 in a configured Linux GPU environment",
+    ),
+]
+
+
+@pytest.mark.parametrize("variant, class_name", [("mamba", "Mamba"), ("mamba2", "Mamba2")])
+def test_forward_backward_and_checkpoint_match_upstream(variant, class_name):
+    import mamba_ssm
+    import torch
+
+    from flash_mamba import create_mamba
+
+    assert torch.cuda.is_available(), "GPU test explicitly requested but CUDA/ROCm is unavailable"
+    kwargs = {"d_state": 64, "d_conv": 4, "expand": 2, "device": "cuda"}
+    if variant == "mamba2":
+        kwargs["headdim"] = 64
+    torch.manual_seed(42)
+    wrapped = create_mamba(128, variant=variant, **kwargs)
+    official = getattr(mamba_ssm, class_name)(d_model=128, **kwargs)
+    assert type(wrapped) is type(official)
+    assert wrapped.state_dict().keys() == official.state_dict().keys()
+    official.load_state_dict(wrapped.state_dict(), strict=True)
+    wrapped.load_state_dict(official.state_dict(), strict=True)
+    x1 = torch.randn(2, 64, 128, device="cuda", requires_grad=True)
+    x2 = x1.detach().clone().requires_grad_(True)
+    y1, y2 = wrapped(x1), official(x2)
+    assert y1.shape == x1.shape
+    assert torch.isfinite(y1).all()
+    torch.testing.assert_close(y1, y2, rtol=1e-4, atol=1e-5)
+    y1.square().mean().backward()
+    y2.square().mean().backward()
+    assert torch.isfinite(x1.grad).all()
+    torch.testing.assert_close(x1.grad, x2.grad, rtol=1e-4, atol=1e-5)
+    for (name1, param1), (name2, param2) in zip(
+        wrapped.named_parameters(), official.named_parameters(), strict=True
+    ):
+        assert name1 == name2
+        assert param1.grad is not None, f"Missing gradient: {name1}"
+        assert torch.isfinite(param1.grad).all(), f"Nonfinite gradient: {name1}"
+        torch.testing.assert_close(param1.grad, param2.grad, rtol=1e-4, atol=1e-5)