Bazel + cmake(): Building CUDA pybind11 Extensions Inside the Build Graph

A month ago I wrote about how we migrated mai_kernels to Bazel without Bazel knowing about CUDA. Docker compiled the .so files, Bazel managed the Python layer, and MagicMock papered over the gap on CPU. It worked. It shipped.

But it always felt like a trick. The native modules were invisible to Bazel. You couldn’t bazel build //mai_kernels:mai_kernels_cuda and get a .so out. Cache invalidation was Docker’s problem, not Bazel’s. And if someone changed a CUDA kernel, the feedback loop was “push, wait for Docker rebuild, pray.”

PR #30499 closes that gap. Bazel now builds all eight pybind11 CUDA .so files through rules_foreign_cc’s cmake() rule. Same CMakeLists.txt, same nvcc, same output. But now it’s in the build graph.

The trilogy

This is the third approach we’ve tried for native extensions in Bazel. Each one trades a different thing:

The sdist approach works great for packages with simple setup.py or pyproject.toml builds. mai_kernels has a 300-line CMakeLists.txt that fetches CUTLASS from GitHub, finds MPI, and conditionally compiles different kernel sets per architecture. You’re not running that through pip wheel in a sandbox.

The cmake() foreign rule is the right tool here. It runs the existing CMake build system inside Bazel’s sandbox, then captures the output artifacts. You don’t rewrite your build in Starlark. You wrap it.

How cmake() works

rules_foreign_cc provides a cmake() rule that runs a three-phase build inside the Bazel sandbox:

1. Configure — runs cmake with your cache entries and env vars
2. Build — runs cmake --build (which calls make or ninja)
3. Install — runs cmake --install into a Bazel-managed output directory

The rule declares the output shared libraries upfront (out_shared_libs), and Bazel treats them as build artifacts. Downstream targets can depend on them. Caching works. Invalidation is based on input file hashes, like everything else in Bazel.

Here’s the actual target from the PR:

cmake(
    name = "mai_kernels_cuda",
    cache_entries = {
        "BUILD_TESTING": "OFF",
        "CMAKE_CUDA_COMPILER_LAUNCHER": "ccache",
        "CMAKE_CXX_COMPILER_LAUNCHER": "ccache",
        "CMAKE_C_COMPILER_LAUNCHER": "ccache",
        # MPI bypass — FindMPI fails in sandbox
        "MPIEXEC_EXECUTABLE": "/opt/hpcx/ompi/bin/mpiexec",
        "MPI_CXX_COMPILER": "/opt/hpcx/ompi/bin/mpicxx",
        "MPI_CXX_HEADER_DIR": "/opt/hpcx/ompi/include",
        "MPI_CXX_LIB_NAMES": "mpi",
        "MPI_CXX_WORKS": "TRUE",
        "MPI_mpi_LIBRARY": "/opt/hpcx/ompi/lib/libmpi.so",
    },
    env = {
        "CCACHE_BASEDIR": "$$EXT_BUILD_ROOT$$",
        "CCACHE_COMPILERCHECK": "content",
        "CCACHE_DIR": "/mnt/ci/ccache/gpu",
        "CCACHE_IGNOREOPTIONS": "--time=*",
        "CCACHE_LOGFILE": "/mnt/ci/ccache/gpu/ccache.log",
        "CCACHE_NOHASHDIR": "1",
        "LD_LIBRARY_PATH": "/opt/hpcx/ompi/lib:...",
        "PATH": "/opt/hpcx/ompi/bin:/usr/local/cuda/bin:...",
    },
    generate_crosstool_file = False,
    lib_source = ":csrc_srcs",
    out_lib_dir = "mai_kernels",
    out_shared_libs = select({
        "@platforms//cpu:aarch64": _AARCH64_SHARED_LIBS,
        "//conditions:default": _X86_64_SHARED_LIBS,
    }),
    tags = ["gpu", "manual", "requires-network"],
    target_compatible_with = ["@platforms//os:linux"],
)

Lot going on. Let’s unpack the gotchas.

Gotcha 1: FindMPI dies in the sandbox

CMake’s FindMPI module works by running mpicc --showme:compile and parsing the output. Inside Bazel’s sandbox, the PATH is restricted. mpicc isn’t there. FindMPI fails, the configure step dies, and you get a cryptic error about MPI not being found even though it’s definitely installed in the Docker image.

The fix: bypass FindMPI entirely by setting every MPI variable as a CMake cache entry.

cache_entries = {
    "MPIEXEC_EXECUTABLE": "/opt/hpcx/ompi/bin/mpiexec",
    "MPI_CXX_COMPILER": "/opt/hpcx/ompi/bin/mpicxx",
    "MPI_CXX_HEADER_DIR": "/opt/hpcx/ompi/include",
    "MPI_CXX_LIB_NAMES": "mpi",
    "MPI_CXX_WORKS": "TRUE",
    "MPI_mpi_LIBRARY": "/opt/hpcx/ompi/lib/libmpi.so",
}

MPI_CXX_WORKS=TRUE is the key. It tells CMake “don’t bother testing whether MPI works, just trust me.” Without it, CMake tries to compile and run a test program, which also fails in the sandbox.

This is a general pattern with rules_foreign_cc: any CMake find_package() that probes the system will probably break. The sandbox is not the system. You either pre-populate the cache entries or restructure the CMakeLists.txt to accept paths directly.

Gotcha 2: ccache and sandbox paths

ccache hashes the compiler invocation to decide whether it’s a cache hit. One of the things it hashes is the file path. Bazel sandboxes use paths like /tmp/bazel-sandbox/12345/execroot/yolo/mai_kernels/csrc/foo.cu. The 12345 changes every build. Every ccache lookup misses. Every build is cold.

The fix: CCACHE_BASEDIR.

env = {
    "CCACHE_BASEDIR": "$$EXT_BUILD_ROOT$$",
}

$$EXT_BUILD_ROOT$$ is a rules_foreign_cc variable that expands to the sandbox root. ccache strips this prefix before hashing, so /tmp/bazel-sandbox/12345/execroot/yolo/foo.cu and /tmp/bazel-sandbox/67890/execroot/yolo/foo.cu hash the same.

Two more settings make this work:

• CCACHE_NOHASHDIR=1 — don’t hash the CWD at all (sandbox CWD also changes)
• CCACHE_COMPILERCHECK=content — hash the compiler binary content, not its path (the nvcc path is stable, but this is defensive)

Gotcha 3: nvcc –time and ccache

This one was subtle. The CMakeLists.txt had:

set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --time nvcc_profile.txt")

The space between --time and nvcc_profile.txt means CMake passes them as two separate arguments. nvcc handles this fine. But ccache sees --time as one option and nvcc_profile.txt as a positional argument, gets confused, and falls back to running nvcc directly (no caching).

The fix: use the equals form.

set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --time=nvcc_profile.txt")

One character. = instead of . The difference between ccache working and not working for every nvcc invocation.

We also set CCACHE_IGNOREOPTIONS="--time=*" so ccache doesn’t include the profiling flag in its hash. The --time output file changes per compilation unit, so including it would bust the cache for no reason.

Gotcha 4: generate_crosstool_file = False

By default, rules_foreign_cc generates a CMake toolchain file that tells CMake to use Bazel’s registered CC toolchain. This is great for plain C/C++ projects. It’s terrible for CUDA projects.

The generated toolchain file sets CMAKE_C_COMPILER and CMAKE_CXX_COMPILER but knows nothing about CMAKE_CUDA_COMPILER. CMake then tries to find nvcc on its own, fails (sandbox PATH), and you’re stuck.

Setting generate_crosstool_file = False tells the rule: don’t generate a toolchain file, let CMake use its own compiler discovery. Since we set PATH in env to include /usr/local/cuda/bin, CMake finds nvcc there. Not hermetic. But it works, and for CUDA, “works” is the right bar right now.

Gotcha 5: architecture-specific outputs

mai_kernels produces different .so files depending on architecture. x86_64 gets SM80/SM90a kernels. aarch64 gets SM100a/SM103a (Blackwell) too. The out_shared_libs attribute needs to list every .so file that the CMake install step produces, and it’s different per platform.

out_shared_libs = select({
    "@platforms//cpu:aarch64": _AARCH64_SHARED_LIBS,
    "//conditions:default": _X86_64_SHARED_LIBS,
}),

If you list a .so that CMake didn’t produce, Bazel fails with “expected output not found.” If you miss one, it’s silently dropped. Both are bad. The lists need to match exactly what CMake installs, per architecture.

The CMakeLists.txt changes

The existing CMakeLists.txt needed two small patches to work inside the Bazel sandbox.

CMake 3.31 compatibility. Newer CMake requires explicit enable_language(CUDA) before you can use CUDA features. The old file relied on implicit enablement through project(). One line fix.

enable_language(CUDA)
find_package(CUDAToolkit REQUIRED)

Test guards. The CMakeLists.txt built GTest and unit tests unconditionally. Inside Bazel, we don’t want CMake running its own tests (Bazel runs tests). Wrapping in if(BUILD_TESTING) and setting BUILD_TESTING=OFF in cache entries skips the test build entirely.

Tags and scheduling

The target is tagged manual, gpu, and requires-network.

• manual — excluded from bazel build //.... You don’t want every CPU CI job trying to build CUDA kernels.
• gpu — CI can filter: “only run targets tagged gpu on GPU runners.”
• requires-network — CMake fetches CUTLASS and MathDx via FetchContent. Sandbox needs network.

This means the target never runs accidentally. You either ask for it explicitly (bazel build //mai_kernels:mai_kernels_cuda) or CI selects it based on tags.

The cache stack: 47 min → 44s → 14s

Three tiers of caching, each building on the last.

ccache catches the compiler-level work. If the .cu source didn’t change, nvcc’s output is cached. This takes the build from 47 minutes (cold CUDA compilation of 8 targets across two SM architectures) down to 44 seconds (ccache hits for everything, CMake still runs configure + install).

Bazel disk cache catches the entire action. If none of the inputs changed (source files, cache entries, env vars), Bazel skips the cmake() action entirely and serves the output .so files from its own cache. 14 seconds.

The 44s → 14s gap is the overhead of CMake configure + ccache lookups + install. Not huge, but noticeable when you’re iterating on the BUILD file rather than the source.

When to use which approach

After three posts and three approaches, here’s how I’d decide:

Use the Docker/mock pattern when:

• The native code is stable and rarely changes
• You only need Bazel for the Python layer (testing, dependency management)
• You don’t want to deal with CUDA toolchains in Bazel at all
• The package has a mock/fallback pattern that lets Bazel work on CPU

Use the sdist/pip pattern when:

• The package has a standard pyproject.toml or setup.py
• The C extensions are simple (pybind11 wrapper, Cython, no CUDA)
• You want full hermeticity (hermetic CC toolchain, no system deps)

Use the cmake() foreign rule when:

• The package has a complex CMakeLists.txt you don’t want to rewrite
• CUDA is involved (no hermetic CUDA toolchain in Bazel yet)
• You want the native build in Bazel’s dependency graph (cache invalidation, CI scheduling)
• You’re willing to trade hermeticity for pragmatism (generate_crosstool_file = False, system CUDA)

The docker/mock pattern is the fastest to set up. The cmake() pattern gives you the most control. The sdist pattern is the most hermetic. Pick the one that matches what you actually need.

What’s still not great

Honesty section. Things I’m not thrilled about.

Not hermetic. The build depends on system CUDA, system MPI, system libtorch. If the Docker image changes compilers, the cache is invalid and you won’t know until the build fails. This is the same problem we had before, just with better caching on top.

FetchContent needs network. CMake downloads CUTLASS and MathDx at configure time. requires-network tag handles this in CI, but it means you can’t build this target in a fully offline sandbox. A future improvement would be to vendor these or pre-fetch them as Bazel deps.

Hardcoded paths. /opt/hpcx/ompi/, /usr/local/cuda/, /mnt/ci/ccache/. These are correct for our GPU Docker images. They’ll break anywhere else. If we ever change the base image, every path needs updating.

47-minute cold build. Caching makes this tolerable, but the first build after a cache bust is painful. That’s just the reality of compiling CUDA kernels for multiple SM architectures.

Wrapping up

rules_foreign_cc’s cmake() rule is the duct tape between “CMake builds this” and “Bazel needs to know about this.” You’re not rewriting your build system. You’re wrapping it. The gotchas are all about sandbox isolation: paths that move, tools that aren’t in PATH, compiler flags that interact badly with ccache.

For mai_kernels, the payoff is concrete. The build is in Bazel’s graph. Changes to .cu files invalidate the right targets. CI can cache the output and skip the build entirely when nothing changed. 47 minutes becomes 14 seconds.

If you’re curious about the layers underneath all of this (compilers, linkers, shared libraries, nvcc), the Native Build Tooling post covers the conceptual foundation. If you want to see the opposite approach (Docker does the work, Bazel doesn’t know), that’s the CUDA without CUDA post. This post is the middle ground.