build(docker): bump SGLang to v0.5.23

.pre-commit-config.yaml

@@ -3,7 +3,7 @@ default_language_version:
 
 # The vendored upstream module stays byte-diffable against its source; keep
 # formatters and linters away so upstream-tracking diffs remain minimal.
-exclude: '^miles/utils/chat_template_utils/templates/encoding_dsv32\.py$'
+exclude: '^miles/utils/chat_template_utils/templates/encoding_dsv(32|4|41)\.py$'
 
 ci:
   autofix_prs: true

docker/Dockerfile

@@ -10,7 +10,7 @@
 # release for the current arch picked from TARGETARCH — WHEELS_TAG_X86 on amd64,
 # WHEELS_TAG_ARM64 on arm64.
 
-ARG SGLANG_IMAGE_TAG=v0.5.22
+ARG SGLANG_IMAGE_TAG=v0.5.23
 FROM lmsysorg/sglang:${SGLANG_IMAGE_TAG} AS sglang
 
 # ======================================== Arguments =============================================
@@ -136,8 +136,21 @@ RUN if ls /tmp/wheels/fast_hadamard_transform-*.whl 2>/dev/null | grep -q .; the
 # Megatron #5992 requires on SM100; 1.26.0 has no compact wrapper and fails
 # TransformerConfig validation once that change is in miles-main. A floor, not a
 # pin: the base image may already ship a newer one.
-RUN pip install /tmp/wheels/flash_mla-*.whl && \
-    pip install --no-deps "nvidia-cudnn-frontend>=1.28"
+RUN pip install --no-deps "nvidia-cudnn-frontend>=1.28"
+
+# Two FlashMLA builds, both packaged as `flash_mla`. SGLang pins sgl-flash-mla for decode
+# APIs the nv_dev build lacks; Megatron's CSA forward and miles' DSA kernel need nv_dev
+# (`indexer_topk`, SM90 and 576-dim sparse prefill). sglang-miles imports `sgl_flash_mla`
+# first, so move the base image's package there and install nv_dev as `flash_mla`.
+RUN FLASH_MLA_DIR=$(python3 -c 'import importlib.util; print(importlib.util.find_spec("flash_mla").submodule_search_locations[0])') && \
+    SGL_FLASH_MLA_DIR="$(dirname "$FLASH_MLA_DIR")/sgl_flash_mla" && \
+    mv "$FLASH_MLA_DIR" "$SGL_FLASH_MLA_DIR" && \
+    rm -rf "$SGL_FLASH_MLA_DIR/__pycache__" && \
+    sed -i 's/\bfrom flash_mla\b/from sgl_flash_mla/' "$SGL_FLASH_MLA_DIR"/*.py && \
+    if grep -nE '^\s*(from|import) flash_mla\b' "$SGL_FLASH_MLA_DIR"/*.py; then \
+      echo "sgl_flash_mla still imports flash_mla; extend the rename above" >&2; exit 1; \
+    fi && \
+    pip install /tmp/wheels/flash_mla-*.whl
 
 # Mamba kernels for nemotron_h hybrid (mamba+attention) models.
 RUN if ls /tmp/wheels/causal_conv1d-*.whl 2>/dev/null | grep -q . && \

docs/developer/versions.md

@@ -37,7 +37,7 @@ The default build-args are the version surface:
 
 | Build-arg | Default | What it selects |
 |---|---|---|
-| `SGLANG_IMAGE_TAG` | `v0.5.22` | The `lmsysorg/sglang` base image, which brings torch, CUDA, cuDNN and Mooncake |
+| `SGLANG_IMAGE_TAG` | `v0.5.23` | The `lmsysorg/sglang` base image, which brings torch, CUDA, cuDNN and Mooncake |
 | `SGLANG_BRANCH` | `sglang-miles` | The branch fetched into the base image's SGLang checkout |
 | `SGLANG_COMMIT` | empty | Empty means the branch HEAD at build time; set it to freeze one commit |
 | `MEGATRON_REPO` / `MEGATRON_BRANCH` / `MEGATRON_COMMIT` | `radixark/Megatron-LM` / `miles-main` / empty | The Megatron-LM checkout; an empty commit follows branch HEAD, while a release build supplies the locked commit |

miles/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py

@@ -179,13 +179,15 @@ def _quantize_param(args, name, weight, weight_block_size, is_skipped):
 
 
 def _get_scale_format(args, name, weight_block_size):
+    is_moe_expert = ".experts." in name
+    # grouped_moe mirrors SGLang's own MoE requant call: SM120 requantizes only grouped MoE scales.
     if not (
         should_deepgemm_weight_requant_ue8m0
-        and should_deepgemm_weight_requant_ue8m0(weight_block_size=weight_block_size)
+        and should_deepgemm_weight_requant_ue8m0(weight_block_size=weight_block_size, grouped_moe=is_moe_expert)
     ):
         return None  # use default fp32 scale format
 
-    if ".experts." not in name:
+    if not is_moe_expert:
         # Non-MoE linear weights: ue8m0 when deepgemm is enabled
         return "ue8m0"
 

miles/backends/sglang_utils/sglang_engine.py

@@ -145,6 +145,11 @@ def _compute_server_args(
         kwargs["enable_return_routed_experts"] = True
     if args.use_rollout_indexer_replay:
         kwargs["enable_return_indexer_topk"] = True
+    # SGLang v0.5.23 turns DeepSeek-V4.1 decoder SWA bounded replay on when this is unset; it is not
+    # numerically equivalent to full prefill, so rollout logprobs and replayed routing would drift
+    # from the trainer. Keep full prefill unless --sglang-enable-decoder-swa-bounded-replay is set.
+    if base_args.get("enable_decoder_swa_bounded_replay") is None:
+        kwargs["enable_decoder_swa_bounded_replay"] = False
     if args.use_output_store:
         # Every engine, eval included: an eval request never asks for the store.
         kwargs.update(sglang_output_store_server_args(args, local_hostname=host.strip("[]")))

miles/utils/chat_template_utils/deepseek.py

@@ -1,11 +1,11 @@
 """Shared bridge for the DeepSeek official-encoder families (V3.2, V4, V4.1).
 
-None of the families ships a jinja chat_template: V4 and V4.1 render through
-sglang's ``encoding_dsv4`` / ``encoding_dsv41``, while V3.2 uses miles' vendored
-``templates.encoding_dsv32``.  Both modules share one calling convention,
-and miles' ``apply_chat_template`` routes any matching tokenizer here.  Each
-family is one ``DeepSeekFamily`` instance wrapping its encoder module;
-everything else is shared.
+None of the families ships a jinja chat_template: each renders through miles'
+vendored official encoder (``templates.encoding_dsv32`` / ``encoding_dsv4`` /
+``encoding_dsv41``).  The modules share one calling convention, and miles'
+``apply_chat_template`` routes any matching tokenizer here.  Each family is one
+``DeepSeekFamily`` instance wrapping its encoder module; everything else is
+shared.
 """
 
 from __future__ import annotations
@@ -17,10 +17,9 @@
 from collections.abc import Callable
 from typing import Any
 
-from sglang.srt.entrypoints.openai import chat_encoding, encoding_dsv4, encoding_dsv41
 from sglang.srt.entrypoints.openai.protocol import Tool
 
-from miles.utils.chat_template_utils.templates import encoding_dsv32
+from miles.utils.chat_template_utils.templates import encoding_dsv4, encoding_dsv32, encoding_dsv41
 
 _ASSISTANT_SP_TOKEN = "<｜Assistant｜>"
 
@@ -51,9 +50,10 @@ def _inject_tools_into_system(
     """Put *tools* in the system message, where ``encode_messages`` reads them.
 
     The encoder serializes each tool dict verbatim into ``<functions>``, so they
-    must round-trip through the same pydantic payload sglang's ``serving_chat``
-    builds (``Tool.model_dump()`` fills defaults / orders fields; V4.1 keeps
-    only the fields the client sent) or the token ids drift from what sglang serves.
+    must round-trip through the payload the encoder is paired with or the token
+    ids drift: V3.2 uses ``Tool.model_dump()`` as sglang's dsv32 serving path
+    does, V4 keeps the v0.5.22 payload (``dsv4_tool_payload``), and V4.1 keeps
+    only the fields the client sent (``dsv41_tool_payload``).
     """
     out = copy.deepcopy(messages)
     if not out or out[0].get("role") != "system":
@@ -97,7 +97,7 @@ def render_messages(
         add_generation_prompt: bool = True,
         **kwargs: Any,
     ) -> str:
-        """Render *messages* into this family's prompt via sglang ``encode_messages``.
+        """Render *messages* into this family's prompt via its encoder's ``encode_messages``.
 
         Tool_call ``arguments`` must already be JSON strings; *tools*, if given, are
         injected into the system message (see ``_inject_tools_into_system``).  The
@@ -133,34 +133,105 @@ def _generation_prompt_suffix(self, tail_role: str | None, thinking_token: str)
         return None
 
 
+def dsv4_tool_payload(tool: Tool) -> dict[str, Any]:
+    """The tool dict sglang v0.5.22's ``Tool.model_dump()`` gave the V4 encoder.
+
+    v0.5.23 keeps the client's key order and drops an unset ``strict``
+    (sgl-project/sglang#46133); V4 renders the dict verbatim, so rebuild the
+    declared field order with ``strict`` defaulting to False.
+    """
+    payload = tool.model_dump()
+    function = payload["function"]
+    declared = {
+        "description": function.get("description"),
+        "name": function["name"],
+        "parameters": function.get("parameters"),
+        "strict": bool(function.get("strict", False)),
+    }
+    if "defer_loading" in function:
+        declared["defer_loading"] = function["defer_loading"]
+    payload["function"] = declared
+    return payload
+
+
 class DeepSeekV4Family(DeepSeekFamily):
     template = encoding_dsv4
 
+    def _inject_tools(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        return _inject_tools_into_system(messages, tools, dsv4_tool_payload)
+
     def _generation_prompt_suffix(self, tail_role: str | None, thinking_token: str) -> str | None:
         if tail_role in {"user", "developer", "tool"}:
             return _ASSISTANT_SP_TOKEN + thinking_token
         return None
 
 
+def parse_dsv41_reasoning_effort(value: Any) -> str | int | None:
+    """Map an API ``reasoning_effort`` onto what the V4.1 encoder accepts.
+
+    An int budget only reaches here through ``chat_template_kwargs``; None
+    means unsupported, and the caller applies its default.
+    """
+    if isinstance(value, bool):
+        return None
+    if isinstance(value, int):
+        return value if 1 <= value <= 100 else None
+    if isinstance(value, float):
+        return max(1, round(value * 100)) if 0.0 <= value <= 0.99 else None
+    if value in encoding_dsv41.REASONING_EFFORT_MAPPINGS:
+        return value
+    return None
+
+
+_OPENAI_FUNCTION_FIELD_ORDER = ("name", "description", "parameters")
+
+
+def dsv41_tool_payload(tool: Tool) -> dict[str, Any]:
+    """The tool dict the V4.1 encoder renders verbatim into the prompt.
+
+    Only fields the client sent, in the OpenAI field order; pydantic would
+    otherwise add defaults (strict=false) and reorder keys by declaration.
+    """
+    payload = tool.model_dump(exclude_unset=True, exclude_none=True)
+    function = dict(payload.get("function") or {})
+    ordered = {k: function.pop(k) for k in _OPENAI_FUNCTION_FIELD_ORDER if k in function}
+    ordered.update(function)
+    payload["function"] = ordered
+    return payload
+
+
+def default_dsv41_reasoning_effort_from_env(raw: str | None) -> str | int:
+    """Parse ``SGLANG_DSV41_REASONING_EFFORT``; raises so a bad value fails at boot."""
+    if raw is None or not raw.strip():
+        return encoding_dsv41.DEFAULT_REASONING_EFFORT
+    value: Any = int(raw) if raw.strip().isdigit() else raw.strip()
+    effort = parse_dsv41_reasoning_effort(value)
+    if effort is None:
+        raise ValueError(
+            f"Invalid SGLANG_DSV41_REASONING_EFFORT={raw!r}; expected one of "
+            f"{list(encoding_dsv41.REASONING_EFFORT_MAPPINGS)} or an integer in [1, 100]"
+        )
+    return effort
+
+
 class DeepSeekV41Family(DeepSeekV4Family):
-    """V4.1 mirrors sglang's ``serving_chat`` dsv41 branch: tools keep only the
-    fields the client sent, ``reasoning_effort`` always reaches the encoder (the
-    server default comes from ``SGLANG_DSV41_REASONING_EFFORT``), and a
-    mid-conversation system message also opens an assistant turn."""
+    """V4.1 keeps the rendering of sglang's ``serving_chat`` dsv41 branch from
+    before v0.5.23 moved it to deepseek-recipe: tools keep only the fields the
+    client sent, ``reasoning_effort`` always reaches the encoder (the default
+    comes from ``SGLANG_DSV41_REASONING_EFFORT``), and a mid-conversation system
+    message also opens an assistant turn."""
 
     template = encoding_dsv41
 
     def _build_encode_config(self, kwargs: dict) -> dict:
         kwargs = super()._build_encode_config(kwargs)
-        default = chat_encoding.default_dsv41_reasoning_effort_from_env(
-            os.environ.get("SGLANG_DSV41_REASONING_EFFORT", "high")
-        )
-        effort = chat_encoding.parse_dsv41_reasoning_effort(kwargs.get("reasoning_effort"))
+        default = default_dsv41_reasoning_effort_from_env(os.environ.get("SGLANG_DSV41_REASONING_EFFORT", "high"))
+        effort = parse_dsv41_reasoning_effort(kwargs.get("reasoning_effort"))
         kwargs["reasoning_effort"] = default if effort is None else effort
         return kwargs
 
     def _inject_tools(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
-        return _inject_tools_into_system(messages, tools, chat_encoding.dsv41_tool_payload)
+        return _inject_tools_into_system(messages, tools, dsv41_tool_payload)
 
     def _generation_prompt_suffix(self, tail_role: str | None, thinking_token: str) -> str | None:
         if tail_role in {"user", "developer", "tool", "system"}:

miles/utils/chat_template_utils/templates/__init__.py

@@ -1 +1 @@
-"""Bundled fixed chat templates and the vendored official DeepSeek V3.2 encoder."""
+"""Bundled fixed chat templates and the vendored official DeepSeek V3.2 / V4 / V4.1 encoders."""