From 13d8814e2a6e8f66315669f32c51a80a37efedca Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Sat, 25 Jul 2026 02:55:19 -0700
Subject: [PATCH 01/26] Bump sglang to v0.5.16

---
 docker/Dockerfile | 4 ++--
 docker/build.py   | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index f8ce794fef..2ef779677f 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -10,12 +10,12 @@
 # release for the current arch picked from TARGETARCH — WHEELS_TAG_X86 on amd64,
 # WHEELS_TAG_ARM64 on arm64 (cu12-x86 overrides WHEELS_TAG_X86).
 
-ARG SGLANG_IMAGE_TAG=v0.5.15
+ARG SGLANG_IMAGE_TAG=v0.5.16
 FROM lmsysorg/sglang:${SGLANG_IMAGE_TAG} AS sglang
 
 # ======================================== Arguments =============================================
 
-ARG SGLANG_BRANCH=sglang-miles
+ARG SGLANG_BRANCH=sglang-miles-v0.5.16
 ARG SGLANG_COMMIT=""
 
 ARG MEGATRON_REPO=radixark/Megatron-LM
diff --git a/docker/build.py b/docker/build.py
index 207746417c..678d9add1d 100644
--- a/docker/build.py
+++ b/docker/build.py
@@ -45,7 +45,7 @@
         "tag_postfix": "-cu12",
         "build_args": {
             "ENABLE_CUDA_13": "0",
-            "SGLANG_IMAGE_TAG": "v0.5.15-cu129",
+            "SGLANG_IMAGE_TAG": "v0.5.16-cu129",
             "WHEELS_TAG_X86": "cu129-x86_64",
         },
     },

From db44327f6c402b696927943be46ad562e4176ed0 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Sat, 25 Jul 2026 02:58:06 -0700
Subject: [PATCH 02/26] Re-trigger CI to pick up ci-sglang-pr in the PR
 description

pr-test.yml does not listen for the 'edited' event, so the run started by the
label additions captured the PR body before 'ci-sglang-pr: sglang-miles-v0.5.16'
was added and would have resolved sglang to the default 'sglang-miles' branch.

From 25a2b9d2a7e3e9a7e091a9a80b61ce30f70ae886 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Sat, 25 Jul 2026 03:14:45 -0700
Subject: [PATCH 03/26] ci: add xgrammar to the CPU suite requirements

sglang v0.5.16 added `function_call/inkling_detector.py`, which imports
`from xgrammar import StructuralTag` at module level, and
`function_call_parser.py` imports that detector unconditionally. miles' conftest
reaches the parser through `mock_sglang_server`, so every CPU stage failed at
collection with `ModuleNotFoundError: No module named 'xgrammar'`.

The hosted CPU runner puts sglang on PYTHONPATH from source without installing
sglang's own dependencies, which is exactly what this file exists to mirror.
xgrammar is a hard sglang dependency (`xgrammar==0.2.1` in its pyproject, not an
extra); pinned to the same version. GPU/ROCm stages already get it from the
sglang base image.

Note `function_call/base_format_detector.py` guards the same import with
try/except ImportError, so the rest of that module family tolerates a missing
xgrammar; the new inkling detector does not.

Checked the whole delta rather than fixing one symbol at a time: walking every
unguarded module-level third-party import reachable from the CPU suite's sglang
entry points, xgrammar is the only one v0.5.16 adds over v0.5.15.
---
 tests/ci/requirements-ci-cpu.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/ci/requirements-ci-cpu.txt b/tests/ci/requirements-ci-cpu.txt
index 766fc9e1ee..e812c695af 100644
--- a/tests/ci/requirements-ci-cpu.txt
+++ b/tests/ci/requirements-ci-cpu.txt
@@ -11,3 +11,4 @@ sentencepiece==0.2.1
 tiktoken==0.13.0
 torch==2.11.0
 torchvision==0.26.0
+xgrammar==0.2.1

From 65dfeca5fc27314371a6dba13e15669349c73823 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Sat, 25 Jul 2026 03:30:32 -0700
Subject: [PATCH 04/26] test: mirror sglang v0.5.16's
 default_chat_template_kwargs in the serving double

`_make_serving` builds an OpenAIServingChat via `object.__new__` and sets only the
attributes `_process_messages` reads, so every attribute `__init__` gains has to be
mirrored here. v0.5.16 added server-level `default_chat_template_kwargs`, merged into
the request's `chat_template_kwargs` at the top of `_process_messages`, which made all
174 TestAlignWithSGLang / TestDatasetRouting cases fail with

    AttributeError: 'OpenAIServingChat' object has no attribute
    'default_chat_template_kwargs'

Mirrored as `{}` (what `__init__` produces when the server arg is unset) so the
request's own kwargs stay the only ones applied. Same pattern as the existing
`_tokenizer_auto_adds_specials` line added for v0.5.13.

Diffed every `self.X =` in OpenAIServingChat.__init__ between v0.5.15 and v0.5.16:
`default_chat_template_kwargs` is the only addition, and `_make_serving` is the only
hand-built construction site in the repo.
---
 tests/fast/utils/chat_template_utils/test_template.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/fast/utils/chat_template_utils/test_template.py b/tests/fast/utils/chat_template_utils/test_template.py
index 52e54ad85d..dfc5b6f638 100644
--- a/tests/fast/utils/chat_template_utils/test_template.py
+++ b/tests/fast/utils/chat_template_utils/test_template.py
@@ -67,6 +67,11 @@ def _make_serving(tokenizer) -> OpenAIServingChat:
     serving.is_gemma4 = False
     serving.tool_call_parser = None
     serving.reasoning_parser = None
+    # sglang v0.5.16 added server-level default chat-template kwargs, merged into
+    # the request's chat_template_kwargs at the top of _process_messages.
+    # __init__ always sets it (to `... or {}`); mirror the empty default so the
+    # request's own kwargs are the only ones applied.
+    serving.default_chat_template_kwargs = {}
     # sglang v0.5.13 probes whether the tokenizer auto-adds special tokens
     # (encode("") non-empty) to decide add_special_tokens at the chat-template
     # encode site. __init__ always sets this; mirror it here so _process_messages

From ef9a1254aba61c50650b6d954dbeb320aa4c45d3 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Sat, 25 Jul 2026 04:09:19 -0700
Subject: [PATCH 05/26] docker: reconcile apache-tvm-ffi with sglang's pin

3 of the 4 failing GPU stages died during CUDA-graph capture with

    File nvidia_cutlass_dsl/.../cutlass_dsl/tvm_ffi_provider.py, line 659
      self._kwargs_wrapper = kwargs_wrapper.make_kwargs_wrapper(
    TypeError: make_kwargs_wrapper() got an unexpected keyword argument
               'map_dataclass_to_tuple'

reached from layernorm.forward_cuda -> sgl_kernel rmsnorm -> flashinfer's CuTe
rmsnorm kernel. Not a migration regression: these jobs are green on main, and the
trigger is a dependency bump inside sglang v0.5.16 itself --

    nvidia-cutlass-dsl  4.5.2 -> 4.6.0
    flashinfer_python   0.6.12 -> 0.6.14
    apache-tvm-ffi      0.1.11 -> 0.1.11  (unchanged)

`map_dataclass_to_tuple` was added to apache-tvm-ffi in 0.1.10 (verified by
diffing the 0.1.9/0.1.10/0.1.11/0.1.12 sdists), so the TypeError means the
apache-tvm-ffi actually resolved at runtime is <= 0.1.9 -- older than the 0.1.11
sglang pins. cutlass-dsl 4.5.2 did not pass that kwarg, which is why v0.5.15 was
unaffected. Something among this file's many later pip installs drags in the
older build.

Restores sglang's own pin after every other install. The step verifies the
signature it depends on and fails the build if the pin does not stick, so a
future downgrade surfaces here instead of three hours later in a GPU stage.

I could not inspect the built image from this host, so the exact package doing
the downgrade is not identified; the assert is there to make that assumption
self-checking.
---
 docker/Dockerfile | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 2ef779677f..bd30b4434b 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -230,6 +230,22 @@ RUN --mount=type=cache,target=/root/.cache/pip \
       rm -rf /root/.cargo /root/.rustup /build/sgl-model-gateway /build/gateway_wheels; \
     fi
 
+# ==================== Reconcile apache-tvm-ffi with sglang's pin ====================
+# sglang v0.5.16 bumped nvidia-cutlass-dsl 4.5.2 -> 4.6.0 while leaving its
+# apache-tvm-ffi pin at 0.1.11. cutlass-dsl 4.6.0's tvm_ffi provider calls
+#   tvm_ffi.utils.kwargs_wrapper.make_kwargs_wrapper(..., map_dataclass_to_tuple=...)
+# and that parameter only exists in apache-tvm-ffi >= 0.1.10. An earlier install in
+# this file drags in an older apache-tvm-ffi, so flashinfer's CuTe rmsnorm path dies
+# during CUDA-graph capture with
+#   TypeError: make_kwargs_wrapper() got an unexpected keyword argument 'map_dataclass_to_tuple'
+# Restore sglang's own pin after every other install, and fail the build loudly if it
+# does not take effect rather than shipping an image that only breaks under CI.
+RUN pip install --force-reinstall --no-deps "apache-tvm-ffi==0.1.11" && \
+    python3 -c "import inspect; from tvm_ffi.utils import kwargs_wrapper as k; \
+p = inspect.signature(k.make_kwargs_wrapper).parameters; \
+assert 'map_dataclass_to_tuple' in p, \
+  'apache-tvm-ffi resolves to a build too old for nvidia-cutlass-dsl 4.6.0: %s' % list(p)"
+
 RUN rm -rf /tmp/wheels
 
 # The cu130 sglang base ships its baked Rust toolchain outside the default PATH.

From 635943f55bb499f9fc2be5b6d9613df3e89891df Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Sat, 25 Jul 2026 05:21:59 -0700
Subject: [PATCH 06/26] Re-trigger CI for sglang-miles-v0.5.16 [25/25]

Picks up the sglang-side fixes for three symbols v0.5.16 renamed or moved:
DecodeCudaGraphRunner.num_tokens_per_bs -> captured_req_width,
dp_attention.get_attention_tp_size -> ParallelState.attn_tp_size, and
layers.utils.cp_utils -> layers.cp.padding. All eight completed stage-c
failures in the previous run were the first two.

From 084aaff89625b47d595ad7910a0df66d8e9e0255 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Sat, 25 Jul 2026 06:40:07 -0700
Subject: [PATCH 07/26] Re-trigger CI for sglang-miles-v0.5.16 [26-27/27]

TpModelWorker.tp_rank -> self.ps.tp_rank, ModelRunner._model_update_group ->
weight_updater._model_update_group, and iter_runners reaching draft_runner /
draft_runner_list through the BaseSpecWorker.draft_worker property instead of
off the *V2 wrapper.

Round 4 got 3 stage-c jobs green (round 3: 0). The remaining three failures were
these bugs, plus 'window_size_left' which is still entangled with them -- this
run should separate the two.

From 1998911168824e986a2c4ea4dee2b1e7613f2d3b Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Sat, 25 Jul 2026 07:25:41 -0700
Subject: [PATCH 08/26] docker: fail the build when a transformer_engine patch
 does not apply

A `for` loop exits with the status of its last iteration, so a patch that failed
to apply was silently swallowed as long as a later one succeeded. With two
patches, a failure in the alphabetically-first
`te_dequantized_backward_override.patch` left `grouped_linear.py` unpatched and
the image shipped anyway.

That matters right now because both remaining GPU failures are inside the two
files these patches touch:

  ValueError: not enough values to unpack (expected 21, got 16)
    transformer_engine/pytorch/module/grouped_linear.py:405
      <- te_dequantized_backward_override.patch

  TypeError: _flash_attn_forward() got an unexpected keyword argument
             'window_size_left'
    transformer_engine/pytorch/attention/dot_product_attention/...
      <- te_fa2_sm103_whitelist.patch

Both jobs are green on main, so something under them changed with the v0.5.16
base image. This commit does not claim to fix that -- it makes the build say so
if a patch stopped applying, instead of producing an image that only fails hours
later inside a GPU stage.
---
 docker/Dockerfile | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index bd30b4434b..fb589f80a0 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -119,11 +119,17 @@ RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \
 # te_dequantized_backward_override.patch is a hot fix from
 # https://github.com/NVIDIA/TransformerEngine/pull/3141; drop it after TE v2.18.
 COPY docker/patch/ /tmp/patches/
+# A `for` loop exits with the status of its LAST iteration, so a patch that fails
+# to apply here used to be swallowed whenever a later one succeeded -- the image
+# shipped with TE silently unpatched. Fail the build instead.
 RUN if [ "${ENABLE_CUDA_13}" = "1" ] && [ -d /tmp/patches/cu13 ]; then \
       TE_DIR=$(python -c 'import importlib.util; print(importlib.util.find_spec("transformer_engine").submodule_search_locations[0])') && \
       for p in /tmp/patches/cu13/*.patch; do \
-        echo "Applying $(basename $p) to $TE_DIR" && \
-        patch -d "$TE_DIR" -p1 < "$p"; \
+        echo "Applying $(basename $p) to $TE_DIR"; \
+        patch -d "$TE_DIR" -p1 < "$p" || { \
+          echo "TE patch $(basename $p) did not apply cleanly against the installed transformer_engine" >&2; \
+          exit 1; \
+        }; \
       done; \
     fi && rm -rf /tmp/patches
 

From 90c263ab772a3ffa18dc07119bab634b7ee5e153 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Sat, 25 Jul 2026 23:09:12 -0700
Subject: [PATCH 09/26] docker: temporarily build megatron-bridge from
 radixark/Megatron-Bridge#24

Validates the TE 2.17 grouped-linear contract fix end to end. The
`not enough values to unpack (expected 21, got 16)` failures in the
test_glm5_*_lora_ci tests come from megatron-bridge's
`_forward_te_grouped_linear` calling TE's private `_GroupedLinear` with TE
<=2.14's positional layout; miles picked that up when it bumped
TransformerEngine 2.12 -> 2.17 in #1781.

Revert to `@bridge` -- ideally pinned to the merged commit rather than a branch,
as line 82 does for mbridge -- once #24 lands.

Note this does not address `window_size_left`, which is a separate TE 2.17 vs
flash-attn skew inside TE's own context-parallel path.
---
 docker/Dockerfile | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index fb589f80a0..07d606ca23 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -149,7 +149,9 @@ RUN pip install "git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.1
 RUN TMS_CUDA_MAJOR=$(python3 -c "import torch; print(torch.version.cuda.split('.')[0])") \
     pip install git+https://github.com/fzyzcjy/torch_memory_saver.git --no-cache-dir --force-reinstall
 RUN pip install "nvidia-modelopt[torch]>=0.37.0" --no-build-isolation
-RUN pip install git+https://github.com/radixark/Megatron-Bridge.git@bridge --no-deps --no-build-isolation
+# TEMPORARY: validating radixark/Megatron-Bridge#24 (TE 2.17 grouped-linear contract).
+# Point back at @bridge -- ideally pinned to the merged SHA -- once that PR lands.
+RUN pip install git+https://github.com/radixark/Megatron-Bridge.git@radixark/bridge/te-grouped-linear-2.17 --no-deps --no-build-isolation
 RUN pip install megatron-energon --no-deps
 RUN pip install multi-storage-client --no-deps
 

From 915c714895f12d00501423980ae7aeb775d6297d Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Sun, 26 Jul 2026 00:57:39 -0700
Subject: [PATCH 10/26] docker: use FA3 3.0.0 for transformer_engine 2.17

TE 2.17 moved use_flash_attn_3 out of the window_size-tuple branch in
cp_p2p_fwd/bwd_flash_attn, so it passes window_size_left/right to FA3. The
flash_attn_3 wheel in miles-wheels was 3.0.0b1, whose _flash_attn_forward
only accepts a window_size tuple, so every context-parallel path raises
TypeError. TE has no FA3-only kill switch -- NVTE_FLASH_ATTN disables FA2 too.

miles-wheels#10 rebuilt the wheel at flash-attention 00756db (3.0.0) and
replaced the release asset, so the pip install here already picks it up. That
wheel ships flash_attn_interface top-level and a flash_attn_3 re-export shim,
making the separate interface fetch unnecessary -- and harmful, since it
pinned a second revision that can drift from the one the extension module was
built at, and overwriting the shim with the full module double-registers the
flash_attn_3:: custom ops.
---
 docker/Dockerfile | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 07d606ca23..19b383fa52 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -73,11 +73,11 @@ RUN case "${TARGETARCH}" in \
 RUN pip install /tmp/wheels/flash_attn-*.whl
 
 # flash-attn hopper (FA3): Hopper-only (sm_90a). Coexists with FA2.
-RUN pip install /tmp/wheels/flash_attn_3-*.whl && \
-    python_path=$(python -c "import site; print(site.getsitepackages()[0])") && \
-    mkdir -p $python_path/flash_attn_3 && \
-    curl -fSL https://raw.githubusercontent.com/Dao-AILab/flash-attention/fbf24f67cf7f6442c5cfb2c1057f4bfc57e72d89/hopper/flash_attn_interface.py \
-      -o $python_path/flash_attn_3/flash_attn_interface.py
+# The wheel ships flash_attn_interface top-level (what transformer_engine imports)
+# plus a flash_attn_3.flash_attn_interface re-export shim, so there is nothing to
+# fetch separately. Do not overwrite the shim with a copy of the full module: it
+# re-runs @torch.library.custom_op("flash_attn_3::...") and double-registers.
+RUN pip install /tmp/wheels/flash_attn_3-*.whl
 
 RUN pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c --no-deps
 

From 2b07b98695eae9d94dbb9333123f8d6d2bf9a53e Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Sun, 26 Jul 2026 01:03:02 -0700
Subject: [PATCH 11/26] docker: install transformer_engine_torch runtime deps

#1796 switched TransformerEngine to the rolling wheels installed with
--no-deps. That pins the three TE dists to exactly those wheels, which is the
point, but it also drops transformer_engine_torch's own runtime requirements:
einops, onnx, onnxscript, pydantic, nvdlfw-inspect.

transformer_engine.pytorch imports onnxscript unconditionally on its core
import path (module/__init__ -> layernorm_linear -> base -> _common -> export
-> onnx_extensions), so every image built since #1796 fails at

    ModuleNotFoundError: No module named 'onnxscript'

as soon as anything touches TE -- which is every GPU test.

verify_transformer_engine.py did not catch it because it only compared
versions and metadata and called find_spec('transformer_engine'), which
resolves the package directory without importing it. Check that every non-TE
requirement of transformer_engine_torch is actually installed, so the build
fails at the layer that introduced the gap instead of shipping a green image.
---
 docker/Dockerfile                   |  5 +++++
 docker/verify_transformer_engine.py | 27 +++++++++++++++++++++++----
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 7faceb2e62..79e46a2cad 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -105,6 +105,10 @@ RUN if ls /tmp/wheels/causal_conv1d-*.whl 2>/dev/null | grep -q . && \
     fi
 
 # transformer_engine
+# --no-deps pins the three TE dists to exactly the wheels above, so pip cannot
+# pull a conflicting core dist. That also drops transformer_engine_torch's own
+# runtime deps, which still have to be installed: transformer_engine.pytorch
+# imports onnxscript unconditionally (module/_common -> export -> onnx_extensions).
 RUN --mount=type=bind,source=docker/verify_transformer_engine.py,target=/tmp/verify_transformer_engine.py \
     if [ "${ENABLE_CUDA_13}" = "1" ]; then \
       TE_CORE_DIST=transformer_engine_cu13; \
@@ -120,6 +124,7 @@ RUN --mount=type=bind,source=docker/verify_transformer_engine.py,target=/tmp/ver
     fi && \
     pip uninstall -y transformer-engine transformer-engine-cu12 transformer-engine-cu13 transformer-engine-torch && \
     pip install --force-reinstall --no-deps "$@" && \
+    pip install einops onnx onnxscript pydantic nvdlfw-inspect && \
     python3 /tmp/verify_transformer_engine.py "${TE_CORE_DIST}"
 
 # TE patches (cu13): B300/GB300 FA2 and backward override fixes
diff --git a/docker/verify_transformer_engine.py b/docker/verify_transformer_engine.py
index cde1fd1695..8735b7aea1 100644
--- a/docker/verify_transformer_engine.py
+++ b/docker/verify_transformer_engine.py
@@ -2,12 +2,18 @@
 
 import importlib.metadata as metadata
 import importlib.util
+import re
 import sys
 
 
 VERSION = "2.17.0"
 
 
+def _dist_name(requirement: str) -> str:
+    """'onnxscript', 'torch>=2.1', 'foo; extra == \"bar\"' -> normalised dist name."""
+    return re.split(r"[<>=!~;\[ ]", requirement.strip(), maxsplit=1)[0].lower().replace("_", "-")
+
+
 def verify(core_dist: str) -> None:
     core = core_dist.replace("_", "-")
     expected = {
@@ -16,15 +22,28 @@ def verify(core_dist: str) -> None:
         "transformer-engine-torch": VERSION,
     }
     actual = {name: metadata.version(name) for name in expected}
-    requires = {
-        requirement.lower().replace("_", "-").replace(" ", "")
-        for requirement in (metadata.requires("transformer-engine-torch") or [])
-    }
+    raw_requires = metadata.requires("transformer-engine-torch") or []
+    requires = {requirement.lower().replace("_", "-").replace(" ", "") for requirement in raw_requires}
 
     assert actual == expected, f"unexpected TransformerEngine versions: {actual}"
     assert f"{core}=={VERSION}" in requires, f"unexpected TransformerEngine torch requirements: {requires}"
     assert importlib.util.find_spec("transformer_engine") is not None, "transformer_engine package not found"
 
+    # The triplet is installed with --no-deps, so nothing makes transformer_engine_torch's
+    # own runtime deps present. Missing ones do not surface until something imports
+    # transformer_engine.pytorch, which is far too late -- a missing onnxscript shipped
+    # a green image that failed every GPU test.
+    missing = []
+    for requirement in raw_requires:
+        name = _dist_name(requirement)
+        if name.startswith("transformer-engine"):
+            continue  # pinned above
+        try:
+            metadata.version(name)
+        except metadata.PackageNotFoundError:
+            missing.append(name)
+    assert not missing, f"transformer_engine_torch runtime deps not installed: {missing}"
+
 
 if __name__ == "__main__":
     verify(sys.argv[1])

From 81f52f8e2a9f7e955883f7d611e5e6afb32bb721 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Sun, 26 Jul 2026 02:55:10 -0700
Subject: [PATCH 12/26] docker: pin Megatron-Bridge to the m_splits fix

Megatron-Bridge#24's first revision only handled half of TE 2.17's
grouped-linear contract change: m_splits was hoisted to a positional
parameter *and* retyped to torch.Tensor. Passing the Python list made
_GroupedLinear.forward raise 'list' object has no attribute 'tolist',
which failed every LoRA test in stage-c.

Pin by SHA rather than branch. buildkit keys this layer on the instruction
text, so moving the branch would have left the old revision in the image
with nothing to signal it.
---
 docker/Dockerfile | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 79e46a2cad..df49a66faf 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -162,8 +162,10 @@ RUN TMS_CUDA_MAJOR=$(python3 -c "import torch; print(torch.version.cuda.split('.
     pip install git+https://github.com/fzyzcjy/torch_memory_saver.git --no-cache-dir --force-reinstall
 RUN pip install "nvidia-modelopt[torch]>=0.37.0" --no-build-isolation
 # TEMPORARY: validating radixark/Megatron-Bridge#24 (TE 2.17 grouped-linear contract).
-# Point back at @bridge -- ideally pinned to the merged SHA -- once that PR lands.
-RUN pip install git+https://github.com/radixark/Megatron-Bridge.git@radixark/bridge/te-grouped-linear-2.17 --no-deps --no-build-isolation
+# Point back at @bridge -- pinned to the merged SHA -- once that PR lands.
+# Pinned by SHA, not branch: buildkit caches this layer on the instruction text,
+# so a branch that moves leaves the old revision baked into the image.
+RUN pip install git+https://github.com/radixark/Megatron-Bridge.git@36d20288b9f064b329cbbe49ee27f85454453668 --no-deps --no-build-isolation
 RUN pip install megatron-energon --no-deps
 RUN pip install multi-storage-client --no-deps
 

From bf5853b2af7093925cd180d5bcd7247b796b0f42 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Sun, 26 Jul 2026 21:34:37 -0700
Subject: [PATCH 13/26] docker: pin Megatron-Bridge to the merged #24

radixark/Megatron-Bridge#24 landed on @bridge as 7f0fb345 (squash). Its
src/megatron/bridge/peft/utils.py is byte-identical to 36d20288, the revision
this PR already exercised, so the pin moves off the PR branch without changing
what runs.

Kept as a SHA rather than @bridge. Line 82 pins mbridge the same way, and the
alternative is the failure this PR already hit once: buildkit keys the layer on
the instruction text, so a moving branch silently reuses the cached layer and
ships the previous revision.
---
 docker/Dockerfile | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index df49a66faf..db54fb4931 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -161,11 +161,10 @@ RUN pip install "git+https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git@v0.1
 RUN TMS_CUDA_MAJOR=$(python3 -c "import torch; print(torch.version.cuda.split('.')[0])") \
     pip install git+https://github.com/fzyzcjy/torch_memory_saver.git --no-cache-dir --force-reinstall
 RUN pip install "nvidia-modelopt[torch]>=0.37.0" --no-build-isolation
-# TEMPORARY: validating radixark/Megatron-Bridge#24 (TE 2.17 grouped-linear contract).
-# Point back at @bridge -- pinned to the merged SHA -- once that PR lands.
-# Pinned by SHA, not branch: buildkit caches this layer on the instruction text,
-# so a branch that moves leaves the old revision baked into the image.
-RUN pip install git+https://github.com/radixark/Megatron-Bridge.git@36d20288b9f064b329cbbe49ee27f85454453668 --no-deps --no-build-isolation
+# radixark/Megatron-Bridge#24 (TE 2.17 grouped-linear contract), merged to @bridge.
+# Pinned by SHA, not branch: buildkit caches this layer on the instruction text, so a
+# branch that moves leaves the old revision baked into the image with nothing to show it.
+RUN pip install git+https://github.com/radixark/Megatron-Bridge.git@7f0fb3456f8ffe47599b5fd167b454605d85f932 --no-deps --no-build-isolation
 RUN pip install megatron-energon --no-deps
 RUN pip install multi-storage-client --no-deps
 

From c24e7dcec2a678d6d7a4e27ecd7b590758bed900 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Sun, 26 Jul 2026 18:07:03 -0700
Subject: [PATCH 14/26] docker: raise cuDNN to 9.22.0 for transformer_engine
 2.17

TE 2.17's fused attention calls cudnn-frontend's
Reshape_attributes::set_reshape_mode, which sets the backend attribute
CUDNN_ATTR_OPERATION_RESHAPE_MODE. That attribute exists from cuDNN 9.22.0.
Against the pinned 9.16.0.29 it comes back BAD_PARAM and TE raises out of
fused_attn_bwd:

    fused_attn_f16_arbitrary_seqlen.cu:934 cuDNN Error:
      set_attribute(..., CUDNN_ATTR_OPERATION_RESHAPE_MODE, ...)
      ... CUDNN_STATUS_BAD_PARAM

It killed test_qwen3_5_35b_a3b_lora_ci (h100) and test_lora_qwen2.5_0.5B
(h200) in #1795, and reaches only configs where TE does not select FA3, so
most tests never touch it.

The old pin is obsolete in both of its purposes. It existed for
pytorch/pytorch#168167 (Conv3D bf16 collapse on H200), whose maintainers say
the fix landed in PyTorch 2.10 and whose stated workaround for 2.9.x was to
install a *newer* cuDNN, 9.15+. This image runs torch 2.11.0, and torch 2.11
already ships 9.19.0.56, so the pin now only drags cuDNN backwards.

Dropping the pin is not enough: 9.19.0.56 is still below 9.22.0. 9.22.0.52 is
the lowest release that satisfies TE, chosen over the newer 9.23-9.25 to keep
the change minimal.
---
 docker/Dockerfile | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index db54fb4931..e80e3f3c58 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -171,11 +171,16 @@ RUN pip install multi-storage-client --no-deps
 COPY requirements.txt /tmp/requirements.txt
 RUN rm -rf /usr/lib/python3/dist-packages/jwt /usr/lib/python3/dist-packages/PyJWT* && pip install -r /tmp/requirements.txt
 
-# https://github.com/pytorch/pytorch/issues/168167
+# transformer_engine 2.17's fused attention emits cudnn-frontend's
+# Reshape_attributes::set_reshape_mode, i.e. the backend attribute
+# CUDNN_ATTR_OPERATION_RESHAPE_MODE, which cuDNN only understands from 9.22.0.
+# torch 2.11 ships 9.19.0.56, so this has to raise it rather than just get out of
+# the way. pip reports a conflict against torch's exact ==9.19.0.56 pin and then
+# installs anyway; the previous pin here produced the same conflict downwards.
 RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \
-    pip install nvidia-cudnn-cu13==9.16.0.29; \
+    pip install nvidia-cudnn-cu13==9.22.0.52; \
   else \
-    pip install nvidia-cudnn-cu12==9.16.0.29; \
+    pip install nvidia-cudnn-cu12==9.22.0.52; \
   fi
 
 

From 17172194f3c08a0e32e41288aa033aa80fa4db36 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Mon, 27 Jul 2026 07:52:58 -0700
Subject: [PATCH 15/26] docker: remove the base image's apt cuDNN so the pip
 one is what loads

The apt copy shadows the pip copy in ldconfig, and the loader interleaves the
two: libcudnn_heuristic.so.9 comes from apt while libcudnn_graph.so.9 comes
from pip. transformer_engine then fails either on import with an undefined
symbol or inside fused attention with CUDNN_STATUS_BAD_PARAM, which is why
moving the pip pin alone changed nothing.

Verified on radixark/miles:pr-1795@sha256:792974ed with 8 DotProductAttention
forward+backward configs (d=64/128 x MHA/GQA x sbhd/bshd):

  apt present  + pip 9.16.0.29  -> torch refuses to init cuDNN
  apt present  + pip 9.22.0.52  -> 8/8 fail
  apt removed  + pip 9.19.0.56  -> 8/8 fail (the version floor is real)
  apt removed  + pip 9.22.0.52  -> 8/8 pass
---
 docker/Dockerfile | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index e80e3f3c58..f7981aeddf 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -171,12 +171,23 @@ RUN pip install multi-storage-client --no-deps
 COPY requirements.txt /tmp/requirements.txt
 RUN rm -rf /usr/lib/python3/dist-packages/jwt /usr/lib/python3/dist-packages/PyJWT* && pip install -r /tmp/requirements.txt
 
+# The base image also carries cuDNN as apt packages, in /usr/lib/<triple> where
+# ldconfig finds it ahead of the pip copy. The two then get interleaved at dlopen
+# time -- libcudnn_heuristic.so.9 resolves to the apt build while libcudnn_graph.so.9
+# resolves to the pip one -- and transformer_engine dies either on import
+# (undefined symbol _ZTVN5cudnn7backend12OperationSetE) or inside fused attention.
+# The apt copy is held, so removing it needs --allow-change-held-packages; without
+# that flag apt skips it silently and the pip pin below stays inert.
+RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then CU=13; else CU=12; fi; \
+  apt-get remove -y --purge --allow-change-held-packages \
+    libcudnn9-cuda-${CU} libcudnn9-dev-cuda-${CU} libcudnn9-headers-cuda-${CU} \
+  && rm -rf /var/lib/apt/lists/*
+
 # transformer_engine 2.17's fused attention emits cudnn-frontend's
 # Reshape_attributes::set_reshape_mode, i.e. the backend attribute
 # CUDNN_ATTR_OPERATION_RESHAPE_MODE, which cuDNN only understands from 9.22.0.
 # torch 2.11 ships 9.19.0.56, so this has to raise it rather than just get out of
-# the way. pip reports a conflict against torch's exact ==9.19.0.56 pin and then
-# installs anyway; the previous pin here produced the same conflict downwards.
+# the way: with the apt copy gone, 9.19.0.56 still fails and 9.22.0.52 passes.
 RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \
     pip install nvidia-cudnn-cu13==9.22.0.52; \
   else \

From 7d153f3a277a1d0f5912c8b8b83470fc8261fe70 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Mon, 27 Jul 2026 09:08:44 -0700
Subject: [PATCH 16/26] test: retarget the dumper mlp_output patch at sglang
 v0.5.16

v0.5.16 moved should_allreduce_fusion/use_reduce_scatter out of the
self.mlp(...) call args and into a surrounding get_forward().scoped(...)
context manager, so the match text found 0 occurrences. source_patcher
raises on a non-unique match, which killed the sglang scheduler during
load_model and left test_dumper hanging to its 1800s timeout.

The other seven matches in the sglang config still resolve uniquely
against v0.5.16.
---
 tests/e2e/conftest_dumper.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/tests/e2e/conftest_dumper.py b/tests/e2e/conftest_dumper.py
index dfaa1f0998..4c8b51c5c1 100644
--- a/tests/e2e/conftest_dumper.py
+++ b/tests/e2e/conftest_dumper.py
@@ -147,10 +147,9 @@
         append: |
           dumper.dump('pre_mlp_residual', residual, dims='t h # tp:replicated dp:=attn_dp')
           dumper.dump('pre_mlp_layernorm_output', hidden_states, dims='t h # tp:replicated')
-      - match: |
-          hidden_states = self.mlp(
-              hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter
-          )
+      # v0.5.16 moved should_allreduce_fusion/use_reduce_scatter out of the call
+      # args and into the surrounding get_forward().scoped(...) context manager.
+      - match: "hidden_states = self.mlp(hidden_states, forward_batch)"
         append: "dumper.dump('mlp_output', hidden_states, dims='t h # tp:replicated')"
 
   # --- attention internals ---

From 318346caaf7e5654d0c3c67aa07cf6b801b28fc0 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Mon, 27 Jul 2026 09:12:53 -0700
Subject: [PATCH 17/26] docker,test: drop the version-history comments

---
 docker/Dockerfile            | 15 +++------------
 tests/e2e/conftest_dumper.py |  2 --
 2 files changed, 3 insertions(+), 14 deletions(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index f7981aeddf..91667285c7 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -171,23 +171,14 @@ RUN pip install multi-storage-client --no-deps
 COPY requirements.txt /tmp/requirements.txt
 RUN rm -rf /usr/lib/python3/dist-packages/jwt /usr/lib/python3/dist-packages/PyJWT* && pip install -r /tmp/requirements.txt
 
-# The base image also carries cuDNN as apt packages, in /usr/lib/<triple> where
-# ldconfig finds it ahead of the pip copy. The two then get interleaved at dlopen
-# time -- libcudnn_heuristic.so.9 resolves to the apt build while libcudnn_graph.so.9
-# resolves to the pip one -- and transformer_engine dies either on import
-# (undefined symbol _ZTVN5cudnn7backend12OperationSetE) or inside fused attention.
-# The apt copy is held, so removing it needs --allow-change-held-packages; without
-# that flag apt skips it silently and the pip pin below stays inert.
+# The apt copy shadows the pip one in ldconfig and the loader interleaves the two,
+# so transformer_engine gets a mixed set of libcudnn sub-libraries. The packages are
+# held, hence --allow-change-held-packages.
 RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then CU=13; else CU=12; fi; \
   apt-get remove -y --purge --allow-change-held-packages \
     libcudnn9-cuda-${CU} libcudnn9-dev-cuda-${CU} libcudnn9-headers-cuda-${CU} \
   && rm -rf /var/lib/apt/lists/*
 
-# transformer_engine 2.17's fused attention emits cudnn-frontend's
-# Reshape_attributes::set_reshape_mode, i.e. the backend attribute
-# CUDNN_ATTR_OPERATION_RESHAPE_MODE, which cuDNN only understands from 9.22.0.
-# torch 2.11 ships 9.19.0.56, so this has to raise it rather than just get out of
-# the way: with the apt copy gone, 9.19.0.56 still fails and 9.22.0.52 passes.
 RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \
     pip install nvidia-cudnn-cu13==9.22.0.52; \
   else \
diff --git a/tests/e2e/conftest_dumper.py b/tests/e2e/conftest_dumper.py
index 4c8b51c5c1..867ec70d87 100644
--- a/tests/e2e/conftest_dumper.py
+++ b/tests/e2e/conftest_dumper.py
@@ -147,8 +147,6 @@
         append: |
           dumper.dump('pre_mlp_residual', residual, dims='t h # tp:replicated dp:=attn_dp')
           dumper.dump('pre_mlp_layernorm_output', hidden_states, dims='t h # tp:replicated')
-      # v0.5.16 moved should_allreduce_fusion/use_reduce_scatter out of the call
-      # args and into the surrounding get_forward().scoped(...) context manager.
       - match: "hidden_states = self.mlp(hidden_states, forward_batch)"
         append: "dumper.dump('mlp_output', hidden_states, dims='t h # tp:replicated')"
 

From b45c3d7f4b26520c50b8f6d8fd07643b62c77bf1 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Mon, 27 Jul 2026 17:13:21 -0700
Subject: [PATCH 18/26] TEMP diag: dump host/driver/cuDNN/env facts from the
 two failing LoRA shards

Revert once the cuDNN BAD_PARAM in fused_attn_bwd is understood. A devbox on the
same image digest runs the qwen2.5 LoRA test to completion through FusedAttention
sub-backend 1 with zero cuDNN errors, so the difference is host-level and not
recoverable from an ordinary CI log.

Captures per shard: driver version and GPU state (including memory already in use
by earlier tests in the same container), kernel, ldconfig/dpkg/pip cuDNN, the
libcudnn actually mapped into a torch+TE process and its cudnnGetVersion(), the
mounted model snapshot's config.json, ulimits, and the full container env.
NVTE_DEBUG=2 is injected into the training process so TE logs the available and
selected attention (sub-)backend per call.
---
 tests/e2e/lora/test_lora_qwen2.5_0.5B.py      | 70 ++++++++++++++++++-
 .../megatron/test_qwen3_5_35b_a3b_lora_ci.py  | 33 +++++++++
 2 files changed, 102 insertions(+), 1 deletion(-)

diff --git a/tests/e2e/lora/test_lora_qwen2.5_0.5B.py b/tests/e2e/lora/test_lora_qwen2.5_0.5B.py
index 94f8293eff..d985c0dbab 100644
--- a/tests/e2e/lora/test_lora_qwen2.5_0.5B.py
+++ b/tests/e2e/lora/test_lora_qwen2.5_0.5B.py
@@ -28,10 +28,70 @@
 NUM_GPUS = 4
 
 
+# TEMPORARY diagnostic (revert once the cuDNN BAD_PARAM in fused_attn_bwd is understood).
+# Everything here is read-only: it dumps the host/driver/cuDNN/env facts that cannot be
+# recovered from an ordinary CI log, so the CI container can be compared field-by-field
+# against a devbox reproduction that passes on the same image digest.
+_PROBE = r"""
+set +e
+echo '===== DIAG begin ====='
+echo '--- [1] host / driver / kernel ---'
+nvidia-smi --query-gpu=index,name,driver_version,memory.total,memory.used,persistence_mode \
+           --format=csv 2>&1
+nvidia-smi | sed -n '1,12p' 2>&1
+uname -a; cat /proc/version
+echo '--- [2] gpu memory already in use BEFORE this test (prior-test residue) ---'
+nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv 2>&1
+echo '--- [3] cuDNN on disk ---'
+ldconfig -p | grep -i cudnn || echo '(none in ldconfig)'
+for d in $(echo "${LD_LIBRARY_PATH}" | tr ':' ' ') /usr/lib/x86_64-linux-gnu /usr/local/cuda/lib64 \
+         /usr/local/lib/python3.12/dist-packages/nvidia/cudnn/lib; do
+  echo "  dir $d -> $(ls "$d"/libcudnn* 2>/dev/null | wc -l) libcudnn file(s)"
+done
+dpkg -l 2>/dev/null | grep -i cudnn || echo '(no cudnn apt package)'
+echo '--- [4] package versions ---'
+python -m pip list 2>/dev/null | grep -iE 'cudnn|^torch |transformer.engine|flash|nvidia-cublas'
+echo '--- [5] what a torch+TE process actually loads ---'
+python - <<'PY' 2>&1
+import ctypes, torch
+try:
+    print("torch", torch.__version__, "| torch.backends.cudnn.version() =", torch.backends.cudnn.version())
+except Exception as e:
+    print("torch cudnn init FAILED:", e)
+import transformer_engine, transformer_engine.pytorch  # noqa: F401
+print("transformer_engine", transformer_engine.__version__)
+try:
+    import flash_attn_interface
+    print("flash_attn_interface import OK")
+except Exception as e:
+    print("flash_attn_interface import FAILED:", type(e).__name__, e)
+maps = sorted({l.split()[-1] for l in open("/proc/self/maps") if "libcudnn" in l})
+for m in maps:
+    print("  mapped:", m)
+for m in maps:
+    if "libcudnn.so" in m:
+        lib = ctypes.CDLL(m); lib.cudnnGetVersion.restype = ctypes.c_size_t
+        v = lib.cudnnGetVersion()
+        print("  cudnnGetVersion() =", v, "->", "%d.%d.%d" % (v // 10000, (v % 10000) // 100, v % 100))
+print("device", torch.cuda.get_device_name(0), "| capability", torch.cuda.get_device_capability(0))
+print("cuda driver/runtime:", torch.version.cuda, torch.cuda_version if hasattr(torch, "cuda_version") else "")
+PY
+echo '--- [6] model snapshot identity (CI mounts a host cache; a devbox downloads fresh) ---'
+cat /root/models/Qwen2.5-0.5B-Instruct/config.json 2>&1
+ls -la /root/models/Qwen2.5-0.5B-Instruct/ 2>&1
+echo '--- [7] container limits ---'
+ulimit -a; df -h /dev/shm
+echo '--- [8] full container env ---'
+env | sort
+echo '===== DIAG end ====='
+"""
+
+
 def prepare():
     U.exec_command("mkdir -p /root/models /root/datasets")
     U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}")
     U.exec_command("hf download --repo-type dataset zhuzilin/gsm8k --local-dir /root/datasets/gsm8k")
+    U.exec_command(f"bash -c {_PROBE!r}")
 
 
 def execute():
@@ -130,7 +190,15 @@ def execute():
         train_args=train_args,
         num_gpus_per_node=NUM_GPUS,
         megatron_model_type=MODEL_TYPE,
-        extra_env_vars={"MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1"},
+        # NVTE_DEBUG here is the other half of the TEMPORARY diagnostic above: it makes TE
+        # log the available backends and the selected (sub-)backend for every attention
+        # call, which is the only way to line the CI run up against a devbox run that
+        # takes the same code path and passes.
+        extra_env_vars={
+            "MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1",
+            "NVTE_DEBUG": "1",
+            "NVTE_DEBUG_LEVEL": "2",
+        },
     )
 
 
diff --git a/tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py b/tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py
index 583fb6f5fc..de3ce7e96f 100644
--- a/tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py
+++ b/tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py
@@ -36,8 +36,41 @@ def _args(shared_outer: bool, virtual_experts: bool) -> ScriptArgs:
     )
 
 
+# TEMPORARY diagnostic (revert with the one in tests/e2e/lora/test_lora_qwen2.5_0.5B.py).
+# This shard runs on the H100 hosts, so it gives the second host's driver/cuDNN facts to
+# compare against both the H200 shard and the devbox reproduction.
+_PROBE = r"""
+set +e
+echo '===== DIAG begin (h100 shard) ====='
+nvidia-smi --query-gpu=index,name,driver_version,memory.total,memory.used --format=csv 2>&1
+nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv 2>&1
+uname -a
+ldconfig -p | grep -i cudnn || echo '(none in ldconfig)'
+dpkg -l 2>/dev/null | grep -i cudnn || echo '(no cudnn apt package)'
+python -m pip list 2>/dev/null | grep -iE 'cudnn|^torch |transformer.engine|flash'
+python - <<'PY' 2>&1
+import ctypes, torch
+try:
+    print("torch", torch.__version__, "| cudnn.version() =", torch.backends.cudnn.version())
+except Exception as e:
+    print("torch cudnn init FAILED:", e)
+import transformer_engine.pytorch  # noqa: F401
+maps = sorted({l.split()[-1] for l in open("/proc/self/maps") if "libcudnn" in l})
+for m in maps:
+    print("  mapped:", m)
+    if "libcudnn.so" in m:
+        lib = ctypes.CDLL(m); lib.cudnnGetVersion.restype = ctypes.c_size_t
+        print("  cudnnGetVersion() =", lib.cudnnGetVersion())
+print("device", torch.cuda.get_device_name(0), torch.cuda.get_device_capability(0))
+PY
+env | sort
+echo '===== DIAG end (h100 shard) ====='
+"""
+
+
 def prepare(args: ScriptArgs):
     _prepare_download(args)
+    U.exec_command(f"bash -c {_PROBE!r}")
 
 
 def execute(args: ScriptArgs):

From afb806634c7a4da3984944d3651dafa9c0e2500f Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Mon, 27 Jul 2026 22:12:08 -0700
Subject: [PATCH 19/26] TEMP diag: fix probe quoting (base64 transport)

f-string !r emitted literal backslash escapes, which exec_command then wrapped in
another bash -c, so the probe died with a syntax error and took the test with it
(qwen3.5 lora failed at 49s). Ship the script as base64 instead: the emitted
command contains no quotes or newlines at all. Verified locally that the payload
round-trips and that the decoded script passes bash -n.
---
 tests/e2e/lora/test_lora_qwen2.5_0.5B.py           | 3 ++-
 tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/e2e/lora/test_lora_qwen2.5_0.5B.py b/tests/e2e/lora/test_lora_qwen2.5_0.5B.py
index d985c0dbab..c8404827b5 100644
--- a/tests/e2e/lora/test_lora_qwen2.5_0.5B.py
+++ b/tests/e2e/lora/test_lora_qwen2.5_0.5B.py
@@ -11,6 +11,7 @@
 Triggered by label: run-ci-lora
 """
 
+import base64
 import os
 
 from tests.ci.ci_register import register_cuda_ci, register_rocm_ci
@@ -91,7 +92,7 @@ def prepare():
     U.exec_command("mkdir -p /root/models /root/datasets")
     U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}")
     U.exec_command("hf download --repo-type dataset zhuzilin/gsm8k --local-dir /root/datasets/gsm8k")
-    U.exec_command(f"bash -c {_PROBE!r}")
+    U.exec_command(f"echo {base64.b64encode(_PROBE.encode()).decode()} | base64 -d | bash")
 
 
 def execute():
diff --git a/tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py b/tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py
index de3ce7e96f..8c4eeb8745 100644
--- a/tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py
+++ b/tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py
@@ -1,3 +1,4 @@
+import base64
 import os
 
 from scripts.run_qwen3_5_35b_a3b_lora import ScriptArgs, _prepare_download, _train
@@ -70,7 +71,7 @@ def _args(shared_outer: bool, virtual_experts: bool) -> ScriptArgs:
 
 def prepare(args: ScriptArgs):
     _prepare_download(args)
-    U.exec_command(f"bash -c {_PROBE!r}")
+    U.exec_command(f"echo {base64.b64encode(_PROBE.encode()).decode()} | base64 -d | bash")
 
 
 def execute(args: ScriptArgs):

From bc1b5a56edc486c2bf05f200d3337309876bd326 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Tue, 28 Jul 2026 01:03:08 -0700
Subject: [PATCH 20/26] Revert the temporary cuDNN diagnostics

Root cause found: the CI job's requirements resolve downgraded the image's cuDNN.
Fixed in radixark/miles#1836, cherry-picked below.
---
 tests/e2e/lora/test_lora_qwen2.5_0.5B.py      | 71 +------------------
 .../megatron/test_qwen3_5_35b_a3b_lora_ci.py  | 34 ---------
 2 files changed, 1 insertion(+), 104 deletions(-)

diff --git a/tests/e2e/lora/test_lora_qwen2.5_0.5B.py b/tests/e2e/lora/test_lora_qwen2.5_0.5B.py
index c8404827b5..94f8293eff 100644
--- a/tests/e2e/lora/test_lora_qwen2.5_0.5B.py
+++ b/tests/e2e/lora/test_lora_qwen2.5_0.5B.py
@@ -11,7 +11,6 @@
 Triggered by label: run-ci-lora
 """
 
-import base64
 import os
 
 from tests.ci.ci_register import register_cuda_ci, register_rocm_ci
@@ -29,70 +28,10 @@
 NUM_GPUS = 4
 
 
-# TEMPORARY diagnostic (revert once the cuDNN BAD_PARAM in fused_attn_bwd is understood).
-# Everything here is read-only: it dumps the host/driver/cuDNN/env facts that cannot be
-# recovered from an ordinary CI log, so the CI container can be compared field-by-field
-# against a devbox reproduction that passes on the same image digest.
-_PROBE = r"""
-set +e
-echo '===== DIAG begin ====='
-echo '--- [1] host / driver / kernel ---'
-nvidia-smi --query-gpu=index,name,driver_version,memory.total,memory.used,persistence_mode \
-           --format=csv 2>&1
-nvidia-smi | sed -n '1,12p' 2>&1
-uname -a; cat /proc/version
-echo '--- [2] gpu memory already in use BEFORE this test (prior-test residue) ---'
-nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv 2>&1
-echo '--- [3] cuDNN on disk ---'
-ldconfig -p | grep -i cudnn || echo '(none in ldconfig)'
-for d in $(echo "${LD_LIBRARY_PATH}" | tr ':' ' ') /usr/lib/x86_64-linux-gnu /usr/local/cuda/lib64 \
-         /usr/local/lib/python3.12/dist-packages/nvidia/cudnn/lib; do
-  echo "  dir $d -> $(ls "$d"/libcudnn* 2>/dev/null | wc -l) libcudnn file(s)"
-done
-dpkg -l 2>/dev/null | grep -i cudnn || echo '(no cudnn apt package)'
-echo '--- [4] package versions ---'
-python -m pip list 2>/dev/null | grep -iE 'cudnn|^torch |transformer.engine|flash|nvidia-cublas'
-echo '--- [5] what a torch+TE process actually loads ---'
-python - <<'PY' 2>&1
-import ctypes, torch
-try:
-    print("torch", torch.__version__, "| torch.backends.cudnn.version() =", torch.backends.cudnn.version())
-except Exception as e:
-    print("torch cudnn init FAILED:", e)
-import transformer_engine, transformer_engine.pytorch  # noqa: F401
-print("transformer_engine", transformer_engine.__version__)
-try:
-    import flash_attn_interface
-    print("flash_attn_interface import OK")
-except Exception as e:
-    print("flash_attn_interface import FAILED:", type(e).__name__, e)
-maps = sorted({l.split()[-1] for l in open("/proc/self/maps") if "libcudnn" in l})
-for m in maps:
-    print("  mapped:", m)
-for m in maps:
-    if "libcudnn.so" in m:
-        lib = ctypes.CDLL(m); lib.cudnnGetVersion.restype = ctypes.c_size_t
-        v = lib.cudnnGetVersion()
-        print("  cudnnGetVersion() =", v, "->", "%d.%d.%d" % (v // 10000, (v % 10000) // 100, v % 100))
-print("device", torch.cuda.get_device_name(0), "| capability", torch.cuda.get_device_capability(0))
-print("cuda driver/runtime:", torch.version.cuda, torch.cuda_version if hasattr(torch, "cuda_version") else "")
-PY
-echo '--- [6] model snapshot identity (CI mounts a host cache; a devbox downloads fresh) ---'
-cat /root/models/Qwen2.5-0.5B-Instruct/config.json 2>&1
-ls -la /root/models/Qwen2.5-0.5B-Instruct/ 2>&1
-echo '--- [7] container limits ---'
-ulimit -a; df -h /dev/shm
-echo '--- [8] full container env ---'
-env | sort
-echo '===== DIAG end ====='
-"""
-
-
 def prepare():
     U.exec_command("mkdir -p /root/models /root/datasets")
     U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}")
     U.exec_command("hf download --repo-type dataset zhuzilin/gsm8k --local-dir /root/datasets/gsm8k")
-    U.exec_command(f"echo {base64.b64encode(_PROBE.encode()).decode()} | base64 -d | bash")
 
 
 def execute():
@@ -191,15 +130,7 @@ def execute():
         train_args=train_args,
         num_gpus_per_node=NUM_GPUS,
         megatron_model_type=MODEL_TYPE,
-        # NVTE_DEBUG here is the other half of the TEMPORARY diagnostic above: it makes TE
-        # log the available backends and the selected (sub-)backend for every attention
-        # call, which is the only way to line the CI run up against a devbox run that
-        # takes the same code path and passes.
-        extra_env_vars={
-            "MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1",
-            "NVTE_DEBUG": "1",
-            "NVTE_DEBUG_LEVEL": "2",
-        },
+        extra_env_vars={"MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1"},
     )
 
 
diff --git a/tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py b/tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py
index 8c4eeb8745..583fb6f5fc 100644
--- a/tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py
+++ b/tests/e2e/megatron/test_qwen3_5_35b_a3b_lora_ci.py
@@ -1,4 +1,3 @@
-import base64
 import os
 
 from scripts.run_qwen3_5_35b_a3b_lora import ScriptArgs, _prepare_download, _train
@@ -37,41 +36,8 @@ def _args(shared_outer: bool, virtual_experts: bool) -> ScriptArgs:
     )
 
 
-# TEMPORARY diagnostic (revert with the one in tests/e2e/lora/test_lora_qwen2.5_0.5B.py).
-# This shard runs on the H100 hosts, so it gives the second host's driver/cuDNN facts to
-# compare against both the H200 shard and the devbox reproduction.
-_PROBE = r"""
-set +e
-echo '===== DIAG begin (h100 shard) ====='
-nvidia-smi --query-gpu=index,name,driver_version,memory.total,memory.used --format=csv 2>&1
-nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv 2>&1
-uname -a
-ldconfig -p | grep -i cudnn || echo '(none in ldconfig)'
-dpkg -l 2>/dev/null | grep -i cudnn || echo '(no cudnn apt package)'
-python -m pip list 2>/dev/null | grep -iE 'cudnn|^torch |transformer.engine|flash'
-python - <<'PY' 2>&1
-import ctypes, torch
-try:
-    print("torch", torch.__version__, "| cudnn.version() =", torch.backends.cudnn.version())
-except Exception as e:
-    print("torch cudnn init FAILED:", e)
-import transformer_engine.pytorch  # noqa: F401
-maps = sorted({l.split()[-1] for l in open("/proc/self/maps") if "libcudnn" in l})
-for m in maps:
-    print("  mapped:", m)
-    if "libcudnn.so" in m:
-        lib = ctypes.CDLL(m); lib.cudnnGetVersion.restype = ctypes.c_size_t
-        print("  cudnnGetVersion() =", lib.cudnnGetVersion())
-print("device", torch.cuda.get_device_name(0), torch.cuda.get_device_capability(0))
-PY
-env | sort
-echo '===== DIAG end (h100 shard) ====='
-"""
-
-
 def prepare(args: ScriptArgs):
     _prepare_download(args)
-    U.exec_command(f"echo {base64.b64encode(_PROBE.encode()).decode()} | base64 -d | bash")
 
 
 def execute(args: ScriptArgs):

From 0e1db096c155f822ed5a199b628a45367a9411b2 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Tue, 28 Jul 2026 01:02:30 -0700
Subject: [PATCH 21/26] ci: restore the image's cuDNN pin after reconciling
 requirements

The GPU jobs install requirements.txt inside the test container. torchft-nightly
pulls torch, torch pins nvidia-cudnn-cu13 exactly, and that resolve uninstalls the
cuDNN the image deliberately installed last:

  Collecting nvidia-cudnn-cu13==9.19.0.56 (from torch>=2.7->torchft-nightly->-r requirements.txt)
  Uninstalling nvidia-cudnn-cu13-9.22.0.52
  Successfully installed nvidia-cudnn-cu13-9.19.0.56

transformer_engine 2.17's fused attention backward then fails with
CUDNN_STATUS_BAD_PARAM on CUDNN_ATTR_OPERATION_RESHAPE_MODE, which is what has been
failing test_lora_qwen2.5_0.5B, test_gpt_oss_20b_moe_lora_ci and
test_qwen3_5_35b_a3b_lora_ci on radixark/miles#1795.

Record the version the image shipped and reinstall it after the resolve, rather than
hardcoding it here: the Dockerfile stays the single place the version is chosen, and
--no-deps keeps the rest of the resolved set untouched.

Declaring the pin in requirements.txt instead does not work -- pip resolves the
conflict against torch's exact pin by downgrading torch itself to 2.10.0 and swapping
the whole CUDA stack to cu12.

Verified on an H200 devbox running the exact CI image digest: the requirements resolve
reproduces the downgrade (9.22.0.52 -> 9.19.0.56, torch untouched), the restore puts
9.22.0.52 back, and TE fused attention backward then passes 8/8 configs that fail at
9.19.0.56.

(cherry picked from commit bad0c0f6ac77e8c20ace40ce029446776fd8a34d)
---
 .github/workflows/_run-ci.yml | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/_run-ci.yml b/.github/workflows/_run-ci.yml
index 46d7dfd8ff..e730fc3be3 100644
--- a/.github/workflows/_run-ci.yml
+++ b/.github/workflows/_run-ci.yml
@@ -134,7 +134,27 @@ jobs:
 
       - name: Reconcile Miles dependencies
         shell: bash
-        run: python -m pip install -r "$GITHUB_WORKSPACE/requirements.txt" --break-system-packages
+        run: |
+          # The image pins cuDNN deliberately: transformer_engine's fused attention needs a
+          # newer build than torch's own exact pin, so the Dockerfile installs it last. This
+          # resolve drags it back down to whatever torch asks for, which makes fused attention
+          # backward fail with CUDNN_STATUS_BAD_PARAM. Record what the image shipped and put it
+          # back afterwards; --no-deps so nothing else in the resolved set moves.
+          pinned=""
+          for pkg in nvidia-cudnn-cu13 nvidia-cudnn-cu12; do
+            v=$(python -m pip show "$pkg" 2>/dev/null | awk '/^Version:/{print $2}')
+            [ -n "$v" ] && pinned="$pinned $pkg==$v"
+          done
+          python -m pip install -r "$GITHUB_WORKSPACE/requirements.txt" --break-system-packages
+          for spec in $pinned; do
+            pkg=${spec%%==*}
+            want=${spec##*==}
+            got=$(python -m pip show "$pkg" 2>/dev/null | awk '/^Version:/{print $2}')
+            if [ -n "$got" ] && [ "$got" != "$want" ]; then
+              echo "Restoring image pin: $pkg $got -> $want"
+              python -m pip install --no-deps --break-system-packages "$spec"
+            fi
+          done
 
       - name: Sync dependency sources
         if: ${{ !inputs.skip_dependency_install }}

From d143d2299c0e41c110a024dec89b0ac108361c23 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Tue, 28 Jul 2026 02:08:41 -0700
Subject: [PATCH 22/26] ci: do not let the empty-package probe trip errexit

GitHub runs 'shell: bash' as 'bash -e {0}'. '[ -n "$v" ] && pinned=...' as the last
statement of the loop body makes the loop exit non-zero whenever a package is absent
(cu12 on a cu13 image), which aborted the whole step and every GPU job with it. Use an
if-block so the loop always ends cleanly.

Verified by extracting the step from the workflow and running it under 'bash -e' with a
stub pip: exits 0 both when the version is untouched and when it has to be restored.

(cherry picked from commit 47ecd5dba4db65233691c2a55fa120ab90855548)
---
 .github/workflows/_run-ci.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/_run-ci.yml b/.github/workflows/_run-ci.yml
index e730fc3be3..aa4ce50483 100644
--- a/.github/workflows/_run-ci.yml
+++ b/.github/workflows/_run-ci.yml
@@ -143,7 +143,7 @@ jobs:
           pinned=""
           for pkg in nvidia-cudnn-cu13 nvidia-cudnn-cu12; do
             v=$(python -m pip show "$pkg" 2>/dev/null | awk '/^Version:/{print $2}')
-            [ -n "$v" ] && pinned="$pinned $pkg==$v"
+            if [ -n "$v" ]; then pinned="$pinned $pkg==$v"; fi
           done
           python -m pip install -r "$GITHUB_WORKSPACE/requirements.txt" --break-system-packages
           for spec in $pinned; do

From 1f35f2c86932bb5b7a35095f646a51e0392adb03 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Tue, 28 Jul 2026 03:01:59 -0700
Subject: [PATCH 23/26] ci: guard the pip-show probe against pipefail

GitHub runs steps as 'bash --noprofile --norc -e -o pipefail'. 'pip show' exits 1 for a
package that is not installed (cu12 on a cu13 image), pipefail propagates that through
the awk pipe, and -e then aborts the step -- taking every GPU job with it. Terminate
both command substitutions with '|| true', matching how this file already reads optional
values out of the PR body.

Verified by extracting the step from the workflow and running it under the exact shell
line GitHub logs, with a stub pip where cu12 is absent: exits 0, and still restores the
pin when the resolve moves it.

(cherry picked from commit 20e30ced19029972a089d05f437baa202de1110b)
---
 .github/workflows/_run-ci.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/_run-ci.yml b/.github/workflows/_run-ci.yml
index aa4ce50483..dd420cc562 100644
--- a/.github/workflows/_run-ci.yml
+++ b/.github/workflows/_run-ci.yml
@@ -142,14 +142,14 @@ jobs:
           # back afterwards; --no-deps so nothing else in the resolved set moves.
           pinned=""
           for pkg in nvidia-cudnn-cu13 nvidia-cudnn-cu12; do
-            v=$(python -m pip show "$pkg" 2>/dev/null | awk '/^Version:/{print $2}')
+            v=$(python -m pip show "$pkg" 2>/dev/null | awk '/^Version:/{print $2}' || true)
             if [ -n "$v" ]; then pinned="$pinned $pkg==$v"; fi
           done
           python -m pip install -r "$GITHUB_WORKSPACE/requirements.txt" --break-system-packages
           for spec in $pinned; do
             pkg=${spec%%==*}
             want=${spec##*==}
-            got=$(python -m pip show "$pkg" 2>/dev/null | awk '/^Version:/{print $2}')
+            got=$(python -m pip show "$pkg" 2>/dev/null | awk '/^Version:/{print $2}' || true)
             if [ -n "$got" ] && [ "$got" != "$want" ]; then
               echo "Restoring image pin: $pkg $got -> $want"
               python -m pip install --no-deps --break-system-packages "$spec"

From 8c2cf16406f24a215061217ab50c1a5e12eae1f9 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Tue, 28 Jul 2026 15:45:44 -0700
Subject: [PATCH 24/26] megatron: gather expert weights TEGroupedLinear leaves
 unmarked

TEGroupedLinear never sets the tensor-model-parallel attributes on its per-expert
weight0..weightN, so Megatron fills in the defaults and the params claim to be
unsharded (tensor_model_parallel=False, partition_dim=-1). They are in fact
expert-TP sharded, so all_gather_param(s) returned the local shard untouched and
the rollout engine received half-size expert weights: ranks past the first two
raise 'start (384) + length (192) exceeds dimension size (384)' in _load_w13, and
the earlier ranks silently load the wrong rows.

Gather these weights when etp > 1 and normalise the partition dim the same way
_check_and_fix_partition already compensates for this module's stride. ETP=1 is
unaffected, which is why only test_dumper (etp=2) hit it.

Verified on 8xH200: test_dumper passes, comparator reports failed=0 over
tp2_pp2_cp2_ep2_etp2.
---
 .../megatron_utils/update_weight/common.py    | 25 ++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/miles/backends/megatron_utils/update_weight/common.py b/miles/backends/megatron_utils/update_weight/common.py
index 2e289acbc6..379dc7a9c5 100644
--- a/miles/backends/megatron_utils/update_weight/common.py
+++ b/miles/backends/megatron_utils/update_weight/common.py
@@ -138,6 +138,19 @@ def _gather_with_stride(
     return torch.cat(interleaved, dim=partition_dim)
 
 
+def _is_unmarked_grouped_expert_weight(name: str, param: torch.nn.Parameter) -> bool:
+    """TEGroupedLinear never marks its per-expert weight0..weightN, so Megatron fills in
+    the defaults (tensor_model_parallel=False, partition_dim=-1) and the tensor claims to
+    be unsharded. It is expert-TP sharded whenever etp > 1, so the gather must still run.
+    """
+    return (
+        ".experts." in name
+        and ("linear_fc1.weight" in name or "linear_fc2.weight" in name)
+        and not param.tensor_model_parallel
+        and get_parallel_state().etp.size > 1
+    )
+
+
 def _check_and_fix_partition(args: Namespace, name: str, partition_stride: int, partition_dim: int) -> tuple[int, int]:
     """Validate partition_stride values for known parameter patterns.
 
@@ -147,9 +160,11 @@ def _check_and_fix_partition(args: Namespace, name: str, partition_stride: int,
     """
     if "linear_fc1.weight" in name and args.swiglu:
         partition_stride = 2
+        if partition_dim < 0:
+            partition_dim = 0
     elif "linear_fc2.weight" in name:
         assert partition_stride == 1, f"Expected partition_stride=1 for {name}, got {partition_stride}"
-        if partition_dim == 0:
+        if partition_dim <= 0:
             partition_dim = 1
     else:
         assert partition_stride == 1, f"Expected partition_stride=1 for {name}, got {partition_stride}"
@@ -165,7 +180,9 @@ def all_gather_param(args: Namespace, name: str, param: torch.nn.Parameter) -> t
         return param
 
     assert hasattr(param, "tensor_model_parallel"), f"{name} does not have tensor_model_parallel attribute"
-    if not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated":
+    if getattr(param, "parallel_mode", None) == "duplicated":
+        return param.data
+    if not param.tensor_model_parallel and not _is_unmarked_grouped_expert_weight(name, param):
         return param.data
 
     if ".experts." in name:
@@ -203,7 +220,9 @@ def all_gather_params_async(
         if "expert_bias" in info.name:
             gather_tasks.append((info, param, None, None, None, None))
             handles.append(None)
-        elif not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated":
+        elif getattr(param, "parallel_mode", None) == "duplicated" or (
+            not param.tensor_model_parallel and not _is_unmarked_grouped_expert_weight(info.name, param)
+        ):
             gather_tasks.append((info, param.data, None, None, None, None))
             handles.append(None)
         else:

From 3031e1e4ce8f62524d6fea374d975123215f292f Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Tue, 28 Jul 2026 15:45:50 -0700
Subject: [PATCH 25/26] fsdp: open a weight-update session before sending
 tensors

sglang requires update_weights_from_tensor to run inside a begin/end_weight_update
session; it restores the packed weights on begin and runs the post-load and quant
post-process on the full model at end. The megatron paths already do this, the FSDP
IPC path did not, so the scheduler died with 'update_weights_from_tensor requires an
open begin_weight_update session' and the job hung until the CI timeout.

Verified on 2xH200: 29 rollouts of test_qwen3_0.6B_fsdp_colocated_2xGPU with zero
session errors, where it previously failed on the first weight update.
---
 .../backends/experimental/fsdp_utils/update_weight_utils.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/miles/backends/experimental/fsdp_utils/update_weight_utils.py b/miles/backends/experimental/fsdp_utils/update_weight_utils.py
index 98000e7d34..7ba33aef9c 100644
--- a/miles/backends/experimental/fsdp_utils/update_weight_utils.py
+++ b/miles/backends/experimental/fsdp_utils/update_weight_utils.py
@@ -170,6 +170,11 @@ def update_bucket_weights(self, named_tensors, weight_version=None) -> None:
         )
 
         if dist.get_rank() == self._ipc_gather_src:
+            # sglang requires update_weights_from_tensor to run inside a weight-update
+            # session; it restores the packed weights on begin and runs the post-load /
+            # quant post-process on the full model at end.
+            ray.get(self._ipc_engine.begin_weight_update.remote())
+
             # Handle flattened bucket format (same as Megatron approach)
             # Each rank may have multiple dtype buckets
             # TODO: here we assume all ranks have the same number of dtypes
@@ -196,6 +201,7 @@ def update_bucket_weights(self, named_tensors, weight_version=None) -> None:
                     )
 
         if dist.get_rank() == self._ipc_gather_src:
+            ray.get(self._ipc_engine.end_weight_update.remote())
             ref = self._ipc_engine.flush_cache.remote()
             ray.get(ref)
 

From f27d4510fb14dc216a4840b1bcde51c63b4daa61 Mon Sep 17 00:00:00 2001
From: yueming-yuan <yym022502@gmail.com>
Date: Wed, 29 Jul 2026 13:58:37 -0700
Subject: [PATCH 26/26] docker: point at the sglang-miles branch now that it
 carries v0.5.16

sglang-miles was moved to the v0.5.16 stack, so the bump no longer needs the
versioned branch name. Two fixes that had landed on sglang-miles after the
v0.5.16 branch was cut (#30742, #32376) were carried over first, and
sglang-miles-v0.5.15-final now records the pre-move tip.
---
 docker/Dockerfile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index 41956d53fe..9612d935c5 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -15,7 +15,7 @@ FROM lmsysorg/sglang:${SGLANG_IMAGE_TAG} AS sglang
 
 # ======================================== Arguments =============================================
 
-ARG SGLANG_BRANCH=sglang-miles-v0.5.16
+ARG SGLANG_BRANCH=sglang-miles
 ARG SGLANG_COMMIT=""
 
 ARG MEGATRON_REPO=radixark/Megatron-LM