chore(customizer): add magi_attention dependency in automodel

docker-bake.hcl

@@ -76,6 +76,10 @@ variable "MAMBA_SSM_WHEEL_CONTEXT" {
   default = ""
 }
 
+variable "MAGI_ATTENTION_WHEEL_CONTEXT" {
+  default = ""
+}
+
 variable "FFMPEG_VLM_WHEEL_CONTEXT" {
   default = ""
 }
@@ -138,7 +142,7 @@ variable "RL_BASE_CONTEXT" {
 
 # The tag for base images if needed
 variable "WHEELS_TAG" {
-  default = "a696c78488b10d15d07fab80c75b42c2983c7d8c"
+  default = "67cc126a137e475f22c4925be21cfc76f99847f2"
 }
 
 variable "BAKE_CACHE_SOURCE_BRANCH" {
@@ -186,6 +190,19 @@ variable "CAUSAL_CONV1D_VERSION" {
   default = "v1.5.3"
 }
 
+# Keep in sync with the magi_attention rev in Automodel's [tool.uv.sources]
+# (AUTOMODEL_COMMIT in docker/automodel/Dockerfile.nhx-automodel-base).
+variable "MAGI_ATTENTION_COMMIT" {
+  default = "d7ea8afd44c790b65fab68a04a6a0fdd5adbf182"
+}
+
+# nvidia-nvshmem-cu13 that magi-attention-wheel builds against and nhx-automodel-base ships; the
+# wheel's RPATH points at it, so both targets take this one value. Keep in sync with the NVSHMEM
+# pin in Automodel's docker/Dockerfile.
+variable "NVSHMEM_VERSION" {
+  default = "3.6.5"
+}
+
 function "get_causal_conv1d_wheel_image" {
   params = []
   result = "${WHEELS_REGISTRY}/causal-conv1d-wheel:${WHEELS_TAG}"
@@ -196,6 +213,11 @@ function "get_mamba_ssm_wheel_image" {
   result = "${WHEELS_REGISTRY}/mamba-ssm-wheel:${WHEELS_TAG}"
 }
 
+function "get_magi_attention_wheel_image" {
+  params = []
+  result = "${WHEELS_REGISTRY}/magi-attention-wheel:${WHEELS_TAG}"
+}
+
 function "get_ffmpeg_vlm_wheel_image" {
   params = []
   result = "${WHEELS_REGISTRY}/ffmpeg-vlm-wheel:${WHEELS_TAG}"
@@ -238,6 +260,11 @@ function "mamba_ssm_wheel_context" {
   result = notequal(MAMBA_SSM_WHEEL_CONTEXT, "") ? MAMBA_SSM_WHEEL_CONTEXT : notequal(USE_LOCAL_WHEELS, "") ? "target:mamba-ssm-wheel" : "docker-image://${get_mamba_ssm_wheel_image()}"
 }
 
+function "magi_attention_wheel_context" {
+  params = []
+  result = notequal(MAGI_ATTENTION_WHEEL_CONTEXT, "") ? MAGI_ATTENTION_WHEEL_CONTEXT : notequal(USE_LOCAL_WHEELS, "") ? "target:magi-attention-wheel" : "docker-image://${get_magi_attention_wheel_image()}"
+}
+
 function "ffmpeg_vlm_wheel_context" {
   params = []
   result = notequal(FFMPEG_VLM_WHEEL_CONTEXT, "") ? FFMPEG_VLM_WHEEL_CONTEXT : notequal(USE_LOCAL_WHEELS, "") ? "target:ffmpeg-vlm-wheel" : "docker-image://${get_ffmpeg_vlm_wheel_image()}"
@@ -378,6 +405,7 @@ group "nhx-automodel-gpu-wheels" {
     "causal-conv1d-wheel",
     "mamba-ssm-wheel",
     "rl-cuda-ext-wheel",
+    "magi-attention-wheel",
   ]
 }
 
@@ -840,8 +868,8 @@ target "nhx-gym-host-smoke-test" {
   platforms  = get_platforms()
 }
 
-# Python wheel builders (causal-conv1d, mamba-ssm, av, opencv-python-headless).
-# CUDA extensions only ship source on PyPI; av/opencv bundle FFmpeg. Pre-built for
+# Python wheel builders (causal-conv1d, mamba-ssm, magi-attention, av, opencv-python-headless).
+# CUDA extensions only ship source on PyPI (or git); av/opencv bundle FFmpeg. Pre-built for
 # amd64 and arm64. Wheels live at /wheels/*.whl inside each image.
 
 target "causal-conv1d-wheel" {
@@ -875,6 +903,21 @@ target "mamba-ssm-wheel" {
   platforms = get_platforms()
 }
 
+target "magi-attention-wheel" {
+  target     = "magi-attention-wheel"
+  context    = "."
+  dockerfile = "docker/base/Dockerfile.python-wheels"
+  cache-to   = maybe_registry_cache_to("magi-attention-wheel")
+  cache-from = maybe_registry_cache_from("magi-attention-wheel")
+  tags       = wheel_tags("magi-attention-wheel")
+  output     = image_output()
+  args = {
+    MAGI_ATTENTION_COMMIT = MAGI_ATTENTION_COMMIT
+    NVSHMEM_VERSION       = NVSHMEM_VERSION
+  }
+  platforms = get_platforms()
+}
+
 target "ffmpeg-vlm-wheel" {
   target     = "ffmpeg-vlm-wheel"
   context    = "."
@@ -1080,11 +1123,13 @@ target "nhx-automodel-base-builder" {
   tags            = base_tags("nhx-automodel-base")
   output          = image_output()
   contexts = {
-    causal-conv1d-wheel-image = causal_conv1d_wheel_context()
-    mamba-ssm-wheel-image     = mamba_ssm_wheel_context()
+    causal-conv1d-wheel-image  = causal_conv1d_wheel_context()
+    mamba-ssm-wheel-image      = mamba_ssm_wheel_context()
+    magi-attention-wheel-image = magi_attention_wheel_context()
   }
   args = {
     NHX_COLLECT_SOURCES = NHX_COLLECT_SOURCES
+    NVSHMEM_VERSION     = NVSHMEM_VERSION
   }
   platforms = get_platforms()
 }

docker/automodel/Dockerfile.nhx-automodel-base

@@ -17,6 +17,7 @@ ARG PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.07-py3
 
 FROM causal-conv1d-wheel-image AS causal-conv1d-wheel-src
 FROM mamba-ssm-wheel-image AS mamba-ssm-wheel-src
+FROM magi-attention-wheel-image AS magi-attention-wheel-src
 
 FROM alpine/git AS automodel-clone
 ARG AUTOMODEL_BRANCH
@@ -51,8 +52,8 @@ RUN cd /opt/Automodel && \
 # OpenCV and FFmpeg-bearing wheels (av, imageio-ffmpeg, decord) in the opt-in
 # vlm-media / diffusion-media extras; skipping them keeps those wheels out. We'll need
 # to re-examine when Customizer supports VLM and FFmpeg 8.xx only releases every 6 months
-# The magi group (MagiAttention) is a git-only CUDA build that Customizer does not
-# expose yet, so skip it rather than compiling it here.
+# The magi group is a git-only CUDA build; it is installed below from the prebuilt
+# magi-attention-wheel image instead.
 RUN --mount=type=cache,target=/root/.cache/uv \
     cd /opt/Automodel && \
     UV_HTTP_TIMEOUT=120 uv sync --locked \
@@ -73,9 +74,11 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 # Install AFTER Automodel sync - uv sync drops packages not in its lockfile.
 RUN --mount=from=causal-conv1d-wheel-src,target=/tmp/causal-conv1d-wheel-src,readonly \
     --mount=from=mamba-ssm-wheel-src,target=/tmp/mamba-ssm-wheel-src,readonly \
+    --mount=from=magi-attention-wheel-src,target=/tmp/magi-attention-wheel-src,readonly \
     uv pip install --no-cache-dir --no-deps \
     /tmp/causal-conv1d-wheel-src/wheels/cu13.3/causal_conv1d-*cp312*.whl \
-    /tmp/mamba-ssm-wheel-src/wheels/cu13.3/mamba_ssm-2.3.0-cp312*.whl
+    /tmp/mamba-ssm-wheel-src/wheels/cu13.3/mamba_ssm-2.3.0-cp312*.whl \
+    /tmp/magi-attention-wheel-src/wheels/cu13.3/magi_attention-*cp312*.whl
 
 RUN --mount=type=cache,target=/root/.cache/uv \
     git clone https://github.com/fanshiqing/grouped_gemm.git /tmp/grouped_gemm && \
@@ -178,6 +181,12 @@ RUN --mount=type=bind,source=docker/scripts/remove-listed-files.py,target=/tmp/r
     --mount=type=bind,source=docker/removals/files/base/pytorch-ngc-common.txt,target=/tmp/common-file-removals.txt,readonly \
     python /tmp/remove-listed-files.py /tmp/common-file-removals.txt
 
+# magi_attn_comm (magi-attention-wheel) loads libnvshmem_host through an RPATH into the system
+# site-packages it was built against, so install the same NVSHMEM with the system pip, before
+# PATH puts /opt/venv first. Keep NVSHMEM_VERSION equal to the magi-attention-wheel build.
+ARG NVSHMEM_VERSION=3.6.5
+RUN pip install --no-cache-dir "nvidia-nvshmem-cu13==${NVSHMEM_VERSION}"
+
 ARG NHX_COLLECT_SOURCES=0
 ENV NHX_COLLECT_SOURCES=${NHX_COLLECT_SOURCES}
 RUN bash /usr/local/bin/collect-apt-sources.sh \

docker/automodel/README.md

@@ -9,7 +9,7 @@ Two images for the **nhx-automodel** customization backend, plus the shared **`n
 
 | Image | Dockerfile | Role |
 |-------|------------|------|
-| `nhx-automodel-base` | `docker/automodel/Dockerfile.nhx-automodel-base` | PyTorch 26.05 + Automodel + `mamba-ssm` / `causal-conv1d` wheels |
+| `nhx-automodel-base` | `docker/automodel/Dockerfile.nhx-automodel-base` | PyTorch 26.07 + Automodel + `mamba-ssm` / `causal-conv1d` / `magi-attention` wheels |
 | `nhx-customizer-tasks` | `docker/Dockerfile.nhx-customizer-tasks` | Shared CPU tasks (`file_io`, `model_entity`, `model_spec`, LoRA sidecar) |
 | `nhx-automodel-training` | `docker/automodel/Dockerfile.nhx-automodel-training` | Training (`nhx.automodel.tasks.training`) and retrieval mining (`nhx.automodel.tasks.retrieval_mine`) |
 
@@ -27,7 +27,7 @@ Bake file: **`docker-bake.hcl`** at the NeMo Helix repo root (`context = "."`).
 
 ## Prerequisites
 
-1. **CUDA extension wheels** (`causal-conv1d-wheel`, `mamba-ssm-wheel`) - built from this directory or pulled from NGC. The wheel Dockerfile and uv locks live under `docker/locks/` (ported from `nhx`).
+1. **CUDA extension wheels** (`causal-conv1d-wheel`, `mamba-ssm-wheel`, `magi-attention-wheel`) - built from this directory or pulled from NGC. The wheel Dockerfile and uv locks live under `docker/locks/` (ported from `nhx`).
 
 2. **Base image tag** - after building the base, set `BASE_TAG_AUTOMODEL` (or push to `BASE_REGISTRY`) before building tasks/training.
 
@@ -82,7 +82,7 @@ Override registry: `export WHEELS_REGISTRY=...` and `export IMAGE_REGISTRY=...`
 
 ## Tasks / training runtime (platform glue)
 
-**Base (`nhx-automodel-base`):** NGC PyTorch 26.05, Automodel `uv sync --locked`, pinned `transformers`/`torch`.
+**Base (`nhx-automodel-base`):** NGC PyTorch 26.07, Automodel `uv sync --locked`, pinned `transformers`/`torch`.
 
 **Media dependencies:** The base syncs only the Automodel extras text training needs (`extra`, `delta-databricks`). Since r0.6.0, Automodel keeps `opencv-python-headless` and FFmpeg-bearing wheels (`av`, `imageio-ffmpeg`, `decord`) in the opt-in `vlm-media` / `diffusion-media` extras, so they are not installed. Earlier releases needed a local cherry-pick (`3d98f6e3.diff`) to get the same result.
 

docker/base/Dockerfile.python-wheels

@@ -8,10 +8,12 @@
 # Builds Python wheels for:
 #   - causal-conv1d (CUDA extension required by mamba-ssm)
 #   - mamba-ssm     (selective state space model)
+#   - magi-attention (context-parallel attention, Automodel `magi` dependency group)
 #   - av + opencv-python-headless + decord2 (FFmpeg 8.1.2 bundled wheels for nspect / VLM)
 #   - rl-cuda-ext-wheel (CPython 3.13 / torch 2.13.0+cu130 extensions for nhx-rl-base)
 #
-# causal-conv1d / mamba-ssm only ship source on PyPI and require nvcc.
+# causal-conv1d / mamba-ssm only ship source on PyPI and require nvcc; magi-attention is
+# git-only and compiles CUDA kernels against the NGC torch.
 # av / opencv wheels bundle FFmpeg; upstream PyPI wheels lag nspect 8.1.2.
 # The cp311/cp312 causal-conv1d and mamba-ssm targets stay for Automodel and Unsloth.
 # Build those via the nhx-automodel-gpu-wheels bake group.
@@ -20,6 +22,8 @@
 # Build args:
 #   CAUSAL_CONV1D_VERSION  - git tag to build              (default: v1.5.3)
 #   MAMBA_22_COMMIT        - git SHA or tag to build
+#   MAGI_ATTENTION_COMMIT  - MagiAttention git SHA (match Automodel's [tool.uv.sources] rev)
+#   NVSHMEM_VERSION        - nvidia-nvshmem-cu13 for magi-attention (match nhx-automodel-base)
 #   CUDA_VERSION           - CUDA devel image version      (default: 12.8.1)
 #   TORCH_CUDA_ARCH_LIST   - semicolon-separated SM targets (default: "8.0;8.6;9.0")
 #   FFMPEG_VERSION         - FFmpeg for av/opencv wheels   (default: 8.1.2)
@@ -287,6 +291,46 @@ COPY --from=mamba-ssm-23-wheel-builder-py312 /wheels /wheels
 COPY --from=mamba-ssm-23-wheel-builder-py312-cu13.1.1 /wheels /wheels/cu13.1.1
 COPY --from=mamba-ssm-23-wheel-builder-py312-cu13.3 /wheels /wheels/cu13.3
 
+# =============================================================================
+# magi-attention wheel — Python 3.12 (for Python 3.12 consumers) - Using CUDA 13.3
+# =============================================================================
+FROM mamba-wheel-base-py312-cu13.3 AS magi-attention-wheel-builder-py312-cu13.3
+
+ARG MAGI_ATTENTION_COMMIT=d7ea8afd44c790b65fab68a04a6a0fdd5adbf182
+ARG NVSHMEM_VERSION=3.6.5
+
+# Automodel's docker/Dockerfile installs the RDMA headers and NVSHMEM before it builds
+# MagiAttention. Without NVSHMEM, magi_attn_comm still compiles its internode kernels and
+# fails. MagiAttention finds NVSHMEM through the nvidia.nvshmem package and records its lib
+# dir as the extension's RPATH, so nhx-automodel-base installs the same NVSHMEM_VERSION with
+# the system pip, at the same path.
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends --allow-change-held-packages rdma-core libibverbs-dev && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/* && \
+    pip install --no-cache-dir "nvidia-nvshmem-cu13==${NVSHMEM_VERSION}" && \
+    python -c "import nvidia.nvshmem"
+
+# Build settings mirror Automodel's docker/Dockerfile: Hopper + Blackwell kernels only,
+# FFA kernels JIT-compiled at runtime instead of prebuilt. versioningit stamps the wheel
+# version from git tags (Automodel lists it as an extra build dependency). uv checks out a
+# git dependency's submodules (CUTLASS, the flash-attention fork), so this clone does too.
+RUN mkdir -p /wheels && \
+    pip install --no-cache-dir versioningit && \
+    git clone https://github.com/SandAI-org/MagiAttention.git /src/magi-attention && \
+    cd /src/magi-attention && \
+    git checkout ${MAGI_ATTENTION_COMMIT} && \
+    git submodule update --init --recursive && \
+    MAGI_ATTENTION_BUILD_COMPUTE_CAPABILITY="90,100" \
+    MAGI_ATTENTION_PREBUILD_FFA=0 \
+    uv build --wheel --no-build-isolation --out-dir=/wheels . && \
+    rm -rf /src/magi-attention
+
+# The final magi-attention-wheel image contains:
+#   - cu13.3/magi_attention-*-cp312-*.whl  (for Automodel on CUDA 13.3)
+FROM scratch AS magi-attention-wheel
+COPY --from=magi-attention-wheel-builder-py312-cu13.3 /wheels /wheels/cu13.3
+
 # =============================================================================
 # FFmpeg / VLM wheels (manylinux_2_28) — av + opencv-python-headless + decord2
 # =============================================================================

tests/smoke_gpu/test_customizer_automodel.py

@@ -72,6 +72,15 @@ def test_bitsandbytes_importable():
     import bitsandbytes  # noqa: F401
 
 
+@pytest.mark.smoke_nhx_automodel_training
+def test_magi_attention_extensions_importable():
+    # magi_attention's __init__ only warns when its compiled extensions fail to load,
+    # so import them directly to surface a missing .so or torch ABI mismatch. Loading
+    # magi_attn_comm also resolves libnvshmem_host through the RPATH recorded at build time,
+    # so this fails if the base image's NVSHMEM is missing or at another path.
+    from magi_attention import magi_attn_comm, magi_attn_ext  # noqa: F401
+
+
 @pytest.mark.smoke_nhx_automodel_training
 def test_nhx_automodel_training_importable():
     import nemo_automodel  # noqa: F401