Bump sglang to v0.5.16

.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}' || 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}' || true)
+            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 }}

docker/Dockerfile

@@ -10,7 +10,7 @@
 # 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 =============================================
@@ -131,11 +131,17 @@ RUN --mount=type=bind,source=docker/verify_transformer_engine.py,target=/tmp/ver
 # 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
 
@@ -155,13 +161,24 @@ 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@74d68c5e4bedf2b6774f2c92ed0f81b7c8d91ed0 --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
+# 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
 
 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 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/*
+
 RUN if [ "${ENABLE_CUDA_13}" = "1" ]; then \
     pip install nvidia-cudnn-cu13==9.22.0.52; \
   else \
@@ -241,6 +258,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.

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",
         },
     },

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:

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

tests/e2e/conftest_dumper.py

@@ -147,10 +147,7 @@
         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
-          )
+      - match: "hidden_states = self.mlp(hidden_states, forward_batch)"
         append: "dumper.dump('mlp_output', hidden_states, dims='t h # tp:replicated')"
 
   # --- attention internals ---

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