flash-mamba 0.1.0__py3-none-any.whl

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/__main__.py

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

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/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/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/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.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,12 @@
+flash_mamba/__init__.py,sha256=rPxEC-31wRl9qKyIxvDPRXMYy_r7er2XDzzTREpPcJA,607
+flash_mamba/__main__.py,sha256=k1ocEWawweo1qCJWNFAAvyxz3tcY13dzvCenHszij30,48
+flash_mamba/_backend.py,sha256=7U30V_e4LvLSaI0GLvFpgvq26MIzsnuoY97ojhzYTro,1179
+flash_mamba/cli.py,sha256=M6h3nUK3pET1099KkBWpjQm1_2S8mODNqMHGc6COt9A,1494
+flash_mamba/diagnostics.py,sha256=VMjI6wWSzdKiNPQjk5pF2Cz-_m22rq_nWFFplY1hKIs,2666
+flash_mamba/factory.py,sha256=9lqsONABo8whLP8aICeIp8zQpSj6x_1hd7zFIeuhhzo,966
+flash_mamba-0.1.0.dist-info/licenses/LICENSE,sha256=XTpZyKa-FwcK2JNnho79ygfkWIRVLrDji8uJbmyQXrk,1081
+flash_mamba-0.1.0.dist-info/METADATA,sha256=Pe6eWOOAmX3fQzHn7uw6_H71i8SnRmaV-PahDw1Oe1U,6443
+flash_mamba-0.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+flash_mamba-0.1.0.dist-info/entry_points.txt,sha256=DoWzMne10GmUCsrvXV2qQ2NDLh9jBGgv5imPq_vDBeo,53
+flash_mamba-0.1.0.dist-info/top_level.txt,sha256=oIYSQV-R2UScF0N2qI48jmsb7pS8hqR0JtnpdDCFxjo,12
+flash_mamba-0.1.0.dist-info/RECORD,,

flash_mamba-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flash_mamba-0.1.0.dist-info/entry_points.txt

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

flash_mamba-0.1.0.dist-info/licenses/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.dist-info/top_level.txt

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