Bazel Migration: A CUDA Package Without Bazel Knowing About CUDA

Everyone on the team agreed: mai_kernels can’t be migrated to Bazel. It’s a CUDA/C++ extension package. Nine pybind11 native modules. CMake build system. nvcc. The whole nightmare. Bazel doesn’t know how to build that, and teaching it would be – to use a technical term – insane.

They were right about the last part. But wrong about the conclusion.

The Package

mai_kernels is the kind of package that shows up in dependency graphs like a load-bearing wall. You can’t remove it. You can’t go around it. And until last week, you couldn’t migrate it either.

Here’s what’s inside:

The native code lives in csrc/, organized into nine directories that each produce a shared library:

The build system is scikit-build-core -> CMake -> nvcc. Not setup.py. The CMakeLists.txt alone pulls three external dependencies at configure time:

# CUTLASS v4.3.5 -- used by cutlass kernels.
FetchContent_Declare(
    cutlass_4_3_5
    GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git
    GIT_TAG        v4.3.5
)

# MathDx 25.12.1 -- used by conversion kernels
FetchContent_Declare(
    mathdx_25_12_1_cuda13
    URL https://developer.nvidia.com/downloads/compute/cuRANDDx/redist/...
)

Plus pybind11, libtorch C++ headers, MPI, GTest, and the CUDA toolkit itself. This is not “pip install and hope.” This is a multi-stage compilation pipeline that takes 5-10 minutes with a warm ccache and considerably longer without.

And then there’s the architecture split. On x86_64, you get 7 modules targeting SM80 and SM90a. On aarch64, you get 8 modules including SM100a and SM103a (Blackwell). Different architectures, different SM targets, different .so files. The CMake build reads CMAKE_SYSTEM_PROCESSOR and conditionally includes or excludes entire kernel directories.

The Python layer depends on torch (obviously), numpy, plus exotics: nvidia-cutlass-dsl, apache-tvm-ffi, torch-c-dlpack-ext, quack-kernels, nvidia-cudnn-cu12, nvidia-cudnn-frontend. Some of these are Linux-only (sys_platform == 'linux'). All of them come preinstalled in the Docker image.

This package has 10 direct downstream dependents: mai_distributed, mai_layers, mai_trainer, mai_job, mai_evaluator, mai_multimodal, rocket, sgl_server, simplertransformer, torchlab. Every one of those is either a core training package or a serving stack. Blocking mai_kernels means blocking the center of the monorepo’s dependency graph.

The Mock Pattern

Here’s the thing that makes the whole approach possible. Open mai_kernels/__init__.py:

from mai_nano.mocks.magic_mock_packages import magic_mock_packages

missing_packages = magic_mock_packages(
    "mai_kernels.conversion_kernels",
    "mai_kernels.cublaslt_block_scaled_gemm_kernels",
    "mai_kernels.imha_kernels",
    "mai_kernels.imha_train_kernels",
    "mai_kernels.cutlass_gemm_kernels",
    "mai_kernels.cutlass_group_gemm_kernels",
    "mai_kernels.cutlass_group_gemm_sm100_kernels",
    "mai_kernels.de",
    "mai_kernels.dequant_group_gemm_kernels",
    "triton.compiler",
    "triton.language",
    # ... and more
)

All nine .so modules plus triton. When the native libraries aren’t available (CPU machine, wrong architecture, no CUDA toolkit), they get mocked with MagicMock. The package still imports. Type checking works. Dependency resolution works. You just can’t run anything that calls into the native code.

This is the critical insight. The mock pattern means Bazel can manage mai_kernels as a py_library on CPU – builds, type-checks, validates dependencies – and only needs actual GPU hardware for tests.

The Insight

Here’s the thing nobody stopped to ask: does Bazel actually need to compile the C++?

No. It doesn’t.

The pretrain_base Docker image already builds all nine .so modules during image build via uv sync + CMake/nvcc. Those shared libraries exist at runtime inside the container. Bazel only needs to manage the Python source layer – the .py files, the imports, the test targets. The native modules appear at runtime through include_system_site_packages=True, same trick we used for system torch.

Don’t fight the build system boundary. Lean on Docker.

What I Built

PR #23478. Four pieces.

1. py_gpu_test macro

New macro in bazel/rules/defs.bzl, parallel to the existing py_torch_test:

py_gpu_test(
    name = "test_something",
    srcs = ["test_something.py"],
    deps = ["//mai_kernels"],
)

Under the hood it wraps py_venv_test with include_system_site_packages=True, auto-adds ["gpu", "torch"] tags, restricts to Linux via target_compatible_with, and defaults to size = "large" because GPU tests are heavyweight. Uses the same pytest_runner.py pattern as everything else.

2. mai_kernels/BUILD.bazel

Manually maintained (gazelle:ignore – Gazelle has no idea what to do with a csrc/ directory full of CUDA). The whole thing is surprisingly small:

py_library(
    name = "mai_kernels",
    srcs = glob(["src/mai_kernels/**/*.py"], exclude = ["src/mai_kernels/**/tests/**"]),
    imports = ["src"],
    deps = ["//lib/mai_nano/src/mai_nano/mocks", "@pypi//numpy"],
)

That’s it. Two deps. numpy and the mock utilities. Torch, triton, CUDA libs, the nine native .so files – all come from system site-packages in the Docker container. Bazel doesn’t know they exist. Doesn’t need to.

The .bazelignore was updated to ignore only mai_kernels/csrc instead of all of mai_kernels. Bazel sees 196 Python files. Ignores 81 C++/CUDA/header files. Everyone’s happy.

3. Smoke tests

Four py_gpu_test targets: test_numerics_config, test_packing, test_padding, test_sampling. Conservative start out of 78 total test files.

Why only 4? Because the full mai_kernels test suite includes tests that need multi-GPU setups, specific model weights, or 10+ minutes of runtime. Start with smoke tests that validate “can we even run py_gpu_test in this pipeline at all,” then expand once the infrastructure is proven.

4. CI: bazel-gpu-docker job

New job in bazel_build_test.yml. Runs on gpu-condor2-rss-dind runners using the pretrain_base image with Bazelisk + uv installed at runtime. The tag filter matrix now looks like:

Clean separation. Each job runs exactly what it should, nothing more.

The ccache Story

One detail worth calling out: the pretrain_base image already optimizes mai_kernels compilation with ccache. From pyproject.toml:

[[tool.scikit-build.overrides]]
if.env.UV_PROJECT_ENVIRONMENT = "/venv"
build-dir = "/root/.cache/skbuild/mai_kernels"
build.tool-args = ["-j32"]
cmake.args = [
  "-DPython_EXECUTABLE=/venv/bin/python3",
  "-DCMAKE_C_COMPILER_LAUNCHER=ccache",
  "-DCMAKE_CXX_COMPILER_LAUNCHER=ccache",
  "-DCMAKE_CUDA_COMPILER_LAUNCHER=ccache",
]

When the Docker build cache is warm, mai_kernels compilation drops from ~10 minutes to ~1 minute. The build directory is cached at /root/.cache/skbuild/mai_kernels, and all three compilers (C, C++, CUDA) go through ccache. This is why re-using the existing Docker build pipeline is so much better than trying to replicate it in Bazel – the optimization work is already done. Duplicating it would mean rebuilding the ccache integration, the build directory persistence, and the parallel -j32 tuning from scratch.

The Bumps

I pushed the PR expecting the usual “one CI run to shake out the typos” situation. It took three rounds. Each one taught me something I should have known.

Round 1: actionlint says no

First push. Immediate failure. Not even a runtime failure – a linting failure.

hashFiles() is not allowed in jobs.<job_id>.container.image context. I had the image tag computed inline using hashFiles('Dockerfile', ...) right in the container: block. Turns out GitHub Actions has strict rules about which expressions are available where. hashFiles() works in step-level run: and if: contexts. It does not work in job-level container.image. The error message is clear once you see it. The documentation is not.

Fix: split into a separate bazel-gpu-image-tag job that runs on ubuntu-latest, computes the hash in a step, and passes it via outputs to the bazel-gpu-docker job. More YAML. Less elegance. Works.

The lesson is boring but real: GitHub Actions context availability rules are surprisingly strict. What works in steps.run doesn’t work in container.image. You find this out when actionlint yells at you, not when you read the docs.

Round 2: someone else’s broken windows

Second push. Different failure. Not my code this time.

Gazelle freshness check failed. Someone had merged splits.py and test_magic_mock_packages.py into lib/mai_nano on main without updating the BUILD files. The CI runs bazel run //bazel:gazelle and diffs – if anything changed, you fail. Doesn’t matter whose fault it is. Your PR, your problem.

Had to add splits.py to mai_nano/data/BUILD.bazel and a test_magic_mock_packages target to mai_nano/tests/BUILD.bazel. Three lines of Starlark to fix someone else’s oversight.

This is monorepo life. Shared CI means you inherit other people’s broken windows. You can be mad about it, or you can fix it in 30 seconds and move on. I was mad about it for about 45 seconds, then fixed it.

Round 3: the venv/system Python split

Third push. The real one.

All four py_gpu_test targets failed with the same error: ModuleNotFoundError: No module named 'torch'.

This was confusing for about five minutes. The pretrain_base image has torch. The bazel-torch-docker job (which uses a different image) can import torch fine. Why can’t the GPU job find it?

Here’s what’s happening. The pretrain_base Docker image installs all its packages into a venv at /venv. The VIRTUAL_ENV=/venv env var is set, PATH includes /venv/bin. Standard Docker Python setup. But Bazel’s local_runtime_repo rule registers /usr/local/bin/python3 – the system Python. And include_system_site_packages=True exposes the system Python’s site-packages. Not the venv’s.

In the bazel-torch-docker job this never came up because that image’s Dockerfile installs torch via pip install --break-system-packages directly into the system Python’s site-packages. No venv indirection. pretrain_base uses a venv. Same concept, different plumbing, totally different result.

The fix is a .pth file – one of Python’s oldest and least-known tricks:

- name: Expose venv packages to system Python
  run: |
    SYSTEM_SITE=$(/usr/local/bin/python3 -c "import site; print(site.getsitepackages()[0])")
    VENV_SITE=$(/venv/bin/python3 -c "import site; print(site.getsitepackages()[0])")
    echo "$VENV_SITE" > "$SYSTEM_SITE/venv_packages.pth"
    /usr/local/bin/python3 -c "import torch; print(f'torch {torch.__version__} visible to system Python')"

Python’s site module reads .pth files from site-packages directories at startup and adds each line as an additional search path. By writing the venv’s site-packages path into a .pth file in the system site-packages, every venv package becomes visible to the system Python. No copying. No symlinking. One file, one line, two Python environments bridged.

The lesson here cuts deeper than the fix. “System site-packages” means exactly what it says – the system Python’s site-packages. Not whatever virtualenv the container happens to have activated. When you’re bridging two Python environments – Bazel’s toolchain pointing at system Python, Docker’s venv holding all the actual packages – you need to understand which Python is which. The error message says torch isn’t installed. What it means is “torch isn’t installed where I’m looking.”

Current CI Status

After three rounds:

All Bazel-related checks green. The st-tests failure is an azcopy download timeout in the SimplerTransformer Docker build – nothing to do with this PR. Ryan approved it after a round of questions about ACR probing, DinD network race conditions, and whether to use setup-uv-wrapper (yes). All four review threads resolved.

What This Unblocks

The monorepo has 104 workspace packages. Before this work:

mai_kernels was listed under “Blocked by Non-Python Code” alongside lib/bus, simplertransformer, and caas_user_images.

With mai_kernels migrated, its 10 direct dependents get closer to migration:

Most of these were already torch-blocked, and torch is already solved. With mai_kernels also solved, the remaining blocker for these packages is just the torch migration wave – a known, mechanical process.

The bigger picture: mai_kernels was one of the four “non-Python code” packages. This PR proves the pattern works for CUDA packages. simplertransformer (the other CUDA package) can use the exact same approach – py_library for Python, py_gpu_test for tests, Docker for native code.

The FA4 Story

One thing I glossed over: mai_kernels/src/mai_kernels/fa4/ contains 34 vendored Python files from Flash Attention CUTE (Flash Attention 4). It has its own pyproject.toml. Ruff is configured to exclude it entirely (extend-exclude = ["src/mai_kernels/fa4"]).

In the BUILD file, these files are included in the glob() pattern – they’re just Python, so Bazel handles them fine. But they’re worth mentioning because they represent another pattern: vendored code that the team doesn’t maintain, included as a source directory rather than a pip package. Bazel doesn’t care about the distinction. It sees .py files, it builds them.

The Conversation

Here’s how I expect this to go in review.

“Why not have Bazel compile the C++?”

Because the pretrain_base Docker image already does this. The entire CMake/nvcc build pipeline exists and works. The ccache integration alone took significant tuning to get right. Duplicating it in Starlark would be massive effort – custom cc_library rules for CUDA, managing nvcc toolchains, handling pybind11 binding generation, replicating the FetchContent for CUTLASS v4.3.5 and MathDx 25.12.1, dealing with architecture-conditional SM targets – all to produce the exact same .so files that Docker already produces. Zero benefit, months of work.

“Isn’t this cheating? The build isn’t hermetic.”

It’s pragmatic. Bazel hermeticity matters for reproducibility. For CUDA kernels that only run inside a specific Docker image, the image IS the hermetic boundary. You’re not going to bazel run a CUDA kernel on your MacBook. The runtime environment is always the container. This is the same approach that worked for system torch – and nobody calls that cheating anymore.

“Does the Python code build on CPU?”

Yes. bazel build //mai_kernels/... works in the regular bazel job on CPU runners. The mock pattern in __init__.py handles missing native modules gracefully. The py_library target builds, pyright can type-check it, deps are validated. Only the tests require GPU – which is why they’re tagged gpu and restricted to Linux.

“Why a new CI job? Can’t you use existing GPU infra?”

Existing GPU tests (run_gpu_tests.yml) use pytest directly with uv. No Bazel involved. Bazel GPU tests need Bazel + the pretrain_base image + GPU runners. Different enough to warrant a dedicated job. They share the same runner pool (gpu-condor2-rss-dind) and ACR credentials – just different execution engines.

“Will this slow down CI?”

The bazel-gpu-docker job only triggers when **/*.bazel or bazel/** files change – same path filter as the other Bazel jobs. It doesn’t affect the main GPU test pipeline. If you’re not touching BUILD files, this job doesn’t exist for you.

“What about cpu_tests.yml?”

mai_kernels has zero CPU tests. The entry in cpu_tests.yml is literally mai_kernels: [] – an empty list. Every test in this package requires CUDA. That’s why we need py_gpu_test and the GPU CI job. There’s nothing to run on CPU.

The Pattern

Here’s what I want people to take away from this.

If you have a package with native extensions – C++, CUDA, Cython, Fortran, whatever – and you’re migrating to Bazel, and someone says “we can’t migrate that, Bazel doesn’t know how to build C extensions” – they’re probably right. And it probably doesn’t matter.

The question isn’t “can Bazel build the native code.” The question is “does Bazel need to build the native code.” If your CI already builds those .so files in a Docker image, and your tests already run inside that image, then Bazel’s job is managing the Python dependency graph. The native modules are runtime artifacts, same as a GPU driver or a system library. You don’t ask Bazel to compile libcuda.so. Don’t ask it to compile your pybind11 modules either.

The recipe:

1. py_library for the Python source. Two deps: numpy and the mock utilities.
2. include_system_site_packages=True for the native runtime.
3. gazelle:ignore for the directory full of .cu files.
4. .pth file if your Docker image uses a venv (the system Python bridge).
5. py_gpu_test macro for tests that need actual hardware.

Five pieces that unblock a package everyone said was impossible. The hardest part wasn’t the code – it was convincing people (including myself) that “migration” doesn’t have to mean “Bazel compiles everything.”