Merge origin/main into ModelOpt real-quant export

Signed-off-by: Meng Xin <164952785+mxinO@users.noreply.github.com>

Author: mxinO

.main.commit

@@ -1 +1 @@
-ef8cd10b8e34a94c9ad7452e7851f58337f353a0
+441a987409b8fdd2b90b463f7999bac970109a43

3rdparty/Megatron-LM

@@ -229,6 +229,7 @@ Megatron Bridge provides out-of-the-box bridges and training recipes for a wide
 
 | Family | Supported variants |
 |----------------|--------------------|
+| [**BAGEL**](docs/models/bagel/bagel.md) | BAGEL-7B-MoT (requires the tested MCore `dev` revision and `bagel` extra) |
 | [**Bailing**](docs/models/bailing/index.md) | Ling 2.0 / Ling MoE V2 (Bailing) |
 | [**BERT**](docs/models/bert/index.md) | Megatron-Style BERT (`MegatronBertForMaskedLM`) |
 | [**DeepSeek**](docs/models/deepseek/index.md) | DeepSeek V2 / V2 Lite (deprecated), DeepSeek V3, DeepSeek V4 / V4 Flash |
@@ -246,6 +247,7 @@ Megatron Bridge provides out-of-the-box bridges and training recipes for a wide
 | [**Mistral**](docs/models/mistral/index.md) | Mistral 7B / Small 3 24B (deprecated), Ministral 3 (3B/8B/14B) |
 | [**Xiaomi-MiMo**](docs/models/mimo/index.md) | Xiaomi-MiMo, MiMo-V2-Flash |
 | [**Moonlight**](docs/models/moonlight/index.md) | Moonlight |
+| [**Muse Glimmer**](docs/models/muse_glimmer/muse-glimmer.md) | Muse Glimmer 30B |
 | [**Nemotron**](docs/models/nemotron/index.md) | Nemotron H v1 (deprecated), Nemotron Nano v2 (deprecated), Nemotron-3 Nano, Nemotron-3 Super, Nemotron-3 Ultra, Llama Nemotron (deprecated), Nemotron Nano v2 VL (deprecated), Nemotron-3 Nano Omni |
 | [**OLMoE**](docs/models/olmoe/index.md) | OLMoE |
 | [**Qwen**](docs/models/qwen/index.md) | Qwen2 / Qwen2.5, Qwen3, Qwen3-MoE, Qwen3 Next, Qwen3.5 (dense/MoE), Qwen2.5-VL, Qwen3-VL, Qwen3.5-VL, Qwen3.6-VL, Qwen2 Audio, Qwen2.5-Omni, Qwen3-Omni, Qwen3-ASR |

README.md

@@ -1,4 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,7 +13,29 @@
 # limitations under the License.
 
 ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:26.08-py3
-FROM ${BASE_IMAGE} AS megatron_bridge
+FROM ${BASE_IMAGE} AS base
+
+ENV PATH="/opt/venv/bin:/usr/local/bin:$PATH" \
+    UV_PROJECT_ENVIRONMENT=/opt/venv \
+    VIRTUAL_ENV=/opt/venv \
+    UV_LINK_MODE=copy \
+    UV_VERSION="0.7.2" \
+    UV_HTTP_TIMEOUT=120
+
+RUN curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh && \
+    install -m 0755 /root/.local/bin/uv /usr/local/bin/uv && \
+    uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages
+
+# Cache the native wheel independently of Bridge's other dependencies. Both CPU
+# platforms use the same base runtime as the final image and build both GPU targets.
+FROM base AS mok-wheel
+ARG MOK_COMMIT=a2ad2d0ce2366c1817d15498cdbaf8ead5966117
+ARG MOK_ARCH=ALL
+RUN --mount=type=bind,source=docker/common/install_mok.sh,target=/opt/install_mok.sh \
+    --mount=type=bind,source=docker/patches/mok.patch,target=/opt/mok.patch \
+    MOK_ARCH=${MOK_ARCH} bash /opt/install_mok.sh "${MOK_COMMIT}" /opt/mok.patch /opt/mok-wheels
+
+FROM base AS megatron_bridge
 WORKDIR /opt/Megatron-Bridge
 
 ARG APPLY_PYTORCH_LIBRARY_FINALIZER_PATCH=False
@@ -50,12 +72,6 @@ ARG REINSTALL_NVSHMEM=True
 ARG PRESERVE_BASE_RUNTIME=False
 
 ENV MBRIDGE_PRESERVE_BASE_RUNTIME=${PRESERVE_BASE_RUNTIME} \
-    PATH="/opt/venv/bin:/usr/local/bin:$PATH" \
-    UV_PROJECT_ENVIRONMENT=/opt/venv \
-    VIRTUAL_ENV=/opt/venv \
-    UV_LINK_MODE=copy \
-    UV_VERSION="0.7.2" \
-    UV_HTTP_TIMEOUT=120 \
     NVTE_BUILD_NUM_PHILOX_ROUNDS=3 \
     NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD=1 \
     HYBRID_EP_MULTINODE=1 \
@@ -97,9 +113,6 @@ EOF
 ARG DEEPEP_NVCC_THREADS
 
 RUN --mount=type=bind,source=docker/patches/deepep.patch,target=/opt/deepep.patch \
-    curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh && \
-    install -m 0755 /root/.local/bin/uv /usr/local/bin/uv && \
-    uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages && \
     apt-get update && apt install -y --only-upgrade gnupg openssl libssl3t64 && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/* && \
@@ -264,6 +277,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     fi && \
     chown -R 1000:1000 ${UV_PROJECT_ENVIRONMENT}
 
+# Install after every uv sync so the Docker-only package is retained. Bind-mount
+# the wheel to avoid retaining a second archive in the final image.
+RUN --mount=type=bind,from=mok-wheel,source=/opt/mok-wheels,target=/opt/mok-wheels \
+    uv pip install --python "${UV_PROJECT_ENVIRONMENT}/bin/python" --no-deps --no-cache /opt/mok-wheels/*.whl && \
+    chown -R 1000:1000 ${UV_PROJECT_ENVIRONMENT}
+
 COPY --chown=1000:1000 . /opt/Megatron-Bridge
 
 RUN python /opt/preserve_base_runtime.py check /opt/base-runtime-packages.json
@@ -274,5 +293,8 @@ RUN python /opt/preserve_base_runtime.py check /opt/base-runtime-packages.json
 RUN install -d -o 1000 -g 1000 -m 0755 /nemo_run /home/ubuntu/.cache && \
     chown -R 1000:1000 /opt/Megatron-Bridge
 
+# Remove the Nsight Systems EFA metrics plugin to address CVE-2026-27143.
+RUN rm -rf /usr/local/cuda-*/NsightSystems-cli-*/target-linux-x64/plugins/efa_metrics/
+
 ENV HOME=/home/ubuntu
 USER 1000:1000

docker/Dockerfile.ci

@@ -27,6 +27,35 @@ docker build \
   .
 ```
 
+The image installs [Mixture of Kittens](https://github.com/QiZhangNV/mixture-of-kittens)
+on both `linux/amd64` and `linux/arm64`. Each platform's native wheel includes
+GB200/B200 (`sm_100a`) and GB300/B300 (`sm_103a`) kernels by default. This does not
+enable MoK kernels on other GPUs or select MoK as a training backend.
+
+The `mok-wheel` stage builds against the same base image's Python, PyTorch, and CUDA
+as the final image, independently of the other Bridge dependencies. BuildKit caches
+it by platform, base image, source revision, GPU targets, installer, and patch; registry
+caches must export intermediate stages with `mode=max` (as Bridge CI already does).
+The build checks that the wheel contains every requested GPU target without requiring
+a GPU. The pinned source needs `scripts/prepare_thunderkittens.sh` and the compatibility
+patch in `docker/patches/mok.patch`; revisit both workarounds when updating the source.
+
+Build only the MoK wheel stage on either native builder, for example:
+
+```bash
+docker buildx build \
+  --platform linux/arm64 \
+  -f docker/Dockerfile.ci \
+  --target mok-wheel \
+  -t megatron-bridge-mok-wheel:arm64 \
+  --load \
+  .
+```
+
+Use `--platform linux/amd64` on an AMD64 builder. Full image builds use the existing
+`megatron_bridge` target and install the wheel after dependency syncing. MoK is a
+Docker-only dependency, outside `pyproject.toml` and `uv.lock`.
+
 | Argument | Description |
 |---|---|
 | `BASE_IMAGE` | Base container |
@@ -204,6 +233,8 @@ docker build \
 | `APPLY_PYTORCH_LIBRARY_FINALIZER_PATCH` | Set to `True` to apply the torch.library finalizer patch; `False` for base images it does not apply against (e.g. Rubin) |
 | `INSTALL_DEEPEP` | Set to `True` to build and install DeepEP and nvshmem |
 | `DEEPEP_COMMIT` | DeepEP git commit SHA |
+| `MOK_COMMIT` | Full Mixture of Kittens source commit SHA; defaults to `a2ad2d0ce2366c1817d15498cdbaf8ead5966117` |
+| `MOK_ARCH` | GPU targets in each CPU platform's MoK wheel: `ALL` (default, GB200 and GB300), `SM100` (GB200/B200), or `SM103` (GB300/B300) |
 | `REINSTALL_NVSHMEM` | Set to `True` to reinstall nvshmem (`nvidia-nvshmem-cu13`) over the base image version; only applied when `INSTALL_DEEPEP=True` |
 | `MCORE_TRIGGERED_TESTING` | Skip uv lockfile check for cross-version Megatron-LM testing |
 | `INSTALL_DIFFUSION_DEPS` | Install the test-only WAN diffusion dependencies; defaults to `false` for Framework/release images |
@@ -230,9 +261,11 @@ docker build \
 | `common/install_cublas.sh` | Reinstall cuBLAS and cuBLASLt from the public NVIDIA CUDA apt repo |
 | `common/install_nccl.sh` | Reinstall NCCL from the public NVIDIA CUDA apt repo |
 | `common/install_cudnn.sh` | Reinstall cuDNN from the public NVIDIA CUDA apt repo |
+| `common/install_mok.sh` | Build and verify the pinned Mixture of Kittens wheel |
 | `common/install_nsys.sh` | Reinstall Nsight Systems from the public NVIDIA CUDA apt repo |
 | `common/install_mamba.sh` | Build the locked Mamba sdist with the local C++ standard patch |
 | `common/preserve_base_runtime.py` | Record and check selected base package metadata for opt-in runtime preservation |
 | `patches/deepep.patch` | Patch applied to DeepEP during CI image build |
+| `patches/mok.patch` | MoK dual GPU target build support and PyTorch/CUDA namespace compatibility fix |
 | `patches/mamba.patch` | Let Mamba's CUDA extension inherit PyTorch's C++ standard |
 | `patches/vllm.patch` | Patch applied to vLLM after install in fw-base |

docker/README.md

@@ -67,6 +67,7 @@ override-dependencies = [
     "packaging>=25.0",
     "nvidia-pytriton>=0.7.0",
     "flashinfer-python==0.6.18.post1",          # Match flashinfer-cubin installed by MBridge
+    "nvidia-cudnn-frontend==1.29.0",            # Match MBridge; override PyPI MCore's older dev pin
     "urllib3>=2.0.7",
     "sh==2.2.4",
     "jupyter-core==5.8.1",

docker/common/fw_pyproject.toml

@@ -22,6 +22,7 @@ PyNvVideoCodec  # dlopens libnvidia-encode.so.1 (NVENC)
 deep_ep
 deep_ep_cpp
 hybrid_ep_cpp
+mok  # native extension links libcuda.so.1
 torch_tensorrt
 quack  # imports nvidia-cutlass-dsl MLIR libs, which dlopen libcuda at import
 peft # transitively imports optional CUDA-backed modules that dlopen libcuda

docker/common/import_check_skip.txt

@@ -0,0 +1,73 @@
+#!/usr/bin/env bash
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -euxo pipefail
+
+MOK_REVISION="${1:?Usage: install_mok.sh <git commit> <patch> <wheel directory>}"
+PATCH_FILE=$(realpath "${2:?Usage: install_mok.sh <git commit> <patch> <wheel directory>}")
+WHEEL_DIR=$(realpath -m "${3:?Usage: install_mok.sh <git commit> <patch> <wheel directory>}")
+PYTHON="${UV_PROJECT_ENVIRONMENT:-/opt/venv}/bin/python"
+
+if [[ ! "${MOK_REVISION}" =~ ^[0-9a-f]{40}$ ]]; then
+    echo "MOK_COMMIT must be a full Git commit" >&2
+    exit 1
+fi
+
+# ALL includes both GPU architectures supported by the pinned source.
+export MOK_ARCH="${MOK_ARCH:-ALL}"
+case "${MOK_ARCH}" in
+    ALL) MOK_CUBINS=(sm_100a sm_103a) ;;
+    SM100) MOK_CUBINS=(sm_100a) ;;
+    SM103) MOK_CUBINS=(sm_103a) ;;
+    *)
+        echo "Unsupported MOK_ARCH '${MOK_ARCH}'; expected ALL, SM100 or SM103" >&2
+        exit 1
+        ;;
+esac
+
+WORK_DIR=$(mktemp -d)
+trap 'rm -rf "${WORK_DIR}"' EXIT
+
+git init "${WORK_DIR}"
+git -C "${WORK_DIR}" remote add origin https://github.com/QiZhangNV/mixture-of-kittens.git
+git -C "${WORK_DIR}" fetch --depth 1 origin "${MOK_REVISION}"
+git -C "${WORK_DIR}" checkout --detach "${MOK_REVISION}"
+git -C "${WORK_DIR}" apply --check "${PATCH_FILE}"
+git -C "${WORK_DIR}" apply "${PATCH_FILE}"
+bash "${WORK_DIR}/scripts/prepare_thunderkittens.sh"
+
+export NVCC="${CUDA_HOME:-/usr/local/cuda}/bin/nvcc"
+export MOK_NVCC="${NVCC}"
+# Docker builds have no GPU driver. Stubs are link inputs, never runtime libraries.
+export LIBRARY_PATH="${CUDA_HOME:-/usr/local/cuda}/lib64/stubs${LIBRARY_PATH:+:${LIBRARY_PATH}}"
+uv build --python "${PYTHON}" --wheel --no-build-isolation --no-cache --out-dir "${WHEEL_DIR}" "${WORK_DIR}"
+uv pip install --python "${PYTHON}" --no-deps --no-cache "${WHEEL_DIR}"/*.whl
+
+# Inspect the installed wheel without loading libcuda on GPU-less image builders.
+MOK_EXTENSION=$(uv run --no-project --no-sync --python "${PYTHON}" python - <<'PY'
+import sys
+from importlib.metadata import distribution
+from sysconfig import get_config_var
+
+extension = distribution("mixture-of-kittens").locate_file("mok/_C" + get_config_var("EXT_SUFFIX"))
+assert extension.is_file(), f"MoK native extension missing: {extension}"
+sys.stdout.write(str(extension))
+PY
+)
+"${CUDA_HOME:-/usr/local/cuda}/bin/cuobjdump" --list-elf "${MOK_EXTENSION}" > "${WORK_DIR}/cubins.txt"
+cat "${WORK_DIR}/cubins.txt"
+for MOK_CUBIN in "${MOK_CUBINS[@]}"; do
+    grep -Fq ".${MOK_CUBIN}.cubin" "${WORK_DIR}/cubins.txt"
+done

docker/common/install_mok.sh

@@ -0,0 +1,38 @@
+diff --git a/Makefile b/Makefile
+index 58bebae..1c8eb2a 100644
+--- a/Makefile
++++ b/Makefile
+@@ -29,12 +29,18 @@ NVCCFLAGS += -ltorch_python -ltorch_cuda -ltorch_cpu -ltorch -lc10_cuda -lc10
+ NVCCFLAGS += -shared -fPIC
+ NVCCFLAGS += -diag-suppress 3189 # for warnings generated by PyTorch C++ headers
+
+-ifeq ($(ARCH),SM103)
++ifeq ($(ARCH),ALL)
++# MoK uses the BF16/MXFP8 features shared by SM100 and SM103.
++NVCCFLAGS += --threads=2
++NVCCFLAGS += -DKITTENS_SM100
++NVCCFLAGS += -gencode arch=compute_100a,code=sm_100a
++NVCCFLAGS += -gencode arch=compute_103a,code=sm_103a
++else ifeq ($(ARCH),SM103)
+ NVCCFLAGS += -DKITTENS_SM103 -gencode arch=compute_103a,code=sm_103a
+ else ifeq ($(ARCH),SM100)
+ NVCCFLAGS += -DKITTENS_SM100 -gencode arch=compute_100a,code=sm_100a
+ else
+-$(error Unsupported ARCH '$(ARCH)'; expected SM100 or SM103)
++$(error Unsupported ARCH '$(ARCH)'; expected ALL, SM100 or SM103)
+ endif
+
+ all: $(OUT)
+diff --git a/csrc/megakernel/grouped_gemm.cuh b/csrc/megakernel/grouped_gemm.cuh
+index 02f3899..07ba154 100644
+--- a/csrc/megakernel/grouped_gemm.cuh
++++ b/csrc/megakernel/grouped_gemm.cuh
+@@ -291,7 +291,7 @@ static __device__ __forceinline__ void expert_grouped_gemm_kernel(
+                 update_phasebit<0>(gemm_bitfield, input_ring);
+                 input_ring = ring_advance<config::MLP_LOAD_PIPE_DEPTH>(input_ring);
+             }
+-            detail::tcgen05::commit<config::CLUSTER_SIZE>(gemm_outputs_arrived);
++            kittens::detail::tcgen05::commit<config::CLUSTER_SIZE>(gemm_outputs_arrived);
+         }
+     } else {
+         using epilogue_group = group<WARPGROUP_WARPS>;