setup.py

"""Build the flashboot Python extension (flashboot._C) via torch's CUDA toolchain.

    pip install -e .                        # editable, builds the extension in place
    python setup.py build_ext --inplace     # rebuild just the extension

Requires: a working CUDA toolkit (nvcc) and torch already installed (it provides the
CUDA build helpers and the pybind11 headers). The raw-ibverbs RDMA engine additionally
needs the libibverbs headers; FB_BUILD_RDMA (auto/1/0, default auto) probes for them
and drops the engine when they are absent (GB300 images ship without libibverbs-dev —
only the rdma transport is lost, fabric/ipc are unaffected).
"""

import os

from setuptools import find_packages, setup

from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CUDAExtension

HERE = os.path.abspath(os.path.dirname(__file__))


def _rdma_enabled() -> bool:
    """Whether to build the raw-ibverbs RDMA engine (rdma_read.cpp + -libverbs).

    FB_BUILD_RDMA=1 forces it on (build fails loudly if verbs.h is missing),
    FB_BUILD_RDMA=0 forces it off, and auto (the default) probes the standard
    include dirs for infiniband/verbs.h — so H100 images (libibverbs-dev
    installed) get the rdma transport and GB300 images (no verbs headers)
    still build, with transport auto-selection falling through to fabric/ipc."""
    requested = os.environ.get("FB_BUILD_RDMA", "auto").strip().lower()
    if requested in ("1", "true", "yes", "on"):
        return True
    if requested in ("0", "false", "no", "off"):
        return False
    if requested not in ("", "auto"):
        raise ValueError(f"FB_BUILD_RDMA={requested!r} is not supported: "
                         f"expected auto, 1 or 0")
    include_dirs = ["/usr/include", "/usr/local/include"]
    include_dirs += os.environ.get("CPATH", "").split(os.pathsep)
    include_dirs += os.environ.get("C_INCLUDE_PATH", "").split(os.pathsep)
    # rdma_read.cpp is C++, and conda/toolchain installs land verbs.h outside the
    # two standard prefixes — a probe miss silently ships a build without the rdma
    # transport, so cover every dir the compiler itself would search.
    include_dirs += os.environ.get("CPLUS_INCLUDE_PATH", "").split(os.pathsep)
    for prefix_var in ("CONDA_PREFIX", "PREFIX"):
        prefix = os.environ.get(prefix_var, "")
        if prefix:
            include_dirs.append(os.path.join(prefix, "include"))
    found = next((d for d in include_dirs
                  if d and os.path.exists(os.path.join(d, "infiniband", "verbs.h"))),
                 None)
    print(f"[flashboot setup] FB_BUILD_RDMA=auto -> "
          f"{'enabled (verbs.h in ' + found + ')' if found else 'DISABLED (infiniband/verbs.h not found; the rdma transport will be unavailable — set FB_BUILD_RDMA=1 to force)'}")
    return found is not None


_BUILD_RDMA = _rdma_enabled()

# The driver library (libcuda) is not installed on GPU-less build hosts; link against
# the toolkit stub there. -L only affects link time — at runtime the loader resolves
# libcuda.so.1 from the node's real driver.
_driver_stub_dirs = [os.path.join(CUDA_HOME, "lib64", "stubs")] if CUDA_HOME else []

SOURCES = [
    "csrc/python_bindings.cc",
    "csrc/device_arena.cu",  # the per-rank contiguous cudaMalloc weight arena
    # NUMA-pinned host buffer + parallel file read + multi-stream H2D (seed arena fill
    # and the weight_preload GPFS->tmpfs prestager).
    "csrc/pinned_stage.cpp",
    # CUDA-IPC / NVLink-fabric handle-import engine (flashboot._C.PeerArenaImporter):
    # the clone's same-node / same-IMEX-domain weight pull over NVLink.
    "csrc/peer_arena_import.cu",
    # NVLink-fabric / IMEX preflight (flashboot._C.check_imex) — lets transport
    # auto-selection prefer fabric only on IMEX-configured systems (GB300 NVL72).
    "csrc/imex_check.cc",
    # Pipelined chunked chain broadcast (flashboot._C.ChainReceiver), ported from
    # the GB300-verified main-branch RingReceiver: per-chunk flag-gated forwarding
    # so all chain hops stream concurrently.
    "csrc/chain_broadcast.cu",
]

if _BUILD_RDMA:
    # raw-ibverbs RDMA-READ engine — registered into flashboot._C via register_rdma(),
    # so it ships as flashboot._C.Endpoint (the cross-node rdma chain broadcast and
    # all-gather collectives; see python/flashboot/rdma_collectives.py).
    SOURCES.append("csrc/rdma_read.cpp")

# Default covers Hopper (sm_90, H100) and Blackwell (sm_100, GB200/GB300 — needs
# CUDA >= 12.8). Override with a semicolon-separated list, e.g.
# FLASHBOOT_CUDA_ARCH=9.0 or FLASHBOOT_CUDA_ARCH="10.0;10.3".
_arches = os.environ.get("FLASHBOOT_CUDA_ARCH", "9.0;10.0").split(";")
_nvcc_arch = []
for arch in _arches:
    arch = arch.strip().replace(".", "")
    if arch:
        _nvcc_arch += [f"-gencode=arch=compute_{arch},code=sm_{arch}"]

# -DFB_NO_RDMA tells python_bindings.cc to skip the register_rdma() declaration and
# call when rdma_read.cpp is out of the build.
_rdma_macros = [] if _BUILD_RDMA else ["-DFB_NO_RDMA"]

ext = CUDAExtension(
    name="flashboot._C",
    sources=SOURCES,
    include_dirs=[os.path.join(HERE, "include")],
    # libibverbs for the raw-ibverbs RDMA engine (rdma_read.cpp, when built); libcuda
    # for the driver VMM API (cuMem*) behind the fabric arena + PeerArenaImporter;
    # everything else uses only the CUDA runtime, which the torch extension links by
    # default.
    libraries=(["ibverbs"] if _BUILD_RDMA else []) + ["cuda"],
    library_dirs=_driver_stub_dirs,
    extra_compile_args={
        # -fvisibility=default is required: flashboot._C coexists with torch's own
        # pybind11 in one process, and hidden visibility broke cross-module pybind
        # type/GIL handling.
        "cxx": ["-O3", "-std=c++17", "-fvisibility=default", "-pthread"] + _rdma_macros,
        "nvcc": ["-O3", "-std=c++17"] + _rdma_macros + _nvcc_arch,
    },
)

ext_modules = [ext]

setup(
    name="flashboot",
    version="0.2.0",
    description="Fast sglang startup: sharded-checkpoint GPU arena loader + NVLink clone",
    packages=find_packages(where="python"),
    package_dir={"": "python"},
    ext_modules=ext_modules,
    cmdclass={"build_ext": BuildExtension},
    python_requires=">=3.9",
    install_requires=["torch>=2.3"],
    zip_safe=False,
)