| """Development loader: JIT-builds the kernel and exposes the packaged API. |
| |
| Binds torch.ops to the package's `_ops` slot and imports torch-ext/mamba3 |
| unchanged, so local runs exercise the code that ships. |
| """ |
|
|
| import importlib.util |
| import sys |
| import types |
| from pathlib import Path |
|
|
| import torch |
| from torch.utils.cpp_extension import load |
|
|
| _ROOT = Path(__file__).parent |
| _NAME = "mamba3_cuda" |
|
|
| load( |
| name=_NAME, |
| sources=[str(_ROOT / "torch-ext" / "torch_binding.cpp"), |
| str(_ROOT / "mamba3_cuda" / "mamba3_step.cu"), |
| str(_ROOT / "mamba3_cuda" / "mamba3_fwd.cu"), |
| str(_ROOT / "mamba3_cuda" / "mamba3_bwd.cu")], |
| extra_include_paths=[str(_ROOT / "torch-ext")], |
| |
| |
| extra_cuda_cflags=["-O3"], |
| is_python_module=False, |
| verbose=False, |
| ) |
|
|
| _ops_mod = types.ModuleType("_m3pkg._ops") |
| _ops_mod.ops = getattr(torch.ops, _NAME) |
| sys.modules["_m3pkg._ops"] = _ops_mod |
|
|
| _spec = importlib.util.spec_from_file_location( |
| "_m3pkg", _ROOT / "torch-ext" / "mamba3" / "__init__.py", |
| submodule_search_locations=[str(_ROOT / "torch-ext" / "mamba3")]) |
| _pkg = importlib.util.module_from_spec(_spec) |
| sys.modules["_m3pkg"] = _pkg |
| _spec.loader.exec_module(_pkg) |
|
|
| def ext_ops(): |
| return getattr(torch.ops, _NAME) |
|
|
|
|
| forward = _pkg.forward |
| DecodeState = _pkg.DecodeState |
| cumulative_angles = _pkg.cumulative_angles |
| chunk_decay = _pkg.chunk_decay |
| dispatch_paths = _pkg.dispatch_paths |
| __all__ = ["forward", "DecodeState", "cumulative_angles", "chunk_decay", |
| "dispatch_paths"] |
|
|
