Sync mad-rccl with develop

.dockerignore

@@ -0,0 +1,22 @@
+# Runtime artifacts from madengine runs. The primus dockerfiles build with the repo
+# root as context and copy only scripts/Primus, so none of this belongs in the build.
+# hf_cache matters most: scripts/jax-maxdiffusion/run.sh caches weights there, and
+# flux/wan checkpoints push it past 100GB, which docker would otherwise send to the
+# daemon on every rebuild.
+.git
+.venv
+hf_cache/
+run_directory/
+output/
+logs/
+.*_env/
+.jax-*_env/
+*.csv
+*.log
+*.jsonl
+perf_metrics_*.jsonl
+scripts/jax-maxtext/output/
+scripts/jax-maxdiffusion/output/
+**/.jax_cache/
+**/core.*
+scripts/Primus/data/

.gitignore

@@ -82,3 +82,18 @@ venv.bak/
 !.claude/skills/mad-profiling/assets/**/*.json
 __pycache__/
 .pytest_cache/
+# Primus checkout for JAX/MaxText discovery + image build. Use either:
+#   git submodule update --init scripts/Primus
+#   tools/fetch_primus.sh
+# See benchmark/jax_maxtext/README.md for details.
+scripts/Primus/
+
+# Artifacts left in the repo root by a madengine run. hf_cache matters most: the
+# jax-maxdiffusion wrapper points HF_HOME here so gated FLUX/WAN weights survive between
+# runs, which puts 100+GB one careless `git add -A` away from being staged.
+hf_cache/
+run_directory/
+.jax-*_env/
+perf_metrics_*.jsonl
+scripts/jax-maxtext/output/
+scripts/jax-maxdiffusion/output/

docker/primus_maxdiffusion.ubuntu.amd.Dockerfile

@@ -0,0 +1,84 @@
+# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'}
+###############################################################################
+#
+# MIT License
+#
+# Copyright (c) Advanced Micro Devices, Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################
+
+# Primus JAX/MaxDiffusion launcher image for MAD (WAN 2.1 + FLUX.1-dev): bakes the Primus
+# repo onto a JAX training base so scripts/jax-maxdiffusion/run.sh can run
+# `train pretrain --config ...`.
+#
+# The base image owns the maxdiffusion stack: maxdiffusion is installed and patched at
+# /workspace/maxdiffusion, at the same commit as Primus's third_party/maxdiffusion pin. So
+# this image runs no setup_maxdiffusion_env.sh and installs no maxdiffusion deps. For a
+# base without the stack, use the setup-script build from git history before this commit.
+#
+# Check Primus out first with tools/fetch_primus.sh. It is gitignored here and baked from
+# the build context, which keeps git auth for a private repo out of the build. That script
+# initializes no submodules: third_party/maxdiffusion at the same commit is unpatched, and
+# run_pretrain.sh would select it over the base's tree if MAXDIFFUSION_PATH were ever unset.
+#
+# Build from the repo root, as madengine does for dockerfile paths containing "primus":
+#   docker build -f docker/primus_maxdiffusion.ubuntu.amd.Dockerfile .
+
+# madengine passes the base via docker_build_arg, which is how the v26.6 sweep put both
+# maxtext and maxdiffusion on one unified CI image so their numbers share a toolchain.
+ARG BASE_DOCKER=rocm/jax-training:maxtext-v26.6
+FROM $BASE_DOCKER
+
+USER root
+ENV WORKSPACE_DIR=/workspace
+# The Primus repo root, not /workspace: run.sh resolves examples/ relative to it.
+ENV PRIMUS_ROOT=/workspace/Primus
+# Pin the base's patched tree; run_pretrain.sh would otherwise default to
+# $PRIMUS_ROOT/third_party/maxdiffusion and insert it at sys.path[0].
+ENV MAXDIFFUSION_PATH=/workspace/maxdiffusion
+# Transformer Engine must load only its JAX extension (torch is present too).
+ENV NVTE_FRAMEWORK=jax
+RUN mkdir -p $WORKSPACE_DIR
+WORKDIR $WORKSPACE_DIR
+
+LABEL mad.launcher=primus
+
+# The base may ship /workspace/Primus as a git clone, and COPY cannot replace a
+# .git directory with a submodule checkout's .git file.
+RUN rm -rf /workspace/Primus
+COPY scripts/Primus/ /workspace/Primus/
+
+RUN test -f /workspace/Primus/examples/run_pretrain.sh
+RUN test -d /workspace/Primus/primus/backends/maxdiffusion \
+    || (echo "ERROR: Primus checkout lacks primus/backends/maxdiffusion; use Primus main branch." >&2 && exit 1)
+
+# Prove the base's stack is really there, so a wrong base fails the build instead
+# of step 0 of a training run. The patch fixes a segfault on TE import order.
+RUN python3 -c "import maxdiffusion, os; print('maxdiffusion ->', os.path.dirname(maxdiffusion.__file__))"
+RUN grep -q "preload before Transformer Engine" /workspace/maxdiffusion/src/maxdiffusion/train_utils.py \
+    || (echo "ERROR: /workspace/maxdiffusion is missing or lacks the TF-preload patch." >&2 && exit 1)
+
+# Primus's own requirements, not maxdiffusion's, which the base already covers.
+# Installed here rather than on every run: run.sh sets PRIMUS_SKIP_PIP=1 so a
+# launch stays off the network. On this base it adds loguru.
+RUN pip3 install --no-cache-dir -r /workspace/Primus/requirements-maxdiffusion.txt
+
+RUN pip3 list 2>/dev/null || true

docker/primus_maxtext.ubuntu.amd.Dockerfile

@@ -0,0 +1,75 @@
+# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'}
+###############################################################################
+#
+# MIT License
+#
+# Copyright (c) 2026 Advanced Micro Devices, Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#################################################################################
+
+# Primus JAX/MaxText launcher image for MAD: bakes the Primus repo onto a JAX training
+# base so scripts/jax-maxtext/run.sh can run `train pretrain --config ...`.
+#
+# The base image owns the MaxText stack: it is installed at /workspace/maxtext, at the same
+# commit as Primus's third_party/maxtext pin, so no submodule is needed here. This mirrors
+# how primus_maxdiffusion takes maxdiffusion from the base.
+#
+# Check Primus out first with tools/fetch_primus.sh. It is gitignored here and baked from
+# the build context, which keeps git auth for a private repo out of the build.
+#
+# Build from the repo root, as madengine does for dockerfile paths containing "primus":
+#   docker build -f docker/primus_maxtext.ubuntu.amd.Dockerfile .
+
+# madengine passes the base via docker_build_arg, which is how the v26.6 sweep put both
+# maxtext and maxdiffusion on one unified CI image so their numbers share a toolchain.
+ARG BASE_DOCKER=rocm/jax-training:maxtext-v26.6
+FROM $BASE_DOCKER
+
+USER root
+ENV WORKSPACE_DIR=/workspace
+# The Primus repo root, not /workspace: run.sh resolves examples/ relative to it.
+ENV PRIMUS_ROOT=/workspace/Primus
+# Pin the base's tree; prepare.py would otherwise default to
+# $PRIMUS_ROOT/third_party/maxtext, which this image does not ship.
+ENV MAXTEXT_PATH=/workspace/maxtext
+RUN mkdir -p "$WORKSPACE_DIR"
+WORKDIR $WORKSPACE_DIR
+
+LABEL mad.launcher=primus
+
+# The base may ship /workspace/Primus as a git clone, and COPY cannot replace a
+# .git directory with a submodule checkout's .git file.
+RUN rm -rf /workspace/Primus
+COPY scripts/Primus/ /workspace/Primus/
+
+RUN test -f /workspace/Primus/examples/run_pretrain.sh
+RUN test -f /workspace/Primus/requirements-jax.txt
+
+# Prove the base's stack is really there, so a wrong base fails the build instead
+# of step 0 of a training run.
+RUN test -f /workspace/maxtext/pyproject.toml \
+    || (echo "ERROR: no MaxText at /workspace/maxtext. Use a base that bakes it, or point MAXTEXT_PATH at a checkout." >&2 && exit 1)
+
+# Installed here rather than on every run: run.sh sets PRIMUS_SKIP_PIP=1 so a
+# launch stays off the network. On this base it adds loguru, wandb, pre-commit.
+RUN pip3 install --no-cache-dir -r /workspace/Primus/requirements-jax.txt
+
+RUN pip3 list 2>/dev/null || true

docker/sglang_disagg_inference_full_overlay.ubuntu.amd.Dockerfile

@@ -214,6 +214,20 @@ ARG MORI_WHEEL_URL=
 # build MoRI for a different arch than the rest of the image.
 ARG MORI_GPU_ARCHS=${BUILD_GPU_TARGETS}
 ARG MORI_VERSION=1.2.0
+# NOTE (UMBP unit-test build guard): mori's src/umbp/CMakeLists.txt calls
+# `add_subdirectory(tests)` UNCONDITIONALLY, ignoring the top-level BUILD_TESTS
+# option (which setup.py already defaults to OFF, and which correctly gates
+# tests/cpp). Those UMBP tests use gtest_discover_tests() in its default
+# POST_BUILD discovery mode, which EXECUTES each freshly-linked test binary at
+# build time to enumerate test cases. The binaries link the ROCm/HIP runtime
+# (libamdhip64 / libhsa-runtime64 / librocm_smi64 / libhsakmt), so they abort on
+# startup inside `docker build`, which has NO GPU (RUN steps cannot pass
+# --gpus/--device, so /dev/kfd and /dev/dri are absent). The empty output then
+# breaks CMake's string(JSON ...) in GoogleTest/ParseTestList.cmake and fails the
+# amd_mori wheel build. Nothing is "pulled from the build host" -- CMake just
+# tries to RUN the test binaries during the build. Re-gated behind BUILD_TESTS in
+# the RUN block below; idempotent no-op once upstream fixes it (already fixed on
+# mori main, still present at the pinned MORI_COMMIT).
 ARG MORI_SRC_DIR=/sgl-workspace/mori
 
 ENV MORI_GPU_ARCHS=${MORI_GPU_ARCHS} \
@@ -253,6 +267,12 @@ RUN set -e; \
         git clone --depth 1 --branch "${MORI_BRANCH}" "${MORI_REPO}" "${MORI_SRC_DIR}"; \
       fi; \
       cd "${MORI_SRC_DIR}" && git submodule update --init --recursive || true; \
+      if grep -qE '^add_subdirectory\(tests\)$' src/umbp/CMakeLists.txt; then \
+        sed -i 's|^add_subdirectory(tests)$|if(BUILD_TESTS)\n  add_subdirectory(tests)\nendif()|' src/umbp/CMakeLists.txt; \
+        echo "MORI_UMBP_TESTS_GUARDED (add_subdirectory(tests) now gated by BUILD_TESTS)"; \
+      else \
+        echo "MORI_UMBP_TESTS_GUARD_SKIPPED (unguarded pattern not found; upstream may have fixed it)"; \
+      fi; \
       pip install --no-cache-dir "cmake<4.0" setuptools wheel pybind11 ninja; \
       pip install --no-build-isolation --no-cache-dir --force-reinstall .; \
     fi