flash-rt 0.2.0__py3-none-any.whl

flash_rt/__init__.py

@@ -0,0 +1,107 @@
+"""
+FlashRT — High-performance VLA inference engine.
+
+Public exports (stable API — see ``docs/stable_api.md``):
+
+    flash_rt.load_model(...)   → VLAModel
+    flash_rt.VLAModel          — unified inference wrapper
+
+Supported models: Pi0.5, Pi0, Pi0-FAST, GROOT N1.6, GROOT N1.7.
+Supported hardware: Jetson Thor (SM110), RTX 5090 (SM120), RTX 4090
+(SM89), AMD Instinct MI350 series (ROCm gfx950, pi05).
+
+Extending with new models: see ``docs/plugin_model_template.md``.
+
+Usage::
+
+    import flash_rt
+
+    model = flash_rt.load_model(
+        checkpoint="/path/to/checkpoint",
+        framework="torch",
+        autotune=3,
+    )
+
+    actions = model.predict(images=[base_img, wrist_img],
+                            prompt="pick up the red block")
+"""
+
+__version__ = "0.2.0"
+
+# ── Windows: register CUDA / cuDNN DLL search paths ──
+# Python 3.8+ on Windows ignores PATH for C-extension dependencies
+# (security hardening). The compiled .pyd needs cudart64_*.dll,
+# cublas64_*.dll, cublasLt, cudnn — we add their canonical install
+# directories to the secure DLL loader so `import flash_rt` works
+# without the user pre-loading them. Linux is unaffected: this whole
+# block is skipped via the sys.platform guard.
+import os as _os
+import sys as _sys
+if _sys.platform == 'win32':
+    _cuda_roots = [
+        _os.environ.get('CUDA_PATH'),
+        _os.environ.get('CUDA_PATH_V13_0'),
+        _os.environ.get('CUDA_PATH_V12_9'),
+        _os.environ.get('CUDA_PATH_V12_8'),
+        _os.environ.get('CUDA_PATH_V12_4'),
+        r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.0',
+        r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.9',
+        r'C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8',
+        _os.environ.get('CUDNN_PATH'),
+    ]
+    _seen = set()
+    for _root in filter(None, _cuda_roots):
+        for _sub in ('bin', r'extras\CUPTI\lib64', ''):
+            _p = _os.path.join(_root, _sub) if _sub else _root
+            if _p in _seen:
+                continue
+            _seen.add(_p)
+            if _os.path.isdir(_p):
+                try:
+                    _os.add_dll_directory(_p)
+                except (OSError, ValueError):
+                    pass
+    del _root, _sub, _p, _seen, _cuda_roots
+del _os, _sys
+
+from flash_rt import _extensions as _ext  # noqa: E402
+
+__all__ = ["load_model", "VLAModel", "catalog"]
+
+
+def __getattr__(name):
+    """PEP 562. ``import flash_rt`` stays free of torch and of the
+    compiled extensions.
+
+    The structure catalog is usable without either, and a consumer that
+    only wants ``flash_rt.catalog`` should not pay for the VLA API to
+    reach it. Naming ``load_model`` still loads everything it needs.
+
+    ``flash_rt.structures`` moved to its own distribution
+    (``flashrt-structures``); asking for it here answers with that
+    pointer rather than an AttributeError.
+
+    An extension name reaching here means the import machinery did not
+    find it beside the package — this distribution ships no ``.so`` — so
+    answer with the build instructions rather than an AttributeError.
+    """
+    if name in ("load_model", "VLAModel"):
+        from flash_rt import api
+        return getattr(api, name)
+    if name == "catalog":
+        import flash_rt.catalog as mod
+        return mod
+    if name == "structures":
+        raise ImportError(
+            "flash_rt.structures moved to the flashrt-structures "
+            "distribution: pip install flashrt-structures, then "
+            "`import flashrt_structures as structures`. The structure "
+            "catalog itself stayed here as flash_rt.catalog. "
+            "See https://github.com/flashrt-project/FlashRT-Structures")
+    if name in _ext.EXTENSIONS:
+        return _ext.require(name)
+    raise AttributeError("module %r has no attribute %r" % (__name__, name))
+
+
+def __dir__():
+    return sorted(set(globals()) | set(__all__))

flash_rt/_extensions.py

@@ -0,0 +1,119 @@
+"""What to say when the compiled half is not there.
+
+This distribution is pure Python. It carries the frontends, the structure
+catalog, the host adapters and the kernel sources, but no ``.so``: kernels
+either arrive through the kernel hub, which the structures layer talks to,
+or are built locally for the models actually being run. So the first thing
+a fresh install meets is an absent extension, and that moment has to carry
+its own instructions — a bare ``ModuleNotFoundError: No module named
+'flash_rt.flash_rt_kernels'`` reads as a broken package rather than as a
+step the user has not taken yet.
+
+The structures layer already answers this way: a refusal names its reason
+and the rung below it. This module gives the native path the same manners.
+"""
+
+from __future__ import annotations
+
+import importlib
+import importlib.util
+
+#: extension module → what it serves. ``flash_rt_kernels`` is the core
+#: library every native frontend needs; the rest are per-family.
+EXTENSIONS = {
+    "flash_rt_kernels": "core kernels — every native frontend needs these",
+    "flash_rt_fa2": "FlashAttention-2, the RTX attention path",
+    "flash_rt_qwen3_vl_kernels": "Qwen3-VL SM89 FP8 kernels",
+}
+
+#: model family → the CMake switch that builds its kernels and nothing
+#: else. Pair any of these with ``-DFLASHRT_SLIM_BUILD=ON``.
+BUILD_SWITCHES = {
+    "lingbot": "FLASHRT_ENABLE_LINGBOT",
+    "motus": "FLASHRT_ENABLE_MOTUS",
+    "qwen3_5_moe": "FLASHRT_ENABLE_QWEN35MOE",
+    "nexn2": "FLASHRT_ENABLE_QWEN35MOE",
+    "qwen3_vl": "FLASHRT_BUILD_QWEN3_VL",
+    "melband_roformer": "FLASHRT_ENABLE_MELBAND_ROFORMER",
+    "omnivoice": "FLASHRT_ENABLE_OMNIVOICE",
+    "audio_codebook": "FLASHRT_ENABLE_AUDIO_CODEBOOK",
+    "minimax_remover": "FLASHRT_ENABLE_MINIMAX_REMOVER",
+}
+
+_REPO = "https://github.com/flashrt-project/FlashRT"
+
+
+def present(name: str = "flash_rt_kernels") -> bool:
+    """True if the named extension is importable beside this package."""
+    try:
+        return importlib.util.find_spec("flash_rt." + name) is not None
+    except (ImportError, AttributeError, ValueError):
+        return False
+
+
+def missing() -> list[str]:
+    """The extensions this install does not have."""
+    return [n for n in EXTENSIONS if not present(n)]
+
+
+def build_command(config: str | None = None) -> str:
+    """The shortest build that produces what ``config`` needs."""
+    switch = BUILD_SWITCHES.get(config or "")
+    if switch:
+        return ("cmake -B build -S . -DFLASHRT_SLIM_BUILD=ON "
+                "-D%s=ON && cmake --build build -j$(nproc)" % switch)
+    return "cmake -B build -S . && cmake --build build -j$(nproc)"
+
+
+def require(name: str = "flash_rt_kernels", *, config: str | None = None):
+    """Return the extension module, or refuse with reason and next step.
+
+    Raises :class:`ImportError` so call sites that already guard imports
+    keep behaving as they did when the extension shipped in the wheel.
+    """
+    if present(name):
+        try:
+            return importlib.import_module("flash_rt." + name)
+        except ImportError as exc:
+            # Built, but not for this environment. Worth saying plainly:
+            # the fix is a rebuild, not a build, and the two look alike
+            # from the traceback alone.
+            raise ImportError(
+                "flash_rt.{name} is present but will not load here:\n"
+                "    {exc}\n"
+                "The extension is compiled against a specific Python, torch\n"
+                "and CUDA ABI. Rebuild it in the environment you are running:\n"
+                "    {build}".format(name=name, exc=exc,
+                                     build=build_command(config))) from exc
+    raise ImportError(
+        "flash_rt.{name} is not built ({what}).\n"
+        "\n"
+        "This distribution ships pure Python; the CUDA extensions are built\n"
+        "from the source tree, for the hardware and the models you run:\n"
+        "    git clone {repo} && cd FlashRT\n"
+        "    pip install -e .\n"
+        "    {build}\n"
+        "The editable install matters: the build writes the extensions into\n"
+        "the clone's own flash_rt/ directory, so that clone has to be what\n"
+        "your interpreter imports.\n"
+        "Add -DGPU_ARCH=120 for RTX 5090, 110 for Jetson Thor, 89 for RTX\n"
+        "4090; see the Build section of the README for the full table.\n"
+        "\n"
+        "To run without a local build, use the structures layer instead: it\n"
+        "obtains kernels from the kernel hub and refuses legibly when one is\n"
+        "unavailable.\n"
+        "    from flash_rt import structures".format(
+            name=name, what=EXTENSIONS.get(name, "compiled extension"),
+            repo=_REPO, build=build_command(config)))
+
+
+def report() -> str:
+    """One line per extension: present, or the switch that builds it."""
+    lines = []
+    for name, what in EXTENSIONS.items():
+        mark = "present" if present(name) else "absent"
+        lines.append("  %-28s %-8s  %s" % (name, mark, what))
+    if missing():
+        lines.append("")
+        lines.append("  build:  " + build_command())
+    return "\n".join(lines)

flash_rt/amd/__init__.py

@@ -0,0 +1,9 @@
+"""FlashRT AMD (ROCm/HIP) backend.
+
+Self-contained AMD tree: HIP runtime twins of core/cuda_buffer and
+core/cuda_graph, the flash_rt_amd_kernels extension (built by
+csrc/amd/CMakeLists.txt, dropped into this directory), and — as the
+port progresses — CDNA attention backends, pi05 pipeline, frontends.
+
+The NVIDIA package tree never imports from here and vice versa.
+"""

flash_rt/amd/core/hip_buffer.py

@@ -0,0 +1,176 @@
+"""FlashRT AMD — HipBuffer: hipMalloc/managed wrapper for engine-facing GPU buffers.
+
+Twin of flash_rt/core/cuda_buffer.py over libamdhip64. Memcpy-kind and
+attach-flag enums are numerically identical to CUDA's (verified at
+bring-up): H2D=1, D2H=2, D2D=3, hipMemAttachGlobal=1.
+"""
+
+import ctypes
+import logging
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+try:
+    _hip = ctypes.CDLL("libamdhip64.so")
+except OSError as exc:  # no ROCm runtime on this machine
+    # Surface the conventional "optional backend unavailable" signal so
+    # callers (flash_rt.api's hardware gate, test suites) can guard with
+    # a plain ImportError instead of a platform-specific OSError.
+    raise ImportError(
+        "the AMD backend requires the ROCm runtime (libamdhip64.so), "
+        f"which could not be loaded: {exc}") from exc
+
+
+def _configure_hip_signatures() -> None:
+    """Declare ctypes signatures — host pointers ≥2GiB truncate under the
+    default c_int argtype (see cuda_buffer.py download() docstring)."""
+    ptr_p = ctypes.POINTER(ctypes.c_void_p)
+    signatures = {
+        "hipMallocManaged": (
+            [ptr_p, ctypes.c_size_t, ctypes.c_uint], ctypes.c_int),
+        "hipMalloc": ([ptr_p, ctypes.c_size_t], ctypes.c_int),
+        "hipFree": ([ctypes.c_void_p], ctypes.c_int),
+        "hipMemcpy": (
+            [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t,
+             ctypes.c_int],
+            ctypes.c_int,
+        ),
+        "hipMemcpyAsync": (
+            [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t,
+             ctypes.c_int, ctypes.c_void_p],
+            ctypes.c_int,
+        ),
+        "hipMemset": (
+            [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t], ctypes.c_int),
+        "hipMemsetAsync": (
+            [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t,
+             ctypes.c_void_p],
+            ctypes.c_int,
+        ),
+        "hipDeviceSynchronize": ([], ctypes.c_int),
+        "hipStreamSynchronize": ([ctypes.c_void_p], ctypes.c_int),
+    }
+    for name, (argtypes, restype) in signatures.items():
+        fn = getattr(_hip, name)
+        fn.argtypes = argtypes
+        fn.restype = restype
+
+
+_configure_hip_signatures()
+
+
+def _check(ret, msg=""):
+    if ret != 0:
+        raise RuntimeError(f"HIP error {ret}: {msg}")
+
+
+class HipBuffer:
+    """GPU buffer — managed or device memory."""
+
+    def __init__(self, nbytes: int, managed: bool = True):
+        self._ptr = ctypes.c_void_p()
+        self._managed = managed
+        if managed:
+            _check(_hip.hipMallocManaged(ctypes.byref(self._ptr), nbytes, 1),
+                   "hipMallocManaged")
+        else:
+            _check(_hip.hipMalloc(ctypes.byref(self._ptr), nbytes), "hipMalloc")
+        self._nbytes = nbytes
+
+    @property
+    def ptr(self) -> ctypes.c_void_p:
+        return self._ptr
+
+    @property
+    def nbytes(self) -> int:
+        return self._nbytes
+
+    @classmethod
+    def from_numpy(cls, arr: np.ndarray) -> 'HipBuffer':
+        """Create device buffer, upload H2D (device memory for replay bandwidth)."""
+        arr = np.ascontiguousarray(arr)
+        buf = cls(arr.nbytes, managed=False)
+        _check(_hip.hipMemcpy(
+            buf._ptr, ctypes.c_void_p(arr.ctypes.data), arr.nbytes, 1), "H2D")
+        return buf
+
+    @classmethod
+    def from_numpy_managed(cls, arr: np.ndarray) -> 'HipBuffer':
+        """Create managed buffer, upload via memmove. Use for buffers needing D2H readback."""
+        arr = np.ascontiguousarray(arr)
+        buf = cls(arr.nbytes, managed=True)
+        ctypes.memmove(buf._ptr, arr.ctypes.data, arr.nbytes)
+        return buf
+
+    @classmethod
+    def zeros(cls, count: int, dtype, managed: bool = True) -> 'HipBuffer':
+        nbytes = count * np.dtype(dtype).itemsize
+        buf = cls(nbytes, managed=managed)
+        _check(_hip.hipMemset(buf._ptr, 0, nbytes), "hipMemset")
+        return buf
+
+    @classmethod
+    def empty(cls, count: int, dtype, managed: bool = True) -> 'HipBuffer':
+        return cls(count * np.dtype(dtype).itemsize, managed=managed)
+
+    @classmethod
+    def device_zeros(cls, count: int, dtype) -> 'HipBuffer':
+        return cls.zeros(count, dtype, managed=False)
+
+    @classmethod
+    def device_empty(cls, count: int, dtype) -> 'HipBuffer':
+        return cls.empty(count, dtype, managed=False)
+
+    def upload(self, arr: np.ndarray):
+        """Upload numpy → buffer."""
+        assert arr.nbytes <= self._nbytes
+        arr = np.ascontiguousarray(arr)
+        if self._managed:
+            ctypes.memmove(self._ptr, arr.ctypes.data, arr.nbytes)
+        else:
+            _check(_hip.hipMemcpy(
+                self._ptr, ctypes.c_void_p(arr.ctypes.data), arr.nbytes, 1), "H2D")
+
+    def download(self, arr: np.ndarray):
+        """Download buffer → numpy."""
+        assert arr.nbytes <= self._nbytes
+        _check(_hip.hipDeviceSynchronize(), "hipDeviceSynchronize")
+        if self._managed:
+            ctypes.memmove(arr.ctypes.data, self._ptr, arr.nbytes)
+        else:
+            _check(_hip.hipMemcpy(
+                ctypes.c_void_p(arr.ctypes.data),
+                self._ptr, arr.nbytes, 2), "D2H")
+
+    def download_new(self, shape, dtype) -> np.ndarray:
+        arr = np.empty(shape, dtype=dtype)
+        self.download(arr)
+        return arr
+
+    def zero_(self, stream=None):
+        if stream is not None:
+            _check(_hip.hipMemsetAsync(self._ptr, 0, self._nbytes, stream),
+                   "hipMemsetAsync")
+        else:
+            _check(_hip.hipMemset(self._ptr, 0, self._nbytes), "hipMemset")
+
+    def __del__(self):
+        try:
+            if _hip is not None and hasattr(self, '_ptr') and self._ptr.value:
+                ret = _hip.hipFree(self._ptr)
+                self._ptr = ctypes.c_void_p()
+                if ret != 0:
+                    # Raising in __del__ is unraisable; a failed free cannot
+                    # corrupt results, so log instead of _check here.
+                    logger.warning("hipFree failed with HIP error %d", ret)
+        except Exception:
+            pass
+
+    def __repr__(self):
+        t = "managed" if self._managed else "device"
+        return f"HipBuffer({self._nbytes}B, {t}, ptr=0x{self._ptr.value:x})"
+
+
+def sync():
+    _check(_hip.hipDeviceSynchronize(), "hipDeviceSynchronize")

flash_rt/amd/core/hip_graph.py

@@ -0,0 +1,102 @@
+"""FlashRT AMD — Framework-agnostic HIP Graph capture/replay.
+
+Twin of flash_rt/core/cuda_graph.py over libamdhip64. Same usage:
+
+    graph = HipGraph()
+    stream = graph.create_stream()
+    # warmup ... then:
+    graph.begin_capture(stream)
+    my_kernel(args..., stream)
+    graph.end_capture(stream)
+    graph.replay(stream)
+
+HIP-vs-CUDA deltas handled here:
+  - the 3-arg instantiate is hipGraphInstantiateWithFlags
+    (plain hipGraphInstantiate is the 5-arg errorNode/logBuffer form)
+  - capture-mode enum verified on hardware: hipStreamCaptureModeRelaxed == 2
+"""
+
+import ctypes
+import logging
+
+logger = logging.getLogger(__name__)
+
+try:
+    _hip = ctypes.CDLL("libamdhip64.so")
+except OSError as exc:  # no ROCm runtime on this machine
+    # Same contract as hip_buffer: report an unavailable optional
+    # backend as ImportError, not a platform-specific OSError.
+    raise ImportError(
+        "the AMD backend requires the ROCm runtime (libamdhip64.so), "
+        f"which could not be loaded: {exc}") from exc
+
+
+def _configure_signatures() -> None:
+    """Declare ctypes signatures — pointer args must never fall back to
+    the 32-bit c_int default (see cuda_buffer.py download() incident)."""
+    p = ctypes.c_void_p
+    pp = ctypes.POINTER(ctypes.c_void_p)
+    signatures = {
+        "hipStreamCreate": ([pp], ctypes.c_int),
+        "hipStreamBeginCapture": ([p, ctypes.c_uint], ctypes.c_int),
+        "hipStreamEndCapture": ([p, pp], ctypes.c_int),
+        "hipGraphInstantiateWithFlags": ([pp, p, ctypes.c_ulonglong], ctypes.c_int),
+        "hipGraphLaunch": ([p, p], ctypes.c_int),
+        "hipStreamSynchronize": ([p], ctypes.c_int),
+    }
+    for name, (argtypes, restype) in signatures.items():
+        fn = getattr(_hip, name)
+        fn.argtypes = argtypes
+        fn.restype = restype
+
+
+_configure_signatures()
+
+
+def _check(status, msg=""):
+    if status != 0:
+        raise RuntimeError(f"HIP error {status}: {msg}")
+
+
+class HipGraph:
+    """Framework-agnostic HIP Graph using raw HIP Runtime API."""
+
+    def __init__(self):
+        self._graph = ctypes.c_void_p()
+        self._graph_exec = ctypes.c_void_p()
+        self._captured = False
+
+    def create_stream(self) -> ctypes.c_void_p:
+        stream = ctypes.c_void_p()
+        _check(_hip.hipStreamCreate(ctypes.byref(stream)), "hipStreamCreate")
+        return stream
+
+    def begin_capture(self, stream: ctypes.c_void_p):
+        """Begin HIP Graph capture on the given stream.
+
+        hipStreamCaptureModeRelaxed=2: only capture ops on THIS stream,
+        same rationale as the CUDA path (don't block framework streams).
+        """
+        _check(_hip.hipStreamBeginCapture(stream, 2), "hipStreamBeginCapture")
+
+    def end_capture(self, stream: ctypes.c_void_p):
+        """End capture and instantiate the graph for replay."""
+        _check(_hip.hipStreamEndCapture(stream, ctypes.byref(self._graph)),
+               "hipStreamEndCapture")
+        _check(_hip.hipGraphInstantiateWithFlags(
+            ctypes.byref(self._graph_exec), self._graph, 0),
+               "hipGraphInstantiateWithFlags")
+        self._captured = True
+
+    def replay(self, stream: ctypes.c_void_p):
+        """Replay the captured graph (single CPU call → full GPU replay)."""
+        if not self._captured:
+            raise RuntimeError("No graph captured")
+        _check(_hip.hipGraphLaunch(self._graph_exec, stream), "hipGraphLaunch")
+
+    def sync(self, stream: ctypes.c_void_p):
+        _check(_hip.hipStreamSynchronize(stream), "hipStreamSynchronize")
+
+    @property
+    def captured(self) -> bool:
+        return self._captured

flash_rt/amd/frontends/__init__.py

@@ -0,0 +1 @@
+"""FlashRT AMD — framework frontends (weights + preprocessing)."""

flash_rt/amd/frontends/torch/__init__.py

@@ -0,0 +1 @@
+"""FlashRT AMD — torch (ROCm) frontends."""