File size: 5,638 Bytes
e92f76f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | #!/usr/bin/env python3
"""
build_mamba2.py — Build libmamba2.so from mamba2.cu
Run this on bbqbaddie (where nvcc lives):
python build_mamba2.py # auto-detect arch
python build_mamba2.py --arch sm_86 # RTX 3080
python build_mamba2.py --arch sm_89 # bbqbaddie RTX 5000 (Ada)
python build_mamba2.py --arch sm_80 # A100
Output: kernels/libmamba2.so
Then scp to dev machine or bundle with the package.
The .so exposes:
mamba2_step_fp8()
mamba2_forward_fp8()
mamba2_get_version()
Haskell links via:
ghc -L<kernels_dir> -lmamba2 -rpath <kernels_dir> BOB/Mamba2FFI.hs
"""
from __future__ import annotations
import argparse
import os
import subprocess
import sys
from pathlib import Path
KERNELS_DIR = Path(__file__).parent.resolve()
CUDA_SRC = KERNELS_DIR / "mamba2.cu"
OUT_SO = KERNELS_DIR / "libmamba2.so"
OUT_OBJ = KERNELS_DIR / "mamba2.o"
def detect_arch() -> str:
"""Detect GPU compute capability via torch."""
try:
import torch
if torch.cuda.is_available():
cap = torch.cuda.get_device_capability(0)
arch = f"sm_{cap[0]*10 + cap[1]}"
print(f"[build_mamba2] detected GPU arch: {arch}")
return arch
except ImportError:
pass
print("[build_mamba2] WARNING: torch not available, defaulting to sm_86")
return "sm_86"
def find_nvcc() -> str:
"""Return path to nvcc binary."""
# 1. On PATH
r = subprocess.run(["which", "nvcc"], capture_output=True, text=True)
if r.returncode == 0 and r.stdout.strip():
return r.stdout.strip()
# 2. Via torch CUDA_HOME
try:
from torch.utils.cpp_extension import CUDA_HOME
if CUDA_HOME:
candidate = Path(CUDA_HOME) / "bin" / "nvcc"
if candidate.exists():
return str(candidate)
except ImportError:
pass
# 3. Common Linux paths
for p in ["/usr/local/cuda/bin/nvcc", "/usr/bin/nvcc"]:
if Path(p).exists():
return p
raise FileNotFoundError(
"nvcc not found. Run this script on bbqbaddie where CUDA toolkit is installed.\n"
"On dev machine (no nvcc) use the pure-PyTorch fallback in mamba2_torch.py."
)
def build(arch: str, debug: bool = False):
nvcc = find_nvcc()
print(f"[build_mamba2] nvcc: {nvcc}")
print(f"[build_mamba2] arch: {arch}")
print(f"[build_mamba2] src: {CUDA_SRC}")
print(f"[build_mamba2] out: {OUT_SO}")
if not CUDA_SRC.exists():
raise FileNotFoundError(f"Source not found: {CUDA_SRC}")
opt_flags = ["-G", "-g"] if debug else ["-O3", "--use_fast_math"]
# Step 1: compile to relocatable device code object
compile_cmd = [
nvcc,
str(CUDA_SRC),
f"-arch={arch}",
"--compiler-options", "-fPIC",
"-dc", # device code compilation (relocatable)
"-o", str(OUT_OBJ),
*opt_flags,
"-I", str(KERNELS_DIR),
]
# Step 2: link into shared library
link_cmd = [
nvcc,
str(OUT_OBJ),
f"-arch={arch}",
"--shared",
"-o", str(OUT_SO),
*opt_flags,
]
print("\n[build_mamba2] Compiling...")
print(" ".join(compile_cmd))
r = subprocess.run(compile_cmd, capture_output=False)
if r.returncode != 0:
print("[build_mamba2] COMPILE FAILED")
sys.exit(r.returncode)
print("\n[build_mamba2] Linking...")
print(" ".join(link_cmd))
r = subprocess.run(link_cmd, capture_output=False)
if r.returncode != 0:
print("[build_mamba2] LINK FAILED")
sys.exit(r.returncode)
# Verify symbols
nm_r = subprocess.run(["nm", "-D", str(OUT_SO)], capture_output=True, text=True)
required_syms = ["mamba2_step_fp8", "mamba2_forward_fp8", "mamba2_get_version"]
missing = [s for s in required_syms if s not in nm_r.stdout]
if missing:
print(f"[build_mamba2] WARNING: missing symbols in .so: {missing}")
else:
print("[build_mamba2] All required symbols present.")
so_size = OUT_SO.stat().st_size
print(f"\n[build_mamba2] SUCCESS: {OUT_SO} ({so_size // 1024} KB)")
print("\nTo use from Python:")
print(f" import ctypes")
print(f" lib = ctypes.CDLL('{OUT_SO}')")
print(f" print(lib.mamba2_get_version().decode())")
print("\nTo link from Haskell:")
print(f" ghc -L{KERNELS_DIR} -lmamba2 -rpath {KERNELS_DIR} BOB/Mamba2FFI.hs")
def verify_so():
"""Quick sanity check: load the .so and call mamba2_get_version."""
if not OUT_SO.exists():
print(f"[verify] {OUT_SO} not found — run build first")
return False
import ctypes, ctypes.util
try:
lib = ctypes.CDLL(str(OUT_SO))
lib.mamba2_get_version.restype = ctypes.c_char_p
version = lib.mamba2_get_version().decode()
print(f"[verify] mamba2_get_version() = '{version}'")
return True
except OSError as e:
print(f"[verify] Failed to load {OUT_SO}: {e}")
return False
def main():
parser = argparse.ArgumentParser(description="Build libmamba2.so from mamba2.cu")
parser.add_argument("--arch", default=None, help="CUDA arch (e.g. sm_86, sm_89)")
parser.add_argument("--debug", action="store_true", help="Debug build (-G -g)")
parser.add_argument("--verify", action="store_true", help="Verify existing .so only")
args = parser.parse_args()
if args.verify:
ok = verify_so()
sys.exit(0 if ok else 1)
arch = args.arch or detect_arch()
build(arch, debug=args.debug)
verify_so()
if __name__ == "__main__":
main()
|